update OmniGenTransformerModel

Author: staoxiao

src/diffusers/__init__.py

@@ -108,6 +108,7 @@
             "MotionAdapter",
             "MultiAdapter",
             "MultiControlNetModel",
+            "OmniGenTransformerModel",
             "PixArtTransformer2DModel",
             "PriorTransformer",
             "SD3ControlNetModel",
@@ -599,6 +600,7 @@
             MotionAdapter,
             MultiAdapter,
             MultiControlNetModel,
+            OmniGenTransformerModel,
             PixArtTransformer2DModel,
             PriorTransformer,
             SD3ControlNetModel,

src/diffusers/models/__init__.py

@@ -66,6 +66,7 @@
     _import_structure["transformers.transformer_mochi"] = ["MochiTransformer3DModel"]
     _import_structure["transformers.transformer_sd3"] = ["SD3Transformer2DModel"]
     _import_structure["transformers.transformer_temporal"] = ["TransformerTemporalModel"]
+    _import_structure["transformers.transformer_omnigen"] = ["OmniGenTransformerModel"]
     _import_structure["unets.unet_1d"] = ["UNet1DModel"]
     _import_structure["unets.unet_2d"] = ["UNet2DModel"]
     _import_structure["unets.unet_2d_condition"] = ["UNet2DConditionModel"]
@@ -125,6 +126,7 @@
             LatteTransformer3DModel,
             LuminaNextDiT2DModel,
             MochiTransformer3DModel,
+            OmniGenTransformerModel,
             PixArtTransformer2DModel,
             PriorTransformer,
             SD3Transformer2DModel,

src/diffusers/models/transformers/__init__.py

@@ -20,3 +20,4 @@
     from .transformer_mochi import MochiTransformer3DModel
     from .transformer_sd3 import SD3Transformer2DModel
     from .transformer_temporal import TransformerTemporalModel
+    from .transformer_omnigen import OmniGenTransformerModel

src/diffusers/models/transformers/transformer_omnigen.py

@@ -26,7 +26,7 @@
 from ...loaders import PeftAdapterMixin
 from ...utils import logging
 from ..attention_processor import AttentionProcessor
-from ..normalization import AdaLayerNorm, CogVideoXLayerNormZero
+from ..normalization import AdaLayerNorm
 from ..embeddings import OmniGenPatchEmbed, OmniGenTimestepEmbed
 from ..modeling_utils import ModelMixin
 
@@ -162,7 +162,7 @@ def forward(
         )
 
 
-class OmniGenTransformer(ModelMixin, ConfigMixin, PeftAdapterMixin):
+class OmniGenTransformerModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
     """
     The Transformer model introduced in OmniGen.
 
@@ -343,3 +343,4 @@ def forward(self,
 
 
 
+

test.py

@@ -0,0 +1,52 @@
+import os
+os.environ['HF_HUB_CACHE'] = '/share/shitao/downloaded_models2'
+
+from huggingface_hub import snapshot_download
+
+from diffusers.models import OmniGenTransformerModel
+from transformers import Phi3Model, Phi3Config
+
+
+from safetensors.torch import load_file
+
+model_name = "Shitao/OmniGen-v1"
+config = Phi3Config.from_pretrained("Shitao/OmniGen-v1")
+model = OmniGenTransformerModel(transformer_config=config)
+cache_folder = os.getenv('HF_HUB_CACHE')
+model_name = snapshot_download(repo_id=model_name,
+                                cache_dir=cache_folder,
+                                ignore_patterns=['flax_model.msgpack', 'rust_model.ot', 'tf_model.h5'])
+print(model_name)
+model_path = os.path.join(model_name, 'model.safetensors')
+ckpt = load_file(model_path, 'cpu')
+
+
+mapping_dict = {
+    "pos_embed": "patch_embedding.pos_embed",
+    "x_embedder.proj.weight": "patch_embedding.output_image_proj.weight",
+    "x_embedder.proj.bias": "patch_embedding.output_image_proj.bias",
+    "input_x_embedder.proj.weight": "patch_embedding.input_image_proj.weight",
+    "input_x_embedder.proj.bias": "patch_embedding.input_image_proj.bias",
+    "final_layer.adaLN_modulation.1.weight": "norm_out.linear.weight",
+    "final_layer.adaLN_modulation.1.bias": "norm_out.linear.bias",
+    "final_layer.linear.weight": "proj_out.weight",
+    "final_layer.linear.bias": "proj_out.bias",
+
+}
+
+new_ckpt = {}
+for k, v in ckpt.items():
+    # new_ckpt[k] = v
+    if k in mapping_dict:
+        new_ckpt[mapping_dict[k]] = v
+    else:
+        new_ckpt[k] = v
+    
+
+
+model.load_state_dict(new_ckpt)
+
+
+
+
+