CogVideoX: remove per-step host syncs from the denoising loop, enable regional compilation

examples/profiling/README.md

@@ -105,6 +105,7 @@ This is non-invasive — it monkey-patches bound methods without modifying sourc
 | Wan | 480x832 | 81 | 2 | — |
 | LTX2 | 768x512 | 121 | 2 | `guidance_scale=4.0` |
 | QwenImage | 1024x1024 | — | 2 | `true_cfg_scale=4.0` |
+| CogVideoX | 720x480 | 49 | 4 | `guidance_scale=6.0`, `use_dynamic_cfg=True` |
 
 All configs use `output_type="latent"` by default (skip VAE decode for cleaner denoising-loop traces).
 

examples/profiling/profiling_pipelines.py

@@ -30,7 +30,14 @@
 
 def build_registry():
     """Build the pipeline config registry. Imports are deferred to avoid loading all pipelines upfront."""
-    from diffusers import Flux2KleinPipeline, FluxPipeline, LTX2Pipeline, QwenImagePipeline, WanPipeline
+    from diffusers import (
+        CogVideoXPipeline,
+        Flux2KleinPipeline,
+        FluxPipeline,
+        LTX2Pipeline,
+        QwenImagePipeline,
+        WanPipeline,
+    )
 
     return {
         "flux": PipelineProfilingConfig(
@@ -100,6 +107,24 @@ def build_registry():
                 "output_type": "latent",
             },
         ),
+        "cogvideox": PipelineProfilingConfig(
+            name="cogvideox",
+            pipeline_cls=CogVideoXPipeline,
+            pipeline_init_kwargs={
+                "pretrained_model_name_or_path": "THUDM/CogVideoX-2b",
+                "dtype": torch.float16,
+            },
+            pipeline_call_kwargs={
+                "prompt": PROMPT,
+                "height": 480,
+                "width": 720,
+                "num_frames": 49,
+                "num_inference_steps": 4,
+                "guidance_scale": 6.0,
+                "use_dynamic_cfg": True,
+                "output_type": "latent",
+            },
+        ),
         "qwenimage": PipelineProfilingConfig(
             name="qwenimage",
             pipeline_cls=QwenImagePipeline,
@@ -124,7 +149,7 @@ def main():
     parser = argparse.ArgumentParser(description="Profile diffusers pipelines with torch.profiler")
     parser.add_argument(
         "--pipeline",
-        choices=["flux", "flux2", "wan", "ltx2", "qwenimage", "all"],
+        choices=["flux", "flux2", "wan", "ltx2", "qwenimage", "cogvideox", "all"],
         required=True,
         help="Which pipeline to profile",
     )

src/diffusers/models/transformers/cogvideox_transformer_3d.py

@@ -216,6 +216,7 @@ class CogVideoXTransformer3DModel(ModelMixin, AttentionMixin, ConfigMixin, PeftA
     _skip_layerwise_casting_patterns = ["patch_embed", "norm"]
     _supports_gradient_checkpointing = True
     _no_split_modules = ["CogVideoXBlock", "CogVideoXPatchEmbed"]
+    _repeated_blocks = ["CogVideoXBlock"]
 
     @register_to_config
     def __init__(

src/diffusers/pipelines/cogvideo/pipeline_cogvideox.py

@@ -713,6 +713,8 @@ def __call__(
         num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
 
         with self.progress_bar(total=num_inference_steps) as progress_bar:
+            # read timesteps on the host; a device scalar would sync every step
+            timesteps_cpu = timesteps.tolist()
             # for DPM-solver++
             old_pred_original_sample = None
             for i, t in enumerate(timesteps):
@@ -740,22 +742,25 @@ def __call__(
 
                 # perform guidance
                 if use_dynamic_cfg:
+                    t_cpu = timesteps_cpu[i]
                     self._guidance_scale = 1 + guidance_scale * (
-                        (1 - math.cos(math.pi * ((num_inference_steps - t.item()) / num_inference_steps) ** 5.0)) / 2
+                        (1 - math.cos(math.pi * ((num_inference_steps - t_cpu) / num_inference_steps) ** 5.0)) / 2
                     )
                 if do_classifier_free_guidance:
                     noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                     noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
 
                 # compute the previous noisy sample x_t -> x_t-1
                 if not isinstance(self.scheduler, CogVideoXDPMScheduler):
-                    latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+                    latents = self.scheduler.step(
+                        noise_pred, timesteps_cpu[i], latents, **extra_step_kwargs, return_dict=False
+                    )[0]
                 else:
                     latents, old_pred_original_sample = self.scheduler.step(
                         noise_pred,
                         old_pred_original_sample,
-                        t,
-                        timesteps[i - 1] if i > 0 else None,
+                        timesteps_cpu[i],
+                        timesteps_cpu[i - 1] if i > 0 else None,
                         latents,
                         **extra_step_kwargs,
                         return_dict=False,

src/diffusers/pipelines/cogvideo/pipeline_cogvideox_fun_control.py

@@ -774,6 +774,8 @@ def __call__(
         num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
 
         with self.progress_bar(total=num_inference_steps) as progress_bar:
+            # read timesteps on the host; a device scalar would sync every step
+            timesteps_cpu = timesteps.tolist()
             # for DPM-solver++
             old_pred_original_sample = None
             for i, t in enumerate(timesteps):
@@ -806,15 +808,18 @@ def __call__(
 
                 # perform guidance
                 if use_dynamic_cfg:
+                    t_cpu = timesteps_cpu[i]
                     self._guidance_scale = 1 + guidance_scale * (
-                        (1 - math.cos(math.pi * ((num_inference_steps - t.item()) / num_inference_steps) ** 5.0)) / 2
+                        (1 - math.cos(math.pi * ((num_inference_steps - t_cpu) / num_inference_steps) ** 5.0)) / 2
                     )
                 if do_classifier_free_guidance:
                     noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                     noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
 
                 # compute the previous noisy sample x_t -> x_t-1
-                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+                latents = self.scheduler.step(
+                    noise_pred, timesteps_cpu[i], latents, **extra_step_kwargs, return_dict=False
+                )[0]
                 latents = latents.to(prompt_embeds.dtype)
 
                 # call the callback, if provided

src/diffusers/pipelines/cogvideo/pipeline_cogvideox_image2video.py

@@ -819,6 +819,8 @@ def __call__(
         num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
 
         with self.progress_bar(total=num_inference_steps) as progress_bar:
+            # read timesteps on the host; a device scalar would sync every step
+            timesteps_cpu = timesteps.tolist()
             # for DPM-solver++
             old_pred_original_sample = None
             for i, t in enumerate(timesteps):
@@ -850,22 +852,25 @@ def __call__(
 
                 # perform guidance
                 if use_dynamic_cfg:
+                    t_cpu = timesteps_cpu[i]
                     self._guidance_scale = 1 + guidance_scale * (
-                        (1 - math.cos(math.pi * ((num_inference_steps - t.item()) / num_inference_steps) ** 5.0)) / 2
+                        (1 - math.cos(math.pi * ((num_inference_steps - t_cpu) / num_inference_steps) ** 5.0)) / 2
                     )
                 if do_classifier_free_guidance:
                     noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                     noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
 
                 # compute the previous noisy sample x_t -> x_t-1
                 if not isinstance(self.scheduler, CogVideoXDPMScheduler):
-                    latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+                    latents = self.scheduler.step(
+                        noise_pred, timesteps_cpu[i], latents, **extra_step_kwargs, return_dict=False
+                    )[0]
                 else:
                     latents, old_pred_original_sample = self.scheduler.step(
                         noise_pred,
                         old_pred_original_sample,
-                        t,
-                        timesteps[i - 1] if i > 0 else None,
+                        timesteps_cpu[i],
+                        timesteps_cpu[i - 1] if i > 0 else None,
                         latents,
                         **extra_step_kwargs,
                         return_dict=False,

src/diffusers/pipelines/cogvideo/pipeline_cogvideox_video2video.py

@@ -794,6 +794,8 @@ def __call__(
         num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
 
         with self.progress_bar(total=num_inference_steps) as progress_bar:
+            # read timesteps on the host; a device scalar would sync every step
+            timesteps_cpu = timesteps.tolist()
             # for DPM-solver++
             old_pred_original_sample = None
             for i, t in enumerate(timesteps):
@@ -821,22 +823,25 @@ def __call__(
 
                 # perform guidance
                 if use_dynamic_cfg:
+                    t_cpu = timesteps_cpu[i]
                     self._guidance_scale = 1 + guidance_scale * (
-                        (1 - math.cos(math.pi * ((num_inference_steps - t.item()) / num_inference_steps) ** 5.0)) / 2
+                        (1 - math.cos(math.pi * ((num_inference_steps - t_cpu) / num_inference_steps) ** 5.0)) / 2
                     )
                 if do_classifier_free_guidance:
                     noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                     noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
 
                 # compute the previous noisy sample x_t -> x_t-1
                 if not isinstance(self.scheduler, CogVideoXDPMScheduler):
-                    latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+                    latents = self.scheduler.step(
+                        noise_pred, timesteps_cpu[i], latents, **extra_step_kwargs, return_dict=False
+                    )[0]
                 else:
                     latents, old_pred_original_sample = self.scheduler.step(
                         noise_pred,
                         old_pred_original_sample,
-                        t,
-                        timesteps[i - 1] if i > 0 else None,
+                        timesteps_cpu[i],
+                        timesteps_cpu[i - 1] if i > 0 else None,
                         latents,
                         **extra_step_kwargs,
                         return_dict=False,

tests/pipelines/cogvideo/test_cogvideox.py

@@ -26,6 +26,7 @@
     nightly,
     numpy_cosine_similarity_distance,
     require_torch_accelerator,
+    require_torch_gpu,
     torch_device,
 )
 from ..testing_utils import (
@@ -127,6 +128,31 @@ def get_dummy_inputs(self):
         }
 
 
+class TestCogVideoXPipelineHostSync(CogVideoXPipelineTesterConfig):
+    @require_torch_gpu
+    def test_denoising_loop_does_not_sync_with_host(self):
+        # A device-to-host copy inside the loop stalls the CPU every step and leaves the
+        # step uncapturable by a CUDA graph.
+        pipe = CogVideoXPipeline(**self.get_dummy_components()).to(torch_device)
+        pipe.set_progress_bar_config(disable=True)
+
+        inputs = self.get_dummy_inputs()
+        inputs["num_inference_steps"] = 3
+        inputs["use_dynamic_cfg"] = True
+        last_step = inputs["num_inference_steps"] - 1
+
+        def toggle_sync_debug(pipe, i, t, callback_kwargs):
+            # Arm once the one-off setup copies are done, disarm before decoding.
+            torch.cuda.set_sync_debug_mode("error" if i < last_step else "default")
+            return callback_kwargs
+
+        inputs["callback_on_step_end"] = toggle_sync_debug
+        try:
+            pipe(**inputs)
+        finally:
+            torch.cuda.set_sync_debug_mode("default")
+
+
 class TestCogVideoXPipeline(CogVideoXPipelineTesterConfig, PipelineTesterMixin):
     def test_inference(self):
         # Run on CPU: the expected slice below is CPU-specific.