cpu_offload

nodes.py

@@ -123,7 +123,8 @@ class PyramidFlowModelLoader:
         return {
             "required": {
                 "model": (folder_paths.get_filename_list("diffusion_models"), {"tooltip": "The name of the checkpoint (model) to load.",}),
-                    "precision": (["fp8_e4m3fn","fp8_e4m3fn_fast","fp16", "fp32", "bf16"], {"default": "bf16"}),
+                "precision": (["fp8_e4m3fn","fp8_e4m3fn_fast","fp16", "fp32", "bf16"], {"default": "bf16"}),
+                "enable_sequential_cpu_offload": ("BOOLEAN", {"default": False, "tooltip": "Enable sequential cpu offload, saves VRAM but is MUCH slower, do not use unless you have to"}),
             },
             "optional": {
                 "compile_args": ("MOCHICOMPILEARGS", {"tooltip": "Optional torch.compile arguments",}),
@@ -135,7 +136,7 @@ class PyramidFlowModelLoader:
     FUNCTION = "loadmodel"
     CATEGORY = "PyramidFlowWrapper"
 
-    def loadmodel(self, model, precision, compile_args=None):
+    def loadmodel(self, model, precision,enable_sequential_cpu_offload, compile_args=None):
 
         device = mm.get_torch_device()
         offload_device = mm.unet_offload_device()
@@ -216,7 +217,10 @@ class PyramidFlowModelLoader:
                     transformer.proj_out = torch.compile(transformer.proj_out, fullgraph=compile_args["fullgraph"], dynamic=dynamic, backend=compile_args["backend"])
  
 
-        pyramid_model = PyramidDiTForVideoGeneration(transformer, dtype, model_name)
+        pyramid_model = PyramidDiTForVideoGeneration(transformer, dtype, model_name, device)
+
+        if enable_sequential_cpu_offload:
+            pyramid_model.enable_sequential_cpu_offload()
 
         return (pyramid_model,)
 
@@ -301,7 +305,7 @@ class PyramidFlowSampler:
                     output_type="latent",
                 )
 
-        if not keep_model_loaded:
+        if not keep_model_loaded and not pyramid_model.sequential_offload_enabled:
             pyramid_model.dit.to(offload_device)      
 
         return ({"samples": latents},)

pyramid_dit/pyramid_dit_for_video_gen_pipeline.py

@@ -9,7 +9,7 @@ from tqdm import tqdm
 
 from typing import List, Optional, Union
 from ..diffusion_schedulers import PyramidFlowMatchEulerDiscreteScheduler
-
+from accelerate import cpu_offload
 from comfy.utils import ProgressBar
 
 def compute_density_for_timestep_sampling(
@@ -36,6 +36,7 @@ class PyramidDiTForVideoGeneration:
             transformer, 
             model_dtype, 
             model_name, 
+            main_device,
             return_log=True, 
             timestep_shift=1.0, 
             stage_range=[0, 1/3, 2/3, 1],
@@ -52,6 +53,9 @@ class PyramidDiTForVideoGeneration:
         super().__init__()
 
         self.dit = transformer
+        self.device = main_device
+        self.sequential_offload_enabled = False
+
         if model_dtype in [torch.float8_e4m3fn, torch.float8_e5m2]:
             self.dtype = torch.bfloat16
         else:
@@ -110,6 +114,14 @@ class PyramidDiTForVideoGeneration:
         vae_latent_list = list(reversed(vae_latent_list))
         return vae_latent_list
 
+    def _enable_sequential_cpu_offload(self, model, device):
+        self.sequential_offload_enabled = True
+        offload_buffers = len(model._parameters) > 0
+        cpu_offload(model, device, offload_buffers=offload_buffers)
+    
+    def enable_sequential_cpu_offload(self):
+        self._enable_sequential_cpu_offload(self.dit, device=self.device)
+
     def prepare_latents(
         self,
         batch_size,
@@ -330,7 +342,8 @@ class PyramidDiTForVideoGeneration:
         generated_latents_list = [input_image_latent]    # The generated results
         #last_generated_latents = input_image_latent
 
-        self.dit.to(device)
+        if not self.sequential_offload_enabled:
+            self.dit.to(device)
         comfy_pbar = ProgressBar(num_units)
 
         for unit_index in tqdm(range(1, num_units + 1)):
@@ -500,8 +513,8 @@ class PyramidDiTForVideoGeneration:
 
         generated_latents_list = []    # The generated results
         #last_generated_latents = None
-
-        self.dit.to(device)
+        if not self.sequential_offload_enabled:
+            self.dit.to(device)
         comfy_pbar = ProgressBar(num_units)
 
         for unit_index in tqdm(range(num_units)):
@@ -584,9 +597,9 @@ class PyramidDiTForVideoGeneration:
 
         return image
 
-    @property
-    def device(self):
-        return next(self.dit.parameters()).device
+    # @property
+    # def device(self):
+    #     return next(self.dit.parameters()).device
     
     @property
     def guidance_scale(self):