Use Comfy API for MoneyPrinter video generation

This commit is contained in:
boos
2026-06-02 16:47:48 +08:00
parent b29c7ffa7a
commit 1293ebae68
3 changed files with 505 additions and 43 deletions
+460 -43
View File
@@ -7,7 +7,7 @@ use crate::{
};
use axum::{
extract::{Path, Query, State},
http::{header, HeaderMap},
http::{header, HeaderMap, Method},
routing::{get, post},
Json, Router,
};
@@ -19,6 +19,8 @@ use std::time::Duration;
const TASK_TYPE_MONEY_PRINTER: &str = "money-printer-video";
const VIDEO_POLL_INTERVAL_SECS: u64 = 6;
const VIDEO_MAX_POLLS: usize = 120;
const COMFY_POLL_INTERVAL_SECS: u64 = 5;
const COMFY_MAX_POLLS: usize = 180;
#[derive(Debug, Clone)]
struct MoneyPrinterAuth {
@@ -130,10 +132,16 @@ async fn status(State(state): State<AppState>, headers: HeaderMap) -> AppResult<
"capabilities": [
"topic_to_script",
"boosapi_chat_completion",
"comfy_prompt",
"boosapi_text_to_video",
"workspace_archive",
"background_jobs"
]
],
"comfy": {
"enabled": state.config.comfyui.uses_native_backend(),
"backend": format!("{:?}", state.config.comfyui.backend).to_ascii_lowercase(),
"baseUrl": state.config.comfyui.target_base_url()
}
})))
}
@@ -351,6 +359,355 @@ fn video_task_id_from_response(value: &Value) -> Option<String> {
})
}
fn scalar_string(value: Option<&Value>) -> Option<String> {
match value? {
Value::String(value) => Some(value.trim().to_string()).filter(|value| !value.is_empty()),
Value::Number(value) => Some(value.to_string()),
Value::Bool(value) => Some(value.to_string()),
_ => None,
}
}
fn native_path(state: &AppState, path: &str) -> String {
if state.config.comfyui.uses_cloud_backend() {
format!(
"/api{}",
path.strip_prefix('/')
.map(|value| format!("/{value}"))
.unwrap_or_else(|| path.to_string())
)
} else {
path.to_string()
}
}
fn native_url(state: &AppState, path: &str) -> String {
format!("{}{}", state.config.comfyui.target_base_url(), native_path(state, path))
}
fn native_request(state: &AppState, method: Method, path: &str) -> reqwest::RequestBuilder {
let client = Client::builder()
.timeout(Duration::from_secs(900))
.build()
.unwrap_or_else(|_| Client::new());
let reqwest_method =
reqwest::Method::from_bytes(method.as_str().as_bytes()).unwrap_or(reqwest::Method::GET);
let mut request = client.request(reqwest_method, native_url(state, path));
if let Some(api_key) = state
.config
.comfyui
.api_key
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{
request = request
.header("X-API-Key", api_key)
.header("Authorization", format!("Bearer {api_key}"));
}
request
}
async fn native_json(
state: &AppState,
method: Method,
path: &str,
body: Option<Value>,
) -> AppResult<Value> {
let mut request = native_request(state, method, path);
if let Some(body) = body {
request = request.json(&body);
}
let response = request
.send()
.await
.map_err(|err| AppError::UpstreamError(format!("Comfy backend unreachable: {err}")))?;
let status = response.status();
let text = response
.text()
.await
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
let value = serde_json::from_str::<Value>(&text).unwrap_or_else(|_| json!({ "body": text }));
if !status.is_success() {
return Err(AppError::UpstreamError(value.to_string()));
}
Ok(value)
}
fn encode_query(value: &str) -> String {
let mut out = String::new();
for byte in value.bytes() {
if byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.' | b'~') {
out.push(byte as char);
} else {
out.push_str(&format!("%{byte:02X}"));
}
}
out
}
fn mapping_value<'a>(mapping: &'a Value, keys: &[&str]) -> Option<&'a Value> {
keys.iter().find_map(|key| mapping.get(*key))
}
fn mapping_node_id(mapping: &Value, keys: &[&str]) -> Option<String> {
mapping_value(mapping, keys).and_then(|value| {
if value.is_object() {
scalar_string(value.get("nodeId").or_else(|| value.get("node_id")).or(Some(value)))
} else {
scalar_string(Some(value))
}
})
}
fn mapping_field(mapping: &Value, keys: &[&str], default_field: &str) -> String {
mapping_value(mapping, keys)
.and_then(|value| {
if value.is_object() {
scalar_string(value.get("field").or_else(|| value.get("fieldName")))
} else {
None
}
})
.unwrap_or_else(|| default_field.to_string())
}
fn default_comfy_mapping() -> Value {
json!({
"prompt": { "nodeId": "", "field": "prompt" },
"negativePrompt": { "nodeId": "", "field": "text" },
"ratio": { "nodeId": "", "field": "ratio" },
"resolution": { "nodeId": "", "field": "resolution" },
"duration": { "nodeId": "", "field": "duration" },
"output": { "nodeId": "", "field": "video" }
})
}
fn comfy_mapping_from_body(body: &Value) -> Value {
body.get("comfyNodeMapping")
.or_else(|| body.get("comfy_node_mapping"))
.or_else(|| body.get("nodeMapping"))
.or_else(|| body.get("node_mapping"))
.cloned()
.unwrap_or_else(default_comfy_mapping)
}
fn workflow_prompt_from_body(body: &Value) -> AppResult<Value> {
body.get("comfyPrompt")
.or_else(|| body.get("comfy_prompt"))
.or_else(|| body.get("promptWorkflow"))
.or_else(|| body.get("prompt_workflow"))
.or_else(|| body.get("workflowPrompt"))
.or_else(|| body.get("workflow_prompt"))
.cloned()
.ok_or_else(|| {
AppError::BadRequest(
"Comfy API 模式需要传入 comfyPrompt/promptWorkflow 和节点映射".into(),
)
})
}
fn apply_node_input(prompt: &mut Value, node_id: &str, field: &str, value: Value) -> AppResult<()> {
if node_id.trim().is_empty() {
return Err(AppError::BadRequest("Comfy 节点映射包含空 nodeId".into()));
}
let Some(node) = prompt.get_mut(node_id) else {
return Err(AppError::BadRequest(format!("Comfy prompt missing node {node_id}")));
};
let Some(inputs) = node.get_mut("inputs").and_then(Value::as_object_mut) else {
return Err(AppError::BadRequest(format!("Comfy node {node_id} has no inputs")));
};
inputs.insert(field.to_string(), value);
Ok(())
}
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")
.and_then(Value::as_str)
.or_else(|| plan.get("script").and_then(Value::as_str))
.unwrap_or_default()
.to_string()
});
if video_prompt.trim().is_empty() {
return Err(AppError::BadRequest("video prompt is empty".into()));
}
let prompt_node = mapping_node_id(mapping, &["prompt", "positive", "text", "promptNode"])
.ok_or_else(|| AppError::BadRequest("Comfy API 模式缺少提示词节点映射".into()))?;
apply_node_input(
prompt,
&prompt_node,
&mapping_field(mapping, &["prompt", "positive", "text", "promptNode"], "prompt"),
json!(video_prompt),
)?;
if let Some(negative_node) = mapping_node_id(mapping, &["negativePrompt", "negative", "negativeNode"]) {
let negative = string_from(body, &["negativePrompt", "negative_prompt", "negative"])
.unwrap_or_else(|| "low quality, blurry, distorted, watermark, text artifacts".to_string());
apply_node_input(
prompt,
&negative_node,
&mapping_field(mapping, &["negativePrompt", "negative", "negativeNode"], "text"),
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)?;
}
}
}
Ok(video_prompt)
}
async fn submit_comfy_prompt(state: &AppState, prompt: Value) -> AppResult<(String, Value)> {
let body = json!({
"client_id": format!("boosapi-money-printer-{}", uuid::Uuid::new_v4().simple()),
"prompt": prompt,
});
let value = native_json(state, Method::POST, "/prompt", Some(body)).await?;
let prompt_id = scalar_string(
value
.get("prompt_id")
.or_else(|| value.get("promptId"))
.or_else(|| value.pointer("/data/prompt_id"))
.or_else(|| value.pointer("/data/promptId")),
)
.ok_or_else(|| AppError::UpstreamError(format!("Comfy response missing prompt_id: {value}")))?;
Ok((prompt_id, value))
}
async fn comfy_history(state: &AppState, prompt_id: &str) -> AppResult<Value> {
native_json(state, Method::GET, &format!("/history/{prompt_id}"), None).await
}
fn comfy_output_objects(value: &Value) -> Vec<Value> {
let mut out = Vec::new();
fn walk(value: &Value, out: &mut Vec<Value>) {
match value {
Value::Object(object) => {
for key in ["videos", "gifs", "images", "audio"] {
if let Some(items) = object.get(key).and_then(Value::as_array) {
out.extend(items.iter().filter(|item| item.is_object()).cloned());
}
}
for value in object.values() {
walk(value, out);
}
}
Value::Array(items) => {
for item in items {
walk(item, out);
}
}
_ => {}
}
}
walk(value, &mut out);
out
}
fn output_kind(item: &Value) -> &'static str {
let filename = scalar_string(item.get("filename")).unwrap_or_default().to_ascii_lowercase();
if filename.ends_with(".mp4") || filename.ends_with(".webm") || filename.ends_with(".mov") {
"video"
} else if filename.ends_with(".mp3") || filename.ends_with(".wav") || filename.ends_with(".m4a") {
"audio"
} else {
"image"
}
}
fn output_query(item: &Value) -> Option<String> {
let filename = scalar_string(item.get("filename"))?;
let file_type = scalar_string(item.get("type")).unwrap_or_else(|| "output".to_string());
let subfolder = scalar_string(item.get("subfolder")).unwrap_or_default();
let mut parts = vec![
format!("filename={}", encode_query(&filename)),
format!("type={}", encode_query(&file_type)),
];
if !subfolder.is_empty() {
parts.push(format!("subfolder={}", encode_query(&subfolder)));
}
Some(format!("?{}", parts.join("&")))
}
async fn archive_comfy_output(
state: &AppState,
auth: &MoneyPrinterAuth,
item: &Value,
prompt_id: &str,
prompt: &str,
title: &str,
) -> Option<Value> {
let query = output_query(item)?;
let kind = output_kind(item).to_string();
let path = if kind == "video" { "/viewvideo" } else { "/view" };
let response = native_request(state, Method::GET, &format!("{path}{query}"))
.send()
.await
.ok()?;
if !response.status().is_success() {
return None;
}
let content_type = response
.headers()
.get(reqwest::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok())
.unwrap_or("application/octet-stream")
.to_string();
let bytes = response.bytes().await.ok()?;
let input = WorkspaceAssetInput {
user_id: auth.user_id,
source: "money-printer".to_string(),
kind: kind.clone(),
title: Some(title.chars().take(80).collect()),
prompt: Some(prompt.chars().take(1000).collect()),
model: Some("comfy-api".to_string()),
remote_url: None,
metadata: json!({ "moneyPrinter": { "backend": "comfy", "promptId": prompt_id, "raw": item } }),
};
match workspace::record_bytes_asset(state, input, &bytes, &content_type).await {
Ok(asset) => Some(json!({
"publicId": asset.public_id,
"kind": kind,
"status": asset.status,
})),
Err(err) => Some(json!({
"kind": kind,
"status": "failed",
"error": err.to_string(),
})),
}
}
async fn submit_comfy_video(
state: &AppState,
body: &Value,
plan: &Value,
) -> AppResult<Value> {
if !state.config.comfyui.uses_native_backend() {
return Err(AppError::BadRequest(
"Comfy API 视频生成需要配置 comfyui.backend = cloud 或 external".into(),
));
}
let mut prompt = workflow_prompt_from_body(body)?;
let mapping = comfy_mapping_from_body(body);
let video_prompt = inject_comfy_video_inputs(&mut prompt, &mapping, body, plan)?;
let (prompt_id, response) = submit_comfy_prompt(state, prompt).await?;
Ok(json!({
"backend": "comfy",
"promptId": prompt_id,
"prompt": video_prompt,
"mapping": mapping,
"response": response,
}))
}
async fn submit_video(
state: &AppState,
auth: &MoneyPrinterAuth,
@@ -540,54 +897,111 @@ async fn create_job(
.execute(&*state_for_task.db)
.await?;
update_progress(&state_for_task, &task_id_for_spawn, 3, "正在提交 BoosAPI 视频模型").await?;
let submitted = submit_video(&state_for_task, &auth_for_task, &body, &plan).await?;
if is_cancelled(&state_for_task, &task_id_for_spawn).await? {
return Ok(());
}
let model = submitted
.get("model")
.and_then(Value::as_str)
.unwrap_or("default")
.to_string();
let first_response = submitted.get("response").cloned().unwrap_or(Value::Null);
let final_video_response = if video_url_from_response(&first_response).is_some() {
first_response
} else if let Some(video_task_id) = video_task_id_from_response(&first_response) {
update_progress(&state_for_task, &task_id_for_spawn, 4, "视频生成中,正在轮询上游任务").await?;
poll_video(&state_for_task, &auth_for_task, &model, &video_task_id).await?
} else {
first_response
};
let video_url = video_url_from_response(&final_video_response);
let prompt = submitted
.get("request")
.and_then(|request| request.get("prompt"))
.and_then(Value::as_str)
.unwrap_or_default()
.to_string();
let video_backend = string_from(&body, &["videoBackend", "video_backend"])
.unwrap_or_else(|| "comfy".to_string())
.to_ascii_lowercase();
let title = plan
.get("title")
.and_then(Value::as_str)
.unwrap_or("短视频工厂成片")
.to_string();
let mut assets = Vec::new();
if let Some(url) = video_url.as_deref() {
update_progress(&state_for_task, &task_id_for_spawn, 5, "视频完成,正在归档到工作空间").await?;
if let Some(asset) = archive_video(
&state_for_task,
&auth_for_task,
&title,
&prompt,
&model,
url,
json!({ "moneyPrinter": { "topic": topic, "plan": plan, "response": final_video_response } }),
)
.await
{
assets.push(asset);
let (submitted, final_video_response, video_url, prompt, model) = if video_backend == "comfy" {
update_progress(&state_for_task, &task_id_for_spawn, 3, "正在提交 Comfy API 工作流").await?;
let submitted = submit_comfy_video(&state_for_task, &body, &plan).await?;
if is_cancelled(&state_for_task, &task_id_for_spawn).await? {
return Ok(());
}
}
let prompt_id = submitted
.get("promptId")
.and_then(Value::as_str)
.unwrap_or_default()
.to_string();
let prompt = submitted
.get("prompt")
.and_then(Value::as_str)
.unwrap_or_default()
.to_string();
let mut final_history = Value::Null;
for _ in 0..COMFY_MAX_POLLS {
if is_cancelled(&state_for_task, &task_id_for_spawn).await? {
return Ok(());
}
update_progress(&state_for_task, &task_id_for_spawn, 4, "Comfy 正在生成视频").await?;
let history = comfy_history(&state_for_task, &prompt_id).await?;
let outputs = comfy_output_objects(&history);
if !outputs.is_empty() {
update_progress(&state_for_task, &task_id_for_spawn, 5, "Comfy 视频完成,正在归档").await?;
for output in &outputs {
if output_kind(output) == "video" {
if let Some(asset) = archive_comfy_output(
&state_for_task,
&auth_for_task,
output,
&prompt_id,
&prompt,
&title,
)
.await
{
assets.push(asset);
}
}
}
final_history = history;
break;
}
tokio::time::sleep(Duration::from_secs(COMFY_POLL_INTERVAL_SECS)).await;
}
if final_history.is_null() {
return Err(AppError::GatewayTimeout("Comfy 视频任务等待超时".into()));
}
(submitted, final_history, None, prompt, "comfy-api".to_string())
} else {
update_progress(&state_for_task, &task_id_for_spawn, 3, "正在提交 BoosAPI 视频模型").await?;
let submitted = submit_video(&state_for_task, &auth_for_task, &body, &plan).await?;
if is_cancelled(&state_for_task, &task_id_for_spawn).await? {
return Ok(());
}
let model = submitted
.get("model")
.and_then(Value::as_str)
.unwrap_or("default")
.to_string();
let first_response = submitted.get("response").cloned().unwrap_or(Value::Null);
let final_video_response = if video_url_from_response(&first_response).is_some() {
first_response
} else if let Some(video_task_id) = video_task_id_from_response(&first_response) {
update_progress(&state_for_task, &task_id_for_spawn, 4, "视频生成中,正在轮询上游任务").await?;
poll_video(&state_for_task, &auth_for_task, &model, &video_task_id).await?
} else {
first_response
};
let video_url = video_url_from_response(&final_video_response);
let prompt = submitted
.get("request")
.and_then(|request| request.get("prompt"))
.and_then(Value::as_str)
.unwrap_or_default()
.to_string();
if let Some(url) = video_url.as_deref() {
update_progress(&state_for_task, &task_id_for_spawn, 5, "视频完成,正在归档到工作空间").await?;
if let Some(asset) = archive_video(
&state_for_task,
&auth_for_task,
&title,
&prompt,
&model,
url,
json!({ "moneyPrinter": { "topic": topic, "plan": plan, "response": final_video_response } }),
)
.await
{
assets.push(asset);
}
}
(submitted, final_video_response, video_url, prompt, model)
};
let final_result = json!({
"userId": auth_for_task.user_id,
"role": auth_for_task.role,
@@ -595,6 +1009,9 @@ async fn create_job(
"topic": topic,
"request": body,
"script": script,
"videoBackend": video_backend,
"model": model,
"prompt": prompt,
"video": submitted,
"finalVideoResponse": final_video_response,
"videoUrl": video_url,