refactor: use a distinct version id for Chroma (#2237)

src/model.h

@@ -69,6 +69,7 @@ enum SDVersion {
     VERSION_FLUX3,
     VERSION_MOSS_IMAGE,
     VERSION_FIBO2,
+    VERSION_CHROMA,
     VERSION_COUNT,
 };
 
@@ -139,6 +140,7 @@ static inline bool sd_version_is_flux(SDVersion version) {
         version == VERSION_FLUX_CONTROLS ||
         version == VERSION_FLEX_2 ||
         version == VERSION_OVIS_IMAGE ||
+        version == VERSION_CHROMA ||
         version == VERSION_CHROMA_RADIANCE) {
         return true;
     }

src/model/diffusion/flux.hpp

@@ -77,7 +77,10 @@ namespace Flux {
                 config.in_channels = 128;
             } else if (version == VERSION_FLEX_2) {
                 config.in_channels = 196;
+            } else if (version == VERSION_CHROMA) {
+                config.is_chroma = true;
             } else if (version == VERSION_CHROMA_RADIANCE) {
+                config.is_chroma   = true;
                 config.in_channels = 3;
                 config.patch_size  = 16;
             } else if (version == VERSION_OVIS_IMAGE) {
@@ -138,9 +141,6 @@ namespace Flux {
                     actual_radiance_patch_size = tensor_storage.ne[0];
                     LOG_VERBOSE("actual radiance patch size: %" PRId64, actual_radiance_patch_size);
                 }
-                if (name.find("distilled_guidance_layer.in_proj.weight") != std::string::npos) {
-                    config.is_chroma = true;
-                }
                 size_t db = name.find("double_blocks.");
                 if (db != std::string::npos) {
                     std::string block_name = name.substr(db);

src/model/integrations/flux.cpp

@@ -9,14 +9,7 @@ static bool build_flux(const sd::model_builders::Context& ctx, const TokenizerCo
     const auto& tensor_storage_map = ctx.tensor_storage_map;
     const auto version             = ctx.version;
     const auto& weight_manager     = ctx.weight_manager;
-    bool is_chroma                 = false;
+    if (version == VERSION_CHROMA || version == VERSION_CHROMA_RADIANCE) {
-    for (auto pair : tensor_storage_map) {
-        if (pair.first.find("distilled_guidance_layer.in_proj.weight") != std::string::npos) {
-            is_chroma = true;
-            break;
-        }
-    }
-    if (is_chroma) {
         result.conditioner = std::make_shared<T5CLIPEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
                                                               tensor_storage_map,
                                                               false,
@@ -130,6 +123,21 @@ const std::vector<ModelDefinition>& flux_definitions() {
             definition.sampling.flow_shift_from_weights = flux_guidance_flow_shift;
             result.push_back(std::move(definition));
         }
+        {
+            ModelDefinition definition;
+            definition.version                          = VERSION_CHROMA;
+            definition.geometry                         = {1, 16, 3};
+            definition.ref_image_preset                 = "flux_kontext";
+            definition.name                             = "Chroma";
+            definition.traits.is_dit                    = true;
+            definition.traits.vae_kind                  = ModelVAEKind::FLUX;
+            definition.sampling                         = {FLUX_FLOW_PRED, 1.0f};
+            definition.sampling.default_sample_method   = EULER_SAMPLE_METHOD;
+            definition.sampling.default_scheduler       = FLUX_SCHEDULER;
+            definition.build                            = build_flux;
+            definition.sampling.flow_shift_from_weights = flux_guidance_flow_shift;
+            result.push_back(std::move(definition));
+        }
         {
             ModelDefinition definition;
             definition.version                          = VERSION_OVIS_IMAGE;

src/model_loader.cpp

@@ -468,6 +468,7 @@ SDVersion ModelLoader::get_sd_version() const {
     bool has_output_block_311        = false;
     bool has_output_block_71         = false;
     bool has_attn_1024               = false;
+    bool has_distilled_guidance      = false;
 
     for (auto& [name, tensor_storage] : tensor_storage_map) {
         if (tensor_storage.name.find("model.diffusion_model.double_blocks.") != std::string::npos ||
@@ -487,6 +488,9 @@ SDVersion ModelLoader::get_sd_version() const {
         if (tensor_storage.name.find("model.diffusion_model.nerf_final_layer_conv.") != std::string::npos) {
             return VERSION_CHROMA_RADIANCE;
         }
+        if (tensor_storage.name.find("distilled_guidance_layer.in_proj.weight") != std::string::npos) {
+            has_distilled_guidance = true;
+        }
         if (tensor_storage.name.find("model.diffusion_model.joint_blocks.") != std::string::npos) {
             return VERSION_SD3;
         }
@@ -676,6 +680,9 @@ SDVersion ModelLoader::get_sd_version() const {
     }
 
     if (is_flux && !is_flux2) {
+        if (has_distilled_guidance) {
+            return VERSION_CHROMA;
+        }
         if (context_ebedding_weight.ne[0] == 3584) {
             return VERSION_LONGCAT;
         } else {

src/pipeline/diffusion_engine.h

@@ -36,6 +36,7 @@ namespace Wav2Vec2 {
 
 static inline bool sd_version_supports_ref_latent_img_cfg(SDVersion version) {
     return version == VERSION_FLUX ||
+           version == VERSION_CHROMA ||
            sd_version_is_flux2(version) ||
            (sd_version_is_qwen_image(version) && version != VERSION_QWEN_IMAGE_2_1) ||
            sd_version_is_mage_flow(version) ||