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    pub repository: &'static str,
198    pub kind: GenerationKind,
199    pub workflow: Workflow,
200    pub source: &'static str,
201    /// Diffusion files the example uses (and same-architecture precision siblings).
202    pub diffusion: &'static [&'static str],
203    pub text_encoder: &'static str,
204    pub vae: &'static str,
205}
206
207pub const RECIPES: [Recipe; 2] = [
208    Recipe {
209        repository: "Comfy-Org/Wan_2.2_ComfyUI_Repackaged",
210        kind: GenerationKind::Video,
211        workflow: Workflow::Wan22Ti2v,
212        source: "https://comfyanonymous.github.io/ComfyUI_examples/wan22/",
213        diffusion: &["wan2.2_ti2v_5B_fp16.safetensors"],
214        text_encoder: "umt5_xxl_fp8_e4m3fn_scaled.safetensors",
215        vae: "wan2.2_vae.safetensors",
216    },
217    Recipe {
218        repository: "Comfy-Org/Qwen-Image_ComfyUI",
219        kind: GenerationKind::Image,
220        workflow: Workflow::QwenImage,
221        source: "https://comfyanonymous.github.io/ComfyUI_examples/qwen_image/",
222        diffusion: &[
223            "qwen_image_fp8_e4m3fn.safetensors",
224            "qwen_image_bf16.safetensors",
225        ],
226        text_encoder: "qwen_2.5_vl_7b_fp8_scaled.safetensors",
227        vae: "qwen_image_vae.safetensors",
228    },
229];
230
231pub fn recipe(repository: &str) -> Option<&'static Recipe> {
232    RECIPES
233        .iter()
234        .find(|recipe| recipe.repository == repository)
235}
236
237/// The model family from a Comfy-Org repository name, without packaging suffixes.
238pub fn family(repository: &str) -> String {
239    let mut name = repository
240        .rsplit('/')
241        .next()
242        .unwrap_or_default()
243        .to_ascii_lowercase();
244    loop {
245        let trimmed = ["_repackaged", "-repackaged", "_comfyui", "-comfyui"]
246            .iter()
247            .find_map(|suffix| name.strip_suffix(suffix).map(str::to_string));
248        match trimmed {
249            Some(shorter) if !shorter.is_empty() => name = shorter,
250            _ => return name,
251        }
252    }
253}
254
255fn folder(path: &str) -> Option<FileRole> {
256    let segments: Vec<&str> = path.split('/').collect();
257    let parent = segments.get(segments.len().checked_sub(2)?)?;
258    match *parent {
259        "diffusion_models" => Some(FileRole::Diffusion),
260        "text_encoders" => Some(FileRole::TextEncoder),
261        "vae" => Some(FileRole::Vae),
262        _ => None,
263    }
264}
265
266fn stem(path: &str) -> &str {
267    let name = path.rsplit('/').next().unwrap_or(path);
268    name.strip_suffix(".safetensors").unwrap_or(name)
269}
270
271/// Variants that condition on something other than a text prompt.
272fn text_prompted(path: &str) -> bool {
273    let name = stem(path).to_ascii_lowercase();
274    // `ti2v` stays: it takes a text-only prompt.
275    ![
276        "edit",
277        "lora",
278        "inpaint",
279        "controlnet",
280        "control_",
281        "fill",
282        "depth",
283        "canny",
284        "animate",
285        "fun_",
286        "_i2v",
287        "s2v",
288        "vace",
289        "camera",
290    ]
291    .iter()
292    .any(|marker| name.contains(marker))
293}
294
295fn generation_file(model: &HfModel, sibling: &Sibling, role: FileRole) -> Option<GenerationFile> {
296    let lfs = sibling.lfs.as_ref()?;
297    (lfs.sha256.len() == 64
298        && lfs.sha256.bytes().all(|byte| byte.is_ascii_hexdigit())
299        && lfs.size > 0)
300        .then(|| GenerationFile {
301            role,
302            folder: role.folder().into(),
303            repo: model.id.clone(),
304            revision: model.sha.clone(),
305            file: sibling.rfilename.clone(),
306            sha256: lfs.sha256.to_ascii_lowercase(),
307            bytes: lfs.size,
308        })
309}
310
311/// Text-prompted diffusion files whose parameter count an entry needs.
312pub fn diffusion_files(model: &HfModel) -> Vec<&Sibling> {
313    model
314        .siblings
315        .iter()
316        .filter(|sibling| {
317            sibling.rfilename.ends_with(".safetensors")
318                && folder(&sibling.rfilename) == Some(FileRole::Diffusion)
319                && text_prompted(&sibling.rfilename)
320                && sibling.lfs.is_some()
321        })
322        .collect()
323}
324
325/// One fit-only entry per text-prompted diffusion file, paired with the smallest published
326/// text encoder and VAE (the least memory the release can run with).
327pub fn fit_only_entries(input: &FitOnlyInput) -> Result<Vec<GenerationModel>, ValidationError> {
328    let model = input.model;
329    require(model.pinned(), "repository is not a public pinned revision")?;
330    let family = family(&model.id);
331    let files: Vec<(&Sibling, FileRole)> = model
332        .siblings
333        .iter()
334        .filter(|sibling| sibling.rfilename.ends_with(".safetensors"))
335        .filter_map(|sibling| folder(&sibling.rfilename).map(|role| (sibling, role)))
336        .collect();
337    let smallest = |role: FileRole| {
338        files
339            .iter()
340            .filter(|(_, candidate)| *candidate == role)
341            .filter_map(|(sibling, _)| generation_file(model, sibling, role))
342            .min_by(|left, right| {
343                left.bytes
344                    .cmp(&right.bytes)
345                    .then_with(|| left.file.cmp(&right.file))
346            })
347    };
348    let companions: Vec<GenerationFile> = [FileRole::TextEncoder, FileRole::Vae]
349        .into_iter()
350        .filter_map(smallest)
351        .collect();
352    let mut entries = Vec::new();
353    for (sibling, role) in &files {
354        if *role != FileRole::Diffusion || !text_prompted(&sibling.rfilename) {
355            continue;
356        }
357        let Some(diffusion) = generation_file(model, sibling, FileRole::Diffusion) else {
358            continue;
359        };
360        let Some(params) = (input.params)(&sibling.rfilename) else {
361            continue;
362        };
363        let id = format!("{family}:{}", stem(&sibling.rfilename).to_ascii_lowercase());
364        let file_name = sibling.rfilename.rsplit('/').next().unwrap_or_default();
365        let exact = |role: FileRole, name: &str| {
366            files
367                .iter()
368                .find(|(candidate, candidate_role)| {
369                    *candidate_role == role && candidate.rfilename.rsplit('/').next() == Some(name)
370                })
371                .and_then(|(candidate, _)| generation_file(model, candidate, role))
372        };
373        // Runnable only when every file the official example names is published here.
374        let runnable = recipe(&model.id)
375            .filter(|recipe| recipe.kind == input.kind && recipe.diffusion.contains(&file_name))
376            .and_then(|recipe| {
377                Some((
378                    recipe,
379                    exact(FileRole::TextEncoder, recipe.text_encoder)?,
380                    exact(FileRole::Vae, recipe.vae)?,
381                ))
382            });
383        let mut entry_files = vec![diffusion];
384        let (workflow, workflow_source) = match runnable {
385            Some((recipe, encoder, vae)) => {
386                entry_files.extend([encoder, vae]);
387                (Some(recipe.workflow), Some(recipe.source.to_string()))
388            }
389            None => {
390                entry_files.extend(companions.iter().cloned());
391                (None, None)
392            }
393        };
394        entries.push(GenerationModel {
395            id,
396            family: family.clone(),
397            kind: input.kind,
398            params: crate::admission::parameter_label(params as f64),
399            license: input.license.into(),
400            open_weight: true,
401            release_date: None,
402            added_at: Some(input.today.into()),
403            provenance: EntryProvenance::Auto,
404            default: false,
405            source: format!("https://huggingface.co/{}", model.id),
406            workflow,
407            workflow_source,
408            files: entry_files,
409        });
410    }
411    entries.sort_by(|left, right| left.id.cmp(&right.id));
412    Ok(entries)
413}
rigspark_core/generation_admission.rs

rigspark_core/
generation_admission.rs
