convert passed

Author: hlky

scripts/convert_z_image_controlnet_to_diffusers.py

@@ -1,19 +1,17 @@
 import argparse
 from contextlib import nullcontext
 
-import safetensors.torch
 import torch
+import safetensors.torch
 from accelerate import init_empty_weights
 from huggingface_hub import hf_hub_download
 
-from diffusers.models import ZImageTransformer2DModel
 from diffusers.models.controlnets.controlnet_z_image import ZImageControlNetModel
 from diffusers.utils.import_utils import is_accelerate_available
 
 
 """
 python scripts/convert_z_image_controlnet_to_diffusers.py  \
---original_z_image_repo_id "Tongyi-MAI/Z-Image-Turbo" \
 --original_controlnet_repo_id "alibaba-pai/Z-Image-Turbo-Fun-Controlnet-Union" \
 --filename "Z-Image-Turbo-Fun-Controlnet-Union.safetensors"
 --output_path "z-image-controlnet-hf/"
@@ -23,7 +21,6 @@
 CTX = init_empty_weights if is_accelerate_available else nullcontext
 
 parser = argparse.ArgumentParser()
-parser.add_argument("--original_z_image_repo_id", default="Tongyi-MAI/Z-Image-Turbo", type=str)
 parser.add_argument("--original_controlnet_repo_id", default=None, type=str)
 parser.add_argument("--filename", default="Z-Image-Turbo-Fun-Controlnet-Union.safetensors", type=str)
 parser.add_argument("--checkpoint_path", default=None, type=str)
@@ -44,72 +41,29 @@ def load_original_checkpoint(args):
     return original_state_dict
 
 
-def load_z_image(args):
-    model = ZImageTransformer2DModel.from_pretrained(
-        args.original_z_image_repo_id, subfolder="transformer", torch_dtype=torch.bfloat16
-    )
-    return model.state_dict(), model.config
-
-
-def convert_z_image_controlnet_checkpoint_to_diffusers(z_image, original_state_dict):
+def convert_z_image_controlnet_checkpoint_to_diffusers(original_state_dict):
     converted_state_dict = {}
 
     converted_state_dict.update(original_state_dict)
 
-    to_copy = {
-        "all_x_embedder.",
-        "noise_refiner.",
-        "context_refiner.",
-        "t_embedder.",
-        "cap_embedder.",
-        "x_pad_token",
-        "cap_pad_token",
-    }
-
-    for key in z_image.keys():
-        for copy_key in to_copy:
-            if key.startswith(copy_key):
-                converted_state_dict[key] = z_image[key]
-
     return converted_state_dict
 
 
 def main(args):
     original_ckpt = load_original_checkpoint(args)
-    z_image, config = load_z_image(args)
 
     control_in_dim = 16
     control_layers_places = [0, 5, 10, 15, 20, 25]
 
-    converted_controlnet_state_dict = convert_z_image_controlnet_checkpoint_to_diffusers(z_image, original_ckpt)
-
-    for key, tensor in converted_controlnet_state_dict.items():
-        print(f"{key} - {tensor.dtype}")
+    converted_controlnet_state_dict = convert_z_image_controlnet_checkpoint_to_diffusers(original_ckpt)
 
     controlnet = ZImageControlNetModel(
-        all_patch_size=config["all_patch_size"],
-        all_f_patch_size=config["all_f_patch_size"],
-        in_channels=config["in_channels"],
-        dim=config["dim"],
-        n_layers=config["n_layers"],
-        n_refiner_layers=config["n_refiner_layers"],
-        n_heads=config["n_heads"],
-        n_kv_heads=config["n_kv_heads"],
-        norm_eps=config["norm_eps"],
-        qk_norm=config["qk_norm"],
-        cap_feat_dim=config["cap_feat_dim"],
-        rope_theta=config["rope_theta"],
-        t_scale=config["t_scale"],
-        axes_dims=config["axes_dims"],
-        axes_lens=config["axes_lens"],
         control_layers_places=control_layers_places,
         control_in_dim=control_in_dim,
-    )
-    missing, unexpected = controlnet.load_state_dict(converted_controlnet_state_dict)
-    print(f"{missing=}")
-    print(f"{unexpected=}")
+    ).to(torch.bfloat16)
+    controlnet.load_state_dict(converted_controlnet_state_dict)
     print("Saving Z-Image ControlNet in Diffusers format")
-    controlnet.save_pretrained(args.output_path, max_shard_size="5GB")
+    controlnet.save_pretrained(args.output_path)
 
 
 if __name__ == "__main__":