Fix Diffusers CPU offload lifecycle

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,39 +342,41 @@ 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()
     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:

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(

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,