Skip to content

cogvideo model/pipeline review #13622

Description

@hlky

cogvideo model/pipeline review

Commit tested: 0f1abc4ae8b0eb2a3b40e82a310507281144c423

Review performed against the repository review rules.

Reviewed target pipelines, model files, lazy exports, docs/tests/examples, dtype/device paths, offload-facing behavior, attention processors, and coverage. Public imports/lazy loading looked consistent. I did not find separate actionable issues in pipeline_output.py or autoencoder_kl_cogvideox.py.

Execution: standalone repros were run with .venv/Scripts/python.exe; no full pytest suite was run.

Duplicate search: searched GitHub Issues and PRs for cogvideo, affected class/file names, and the specific failure modes. Exact duplicate found only for Issue 5: #9641. Related but not exact: #11133, #9972, #13586, PR #11368, PR #9333.

Issue 1: num_videos_per_prompt is accepted but ignored

Affected code:








Problem:
All four pipelines expose num_videos_per_prompt, but each __call__ resets it to 1 before prompt encoding and latent preparation. The text-only pipeline can already use the parameter correctly if that reset is removed. The conditioned pipelines also need image/video/control latents expanded to the effective batch, or they should reject values above 1.

Impact:
Users requesting multiple videos per prompt silently get one video per prompt. Batch behavior and callback tensor shapes are also misleading.

Reproduction:

importtorchfromdiffusersimportAutoencoderKLCogVideoX, CogVideoXDDIMScheduler, CogVideoXPipeline, CogVideoXTransformer3DModeldeftiny_pipe():
transformer=CogVideoXTransformer3DModel(
num_attention_heads=4, attention_head_dim=8, in_channels=4, out_channels=4,
time_embed_dim=2, text_embed_dim=32, num_layers=1,
sample_width=2, sample_height=2, sample_frames=9, patch_size=2,
temporal_compression_ratio=4, max_text_seq_length=16,
)
vae=AutoencoderKLCogVideoX(
in_channels=3, out_channels=3,
down_block_types=("CogVideoXDownBlock3D",) *4,
up_block_types=("CogVideoXUpBlock3D",) *4,
block_out_channels=(8, 8, 8, 8), latent_channels=4,
layers_per_block=1, norm_num_groups=2, temporal_compression_ratio=4,
)
returnCogVideoXPipeline(None, None, transformer, vae, CogVideoXDDIMScheduler())
pipe=tiny_pipe()
frames=pipe(
prompt_embeds=torch.zeros(1, 16, 32),
height=16, width=16, num_frames=5,
num_inference_steps=1, guidance_scale=1,
num_videos_per_prompt=2, output_type="latent",
).framesprint(frames.shape)
assertframes.shape[0] ==2, f"expected 2 videos, got {frames.shape[0]}"

Relevant precedent:

prompt_embeds, negative_prompt_embeds=self.encode_prompt(
prompt=prompt,
negative_prompt=negative_prompt,
do_classifier_free_guidance=self.do_classifier_free_guidance,
num_videos_per_prompt=num_videos_per_prompt,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
max_sequence_length=max_sequence_length,
device=device,
)
transformer_dtype=self.transformer.dtypeifself.transformerisnotNoneelseself.transformer_2.dtype
prompt_embeds=prompt_embeds.to(transformer_dtype)
ifnegative_prompt_embedsisnotNone:
negative_prompt_embeds=negative_prompt_embeds.to(transformer_dtype)
# 4. Prepare timesteps
self.scheduler.set_timesteps(num_inference_steps, device=device)
timesteps=self.scheduler.timesteps
# 5. Prepare latent variables
num_channels_latents= (
self.transformer.config.in_channels
ifself.transformerisnotNone
elseself.transformer_2.config.in_channels
)
latents=self.prepare_latents(
batch_size*num_videos_per_prompt,

) =self.encode_prompt(
prompt=prompt,
negative_prompt=negative_prompt,
do_classifier_free_guidance=self.do_classifier_free_guidance,
num_videos_per_prompt=num_videos_per_prompt,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
prompt_attention_mask=prompt_attention_mask,
negative_prompt_attention_mask=negative_prompt_attention_mask,
max_sequence_length=max_sequence_length,
device=device,
)
# 4. Prepare latent variables
num_channels_latents=self.transformer.config.in_channels
latents=self.prepare_latents(
batch_size*num_videos_per_prompt,

Suggested fix:

# Text-to-video: remove the forced reset.# num_videos_per_prompt = 1# Conditioned pipelines should either expand conditioning latents:def_repeat_to_effective_batch(tensor, batch_size, num_videos_per_prompt):
iftensor.shape[0] ==1:
returntensor.repeat_interleave(batch_size*num_videos_per_prompt, dim=0)
iftensor.shape[0] ==batch_size:
returntensor.repeat_interleave(num_videos_per_prompt, dim=0)
returntensor# Or reject unsupported requests until expansion is implemented:ifnum_videos_per_prompt!=1:
raiseValueError("`num_videos_per_prompt > 1` is not currently supported by this conditioned CogVideoX pipeline.")

Issue 2: CogVideoXFunControlPipeline crashes when control_video_latents is supplied

Affected code:

control_video: list[Image.Image] |None=None,
height: int|None=None,
width: int|None=None,
num_inference_steps: int=50,
timesteps: list[int] |None=None,
guidance_scale: float=6,
use_dynamic_cfg: bool=False,
num_videos_per_prompt: int=1,
eta: float=0.0,
generator: torch.Generator|list[torch.Generator] |None=None,
latents: torch.Tensor|None=None,
control_video_latents: torch.Tensor|None=None,

ifcontrol_video_latentsisNone:
control_video=self.video_processor.preprocess_video(control_video, height=height, width=width)
control_video=control_video.to(device=device, dtype=prompt_embeds.dtype)
_, control_video_latents=self.prepare_control_latents(None, control_video)
control_video_latents=control_video_latents.permute(0, 2, 1, 3, 4)

Problem:
The API documents control_video_latents, and check_inputs only rejects passing both raw control video and latents. But __call__ unconditionally runs prepare_control_latents(None, control_video) after the preprocessing branch. When only control_video_latents is supplied, control_video is None, so the user tensor is discarded and the next .permute() crashes.

Impact:
The precomputed control-latent path is unusable.

Reproduction:

importtorchfromdiffusersimportAutoencoderKLCogVideoX, CogVideoXDDIMScheduler, CogVideoXFunControlPipeline, CogVideoXTransformer3DModeltransformer=CogVideoXTransformer3DModel(
num_attention_heads=4, attention_head_dim=8, in_channels=8, out_channels=4,
time_embed_dim=2, text_embed_dim=32, num_layers=1,
sample_width=2, sample_height=2, sample_frames=9, patch_size=2,
temporal_compression_ratio=4, max_text_seq_length=16,
)
vae=AutoencoderKLCogVideoX(
in_channels=3, out_channels=3,
down_block_types=("CogVideoXDownBlock3D",) *4,
up_block_types=("CogVideoXUpBlock3D",) *4,
block_out_channels=(8, 8, 8, 8), latent_channels=4,
layers_per_block=1, norm_num_groups=2, temporal_compression_ratio=4,
)
pipe=CogVideoXFunControlPipeline(None, None, transformer, vae, CogVideoXDDIMScheduler())
try:
pipe(
prompt_embeds=torch.zeros(1, 16, 32),
control_video_latents=torch.zeros(1, 4, 2, 2, 2),
height=16, width=16, num_inference_steps=1,
guidance_scale=1, output_type="latent",
)
exceptExceptionase:
print(type(e).__name__, e)

Relevant precedent:
The raw-image/video latent paths in the other CogVideoX pipelines keep the precomputed latents branch separate from preprocessing.

Suggested fix:

ifcontrol_video_latentsisNone:
ifcontrol_videoisNone:
raiseValueError("Provide either `control_video` or `control_video_latents`.")
control_video=self.video_processor.preprocess_video(control_video, height=height, width=width)
control_video=control_video.to(device=device, dtype=prompt_embeds.dtype)
_, control_video_latents=self.prepare_control_latents(None, control_video)
else:
control_video_latents=control_video_latents.to(device=device, dtype=prompt_embeds.dtype)
control_video_latents=control_video_latents.permute(0, 2, 1, 3, 4)

Issue 3: Supplied prompt_embeds and latents are not cast to execution dtype

Affected code:

device=deviceorself._execution_device
prompt= [prompt] ifisinstance(prompt, str) elseprompt
ifpromptisnotNone:
batch_size=len(prompt)
else:
batch_size=prompt_embeds.shape[0]
ifprompt_embedsisNone:
prompt_embeds=self._get_t5_prompt_embeds(
prompt=prompt,
num_videos_per_prompt=num_videos_per_prompt,
max_sequence_length=max_sequence_length,
device=device,
dtype=dtype,
)
ifdo_classifier_free_guidanceandnegative_prompt_embedsisNone:
negative_prompt=negative_promptor""
negative_prompt=batch_size* [negative_prompt] ifisinstance(negative_prompt, str) elsenegative_prompt
ifpromptisnotNoneandtype(prompt) isnottype(negative_prompt):
raiseTypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
f" {type(prompt)}."
)
elifbatch_size!=len(negative_prompt):
raiseValueError(
f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
" the batch size of `prompt`."
)
negative_prompt_embeds=self._get_t5_prompt_embeds(
prompt=negative_prompt,
num_videos_per_prompt=num_videos_per_prompt,
max_sequence_length=max_sequence_length,
device=device,
dtype=dtype,
)
returnprompt_embeds, negative_prompt_embeds

iflatentsisNone:
latents=randn_tensor(shape, generator=generator, device=device, dtype=dtype)
else:
latents=latents.to(device)
# scale the initial noise by the standard deviation required by the scheduler
latents=latents*self.scheduler.init_noise_sigma

defencode_prompt(
self,
prompt: str|list[str],
negative_prompt: str|list[str] |None=None,
do_classifier_free_guidance: bool=True,
num_videos_per_prompt: int=1,
prompt_embeds: torch.Tensor|None=None,
negative_prompt_embeds: torch.Tensor|None=None,
max_sequence_length: int=226,
device: torch.device|None=None,
dtype: torch.dtype|None=None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
Whether to use classifier free guidance or not.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
Number of videos that should be generated per prompt. torch device to place the resulting embeddings on
prompt_embeds (`torch.Tensor`, *optional*):
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
provided, text embeddings will be generated from `prompt` input argument.
negative_prompt_embeds (`torch.Tensor`, *optional*):
Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
argument.
device: (`torch.device`, *optional*):
torch device
dtype: (`torch.dtype`, *optional*):
torch dtype
"""
device=deviceorself._execution_device
prompt= [prompt] ifisinstance(prompt, str) elseprompt
ifpromptisnotNone:
batch_size=len(prompt)
else:
batch_size=prompt_embeds.shape[0]
ifprompt_embedsisNone:
prompt_embeds=self._get_t5_prompt_embeds(
prompt=prompt,
num_videos_per_prompt=num_videos_per_prompt,
max_sequence_length=max_sequence_length,
device=device,
dtype=dtype,
)
ifdo_classifier_free_guidanceandnegative_prompt_embedsisNone:
negative_prompt=negative_promptor""
negative_prompt=batch_size* [negative_prompt] ifisinstance(negative_prompt, str) elsenegative_prompt
ifpromptisnotNoneandtype(prompt) isnottype(negative_prompt):
raiseTypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
f" {type(prompt)}."
)
elifbatch_size!=len(negative_prompt):
raiseValueError(
f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
" the batch size of `prompt`."
)
negative_prompt_embeds=self._get_t5_prompt_embeds(
prompt=negative_prompt,
num_videos_per_prompt=num_videos_per_prompt,
max_sequence_length=max_sequence_length,
device=device,
dtype=dtype,
)
returnprompt_embeds, negative_prompt_embeds

iflatentsisNone:
latents=randn_tensor(shape, generator=generator, device=device, dtype=dtype)
else:
latents=latents.to(device)
# scale the initial noise by the standard deviation required by the scheduler
latents=latents*self.scheduler.init_noise_sigma

defencode_prompt(
self,
prompt: str|list[str],
negative_prompt: str|list[str] |None=None,
do_classifier_free_guidance: bool=True,
num_videos_per_prompt: int=1,
prompt_embeds: torch.Tensor|None=None,
negative_prompt_embeds: torch.Tensor|None=None,
max_sequence_length: int=226,
device: torch.device|None=None,
dtype: torch.dtype|None=None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
Whether to use classifier free guidance or not.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
Number of videos that should be generated per prompt. torch device to place the resulting embeddings on
prompt_embeds (`torch.Tensor`, *optional*):
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
provided, text embeddings will be generated from `prompt` input argument.
negative_prompt_embeds (`torch.Tensor`, *optional*):
Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
argument.
device: (`torch.device`, *optional*):
torch device
dtype: (`torch.dtype`, *optional*):
torch dtype
"""
device=deviceorself._execution_device
prompt= [prompt] ifisinstance(prompt, str) elseprompt
ifpromptisnotNone:
batch_size=len(prompt)
else:
batch_size=prompt_embeds.shape[0]
ifprompt_embedsisNone:
prompt_embeds=self._get_t5_prompt_embeds(
prompt=prompt,
num_videos_per_prompt=num_videos_per_prompt,
max_sequence_length=max_sequence_length,
device=device,
dtype=dtype,
)
ifdo_classifier_free_guidanceandnegative_prompt_embedsisNone:
negative_prompt=negative_promptor""
negative_prompt=batch_size* [negative_prompt] ifisinstance(negative_prompt, str) elsenegative_prompt
ifpromptisnotNoneandtype(prompt) isnottype(negative_prompt):
raiseTypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
f" {type(prompt)}."
)
elifbatch_size!=len(negative_prompt):
raiseValueError(
f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
" the batch size of `prompt`."
)
negative_prompt_embeds=self._get_t5_prompt_embeds(
prompt=negative_prompt,
num_videos_per_prompt=num_videos_per_prompt,
max_sequence_length=max_sequence_length,
device=device,
dtype=dtype,
)
returnprompt_embeds, negative_prompt_embeds

iflatentsisNone:
latents=randn_tensor(shape, generator=generator, device=device, dtype=dtype)
else:
latents=latents.to(device)
# scale the initial noise by the standard deviation required by the scheduler
latents=latents*self.scheduler.init_noise_sigma

defencode_prompt(
self,
prompt: str|list[str],
negative_prompt: str|list[str] |None=None,
do_classifier_free_guidance: bool=True,
num_videos_per_prompt: int=1,
prompt_embeds: torch.Tensor|None=None,
negative_prompt_embeds: torch.Tensor|None=None,
max_sequence_length: int=226,
device: torch.device|None=None,
dtype: torch.dtype|None=None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
prompt (`str` or `list[str]`, *optional*):
prompt to be encoded
negative_prompt (`str` or `list[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
Whether to use classifier free guidance or not.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
Number of videos that should be generated per prompt. torch device to place the resulting embeddings on
prompt_embeds (`torch.Tensor`, *optional*):
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
provided, text embeddings will be generated from `prompt` input argument.
negative_prompt_embeds (`torch.Tensor`, *optional*):
Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
argument.
device: (`torch.device`, *optional*):
torch device
dtype: (`torch.dtype`, *optional*):
torch dtype
"""
device=deviceorself._execution_device
prompt= [prompt] ifisinstance(prompt, str) elseprompt
ifpromptisnotNone:
batch_size=len(prompt)
else:
batch_size=prompt_embeds.shape[0]
ifprompt_embedsisNone:
prompt_embeds=self._get_t5_prompt_embeds(
prompt=prompt,
num_videos_per_prompt=num_videos_per_prompt,
max_sequence_length=max_sequence_length,
device=device,
dtype=dtype,
)
ifdo_classifier_free_guidanceandnegative_prompt_embedsisNone:
negative_prompt=negative_promptor""
negative_prompt=batch_size* [negative_prompt] ifisinstance(negative_prompt, str) elsenegative_prompt
ifpromptisnotNoneandtype(prompt) isnottype(negative_prompt):
raiseTypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
f" {type(prompt)}."
)
elifbatch_size!=len(negative_prompt):
raiseValueError(
f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
" the batch size of `prompt`."
)
negative_prompt_embeds=self._get_t5_prompt_embeds(
prompt=negative_prompt,
num_videos_per_prompt=num_videos_per_prompt,
max_sequence_length=max_sequence_length,
device=device,
dtype=dtype,
)
returnprompt_embeds, negative_prompt_embeds

noise=randn_tensor(shape, generator=generator, device=device, dtype=dtype)
latents=self.scheduler.add_noise(init_latents, noise, timestep)
else:
latents=latents.to(device)
# scale the initial noise by the standard deviation required by the scheduler
latents=latents*self.scheduler.init_noise_sigma

Problem:
Generated prompt embeddings are cast, but user-supplied prompt_embeds and negative_prompt_embeds are returned unchanged. User-supplied latents are moved to device but not dtype. With a bf16/fp16 pipeline and fp32 tensors, transformer projections hit dtype mismatches.

Impact:
Documented advanced inputs break mixed precision inference.

Reproduction:

importtorchfromdiffusersimportAutoencoderKLCogVideoX, CogVideoXDDIMScheduler, CogVideoXPipeline, CogVideoXTransformer3DModeltransformer=CogVideoXTransformer3DModel(
num_attention_heads=4, attention_head_dim=8, in_channels=4, out_channels=4,
time_embed_dim=2, text_embed_dim=32, num_layers=1,
sample_width=2, sample_height=2, sample_frames=9, patch_size=2,
temporal_compression_ratio=4, max_text_seq_length=16,
)
vae=AutoencoderKLCogVideoX(
in_channels=3, out_channels=3,
down_block_types=("CogVideoXDownBlock3D",) *4,
up_block_types=("CogVideoXUpBlock3D",) *4,
block_out_channels=(8, 8, 8, 8), latent_channels=4,
layers_per_block=1, norm_num_groups=2, temporal_compression_ratio=4,
)
pipe=CogVideoXPipeline(None, None, transformer, vae, CogVideoXDDIMScheduler()).to(dtype=torch.bfloat16)
try:
pipe(
prompt_embeds=torch.zeros(1, 16, 32, dtype=torch.float32),
height=16, width=16, num_frames=5,
num_inference_steps=1, guidance_scale=1, output_type="latent",
)
exceptRuntimeErrorase:
print(e)

Relevant precedent:

transformer_dtype=self.transformer.dtypeifself.transformerisnotNoneelseself.transformer_2.dtype
prompt_embeds=prompt_embeds.to(transformer_dtype)
ifnegative_prompt_embedsisnotNone:
negative_prompt_embeds=negative_prompt_embeds.to(transformer_dtype)

iflatentsisnotNone:
returnlatents.to(device=device, dtype=dtype)

Suggested fix:

# After prompt embedding selection in encode_prompt:prompt_embeds=prompt_embeds.to(device=device, dtype=dtype)
ifnegative_prompt_embedsisnotNone:
negative_prompt_embeds=negative_prompt_embeds.to(device=device, dtype=dtype)
# In prepare_latents branches:latents=latents.to(device=device, dtype=dtype)

Issue 4: Spatial validation accepts sizes that later fail patchification

Affected code:

ifheight%8!=0orwidth%8!=0:
raiseValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")

ifheight%8!=0orwidth%8!=0:
raiseValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")

ifheight%8!=0orwidth%8!=0:
raiseValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")

ifheight%8!=0orwidth%8!=0:
raiseValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")

p=self.config.patch_size
p_t=self.config.patch_size_t
ifp_tisNone:
output=hidden_states.reshape(batch_size, num_frames, height//p, width//p, -1, p, p)
output=output.permute(0, 1, 4, 2, 5, 3, 6).flatten(5, 6).flatten(3, 4)
else:
output=hidden_states.reshape(
batch_size, (num_frames+p_t-1) //p_t, height//p, width//p, -1, p_t, p, p
)
output=output.permute(0, 1, 5, 4, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(1, 2)

Problem:
Pipelines only require height and width to be divisible by 8, the VAE scale factor. The transformer then patchifies latents with patch_size=2, so the original size must usually be divisible by 8 * 2 = 16. For example, 24x24 passes validation but produces latent 3x3, which fails later.

Impact:
Users get a late low-level tensor shape error instead of an actionable validation error.

Reproduction:

importtorchfromdiffusersimportAutoencoderKLCogVideoX, CogVideoXDDIMScheduler, CogVideoXPipeline, CogVideoXTransformer3DModeltransformer=CogVideoXTransformer3DModel(
num_attention_heads=4, attention_head_dim=8, in_channels=4, out_channels=4,
time_embed_dim=2, text_embed_dim=32, num_layers=1,
sample_width=2, sample_height=2, sample_frames=9, patch_size=2,
temporal_compression_ratio=4, max_text_seq_length=16,
)
vae=AutoencoderKLCogVideoX(
in_channels=3, out_channels=3,
down_block_types=("CogVideoXDownBlock3D",) *4,
up_block_types=("CogVideoXUpBlock3D",) *4,
block_out_channels=(8, 8, 8, 8), latent_channels=4,
layers_per_block=1, norm_num_groups=2, temporal_compression_ratio=4,
)
pipe=CogVideoXPipeline(None, None, transformer, vae, CogVideoXDDIMScheduler())
try:
pipe(
prompt_embeds=torch.zeros(1, 16, 32),
height=24, width=24, num_frames=5,
num_inference_steps=1, guidance_scale=1, output_type="latent",
)
exceptExceptionase:
print(type(e).__name__, e)

Relevant precedent:

patch_size= (
self.transformer.config.patch_size
ifself.transformerisnotNone
elseself.transformer_2.config.patch_size
)
h_multiple_of=self.vae_scale_factor_spatial*patch_size[1]
w_multiple_of=self.vae_scale_factor_spatial*patch_size[2]
calc_height=height//h_multiple_of*h_multiple_of
calc_width=width//w_multiple_of*w_multiple_of
ifheight!=calc_heightorwidth!=calc_width:
logger.warning(
f"`height` and `width` must be multiples of ({h_multiple_of}, {w_multiple_of}) for proper patchification. "
f"Adjusting ({height}, {width}) -> ({calc_height}, {calc_width})."
)
height, width=calc_height, calc_width

raiseValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.")

Suggested fix:

spatial_multiple=self.vae_scale_factor_spatial*self.transformer.config.patch_sizeifheight%spatial_multiple!=0orwidth%spatial_multiple!=0:
raiseValueError(
f"`height` and `width` have to be divisible by {spatial_multiple} "f"because CogVideoX patchifies VAE latents with patch_size={self.transformer.config.patch_size}; "f"got {height} and {width}."
)

Issue 5: Attention backend selection cannot affect CogVideoX attention

Affected code:

from ..attention_processorimportCogVideoXAttnProcessor2_0, FusedCogVideoXAttnProcessor2_0

self.attn1=Attention(
query_dim=dim,
dim_head=attention_head_dim,
heads=num_attention_heads,
qk_norm="layer_norm"ifqk_normelseNone,
eps=1e-6,
bias=attention_bias,
out_bias=attention_out_bias,
processor=CogVideoXAttnProcessor2_0(),
)

# Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.fuse_qkv_projections with FusedAttnProcessor2_0->FusedCogVideoXAttnProcessor2_0
deffuse_qkv_projections(self):
"""
Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value)
are fused. For cross-attention modules, key and value projection matrices are fused.
> [!WARNING] > This API is 🧪 experimental.
"""
self.original_attn_processors=None
for_, attn_processorinself.attn_processors.items():
if"Added"instr(attn_processor.__class__.__name__):
raiseValueError("`fuse_qkv_projections()` is not supported for models having added KV projections.")
self.original_attn_processors=self.attn_processors
formoduleinself.modules():
ifisinstance(module, Attention):
module.fuse_projections(fuse=True)
self.set_attn_processor(FusedCogVideoXAttnProcessor2_0())

classCogVideoXAttnProcessor2_0:
r"""
Processor for implementing scaled dot-product attention for the CogVideoX model. It applies a rotary embedding on
query and key vectors, but does not include spatial normalization.
"""
def__init__(self):
ifnothasattr(F, "scaled_dot_product_attention"):
raiseImportError("CogVideoXAttnProcessor requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
def__call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
attention_mask: torch.Tensor|None=None,
image_rotary_emb: torch.Tensor|None=None,
) ->torch.Tensor:
text_seq_length=encoder_hidden_states.size(1)
hidden_states=torch.cat([encoder_hidden_states, hidden_states], dim=1)
batch_size, sequence_length, _=hidden_states.shape
ifattention_maskisnotNone:
attention_mask=attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
attention_mask=attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
query=attn.to_q(hidden_states)
key=attn.to_k(hidden_states)
value=attn.to_v(hidden_states)
inner_dim=key.shape[-1]
head_dim=inner_dim//attn.heads
query=query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
key=key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
value=value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
ifattn.norm_qisnotNone:
query=attn.norm_q(query)
ifattn.norm_kisnotNone:
key=attn.norm_k(key)
# Apply RoPE if needed
ifimage_rotary_embisnotNone:
from .embeddingsimportapply_rotary_emb
query[:, :, text_seq_length:] =apply_rotary_emb(query[:, :, text_seq_length:], image_rotary_emb)
ifnotattn.is_cross_attention:
key[:, :, text_seq_length:] =apply_rotary_emb(key[:, :, text_seq_length:], image_rotary_emb)
hidden_states=F.scaled_dot_product_attention(
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False

classFusedCogVideoXAttnProcessor2_0:
r"""
Processor for implementing scaled dot-product attention for the CogVideoX model. It applies a rotary embedding on
query and key vectors, but does not include spatial normalization.
"""
def__init__(self):
ifnothasattr(F, "scaled_dot_product_attention"):
raiseImportError("CogVideoXAttnProcessor requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
def__call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
attention_mask: torch.Tensor|None=None,
image_rotary_emb: torch.Tensor|None=None,
) ->torch.Tensor:
text_seq_length=encoder_hidden_states.size(1)
hidden_states=torch.cat([encoder_hidden_states, hidden_states], dim=1)
batch_size, sequence_length, _= (
hidden_states.shapeifencoder_hidden_statesisNoneelseencoder_hidden_states.shape
)
ifattention_maskisnotNone:
attention_mask=attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
attention_mask=attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
qkv=attn.to_qkv(hidden_states)
split_size=qkv.shape[-1] //3
query, key, value=torch.split(qkv, split_size, dim=-1)
inner_dim=key.shape[-1]
head_dim=inner_dim//attn.heads
query=query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
key=key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
value=value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
ifattn.norm_qisnotNone:
query=attn.norm_q(query)
ifattn.norm_kisnotNone:
key=attn.norm_k(key)
# Apply RoPE if needed
ifimage_rotary_embisnotNone:
from .embeddingsimportapply_rotary_emb
query[:, :, text_seq_length:] =apply_rotary_emb(query[:, :, text_seq_length:], image_rotary_emb)
ifnotattn.is_cross_attention:
key[:, :, text_seq_length:] =apply_rotary_emb(key[:, :, text_seq_length:], image_rotary_emb)
hidden_states=F.scaled_dot_product_attention(
query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False

Problem:
CogVideoXAttnProcessor2_0 and FusedCogVideoXAttnProcessor2_0 are shared legacy processors without _attention_backend / _parallel_config, and they call F.scaled_dot_product_attention directly. CogVideoXTransformer3DModel.set_attention_backend(...) therefore leaves them unchanged and cannot route through the dispatcher.

Impact:
CogVideoX cannot use the model-level attention backend infrastructure consistently, including alternate kernels and parallel attention integrations covered by the review rules.

Reproduction:

fromdiffusersimportCogVideoXTransformer3DModelmodel=CogVideoXTransformer3DModel(
num_attention_heads=2, attention_head_dim=8, in_channels=4, out_channels=4,
time_embed_dim=2, text_embed_dim=8, num_layers=1,
sample_width=8, sample_height=8, sample_frames=8,
patch_size=2, temporal_compression_ratio=4, max_text_seq_length=8,
)
model.set_attention_backend("native")
processors=list(model.attn_processors.values())
print([type(p).__name__forpinprocessors])
print([hasattr(p, "_attention_backend") forpinprocessors])
assertall(hasattr(p, "_attention_backend") forpinprocessors)

Relevant precedent:

_attention_backend=None
_parallel_config=None
def__init__(self):
ifnothasattr(F, "scaled_dot_product_attention"):
raiseImportError(
"WanAttnProcessor requires PyTorch 2.0. To use it, please upgrade PyTorch to version 2.0 or higher."
)
def__call__(
self,
attn: "WanAttention",
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor|None=None,
attention_mask: torch.Tensor|None=None,
rotary_emb: tuple[torch.Tensor, torch.Tensor] |None=None,
) ->torch.Tensor:
encoder_hidden_states_img=None
ifattn.add_k_projisnotNone:
# 512 is the context length of the text encoder, hardcoded for now
image_context_length=encoder_hidden_states.shape[1] -512
encoder_hidden_states_img=encoder_hidden_states[:, :image_context_length]
encoder_hidden_states=encoder_hidden_states[:, image_context_length:]
query, key, value=_get_qkv_projections(attn, hidden_states, encoder_hidden_states)
query=attn.norm_q(query)
key=attn.norm_k(key)
query=query.unflatten(2, (attn.heads, -1))
key=key.unflatten(2, (attn.heads, -1))
value=value.unflatten(2, (attn.heads, -1))
ifrotary_embisnotNone:
defapply_rotary_emb(
hidden_states: torch.Tensor,
freqs_cos: torch.Tensor,
freqs_sin: torch.Tensor,
):
x1, x2=hidden_states.unflatten(-1, (-1, 2)).unbind(-1)
cos=freqs_cos[..., 0::2]
sin=freqs_sin[..., 1::2]
out=torch.empty_like(hidden_states)
out[..., 0::2] =x1*cos-x2*sin
out[..., 1::2] =x1*sin+x2*cos
returnout.type_as(hidden_states)
query=apply_rotary_emb(query, *rotary_emb)
key=apply_rotary_emb(key, *rotary_emb)
# I2V task
hidden_states_img=None
ifencoder_hidden_states_imgisnotNone:
key_img, value_img=_get_added_kv_projections(attn, encoder_hidden_states_img)
key_img=attn.norm_added_k(key_img)
key_img=key_img.unflatten(2, (attn.heads, -1))
value_img=value_img.unflatten(2, (attn.heads, -1))
hidden_states_img=dispatch_attention_fn(
query,
key_img,
value_img,
attn_mask=None,
dropout_p=0.0,
is_causal=False,
backend=self._attention_backend,
# Reference: https://github.com/huggingface/diffusers/pull/12909
parallel_config=None,
)
hidden_states_img=hidden_states_img.flatten(2, 3)
hidden_states_img=hidden_states_img.type_as(query)
hidden_states=dispatch_attention_fn(

from ..attention_dispatchimportdispatch_attention_fn
from ..cache_utilsimportCacheMixin
from ..embeddingsimport (
CombinedTimestepGuidanceTextProjEmbeddings,
CombinedTimestepTextProjEmbeddings,
apply_rotary_emb,
get_1d_rotary_pos_embed,
)
from ..modeling_outputsimportTransformer2DModelOutput
from ..modeling_utilsimportModelMixin
from ..normalizationimportAdaLayerNormContinuous, AdaLayerNormZero, AdaLayerNormZeroSingle
logger=logging.get_logger(__name__) # pylint: disable=invalid-name
def_get_projections(attn: "FluxAttention", hidden_states, encoder_hidden_states=None):
query=attn.to_q(hidden_states)
key=attn.to_k(hidden_states)
value=attn.to_v(hidden_states)
encoder_query=encoder_key=encoder_value=None
ifencoder_hidden_statesisnotNoneandattn.added_kv_proj_dimisnotNone:
encoder_query=attn.add_q_proj(encoder_hidden_states)
encoder_key=attn.add_k_proj(encoder_hidden_states)
encoder_value=attn.add_v_proj(encoder_hidden_states)
returnquery, key, value, encoder_query, encoder_key, encoder_value
def_get_fused_projections(attn: "FluxAttention", hidden_states, encoder_hidden_states=None):
query, key, value=attn.to_qkv(hidden_states).chunk(3, dim=-1)
encoder_query=encoder_key=encoder_value= (None,)
ifencoder_hidden_statesisnotNoneandhasattr(attn, "to_added_qkv"):
encoder_query, encoder_key, encoder_value=attn.to_added_qkv(encoder_hidden_states).chunk(3, dim=-1)
returnquery, key, value, encoder_query, encoder_key, encoder_value
def_get_qkv_projections(attn: "FluxAttention", hidden_states, encoder_hidden_states=None):
ifattn.fused_projections:
return_get_fused_projections(attn, hidden_states, encoder_hidden_states)
return_get_projections(attn, hidden_states, encoder_hidden_states)
classFluxAttnProcessor:
_attention_backend=None
_parallel_config=None
def__init__(self):
ifnothasattr(F, "scaled_dot_product_attention"):
raiseImportError(f"{self.__class__.__name__} requires PyTorch 2.0. Please upgrade your pytorch version.")
def__call__(
self,
attn: "FluxAttention",
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor=None,
attention_mask: torch.Tensor|None=None,
image_rotary_emb: torch.Tensor|None=None,
) ->torch.Tensor:
query, key, value, encoder_query, encoder_key, encoder_value=_get_qkv_projections(
attn, hidden_states, encoder_hidden_states
)
query=query.unflatten(-1, (attn.heads, -1))
key=key.unflatten(-1, (attn.heads, -1))
value=value.unflatten(-1, (attn.heads, -1))
query=attn.norm_q(query)
key=attn.norm_k(key)
ifattn.added_kv_proj_dimisnotNone:
encoder_query=encoder_query.unflatten(-1, (attn.heads, -1))
encoder_key=encoder_key.unflatten(-1, (attn.heads, -1))
encoder_value=encoder_value.unflatten(-1, (attn.heads, -1))
encoder_query=attn.norm_added_q(encoder_query)
encoder_key=attn.norm_added_k(encoder_key)
query=torch.cat([encoder_query, query], dim=1)
key=torch.cat([encoder_key, key], dim=1)
value=torch.cat([encoder_value, value], dim=1)
ifimage_rotary_embisnotNone:
query=apply_rotary_emb(query, image_rotary_emb, sequence_dim=1)
key=apply_rotary_emb(key, image_rotary_emb, sequence_dim=1)
hidden_states=dispatch_attention_fn(
query,
key,
value,
attn_mask=attention_mask,
backend=self._attention_backend,

Suggested fix:
Move CogVideoX attention processors into cogvideox_transformer_3d.py as model-local processors, add _attention_backend and _parallel_config, and replace direct SDPA calls with dispatcher calls while preserving text/video split and RoPE behavior:

hidden_states=dispatch_attention_fn(
query,
key,
value,
attn_mask=attention_mask,
dropout_p=0.0,
is_causal=False,
backend=self._attention_backend,
parallel_config=self._parallel_config,
)

Issue 6: Dynamic CFG uses scheduler timestep values as denoising progress

Affected code:

ifuse_dynamic_cfg:
self._guidance_scale=1+guidance_scale* (
(1-math.cos(math.pi* ((num_inference_steps-t.item()) /num_inference_steps) **5.0)) /2
)

ifuse_dynamic_cfg:
self._guidance_scale=1+guidance_scale* (
(1-math.cos(math.pi* ((num_inference_steps-t.item()) /num_inference_steps) **5.0)) /2
)

ifuse_dynamic_cfg:
self._guidance_scale=1+guidance_scale* (
(1-math.cos(math.pi* ((num_inference_steps-t.item()) /num_inference_steps) **5.0)) /2
)

ifuse_dynamic_cfg:
self._guidance_scale=1+guidance_scale* (
(1-math.cos(math.pi* ((num_inference_steps-t.item()) /num_inference_steps) **5.0)) /2
)

Problem:
This is already tracked by #9641, so I am not presenting it as new. The dynamic CFG formula uses t.item() from scheduler timesteps, not the denoising loop index/progress. With standard timesteps like 980, 960, ..., the expression does not represent normalized progress and can exceed the requested guidance scale because it computes 1 + guidance_scale * ....

Impact:
use_dynamic_cfg=True applies an unintuitive and scheduler-dependent guidance schedule across all CogVideoX pipelines.

Reproduction:

importmathfromdiffusersimportCogVideoXDDIMSchedulerguidance_scale=6num_inference_steps=50scheduler=CogVideoXDDIMScheduler()
scheduler.set_timesteps(num_inference_steps)
scales= [
1+guidance_scale* (
(1-math.cos(math.pi* ((num_inference_steps-t.item()) /num_inference_steps) **5.0)) /2
)
fortinscheduler.timesteps
]
print(scheduler.timesteps[:5].tolist())
print([round(x, 3) forxinscales[:12]], max(scales))
assertmax(scales) <=guidance_scale

Relevant precedent:
The intended behavior should be resolved in the existing issue: #9641

Suggested fix:
Use loop progress instead of scheduler timestep value, and confirm the intended max scale against the CogVideoX implementation:

progress= (num_inference_steps-i) /num_inference_stepsself._guidance_scale=1+ (guidance_scale-1) * ((1-math.cos(math.pi*progress**5.0)) /2)

Issue 7: Slow tests are missing for FunControl and Video-to-Video

Affected code:

# Copyright 2025 The HuggingFace Team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
importinspect
importunittest
importnumpyasnp
importtorch
fromPILimportImage
fromtransformersimportAutoConfig, AutoTokenizer, T5EncoderModel
fromdiffusersimportAutoencoderKLCogVideoX, CogVideoXFunControlPipeline, CogVideoXTransformer3DModel, DDIMScheduler
from ...testing_utilsimport (
enable_full_determinism,
torch_device,
)
from ..pipeline_paramsimportTEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_commonimport (
PipelineTesterMixin,
check_qkv_fusion_matches_attn_procs_length,
check_qkv_fusion_processors_exist,
to_np,
)
enable_full_determinism()
classCogVideoXFunControlPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
pipeline_class=CogVideoXFunControlPipeline
params=TEXT_TO_IMAGE_PARAMS- {"cross_attention_kwargs"}
batch_params=TEXT_TO_IMAGE_BATCH_PARAMS.union({"control_video"})
image_params=TEXT_TO_IMAGE_IMAGE_PARAMS
image_latents_params=TEXT_TO_IMAGE_IMAGE_PARAMS
required_optional_params=frozenset(
[
"num_inference_steps",
"generator",
"latents",
"return_dict",
"callback_on_step_end",
"callback_on_step_end_tensor_inputs",
]
)
test_xformers_attention=False
test_layerwise_casting=True
test_group_offloading=True
defget_dummy_components(self):
torch.manual_seed(0)
transformer=CogVideoXTransformer3DModel(
# Product of num_attention_heads * attention_head_dim must be divisible by 16 for 3D positional embeddings
# But, since we are using tiny-random-t5 here, we need the internal dim of CogVideoXTransformer3DModel
# to be 32. The internal dim is product of num_attention_heads and attention_head_dim
num_attention_heads=4,
attention_head_dim=8,
in_channels=8,
out_channels=4,
time_embed_dim=2,
text_embed_dim=32, # Must match with tiny-random-t5
num_layers=1,
sample_width=2, # latent width: 2 -> final width: 16
sample_height=2, # latent height: 2 -> final height: 16
sample_frames=9, # latent frames: (9 - 1) / 4 + 1 = 3 -> final frames: 9
patch_size=2,
temporal_compression_ratio=4,
max_text_seq_length=16,
)
torch.manual_seed(0)
vae=AutoencoderKLCogVideoX(
in_channels=3,
out_channels=3,
down_block_types=(
"CogVideoXDownBlock3D",
"CogVideoXDownBlock3D",
"CogVideoXDownBlock3D",
"CogVideoXDownBlock3D",
),
up_block_types=(
"CogVideoXUpBlock3D",
"CogVideoXUpBlock3D",
"CogVideoXUpBlock3D",
"CogVideoXUpBlock3D",
),
block_out_channels=(8, 8, 8, 8),
latent_channels=4,
layers_per_block=1,
norm_num_groups=2,
temporal_compression_ratio=4,
)
torch.manual_seed(0)
scheduler=DDIMScheduler()
config=AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
text_encoder=T5EncoderModel(config)
tokenizer=AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
components= {
"transformer": transformer,
"vae": vae,
"scheduler": scheduler,
"text_encoder": text_encoder,
"tokenizer": tokenizer,
}
returncomponents
defget_dummy_inputs(self, device, seed: int=0, num_frames: int=8):
ifstr(device).startswith("mps"):
generator=torch.manual_seed(seed)
else:
generator=torch.Generator(device=device).manual_seed(seed)
# Cannot reduce because convolution kernel becomes bigger than sample
height=16
width=16
control_video= [Image.new("RGB", (width, height))] *num_frames
inputs= {
"prompt": "dance monkey",
"negative_prompt": "",
"control_video": control_video,
"generator": generator,
"num_inference_steps": 2,
"guidance_scale": 6.0,
"height": height,
"width": width,
"max_sequence_length": 16,
"output_type": "pt",
}
returninputs
deftest_inference(self):
device="cpu"
components=self.get_dummy_components()
pipe=self.pipeline_class(**components)
pipe.to(device)
pipe.set_progress_bar_config(disable=None)
inputs=self.get_dummy_inputs(device)
video=pipe(**inputs).frames
generated_video=video[0]
self.assertEqual(generated_video.shape, (8, 3, 16, 16))
expected_video=torch.randn(8, 3, 16, 16)
max_diff=np.abs(generated_video-expected_video).max()
self.assertLessEqual(max_diff, 1e10)
deftest_callback_inputs(self):
sig=inspect.signature(self.pipeline_class.__call__)
has_callback_tensor_inputs="callback_on_step_end_tensor_inputs"insig.parameters
has_callback_step_end="callback_on_step_end"insig.parameters
ifnot (has_callback_tensor_inputsandhas_callback_step_end):
return
components=self.get_dummy_components()
pipe=self.pipeline_class(**components)
pipe=pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)
self.assertTrue(
hasattr(pipe, "_callback_tensor_inputs"),
f" {self.pipeline_class} should have `_callback_tensor_inputs` that defines a list of tensor variables its callback function can use as inputs",
)
defcallback_inputs_subset(pipe, i, t, callback_kwargs):
# iterate over callback args
fortensor_name, tensor_valueincallback_kwargs.items():
# check that we're only passing in allowed tensor inputs
asserttensor_nameinpipe._callback_tensor_inputs
returncallback_kwargs
defcallback_inputs_all(pipe, i, t, callback_kwargs):
fortensor_nameinpipe._callback_tensor_inputs:
asserttensor_nameincallback_kwargs
# iterate over callback args
fortensor_name, tensor_valueincallback_kwargs.items():
# check that we're only passing in allowed tensor inputs
asserttensor_nameinpipe._callback_tensor_inputs
returncallback_kwargs
inputs=self.get_dummy_inputs(torch_device)
# Test passing in a subset
inputs["callback_on_step_end"] =callback_inputs_subset
inputs["callback_on_step_end_tensor_inputs"] = ["latents"]
output=pipe(**inputs)[0]
# Test passing in a everything
inputs["callback_on_step_end"] =callback_inputs_all
inputs["callback_on_step_end_tensor_inputs"] =pipe._callback_tensor_inputs
output=pipe(**inputs)[0]
defcallback_inputs_change_tensor(pipe, i, t, callback_kwargs):
is_last=i== (pipe.num_timesteps-1)
ifis_last:
callback_kwargs["latents"] =torch.zeros_like(callback_kwargs["latents"])
returncallback_kwargs
inputs["callback_on_step_end"] =callback_inputs_change_tensor
inputs["callback_on_step_end_tensor_inputs"] =pipe._callback_tensor_inputs
output=pipe(**inputs)[0]
assertoutput.abs().sum() <1e10
deftest_inference_batch_single_identical(self):
self._test_inference_batch_single_identical(batch_size=3, expected_max_diff=1e-3)
deftest_attention_slicing_forward_pass(
self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3
):
ifnotself.test_attention_slicing:
return
components=self.get_dummy_components()
forkeyincomponents:
if"text_encoder"inkeyandhasattr(components[key], "eval"):
components[key].eval()
pipe=self.pipeline_class(**components)
forcomponentinpipe.components.values():
ifhasattr(component, "set_default_attn_processor"):
component.set_default_attn_processor()
pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)
generator_device="cpu"
inputs=self.get_dummy_inputs(generator_device)
output_without_slicing=pipe(**inputs)[0]
pipe.enable_attention_slicing(slice_size=1)
inputs=self.get_dummy_inputs(generator_device)
output_with_slicing1=pipe(**inputs)[0]
pipe.enable_attention_slicing(slice_size=2)
inputs=self.get_dummy_inputs(generator_device)
output_with_slicing2=pipe(**inputs)[0]
iftest_max_difference:
max_diff1=np.abs(to_np(output_with_slicing1) -to_np(output_without_slicing)).max()
max_diff2=np.abs(to_np(output_with_slicing2) -to_np(output_without_slicing)).max()
self.assertLess(
max(max_diff1, max_diff2),
expected_max_diff,
"Attention slicing should not affect the inference results",
)
deftest_vae_tiling(self, expected_diff_max: float=0.5):
# NOTE(aryan): This requires a higher expected_max_diff than other CogVideoX pipelines
generator_device="cpu"
components=self.get_dummy_components()
pipe=self.pipeline_class(**components)
pipe.to("cpu")
pipe.set_progress_bar_config(disable=None)
# Without tiling
inputs=self.get_dummy_inputs(generator_device)
inputs["height"] =inputs["width"] =128
output_without_tiling=pipe(**inputs)[0]
# With tiling
pipe.vae.enable_tiling(
tile_sample_min_height=96,
tile_sample_min_width=96,
tile_overlap_factor_height=1/12,
tile_overlap_factor_width=1/12,
)
inputs=self.get_dummy_inputs(generator_device)
inputs["height"] =inputs["width"] =128
output_with_tiling=pipe(**inputs)[0]
self.assertLess(
(to_np(output_without_tiling) -to_np(output_with_tiling)).max(),
expected_diff_max,
"VAE tiling should not affect the inference results",
)
deftest_fused_qkv_projections(self):
device="cpu"# ensure determinism for the device-dependent torch.Generator
components=self.get_dummy_components()
pipe=self.pipeline_class(**components)
pipe=pipe.to(device)
pipe.set_progress_bar_config(disable=None)
inputs=self.get_dummy_inputs(device)
frames=pipe(**inputs).frames# [B, F, C, H, W]
original_image_slice=frames[0, -2:, -1, -3:, -3:]
pipe.fuse_qkv_projections()
assertcheck_qkv_fusion_processors_exist(pipe.transformer), (
"Something wrong with the fused attention processors. Expected all the attention processors to be fused."
)
assertcheck_qkv_fusion_matches_attn_procs_length(
pipe.transformer, pipe.transformer.original_attn_processors
), "Something wrong with the attention processors concerning the fused QKV projections."
inputs=self.get_dummy_inputs(device)
frames=pipe(**inputs).frames
image_slice_fused=frames[0, -2:, -1, -3:, -3:]
pipe.transformer.unfuse_qkv_projections()
inputs=self.get_dummy_inputs(device)
frames=pipe(**inputs).frames
image_slice_disabled=frames[0, -2:, -1, -3:, -3:]
assertnp.allclose(original_image_slice, image_slice_fused, atol=1e-3, rtol=1e-3), (
"Fusion of QKV projections shouldn't affect the outputs."
)
assertnp.allclose(image_slice_fused, image_slice_disabled, atol=1e-3, rtol=1e-3), (
"Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled."
)
assertnp.allclose(original_image_slice, image_slice_disabled, atol=1e-2, rtol=1e-2), (
"Original outputs should match when fused QKV projections are disabled."
)

# Copyright 2025 The HuggingFace Team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
importinspect
importunittest
importnumpyasnp
importtorch
fromPILimportImage
fromtransformersimportAutoConfig, AutoTokenizer, T5EncoderModel
fromdiffusersimportAutoencoderKLCogVideoX, CogVideoXTransformer3DModel, CogVideoXVideoToVideoPipeline, DDIMScheduler
from ...testing_utilsimportenable_full_determinism, torch_device
from ..pipeline_paramsimportTEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_commonimport (
PipelineTesterMixin,
check_qkv_fusion_matches_attn_procs_length,
check_qkv_fusion_processors_exist,
to_np,
)
enable_full_determinism()
classCogVideoXVideoToVideoPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
pipeline_class=CogVideoXVideoToVideoPipeline
params=TEXT_TO_IMAGE_PARAMS- {"cross_attention_kwargs"}
batch_params=TEXT_TO_IMAGE_BATCH_PARAMS.union({"video"})
image_params=TEXT_TO_IMAGE_IMAGE_PARAMS
image_latents_params=TEXT_TO_IMAGE_IMAGE_PARAMS
required_optional_params=frozenset(
[
"num_inference_steps",
"generator",
"latents",
"return_dict",
"callback_on_step_end",
"callback_on_step_end_tensor_inputs",
]
)
test_xformers_attention=False
defget_dummy_components(self):
torch.manual_seed(0)
transformer=CogVideoXTransformer3DModel(
# Product of num_attention_heads * attention_head_dim must be divisible by 16 for 3D positional embeddings
# But, since we are using tiny-random-t5 here, we need the internal dim of CogVideoXTransformer3DModel
# to be 32. The internal dim is product of num_attention_heads and attention_head_dim
num_attention_heads=4,
attention_head_dim=8,
in_channels=4,
out_channels=4,
time_embed_dim=2,
text_embed_dim=32, # Must match with tiny-random-t5
num_layers=1,
sample_width=2, # latent width: 2 -> final width: 16
sample_height=2, # latent height: 2 -> final height: 16
sample_frames=9, # latent frames: (9 - 1) / 4 + 1 = 3 -> final frames: 9
patch_size=2,
temporal_compression_ratio=4,
max_text_seq_length=16,
)
torch.manual_seed(0)
vae=AutoencoderKLCogVideoX(
in_channels=3,
out_channels=3,
down_block_types=(
"CogVideoXDownBlock3D",
"CogVideoXDownBlock3D",
"CogVideoXDownBlock3D",
"CogVideoXDownBlock3D",
),
up_block_types=(
"CogVideoXUpBlock3D",
"CogVideoXUpBlock3D",
"CogVideoXUpBlock3D",
"CogVideoXUpBlock3D",
),
block_out_channels=(8, 8, 8, 8),
latent_channels=4,
layers_per_block=1,
norm_num_groups=2,
temporal_compression_ratio=4,
)
torch.manual_seed(0)
scheduler=DDIMScheduler()
config=AutoConfig.from_pretrained("hf-internal-testing/tiny-random-t5")
text_encoder=T5EncoderModel(config)
tokenizer=AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5")
components= {
"transformer": transformer,
"vae": vae,
"scheduler": scheduler,
"text_encoder": text_encoder,
"tokenizer": tokenizer,
}
returncomponents
defget_dummy_inputs(self, device, seed: int=0, num_frames: int=8):
ifstr(device).startswith("mps"):
generator=torch.manual_seed(seed)
else:
generator=torch.Generator(device=device).manual_seed(seed)
video_height=16
video_width=16
video= [Image.new("RGB", (video_width, video_height))] *num_frames
inputs= {
"video": video,
"prompt": "dance monkey",
"negative_prompt": "",
"generator": generator,
"num_inference_steps": 2,
"strength": 0.5,
"guidance_scale": 6.0,
# Cannot reduce because convolution kernel becomes bigger than sample
"height": video_height,
"width": video_width,
"max_sequence_length": 16,
"output_type": "pt",
}
returninputs
deftest_inference(self):
device="cpu"
components=self.get_dummy_components()
pipe=self.pipeline_class(**components)
pipe.to(device)
pipe.set_progress_bar_config(disable=None)
inputs=self.get_dummy_inputs(device)
video=pipe(**inputs).frames
generated_video=video[0]
self.assertEqual(generated_video.shape, (8, 3, 16, 16))
expected_video=torch.randn(8, 3, 16, 16)
max_diff=np.abs(generated_video-expected_video).max()
self.assertLessEqual(max_diff, 1e10)
deftest_callback_inputs(self):
sig=inspect.signature(self.pipeline_class.__call__)
has_callback_tensor_inputs="callback_on_step_end_tensor_inputs"insig.parameters
has_callback_step_end="callback_on_step_end"insig.parameters
ifnot (has_callback_tensor_inputsandhas_callback_step_end):
return
components=self.get_dummy_components()
pipe=self.pipeline_class(**components)
pipe=pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)
self.assertTrue(
hasattr(pipe, "_callback_tensor_inputs"),
f" {self.pipeline_class} should have `_callback_tensor_inputs` that defines a list of tensor variables its callback function can use as inputs",
)
defcallback_inputs_subset(pipe, i, t, callback_kwargs):
# iterate over callback args
fortensor_name, tensor_valueincallback_kwargs.items():
# check that we're only passing in allowed tensor inputs
asserttensor_nameinpipe._callback_tensor_inputs
returncallback_kwargs
defcallback_inputs_all(pipe, i, t, callback_kwargs):
fortensor_nameinpipe._callback_tensor_inputs:
asserttensor_nameincallback_kwargs
# iterate over callback args
fortensor_name, tensor_valueincallback_kwargs.items():
# check that we're only passing in allowed tensor inputs
asserttensor_nameinpipe._callback_tensor_inputs
returncallback_kwargs
inputs=self.get_dummy_inputs(torch_device)
# Test passing in a subset
inputs["callback_on_step_end"] =callback_inputs_subset
inputs["callback_on_step_end_tensor_inputs"] = ["latents"]
output=pipe(**inputs)[0]
# Test passing in a everything
inputs["callback_on_step_end"] =callback_inputs_all
inputs["callback_on_step_end_tensor_inputs"] =pipe._callback_tensor_inputs
output=pipe(**inputs)[0]
defcallback_inputs_change_tensor(pipe, i, t, callback_kwargs):
is_last=i== (pipe.num_timesteps-1)
ifis_last:
callback_kwargs["latents"] =torch.zeros_like(callback_kwargs["latents"])
returncallback_kwargs
inputs["callback_on_step_end"] =callback_inputs_change_tensor
inputs["callback_on_step_end_tensor_inputs"] =pipe._callback_tensor_inputs
output=pipe(**inputs)[0]
assertoutput.abs().sum() <1e10
deftest_inference_batch_single_identical(self):
self._test_inference_batch_single_identical(batch_size=3, expected_max_diff=1e-3)
deftest_attention_slicing_forward_pass(
self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3
):
ifnotself.test_attention_slicing:
return
components=self.get_dummy_components()
pipe=self.pipeline_class(**components)
forcomponentinpipe.components.values():
ifhasattr(component, "set_default_attn_processor"):
component.set_default_attn_processor()
pipe.to(torch_device)
pipe.set_progress_bar_config(disable=None)
generator_device="cpu"
inputs=self.get_dummy_inputs(generator_device)
output_without_slicing=pipe(**inputs)[0]
pipe.enable_attention_slicing(slice_size=1)
inputs=self.get_dummy_inputs(generator_device)
output_with_slicing1=pipe(**inputs)[0]
pipe.enable_attention_slicing(slice_size=2)
inputs=self.get_dummy_inputs(generator_device)
output_with_slicing2=pipe(**inputs)[0]
iftest_max_difference:
max_diff1=np.abs(to_np(output_with_slicing1) -to_np(output_without_slicing)).max()
max_diff2=np.abs(to_np(output_with_slicing2) -to_np(output_without_slicing)).max()
self.assertLess(
max(max_diff1, max_diff2),
expected_max_diff,
"Attention slicing should not affect the inference results",
)
deftest_vae_tiling(self, expected_diff_max: float=0.2):
# Since VideoToVideo uses both encoder and decoder tiling, there seems to be much more numerical
# difference. We seem to need a higher tolerance here...
# TODO(aryan): Look into this more deeply
expected_diff_max=0.4
generator_device="cpu"
components=self.get_dummy_components()
pipe=self.pipeline_class(**components)
pipe.to("cpu")
pipe.set_progress_bar_config(disable=None)
# Without tiling
inputs=self.get_dummy_inputs(generator_device)
inputs["height"] =inputs["width"] =128
output_without_tiling=pipe(**inputs)[0]
# With tiling
pipe.vae.enable_tiling(
tile_sample_min_height=96,
tile_sample_min_width=96,
tile_overlap_factor_height=1/12,
tile_overlap_factor_width=1/12,
)
inputs=self.get_dummy_inputs(generator_device)
inputs["height"] =inputs["width"] =128
output_with_tiling=pipe(**inputs)[0]
self.assertLess(
(to_np(output_without_tiling) -to_np(output_with_tiling)).max(),
expected_diff_max,
"VAE tiling should not affect the inference results",
)
deftest_fused_qkv_projections(self):
device="cpu"# ensure determinism for the device-dependent torch.Generator
components=self.get_dummy_components()
pipe=self.pipeline_class(**components)
pipe=pipe.to(device)
pipe.set_progress_bar_config(disable=None)
inputs=self.get_dummy_inputs(device)
frames=pipe(**inputs).frames# [B, F, C, H, W]
original_image_slice=frames[0, -2:, -1, -3:, -3:]
pipe.fuse_qkv_projections()
assertcheck_qkv_fusion_processors_exist(pipe.transformer), (
"Something wrong with the fused attention processors. Expected all the attention processors to be fused."
)
assertcheck_qkv_fusion_matches_attn_procs_length(
pipe.transformer, pipe.transformer.original_attn_processors
), "Something wrong with the attention processors concerning the fused QKV projections."
inputs=self.get_dummy_inputs(device)
frames=pipe(**inputs).frames
image_slice_fused=frames[0, -2:, -1, -3:, -3:]
pipe.transformer.unfuse_qkv_projections()
inputs=self.get_dummy_inputs(device)
frames=pipe(**inputs).frames
image_slice_disabled=frames[0, -2:, -1, -3:, -3:]
assertnp.allclose(original_image_slice, image_slice_fused, atol=1e-3, rtol=1e-3), (
"Fusion of QKV projections shouldn't affect the outputs."
)
assertnp.allclose(image_slice_fused, image_slice_disabled, atol=1e-3, rtol=1e-3), (
"Outputs, with QKV projection fusion enabled, shouldn't change when fused QKV projections are disabled."
)
assertnp.allclose(original_image_slice, image_slice_disabled, atol=1e-2, rtol=1e-2), (
"Original outputs should match when fused QKV projections are disabled."
)

Problem:
Fast tests exist for these pipelines, but there are no @slow integration tests for CogVideoXFunControlPipeline or CogVideoXVideoToVideoPipeline. Text-to-video and image-to-video do have slow coverage.

Impact:
Checkpoint compatibility, preprocessing with real media, and end-to-end output regressions are not covered for two public CogVideoX pipelines.

Reproduction:

frompathlibimportPathforpathin [
"tests/pipelines/cogvideo/test_cogvideox_fun_control.py",
"tests/pipelines/cogvideo/test_cogvideox_video2video.py",
]:
text=Path(path).read_text()
print(path, "@slow"intext, "IntegrationTests"intext)
assert"@slow"intextand"IntegrationTests"intext

Relevant precedent:

@slow
@require_torch_accelerator
classCogVideoXPipelineIntegrationTests(unittest.TestCase):
prompt="A painting of a squirrel eating a burger."
defsetUp(self):
super().setUp()
gc.collect()
backend_empty_cache(torch_device)
deftearDown(self):
super().tearDown()
gc.collect()
backend_empty_cache(torch_device)
deftest_cogvideox(self):
generator=torch.Generator("cpu").manual_seed(0)
pipe=CogVideoXPipeline.from_pretrained("THUDM/CogVideoX-2b", torch_dtype=torch.float16)

@slow
@require_torch_accelerator
classCogVideoXImageToVideoPipelineIntegrationTests(unittest.TestCase):
prompt="A painting of a squirrel eating a burger."
defsetUp(self):
super().setUp()
gc.collect()
backend_empty_cache(torch_device)
deftearDown(self):
super().tearDown()
gc.collect()
backend_empty_cache(torch_device)
deftest_cogvideox(self):
generator=torch.Generator("cpu").manual_seed(0)
pipe=CogVideoXImageToVideoPipeline.from_pretrained("THUDM/CogVideoX-5b-I2V", torch_dtype=torch.bfloat16)

Suggested fix:
Add slow integration classes for CogVideoXVideoToVideoPipeline and CogVideoXFunControlPipeline with published checkpoints, fixed seeds, small media fixtures, and output slice assertions. Also add a fast regression test for the control_video_latents path from Issue 2.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions