cosmos model/pipeline review

[hlky]

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.