omnigen model/pipeline review

[hlky]

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:

diffusers/src/diffusers/pipelines/omnigen/pipeline_omnigen.py

Line 340 in 0f1abc4

timesteps: list[int] = None,

diffusers/src/diffusers/pipelines/omnigen/pipeline_omnigen.py

Lines 461 to 468 in 0f1abc4

	sigmas = np.linspace(1, 0, num_inference_steps + 1)[:num_inference_steps]
	if XLA_AVAILABLE:
	timestep_device = "cpu"
	else:
	timestep_device = device
	timesteps, num_inference_steps = retrieve_timesteps(
	self.scheduler, num_inference_steps, timestep_device, timesteps, sigmas=sigmas
	)

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:

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:

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:

diffusers/src/diffusers/pipelines/omnigen/__init__.py

Lines 17 to 25 in 0f1abc4

	try:
	if not (is_transformers_available() and is_torch_available()):
	raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
	from ...utils import dummy_torch_and_transformers_objects # noqa F403
	_dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
	else:
	_import_structure["pipeline_omnigen"] = ["OmniGenPipeline"]

diffusers/src/diffusers/pipelines/omnigen/pipeline_omnigen.py

Lines 31 to 32 in 0f1abc4

	if is_torchvision_available():
	from .processor_omnigen import OmniGenMultiModalProcessor

diffusers/src/diffusers/pipelines/omnigen/pipeline_omnigen.py

Line 165 in 0f1abc4

self.multimodal_processor = OmniGenMultiModalProcessor(tokenizer, max_image_size=1024)

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:

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:

diffusers/src/diffusers/pipelines/omnigen/pipeline_omnigen.py

Lines 171 to 190 in 0f1abc4

	def encode_input_images(
	self,
	input_pixel_values: list[torch.Tensor],
	device: torch.device | None = None,
	dtype: torch.dtype | None = None,
	):
	"""
	get the continue embedding of input images by VAE
	Args:
	input_pixel_values: normalized pixel of input images
	device:
	Returns: torch.Tensor
	"""
	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().mul_(self.vae.config.scaling_factor)

diffusers/src/diffusers/pipelines/omnigen/pipeline_omnigen.py

Line 458 in 0f1abc4

input_img_latents = self.encode_input_images(processed_data["input_pixel_values"], device=device)

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:

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:

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:

diffusers/src/diffusers/pipelines/omnigen/pipeline_omnigen.py

Lines 335 to 360 in 0f1abc4

	input_images: PipelineImageInput | list[PipelineImageInput] = None,
	height: int | None = None,
	width: int | None = None,
	num_inference_steps: int = 50,
	max_input_image_size: int = 1024,
	timesteps: list[int] = None,
	guidance_scale: float = 2.5,
	img_guidance_scale: float = 1.6,
	use_input_image_size_as_output: bool = False,
	num_images_per_prompt: int | None = 1,
	generator: torch.Generator | list[torch.Generator] | None = None,
	latents: torch.Tensor | None = None,
	output_type: str | None = "pil",
	return_dict: bool = True,
	callback_on_step_end: Callable[[int, int], None] | None = None,
	callback_on_step_end_tensor_inputs: list[str] = ["latents"],
	):
	r"""
	Function invoked when calling the pipeline for generation.
	Args:
	prompt (`str` or `list[str]`, *optional*):
	The prompt or prompts to guide the image generation. If the input includes images, need to add
	placeholders `<img><|image_i|></img>` in the prompt to indicate the position of the i-th images.
	input_images (`PipelineImageInput` or `list[PipelineImageInput]`, *optional*):
	The list of input images. We will replace the "<|image_i|>" in prompt with the i-th image in list.

diffusers/src/diffusers/pipelines/omnigen/pipeline_omnigen.py

Lines 203 to 210 in 0f1abc4

	if input_images is not None:
	if len(input_images) != len(prompt):
	raise ValueError(
	f"The number of prompts: {len(prompt)} does not match the number of input images: {len(input_images)}."
	)
	for i in range(len(input_images)):
	if input_images[i] is not None:
	if not all(f"<img><|image_{k + 1}|></img>" in prompt[i] for k in range(len(input_images[i]))):

diffusers/tests/pipelines/omnigen/test_pipeline_omnigen.py

Lines 71 to 88 in 0f1abc4

	inputs = {
	"prompt": "A painting of a squirrel eating a burger",
	"generator": generator,
	"num_inference_steps": 1,
	"guidance_scale": 3.0,
	"output_type": "np",
	"height": 16,
	"width": 16,
	}
	return inputs
	def test_inference(self):
	pipe = self.pipeline_class(**self.get_dummy_components()).to(torch_device)
	inputs = self.get_dummy_inputs(torch_device)
	generated_image = pipe(**inputs).images[0]
	self.assertEqual(generated_image.shape, (16, 16, 3))

diffusers/tests/pipelines/omnigen/test_pipeline_omnigen.py

Lines 113 to 127 in 0f1abc4

	return {
	"prompt": "A photo of a cat",
	"num_inference_steps": 2,
	"guidance_scale": 2.5,
	"output_type": "np",
	"generator": generator,
	}
	def test_omnigen_inference(self):
	pipe = self.pipeline_class.from_pretrained(self.repo_id, torch_dtype=torch.bfloat16)
	pipe.enable_model_cpu_offload()
	inputs = self.get_inputs(torch_device)
	image = pipe(**inputs).images[0]

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:

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(
 ["<img><|image_1|></img> a", "<img><|image_1|></img> 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:

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:

diffusers/src/diffusers/models/transformers/transformer_omnigen.py

Lines 187 to 230 in 0f1abc4

	class OmniGenAttnProcessor2_0:
	r"""
	Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). This is
	used in the OmniGen model.
	"""
	def __init__(self):
	if not hasattr(F, "scaled_dot_product_attention"):
	raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
	def __call__(
	self,
	attn: Attention,
	hidden_states: torch.Tensor,
	encoder_hidden_states: torch.Tensor,
	attention_mask: torch.Tensor | None = None,
	image_rotary_emb: torch.Tensor | None = None,
	) -> torch.Tensor:
	batch_size, sequence_length, _ = hidden_states.shape
	# Get Query-Key-Value Pair
	query = attn.to_q(hidden_states)
	key = attn.to_k(encoder_hidden_states)
	value = attn.to_v(encoder_hidden_states)
	bsz, q_len, query_dim = query.size()
	inner_dim = key.shape[-1]
	head_dim = query_dim // attn.heads
	# Get key-value heads
	kv_heads = inner_dim // head_dim
	query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
	key = key.view(batch_size, -1, kv_heads, head_dim).transpose(1, 2)
	value = value.view(batch_size, -1, kv_heads, head_dim).transpose(1, 2)
	# Apply RoPE if needed
	if image_rotary_emb is not None:
	from ..embeddings import apply_rotary_emb
	query = apply_rotary_emb(query, image_rotary_emb, use_real_unbind_dim=-2)
	key = apply_rotary_emb(key, image_rotary_emb, use_real_unbind_dim=-2)
	hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask)

diffusers/tests/models/transformers/test_models_transformer_omnigen.py

Lines 29 to 32 in 0f1abc4

	class OmniGenTransformerTests(ModelTesterMixin, unittest.TestCase):
	model_class = OmniGenTransformer2DModel
	main_input_name = "hidden_states"
	uses_custom_attn_processor = True

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:

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:

class OmniGenAttnProcessor2_0:
 _attention_backend = None
 _parallel_config = None

Issue 6: rope_scaling=None default crashes model construction

Affected code:

diffusers/src/diffusers/models/transformers/transformer_omnigen.py

Lines 137 to 150 in 0f1abc4

	def __init__(
	self, dim, max_position_embeddings=131072, original_max_position_embeddings=4096, base=10000, rope_scaling=None
	):
	super().__init__()
	self.dim = dim
	self.max_position_embeddings = max_position_embeddings
	self.base = base
	inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float() / self.dim))
	self.register_buffer("inv_freq", tensor=inv_freq, persistent=False)
	self.short_factor = rope_scaling["short_factor"]
	self.long_factor = rope_scaling["long_factor"]

diffusers/src/diffusers/models/transformers/transformer_omnigen.py

Lines 329 to 344 in 0f1abc4

	def __init__(
	self,
	in_channels: int = 4,
	patch_size: int = 2,
	hidden_size: int = 3072,
	rms_norm_eps: float = 1e-5,
	num_attention_heads: int = 32,
	num_key_value_heads: int = 32,
	intermediate_size: int = 8192,
	num_layers: int = 32,
	pad_token_id: int = 32000,
	vocab_size: int = 32064,
	max_position_embeddings: int = 131072,
	original_max_position_embeddings: int = 4096,
	rope_base: int = 10000,
	rope_scaling: dict = None,

diffusers/src/diffusers/models/transformers/transformer_omnigen.py

Lines 367 to 373 in 0f1abc4

	self.rope = OmniGenSuScaledRotaryEmbedding(
	hidden_size // num_attention_heads,
	max_position_embeddings=max_position_embeddings,
	original_max_position_embeddings=original_max_position_embeddings,
	base=rope_base,
	rope_scaling=rope_scaling,
	)

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:

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:

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:

diffusers/src/diffusers/models/transformers/transformer_omnigen.py

Lines 153 to 163 in 0f1abc4

	def forward(self, hidden_states, position_ids):
	seq_len = torch.max(position_ids) + 1
	if seq_len > self.original_max_position_embeddings:
	ext_factors = torch.tensor(self.long_factor, dtype=torch.float32, device=hidden_states.device)
	else:
	ext_factors = torch.tensor(self.short_factor, dtype=torch.float32, device=hidden_states.device)
	inv_freq_shape = (
	torch.arange(0, self.dim, 2, dtype=torch.int64, device=hidden_states.device).float() / self.dim
	)
	self.inv_freq = 1.0 / (ext_factors * self.base**inv_freq_shape)

diffusers/src/diffusers/models/transformers/transformer_omnigen.py

Lines 447 to 455 in 0f1abc4

	image_rotary_emb = self.rope(hidden_states, position_ids)
	# 4. Transformer blocks
	for block in self.layers:
	if torch.is_grad_enabled() and self.gradient_checkpointing:
	hidden_states = self._gradient_checkpointing_func(
	block, hidden_states, attention_mask, image_rotary_emb
	)
	else:

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:

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:

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)