ImageApi/src/enhance/template.rs

//! Workflow templates (#192): load, validate, and fill in an API-format
//! ComfyUI workflow for one photo.
//!
//! 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::{LoraConfig, LoraRequest, 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:";
/// `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
/// (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"];

/// What a template produces, from its `imageapi:save` node's class.
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
#[serde(rename_all = "lowercase")]
pub enum OutputKind {
    /// Save Image Version: PNG versions.
    Image,
    /// Save Video Version: an MP4 (e.g. a clip animated from the photo).
    Video,
}

/// The save node class that makes a template's output a video.
pub const VIDEO_SAVE_CLASS: &str = "ImageApiSaveVideoVersion";

/// Output nodes that write into ComfyUI's output folder on every API run.
const STRAY_OUTPUT_CLASSES: &[&str] = &[
    "SaveImage",
    "SaveAnimatedWEBP",
    "SaveAnimatedPNG",
    "SaveWEBM",
    "SaveVideo",
    "SaveGLB",
];

#[derive(Debug, Clone)]
pub struct Template {
    pub cfg: TemplateConfig,
    workflow: Map<String, Value>,
    source_id: String,
    save_id: String,
    /// `(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>,
    /// 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).
    pub warnings: Vec<String>,
}

#[derive(Debug, Clone)]
struct LoraSwitch {
    cfg: LoraConfig,
    /// 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)>)>,
}

/// 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)]
pub struct FillError(pub String);

impl std::fmt::Display for FillError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(&self.0)
    }
}

/// A filled-in prompt plus whether the reference node had to be removed
/// (so a ComfyUI rejection can be reported as "needs a reference photo").
#[derive(Debug)]
pub struct Filled {
    pub prompt: Value,
    pub dropped_reference: 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 {
    pub fn load(cfg: &TemplateConfig, base_dir: &Path) -> Result<Template> {
        let path = if cfg.file.is_absolute() {
            cfg.file.clone()
        } else {
            base_dir.join(&cfg.file)
        };
        let text = std::fs::read_to_string(&path)
            .with_context(|| format!("reading workflow {}", path.display()))?;
        let value: Value = serde_json::from_str(&text)
            .with_context(|| format!("parsing workflow {}", path.display()))?;
        Self::from_value(cfg.clone(), value)
    }

    pub fn from_value(cfg: TemplateConfig, value: Value) -> Result<Template> {
        let Value::Object(workflow) = value else {
            bail!("workflow is not a JSON object (is it the API-format export?)");
        };
        for (id, node) in &workflow {
            if node.get("class_type").and_then(Value::as_str).is_none()
                || !node.get("inputs").is_some_and(Value::is_object)
            {
                bail!(
                    "node {id} has no class_type/inputs: this looks like a UI-format workflow, \
                     export with Workflow -> Export (API)"
                );
            }
        }
        let find = |title: &str| -> Result<Option<String>> {
            let ids: Vec<&String> = workflow
                .iter()
                .filter(|(_, n)| title_of(n) == Some(title))
                .map(|(id, _)| id)
                .collect();
            match ids.len() {
                0 => Ok(None),
                1 => Ok(Some(ids[0].clone())),
                _ => bail!("more than one node is titled {title:?}"),
            }
        };
        let Some(source_id) = find(TITLE_SOURCE)? else {
            bail!("no node titled {TITLE_SOURCE:?} (the Load Image From Path that gets the photo)");
        };
        let Some(save_id) = find(TITLE_SAVE)? else {
            bail!("no node titled {TITLE_SAVE:?} (the Save Image Version node)");
        };
        require_input(&workflow, &source_id, "path", TITLE_SOURCE)?;

        let prompt = match find(TITLE_PROMPT)? {
            None => None,
            Some(id) => {
                let inputs = inputs_of(&workflow, &id);
                let Some(name) = PROMPT_INPUTS
                    .iter()
                    .copied()
                    .find(|n| inputs.get(*n).is_some_and(Value::is_string))
                else {
                    bail!(
                        "node {TITLE_PROMPT:?} has none of the inputs {PROMPT_INPUTS:?} \
                         as a plain string (a linked input can't take the user's text)"
                    );
                };
                Some((id, name))
            }
        };
        if cfg.prompt_mode == PromptMode::Replace && prompt.is_none() {
            bail!("prompt_mode = \"replace\" needs a node titled {TITLE_PROMPT:?}");
        }
        let reference_id = find(TITLE_REFERENCE)?;
        if let Some(id) = &reference_id {
            require_input(&workflow, id, "path", TITLE_REFERENCE)?;
        }

        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 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,
                    });
                }
                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}")
            ));
        }

        // 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) {
                warnings.push(format!(
                    "node {id} ({class}) also writes into ComfyUI's output folder on every run"
                ));
            }
        }
        let output = if workflow[&save_id]["class_type"].as_str() == Some(VIDEO_SAVE_CLASS) {
            OutputKind::Video
        } else {
            OutputKind::Image
        };
        Ok(Template {
            cfg,
            workflow,
            source_id,
            save_id,
            prompt,
            reference_id,
            loras,
            slots,
            fixed_loras,
            output,
            warnings,
        })
    }

    pub fn accepts_prompt(&self) -> bool {
        self.prompt.is_some()
    }

    pub fn accepts_reference(&self) -> bool {
        self.reference_id.is_some()
    }

    pub fn prompt_required(&self) -> bool {
        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> {
        self.loras.iter().map(|l| &l.cfg)
    }

    /// Fill the workflow for one photo.
    ///
    /// - `source` / `reference` are paths as ComfyUI sees them.
    /// - Without a reference the reference node and every link to it are
    ///   removed: its leftover export path must never be used.
    /// - 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` 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, 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))
        {
            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(_)) => {
                return Err(FillError("this template doesn't take a prompt".to_string()));
            }
            (Some(_), None) if self.prompt_required() => {
                return Err(FillError(
                    "this template needs an instruction (the prompt can't be empty)".to_string(),
                ));
            }
            _ => {}
        }
        if reference.is_some() && self.reference_id.is_none() {
            return Err(FillError(
                "this template doesn't take a reference photo".to_string(),
            ));
        }

        let mut wf = self.workflow.clone();
        set_input(&mut wf, &self.source_id, "path", json!(source));
        set_input(&mut wf, &self.save_id, "require_approval", json!(true));

        if let (Some((id, input)), Some(text)) = (&self.prompt, user_text) {
            let base = inputs_of(&wf, id)
                .get(*input)
                .and_then(Value::as_str)
                .unwrap_or_default()
                .to_string();
            let combined = match self.cfg.prompt_mode {
                PromptMode::Replace => text.to_string(),
                PromptMode::Append if base.trim().is_empty() => text.to_string(),
                PromptMode::Append => format!("{base}{}{text}", self.cfg.prompt_join),
            };
            set_input(&mut wf, id, input, json!(combined));
        }

        let mut dropped_reference = false;
        if let Some(ref_id) = &self.reference_id {
            match reference {
                Some(path) => set_input(&mut wf, ref_id, "path", json!(path)),
                None => {
                    wf.remove(ref_id);
                    for node in wf.values_mut() {
                        if let Some(inputs) = node.get_mut("inputs").and_then(Value::as_object_mut)
                        {
                            inputs.retain(|_, v| !links_to(v, ref_id));
                        }
                    }
                    dropped_reference = true;
                }
            }
        }

        let mut resolved = BTreeMap::new();
        for switch in &self.loras {
            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));
                }
            }
        }

        // 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;
            };
            for (key, value) in inputs.iter_mut() {
                if is_seed_input(key) && value.is_number() {
                    *value = json!(next_seed());
                }
            }
        }
        Ok(Filled {
            prompt: Value::Object(wf),
            dropped_reference,
            loras: resolved,
        })
    }
}

fn title_of(node: &Value) -> Option<&str> {
    node.get("_meta")?.get("title")?.as_str()
}

fn inputs_of<'a>(wf: &'a Map<String, Value>, id: &str) -> &'a Map<String, Value> {
    static EMPTY: std::sync::LazyLock<Map<String, Value>> = std::sync::LazyLock::new(Map::new);
    wf.get(id)
        .and_then(|n| n.get("inputs"))
        .and_then(Value::as_object)
        .unwrap_or(&EMPTY)
}

fn require_input(wf: &Map<String, Value>, id: &str, input: &str, title: &str) -> Result<()> {
    if !inputs_of(wf, id).get(input).is_some_and(Value::is_string) {
        bail!("node {title:?} has no string input {input:?}");
    }
    Ok(())
}

fn set_input(wf: &mut Map<String, Value>, id: &str, input: &str, value: Value) {
    if let Some(inputs) = wf
        .get_mut(id)
        .and_then(|n| n.get_mut("inputs"))
        .and_then(Value::as_object_mut)
    {
        inputs.insert(input.to_string(), value);
    }
}

/// `["<id>", n]`: a link to output `n` of node `id`.
fn links_to(v: &Value, id: &str) -> bool {
    matches!(v.as_array().map(Vec::as_slice), Some([Value::String(s), Value::Number(_)]) if s == id)
}

/// A fresh seed for one job. 32 bits: some nodes cap their seed there
/// (SeedVR2VideoUpscaler's max is 2^32 - 1), and it's plenty of variety for
/// the ones that take 64.
pub fn random_seed() -> u64 {
    u64::from(rand::random::<u32>())
}

fn is_seed_input(key: &str) -> bool {
    key == "seed" || key == "noise_seed" || key.ends_with(".seed")
}

#[cfg(test)]
mod tests {
    use super::*;

    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");
    /// 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 {
            id: id.into(),
            label: id.into(),
            default,
            strength: 1.0,
            min: 0.0,
            max: 1.5,
        }
    }

    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 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(),
            label: name.into(),
            description: String::new(),
            file: "x.json".into(),
            prompt_mode: mode,
            prompt_join: " ".into(),
            output_megapixels: 1.0,
            loras: Vec::new(),
            catalog: Default::default(),
        }
    }

    fn load(json: &str, mode: PromptMode) -> Result<Template> {
        Template::from_value(cfg("t", mode), serde_json::from_str(json).unwrap())
    }

    fn seeds() -> impl FnMut() -> u64 {
        let mut n = 1000;
        move || {
            n += 1;
            n
        }
    }

    fn input<'a>(p: &'a Value, id: &str, name: &str) -> &'a Value {
        &p[id]["inputs"][name]
    }

    #[test]
    fn real_templates_load() {
        let restore = load(RESTORE, PromptMode::Append).unwrap();
        assert!(restore.accepts_prompt());
        assert!(!restore.accepts_reference());
        assert!(!restore.prompt_required());
        let edit = load(EDIT, PromptMode::Replace).unwrap();
        assert!(edit.accepts_prompt() && edit.accepts_reference() && edit.prompt_required());
        assert!(load(DEV, PromptMode::Append).is_ok());
        // The prompt node is a String Multiline primitive: input `value`.
        assert_eq!(restore.prompt.as_ref().unwrap().1, "value");
        assert!(restore.warnings.is_empty());
    }

    #[test]
    fn missing_required_titles_are_rejected_with_a_clear_message() {
        let no_source = RESTORE.replace("imageapi:source", "something else");
        let err = load(&no_source, PromptMode::Append)
            .unwrap_err()
            .to_string();
        assert!(err.contains("imageapi:source"), "{err}");
        let no_save = RESTORE.replace("imageapi:save", "something else");
        let err = load(&no_save, PromptMode::Append).unwrap_err().to_string();
        assert!(err.contains("imageapi:save"), "{err}");
    }

    #[test]
    fn rejects_ui_format_duplicates_and_bad_prompt_nodes() {
        let ui = r#"{"nodes": [], "links": []}"#;
        assert!(load(ui, PromptMode::Append).is_err());
        let dup = RESTORE.replace("\"Load CLIP\"", "\"imageapi:source\"");
        assert!(load(&dup, PromptMode::Append).is_err());
        // Prompt node with none of text / prompt / value.
        let mut v: Value = serde_json::from_str(RESTORE).unwrap();
        v["13"]["inputs"] = json!({"other": "x"});
        assert!(Template::from_value(cfg("t", PromptMode::Append), v).is_err());
        // replace without a prompt node
        let no_prompt = RESTORE.replace("imageapi:prompt", "nope");
        assert!(load(&no_prompt, PromptMode::Replace).is_err());
        assert!(load(&no_prompt, PromptMode::Append).is_ok());
    }

    #[test]
    fn prompt_node_accepts_text_or_prompt_inputs() {
        let mut v: Value = serde_json::from_str(RESTORE).unwrap();
        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,
                &BTreeMap::new(),
                &[],
                &mut seeds(),
            )
            .unwrap();
        assert_eq!(input(&f.prompt, "13", "prompt"), "base more");
    }

    #[test]
    fn stray_save_image_is_a_warning_not_an_error() {
        let mut v: Value = serde_json::from_str(RESTORE).unwrap();
        v["99"] = json!({"class_type": "SaveImage", "inputs": {"images": ["6", 0]}});
        let t = Template::from_value(cfg("t", PromptMode::Append), v).unwrap();
        assert_eq!(t.warnings.len(), 1);
        assert!(t.warnings[0].contains("SaveImage"));
    }

    #[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, &none(), &[], &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.
        let a = input(&f.prompt, "5", "seed").as_u64().unwrap();
        let b = input(&f.prompt, "14", "sampling_mode.seed")
            .as_u64()
            .unwrap();
        assert!(a > 1000 && b > 1000 && a != b);
        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();
        assert_ne!(input(&x.prompt, "5", "seed"), input(&y.prompt, "5", "seed"));
    }

    #[test]
    fn output_kind_follows_the_save_node() {
        assert_eq!(
            load(RESTORE, PromptMode::Append).unwrap().output,
            OutputKind::Image
        );
        let video = RESTORE.replace(
            "\"ImageApiSaveImageVersion\"",
            "\"ImageApiSaveVideoVersion\"",
        );
        assert_ne!(video, RESTORE);
        assert_eq!(
            load(&video, PromptMode::Append).unwrap().output,
            OutputKind::Video
        );
    }

    #[test]
    fn random_seeds_fit_seedvr2s_u32_limit() {
        // 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();
            assert!(input(&f.prompt, "5", "seed").as_u64().unwrap() <= u64::from(u32::MAX));
        }
    }

    #[test]
    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();
        assert_eq!(input(&f.prompt, "36", "require_approval"), true);
    }

    #[test]
    fn append_uses_the_join_text_and_empty_leaves_the_base() {
        let base = "Increase the clarity";
        let mut c = cfg("t", PromptMode::Append);
        c.prompt_join = "\n\nAdditional instructions: ".into();
        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,
                &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, &none(), &[], &mut seeds())
            .unwrap();
        assert_eq!(input(&f.prompt, "13", "value"), base);
        let f = t.fill("p", None, None, &none(), &[], &mut seeds()).unwrap();
        assert_eq!(input(&f.prompt, "13", "value"), base);
    }

    #[test]
    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,
                &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."));
    }

    #[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();
        assert!(err.0.contains("instruction"));
        assert!(
            t.fill("p", Some("  "), None, &none(), &[], &mut seeds())
                .is_err()
        );
        let f = t
            .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");
    }

    #[test]
    fn prompt_or_reference_the_template_cant_take_are_refused() {
        let mut v: Value = serde_json::from_str(RESTORE).unwrap();
        v.as_object_mut().unwrap().remove("13");
        // 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());

        let restore = load(RESTORE, PromptMode::Append).unwrap();
        assert!(
            restore
                .fill("p", None, Some("r"), &none(), &[], &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"),
                &BTreeMap::new(),
                &[],
                &mut seeds(),
            )
            .unwrap();
        assert!(!f.dropped_reference);
        assert_eq!(input(&f.prompt, "38", "path"), "Z:\\ref.jpg");
        assert_eq!(input(&f.prompt, "36", "reference_1"), &json!(["38", 1]));
    }

    #[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, &none(), &[], &mut seeds())
            .unwrap();
        assert!(f.dropped_reference);
        let wf = f.prompt.as_object().unwrap();
        assert!(!wf.contains_key("38"));
        // No input anywhere still links to node 38, and the leftover export
        // path can't leak.
        assert!(
            !serde_json::to_string(&f.prompt)
                .unwrap()
                .contains("[\"38\"")
        );
        assert!(wf["36"]["inputs"].get("reference_1").is_none());
        // 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, &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));
        }
        // 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 configured strength (1.0 = as exported).
        let on = t
            .fill(
                "p",
                None,
                None,
                &reqs(&[("style", LoraRequest::On(true))]),
                &[],
                &mut seeds(),
            )
            .unwrap();
        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 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,
                &reqs(&[("style", LoraRequest::Scale(0.5))]),
                &[],
                &mut seeds(),
            )
            .unwrap();
        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,
                &reqs(&[("nope", LoraRequest::On(true))]),
                &[],
                &mut seeds(),
            )
            .unwrap_err();
        assert!(err.0.contains("nope"), "{err}");
    }

    #[test]
    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 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(&half.prompt, "50", "strength_clip"), &json!(0.4));
    }

    #[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, &none(), &[], &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,
                &reqs(&[("style", LoraRequest::On(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}");
    }

    #[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());
    }
}