Skip to content

omnigen model/pipeline review #13633

Open
Open
omnigen model/pipeline review#13633

Description

@hlky
hlky
opened on Apr 29, 2026

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)

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions