kernel/resolution/
pipelines.rs

1//! The diffusers pipeline-family registry: maps a `model_index.json`
2//! `_class_name` to a modality/capability/parameter profile, refined by the
3//! model's scheduler and repo name. This is how a discovered diffusers bundle
4//! gets its image/video/audio shape and its sampler parameter schema.
5
6use std::collections::HashSet;
7use std::path::Path;
8use std::sync::LazyLock;
9
10use crate::records::{Capability, JsonValue, Modality, ParamSpec, ParamType};
11
12/// The `_class_name` declared in a diffusers `model_index.json` (or scheduler
13/// config), if the file reads and parses as an object.
14pub(crate) fn diffusers_pipeline_class(path: &Path) -> Option<String> {
15    let bytes = std::fs::read(path).ok()?;
16    let JsonValue::Object(index) = serde_json::from_slice::<JsonValue>(&bytes).ok()? else {
17        return None;
18    };
19    index
20        .get("_class_name")
21        .and_then(JsonValue::as_str)
22        .map(str::to_owned)
23}
24
25/// The resolved shape of a diffusers pipeline: what it produces and the sampler
26/// parameters it exposes.
27#[derive(Debug, Clone, PartialEq)]
28pub struct DiffusersPipelineProfile {
29    /// What the pipeline generates.
30    pub modality: Modality,
31    /// The capabilities it serves.
32    pub capabilities: Vec<Capability>,
33    /// Its tunable parameter schema.
34    pub params: Vec<ParamSpec>,
35}
36
37/// The facts read from a pipeline's `scheduler/scheduler_config.json` that a
38/// refinement matches against.
39#[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
40pub struct SchedulerFacts {
41    /// The scheduler's `_class_name`.
42    pub class_name: Option<String>,
43    /// Its `timestep_spacing`.
44    pub timestep_spacing: Option<String>,
45}
46
47impl SchedulerFacts {
48    /// Scheduler facts from a class name and timestep spacing.
49    pub fn new(class_name: Option<String>, timestep_spacing: Option<String>) -> Self {
50        Self {
51            class_name,
52            timestep_spacing,
53        }
54    }
55}
56
57/// A conditional parameter override applied to a family when the model's
58/// scheduler (and optionally its repo name) matches — e.g. the SDXL "turbo"
59/// low-step preset.
60#[derive(Debug, Clone, PartialEq)]
61pub struct PipelineRefinement {
62    /// Scheduler `_class_name`s this refinement applies to.
63    pub scheduler_classes: HashSet<String>,
64    /// A required `timestep_spacing`, if the refinement is spacing-specific.
65    pub timestep_spacing: Option<String>,
66    /// Repo-name tokens that must be present (any one), if name-gated.
67    pub name_signals: HashSet<String>,
68    /// The parameter specs to overlay (replacing by key, or appending).
69    pub param_overrides: Vec<ParamSpec>,
70}
71
72impl PipelineRefinement {
73    /// Whether this refinement applies given the scheduler `facts` and an optional
74    /// `repo_hint`. The scheduler class must match; a set `timestep_spacing` must
75    /// match; and if any `name_signals` are declared, a token of `repo_hint` must
76    /// be among them.
77    pub fn matches(&self, facts: &SchedulerFacts, repo_hint: Option<&str>) -> bool {
78        let Some(class_name) = &facts.class_name else {
79            return false;
80        };
81        if !self.scheduler_classes.contains(class_name) {
82            return false;
83        }
84        if let Some(spacing) = &self.timestep_spacing
85            && facts.timestep_spacing.as_deref() != Some(spacing.as_str())
86        {
87            return false;
88        }
89        if self.name_signals.is_empty() {
90            return true;
91        }
92        let Some(repo_hint) = repo_hint else {
93            return false;
94        };
95        let tokens = name_tokens(repo_hint);
96        self.name_signals
97            .iter()
98            .any(|signal| tokens.contains(signal))
99    }
100}
101
102/// The lowercase alphanumeric tokens of a repo name (`"SDXL-Turbo/v1"` →
103/// `{"sdxl", "turbo", "v1"}`).
104fn name_tokens(repo_hint: &str) -> HashSet<String> {
105    repo_hint
106        .to_lowercase()
107        .split(|ch: char| !ch.is_alphanumeric())
108        .filter(|token| !token.is_empty())
109        .map(str::to_owned)
110        .collect()
111}
112
113/// A family of diffusers pipeline classes that share a modality/capability/
114/// parameter profile (e.g. all SDXL pipelines), plus any scheduler refinements.
115#[derive(Debug, Clone, PartialEq)]
116pub struct PipelineFamily {
117    /// A stable family id (`"stable-diffusion-xl"`).
118    pub id: String,
119    /// The `_class_name`s that belong to this family.
120    pub class_names: HashSet<String>,
121    /// What the family generates.
122    pub modality: Modality,
123    /// The capabilities it serves (empty for edit/upscale/video/audio families).
124    pub capabilities: Vec<Capability>,
125    /// The base parameter schema.
126    pub params: Vec<ParamSpec>,
127    /// Scheduler-conditional parameter refinements.
128    pub refinements: Vec<PipelineRefinement>,
129}
130
131impl PipelineFamily {
132    fn new(
133        id: &str,
134        class_names: &[&str],
135        modality: Modality,
136        capabilities: Vec<Capability>,
137        params: Vec<ParamSpec>,
138        refinements: Vec<PipelineRefinement>,
139    ) -> Self {
140        Self {
141            id: id.to_owned(),
142            class_names: class_names.iter().map(|name| (*name).to_owned()).collect(),
143            modality,
144            capabilities,
145            params,
146            refinements,
147        }
148    }
149}
150
151/// The registry of diffusers pipeline families. Look up a `_class_name` to get
152/// its [`DiffusersPipelineProfile`].
153#[derive(Debug, Clone)]
154pub struct PipelineFamilyRegistry {
155    /// The known families.
156    pub families: Vec<PipelineFamily>,
157}
158
159impl PipelineFamilyRegistry {
160    /// A registry over `families`.
161    pub fn new(families: Vec<PipelineFamily>) -> Self {
162        Self { families }
163    }
164
165    /// The process-wide built-in registry, built once on first use. Prefer this
166    /// over [`builtin`](Self::builtin) at call sites — identification runs it per
167    /// model, and the table is immutable.
168    pub fn shared() -> &'static Self {
169        static BUILTIN: LazyLock<PipelineFamilyRegistry> =
170            LazyLock::new(PipelineFamilyRegistry::builtin);
171        &BUILTIN
172    }
173
174    /// The family owning `class_name`, if any.
175    pub fn family(&self, class_name: &str) -> Option<&PipelineFamily> {
176        self.families
177            .iter()
178            .find(|family| family.class_names.contains(class_name))
179    }
180
181    /// The resolved profile for `class_name`, applying the first scheduler
182    /// refinement that matches `scheduler`/`repo_hint`. `None` if the class name
183    /// is unknown.
184    pub fn profile(
185        &self,
186        class_name: &str,
187        scheduler: Option<&SchedulerFacts>,
188        repo_hint: Option<&str>,
189    ) -> Option<DiffusersPipelineProfile> {
190        let family = self.family(class_name)?;
191        let mut params = family.params.clone();
192        if let Some(scheduler) = scheduler
193            && let Some(refinement) = family
194                .refinements
195                .iter()
196                .find(|refinement| refinement.matches(scheduler, repo_hint))
197        {
198            for overlay in &refinement.param_overrides {
199                if let Some(index) = params.iter().position(|spec| spec.key == overlay.key) {
200                    params[index] = overlay.clone();
201                } else {
202                    params.push(overlay.clone());
203                }
204            }
205        }
206        Some(DiffusersPipelineProfile {
207            modality: family.modality.clone(),
208            capabilities: family.capabilities.clone(),
209            params,
210        })
211    }
212
213    /// The built-in family table covering the common diffusers pipelines.
214    pub fn builtin() -> Self {
215        Self::new(vec![
216            PipelineFamily::new(
217                "flux",
218                &["FluxPipeline"],
219                Modality::image(),
220                vec![Capability::image()],
221                flux_params(),
222                Vec::new(),
223            ),
224            PipelineFamily::new(
225                "stable-diffusion",
226                &["StableDiffusionPipeline"],
227                Modality::image(),
228                vec![Capability::image()],
229                sd1_params(),
230                vec![turbo_refinement()],
231            ),
232            PipelineFamily::new(
233                "stable-diffusion-xl",
234                &["StableDiffusionXLPipeline"],
235                Modality::image(),
236                vec![Capability::image()],
237                sdxl_params(),
238                vec![turbo_refinement()],
239            ),
240            PipelineFamily::new(
241                "stable-diffusion-3",
242                &["StableDiffusion3Pipeline"],
243                Modality::image(),
244                vec![Capability::image()],
245                sd3_params(),
246                Vec::new(),
247            ),
248            PipelineFamily::new(
249                "pixart",
250                &["PixArtAlphaPipeline", "PixArtSigmaPipeline"],
251                Modality::image(),
252                vec![Capability::image()],
253                pixart_params(),
254                Vec::new(),
255            ),
256            PipelineFamily::new(
257                "kandinsky",
258                &["KandinskyV22Pipeline", "KandinskyV22CombinedPipeline"],
259                Modality::image(),
260                vec![Capability::image()],
261                kandinsky_params(),
262                Vec::new(),
263            ),
264            PipelineFamily::new(
265                "latent-consistency",
266                &["LatentConsistencyModelPipeline"],
267                Modality::image(),
268                vec![Capability::image()],
269                lcm_params(),
270                Vec::new(),
271            ),
272            PipelineFamily::new(
273                "image-edit",
274                &[
275                    "StableDiffusionImg2ImgPipeline",
276                    "StableDiffusionInpaintPipeline",
277                    "StableDiffusionXLImg2ImgPipeline",
278                    "StableDiffusionXLInpaintPipeline",
279                    "StableDiffusion3Img2ImgPipeline",
280                    "StableDiffusion3InpaintPipeline",
281                    "FluxImg2ImgPipeline",
282                    "FluxInpaintPipeline",
283                    "KandinskyV22Img2ImgPipeline",
284                    "LatentConsistencyModelImg2ImgPipeline",
285                ],
286                Modality::image(),
287                Vec::new(),
288                Vec::new(),
289                Vec::new(),
290            ),
291            PipelineFamily::new(
292                "image-upscale",
293                &[
294                    "StableDiffusionUpscalePipeline",
295                    "StableDiffusionLatentUpscalePipeline",
296                ],
297                Modality::image(),
298                Vec::new(),
299                Vec::new(),
300                Vec::new(),
301            ),
302            PipelineFamily::new(
303                "video",
304                &[
305                    "TextToVideoSDPipeline",
306                    "AnimateDiffPipeline",
307                    "CogVideoXPipeline",
308                    "StableVideoDiffusionPipeline",
309                    "HunyuanVideoPipeline",
310                    "LTXPipeline",
311                    "WanPipeline",
312                ],
313                Modality::video(),
314                Vec::new(),
315                Vec::new(),
316                Vec::new(),
317            ),
318            PipelineFamily::new(
319                "audio",
320                &[
321                    "AudioLDMPipeline",
322                    "AudioLDM2Pipeline",
323                    "MusicLDMPipeline",
324                    "StableAudioPipeline",
325                ],
326                Modality::audio(),
327                Vec::new(),
328                Vec::new(),
329                Vec::new(),
330            ),
331        ])
332    }
333}
334
335fn int_param(key: &str, default: i64, min: i64, max: i64) -> ParamSpec {
336    ParamSpec {
337        key: key.to_owned(),
338        param_type: ParamType::Int,
339        default_value: Some(JsonValue::Int(default)),
340        range: Some(vec![JsonValue::Int(min), JsonValue::Int(max)]),
341        values: None,
342    }
343}
344
345fn float_param(key: &str, default: f64, min: f64, max: f64) -> ParamSpec {
346    ParamSpec {
347        key: key.to_owned(),
348        param_type: ParamType::Float,
349        default_value: Some(JsonValue::Double(default)),
350        range: Some(vec![JsonValue::Double(min), JsonValue::Double(max)]),
351        values: None,
352    }
353}
354
355fn size_param(default: &str, values: &[&str]) -> ParamSpec {
356    ParamSpec {
357        key: "size".to_owned(),
358        param_type: ParamType::Enum,
359        default_value: Some(JsonValue::String(default.to_owned())),
360        range: None,
361        values: Some(values.iter().map(|value| (*value).to_owned()).collect()),
362    }
363}
364
365/// A bare parameter with only a key and type (no default/range/values) — `seed`
366/// and `negative_prompt`.
367fn bare_param(key: &str, param_type: ParamType) -> ParamSpec {
368    ParamSpec {
369        key: key.to_owned(),
370        param_type,
371        default_value: None,
372        range: None,
373        values: None,
374    }
375}
376
377fn seed() -> ParamSpec {
378    bare_param("seed", ParamType::Int)
379}
380
381fn negative_prompt() -> ParamSpec {
382    bare_param("negative_prompt", ParamType::String)
383}
384
385fn flux_params() -> Vec<ParamSpec> {
386    vec![
387        int_param("steps", 4, 1, 50),
388        float_param("guidance", 4.0, 0.0, 10.0),
389        size_param("1024x1024", &["512x512", "768x768", "1024x1024"]),
390        seed(),
391    ]
392}
393
394fn sd1_params() -> Vec<ParamSpec> {
395    vec![
396        int_param("steps", 30, 1, 75),
397        float_param("guidance", 7.5, 0.0, 15.0),
398        size_param("512x512", &["512x512", "576x576", "640x640", "768x768"]),
399        seed(),
400        negative_prompt(),
401    ]
402}
403
404fn sdxl_params() -> Vec<ParamSpec> {
405    vec![
406        int_param("steps", 30, 1, 75),
407        float_param("guidance", 7.0, 0.0, 15.0),
408        size_param(
409            "1024x1024",
410            &["768x768", "1024x1024", "1152x896", "896x1152"],
411        ),
412        seed(),
413        negative_prompt(),
414    ]
415}
416
417fn sd3_params() -> Vec<ParamSpec> {
418    vec![
419        int_param("steps", 28, 1, 75),
420        float_param("guidance", 7.0, 0.0, 15.0),
421        size_param(
422            "1024x1024",
423            &["768x768", "1024x1024", "1152x896", "896x1152"],
424        ),
425        seed(),
426        negative_prompt(),
427    ]
428}
429
430fn pixart_params() -> Vec<ParamSpec> {
431    vec![
432        int_param("steps", 20, 1, 75),
433        float_param("guidance", 4.5, 0.0, 15.0),
434        size_param("1024x1024", &["512x512", "768x768", "1024x1024"]),
435        seed(),
436        negative_prompt(),
437    ]
438}
439
440fn kandinsky_params() -> Vec<ParamSpec> {
441    vec![
442        int_param("steps", 30, 1, 75),
443        float_param("guidance", 4.0, 0.0, 15.0),
444        size_param("768x768", &["512x512", "768x768", "1024x1024"]),
445        seed(),
446        negative_prompt(),
447    ]
448}
449
450fn lcm_params() -> Vec<ParamSpec> {
451    vec![
452        int_param("steps", 4, 1, 8),
453        float_param("guidance", 1.5, 0.0, 2.0),
454        size_param("512x512", &["512x512", "768x768"]),
455        seed(),
456    ]
457}
458
459/// The "turbo/lightning/lcm" low-step refinement shared by the SD and SDXL
460/// families: fewer steps and (near-)zero guidance under the trailing-spacing
461/// ancestral scheduler.
462fn turbo_refinement() -> PipelineRefinement {
463    PipelineRefinement {
464        scheduler_classes: ["EulerAncestralDiscreteScheduler"]
465            .into_iter()
466            .map(str::to_owned)
467            .collect(),
468        timestep_spacing: Some("trailing".to_owned()),
469        name_signals: ["turbo", "lightning", "lcm"]
470            .into_iter()
471            .map(str::to_owned)
472            .collect(),
473        param_overrides: vec![
474            int_param("steps", 2, 1, 8),
475            float_param("guidance", 0.0, 0.0, 2.0),
476        ],
477    }
478}
kernel/resolution/pipelines.rs

kernel/resolution/
pipelines.rs
