studio_worker/
catalog.rs

1//! Local model catalog — the offline equivalent of the studio's
2//! `studioModels` registry.
3//!
4//! The studio is normally the single source of truth for a model's
5//! [`ModelSource`] (which files to download + the CLI defaults). When generating
6//! locally there is no studio, so the worker keeps a small JSON catalog the
7//! operator can edit and extend exactly the way they would add a model in the
8//! studio. It ships seeded with Z-Image-Turbo (the studio's default image
9//! model) so a fresh install can generate out of the box.
10
11use std::path::Path;
12
13use serde::{Deserialize, Serialize};
14
15use crate::types::{
16    ModelCliDefaults, ModelEngine, ModelFile, ModelFileRole, ModelSource, TaskKind,
17};
18
19/// One catalog entry: a model id plus everything needed to run it. Mirrors the
20/// columns of the studio's `studioModels` row.
21#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
22#[serde(rename_all = "camelCase")]
23pub struct CatalogModel {
24    /// The model id the operator references (e.g. `z-image-turbo-q4_k_m.gguf`).
25    pub id: String,
26    /// Human-readable name shown in the UI.
27    pub display_name: String,
28    /// Task kind this model serves.
29    pub kind: TaskKind,
30    /// Rough VRAM requirement in GB (informational).
31    #[serde(default)]
32    pub vram_gb_estimate: f32,
33    /// Optional human description.
34    #[serde(default, skip_serializing_if = "Option::is_none")]
35    pub description: Option<String>,
36    /// Download spec + engine + CLI defaults (same shape the studio sends).
37    pub source: ModelSource,
38    /// Whether the model is selectable.
39    #[serde(default = "default_true")]
40    pub enabled: bool,
41}
42
43fn default_true() -> bool {
44    true
45}
46
47/// A collection of locally-available models.
48#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
49pub struct Catalog {
50    #[serde(default)]
51    pub models: Vec<CatalogModel>,
52}
53
54impl Catalog {
55    /// The built-in catalog: every model the worker ships seeded with.
56    pub fn seed() -> Self {
57        Catalog {
58            models: vec![zimage_turbo()],
59        }
60    }
61
62    /// Parse a catalog from a JSON string.
63    pub fn from_json(json: &str) -> serde_json::Result<Self> {
64        serde_json::from_str(json)
65    }
66
67    /// Serialise to pretty JSON.
68    pub fn to_json(&self) -> serde_json::Result<String> {
69        serde_json::to_string_pretty(self)
70    }
71
72    /// Load the catalog from `path`. If the file does not exist it is seeded
73    /// with the built-in defaults and written to `path`.
74    pub fn load_or_seed(path: &Path) -> std::io::Result<Self> {
75        match std::fs::read_to_string(path) {
76            Ok(contents) => Self::from_json(&contents)
77                .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e)),
78            Err(err) if err.kind() == std::io::ErrorKind::NotFound => {
79                let seeded = Self::seed();
80                seeded.save(path)?;
81                Ok(seeded)
82            }
83            Err(err) => Err(err),
84        }
85    }
86
87    /// Write the catalog to `path` (creating parent dirs).
88    pub fn save(&self, path: &Path) -> std::io::Result<()> {
89        if let Some(parent) = path.parent() {
90            std::fs::create_dir_all(parent)?;
91        }
92        let json = self
93            .to_json()
94            .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
95        std::fs::write(path, json)
96    }
97
98    /// Look up a model by id.
99    pub fn get(&self, id: &str) -> Option<&CatalogModel> {
100        self.models.iter().find(|m| m.id == id)
101    }
102
103    /// All catalog entries.
104    pub fn list(&self) -> &[CatalogModel] {
105        &self.models
106    }
107
108    /// Insert a model, replacing any existing entry with the same id.
109    pub fn upsert(&mut self, model: CatalogModel) {
110        if let Some(existing) = self.models.iter_mut().find(|m| m.id == model.id) {
111            *existing = model;
112        } else {
113            self.models.push(model);
114        }
115    }
116
117    /// Remove a model by id. Returns whether it existed.
118    pub fn remove(&mut self, id: &str) -> bool {
119        let before = self.models.len();
120        self.models.retain(|m| m.id != id);
121        self.models.len() != before
122    }
123
124    /// The first enabled image model — used when a request names no model.
125    pub fn default_image_model(&self) -> Option<&CatalogModel> {
126        self.models
127            .iter()
128            .find(|m| m.enabled && m.kind == TaskKind::Image)
129    }
130}
131
132/// The canonical Z-Image-Turbo entry, mirroring the studio seed
133/// (`migrations/graphics/0017_seed_registry.sql`).
134fn zimage_turbo() -> CatalogModel {
135    CatalogModel {
136        id: "z-image-turbo-q4_k_m.gguf".into(),
137        display_name: "Z-Image Turbo (Q4_K_M)".into(),
138        kind: TaskKind::Image,
139        vram_gb_estimate: 12.0,
140        description: Some(
141            "Distilled 8-step diffusion model packaged for sd.cpp. Diffusion (Q4_K), \
142             Qwen3-4B text encoder, Flux ae.safetensors VAE."
143                .into(),
144        ),
145        source: ModelSource {
146            engine: ModelEngine::SdCpp,
147            files: vec![
148                ModelFile {
149                    role: ModelFileRole::DiffusionModel,
150                    url: "https://huggingface.co/leejet/Z-Image-Turbo-GGUF/resolve/main/z_image_turbo-Q4_K.gguf".into(),
151                    filename: "z_image_turbo-Q4_K.gguf".into(),
152                    approx_bytes: Some(2_700_000_000),
153                    sha256: None,
154                },
155                ModelFile {
156                    role: ModelFileRole::TextEncoder,
157                    url: "https://huggingface.co/unsloth/Qwen3-4B-Instruct-2507-GGUF/resolve/main/Qwen3-4B-Instruct-2507-Q4_K_M.gguf".into(),
158                    filename: "Qwen3-4B-Instruct-2507-Q4_K_M.gguf".into(),
159                    approx_bytes: Some(2_500_000_000),
160                    sha256: None,
161                },
162                ModelFile {
163                    role: ModelFileRole::Vae,
164                    url: "https://huggingface.co/Comfy-Org/Lumina_Image_2.0_Repackaged/resolve/main/split_files/vae/ae.safetensors".into(),
165                    filename: "ae.safetensors".into(),
166                    approx_bytes: Some(335_000_000),
167                    sha256: None,
168                },
169            ],
170            cli_defaults: ModelCliDefaults {
171                cfg_scale: 1.0,
172                steps: 8,
173                width: 1024,
174                height: 1024,
175                sampling_method: Some("euler".into()),
176                flow_shift: None,
177                zero_cond_t: None,
178                offload_to_cpu: None,
179            },
180        },
181        enabled: true,
182    }
183}
184
185#[cfg(test)]
186mod tests {
187    use super::*;
188
189    #[test]
190    fn seed_contains_zimage_with_three_files() {
191        let catalog = Catalog::seed();
192        let model = catalog
193            .get("z-image-turbo-q4_k_m.gguf")
194            .expect("z-image seeded");
195        assert_eq!(model.kind, TaskKind::Image);
196        assert_eq!(model.source.engine, ModelEngine::SdCpp);
197        assert_eq!(model.source.files.len(), 3);
198        assert_eq!(model.source.cli_defaults.steps, 8);
199        assert!(model.enabled);
200    }
201
202    #[test]
203    fn default_image_model_is_zimage() {
204        let catalog = Catalog::seed();
205        assert_eq!(
206            catalog.default_image_model().map(|m| m.id.as_str()),
207            Some("z-image-turbo-q4_k_m.gguf")
208        );
209    }
210
211    #[test]
212    fn json_round_trips() {
213        let catalog = Catalog::seed();
214        let json = catalog.to_json().unwrap();
215        // camelCase wire keys, mirroring the studio.
216        assert!(json.contains("\"displayName\""));
217        assert!(json.contains("\"cliDefaults\""));
218        assert!(json.contains("\"diffusion-model\""));
219        let parsed = Catalog::from_json(&json).unwrap();
220        assert_eq!(parsed, catalog);
221    }
222
223    #[test]
224    fn upsert_adds_then_replaces() {
225        let mut catalog = Catalog::default();
226        let mut model = zimage_turbo();
227        catalog.upsert(model.clone());
228        assert_eq!(catalog.list().len(), 1);
229
230        model.display_name = "Renamed".into();
231        catalog.upsert(model);
232        assert_eq!(catalog.list().len(), 1);
233        assert_eq!(
234            catalog
235                .get("z-image-turbo-q4_k_m.gguf")
236                .unwrap()
237                .display_name,
238            "Renamed"
239        );
240    }
241
242    #[test]
243    fn remove_reports_presence() {
244        let mut catalog = Catalog::seed();
245        assert!(catalog.remove("z-image-turbo-q4_k_m.gguf"));
246        assert!(!catalog.remove("z-image-turbo-q4_k_m.gguf"));
247        assert!(catalog.get("z-image-turbo-q4_k_m.gguf").is_none());
248    }
249
250    #[test]
251    fn load_or_seed_writes_then_reads_back() {
252        let dir = std::env::temp_dir().join(format!("sw-catalog-{}", std::process::id()));
253        let _ = std::fs::remove_dir_all(&dir);
254        let path = dir.join("models.json");
255
256        // Missing -> seeded + persisted.
257        let seeded = Catalog::load_or_seed(&path).unwrap();
258        assert!(path.exists());
259        assert!(seeded.get("z-image-turbo-q4_k_m.gguf").is_some());
260
261        // Existing -> read back unchanged.
262        let reloaded = Catalog::load_or_seed(&path).unwrap();
263        assert_eq!(reloaded, seeded);
264
265        let _ = std::fs::remove_dir_all(&dir);
266    }
267
268    #[test]
269    fn equality_derives_hold_for_catalog_model() {
270        assert_eq!(zimage_turbo(), zimage_turbo());
271    }
272}
studio_worker/catalog.rs

studio_worker/
catalog.rs
