# SPDX-License-Identifier: Apache-2.0
"""Convert a `stable_audio_tools` raw checkpoint into a Diffusers-style
per-component layout consumable by FastVideo's standard
`ComposedPipelineBase` loader.

Why: the published Stability AI Stable Audio Open repos ship a single
monolithic `model.safetensors` containing DiT (`model.model.*`),
VAE / pretransform (`pretransform.model.*`), and NumberConditioner
(`conditioner.conditioners.*`) state dicts under shared prefixes.
FastVideo's standard loader expects the Diffusers per-subfolder layout
(`model_index.json` + per-component subfolders). This script splits the
monolithic checkpoint, drops the host-pipeline prefixes, normalizes
LayerNorm `gamma`/`beta` keys to `weight`/`bias`, and writes a
Diffusers-style tree.

Usage::

    python scripts/checkpoint_conversion/stable_audio_to_diffusers.py \
        --src stabilityai/stable-audio-open-1.0 \
        --dst converted_weights/stable_audio_open_1_0_diffusers

    # Or with a local snapshot directory:
    python scripts/checkpoint_conversion/stable_audio_to_diffusers.py \
        --src /path/to/stable-audio-open-1.0 \
        --dst converted_weights/stable_audio_open_small_diffusers

The output tree is::

    <dst>/
    ├── model_index.json
    ├── transformer/{config.json, diffusion_pytorch_model.safetensors}
    ├── vae/{config.json, diffusion_pytorch_model.safetensors}
    ├── conditioner/{config.json, diffusion_pytorch_model.safetensors}
    ├── text_encoder/  (copied from src if present)
    ├── tokenizer/     (copied from src if present)
    └── scheduler/     (copied from src if present)

Push to the FastVideo HF org with `huggingface-cli upload`. Confirm
visibility (`--private` for staging, omit for public release)::

    huggingface-cli upload --private FastVideo/stable-audio-open-1.0-Diffusers \
        converted_weights/stable-audio-open-1.0-Diffusers
    huggingface-cli upload --private FastVideo/stable-audio-open-small-Diffusers \
        converted_weights/stable-audio-open-small-Diffusers
"""
from __future__ import annotations

import argparse
import json
import os
import shutil
from pathlib import Path
from typing import Any

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

# Map host-pipeline prefix → output subfolder. Order matters only for
# logging.
_COMPONENT_PREFIXES: dict[str, str] = {
    "model.model.": "transformer",
    "pretransform.model.": "vae",
    "conditioner.": "conditioner",
}

# Subfolders we copy verbatim from the source repo. The `vae/` subfolder
# in `stabilityai/stable-audio-open-1.0` already ships in Diffusers
# format with the key naming our `OobleckVAE` expects (`decoder.block.X.
# conv_t1.weight_g` etc.) — preferred over the raw `pretransform.model.*`
# extraction below, which uses a different module nesting (`decoder.
# layers.X.layers.Y.*`) that would need a separate remap.
_PASSTHROUGH_SUBFOLDERS = ("text_encoder", "tokenizer", "scheduler", "vae")

_MODEL_INDEX = {
    "_class_name": "StableAudioPipeline",
    # `ComposedPipelineBase._load_config` requires this field; the value
    # is informational (we don't pin a Diffusers version).
    "_diffusers_version": "0.30.0",
    "_fastvideo_converted_from": None,  # filled in at write time
    "transformer": ["fastvideo.models.dits.stable_audio", "StableAudioDiT"],
    "vae": ["fastvideo.models.vaes.oobleck", "OobleckVAE"],
    "conditioner": [
        "fastvideo.models.encoders.stable_audio_conditioner",
        "StableAudioMultiConditioner",
    ],
    "text_encoder": ["transformers", "T5EncoderModel"],
    "tokenizer": ["transformers", "T5TokenizerFast"],
    "scheduler": ["diffusers", "CosineDPMSolverMultistepScheduler"],
}


def _resolve_src(src: str) -> Path:
    """Accept either a HF repo id or a local directory."""
    if os.path.isdir(src):
        return Path(src)
    from fastvideo.utils import resolve_hf_token
    return Path(snapshot_download(repo_id=src, token=resolve_hf_token()))


def _rename_layernorm_keys(state: dict[str, Any]) -> dict[str, Any]:
    """`nn.LayerNorm` ships `gamma`/`beta` in the official checkpoint;
    torch's `nn.LayerNorm` uses `weight`/`bias`."""
    out: dict[str, Any] = {}
    for k, v in state.items():
        if k.endswith(".gamma"):
            out[k[:-len(".gamma")] + ".weight"] = v
        elif k.endswith(".beta"):
            out[k[:-len(".beta")] + ".bias"] = v
        else:
            out[k] = v
    return out


def _split_state_dict(src_safetensors: Path) -> dict[str, dict[str, Any]]:
    """Read `model.safetensors` and bucket its keys by the host-pipeline
    prefix listed in `_COMPONENT_PREFIXES`."""
    buckets: dict[str, dict[str, Any]] = {name: {} for name in _COMPONENT_PREFIXES.values()}
    skipped = 0
    with safe_open(str(src_safetensors), framework="pt") as f:
        for k in f.keys():
            for prefix, subfolder in _COMPONENT_PREFIXES.items():
                if k.startswith(prefix):
                    buckets[subfolder][k[len(prefix):]] = f.get_tensor(k)
                    break
            else:
                skipped += 1
    if skipped:
        print(f"  skipped {skipped} keys (no matching prefix)")
    for subfolder, state in buckets.items():
        if subfolder == "transformer":
            buckets[subfolder] = _rename_layernorm_keys(state)
        print(f"  {subfolder}: {len(buckets[subfolder])} keys")
    return buckets


def _detect_projection_flags(diff_cfg: dict[str, Any], dit_state: dict[str, Any]) -> tuple[bool, bool]:
    """Decide `project_cond_tokens` / `project_global_cond` by comparing
    the actual `to_cond_embed.0` / `to_global_embed.0` weight shapes
    in the DiT state dict against `cond_token_dim` / `global_cond_dim`.
    Upstream's factory toggles these per-variant in ways that aren't
    derivable from `embed_dim` alone (SA-1.0 keeps `cond_embed_dim` =
    `cond_token_dim` and uses GQA in cross-attn; SA-small projects to
    `embed_dim` and uses MHA).
    """
    cond_dim = diff_cfg.get("cond_token_dim")
    glob_dim = diff_cfg.get("global_cond_dim")
    cond_w = dit_state.get("to_cond_embed.0.weight")
    glob_w = dit_state.get("to_global_embed.0.weight")
    project_cond = (cond_w is not None and cond_w.shape[0] != cond_dim)
    project_glob = (glob_w is not None and glob_w.shape[0] != glob_dim)
    return project_cond, project_glob


def _component_config(model_config: dict[str, Any]) -> dict[str, dict[str, Any]]:
    """Pull per-component sizing out of the official `model_config.json`
    so the converted repo records the architecture authoritatively."""
    diff = model_config["model"]["diffusion"]
    pre = model_config["model"]["pretransform"]["config"]
    cond = model_config["model"]["conditioning"]
    # Translate upstream `model_config.json` field names to FastVideo's
    # `StableAudioArchConfig` fields. The standard `TransformerLoader`
    # calls `dit_config.update_model_arch(json_dict)` which raises on
    # unknown keys.
    diff_cfg = dict(diff["config"])
    if "num_heads" in diff_cfg:
        diff_cfg["num_attention_heads"] = diff_cfg.pop("num_heads")
    diff_cfg.pop("transformer_type", None)  # always "continuous_transformer"
    # `attn_kwargs.qk_norm` (small variant) → top-level `qk_norm`.
    attn_kwargs = diff_cfg.pop("attn_kwargs", None) or {}
    if "qk_norm" in attn_kwargs:
        diff_cfg["qk_norm"] = attn_kwargs["qk_norm"]
    # `cross_attention_cond_ids` / `global_cond_ids` belong to the
    # conditioner, not the DiT — keep them out of transformer/config.json.
    transformer_cfg = {
        "_class_name": "StableAudioDiT",
        **diff_cfg,
    }
    # Mirror Diffusers' `AutoencoderOobleck` config field naming so
    # `OobleckVAEArchConfig.update_model_arch` (called by `VAELoader`)
    # accepts the keys.
    enc_cfg = pre["encoder"]["config"]
    dec_cfg = pre["decoder"]["config"]
    vae_cfg = {
        "_class_name": "AutoencoderOobleck",
        "encoder_hidden_size": enc_cfg["channels"],
        "downsampling_ratios": enc_cfg["strides"],
        "channel_multiples": enc_cfg["c_mults"],
        "decoder_channels": dec_cfg["channels"],
        "decoder_input_channels": dec_cfg["latent_dim"],
        "audio_channels": pre.get("io_channels", 2),
        "sampling_rate": model_config.get("sample_rate", 44100),
    }
    conditioner_cfg = {
        "_class_name": "StableAudioMultiConditioner",
        "cond_dim": cond["cond_dim"],
        "configs": cond["configs"],
        "cross_attention_cond_ids": diff["cross_attention_cond_ids"],
        "global_cond_ids": diff["global_cond_ids"],
    }
    return {
        "transformer": transformer_cfg,
        "vae": vae_cfg,
        "conditioner": conditioner_cfg,
    }


def _write_component(dst_dir: Path, name: str, state: dict[str, Any], config: dict[str, Any]) -> None:
    sub = dst_dir / name
    sub.mkdir(parents=True, exist_ok=True)
    save_file(state, str(sub / "diffusion_pytorch_model.safetensors"))
    with open(sub / "config.json", "w") as f:
        json.dump(config, f, indent=2)
    print(f"  wrote {sub}/")


def _copy_passthrough(src_dir: Path, dst_dir: Path) -> list[str]:
    copied: list[str] = []
    for sub in _PASSTHROUGH_SUBFOLDERS:
        s = src_dir / sub
        if not s.is_dir():
            continue
        d = dst_dir / sub
        if d.exists():
            shutil.rmtree(d)
        shutil.copytree(s, d, ignore=shutil.ignore_patterns("*.bin", "*.ckpt"))
        copied.append(sub)
        print(f"  copied {sub}/")
    return copied


def convert(src: str, dst: str) -> None:
    src_dir = _resolve_src(src)
    dst_dir = Path(dst)
    dst_dir.mkdir(parents=True, exist_ok=True)
    print(f"src: {src_dir}")
    print(f"dst: {dst_dir}")

    monolithic = src_dir / "model.safetensors"
    if not monolithic.is_file():
        raise FileNotFoundError(f"Expected monolithic safetensors at {monolithic}")
    cfg_path = src_dir / "model_config.json"
    if not cfg_path.is_file():
        raise FileNotFoundError(f"Expected model_config.json at {cfg_path} (Stable Audio's authoritative "
                                f"per-component arch config).")
    with open(cfg_path) as f:
        model_config = json.load(f)

    print("\n[1/4] Splitting monolithic state dict by component prefix:")
    buckets = _split_state_dict(monolithic)

    print("\n[2/4] Copying passthrough subfolders (text_encoder/tokenizer/scheduler/vae):")
    copied = _copy_passthrough(src_dir, dst_dir)

    print("\n[3/4] Writing per-component configs + safetensors:")
    component_cfgs = _component_config(model_config)
    project_cond, project_glob = _detect_projection_flags(component_cfgs["transformer"], buckets["transformer"])
    component_cfgs["transformer"]["project_cond_tokens"] = project_cond
    component_cfgs["transformer"]["project_global_cond"] = project_glob
    for name in ("transformer", "vae", "conditioner"):
        if name in copied:
            print(f"  skipped {name}/ (already copied from source in Diffusers format)")
            continue
        _write_component(dst_dir, name, buckets[name], component_cfgs[name])

    print("\n[4/4] Writing model_index.json:")
    index = dict(_MODEL_INDEX)
    # Prefer the HF repo id over a snapshot directory (which is
    # machine-local and full of hash gunk).
    index["_fastvideo_converted_from"] = src if "/" in src and not os.path.isdir(src) else src_dir.name
    # Drop entries for subfolders we didn't actually populate.
    available = {"transformer", "vae", "conditioner", *copied}
    index = {k: v for k, v in index.items() if not k or k.startswith("_") or k in available}
    with open(dst_dir / "model_index.json", "w") as f:
        json.dump(index, f, indent=2)
    print(f"  wrote {dst_dir}/model_index.json")
    print(f"\nDone. Push to HF with:")
    print(f"  huggingface-cli upload <FastVideo/repo-name> {dst_dir}")


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    parser.add_argument("--src",
                        required=True,
                        help="HF repo id or local directory containing model.safetensors + model_config.json")
    parser.add_argument("--dst", required=True, help="Output directory for the Diffusers-format repo")
    args = parser.parse_args()
    convert(args.src, args.dst)


if __name__ == "__main__":
    main()