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.

Written by the indexing model from the issue text.