omnigen pipeline

Author: staoxiao

scripts/convert_omnigen_to_diffusers.py

@@ -4,13 +4,25 @@
 import torch
 from safetensors.torch import load_file
 from transformers import AutoModel, AutoTokenizer, AutoConfig
+from huggingface_hub import snapshot_download
 
 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")
+
+    if not os.path.exists(args.origin_ckpt_path):
+        print("Model not found, downloading...")
+        cache_folder = os.getenv('HF_HUB_CACHE')
+        args.origin_ckpt_path = snapshot_download(repo_id=args.origin_ckpt_path,
+                                       cache_dir=cache_folder,
+                                       ignore_patterns=['flax_model.msgpack', 'rust_model.ot', 'tf_model.h5',
+                                                        'model.pt'])
+        print(f"Downloaded model to {args.origin_ckpt_path}")
+
+    ckpt = os.path.join(args.origin_ckpt_path, 'model.safetensors')
+    ckpt = load_file(ckpt, device="cpu")
 
     mapping_dict = {
         "pos_embed": "patch_embedding.pos_embed",
@@ -27,15 +39,13 @@ def main(args):
 
     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,
@@ -49,7 +59,7 @@ def main(args):
 
     scheduler = FlowMatchEulerDiscreteScheduler()
 
-    vae = AutoencoderKL.from_pretrained(args.origin_ckpt_path, torch_dtype=torch.float32)
+    vae = AutoencoderKL.from_pretrained(os.path.join(args.origin_ckpt_path, "vae"), torch_dtype=torch.float32)
 
     tokenizer = AutoTokenizer.from_pretrained(args.origin_ckpt_path)
 
@@ -64,10 +74,10 @@ def main(args):
     parser = argparse.ArgumentParser()
 
     parser.add_argument(
-        "--origin_ckpt_path", default=None, type=str, required=False, help="Path to the checkpoint to convert."
+        "--origin_ckpt_path", default="Shitao/OmniGen-v1", 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.")
+    parser.add_argument("--dump_path", default="OmniGen-v1-diffusers", type=str, required=True, help="Path to the output pipeline.")
 
     args = parser.parse_args()
     main(args)