feat(enhance): config, templates, ComfyUI client, jobs, endpoints (#192)
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Cameron CordesandClaude Opus 5.5 authored and 
cameron committed
1 parent
21920e8608
commit
3a21761cfe
19 files changed
+3912
No files matched your search
@@ -0,0 +1,490 @@
|
||||
//! 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::path::Path;
|
||||
|
||||
use anyhow::{Context as _, Result, bail};
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
use super::config::{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";
|
||||
|
||||
/// Prompt-node inputs, in the order they are tried.
|
||||
const PROMPT_INPUTS: &[&str] = &["text", "prompt", "value"];
|
||||
|
||||
/// 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>,
|
||||
/// Startup warnings (stray output nodes).
|
||||
pub warnings: Vec<String>,
|
||||
}
|
||||
|
||||
/// 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,
|
||||
}
|
||||
|
||||
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();
|
||||
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"
|
||||
));
|
||||
}
|
||||
}
|
||||
Ok(Template {
|
||||
cfg,
|
||||
workflow,
|
||||
source_id,
|
||||
save_id,
|
||||
prompt,
|
||||
reference_id,
|
||||
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
|
||||
}
|
||||
|
||||
/// 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.
|
||||
pub fn fill(
|
||||
&self,
|
||||
source: &str,
|
||||
prompt: Option<&str>,
|
||||
reference: Option<&str>,
|
||||
next_seed: &mut dyn FnMut() -> u64,
|
||||
) -> Result<Filled, FillError> {
|
||||
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;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
|
||||
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");
|
||||
|
||||
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,
|
||||
}
|
||||
}
|
||||
|
||||
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, &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, &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 = || rand::random::<u64>() >> 12;
|
||||
let x = t.fill("p", None, None, &mut rng).unwrap();
|
||||
let y = t.fill("p", None, None, &mut rng).unwrap();
|
||||
assert_ne!(input(&x.prompt, "5", "seed"), input(&y.prompt, "5", "seed"));
|
||||
}
|
||||
|
||||
#[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, &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, &mut seeds()).unwrap();
|
||||
assert_eq!(
|
||||
input(&f.prompt, "13", "value"),
|
||||
"Increase the clarity\n\nAdditional instructions: warmer"
|
||||
);
|
||||
let f = t.fill("p", Some(" "), None, &mut seeds()).unwrap();
|
||||
assert_eq!(input(&f.prompt, "13", "value"), base);
|
||||
let f = t.fill("p", 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, &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, &mut seeds()).unwrap_err();
|
||||
assert!(err.0.contains("instruction"));
|
||||
assert!(t.fill("p", Some(" "), None, &mut seeds()).is_err());
|
||||
let f = t.fill("p", Some("make it black and white"), None, &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, &mut seeds()).is_err());
|
||||
assert!(t.fill("p", None, None, &mut seeds()).is_ok());
|
||||
|
||||
let restore = load(RESTORE, PromptMode::Append).unwrap();
|
||||
assert!(restore.fill("p", None, Some("r"), &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"), &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, &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]));
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user