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

feat(enhance): LoRA strength scale alongside the on/off switch (#218)

Browse Source 
So the app builds against the final shape: each [[template.lora]] gains
`strength` (the scale when on, default 1.0) and `min` / `max` (the
slider's range, default 0.0-1.5).

- POST /enhance `loras` takes true / false or a scale: true = strength,
  false or <= 0 = off, a scale is clamped to min..max.
- The scale multiplies each node's exported strength_model /
  strength_clip, so 1.0 = as exported and a high/low pair keeps its
  ratio.
- Jobs and the job view store the resolved scale per id (0 = off);
  GET /enhance/templates lists strength, min and max.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Cameron CordesandClaude Opus 5.5 committed 2026-11-03 14:01:04 -05:00
1 parent 31da21e061
commit cc96873f86
7 files changed
+298 -110
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/enhance/template.rs
+143 -83
View File
@@ -10,7 +10,7 @@ use std::path::Path;
use anyhow::{Context as _, Result, bail};
use serde_json::{Map, Value, json};
use super::config::{LoraConfig, PromptMode, TemplateConfig};
use super::config::{LoraConfig, LoraRequest, PromptMode, TemplateConfig};
pub const TITLE_SOURCE: &str = "imageapi:source";
pub const TITLE_SAVE: &str = "imageapi:save";
@@ -19,8 +19,8 @@ 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
/// 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"];
@@ -70,8 +70,9 @@ pub struct Template {
#[derive(Debug, Clone)]
struct LoraSwitch {
cfg: LoraConfig,
/// Every node titled `imageapi:lora:<id>`.
nodes: Vec<String>,
/// Every node titled `imageapi:lora:<id>`, 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
@@ -91,9 +92,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>,
/// 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<String, f64>,
}
impl Template {
@@ -192,8 +193,21 @@ impl Template {
let mut loras = Vec::new();
for l in &cfg.loras {
match lora_nodes.remove(l.id.as_str()) {
Some(mut nodes) => {
nodes.sort();
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,
@@ -265,14 +279,15 @@ 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.
/// - `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<String, bool>,
loras: &BTreeMap<String, LoraRequest>,
next_seed: &mut dyn FnMut() -> u64,
) -> Result<Filled, FillError> {
if let Some(unknown) = loras
@@ -338,18 +353,11 @@ 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));
}
}
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));
}
}
}
@@ -433,6 +441,9 @@ mod tests {
id: id.into(),
label: id.into(),
default,
strength: 1.0,
min: 0.0,
max: 1.5,
}
}
@@ -442,10 +453,18 @@ mod tests {
Template::from_value(c, serde_json::from_str(json).unwrap())
}
fn switches(pairs: &[(&str, bool)]) -> BTreeMap<String, bool> {
fn reqs(pairs: &[(&str, LoraRequest)]) -> BTreeMap<String, LoraRequest> {
pairs.iter().map(|(k, v)| (k.to_string(), *v)).collect()
}
fn scales(pairs: &[(&str, f64)]) -> BTreeMap<String, f64> {
pairs.iter().map(|(k, v)| (k.to_string(), *v)).collect()
}
fn none() -> BTreeMap<String, LoraRequest> {
BTreeMap::new()
}
fn cfg(name: &str, mode: PromptMode) -> TemplateConfig {
TemplateConfig {
name: name.into(),
@@ -547,7 +566,7 @@ mod tests {
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, &BTreeMap::new(), &mut seeds())
.fill("Z:\\p\\a.jpg", None, None, &none(), &mut seeds())
.unwrap();
assert_eq!(input(&f.prompt, "37", "path"), "Z:\\p\\a.jpg");
assert_eq!(input(&f.prompt, "36", "require_approval"), true);
@@ -560,8 +579,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, &BTreeMap::new(), &mut rng).unwrap();
let y = t.fill("p", None, None, &BTreeMap::new(), &mut rng).unwrap();
let x = t.fill("p", None, None, &none(), &mut rng).unwrap();
let y = t.fill("p", None, None, &none(), &mut rng).unwrap();
assert_ne!(input(&x.prompt, "5", "seed"), input(&y.prompt, "5", "seed"));
}
@@ -587,9 +606,7 @@ 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, &BTreeMap::new(), &mut random_seed)
.unwrap();
let f = t.fill("p", None, None, &none(), &mut random_seed).unwrap();
assert!(input(&f.prompt, "5", "seed").as_u64().unwrap() <= u64::from(u32::MAX));
}
}
@@ -598,9 +615,7 @@ 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, &BTreeMap::new(), &mut seeds())
.unwrap();
let f = t.fill("p", None, None, &none(), &mut seeds()).unwrap();
assert_eq!(input(&f.prompt, "36", "require_approval"), true);
}
@@ -626,12 +641,10 @@ mod tests {
"Increase the clarity\n\nAdditional instructions: warmer"
);
let f = t
.fill("p", Some(" "), None, &BTreeMap::new(), &mut seeds())
.fill("p", Some(" "), None, &none(), &mut seeds())
.unwrap();
assert_eq!(input(&f.prompt, "13", "value"), base);
let f = t
.fill("p", None, None, &BTreeMap::new(), &mut seeds())
.unwrap();
let f = t.fill("p", None, None, &none(), &mut seeds()).unwrap();
assert_eq!(input(&f.prompt, "13", "value"), base);
}
@@ -654,12 +667,10 @@ 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, &BTreeMap::new(), &mut seeds())
.unwrap_err();
let err = t.fill("p", None, None, &none(), &mut seeds()).unwrap_err();
assert!(err.0.contains("instruction"));
assert!(
t.fill("p", Some(" "), None, &BTreeMap::new(), &mut seeds())
t.fill("p", Some(" "), None, &none(), &mut seeds())
.is_err()
);
let f = t
@@ -681,19 +692,13 @@ 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, &BTreeMap::new(), &mut seeds())
.is_err()
);
assert!(
t.fill("p", None, None, &BTreeMap::new(), &mut seeds())
.is_ok()
);
assert!(t.fill("p", Some("x"), None, &none(), &mut seeds()).is_err());
assert!(t.fill("p", None, None, &none(), &mut seeds()).is_ok());
let restore = load(RESTORE, PromptMode::Append).unwrap();
assert!(
restore
.fill("p", None, Some("r"), &BTreeMap::new(), &mut seeds())
.fill("p", None, Some("r"), &none(), &mut seeds())
.is_err()
);
}
@@ -719,7 +724,7 @@ mod tests {
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, &BTreeMap::new(), &mut seeds())
.fill("src", Some("x"), None, &none(), &mut seeds())
.unwrap();
assert!(f.dropped_reference);
let wf = f.prompt.as_object().unwrap();
@@ -746,10 +751,8 @@ mod tests {
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)]));
let off = t.fill("p", None, None, &none(), &mut seeds()).unwrap();
assert_eq!(off.loras, scales(&[("style", 0.0)]));
for id in ["50", "51"] {
assert_eq!(input(&off.prompt, id, "strength_model"), &json!(0.0));
}
@@ -757,49 +760,108 @@ mod tests {
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.
// On: the configured strength (1.0 = as exported).
let on = t
.fill("p", None, None, &switches(&[("style", true)]), &mut seeds())
.fill(
"p",
None,
None,
&reqs(&[("style", LoraRequest::On(true))]),
&mut seeds(),
)
.unwrap();
assert_eq!(on.loras, switches(&[("style", true)]));
assert_eq!(on.loras, scales(&[("style", 1.0)]));
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();
fn a_scale_multiplies_each_nodes_exported_strength() {
// A pair exported at different strengths keeps its ratio.
let mut v: Value = serde_json::from_str(ANIMATE_LORAS).unwrap();
v["51"]["inputs"]["strength_model"] = json!(0.8);
let mut c = cfg("animate", PromptMode::Append);
c.loras = vec![lora("style", false)];
let t = Template::from_value(c, v).unwrap();
let f = t
.fill("p", None, None, &BTreeMap::new(), &mut seeds())
.fill(
"p",
None,
None,
&reqs(&[("style", LoraRequest::Scale(0.5))]),
&mut seeds(),
)
.unwrap();
assert_eq!(f.loras, switches(&[("style", true)]));
assert_eq!(input(&f.prompt, "50", "strength_model"), &json!(1.0));
assert_eq!(f.loras, scales(&[("style", 0.5)]));
assert_eq!(input(&f.prompt, "50", "strength_model"), &json!(0.5));
assert_eq!(input(&f.prompt, "51", "strength_model"), &json!(0.4));
// Clamped to the switch's max; 0 is off.
let f = t
.fill(
"p",
None,
None,
&reqs(&[("style", LoraRequest::Scale(9.0))]),
&mut seeds(),
)
.unwrap();
assert_eq!(f.loras, scales(&[("style", 1.5)]));
let f = t
.fill(
"p",
None,
None,
&reqs(&[("style", LoraRequest::Scale(0.0))]),
&mut seeds(),
)
.unwrap();
assert_eq!(input(&f.prompt, "51", "strength_model"), &json!(0.0));
}
#[test]
fn lora_default_on_and_unknown_ids_refused() {
let mut l = lora("style", true);
l.strength = 0.7;
let t = load_with_loras(ANIMATE_LORAS, vec![l]).unwrap();
let f = t.fill("p", None, None, &none(), &mut seeds()).unwrap();
assert_eq!(f.loras, scales(&[("style", 0.7)]));
assert_eq!(input(&f.prompt, "50", "strength_model"), &json!(0.7));
let err = t
.fill("p", None, None, &switches(&[("nope", true)]), &mut seeds())
.fill(
"p",
None,
None,
&reqs(&[("nope", LoraRequest::On(true))]),
&mut seeds(),
)
.unwrap_err();
assert!(err.0.contains("nope"), "{err}");
}
#[test]
fn lora_clip_strength_is_zeroed_too() {
fn lora_clip_strength_is_scaled_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())
let mut c = cfg("animate", PromptMode::Append);
c.loras = vec![lora("style", false)];
let t = Template::from_value(c, v).unwrap();
let off = t.fill("p", None, None, &none(), &mut seeds()).unwrap();
assert_eq!(input(&off.prompt, "50", "strength_clip"), &json!(0.0));
let half = t
.fill(
"p",
None,
None,
&reqs(&[("style", LoraRequest::Scale(0.5))]),
&mut seeds(),
)
.unwrap();
assert_eq!(input(&f.prompt, "50", "strength_clip"), &json!(0.0));
assert_eq!(input(&half.prompt, "50", "strength_clip"), &json!(0.4));
}
#[test]
@@ -812,9 +874,7 @@ mod tests {
"{:?}",
t.warnings
);
let f = t
.fill("p", None, None, &BTreeMap::new(), &mut seeds())
.unwrap();
let f = t.fill("p", None, None, &none(), &mut seeds()).unwrap();
assert!(f.loras.is_empty());
assert_eq!(input(&f.prompt, "50", "strength_model"), &json!(1.0));
assert!(
@@ -822,7 +882,7 @@ mod tests {
"p",
None,
None,
&switches(&[("style", false)]),
&reqs(&[("style", LoraRequest::On(false))]),
&mut seeds()
)
.is_err()