diffusers/scripts/convert_anima_to_diffusers.py

"""
Convert Anima checkpoints to Diffusers format.

Example:
```bash
python scripts/convert_anima_to_diffusers.py \
    --transformer_ckpt_path anima_model/anima-preview3-base.safetensors \
    --text_encoder_ckpt_path anima_model/qwen_3_06b_base.safetensors \
    --vae_ckpt_path anima_model/qwen_image_vae.safetensors \
    --qwen_tokenizer_path path/to/qwen25_tokenizer \
    --t5_tokenizer_path path/to/t5_tokenizer \
    --output_path anima_model/anima-preview3-diffusers \
    --save_pipeline
```
"""

import argparse
import pathlib
import sys
from typing import Any

import torch
from accelerate import init_empty_weights
from convert_cosmos_to_diffusers import convert_transformer
from safetensors.torch import load_file
from transformers import AutoTokenizer, Qwen3Config, Qwen3Model, T5TokenizerFast

from diffusers import (
    AnimaAutoBlocks,
    AnimaTextConditioner,
    AutoencoderKLQwenImage,
    FlowMatchEulerDiscreteScheduler,
)


DTYPE_MAPPING = {
    "fp32": torch.float32,
    "fp16": torch.float16,
    "bf16": torch.bfloat16,
}


def rename_residual_key(key: str) -> str:
    replacements = {
        ".residual.0.": ".norm1.",
        ".residual.2.": ".conv1.",
        ".residual.3.": ".norm2.",
        ".residual.6.": ".conv2.",
        ".shortcut.": ".conv_shortcut.",
    }
    for old, new in replacements.items():
        key = key.replace(old, new)
    return key


def rename_mid_key(key: str) -> str:
    replacements = {
        ".middle.0.": ".mid_block.resnets.0.",
        ".middle.1.": ".mid_block.attentions.0.",
        ".middle.2.": ".mid_block.resnets.1.",
    }
    for old, new in replacements.items():
        key = key.replace(old, new)
    return rename_residual_key(key)


def rename_decoder_upsample_key(key: str) -> str:
    prefix = "decoder.upsamples."
    suffix = key.removeprefix(prefix)
    index_str, rest = suffix.split(".", 1)
    index = int(index_str)

    if index in (3, 7, 11):
        block_index = (index - 3) // 4
        new_key = f"decoder.up_blocks.{block_index}.upsamplers.0.{rest}"
    else:
        block_index = index // 4
        resnet_index = index % 4
        new_key = f"decoder.up_blocks.{block_index}.resnets.{resnet_index}.{rest}"

    return rename_residual_key(new_key)


def convert_qwen_image_vae_state_dict(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
    converted_state_dict = {}
    for key, value in state_dict.items():
        if key.startswith("conv1."):
            new_key = key.replace("conv1.", "quant_conv.", 1)
        elif key.startswith("conv2."):
            new_key = key.replace("conv2.", "post_quant_conv.", 1)
        elif key.startswith("encoder.conv1."):
            new_key = key.replace("encoder.conv1.", "encoder.conv_in.", 1)
        elif key.startswith("decoder.conv1."):
            new_key = key.replace("decoder.conv1.", "decoder.conv_in.", 1)
        elif key.startswith("encoder.downsamples."):
            new_key = rename_residual_key(key.replace("encoder.downsamples.", "encoder.down_blocks.", 1))
        elif key.startswith("decoder.upsamples."):
            new_key = rename_decoder_upsample_key(key)
        elif key.startswith("encoder.middle.") or key.startswith("decoder.middle."):
            new_key = rename_mid_key(key)
        elif key.startswith("encoder.head.0."):
            new_key = key.replace("encoder.head.0.", "encoder.norm_out.", 1)
        elif key.startswith("encoder.head.2."):
            new_key = key.replace("encoder.head.2.", "encoder.conv_out.", 1)
        elif key.startswith("decoder.head.0."):
            new_key = key.replace("decoder.head.0.", "decoder.norm_out.", 1)
        elif key.startswith("decoder.head.2."):
            new_key = key.replace("decoder.head.2.", "decoder.conv_out.", 1)
        else:
            new_key = rename_residual_key(key)

        if new_key in converted_state_dict:
            raise ValueError(f"Duplicate converted VAE key: {new_key}")
        converted_state_dict[new_key] = value

    return converted_state_dict


def convert_qwen_image_vae(state_dict: dict[str, torch.Tensor]) -> AutoencoderKLQwenImage:
    converted_state_dict = convert_qwen_image_vae_state_dict(state_dict)
    with init_empty_weights():
        vae = AutoencoderKLQwenImage()

    expected_keys = set(vae.state_dict().keys())
    converted_keys = set(converted_state_dict.keys())
    missing_keys = expected_keys - converted_keys
    unexpected_keys = converted_keys - expected_keys
    if missing_keys or unexpected_keys:
        if missing_keys:
            print(f"ERROR: missing VAE keys ({len(missing_keys)}):", file=sys.stderr)
            for key in sorted(missing_keys):
                print(key, file=sys.stderr)
        if unexpected_keys:
            print(f"ERROR: unexpected VAE keys ({len(unexpected_keys)}):", file=sys.stderr)
            for key in sorted(unexpected_keys):
                print(key, file=sys.stderr)
        sys.exit(1)

    vae.load_state_dict(converted_state_dict, strict=True, assign=True)
    return vae


def infer_text_conditioner_config(state_dict: dict[str, torch.Tensor]) -> dict[str, Any]:
    model_dim = state_dict["blocks.0.self_attn.q_proj.weight"].shape[0]
    source_dim = state_dict["blocks.0.cross_attn.k_proj.weight"].shape[1]
    target_vocab_size, target_dim = state_dict["embed.weight"].shape
    attention_head_dim = state_dict["blocks.0.self_attn.q_norm.weight"].shape[0]
    num_layers = 1 + max(int(key.split(".")[1]) for key in state_dict if key.startswith("blocks."))

    return {
        "source_dim": source_dim,
        "target_dim": target_dim,
        "model_dim": model_dim,
        "num_layers": num_layers,
        "num_attention_heads": model_dim // attention_head_dim,
        "target_vocab_size": target_vocab_size,
    }


def convert_text_conditioner(state_dict: dict[str, torch.Tensor]) -> AnimaTextConditioner:
    config = infer_text_conditioner_config(state_dict)
    with init_empty_weights():
        text_conditioner = AnimaTextConditioner(**config)

    expected_keys = set(text_conditioner.state_dict().keys())
    converted_keys = set(state_dict.keys())
    missing_keys = expected_keys - converted_keys
    unexpected_keys = converted_keys - expected_keys
    if missing_keys or unexpected_keys:
        if missing_keys:
            print(f"ERROR: missing text conditioner keys ({len(missing_keys)}):", file=sys.stderr)
            for key in sorted(missing_keys):
                print(key, file=sys.stderr)
        if unexpected_keys:
            print(f"ERROR: unexpected text conditioner keys ({len(unexpected_keys)}):", file=sys.stderr)
            for key in sorted(unexpected_keys):
                print(key, file=sys.stderr)
        sys.exit(1)

    text_conditioner.load_state_dict(state_dict, strict=True, assign=True)
    return text_conditioner


def infer_qwen3_config(state_dict: dict[str, torch.Tensor]) -> Qwen3Config:
    vocab_size, hidden_size = state_dict["embed_tokens.weight"].shape
    intermediate_size = state_dict["layers.0.mlp.gate_proj.weight"].shape[0]
    num_hidden_layers = 1 + max(int(key.split(".")[1]) for key in state_dict if key.startswith("layers."))
    head_dim = state_dict["layers.0.self_attn.q_norm.weight"].shape[0]
    num_attention_heads = state_dict["layers.0.self_attn.q_proj.weight"].shape[0] // head_dim
    num_key_value_heads = state_dict["layers.0.self_attn.k_proj.weight"].shape[0] // head_dim

    return Qwen3Config(
        vocab_size=vocab_size,
        hidden_size=hidden_size,
        intermediate_size=intermediate_size,
        num_hidden_layers=num_hidden_layers,
        num_attention_heads=num_attention_heads,
        num_key_value_heads=num_key_value_heads,
        max_position_embeddings=32768,
        rms_norm_eps=1e-6,
        rope_theta=1000000.0,
        head_dim=head_dim,
        attention_bias=False,
        tie_word_embeddings=False,
    )


def convert_text_encoder(state_dict: dict[str, torch.Tensor]) -> Qwen3Model:
    state_dict = {key.removeprefix("model."): value for key, value in state_dict.items()}
    config = infer_qwen3_config(state_dict)
    with init_empty_weights():
        text_encoder = Qwen3Model(config)

    expected_keys = set(text_encoder.state_dict().keys())
    converted_keys = set(state_dict.keys())
    missing_keys = expected_keys - converted_keys
    unexpected_keys = converted_keys - expected_keys
    if missing_keys or unexpected_keys:
        if missing_keys:
            print(f"ERROR: missing Qwen3 keys ({len(missing_keys)}):", file=sys.stderr)
            for key in sorted(missing_keys):
                print(key, file=sys.stderr)
        if unexpected_keys:
            print(f"ERROR: unexpected Qwen3 keys ({len(unexpected_keys)}):", file=sys.stderr)
            for key in sorted(unexpected_keys):
                print(key, file=sys.stderr)
        sys.exit(1)

    text_encoder.load_state_dict(state_dict, strict=True, assign=True)
    return text_encoder


def split_anima_transformer_checkpoint(
    state_dict: dict[str, torch.Tensor],
) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
    transformer_state_dict = {}
    text_conditioner_state_dict = {}
    adapter_prefix = "net.llm_adapter."

    for key, value in state_dict.items():
        if key.startswith(adapter_prefix):
            text_conditioner_state_dict[key.removeprefix(adapter_prefix)] = value
        else:
            transformer_state_dict[key] = value

    return transformer_state_dict, text_conditioner_state_dict


def save_pipeline(args, transformer, text_conditioner, text_encoder, vae):
    tokenizer = AutoTokenizer.from_pretrained(args.qwen_tokenizer_path)
    t5_tokenizer = T5TokenizerFast.from_pretrained(args.t5_tokenizer_path)
    scheduler = FlowMatchEulerDiscreteScheduler(shift=3.0)

    pipe = AnimaAutoBlocks().init_pipeline()
    pipe.update_components(
        text_encoder=text_encoder,
        tokenizer=tokenizer,
        t5_tokenizer=t5_tokenizer,
        text_conditioner=text_conditioner,
        transformer=transformer,
        vae=vae,
        scheduler=scheduler,
    )
    pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size=args.max_shard_size)


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--transformer_ckpt_path", type=str, required=True, help="Path to Anima DiT safetensors")
    parser.add_argument("--text_encoder_ckpt_path", type=str, required=True, help="Path to Qwen3 text encoder")
    parser.add_argument("--vae_ckpt_path", type=str, required=True, help="Path to Qwen-Image VAE safetensors")
    parser.add_argument("--qwen_tokenizer_path", type=str, default=None)
    parser.add_argument("--t5_tokenizer_path", type=str, default=None)
    parser.add_argument("--output_path", type=str, required=True)
    parser.add_argument("--save_pipeline", action="store_true")
    parser.add_argument("--dtype", default="bf16", choices=list(DTYPE_MAPPING.keys()))
    parser.add_argument("--max_shard_size", default="5GB")
    return parser.parse_args()


if __name__ == "__main__":
    args = get_args()
    output_path = pathlib.Path(args.output_path)
    dtype = DTYPE_MAPPING[args.dtype]

    raw_transformer_state_dict = load_file(args.transformer_ckpt_path, device="cpu")
    transformer_state_dict, text_conditioner_state_dict = split_anima_transformer_checkpoint(
        raw_transformer_state_dict
    )
    transformer = convert_transformer(
        "Cosmos-2.0-Diffusion-2B-Text2Image", state_dict=transformer_state_dict, weights_only=True
    ).to(dtype=dtype)
    text_conditioner = convert_text_conditioner(text_conditioner_state_dict).to(dtype=dtype)

    text_encoder_state_dict = load_file(args.text_encoder_ckpt_path, device="cpu")
    text_encoder = convert_text_encoder(text_encoder_state_dict).to(dtype=dtype)

    vae_state_dict = load_file(args.vae_ckpt_path, device="cpu")
    vae = convert_qwen_image_vae(vae_state_dict).to(dtype=dtype)

    if args.save_pipeline:
        if args.qwen_tokenizer_path is None or args.t5_tokenizer_path is None:
            raise ValueError("`--qwen_tokenizer_path` and `--t5_tokenizer_path` are required with `--save_pipeline`.")
        save_pipeline(args, transformer, text_conditioner, text_encoder, vae)
    else:
        output_path.mkdir(parents=True, exist_ok=True)
        transformer.save_pretrained(
            output_path / "transformer", safe_serialization=True, max_shard_size=args.max_shard_size
        )
        text_conditioner.save_pretrained(
            output_path / "text_conditioner", safe_serialization=True, max_shard_size=args.max_shard_size
        )
        text_encoder.save_pretrained(
            output_path / "text_encoder", safe_serialization=True, max_shard_size=args.max_shard_size
        )
        vae.save_pretrained(output_path / "vae", safe_serialization=True, max_shard_size=args.max_shard_size)