#!/usr/bin/env python
# Copyright 2026 The HuggingFace Team and SANA-WM Authors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Convert the public SANA-WM release into a diffusers-loadable directory.

Reads the ``Efficient-Large-Model/SANA-WM_bidirectional`` HF repo (or a local
mirror) and writes a directory ready for ``SanaWMPipeline.from_pretrained(path)``:

    <output_dir>/
    ├── model_index.json
    ├── tokenizer/
    ├── text_encoder/
    ├── vae/
    ├── transformer/
    ├── scheduler/
    └── refiner/
        ├── transformer/
        ├── connectors/
        ├── text_encoder/
        └── tokenizer/

Usage:
    python scripts/convert_sana_wm_to_diffusers.py \\
        --src Efficient-Large-Model/SANA-WM_bidirectional \\
        --dst /path/to/SANA-WM_bidirectional-diffusers \\
        [--no-refiner]

The output is local-only; no upload to the Hub.
"""

from __future__ import annotations

import argparse
import json
import shutil
from pathlib import Path

import torch
from huggingface_hub import snapshot_download
from safetensors import safe_open
from safetensors.torch import save_file


def _copy_subdir(src: Path, dst: Path) -> None:
    if dst.exists():
        shutil.rmtree(dst)
    shutil.copytree(src, dst, symlinks=False)


def _cast_to_bfloat16(component: Path) -> None:
    """Rewrite a component's safetensors shards in bfloat16.

    SANA-WM is served in `torch.bfloat16`, so `from_pretrained(torch_dtype=torch.bfloat16)` casts these weights
    anyway; storing them cast halves the download without changing a single value. Non-float32 tensors are left
    alone, and the VAE is deliberately not passed through here since it runs in `torch.float32`.
    """
    shards = sorted(component.glob("*.safetensors"))
    if not shards:
        return
    total = 0
    for shard in shards:
        with safe_open(shard, framework="pt") as f:
            metadata = f.metadata()
            tensors = {
                k: (lambda t: t.to(torch.bfloat16) if t.dtype == torch.float32 else t)(f.get_tensor(k))
                for k in f.keys()
            }
        tmp = shard.with_suffix(shard.suffix + ".tmp")
        save_file(tensors, tmp, metadata=metadata)
        tmp.replace(shard)
        total += shard.stat().st_size
        del tensors
    index = component / "model.safetensors.index.json"
    if index.exists():
        payload = json.loads(index.read_text())
        payload["metadata"]["total_size"] = total
        index.write_text(json.dumps(payload, indent=2))


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--src", default="Efficient-Large-Model/SANA-WM_bidirectional", help="HF repo or local dir")
    parser.add_argument("--dst", required=True, type=Path, help="Output directory")
    parser.add_argument("--no-refiner", action="store_true", help="Skip refiner export")
    parser.add_argument(
        "--dst-refiner",
        type=Path,
        default=None,
        help="Output directory for the stage-2 refiner pipeline (default: `<dst>-refiner`)",
    )
    parser.add_argument(
        "--torch-dtype", default="bfloat16", choices=["bfloat16", "float16", "float32"], help="Weight dtype"
    )
    args = parser.parse_args()

    torch_dtype = getattr(torch, args.torch_dtype)
    dst: Path = args.dst.absolute()
    dst.mkdir(parents=True, exist_ok=True)

    # Resolve the source on disk (snapshot_download for HF repos, otherwise use as-is).
    src_path = Path(args.src)
    if not src_path.is_dir():
        print(f"[convert] snapshot_download({args.src}) …")
        src_path = Path(snapshot_download(args.src))
    print(f"[convert] source: {src_path}")

    # 1. VAE (already diffusers format under <src>/vae).
    print("[convert] vae …")
    _copy_subdir(src_path / "vae", dst / "vae")

    # 2. Tokenizer + text encoder — fetch via the configured Gemma-2 repo.
    # We save the full ``Gemma2ForCausalLM``; the pipeline grabs the decoder
    # at runtime via ``self.text_encoder.model(...)``. This matches the sana
    # inference recipe of ``AutoModelForCausalLM.from_pretrained(...).get_decoder()``
    # and avoids subtle state-dict prefix differences when saving just the
    # decoder submodule.
    print("[convert] tokenizer + text_encoder (gemma-2-2b-it) …")
    from transformers import AutoModelForCausalLM, AutoTokenizer

    gemma_repo = "Efficient-Large-Model/gemma-2-2b-it"
    tokenizer = AutoTokenizer.from_pretrained(gemma_repo)
    tokenizer.padding_side = "right"
    tokenizer.save_pretrained(dst / "tokenizer")
    text_encoder = AutoModelForCausalLM.from_pretrained(gemma_repo, torch_dtype=torch_dtype)
    text_encoder.save_pretrained(dst / "text_encoder")
    del text_encoder

    # 3. Transformer (SanaWMTransformer3DModel) — load the public DiT, save in diffusers format.
    print("[convert] transformer (SanaWMTransformer3DModel) …")
    from diffusers import SanaWMTransformer3DModel

    transformer = SanaWMTransformer3DModel().to(torch_dtype).eval()
    dit_ckpt = src_path / "dit" / "sana_wm_1600m_720p.safetensors"
    if not dit_ckpt.is_file():
        raise FileNotFoundError(f"DiT checkpoint not found at {dit_ckpt}")
    from safetensors.torch import load_file

    sd = load_file(str(dit_ckpt))
    sd.pop("pos_embed", None)  # unused at inference (wan_rope is computed on-the-fly)
    # `y_embedding` is the training-time null-caption embedding; inference never reads it.
    sd.pop("y_embedder.y_embedding", None)
    # The timestep MLP is diffusers' `TimestepEmbedding`, which names its layers `linear_1` / `linear_2`.
    sd = {
        k.replace("t_embedder.mlp.0.", "t_embedder.linear_1.").replace("t_embedder.mlp.2.", "t_embedder.linear_2."): v
        for k, v in sd.items()
    }
    # `.pos_embed` entries are non-persistent buffers rebuilt at construction, so they are
    # expected to be absent from the converted state dict; anything else means the mapping
    # is wrong and would silently produce a broken transformer.
    missing, unexpected = transformer.load_state_dict(sd, strict=False)
    missing = [k for k in missing if not k.endswith(".pos_embed")]
    if missing or unexpected:
        raise RuntimeError(
            "State dict does not match `SanaWMTransformer3DModel`.\n"
            f"  missing keys ({len(missing)}): {missing[:10]}{' …' if len(missing) > 10 else ''}\n"
            f"  unexpected keys ({len(unexpected)}): {unexpected[:10]}{' …' if len(unexpected) > 10 else ''}"
        )
    transformer.save_pretrained(dst / "transformer")
    del transformer, sd

    # 4. Scheduler — FlowMatchEulerDiscreteScheduler config.
    print("[convert] scheduler …")
    from diffusers import FlowMatchEulerDiscreteScheduler

    FlowMatchEulerDiscreteScheduler(shift=9.8).save_pretrained(dst / "scheduler")

    # 5. Refiner (LTX-2): a standalone DiffusionPipeline written to its own output
    #    directory, the way SDXL ships base and refiner as separate repos.
    #    `DiffusionPipeline.from_pretrained` has no `subfolder` argument, so a nested
    #    folder would silently load the *base* pipeline's components instead. Copy the LTX-2 sub-model folders as-is, split out a ``tokenizer/``
    #    folder, add a ``scheduler/`` (FlowMatchEulerDiscreteScheduler), and write the
    #    manifest. The transformer weights are LTX-2's, but the refiner drives them
    #    through its own model class, so the manifest names that class.
    if not args.no_refiner:
        print("[convert] refiner …")
        from transformers import AutoTokenizer

        refiner_src = src_path / "refiner"
        refiner_dst = args.dst_refiner or dst.with_name(dst.name + "-refiner")
        refiner_dst = Path(refiner_dst)
        refiner_dst.mkdir(parents=True, exist_ok=True)
        for sub in ("transformer", "connectors", "text_encoder"):
            if (refiner_src / sub).is_dir():
                _copy_subdir(refiner_src / sub, refiner_dst / sub)
                _cast_to_bfloat16(refiner_dst / sub)

        # The weights are LTX-2's, but they are driven by `SanaWMLTX2RefinerTransformer3DModel`
        # (same submodule layout, streaming-attention forward), so point the config at that class.
        transformer_config_path = refiner_dst / "transformer" / "config.json"
        transformer_config = json.loads(transformer_config_path.read_text())
        transformer_config["_class_name"] = "SanaWMLTX2RefinerTransformer3DModel"
        transformer_config_path.write_text(json.dumps(transformer_config, indent=2))

        # Tokenizer lives co-located with the Gemma-3 text encoder in the release;
        # re-save it into its own subfolder so it registers as a pipeline component.
        refiner_tokenizer = AutoTokenizer.from_pretrained(refiner_src / "text_encoder")
        refiner_tokenizer.save_pretrained(refiner_dst / "tokenizer")

        # Scheduler carries the distilled sigma schedule; shift=1.0 leaves the
        # explicit sigmas passed at inference time unmodified.
        FlowMatchEulerDiscreteScheduler(shift=1.0).save_pretrained(refiner_dst / "scheduler")

        refiner_index = {
            "_class_name": "SanaWMLTX2Refiner",
            "_diffusers_version": "0.38.0",
            "transformer": ["diffusers", "SanaWMLTX2RefinerTransformer3DModel"],
            # LTX2TextConnectors lives in diffusers.pipelines.ltx2 (not top-level),
            # so the loader resolves it via the pipeline-module path ("ltx2", ...).
            "connectors": ["ltx2", "LTX2TextConnectors"],
            "tokenizer": ["transformers", type(refiner_tokenizer).__name__],
            "text_encoder": ["transformers", "Gemma3ForConditionalGeneration"],
            "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"],
            "text_max_sequence_length": 1024,
        }
        (refiner_dst / "model_index.json").write_text(json.dumps(refiner_index, indent=2))

    # 6. model_index.json — the top-level diffusers manifest.
    print("[convert] model_index.json …")
    index = {
        "_class_name": "SanaWMPipeline",
        "_diffusers_version": "0.38.0",
        "tokenizer": ["transformers", "GemmaTokenizerFast"],
        "text_encoder": ["transformers", "Gemma2ForCausalLM"],
        "vae": ["diffusers", "AutoencoderKLLTX2Video"],
        "transformer": ["diffusers", "SanaWMTransformer3DModel"],
        "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"],
    }
    (dst / "model_index.json").write_text(json.dumps(index, indent=2))

    print(f"[convert] done — wrote {dst}")
    if not args.no_refiner:
        print(f"[convert] done — wrote {refiner_dst}")


if __name__ == "__main__":
    main()