Allow --clip-on-cpu and --control-net-cpu

Author: cmdr2

include/image_generator.h

@@ -115,6 +115,8 @@ class ImageGenerator {
     bool vae_tiling_;
     bool offload_to_cpu_;
     bool diffusion_fa_;
+    bool control_net_cpu_;
+    bool clip_on_cpu_;
 };
 
 #endif  // __IMAGE_GENERATOR_H__

include/server.h

@@ -18,6 +18,8 @@ struct ServerParams {
     bool vae_tiling = false;
     bool offload_to_cpu = false;
     bool diffusion_fa = false;
+    bool control_net_cpu = false;
+    bool clip_on_cpu = false;
 };
 
 class Server {

src/image_generator.cpp

@@ -85,7 +85,9 @@ ImageGenerator::ImageGenerator(std::shared_ptr<TaskStateManager> task_state_mana
       vae_on_cpu_(server_params.vae_on_cpu),
       vae_tiling_(server_params.vae_tiling),
       offload_to_cpu_(server_params.offload_to_cpu),
-      diffusion_fa_(server_params.diffusion_fa) {
+      diffusion_fa_(server_params.diffusion_fa),
+      control_net_cpu_(server_params.control_net_cpu),
+      clip_on_cpu_(server_params.clip_on_cpu) {
     LOG_INFO("ImageGenerator created");
 }
 
@@ -546,6 +548,8 @@ bool ImageGenerator::ensureModelLoaded(const std::string& controlnet_model) {
     params.keep_vae_on_cpu = vae_on_cpu_;
     params.offload_params_to_cpu = offload_to_cpu_;
     params.diffusion_flash_attn = diffusion_fa_;
+    params.keep_control_net_on_cpu = control_net_cpu_;
+    params.keep_clip_on_cpu = clip_on_cpu_;
 
     if (vae_on_cpu_) {
         LOG_INFO("VAE will be kept on CPU");
@@ -556,6 +560,12 @@ bool ImageGenerator::ensureModelLoaded(const std::string& controlnet_model) {
     if (diffusion_fa_) {
         LOG_INFO("Diffusion flash attention enabled");
     }
+    if (control_net_cpu_) {
+        LOG_INFO("ControlNet will be kept on CPU");
+    }
+    if (clip_on_cpu_) {
+        LOG_INFO("CLIP will be kept on CPU");
+    }
 
     // Create SD context
     sd_ctx_ = new_sd_ctx(&params);

src/main.cpp

@@ -93,6 +93,8 @@ struct CommandLineArgs {
     bool vae_tiling = false;
     bool offload_to_cpu = false;
     bool diffusion_fa = false;
+    bool control_net_cpu = false;
+    bool clip_on_cpu = false;
 };
 
 void print_usage(const char* program_name) {
@@ -117,6 +119,8 @@ void print_usage(const char* program_name) {
     std::cerr << "  --vae-tiling                       Enable VAE tiling (default: false)" << std::endl;
     std::cerr << "  --offload-to-cpu                   Offload parameters to CPU (default: false)" << std::endl;
     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;
 }
 
 CommandLineArgs parse_args(int argc, char* argv[]) {
@@ -171,6 +175,10 @@ CommandLineArgs parse_args(int argc, char* argv[]) {
             args.offload_to_cpu = true;
         } else if (arg == "--diffusion-fa") {
             args.diffusion_fa = true;
+        } else if (arg == "--control-net-cpu") {
+            args.control_net_cpu = true;
+        } else if (arg == "--clip-on-cpu") {
+            args.clip_on_cpu = true;
         } else if (arg == "--help" || arg == "-h") {
             print_usage(argv[0]);
             exit(0);
@@ -249,6 +257,8 @@ int main(int argc, char* argv[]) {
         server_params.vae_tiling = args.vae_tiling;
         server_params.offload_to_cpu = args.offload_to_cpu;
         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;
 
         // Create and start the server
         g_server = std::make_unique<Server>(server_params);