refactor: centralize model factories and defaults (#2198)

2026-11-11 11:26:34 +00:00 · 2026-10-24 01:03:05 +08:00
1 parent baf45b49e2
commit cd7f986058
@@ -0,0 +1,119 @@
+#include "model/diffusion/wan.hpp"
+#include "conditioning/conditioner.hpp"
+#include "model/audio/wav2vec2.hpp"
+#include "model_definition.h"
+#include "pipeline/model_builders.h"
+
+static bool build_wan(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) {
+    const auto* sd_ctx_params      = &ctx.params;
+    const auto& tensor_storage_map = ctx.tensor_storage_map;
+    const auto version             = ctx.version;
+    const auto& weight_manager     = ctx.weight_manager;
+    result.conditioner             = std::make_shared<T5CLIPEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                          tensor_storage_map,
+                                                          true,
+                                                          0,
+                                                          true,
+                                                          weight_manager);
+    result.diffusion               = std::make_shared<WAN::WanRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                        tensor_storage_map,
+                                                        "model.diffusion_model",
+                                                        version,
+                                                        weight_manager);
+    if (strlen(SAFE_STR(sd_ctx_params->high_noise_diffusion_model_path)) > 0) {
+        result.high_noise_diffusion = std::make_shared<WAN::WanRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                       tensor_storage_map,
+                                                                       "model.high_noise_diffusion_model",
+                                                                       version,
+                                                                       weight_manager);
+    }
+    if (result.diffusion->get_desc() == "Wan2.1-I2V-14B" ||
+        result.diffusion->get_desc() == "Wan2.1-FLF2V-14B" ||
+        result.diffusion->get_desc() == "Wan2.1-I2V-1.3B") {
+        if (!sd::model_builders::ensure_backend_pair(ctx.backends, SDBackendModule::CLIP_VISION)) {
+            return false;
+        }
+        result.clip_vision = std::make_shared<FrozenCLIPVisionEmbedder>(ctx.backends.runtime_backend(SDBackendModule::CLIP_VISION),
+                                                                        tensor_storage_map,
+                                                                        weight_manager);
+    }
+    if (version == VERSION_WAN2_2_S2V &&
+        tensor_storage_map.count("wav2vec2.encoder.layer_norm.bias") > 0) {
+        if (!sd::model_builders::ensure_backend_pair(ctx.backends, SDBackendModule::AUDIO_ENCODER)) {
+            return false;
+        }
+        result.audio_encoder = std::make_shared<Wav2Vec2::Wav2Vec2ModelRunner>(ctx.backends.runtime_backend(SDBackendModule::AUDIO_ENCODER),
+                                                                               tensor_storage_map,
+                                                                               "wav2vec2.",
+                                                                               weight_manager);
+    }
+    return true;
+}
+
+const std::vector<ModelDefinition>& wan_definitions() {
+    static const std::vector<ModelDefinition> definitions = [] {
+        std::vector<ModelDefinition> result;
+        {
+            ModelDefinition definition;
+            definition.version                          = VERSION_WAN2;
+            definition.geometry                         = {2, 16, 3};
+            definition.name                             = "Wan 2.x";
+            definition.traits.is_dit                    = true;
+            definition.traits.supports_video_generation = true;
+            definition.traits.vae_kind                  = ModelVAEKind::WAN;
+            definition.traits.video_tae                 = true;
+            definition.sampling                         = {FLOW_PRED, 5.f};
+            definition.sampling.default_sample_method   = EULER_SAMPLE_METHOD;
+            definition.sampling.default_scheduler       = DISCRETE_SCHEDULER;
+            definition.build                            = build_wan;
+            result.push_back(std::move(definition));
+        }
+        {
+            ModelDefinition definition;
+            definition.version                          = VERSION_WAN2_2_I2V;
+            definition.geometry                         = {2, 16, 3};
+            definition.name                             = "Wan 2.2 I2V";
+            definition.traits.is_dit                    = true;
+            definition.traits.supports_video_generation = true;
+            definition.traits.vae_kind                  = ModelVAEKind::WAN;
+            definition.traits.video_tae                 = true;
+            definition.sampling                         = {FLOW_PRED, 5.f};
+            definition.sampling.default_sample_method   = EULER_SAMPLE_METHOD;
+            definition.sampling.default_scheduler       = DISCRETE_SCHEDULER;
+            definition.build                            = build_wan;
+            result.push_back(std::move(definition));
+        }
+        {
+            ModelDefinition definition;
+            definition.version                          = VERSION_WAN2_2_TI2V;
+            definition.geometry                         = {2, 48, 3};
+            definition.name                             = "Wan 2.2 TI2V";
+            definition.traits.is_dit                    = true;
+            definition.traits.supports_video_generation = true;
+            definition.traits.vae_kind                  = ModelVAEKind::WAN;
+            definition.traits.video_tae                 = true;
+            definition.sampling                         = {FLOW_PRED, 5.f};
+            definition.sampling.default_sample_method   = EULER_SAMPLE_METHOD;
+            definition.sampling.default_scheduler       = DISCRETE_SCHEDULER;
+            definition.build                            = build_wan;
+            result.push_back(std::move(definition));
+        }
+        {
+            ModelDefinition definition;
+            definition.version                          = VERSION_WAN2_2_S2V;
+            definition.geometry                         = {2, 16, 3};
+            definition.name                             = "Wan 2.2 S2V";
+            definition.traits.is_dit                    = true;
+            definition.traits.supports_video_generation = true;
+            definition.traits.vae_kind                  = ModelVAEKind::WAN;
+            definition.traits.video_tae                 = true;
+            definition.sampling                         = {FLOW_PRED, 5.f};
+            definition.sampling.default_sample_method   = EULER_SAMPLE_METHOD;
+            definition.sampling.default_scheduler       = DISCRETE_SCHEDULER;
+            definition.build                            = build_wan;
+            result.push_back(std::move(definition));
+        }
+        return result;
+    }();
+    return definitions;
+}