update conversion script

Author: a-r-r-o-w

scripts/convert_mochi_to_diffusers.py

@@ -4,9 +4,9 @@
 import torch
 from accelerate import init_empty_weights
 from safetensors.torch import load_file
+from transformers import T5EncoderModel, T5Tokenizer
 
-# from transformers import T5EncoderModel, T5Tokenizer
-from diffusers import MochiTransformer3DModel
+from diffusers import AutoencoderKLMochi, FlowMatchEulerDiscreteScheduler, MochiPipeline, MochiTransformer3DModel
 from diffusers.utils.import_utils import is_accelerate_available
 
 
@@ -16,7 +16,7 @@
 
 parser = argparse.ArgumentParser()
 parser.add_argument("--transformer_checkpoint_path", default=None, type=str)
-# parser.add_argument("--vae_checkpoint_path", default=None, type=str)
+parser.add_argument("--vae_checkpoint_path", default=None, type=str)
 parser.add_argument("--output_path", required=True, type=str)
 parser.add_argument("--push_to_hub", action="store_true", default=False, help="Whether to push to HF Hub after saving")
 parser.add_argument("--text_encoder_cache_dir", type=str, default=None, help="Path to text encoder cache directory")
@@ -144,9 +144,106 @@ def convert_mochi_transformer_checkpoint_to_diffusers(ckpt_path):
     return new_state_dict
 
 
-# def convert_mochi_vae_checkpoint_to_diffusers(ckpt_path, vae_config):
-#     original_state_dict = torch.load(ckpt_path, map_location="cpu")["state_dict"]
-#     return convert_ldm_vae_checkpoint(original_state_dict, vae_config)
+def convert_mochi_decoder_state_dict_to_diffusers(ckpt_path):
+    original_state_dict = load_file(ckpt_path, device="cpu")
+
+    new_state_dict = {}
+    prefix = "decoder."
+
+    # Convert conv_in
+    new_state_dict[f"{prefix}conv_in.weight"] = original_state_dict["blocks.0.0.weight"]
+    new_state_dict[f"{prefix}conv_in.bias"] = original_state_dict["blocks.0.0.bias"]
+
+    # Convert block_in (MochiMidBlock3D)
+    for i in range(3):  # layers_per_block[-1] = 3
+        new_state_dict[f"{prefix}block_in.resnets.{i}.norm1.norm_layer.weight"] = original_state_dict[
+            f"blocks.0.{i+1}.stack.0.weight"
+        ]
+        new_state_dict[f"{prefix}block_in.resnets.{i}.norm1.norm_layer.bias"] = original_state_dict[
+            f"blocks.0.{i+1}.stack.0.bias"
+        ]
+        new_state_dict[f"{prefix}block_in.resnets.{i}.conv1.conv.weight"] = original_state_dict[
+            f"blocks.0.{i+1}.stack.2.weight"
+        ]
+        new_state_dict[f"{prefix}block_in.resnets.{i}.conv1.conv.bias"] = original_state_dict[
+            f"blocks.0.{i+1}.stack.2.bias"
+        ]
+        new_state_dict[f"{prefix}block_in.resnets.{i}.norm2.norm_layer.weight"] = original_state_dict[
+            f"blocks.0.{i+1}.stack.3.weight"
+        ]
+        new_state_dict[f"{prefix}block_in.resnets.{i}.norm2.norm_layer.bias"] = original_state_dict[
+            f"blocks.0.{i+1}.stack.3.bias"
+        ]
+        new_state_dict[f"{prefix}block_in.resnets.{i}.conv2.conv.weight"] = original_state_dict[
+            f"blocks.0.{i+1}.stack.5.weight"
+        ]
+        new_state_dict[f"{prefix}block_in.resnets.{i}.conv2.conv.bias"] = original_state_dict[
+            f"blocks.0.{i+1}.stack.5.bias"
+        ]
+
+    # Convert up_blocks (MochiUpBlock3D)
+    up_block_layers = [6, 4, 3]  # layers_per_block[-2], layers_per_block[-3], layers_per_block[-4]
+    for block in range(3):
+        for i in range(up_block_layers[block]):
+            new_state_dict[f"{prefix}up_blocks.{block}.resnets.{i}.norm1.norm_layer.weight"] = original_state_dict[
+                f"blocks.{block+1}.blocks.{i}.stack.0.weight"
+            ]
+            new_state_dict[f"{prefix}up_blocks.{block}.resnets.{i}.norm1.norm_layer.bias"] = original_state_dict[
+                f"blocks.{block+1}.blocks.{i}.stack.0.bias"
+            ]
+            new_state_dict[f"{prefix}up_blocks.{block}.resnets.{i}.conv1.conv.weight"] = original_state_dict[
+                f"blocks.{block+1}.blocks.{i}.stack.2.weight"
+            ]
+            new_state_dict[f"{prefix}up_blocks.{block}.resnets.{i}.conv1.conv.bias"] = original_state_dict[
+                f"blocks.{block+1}.blocks.{i}.stack.2.bias"
+            ]
+            new_state_dict[f"{prefix}up_blocks.{block}.resnets.{i}.norm2.norm_layer.weight"] = original_state_dict[
+                f"blocks.{block+1}.blocks.{i}.stack.3.weight"
+            ]
+            new_state_dict[f"{prefix}up_blocks.{block}.resnets.{i}.norm2.norm_layer.bias"] = original_state_dict[
+                f"blocks.{block+1}.blocks.{i}.stack.3.bias"
+            ]
+            new_state_dict[f"{prefix}up_blocks.{block}.resnets.{i}.conv2.conv.weight"] = original_state_dict[
+                f"blocks.{block+1}.blocks.{i}.stack.5.weight"
+            ]
+            new_state_dict[f"{prefix}up_blocks.{block}.resnets.{i}.conv2.conv.bias"] = original_state_dict[
+                f"blocks.{block+1}.blocks.{i}.stack.5.bias"
+            ]
+        new_state_dict[f"{prefix}up_blocks.{block}.proj.weight"] = original_state_dict[f"blocks.{block+1}.proj.weight"]
+        new_state_dict[f"{prefix}up_blocks.{block}.proj.bias"] = original_state_dict[f"blocks.{block+1}.proj.bias"]
+
+    # Convert block_out (MochiMidBlock3D)
+    for i in range(3):  # layers_per_block[0] = 3
+        new_state_dict[f"{prefix}block_out.resnets.{i}.norm1.norm_layer.weight"] = original_state_dict[
+            f"blocks.4.{i}.stack.0.weight"
+        ]
+        new_state_dict[f"{prefix}block_out.resnets.{i}.norm1.norm_layer.bias"] = original_state_dict[
+            f"blocks.4.{i}.stack.0.bias"
+        ]
+        new_state_dict[f"{prefix}block_out.resnets.{i}.conv1.conv.weight"] = original_state_dict[
+            f"blocks.4.{i}.stack.2.weight"
+        ]
+        new_state_dict[f"{prefix}block_out.resnets.{i}.conv1.conv.bias"] = original_state_dict[
+            f"blocks.4.{i}.stack.2.bias"
+        ]
+        new_state_dict[f"{prefix}block_out.resnets.{i}.norm2.norm_layer.weight"] = original_state_dict[
+            f"blocks.4.{i}.stack.3.weight"
+        ]
+        new_state_dict[f"{prefix}block_out.resnets.{i}.norm2.norm_layer.bias"] = original_state_dict[
+            f"blocks.4.{i}.stack.3.bias"
+        ]
+        new_state_dict[f"{prefix}block_out.resnets.{i}.conv2.conv.weight"] = original_state_dict[
+            f"blocks.4.{i}.stack.5.weight"
+        ]
+        new_state_dict[f"{prefix}block_out.resnets.{i}.conv2.conv.bias"] = original_state_dict[
+            f"blocks.4.{i}.stack.5.bias"
+        ]
+
+    # Convert conv_out (Conv1x1)
+    new_state_dict[f"{prefix}conv_out.weight"] = original_state_dict["output_proj.weight"]
+    new_state_dict[f"{prefix}conv_out.bias"] = original_state_dict["output_proj.bias"]
+
+    return new_state_dict
 
 
 def main(args):
@@ -162,7 +259,7 @@ def main(args):
         raise ValueError(f"Unsupported dtype: {args.dtype}")
 
     transformer = None
-    # vae = None
+    vae = None
 
     if args.transformer_checkpoint_path is not None:
         converted_transformer_state_dict = convert_mochi_transformer_checkpoint_to_diffusers(
@@ -171,18 +268,31 @@ def main(args):
         transformer = MochiTransformer3DModel()
         transformer.load_state_dict(converted_transformer_state_dict, strict=True)
         if dtype is not None:
-            # Original checkpoint data type will be preserved
             transformer = transformer.to(dtype=dtype)
 
-    # text_encoder_id = "google/t5-v1_1-xxl"
-    # tokenizer = T5Tokenizer.from_pretrained(text_encoder_id, model_max_length=TOKENIZER_MAX_LENGTH)
-    # text_encoder = T5EncoderModel.from_pretrained(text_encoder_id, cache_dir=args.text_encoder_cache_dir)
-
-    # # Apparently, the conversion does not work anymore without this :shrug:
-    # for param in text_encoder.parameters():
-    #     param.data = param.data.contiguous()
-
-    transformer.save_pretrained("/raid/aryan/mochi-diffusers", subfolder="transformer")
+    if args.vae_checkpoint_path is not None:
+        vae = AutoencoderKLMochi(latent_channels=12, out_channels=3)
+        converted_vae_state_dict = convert_mochi_decoder_state_dict_to_diffusers(args.vae_checkpoint_path)
+        vae.load_state_dict(converted_vae_state_dict, strict=True)
+        if dtype is not None:
+            vae = vae.to(dtype=dtype)
+
+    text_encoder_id = "google/t5-v1_1-xxl"
+    tokenizer = T5Tokenizer.from_pretrained(text_encoder_id, model_max_length=TOKENIZER_MAX_LENGTH)
+    text_encoder = T5EncoderModel.from_pretrained(text_encoder_id, cache_dir=args.text_encoder_cache_dir)
+
+    # Apparently, the conversion does not work anymore without this :shrug:
+    for param in text_encoder.parameters():
+        param.data = param.data.contiguous()
+
+    pipe = MochiPipeline(
+        scheduler=FlowMatchEulerDiscreteScheduler(),
+        vae=vae,
+        text_encoder=text_encoder,
+        tokenizer=tokenizer,
+        transformer=transformer,
+    )
+    pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB", push_to_hub=args.push_to_hub)
 
 
 if __name__ == "__main__":

src/diffusers/pipelines/mochi/pipeline_mochi.py

@@ -204,7 +204,7 @@ def __init__(
         self.vae_temporal_scale_factor = 6
         self.patch_size = 2
 
-        self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor)
+        self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_scale_factor)
         self.tokenizer_max_length = (
             self.tokenizer.model_max_length if hasattr(self, "tokenizer") and self.tokenizer is not None else 77
         )