refactor(workflow_ui): 抽出 model_inject 子模块

Co-Authored-By: Claude <noreply@anthropic.com>

src/routes/workflow_ui/mod.rs

@@ -4,7 +4,6 @@ use crate::{
     services::workspace::{self, WorkspaceAssetInput},
     workflow::{
         models::{RunContext, WorkflowGraph},
-        registry,
     },
     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 OFFICIAL_WORKFLOW_TEMPLATES_BASE: &str =
     "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_EXTERNAL: Dir<'_> =
     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};
 mod 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;
 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;
 use assets::{
@@ -149,6 +146,12 @@ use 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,
 };
+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> {
     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 {
     prompt
         .as_object()

src/routes/workflow_ui/model_inject.rs

@@ -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 });
+        }
+    }
+}
+