always_offload

dkackman · Feb 5, 2024 · ae822a2 · ae822a2
1 parent 48793e8
commit ae822a2
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/swarm/diffusion/diffusion_func.py b/swarm/diffusion/diffusion_func.py
@@ -145,10 +145,8 @@ def diffusion_callback(device_identifier, model_name, **kwargs):
             main_pipeline.enable_vae_slicing()
         if has_method(main_pipeline, "enable_vae_tiling"):
             main_pipeline.enable_vae_tiling()
-        if has_method(main_pipeline, "enable_model_cpu_offload"):
-            main_pipeline.enable_model_cpu_offload()
 
-    if kwargs.pop("always_offload", False) and has_method(main_pipeline, "enable_model_cpu_offload"):
+    if (preserve_vram or kwargs.pop("always_offload", False)) and has_method(main_pipeline, "enable_model_cpu_offload"):
         main_pipeline.enable_model_cpu_offload()
 
     # prior pipeline is used by the Kandinsky v2 and others