rigspark_core/
generation.rs

1//! Curated, offline catalog of local image/video generation models served by ComfyUI.
2//!
3//! Kept separate from the LLM catalog so LLM ranking, signing, and refresh are unaffected.
4use crate::{
5    catalog::{
6        EntryProvenance, LICENSES, RecencyBasis, coordinates, date, digest, model_file,
7        parameter_label, parse_document, recency, recency_label, require, timestamp,
8    },
9    reports::{strip_control, table},
10    sizing::{HEADROOM, Hardware, ValidationError, memory_capacity},
11};
12use serde::{Deserialize, Serialize};
13use std::collections::HashSet;
14
15const MAX_SAFE_INTEGER: u64 = 9_007_199_254_740_991;
16
17#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Deserialize, Serialize)]
18#[serde(rename_all = "lowercase")]
19pub enum GenerationKind {
20    Image,
21    Video,
22}
23impl GenerationKind {
24    pub fn name(self) -> &'static str {
25        match self {
26            Self::Image => "image",
27            Self::Video => "video",
28        }
29    }
30    pub fn parse(value: &str) -> Option<Self> {
31        match value {
32            "image" => Some(Self::Image),
33            "video" => Some(Self::Video),
34            _ => None,
35        }
36    }
37}
38
39/// Built-in ComfyUI workflow graphs implemented by `rigspark-runtime`.
40#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, Serialize)]
41#[serde(rename_all = "kebab-case")]
42pub enum Workflow {
43    FluxCheckpoint,
44    FluxSplit,
45    WanT2v,
46    /// Wan 2.2 TI2V-5B text-to-video (https://comfyanonymous.github.io/ComfyUI_examples/wan22/).
47    #[serde(rename = "wan22-ti2v")]
48    Wan22Ti2v,
49    /// Qwen-Image text-to-image (https://comfyanonymous.github.io/ComfyUI_examples/qwen_image/).
50    QwenImage,
51    /// Candidate only; bundled entries stay fit-only until live generation is verified.
52    MingDesign,
53}
54impl Workflow {
55    pub fn kind(self) -> GenerationKind {
56        match self {
57            Self::FluxCheckpoint | Self::FluxSplit | Self::QwenImage | Self::MingDesign => {
58                GenerationKind::Image
59            }
60            Self::WanT2v | Self::Wan22Ti2v => GenerationKind::Video,
61        }
62    }
63    /// Exact file roles the workflow loads, in a stable order.
64    pub fn roles(self) -> &'static [FileRole] {
65        match self {
66            Self::FluxCheckpoint => &[FileRole::Checkpoint],
67            Self::FluxSplit => &[
68                FileRole::Diffusion,
69                FileRole::TextEncoder,
70                FileRole::Clip,
71                FileRole::Vae,
72            ],
73            Self::WanT2v | Self::Wan22Ti2v | Self::QwenImage | Self::MingDesign => {
74                &[FileRole::Diffusion, FileRole::TextEncoder, FileRole::Vae]
75            }
76        }
77    }
78}
79
80#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Deserialize, Serialize)]
81#[serde(rename_all = "kebab-case")]
82pub enum FileRole {
83    Checkpoint,
84    Diffusion,
85    TextEncoder,
86    Clip,
87    Vae,
88}
89impl FileRole {
90    /// ComfyUI `models/` subfolder the loader node reads this role from.
91    pub fn folder(self) -> &'static str {
92        match self {
93            Self::Checkpoint => "checkpoints",
94            Self::Diffusion => "diffusion_models",
95            Self::TextEncoder => "text_encoders",
96            Self::Clip => "text_encoders",
97            Self::Vae => "vae",
98        }
99    }
100}
101
102#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
103#[serde(rename_all = "camelCase", deny_unknown_fields)]
104pub struct GenerationFile {
105    pub role: FileRole,
106    pub folder: String,
107    pub repo: String,
108    pub revision: String,
109    pub file: String,
110    pub sha256: String,
111    pub bytes: u64,
112}
113
114#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
115#[serde(rename_all = "camelCase", deny_unknown_fields)]
116pub struct GenerationModel {
117    pub id: String,
118    pub family: String,
119    pub kind: GenerationKind,
120    pub params: String,
121    pub license: String,
122    pub open_weight: bool,
123    #[serde(default, skip_serializing_if = "Option::is_none")]
124    pub release_date: Option<String>,
125    #[serde(default, skip_serializing_if = "Option::is_none")]
126    pub added_at: Option<String>,
127    #[serde(default, skip_serializing_if = "EntryProvenance::is_curated")]
128    pub provenance: EntryProvenance,
129    #[serde(default)]
130    pub default: bool,
131    pub source: String,
132    /// Absent for fit-only entries: sourced weights without a built-in workflow yet.
133    #[serde(default, skip_serializing_if = "Option::is_none")]
134    pub workflow: Option<Workflow>,
135    #[serde(default, skip_serializing_if = "Option::is_none")]
136    pub workflow_source: Option<String>,
137    pub files: Vec<GenerationFile>,
138}
139
140pub const GENERATION_SCHEMA_VERSION: u8 = 2;
141
142#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
143#[serde(rename_all = "camelCase", deny_unknown_fields)]
144pub struct GenerationCatalog {
145    pub schema_version: u8,
146    pub generated_at: String,
147    pub models: Vec<GenerationModel>,
148}
149
150fn https(value: &str) -> Result<(), ValidationError> {
151    require(
152        value.len() <= 2048
153            && value.starts_with("https://")
154            && value.len() > "https://".len()
155            && !value
156                .chars()
157                .any(|ch| ch.is_control() || ch.is_whitespace()),
158        "source links must be HTTPS",
159    )
160}
161fn model_id(value: &str) -> bool {
162    !value.is_empty()
163        && value.len() <= 128
164        && value.as_bytes()[0].is_ascii_alphanumeric()
165        && value.bytes().all(|byte| {
166            byte.is_ascii_lowercase() || byte.is_ascii_digit() || b"._-:".contains(&byte)
167        })
168        && GenerationKind::parse(value).is_none()
169}
170
171impl GenerationModel {
172    pub fn recency(&self) -> Option<(&str, RecencyBasis)> {
173        recency(self.release_date.as_deref(), self.added_at.as_deref())
174    }
175    pub fn recency_label(&self) -> String {
176        recency_label(self.recency())
177    }
178    /// The built-in workflow that runs this model, or why it cannot run yet.
179    pub fn runnable(&self) -> Result<Workflow, ValidationError> {
180        self.workflow.ok_or_else(|| {
181            ValidationError(format!(
182                "{} is fit-only (workflow coming): RigSpark has no built-in ComfyUI workflow for this architecture yet",
183                self.id
184            ))
185        })
186    }
187    pub fn total_bytes(&self) -> u64 {
188        self.files.iter().map(|file| file.bytes).sum()
189    }
190    pub fn largest_bytes(&self) -> u64 {
191        self.files.iter().map(|file| file.bytes).max().unwrap_or(0)
192    }
193    pub fn file(&self, role: FileRole) -> Option<&GenerationFile> {
194        self.files.iter().find(|file| file.role == role)
195    }
196    fn validate(&self) -> Result<(), ValidationError> {
197        require(model_id(&self.id), "invalid generation model id")?;
198        require(
199            !self.family.is_empty() && self.family.len() <= 128,
200            "empty text field",
201        )?;
202        require(parameter_label(&self.params), "invalid parameter label")?;
203        require(
204            self.open_weight && LICENSES.contains(&self.license.as_str()),
205            "generation catalog requires an allowlisted open-weight license",
206        )?;
207        if let Some(released) = &self.release_date {
208            date(released)?;
209        }
210        if let Some(added) = &self.added_at {
211            date(added)?;
212        }
213        match self.provenance {
214            EntryProvenance::Curated => require(
215                self.release_date.is_some() && self.workflow.is_some(),
216                "curated generation entries need releaseDate and a workflow",
217            )?,
218            EntryProvenance::Auto => require(
219                self.added_at.is_some(),
220                "auto generation entries need addedAt",
221            )?,
222        }
223        https(&self.source)?;
224        require(
225            self.workflow.is_some() == self.workflow_source.is_some(),
226            "workflowSource must accompany a workflow",
227        )?;
228        let roles: Vec<_> = self.files.iter().map(|file| file.role).collect();
229        match (self.workflow, &self.workflow_source) {
230            (Some(workflow), Some(source)) => {
231                https(source)?;
232                require(
233                    roles.len() == workflow.roles().len()
234                        && workflow.roles().iter().all(|role| roles.contains(role)),
235                    "files must match the workflow roles exactly",
236                )?;
237                require(
238                    workflow.kind() == self.kind,
239                    "workflow does not produce this model kind",
240                )?;
241            }
242            _ => require(
243                !roles.is_empty()
244                    && roles.len() <= 8
245                    && roles.iter().collect::<HashSet<_>>().len() == roles.len(),
246                "fit-only entries need sourced files with distinct roles",
247            )?,
248        }
249        let mut total = 0u64;
250        for file in &self.files {
251            require(
252                file.folder == file.role.folder(),
253                "file folder does not match its role",
254            )?;
255            coordinates(&file.repo, &file.revision)?;
256            model_file(&file.file, 512)?;
257            require(
258                file.file.ends_with(".safetensors"),
259                "generation weights must be .safetensors files",
260            )?;
261            digest(&file.sha256)?;
262            require(
263                file.bytes > 0 && file.bytes <= MAX_SAFE_INTEGER,
264                "invalid weight file size",
265            )?;
266            total = total
267                .checked_add(file.bytes)
268                .filter(|total| *total <= MAX_SAFE_INTEGER)
269                .ok_or_else(|| ValidationError("invalid weight file size".into()))?;
270        }
271        Ok(())
272    }
273}
274
275impl GenerationCatalog {
276    pub fn parse(raw: &str) -> Result<Self, ValidationError> {
277        let mut result: Self = parse_document(raw)?;
278        require(
279            (1..=GENERATION_SCHEMA_VERSION).contains(&result.schema_version)
280                && !result.models.is_empty()
281                && result.models.len() <= 1024,
282            "unsupported or empty generation catalog",
283        )?;
284        timestamp(&result.generated_at)?;
285        let mut ids = HashSet::new();
286        let mut defaults = HashSet::new();
287        let mut kinds = HashSet::new();
288        for model in &result.models {
289            require(
290                result.schema_version >= 2
291                    || (model.provenance.is_curated() && model.added_at.is_none()),
292                "schema v1 generation catalogs cannot carry v2 fields",
293            )?;
294            model.validate()?;
295            require(
296                !model.default || model.workflow.is_some(),
297                "a default model must be runnable",
298            )?;
299            require(ids.insert(&model.id), "duplicate generation model id")?;
300            kinds.insert(model.kind);
301            if model.default {
302                require(
303                    defaults.insert(model.kind),
304                    "each kind needs exactly one default model",
305                )?;
306            }
307        }
308        require(
309            kinds == defaults,
310            "each kind needs exactly one default model",
311        )?;
312        for model in &mut result.models {
313            model.family = strip_control(&model.family);
314        }
315        Ok(result)
316    }
317    pub fn bundled() -> Result<Self, ValidationError> {
318        Self::parse(crate::GENERATION_JSON)
319    }
320    /// Resolves `image`/`video` to that kind's default model, otherwise an exact id.
321    pub fn resolve(&self, query: &str) -> Result<&GenerationModel, ValidationError> {
322        let found = match GenerationKind::parse(query) {
323            Some(kind) => self
324                .models
325                .iter()
326                .find(|model| model.kind == kind && model.default),
327            None => self.models.iter().find(|model| model.id == query),
328        };
329        found.ok_or_else(|| {
330            ValidationError(
331                "unknown generation model; use image, video, or an id from `rigspark catalog --generation`"
332                    .into(),
333            )
334        })
335    }
336}
337
338#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
339#[serde(rename_all = "camelCase")]
340pub struct RejectedGenerationCandidate {
341    pub id: String,
342    pub reason: String,
343}
344
345#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
346#[serde(rename_all = "camelCase")]
347pub struct GenerationEnrichment {
348    pub catalog: GenerationCatalog,
349    pub updated: Vec<String>,
350    pub rejected: Vec<RejectedGenerationCandidate>,
351}
352
353fn built_in_family(model: &GenerationModel) -> bool {
354    matches!(
355        (model.kind, model.workflow, model.family.as_str()),
356        (
357            GenerationKind::Image,
358            Some(Workflow::FluxCheckpoint | Workflow::FluxSplit),
359            "flux1-schnell"
360        ) | (GenerationKind::Video, Some(Workflow::WanT2v), "wan2.1-t2v")
361    )
362}
363
364/// Applies already-collected, pinned metadata to the bundled generation catalog.
365///
366/// Collection is intentionally outside `rigspark-core`; this function remains pure and
367/// accepts only complete models for workflow families implemented by the runtime.
368pub fn enrich_generation_catalog(
369    catalog: &GenerationCatalog,
370    candidates: Vec<GenerationModel>,
371    kind: GenerationKind,
372    now: &str,
373) -> Result<GenerationEnrichment, ValidationError> {
374    timestamp(now)?;
375    let mut next = catalog.clone();
376    let mut updated = Vec::new();
377    let mut rejected = Vec::new();
378
379    for mut candidate in candidates {
380        let reason = if candidate.kind != kind {
381            Some("candidate does not match the requested kind".into())
382        } else if !built_in_family(&candidate) {
383            Some("candidate has no built-in workflow family".into())
384        } else {
385            candidate.validate().err().map(|error| error.to_string())
386        };
387        if let Some(reason) = reason {
388            rejected.push(RejectedGenerationCandidate {
389                id: strip_control(&candidate.id),
390                reason,
391            });
392            continue;
393        }
394
395        if let Some(existing) = next
396            .models
397            .iter_mut()
398            .find(|model| model.id == candidate.id)
399        {
400            candidate.default = existing.default;
401            if *existing != candidate {
402                updated.push(candidate.id.clone());
403                *existing = candidate;
404            }
405        } else {
406            updated.push(candidate.id.clone());
407            next.models.push(candidate);
408        }
409    }
410
411    updated.sort();
412    updated.dedup();
413    rejected.sort_by(|left, right| left.id.cmp(&right.id));
414    if !updated.is_empty() {
415        next.generated_at = now.into();
416        next.models.sort_by(|left, right| left.id.cmp(&right.id));
417        next = GenerationCatalog::parse(
418            &serde_json::to_string(&next)
419                .map_err(|_| ValidationError("invalid generation catalog".into()))?,
420        )?;
421    }
422    Ok(GenerationEnrichment {
423        catalog: next,
424        updated,
425        rejected,
426    })
427}
428
429#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
430#[serde(rename_all = "lowercase")]
431pub enum FitVerdict {
432    Yes,
433    Slow,
434    No,
435}
436impl FitVerdict {
437    pub fn name(self) -> &'static str {
438        match self {
439            Self::Yes => "yes",
440            Self::Slow => "slow",
441            Self::No => "no",
442        }
443    }
444}
445
446/// Memory-only fit of a model's weights; generation speed is never estimated.
447#[derive(Debug, Clone, Serialize)]
448#[serde(rename_all = "camelCase")]
449pub struct GenerationFit {
450    pub verdict: FitVerdict,
451    pub reason: &'static str,
452    pub memory_kind: &'static str,
453    pub usable_bytes: f64,
454    pub budget_bytes: f64,
455    pub total_bytes: u64,
456    pub largest_bytes: u64,
457    /// Always `"unknown"`: no sourced generation throughput dataset exists.
458    pub speed: &'static str,
459}
460
461fn system_ram_budget(hardware: &Hardware) -> f64 {
462    let mut cpu = hardware.clone();
463    cpu.gpu.clear();
464    cpu.unified_memory = Some(false);
465    memory_capacity(&cpu).1 * (1.0 - HEADROOM)
466}
467
468pub fn fit(model: &GenerationModel, hardware: &Hardware) -> GenerationFit {
469    let (memory_kind, usable_bytes) = memory_capacity(hardware);
470    let budget_bytes = usable_bytes * (1.0 - HEADROOM);
471    let total_bytes = model.total_bytes();
472    let largest_bytes = model.largest_bytes();
473    let (verdict, reason) = if total_bytes as f64 <= budget_bytes {
474        (FitVerdict::Yes, "all weights fit in memory together")
475    } else if largest_bytes as f64 <= budget_bytes {
476        (
477            FitVerdict::Slow,
478            "weights fit one stage at a time; ComfyUI swaps models between stages",
479        )
480    } else if memory_kind == "vram" && largest_bytes as f64 <= system_ram_budget(hardware) {
481        (
482            FitVerdict::Slow,
483            "largest weight file exceeds VRAM; ComfyUI offloads to system RAM",
484        )
485    } else {
486        (FitVerdict::No, "largest weight file exceeds usable memory")
487    };
488    GenerationFit {
489        verdict,
490        reason,
491        memory_kind,
492        usable_bytes,
493        budget_bytes,
494        total_bytes,
495        largest_bytes,
496        speed: "unknown",
497    }
498}
499
500pub fn catalog_text(catalog: &GenerationCatalog, hardware: &Hardware) -> String {
501    let mut models: Vec<_> = catalog.models.iter().collect();
502    models.sort_by(|left, right| {
503        left.kind
504            .name()
505            .cmp(right.kind.name())
506            .then_with(|| right.release_date.cmp(&left.release_date))
507            .then_with(|| left.id.cmp(&right.id))
508    });
509    let rows = models
510        .into_iter()
511        .map(|model| {
512            let fit = fit(model, hardware);
513            vec![
514                model.id.clone(),
515                model.kind.name().into(),
516                model.params.clone(),
517                format!("{:.1}", model.total_bytes() as f64 / 1073741824.0),
518                fit.verdict.name().into(),
519                fit.speed.into(),
520                model.license.clone(),
521                model.recency_label(),
522                if model.workflow.is_some() {
523                    "runs".into()
524                } else {
525                    "workflow coming".into()
526                },
527            ]
528        })
529        .collect();
530    format!(
531        "Generation catalog (runtime: ComfyUI, local only; models: {})\n{}\n\
532         Fit covers weight memory only; generation speed is unknown (no sourced data).\n\
533         Run: rigspark generate image --prompt \"...\" --comfyui-dir <ComfyUI directory>\n",
534        catalog.models.len(),
535        table(
536            &[
537                ("Model", false),
538                ("Kind", false),
539                ("Params", false),
540                ("Weights GiB", true),
541                ("Fit", false),
542                ("Speed", false),
543                ("License", false),
544                ("Release", false),
545                ("Generate", false),
546            ],
547            rows,
548        )
549    )
550}
rigspark_core/generation.rs

rigspark_core/
generation.rs
