diff --git a/lingbot_video/pipeline_lingbot_video.py b/lingbot_video/pipeline_lingbot_video.py index 183b0a0..c733ab8 100644 --- a/lingbot_video/pipeline_lingbot_video.py +++ b/lingbot_video/pipeline_lingbot_video.py @@ -1,6 +1,6 @@ from __future__ import annotations -from contextlib import nullcontext +from contextlib import contextmanager, nullcontext from dataclasses import dataclass from typing import Any, Dict, List, Optional, Tuple, Union @@ -127,6 +127,29 @@ def __init__(self, transformer, vae, text_encoder, processor, scheduler): self.video_prompt_template = VIDEO_PROMPT_TEMPLATE self._crop_start: Optional[int] = None + def enable_sequential_cpu_offload( + self, + gpu_id: Optional[int] = None, + device: Optional[Union[str, torch.device]] = None, + ) -> None: + super().enable_sequential_cpu_offload(gpu_id=gpu_id, device=device) + + visual = getattr(getattr(self.text_encoder, "model", None), "visual", None) + if visual is None: + return + + # Qwen reads a child embedding weight directly from its visual parent. + from accelerate import cpu_offload + from accelerate.hooks import remove_hook_from_module + + remove_hook_from_module(self.text_encoder, recurse=True) + cpu_offload( + self.text_encoder, + execution_device=self._offload_device, + offload_buffers=bool(self.text_encoder._parameters), + preload_module_classes=[visual.__class__.__name__], + ) + @staticmethod def check_inputs(height: int, width: int, num_frames: int) -> None: if num_frames != 1 and (num_frames - 1) % 4 != 0: @@ -290,6 +313,27 @@ def _vae_latent_to_dit(self, latents: torch.Tensor) -> torch.Tensor: std_inv = std_inv.view(1, -1, 1, 1, 1) return (latents.float() - mean) * std_inv + def _vae_execution_device(self) -> torch.device: + """Return the VAE hook device, falling back to its parameter device.""" + for module in self.vae.modules(): + hook = getattr(module, "_hf_hook", None) + execution_device = getattr(hook, "execution_device", None) + if execution_device is not None: + return torch.device(execution_device) + return _module_device(self.vae) + + @contextmanager + def _vae_encode_context(self): + """Isolate an out-of-order VAE encode from the linear offload chain.""" + model_offload = bool(getattr(self, "_all_hooks", None)) + if model_offload: + self.maybe_free_model_hooks() + try: + yield self._vae_execution_device() + finally: + if model_offload: + self.maybe_free_model_hooks() + @torch.no_grad() def encode_video_latent( self, @@ -298,31 +342,33 @@ def encode_video_latent( ) -> torch.Tensor: if self.vae is None: raise ValueError("`vae` is required to encode video latents.") - vae_device = _module_device(self.vae) - video = video.to(device=vae_device, dtype=torch.float32) - bsz, channels, frames, height, width = video.shape - flat_video = video.permute(0, 2, 1, 3, 4).reshape(bsz * frames, channels, height, width) - norm_flat_video = normalize_image_tensor( - flat_video, - [0.5, 0.5, 0.5], - [0.5, 0.5, 0.5], - inplace=False, - ) - norm_video = ( - norm_flat_video.reshape(bsz, frames, channels, height, width) - .permute(0, 2, 1, 3, 4) - .contiguous() - ) - with torch.autocast( - "cuda", - dtype=torch.bfloat16, - enabled=vae_device.type == "cuda", - ): - encoded = self.vae.encode(norm_video) - if hasattr(encoded, "latent_dist"): - latents = encoded.latent_dist.sample(generator) - else: - latents = encoded[0] if isinstance(encoded, tuple) else encoded + with self._vae_encode_context() as vae_device: + video = video.to(device=vae_device, dtype=torch.float32) + bsz, channels, frames, height, width = video.shape + flat_video = video.permute(0, 2, 1, 3, 4).reshape( + bsz * frames, channels, height, width + ) + norm_flat_video = normalize_image_tensor( + flat_video, + [0.5, 0.5, 0.5], + [0.5, 0.5, 0.5], + inplace=False, + ) + norm_video = ( + norm_flat_video.reshape(bsz, frames, channels, height, width) + .permute(0, 2, 1, 3, 4) + .contiguous() + ) + with torch.autocast( + "cuda", + dtype=torch.bfloat16, + enabled=vae_device.type == "cuda", + ): + encoded = self.vae.encode(norm_video) + if hasattr(encoded, "latent_dist"): + latents = encoded.latent_dist.sample(generator) + else: + latents = encoded[0] if isinstance(encoded, tuple) else encoded return self._vae_latent_to_dit(latents).to(latents) @torch.no_grad() @@ -330,7 +376,7 @@ def _decode_latents( self, latents: torch.Tensor, ) -> List[np.ndarray]: - vae_device = _module_device(self.vae) + vae_device = self._vae_execution_device() vae_dtype = _module_dtype(self.vae) vae_latents = self._dit_latent_to_vae(latents).to(device=vae_device, dtype=torch.float32) if vae_latents.ndim == 5: diff --git a/lingbot_video/pipeline_lingbot_video_i2v.py b/lingbot_video/pipeline_lingbot_video_i2v.py index 7216cb7..95f1a9a 100644 --- a/lingbot_video/pipeline_lingbot_video_i2v.py +++ b/lingbot_video/pipeline_lingbot_video_i2v.py @@ -15,7 +15,6 @@ LingBotVideoPipeline, LingBotVideoPipelineOutput, _group_global_rank, - _module_device, _module_dtype, _transformer_autocast, _transformer_timestep, @@ -123,11 +122,11 @@ def encode_image_latent( ) -> torch.Tensor: if self.vae is None: raise ValueError("`vae` is required to encode image latents.") - device = _module_device(self.vae) - pixel = pixel.to(device=device, dtype=torch.float32) - norm_pixel = (pixel - 0.5) / 0.5 - with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"): - latents = self.vae.encode(norm_pixel).latent_dist.sample(generator) + with self._vae_encode_context() as device: + pixel = pixel.to(device=device, dtype=torch.float32) + norm_pixel = (pixel - 0.5) / 0.5 + with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"): + latents = self.vae.encode(norm_pixel).latent_dist.sample(generator) mean = torch.tensor(self.vae.config.latents_mean, device=latents.device, dtype=torch.float32) std_inv = 1.0 / torch.tensor( diff --git a/lingbot_video/runner.py b/lingbot_video/runner.py index 99657b1..8b296ea 100644 --- a/lingbot_video/runner.py +++ b/lingbot_video/runner.py @@ -692,6 +692,7 @@ def _load_diffusers_pipe( mode: str, transformer_subfolder: str, defer_transformer_to_device: bool = False, + cpu_offload: str = "none", ) -> Any: pipeline_class = _pipeline_class_for_mode(mode) transformer = _load_transformer_component(model_dir, transformer_subfolder, dtype_map) @@ -707,6 +708,15 @@ def _load_diffusers_pipe( ) _log_progress(f"loaded pipeline mode={mode}") device = _default_device() + if cpu_offload != "none": + if defer_transformer_to_device: + raise ValueError("--cpu_offload cannot be combined with deferred transformer placement.") + _log_progress(f"enabling {cpu_offload} CPU offload to {device}") + if cpu_offload == "model": + pipe.enable_model_cpu_offload(device=device) + else: + pipe.enable_sequential_cpu_offload(device=device) + return pipe if defer_transformer_to_device: return _move_pipeline_aux_modules_to_device(pipe, device) _log_progress(f"moving pipeline to {device}") @@ -757,6 +767,7 @@ def _load_pipe( mode=args.mode, transformer_subfolder=args.transformer_subfolder, defer_transformer_to_device=defer_transformer_to_device, + cpu_offload=args.cpu_offload, ), "diffusers-reference", ) @@ -1009,6 +1020,11 @@ def main() -> None: parser.add_argument("--vae_dtype", default="fp32") parser.add_argument("--diffusers_attn_backend", default=os.environ.get("DIFFUSERS_ATTN_BACKEND", "")) parser.add_argument("--allow_tf32", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument( + "--cpu_offload", + choices=["none", "model", "sequential"], + default="none", + ) parser.add_argument( "--quiet_progress", action="store_true", @@ -1051,6 +1067,21 @@ def main() -> None: backend=args.backend, stderr=sys.stderr, ) + if args.cpu_offload != "none": + _, _, world_size = _distributed_env() + if args.engine != "diffusers": + raise ValueError("--cpu_offload supports the diffusers engine only.") + if not torch.cuda.is_available(): + raise RuntimeError("--cpu_offload requires CUDA.") + if ( + world_size != 1 + or args.enable_fsdp_inference + or args.cfg_parallel_degree > 1 + or args.context_parallel_degree > 1 + ): + raise ValueError( + "--cpu_offload requires single-process inference without FSDP, CFG, or CP." + ) args.negative_prompt = resolve_negative_prompt_arg( args.negative_prompt, args.negative_prompt_json,