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:<id>`, 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 <noreply@anthropic.com>
This commit is contained in:
Cameron CordesandClaude Opus 5.5 committed 2026-11-03 13:53:53 -05:00
1 parent e317bfee98
commit 31da21e061
12 files changed
+910 -33

No files matched your search

+105 -1
View File
@@ -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<String>, 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<String>, 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}");
}