refactor(workflow_ui): 抽出 model_inject 子模块

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
boos
2026-06-18 21:55:27 +08:00
parent e296c01b4b
commit 35f63c46f1
2 changed files with 290 additions and 280 deletions
+10 -280
View File
@@ -4,7 +4,6 @@ use crate::{
services::workspace::{self, WorkspaceAssetInput}, services::workspace::{self, WorkspaceAssetInput},
workflow::{ workflow::{
models::{RunContext, WorkflowGraph}, models::{RunContext, WorkflowGraph},
registry,
}, },
AppState, AppState,
}; };
@@ -42,7 +41,7 @@ const NATIVE_WORKFLOW_CACHE_MAX_AGE: Duration = Duration::from_secs(21600);
const WORKFLOW_UPLOAD_BODY_LIMIT_BYTES: usize = 256 * 1024 * 1024; const WORKFLOW_UPLOAD_BODY_LIMIT_BYTES: usize = 256 * 1024 * 1024;
const OFFICIAL_WORKFLOW_TEMPLATES_BASE: &str = const OFFICIAL_WORKFLOW_TEMPLATES_BASE: &str =
"https://raw.githubusercontent.com/Comfy-Org/workflow_templates/main/templates"; "https://raw.githubusercontent.com/Comfy-Org/workflow_templates/main/templates";
const DEFAULT_TEXT_TO_VIDEO_MODEL: &str = "seedance-1-0-lite-t2v-250428"; pub(super) const DEFAULT_TEXT_TO_VIDEO_MODEL: &str = "seedance-1-0-lite-t2v-250428";
static EMBEDDED_WORKFLOW_UI: Dir<'_> = include_dir!("$CARGO_MANIFEST_DIR/dist/workflow-ui"); static EMBEDDED_WORKFLOW_UI: Dir<'_> = include_dir!("$CARGO_MANIFEST_DIR/dist/workflow-ui");
static EMBEDDED_WORKFLOW_UI_EXTERNAL: Dir<'_> = static EMBEDDED_WORKFLOW_UI_EXTERNAL: Dir<'_> =
include_dir!("$CARGO_MANIFEST_DIR/dist/workflow-ui-external"); include_dir!("$CARGO_MANIFEST_DIR/dist/workflow-ui-external");
@@ -122,14 +121,12 @@ mod auth;
use auth::{auth_caller, caller_from_headers, caller_is_admin, caller_user_id}; use auth::{auth_caller, caller_from_headers, caller_is_admin, caller_user_id};
mod node_defs; mod node_defs;
use node_defs::{ use node_defs::{
core_node_definitions, core_workflow_template_index, default_boosapi_workflow_graph, core_workflow_template_index, default_boosapi_workflow_graph, text_to_video_workflow_graph,
node_type_def_to_comfy, text_to_video_workflow_graph,
}; };
mod models; mod models;
use models::{ use models::{
image_model_choices, is_video_workflow_model, native_folder_models_from_object_info, native_folder_models_from_object_info, workflow_folder_has_route_models,
workflow_folder_has_route_models, workflow_folder_is_api_only, workflow_folder_route_model_items, workflow_folder_is_api_only, workflow_folder_route_model_items, workflow_route_models,
workflow_route_models,
}; };
mod assets; mod assets;
use assets::{ use assets::{
@@ -149,6 +146,12 @@ use queue::{
history, history_by_id, history_mutate, interrupt, internal_logs, internal_logs_raw, queue, history, history_by_id, history_mutate, interrupt, internal_logs, internal_logs_raw, queue,
queue_mutate, queue_prompt, queue_status, repair_native_prompt_payload, view, queue_mutate, queue_prompt, queue_status, repair_native_prompt_payload, view,
}; };
mod model_inject;
use model_inject::{
inject_boosapi_object_info, merge_model_payload_items, model_payload_items_from_names,
native_model_folder_payload_or_fallback, workflow_model_folder_payload,
workflow_node_defs_with_models,
};
pub fn routes() -> Router<AppState> { pub fn routes() -> Router<AppState> {
Router::new() Router::new()
@@ -303,279 +306,6 @@ pub(super) fn parse_boolish(input: Option<&str>) -> bool {
) )
} }
fn comfy_combo_input(options: &[String], fallback: &str) -> Value {
let choices = if options.is_empty() {
vec![fallback.to_string()]
} else {
options.to_vec()
};
let default = choices
.iter()
.find(|model| model.as_str() == fallback)
.cloned()
.unwrap_or_else(|| choices[0].clone());
json!([choices, { "default": default }])
}
fn inject_model_choices(defs: &mut Map<String, Value>, models: &[String]) {
let image_models = image_model_choices(models);
let video_models = models
.iter()
.filter(|model| is_video_workflow_model(model))
.cloned()
.collect::<Vec<_>>();
if let Some(node) = defs.get_mut("boosapi.text_to_image") {
if let Some(optional) = node
.pointer_mut("/input/optional")
.and_then(Value::as_object_mut)
{
optional.insert(
"model".to_string(),
comfy_combo_input(&image_models, "gpt-image-2"),
);
}
}
for node_type in [
"boosapi.text_to_video",
"boosapi.image_to_video",
"boosapi.video_status",
] {
if let Some(node) = defs.get_mut(node_type) {
if let Some(optional) = node
.pointer_mut("/input/optional")
.and_then(Value::as_object_mut)
{
optional.insert(
"model".to_string(),
comfy_combo_input(&video_models, DEFAULT_TEXT_TO_VIDEO_MODEL),
);
}
}
}
for node_type in ["boosapi.chat_completion", "boosapi.comfyui_workflow"] {
if let Some(node) = defs.get_mut(node_type) {
if let Some(optional) = node
.pointer_mut("/input/optional")
.and_then(Value::as_object_mut)
{
optional.insert("model".to_string(), comfy_combo_input(models, "gpt-5.4"));
}
}
}
}
fn is_boosapi_extension_node_type(node_type: &str) -> bool {
node_type.starts_with("boosapi.")
|| matches!(
node_type,
"input.text"
| "output.preview_image"
| "output.preview_video"
| "output.save_file"
| "util.text_join"
| "util.json_extract"
| "comfyui.submit_workflow"
)
}
pub(super) fn workflow_node_defs_with_models(
models: &[String],
include_compat_core: bool,
) -> Map<String, Value> {
let core_defs = include_compat_core
.then(core_node_definitions)
.unwrap_or_default();
let mut defs = core_defs
.into_iter()
.chain(
registry::node_types()
.into_iter()
.filter(|def| include_compat_core || is_boosapi_extension_node_type(def.node_type)),
)
.map(|def| (def.node_type.to_string(), node_type_def_to_comfy(&def)))
.collect::<Map<_, _>>();
inject_model_choices(&mut defs, models);
defs
}
pub(super) fn inject_boosapi_object_info(value: &mut Value, models: &[String]) {
let Some(object) = value.as_object_mut() else {
return;
};
for (key, def) in workflow_node_defs_with_models(models, false) {
object.entry(key).or_insert(def);
}
}
fn builtin_model_folders() -> Vec<&'static str> {
vec![
"boosapi_models",
"checkpoints",
"chat_models",
"video_models",
"audio_tts_models",
"embeddings",
"vision_models",
"3d_models",
"loras",
"vae",
]
}
fn inject_boosapi_model_folders(value: &mut Value) {
let Some(items) = value.as_array_mut() else {
*value = workflow_model_folder_payload();
return;
};
let mut seen = items
.iter()
.filter_map(|item| {
item.get("name")
.and_then(Value::as_str)
.map(|name| name.to_ascii_lowercase())
})
.collect::<HashSet<_>>();
for name in builtin_model_folders() {
if seen.insert(name.to_ascii_lowercase()) {
items.push(json!({ "name": name }));
}
}
}
fn workflow_model_folder_payload() -> Value {
Value::Array(
builtin_model_folders()
.into_iter()
.map(|name| json!({ "name": name }))
.collect(),
)
}
fn native_model_folder_payload_or_fallback(value: Value) -> Value {
match value {
Value::Array(_) => {
let mut value = value;
normalize_model_folder_payload(&mut value);
inject_boosapi_model_folders(&mut value);
value
}
Value::Object(object)
if object
.get("code")
.and_then(Value::as_i64)
.is_some_and(|code| code == 404) =>
{
workflow_model_folder_payload()
}
_ => workflow_model_folder_payload(),
}
}
pub(super) fn model_payload_items_from_names(names: Vec<String>, path_index: i64) -> Vec<Value> {
names
.into_iter()
.map(|name| json!({ "name": name, "pathIndex": path_index }))
.collect()
}
fn normalize_model_item_name(item: &Value) -> Option<String> {
item.get("name")
.and_then(Value::as_str)
.or_else(|| item.as_str())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
fn workflow_model_display_aliases(name: &str) -> Vec<String> {
let trimmed = name.trim();
if trimmed.is_empty() {
return Vec::new();
}
let mut aliases = Vec::new();
let lower = trimmed.to_ascii_lowercase();
if lower == "ltx-2.3-22b-distilled-1.1_lora-dynamic_fro09_avg_rank_111_bf16.safetensors" {
aliases.push(
"ltx_2.3_22b_distilled_1.1_lora_dynamic_fro09_avg_rank_111_bf16.safetensors"
.to_string(),
);
aliases.push(
"ltx_2.3_22b_distilled_1.1_lora_dynamic_fr009_avg_rank_111_bf16.safetensors"
.to_string(),
);
}
if lower.contains("ltx") && lower.contains("lora") {
aliases.push(trimmed.replace('-', "_"));
}
aliases.sort();
aliases.dedup_by(|a, b| a.eq_ignore_ascii_case(b));
aliases
}
fn model_payload_item_with_name(source: &Value, name: String) -> Value {
match source {
Value::Object(object) => {
let mut object = object.clone();
object.insert("name".to_string(), Value::String(name));
object
.entry("pathIndex".to_string())
.or_insert_with(|| Value::Number(0.into()));
Value::Object(object)
}
_ => json!({ "name": name, "pathIndex": 0 }),
}
}
fn merge_model_payload_items(primary: Vec<Value>, extra: Vec<Value>) -> Value {
let mut seen = HashSet::<String>::new();
let mut items = Vec::with_capacity(primary.len() + extra.len());
for item in primary.into_iter().chain(extra) {
let Some(name) = normalize_model_item_name(&item) else {
continue;
};
if seen.insert(name.to_ascii_lowercase()) {
let aliases = workflow_model_display_aliases(&name);
items.push(match item {
Value::String(_) => json!({ "name": name.clone(), "pathIndex": 0 }),
Value::Object(mut object) => {
object
.entry("name".to_string())
.or_insert_with(|| Value::String(name.clone()));
object
.entry("pathIndex".to_string())
.or_insert_with(|| Value::Number(0.into()));
Value::Object(object)
}
_ => json!({ "name": name.clone(), "pathIndex": 0 }),
});
let source = items.last().cloned().unwrap_or_else(|| json!({}));
for alias in aliases {
if seen.insert(alias.to_ascii_lowercase()) {
items.push(model_payload_item_with_name(&source, alias));
}
}
}
}
Value::Array(items)
}
fn normalize_model_folder_payload(value: &mut Value) {
let Some(items) = value.as_array_mut() else {
return;
};
for item in items {
if let Some(name) = item.as_str().map(str::to_string) {
*item = json!({ "name": name });
}
}
}
pub(super) fn prompt_contains_boosapi_nodes(prompt: &Value) -> bool { pub(super) fn prompt_contains_boosapi_nodes(prompt: &Value) -> bool {
prompt prompt
.as_object() .as_object()
+280
View File
@@ -0,0 +1,280 @@
use super::models::{image_model_choices, is_video_workflow_model};
use super::node_defs::{core_node_definitions, node_type_def_to_comfy};
use super::DEFAULT_TEXT_TO_VIDEO_MODEL;
use crate::workflow::registry;
use serde_json::{json, Map, Value};
use std::collections::HashSet;
pub(super) fn comfy_combo_input(options: &[String], fallback: &str) -> Value {
let choices = if options.is_empty() {
vec![fallback.to_string()]
} else {
options.to_vec()
};
let default = choices
.iter()
.find(|model| model.as_str() == fallback)
.cloned()
.unwrap_or_else(|| choices[0].clone());
json!([choices, { "default": default }])
}
pub(super) fn inject_model_choices(defs: &mut Map<String, Value>, models: &[String]) {
let image_models = image_model_choices(models);
let video_models = models
.iter()
.filter(|model| is_video_workflow_model(model))
.cloned()
.collect::<Vec<_>>();
if let Some(node) = defs.get_mut("boosapi.text_to_image") {
if let Some(optional) = node
.pointer_mut("/input/optional")
.and_then(Value::as_object_mut)
{
optional.insert(
"model".to_string(),
comfy_combo_input(&image_models, "gpt-image-2"),
);
}
}
for node_type in [
"boosapi.text_to_video",
"boosapi.image_to_video",
"boosapi.video_status",
] {
if let Some(node) = defs.get_mut(node_type) {
if let Some(optional) = node
.pointer_mut("/input/optional")
.and_then(Value::as_object_mut)
{
optional.insert(
"model".to_string(),
comfy_combo_input(&video_models, DEFAULT_TEXT_TO_VIDEO_MODEL),
);
}
}
}
for node_type in ["boosapi.chat_completion", "boosapi.comfyui_workflow"] {
if let Some(node) = defs.get_mut(node_type) {
if let Some(optional) = node
.pointer_mut("/input/optional")
.and_then(Value::as_object_mut)
{
optional.insert("model".to_string(), comfy_combo_input(models, "gpt-5.4"));
}
}
}
}
pub(super) fn is_boosapi_extension_node_type(node_type: &str) -> bool {
node_type.starts_with("boosapi.")
|| matches!(
node_type,
"input.text"
| "output.preview_image"
| "output.preview_video"
| "output.save_file"
| "util.text_join"
| "util.json_extract"
| "comfyui.submit_workflow"
)
}
pub(super) fn workflow_node_defs_with_models(
models: &[String],
include_compat_core: bool,
) -> Map<String, Value> {
let core_defs = include_compat_core
.then(core_node_definitions)
.unwrap_or_default();
let mut defs = core_defs
.into_iter()
.chain(
registry::node_types()
.into_iter()
.filter(|def| include_compat_core || is_boosapi_extension_node_type(def.node_type)),
)
.map(|def| (def.node_type.to_string(), node_type_def_to_comfy(&def)))
.collect::<Map<_, _>>();
inject_model_choices(&mut defs, models);
defs
}
pub(super) fn inject_boosapi_object_info(value: &mut Value, models: &[String]) {
let Some(object) = value.as_object_mut() else {
return;
};
for (key, def) in workflow_node_defs_with_models(models, false) {
object.entry(key).or_insert(def);
}
}
pub(super) fn builtin_model_folders() -> Vec<&'static str> {
vec![
"boosapi_models",
"checkpoints",
"chat_models",
"video_models",
"audio_tts_models",
"embeddings",
"vision_models",
"3d_models",
"loras",
"vae",
]
}
pub(super) fn inject_boosapi_model_folders(value: &mut Value) {
let Some(items) = value.as_array_mut() else {
*value = workflow_model_folder_payload();
return;
};
let mut seen = items
.iter()
.filter_map(|item| {
item.get("name")
.and_then(Value::as_str)
.map(|name| name.to_ascii_lowercase())
})
.collect::<HashSet<_>>();
for name in builtin_model_folders() {
if seen.insert(name.to_ascii_lowercase()) {
items.push(json!({ "name": name }));
}
}
}
pub(super) fn workflow_model_folder_payload() -> Value {
Value::Array(
builtin_model_folders()
.into_iter()
.map(|name| json!({ "name": name }))
.collect(),
)
}
pub(super) fn native_model_folder_payload_or_fallback(value: Value) -> Value {
match value {
Value::Array(_) => {
let mut value = value;
normalize_model_folder_payload(&mut value);
inject_boosapi_model_folders(&mut value);
value
}
Value::Object(object)
if object
.get("code")
.and_then(Value::as_i64)
.is_some_and(|code| code == 404) =>
{
workflow_model_folder_payload()
}
_ => workflow_model_folder_payload(),
}
}
pub(super) fn model_payload_items_from_names(names: Vec<String>, path_index: i64) -> Vec<Value> {
names
.into_iter()
.map(|name| json!({ "name": name, "pathIndex": path_index }))
.collect()
}
pub(super) fn normalize_model_item_name(item: &Value) -> Option<String> {
item.get("name")
.and_then(Value::as_str)
.or_else(|| item.as_str())
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
pub(super) fn workflow_model_display_aliases(name: &str) -> Vec<String> {
let trimmed = name.trim();
if trimmed.is_empty() {
return Vec::new();
}
let mut aliases = Vec::new();
let lower = trimmed.to_ascii_lowercase();
if lower == "ltx-2.3-22b-distilled-1.1_lora-dynamic_fro09_avg_rank_111_bf16.safetensors" {
aliases.push(
"ltx_2.3_22b_distilled_1.1_lora_dynamic_fro09_avg_rank_111_bf16.safetensors"
.to_string(),
);
aliases.push(
"ltx_2.3_22b_distilled_1.1_lora_dynamic_fr009_avg_rank_111_bf16.safetensors"
.to_string(),
);
}
if lower.contains("ltx") && lower.contains("lora") {
aliases.push(trimmed.replace('-', "_"));
}
aliases.sort();
aliases.dedup_by(|a, b| a.eq_ignore_ascii_case(b));
aliases
}
pub(super) fn model_payload_item_with_name(source: &Value, name: String) -> Value {
match source {
Value::Object(object) => {
let mut object = object.clone();
object.insert("name".to_string(), Value::String(name));
object
.entry("pathIndex".to_string())
.or_insert_with(|| Value::Number(0.into()));
Value::Object(object)
}
_ => json!({ "name": name, "pathIndex": 0 }),
}
}
pub(super) fn merge_model_payload_items(primary: Vec<Value>, extra: Vec<Value>) -> Value {
let mut seen = HashSet::<String>::new();
let mut items = Vec::with_capacity(primary.len() + extra.len());
for item in primary.into_iter().chain(extra) {
let Some(name) = normalize_model_item_name(&item) else {
continue;
};
if seen.insert(name.to_ascii_lowercase()) {
let aliases = workflow_model_display_aliases(&name);
items.push(match item {
Value::String(_) => json!({ "name": name.clone(), "pathIndex": 0 }),
Value::Object(mut object) => {
object
.entry("name".to_string())
.or_insert_with(|| Value::String(name.clone()));
object
.entry("pathIndex".to_string())
.or_insert_with(|| Value::Number(0.into()));
Value::Object(object)
}
_ => json!({ "name": name.clone(), "pathIndex": 0 }),
});
let source = items.last().cloned().unwrap_or_else(|| json!({}));
for alias in aliases {
if seen.insert(alias.to_ascii_lowercase()) {
items.push(model_payload_item_with_name(&source, alias));
}
}
}
}
Value::Array(items)
}
pub(super) fn normalize_model_folder_payload(value: &mut Value) {
let Some(items) = value.as_array_mut() else {
return;
};
for item in items {
if let Some(name) = item.as_str().map(str::to_string) {
*item = json!({ "name": name });
}
}
}