diff --git a/nodes.py b/nodes.py index c4a05db..a63fe87 100644 --- a/nodes.py +++ b/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},) diff --git a/pyramid_dit/pyramid_dit_for_video_gen_pipeline.py b/pyramid_dit/pyramid_dit_for_video_gen_pipeline.py index 443e8d4..cfb7dff 100644 --- a/pyramid_dit/pyramid_dit_for_video_gen_pipeline.py +++ b/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):