ImageApi/src/enhance/lora_catalog.rs

//! 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());
    }
}