forked from boos/BoosAPI
Make short video factory use Comfy templates
This commit is contained in:
+177
-15
@@ -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)> {
|
||||
|
||||
Reference in New Issue
Block a user