feat: add Chroma support (#696)

---------

Co-authored-by: Green Sky <Green-Sky@users.noreply.github.com>
Co-authored-by: leejet <leejet714@gmail.com>
This commit is contained in:
stduhpf
2025-06-29 17:36:42 +02:00
committed by GitHub
parent 884e23eeeb
commit b1cc40c35c
8 changed files with 563 additions and 112 deletions

View File

@@ -159,7 +159,10 @@ public:
bool clip_on_cpu,
bool control_net_cpu,
bool vae_on_cpu,
bool diffusion_flash_attn) {
bool diffusion_flash_attn,
bool chroma_use_dit_mask,
bool chroma_use_t5_mask,
int chroma_t5_mask_pad) {
use_tiny_autoencoder = taesd_path.size() > 0;
#ifdef SD_USE_CUDA
LOG_DEBUG("Using CUDA backend");
@@ -334,8 +337,19 @@ public:
cond_stage_model = std::make_shared<SD3CLIPEmbedder>(clip_backend, model_loader.tensor_storages_types);
diffusion_model = std::make_shared<MMDiTModel>(backend, model_loader.tensor_storages_types);
} else if (sd_version_is_flux(version)) {
cond_stage_model = std::make_shared<FluxCLIPEmbedder>(clip_backend, model_loader.tensor_storages_types);
diffusion_model = std::make_shared<FluxModel>(backend, model_loader.tensor_storages_types, version, diffusion_flash_attn);
bool is_chroma = false;
for (auto pair : model_loader.tensor_storages_types) {
if (pair.first.find("distilled_guidance_layer.in_proj.weight") != std::string::npos) {
is_chroma = true;
break;
}
}
if (is_chroma) {
cond_stage_model = std::make_shared<PixArtCLIPEmbedder>(clip_backend, model_loader.tensor_storages_types, -1, chroma_use_t5_mask, chroma_t5_mask_pad);
} else {
cond_stage_model = std::make_shared<FluxCLIPEmbedder>(clip_backend, model_loader.tensor_storages_types);
}
diffusion_model = std::make_shared<FluxModel>(backend, model_loader.tensor_storages_types, version, diffusion_flash_attn, chroma_use_dit_mask);
} else {
if (id_embeddings_path.find("v2") != std::string::npos) {
cond_stage_model = std::make_shared<FrozenCLIPEmbedderWithCustomWords>(clip_backend, model_loader.tensor_storages_types, embeddings_path, version, PM_VERSION_2);
@@ -1135,7 +1149,10 @@ sd_ctx_t* new_sd_ctx(const char* model_path_c_str,
bool keep_clip_on_cpu,
bool keep_control_net_cpu,
bool keep_vae_on_cpu,
bool diffusion_flash_attn) {
bool diffusion_flash_attn,
bool chroma_use_dit_mask,
bool chroma_use_t5_mask,
int chroma_t5_mask_pad) {
sd_ctx_t* sd_ctx = (sd_ctx_t*)malloc(sizeof(sd_ctx_t));
if (sd_ctx == NULL) {
return NULL;
@@ -1177,7 +1194,10 @@ sd_ctx_t* new_sd_ctx(const char* model_path_c_str,
keep_clip_on_cpu,
keep_control_net_cpu,
keep_vae_on_cpu,
diffusion_flash_attn)) {
diffusion_flash_attn,
chroma_use_dit_mask,
chroma_use_t5_mask,
chroma_t5_mask_pad)) {
delete sd_ctx->sd;
sd_ctx->sd = NULL;
free(sd_ctx);