[BugFix]: Fix OmniGen2 Model Loading (#1711)

Signed-off-by: Yupu <feng.yu.pu0330@gmail.com>

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__()