From 31da21e0619bfa8e8c5be2d106db2693e036af27 Mon Sep 17 00:00:00 2001 From: Cameron Cordes Date: Tue, 3 Nov 2026 13:53:53 -0500 Subject: [PATCH 1/2] feat(enhance): LoRA switches per template (#218) A template can offer optional LoRAs the app turns on or off. - enhance.toml: [[template.lora]] { id, label, default } per template. Each id drives every workflow node titled `imageapi:lora:`, so Wan's high- and low-noise pair share one switch. - POST /enhance takes `loras: {id: bool}`; a missing id uses its default, an id the template doesn't offer is a 400. Off sets strength_model (and strength_clip when present) to 0, so the graph is never rewired. - The resolved switches are stored on the job (new enhance_jobs.loras JSON column) and returned in the job view, so a retry or adjust can repeat them; GET /enhance/templates lists each template's switches. - A lora-titled node without a declaration keeps its exported strength, and a declaration without nodes isn't offered; both are startup warnings. A lora-titled node without a numeric strength_model fails the template. Co-Authored-By: Claude Opus 5.5 --- CLAUDE.md | 18 +- enhance.example.toml | 18 + .../down.sql | 1 + .../up.sql | 2 + src/database/enhance_dao.rs | 10 +- src/database/schema.rs | 1 + src/enhance/config.rs | 72 +++- src/enhance/handlers.rs | 15 + src/enhance/integration_tests.rs | 106 +++++- src/enhance/service.rs | 11 +- src/enhance/template.rs | 341 ++++++++++++++++- tests/fixtures/enhance/animate_loras.json | 348 ++++++++++++++++++ 12 files changed, 910 insertions(+), 33 deletions(-) create mode 100644 migrations/2026-11-03-100000_add_enhance_loras/down.sql create mode 100644 migrations/2026-11-03-100000_add_enhance_loras/up.sql create mode 100644 tests/fixtures/enhance/animate_loras.json diff --git a/CLAUDE.md b/CLAUDE.md index 1a19348..eaf17bf 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -459,14 +459,19 @@ 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`)). - **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; off = `strength_model`/`strength_clip` set to 0, 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 +517,10 @@ 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"}]}]} +POST /enhance {"path","template","prompt"?,"reference_path"?,"loras"?: {id: bool},"library"?} -> job + (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 +537,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: bool} (resolved; 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..5d7a787 100644 --- a/enhance.example.toml +++ b/enhance.example.toml @@ -63,3 +63,21 @@ 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 (#218). +# Each [[template.lora]] matches the workflow's LoRA loader nodes titled +# "imageapi:lora:" (Wan's high- and low-noise pair share one id). Off sets +# their strength_model (and strength_clip) to 0; the graph isn't rewired. +# 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 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..0265c99 100644 --- a/src/enhance/config.rs +++ b/src/enhance/config.rs @@ -100,6 +100,30 @@ 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, +} + +/// 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 +156,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 +171,22 @@ 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); + } + } } if let Some(d) = &self.default_template && !self.templates.iter().any(|t| &t.name == d) @@ -252,6 +288,31 @@ 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 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)); + + let base = "[[template]]\nname='a'\nlabel='A'\nfile='a.json'\n"; + 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'\nstrength=1" + )) + .is_err() + ); + } + #[test] fn template_defaults() { let cfg = parse("[[template]]\nname='a'\nlabel='A'\nfile='a.json'").unwrap(); @@ -259,5 +320,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..591f142 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 LoRA switches this job ran with (#218); 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,9 @@ pub struct SubmitBody { pub prompt: Option, #[serde(default)] pub reference_path: Option, + /// LoRA switches by id (#218); a missing id uses its default. + #[serde(default)] + pub loras: std::collections::BTreeMap, #[serde(default)] pub library: Option, } @@ -325,6 +339,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..74f8f5f 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,24 @@ 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, + }]; + 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 +218,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 +347,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 +1006,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 +1116,92 @@ 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}]) + ); + let restore = templates.iter().find(|t| t["name"] == "restore").unwrap(); + assert_eq!(restore["loras"], json!([])); + + // On: the pair keeps its strength, and the job remembers the choice. + 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": true})); + let (_, again, _) = call( + &fx, + "GET", + &format!("/enhance/jobs/{}", job["id"].as_str().unwrap()), + None, + ) + .await; + assert_eq!(again["loras"], json!({"style": true})); + + // 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": false})); + { + let sent = fx.mock.submitted.lock().unwrap(); + assert_eq!( + sent[0]["prompt"]["50"]["inputs"]["strength_model"], + json!(1.0) + ); + 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..a69d358 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..6fc1e7a 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, 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 switched-off LoRA sets to 0. `strength_model` is +/// required (LoraLoaderModelOnly); `strength_clip` is zeroed 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,21 @@ 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:`. + nodes: Vec, +} + /// 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 +91,9 @@ impl std::fmt::Display for FillError { pub struct Filled { pub prompt: Value, pub dropped_reference: bool, + /// Every switch's resolved state (request, else its default); stored on + /// the job so a retry repeats it. Empty for templates without switches. + pub loras: BTreeMap, } impl Template { @@ -150,6 +170,49 @@ 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 nodes) => { + nodes.sort(); + 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 +233,7 @@ impl Template { save_id, prompt, reference_id, + loras, output, warnings, }) @@ -187,6 +251,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 +265,24 @@ 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` turns switches on or off (missing = the switch's default); + /// an off LoRA's strengths are set to 0, 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 +336,24 @@ impl Template { } } + let mut resolved = BTreeMap::new(); + for switch in &self.loras { + let on = loras + .get(&switch.cfg.id) + .copied() + .unwrap_or(switch.cfg.default); + resolved.insert(switch.cfg.id.clone(), on); + if !on { + for id in &switch.nodes { + for input in LORA_STRENGTH_INPUTS { + if inputs_of(&wf, id).get(*input).is_some_and(Value::is_number) { + set_input(&mut wf, id, input, json!(0.0)); + } + } + } + } + } + for node in wf.values_mut() { let Some(inputs) = node.get_mut("inputs").and_then(Value::as_object_mut) else { continue; @@ -268,6 +367,7 @@ impl Template { Ok(Filled { prompt: Value::Object(wf), dropped_reference, + loras: resolved, }) } } @@ -324,6 +424,27 @@ 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, + } + } + + fn load_with_loras(json: &str, loras: Vec) -> Result