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

This commit is contained in:
Wagner Bruna authored and GitHub committed 2026-10-28 22:02:18 +08:00
1 parent 687915518e
commit 8cbe905fdd
5 files changed
+29 -11

No files matched your search

+3 -3
View File
@@ -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);
+16 -8
View File
@@ -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<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;