| import argparse |
| import os |
|
|
| import torch |
| from huggingface_hub import snapshot_download |
| from safetensors.torch import load_file |
| from transformers import AutoTokenizer |
|
|
| from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, OmniGenPipeline, OmniGenTransformer2DModel |
|
|
|
|
| def main(args): |
| |
|
|
| 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", |
| "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", |
| "time_token.mlp.0.weight": "time_token.linear_1.weight", |
| "time_token.mlp.0.bias": "time_token.linear_1.bias", |
| "time_token.mlp.2.weight": "time_token.linear_2.weight", |
| "time_token.mlp.2.bias": "time_token.linear_2.bias", |
| "t_embedder.mlp.0.weight": "t_embedder.linear_1.weight", |
| "t_embedder.mlp.0.bias": "t_embedder.linear_1.bias", |
| "t_embedder.mlp.2.weight": "t_embedder.linear_2.weight", |
| "t_embedder.mlp.2.bias": "t_embedder.linear_2.bias", |
| "llm.embed_tokens.weight": "embed_tokens.weight", |
| } |
|
|
| converted_state_dict = {} |
| for k, v in ckpt.items(): |
| if k in mapping_dict: |
| converted_state_dict[mapping_dict[k]] = v |
| elif "qkv" in k: |
| to_q, to_k, to_v = v.chunk(3) |
| converted_state_dict[f"layers.{k.split('.')[2]}.self_attn.to_q.weight"] = to_q |
| converted_state_dict[f"layers.{k.split('.')[2]}.self_attn.to_k.weight"] = to_k |
| converted_state_dict[f"layers.{k.split('.')[2]}.self_attn.to_v.weight"] = to_v |
| elif "o_proj" in k: |
| converted_state_dict[f"layers.{k.split('.')[2]}.self_attn.to_out.0.weight"] = v |
| else: |
| converted_state_dict[k[4:]] = v |
|
|
| transformer = OmniGenTransformer2DModel( |
| rope_scaling={ |
| "long_factor": [ |
| 1.0299999713897705, |
| 1.0499999523162842, |
| 1.0499999523162842, |
| 1.0799999237060547, |
| 1.2299998998641968, |
| 1.2299998998641968, |
| 1.2999999523162842, |
| 1.4499999284744263, |
| 1.5999999046325684, |
| 1.6499998569488525, |
| 1.8999998569488525, |
| 2.859999895095825, |
| 3.68999981880188, |
| 5.419999599456787, |
| 5.489999771118164, |
| 5.489999771118164, |
| 9.09000015258789, |
| 11.579999923706055, |
| 15.65999984741211, |
| 15.769999504089355, |
| 15.789999961853027, |
| 18.360000610351562, |
| 21.989999771118164, |
| 23.079999923706055, |
| 30.009998321533203, |
| 32.35000228881836, |
| 32.590003967285156, |
| 35.56000518798828, |
| 39.95000457763672, |
| 53.840003967285156, |
| 56.20000457763672, |
| 57.95000457763672, |
| 59.29000473022461, |
| 59.77000427246094, |
| 59.920005798339844, |
| 61.190006256103516, |
| 61.96000671386719, |
| 62.50000762939453, |
| 63.3700065612793, |
| 63.48000717163086, |
| 63.48000717163086, |
| 63.66000747680664, |
| 63.850006103515625, |
| 64.08000946044922, |
| 64.760009765625, |
| 64.80001068115234, |
| 64.81001281738281, |
| 64.81001281738281, |
| ], |
| "short_factor": [ |
| 1.05, |
| 1.05, |
| 1.05, |
| 1.1, |
| 1.1, |
| 1.1, |
| 1.2500000000000002, |
| 1.2500000000000002, |
| 1.4000000000000004, |
| 1.4500000000000004, |
| 1.5500000000000005, |
| 1.8500000000000008, |
| 1.9000000000000008, |
| 2.000000000000001, |
| 2.000000000000001, |
| 2.000000000000001, |
| 2.000000000000001, |
| 2.000000000000001, |
| 2.000000000000001, |
| 2.000000000000001, |
| 2.000000000000001, |
| 2.000000000000001, |
| 2.000000000000001, |
| 2.000000000000001, |
| 2.000000000000001, |
| 2.000000000000001, |
| 2.000000000000001, |
| 2.000000000000001, |
| 2.000000000000001, |
| 2.000000000000001, |
| 2.000000000000001, |
| 2.000000000000001, |
| 2.1000000000000005, |
| 2.1000000000000005, |
| 2.2, |
| 2.3499999999999996, |
| 2.3499999999999996, |
| 2.3499999999999996, |
| 2.3499999999999996, |
| 2.3999999999999995, |
| 2.3999999999999995, |
| 2.6499999999999986, |
| 2.6999999999999984, |
| 2.8999999999999977, |
| 2.9499999999999975, |
| 3.049999999999997, |
| 3.049999999999997, |
| 3.049999999999997, |
| ], |
| "type": "su", |
| }, |
| patch_size=2, |
| in_channels=4, |
| pos_embed_max_size=192, |
| ) |
| transformer.load_state_dict(converted_state_dict, strict=True) |
| transformer.to(torch.bfloat16) |
|
|
| num_model_params = sum(p.numel() for p in transformer.parameters()) |
| print(f"Total number of transformer parameters: {num_model_params}") |
|
|
| scheduler = FlowMatchEulerDiscreteScheduler(invert_sigmas=True, num_train_timesteps=1) |
|
|
| vae = AutoencoderKL.from_pretrained(os.path.join(args.origin_ckpt_path, "vae"), 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="Shitao/OmniGen-v1", |
| type=str, |
| required=False, |
| help="Path to the checkpoint to convert.", |
| ) |
|
|
| parser.add_argument( |
| "--dump_path", default="OmniGen-v1-diffusers", type=str, required=False, help="Path to the output pipeline." |
| ) |
|
|
| args = parser.parse_args() |
| main(args) |
|
|
