refactor: centralize model factories and defaults (#2198)

src/model/integrations/sd.cpp

@@ -0,0 +1,261 @@
+#include "conditioning/conditioner.hpp"
+#include "model/diffusion/unet.hpp"
+#include "model/vae/auto_encoder_kl.hpp"
+#include "model_definition.h"
+#include "pipeline/model_builders.h"
+
+static void configure_sdxl_vae(const sd::model_builders::Context& ctx,
+                               const sd::model_builders::VAEOptions& options,
+                               VAE& vae) {
+    if (strlen(SAFE_STR(ctx.params.vae_path)) == 0 || ctx.params.force_sdxl_vae_conv_scale || options.external_vae_is_invalid) {
+        float vae_conv_2d_scale = 1.f / 32.f;
+        LOG_WARN(
+            "No valid VAE specified with --vae or --force-sdxl-vae-conv-scale flag set, "
+            "using Conv2D scale %.3f",
+            vae_conv_2d_scale);
+        static_cast<AutoEncoderKL&>(vae).set_conv2d_scale(vae_conv_2d_scale);
+    }
+}
+
+static bool build_sd(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;
+    // SD1.x SD2.x SDXL
+    std::map<std::string, std::string> embbeding_map;
+    for (uint32_t i = 0; i < sd_ctx_params->embedding_count; i++) {
+        embbeding_map.emplace(SAFE_STR(sd_ctx_params->embeddings[i].name), SAFE_STR(sd_ctx_params->embeddings[i].path));
+    }
+    result.conditioner = std::make_shared<FrozenCLIPEmbedderWithCustomWords>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                                             tensor_storage_map,
+                                                                             embbeding_map,
+                                                                             version,
+                                                                             weight_manager,
+                                                                             tokenizers);
+    result.diffusion   = std::make_shared<UNetModelRunner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                         tensor_storage_map,
+                                                         "model.diffusion_model",
+                                                         version,
+                                                         weight_manager);
+    return true;
+}
+static prediction_t detect_sdxl_prediction(const String2TensorStorage& weights) {
+    if (weights.count("edm_vpred.sigma_max"))
+        return EDM_V_PRED;
+    return weights.count("v_pred") ? V_PRED : EPS_PRED;
+}
+
+const std::vector<ModelDefinition>& sd_definitions() {
+    static const std::vector<ModelDefinition> definitions = [] {
+        std::vector<ModelDefinition> result;
+        {
+            ModelDefinition definition;
+            definition.version  = VERSION_SD1;
+            definition.geometry = {8, 4, 3};
+            definition.name     = "SD 1.x";
+
+            definition.sampling                       = {EPS_PRED, INFINITY};
+            definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD;
+            definition.sampling.default_scheduler     = DISCRETE_SCHEDULER;
+            definition.build                          = build_sd;
+            result.push_back(std::move(definition));
+        }
+        {
+            ModelDefinition definition;
+            definition.version                        = VERSION_SD1_INPAINT;
+            definition.geometry                       = {8, 4, 3};
+            definition.name                           = "SD 1.x Inpaint";
+            definition.traits.is_inpaint              = true;
+            definition.sampling                       = {EPS_PRED, INFINITY};
+            definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD;
+            definition.sampling.default_scheduler     = DISCRETE_SCHEDULER;
+            definition.build                          = build_sd;
+            result.push_back(std::move(definition));
+        }
+        {
+            ModelDefinition definition;
+            definition.version                        = VERSION_SD1_PIX2PIX;
+            definition.geometry                       = {8, 4, 3};
+            definition.name                           = "Instruct-Pix2Pix";
+            definition.traits.is_unet_edit            = true;
+            definition.sampling                       = {EPS_PRED, INFINITY};
+            definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD;
+            definition.sampling.default_scheduler     = DISCRETE_SCHEDULER;
+            definition.build                          = build_sd;
+            result.push_back(std::move(definition));
+        }
+        {
+            ModelDefinition definition;
+            definition.version  = VERSION_SD1_TINY_UNET;
+            definition.geometry = {8, 4, 3};
+            definition.name     = "SD 1.x Tiny UNet";
+
+            definition.sampling                       = {EPS_PRED, INFINITY};
+            definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD;
+            definition.sampling.default_scheduler     = DISCRETE_SCHEDULER;
+            definition.build                          = build_sd;
+            result.push_back(std::move(definition));
+        }
+        {
+            ModelDefinition definition;
+            definition.version  = VERSION_SD2;
+            definition.geometry = {8, 4, 3};
+            definition.name     = "SD 2.x";
+
+            definition.sampling                       = {EPS_PRED, INFINITY};
+            definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD;
+            definition.sampling.default_scheduler     = DISCRETE_SCHEDULER;
+            definition.build                          = build_sd;
+            definition.sampling.probe_v_prediction    = true;
+            result.push_back(std::move(definition));
+        }
+        {
+            ModelDefinition definition;
+            definition.version                        = VERSION_SD2_INPAINT;
+            definition.geometry                       = {8, 4, 3};
+            definition.name                           = "SD 2.x Inpaint";
+            definition.traits.is_inpaint              = true;
+            definition.sampling                       = {EPS_PRED, INFINITY};
+            definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD;
+            definition.sampling.default_scheduler     = DISCRETE_SCHEDULER;
+            definition.build                          = build_sd;
+            definition.sampling.probe_v_prediction    = true;
+            result.push_back(std::move(definition));
+        }
+        {
+            ModelDefinition definition;
+            definition.version  = VERSION_SD2_TINY_UNET;
+            definition.geometry = {8, 4, 3};
+            definition.name     = "SD 2.x Tiny UNet";
+
+            definition.sampling                       = {EPS_PRED, INFINITY};
+            definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD;
+            definition.sampling.default_scheduler     = DISCRETE_SCHEDULER;
+            definition.build                          = build_sd;
+            definition.sampling.probe_v_prediction    = true;
+            result.push_back(std::move(definition));
+        }
+        {
+            ModelDefinition definition;
+            definition.version  = VERSION_SDXS_512_DS;
+            definition.geometry = {8, 4, 3};
+            definition.name     = "SDXS (512-DS)";
+
+            definition.sampling                       = {EPS_PRED, INFINITY};
+            definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD;
+            definition.sampling.default_scheduler     = DISCRETE_SCHEDULER;
+            definition.build                          = build_sd;
+            result.push_back(std::move(definition));
+        }
+        {
+            ModelDefinition definition;
+            definition.version  = VERSION_SDXS_09;
+            definition.geometry = {8, 4, 3};
+            definition.name     = "SDXS (09)";
+
+            definition.sampling                       = {EPS_PRED, INFINITY};
+            definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD;
+            definition.sampling.default_scheduler     = DISCRETE_SCHEDULER;
+            definition.build                          = build_sd;
+            definition.sampling.probe_v_prediction    = true;
+            result.push_back(std::move(definition));
+        }
+        {
+            ModelDefinition definition;
+            definition.version  = VERSION_SDXL;
+            definition.geometry = {8, 4, 3};
+            definition.name     = "SDXL";
+
+            definition.sampling                         = {EPS_PRED, INFINITY};
+            definition.sampling.default_sample_method   = EULER_A_SAMPLE_METHOD;
+            definition.sampling.default_scheduler       = DISCRETE_SCHEDULER;
+            definition.build                            = build_sd;
+            definition.sampling.prediction_from_weights = detect_sdxl_prediction;
+            definition.configure_vae                    = configure_sdxl_vae;
+            result.push_back(std::move(definition));
+        }
+        {
+            ModelDefinition definition;
+            definition.version                          = VERSION_SDXL_INPAINT;
+            definition.geometry                         = {8, 4, 3};
+            definition.name                             = "SDXL Inpaint";
+            definition.traits.is_inpaint                = true;
+            definition.sampling                         = {EPS_PRED, INFINITY};
+            definition.sampling.default_sample_method   = EULER_A_SAMPLE_METHOD;
+            definition.sampling.default_scheduler       = DISCRETE_SCHEDULER;
+            definition.build                            = build_sd;
+            definition.sampling.prediction_from_weights = detect_sdxl_prediction;
+            definition.configure_vae                    = configure_sdxl_vae;
+            result.push_back(std::move(definition));
+        }
+        {
+            ModelDefinition definition;
+            definition.version                          = VERSION_SDXL_PIX2PIX;
+            definition.geometry                         = {8, 4, 3};
+            definition.name                             = "SDXL Instruct-Pix2Pix";
+            definition.traits.is_unet_edit              = true;
+            definition.sampling                         = {EPS_PRED, INFINITY};
+            definition.sampling.default_sample_method   = EULER_A_SAMPLE_METHOD;
+            definition.sampling.default_scheduler       = DISCRETE_SCHEDULER;
+            definition.build                            = build_sd;
+            definition.sampling.prediction_from_weights = detect_sdxl_prediction;
+            definition.configure_vae                    = configure_sdxl_vae;
+            result.push_back(std::move(definition));
+        }
+        {
+            ModelDefinition definition;
+            definition.version  = VERSION_SDXL_VEGA;
+            definition.geometry = {8, 4, 3};
+            definition.name     = "SDXL (Vega)";
+
+            definition.sampling                         = {EPS_PRED, INFINITY};
+            definition.sampling.default_sample_method   = EULER_A_SAMPLE_METHOD;
+            definition.sampling.default_scheduler       = DISCRETE_SCHEDULER;
+            definition.build                            = build_sd;
+            definition.sampling.prediction_from_weights = detect_sdxl_prediction;
+            definition.configure_vae                    = configure_sdxl_vae;
+            result.push_back(std::move(definition));
+        }
+        {
+            ModelDefinition definition;
+            definition.version  = VERSION_SDXL_SSD1B;
+            definition.geometry = {8, 4, 3};
+            definition.name     = "SDXL (SSD1B)";
+
+            definition.sampling                         = {EPS_PRED, INFINITY};
+            definition.sampling.default_sample_method   = EULER_A_SAMPLE_METHOD;
+            definition.sampling.default_scheduler       = DISCRETE_SCHEDULER;
+            definition.build                            = build_sd;
+            definition.sampling.prediction_from_weights = detect_sdxl_prediction;
+            definition.configure_vae                    = configure_sdxl_vae;
+            result.push_back(std::move(definition));
+        }
+        {
+            ModelDefinition definition;
+            definition.version                          = VERSION_SVD;
+            definition.geometry                         = {8, 4, 3};
+            definition.name                             = "SVD";
+            definition.traits.supports_video_generation = true;
+            definition.sampling                         = {EPS_PRED, INFINITY};
+            definition.sampling.default_sample_method   = EULER_A_SAMPLE_METHOD;
+            definition.sampling.default_scheduler       = DISCRETE_SCHEDULER;
+            definition.build                            = build_sd;
+            result.push_back(std::move(definition));
+        }
+        {
+            ModelDefinition definition;
+            definition.version  = VERSION_ESRGAN;
+            definition.geometry = {8, 4, 3};
+            definition.name     = "ESRGAN";
+
+            definition.sampling                       = {EPS_PRED, INFINITY};
+            definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD;
+            definition.sampling.default_scheduler     = DISCRETE_SCHEDULER;
+            definition.build                          = build_sd;
+            result.push_back(std::move(definition));
+        }
+        return result;
+    }();
+    return definitions;
+}