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:
1 parent
e317bfee98
commit
31da21e061
12 files changed
+910
-33
No files matched your search
@@ -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}");
|
||||
}
|
||||
Reference in new issue
Block a user