add cpu offload option

app.py

@@ -17,6 +17,12 @@ def model_refiner_change(refiner, config):
     
     return config, str(config), assemble_code(config)
 
+def cpu_offload_change(cpu_offload, config):
+    
+    config = set_config(config, 'cpu_offload', cpu_offload)
+    
+    return config, str(config), assemble_code(config)
+
 def models_change(model, scheduler, config):
 
     config = set_config(config, 'model', model)
@@ -171,6 +177,9 @@ def run_inference(config, config_history, progress=gr.Progress(track_tqdm=True))
             num_inference_steps = int(config["inference_steps"]), 
             guidance_scale = float(config["guidance_scale"])).images
 
+        if str(config["use_safetensors"]).lower() != 'false':
+            pipeline.enable_model_cpu_offload()
+
         if config['refiner'] != '':
             image = refiner(
                 prompt = config["prompt"],
@@ -178,6 +187,9 @@ def run_inference(config, config_history, progress=gr.Progress(track_tqdm=True))
                 image=image,
             ).images
             
+            if str(config["use_safetensors"]).lower() != 'false':
+                refiner.enable_model_cpu_offload()
+
         config_history.append(config.copy())
 
         return image[0], dict_list_to_markdown_table(config_history), config_history
@@ -216,6 +228,7 @@ with gr.Blocks(analytics_enabled=False) as demo:
     with gr.Row():
         with gr.Column(scale=1):
             in_use_safetensors = gr.Radio(label="Use safe tensors:", choices=["True", "False"], interactive=False)
+            in_cpu_offload = gr.Radio(label="CPU Offload:", choices=["True", "False"], interactive=False, info="This may increase performance, as it offloads computations from the GPU to the CPU. But this can also lead to slower executions and lower effectiveness. Compare running time and outputs before making sure, that this setting will help you")
             in_model_refiner = gr.Dropdown(value="", choices=[""], label="Refiner", allow_custom_value=True, multiselect=False)
         with gr.Column(scale=1):
             in_safety_checker = gr.Radio(label="Enable safety checker:", value=config.value["safety_checker"], choices=["True", "False"])
@@ -223,7 +236,7 @@ with gr.Blocks(analytics_enabled=False) as demo:
 
     gr.Markdown("### Scheduler")
     with gr.Row():
-        in_schedulers = gr.Dropdown(choices=list(schedulers.keys()), label="Scheduler/Solver", info="the scheduler controls parameter adaption between each inference step, depending on the right scheduler for your model, it may only take 10 or 20 steps to achieve very good results, see https://huggingface.co/docs/diffusers/using-diffusers/loading#schedulers" )
+        in_schedulers = gr.Dropdown(choices=list(schedulers.keys()), label="Scheduler/Solver", info="schedulers employ various strategies for noise control, the scheduler controls parameter adaption between each inference step, depending on the right scheduler for your model, it may only take 10 or 20 steps to achieve very good results, see https://huggingface.co/docs/diffusers/using-diffusers/loading#schedulers" )
         out_scheduler_description = gr.Textbox(value="", label="Description")
     
     gr.Markdown("### Adapters")
@@ -258,6 +271,7 @@ with gr.Blocks(analytics_enabled=False) as demo:
     in_variant.change(variant_change, inputs=[in_variant, config], outputs=[config, out_config, out_code])
     in_models.change(models_change, inputs=[in_models, in_schedulers, config], outputs=[out_model_description, in_model_refiner, in_use_safetensors, in_schedulers, config, out_config, out_code])
     in_model_refiner.change(model_refiner_change, inputs=[in_model_refiner, config], outputs=[config, out_config, out_code])
+    in_cpu_offload.change(cpu_offload_change, inputs=[in_cpu_offload, config], outputs=[config, out_config, out_code])
     in_safety_checker.change(safety_checker_change, inputs=[in_safety_checker, config], outputs=[config, out_config, out_code])
     in_requires_safety_checker.change(requires_safety_checker_change, inputs=[in_requires_safety_checker, config], outputs=[config, out_config, out_code])
     in_schedulers.change(schedulers_change, inputs=[in_schedulers, config], outputs=[out_scheduler_description, config, out_config, out_code])

config.py

@@ -37,6 +37,7 @@ def get_initial_config():
     config = {
         "device": device,
         "model": None,
+        "cpu_offload": "False",
         "scheduler": None,
         "variant": None,
         "allow_tensorfloat32": allow_tensorfloat32,
@@ -141,8 +142,10 @@ def assemble_code(str_config):
             torch_dtype=data_type, 
             variant=variant).to(device)'''
     
+    if str(config["use_safetensors"]).lower() == 'false': code['051_cpu_offload'] = "pipeline.enable_model_cpu_offload()"
+    
     if config['refiner'] != '':
-        code['051_refiner'] = f'''refiner = DiffusionPipeline.from_pretrained(
+        code['052_refiner'] = f'''refiner = DiffusionPipeline.from_pretrained(
                 "{config['refiner']}",
                 text_encoder_2 = base.text_encoder_2,
                 vae = base.vae,
@@ -151,6 +154,8 @@ def assemble_code(str_config):
                 variant=variant,
             ).to(device)'''
         
+        if str(config["use_safetensors"]).lower() == 'false': code['053_cpu_offload'] = "refiner.enable_model_cpu_offload()"
+        
     code['054_requires_safety_checker'] = f'pipeline.requires_safety_checker = {config["requires_safety_checker"]}'
 
     if str(config["safety_checker"]).lower() == 'false':