Apps/ImageApi
1
0 Fork 0
Code Issues 8 Pull Requests 1 Releases Wiki Activity

feat(enhance): LoRA switches per template (#218)

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/enhance/template.rs
+322 -19
View File
@@ -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:<id>`: 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<String>,
/// Declared LoRA switches that have nodes, in `enhance.toml` order.
loras: Vec<LoraSwitch>,
/// Image or video, from the save node's class.
pub output: OutputKind,
/// Startup warnings (stray output nodes).
pub warnings: Vec<String>,
}
#[derive(Debug, Clone)]
struct LoraSwitch {
cfg: LoraConfig,
/// Every node titled `imageapi:lora:<id>`.
nodes: Vec<String>,
}
/// 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<String, bool>,
}
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<String>> = 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<Item = &LoraConfig> {
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<String, bool>,
next_seed: &mut dyn FnMut() -> u64,
) -> Result<Filled, FillError> {
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<LoraConfig>) -> Result<Template> {
let mut c = cfg("animate", PromptMode::Append);
c.loras = loras;
Template::from_value(c, serde_json::from_str(json).unwrap())
}
fn switches(pairs: &[(&str, bool)]) -> BTreeMap<String, bool> {
pairs.iter().map(|(k, v)| (k.to_string(), *v)).collect()
}
fn cfg(name: &str, mode: PromptMode) -> TemplateConfig {
TemplateConfig {
@@ -334,6 +455,7 @@ mod tests {
prompt_mode: mode,
prompt_join: " ".into(),
output_megapixels: 1.0,
loras: Vec::new(),
}
}
@@ -401,7 +523,13 @@ mod tests {
v["13"]["inputs"] = json!({"prompt": "base"});
let t = Template::from_value(cfg("t", PromptMode::Append), v).unwrap();
let f = t
.fill("Z:\\a.jpg", Some("more"), None, &mut seeds())
.fill(
"Z:\\a.jpg",
Some("more"),
None,
&BTreeMap::new(),
&mut seeds(),
)
.unwrap();
assert_eq!(input(&f.prompt, "13", "prompt"), "base more");
}
@@ -418,7 +546,9 @@ mod tests {
#[test]
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, &mut seeds()).unwrap();
let f = t
.fill("Z:\\p\\a.jpg", None, None, &BTreeMap::new(), &mut seeds())
.unwrap();
assert_eq!(input(&f.prompt, "37", "path"), "Z:\\p\\a.jpg");
assert_eq!(input(&f.prompt, "36", "require_approval"), true);
// KSampler.seed and TextGenerate's sampling_mode.seed, both fresh.
@@ -430,8 +560,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, &mut rng).unwrap();
let y = t.fill("p", None, None, &mut rng).unwrap();
let x = t.fill("p", None, None, &BTreeMap::new(), &mut rng).unwrap();
let y = t.fill("p", None, None, &BTreeMap::new(), &mut rng).unwrap();
assert_ne!(input(&x.prompt, "5", "seed"), input(&y.prompt, "5", "seed"));
}
@@ -457,7 +587,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, &mut random_seed).unwrap();
let f = t
.fill("p", None, None, &BTreeMap::new(), &mut random_seed)
.unwrap();
assert!(input(&f.prompt, "5", "seed").as_u64().unwrap() <= u64::from(u32::MAX));
}
}
@@ -466,7 +598,9 @@ 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, &mut seeds()).unwrap();
let f = t
.fill("p", None, None, &BTreeMap::new(), &mut seeds())
.unwrap();
assert_eq!(input(&f.prompt, "36", "require_approval"), true);
}
@@ -478,14 +612,26 @@ mod tests {
let mut v: Value = serde_json::from_str(RESTORE).unwrap();
v["13"]["inputs"]["value"] = json!(base);
let t = Template::from_value(c, v).unwrap();
let f = t.fill("p", Some(" warmer "), None, &mut seeds()).unwrap();
let f = t
.fill(
"p",
Some(" warmer "),
None,
&BTreeMap::new(),
&mut seeds(),
)
.unwrap();
assert_eq!(
input(&f.prompt, "13", "value"),
"Increase the clarity\n\nAdditional instructions: warmer"
);
let f = t.fill("p", Some(" "), None, &mut seeds()).unwrap();
let f = t
.fill("p", Some(" "), None, &BTreeMap::new(), &mut seeds())
.unwrap();
assert_eq!(input(&f.prompt, "13", "value"), base);
let f = t.fill("p", None, None, &mut seeds()).unwrap();
let f = t
.fill("p", None, None, &BTreeMap::new(), &mut seeds())
.unwrap();
assert_eq!(input(&f.prompt, "13", "value"), base);
}
@@ -493,7 +639,13 @@ mod tests {
fn default_join_is_a_single_space_sentence() {
let t = load(RESTORE, PromptMode::Append).unwrap();
let f = t
.fill("p", Some("Make it warmer."), None, &mut seeds())
.fill(
"p",
Some("Make it warmer."),
None,
&BTreeMap::new(),
&mut seeds(),
)
.unwrap();
let text = input(&f.prompt, "13", "value").as_str().unwrap();
assert!(text.ends_with("Remove any phone interface overlays. Make it warmer."));
@@ -502,11 +654,22 @@ 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, &mut seeds()).unwrap_err();
let err = t
.fill("p", None, None, &BTreeMap::new(), &mut seeds())
.unwrap_err();
assert!(err.0.contains("instruction"));
assert!(t.fill("p", Some(" "), None, &mut seeds()).is_err());
assert!(
t.fill("p", Some(" "), None, &BTreeMap::new(), &mut seeds())
.is_err()
);
let f = t
.fill("p", Some("make it black and white"), None, &mut seeds())
.fill(
"p",
Some("make it black and white"),
None,
&BTreeMap::new(),
&mut seeds(),
)
.unwrap();
assert_eq!(input(&f.prompt, "13", "value"), "make it black and white");
}
@@ -518,18 +681,34 @@ 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, &mut seeds()).is_err());
assert!(t.fill("p", None, None, &mut seeds()).is_ok());
assert!(
t.fill("p", Some("x"), None, &BTreeMap::new(), &mut seeds())
.is_err()
);
assert!(
t.fill("p", None, None, &BTreeMap::new(), &mut seeds())
.is_ok()
);
let restore = load(RESTORE, PromptMode::Append).unwrap();
assert!(restore.fill("p", None, Some("r"), &mut seeds()).is_err());
assert!(
restore
.fill("p", None, Some("r"), &BTreeMap::new(), &mut seeds())
.is_err()
);
}
#[test]
fn reference_given_sets_its_path_and_keeps_links() {
let t = load(EDIT, PromptMode::Replace).unwrap();
let f = t
.fill("src", Some("x"), Some("Z:\\ref.jpg"), &mut seeds())
.fill(
"src",
Some("x"),
Some("Z:\\ref.jpg"),
&BTreeMap::new(),
&mut seeds(),
)
.unwrap();
assert!(!f.dropped_reference);
assert_eq!(input(&f.prompt, "38", "path"), "Z:\\ref.jpg");
@@ -539,7 +718,9 @@ mod tests {
#[test]
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, &mut seeds()).unwrap();
let f = t
.fill("src", Some("x"), None, &BTreeMap::new(), &mut seeds())
.unwrap();
assert!(f.dropped_reference);
let wf = f.prompt.as_object().unwrap();
assert!(!wf.contains_key("38"));
@@ -554,4 +735,126 @@ mod tests {
// Links to other nodes survive.
assert_eq!(input(&f.prompt, "36", "source"), &json!(["37", 1]));
}
#[test]
fn lora_switch_drives_every_node_with_its_id() {
let t = load_with_loras(ANIMATE_LORAS, vec![lora("style", false)]).unwrap();
assert_eq!(
t.lora_switches().map(|l| l.id.as_str()).collect::<Vec<_>>(),
["style"]
);
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, &BTreeMap::new(), &mut seeds())
.unwrap();
assert_eq!(off.loras, switches(&[("style", false)]));
for id in ["50", "51"] {
assert_eq!(input(&off.prompt, id, "strength_model"), &json!(0.0));
}
// The graph isn't rewired, and the fixed lightx2v LoRAs are untouched.
assert_eq!(input(&off.prompt, "14", "model"), &json!(["50", 0]));
assert_eq!(input(&off.prompt, "12", "strength_model"), &json!(1.0));
// On: the exported strength stays.
let on = t
.fill("p", None, None, &switches(&[("style", true)]), &mut seeds())
.unwrap();
assert_eq!(on.loras, switches(&[("style", true)]));
for id in ["50", "51"] {
assert_eq!(input(&on.prompt, id, "strength_model"), &json!(1.0));
}
}
#[test]
fn lora_default_on_and_unknown_ids_refused() {
let t = load_with_loras(ANIMATE_LORAS, vec![lora("style", true)]).unwrap();
let f = t
.fill("p", None, None, &BTreeMap::new(), &mut seeds())
.unwrap();
assert_eq!(f.loras, switches(&[("style", true)]));
assert_eq!(input(&f.prompt, "50", "strength_model"), &json!(1.0));
let err = t
.fill("p", None, None, &switches(&[("nope", true)]), &mut seeds())
.unwrap_err();
assert!(err.0.contains("nope"), "{err}");
}
#[test]
fn lora_clip_strength_is_zeroed_too() {
let mut v: Value = serde_json::from_str(ANIMATE_LORAS).unwrap();
v["50"]["class_type"] = json!("LoraLoader");
v["50"]["inputs"]["strength_clip"] = json!(0.8);
let t = Template::from_value(
{
let mut c = cfg("animate", PromptMode::Append);
c.loras = vec![lora("style", false)];
c
},
v,
)
.unwrap();
let f = t
.fill("p", None, None, &BTreeMap::new(), &mut seeds())
.unwrap();
assert_eq!(input(&f.prompt, "50", "strength_clip"), &json!(0.0));
}
#[test]
fn undeclared_or_missing_lora_switches_are_warnings() {
// Nodes but no [[template.lora]]: fixed at their strength, not offered.
let t = load_with_loras(ANIMATE_LORAS, Vec::new()).unwrap();
assert_eq!(t.lora_switches().count(), 0);
assert!(
t.warnings.iter().any(|w| w.contains("imageapi:lora:style")),
"{:?}",
t.warnings
);
let f = t
.fill("p", None, None, &BTreeMap::new(), &mut seeds())
.unwrap();
assert!(f.loras.is_empty());
assert_eq!(input(&f.prompt, "50", "strength_model"), &json!(1.0));
assert!(
t.fill(
"p",
None,
None,
&switches(&[("style", false)]),
&mut seeds()
)
.is_err()
);
// Declared but no node: skipped with a warning, the rest load.
let t = load_with_loras(
ANIMATE_LORAS,
vec![lora("style", false), lora("detail", true)],
)
.unwrap();
assert_eq!(
t.lora_switches().map(|l| l.id.as_str()).collect::<Vec<_>>(),
["style"]
);
assert!(
t.warnings.iter().any(|w| w.contains("\"detail\"")),
"{:?}",
t.warnings
);
}
#[test]
fn a_lora_node_without_strength_model_is_rejected() {
let bad = ANIMATE_LORAS.replace(
"\"strength_model\": 1.0,\n \"model\": [\n \"12\"",
"\"weight\": 1.0,\n \"model\": [\n \"12\"",
);
assert_ne!(bad, ANIMATE_LORAS, "fixture layout changed");
let err = load_with_loras(&bad, vec![lora("style", false)])
.unwrap_err()
.to_string();
assert!(err.contains("strength_model"), "{err}");
}
}