"""The released AudioSR checkpoint, trimmed to what super-resolution uses, as safetensors.

    python convert.py basic audiosr-basic.safetensors

Kept: the diffusion UNet (`model.`), the VAE with its HiFi-GAN vocoder
(`first_stage_model.`) and the scalar/schedule buffers. Dropped: CLAP
(`clap.`, unused for super-resolution), the EMA copy of the UNet (`model_ema.`,
inference uses the live weights) and the condition encoder's VAE
(`cond_stage_models.0.vae.`), after checking it equals `first_stage_model`.
"""

import sys

import torch
from audiosr.utils import download_checkpoint
from safetensors.torch import save_file

name, out = sys.argv[1], sys.argv[2]
state = torch.load(download_checkpoint(name), map_location="cpu", weights_only=False)["state_dict"]

cond_prefix = "cond_stage_models.0.vae."
mismatched = [
    k for k in state
    if k.startswith(cond_prefix) and not torch.equal(state[k], state["first_stage_model." + k[len(cond_prefix):]])
]
assert not mismatched, f"condition VAE differs from first_stage_model: {mismatched[:5]}"

dropped = ("clap.", "model_ema.", "cond_stage_models.")
kept = {k: v.contiguous() for k, v in state.items() if not k.startswith(dropped)}
total = sum(v.numel() for v in kept.values())
print(f"kept {len(kept)} of {len(state)} tensors, {total / 1e6:.0f}M values, {total * 4 / 1e9:.2f} GB at fp32")
for prefix in ("model.", "first_stage_model.encoder", "first_stage_model.decoder", "first_stage_model.vocoder"):
    n = sum(v.numel() for k, v in kept.items() if k.startswith(prefix))
    print(f"  {prefix:30} {n / 1e6:7.1f}M")
save_file(kept, out, metadata={"source": f"audiosr {name}", "dtype": "float32"})