diff --git a/examples/profiling/README.md b/examples/profiling/README.md index b38979577b60..ae850ae03e77 100644 --- a/examples/profiling/README.md +++ b/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). diff --git a/examples/profiling/profiling_pipelines.py b/examples/profiling/profiling_pipelines.py index 5a0b4bfe938b..524bb2c545ae 100644 --- a/examples/profiling/profiling_pipelines.py +++ b/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", ) diff --git a/src/diffusers/models/transformers/cogvideox_transformer_3d.py b/src/diffusers/models/transformers/cogvideox_transformer_3d.py index 08299f05e1b8..57a24f2e5cfb 100644 --- a/src/diffusers/models/transformers/cogvideox_transformer_3d.py +++ b/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__( diff --git a/src/diffusers/pipelines/cogvideo/pipeline_cogvideox.py b/src/diffusers/pipelines/cogvideo/pipeline_cogvideox.py index 9043abcab65e..ba785be896a5 100644 --- a/src/diffusers/pipelines/cogvideo/pipeline_cogvideox.py +++ b/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,8 +742,9 @@ 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) @@ -749,13 +752,15 @@ def __call__( # 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, diff --git a/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_fun_control.py b/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_fun_control.py index e2b45a08ee90..df2d990fd57b 100644 --- a/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_fun_control.py +++ b/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 diff --git a/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_image2video.py b/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_image2video.py index 42f5109bb877..ba1a50e24088 100644 --- a/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_image2video.py +++ b/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,8 +852,9 @@ 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) @@ -859,13 +862,15 @@ def __call__( # 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, diff --git a/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_video2video.py b/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_video2video.py index 3cd72b0c2126..bd788a2b2242 100644 --- a/src/diffusers/pipelines/cogvideo/pipeline_cogvideox_video2video.py +++ b/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,8 +823,9 @@ 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) @@ -830,13 +833,15 @@ def __call__( # 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, diff --git a/tests/pipelines/cogvideo/test_cogvideox.py b/tests/pipelines/cogvideo/test_cogvideox.py index 5bc34896e194..3e8e3fa828be 100644 --- a/tests/pipelines/cogvideo/test_cogvideox.py +++ b/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.