omnigen pipeline

Author: staoxiao

scripts/convert_omnigen_to_diffusers.py

@@ -0,0 +1,73 @@
+import argparse
+import os
+
+import torch
+from safetensors.torch import load_file
+from transformers import AutoModel, AutoTokenizer, AutoConfig
+
+from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, OmniGenTransformer2DModel, OmniGenPipeline
+
+
+def main(args):
+    # checkpoint from https://huggingface.co/Shitao/OmniGen-v1
+    ckpt = load_file(args.origin_ckpt_path, device="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",
+
+    }
+
+    converted_state_dict = {}
+    for k, v in ckpt.items():
+        # new_ckpt[k] = v
+        if k in mapping_dict:
+            converted_state_dict[mapping_dict[k]] = v
+        else:
+            converted_state_dict[k] = v
+
+    transformer_config = AutoConfig.from_pretrained(args.origin_ckpt_path)
+
+    # Lumina-Next-SFT 2B
+    transformer = OmniGenTransformer2DModel(
+        transformer_config=transformer_config,
+        patch_size=2,
+        in_channels=4,
+        pos_embed_max_size=192,
+    )
+    transformer.load_state_dict(converted_state_dict, strict=True)
+
+    num_model_params = sum(p.numel() for p in transformer.parameters())
+    print(f"Total number of transformer parameters: {num_model_params}")
+
+    scheduler = FlowMatchEulerDiscreteScheduler()
+
+    vae = AutoencoderKL.from_pretrained(args.origin_ckpt_path, torch_dtype=torch.float32)
+
+    tokenizer = AutoTokenizer.from_pretrained(args.origin_ckpt_path)
+
+
+    pipeline = OmniGenPipeline(
+        tokenizer=tokenizer, transformer=transformer, vae=vae, scheduler=scheduler
+    )
+    pipeline.save_pretrained(args.dump_path)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--origin_ckpt_path", default=None, type=str, required=False, help="Path to the checkpoint to convert."
+    )
+
+    parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output pipeline.")
+
+    args = parser.parse_args()
+    main(args)

src/diffusers/__init__.py

@@ -108,7 +108,7 @@
             "MotionAdapter",
             "MultiAdapter",
             "MultiControlNetModel",
-            "OmniGenTransformerModel",
+            "OmniGenTransformer2DModel",
             "PixArtTransformer2DModel",
             "PriorTransformer",
             "SD3ControlNetModel",
@@ -321,6 +321,7 @@
             "MarigoldNormalsPipeline",
             "MochiPipeline",
             "MusicLDMPipeline",
+            "OmniGenPipeline",
             "PaintByExamplePipeline",
             "PIAPipeline",
             "PixArtAlphaPipeline",
@@ -600,7 +601,7 @@
             MotionAdapter,
             MultiAdapter,
             MultiControlNetModel,
-            OmniGenTransformerModel,
+            OmniGenTransformer2DModel,
             PixArtTransformer2DModel,
             PriorTransformer,
             SD3ControlNetModel,
@@ -792,6 +793,7 @@
             MarigoldNormalsPipeline,
             MochiPipeline,
             MusicLDMPipeline,
+            OmniGenPipeline,
             PaintByExamplePipeline,
             PIAPipeline,
             PixArtAlphaPipeline,

src/diffusers/models/__init__.py

@@ -66,7 +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["transformers.transformer_omnigen"] = ["OmniGenTransformer2DModel"]
     _import_structure["unets.unet_1d"] = ["UNet1DModel"]
     _import_structure["unets.unet_2d"] = ["UNet2DModel"]
     _import_structure["unets.unet_2d_condition"] = ["UNet2DConditionModel"]
@@ -126,7 +126,7 @@
             LatteTransformer3DModel,
             LuminaNextDiT2DModel,
             MochiTransformer3DModel,
-            OmniGenTransformerModel,
+            OmniGenTransformer2DModel,
             PixArtTransformer2DModel,
             PriorTransformer,
             SD3Transformer2DModel,

src/diffusers/models/embeddings.py

@@ -351,7 +351,20 @@ def patch_embeddings(self, latent, is_input_image: bool):
         latent = latent.flatten(2).transpose(1, 2)
         return latent
 
-    def forward(self, latent, is_input_image: bool, padding_latent=None):
+    def forward(self,
+                latent: torch.Tensor,
+                is_input_image: bool,
+                padding_latent: torch.Tensor = None
+                ):
+        """
+        Args:
+            latent:
+            is_input_image:
+            padding_latent: When sizes of target images are inconsistent, use `padding_latent` to maintain consistent sequence length.
+
+        Returns: torch.Tensor
+
+        """
         if isinstance(latent, list):
             if padding_latent is None:
                 padding_latent = [None] * len(latent)
@@ -362,7 +375,7 @@ def forward(self, latent, is_input_image: bool, padding_latent=None):
                 pos_embed = self.cropped_pos_embed(height, width)
                 sub_latent = sub_latent + pos_embed
                 if padding is not None:
-                    sub_latent = torch.cat([sub_latent, padding], dim=-2)
+                    sub_latent = torch.cat([sub_latent, padding.to(sub_latent.device)], dim=-2)
                 patched_latents.append(sub_latent)
         else:
             height, width = latent.shape[-2:]

src/diffusers/models/transformers/__init__.py

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