rigspark_core/
generation_admission.rs

1//! Pure parsers and builders for admitting Comfy-Org image and video releases as
2//! fit-only generation entries. Network access lives in `rigspark-runtime`.
3
4use crate::{
5    catalog::{EntryProvenance, LICENSES, require},
6    generation::{FileRole, GenerationFile, GenerationKind, GenerationModel, Workflow},
7    sizing::ValidationError,
8};
9use serde::Deserialize;
10use serde_json::Value;
11
12pub const PUBLISHER: &str = "Comfy-Org";
13pub const MAX_METADATA_BYTES: usize = 4 * 1024 * 1024;
14pub const MAX_SAFETENSORS_HEADER_BYTES: u64 = 100 * 1024 * 1024;
15
16/// Repositories from `huggingface.co/api/models?author=Comfy-Org&sort=createdAt`, newest first.
17pub fn parse_listing(raw: &str) -> Result<Vec<String>, ValidationError> {
18    require(
19        raw.len() <= MAX_METADATA_BYTES,
20        "model listing exceeds 4 MiB",
21    )?;
22    #[derive(Deserialize)]
23    struct Item {
24        id: String,
25    }
26    let items: Vec<Item> =
27        serde_json::from_str(raw).map_err(|error| ValidationError(error.to_string()))?;
28    require(items.len() <= 1000, "model listing has too many entries")?;
29    Ok(items
30        .into_iter()
31        .map(|item| item.id)
32        .filter(|id| repository(id) && id.starts_with(&format!("{PUBLISHER}/")))
33        .collect())
34}
35
36fn repository(id: &str) -> bool {
37    let mut parts = id.split('/');
38    let valid = |part: &str| {
39        !part.is_empty()
40            && part.len() <= 96
41            && part.as_bytes()[0].is_ascii_alphanumeric()
42            && part
43                .bytes()
44                .all(|byte| byte.is_ascii_alphanumeric() || b"._-".contains(&byte))
45    };
46    matches!((parts.next(), parts.next(), parts.next()), (Some(owner), Some(name), None) if valid(owner) && valid(name))
47}
48
49#[derive(Debug, Clone, Deserialize)]
50pub struct Lfs {
51    pub sha256: String,
52    pub size: u64,
53}
54#[derive(Debug, Clone, Deserialize)]
55pub struct Sibling {
56    pub rfilename: String,
57    #[serde(default)]
58    pub lfs: Option<Lfs>,
59}
60#[derive(Debug, Clone, Deserialize)]
61pub struct HfModel {
62    pub id: String,
63    pub sha: String,
64    #[serde(default)]
65    pub tags: Vec<String>,
66    #[serde(default)]
67    pub gated: Value,
68    #[serde(default)]
69    pub private: bool,
70    #[serde(default)]
71    pub siblings: Vec<Sibling>,
72    #[serde(default)]
73    pub pipeline_tag: Option<String>,
74}
75impl HfModel {
76    pub fn parse(raw: &str) -> Result<Self, ValidationError> {
77        require(
78            raw.len() <= MAX_METADATA_BYTES,
79            "model metadata exceeds 4 MiB",
80        )?;
81        let model: Self =
82            serde_json::from_str(raw).map_err(|error| ValidationError(error.to_string()))?;
83        require(repository(&model.id), "invalid repository id")?;
84        require(
85            model.siblings.len() <= 10_000 && model.tags.len() <= 1_000,
86            "model metadata has too many entries",
87        )?;
88        Ok(model)
89    }
90    /// Pinned revisions only: a 40-hex commit, not gated and not private.
91    pub fn pinned(&self) -> bool {
92        self.sha.len() == 40
93            && self.sha.bytes().all(|byte| byte.is_ascii_hexdigit())
94            && !self.private
95            && matches!(self.gated, Value::Null | Value::Bool(false))
96    }
97    pub fn base_model(&self) -> Option<String> {
98        self.base_models().into_iter().next()
99    }
100    /// Direct base models (not `base_model:finetune:…` relations), at most four.
101    pub fn base_models(&self) -> Vec<String> {
102        self.tags
103            .iter()
104            .filter_map(|tag| tag.strip_prefix("base_model:"))
105            .filter(|value| repository(value))
106            .take(4)
107            .map(str::to_string)
108            .collect()
109    }
110}
111
112pub fn license_from_tags(tags: &[String]) -> Option<&'static str> {
113    let mut licenses = tags.iter().filter_map(|tag| tag.strip_prefix("license:"));
114    let license = licenses.next()?;
115    if licenses.next().is_some() {
116        return None;
117    }
118    LICENSES.iter().copied().find(|known| *known == license)
119}
120
121/// Text-prompted generation only; editing, audio, 3D and depth pipelines are not admitted.
122pub fn kind_from_pipeline(pipeline: Option<&str>) -> Option<GenerationKind> {
123    match pipeline? {
124        "text-to-image" => Some(GenerationKind::Image),
125        "text-to-video" => Some(GenerationKind::Video),
126        _ => None,
127    }
128}
129
130#[derive(Debug, Clone, Copy, PartialEq, Eq)]
131pub enum SafetensorsError {
132    /// More header bytes are needed than were read.
133    Truncated(u64),
134    Invalid(&'static str),
135}
136
137/// Exact parameter count from a safetensors header (8-byte length + JSON tensor table).
138pub fn safetensors_params(bytes: &[u8]) -> Result<u64, SafetensorsError> {
139    let length = u64::from_le_bytes(
140        bytes
141            .get(..8)
142            .ok_or(SafetensorsError::Truncated(8))?
143            .try_into()
144            .expect("eight bytes"),
145    );
146    if length == 0 || length > MAX_SAFETENSORS_HEADER_BYTES {
147        return Err(SafetensorsError::Invalid(
148            "implausible safetensors header length",
149        ));
150    }
151    let end = 8 + length as usize;
152    let header = bytes
153        .get(8..end)
154        .ok_or(SafetensorsError::Truncated(end as u64))?;
155    let table: serde_json::Map<String, Value> = serde_json::from_slice(header)
156        .map_err(|_| SafetensorsError::Invalid("safetensors header is not a JSON object"))?;
157    let mut total: u64 = 0;
158    for (name, tensor) in table {
159        if name == "__metadata__" {
160            continue;
161        }
162        let shape = tensor["shape"]
163            .as_array()
164            .ok_or(SafetensorsError::Invalid("tensor without shape"))?;
165        let mut elements: u64 = 1;
166        for dim in shape {
167            elements = elements
168                .checked_mul(
169                    dim.as_u64()
170                        .ok_or(SafetensorsError::Invalid("invalid tensor dimension"))?,
171                )
172                .ok_or(SafetensorsError::Invalid("tensor size overflow"))?;
173        }
174        total = total
175            .checked_add(elements)
176            .ok_or(SafetensorsError::Invalid("tensor size overflow"))?;
177    }
178    if total == 0 {
179        return Err(SafetensorsError::Invalid(
180            "safetensors header has no tensors",
181        ));
182    }
183    Ok(total)
184}
185
186pub struct FitOnlyInput<'a> {
187    pub model: &'a HfModel,
188    pub kind: GenerationKind,
189    pub license: &'static str,
190    /// Exact parameter count of a diffusion file (from its safetensors header), if read.
191    pub params: &'a dyn Fn(&str) -> Option<u64>,
192    pub today: &'a str,
193}
194
195/// A built-in workflow and the exact files ComfyUI's official example runs it with.
196pub struct Recipe {
197    /// Only verified recipes may promote production entries to runnable.
198    pub verified: bool,
199    pub repository: &'static str,
200    pub kind: GenerationKind,
201    pub workflow: Workflow,
202    pub source: &'static str,
203    /// Diffusion files the example uses (and same-architecture precision siblings).
204    pub diffusion: &'static [&'static str],
205    pub text_encoder: &'static str,
206    pub vae: &'static str,
207}
208
209pub const RECIPES: [Recipe; 3] = [
210    Recipe {
211        verified: true,
212        repository: "Comfy-Org/Wan_2.2_ComfyUI_Repackaged",
213        kind: GenerationKind::Video,
214        workflow: Workflow::Wan22Ti2v,
215        source: "https://comfyanonymous.github.io/ComfyUI_examples/wan22/",
216        diffusion: &["wan2.2_ti2v_5B_fp16.safetensors"],
217        text_encoder: "umt5_xxl_fp8_e4m3fn_scaled.safetensors",
218        vae: "wan2.2_vae.safetensors",
219    },
220    Recipe {
221        verified: true,
222        repository: "Comfy-Org/Qwen-Image_ComfyUI",
223        kind: GenerationKind::Image,
224        workflow: Workflow::QwenImage,
225        source: "https://comfyanonymous.github.io/ComfyUI_examples/qwen_image/",
226        diffusion: &[
227            "qwen_image_fp8_e4m3fn.safetensors",
228            "qwen_image_bf16.safetensors",
229        ],
230        text_encoder: "qwen_2.5_vl_7b_fp8_scaled.safetensors",
231        vae: "qwen_image_vae.safetensors",
232    },
233    Recipe {
234        verified: true,
235        repository: "Comfy-Org/Ming-Image",
236        kind: GenerationKind::Image,
237        workflow: Workflow::MingDesign,
238        source: "https://github.com/Comfy-Org/workflow_templates/blob/99e3d43745926b78466d99f937c0cd2bb622423a/templates/image_ming_image_01_design_t2i.json",
239        diffusion: &["ming_image_0.1_design_int8_convrot.safetensors"],
240        text_encoder: "ming_image_0.1_ling_mini_2.0_int8_convrot.safetensors",
241        vae: "ming_image_vae_bf16.safetensors",
242    },
243];
244
245pub fn recipe(repository: &str) -> Option<&'static Recipe> {
246    RECIPES
247        .iter()
248        .find(|recipe| recipe.verified && recipe.repository == repository)
249}
250
251/// The model family from a Comfy-Org repository name, without packaging suffixes.
252pub fn family(repository: &str) -> String {
253    let mut name = repository
254        .rsplit('/')
255        .next()
256        .unwrap_or_default()
257        .to_ascii_lowercase();
258    loop {
259        let trimmed = ["_repackaged", "-repackaged", "_comfyui", "-comfyui"]
260            .iter()
261            .find_map(|suffix| name.strip_suffix(suffix).map(str::to_string));
262        match trimmed {
263            Some(shorter) if !shorter.is_empty() => name = shorter,
264            _ => return name,
265        }
266    }
267}
268
269fn folder(path: &str) -> Option<FileRole> {
270    let segments: Vec<&str> = path.split('/').collect();
271    let parent = segments.get(segments.len().checked_sub(2)?)?;
272    match *parent {
273        "diffusion_models" => Some(FileRole::Diffusion),
274        "text_encoders" => Some(FileRole::TextEncoder),
275        "vae" => Some(FileRole::Vae),
276        _ => None,
277    }
278}
279
280fn stem(path: &str) -> &str {
281    let name = path.rsplit('/').next().unwrap_or(path);
282    name.strip_suffix(".safetensors").unwrap_or(name)
283}
284
285/// Variants that condition on something other than a text prompt.
286fn text_prompted(path: &str) -> bool {
287    let name = stem(path).to_ascii_lowercase();
288    // `ti2v` stays: it takes a text-only prompt.
289    ![
290        "edit",
291        "lora",
292        "inpaint",
293        "controlnet",
294        "control_",
295        "fill",
296        "depth",
297        "canny",
298        "animate",
299        "fun_",
300        "_i2v",
301        "s2v",
302        "vace",
303        "camera",
304    ]
305    .iter()
306    .any(|marker| name.contains(marker))
307}
308
309fn generation_file(model: &HfModel, sibling: &Sibling, role: FileRole) -> Option<GenerationFile> {
310    let lfs = sibling.lfs.as_ref()?;
311    (lfs.sha256.len() == 64
312        && lfs.sha256.bytes().all(|byte| byte.is_ascii_hexdigit())
313        && lfs.size > 0)
314        .then(|| GenerationFile {
315            role,
316            folder: role.folder().into(),
317            repo: model.id.clone(),
318            revision: model.sha.clone(),
319            file: sibling.rfilename.clone(),
320            sha256: lfs.sha256.to_ascii_lowercase(),
321            bytes: lfs.size,
322        })
323}
324
325/// Text-prompted diffusion files whose parameter count an entry needs.
326pub fn diffusion_files(model: &HfModel) -> Vec<&Sibling> {
327    model
328        .siblings
329        .iter()
330        .filter(|sibling| {
331            sibling.rfilename.ends_with(".safetensors")
332                && folder(&sibling.rfilename) == Some(FileRole::Diffusion)
333                && text_prompted(&sibling.rfilename)
334                && sibling.lfs.is_some()
335        })
336        .collect()
337}
338
339/// One fit-only entry per text-prompted diffusion file, paired with the smallest published
340/// text encoder and VAE (the least memory the release can run with).
341pub fn fit_only_entries(input: &FitOnlyInput) -> Result<Vec<GenerationModel>, ValidationError> {
342    let model = input.model;
343    require(model.pinned(), "repository is not a public pinned revision")?;
344    let family = family(&model.id);
345    let files: Vec<(&Sibling, FileRole)> = model
346        .siblings
347        .iter()
348        .filter(|sibling| sibling.rfilename.ends_with(".safetensors"))
349        .filter_map(|sibling| folder(&sibling.rfilename).map(|role| (sibling, role)))
350        .collect();
351    let smallest = |role: FileRole| {
352        files
353            .iter()
354            .filter(|(_, candidate)| *candidate == role)
355            .filter_map(|(sibling, _)| generation_file(model, sibling, role))
356            .min_by(|left, right| {
357                left.bytes
358                    .cmp(&right.bytes)
359                    .then_with(|| left.file.cmp(&right.file))
360            })
361    };
362    let companions: Vec<GenerationFile> = [FileRole::TextEncoder, FileRole::Vae]
363        .into_iter()
364        .filter_map(smallest)
365        .collect();
366    let mut entries = Vec::new();
367    for (sibling, role) in &files {
368        if *role != FileRole::Diffusion || !text_prompted(&sibling.rfilename) {
369            continue;
370        }
371        let Some(diffusion) = generation_file(model, sibling, FileRole::Diffusion) else {
372            continue;
373        };
374        let Some(params) = (input.params)(&sibling.rfilename) else {
375            continue;
376        };
377        let id = format!("{family}:{}", stem(&sibling.rfilename).to_ascii_lowercase());
378        let file_name = sibling.rfilename.rsplit('/').next().unwrap_or_default();
379        let exact = |role: FileRole, name: &str| {
380            files
381                .iter()
382                .find(|(candidate, candidate_role)| {
383                    *candidate_role == role && candidate.rfilename.rsplit('/').next() == Some(name)
384                })
385                .and_then(|(candidate, _)| generation_file(model, candidate, role))
386        };
387        // Runnable only when every file the official example names is published here.
388        let runnable = recipe(&model.id)
389            .filter(|recipe| recipe.kind == input.kind && recipe.diffusion.contains(&file_name))
390            .and_then(|recipe| {
391                Some((
392                    recipe,
393                    exact(FileRole::TextEncoder, recipe.text_encoder)?,
394                    exact(FileRole::Vae, recipe.vae)?,
395                ))
396            });
397        let mut entry_files = vec![diffusion];
398        let (workflow, workflow_source) = match runnable {
399            Some((recipe, encoder, vae)) => {
400                entry_files.extend([encoder, vae]);
401                (Some(recipe.workflow), Some(recipe.source.to_string()))
402            }
403            None => {
404                entry_files.extend(companions.iter().cloned());
405                (None, None)
406            }
407        };
408        entries.push(GenerationModel {
409            id,
410            family: family.clone(),
411            kind: input.kind,
412            params: crate::admission::parameter_label(params as f64),
413            license: input.license.into(),
414            open_weight: true,
415            release_date: None,
416            added_at: Some(input.today.into()),
417            provenance: EntryProvenance::Auto,
418            default: false,
419            source: format!("https://huggingface.co/{}", model.id),
420            workflow,
421            workflow_source,
422            files: entry_files,
423        });
424    }
425    entries.sort_by(|left, right| left.id.cmp(&right.id));
426    Ok(entries)
427}
rigspark_core/generation_admission.rs

rigspark_core/
generation_admission.rs
