Add --chroma-disable-dit-mask

Author: cmdr2

include/image_generator.h

@@ -118,6 +118,7 @@ class ImageGenerator {
     bool diffusion_fa_;
     bool control_net_cpu_;
     bool clip_on_cpu_;
+    bool chroma_disable_dit_mask_;
 };
 
 #endif  // __IMAGE_GENERATOR_H__

include/server.h

@@ -21,6 +21,7 @@ struct ServerParams {
     bool diffusion_fa = false;
     bool control_net_cpu = false;
     bool clip_on_cpu = false;
+    bool chroma_disable_dit_mask = false;
 };
 
 class Server {

src/image_generator.cpp

@@ -89,7 +89,8 @@ ImageGenerator::ImageGenerator(std::shared_ptr<TaskStateManager> task_state_mana
       offload_to_cpu_(server_params.offload_to_cpu),
       diffusion_fa_(server_params.diffusion_fa),
       control_net_cpu_(server_params.control_net_cpu),
-      clip_on_cpu_(server_params.clip_on_cpu) {
+      clip_on_cpu_(server_params.clip_on_cpu),
+      chroma_disable_dit_mask_(server_params.chroma_disable_dit_mask) {
     LOG_INFO("ImageGenerator created");
 }
 
@@ -578,6 +579,7 @@ bool ImageGenerator::ensureModelLoaded(const std::string& controlnet_model) {
     params.diffusion_flash_attn = diffusion_fa_;
     params.keep_control_net_on_cpu = control_net_cpu_;
     params.keep_clip_on_cpu = clip_on_cpu_;
+    params.chroma_use_dit_mask = !chroma_disable_dit_mask_;
 
     if (vae_on_cpu_) {
         LOG_INFO("VAE will be kept on CPU");
@@ -594,6 +596,9 @@ bool ImageGenerator::ensureModelLoaded(const std::string& controlnet_model) {
     if (clip_on_cpu_) {
         LOG_INFO("CLIP will be kept on CPU");
     }
+    if (chroma_disable_dit_mask_) {
+        LOG_INFO("DiT mask disabled for Chroma models");
+    }
 
     // Create SD context
     sd_ctx_ = new_sd_ctx(&params);

src/main.cpp

@@ -96,6 +96,7 @@ struct CommandLineArgs {
     bool diffusion_fa = false;
     bool control_net_cpu = false;
     bool clip_on_cpu = false;
+    bool chroma_disable_dit_mask = false;
 };
 
 void print_usage(const char* program_name) {
@@ -124,6 +125,8 @@ void print_usage(const char* program_name) {
     std::cerr << "  --diffusion-fa                     Enable diffusion flash attention (default: false)" << std::endl;
     std::cerr << "  --control-net-cpu                  Keep ControlNet on CPU (default: false)" << std::endl;
     std::cerr << "  --clip-on-cpu                      Keep CLIP on CPU (default: false)" << std::endl;
+    std::cerr << "  --chroma-disable-dit-mask          Disable DiT mask for Chroma models (default: false)"
+              << std::endl;
 }
 
 CommandLineArgs parse_args(int argc, char* argv[]) {
@@ -184,6 +187,8 @@ CommandLineArgs parse_args(int argc, char* argv[]) {
             args.control_net_cpu = true;
         } else if (arg == "--clip-on-cpu") {
             args.clip_on_cpu = true;
+        } else if (arg == "--chroma-disable-dit-mask") {
+            args.chroma_disable_dit_mask = true;
         } else if (arg == "--help" || arg == "-h") {
             print_usage(argv[0]);
             exit(0);
@@ -265,6 +270,7 @@ int main(int argc, char* argv[]) {
         server_params.diffusion_fa = args.diffusion_fa;
         server_params.control_net_cpu = args.control_net_cpu;
         server_params.clip_on_cpu = args.clip_on_cpu;
+        server_params.chroma_disable_dit_mask = args.chroma_disable_dit_mask;
 
         // Create and start the server
         g_server = std::make_unique<Server>(server_params);