feat(enhance): pick LoRAs from ComfyUI per job (#220) #222
No files matched your search
@@ -460,7 +460,8 @@ auth and stays LAN-only, so ImageApi relays previews too.
|
||||
`unreachable_grace_secs`, `default_template`, `warn_above_megapixels`,
|
||||
`[[path_map]]` (`from`, `to`), `[[template]]` (`name`, `label`, `description`,
|
||||
`file`, `prompt_mode`, `prompt_join`, `output_megapixels`, `[[template.lora]]`
|
||||
(`id`, `label`, `default`, `strength`, `min`, `max`)).
|
||||
(`id`, `label`, `default`, `strength`, `min`, `max`), `[template.catalog]` (`strength`,
|
||||
`min`, `max`, `exclude`)).
|
||||
- **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
|
||||
@@ -472,7 +473,12 @@ auth and stays LAN-only, so ImageApi relays previews too.
|
||||
low pair; the resolved scale multiplies each node's exported `strength_model` /
|
||||
`strength_clip` (1.0 = as exported, 0 = off, a pair keeps its ratio), 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
|
||||
both log a startup warning), `imageapi:lora-slot:<n>:high|low` (pickable LoRAs, #220:
|
||||
each slot needs a high and a low node, numbered 1.. without gaps; picks fill them in order,
|
||||
setting `lora_name` and scaling the exported strength, unused slots run at 0; the catalog is
|
||||
ComfyUI's `/object_info/LoraLoaderModelOnly` list (cached 60 s, re-read when a submit names
|
||||
an unknown id) paired by High/Low name tokens (`src/enhance/lora_catalog.rs`), minus every
|
||||
`lora_name` the workflow loads outside the slots and `catalog.exclude` substrings). `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
|
||||
@@ -519,8 +525,12 @@ GET /enhance/templates
|
||||
"prompt_mode": "append"|"replace", "prompt_required",
|
||||
"accepts_reference","default","output_megapixels",
|
||||
"warn_above_megapixels", "output": "image"|"video",
|
||||
"loras": [{"id","label","default","strength","min","max"}]}]}
|
||||
POST /enhance {"path","template","prompt"?,"reference_path"?,"loras"?: {id: bool|scale},"library"?} -> job
|
||||
"loras": [{"id","label","default","strength","min","max"}],
|
||||
"slots": n, "catalog"?: {"loras": [{"id","label"}],"strength","min","max"}}]}
|
||||
POST /enhance {"path","template","prompt"?,"reference_path"?,"loras"?: {id: bool|scale},
|
||||
"slots"?: [{"lora": catalog id, "scale"?}],"library"?} -> job
|
||||
(slots: at most the template's slot count, no duplicates, unknown id = 400; no scale =
|
||||
catalog.strength, <= 0 = slot unused)
|
||||
(true = the switch's strength, false or scale <= 0 = off, a scale is clamped to min..max;
|
||||
a missing switch uses its default; an id the template doesn't offer = 400)
|
||||
GET /enhance/jobs/{id} -> job
|
||||
@@ -540,6 +550,7 @@ POST /enhance/jobs/{id}/discard -> job (hidden from the list; ComfyUI untou
|
||||
|
||||
job = {"id","status","queue_position"?,"error"?,"path","library_id","template",
|
||||
"prompt"?,"reference_path"?,"loras"?: {id: scale} (resolved, 0 = off; absent without switches),
|
||||
"slots"?: [{"lora","label","scale"}] (picked LoRAs; absent when none),
|
||||
"created_at","updated_at",
|
||||
"run_ms": <ms ComfyUI ran it, or null until finished / when ComfyUI gave no timings>,
|
||||
"pending": [{"n","preview": "/enhance/jobs/{id}/preview/{n}","approved"}],
|
||||
|
||||
@@ -87,3 +87,17 @@ default = false # default: false; used when the request omi
|
||||
strength = 1.0 # default: 1.0; the scale when on without a number
|
||||
min = 0.0 # default: 0.0; slider range while on (a requested
|
||||
max = 1.5 # default: 1.5; scale is clamped; <= 0 is off)
|
||||
|
||||
# Picked LoRAs (#220): LoRA loader nodes titled "imageapi:lora-slot:<n>:high"
|
||||
# and ":low" (n = 1, 2, ... without gaps) are slots the app fills from a
|
||||
# catalog, one picked LoRA per slot. The catalog is every LoRA file ComfyUI
|
||||
# lists whose name pairs as High/Low (StyleHigh/StyleLow, x-HIGH/x-LOW,
|
||||
# x_high_noise/x_low_noise), minus the LoRAs the workflow already loads
|
||||
# (lightx2v, switches); a new download shows up within a minute. A pick's
|
||||
# scale multiplies the slot's exported strength; unused slots run at 0.
|
||||
# Settings for the template above (all optional):
|
||||
[template.catalog]
|
||||
strength = 1.0 # default: 1.0; a pick's scale when the request gives none
|
||||
min = 0.0 # default: 0.0; slider range (a requested scale is
|
||||
max = 1.5 # default: 1.5; clamped; <= 0 leaves the slot unused)
|
||||
exclude = [] # default: []; file-name substrings never offered
|
||||
@@ -0,0 +1 @@
|
||||
ALTER TABLE enhance_jobs DROP COLUMN lora_slots;
|
||||
@@ -0,0 +1,3 @@
|
||||
-- LoRAs picked into the template's slots (#220): JSON [{lora, label, scale}];
|
||||
-- NULL = none picked.
|
||||
ALTER TABLE enhance_jobs ADD COLUMN lora_slots TEXT;
|
||||
@@ -66,11 +66,14 @@ pub struct EnhanceJobRow {
|
||||
/// template had none. Owned by `crate::enhance` like `pending`.
|
||||
#[diesel(sql_type = Nullable<Text>)]
|
||||
pub loras: Option<String>,
|
||||
/// LoRAs picked into slots (#220), JSON `[{lora, label, scale}]`.
|
||||
#[diesel(sql_type = Nullable<Text>)]
|
||||
pub lora_slots: Option<String>,
|
||||
}
|
||||
|
||||
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, loras";
|
||||
discarded, created_at, updated_at, run_ms, loras, lora_slots";
|
||||
|
||||
/// One distinct prompt used across jobs, newest use first.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, QueryableByName)]
|
||||
@@ -175,7 +178,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::<Text, _>(&job.id)
|
||||
.bind::<Text, _>(&job.original_hash)
|
||||
@@ -196,6 +199,7 @@ impl EnhanceJobDao for SqliteEnhanceJobDao {
|
||||
.bind::<BigInt, _>(job.updated_at)
|
||||
.bind::<Nullable<BigInt>, _>(job.run_ms)
|
||||
.bind::<Nullable<Text>, _>(&job.loras)
|
||||
.bind::<Nullable<Text>, _>(&job.lora_slots)
|
||||
.execute(&mut *conn)
|
||||
.map(|_| ())
|
||||
.map_err(|e| anyhow::anyhow!("Insert error: {}", e))
|
||||
@@ -376,6 +380,7 @@ mod tests {
|
||||
updated_at: created,
|
||||
run_ms: None,
|
||||
loras: None,
|
||||
lora_slots: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -74,6 +74,7 @@ diesel::table! {
|
||||
updated_at -> BigInt,
|
||||
run_ms -> Nullable<BigInt>,
|
||||
loras -> Nullable<Text>,
|
||||
lora_slots -> Nullable<Text>,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -330,6 +330,13 @@ impl ComfyClient {
|
||||
}
|
||||
}
|
||||
|
||||
/// Every LoRA file ComfyUI can load (`lora_name` choices of
|
||||
/// LoraLoaderModelOnly), for the pickable-LoRA catalog (#220).
|
||||
pub async fn lora_names(&self) -> Result<Vec<String>, ComfyError> {
|
||||
let body = self.get_json("/object_info/LoraLoaderModelOnly").await?;
|
||||
Ok(super::lora_catalog::names_from_object_info(&body))
|
||||
}
|
||||
|
||||
async fn get_json(&self, path: &str) -> Result<Value, ComfyError> {
|
||||
let resp = self.http.get(self.url(path)).send().await?;
|
||||
if !resp.status().is_success() {
|
||||
|
||||
+82
-6
@@ -104,6 +104,55 @@ pub struct TemplateConfig {
|
||||
/// workflow nodes titled `imageapi:lora:<id>`.
|
||||
#[serde(rename = "lora", default)]
|
||||
pub loras: Vec<LoraConfig>,
|
||||
/// Picked LoRAs (#220), for templates with `imageapi:lora-slot:<n>`
|
||||
/// nodes: `[template.catalog]`.
|
||||
#[serde(default)]
|
||||
pub catalog: CatalogConfig,
|
||||
}
|
||||
|
||||
/// How LoRAs picked from the catalog are scaled, and which never show.
|
||||
#[derive(Debug, Clone, PartialEq, Deserialize, serde::Serialize)]
|
||||
#[serde(default, deny_unknown_fields)]
|
||||
pub struct CatalogConfig {
|
||||
/// Scale a pick gets when the request gives none (1.0 = as exported).
|
||||
pub strength: f64,
|
||||
/// Slider range; a requested scale is clamped to it (<= 0 = slot unused).
|
||||
pub min: f64,
|
||||
pub max: f64,
|
||||
/// LoRA file names containing any of these (case-insensitive) are never
|
||||
/// offered, on top of the LoRAs the workflow already loads.
|
||||
#[serde(skip_serializing)]
|
||||
pub exclude: Vec<String>,
|
||||
}
|
||||
|
||||
impl Default for CatalogConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
strength: 1.0,
|
||||
min: 0.0,
|
||||
max: 1.5,
|
||||
exclude: Vec::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl CatalogConfig {
|
||||
/// A pick's scale: absent = `strength`, <= 0 = unused, else clamped.
|
||||
pub fn resolve(&self, scale: Option<f64>) -> f64 {
|
||||
match scale {
|
||||
None => self.strength,
|
||||
Some(s) if s <= 0.0 => 0.0,
|
||||
Some(s) => s.clamp(self.min, self.max),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `0 <= min < max`, `strength` within them, all finite.
|
||||
fn check_range(strength: f64, min: f64, max: f64) -> bool {
|
||||
[strength, min, max].iter().all(|v| v.is_finite())
|
||||
&& min >= 0.0
|
||||
&& min < max
|
||||
&& (min..=max).contains(&strength)
|
||||
}
|
||||
|
||||
/// One LoRA switch (`[[template.lora]]`).
|
||||
@@ -228,12 +277,7 @@ impl EnhanceConfig {
|
||||
if l.label.trim().is_empty() {
|
||||
bail!("template {:?}: lora {:?} needs a label", t.name, l.id);
|
||||
}
|
||||
let finite = [l.strength, l.min, l.max].iter().all(|v| v.is_finite());
|
||||
if !finite
|
||||
|| l.min < 0.0
|
||||
|| l.min >= l.max
|
||||
|| !(l.min..=l.max).contains(&l.strength)
|
||||
{
|
||||
if !check_range(l.strength, l.min, l.max) {
|
||||
bail!(
|
||||
"template {:?}: lora {:?} needs 0 <= min < max and strength within them \
|
||||
(got min {}, strength {}, max {})",
|
||||
@@ -245,6 +289,17 @@ impl EnhanceConfig {
|
||||
);
|
||||
}
|
||||
}
|
||||
let c = &t.catalog;
|
||||
if !check_range(c.strength, c.min, c.max) {
|
||||
bail!(
|
||||
"template {:?}: [template.catalog] needs 0 <= min < max and strength within them \
|
||||
(got min {}, strength {}, max {})",
|
||||
t.name,
|
||||
c.min,
|
||||
c.strength,
|
||||
c.max
|
||||
);
|
||||
}
|
||||
}
|
||||
if let Some(d) = &self.default_template
|
||||
&& !self.templates.iter().any(|t| &t.name == d)
|
||||
@@ -414,6 +469,27 @@ mod tests {
|
||||
assert_eq!(parsed["d"], LoraRequest::Scale(1.0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn catalog_settings_parse_validate_and_resolve() {
|
||||
let base = "[[template]]\nname='a'\nlabel='A'\nfile='a.json'\n";
|
||||
let c = parse(base).unwrap().templates[0].catalog.clone();
|
||||
assert_eq!(c, CatalogConfig::default());
|
||||
let set = parse(&format!(
|
||||
"{base}[template.catalog]\nstrength=0.7\nmax=1.2\nexclude=['x']"
|
||||
))
|
||||
.unwrap();
|
||||
let c = &set.templates[0].catalog;
|
||||
assert_eq!(
|
||||
(c.strength, c.min, c.max, c.exclude.clone()),
|
||||
(0.7, 0.0, 1.2, vec!["x".to_string()])
|
||||
);
|
||||
assert_eq!(c.resolve(None), 0.7);
|
||||
assert_eq!(c.resolve(Some(0.0)), 0.0);
|
||||
assert_eq!(c.resolve(Some(5.0)), 1.2);
|
||||
assert!(parse(&format!("{base}[template.catalog]\nstrength=2.0")).is_err());
|
||||
assert!(parse(&format!("{base}[template.catalog]\nslots=2")).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn template_defaults() {
|
||||
let cfg = parse("[[template]]\nname='a'\nlabel='A'\nfile='a.json'").unwrap();
|
||||
|
||||
+38
-1
@@ -95,6 +95,9 @@ pub struct JobView {
|
||||
/// template had none.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
loras: Option<std::collections::BTreeMap<String, f64>>,
|
||||
/// The LoRAs picked into slots (#220); absent when none were.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
slots: Option<Vec<super::service::SlotChoice>>,
|
||||
created_at: i64,
|
||||
updated_at: i64,
|
||||
/// Milliseconds ComfyUI ran the job; null until it has finished.
|
||||
@@ -134,6 +137,10 @@ pub fn job_view(job: &EnhanceJobRow) -> JobView {
|
||||
.loras
|
||||
.as_deref()
|
||||
.and_then(|j| serde_json::from_str(j).ok()),
|
||||
slots: job
|
||||
.lora_slots
|
||||
.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,
|
||||
@@ -233,6 +240,19 @@ struct TemplateView {
|
||||
output: super::template::OutputKind,
|
||||
/// Optional LoRAs the client can switch on or off (#218).
|
||||
loras: Vec<super::config::LoraConfig>,
|
||||
/// How many LoRAs a job can pick (#220); 0 = none.
|
||||
slots: usize,
|
||||
/// The LoRAs it can pick from (`[{id, label}]`), and how they scale.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
catalog: Option<CatalogView>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct CatalogView {
|
||||
loras: Vec<super::lora_catalog::CatalogEntry>,
|
||||
strength: f64,
|
||||
min: f64,
|
||||
max: f64,
|
||||
}
|
||||
|
||||
#[get("/enhance/templates")]
|
||||
@@ -246,10 +266,15 @@ pub async fn get_templates(_: Claims, svc: web::Data<EnhanceService>) -> impl Re
|
||||
None
|
||||
};
|
||||
let default = svc.default_template().map(str::to_string);
|
||||
let mut catalogs = Vec::new();
|
||||
for t in svc.templates() {
|
||||
catalogs.push((t.slot_count() > 0).then_some(svc.catalog(t).await));
|
||||
}
|
||||
let templates: Vec<TemplateView> = svc
|
||||
.templates()
|
||||
.iter()
|
||||
.map(|t| TemplateView {
|
||||
.zip(catalogs)
|
||||
.map(|(t, catalog)| TemplateView {
|
||||
name: t.cfg.name.clone(),
|
||||
label: t.cfg.label.clone(),
|
||||
description: t.cfg.description.clone(),
|
||||
@@ -262,6 +287,13 @@ pub async fn get_templates(_: Claims, svc: web::Data<EnhanceService>) -> impl Re
|
||||
warn_above_megapixels: svc.config.warn_above_megapixels,
|
||||
output: t.output,
|
||||
loras: t.lora_switches().cloned().collect(),
|
||||
slots: t.slot_count(),
|
||||
catalog: catalog.map(|loras| CatalogView {
|
||||
loras,
|
||||
strength: t.cfg.catalog.strength,
|
||||
min: t.cfg.catalog.min,
|
||||
max: t.cfg.catalog.max,
|
||||
}),
|
||||
})
|
||||
.collect();
|
||||
let mut body = json!({
|
||||
@@ -295,6 +327,10 @@ pub struct SubmitBody {
|
||||
/// the switch's range; <= 0 is off). A missing id uses its default.
|
||||
#[serde(default)]
|
||||
pub loras: std::collections::BTreeMap<String, super::config::LoraRequest>,
|
||||
/// LoRAs picked from the template's catalog (#220), one per slot in
|
||||
/// order: `[{lora: id, scale?}]`; no scale = the catalog's strength.
|
||||
#[serde(default)]
|
||||
pub slots: Vec<super::service::SlotRequest>,
|
||||
#[serde(default)]
|
||||
pub library: Option<String>,
|
||||
}
|
||||
@@ -341,6 +377,7 @@ pub async fn submit(
|
||||
body.prompt.as_deref(),
|
||||
reference.as_ref(),
|
||||
&body.loras,
|
||||
&body.slots,
|
||||
)
|
||||
.await
|
||||
{
|
||||
|
||||
@@ -27,6 +27,7 @@ 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 ANIMATE_SLOTS: &str = include_str!("../../tests/fixtures/enhance/animate_slots.json");
|
||||
const EDIT: &str = include_str!("../../tests/fixtures/enhance/edit_turbo.json");
|
||||
|
||||
// ---- mock ComfyUI ---------------------------------------------------------
|
||||
@@ -45,6 +46,9 @@ struct Mock {
|
||||
free_calls: AtomicUsize,
|
||||
reject_prompts: Mutex<Option<Value>>,
|
||||
view_404: Mutex<bool>,
|
||||
/// What `/object_info/LoraLoaderModelOnly` lists.
|
||||
lora_names: Mutex<Vec<String>>,
|
||||
lora_list_calls: AtomicUsize,
|
||||
}
|
||||
|
||||
async fn start_mock() -> (Arc<Mock>, String) {
|
||||
@@ -54,6 +58,16 @@ async fn start_mock() -> (Arc<Mock>, String) {
|
||||
App::new()
|
||||
.app_data(Data::from(m.clone()))
|
||||
.route("/system_stats", web::get().to(|| async { HttpResponse::Ok().json(json!({})) }))
|
||||
.route(
|
||||
"/object_info/LoraLoaderModelOnly",
|
||||
web::get().to(|m: Data<Mock>| async move {
|
||||
m.lora_list_calls.fetch_add(1, Ordering::SeqCst);
|
||||
let names = m.lora_names.lock().unwrap().clone();
|
||||
HttpResponse::Ok().json(json!({"LoraLoaderModelOnly": {"input": {"required": {
|
||||
"lora_name": [names], "strength_model": ["FLOAT", {}]
|
||||
}}}}))
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/prompt",
|
||||
web::post().to(|m: Data<Mock>, body: web::Json<Value>| async move {
|
||||
@@ -170,6 +184,7 @@ fn template(name: &str, mode: PromptMode, json: &str) -> Template {
|
||||
prompt_join: " ".into(),
|
||||
output_megapixels: 1.0,
|
||||
loras: Vec::new(),
|
||||
catalog: Default::default(),
|
||||
},
|
||||
serde_json::from_str(json).unwrap(),
|
||||
)
|
||||
@@ -222,6 +237,7 @@ async fn fx_with(url: Option<String>, tweak: impl FnOnce(&mut EnhanceConfig)) ->
|
||||
template("restore", PromptMode::Append, RESTORE),
|
||||
template("edit", PromptMode::Replace, EDIT),
|
||||
animate_with_style(),
|
||||
template("animate-picks", PromptMode::Append, ANIMATE_SLOTS),
|
||||
],
|
||||
dao,
|
||||
)
|
||||
@@ -350,7 +366,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(), 3);
|
||||
assert_eq!(body["templates"].as_array().unwrap().len(), 4);
|
||||
|
||||
let (st, body, _) = call(
|
||||
&fx,
|
||||
@@ -1010,6 +1026,7 @@ async fn unreachable_active_jobs_become_unavailable_after_the_grace_period() {
|
||||
updated_at: now,
|
||||
run_ms: None,
|
||||
loras: None,
|
||||
lora_slots: None,
|
||||
};
|
||||
SqliteEnhanceJobDao::from_shared(fx.conn.clone())
|
||||
.insert(&dao_job)
|
||||
@@ -1223,3 +1240,114 @@ async fn lora_switches_are_listed_sent_and_stored_on_the_job() {
|
||||
.await;
|
||||
assert!(plain.get("loras").is_none(), "{plain}");
|
||||
}
|
||||
|
||||
#[actix_rt::test]
|
||||
async fn picked_loras_are_listed_filled_into_slots_and_stored() {
|
||||
let fx = fx().await;
|
||||
fx.put(PHOTO, b"photo");
|
||||
*fx.mock.lora_names.lock().unwrap() = [
|
||||
"wan2.2_i2v_lightx2v_4steps_lora_v1_high_noise.safetensors",
|
||||
"wan2.2_i2v_lightx2v_4steps_lora_v1_low_noise.safetensors",
|
||||
"example_style_high.safetensors",
|
||||
"example_style_low.safetensors",
|
||||
"MotionHigh.safetensors",
|
||||
"MotionLow.safetensors",
|
||||
"Lonely_High.safetensors",
|
||||
]
|
||||
.map(String::from)
|
||||
.to_vec();
|
||||
|
||||
// Only the slot template lists a catalog; the workflow's own LoRAs
|
||||
// (lightx2v, the style switch) and unpaired files are left out.
|
||||
let (_, body, _) = call(&fx, "GET", "/enhance/templates", None).await;
|
||||
let templates = body["templates"].as_array().unwrap();
|
||||
let picks = templates
|
||||
.iter()
|
||||
.find(|t| t["name"] == "animate-picks")
|
||||
.unwrap();
|
||||
assert_eq!(picks["slots"], 2);
|
||||
assert_eq!(
|
||||
picks["catalog"],
|
||||
json!({"loras": [{"id": "motion", "label": "Motion"}], "strength": 1.0, "min": 0.0, "max": 1.5})
|
||||
);
|
||||
let restore = templates.iter().find(|t| t["name"] == "restore").unwrap();
|
||||
assert_eq!(restore["slots"], 0);
|
||||
assert!(restore.get("catalog").is_none());
|
||||
|
||||
let (st, job, _) = call(
|
||||
&fx,
|
||||
"POST",
|
||||
"/enhance",
|
||||
Some(json!({"path": PHOTO, "template": "animate-picks", "slots": [{"lora": "motion", "scale": 0.5}]})),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(st, 200, "{job}");
|
||||
assert_eq!(
|
||||
job["slots"],
|
||||
json!([{"lora": "motion", "label": "Motion", "scale": 0.5}])
|
||||
);
|
||||
let (_, again, _) = call(
|
||||
&fx,
|
||||
"GET",
|
||||
&format!("/enhance/jobs/{}", job["id"].as_str().unwrap()),
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
assert_eq!(again["slots"], job["slots"]);
|
||||
{
|
||||
let sent = fx.mock.submitted.lock().unwrap();
|
||||
let p = &sent[0]["prompt"];
|
||||
assert_eq!(p["60"]["inputs"]["lora_name"], "MotionHigh.safetensors");
|
||||
assert_eq!(p["61"]["inputs"]["lora_name"], "MotionLow.safetensors");
|
||||
assert_eq!(p["60"]["inputs"]["strength_model"], json!(0.5));
|
||||
assert_eq!(p["62"]["inputs"]["strength_model"], json!(0.0));
|
||||
}
|
||||
|
||||
// No scale: the catalog's strength.
|
||||
let (_, job, _) = call(
|
||||
&fx,
|
||||
"POST",
|
||||
"/enhance",
|
||||
Some(json!({"path": PHOTO, "template": "animate-picks", "slots": [{"lora": "motion"}]})),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(job["slots"][0]["scale"], json!(1.0));
|
||||
|
||||
// A LoRA downloaded since the last listing is found by re-reading it.
|
||||
fx.mock.lora_names.lock().unwrap().extend([
|
||||
"NewHigh.safetensors".to_string(),
|
||||
"NewLow.safetensors".to_string(),
|
||||
]);
|
||||
let (st, job, _) = call(
|
||||
&fx,
|
||||
"POST",
|
||||
"/enhance",
|
||||
Some(json!({"path": PHOTO, "template": "animate-picks", "slots": [{"lora": "new"}, {"lora": "motion"}]})),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(st, 200, "{job}");
|
||||
assert_eq!(job["slots"][0]["lora"], "new");
|
||||
|
||||
// Unknown, duplicate, too many, or on a template without slots: 400.
|
||||
for (template, slots) in [
|
||||
("animate-picks", json!([{"lora": "nope"}])),
|
||||
(
|
||||
"animate-picks",
|
||||
json!([{"lora": "motion"}, {"lora": "motion"}]),
|
||||
),
|
||||
(
|
||||
"animate-picks",
|
||||
json!([{"lora": "motion"}, {"lora": "new"}, {"lora": "lonely"}]),
|
||||
),
|
||||
("restore", json!([{"lora": "motion"}])),
|
||||
] {
|
||||
let (st, err, _) = call(
|
||||
&fx,
|
||||
"POST",
|
||||
"/enhance",
|
||||
Some(json!({"path": PHOTO, "template": template, "slots": slots})),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(st, 400, "{template} {err}");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,336 @@
|
||||
//! Pickable LoRAs (#220): the LoRA files ComfyUI can see, paired into the
|
||||
//! high-noise / low-noise sets Wan 2.2's two experts need.
|
||||
//!
|
||||
//! Nothing is configured per LoRA: a file whose name has a `high` or `low`
|
||||
//! token pairs with the file that has the other token and is otherwise the
|
||||
//! same name (`BounceHighWan2_2` + `BounceLowWan2_2`, `x-HIGH` + `x-LOW`,
|
||||
//! `…_high_noise` + `…_low_noise`). A file without a partner isn't offered.
|
||||
//! The ids are derived from the names, so they stay stable across restarts
|
||||
//! and refreshes.
|
||||
|
||||
use std::collections::{BTreeMap, HashSet};
|
||||
|
||||
use serde::Serialize;
|
||||
|
||||
/// One pickable LoRA: the pair of files and how the app shows it.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
|
||||
pub struct CatalogEntry {
|
||||
/// Stable id clients send (`POST /enhance { slots: [{lora}] }`).
|
||||
pub id: String,
|
||||
pub label: String,
|
||||
/// `lora_name` values as ComfyUI lists them (folder prefix included).
|
||||
#[serde(skip)]
|
||||
pub high: String,
|
||||
#[serde(skip)]
|
||||
pub low: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
enum Noise {
|
||||
High,
|
||||
Low,
|
||||
}
|
||||
|
||||
/// Pair `names` (ComfyUI's `lora_name` choices) into catalog entries, sorted
|
||||
/// by label. `exclude` holds names that never appear (the template's own
|
||||
/// fixed LoRAs), compared exactly; `exclude_substrings` drops any name that
|
||||
/// contains one of them (case-insensitive).
|
||||
pub fn build(
|
||||
names: &[String],
|
||||
exclude: &HashSet<String>,
|
||||
exclude_substrings: &[String],
|
||||
) -> Vec<CatalogEntry> {
|
||||
let lowered: Vec<String> = exclude_substrings
|
||||
.iter()
|
||||
.map(|s| s.to_lowercase())
|
||||
.filter(|s| !s.is_empty())
|
||||
.collect();
|
||||
// key (the name with its noise token masked) -> (high, low, masked stem)
|
||||
let mut sides: BTreeMap<String, (Option<&String>, Option<&String>, String)> = BTreeMap::new();
|
||||
for name in names {
|
||||
if exclude.contains(name) {
|
||||
continue;
|
||||
}
|
||||
let lower = name.to_lowercase();
|
||||
if lowered.iter().any(|s| lower.contains(s)) {
|
||||
continue;
|
||||
}
|
||||
let Some((noise, masked)) = classify(name) else {
|
||||
continue;
|
||||
};
|
||||
let entry = sides
|
||||
.entry(masked.to_lowercase())
|
||||
.or_insert((None, None, masked));
|
||||
match noise {
|
||||
Noise::High => entry.0 = entry.0.or(Some(name)),
|
||||
Noise::Low => entry.1 = entry.1.or(Some(name)),
|
||||
}
|
||||
}
|
||||
|
||||
let mut taken = HashSet::new();
|
||||
let mut out: Vec<CatalogEntry> = sides
|
||||
.into_values()
|
||||
.filter_map(|(high, low, masked)| {
|
||||
let (high, low) = (high?, low?);
|
||||
let base = slug(&masked);
|
||||
let mut id = base.clone();
|
||||
let mut n = 2;
|
||||
while !taken.insert(id.clone()) {
|
||||
id = format!("{base}-{n}");
|
||||
n += 1;
|
||||
}
|
||||
Some(CatalogEntry {
|
||||
id,
|
||||
label: label(&masked),
|
||||
high: high.clone(),
|
||||
low: low.clone(),
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
out.sort_by(|a, b| {
|
||||
a.label
|
||||
.to_lowercase()
|
||||
.cmp(&b.label.to_lowercase())
|
||||
.then(a.id.cmp(&b.id))
|
||||
});
|
||||
out
|
||||
}
|
||||
|
||||
/// The name's noise side and the name with that token masked as `{}`
|
||||
/// (folder and extension kept, so two folders never pair with each other).
|
||||
///
|
||||
/// A token counts when it starts the name, follows a non-letter, or is
|
||||
/// capitalised (camel case), so `slow` or `Glowing` are never a `low`. The
|
||||
/// last such token wins (`High…_low_noise` is low).
|
||||
fn classify(name: &str) -> Option<(Noise, String)> {
|
||||
let (dir, file) = match name.rfind(['/', '\\']) {
|
||||
Some(i) => name.split_at(i + 1),
|
||||
None => ("", name),
|
||||
};
|
||||
let stem_end = file.rfind('.').unwrap_or(file.len());
|
||||
let stem = &file[..stem_end];
|
||||
let lower = stem.to_lowercase();
|
||||
let mut best: Option<(usize, usize, Noise)> = None;
|
||||
for (token, noise) in [("high", Noise::High), ("low", Noise::Low)] {
|
||||
for (at, _) in lower.match_indices(token) {
|
||||
let original = &stem[at..at + token.len()];
|
||||
let boundary = at == 0
|
||||
|| !stem[..at]
|
||||
.chars()
|
||||
.next_back()
|
||||
.is_some_and(|c| c.is_alphabetic())
|
||||
|| original.starts_with(char::is_uppercase);
|
||||
// Not the start of a longer lowercase word (`lower`, `highway`).
|
||||
let after = stem[at + token.len()..].chars().next();
|
||||
let ends = match after {
|
||||
None => true,
|
||||
Some(c) => {
|
||||
!c.is_alphabetic()
|
||||
|| c.is_uppercase()
|
||||
|| original.chars().all(char::is_uppercase)
|
||||
}
|
||||
};
|
||||
if boundary && ends && best.is_none_or(|(b, _, _)| at > b) {
|
||||
best = Some((at, token.len(), noise));
|
||||
}
|
||||
}
|
||||
}
|
||||
let (at, len, noise) = best?;
|
||||
let masked = format!(
|
||||
"{dir}{}{{}}{}{}",
|
||||
&stem[..at],
|
||||
&stem[at + len..],
|
||||
&file[stem_end..]
|
||||
);
|
||||
Some((noise, masked))
|
||||
}
|
||||
|
||||
/// `wan22/Bounce{}Wan2_2.safetensors` -> `bounce-wan2-2`.
|
||||
fn slug(masked: &str) -> String {
|
||||
let file = masked.rsplit(['/', '\\']).next().unwrap_or(masked);
|
||||
let stem = file.rsplit_once('.').map(|(s, _)| s).unwrap_or(file);
|
||||
let mut out = String::new();
|
||||
for c in split_words(stem).join(" ").chars() {
|
||||
if c.is_ascii_alphanumeric() {
|
||||
out.push(c.to_ascii_lowercase());
|
||||
} else if !out.ends_with('-') {
|
||||
out.push('-');
|
||||
}
|
||||
}
|
||||
let out = out.trim_matches('-').to_string();
|
||||
if out.is_empty() { "lora".into() } else { out }
|
||||
}
|
||||
|
||||
/// `wan22/Bounce{}Wan2_2.safetensors` -> `Bounce Wan2 2`.
|
||||
fn label(masked: &str) -> String {
|
||||
let file = masked.rsplit(['/', '\\']).next().unwrap_or(masked);
|
||||
let stem = file.rsplit_once('.').map(|(s, _)| s).unwrap_or(file);
|
||||
let words = split_words(stem);
|
||||
if words.is_empty() {
|
||||
stem.to_string()
|
||||
} else {
|
||||
words.join(" ")
|
||||
}
|
||||
}
|
||||
|
||||
/// Words of a stem: split on separators, the noise mask and camel case.
|
||||
fn split_words(stem: &str) -> Vec<String> {
|
||||
let mut words = Vec::new();
|
||||
let mut cur = String::new();
|
||||
let mut prev: Option<char> = None;
|
||||
for c in stem.replace("{}", " ").chars() {
|
||||
if c == '_' || c == '-' || c == ' ' {
|
||||
if !cur.is_empty() {
|
||||
words.push(std::mem::take(&mut cur));
|
||||
}
|
||||
} else {
|
||||
if c.is_uppercase() && prev.is_some_and(|p| p.is_lowercase()) && !cur.is_empty() {
|
||||
words.push(std::mem::take(&mut cur));
|
||||
}
|
||||
cur.push(c);
|
||||
}
|
||||
prev = Some(c);
|
||||
}
|
||||
if !cur.is_empty() {
|
||||
words.push(cur);
|
||||
}
|
||||
words
|
||||
}
|
||||
|
||||
/// ComfyUI's `lora_name` choices from `GET /object_info/LoraLoaderModelOnly`,
|
||||
/// in either shape: `[[..names]]` (classic) or `["COMBO", {"options": [..]}]`.
|
||||
pub fn names_from_object_info(body: &serde_json::Value) -> Vec<String> {
|
||||
let spec = &body["LoraLoaderModelOnly"]["input"]["required"]["lora_name"];
|
||||
let list = if spec[0].is_array() {
|
||||
&spec[0]
|
||||
} else {
|
||||
&spec[1]["options"]
|
||||
};
|
||||
list.as_array()
|
||||
.map(|a| {
|
||||
a.iter()
|
||||
.filter_map(|v| v.as_str().map(str::to_string))
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
|
||||
fn names(v: &[&str]) -> Vec<String> {
|
||||
v.iter().map(|s| s.to_string()).collect()
|
||||
}
|
||||
|
||||
fn ids(c: &[CatalogEntry]) -> Vec<(&str, &str)> {
|
||||
c.iter()
|
||||
.map(|e| (e.id.as_str(), e.label.as_str()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pairs_the_usual_naming_styles() {
|
||||
let c = build(
|
||||
&names(&[
|
||||
"StyleHighWan2_2.safetensors",
|
||||
"StyleLowWan2_2.safetensors",
|
||||
"Motionv2.5-HIGH.safetensors",
|
||||
"Motionv2.5-LOW.safetensors",
|
||||
"Detail_High.safetensors",
|
||||
"Detail_Low.safetensors",
|
||||
"wan_effect_high_noise.safetensors",
|
||||
"wan_effect_low_noise.safetensors",
|
||||
]),
|
||||
&HashSet::new(),
|
||||
&[],
|
||||
);
|
||||
assert_eq!(
|
||||
ids(&c),
|
||||
[
|
||||
("detail", "Detail"),
|
||||
("motionv2-5", "Motionv2.5"),
|
||||
("style-wan2-2", "Style Wan2 2"),
|
||||
("wan-effect-noise", "wan effect noise"),
|
||||
]
|
||||
);
|
||||
let style = c.iter().find(|e| e.id == "style-wan2-2").unwrap();
|
||||
assert_eq!(style.high, "StyleHighWan2_2.safetensors");
|
||||
assert_eq!(style.low, "StyleLowWan2_2.safetensors");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unpaired_and_tokenless_files_are_left_out() {
|
||||
let c = build(
|
||||
&names(&[
|
||||
"OnlyHigh_High.safetensors",
|
||||
"distill_rank64_bf16.safetensors",
|
||||
"slow_motion.safetensors",
|
||||
"Glowing_low.safetensors",
|
||||
"Glowing_high.safetensors",
|
||||
"lower_body_high.safetensors",
|
||||
]),
|
||||
&HashSet::new(),
|
||||
&[],
|
||||
);
|
||||
// `slow`, `Glowing` and `lower` hold no token; Glowing pairs on its suffix.
|
||||
assert_eq!(ids(&c), [("glowing", "Glowing")]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn excludes_fixed_loras_and_substrings() {
|
||||
let all = names(&[
|
||||
"wan2.2_i2v_lightx2v_4steps_lora_v1_high_noise.safetensors",
|
||||
"wan2.2_i2v_lightx2v_4steps_lora_v1_low_noise.safetensors",
|
||||
"StyleHigh.safetensors",
|
||||
"StyleLow.safetensors",
|
||||
"SecretHigh.safetensors",
|
||||
"SecretLow.safetensors",
|
||||
]);
|
||||
let fixed: HashSet<String> = all[..2].iter().cloned().collect();
|
||||
let c = build(&all, &fixed, &["secret".into()]);
|
||||
assert_eq!(ids(&c), [("style", "Style")]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn folders_keep_pairs_apart_and_ids_stay_unique() {
|
||||
let c = build(
|
||||
&names(&[
|
||||
"a/StyleHigh.safetensors",
|
||||
"a/StyleLow.safetensors",
|
||||
"b\\StyleHigh.safetensors",
|
||||
"b\\StyleLow.safetensors",
|
||||
"a/StyleHigh.safetensors",
|
||||
]),
|
||||
&HashSet::new(),
|
||||
&[],
|
||||
);
|
||||
assert_eq!(c.len(), 2);
|
||||
let mut got: Vec<_> = c.iter().map(|e| (e.id.as_str(), e.high.as_str())).collect();
|
||||
got.sort();
|
||||
assert_eq!(
|
||||
got,
|
||||
[
|
||||
("style", "a/StyleHigh.safetensors"),
|
||||
("style-2", "b\\StyleHigh.safetensors")
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reads_both_object_info_shapes() {
|
||||
let classic = json!({"LoraLoaderModelOnly": {"input": {"required": {
|
||||
"lora_name": [["a.safetensors", "b.safetensors"]], "strength_model": ["FLOAT", {}]
|
||||
}}}});
|
||||
assert_eq!(
|
||||
names_from_object_info(&classic),
|
||||
names(&["a.safetensors", "b.safetensors"])
|
||||
);
|
||||
let combo = json!({"LoraLoaderModelOnly": {"input": {"required": {
|
||||
"lora_name": ["COMBO", {"options": ["c.safetensors"]}]
|
||||
}}}});
|
||||
assert_eq!(names_from_object_info(&combo), names(&["c.safetensors"]));
|
||||
assert!(names_from_object_info(&json!({})).is_empty());
|
||||
}
|
||||
}
|
||||
@@ -14,6 +14,7 @@
|
||||
pub mod comfy;
|
||||
pub mod config;
|
||||
pub mod handlers;
|
||||
pub mod lora_catalog;
|
||||
pub mod path_map;
|
||||
pub mod service;
|
||||
pub mod template;
|
||||
|
||||
+135
-1
@@ -15,8 +15,9 @@ use serde::{Deserialize, Serialize};
|
||||
|
||||
use super::comfy::{ApproveError, ComfyClient, ComfyError, HistoryEntry, Staged, SubmitError};
|
||||
use super::config::{EnhanceConfig, Loaded, PromptMode};
|
||||
use super::lora_catalog::{self, CatalogEntry};
|
||||
use super::path_map::PathMap;
|
||||
use super::template::{Template, random_seed};
|
||||
use super::template::{SlotPick, Template, random_seed};
|
||||
use crate::database::enhance_dao::{
|
||||
EnhanceJobDao, EnhanceJobRow, JobUpdate, RecentPrompt, STATUS_APPROVED, STATUS_EXPIRED,
|
||||
STATUS_FAILED, STATUS_QUEUED, STATUS_READY, STATUS_RUNNING, STATUS_UNAVAILABLE,
|
||||
@@ -165,6 +166,27 @@ pub struct EnhanceService {
|
||||
idle: Mutex<IdleState>,
|
||||
/// Serialises approvals (rare; keeps "approve twice" simple).
|
||||
approve_lock: tokio::sync::Mutex<()>,
|
||||
/// ComfyUI's LoRA files and when they were read (#220).
|
||||
lora_names: Mutex<Option<(Instant, Vec<String>)>>,
|
||||
}
|
||||
|
||||
/// How long ComfyUI's LoRA list is reused; a new download shows up after it.
|
||||
const LORA_LIST_TTL: Duration = Duration::from_secs(60);
|
||||
|
||||
/// One `POST /enhance { slots }` entry: a catalog id and an optional scale.
|
||||
#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
|
||||
pub struct SlotRequest {
|
||||
pub lora: String,
|
||||
#[serde(default)]
|
||||
pub scale: Option<f64>,
|
||||
}
|
||||
|
||||
/// A resolved pick, as stored on the job and returned in its view.
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct SlotChoice {
|
||||
pub lora: String,
|
||||
pub label: String,
|
||||
pub scale: f64,
|
||||
}
|
||||
|
||||
impl EnhanceService {
|
||||
@@ -233,9 +255,116 @@ impl EnhanceService {
|
||||
unreachable_since: Mutex::new(HashMap::new()),
|
||||
idle: Mutex::new(IdleState::default()),
|
||||
approve_lock: tokio::sync::Mutex::new(()),
|
||||
lora_names: Mutex::new(None),
|
||||
})
|
||||
}
|
||||
|
||||
/// ComfyUI's LoRA files, cached for [`LORA_LIST_TTL`]. `fresh` re-reads
|
||||
/// them (a submit naming an id the cache doesn't know). An unreachable
|
||||
/// ComfyUI keeps the last list.
|
||||
async fn lora_names(&self, fresh: bool) -> Vec<String> {
|
||||
let Some(client) = &self.client else {
|
||||
return Vec::new();
|
||||
};
|
||||
if !fresh
|
||||
&& let Some((at, names)) = &*self.lora_names.lock().unwrap()
|
||||
&& at.elapsed() < LORA_LIST_TTL
|
||||
{
|
||||
return names.clone();
|
||||
}
|
||||
match client.lora_names().await {
|
||||
Ok(names) => {
|
||||
*self.lora_names.lock().unwrap() = Some((Instant::now(), names.clone()));
|
||||
names
|
||||
}
|
||||
Err(e) => {
|
||||
warn!("enhance: couldn't list ComfyUI's LoRAs: {e}");
|
||||
self.lora_names
|
||||
.lock()
|
||||
.unwrap()
|
||||
.as_ref()
|
||||
.map(|(_, n)| n.clone())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The LoRAs a template's slots can take; empty without slots.
|
||||
pub async fn catalog(&self, template: &Template) -> Vec<CatalogEntry> {
|
||||
if template.slot_count() == 0 {
|
||||
return Vec::new();
|
||||
}
|
||||
let names = self.lora_names(false).await;
|
||||
lora_catalog::build(
|
||||
&names,
|
||||
template.fixed_loras(),
|
||||
&template.cfg.catalog.exclude,
|
||||
)
|
||||
}
|
||||
|
||||
/// Turn `slots` into picks for the template's slots, checking the ids
|
||||
/// against the catalog (re-read once if one is unknown).
|
||||
async fn resolve_slots(
|
||||
&self,
|
||||
template: &Template,
|
||||
slots: &[SlotRequest],
|
||||
) -> Result<(Vec<SlotPick>, Vec<SlotChoice>), EnhanceError> {
|
||||
if slots.is_empty() {
|
||||
return Ok((Vec::new(), Vec::new()));
|
||||
}
|
||||
if slots.len() > template.slot_count() {
|
||||
return Err(EnhanceError::BadRequest(if template.slot_count() == 0 {
|
||||
"this template doesn't take picked LoRAs".to_string()
|
||||
} else {
|
||||
format!(
|
||||
"this template takes at most {} picked LoRAs",
|
||||
template.slot_count()
|
||||
)
|
||||
}));
|
||||
}
|
||||
let mut seen = std::collections::HashSet::new();
|
||||
if let Some(dup) = slots.iter().find(|s| !seen.insert(s.lora.as_str())) {
|
||||
return Err(EnhanceError::BadRequest(format!(
|
||||
"LoRA {:?} is picked twice",
|
||||
dup.lora
|
||||
)));
|
||||
}
|
||||
let mut catalog = self.catalog(template).await;
|
||||
if slots
|
||||
.iter()
|
||||
.any(|s| !catalog.iter().any(|e| e.id == s.lora))
|
||||
{
|
||||
let names = self.lora_names(true).await;
|
||||
catalog = lora_catalog::build(
|
||||
&names,
|
||||
template.fixed_loras(),
|
||||
&template.cfg.catalog.exclude,
|
||||
);
|
||||
}
|
||||
let mut picks = Vec::new();
|
||||
let mut choices = Vec::new();
|
||||
for s in slots {
|
||||
let Some(entry) = catalog.iter().find(|e| e.id == s.lora) else {
|
||||
return Err(EnhanceError::BadRequest(format!(
|
||||
"no LoRA {:?} (it may have been removed from ComfyUI)",
|
||||
s.lora
|
||||
)));
|
||||
};
|
||||
let scale = template.cfg.catalog.resolve(s.scale);
|
||||
picks.push(SlotPick {
|
||||
high: entry.high.clone(),
|
||||
low: entry.low.clone(),
|
||||
scale,
|
||||
});
|
||||
choices.push(SlotChoice {
|
||||
lora: entry.id.clone(),
|
||||
label: entry.label.clone(),
|
||||
scale,
|
||||
});
|
||||
}
|
||||
Ok((picks, choices))
|
||||
}
|
||||
|
||||
pub fn is_configured(&self) -> bool {
|
||||
self.client.is_some()
|
||||
}
|
||||
@@ -293,6 +422,7 @@ impl EnhanceService {
|
||||
prompt: Option<&str>,
|
||||
reference: Option<&SourceFile>,
|
||||
loras: &std::collections::BTreeMap<String, super::config::LoraRequest>,
|
||||
slots: &[SlotRequest],
|
||||
) -> Result<EnhanceJobRow, EnhanceError> {
|
||||
let Some(client) = &self.client else {
|
||||
return Err(EnhanceError::Disabled);
|
||||
@@ -319,12 +449,14 @@ impl EnhanceService {
|
||||
),
|
||||
None => None,
|
||||
};
|
||||
let (picks, choices) = self.resolve_slots(template, slots).await?;
|
||||
let filled = template
|
||||
.fill(
|
||||
&source,
|
||||
prompt,
|
||||
reference_path.as_deref(),
|
||||
loras,
|
||||
&picks,
|
||||
&mut random_seed,
|
||||
)
|
||||
.map_err(|e| EnhanceError::BadRequest(e.0))?;
|
||||
@@ -374,6 +506,8 @@ impl EnhanceService {
|
||||
run_ms: None,
|
||||
loras: (!filled.loras.is_empty())
|
||||
.then(|| serde_json::to_string(&filled.loras).unwrap_or_default()),
|
||||
lora_slots: (!choices.is_empty())
|
||||
.then(|| serde_json::to_string(&choices).unwrap_or_default()),
|
||||
};
|
||||
self.dao
|
||||
.lock()
|
||||
|
||||
+281
-18
@@ -18,6 +18,9 @@ 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:";
|
||||
/// `imageapi:lora-slot:<n>:high|low`: a LoRA loader the app fills with a
|
||||
/// LoRA picked from the catalog (#220); unused slots run at strength 0.
|
||||
pub const TITLE_SLOT_PREFIX: &str = "imageapi:lora-slot:";
|
||||
|
||||
/// LoRA loader inputs a switch scales (0 = off). `strength_model` is
|
||||
/// required (LoraLoaderModelOnly); `strength_clip` is scaled when present
|
||||
@@ -61,6 +64,11 @@ pub struct Template {
|
||||
reference_id: Option<String>,
|
||||
/// Declared LoRA switches that have nodes, in `enhance.toml` order.
|
||||
loras: Vec<LoraSwitch>,
|
||||
/// Pickable-LoRA slots, by slot number.
|
||||
slots: Vec<Slot>,
|
||||
/// `lora_name`s the workflow loads outside the slots (lightx2v, switch
|
||||
/// LoRAs): never offered in the catalog.
|
||||
fixed_loras: std::collections::HashSet<String>,
|
||||
/// Image or video, from the save node's class.
|
||||
pub output: OutputKind,
|
||||
/// Startup warnings (stray output nodes).
|
||||
@@ -75,6 +83,22 @@ struct LoraSwitch {
|
||||
nodes: Vec<(String, Vec<(&'static str, f64)>)>,
|
||||
}
|
||||
|
||||
/// One pickable-LoRA slot: its high- and low-noise loader nodes, each with
|
||||
/// its exported strengths (a pick's scale multiplies them).
|
||||
#[derive(Debug, Clone)]
|
||||
struct Slot {
|
||||
high: Vec<(String, Vec<(&'static str, f64)>)>,
|
||||
low: Vec<(String, Vec<(&'static str, f64)>)>,
|
||||
}
|
||||
|
||||
/// A LoRA picked for one slot: its two files and the resolved scale.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct SlotPick {
|
||||
pub high: String,
|
||||
pub low: String,
|
||||
pub scale: f64,
|
||||
}
|
||||
|
||||
/// 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)]
|
||||
@@ -227,6 +251,69 @@ impl Template {
|
||||
));
|
||||
}
|
||||
|
||||
// Pickable-LoRA slots, and the LoRAs loaded outside them.
|
||||
let exported = |id: &str| -> Vec<(&'static str, f64)> {
|
||||
let inputs = inputs_of(&workflow, id);
|
||||
LORA_STRENGTH_INPUTS
|
||||
.iter()
|
||||
.filter_map(|k| inputs.get(*k).and_then(Value::as_f64).map(|v| (*k, v)))
|
||||
.collect()
|
||||
};
|
||||
let mut slot_map: BTreeMap<u32, Slot> = BTreeMap::new();
|
||||
let mut fixed_loras = std::collections::HashSet::new();
|
||||
for (id, node) in &workflow {
|
||||
let Some(rest) = title_of(node).and_then(|t| t.strip_prefix(TITLE_SLOT_PREFIX)) else {
|
||||
if let Some(name) = inputs_of(&workflow, id)
|
||||
.get("lora_name")
|
||||
.and_then(Value::as_str)
|
||||
{
|
||||
fixed_loras.insert(name.to_string());
|
||||
}
|
||||
continue;
|
||||
};
|
||||
let title = format!("{TITLE_SLOT_PREFIX}{rest}");
|
||||
let (n, role) = rest
|
||||
.split_once(':')
|
||||
.and_then(|(n, role)| Some((n.parse::<u32>().ok().filter(|n| *n > 0)?, role)))
|
||||
.ok_or_else(|| {
|
||||
anyhow::anyhow!(
|
||||
"node {id} titled {title:?}: expected {TITLE_SLOT_PREFIX}<n>:high or :low (n from 1)"
|
||||
)
|
||||
})?;
|
||||
let inputs = inputs_of(&workflow, id);
|
||||
if !inputs.get("lora_name").is_some_and(Value::is_string)
|
||||
|| !inputs.get("strength_model").is_some_and(Value::is_number)
|
||||
{
|
||||
bail!(
|
||||
"node {id} titled {title:?} needs a lora_name and a numeric strength_model \
|
||||
(use a LoraLoaderModelOnly or LoraLoader)"
|
||||
);
|
||||
}
|
||||
let slot = slot_map.entry(n).or_insert(Slot {
|
||||
high: Vec::new(),
|
||||
low: Vec::new(),
|
||||
});
|
||||
match role {
|
||||
"high" => slot.high.push((id.clone(), exported(id))),
|
||||
"low" => slot.low.push((id.clone(), exported(id))),
|
||||
_ => bail!("node {id} titled {title:?}: the role must be high or low"),
|
||||
}
|
||||
}
|
||||
for (n, slot) in &slot_map {
|
||||
if slot.high.is_empty() || slot.low.is_empty() {
|
||||
bail!("LoRA slot {n} needs both a :high and a :low node");
|
||||
}
|
||||
}
|
||||
if let Some((last, _)) = slot_map.last_key_value()
|
||||
&& *last as usize != slot_map.len()
|
||||
{
|
||||
bail!(
|
||||
"LoRA slots must be numbered 1..{} without gaps",
|
||||
slot_map.len()
|
||||
);
|
||||
}
|
||||
let slots: Vec<Slot> = slot_map.into_values().collect();
|
||||
|
||||
for (id, node) in &workflow {
|
||||
let class = node["class_type"].as_str().unwrap_or_default();
|
||||
if STRAY_OUTPUT_CLASSES.contains(&class) {
|
||||
@@ -248,6 +335,8 @@ impl Template {
|
||||
prompt,
|
||||
reference_id,
|
||||
loras,
|
||||
slots,
|
||||
fixed_loras,
|
||||
output,
|
||||
warnings,
|
||||
})
|
||||
@@ -265,6 +354,16 @@ impl Template {
|
||||
self.cfg.prompt_mode == PromptMode::Replace
|
||||
}
|
||||
|
||||
/// How many pickable-LoRA slots the workflow has (#220).
|
||||
pub fn slot_count(&self) -> usize {
|
||||
self.slots.len()
|
||||
}
|
||||
|
||||
/// `lora_name`s loaded outside the slots; kept out of the catalog.
|
||||
pub fn fixed_loras(&self) -> &std::collections::HashSet<String> {
|
||||
&self.fixed_loras
|
||||
}
|
||||
|
||||
/// 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> {
|
||||
@@ -288,8 +387,19 @@ impl Template {
|
||||
prompt: Option<&str>,
|
||||
reference: Option<&str>,
|
||||
loras: &BTreeMap<String, LoraRequest>,
|
||||
picks: &[SlotPick],
|
||||
next_seed: &mut dyn FnMut() -> u64,
|
||||
) -> Result<Filled, FillError> {
|
||||
if picks.len() > self.slots.len() {
|
||||
return Err(FillError(if self.slots.is_empty() {
|
||||
"this template doesn't take picked LoRAs".to_string()
|
||||
} else {
|
||||
format!(
|
||||
"this template takes at most {} picked LoRAs",
|
||||
self.slots.len()
|
||||
)
|
||||
}));
|
||||
}
|
||||
if let Some(unknown) = loras
|
||||
.keys()
|
||||
.find(|id| !self.loras.iter().any(|l| &l.cfg.id == *id))
|
||||
@@ -362,6 +472,26 @@ impl Template {
|
||||
}
|
||||
}
|
||||
|
||||
// Picked LoRAs fill the slots in order; the rest run at 0 (their
|
||||
// exported lora_name stays, a file ComfyUI knows).
|
||||
for (i, slot) in self.slots.iter().enumerate() {
|
||||
let pick = picks.get(i).filter(|p| p.scale > 0.0);
|
||||
for (nodes, file) in [
|
||||
(&slot.high, pick.map(|p| p.high.as_str())),
|
||||
(&slot.low, pick.map(|p| p.low.as_str())),
|
||||
] {
|
||||
for (id, exported) in nodes {
|
||||
if let Some(file) = file {
|
||||
set_input(&mut wf, id, "lora_name", json!(file));
|
||||
}
|
||||
let scale = pick.map_or(0.0, |p| p.scale);
|
||||
for (input, value) in exported {
|
||||
set_input(&mut wf, id, input, json!(value * scale));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for node in wf.values_mut() {
|
||||
let Some(inputs) = node.get_mut("inputs").and_then(Value::as_object_mut) else {
|
||||
continue;
|
||||
@@ -435,6 +565,17 @@ mod tests {
|
||||
/// 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");
|
||||
/// ANIMATE_LORAS plus two pickable slots chained after the switch:
|
||||
/// 60/61 (slot 1 high/low), 62/63 (slot 2).
|
||||
pub const ANIMATE_SLOTS: &str = include_str!("../../tests/fixtures/enhance/animate_slots.json");
|
||||
|
||||
fn pick(name: &str, scale: f64) -> SlotPick {
|
||||
SlotPick {
|
||||
high: format!("{name}_high.safetensors"),
|
||||
low: format!("{name}_low.safetensors"),
|
||||
scale,
|
||||
}
|
||||
}
|
||||
|
||||
fn lora(id: &str, default: bool) -> LoraConfig {
|
||||
LoraConfig {
|
||||
@@ -475,6 +616,7 @@ mod tests {
|
||||
prompt_join: " ".into(),
|
||||
output_megapixels: 1.0,
|
||||
loras: Vec::new(),
|
||||
catalog: Default::default(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -547,6 +689,7 @@ mod tests {
|
||||
Some("more"),
|
||||
None,
|
||||
&BTreeMap::new(),
|
||||
&[],
|
||||
&mut seeds(),
|
||||
)
|
||||
.unwrap();
|
||||
@@ -566,7 +709,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, &none(), &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);
|
||||
@@ -579,8 +722,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, &none(), &mut rng).unwrap();
|
||||
let y = t.fill("p", None, None, &none(), &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"));
|
||||
}
|
||||
|
||||
@@ -606,7 +749,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, &none(), &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));
|
||||
}
|
||||
}
|
||||
@@ -615,7 +760,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, &none(), &mut seeds()).unwrap();
|
||||
let f = t.fill("p", None, None, &none(), &[], &mut seeds()).unwrap();
|
||||
assert_eq!(input(&f.prompt, "36", "require_approval"), true);
|
||||
}
|
||||
|
||||
@@ -633,6 +778,7 @@ mod tests {
|
||||
Some(" warmer "),
|
||||
None,
|
||||
&BTreeMap::new(),
|
||||
&[],
|
||||
&mut seeds(),
|
||||
)
|
||||
.unwrap();
|
||||
@@ -641,10 +787,10 @@ mod tests {
|
||||
"Increase the clarity\n\nAdditional instructions: warmer"
|
||||
);
|
||||
let f = t
|
||||
.fill("p", Some(" "), None, &none(), &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, &none(), &mut seeds()).unwrap();
|
||||
let f = t.fill("p", None, None, &none(), &[], &mut seeds()).unwrap();
|
||||
assert_eq!(input(&f.prompt, "13", "value"), base);
|
||||
}
|
||||
|
||||
@@ -657,6 +803,7 @@ mod tests {
|
||||
Some("Make it warmer."),
|
||||
None,
|
||||
&BTreeMap::new(),
|
||||
&[],
|
||||
&mut seeds(),
|
||||
)
|
||||
.unwrap();
|
||||
@@ -667,10 +814,12 @@ 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, &none(), &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, &none(), &mut seeds())
|
||||
t.fill("p", Some(" "), None, &none(), &[], &mut seeds())
|
||||
.is_err()
|
||||
);
|
||||
let f = t
|
||||
@@ -679,6 +828,7 @@ mod tests {
|
||||
Some("make it black and white"),
|
||||
None,
|
||||
&BTreeMap::new(),
|
||||
&[],
|
||||
&mut seeds(),
|
||||
)
|
||||
.unwrap();
|
||||
@@ -692,13 +842,16 @@ 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, &none(), &mut seeds()).is_err());
|
||||
assert!(t.fill("p", None, None, &none(), &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"), &none(), &mut seeds())
|
||||
.fill("p", None, Some("r"), &none(), &[], &mut seeds())
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
@@ -712,6 +865,7 @@ mod tests {
|
||||
Some("x"),
|
||||
Some("Z:\\ref.jpg"),
|
||||
&BTreeMap::new(),
|
||||
&[],
|
||||
&mut seeds(),
|
||||
)
|
||||
.unwrap();
|
||||
@@ -724,7 +878,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, &none(), &mut seeds())
|
||||
.fill("src", Some("x"), None, &none(), &[], &mut seeds())
|
||||
.unwrap();
|
||||
assert!(f.dropped_reference);
|
||||
let wf = f.prompt.as_object().unwrap();
|
||||
@@ -751,7 +905,7 @@ 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, &none(), &mut seeds()).unwrap();
|
||||
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));
|
||||
@@ -767,6 +921,7 @@ mod tests {
|
||||
None,
|
||||
None,
|
||||
&reqs(&[("style", LoraRequest::On(true))]),
|
||||
&[],
|
||||
&mut seeds(),
|
||||
)
|
||||
.unwrap();
|
||||
@@ -791,6 +946,7 @@ mod tests {
|
||||
None,
|
||||
None,
|
||||
&reqs(&[("style", LoraRequest::Scale(0.5))]),
|
||||
&[],
|
||||
&mut seeds(),
|
||||
)
|
||||
.unwrap();
|
||||
@@ -805,6 +961,7 @@ mod tests {
|
||||
None,
|
||||
None,
|
||||
&reqs(&[("style", LoraRequest::Scale(9.0))]),
|
||||
&[],
|
||||
&mut seeds(),
|
||||
)
|
||||
.unwrap();
|
||||
@@ -815,6 +972,7 @@ mod tests {
|
||||
None,
|
||||
None,
|
||||
&reqs(&[("style", LoraRequest::Scale(0.0))]),
|
||||
&[],
|
||||
&mut seeds(),
|
||||
)
|
||||
.unwrap();
|
||||
@@ -826,7 +984,7 @@ mod tests {
|
||||
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();
|
||||
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));
|
||||
|
||||
@@ -836,6 +994,7 @@ mod tests {
|
||||
None,
|
||||
None,
|
||||
&reqs(&[("nope", LoraRequest::On(true))]),
|
||||
&[],
|
||||
&mut seeds(),
|
||||
)
|
||||
.unwrap_err();
|
||||
@@ -850,7 +1009,7 @@ mod tests {
|
||||
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();
|
||||
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(
|
||||
@@ -858,6 +1017,7 @@ mod tests {
|
||||
None,
|
||||
None,
|
||||
&reqs(&[("style", LoraRequest::Scale(0.5))]),
|
||||
&[],
|
||||
&mut seeds(),
|
||||
)
|
||||
.unwrap();
|
||||
@@ -874,7 +1034,7 @@ mod tests {
|
||||
"{:?}",
|
||||
t.warnings
|
||||
);
|
||||
let f = t.fill("p", None, None, &none(), &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!(
|
||||
@@ -883,7 +1043,8 @@ mod tests {
|
||||
None,
|
||||
None,
|
||||
&reqs(&[("style", LoraRequest::On(false))]),
|
||||
&mut seeds()
|
||||
&[],
|
||||
&mut seeds(),
|
||||
)
|
||||
.is_err()
|
||||
);
|
||||
@@ -917,4 +1078,106 @@ mod tests {
|
||||
.to_string();
|
||||
assert!(err.contains("strength_model"), "{err}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn slots_are_found_and_fixed_loras_listed() {
|
||||
let t = load(ANIMATE_SLOTS, PromptMode::Append).unwrap();
|
||||
assert_eq!(t.slot_count(), 2);
|
||||
let mut fixed: Vec<_> = t.fixed_loras().iter().cloned().collect();
|
||||
fixed.sort();
|
||||
assert_eq!(
|
||||
fixed,
|
||||
[
|
||||
"example_style_high.safetensors",
|
||||
"example_style_low.safetensors",
|
||||
"wan2.2_i2v_lightx2v_4steps_lora_v1_high_noise.safetensors",
|
||||
"wan2.2_i2v_lightx2v_4steps_lora_v1_low_noise.safetensors",
|
||||
]
|
||||
);
|
||||
assert_eq!(
|
||||
load(ANIMATE_LORAS, PromptMode::Append)
|
||||
.unwrap()
|
||||
.slot_count(),
|
||||
0
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn picks_fill_slots_in_order_and_unused_slots_run_at_zero() {
|
||||
let t = load(ANIMATE_SLOTS, PromptMode::Append).unwrap();
|
||||
let f = t
|
||||
.fill(
|
||||
"p",
|
||||
None,
|
||||
None,
|
||||
&none(),
|
||||
&[pick("motion", 0.5)],
|
||||
&mut seeds(),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
input(&f.prompt, "60", "lora_name"),
|
||||
&json!("motion_high.safetensors")
|
||||
);
|
||||
assert_eq!(
|
||||
input(&f.prompt, "61", "lora_name"),
|
||||
&json!("motion_low.safetensors")
|
||||
);
|
||||
assert_eq!(input(&f.prompt, "60", "strength_model"), &json!(0.5));
|
||||
assert_eq!(input(&f.prompt, "61", "strength_model"), &json!(0.5));
|
||||
// Slot 2 unused: its exported file stays, at strength 0.
|
||||
assert_eq!(
|
||||
input(&f.prompt, "62", "lora_name"),
|
||||
&json!("example_slot_high.safetensors")
|
||||
);
|
||||
assert_eq!(input(&f.prompt, "62", "strength_model"), &json!(0.0));
|
||||
assert_eq!(input(&f.prompt, "63", "strength_model"), &json!(0.0));
|
||||
// The chain isn't rewired.
|
||||
assert_eq!(input(&f.prompt, "14", "model"), &json!(["62", 0]));
|
||||
|
||||
// No picks at all: every slot off.
|
||||
let f = t.fill("p", None, None, &none(), &[], &mut seeds()).unwrap();
|
||||
assert_eq!(input(&f.prompt, "60", "strength_model"), &json!(0.0));
|
||||
|
||||
let err = t
|
||||
.fill(
|
||||
"p",
|
||||
None,
|
||||
None,
|
||||
&none(),
|
||||
&[pick("a", 1.0), pick("b", 1.0), pick("c", 1.0)],
|
||||
&mut seeds(),
|
||||
)
|
||||
.unwrap_err();
|
||||
assert!(err.0.contains("at most 2"), "{err}");
|
||||
let plain = load(ANIMATE_LORAS, PromptMode::Append).unwrap();
|
||||
assert!(
|
||||
plain
|
||||
.fill("p", None, None, &none(), &[pick("a", 1.0)], &mut seeds())
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn malformed_slots_are_rejected() {
|
||||
let only_high = ANIMATE_SLOTS.replace("imageapi:lora-slot:2:low", "something else");
|
||||
let err = load(&only_high, PromptMode::Append)
|
||||
.unwrap_err()
|
||||
.to_string();
|
||||
assert!(err.contains("slot 2 needs both"), "{err}");
|
||||
let gap = ANIMATE_SLOTS
|
||||
.replace("imageapi:lora-slot:2:high", "imageapi:lora-slot:3:high")
|
||||
.replace("imageapi:lora-slot:2:low", "imageapi:lora-slot:3:low");
|
||||
let err = load(&gap, PromptMode::Append).unwrap_err().to_string();
|
||||
assert!(err.contains("without gaps"), "{err}");
|
||||
let role = ANIMATE_SLOTS.replace("imageapi:lora-slot:1:high", "imageapi:lora-slot:1:mid");
|
||||
assert!(
|
||||
load(&role, PromptMode::Append)
|
||||
.unwrap_err()
|
||||
.to_string()
|
||||
.contains("high or low")
|
||||
);
|
||||
let no_n = ANIMATE_SLOTS.replace("imageapi:lora-slot:1:high", "imageapi:lora-slot:x:high");
|
||||
assert!(load(&no_n, PromptMode::Append).is_err());
|
||||
}
|
||||
}
|
||||
+404
@@ -0,0 +1,404 @@
|
||||
{
|
||||
"1": {
|
||||
"class_type": "ImageApiLoadImageFromPath",
|
||||
"_meta": {
|
||||
"title": "imageapi:source"
|
||||
},
|
||||
"inputs": {
|
||||
"path": ""
|
||||
}
|
||||
},
|
||||
"2": {
|
||||
"class_type": "ImageApiTargetSize",
|
||||
"_meta": {
|
||||
"title": "Target Size (~0.92 MP, keeps aspect)"
|
||||
},
|
||||
"inputs": {
|
||||
"image": [
|
||||
"1",
|
||||
0
|
||||
],
|
||||
"max_megapixels": 0.92,
|
||||
"min_megapixels": 0.92,
|
||||
"multiple_of": 16
|
||||
}
|
||||
},
|
||||
"10": {
|
||||
"class_type": "UNETLoader",
|
||||
"_meta": {
|
||||
"title": "Wan 2.2 i2v high noise"
|
||||
},
|
||||
"inputs": {
|
||||
"unet_name": "wan2.2_i2v_high_noise_14B_fp8_scaled.safetensors",
|
||||
"weight_dtype": "default"
|
||||
}
|
||||
},
|
||||
"11": {
|
||||
"class_type": "UNETLoader",
|
||||
"_meta": {
|
||||
"title": "Wan 2.2 i2v low noise"
|
||||
},
|
||||
"inputs": {
|
||||
"unet_name": "wan2.2_i2v_low_noise_14B_fp8_scaled.safetensors",
|
||||
"weight_dtype": "default"
|
||||
}
|
||||
},
|
||||
"12": {
|
||||
"class_type": "LoraLoaderModelOnly",
|
||||
"_meta": {
|
||||
"title": "lightx2v 4-step (high)"
|
||||
},
|
||||
"inputs": {
|
||||
"model": [
|
||||
"10",
|
||||
0
|
||||
],
|
||||
"lora_name": "wan2.2_i2v_lightx2v_4steps_lora_v1_high_noise.safetensors",
|
||||
"strength_model": 1.0
|
||||
}
|
||||
},
|
||||
"13": {
|
||||
"class_type": "LoraLoaderModelOnly",
|
||||
"_meta": {
|
||||
"title": "lightx2v 4-step (low)"
|
||||
},
|
||||
"inputs": {
|
||||
"model": [
|
||||
"11",
|
||||
0
|
||||
],
|
||||
"lora_name": "wan2.2_i2v_lightx2v_4steps_lora_v1_low_noise.safetensors",
|
||||
"strength_model": 1.0
|
||||
}
|
||||
},
|
||||
"14": {
|
||||
"class_type": "ModelSamplingSD3",
|
||||
"_meta": {
|
||||
"title": "Shift (high)"
|
||||
},
|
||||
"inputs": {
|
||||
"model": [
|
||||
"62",
|
||||
0
|
||||
],
|
||||
"shift": 5.0
|
||||
}
|
||||
},
|
||||
"15": {
|
||||
"class_type": "ModelSamplingSD3",
|
||||
"_meta": {
|
||||
"title": "Shift (low)"
|
||||
},
|
||||
"inputs": {
|
||||
"model": [
|
||||
"63",
|
||||
0
|
||||
],
|
||||
"shift": 5.0
|
||||
}
|
||||
},
|
||||
"40": {
|
||||
"class_type": "PathchSageAttentionKJ",
|
||||
"_meta": {
|
||||
"title": "Sage attention (high)"
|
||||
},
|
||||
"inputs": {
|
||||
"model": [
|
||||
"14",
|
||||
0
|
||||
],
|
||||
"sage_attention": "auto"
|
||||
}
|
||||
},
|
||||
"41": {
|
||||
"class_type": "PathchSageAttentionKJ",
|
||||
"_meta": {
|
||||
"title": "Sage attention (low)"
|
||||
},
|
||||
"inputs": {
|
||||
"model": [
|
||||
"15",
|
||||
0
|
||||
],
|
||||
"sage_attention": "auto"
|
||||
}
|
||||
},
|
||||
"16": {
|
||||
"class_type": "CLIPLoader",
|
||||
"_meta": {
|
||||
"title": "Load CLIP"
|
||||
},
|
||||
"inputs": {
|
||||
"clip_name": "umt5_xxl_fp8_e4m3fn_scaled.safetensors",
|
||||
"type": "wan",
|
||||
"device": "default"
|
||||
}
|
||||
},
|
||||
"17": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"_meta": {
|
||||
"title": "imageapi:prompt"
|
||||
},
|
||||
"inputs": {
|
||||
"clip": [
|
||||
"16",
|
||||
0
|
||||
],
|
||||
"text": "natural, lifelike motion with a slow, gentle camera movement, realistic, cinematic"
|
||||
}
|
||||
},
|
||||
"18": {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"_meta": {
|
||||
"title": "Negative"
|
||||
},
|
||||
"inputs": {
|
||||
"clip": [
|
||||
"16",
|
||||
0
|
||||
],
|
||||
"text": "static, frozen, blurry, low quality, distorted, warped faces, morphing, flicker, watermark, text"
|
||||
}
|
||||
},
|
||||
"19": {
|
||||
"class_type": "VAELoader",
|
||||
"_meta": {
|
||||
"title": "Load VAE"
|
||||
},
|
||||
"inputs": {
|
||||
"vae_name": "wan_2.1_vae.safetensors"
|
||||
}
|
||||
},
|
||||
"21": {
|
||||
"class_type": "WanImageToVideo",
|
||||
"_meta": {
|
||||
"title": "Wan Image To Video (81 frames)"
|
||||
},
|
||||
"inputs": {
|
||||
"positive": [
|
||||
"17",
|
||||
0
|
||||
],
|
||||
"negative": [
|
||||
"18",
|
||||
0
|
||||
],
|
||||
"vae": [
|
||||
"19",
|
||||
0
|
||||
],
|
||||
"start_image": [
|
||||
"1",
|
||||
0
|
||||
],
|
||||
"width": [
|
||||
"2",
|
||||
0
|
||||
],
|
||||
"height": [
|
||||
"2",
|
||||
1
|
||||
],
|
||||
"length": 81,
|
||||
"batch_size": 1
|
||||
}
|
||||
},
|
||||
"22": {
|
||||
"class_type": "KSamplerAdvanced",
|
||||
"_meta": {
|
||||
"title": "Sampler (high, steps 0-2)"
|
||||
},
|
||||
"inputs": {
|
||||
"model": [
|
||||
"40",
|
||||
0
|
||||
],
|
||||
"positive": [
|
||||
"21",
|
||||
0
|
||||
],
|
||||
"negative": [
|
||||
"21",
|
||||
1
|
||||
],
|
||||
"latent_image": [
|
||||
"21",
|
||||
2
|
||||
],
|
||||
"add_noise": "enable",
|
||||
"noise_seed": 1,
|
||||
"steps": 4,
|
||||
"cfg": 1.0,
|
||||
"sampler_name": "euler",
|
||||
"scheduler": "simple",
|
||||
"start_at_step": 0,
|
||||
"end_at_step": 2,
|
||||
"return_with_leftover_noise": "enable"
|
||||
}
|
||||
},
|
||||
"23": {
|
||||
"class_type": "KSamplerAdvanced",
|
||||
"_meta": {
|
||||
"title": "Sampler (low, steps 2-4)"
|
||||
},
|
||||
"inputs": {
|
||||
"model": [
|
||||
"41",
|
||||
0
|
||||
],
|
||||
"positive": [
|
||||
"21",
|
||||
0
|
||||
],
|
||||
"negative": [
|
||||
"21",
|
||||
1
|
||||
],
|
||||
"latent_image": [
|
||||
"22",
|
||||
0
|
||||
],
|
||||
"add_noise": "disable",
|
||||
"noise_seed": 0,
|
||||
"steps": 4,
|
||||
"cfg": 1.0,
|
||||
"sampler_name": "euler",
|
||||
"scheduler": "simple",
|
||||
"start_at_step": 2,
|
||||
"end_at_step": 10000,
|
||||
"return_with_leftover_noise": "disable"
|
||||
}
|
||||
},
|
||||
"24": {
|
||||
"class_type": "VAEDecode",
|
||||
"_meta": {
|
||||
"title": "VAE Decode"
|
||||
},
|
||||
"inputs": {
|
||||
"samples": [
|
||||
"23",
|
||||
0
|
||||
],
|
||||
"vae": [
|
||||
"19",
|
||||
0
|
||||
]
|
||||
}
|
||||
},
|
||||
"25": {
|
||||
"class_type": "CreateVideo",
|
||||
"_meta": {
|
||||
"title": "Create Video (16 fps)"
|
||||
},
|
||||
"inputs": {
|
||||
"images": [
|
||||
"24",
|
||||
0
|
||||
],
|
||||
"fps": 16.0
|
||||
}
|
||||
},
|
||||
"36": {
|
||||
"class_type": "ImageApiSaveVideoVersion",
|
||||
"_meta": {
|
||||
"title": "imageapi:save"
|
||||
},
|
||||
"inputs": {
|
||||
"video": [
|
||||
"25",
|
||||
0
|
||||
],
|
||||
"source": [
|
||||
"1",
|
||||
1
|
||||
],
|
||||
"kind": "animated",
|
||||
"subfolder": "aiedits",
|
||||
"require_approval": true
|
||||
}
|
||||
},
|
||||
"50": {
|
||||
"inputs": {
|
||||
"lora_name": "example_style_high.safetensors",
|
||||
"strength_model": 1.0,
|
||||
"model": [
|
||||
"12",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "LoraLoaderModelOnly",
|
||||
"_meta": {
|
||||
"title": "imageapi:lora:style"
|
||||
}
|
||||
},
|
||||
"51": {
|
||||
"inputs": {
|
||||
"lora_name": "example_style_low.safetensors",
|
||||
"strength_model": 1.0,
|
||||
"model": [
|
||||
"13",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "LoraLoaderModelOnly",
|
||||
"_meta": {
|
||||
"title": "imageapi:lora:style"
|
||||
}
|
||||
},
|
||||
"60": {
|
||||
"inputs": {
|
||||
"lora_name": "example_slot_high.safetensors",
|
||||
"strength_model": 1.0,
|
||||
"model": [
|
||||
"50",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "LoraLoaderModelOnly",
|
||||
"_meta": {
|
||||
"title": "imageapi:lora-slot:1:high"
|
||||
}
|
||||
},
|
||||
"61": {
|
||||
"inputs": {
|
||||
"lora_name": "example_slot_low.safetensors",
|
||||
"strength_model": 1.0,
|
||||
"model": [
|
||||
"51",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "LoraLoaderModelOnly",
|
||||
"_meta": {
|
||||
"title": "imageapi:lora-slot:1:low"
|
||||
}
|
||||
},
|
||||
"62": {
|
||||
"inputs": {
|
||||
"lora_name": "example_slot_high.safetensors",
|
||||
"strength_model": 1.0,
|
||||
"model": [
|
||||
"60",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "LoraLoaderModelOnly",
|
||||
"_meta": {
|
||||
"title": "imageapi:lora-slot:2:high"
|
||||
}
|
||||
},
|
||||
"63": {
|
||||
"inputs": {
|
||||
"lora_name": "example_slot_low.safetensors",
|
||||
"strength_model": 1.0,
|
||||
"model": [
|
||||
"61",
|
||||
0
|
||||
]
|
||||
},
|
||||
"class_type": "LoraLoaderModelOnly",
|
||||
"_meta": {
|
||||
"title": "imageapi:lora-slot:2:low"
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user