//! Enhance configuration (#192): one `enhance.toml`, next to `llm.toml`. //! //! Path from `ENHANCE_CONFIG_PATH`, else `./enhance.toml`. A missing file at //! the default path means Enhance is simply off (no `url`); a missing //! explicit path is an error. Nothing here is a secret: ComfyUI has no auth //! and stays LAN-only. `enhance.example.toml` is the annotated schema (a //! test parses it). use std::path::{Path, PathBuf}; use anyhow::{Context as _, Result, bail}; use serde::Deserialize; use super::path_map::PathMap; const DEFAULT_PATH: &str = "enhance.toml"; /// How the app's text combines with the template's own prompt. #[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, serde::Serialize, Default)] #[serde(rename_all = "lowercase")] pub enum PromptMode { /// The template's prompt is a full base instruction; the app's text is /// appended after `prompt_join`. Empty text leaves the prompt as is. #[default] Append, /// The app's text *is* the prompt (and is required). Replace, } #[derive(Debug, Clone, PartialEq, Deserialize)] #[serde(default, deny_unknown_fields)] pub struct EnhanceConfig { /// ComfyUI base URL, e.g. `http://192.168.10.218:8188`. Unset = Enhance /// is disabled (`/enhance/templates` reports `available: false`). pub url: Option, /// Per-request timeout for ComfyUI calls. pub timeout_secs: u64, /// How long the `available` probe result is cached. pub probe_ttl_secs: u64, /// Job poll interval. pub poll_secs: u64, /// Idle time (empty queue, newest history entry this old) before /// ImageApi calls `POST /free` once. 0 disables it. pub idle_free_secs: u64, /// An active job whose ComfyUI stays unreachable this long becomes /// `unavailable`. pub unreachable_grace_secs: u64, /// Template preselected by clients; default = the first one. pub default_template: Option, /// Clients warn above this many megapixels in the original. pub warn_above_megapixels: f64, /// Library path prefix <-> ComfyUI path prefix. Longest prefix wins. pub path_map: Vec, #[serde(rename = "template")] pub templates: Vec, } impl Default for EnhanceConfig { fn default() -> Self { Self { url: None, timeout_secs: 15, probe_ttl_secs: 10, poll_secs: 2, idle_free_secs: 300, unreachable_grace_secs: 120, default_template: None, warn_above_megapixels: 1.5, path_map: Vec::new(), templates: Vec::new(), } } } #[derive(Debug, Clone, PartialEq, Deserialize)] #[serde(deny_unknown_fields)] pub struct PathMapping { /// Prefix as ImageApi sees it, e.g. `/media/nas/`. pub from: String, /// Prefix as ComfyUI sees it, e.g. `Z:\` or `\\nas\`. pub to: String, } #[derive(Debug, Clone, PartialEq, Deserialize)] #[serde(deny_unknown_fields)] pub struct TemplateConfig { /// Stable id clients send (`POST /enhance { template }`). pub name: String, pub label: String, #[serde(default)] pub description: String, /// API-format workflow JSON (ComfyUI -> Workflow -> Export (API)). /// Relative paths resolve against the config file's directory. pub file: PathBuf, #[serde(default)] pub prompt_mode: PromptMode, /// Text between the template's prompt and the app's text (append mode). #[serde(default = "default_join")] pub prompt_join: String, /// Approximate output size, for the client's large-original warning. #[serde(default = "default_output_mp")] pub output_megapixels: f64, /// Optional LoRAs the app can switch on or off (#218), each matching the /// workflow nodes titled `imageapi:lora:`. #[serde(rename = "lora", default)] pub loras: Vec, } /// One LoRA switch (`[[template.lora]]`). #[derive(Debug, Clone, PartialEq, Deserialize, serde::Serialize)] #[serde(deny_unknown_fields)] pub struct LoraConfig { /// Matches the nodes titled `imageapi:lora:` (Wan's high and low /// pair share one id); clients send it in `POST /enhance { loras }`. pub id: String, pub label: String, /// On when the request doesn't say. #[serde(default)] pub default: bool, } /// Template names and LoRA ids: non-empty, letters, digits, `_` and `-`. fn valid_id(s: &str) -> bool { !s.is_empty() && s.chars() .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-') } fn default_join() -> String { " ".into() } fn default_output_mp() -> f64 { 1.0 } impl EnhanceConfig { pub fn validate(&self) -> Result<()> { if let Some(url) = &self.url && !(url.starts_with("http://") || url.starts_with("https://")) { bail!("url must start with http:// or https://, got {url:?}"); } for (key, v) in [ ("timeout_secs", self.timeout_secs), ("probe_ttl_secs", self.probe_ttl_secs), ("poll_secs", self.poll_secs), ("unreachable_grace_secs", self.unreachable_grace_secs), ] { if v == 0 { bail!("{key} must be greater than 0"); } } if self.warn_above_megapixels.is_nan() || self.warn_above_megapixels <= 0.0 { bail!("warn_above_megapixels must be greater than 0"); } let mut seen = std::collections::HashSet::new(); for t in &self.templates { if !valid_id(&t.name) { bail!( "template name {:?} must be non-empty and use only letters, digits, '_' and '-'", t.name ); } if !seen.insert(t.name.as_str()) { bail!("duplicate template name {:?}", t.name); } if t.label.trim().is_empty() { bail!("template {:?} needs a label", t.name); } if t.output_megapixels.is_nan() || t.output_megapixels <= 0.0 { bail!("template {:?}: output_megapixels must be > 0", t.name); } let mut ids = std::collections::HashSet::new(); for l in &t.loras { if !valid_id(&l.id) { bail!( "template {:?}: lora id {:?} must be non-empty and use only letters, digits, '_' and '-'", t.name, l.id ); } if !ids.insert(l.id.as_str()) { bail!("template {:?}: duplicate lora id {:?}", t.name, l.id); } if l.label.trim().is_empty() { bail!("template {:?}: lora {:?} needs a label", t.name, l.id); } } } if let Some(d) = &self.default_template && !self.templates.iter().any(|t| &t.name == d) { bail!("default_template {d:?} is not a configured template"); } // Parses every mapping (rejects empty prefixes). PathMap::new(&self.path_map)?; Ok(()) } /// The preselected template's name. pub fn default_name(&self) -> Option<&str> { self.default_template .as_deref() .or_else(|| self.templates.first().map(|t| t.name.as_str())) } } /// Where the config was read from, and what it said. #[derive(Debug)] pub struct Loaded { pub config: EnhanceConfig, /// Directory relative template `file`s resolve against. pub base_dir: PathBuf, /// The file, when one was read. pub path: Option, } /// Read `enhance.toml` (see module docs). A bad file is an error so a typo /// fails startup instead of silently disabling the feature. pub fn load() -> Result { let (path, explicit) = match std::env::var("ENHANCE_CONFIG_PATH") { Ok(p) if !p.trim().is_empty() => (PathBuf::from(p.trim()), true), _ => (PathBuf::from(DEFAULT_PATH), false), }; match std::fs::read_to_string(&path) { Ok(text) => { let config = parse(&text).with_context(|| format!("in {}", path.display()))?; let base_dir = path .parent() .filter(|p| !p.as_os_str().is_empty()) .map(Path::to_path_buf) .unwrap_or_else(|| PathBuf::from(".")); Ok(Loaded { config, base_dir, path: Some(path), }) } Err(e) if e.kind() == std::io::ErrorKind::NotFound && !explicit => Ok(Loaded { config: EnhanceConfig::default(), base_dir: PathBuf::from("."), path: None, }), Err(e) => Err(e).with_context(|| format!("reading enhance config {}", path.display())), } } pub fn parse(text: &str) -> Result { let config: EnhanceConfig = toml::from_str(text)?; config.validate()?; Ok(config) } #[cfg(test)] mod tests { use super::*; #[test] fn example_file_parses() { let cfg = parse(include_str!("../../enhance.example.toml")).unwrap(); assert!(cfg.url.is_some()); assert_eq!(cfg.templates[0].name, "restore"); assert_eq!(cfg.default_name(), Some("restore")); assert!( cfg.templates .iter() .any(|t| t.prompt_mode == PromptMode::Replace) ); } #[test] fn empty_file_is_disabled_defaults() { let cfg = parse("").unwrap(); assert_eq!(cfg, EnhanceConfig::default()); assert!(cfg.url.is_none()); } #[test] fn rejects_unknown_keys_and_bad_values() { assert!(parse("nope = 1").is_err()); assert!(parse("url = \"ftp://x\"").is_err()); assert!(parse("poll_secs = 0").is_err()); assert!(parse("default_template = \"x\"").is_err()); let dup = "[[template]]\nname='a'\nlabel='A'\nfile='a.json'\n[[template]]\nname='a'\nlabel='B'\nfile='b.json'"; assert!(parse(dup).is_err()); assert!(parse("[[template]]\nname='a b'\nlabel='A'\nfile='a.json'").is_err()); assert!(parse("[[path_map]]\nfrom=''\nto='Z:\\\\'").is_err()); } #[test] fn lora_switches_parse_and_validate() { let ok = "[[template]]\nname='a'\nlabel='A'\nfile='a.json'\n\ [[template.lora]]\nid='style'\nlabel='Style'\n\ [[template.lora]]\nid='detail'\nlabel='Detail'\ndefault=true"; let cfg = parse(ok).unwrap(); let l = &cfg.templates[0].loras; assert_eq!((l[0].id.as_str(), l[0].default), ("style", false)); assert_eq!((l[1].id.as_str(), l[1].default), ("detail", true)); let base = "[[template]]\nname='a'\nlabel='A'\nfile='a.json'\n"; let dup = format!( "{base}[[template.lora]]\nid='x'\nlabel='X'\n[[template.lora]]\nid='x'\nlabel='Y'" ); assert!(parse(&dup).is_err()); assert!(parse(&format!("{base}[[template.lora]]\nid='a b'\nlabel='X'")).is_err()); assert!(parse(&format!("{base}[[template.lora]]\nid='x'\nlabel=' '")).is_err()); assert!( parse(&format!( "{base}[[template.lora]]\nid='x'\nlabel='X'\nstrength=1" )) .is_err() ); } #[test] fn template_defaults() { let cfg = parse("[[template]]\nname='a'\nlabel='A'\nfile='a.json'").unwrap(); let t = &cfg.templates[0]; assert_eq!(t.prompt_mode, PromptMode::Append); assert_eq!(t.prompt_join, " "); assert_eq!(t.output_megapixels, 1.0); assert!(t.loras.is_empty()); } }