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

rigspark_core/
generation.rs
