From cb6d012d918ee10d519ba69a45d8576cdd7daebf Mon Sep 17 00:00:00 2001 From: Yupu Date: Sun, 22 Mar 2026 23:13:13 +0800 Subject: [PATCH] [BugFix]: Fix OmniGen2 Model Loading (#1711) Signed-off-by: Yupu --- .../models/omnigen2/omnigen2_transformer.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/vllm_omni/diffusion/models/omnigen2/omnigen2_transformer.py b/vllm_omni/diffusion/models/omnigen2/omnigen2_transformer.py index 6f803410f..e3cac8f5c 100644 --- a/vllm_omni/diffusion/models/omnigen2/omnigen2_transformer.py +++ b/vllm_omni/diffusion/models/omnigen2/omnigen2_transformer.py @@ -740,18 +740,18 @@ class OmniGen2Transformer2DModel(nn.Module): patch_size: int = 2, in_channels: int = 16, out_channels: int | None = None, - hidden_size: int = 2304, - num_layers: int = 26, + hidden_size: int = 2520, + num_layers: int = 32, num_refiner_layers: int = 2, - num_attention_heads: int = 24, - num_kv_heads: int = 8, + num_attention_heads: int = 21, + num_kv_heads: int = 7, multiple_of: int = 256, ffn_dim_multiplier: float | None = None, norm_eps: float = 1e-5, - axes_dim_rope: tuple[int, int, int] = (32, 32, 32), - axes_lens: tuple[int, int, int] = (300, 512, 512), - text_feat_dim: int = 1024, - timestep_scale: float = 1.0, + axes_dim_rope: tuple[int, int, int] = (40, 40, 40), + axes_lens: tuple[int, int, int] = (1024, 1664, 1664), + text_feat_dim: int = 2048, + timestep_scale: float = 1000.0, ) -> None: """Initialize the OmniGen2 transformer model.""" super().__init__()