Make short video factory use Comfy templates

This commit is contained in:
boos
2026-06-05 21:35:12 +08:00
parent 5ce5ea1687
commit 3f22d7b1ab
2 changed files with 239 additions and 39 deletions
+177 -15
View File
@@ -14,7 +14,7 @@ use axum::{
};
use reqwest::Client;
use serde::Deserialize;
use serde_json::{json, Value};
use serde_json::{json, Map, Value};
use std::time::Duration;
const TASK_TYPE_MONEY_PRINTER: &str = "money-printer-video";
@@ -585,6 +585,24 @@ fn first_node_by_class(prompt: &serde_json::Map<String, Value>, class_type: &str
.map(|(id, _)| id.clone())
}
fn first_existing_input_field(
prompt: &serde_json::Map<String, Value>,
node_id: &str,
fields: &[&str],
fallback: &str,
) -> String {
let inputs = prompt
.get(node_id)
.and_then(|node| node.get("inputs"))
.and_then(Value::as_object);
fields
.iter()
.find(|field| inputs.is_some_and(|inputs| inputs.contains_key(**field)))
.copied()
.unwrap_or(fallback)
.to_string()
}
fn prompt_contains_class_type(prompt: &Value, class_type: &str) -> bool {
match prompt {
Value::Object(object) => {
@@ -616,10 +634,14 @@ fn infer_comfy_mapping(prompt: &Value) -> Value {
.or_else(|| first_node_by_class(object, "KSampler"))
.or_else(|| first_node_by_class(object, "EmptyLatentImage"))
.unwrap_or_default();
let prompt_field = first_existing_input_field(object, &prompt_node, &["prompt", "positive", "text"], "text");
let ratio_field = first_existing_input_field(object, &video_node, &["ratio", "model.ratio", "aspect_ratio"], "ratio");
let resolution_field = first_existing_input_field(object, &video_node, &["model.resolution", "resolution", "size"], "resolution");
let duration_field = first_existing_input_field(object, &video_node, &["model.duration", "duration", "length"], "duration");
json!({
"prompt": {
"nodeId": prompt_node,
"field": "text"
"field": prompt_field
},
"negativePrompt": {
"nodeId": negative_node,
@@ -627,15 +649,15 @@ fn infer_comfy_mapping(prompt: &Value) -> Value {
},
"ratio": {
"nodeId": video_node,
"field": "ratio"
"field": ratio_field
},
"resolution": {
"nodeId": video_node,
"field": "resolution"
"field": resolution_field
},
"duration": {
"nodeId": video_node,
"field": "duration"
"field": duration_field
}
})
}
@@ -744,6 +766,47 @@ fn apply_node_input(prompt: &mut Value, node_id: &str, field: &str, value: Value
Ok(())
}
fn node_has_input(prompt: &Value, node_id: &str, field: &str) -> bool {
prompt
.get(node_id)
.and_then(|node| node.get("inputs"))
.and_then(Value::as_object)
.is_some_and(|inputs| inputs.contains_key(field))
}
fn apply_optional_node_input(
prompt: &mut Value,
mapping: &Value,
mapping_keys: &[&str],
default_field: &str,
value: Value,
) {
let Some(node_id) = mapping_node_id(mapping, mapping_keys) else {
return;
};
let field = mapping_field(mapping, mapping_keys, default_field);
if !node_has_input(prompt, &node_id, &field) {
return;
}
let _ = apply_node_input(prompt, &node_id, &field, value);
}
fn value_to_duration_seconds(value: &Value) -> Option<i64> {
value
.as_i64()
.or_else(|| value.as_str().and_then(|text| text.trim().parse::<i64>().ok()))
.map(|duration| duration.clamp(2, 120))
}
fn normalize_resolution_for_comfy(value: &Value) -> Value {
let text = value.as_str().unwrap_or_default().trim().to_ascii_lowercase();
match text.as_str() {
"720p" | "720" => json!("720P"),
"1080p" | "1080" => json!("1080P"),
_ => value.clone(),
}
}
fn inject_comfy_video_inputs(prompt: &mut Value, mapping: &Value, body: &Value, plan: &Value) -> AppResult<String> {
let video_prompt = string_from(body, &["videoPrompt", "video_prompt"]).unwrap_or_else(|| {
plan.get("video_prompt")
@@ -773,18 +836,117 @@ fn inject_comfy_video_inputs(prompt: &mut Value, mapping: &Value, body: &Value,
json!(negative),
)?;
}
for (keys, request_keys, default_field) in [
(&["ratio", "ratioNode"][..], &["ratio", "aspectRatio", "aspect_ratio"][..], "ratio"),
(&["resolution", "resolutionNode"][..], &["resolution", "size"][..], "resolution"),
(&["duration", "durationNode"][..], &["duration"][..], "duration"),
] {
if let Some(node_id) = mapping_node_id(mapping, keys) {
if let Some(value) = request_keys.iter().find_map(|key| body.get(*key).cloned()) {
apply_node_input(prompt, &node_id, &mapping_field(mapping, keys, default_field), value)?;
}
if let Some(value) = ["ratio", "aspectRatio", "aspect_ratio"].iter().find_map(|key| body.get(*key).cloned()) {
apply_optional_node_input(prompt, mapping, &["ratio", "ratioNode"], "ratio", value);
}
if let Some(value) = ["resolution", "size"].iter().find_map(|key| body.get(*key).cloned()) {
apply_optional_node_input(
prompt,
mapping,
&["resolution", "resolutionNode"],
"resolution",
normalize_resolution_for_comfy(&value),
);
}
if let Some(value) = ["duration"].iter().find_map(|key| body.get(*key)).and_then(value_to_duration_seconds) {
apply_optional_node_input(prompt, mapping, &["duration", "durationNode"], "duration", json!(value));
}
repair_money_printer_comfy_prompt(prompt);
Ok(video_prompt)
}
fn repair_money_printer_comfy_prompt(prompt: &mut Value) {
repair_resize_image_mask_nodes(prompt);
repair_text_generate_ltx2_prompt_nodes(prompt);
repair_wan2_api_nodes(prompt);
}
fn repair_resize_image_mask_nodes(prompt: &mut Value) {
let Some(nodes) = prompt.as_object_mut() else {
return;
};
for (_node_id, node) in nodes.iter_mut() {
if node.get("class_type").and_then(Value::as_str) != Some("ResizeImageMaskNode") {
continue;
}
let Some(inputs) = node.get_mut("inputs").and_then(Value::as_object_mut) else {
continue;
};
if inputs
.get("resize_type.crop")
.is_some_and(|value| !matches!(value.as_str(), Some("disabled" | "center")))
{
inputs.insert("resize_type.crop".to_string(), Value::String("center".to_string()));
}
}
Ok(video_prompt)
}
fn repair_text_generate_ltx2_prompt_nodes(prompt: &mut Value) {
let Some(nodes) = prompt.as_object_mut() else {
return;
};
for (_node_id, node) in nodes.iter_mut() {
if node.get("class_type").and_then(Value::as_str) != Some("TextGenerateLTX2Prompt") {
continue;
}
let Some(inputs) = node.get_mut("inputs").and_then(Value::as_object_mut) else {
continue;
};
if inputs
.get("sampling_mode.repetition_penalty")
.and_then(Value::as_f64)
.is_some_and(|value| !(0.0..=5.0).contains(&value))
{
inputs.insert("sampling_mode".to_string(), Value::String("on".to_string()));
inputs.insert("sampling_mode.temperature".to_string(), json!(0.7));
inputs.insert("sampling_mode.top_k".to_string(), json!(64));
inputs.insert("sampling_mode.top_p".to_string(), json!(0.95));
inputs.insert("sampling_mode.min_p".to_string(), json!(0.05));
inputs.insert("sampling_mode.repetition_penalty".to_string(), json!(1.05));
inputs.entry("sampling_mode.seed".to_string()).or_insert_with(|| json!(0));
inputs.insert("sampling_mode.presence_penalty".to_string(), json!(0.0));
}
}
}
fn repair_wan2_api_nodes(prompt: &mut Value) {
let Some(nodes) = prompt.as_object_mut() else {
return;
};
for (_node_id, node) in nodes.iter_mut() {
let class_type = node.get("class_type").and_then(Value::as_str).unwrap_or_default();
if !class_type.starts_with("Wan2") || !class_type.ends_with("Api") {
continue;
}
let Some(inputs) = node.get_mut("inputs").and_then(Value::as_object_mut) else {
continue;
};
repair_wan2_model_group(inputs);
}
}
fn repair_wan2_model_group(inputs: &mut Map<String, Value>) {
let model_value = inputs.get("model").cloned().unwrap_or(Value::Null);
if let Some(items) = model_value.as_array() {
if let Some(model) = items.first().cloned() {
inputs.insert("model".to_string(), model);
}
if let Some(resolution) = items.get(1).cloned() {
inputs.insert("model.resolution".to_string(), normalize_resolution_for_comfy(&resolution));
}
if let Some(duration) = items.get(2).and_then(value_to_duration_seconds) {
inputs.insert("model.duration".to_string(), json!(duration));
}
}
if inputs
.get("model.resolution")
.is_some_and(|value| !matches!(value.as_str(), Some("720P" | "1080P")))
{
inputs.insert("model.resolution".to_string(), json!("1080P"));
}
if let Some(duration) = inputs.get("model.duration").and_then(value_to_duration_seconds) {
inputs.insert("model.duration".to_string(), json!(duration));
}
}
async fn submit_comfy_prompt(state: &AppState, prompt: Value) -> AppResult<(String, Value)> {