diff --git a/CLAUDE.md b/CLAUDE.md index 1a19348..480c9bc 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -459,14 +459,20 @@ auth and stays LAN-only, so ImageApi relays previews too. Keys: `url`, `timeout_secs`, `probe_ttl_secs`, `poll_secs`, `idle_free_secs`, `unreachable_grace_secs`, `default_template`, `warn_above_megapixels`, `[[path_map]]` (`from`, `to`), `[[template]]` (`name`, `label`, `description`, - `file`, `prompt_mode`, `prompt_join`, `output_megapixels`). + `file`, `prompt_mode`, `prompt_join`, `output_megapixels`, `[[template.lora]]` + (`id`, `label`, `default`, `strength`, `min`, `max`)). - **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 `ImageApiSaveVideoVersion` there makes the template's `output` `video`, an MP4 clip, else `image`), `imageapi:prompt` (optional; fills the first of `text` / `prompt` / `value`), - `imageapi:reference` (optional). `prompt_mode = "append"`: the user's text is joined + `imageapi:reference` (optional), `imageapi:lora:` (optional LoRA switches, #218: + a declared `[[template.lora]]` id drives every node with that title, e.g. Wan's high and + 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 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 @@ -512,8 +518,11 @@ GET /enhance/templates "templates": [{"name","label","description","accepts_prompt", "prompt_mode": "append"|"replace", "prompt_required", "accepts_reference","default","output_megapixels", - "warn_above_megapixels", "output": "image"|"video"}]} -POST /enhance {"path","template","prompt"?,"reference_path"?,"library"?} -> job + "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 + (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 GET /enhance/jobs?path=&library= -> {"jobs": [job, ...]} newest first, discarded hidden GET /enhance/jobs[?status=a,b] -> {"jobs": [...]} no path = the jobs manager: every photo's @@ -530,7 +539,8 @@ POST /enhance/jobs/{id}/approve {"n"} -> a /image/versions `versions[]` (or `cli POST /enhance/jobs/{id}/discard -> job (hidden from the list; ComfyUI untouched) job = {"id","status","queue_position"?,"error"?,"path","library_id","template", - "prompt"?,"reference_path"?,"created_at","updated_at", + "prompt"?,"reference_path"?,"loras"?: {id: scale} (resolved, 0 = off; absent without switches), + "created_at","updated_at", "run_ms": , "pending": [{"n","preview": "/enhance/jobs/{id}/preview/{n}","approved"}], "saved": [{"n","path","library_id"}]} diff --git a/enhance.example.toml b/enhance.example.toml index 9cc80a9..b52e0c4 100644 --- a/enhance.example.toml +++ b/enhance.example.toml @@ -63,3 +63,27 @@ prompt_mode = "replace" # name = "restore-dev" # label = "Restore (Dev, slow)" # file = "/media/nas/image_fibo2_restore_imageapi_API.json" + +# A video template with optional LoRAs the app can switch on or off and scale +# (#218). Each [[template.lora]] matches the workflow's LoRA loader nodes titled +# "imageapi:lora:" (Wan's high- and low-noise pair share one id). The +# request's scale multiplies each node's exported strength_model (and +# strength_clip), so 1.0 = as exported, 0 = off and a pair keeps its ratio; the +# graph isn't rewired. Any number of switches can be on at once (LoRA patches +# add up, so their order in the graph doesn't matter). Nodes titled that way +# with no entry here keep their exported strength. +[[template]] +name = "animate" +label = "Animate" +description = "A short clip from the photo." +file = "/media/nas/image_wan22_animate_imageapi_API.json" +prompt_join = ", " +output_megapixels = 0.92 + +[[template.lora]] +id = "style" # node title "imageapi:lora:style" +label = "Style" +default = false # default: false; used when the request omits it +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) diff --git a/migrations/2026-11-03-100000_add_enhance_loras/down.sql b/migrations/2026-11-03-100000_add_enhance_loras/down.sql new file mode 100644 index 0000000..d532b70 --- /dev/null +++ b/migrations/2026-11-03-100000_add_enhance_loras/down.sql @@ -0,0 +1 @@ +ALTER TABLE enhance_jobs DROP COLUMN loras; diff --git a/migrations/2026-11-03-100000_add_enhance_loras/up.sql b/migrations/2026-11-03-100000_add_enhance_loras/up.sql new file mode 100644 index 0000000..196b604 --- /dev/null +++ b/migrations/2026-11-03-100000_add_enhance_loras/up.sql @@ -0,0 +1,2 @@ +-- The job's resolved LoRA switches (#218): JSON {id: bool}; NULL = the template had none. +ALTER TABLE enhance_jobs ADD COLUMN loras TEXT; diff --git a/src/database/enhance_dao.rs b/src/database/enhance_dao.rs index 544594b..a426198 100644 --- a/src/database/enhance_dao.rs +++ b/src/database/enhance_dao.rs @@ -62,11 +62,15 @@ pub struct EnhanceJobRow { /// Milliseconds ComfyUI actually ran the job; set once it finishes. #[diesel(sql_type = Nullable)] pub run_ms: Option, + /// The resolved LoRA switches (#218), JSON `{id: bool}`; `None` when the + /// template had none. Owned by `crate::enhance` like `pending`. + #[diesel(sql_type = Nullable)] + pub loras: 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"; + discarded, created_at, updated_at, run_ms, loras"; /// One distinct prompt used across jobs, newest use first. #[derive(Debug, Clone, PartialEq, Eq, QueryableByName)] @@ -171,7 +175,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) @@ -191,6 +195,7 @@ impl EnhanceJobDao for SqliteEnhanceJobDao { .bind::(job.created_at) .bind::(job.updated_at) .bind::, _>(job.run_ms) + .bind::, _>(&job.loras) .execute(&mut *conn) .map(|_| ()) .map_err(|e| anyhow::anyhow!("Insert error: {}", e)) @@ -370,6 +375,7 @@ mod tests { created_at: created, updated_at: created, run_ms: None, + loras: None, } } diff --git a/src/database/schema.rs b/src/database/schema.rs index 2eb3d0d..6a742fb 100644 --- a/src/database/schema.rs +++ b/src/database/schema.rs @@ -73,6 +73,7 @@ diesel::table! { created_at -> BigInt, updated_at -> BigInt, run_ms -> Nullable, + loras -> Nullable, } } diff --git a/src/enhance/config.rs b/src/enhance/config.rs index 1cac49f..f958fb7 100644 --- a/src/enhance/config.rs +++ b/src/enhance/config.rs @@ -100,6 +100,72 @@ pub struct TemplateConfig { /// Approximate output size, for the client's large-original warning. #[serde(default = "default_output_mp")] pub output_megapixels: f64, + /// Optional LoRAs the app can switch on or off (#218), each matching the + /// workflow nodes titled `imageapi:lora:`. + #[serde(rename = "lora", default)] + pub loras: Vec, +} + +/// One LoRA switch (`[[template.lora]]`). +#[derive(Debug, Clone, PartialEq, Deserialize, serde::Serialize)] +#[serde(deny_unknown_fields)] +pub struct LoraConfig { + /// Matches the nodes titled `imageapi:lora:` (Wan's high and low + /// pair share one id); clients send it in `POST /enhance { loras }`. + pub id: String, + pub label: String, + /// On when the request doesn't say. + #[serde(default)] + pub default: bool, + /// Scale applied when on and the request gives no number: a factor on the + /// strengths exported in the workflow (1.0 = as exported), so a high/low + /// pair keeps its ratio. + #[serde(default = "default_lora_strength")] + pub strength: f64, + /// Slider bounds; a requested scale is clamped to them. + #[serde(default)] + pub min: f64, + #[serde(default = "default_lora_max")] + pub max: f64, +} + +impl LoraConfig { + /// The scale a request resolves to: `true` = `strength`, `false` or a + /// number <= 0 = off (0), another number clamped to `min..=max` (the + /// slider's range while on), absent = the default's. + pub fn resolve(&self, request: Option) -> f64 { + match request { + None if self.default => self.strength, + None => 0.0, + Some(LoraRequest::On(true)) => self.strength, + Some(LoraRequest::On(false)) => 0.0, + Some(LoraRequest::Scale(s)) if s <= 0.0 => 0.0, + Some(LoraRequest::Scale(s)) => s.clamp(self.min, self.max), + } + } +} + +/// One switch in `POST /enhance { loras }`: on/off, or a scale for the slider. +#[derive(Debug, Clone, Copy, PartialEq, Deserialize)] +#[serde(untagged)] +pub enum LoraRequest { + On(bool), + Scale(f64), +} + +fn default_lora_strength() -> f64 { + 1.0 +} + +fn default_lora_max() -> f64 { + 1.5 +} + +/// Template names and LoRA ids: non-empty, letters, digits, `_` and `-`. +fn valid_id(s: &str) -> bool { + !s.is_empty() + && s.chars() + .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-') } fn default_join() -> String { @@ -132,11 +198,7 @@ impl EnhanceConfig { } let mut seen = std::collections::HashSet::new(); for t in &self.templates { - let ok = !t.name.is_empty() - && t.name - .chars() - .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-'); - if !ok { + if !valid_id(&t.name) { bail!( "template name {:?} must be non-empty and use only letters, digits, '_' and '-'", t.name @@ -151,6 +213,38 @@ impl EnhanceConfig { if t.output_megapixels.is_nan() || t.output_megapixels <= 0.0 { bail!("template {:?}: output_megapixels must be > 0", t.name); } + let mut ids = std::collections::HashSet::new(); + for l in &t.loras { + if !valid_id(&l.id) { + bail!( + "template {:?}: lora id {:?} must be non-empty and use only letters, digits, '_' and '-'", + t.name, + l.id + ); + } + if !ids.insert(l.id.as_str()) { + bail!("template {:?}: duplicate lora id {:?}", t.name, l.id); + } + 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) + { + bail!( + "template {:?}: lora {:?} needs 0 <= min < max and strength within them \ + (got min {}, strength {}, max {})", + t.name, + l.id, + l.min, + l.strength, + l.max + ); + } + } } if let Some(d) = &self.default_template && !self.templates.iter().any(|t| &t.name == d) @@ -252,6 +346,74 @@ mod tests { assert!(parse("[[path_map]]\nfrom=''\nto='Z:\\\\'").is_err()); } + #[test] + fn lora_switches_parse_and_validate() { + let ok = "[[template]]\nname='a'\nlabel='A'\nfile='a.json'\n\ + [[template.lora]]\nid='style'\nlabel='Style'\n\ + [[template.lora]]\nid='detail'\nlabel='Detail'\ndefault=true"; + let base = "[[template]]\nname='a'\nlabel='A'\nfile='a.json'\n"; + let cfg = parse(ok).unwrap(); + let l = &cfg.templates[0].loras; + assert_eq!((l[0].id.as_str(), l[0].default), ("style", false)); + assert_eq!((l[1].id.as_str(), l[1].default), ("detail", true)); + assert_eq!((l[0].strength, l[0].min, l[0].max), (1.0, 0.0, 1.5)); + let ranged = + format!("{base}[[template.lora]]\nid='x'\nlabel='X'\nstrength=0.8\nmin=0.2\nmax=2.0"); + let r = &parse(&ranged).unwrap().templates[0].loras[0]; + assert_eq!((r.strength, r.min, r.max), (0.8, 0.2, 2.0)); + for bad in ["strength=2.0", "min=1.5", "min=-0.5", "max=0.5"] { + assert!( + parse(&format!( + "{base}[[template.lora]]\nid='x'\nlabel='X'\n{bad}" + )) + .is_err(), + "{bad}" + ); + } + + let dup = format!( + "{base}[[template.lora]]\nid='x'\nlabel='X'\n[[template.lora]]\nid='x'\nlabel='Y'" + ); + assert!(parse(&dup).is_err()); + assert!(parse(&format!("{base}[[template.lora]]\nid='a b'\nlabel='X'")).is_err()); + assert!(parse(&format!("{base}[[template.lora]]\nid='x'\nlabel=' '")).is_err()); + assert!( + parse(&format!( + "{base}[[template.lora]]\nid='x'\nlabel='X'\nweight=1" + )) + .is_err() + ); + } + + #[test] + fn lora_requests_resolve_to_a_scale() { + let l = LoraConfig { + id: "x".into(), + label: "X".into(), + default: false, + strength: 0.8, + min: 0.2, + max: 1.2, + }; + assert_eq!(l.resolve(None), 0.0); + assert_eq!(l.resolve(Some(LoraRequest::On(true))), 0.8); + assert_eq!(l.resolve(Some(LoraRequest::On(false))), 0.0); + assert_eq!(l.resolve(Some(LoraRequest::Scale(1.0))), 1.0); + // Clamped to the slider's range; 0 (or less) still means off. + assert_eq!(l.resolve(Some(LoraRequest::Scale(5.0))), 1.2); + assert_eq!(l.resolve(Some(LoraRequest::Scale(0.05))), 0.2); + assert_eq!(l.resolve(Some(LoraRequest::Scale(0.0))), 0.0); + assert_eq!(l.resolve(Some(LoraRequest::Scale(-1.0))), 0.0); + let on = LoraConfig { default: true, ..l }; + assert_eq!(on.resolve(None), 0.8); + + let parsed: std::collections::BTreeMap = + serde_json::from_str(r#"{"a": true, "b": false, "c": 0.5, "d": 1}"#).unwrap(); + assert_eq!(parsed["a"], LoraRequest::On(true)); + assert_eq!(parsed["c"], LoraRequest::Scale(0.5)); + assert_eq!(parsed["d"], LoraRequest::Scale(1.0)); + } + #[test] fn template_defaults() { let cfg = parse("[[template]]\nname='a'\nlabel='A'\nfile='a.json'").unwrap(); @@ -259,5 +421,6 @@ mod tests { assert_eq!(t.prompt_mode, PromptMode::Append); assert_eq!(t.prompt_join, " "); assert_eq!(t.output_megapixels, 1.0); + assert!(t.loras.is_empty()); } } diff --git a/src/enhance/handlers.rs b/src/enhance/handlers.rs index a057e9d..55633c0 100644 --- a/src/enhance/handlers.rs +++ b/src/enhance/handlers.rs @@ -91,6 +91,10 @@ pub struct JobView { prompt: Option, #[serde(skip_serializing_if = "Option::is_none")] reference_path: Option, + /// The scale each LoRA switch ran at (#218; 0 = off); absent when the + /// template had none. + #[serde(skip_serializing_if = "Option::is_none")] + loras: Option>, created_at: i64, updated_at: i64, /// Milliseconds ComfyUI ran the job; null until it has finished. @@ -126,6 +130,10 @@ pub fn job_view(job: &EnhanceJobRow) -> JobView { template: job.template.clone(), prompt: job.prompt.clone(), reference_path: job.reference_path.clone(), + loras: job + .loras + .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, @@ -223,6 +231,8 @@ struct TemplateView { warn_above_megapixels: f64, /// `image` (a PNG version) or `video` (an MP4 clip, e.g. Animate). output: super::template::OutputKind, + /// Optional LoRAs the client can switch on or off (#218). + loras: Vec, } #[get("/enhance/templates")] @@ -251,6 +261,7 @@ pub async fn get_templates(_: Claims, svc: web::Data) -> impl Re output_megapixels: t.cfg.output_megapixels, warn_above_megapixels: svc.config.warn_above_megapixels, output: t.output, + loras: t.lora_switches().cloned().collect(), }) .collect(); let mut body = json!({ @@ -280,6 +291,10 @@ pub struct SubmitBody { pub prompt: Option, #[serde(default)] pub reference_path: Option, + /// LoRA switches by id (#218): `true` / `false` or a scale (clamped to + /// the switch's range; <= 0 is off). A missing id uses its default. + #[serde(default)] + pub loras: std::collections::BTreeMap, #[serde(default)] pub library: Option, } @@ -325,6 +340,7 @@ pub async fn submit( &body.template, body.prompt.as_deref(), reference.as_ref(), + &body.loras, ) .await { diff --git a/src/enhance/integration_tests.rs b/src/enhance/integration_tests.rs index e1c7d73..f5670ba 100644 --- a/src/enhance/integration_tests.rs +++ b/src/enhance/integration_tests.rs @@ -26,6 +26,7 @@ use crate::state::AppState; 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 EDIT: &str = include_str!("../../tests/fixtures/enhance/edit_turbo.json"); // ---- mock ComfyUI --------------------------------------------------------- @@ -168,12 +169,27 @@ fn template(name: &str, mode: PromptMode, json: &str) -> Template { prompt_mode: mode, prompt_join: " ".into(), output_megapixels: 1.0, + loras: Vec::new(), }, serde_json::from_str(json).unwrap(), ) .unwrap() } +/// The Animate fixture with its `imageapi:lora:style` pair as a switch. +fn animate_with_style() -> Template { + let mut cfg = template("animate", PromptMode::Append, ANIMATE_LORAS).cfg; + cfg.loras = vec![crate::enhance::config::LoraConfig { + id: "style".into(), + label: "Style".into(), + default: false, + strength: 0.8, + min: 0.0, + max: 1.5, + }]; + Template::from_value(cfg, serde_json::from_str(ANIMATE_LORAS).unwrap()).unwrap() +} + async fn fx_with(url: Option, tweak: impl FnOnce(&mut EnhanceConfig)) -> Fx { let (mock, mock_url) = start_mock().await; let mut state = AppState::test_state(); @@ -205,6 +221,7 @@ async fn fx_with(url: Option, tweak: impl FnOnce(&mut EnhanceConfig)) -> vec![ template("restore", PromptMode::Append, RESTORE), template("edit", PromptMode::Replace, EDIT), + animate_with_style(), ], dao, ) @@ -333,7 +350,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(), 2); + assert_eq!(body["templates"].as_array().unwrap().len(), 3); let (st, body, _) = call( &fx, @@ -992,6 +1009,7 @@ async fn unreachable_active_jobs_become_unavailable_after_the_grace_period() { created_at: now, updated_at: now, run_ms: None, + loras: None, }; SqliteEnhanceJobDao::from_shared(fx.conn.clone()) .insert(&dao_job) @@ -1101,3 +1119,107 @@ async fn recent_prompts_answers_when_enhance_is_off() { let (st, body, _) = call(&fx, "GET", "/enhance/prompts", None).await; assert_eq!((st, body), (200, json!({"prompts": []}))); } + +#[actix_rt::test] +async fn lora_switches_are_listed_sent_and_stored_on_the_job() { + let fx = fx().await; + fx.put(PHOTO, b"photo"); + + let (_, body, _) = call(&fx, "GET", "/enhance/templates", None).await; + let templates = body["templates"].as_array().unwrap(); + let animate = templates.iter().find(|t| t["name"] == "animate").unwrap(); + assert_eq!( + animate["loras"], + json!([{"id": "style", "label": "Style", "default": false, "strength": 0.8, "min": 0.0, "max": 1.5}]) + ); + let restore = templates.iter().find(|t| t["name"] == "restore").unwrap(); + assert_eq!(restore["loras"], json!([])); + + // On: the configured strength, and the job remembers it. + let (st, job, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": PHOTO, "template": "animate", "loras": {"style": true}})), + ) + .await; + assert_eq!(st, 200, "{job}"); + assert_eq!(job["loras"], json!({"style": 0.8})); + let (_, again, _) = call( + &fx, + "GET", + &format!("/enhance/jobs/{}", job["id"].as_str().unwrap()), + None, + ) + .await; + assert_eq!(again["loras"], json!({"style": 0.8})); + + // Omitted: the default (off) zeroes both nodes. + let (st, off, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": PHOTO, "template": "animate"})), + ) + .await; + assert_eq!(st, 200, "{off}"); + assert_eq!(off["loras"], json!({"style": 0.0})); + + // A slider value scales the exported strength. + let (st, half, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": PHOTO, "template": "animate", "loras": {"style": 0.5}})), + ) + .await; + assert_eq!(st, 200, "{half}"); + assert_eq!(half["loras"], json!({"style": 0.5})); + { + let sent = fx.mock.submitted.lock().unwrap(); + assert_eq!( + sent[0]["prompt"]["50"]["inputs"]["strength_model"], + json!(0.8) + ); + assert_eq!( + sent[2]["prompt"]["51"]["inputs"]["strength_model"], + json!(0.5) + ); + assert_eq!( + sent[1]["prompt"]["50"]["inputs"]["strength_model"], + json!(0.0) + ); + assert_eq!( + sent[1]["prompt"]["51"]["inputs"]["strength_model"], + json!(0.0) + ); + } + + // Unknown switch, or one on a template without switches: 400. + let (st, err, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": PHOTO, "template": "animate", "loras": {"nope": true}})), + ) + .await; + assert_eq!(st, 400, "{err}"); + let (st, _, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": PHOTO, "template": "restore", "loras": {"style": true}})), + ) + .await; + assert_eq!(st, 400); + + // A template without switches stores none. + let (_, plain, _) = call( + &fx, + "POST", + "/enhance", + Some(json!({"path": PHOTO, "template": "restore"})), + ) + .await; + assert!(plain.get("loras").is_none(), "{plain}"); +} diff --git a/src/enhance/service.rs b/src/enhance/service.rs index 359b707..a189d89 100644 --- a/src/enhance/service.rs +++ b/src/enhance/service.rs @@ -292,6 +292,7 @@ impl EnhanceService { template_name: &str, prompt: Option<&str>, reference: Option<&SourceFile>, + loras: &std::collections::BTreeMap, ) -> Result { let Some(client) = &self.client else { return Err(EnhanceError::Disabled); @@ -319,7 +320,13 @@ impl EnhanceService { None => None, }; let filled = template - .fill(&source, prompt, reference_path.as_deref(), &mut random_seed) + .fill( + &source, + prompt, + reference_path.as_deref(), + loras, + &mut random_seed, + ) .map_err(|e| EnhanceError::BadRequest(e.0))?; if !self.probe_now().await { @@ -365,6 +372,8 @@ impl EnhanceService { created_at: now, updated_at: now, run_ms: None, + loras: (!filled.loras.is_empty()) + .then(|| serde_json::to_string(&filled.loras).unwrap_or_default()), }; self.dao .lock() diff --git a/src/enhance/template.rs b/src/enhance/template.rs index bc04bfa..acf1549 100644 --- a/src/enhance/template.rs +++ b/src/enhance/template.rs @@ -4,17 +4,25 @@ //! Nodes are found by `_meta.title`, never by id, so re-exporting after //! editing the graph keeps working. +use std::collections::BTreeMap; use std::path::Path; use anyhow::{Context as _, Result, bail}; use serde_json::{Map, Value, json}; -use super::config::{PromptMode, TemplateConfig}; +use super::config::{LoraConfig, LoraRequest, PromptMode, TemplateConfig}; pub const TITLE_SOURCE: &str = "imageapi:source"; pub const TITLE_SAVE: &str = "imageapi:save"; 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:"; + +/// LoRA loader inputs a switch scales (0 = off). `strength_model` is +/// required (LoraLoaderModelOnly); `strength_clip` is scaled when present +/// (LoraLoader). +const LORA_STRENGTH_INPUTS: &[&str] = &["strength_model", "strength_clip"]; /// Prompt-node inputs, in the order they are tried. const PROMPT_INPUTS: &[&str] = &["text", "prompt", "value"]; @@ -51,12 +59,22 @@ pub struct Template { /// `(node id, input name)` of the prompt text. prompt: Option<(String, &'static str)>, reference_id: Option, + /// Declared LoRA switches that have nodes, in `enhance.toml` order. + loras: Vec, /// Image or video, from the save node's class. pub output: OutputKind, /// Startup warnings (stray output nodes). pub warnings: Vec, } +#[derive(Debug, Clone)] +struct LoraSwitch { + cfg: LoraConfig, + /// Every node titled `imageapi:lora:`, with its exported strengths + /// (`(input, value)`), which the switch's scale multiplies. + nodes: Vec<(String, Vec<(&'static str, 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)] @@ -74,6 +92,9 @@ impl std::fmt::Display for FillError { pub struct Filled { pub prompt: Value, pub dropped_reference: bool, + /// Every switch's resolved scale (0 = off; request, else its default); + /// stored on the job so a retry repeats it. Empty without switches. + pub loras: BTreeMap, } impl Template { @@ -150,6 +171,62 @@ impl Template { } let mut warnings = Vec::new(); + + // LoRA switches: declared ids collect their nodes; a declared id with + // none is skipped (not offered), an undeclared node keeps its strength. + let mut lora_nodes: BTreeMap<&str, Vec> = BTreeMap::new(); + for (id, node) in &workflow { + if let Some(lora) = title_of(node).and_then(|t| t.strip_prefix(TITLE_LORA_PREFIX)) { + if !inputs_of(&workflow, id) + .get("strength_model") + .is_some_and(Value::is_number) + { + bail!( + "node {id} titled {:?} has no numeric strength_model input \ + (use a LoraLoaderModelOnly or LoraLoader)", + format!("{TITLE_LORA_PREFIX}{lora}") + ); + } + lora_nodes.entry(lora).or_default().push(id.clone()); + } + } + let mut loras = Vec::new(); + for l in &cfg.loras { + match lora_nodes.remove(l.id.as_str()) { + Some(mut ids) => { + ids.sort(); + let nodes = ids + .into_iter() + .map(|id| { + let inputs = inputs_of(&workflow, &id); + let exported = LORA_STRENGTH_INPUTS + .iter() + .filter_map(|k| { + inputs.get(*k).and_then(Value::as_f64).map(|v| (*k, v)) + }) + .collect(); + (id, exported) + }) + .collect(); + loras.push(LoraSwitch { + cfg: l.clone(), + nodes, + }); + } + None => warnings.push(format!( + "lora {:?} has no node titled {:?}, so it isn't offered", + l.id, + format!("{TITLE_LORA_PREFIX}{}", l.id) + )), + } + } + for lora in lora_nodes.keys() { + warnings.push(format!( + "node titled {:?} has no [[template.lora]] entry, so it stays at its exported strength", + format!("{TITLE_LORA_PREFIX}{lora}") + )); + } + for (id, node) in &workflow { let class = node["class_type"].as_str().unwrap_or_default(); if STRAY_OUTPUT_CLASSES.contains(&class) { @@ -170,6 +247,7 @@ impl Template { save_id, prompt, reference_id, + loras, output, warnings, }) @@ -187,6 +265,12 @@ impl Template { self.cfg.prompt_mode == PromptMode::Replace } + /// The LoRA switches clients can offer (declared and present in the + /// workflow), in `enhance.toml` order. + pub fn lora_switches(&self) -> impl Iterator { + self.loras.iter().map(|l| &l.cfg) + } + /// Fill the workflow for one photo. /// /// - `source` / `reference` are paths as ComfyUI sees them. @@ -195,13 +279,25 @@ impl Template { /// - Every `seed` / `noise_seed` / `*.seed` input gets a fresh value /// (ComfyUI's "randomize after generate" only exists in its page), or a /// retry would return the cached result. + /// - `loras` sets each switch (missing = its default; see + /// [`LoraConfig::resolve`]): the resolved scale multiplies the node's + /// exported strengths, 0 being off, so the graph is never rewired. pub fn fill( &self, source: &str, prompt: Option<&str>, reference: Option<&str>, + loras: &BTreeMap, next_seed: &mut dyn FnMut() -> u64, ) -> Result { + if let Some(unknown) = loras + .keys() + .find(|id| !self.loras.iter().any(|l| &l.cfg.id == *id)) + { + return Err(FillError(format!( + "this template has no LoRA switch {unknown:?}" + ))); + } let user_text = prompt.map(str::trim).filter(|t| !t.is_empty()); match (&self.prompt, user_text) { (None, Some(_)) => { @@ -255,6 +351,17 @@ impl Template { } } + let mut resolved = BTreeMap::new(); + for switch in &self.loras { + let scale = switch.cfg.resolve(loras.get(&switch.cfg.id).copied()); + resolved.insert(switch.cfg.id.clone(), scale); + for (id, exported) in &switch.nodes { + 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; @@ -268,6 +375,7 @@ impl Template { Ok(Filled { prompt: Value::Object(wf), dropped_reference, + loras: resolved, }) } } @@ -324,6 +432,38 @@ mod tests { pub const RESTORE: &str = include_str!("../../tests/fixtures/enhance/restore_turbo.json"); pub const EDIT: &str = include_str!("../../tests/fixtures/enhance/edit_turbo.json"); pub const DEV: &str = include_str!("../../tests/fixtures/enhance/restore_dev.json"); + /// 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"); + + fn lora(id: &str, default: bool) -> LoraConfig { + LoraConfig { + id: id.into(), + label: id.into(), + default, + strength: 1.0, + min: 0.0, + max: 1.5, + } + } + + fn load_with_loras(json: &str, loras: Vec) -> Result