cosmos model/pipeline review

# `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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/cosmos/__init__.py#L25-L34
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:
```python
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:
```python
_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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_cosmos.py#L694-L711
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/controlnets/controlnet_cosmos.py#L154-L218

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:
```python
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:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_cosmos.py#L220-L232
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_cosmos.py#L688-L715

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:
```python
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:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/cosmos/pipeline_cosmos_text2world.py#L528-L591
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/cosmos/pipeline_cosmos_video2world.py#L644-L742
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/cosmos/pipeline_cosmos2_text2image.py#L540-L612
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/cosmos/pipeline_cosmos2_video2world.py#L625-L730
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_predict.py#L690-L817
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_transfer.py#L760-L948

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:
```python
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:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_predict.py#L584-L587
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_predict.py#L705-L711

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:
```python
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:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/cosmos/pipeline_cosmos2_video2world.py#L434-L463
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/cosmos/pipeline_cosmos2_video2world.py#L648-L651
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_predict.py#L494-L523
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/cosmos/pipeline_cosmos2_5_predict.py#L705-L716

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:
```python
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:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/autoencoders/autoencoder_kl_cosmos.py#L971-L1068

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:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/autoencoders/autoencoder_kl_cosmos.py#L416-L443
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/autoencoders/autoencoder_kl_cosmos.py#L474-L513

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:
```python
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:
```python
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:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/models/autoencoders/test_models_autoencoder_cosmos.py#L74-L83

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:
```python
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.

Start with the affected Cosmos files under src/diffusers/pipelines/cosmos, models/transformers/transformer_cosmos.py, and models/controlnets/controlnet_cosmos.py, then run the reproduction snippets. Compare the export, mask, dtype, image-input, and validation paths with the cited Flux, Wan, and other pipeline precedents. Done means the six documented Cosmos cases work consistently without the reported exceptions or silent input handling.

Vom Indexierungsmodell aus dem Issue-Text verfasst.