From e2cf18678f37b9d7a25d7057f043aa9fb9779af6 Mon Sep 17 00:00:00 2001 From: Cameron Cordes Date: Tue, 3 Nov 2026 16:28:23 -0500 Subject: [PATCH] feat(enhance): pick LoRAs from ComfyUI per job (#220) Templates with `imageapi:lora-slot::high|low` nodes take LoRAs picked per job from a catalog built live from ComfyUI. - Catalog: ComfyUI's LoraLoaderModelOnly `lora_name` choices (cached 60 s, re-read when a submit names an unknown id), paired by High/Low name tokens (StyleHigh/StyleLow, x-HIGH/x-LOW, x_high_noise/...). Files the workflow already loads outside the slots (lightx2v, switches) and `[template.catalog] exclude` substrings never show; ids derive from the names so they stay stable. - POST /enhance `slots: [{lora, scale?}]` fills the slots in order: each pick sets the slot nodes' lora_name and scales their exported strength; unused slots run at 0, so the graph is never rewired. More picks than slots, duplicates or unknown ids are 400s. - GET /enhance/templates gains `slots` and `catalog: {loras, strength, min, max}`; jobs store and return the picks (new enhance_jobs.lora_slots JSON column). Co-Authored-By: Claude Opus 5.5 --- CLAUDE.md | 19 +- enhance.example.toml | 14 + .../down.sql | 1 + .../up.sql | 3 + src/database/enhance_dao.rs | 9 +- src/database/schema.rs | 1 + src/enhance/comfy.rs | 7 + src/enhance/config.rs | 88 +++- src/enhance/handlers.rs | 39 +- src/enhance/integration_tests.rs | 130 +++++- src/enhance/lora_catalog.rs | 336 +++++++++++++++ src/enhance/mod.rs | 1 + src/enhance/service.rs | 136 +++++- src/enhance/template.rs | 299 ++++++++++++- tests/fixtures/enhance/animate_slots.json | 404 ++++++++++++++++++ 15 files changed, 1454 insertions(+), 33 deletions(-) create mode 100644 migrations/2026-11-03-200000_add_enhance_lora_slots/down.sql create mode 100644 migrations/2026-11-03-200000_add_enhance_lora_slots/up.sql create mode 100644 src/enhance/lora_catalog.rs create mode 100644 tests/fixtures/enhance/animate_slots.json diff --git a/CLAUDE.md b/CLAUDE.md index 480c9bc..4ed882a 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -460,7 +460,8 @@ auth and stays LAN-only, so ImageApi relays previews too. `unreachable_grace_secs`, `default_template`, `warn_above_megapixels`, `[[path_map]]` (`from`, `to`), `[[template]]` (`name`, `label`, `description`, `file`, `prompt_mode`, `prompt_join`, `output_megapixels`, `[[template.lora]]` - (`id`, `label`, `default`, `strength`, `min`, `max`)). + (`id`, `label`, `default`, `strength`, `min`, `max`), `[template.catalog]` (`strength`, + `min`, `max`, `exclude`)). - **Templates** are API-format workflow exports, read-only (never edit them from here). Nodes are found by `_meta.title`: `imageapi:source` (required, gets the photo path), `imageapi:save` (required, forced to `require_approval = true`; a @@ -472,7 +473,12 @@ auth and stays LAN-only, so ImageApi relays previews too. low pair; the resolved scale multiplies each node's exported `strength_model` / `strength_clip` (1.0 = as exported, 0 = off, a pair keeps its ratio), never rewired; a node with no declaration keeps its strength, a declaration with no node isn't offered; - both log a startup warning). `prompt_mode = "append"`: the user's text is joined + both log a startup warning), `imageapi:lora-slot::high|low` (pickable LoRAs, #220: + each slot needs a high and a low node, numbered 1.. without gaps; picks fill them in order, + setting `lora_name` and scaling the exported strength, unused slots run at 0; the catalog is + ComfyUI's `/object_info/LoraLoaderModelOnly` list (cached 60 s, re-read when a submit names + an unknown id) paired by High/Low name tokens (`src/enhance/lora_catalog.rs`), minus every + `lora_name` the workflow loads outside the slots and `catalog.exclude` substrings). `prompt_mode = "append"`: the user's text is joined after the template's own prompt with `prompt_join` (default one space); empty text leaves the prompt as is. `"replace"`: the text *is* the prompt and is required (empty = 400); it needs a prompt node. With no `reference_path` the reference node **and @@ -519,8 +525,12 @@ GET /enhance/templates "prompt_mode": "append"|"replace", "prompt_required", "accepts_reference","default","output_megapixels", "warn_above_megapixels", "output": "image"|"video", - "loras": [{"id","label","default","strength","min","max"}]}]} -POST /enhance {"path","template","prompt"?,"reference_path"?,"loras"?: {id: bool|scale},"library"?} -> job + "loras": [{"id","label","default","strength","min","max"}], + "slots": n, "catalog"?: {"loras": [{"id","label"}],"strength","min","max"}}]} +POST /enhance {"path","template","prompt"?,"reference_path"?,"loras"?: {id: bool|scale}, + "slots"?: [{"lora": catalog id, "scale"?}],"library"?} -> job + (slots: at most the template's slot count, no duplicates, unknown id = 400; no scale = + catalog.strength, <= 0 = slot unused) (true = the switch's strength, false or scale <= 0 = off, a scale is clamped to min..max; a missing switch uses its default; an id the template doesn't offer = 400) GET /enhance/jobs/{id} -> job @@ -540,6 +550,7 @@ POST /enhance/jobs/{id}/discard -> job (hidden from the list; ComfyUI untou job = {"id","status","queue_position"?,"error"?,"path","library_id","template", "prompt"?,"reference_path"?,"loras"?: {id: scale} (resolved, 0 = off; absent without switches), + "slots"?: [{"lora","label","scale"}] (picked LoRAs; absent when none), "created_at","updated_at", "run_ms": , "pending": [{"n","preview": "/enhance/jobs/{id}/preview/{n}","approved"}], diff --git a/enhance.example.toml b/enhance.example.toml index b52e0c4..c2e3b49 100644 --- a/enhance.example.toml +++ b/enhance.example.toml @@ -87,3 +87,17 @@ default = false # default: false; used when the request omi strength = 1.0 # default: 1.0; the scale when on without a number min = 0.0 # default: 0.0; slider range while on (a requested max = 1.5 # default: 1.5; scale is clamped; <= 0 is off) + +# Picked LoRAs (#220): LoRA loader nodes titled "imageapi:lora-slot::high" +# and ":low" (n = 1, 2, ... without gaps) are slots the app fills from a +# catalog, one picked LoRA per slot. The catalog is every LoRA file ComfyUI +# lists whose name pairs as High/Low (StyleHigh/StyleLow, x-HIGH/x-LOW, +# x_high_noise/x_low_noise), minus the LoRAs the workflow already loads +# (lightx2v, switches); a new download shows up within a minute. A pick's +# scale multiplies the slot's exported strength; unused slots run at 0. +# Settings for the template above (all optional): +[template.catalog] +strength = 1.0 # default: 1.0; a pick's scale when the request gives none +min = 0.0 # default: 0.0; slider range (a requested scale is +max = 1.5 # default: 1.5; clamped; <= 0 leaves the slot unused) +exclude = [] # default: []; file-name substrings never offered diff --git a/migrations/2026-11-03-200000_add_enhance_lora_slots/down.sql b/migrations/2026-11-03-200000_add_enhance_lora_slots/down.sql new file mode 100644 index 0000000..2d52c9d --- /dev/null +++ b/migrations/2026-11-03-200000_add_enhance_lora_slots/down.sql @@ -0,0 +1 @@ +ALTER TABLE enhance_jobs DROP COLUMN lora_slots; diff --git a/migrations/2026-11-03-200000_add_enhance_lora_slots/up.sql b/migrations/2026-11-03-200000_add_enhance_lora_slots/up.sql new file mode 100644 index 0000000..f2472b4 --- /dev/null +++ b/migrations/2026-11-03-200000_add_enhance_lora_slots/up.sql @@ -0,0 +1,3 @@ +-- LoRAs picked into the template's slots (#220): JSON [{lora, label, scale}]; +-- NULL = none picked. +ALTER TABLE enhance_jobs ADD COLUMN lora_slots TEXT; diff --git a/src/database/enhance_dao.rs b/src/database/enhance_dao.rs index a426198..b51a38f 100644 --- a/src/database/enhance_dao.rs +++ b/src/database/enhance_dao.rs @@ -66,11 +66,14 @@ pub struct EnhanceJobRow { /// template had none. Owned by `crate::enhance` like `pending`. #[diesel(sql_type = Nullable)] pub loras: Option, + /// LoRAs picked into slots (#220), JSON `[{lora, label, scale}]`. + #[diesel(sql_type = Nullable)] + pub lora_slots: Option, } const COLUMNS: &str = "id, original_hash, library_id, rel_path, template, prompt, reference_hash, \ reference_path, comfy_prompt_id, status, error, queue_position, pending, saved, \ - discarded, created_at, updated_at, run_ms, loras"; + discarded, created_at, updated_at, run_ms, loras, lora_slots"; /// One distinct prompt used across jobs, newest use first. #[derive(Debug, Clone, PartialEq, Eq, QueryableByName)] @@ -175,7 +178,7 @@ impl EnhanceJobDao for SqliteEnhanceJobDao { trace_db_call("insert", "enhance_job_insert", |_span| { let mut conn = self.connection.lock().expect("Unable to get EnhanceJobDao"); sql_query(format!( - "INSERT INTO enhance_jobs ({COLUMNS}) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)" + "INSERT INTO enhance_jobs ({COLUMNS}) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)" )) .bind::(&job.id) .bind::(&job.original_hash) @@ -196,6 +199,7 @@ impl EnhanceJobDao for SqliteEnhanceJobDao { .bind::(job.updated_at) .bind::, _>(job.run_ms) .bind::, _>(&job.loras) + .bind::, _>(&job.lora_slots) .execute(&mut *conn) .map(|_| ()) .map_err(|e| anyhow::anyhow!("Insert error: {}", e)) @@ -376,6 +380,7 @@ mod tests { updated_at: created, run_ms: None, loras: None, + lora_slots: None, } } diff --git a/src/database/schema.rs b/src/database/schema.rs index 6a742fb..cf7ee1d 100644 --- a/src/database/schema.rs +++ b/src/database/schema.rs @@ -74,6 +74,7 @@ diesel::table! { updated_at -> BigInt, run_ms -> Nullable, loras -> Nullable, + lora_slots -> Nullable, } } diff --git a/src/enhance/comfy.rs b/src/enhance/comfy.rs index 0445265..d6e77fb 100644 --- a/src/enhance/comfy.rs +++ b/src/enhance/comfy.rs @@ -330,6 +330,13 @@ impl ComfyClient { } } + /// Every LoRA file ComfyUI can load (`lora_name` choices of + /// LoraLoaderModelOnly), for the pickable-LoRA catalog (#220). + pub async fn lora_names(&self) -> Result, ComfyError> { + let body = self.get_json("/object_info/LoraLoaderModelOnly").await?; + Ok(super::lora_catalog::names_from_object_info(&body)) + } + async fn get_json(&self, path: &str) -> Result { let resp = self.http.get(self.url(path)).send().await?; if !resp.status().is_success() { diff --git a/src/enhance/config.rs b/src/enhance/config.rs index f958fb7..5ead715 100644 --- a/src/enhance/config.rs +++ b/src/enhance/config.rs @@ -104,6 +104,55 @@ pub struct TemplateConfig { /// workflow nodes titled `imageapi:lora:`. #[serde(rename = "lora", default)] pub loras: Vec, + /// Picked LoRAs (#220), for templates with `imageapi:lora-slot:` + /// nodes: `[template.catalog]`. + #[serde(default)] + pub catalog: CatalogConfig, +} + +/// How LoRAs picked from the catalog are scaled, and which never show. +#[derive(Debug, Clone, PartialEq, Deserialize, serde::Serialize)] +#[serde(default, deny_unknown_fields)] +pub struct CatalogConfig { + /// Scale a pick gets when the request gives none (1.0 = as exported). + pub strength: f64, + /// Slider range; a requested scale is clamped to it (<= 0 = slot unused). + pub min: f64, + pub max: f64, + /// LoRA file names containing any of these (case-insensitive) are never + /// offered, on top of the LoRAs the workflow already loads. + #[serde(skip_serializing)] + pub exclude: Vec, +} + +impl Default for CatalogConfig { + fn default() -> Self { + Self { + strength: 1.0, + min: 0.0, + max: 1.5, + exclude: Vec::new(), + } + } +} + +impl CatalogConfig { + /// A pick's scale: absent = `strength`, <= 0 = unused, else clamped. + pub fn resolve(&self, scale: Option) -> f64 { + match scale { + None => self.strength, + Some(s) if s <= 0.0 => 0.0, + Some(s) => s.clamp(self.min, self.max), + } + } +} + +/// `0 <= min < max`, `strength` within them, all finite. +fn check_range(strength: f64, min: f64, max: f64) -> bool { + [strength, min, max].iter().all(|v| v.is_finite()) + && min >= 0.0 + && min < max + && (min..=max).contains(&strength) } /// One LoRA switch (`[[template.lora]]`). @@ -228,12 +277,7 @@ impl EnhanceConfig { if l.label.trim().is_empty() { bail!("template {:?}: lora {:?} needs a label", t.name, l.id); } - let finite = [l.strength, l.min, l.max].iter().all(|v| v.is_finite()); - if !finite - || l.min < 0.0 - || l.min >= l.max - || !(l.min..=l.max).contains(&l.strength) - { + if !check_range(l.strength, l.min, l.max) { bail!( "template {:?}: lora {:?} needs 0 <= min < max and strength within them \ (got min {}, strength {}, max {})", @@ -245,6 +289,17 @@ impl EnhanceConfig { ); } } + let c = &t.catalog; + if !check_range(c.strength, c.min, c.max) { + bail!( + "template {:?}: [template.catalog] needs 0 <= min < max and strength within them \ + (got min {}, strength {}, max {})", + t.name, + c.min, + c.strength, + c.max + ); + } } if let Some(d) = &self.default_template && !self.templates.iter().any(|t| &t.name == d) @@ -414,6 +469,27 @@ mod tests { assert_eq!(parsed["d"], LoraRequest::Scale(1.0)); } + #[test] + fn catalog_settings_parse_validate_and_resolve() { + let base = "[[template]]\nname='a'\nlabel='A'\nfile='a.json'\n"; + let c = parse(base).unwrap().templates[0].catalog.clone(); + assert_eq!(c, CatalogConfig::default()); + let set = parse(&format!( + "{base}[template.catalog]\nstrength=0.7\nmax=1.2\nexclude=['x']" + )) + .unwrap(); + let c = &set.templates[0].catalog; + assert_eq!( + (c.strength, c.min, c.max, c.exclude.clone()), + (0.7, 0.0, 1.2, vec!["x".to_string()]) + ); + assert_eq!(c.resolve(None), 0.7); + assert_eq!(c.resolve(Some(0.0)), 0.0); + assert_eq!(c.resolve(Some(5.0)), 1.2); + assert!(parse(&format!("{base}[template.catalog]\nstrength=2.0")).is_err()); + assert!(parse(&format!("{base}[template.catalog]\nslots=2")).is_err()); + } + #[test] fn template_defaults() { let cfg = parse("[[template]]\nname='a'\nlabel='A'\nfile='a.json'").unwrap(); diff --git a/src/enhance/handlers.rs b/src/enhance/handlers.rs index 55633c0..47238b1 100644 --- a/src/enhance/handlers.rs +++ b/src/enhance/handlers.rs @@ -95,6 +95,9 @@ pub struct JobView { /// template had none. #[serde(skip_serializing_if = "Option::is_none")] loras: Option>, + /// The LoRAs picked into slots (#220); absent when none were. + #[serde(skip_serializing_if = "Option::is_none")] + slots: Option>, created_at: i64, updated_at: i64, /// Milliseconds ComfyUI ran the job; null until it has finished. @@ -134,6 +137,10 @@ pub fn job_view(job: &EnhanceJobRow) -> JobView { .loras .as_deref() .and_then(|j| serde_json::from_str(j).ok()), + slots: job + .lora_slots + .as_deref() + .and_then(|j| serde_json::from_str(j).ok()), created_at: job.created_at, updated_at: job.updated_at, run_ms: job.run_ms, @@ -233,6 +240,19 @@ struct TemplateView { output: super::template::OutputKind, /// Optional LoRAs the client can switch on or off (#218). loras: Vec, + /// How many LoRAs a job can pick (#220); 0 = none. + slots: usize, + /// The LoRAs it can pick from (`[{id, label}]`), and how they scale. + #[serde(skip_serializing_if = "Option::is_none")] + catalog: Option, +} + +#[derive(Serialize)] +struct CatalogView { + loras: Vec, + strength: f64, + min: f64, + max: f64, } #[get("/enhance/templates")] @@ -246,10 +266,15 @@ pub async fn get_templates(_: Claims, svc: web::Data) -> impl Re None }; let default = svc.default_template().map(str::to_string); + let mut catalogs = Vec::new(); + for t in svc.templates() { + catalogs.push((t.slot_count() > 0).then_some(svc.catalog(t).await)); + } let templates: Vec = svc .templates() .iter() - .map(|t| TemplateView { + .zip(catalogs) + .map(|(t, catalog)| TemplateView { name: t.cfg.name.clone(), label: t.cfg.label.clone(), description: t.cfg.description.clone(), @@ -262,6 +287,13 @@ pub async fn get_templates(_: Claims, svc: web::Data) -> impl Re warn_above_megapixels: svc.config.warn_above_megapixels, output: t.output, loras: t.lora_switches().cloned().collect(), + slots: t.slot_count(), + catalog: catalog.map(|loras| CatalogView { + loras, + strength: t.cfg.catalog.strength, + min: t.cfg.catalog.min, + max: t.cfg.catalog.max, + }), }) .collect(); let mut body = json!({ @@ -295,6 +327,10 @@ pub struct SubmitBody { /// the switch's range; <= 0 is off). A missing id uses its default. #[serde(default)] pub loras: std::collections::BTreeMap, + /// LoRAs picked from the template's catalog (#220), one per slot in + /// order: `[{lora: id, scale?}]`; no scale = the catalog's strength. + #[serde(default)] + pub slots: Vec, #[serde(default)] pub library: Option, } @@ -341,6 +377,7 @@ pub async fn submit( body.prompt.as_deref(), reference.as_ref(), &body.loras, + &body.slots, ) .await { diff --git a/src/enhance/integration_tests.rs b/src/enhance/integration_tests.rs index f5670ba..78ec8e0 100644 --- a/src/enhance/integration_tests.rs +++ b/src/enhance/integration_tests.rs @@ -27,6 +27,7 @@ use crate::versions::xmp::test_png::{producer_xmp, write_png}; const RESTORE: &str = include_str!("../../tests/fixtures/enhance/restore_turbo.json"); const ANIMATE_LORAS: &str = include_str!("../../tests/fixtures/enhance/animate_loras.json"); +const ANIMATE_SLOTS: &str = include_str!("../../tests/fixtures/enhance/animate_slots.json"); const EDIT: &str = include_str!("../../tests/fixtures/enhance/edit_turbo.json"); // ---- mock ComfyUI --------------------------------------------------------- @@ -45,6 +46,9 @@ struct Mock { free_calls: AtomicUsize, reject_prompts: Mutex>, view_404: Mutex, + /// What `/object_info/LoraLoaderModelOnly` lists. + lora_names: Mutex>, + lora_list_calls: AtomicUsize, } async fn start_mock() -> (Arc, String) { @@ -54,6 +58,16 @@ async fn start_mock() -> (Arc, String) { App::new() .app_data(Data::from(m.clone())) .route("/system_stats", web::get().to(|| async { HttpResponse::Ok().json(json!({})) })) + .route( + "/object_info/LoraLoaderModelOnly", + web::get().to(|m: Data| async move { + m.lora_list_calls.fetch_add(1, Ordering::SeqCst); + let names = m.lora_names.lock().unwrap().clone(); + HttpResponse::Ok().json(json!({"LoraLoaderModelOnly": {"input": {"required": { + "lora_name": [names], "strength_model": ["FLOAT", {}] + }}}})) + }), + ) .route( "/prompt", web::post().to(|m: Data, body: web::Json| async move { @@ -170,6 +184,7 @@ fn template(name: &str, mode: PromptMode, json: &str) -> Template { prompt_join: " ".into(), output_megapixels: 1.0, loras: Vec::new(), + catalog: Default::default(), }, serde_json::from_str(json).unwrap(), ) @@ -222,6 +237,7 @@ async fn fx_with(url: Option, tweak: impl FnOnce(&mut EnhanceConfig)) -> template("restore", PromptMode::Append, RESTORE), template("edit", PromptMode::Replace, EDIT), animate_with_style(), + template("animate-picks", PromptMode::Append, ANIMATE_SLOTS), ], dao, ) @@ -350,7 +366,7 @@ async fn comfy_down_means_unavailable_and_a_clear_submit_error() { assert_eq!(st, 200); assert_eq!(body["available"], false); assert_eq!(body["reason"], "ComfyUI isn't running"); - assert_eq!(body["templates"].as_array().unwrap().len(), 3); + assert_eq!(body["templates"].as_array().unwrap().len(), 4); let (st, body, _) = call( &fx, @@ -1010,6 +1026,7 @@ async fn unreachable_active_jobs_become_unavailable_after_the_grace_period() { updated_at: now, run_ms: None, loras: None, + lora_slots: None, }; SqliteEnhanceJobDao::from_shared(fx.conn.clone()) .insert(&dao_job) @@ -1223,3 +1240,114 @@ async fn lora_switches_are_listed_sent_and_stored_on_the_job() { .await; assert!(plain.get("loras").is_none(), "{plain}"); } + +#[actix_rt::test] +async fn picked_loras_are_listed_filled_into_slots_and_stored() { + let fx = fx().await; + fx.put(PHOTO, b"photo"); + *fx.mock.lora_names.lock().unwrap() = [ + "wan2.2_i2v_lightx2v_4steps_lora_v1_high_noise.safetensors", + "wan2.2_i2v_lightx2v_4steps_lora_v1_low_noise.safetensors", + "example_style_high.safetensors", + "example_style_low.safetensors", + "MotionHigh.safetensors", + "MotionLow.safetensors", + "Lonely_High.safetensors", + ] + .map(String::from) + .to_vec(); + + // Only the slot template lists a catalog; the workflow's own LoRAs + // (lightx2v, the style switch) and unpaired files are left out. + let (_, body, _) = call(&fx, "GET", "/enhance/templates", None).await; + let templates = body["templates"].as_array().unwrap(); + let picks = templates + .iter() + .find(|t| t["name"] == "animate-picks") + .unwrap(); + assert_eq!(picks["slots"], 2); + assert_eq!( + picks["catalog"], + json!({"loras": [{"id": "motion", "label": "Motion"}], "strength": 1.0, "min": 0.0, "max": 1.5}) + ); + let restore = templates.iter().find(|t| t["name"] == "restore").unwrap(); + assert_eq!(restore["slots"], 0); + assert!(restore.get("catalog").is_none()); + + let (st, job, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": PHOTO, "template": "animate-picks", "slots": [{"lora": "motion", "scale": 0.5}]})), + ) + .await; + assert_eq!(st, 200, "{job}"); + assert_eq!( + job["slots"], + json!([{"lora": "motion", "label": "Motion", "scale": 0.5}]) + ); + let (_, again, _) = call( + &fx, + "GET", + &format!("/enhance/jobs/{}", job["id"].as_str().unwrap()), + None, + ) + .await; + assert_eq!(again["slots"], job["slots"]); + { + let sent = fx.mock.submitted.lock().unwrap(); + let p = &sent[0]["prompt"]; + assert_eq!(p["60"]["inputs"]["lora_name"], "MotionHigh.safetensors"); + assert_eq!(p["61"]["inputs"]["lora_name"], "MotionLow.safetensors"); + assert_eq!(p["60"]["inputs"]["strength_model"], json!(0.5)); + assert_eq!(p["62"]["inputs"]["strength_model"], json!(0.0)); + } + + // No scale: the catalog's strength. + let (_, job, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": PHOTO, "template": "animate-picks", "slots": [{"lora": "motion"}]})), + ) + .await; + assert_eq!(job["slots"][0]["scale"], json!(1.0)); + + // A LoRA downloaded since the last listing is found by re-reading it. + fx.mock.lora_names.lock().unwrap().extend([ + "NewHigh.safetensors".to_string(), + "NewLow.safetensors".to_string(), + ]); + let (st, job, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": PHOTO, "template": "animate-picks", "slots": [{"lora": "new"}, {"lora": "motion"}]})), + ) + .await; + assert_eq!(st, 200, "{job}"); + assert_eq!(job["slots"][0]["lora"], "new"); + + // Unknown, duplicate, too many, or on a template without slots: 400. + for (template, slots) in [ + ("animate-picks", json!([{"lora": "nope"}])), + ( + "animate-picks", + json!([{"lora": "motion"}, {"lora": "motion"}]), + ), + ( + "animate-picks", + json!([{"lora": "motion"}, {"lora": "new"}, {"lora": "lonely"}]), + ), + ("restore", json!([{"lora": "motion"}])), + ] { + let (st, err, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": PHOTO, "template": template, "slots": slots})), + ) + .await; + assert_eq!(st, 400, "{template} {err}"); + } +} diff --git a/src/enhance/lora_catalog.rs b/src/enhance/lora_catalog.rs new file mode 100644 index 0000000..d64bb54 --- /dev/null +++ b/src/enhance/lora_catalog.rs @@ -0,0 +1,336 @@ +//! Pickable LoRAs (#220): the LoRA files ComfyUI can see, paired into the +//! high-noise / low-noise sets Wan 2.2's two experts need. +//! +//! Nothing is configured per LoRA: a file whose name has a `high` or `low` +//! token pairs with the file that has the other token and is otherwise the +//! same name (`BounceHighWan2_2` + `BounceLowWan2_2`, `x-HIGH` + `x-LOW`, +//! `…_high_noise` + `…_low_noise`). A file without a partner isn't offered. +//! The ids are derived from the names, so they stay stable across restarts +//! and refreshes. + +use std::collections::{BTreeMap, HashSet}; + +use serde::Serialize; + +/// One pickable LoRA: the pair of files and how the app shows it. +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +pub struct CatalogEntry { + /// Stable id clients send (`POST /enhance { slots: [{lora}] }`). + pub id: String, + pub label: String, + /// `lora_name` values as ComfyUI lists them (folder prefix included). + #[serde(skip)] + pub high: String, + #[serde(skip)] + pub low: String, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Noise { + High, + Low, +} + +/// Pair `names` (ComfyUI's `lora_name` choices) into catalog entries, sorted +/// by label. `exclude` holds names that never appear (the template's own +/// fixed LoRAs), compared exactly; `exclude_substrings` drops any name that +/// contains one of them (case-insensitive). +pub fn build( + names: &[String], + exclude: &HashSet, + exclude_substrings: &[String], +) -> Vec { + let lowered: Vec = exclude_substrings + .iter() + .map(|s| s.to_lowercase()) + .filter(|s| !s.is_empty()) + .collect(); + // key (the name with its noise token masked) -> (high, low, masked stem) + let mut sides: BTreeMap, Option<&String>, String)> = BTreeMap::new(); + for name in names { + if exclude.contains(name) { + continue; + } + let lower = name.to_lowercase(); + if lowered.iter().any(|s| lower.contains(s)) { + continue; + } + let Some((noise, masked)) = classify(name) else { + continue; + }; + let entry = sides + .entry(masked.to_lowercase()) + .or_insert((None, None, masked)); + match noise { + Noise::High => entry.0 = entry.0.or(Some(name)), + Noise::Low => entry.1 = entry.1.or(Some(name)), + } + } + + let mut taken = HashSet::new(); + let mut out: Vec = sides + .into_values() + .filter_map(|(high, low, masked)| { + let (high, low) = (high?, low?); + let base = slug(&masked); + let mut id = base.clone(); + let mut n = 2; + while !taken.insert(id.clone()) { + id = format!("{base}-{n}"); + n += 1; + } + Some(CatalogEntry { + id, + label: label(&masked), + high: high.clone(), + low: low.clone(), + }) + }) + .collect(); + out.sort_by(|a, b| { + a.label + .to_lowercase() + .cmp(&b.label.to_lowercase()) + .then(a.id.cmp(&b.id)) + }); + out +} + +/// The name's noise side and the name with that token masked as `{}` +/// (folder and extension kept, so two folders never pair with each other). +/// +/// A token counts when it starts the name, follows a non-letter, or is +/// capitalised (camel case), so `slow` or `Glowing` are never a `low`. The +/// last such token wins (`High…_low_noise` is low). +fn classify(name: &str) -> Option<(Noise, String)> { + let (dir, file) = match name.rfind(['/', '\\']) { + Some(i) => name.split_at(i + 1), + None => ("", name), + }; + let stem_end = file.rfind('.').unwrap_or(file.len()); + let stem = &file[..stem_end]; + let lower = stem.to_lowercase(); + let mut best: Option<(usize, usize, Noise)> = None; + for (token, noise) in [("high", Noise::High), ("low", Noise::Low)] { + for (at, _) in lower.match_indices(token) { + let original = &stem[at..at + token.len()]; + let boundary = at == 0 + || !stem[..at] + .chars() + .next_back() + .is_some_and(|c| c.is_alphabetic()) + || original.starts_with(char::is_uppercase); + // Not the start of a longer lowercase word (`lower`, `highway`). + let after = stem[at + token.len()..].chars().next(); + let ends = match after { + None => true, + Some(c) => { + !c.is_alphabetic() + || c.is_uppercase() + || original.chars().all(char::is_uppercase) + } + }; + if boundary && ends && best.is_none_or(|(b, _, _)| at > b) { + best = Some((at, token.len(), noise)); + } + } + } + let (at, len, noise) = best?; + let masked = format!( + "{dir}{}{{}}{}{}", + &stem[..at], + &stem[at + len..], + &file[stem_end..] + ); + Some((noise, masked)) +} + +/// `wan22/Bounce{}Wan2_2.safetensors` -> `bounce-wan2-2`. +fn slug(masked: &str) -> String { + let file = masked.rsplit(['/', '\\']).next().unwrap_or(masked); + let stem = file.rsplit_once('.').map(|(s, _)| s).unwrap_or(file); + let mut out = String::new(); + for c in split_words(stem).join(" ").chars() { + if c.is_ascii_alphanumeric() { + out.push(c.to_ascii_lowercase()); + } else if !out.ends_with('-') { + out.push('-'); + } + } + let out = out.trim_matches('-').to_string(); + if out.is_empty() { "lora".into() } else { out } +} + +/// `wan22/Bounce{}Wan2_2.safetensors` -> `Bounce Wan2 2`. +fn label(masked: &str) -> String { + let file = masked.rsplit(['/', '\\']).next().unwrap_or(masked); + let stem = file.rsplit_once('.').map(|(s, _)| s).unwrap_or(file); + let words = split_words(stem); + if words.is_empty() { + stem.to_string() + } else { + words.join(" ") + } +} + +/// Words of a stem: split on separators, the noise mask and camel case. +fn split_words(stem: &str) -> Vec { + let mut words = Vec::new(); + let mut cur = String::new(); + let mut prev: Option = None; + for c in stem.replace("{}", " ").chars() { + if c == '_' || c == '-' || c == ' ' { + if !cur.is_empty() { + words.push(std::mem::take(&mut cur)); + } + } else { + if c.is_uppercase() && prev.is_some_and(|p| p.is_lowercase()) && !cur.is_empty() { + words.push(std::mem::take(&mut cur)); + } + cur.push(c); + } + prev = Some(c); + } + if !cur.is_empty() { + words.push(cur); + } + words +} + +/// ComfyUI's `lora_name` choices from `GET /object_info/LoraLoaderModelOnly`, +/// in either shape: `[[..names]]` (classic) or `["COMBO", {"options": [..]}]`. +pub fn names_from_object_info(body: &serde_json::Value) -> Vec { + let spec = &body["LoraLoaderModelOnly"]["input"]["required"]["lora_name"]; + let list = if spec[0].is_array() { + &spec[0] + } else { + &spec[1]["options"] + }; + list.as_array() + .map(|a| { + a.iter() + .filter_map(|v| v.as_str().map(str::to_string)) + .collect() + }) + .unwrap_or_default() +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn names(v: &[&str]) -> Vec { + v.iter().map(|s| s.to_string()).collect() + } + + fn ids(c: &[CatalogEntry]) -> Vec<(&str, &str)> { + c.iter() + .map(|e| (e.id.as_str(), e.label.as_str())) + .collect() + } + + #[test] + fn pairs_the_usual_naming_styles() { + let c = build( + &names(&[ + "StyleHighWan2_2.safetensors", + "StyleLowWan2_2.safetensors", + "Motionv2.5-HIGH.safetensors", + "Motionv2.5-LOW.safetensors", + "Detail_High.safetensors", + "Detail_Low.safetensors", + "wan_effect_high_noise.safetensors", + "wan_effect_low_noise.safetensors", + ]), + &HashSet::new(), + &[], + ); + assert_eq!( + ids(&c), + [ + ("detail", "Detail"), + ("motionv2-5", "Motionv2.5"), + ("style-wan2-2", "Style Wan2 2"), + ("wan-effect-noise", "wan effect noise"), + ] + ); + let style = c.iter().find(|e| e.id == "style-wan2-2").unwrap(); + assert_eq!(style.high, "StyleHighWan2_2.safetensors"); + assert_eq!(style.low, "StyleLowWan2_2.safetensors"); + } + + #[test] + fn unpaired_and_tokenless_files_are_left_out() { + let c = build( + &names(&[ + "OnlyHigh_High.safetensors", + "distill_rank64_bf16.safetensors", + "slow_motion.safetensors", + "Glowing_low.safetensors", + "Glowing_high.safetensors", + "lower_body_high.safetensors", + ]), + &HashSet::new(), + &[], + ); + // `slow`, `Glowing` and `lower` hold no token; Glowing pairs on its suffix. + assert_eq!(ids(&c), [("glowing", "Glowing")]); + } + + #[test] + fn excludes_fixed_loras_and_substrings() { + let all = names(&[ + "wan2.2_i2v_lightx2v_4steps_lora_v1_high_noise.safetensors", + "wan2.2_i2v_lightx2v_4steps_lora_v1_low_noise.safetensors", + "StyleHigh.safetensors", + "StyleLow.safetensors", + "SecretHigh.safetensors", + "SecretLow.safetensors", + ]); + let fixed: HashSet = all[..2].iter().cloned().collect(); + let c = build(&all, &fixed, &["secret".into()]); + assert_eq!(ids(&c), [("style", "Style")]); + } + + #[test] + fn folders_keep_pairs_apart_and_ids_stay_unique() { + let c = build( + &names(&[ + "a/StyleHigh.safetensors", + "a/StyleLow.safetensors", + "b\\StyleHigh.safetensors", + "b\\StyleLow.safetensors", + "a/StyleHigh.safetensors", + ]), + &HashSet::new(), + &[], + ); + assert_eq!(c.len(), 2); + let mut got: Vec<_> = c.iter().map(|e| (e.id.as_str(), e.high.as_str())).collect(); + got.sort(); + assert_eq!( + got, + [ + ("style", "a/StyleHigh.safetensors"), + ("style-2", "b\\StyleHigh.safetensors") + ] + ); + } + + #[test] + fn reads_both_object_info_shapes() { + let classic = json!({"LoraLoaderModelOnly": {"input": {"required": { + "lora_name": [["a.safetensors", "b.safetensors"]], "strength_model": ["FLOAT", {}] + }}}}); + assert_eq!( + names_from_object_info(&classic), + names(&["a.safetensors", "b.safetensors"]) + ); + let combo = json!({"LoraLoaderModelOnly": {"input": {"required": { + "lora_name": ["COMBO", {"options": ["c.safetensors"]}] + }}}}); + assert_eq!(names_from_object_info(&combo), names(&["c.safetensors"])); + assert!(names_from_object_info(&json!({})).is_empty()); + } +} diff --git a/src/enhance/mod.rs b/src/enhance/mod.rs index 1f8c885..9ecf76d 100644 --- a/src/enhance/mod.rs +++ b/src/enhance/mod.rs @@ -14,6 +14,7 @@ pub mod comfy; pub mod config; pub mod handlers; +pub mod lora_catalog; pub mod path_map; pub mod service; pub mod template; diff --git a/src/enhance/service.rs b/src/enhance/service.rs index a189d89..0aa3a4a 100644 --- a/src/enhance/service.rs +++ b/src/enhance/service.rs @@ -15,8 +15,9 @@ use serde::{Deserialize, Serialize}; use super::comfy::{ApproveError, ComfyClient, ComfyError, HistoryEntry, Staged, SubmitError}; use super::config::{EnhanceConfig, Loaded, PromptMode}; +use super::lora_catalog::{self, CatalogEntry}; use super::path_map::PathMap; -use super::template::{Template, random_seed}; +use super::template::{SlotPick, Template, random_seed}; use crate::database::enhance_dao::{ EnhanceJobDao, EnhanceJobRow, JobUpdate, RecentPrompt, STATUS_APPROVED, STATUS_EXPIRED, STATUS_FAILED, STATUS_QUEUED, STATUS_READY, STATUS_RUNNING, STATUS_UNAVAILABLE, @@ -165,6 +166,27 @@ pub struct EnhanceService { idle: Mutex, /// Serialises approvals (rare; keeps "approve twice" simple). approve_lock: tokio::sync::Mutex<()>, + /// ComfyUI's LoRA files and when they were read (#220). + lora_names: Mutex)>>, +} + +/// How long ComfyUI's LoRA list is reused; a new download shows up after it. +const LORA_LIST_TTL: Duration = Duration::from_secs(60); + +/// One `POST /enhance { slots }` entry: a catalog id and an optional scale. +#[derive(Debug, Clone, PartialEq, serde::Deserialize)] +pub struct SlotRequest { + pub lora: String, + #[serde(default)] + pub scale: Option, +} + +/// A resolved pick, as stored on the job and returned in its view. +#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)] +pub struct SlotChoice { + pub lora: String, + pub label: String, + pub scale: f64, } impl EnhanceService { @@ -233,9 +255,116 @@ impl EnhanceService { unreachable_since: Mutex::new(HashMap::new()), idle: Mutex::new(IdleState::default()), approve_lock: tokio::sync::Mutex::new(()), + lora_names: Mutex::new(None), }) } + /// ComfyUI's LoRA files, cached for [`LORA_LIST_TTL`]. `fresh` re-reads + /// them (a submit naming an id the cache doesn't know). An unreachable + /// ComfyUI keeps the last list. + async fn lora_names(&self, fresh: bool) -> Vec { + let Some(client) = &self.client else { + return Vec::new(); + }; + if !fresh + && let Some((at, names)) = &*self.lora_names.lock().unwrap() + && at.elapsed() < LORA_LIST_TTL + { + return names.clone(); + } + match client.lora_names().await { + Ok(names) => { + *self.lora_names.lock().unwrap() = Some((Instant::now(), names.clone())); + names + } + Err(e) => { + warn!("enhance: couldn't list ComfyUI's LoRAs: {e}"); + self.lora_names + .lock() + .unwrap() + .as_ref() + .map(|(_, n)| n.clone()) + .unwrap_or_default() + } + } + } + + /// The LoRAs a template's slots can take; empty without slots. + pub async fn catalog(&self, template: &Template) -> Vec { + if template.slot_count() == 0 { + return Vec::new(); + } + let names = self.lora_names(false).await; + lora_catalog::build( + &names, + template.fixed_loras(), + &template.cfg.catalog.exclude, + ) + } + + /// Turn `slots` into picks for the template's slots, checking the ids + /// against the catalog (re-read once if one is unknown). + async fn resolve_slots( + &self, + template: &Template, + slots: &[SlotRequest], + ) -> Result<(Vec, Vec), EnhanceError> { + if slots.is_empty() { + return Ok((Vec::new(), Vec::new())); + } + if slots.len() > template.slot_count() { + return Err(EnhanceError::BadRequest(if template.slot_count() == 0 { + "this template doesn't take picked LoRAs".to_string() + } else { + format!( + "this template takes at most {} picked LoRAs", + template.slot_count() + ) + })); + } + let mut seen = std::collections::HashSet::new(); + if let Some(dup) = slots.iter().find(|s| !seen.insert(s.lora.as_str())) { + return Err(EnhanceError::BadRequest(format!( + "LoRA {:?} is picked twice", + dup.lora + ))); + } + let mut catalog = self.catalog(template).await; + if slots + .iter() + .any(|s| !catalog.iter().any(|e| e.id == s.lora)) + { + let names = self.lora_names(true).await; + catalog = lora_catalog::build( + &names, + template.fixed_loras(), + &template.cfg.catalog.exclude, + ); + } + let mut picks = Vec::new(); + let mut choices = Vec::new(); + for s in slots { + let Some(entry) = catalog.iter().find(|e| e.id == s.lora) else { + return Err(EnhanceError::BadRequest(format!( + "no LoRA {:?} (it may have been removed from ComfyUI)", + s.lora + ))); + }; + let scale = template.cfg.catalog.resolve(s.scale); + picks.push(SlotPick { + high: entry.high.clone(), + low: entry.low.clone(), + scale, + }); + choices.push(SlotChoice { + lora: entry.id.clone(), + label: entry.label.clone(), + scale, + }); + } + Ok((picks, choices)) + } + pub fn is_configured(&self) -> bool { self.client.is_some() } @@ -293,6 +422,7 @@ impl EnhanceService { prompt: Option<&str>, reference: Option<&SourceFile>, loras: &std::collections::BTreeMap, + slots: &[SlotRequest], ) -> Result { let Some(client) = &self.client else { return Err(EnhanceError::Disabled); @@ -319,12 +449,14 @@ impl EnhanceService { ), None => None, }; + let (picks, choices) = self.resolve_slots(template, slots).await?; let filled = template .fill( &source, prompt, reference_path.as_deref(), loras, + &picks, &mut random_seed, ) .map_err(|e| EnhanceError::BadRequest(e.0))?; @@ -374,6 +506,8 @@ impl EnhanceService { run_ms: None, loras: (!filled.loras.is_empty()) .then(|| serde_json::to_string(&filled.loras).unwrap_or_default()), + lora_slots: (!choices.is_empty()) + .then(|| serde_json::to_string(&choices).unwrap_or_default()), }; self.dao .lock() diff --git a/src/enhance/template.rs b/src/enhance/template.rs index acf1549..81bbea8 100644 --- a/src/enhance/template.rs +++ b/src/enhance/template.rs @@ -18,6 +18,9 @@ pub const TITLE_PROMPT: &str = "imageapi:prompt"; pub const TITLE_REFERENCE: &str = "imageapi:reference"; /// `imageapi:lora:`: a LoRA loader the app can switch off (#218). pub const TITLE_LORA_PREFIX: &str = "imageapi:lora:"; +/// `imageapi:lora-slot::high|low`: a LoRA loader the app fills with a +/// LoRA picked from the catalog (#220); unused slots run at strength 0. +pub const TITLE_SLOT_PREFIX: &str = "imageapi:lora-slot:"; /// LoRA loader inputs a switch scales (0 = off). `strength_model` is /// required (LoraLoaderModelOnly); `strength_clip` is scaled when present @@ -61,6 +64,11 @@ pub struct Template { reference_id: Option, /// Declared LoRA switches that have nodes, in `enhance.toml` order. loras: Vec, + /// Pickable-LoRA slots, by slot number. + slots: Vec, + /// `lora_name`s the workflow loads outside the slots (lightx2v, switch + /// LoRAs): never offered in the catalog. + fixed_loras: std::collections::HashSet, /// Image or video, from the save node's class. pub output: OutputKind, /// Startup warnings (stray output nodes). @@ -75,6 +83,22 @@ struct LoraSwitch { nodes: Vec<(String, Vec<(&'static str, f64)>)>, } +/// One pickable-LoRA slot: its high- and low-noise loader nodes, each with +/// its exported strengths (a pick's scale multiplies them). +#[derive(Debug, Clone)] +struct Slot { + high: Vec<(String, Vec<(&'static str, f64)>)>, + low: Vec<(String, Vec<(&'static str, f64)>)>, +} + +/// A LoRA picked for one slot: its two files and the resolved scale. +#[derive(Debug, Clone, PartialEq)] +pub struct SlotPick { + pub high: String, + pub low: String, + pub scale: f64, +} + /// Why a request can't be turned into a ComfyUI prompt. The message is for /// the caller, so it names the problem in the user's terms. #[derive(Debug, Clone, PartialEq, Eq)] @@ -227,6 +251,69 @@ impl Template { )); } + // Pickable-LoRA slots, and the LoRAs loaded outside them. + let exported = |id: &str| -> Vec<(&'static str, f64)> { + let inputs = inputs_of(&workflow, id); + LORA_STRENGTH_INPUTS + .iter() + .filter_map(|k| inputs.get(*k).and_then(Value::as_f64).map(|v| (*k, v))) + .collect() + }; + let mut slot_map: BTreeMap = BTreeMap::new(); + let mut fixed_loras = std::collections::HashSet::new(); + for (id, node) in &workflow { + let Some(rest) = title_of(node).and_then(|t| t.strip_prefix(TITLE_SLOT_PREFIX)) else { + if let Some(name) = inputs_of(&workflow, id) + .get("lora_name") + .and_then(Value::as_str) + { + fixed_loras.insert(name.to_string()); + } + continue; + }; + let title = format!("{TITLE_SLOT_PREFIX}{rest}"); + let (n, role) = rest + .split_once(':') + .and_then(|(n, role)| Some((n.parse::().ok().filter(|n| *n > 0)?, role))) + .ok_or_else(|| { + anyhow::anyhow!( + "node {id} titled {title:?}: expected {TITLE_SLOT_PREFIX}:high or :low (n from 1)" + ) + })?; + let inputs = inputs_of(&workflow, id); + if !inputs.get("lora_name").is_some_and(Value::is_string) + || !inputs.get("strength_model").is_some_and(Value::is_number) + { + bail!( + "node {id} titled {title:?} needs a lora_name and a numeric strength_model \ + (use a LoraLoaderModelOnly or LoraLoader)" + ); + } + let slot = slot_map.entry(n).or_insert(Slot { + high: Vec::new(), + low: Vec::new(), + }); + match role { + "high" => slot.high.push((id.clone(), exported(id))), + "low" => slot.low.push((id.clone(), exported(id))), + _ => bail!("node {id} titled {title:?}: the role must be high or low"), + } + } + for (n, slot) in &slot_map { + if slot.high.is_empty() || slot.low.is_empty() { + bail!("LoRA slot {n} needs both a :high and a :low node"); + } + } + if let Some((last, _)) = slot_map.last_key_value() + && *last as usize != slot_map.len() + { + bail!( + "LoRA slots must be numbered 1..{} without gaps", + slot_map.len() + ); + } + let slots: Vec = slot_map.into_values().collect(); + for (id, node) in &workflow { let class = node["class_type"].as_str().unwrap_or_default(); if STRAY_OUTPUT_CLASSES.contains(&class) { @@ -248,6 +335,8 @@ impl Template { prompt, reference_id, loras, + slots, + fixed_loras, output, warnings, }) @@ -265,6 +354,16 @@ impl Template { self.cfg.prompt_mode == PromptMode::Replace } + /// How many pickable-LoRA slots the workflow has (#220). + pub fn slot_count(&self) -> usize { + self.slots.len() + } + + /// `lora_name`s loaded outside the slots; kept out of the catalog. + pub fn fixed_loras(&self) -> &std::collections::HashSet { + &self.fixed_loras + } + /// The LoRA switches clients can offer (declared and present in the /// workflow), in `enhance.toml` order. pub fn lora_switches(&self) -> impl Iterator { @@ -288,8 +387,19 @@ impl Template { prompt: Option<&str>, reference: Option<&str>, loras: &BTreeMap, + picks: &[SlotPick], next_seed: &mut dyn FnMut() -> u64, ) -> Result { + if picks.len() > self.slots.len() { + return Err(FillError(if self.slots.is_empty() { + "this template doesn't take picked LoRAs".to_string() + } else { + format!( + "this template takes at most {} picked LoRAs", + self.slots.len() + ) + })); + } if let Some(unknown) = loras .keys() .find(|id| !self.loras.iter().any(|l| &l.cfg.id == *id)) @@ -362,6 +472,26 @@ impl Template { } } + // Picked LoRAs fill the slots in order; the rest run at 0 (their + // exported lora_name stays, a file ComfyUI knows). + for (i, slot) in self.slots.iter().enumerate() { + let pick = picks.get(i).filter(|p| p.scale > 0.0); + for (nodes, file) in [ + (&slot.high, pick.map(|p| p.high.as_str())), + (&slot.low, pick.map(|p| p.low.as_str())), + ] { + for (id, exported) in nodes { + if let Some(file) = file { + set_input(&mut wf, id, "lora_name", json!(file)); + } + let scale = pick.map_or(0.0, |p| p.scale); + for (input, value) in exported { + set_input(&mut wf, id, input, json!(value * scale)); + } + } + } + } + for node in wf.values_mut() { let Some(inputs) = node.get_mut("inputs").and_then(Value::as_object_mut) else { continue; @@ -435,6 +565,17 @@ mod tests { /// The Animate template plus a Style pair titled `imageapi:lora:style` /// (nodes 50 high, 51 low). pub const ANIMATE_LORAS: &str = include_str!("../../tests/fixtures/enhance/animate_loras.json"); + /// ANIMATE_LORAS plus two pickable slots chained after the switch: + /// 60/61 (slot 1 high/low), 62/63 (slot 2). + pub const ANIMATE_SLOTS: &str = include_str!("../../tests/fixtures/enhance/animate_slots.json"); + + fn pick(name: &str, scale: f64) -> SlotPick { + SlotPick { + high: format!("{name}_high.safetensors"), + low: format!("{name}_low.safetensors"), + scale, + } + } fn lora(id: &str, default: bool) -> LoraConfig { LoraConfig { @@ -475,6 +616,7 @@ mod tests { prompt_join: " ".into(), output_megapixels: 1.0, loras: Vec::new(), + catalog: Default::default(), } } @@ -547,6 +689,7 @@ mod tests { Some("more"), None, &BTreeMap::new(), + &[], &mut seeds(), ) .unwrap(); @@ -566,7 +709,7 @@ mod tests { fn fill_sets_path_forces_approval_and_randomises_every_seed() { let t = load(RESTORE, PromptMode::Append).unwrap(); let f = t - .fill("Z:\\p\\a.jpg", None, None, &none(), &mut seeds()) + .fill("Z:\\p\\a.jpg", None, None, &none(), &[], &mut seeds()) .unwrap(); assert_eq!(input(&f.prompt, "37", "path"), "Z:\\p\\a.jpg"); assert_eq!(input(&f.prompt, "36", "require_approval"), true); @@ -579,8 +722,8 @@ mod tests { assert_ne!(a, 774744764437710); // Two fills with a real RNG differ. let mut rng = random_seed; - let x = t.fill("p", None, None, &none(), &mut rng).unwrap(); - let y = t.fill("p", None, None, &none(), &mut rng).unwrap(); + let x = t.fill("p", None, None, &none(), &[], &mut rng).unwrap(); + let y = t.fill("p", None, None, &none(), &[], &mut rng).unwrap(); assert_ne!(input(&x.prompt, "5", "seed"), input(&y.prompt, "5", "seed")); } @@ -606,7 +749,9 @@ mod tests { // SeedVR2VideoUpscaler rejects seeds above 2^32 - 1 at validation. let t = load(RESTORE, PromptMode::Append).unwrap(); for _ in 0..200 { - let f = t.fill("p", None, None, &none(), &mut random_seed).unwrap(); + let f = t + .fill("p", None, None, &none(), &[], &mut random_seed) + .unwrap(); assert!(input(&f.prompt, "5", "seed").as_u64().unwrap() <= u64::from(u32::MAX)); } } @@ -615,7 +760,7 @@ mod tests { fn force_approval_even_if_the_export_had_it_off() { let off = RESTORE.replace("\"require_approval\": true", "\"require_approval\": false"); let t = load(&off, PromptMode::Append).unwrap(); - let f = t.fill("p", None, None, &none(), &mut seeds()).unwrap(); + let f = t.fill("p", None, None, &none(), &[], &mut seeds()).unwrap(); assert_eq!(input(&f.prompt, "36", "require_approval"), true); } @@ -633,6 +778,7 @@ mod tests { Some(" warmer "), None, &BTreeMap::new(), + &[], &mut seeds(), ) .unwrap(); @@ -641,10 +787,10 @@ mod tests { "Increase the clarity\n\nAdditional instructions: warmer" ); let f = t - .fill("p", Some(" "), None, &none(), &mut seeds()) + .fill("p", Some(" "), None, &none(), &[], &mut seeds()) .unwrap(); assert_eq!(input(&f.prompt, "13", "value"), base); - let f = t.fill("p", None, None, &none(), &mut seeds()).unwrap(); + let f = t.fill("p", None, None, &none(), &[], &mut seeds()).unwrap(); assert_eq!(input(&f.prompt, "13", "value"), base); } @@ -657,6 +803,7 @@ mod tests { Some("Make it warmer."), None, &BTreeMap::new(), + &[], &mut seeds(), ) .unwrap(); @@ -667,10 +814,12 @@ mod tests { #[test] fn replace_mode_requires_and_uses_the_instruction() { let t = load(EDIT, PromptMode::Replace).unwrap(); - let err = t.fill("p", None, None, &none(), &mut seeds()).unwrap_err(); + let err = t + .fill("p", None, None, &none(), &[], &mut seeds()) + .unwrap_err(); assert!(err.0.contains("instruction")); assert!( - t.fill("p", Some(" "), None, &none(), &mut seeds()) + t.fill("p", Some(" "), None, &none(), &[], &mut seeds()) .is_err() ); let f = t @@ -679,6 +828,7 @@ mod tests { Some("make it black and white"), None, &BTreeMap::new(), + &[], &mut seeds(), ) .unwrap(); @@ -692,13 +842,16 @@ mod tests { // 14 links to 13; drop that link so the graph stays plausible. v["14"]["inputs"].as_object_mut().unwrap().remove("prompt"); let t = Template::from_value(cfg("t", PromptMode::Append), v).unwrap(); - assert!(t.fill("p", Some("x"), None, &none(), &mut seeds()).is_err()); - assert!(t.fill("p", None, None, &none(), &mut seeds()).is_ok()); + assert!( + t.fill("p", Some("x"), None, &none(), &[], &mut seeds()) + .is_err() + ); + assert!(t.fill("p", None, None, &none(), &[], &mut seeds()).is_ok()); let restore = load(RESTORE, PromptMode::Append).unwrap(); assert!( restore - .fill("p", None, Some("r"), &none(), &mut seeds()) + .fill("p", None, Some("r"), &none(), &[], &mut seeds()) .is_err() ); } @@ -712,6 +865,7 @@ mod tests { Some("x"), Some("Z:\\ref.jpg"), &BTreeMap::new(), + &[], &mut seeds(), ) .unwrap(); @@ -724,7 +878,7 @@ mod tests { fn no_reference_removes_the_node_and_every_link_to_it() { let t = load(EDIT, PromptMode::Replace).unwrap(); let f = t - .fill("src", Some("x"), None, &none(), &mut seeds()) + .fill("src", Some("x"), None, &none(), &[], &mut seeds()) .unwrap(); assert!(f.dropped_reference); let wf = f.prompt.as_object().unwrap(); @@ -751,7 +905,7 @@ mod tests { assert!(t.warnings.is_empty(), "{:?}", t.warnings); // Default off: both the high and the low node go to 0. - let off = t.fill("p", None, None, &none(), &mut seeds()).unwrap(); + let off = t.fill("p", None, None, &none(), &[], &mut seeds()).unwrap(); assert_eq!(off.loras, scales(&[("style", 0.0)])); for id in ["50", "51"] { assert_eq!(input(&off.prompt, id, "strength_model"), &json!(0.0)); @@ -767,6 +921,7 @@ mod tests { None, None, &reqs(&[("style", LoraRequest::On(true))]), + &[], &mut seeds(), ) .unwrap(); @@ -791,6 +946,7 @@ mod tests { None, None, &reqs(&[("style", LoraRequest::Scale(0.5))]), + &[], &mut seeds(), ) .unwrap(); @@ -805,6 +961,7 @@ mod tests { None, None, &reqs(&[("style", LoraRequest::Scale(9.0))]), + &[], &mut seeds(), ) .unwrap(); @@ -815,6 +972,7 @@ mod tests { None, None, &reqs(&[("style", LoraRequest::Scale(0.0))]), + &[], &mut seeds(), ) .unwrap(); @@ -826,7 +984,7 @@ mod tests { let mut l = lora("style", true); l.strength = 0.7; let t = load_with_loras(ANIMATE_LORAS, vec![l]).unwrap(); - let f = t.fill("p", None, None, &none(), &mut seeds()).unwrap(); + let f = t.fill("p", None, None, &none(), &[], &mut seeds()).unwrap(); assert_eq!(f.loras, scales(&[("style", 0.7)])); assert_eq!(input(&f.prompt, "50", "strength_model"), &json!(0.7)); @@ -836,6 +994,7 @@ mod tests { None, None, &reqs(&[("nope", LoraRequest::On(true))]), + &[], &mut seeds(), ) .unwrap_err(); @@ -850,7 +1009,7 @@ mod tests { let mut c = cfg("animate", PromptMode::Append); c.loras = vec![lora("style", false)]; let t = Template::from_value(c, v).unwrap(); - let off = t.fill("p", None, None, &none(), &mut seeds()).unwrap(); + let off = t.fill("p", None, None, &none(), &[], &mut seeds()).unwrap(); assert_eq!(input(&off.prompt, "50", "strength_clip"), &json!(0.0)); let half = t .fill( @@ -858,6 +1017,7 @@ mod tests { None, None, &reqs(&[("style", LoraRequest::Scale(0.5))]), + &[], &mut seeds(), ) .unwrap(); @@ -874,7 +1034,7 @@ mod tests { "{:?}", t.warnings ); - let f = t.fill("p", None, None, &none(), &mut seeds()).unwrap(); + let f = t.fill("p", None, None, &none(), &[], &mut seeds()).unwrap(); assert!(f.loras.is_empty()); assert_eq!(input(&f.prompt, "50", "strength_model"), &json!(1.0)); assert!( @@ -883,7 +1043,8 @@ mod tests { None, None, &reqs(&[("style", LoraRequest::On(false))]), - &mut seeds() + &[], + &mut seeds(), ) .is_err() ); @@ -917,4 +1078,106 @@ mod tests { .to_string(); assert!(err.contains("strength_model"), "{err}"); } + + #[test] + fn slots_are_found_and_fixed_loras_listed() { + let t = load(ANIMATE_SLOTS, PromptMode::Append).unwrap(); + assert_eq!(t.slot_count(), 2); + let mut fixed: Vec<_> = t.fixed_loras().iter().cloned().collect(); + fixed.sort(); + assert_eq!( + fixed, + [ + "example_style_high.safetensors", + "example_style_low.safetensors", + "wan2.2_i2v_lightx2v_4steps_lora_v1_high_noise.safetensors", + "wan2.2_i2v_lightx2v_4steps_lora_v1_low_noise.safetensors", + ] + ); + assert_eq!( + load(ANIMATE_LORAS, PromptMode::Append) + .unwrap() + .slot_count(), + 0 + ); + } + + #[test] + fn picks_fill_slots_in_order_and_unused_slots_run_at_zero() { + let t = load(ANIMATE_SLOTS, PromptMode::Append).unwrap(); + let f = t + .fill( + "p", + None, + None, + &none(), + &[pick("motion", 0.5)], + &mut seeds(), + ) + .unwrap(); + assert_eq!( + input(&f.prompt, "60", "lora_name"), + &json!("motion_high.safetensors") + ); + assert_eq!( + input(&f.prompt, "61", "lora_name"), + &json!("motion_low.safetensors") + ); + assert_eq!(input(&f.prompt, "60", "strength_model"), &json!(0.5)); + assert_eq!(input(&f.prompt, "61", "strength_model"), &json!(0.5)); + // Slot 2 unused: its exported file stays, at strength 0. + assert_eq!( + input(&f.prompt, "62", "lora_name"), + &json!("example_slot_high.safetensors") + ); + assert_eq!(input(&f.prompt, "62", "strength_model"), &json!(0.0)); + assert_eq!(input(&f.prompt, "63", "strength_model"), &json!(0.0)); + // The chain isn't rewired. + assert_eq!(input(&f.prompt, "14", "model"), &json!(["62", 0])); + + // No picks at all: every slot off. + let f = t.fill("p", None, None, &none(), &[], &mut seeds()).unwrap(); + assert_eq!(input(&f.prompt, "60", "strength_model"), &json!(0.0)); + + let err = t + .fill( + "p", + None, + None, + &none(), + &[pick("a", 1.0), pick("b", 1.0), pick("c", 1.0)], + &mut seeds(), + ) + .unwrap_err(); + assert!(err.0.contains("at most 2"), "{err}"); + let plain = load(ANIMATE_LORAS, PromptMode::Append).unwrap(); + assert!( + plain + .fill("p", None, None, &none(), &[pick("a", 1.0)], &mut seeds()) + .is_err() + ); + } + + #[test] + fn malformed_slots_are_rejected() { + let only_high = ANIMATE_SLOTS.replace("imageapi:lora-slot:2:low", "something else"); + let err = load(&only_high, PromptMode::Append) + .unwrap_err() + .to_string(); + assert!(err.contains("slot 2 needs both"), "{err}"); + let gap = ANIMATE_SLOTS + .replace("imageapi:lora-slot:2:high", "imageapi:lora-slot:3:high") + .replace("imageapi:lora-slot:2:low", "imageapi:lora-slot:3:low"); + let err = load(&gap, PromptMode::Append).unwrap_err().to_string(); + assert!(err.contains("without gaps"), "{err}"); + let role = ANIMATE_SLOTS.replace("imageapi:lora-slot:1:high", "imageapi:lora-slot:1:mid"); + assert!( + load(&role, PromptMode::Append) + .unwrap_err() + .to_string() + .contains("high or low") + ); + let no_n = ANIMATE_SLOTS.replace("imageapi:lora-slot:1:high", "imageapi:lora-slot:x:high"); + assert!(load(&no_n, PromptMode::Append).is_err()); + } } diff --git a/tests/fixtures/enhance/animate_slots.json b/tests/fixtures/enhance/animate_slots.json new file mode 100644 index 0000000..4788df3 --- /dev/null +++ b/tests/fixtures/enhance/animate_slots.json @@ -0,0 +1,404 @@ +{ + "1": { + "class_type": "ImageApiLoadImageFromPath", + "_meta": { + "title": "imageapi:source" + }, + "inputs": { + "path": "" + } + }, + "2": { + "class_type": "ImageApiTargetSize", + "_meta": { + "title": "Target Size (~0.92 MP, keeps aspect)" + }, + "inputs": { + "image": [ + "1", + 0 + ], + "max_megapixels": 0.92, + "min_megapixels": 0.92, + "multiple_of": 16 + } + }, + "10": { + "class_type": "UNETLoader", + "_meta": { + "title": "Wan 2.2 i2v high noise" + }, + "inputs": { + "unet_name": "wan2.2_i2v_high_noise_14B_fp8_scaled.safetensors", + "weight_dtype": "default" + } + }, + "11": { + "class_type": "UNETLoader", + "_meta": { + "title": "Wan 2.2 i2v low noise" + }, + "inputs": { + "unet_name": "wan2.2_i2v_low_noise_14B_fp8_scaled.safetensors", + "weight_dtype": "default" + } + }, + "12": { + "class_type": "LoraLoaderModelOnly", + "_meta": { + "title": "lightx2v 4-step (high)" + }, + "inputs": { + "model": [ + "10", + 0 + ], + "lora_name": "wan2.2_i2v_lightx2v_4steps_lora_v1_high_noise.safetensors", + "strength_model": 1.0 + } + }, + "13": { + "class_type": "LoraLoaderModelOnly", + "_meta": { + "title": "lightx2v 4-step (low)" + }, + "inputs": { + "model": [ + "11", + 0 + ], + "lora_name": "wan2.2_i2v_lightx2v_4steps_lora_v1_low_noise.safetensors", + "strength_model": 1.0 + } + }, + "14": { + "class_type": "ModelSamplingSD3", + "_meta": { + "title": "Shift (high)" + }, + "inputs": { + "model": [ + "62", + 0 + ], + "shift": 5.0 + } + }, + "15": { + "class_type": "ModelSamplingSD3", + "_meta": { + "title": "Shift (low)" + }, + "inputs": { + "model": [ + "63", + 0 + ], + "shift": 5.0 + } + }, + "40": { + "class_type": "PathchSageAttentionKJ", + "_meta": { + "title": "Sage attention (high)" + }, + "inputs": { + "model": [ + "14", + 0 + ], + "sage_attention": "auto" + } + }, + "41": { + "class_type": "PathchSageAttentionKJ", + "_meta": { + "title": "Sage attention (low)" + }, + "inputs": { + "model": [ + "15", + 0 + ], + "sage_attention": "auto" + } + }, + "16": { + "class_type": "CLIPLoader", + "_meta": { + "title": "Load CLIP" + }, + "inputs": { + "clip_name": "umt5_xxl_fp8_e4m3fn_scaled.safetensors", + "type": "wan", + "device": "default" + } + }, + "17": { + "class_type": "CLIPTextEncode", + "_meta": { + "title": "imageapi:prompt" + }, + "inputs": { + "clip": [ + "16", + 0 + ], + "text": "natural, lifelike motion with a slow, gentle camera movement, realistic, cinematic" + } + }, + "18": { + "class_type": "CLIPTextEncode", + "_meta": { + "title": "Negative" + }, + "inputs": { + "clip": [ + "16", + 0 + ], + "text": "static, frozen, blurry, low quality, distorted, warped faces, morphing, flicker, watermark, text" + } + }, + "19": { + "class_type": "VAELoader", + "_meta": { + "title": "Load VAE" + }, + "inputs": { + "vae_name": "wan_2.1_vae.safetensors" + } + }, + "21": { + "class_type": "WanImageToVideo", + "_meta": { + "title": "Wan Image To Video (81 frames)" + }, + "inputs": { + "positive": [ + "17", + 0 + ], + "negative": [ + "18", + 0 + ], + "vae": [ + "19", + 0 + ], + "start_image": [ + "1", + 0 + ], + "width": [ + "2", + 0 + ], + "height": [ + "2", + 1 + ], + "length": 81, + "batch_size": 1 + } + }, + "22": { + "class_type": "KSamplerAdvanced", + "_meta": { + "title": "Sampler (high, steps 0-2)" + }, + "inputs": { + "model": [ + "40", + 0 + ], + "positive": [ + "21", + 0 + ], + "negative": [ + "21", + 1 + ], + "latent_image": [ + "21", + 2 + ], + "add_noise": "enable", + "noise_seed": 1, + "steps": 4, + "cfg": 1.0, + "sampler_name": "euler", + "scheduler": "simple", + "start_at_step": 0, + "end_at_step": 2, + "return_with_leftover_noise": "enable" + } + }, + "23": { + "class_type": "KSamplerAdvanced", + "_meta": { + "title": "Sampler (low, steps 2-4)" + }, + "inputs": { + "model": [ + "41", + 0 + ], + "positive": [ + "21", + 0 + ], + "negative": [ + "21", + 1 + ], + "latent_image": [ + "22", + 0 + ], + "add_noise": "disable", + "noise_seed": 0, + "steps": 4, + "cfg": 1.0, + "sampler_name": "euler", + "scheduler": "simple", + "start_at_step": 2, + "end_at_step": 10000, + "return_with_leftover_noise": "disable" + } + }, + "24": { + "class_type": "VAEDecode", + "_meta": { + "title": "VAE Decode" + }, + "inputs": { + "samples": [ + "23", + 0 + ], + "vae": [ + "19", + 0 + ] + } + }, + "25": { + "class_type": "CreateVideo", + "_meta": { + "title": "Create Video (16 fps)" + }, + "inputs": { + "images": [ + "24", + 0 + ], + "fps": 16.0 + } + }, + "36": { + "class_type": "ImageApiSaveVideoVersion", + "_meta": { + "title": "imageapi:save" + }, + "inputs": { + "video": [ + "25", + 0 + ], + "source": [ + "1", + 1 + ], + "kind": "animated", + "subfolder": "aiedits", + "require_approval": true + } + }, + "50": { + "inputs": { + "lora_name": "example_style_high.safetensors", + "strength_model": 1.0, + "model": [ + "12", + 0 + ] + }, + "class_type": "LoraLoaderModelOnly", + "_meta": { + "title": "imageapi:lora:style" + } + }, + "51": { + "inputs": { + "lora_name": "example_style_low.safetensors", + "strength_model": 1.0, + "model": [ + "13", + 0 + ] + }, + "class_type": "LoraLoaderModelOnly", + "_meta": { + "title": "imageapi:lora:style" + } + }, + "60": { + "inputs": { + "lora_name": "example_slot_high.safetensors", + "strength_model": 1.0, + "model": [ + "50", + 0 + ] + }, + "class_type": "LoraLoaderModelOnly", + "_meta": { + "title": "imageapi:lora-slot:1:high" + } + }, + "61": { + "inputs": { + "lora_name": "example_slot_low.safetensors", + "strength_model": 1.0, + "model": [ + "51", + 0 + ] + }, + "class_type": "LoraLoaderModelOnly", + "_meta": { + "title": "imageapi:lora-slot:1:low" + } + }, + "62": { + "inputs": { + "lora_name": "example_slot_high.safetensors", + "strength_model": 1.0, + "model": [ + "60", + 0 + ] + }, + "class_type": "LoraLoaderModelOnly", + "_meta": { + "title": "imageapi:lora-slot:2:high" + } + }, + "63": { + "inputs": { + "lora_name": "example_slot_low.safetensors", + "strength_model": 1.0, + "model": [ + "61", + 0 + ] + }, + "class_type": "LoraLoaderModelOnly", + "_meta": { + "title": "imageapi:lora-slot:2:low" + } + } +} \ No newline at end of file