diff --git a/swarm/diffusion/diffusion_func.py b/swarm/diffusion/diffusion_func.py
index 75f4a7c..3e81de6 100644
--- 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