From 7fd753455cccbc845d25372894109cf7a0c71830 Mon Sep 17 00:00:00 2001 From: Wagner Bruna Date: Tue, 27 Oct 2026 19:39:01 -0300 Subject: [PATCH] refactor: use a distinct version id for Chroma --- src/model.h | 2 ++ src/model/diffusion/flux.hpp | 6 +++--- src/model/integrations/flux.cpp | 24 ++++++++++++++++-------- src/model_loader.cpp | 7 +++++++ src/pipeline/diffusion_engine.h | 1 + 5 files changed, 29 insertions(+), 11 deletions(-) diff --git a/src/model.h b/src/model.h index 3ee9553836..12bc22e0f0 100644 --- a/src/model.h +++ b/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; } diff --git a/src/model/diffusion/flux.hpp b/src/model/diffusion/flux.hpp index 2f4272eb13..3199d0fded 100644 --- a/src/model/diffusion/flux.hpp +++ b/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); diff --git a/src/model/integrations/flux.cpp b/src/model/integrations/flux.cpp index 4f14237aac..fc96da7a38 100644 --- a/src/model/integrations/flux.cpp +++ b/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; - 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) { + if (version == VERSION_CHROMA || version == VERSION_CHROMA_RADIANCE) { result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), tensor_storage_map, false, @@ -130,6 +123,21 @@ const std::vector& 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; diff --git a/src/model_loader.cpp b/src/model_loader.cpp index 63d681de66..bc6265fb9a 100644 --- a/src/model_loader.cpp +++ b/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 { diff --git a/src/pipeline/diffusion_engine.h b/src/pipeline/diffusion_engine.h index c7295d8f78..e067781b45 100644 --- a/src/pipeline/diffusion_engine.h +++ b/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) ||