#!/usr/bin/env python3
"""Convert official Diffusers Wan transformer shards to WanModel key names.

No tensor is numerically changed.  Shards are processed one at a time so the
conversion does not need to materialize the 14B state dict twice in memory.
"""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path

from safetensors.torch import load_file, save_file


REPLACEMENTS = (
    ("attn1.to_q", "self_attn.q"),
    ("attn1.to_k", "self_attn.k"),
    ("attn1.to_v", "self_attn.v"),
    ("attn1.to_out.0", "self_attn.o"),
    ("attn1.norm_q", "self_attn.norm_q"),
    ("attn1.norm_k", "self_attn.norm_k"),
    ("attn2.to_q", "cross_attn.q"),
    ("attn2.to_k", "cross_attn.k"),
    ("attn2.to_v", "cross_attn.v"),
    ("attn2.to_out.0", "cross_attn.o"),
    ("attn2.norm_q", "cross_attn.norm_q"),
    ("attn2.norm_k", "cross_attn.norm_k"),
    ("ffn.net.0.proj", "ffn.0"),
    ("ffn.net.2", "ffn.2"),
    (".norm2.", ".norm3."),
    (".scale_shift_table", ".modulation"),
    ("condition_embedder.text_embedder.linear_1", "text_embedding.0"),
    ("condition_embedder.text_embedder.linear_2", "text_embedding.2"),
    ("condition_embedder.time_embedder.linear_1", "time_embedding.0"),
    ("condition_embedder.time_embedder.linear_2", "time_embedding.2"),
    ("condition_embedder.time_proj", "time_projection.1"),
    ("proj_out", "head.head"),
    ("scale_shift_table", "head.modulation"),
)


def rename(key: str) -> str:
    for source, target in REPLACEMENTS:
        key = key.replace(source, target)
    return key


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    source = args.source.resolve()
    output = args.output.resolve()
    output.mkdir(parents=True, exist_ok=True)
    source_index = json.loads((source / "diffusion_pytorch_model.safetensors.index.json").read_text())
    source_map: dict[str, str] = source_index["weight_map"]
    output_map = {rename(key): filename for key, filename in source_map.items()}
    if len(output_map) != len(source_map):
        raise RuntimeError("Wan key conversion produced a collision")

    for shard in sorted(set(source_map.values())):
        destination = output / shard
        if destination.exists() and destination.stat().st_size == (source / shard).stat().st_size:
            print(f"[skip] {shard}", flush=True)
            continue
        print(f"[convert] {shard}", flush=True)
        tensors = load_file(str(source / shard), device="cpu")
        converted = {rename(key): value.contiguous() for key, value in tensors.items()}
        temporary = destination.with_suffix(destination.suffix + ".tmp")
        save_file(converted, str(temporary))
        os.replace(temporary, destination)
        del tensors, converted

    index = {"metadata": source_index.get("metadata", {}), "weight_map": output_map}
    temporary_index = output / "diffusion_pytorch_model.safetensors.index.json.tmp"
    temporary_index.write_text(json.dumps(index, indent=2, sort_keys=True) + "\n")
    os.replace(temporary_index, output / "diffusion_pytorch_model.safetensors.index.json")
    print(f"[complete] tensors={len(output_map)} shards={len(set(output_map.values()))}", flush=True)


if __name__ == "__main__":
    main()