omnigen model/pipeline review

# `omnigen` model/pipeline review

Commit tested: `0f1abc4ae8b0eb2a3b40e82a310507281144c423`

Review performed against the repository review rules.

Duplicate check: searched GitHub Issues and PRs for `omnigen`, affected class/function names, and specific failure modes. No duplicates found for the actionable issues below. Related but not duplicate: PR `huggingface/diffusers#11799` touched the torchvision guard.

Execution note: direct `.venv` repros were run. Targeted pytest collection failed before test collection because this Windows torch build lacks `torch._C._distributed_c10d`.

## Issue 1: `timesteps` is unusable

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/omnigen/pipeline_omnigen.py#L340
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/omnigen/pipeline_omnigen.py#L461-L468

Problem:
`__call__` exposes `timesteps`, but always also builds and passes `sigmas`. `retrieve_timesteps` rejects receiving both, so any user-provided `timesteps` fails before inference.

Impact:
The documented custom timestep API is broken for OmniGen.

Reproduction:
```python
import numpy as np
from diffusers import FlowMatchEulerDiscreteScheduler
from diffusers.pipelines.omnigen.pipeline_omnigen import retrieve_timesteps

scheduler = FlowMatchEulerDiscreteScheduler(invert_sigmas=True, num_train_timesteps=1)
sigmas = np.linspace(1, 0, 3)[:2]
retrieve_timesteps(scheduler, 2, "cpu", timesteps=[1, 0], sigmas=sigmas)
```

Relevant precedent:
The copied `retrieve_timesteps` helper is designed to receive either `timesteps` or `sigmas`, not both.

Suggested fix:
```python
sigmas = None
if timesteps is None:
sigmas = np.linspace(1, 0, num_inference_steps + 1)[:num_inference_steps]

timesteps, num_inference_steps = retrieve_timesteps(
self.scheduler, num_inference_steps, timestep_device, timesteps=timesteps, sigmas=sigmas
)
```

## Issue 2: Missing torchvision gives `NameError` instead of dependency gating

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/omnigen/__init__.py#L17-L25
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/omnigen/pipeline_omnigen.py#L31-L32
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/omnigen/pipeline_omnigen.py#L165

Problem:
OmniGen requires `torchvision.transforms`, but lazy loading only gates on torch and transformers. If torchvision is absent, `OmniGenPipeline` remains importable and construction fails with `NameError: OmniGenMultiModalProcessor is not defined`.

Impact:
Users get a confusing runtime failure instead of the standard diffusers missing-backend message.

Reproduction:
```python
import diffusers.utils.import_utils as import_utils

import_utils._torchvision_available = False

from diffusers import OmniGenPipeline

OmniGenPipeline(transformer=None, scheduler=None, vae=None, tokenizer=None)
```

Relevant precedent:
PR `huggingface/diffusers#11799` added a guard, but the export/backend gating still does not include torchvision.

Suggested fix:
Treat torchvision as a required OmniGen pipeline backend in `pipelines/omnigen/__init__.py` and in the dummy object backend list, or raise a clear `requires_backends(..., ["torchvision"])` error before constructing the processor.

## Issue 3: Input-image VAE sampling ignores `generator`

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/omnigen/pipeline_omnigen.py#L171-L190
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/omnigen/pipeline_omnigen.py#L458

Problem:
`encode_input_images` calls `latent_dist.sample()` without passing the pipeline `generator`.

Impact:
Image-conditioned OmniGen calls are not fully controlled by the user-provided generator, so repeated runs with the same `generator` can diverge.

Reproduction:
```python
import torch
from diffusers import AutoencoderKL, OmniGenPipeline

pipe = object.__new__(OmniGenPipeline)
pipe.vae = AutoencoderKL(
sample_size=32,
in_channels=3,
out_channels=3,
block_out_channels=(4, 4, 4, 4),
layers_per_block=1,
latent_channels=4,
norm_num_groups=1,
down_block_types=["DownEncoderBlock2D"] * 4,
up_block_types=["UpDecoderBlock2D"] * 4,
)

image = torch.zeros(1, 3, 16, 16)
torch.manual_seed(0)
a = pipe.encode_input_images([image], device=torch.device("cpu"))[0]
torch.manual_seed(1)
b = pipe.encode_input_images([image], device=torch.device("cpu"))[0]
print((a - b).abs().max().item())
```

Relevant precedent:
Other image-conditioning pipelines use `latent_dist.sample(generator=generator)` via `retrieve_latents`.

Suggested fix:
```python
def encode_input_images(self, input_pixel_values, device=None, dtype=None, generator=None):
device = device or self._execution_device
dtype = dtype or self.vae.dtype

input_img_latents = []
for img in input_pixel_values:
img = self.vae.encode(img.to(device, dtype)).latent_dist.sample(generator=generator)
img = img.mul_(self.vae.config.scaling_factor)
input_img_latents.append(img)
return input_img_latents
```

Then call it with `generator=generator`.

## Issue 4: Batched `input_images` crashes despite public type hint

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/omnigen/pipeline_omnigen.py#L335-L360
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/pipelines/omnigen/pipeline_omnigen.py#L203-L210
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/omnigen/test_pipeline_omnigen.py#L71-L88
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/pipelines/omnigen/test_pipeline_omnigen.py#L113-L127

Problem:
The signature/docstring allow `list[PipelineImageInput]`, but for batched prompts the validation treats each image as a list and calls `len(input_images[i])`. A normal batch like two prompts plus two PIL images crashes. Existing fast and slow pipeline tests are text-only, so the multimodal path is not covered.

Impact:
The core OmniGen image-conditioned API is brittle for batched use.

Reproduction:
```python
from PIL import Image
from diffusers import OmniGenPipeline

pipe = object.__new__(OmniGenPipeline)
pipe.vae_scale_factor = 8
pipe._callback_tensor_inputs = ["latents"]

img = Image.new("RGB", (16, 16), "white")
pipe.check_inputs(
["<|image_1|> a", "<|image_1|> b"],
[img, img],
16,
16,
False,
)
```

Relevant precedent:
Most pipeline batch APIs normalize single items and per-prompt lists before validation.

Suggested fix:
Normalize `input_images` after prompt normalization, and add fast plus slow tests for image-conditioned generation:
```python
if isinstance(prompt, str):
prompt = [prompt]
input_images = [input_images]
elif input_images is not None and len(input_images) > 0 and not isinstance(input_images[0], (list, tuple)):
input_images = [[image] if image is not None else None for image in input_images]
```

## Issue 5: Custom attention bypasses diffusers attention dispatch

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_omnigen.py#L187-L230
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/tests/models/transformers/test_models_transformer_omnigen.py#L29-L32

Problem:
`OmniGenAttnProcessor2_0` calls `F.scaled_dot_product_attention` directly and does not define `_attention_backend` / `_parallel_config`. `model.set_attention_backend(...)` therefore no-ops for OmniGen. The same path also fails exposed GQA-style configs where `num_key_value_heads != num_attention_heads`.

Impact:
OmniGen misses diffusers attention backends/context parallel plumbing, and some serialized configs allowed by the constructor fail at runtime.

Reproduction:
```python
import torch
from diffusers import OmniGenTransformer2DModel

model = OmniGenTransformer2DModel(
hidden_size=16,
num_attention_heads=4,
num_key_value_heads=4,
intermediate_size=32,
num_layers=1,
in_channels=4,
time_step_dim=4,
rope_scaling={"long_factor": [1, 1], "short_factor": [1, 1]},
)
processor = model.layers[0].self_attn.processor
model.set_attention_backend("native")
print(getattr(processor, "_attention_backend", None))

gqa_model = OmniGenTransformer2DModel(
hidden_size=16,
num_attention_heads=4,
num_key_value_heads=2,
intermediate_size=32,
num_layers=1,
pad_token_id=0,
vocab_size=100,
in_channels=4,
time_step_dim=4,
rope_scaling={"long_factor": [1, 1], "short_factor": [1, 1]},
)
seq = 4 + 1 + 16
gqa_model(
hidden_states=torch.randn(1, 4, 8, 8),
timestep=torch.tensor([0.5]),
input_ids=torch.randint(0, 100, (1, 4)),
input_img_latents=[],
input_image_sizes={},
attention_mask=torch.ones(1, seq, seq),
position_ids=torch.arange(seq).unsqueeze(0),
)
```

Relevant precedent:
`FluxAttnProcessor` and `QwenDoubleStreamAttnProcessor2_0` use `dispatch_attention_fn`.

Suggested fix:
Refactor the processor to use `[B, S, H, D]` tensors with `dispatch_attention_fn(..., backend=self._attention_backend, parallel_config=self._parallel_config, enable_gqa=kv_heads != attn.heads)` and add processor attributes:
```python
class OmniGenAttnProcessor2_0:
_attention_backend = None
_parallel_config = None
```

## Issue 6: `rope_scaling=None` default crashes model construction

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_omnigen.py#L137-L150
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_omnigen.py#L329-L344
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_omnigen.py#L367-L373

Problem:
`rope_scaling` is documented and typed as optional with default `None`, but `OmniGenSuScaledRotaryEmbedding` immediately subscripts it.

Impact:
The model’s default constructor is invalid, which is bad for config ergonomics and common model tests.

Reproduction:
```python
from diffusers import OmniGenTransformer2DModel

OmniGenTransformer2DModel(
hidden_size=16,
num_attention_heads=4,
num_key_value_heads=4,
intermediate_size=32,
num_layers=1,
in_channels=4,
time_step_dim=4,
)
```

Relevant precedent:
Model defaults should either construct successfully or raise a clear validation error before subcomponent construction.

Suggested fix:
```python
if rope_scaling is None:
rope_dim = hidden_size // num_attention_heads
rope_scaling = {
"short_factor": [1.0] * (rope_dim // 2),
"long_factor": [1.0] * (rope_dim // 2),
}
```

## Issue 7: RoPE forward breaks `torch.compile(fullgraph=True)`

Affected code:
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_omnigen.py#L153-L163
https://github.com/huggingface/diffusers/blob/0f1abc4ae8b0eb2a3b40e82a310507281144c423/src/diffusers/models/transformers/transformer_omnigen.py#L447-L455

Problem:
RoPE computes `seq_len = torch.max(position_ids) + 1`, branches on that tensor in Python, and mutates `self.inv_freq` inside forward.

Impact:
This violates the repo rule to avoid graph breaks in model forward code and prevents fullgraph compilation.

Reproduction:
```python
import torch
from diffusers import OmniGenTransformer2DModel

model = OmniGenTransformer2DModel(
hidden_size=16,
num_attention_heads=4,
num_key_value_heads=4,
intermediate_size=32,
num_layers=1,
pad_token_id=0,
vocab_size=100,
in_channels=4,
time_step_dim=4,
rope_scaling={"long_factor": [1, 1], "short_factor": [1, 1]},
).eval()

seq = 4 + 1 + 16
inputs = dict(
hidden_states=torch.randn(1, 4, 8, 8),
timestep=torch.tensor([0.5]),
input_ids=torch.randint(0, 100, (1, 4)),
input_img_latents=[],
input_image_sizes={},
attention_mask=torch.ones(1, seq, seq),
position_ids=torch.arange(seq).unsqueeze(0),
)

compiled = torch.compile(model, fullgraph=True, backend="eager")
compiled(**inputs)
```

Relevant precedent:
Other transformer processors keep attention/RoPE forward paths tensor-only and avoid mutating module buffers during forward.

Suggested fix:
Register `short_factor` and `long_factor` as tensors, select using a compile-safe shape check or `torch.where`, and keep `inv_freq` local:
```python
seq_len = position_ids.shape[-1]
ext_factors = self.long_factor if seq_len > self.original_max_position_embeddings else self.short_factor
inv_freq = 1.0 / (ext_factors * self.base**inv_freq_shape)
```

Start with the affected sections of src/diffusers/pipelines/omnigen/pipeline_omnigen.py and src/diffusers/pipelines/omnigen/__init__.py, then inspect src/diffusers/models/transformers/transformer_omnigen.py. Use tests/pipelines/omnigen/test_pipeline_omnigen.py and tests/models/transformers/test_models_transformer_omnigen.py as entry points; targeted pytest collection currently fails on the reported Windows torch build. Done means the listed OmniGen failures are fixed and covered by focused tests.

Vom Indexierungsmodell aus dem Issue-Text verfasst.