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        LICENSES, coordinates, date, digest, model_file, parse_document, require, timestamp,
7    },
8    reports::{strip_control, table},
9    sizing::{HEADROOM, Hardware, ValidationError, memory_capacity},
10};
11use serde::{Deserialize, Serialize};
12use std::collections::HashSet;
13
14const MAX_SAFE_INTEGER: u64 = 9_007_199_254_740_991;
15
16#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Deserialize, Serialize)]
17#[serde(rename_all = "lowercase")]
18pub enum GenerationKind {
19    Image,
20    Video,
21}
22impl GenerationKind {
23    pub fn name(self) -> &'static str {
24        match self {
25            Self::Image => "image",
26            Self::Video => "video",
27        }
28    }
29    pub fn parse(value: &str) -> Option<Self> {
30        match value {
31            "image" => Some(Self::Image),
32            "video" => Some(Self::Video),
33            _ => None,
34        }
35    }
36}
37
38/// Built-in ComfyUI workflow graphs implemented by `rigspark-runtime`.
39#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, Serialize)]
40#[serde(rename_all = "kebab-case")]
41pub enum Workflow {
42    FluxCheckpoint,
43    WanT2v,
44}
45impl Workflow {
46    pub fn kind(self) -> GenerationKind {
47        match self {
48            Self::FluxCheckpoint => GenerationKind::Image,
49            Self::WanT2v => GenerationKind::Video,
50        }
51    }
52    /// Exact file roles the workflow loads, in a stable order.
53    pub fn roles(self) -> &'static [FileRole] {
54        match self {
55            Self::FluxCheckpoint => &[FileRole::Checkpoint],
56            Self::WanT2v => &[FileRole::Diffusion, FileRole::TextEncoder, FileRole::Vae],
57        }
58    }
59}
60
61#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Deserialize, Serialize)]
62#[serde(rename_all = "kebab-case")]
63pub enum FileRole {
64    Checkpoint,
65    Diffusion,
66    TextEncoder,
67    Vae,
68}
69impl FileRole {
70    /// ComfyUI `models/` subfolder the loader node reads this role from.
71    pub fn folder(self) -> &'static str {
72        match self {
73            Self::Checkpoint => "checkpoints",
74            Self::Diffusion => "diffusion_models",
75            Self::TextEncoder => "text_encoders",
76            Self::Vae => "vae",
77        }
78    }
79}
80
81#[derive(Debug, Clone, Deserialize, Serialize)]
82#[serde(rename_all = "camelCase", deny_unknown_fields)]
83pub struct GenerationFile {
84    pub role: FileRole,
85    pub folder: String,
86    pub repo: String,
87    pub revision: String,
88    pub file: String,
89    pub sha256: String,
90    pub bytes: u64,
91}
92
93#[derive(Debug, Clone, Deserialize, Serialize)]
94#[serde(rename_all = "camelCase", deny_unknown_fields)]
95pub struct GenerationModel {
96    pub id: String,
97    pub family: String,
98    pub kind: GenerationKind,
99    pub params: String,
100    pub license: String,
101    pub open_weight: bool,
102    pub release_date: String,
103    #[serde(default)]
104    pub default: bool,
105    pub source: String,
106    pub workflow: Workflow,
107    pub workflow_source: String,
108    pub files: Vec<GenerationFile>,
109}
110
111#[derive(Debug, Clone, Deserialize, Serialize)]
112#[serde(rename_all = "camelCase", deny_unknown_fields)]
113pub struct GenerationCatalog {
114    pub schema_version: u8,
115    pub generated_at: String,
116    pub models: Vec<GenerationModel>,
117}
118
119fn https(value: &str) -> Result<(), ValidationError> {
120    require(
121        value.len() <= 2048
122            && value.starts_with("https://")
123            && value.len() > "https://".len()
124            && !value
125                .chars()
126                .any(|ch| ch.is_control() || ch.is_whitespace()),
127        "source links must be HTTPS",
128    )
129}
130fn model_id(value: &str) -> bool {
131    !value.is_empty()
132        && value.len() <= 128
133        && value.as_bytes()[0].is_ascii_alphanumeric()
134        && value.bytes().all(|byte| {
135            byte.is_ascii_lowercase() || byte.is_ascii_digit() || b"._:-".contains(&byte)
136        })
137        && GenerationKind::parse(value).is_none()
138}
139
140impl GenerationModel {
141    pub fn total_bytes(&self) -> u64 {
142        self.files.iter().map(|file| file.bytes).sum()
143    }
144    pub fn largest_bytes(&self) -> u64 {
145        self.files.iter().map(|file| file.bytes).max().unwrap_or(0)
146    }
147    pub fn file(&self, role: FileRole) -> Option<&GenerationFile> {
148        self.files.iter().find(|file| file.role == role)
149    }
150    fn validate(&self) -> Result<(), ValidationError> {
151        require(model_id(&self.id), "invalid generation model id")?;
152        require(
153            !self.family.is_empty() && self.family.len() <= 128,
154            "empty text field",
155        )?;
156        require(
157            regex::Regex::new(r"^\d+(\.\d+)?[BMT]$").is_ok_and(|re| re.is_match(&self.params)),
158            "invalid parameter label",
159        )?;
160        require(
161            self.open_weight && LICENSES.contains(&self.license.as_str()),
162            "generation catalog requires an allowlisted open-weight license",
163        )?;
164        date(&self.release_date)?;
165        https(&self.source)?;
166        https(&self.workflow_source)?;
167        let roles: Vec<_> = self.files.iter().map(|file| file.role).collect();
168        require(
169            roles.len() == self.workflow.roles().len()
170                && self
171                    .workflow
172                    .roles()
173                    .iter()
174                    .all(|role| roles.contains(role)),
175            "files must match the workflow roles exactly",
176        )?;
177        require(
178            self.workflow.kind() == self.kind,
179            "workflow does not produce this model kind",
180        )?;
181        let mut total = 0u64;
182        for file in &self.files {
183            require(
184                file.folder == file.role.folder(),
185                "file folder does not match its role",
186            )?;
187            coordinates(&file.repo, &file.revision)?;
188            model_file(&file.file, 512)?;
189            require(
190                file.file.ends_with(".safetensors"),
191                "generation weights must be .safetensors files",
192            )?;
193            digest(&file.sha256)?;
194            require(
195                file.bytes > 0 && file.bytes <= MAX_SAFE_INTEGER,
196                "invalid weight file size",
197            )?;
198            total = total
199                .checked_add(file.bytes)
200                .filter(|total| *total <= MAX_SAFE_INTEGER)
201                .ok_or_else(|| ValidationError("invalid weight file size".into()))?;
202        }
203        Ok(())
204    }
205}
206
207impl GenerationCatalog {
208    pub fn parse(raw: &str) -> Result<Self, ValidationError> {
209        let mut result: Self = parse_document(raw)?;
210        require(
211            result.schema_version == 1 && !result.models.is_empty() && result.models.len() <= 256,
212            "unsupported or empty generation catalog",
213        )?;
214        timestamp(&result.generated_at)?;
215        let mut ids = HashSet::new();
216        let mut defaults = HashSet::new();
217        let mut kinds = HashSet::new();
218        for model in &result.models {
219            model.validate()?;
220            require(ids.insert(&model.id), "duplicate generation model id")?;
221            kinds.insert(model.kind);
222            if model.default {
223                require(
224                    defaults.insert(model.kind),
225                    "each kind needs exactly one default model",
226                )?;
227            }
228        }
229        require(
230            kinds == defaults,
231            "each kind needs exactly one default model",
232        )?;
233        for model in &mut result.models {
234            model.family = strip_control(&model.family);
235        }
236        Ok(result)
237    }
238    pub fn bundled() -> Result<Self, ValidationError> {
239        Self::parse(crate::GENERATION_JSON)
240    }
241    /// Resolves `image`/`video` to that kind's default model, otherwise an exact id.
242    pub fn resolve(&self, query: &str) -> Result<&GenerationModel, ValidationError> {
243        let found = match GenerationKind::parse(query) {
244            Some(kind) => self
245                .models
246                .iter()
247                .find(|model| model.kind == kind && model.default),
248            None => self.models.iter().find(|model| model.id == query),
249        };
250        found.ok_or_else(|| {
251            ValidationError(
252                "unknown generation model; use image, video, or an id from `rigspark catalog --generation`"
253                    .into(),
254            )
255        })
256    }
257}
258
259#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
260#[serde(rename_all = "lowercase")]
261pub enum FitVerdict {
262    Yes,
263    Slow,
264    No,
265}
266impl FitVerdict {
267    pub fn name(self) -> &'static str {
268        match self {
269            Self::Yes => "yes",
270            Self::Slow => "slow",
271            Self::No => "no",
272        }
273    }
274}
275
276/// Memory-only fit of a model's weights; generation speed is never estimated.
277#[derive(Debug, Clone, Serialize)]
278#[serde(rename_all = "camelCase")]
279pub struct GenerationFit {
280    pub verdict: FitVerdict,
281    pub reason: &'static str,
282    pub memory_kind: &'static str,
283    pub usable_bytes: f64,
284    pub budget_bytes: f64,
285    pub total_bytes: u64,
286    pub largest_bytes: u64,
287    /// Always `"unknown"`: no sourced generation throughput dataset exists.
288    pub speed: &'static str,
289}
290
291fn system_ram_budget(hardware: &Hardware) -> f64 {
292    let mut cpu = hardware.clone();
293    cpu.gpu.clear();
294    cpu.unified_memory = Some(false);
295    memory_capacity(&cpu).1 * (1.0 - HEADROOM)
296}
297
298pub fn fit(model: &GenerationModel, hardware: &Hardware) -> GenerationFit {
299    let (memory_kind, usable_bytes) = memory_capacity(hardware);
300    let budget_bytes = usable_bytes * (1.0 - HEADROOM);
301    let total_bytes = model.total_bytes();
302    let largest_bytes = model.largest_bytes();
303    let (verdict, reason) = if total_bytes as f64 <= budget_bytes {
304        (FitVerdict::Yes, "all weights fit in memory together")
305    } else if largest_bytes as f64 <= budget_bytes {
306        (
307            FitVerdict::Slow,
308            "weights fit one stage at a time; ComfyUI swaps models between stages",
309        )
310    } else if memory_kind == "vram" && largest_bytes as f64 <= system_ram_budget(hardware) {
311        (
312            FitVerdict::Slow,
313            "largest weight file exceeds VRAM; ComfyUI offloads to system RAM",
314        )
315    } else {
316        (FitVerdict::No, "largest weight file exceeds usable memory")
317    };
318    GenerationFit {
319        verdict,
320        reason,
321        memory_kind,
322        usable_bytes,
323        budget_bytes,
324        total_bytes,
325        largest_bytes,
326        speed: "unknown",
327    }
328}
329
330pub fn catalog_text(catalog: &GenerationCatalog, hardware: &Hardware) -> String {
331    let mut models: Vec<_> = catalog.models.iter().collect();
332    models.sort_by(|left, right| {
333        left.kind
334            .name()
335            .cmp(right.kind.name())
336            .then_with(|| right.release_date.cmp(&left.release_date))
337            .then_with(|| left.id.cmp(&right.id))
338    });
339    let rows = models
340        .into_iter()
341        .map(|model| {
342            let fit = fit(model, hardware);
343            vec![
344                model.id.clone(),
345                model.kind.name().into(),
346                model.params.clone(),
347                format!("{:.1}", model.total_bytes() as f64 / 1073741824.0),
348                fit.verdict.name().into(),
349                fit.speed.into(),
350                model.license.clone(),
351                model.release_date.clone(),
352            ]
353        })
354        .collect();
355    format!(
356        "Generation catalog (runtime: ComfyUI, local only; models: {})\n{}\n\
357         Fit covers weight memory only; generation speed is unknown (no sourced data).\n\
358         Run: rigspark generate image --prompt \"...\" --comfyui-dir <ComfyUI directory>\n",
359        catalog.models.len(),
360        table(
361            &[
362                ("Model", false),
363                ("Kind", false),
364                ("Params", false),
365                ("Weights GiB", true),
366                ("Fit", false),
367                ("Speed", false),
368                ("License", false),
369                ("Release", false),
370            ],
371            rows,
372        )
373    )
374}
rigspark_core/generation.rs

rigspark_core/
generation.rs
