Skip to content

cosmos model/pipeline review #13607

Open
Open
cosmos model/pipeline review#13607

Description

@hlky
hlky
opened on Apr 29, 2026

cosmos model/pipeline review

Commit tested: 0f1abc4ae8b0eb2a3b40e82a310507281144c423

Review performed against the repository review rules.

Issue 1: Cosmos pipeline output classes are not exported

Affected code:

diffusers/src/diffusers/pipelines/cosmos/__init__.py

Lines 25 to 34 in 0f1abc4

_import_structure["pipeline_cosmos2_5_predict"] = [
"Cosmos2_5_PredictBasePipeline",
]
_import_structure["pipeline_cosmos2_5_transfer"] = [
"Cosmos2_5_TransferPipeline",
]
_import_structure["pipeline_cosmos2_text2image"] = ["Cosmos2TextToImagePipeline"]
_import_structure["pipeline_cosmos2_video2world"] = ["Cosmos2VideoToWorldPipeline"]
_import_structure["pipeline_cosmos_text2world"] = ["CosmosTextToWorldPipeline"]
_import_structure["pipeline_cosmos_video2world"] = ["CosmosVideoToWorldPipeline"]

https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/cosmos/pipeline_output.py#L14-L42

Problem:
CosmosPipelineOutput and CosmosImagePipelineOutput are public output dataclasses used by the Cosmos pipelines, but diffusers.pipelines.cosmos does not export them through its lazy import structure.

Impact:
Users cannot import the output types from the package namespace, unlike comparable pipeline families. This is a public API consistency and discoverability gap.

Reproduction:

try:
    from diffusers.pipelines.cosmos import CosmosPipelineOutput, CosmosImagePipelineOutput
except Exception as e:
    print(type(e).__name__)
    print(str(e).splitlines()[0])

Relevant precedent:
src/diffusers/pipelines/flux/__init__.py exports its pipeline_output dataclasses through _import_structure.

Suggested fix:

_import_structure["pipeline_output"] = ["CosmosPipelineOutput", "CosmosImagePipelineOutput"]

if TYPE_CHECKING:
    from .pipeline_output import CosmosImagePipelineOutput, CosmosPipelineOutput

Issue 2: padding_mask=None crashes transformer and ControlNet forwards

Affected code:

diffusers/src/diffusers/models/transformers/transformer_cosmos.py

Lines 694 to 711 in 0f1abc4

attention_mask: torch.Tensor | None = None,
fps: int | None = None,
condition_mask: torch.Tensor | None = None,
padding_mask: torch.Tensor | None = None,
return_dict: bool = True,
) -> tuple[torch.Tensor] | Transformer2DModelOutput:
batch_size, num_channels, num_frames, height, width = hidden_states.shape
# 1. Concatenate padding mask if needed & prepare attention mask
if condition_mask is not None:
hidden_states = torch.cat([hidden_states, condition_mask], dim=1)
if self.config.concat_padding_mask:
padding_mask_resized = transforms.functional.resize(
padding_mask, list(hidden_states.shape[-2:]), interpolation=transforms.InterpolationMode.NEAREST
)
hidden_states = torch.cat(
[hidden_states, padding_mask_resized.unsqueeze(2).repeat(batch_size, 1, num_frames, 1, 1)], dim=1

diffusers/src/diffusers/models/controlnets/controlnet_cosmos.py

Lines 154 to 218 in 0f1abc4

padding_mask: torch.Tensor | None = None,
attention_mask: torch.Tensor | None = None,
fps: int | None = None,
return_dict: bool = True,
) -> Union[CosmosControlNetOutput, Tuple[List[torch.Tensor]]]:
"""
Forward pass for the ControlNet.
Args:
controls_latents: Control signal latents [B, C, T, H, W]
latents: Base latents from the noising process [B, C, T, H, W]
timestep: Diffusion timestep tensor
encoder_hidden_states: Tuple of (text_context, img_context) or text_context
condition_mask: Conditioning mask [B, 1, T, H, W]
conditioning_scale: Scale factor(s) for control outputs
padding_mask: Padding mask [B, 1, H, W] or None
attention_mask: Optional attention mask or None
fps: Frames per second for RoPE or None
return_dict: Whether to return a CosmosControlNetOutput or a tuple
Returns:
CosmosControlNetOutput or tuple of control tensors
"""
B, C, T, H, W = controls_latents.shape
# 1. Prepare control latents
control_hidden_states = controls_latents
vace_in_channels = self.config.in_channels - 1
if control_hidden_states.shape[1] < vace_in_channels - 1:
pad_C = vace_in_channels - 1 - control_hidden_states.shape[1]
control_hidden_states = torch.cat(
[
control_hidden_states,
torch.zeros(
(B, pad_C, T, H, W), dtype=control_hidden_states.dtype, device=control_hidden_states.device
),
],
dim=1,
)
if condition_mask is not None:
control_hidden_states = torch.cat([control_hidden_states, condition_mask], dim=1)
else:
control_hidden_states = torch.cat(
[control_hidden_states, torch.zeros_like(controls_latents[:, :1])], dim=1
)
padding_mask_resized = transforms.functional.resize(
padding_mask, list(control_hidden_states.shape[-2:]), interpolation=transforms.InterpolationMode.NEAREST
)
control_hidden_states = torch.cat(
[control_hidden_states, padding_mask_resized.unsqueeze(2).repeat(B, 1, T, 1, 1)], dim=1
)
# 2. Prepare base latents (same processing as transformer.forward)
base_hidden_states = latents
if condition_mask is not None:
base_hidden_states = torch.cat([base_hidden_states, condition_mask], dim=1)
base_padding_mask = transforms.functional.resize(
padding_mask, list(base_hidden_states.shape[-2:]), interpolation=transforms.InterpolationMode.NEAREST
)
base_hidden_states = torch.cat(
[base_hidden_states, base_padding_mask.unsqueeze(2).repeat(B, 1, T, 1, 1)], dim=1
)

Problem:
Both forwards declare padding_mask: torch.Tensor | None = None, but when padding mask concatenation is active they unconditionally pass padding_mask to torchvision.transforms.functional.resize. Omitting the optional argument therefore raises TypeError.

Impact:
The public model API advertises an optional argument that is not actually optional. Direct model use, custom pipelines, and tests that rely on defaults fail before denoising starts.

Reproduction:

import torch
from diffusers import CosmosControlNetModel, CosmosTransformer3DModel

transformer = CosmosTransformer3DModel(
    in_channels=4, out_channels=4, num_attention_heads=2, attention_head_dim=16,
    num_layers=1, mlp_ratio=2, text_embed_dim=16, adaln_lora_dim=4,
    max_size=(1, 16, 16), patch_size=(1, 2, 2), concat_padding_mask=True,
    extra_pos_embed_type=None,
)

try:
    transformer(
        hidden_states=torch.randn(1, 4, 1, 16, 16),
        timestep=torch.tensor([0.5]),
        encoder_hidden_states=torch.randn(1, 8, 16),
    )
except Exception as e:
    print("transformer:", type(e).__name__, str(e).splitlines()[0])

controlnet = CosmosControlNetModel(
    n_controlnet_blocks=1, in_channels=18, latent_channels=18, model_channels=32,
    num_attention_heads=2, attention_head_dim=16, mlp_ratio=2, text_embed_dim=16,
    adaln_lora_dim=4, patch_size=(1, 2, 2), max_size=(1, 16, 16),
    extra_pos_embed_type=None,
)

try:
    controlnet(
        controls_latents=torch.randn(1, 16, 1, 16, 16),
        latents=torch.randn(1, 16, 1, 16, 16),
        timestep=torch.tensor([0.5]),
        encoder_hidden_states=torch.randn(1, 8, 16),
        condition_mask=torch.ones(1, 1, 1, 16, 16),
    )
except Exception as e:
    print("controlnet:", type(e).__name__, str(e).splitlines()[0])

Relevant precedent:
Other model forwards either require masks explicitly or synthesize neutral masks before using them.

Suggested fix:

if self.config.concat_padding_mask:
    if padding_mask is None:
        padding_mask = hidden_states.new_zeros(batch_size, 1, height, width)
    padding_mask = transforms.functional.resize(
        padding_mask,
        list(hidden_states.shape[-2:]),
        interpolation=transforms.InterpolationMode.NEAREST,
    )

Issue 3: Cosmos 2.5 image-context attention crashes on tensor attention masks

Affected code:

diffusers/src/diffusers/models/transformers/transformer_cosmos.py

Lines 220 to 232 in 0f1abc4

def __call__(
self,
attn: Attention,
hidden_states: torch.Tensor,
encoder_hidden_states: tuple[torch.Tensor, torch.Tensor],
attention_mask: tuple[torch.Tensor, torch.Tensor],
image_rotary_emb=None,
) -> torch.Tensor:
if not isinstance(encoder_hidden_states, tuple):
raise ValueError("Expected encoder_hidden_states as (text_context, img_context) tuple.")
text_context, img_context = encoder_hidden_states if encoder_hidden_states else (None, None)
text_mask, img_mask = attention_mask if attention_mask else (None, None)

diffusers/src/diffusers/models/transformers/transformer_cosmos.py

Lines 688 to 715 in 0f1abc4

def forward(
self,
hidden_states: torch.Tensor,
timestep: torch.Tensor,
encoder_hidden_states: torch.Tensor,
block_controlnet_hidden_states: list[torch.Tensor] | None = None,
attention_mask: torch.Tensor | None = None,
fps: int | None = None,
condition_mask: torch.Tensor | None = None,
padding_mask: torch.Tensor | None = None,
return_dict: bool = True,
) -> tuple[torch.Tensor] | Transformer2DModelOutput:
batch_size, num_channels, num_frames, height, width = hidden_states.shape
# 1. Concatenate padding mask if needed & prepare attention mask
if condition_mask is not None:
hidden_states = torch.cat([hidden_states, condition_mask], dim=1)
if self.config.concat_padding_mask:
padding_mask_resized = transforms.functional.resize(
padding_mask, list(hidden_states.shape[-2:]), interpolation=transforms.InterpolationMode.NEAREST
)
hidden_states = torch.cat(
[hidden_states, padding_mask_resized.unsqueeze(2).repeat(batch_size, 1, num_frames, 1, 1)], dim=1
)
if attention_mask is not None:
attention_mask = attention_mask.unsqueeze(1).unsqueeze(1) # [B, 1, 1, S]

Problem:
CosmosAttnProcessor2_5 evaluates attention_mask in boolean context with attention_mask if attention_mask else .... When the public forward passes a tensor mask, PyTorch raises RuntimeError: Boolean value of Tensor with more than one value is ambiguous.

Impact:
Cosmos 2.5 models with image context cannot use normal tensor attention masks. This breaks a standard masking path and makes the processor contract inconsistent with the transformer forward signature.

Reproduction:

import torch
from diffusers import CosmosTransformer3DModel

model = CosmosTransformer3DModel(
    in_channels=5, out_channels=4, num_attention_heads=2, attention_head_dim=16,
    num_layers=1, mlp_ratio=2, text_embed_dim=16, adaln_lora_dim=4,
    max_size=(1, 16, 16), patch_size=(1, 2, 2), concat_padding_mask=True,
    extra_pos_embed_type=None, img_context_dim_in=16, img_context_num_tokens=4,
    img_context_dim_out=16,
)

try:
    model(
        hidden_states=torch.randn(1, 4, 1, 16, 16),
        condition_mask=torch.ones(1, 1, 1, 16, 16),
        timestep=torch.tensor([0.5]),
        encoder_hidden_states=(torch.randn(1, 8, 16), torch.randn(1, 4, 16)),
        attention_mask=torch.ones(1, 8, dtype=torch.bool),
        padding_mask=torch.zeros(1, 1, 16, 16),
    )
except Exception as e:
    print(type(e).__name__)
    print(str(e).splitlines()[0])

Relevant precedent:
Standard attention processors check attention_mask is None explicitly and avoid truthiness checks on tensors.

Suggested fix:

if attention_mask is None:
    text_mask = img_mask = None
elif isinstance(attention_mask, tuple):
    text_mask, img_mask = attention_mask
else:
    text_mask, img_mask = attention_mask, None

Issue 4: Cosmos pipelines do not cast prompt embeddings to transformer dtype

Affected code:

diffusers/src/diffusers/pipelines/cosmos/pipeline_cosmos_text2world.py

Lines 528 to 591 in 0f1abc4

) = 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,
device=device,
max_sequence_length=max_sequence_length,
)
# 4. Prepare timesteps
timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device)
# 5. Prepare latent variables
transformer_dtype = self.transformer.dtype
num_channels_latents = self.transformer.config.in_channels
latents = self.prepare_latents(
batch_size * num_videos_per_prompt,
num_channels_latents,
height,
width,
num_frames,
torch.float32,
device,
generator,
latents,
)
padding_mask = latents.new_zeros(1, 1, height, width, dtype=transformer_dtype)
# 6. Denoising loop
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
self._num_timesteps = len(timesteps)
with self.progress_bar(total=num_inference_steps) as progress_bar:
for i, t in enumerate(timesteps):
if self.interrupt:
continue
self._current_timestep = t
timestep = t.expand(latents.shape[0]).to(transformer_dtype)
latent_model_input = latents
latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
latent_model_input = latent_model_input.to(transformer_dtype)
noise_pred = self.transformer(
hidden_states=latent_model_input,
timestep=timestep,
encoder_hidden_states=prompt_embeds,
fps=fps,
padding_mask=padding_mask,
return_dict=False,
)[0]
sample = latents
if self.do_classifier_free_guidance:
noise_pred_uncond = self.transformer(
hidden_states=latent_model_input,
timestep=timestep,
encoder_hidden_states=negative_prompt_embeds,
fps=fps,
padding_mask=padding_mask,

diffusers/src/diffusers/pipelines/cosmos/pipeline_cosmos_video2world.py

Lines 644 to 742 in 0f1abc4

) = 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,
device=device,
max_sequence_length=max_sequence_length,
)
# 4. Prepare timesteps
timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device)
# 5. Prepare latent variables
vae_dtype = self.vae.dtype
transformer_dtype = self.transformer.dtype
if image is not None:
video = self.video_processor.preprocess(image, height, width).unsqueeze(2)
else:
video = self.video_processor.preprocess_video(video, height, width)
video = video.to(device=device, dtype=vae_dtype)
num_channels_latents = self.transformer.config.in_channels - 1
latents, conditioning_latents, cond_indicator, uncond_indicator, cond_mask, uncond_mask = self.prepare_latents(
video,
batch_size * num_videos_per_prompt,
num_channels_latents,
height,
width,
num_frames,
self.do_classifier_free_guidance,
input_frames_guidance,
torch.float32,
device,
generator,
latents,
)
cond_mask = cond_mask.to(transformer_dtype)
if self.do_classifier_free_guidance:
uncond_mask = uncond_mask.to(transformer_dtype)
augment_sigma = torch.tensor([augment_sigma], device=device, dtype=torch.float32)
padding_mask = latents.new_zeros(1, 1, height, width, dtype=transformer_dtype)
# 6. Denoising loop
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
self._num_timesteps = len(timesteps)
with self.progress_bar(total=num_inference_steps) as progress_bar:
for i, t in enumerate(timesteps):
if self.interrupt:
continue
self._current_timestep = t
timestep = t.expand(latents.shape[0]).to(transformer_dtype)
current_sigma = self.scheduler.sigmas[i]
is_augment_sigma_greater = augment_sigma >= current_sigma
c_in_augment = self.scheduler._get_conditioning_c_in(augment_sigma)
c_in_original = self.scheduler._get_conditioning_c_in(current_sigma)
current_cond_indicator = cond_indicator * 0 if is_augment_sigma_greater else cond_indicator
cond_noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=torch.float32)
cond_latent = conditioning_latents + cond_noise * augment_sigma[:, None, None, None, None]
cond_latent = cond_latent * c_in_augment / c_in_original
cond_latent = current_cond_indicator * cond_latent + (1 - current_cond_indicator) * latents
cond_latent = self.scheduler.scale_model_input(cond_latent, t)
cond_latent = cond_latent.to(transformer_dtype)
noise_pred = self.transformer(
hidden_states=cond_latent,
timestep=timestep,
encoder_hidden_states=prompt_embeds,
fps=fps,
condition_mask=cond_mask,
padding_mask=padding_mask,
return_dict=False,
)[0]
sample = latents
if self.do_classifier_free_guidance:
current_uncond_indicator = uncond_indicator * 0 if is_augment_sigma_greater else uncond_indicator
uncond_noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=torch.float32)
uncond_latent = conditioning_latents + uncond_noise * augment_sigma[:, None, None, None, None]
uncond_latent = uncond_latent * c_in_augment / c_in_original
uncond_latent = current_uncond_indicator * uncond_latent + (1 - current_uncond_indicator) * latents
uncond_latent = self.scheduler.scale_model_input(uncond_latent, t)
uncond_latent = uncond_latent.to(transformer_dtype)
noise_pred_uncond = self.transformer(
hidden_states=uncond_latent,
timestep=timestep,
encoder_hidden_states=negative_prompt_embeds,
fps=fps,
condition_mask=uncond_mask,
padding_mask=padding_mask,

diffusers/src/diffusers/pipelines/cosmos/pipeline_cosmos2_text2image.py

Lines 540 to 612 in 0f1abc4

) = self.encode_prompt(
prompt=prompt,
negative_prompt=negative_prompt,
do_classifier_free_guidance=self.do_classifier_free_guidance,
num_images_per_prompt=num_images_per_prompt,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
device=device,
max_sequence_length=max_sequence_length,
)
# 4. Prepare timesteps
sigmas_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64
sigmas = torch.linspace(0, 1, num_inference_steps, dtype=sigmas_dtype)
timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, device=device, sigmas=sigmas)
if self.scheduler.config.get("final_sigmas_type", "zero") == "sigma_min":
# Replace the last sigma (which is zero) with the minimum sigma value
self.scheduler.sigmas[-1] = self.scheduler.sigmas[-2]
# 5. Prepare latent variables
transformer_dtype = self.transformer.dtype
num_channels_latents = self.transformer.config.in_channels
latents = self.prepare_latents(
batch_size * num_images_per_prompt,
num_channels_latents,
height,
width,
num_frames,
torch.float32,
device,
generator,
latents,
)
padding_mask = latents.new_zeros(1, 1, height, width, dtype=transformer_dtype)
# 6. Denoising loop
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
self._num_timesteps = len(timesteps)
with self.progress_bar(total=num_inference_steps) as progress_bar:
for i, t in enumerate(timesteps):
if self.interrupt:
continue
self._current_timestep = t
current_sigma = self.scheduler.sigmas[i]
current_t = current_sigma / (current_sigma + 1)
c_in = 1 - current_t
c_skip = 1 - current_t
c_out = -current_t
timestep = current_t.expand(latents.shape[0]).to(transformer_dtype) # [B, 1, T, 1, 1]
latent_model_input = latents * c_in
latent_model_input = latent_model_input.to(transformer_dtype)
noise_pred = self.transformer(
hidden_states=latent_model_input,
timestep=timestep,
encoder_hidden_states=prompt_embeds,
padding_mask=padding_mask,
return_dict=False,
)[0]
noise_pred = (c_skip * latents + c_out * noise_pred.float()).to(transformer_dtype)
if self.do_classifier_free_guidance:
noise_pred_uncond = self.transformer(
hidden_states=latent_model_input,
timestep=timestep,
encoder_hidden_states=negative_prompt_embeds,
padding_mask=padding_mask,
return_dict=False,

diffusers/src/diffusers/pipelines/cosmos/pipeline_cosmos2_video2world.py

Lines 625 to 730 in 0f1abc4

) = 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,
device=device,
max_sequence_length=max_sequence_length,
)
# 4. Prepare timesteps
sigmas_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64
sigmas = torch.linspace(0, 1, num_inference_steps, dtype=sigmas_dtype)
timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, device=device, sigmas=sigmas)
if self.scheduler.config.final_sigmas_type == "sigma_min":
# Replace the last sigma (which is zero) with the minimum sigma value
self.scheduler.sigmas[-1] = self.scheduler.sigmas[-2]
# 5. Prepare latent variables
vae_dtype = self.vae.dtype
transformer_dtype = self.transformer.dtype
if image is not None:
video = self.video_processor.preprocess(image, height, width).unsqueeze(2)
else:
video = self.video_processor.preprocess_video(video, height, width)
video = video.to(device=device, dtype=vae_dtype)
num_channels_latents = self.transformer.config.in_channels - 1
latents, conditioning_latents, cond_indicator, uncond_indicator, cond_mask, uncond_mask = self.prepare_latents(
video,
batch_size * num_videos_per_prompt,
num_channels_latents,
height,
width,
num_frames,
self.do_classifier_free_guidance,
torch.float32,
device,
generator,
latents,
)
unconditioning_latents = None
cond_mask = cond_mask.to(transformer_dtype)
if self.do_classifier_free_guidance:
uncond_mask = uncond_mask.to(transformer_dtype)
unconditioning_latents = conditioning_latents
padding_mask = latents.new_zeros(1, 1, height, width, dtype=transformer_dtype)
sigma_conditioning = torch.tensor(sigma_conditioning, dtype=torch.float32, device=device)
t_conditioning = sigma_conditioning / (sigma_conditioning + 1)
# 6. Denoising loop
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
self._num_timesteps = len(timesteps)
with self.progress_bar(total=num_inference_steps) as progress_bar:
for i, t in enumerate(timesteps):
if self.interrupt:
continue
self._current_timestep = t
current_sigma = self.scheduler.sigmas[i]
current_t = current_sigma / (current_sigma + 1)
c_in = 1 - current_t
c_skip = 1 - current_t
c_out = -current_t
timestep = current_t.view(1, 1, 1, 1, 1).expand(
latents.size(0), -1, latents.size(2), -1, -1
) # [B, 1, T, 1, 1]
cond_latent = latents * c_in
cond_latent = cond_indicator * conditioning_latents + (1 - cond_indicator) * cond_latent
cond_latent = cond_latent.to(transformer_dtype)
cond_timestep = cond_indicator * t_conditioning + (1 - cond_indicator) * timestep
cond_timestep = cond_timestep.to(transformer_dtype)
noise_pred = self.transformer(
hidden_states=cond_latent,
timestep=cond_timestep,
encoder_hidden_states=prompt_embeds,
fps=fps,
condition_mask=cond_mask,
padding_mask=padding_mask,
return_dict=False,
)[0]
noise_pred = (c_skip * latents + c_out * noise_pred.float()).to(transformer_dtype)
noise_pred = cond_indicator * conditioning_latents + (1 - cond_indicator) * noise_pred
if self.do_classifier_free_guidance:
uncond_latent = latents * c_in
uncond_latent = uncond_indicator * unconditioning_latents + (1 - uncond_indicator) * uncond_latent
uncond_latent = uncond_latent.to(transformer_dtype)
uncond_timestep = uncond_indicator * t_conditioning + (1 - uncond_indicator) * timestep
uncond_timestep = uncond_timestep.to(transformer_dtype)
noise_pred_uncond = self.transformer(
hidden_states=uncond_latent,
timestep=uncond_timestep,
encoder_hidden_states=negative_prompt_embeds,
fps=fps,
condition_mask=uncond_mask,
padding_mask=padding_mask,

diffusers/src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_predict.py

Lines 690 to 817 in 0f1abc4

) = 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,
device=device,
max_sequence_length=max_sequence_length,
)
vae_dtype = self.vae.dtype
transformer_dtype = self.transformer.dtype
num_frames_in = None
if image is not None:
if batch_size != 1:
raise ValueError(f"batch_size must be 1 for image input (given {batch_size})")
image = torchvision.transforms.functional.to_tensor(image).unsqueeze(0)
video = torch.cat([image, torch.zeros_like(image).repeat(num_frames - 1, 1, 1, 1)], dim=0)
video = video.unsqueeze(0)
num_frames_in = 1
elif video is None:
video = torch.zeros(batch_size, num_frames, 3, height, width, dtype=torch.uint8)
num_frames_in = 0
else:
if batch_size != 1:
raise ValueError(f"batch_size must be 1 for video input (given {batch_size})")
if num_latent_conditional_frames not in [1, 2]:
raise ValueError(
f"num_latent_conditional_frames must be 1 or 2, but got {num_latent_conditional_frames}"
)
frames_to_extract = 4 * (num_latent_conditional_frames - 1) + 1
total_input_frames = len(video)
if total_input_frames < frames_to_extract:
raise ValueError(
f"Input video has only {total_input_frames} frames but Video2World requires at least "
f"{frames_to_extract} frames for conditioning."
)
num_frames_in = frames_to_extract
assert video is not None
video = self.video_processor.preprocess_video(video, height, width)
# For Video2World: extract last frames_to_extract frames from input, then pad
if image is None and num_frames_in > 0 and num_frames_in < video.shape[2]:
video = video[:, :, -num_frames_in:, :, :]
num_frames_out = num_frames
if video.shape[2] < num_frames_out:
n_pad_frames = num_frames_out - video.shape[2]
last_frame = video[:, :, -1:, :, :] # [B, C, T==1, H, W]
pad_frames = last_frame.repeat(1, 1, n_pad_frames, 1, 1) # [B, C, T, H, W]
video = torch.cat((video, pad_frames), dim=2)
assert num_frames_in <= num_frames_out, f"expected ({num_frames_in=}) <= ({num_frames_out=})"
video = video.to(device=device, dtype=vae_dtype)
num_channels_latents = self.transformer.config.in_channels - 1
latents, cond_latent, cond_mask, cond_indicator = self.prepare_latents(
video=video,
batch_size=batch_size * num_videos_per_prompt,
num_channels_latents=num_channels_latents,
height=height,
width=width,
num_frames_in=num_frames_in,
num_frames_out=num_frames,
do_classifier_free_guidance=self.do_classifier_free_guidance,
dtype=torch.float32,
device=device,
generator=generator,
latents=latents,
)
cond_timestep = torch.ones_like(cond_indicator) * conditional_frame_timestep
cond_mask = cond_mask.to(transformer_dtype)
padding_mask = latents.new_zeros(1, 1, height, width, dtype=transformer_dtype)
# Denoising loop
self.scheduler.set_timesteps(num_inference_steps, device=device)
timesteps = self.scheduler.timesteps
self._num_timesteps = len(timesteps)
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
gt_velocity = (latents - cond_latent) * cond_mask
with self.progress_bar(total=num_inference_steps) as progress_bar:
for i, t in enumerate(timesteps):
if self.interrupt:
continue
self._current_timestep = t.cpu().item()
# NOTE: assumes sigma(t) \in [0, 1]
sigma_t = (
torch.tensor(self.scheduler.sigmas[i].item())
.unsqueeze(0)
.to(device=device, dtype=transformer_dtype)
)
in_latents = cond_mask * cond_latent + (1 - cond_mask) * latents
in_latents = in_latents.to(transformer_dtype)
in_timestep = cond_indicator * cond_timestep + (1 - cond_indicator) * sigma_t
noise_pred = self.transformer(
hidden_states=in_latents,
condition_mask=cond_mask,
timestep=in_timestep,
encoder_hidden_states=prompt_embeds,
padding_mask=padding_mask,
return_dict=False,
)[0]
# NOTE: replace velocity (noise_pred) with gt_velocity for conditioning inputs only
noise_pred = gt_velocity + noise_pred * (1 - cond_mask)
if self.do_classifier_free_guidance:
noise_pred_neg = self.transformer(
hidden_states=in_latents,
condition_mask=cond_mask,
timestep=in_timestep,
encoder_hidden_states=negative_prompt_embeds,
padding_mask=padding_mask,

diffusers/src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_transfer.py

Lines 760 to 948 in 0f1abc4

) = 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,
device=device,
max_sequence_length=max_sequence_length,
)
vae_dtype = self.vae.dtype
transformer_dtype = self.transformer.dtype
if getattr(self.transformer.config, "img_context_dim_in", None):
img_context = torch.zeros(
batch_size,
self.transformer.config.img_context_num_tokens,
self.transformer.config.img_context_dim_in,
device=prompt_embeds.device,
dtype=transformer_dtype,
)
if num_videos_per_prompt > 1:
img_context = img_context.repeat_interleave(num_videos_per_prompt, dim=0)
encoder_hidden_states = (prompt_embeds, img_context)
neg_encoder_hidden_states = (negative_prompt_embeds, img_context)
else:
encoder_hidden_states = prompt_embeds
neg_encoder_hidden_states = negative_prompt_embeds
control_video = self.video_processor.preprocess_video(controls, height, width)
if control_video.shape[0] != batch_size:
if control_video.shape[0] == 1:
control_video = control_video.repeat(batch_size, 1, 1, 1, 1)
else:
raise ValueError(
f"Expected controls batch size {batch_size} to match prompt batch size, but got {control_video.shape[0]}."
)
num_frames_out = control_video.shape[2]
if num_frames is not None:
num_frames_out = min(num_frames_out, num_frames)
control_video = _maybe_pad_or_trim_video(control_video, num_frames_out)
# chunk information
num_latent_frames_per_chunk = (num_frames_per_chunk - 1) // self.vae_scale_factor_temporal + 1
chunk_stride = num_frames_per_chunk - num_ar_conditional_frames
chunk_idxs = [
(start_idx, min(start_idx + num_frames_per_chunk, num_frames_out))
for start_idx in range(0, num_frames_out - num_ar_conditional_frames, chunk_stride)
]
video_chunks = []
latents_mean = self.latents_mean.to(dtype=vae_dtype, device=device)
latents_std = self.latents_std.to(dtype=vae_dtype, device=device)
def decode_latents(latents):
latents = latents * latents_std + latents_mean
video = self.vae.decode(latents.to(dtype=self.vae.dtype, device=device), return_dict=False)[0]
return video
latents_arg = latents
initial_num_cond_latent_frames = 0
latent_chunks = []
num_chunks = len(chunk_idxs)
total_steps = num_inference_steps * num_chunks
with self.progress_bar(total=total_steps) as progress_bar:
for chunk_idx, (start_idx, end_idx) in enumerate(chunk_idxs):
if chunk_idx == 0:
prev_output = torch.zeros((batch_size, num_frames_per_chunk, 3, height, width), dtype=vae_dtype)
prev_output = self.video_processor.preprocess_video(prev_output, height, width)
else:
prev_output = video_chunks[-1].clone()
if num_ar_conditional_frames > 0:
prev_output[:, :, :num_ar_conditional_frames] = prev_output[:, :, -num_ar_conditional_frames:]
prev_output[:, :, num_ar_conditional_frames:] = -1 # -1 == 0 in processed video space
else:
prev_output.fill_(-1)
chunk_video = prev_output.to(device=device, dtype=vae_dtype)
chunk_video = _maybe_pad_or_trim_video(chunk_video, num_frames_per_chunk)
latents, cond_latent, cond_mask, cond_indicator = self.prepare_latents(
video=chunk_video,
batch_size=batch_size * num_videos_per_prompt,
num_channels_latents=self.transformer.config.in_channels - 1,
height=height,
width=width,
num_frames_in=chunk_video.shape[2],
num_frames_out=num_frames_per_chunk,
do_classifier_free_guidance=self.do_classifier_free_guidance,
dtype=torch.float32,
device=device,
generator=generator,
num_cond_latent_frames=initial_num_cond_latent_frames
if chunk_idx == 0
else num_cond_latent_frames,
latents=latents_arg,
)
cond_mask = cond_mask.to(transformer_dtype)
cond_timestep = torch.ones_like(cond_indicator) * conditional_frame_timestep
padding_mask = latents.new_zeros(1, 1, height, width, dtype=transformer_dtype)
chunk_control_video = control_video[:, :, start_idx:end_idx, ...].to(
device=device, dtype=self.vae.dtype
)
chunk_control_video = _maybe_pad_or_trim_video(chunk_control_video, num_frames_per_chunk)
if isinstance(generator, list):
controls_latents = [
retrieve_latents(self.vae.encode(chunk_control_video[i].unsqueeze(0)), generator=generator[i])
for i in range(chunk_control_video.shape[0])
]
else:
controls_latents = [
retrieve_latents(self.vae.encode(vid.unsqueeze(0)), generator=generator)
for vid in chunk_control_video
]
controls_latents = torch.cat(controls_latents, dim=0).to(transformer_dtype)
controls_latents = (controls_latents - latents_mean) / latents_std
# Denoising loop
self.scheduler.set_timesteps(num_inference_steps, device=device)
timesteps = self.scheduler.timesteps
self._num_timesteps = len(timesteps)
gt_velocity = (latents - cond_latent) * cond_mask
for i, t in enumerate(timesteps):
if self.interrupt:
continue
self._current_timestep = t.cpu().item()
# NOTE: assumes sigma(t) \in [0, 1]
sigma_t = (
torch.tensor(self.scheduler.sigmas[i].item())
.unsqueeze(0)
.to(device=device, dtype=transformer_dtype)
)
in_latents = cond_mask * cond_latent + (1 - cond_mask) * latents
in_latents = in_latents.to(transformer_dtype)
in_timestep = cond_indicator * cond_timestep + (1 - cond_indicator) * sigma_t
control_output = self.controlnet(
controls_latents=controls_latents,
latents=in_latents,
timestep=in_timestep,
encoder_hidden_states=encoder_hidden_states,
condition_mask=cond_mask,
conditioning_scale=controls_conditioning_scale,
padding_mask=padding_mask,
return_dict=False,
)
control_blocks = control_output[0]
noise_pred = self.transformer(
hidden_states=in_latents,
timestep=in_timestep,
encoder_hidden_states=encoder_hidden_states,
block_controlnet_hidden_states=control_blocks,
condition_mask=cond_mask,
padding_mask=padding_mask,
return_dict=False,
)[0]
noise_pred = gt_velocity + noise_pred * (1 - cond_mask)
if self.do_classifier_free_guidance:
control_output = self.controlnet(
controls_latents=controls_latents,
latents=in_latents,
timestep=in_timestep,
encoder_hidden_states=neg_encoder_hidden_states, # NOTE: negative prompt
condition_mask=cond_mask,
conditioning_scale=controls_conditioning_scale,
padding_mask=padding_mask,
return_dict=False,
)
control_blocks = control_output[0]
noise_pred_neg = self.transformer(
hidden_states=in_latents,
timestep=in_timestep,
encoder_hidden_states=neg_encoder_hidden_states, # NOTE: negative prompt
block_controlnet_hidden_states=control_blocks,
condition_mask=cond_mask,
padding_mask=padding_mask,
return_dict=False,

Problem:
The pipelines pass prompt embeddings from the text encoder, or user-supplied prompt_embeds, directly into the transformer without normalizing to self.transformer.dtype. Mixed precision pipelines can therefore send fp32 embeddings into bf16/fp16 transformer layers.

Impact:
Users running the published Cosmos checkpoints in lower precision can hit dtype mismatch errors, especially when providing precomputed prompt embeddings.

Reproduction:

import torch
from diffusers import CosmosTransformer3DModel

model = CosmosTransformer3DModel(
    in_channels=4, out_channels=4, num_attention_heads=2, attention_head_dim=16,
    num_layers=1, mlp_ratio=2, text_embed_dim=16, adaln_lora_dim=4,
    max_size=(1, 16, 16), patch_size=(1, 2, 2), concat_padding_mask=True,
    extra_pos_embed_type=None,
).to(dtype=torch.bfloat16)

try:
    model(
        hidden_states=torch.randn(1, 4, 1, 16, 16, dtype=torch.bfloat16),
        timestep=torch.tensor([0.5], dtype=torch.bfloat16),
        encoder_hidden_states=torch.randn(1, 8, 16, dtype=torch.float32),
        padding_mask=torch.zeros(1, 1, 16, 16, dtype=torch.bfloat16),
    )
except Exception as e:
    print(type(e).__name__)
    print(str(e).splitlines()[0])

Relevant precedent:
src/diffusers/pipelines/wan/pipeline_wan.py casts prompt embeddings to transformer_dtype after prompt encoding.

Suggested fix:

transformer_dtype = self.transformer.dtype

prompt_embeds, negative_prompt_embeds = self.encode_prompt(...)

prompt_embeds = prompt_embeds.to(device=device, dtype=transformer_dtype)
if negative_prompt_embeds is not None:
    negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=transformer_dtype)

Issue 5: Cosmos2_5_PredictBasePipeline rejects a documented tensor image input

Affected code:

diffusers/src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_predict.py

Lines 584 to 587 in 0f1abc4

image (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, *optional*):
Optional single image for Image2World conditioning. Must be `None` when `video` is provided.
video (`list[PIL.Image.Image]`, `np.ndarray`, `torch.Tensor`, *optional*):
Optional input video for Video2World conditioning. Must be `None` when `image` is provided.

diffusers/src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_predict.py

Lines 705 to 711 in 0f1abc4

if image is not None:
if batch_size != 1:
raise ValueError(f"batch_size must be 1 for image input (given {batch_size})")
image = torchvision.transforms.functional.to_tensor(image).unsqueeze(0)
video = torch.cat([image, torch.zeros_like(image).repeat(num_frames - 1, 1, 1, 1)], dim=0)
video = video.unsqueeze(0)

Problem:
The docstring says image may be a torch.Tensor, but the image path always calls torchvision.transforms.functional.to_tensor(image), which rejects tensor inputs.

Impact:
A documented input type fails before preprocessing. This also diverges from Diffusers pipeline conventions where tensor inputs are normally accepted by image/video processors.

Reproduction:

import torch
from torchvision.transforms.functional import to_tensor

image = torch.zeros(3, 32, 32)

try:
    to_tensor(image)
except Exception as e:
    print(type(e).__name__)
    print(str(e).splitlines()[0])

Relevant precedent:
Other image/video pipelines route accepted tensor, PIL, and NumPy inputs through processor preprocessing instead of forcing to_tensor on every type.

Suggested fix:

if isinstance(image, torch.Tensor):
    image = image if image.ndim == 3 else image.squeeze(0)
else:
    image = torchvision.transforms.functional.to_tensor(image)

Issue 6: Cosmos 2 and Cosmos 2.5 video/image inputs lack validation

Affected code:

diffusers/src/diffusers/pipelines/cosmos/pipeline_cosmos2_video2world.py

Lines 434 to 463 in 0f1abc4

# Copied from diffusers.pipelines.cosmos.pipeline_cosmos_text2world.CosmosTextToWorldPipeline.check_inputs
def check_inputs(
self,
prompt,
height,
width,
prompt_embeds=None,
callback_on_step_end_tensor_inputs=None,
):
if height % 16 != 0 or width % 16 != 0:
raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.")
if callback_on_step_end_tensor_inputs is not None and not all(
k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
):
raise ValueError(
f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
)
if prompt is not None and prompt_embeds is not None:
raise ValueError(
f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
" only forward one of the two."
)
elif prompt is None and prompt_embeds is None:
raise ValueError(
"Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
)
elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")

diffusers/src/diffusers/pipelines/cosmos/pipeline_cosmos2_video2world.py

Lines 648 to 651 in 0f1abc4

if image is not None:
video = self.video_processor.preprocess(image, height, width).unsqueeze(2)
else:
video = self.video_processor.preprocess_video(video, height, width)

diffusers/src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_predict.py

Lines 494 to 523 in 0f1abc4

# Copied from diffusers.pipelines.cosmos.pipeline_cosmos_text2world.CosmosTextToWorldPipeline.check_inputs
def check_inputs(
self,
prompt,
height,
width,
prompt_embeds=None,
callback_on_step_end_tensor_inputs=None,
):
if height % 16 != 0 or width % 16 != 0:
raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.")
if callback_on_step_end_tensor_inputs is not None and not all(
k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
):
raise ValueError(
f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
)
if prompt is not None and prompt_embeds is not None:
raise ValueError(
f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
" only forward one of the two."
)
elif prompt is None and prompt_embeds is None:
raise ValueError(
"Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
)
elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")

diffusers/src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_predict.py

Lines 705 to 716 in 0f1abc4

if image is not None:
if batch_size != 1:
raise ValueError(f"batch_size must be 1 for image input (given {batch_size})")
image = torchvision.transforms.functional.to_tensor(image).unsqueeze(0)
video = torch.cat([image, torch.zeros_like(image).repeat(num_frames - 1, 1, 1, 1)], dim=0)
video = video.unsqueeze(0)
num_frames_in = 1
elif video is None:
video = torch.zeros(batch_size, num_frames, 3, height, width, dtype=torch.uint8)
num_frames_in = 0
else:

Problem:
CosmosVideoToWorldPipeline validates that exactly one of image or video is provided. The Cosmos 2 video pipeline and Cosmos 2.5 predict pipeline do not validate this contract. When both are passed, image silently wins; when neither is passed in video-to-world usage, preprocessing receives None.

Impact:
Invalid user input can be silently ignored or fail later with a less useful error. The behavior is inconsistent across Cosmos pipeline versions.

Reproduction:

from types import SimpleNamespace
from diffusers import Cosmos2VideoToWorldPipeline, CosmosVideoToWorldPipeline

dummy = SimpleNamespace(_callback_tensor_inputs=[])

try:
    CosmosVideoToWorldPipeline.check_inputs(
        dummy, prompt="x", height=16, width=16, image=None, video=None
    )
except Exception as e:
    print("CosmosVideoToWorld:", type(e).__name__, str(e))

print(
    "Cosmos2VideoToWorld:",
    Cosmos2VideoToWorldPipeline.check_inputs(dummy, prompt="x", height=16, width=16),
)

Relevant precedent:
src/diffusers/pipelines/cosmos/pipeline_cosmos_video2world.py already contains the correct image/video validation.

Suggested fix:

if image is not None and video is not None:
    raise ValueError("Only one of `image` or `video` can be provided.")
if image is None and video is None:
    raise ValueError("One of `image` or `video` must be provided.")

Issue 7: AutoencoderKLCosmos.enable_tiling() enables an unused mode

Affected code:

diffusers/src/diffusers/models/autoencoders/autoencoder_kl_cosmos.py

Lines 971 to 1068 in 0f1abc4

# When decoding a batch of video latents at a time, one can save memory by slicing across the batch dimension
# to perform decoding of a single video latent at a time.
self.use_slicing = False
# When decoding spatially large video latents, the memory requirement is very high. By breaking the video latent
# frames spatially into smaller tiles and performing multiple forward passes for decoding, and then blending the
# intermediate tiles together, the memory requirement can be lowered.
self.use_tiling = False
# When decoding temporally long video latents, the memory requirement is very high. By decoding latent frames
# at a fixed frame batch size (based on `self.num_latent_frames_batch_sizes`), the memory requirement can be lowered.
self.use_framewise_encoding = False
self.use_framewise_decoding = False
# This can be configured based on the amount of GPU memory available.
# `16` for sample frames and `2` for latent frames are sensible defaults for consumer GPUs.
# Setting it to higher values results in higher memory usage.
self.num_sample_frames_batch_size = 16
self.num_latent_frames_batch_size = 2
# The minimal tile height and width for spatial tiling to be used
self.tile_sample_min_height = 512
self.tile_sample_min_width = 512
self.tile_sample_min_num_frames = 16
# The minimal distance between two spatial tiles
self.tile_sample_stride_height = 448
self.tile_sample_stride_width = 448
self.tile_sample_stride_num_frames = 8
def enable_tiling(
self,
tile_sample_min_height: int | None = None,
tile_sample_min_width: int | None = None,
tile_sample_min_num_frames: int | None = None,
tile_sample_stride_height: float | None = None,
tile_sample_stride_width: float | None = None,
tile_sample_stride_num_frames: float | None = None,
) -> None:
r"""
Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
processing larger images.
Args:
tile_sample_min_height (`int`, *optional*):
The minimum height required for a sample to be separated into tiles across the height dimension.
tile_sample_min_width (`int`, *optional*):
The minimum width required for a sample to be separated into tiles across the width dimension.
tile_sample_stride_height (`int`, *optional*):
The minimum amount of overlap between two consecutive vertical tiles. This is to ensure that there are
no tiling artifacts produced across the height dimension.
tile_sample_stride_width (`int`, *optional*):
The stride between two consecutive horizontal tiles. This is to ensure that there are no tiling
artifacts produced across the width dimension.
"""
self.use_tiling = True
self.tile_sample_min_height = tile_sample_min_height or self.tile_sample_min_height
self.tile_sample_min_width = tile_sample_min_width or self.tile_sample_min_width
self.tile_sample_min_num_frames = tile_sample_min_num_frames or self.tile_sample_min_num_frames
self.tile_sample_stride_height = tile_sample_stride_height or self.tile_sample_stride_height
self.tile_sample_stride_width = tile_sample_stride_width or self.tile_sample_stride_width
self.tile_sample_stride_num_frames = tile_sample_stride_num_frames or self.tile_sample_stride_num_frames
def _encode(self, x: torch.Tensor) -> torch.Tensor:
x = self.encoder(x)
enc = self.quant_conv(x)
return enc
@apply_forward_hook
def encode(self, x: torch.Tensor, return_dict: bool = True) -> torch.Tensor:
if self.use_slicing and x.shape[0] > 1:
encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)]
h = torch.cat(encoded_slices)
else:
h = self._encode(x)
posterior = IdentityDistribution(h)
if not return_dict:
return (posterior,)
return AutoencoderKLOutput(latent_dist=posterior)
def _decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
z = self.post_quant_conv(z)
dec = self.decoder(z)
if not return_dict:
return (dec,)
return DecoderOutput(sample=dec)
@apply_forward_hook
def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
if self.use_slicing and z.shape[0] > 1:
decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)]
decoded = torch.cat(decoded_slices)
else:
decoded = self._decode(z).sample

Problem:
AutoencoderKLCosmos exposes enable_tiling() and sets self.use_tiling = True, but encode() and decode() only check use_slicing. There are no tiled_encode or tiled_decode implementations.

Impact:
Users can enable a public memory-saving feature that has no effect. For large video VAEs this is particularly misleading because tiling is expected to reduce memory pressure.

Reproduction:

from diffusers import AutoencoderKLCosmos

vae = AutoencoderKLCosmos(
    in_channels=3, out_channels=3, latent_channels=4,
    encoder_block_out_channels=(8, 8, 8, 8),
    decode_block_out_channels=(8, 8, 8, 8),
    attention_resolutions=(8,), resolution=64, num_layers=1,
    patch_size=4, spatial_compression_ratio=4, temporal_compression_ratio=4,
)

vae.enable_tiling(tile_sample_min_height=1, tile_sample_min_width=1, tile_sample_min_num_frames=1)

print(vae.use_tiling)
print(hasattr(vae, "tiled_encode"), hasattr(vae, "tiled_decode"))

Relevant precedent:
AutoencoderKLWan wires use_tiling into encode/decode and implements tiled paths.

Suggested fix:
Implement tiled encode/decode for Cosmos, following AutoencoderKLWan, or temporarily remove/disable enable_tiling() until the mode is functional.

Issue 8: Cosmos VAE attention bypasses the Diffusers attention dispatcher

Affected code:

diffusers/src/diffusers/models/autoencoders/autoencoder_kl_cosmos.py

Lines 416 to 443 in 0f1abc4

class CosmosCausalAttention(nn.Module):
def __init__(
self,
num_attention_heads: int,
attention_head_dim: int,
num_groups: int = 1,
dropout: float = 0.0,
processor: "CosmosSpatialAttentionProcessor2_0" | "CosmosTemporalAttentionProcessor2_0" = None,
) -> None:
super().__init__()
self.num_attention_heads = num_attention_heads
self.norm = CosmosCausalGroupNorm(attention_head_dim, num_groups=num_groups)
self.to_q = CosmosCausalConv3d(attention_head_dim, attention_head_dim, kernel_size=1, stride=1, padding=0)
self.to_k = CosmosCausalConv3d(attention_head_dim, attention_head_dim, kernel_size=1, stride=1, padding=0)
self.to_v = CosmosCausalConv3d(attention_head_dim, attention_head_dim, kernel_size=1, stride=1, padding=0)
self.to_out = nn.ModuleList([])
self.to_out.append(
CosmosCausalConv3d(attention_head_dim, attention_head_dim, kernel_size=1, stride=1, padding=0)
)
self.to_out.append(nn.Dropout(dropout))
self.processor = processor
if self.processor is None:
raise ValueError("CosmosCausalAttention requires a processor.")
def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor:
return self.processor(self, hidden_states=hidden_states, attention_mask=attention_mask)

diffusers/src/diffusers/models/autoencoders/autoencoder_kl_cosmos.py

Lines 474 to 513 in 0f1abc4

hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask)
hidden_states = hidden_states.transpose(1, 2).flatten(2, 3).type_as(query)
hidden_states = hidden_states.unflatten(1, (height, width)).unflatten(0, (batch_size, num_frames))
hidden_states = hidden_states.permute(0, 4, 1, 2, 3)
hidden_states = attn.to_out[0](hidden_states)
hidden_states = attn.to_out[1](hidden_states)
return hidden_states + residual
class CosmosTemporalAttentionProcessor2_0:
def __init__(self):
if not hasattr(F, "scaled_dot_product_attention"):
raise ImportError(
"CosmosSpatialAttentionProcessor2_0 requires PyTorch 2.0 or higher. To use it, please upgrade PyTorch."
)
def __call__(
self, attn: CosmosCausalAttention, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None
) -> torch.Tensor:
batch_size, num_channels, num_frames, height, width = hidden_states.shape
residual = hidden_states
hidden_states = attn.norm(hidden_states)
query = attn.to_q(hidden_states)
key = attn.to_k(hidden_states)
value = attn.to_v(hidden_states)
# [B, C, T, H, W] -> [B * T, H * W, C]
query = query.permute(0, 3, 4, 2, 1).flatten(0, 2)
key = key.permute(0, 3, 4, 2, 1).flatten(0, 2)
value = value.permute(0, 3, 4, 2, 1).flatten(0, 2)
# [B * T, H * W, C] -> [B * T, N, H * W, C // N]
query = query.unflatten(2, (attn.num_attention_heads, -1)).transpose(1, 2)
key = key.unflatten(2, (attn.num_attention_heads, -1)).transpose(1, 2)
value = value.unflatten(2, (attn.num_attention_heads, -1)).transpose(1, 2)
hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask)

Problem:
The VAE attention processors call torch.nn.functional.scaled_dot_product_attention directly instead of dispatch_attention_fn.

Impact:
This violates the model review rule for attention processors and means the VAE does not honor Diffusers attention backend dispatch behavior.

Reproduction:

import inspect
from diffusers.models.autoencoders.autoencoder_kl_cosmos import (
    CosmosSpatialAttentionProcessor2_0,
    CosmosTemporalAttentionProcessor2_0,
)

for cls in (CosmosSpatialAttentionProcessor2_0, CosmosTemporalAttentionProcessor2_0):
    src = inspect.getsource(cls.__call__)
    print(cls.__name__, "dispatch_attention_fn" in src, "scaled_dot_product_attention" in src)

Relevant precedent:
CosmosAttnProcessor2_0 in src/diffusers/models/transformers/transformer_cosmos.py already uses dispatch_attention_fn.

Suggested fix:

from ..attention_dispatch import dispatch_attention_fn

hidden_states = dispatch_attention_fn(
    query.transpose(1, 2),
    key.transpose(1, 2),
    value.transpose(1, 2),
    attn_mask=attention_mask,
    dropout_p=0.0,
    is_causal=False,
)
hidden_states = hidden_states.flatten(2, 3).type_as(query)

Issue 9: Cosmos has no slow tests and one skipped test has an unresolved reason

Affected code:

diffusers/tests/models/autoencoders/test_models_autoencoder_cosmos.py

Lines 74 to 83 in 0f1abc4

def test_gradient_checkpointing_is_applied(self):
expected_set = {
"CosmosEncoder3d",
"CosmosDecoder3d",
}
super().test_gradient_checkpointing_is_applied(expected_set=expected_set)
@unittest.skip("Not sure why this test fails. Investigate later.")
def test_effective_gradient_checkpointing(self):
pass

Problem:
Fast tests exist for the Cosmos model and pipeline files, but there are no Cosmos @slow tests under tests/. The autoencoder test also skips effective gradient checkpointing with the temporary reason "Not sure why this test fails. Investigate later."

Impact:
The public checkpoints and integration paths are not covered by slow smoke tests, and a supported model capability remains skipped without a concrete tracked reason.

Reproduction:

from pathlib import Path

files = list(Path("tests/pipelines/cosmos").glob("test_*.py")) + [
    Path("tests/models/autoencoders/test_models_autoencoder_cosmos.py"),
    Path("tests/models/controlnets/test_models_controlnet_cosmos.py"),
    Path("tests/models/transformers/test_models_transformer_cosmos.py"),
]

slow_hits = [
    (p.as_posix(), i + 1)
    for p in files
    for i, line in enumerate(p.read_text().splitlines())
    if "@slow" in line
]

skip_hits = [
    (p.as_posix(), i + 1, line.strip())
    for p in files
    for i, line in enumerate(p.read_text().splitlines())
    if "Not sure why this test fails" in line
]

print("slow_hits:", slow_hits)
print("ephemeral_skip:", skip_hits)

Relevant precedent:
Other pipeline families include slow checkpoint smoke tests for public model loading and minimal inference.

Suggested fix:
Add slow tests for the public Cosmos pipeline classes using published checkpoints or tiny published fixtures, covering from_pretrained, minimal inference, and dtype/offload where feasible. Replace the skipped gradient-checkpointing reason with a concrete fix or tracked failure reference.

Duplicate Search

I searched existing huggingface/diffusers issues and PRs for cosmos, the affected class/file names, and the specific failure modes above. I found related Cosmos activity, including PR #13573 and issue #12025, but no exact duplicate for these findings.

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