| """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"}) |
|
|
