| """Initialize the decoder student by channel-slicing a frozen pretrained vocoder. |
| |
| Training a decoder this small from scratch does not converge to intelligible |
| audio; slicing the top-left corner of a pretrained neural vocoder's decoder |
| weights does. Pass any ConvNeXt-style vocoder state dict with --vocoder-ckpt; |
| this takes the first `--dim` channels, the first `--pw` hidden units and the |
| first `--blocks` blocks (dim 130 / pw 337 / blocks 4 for the shipped model). |
| |
| The sliced checkpoint still needs decoder training (train_decoder.py) before |
| it is usable: a fresh slice renders noise, and only fine-tuning brings it back. |
| |
| python init_decoder.py --dim 130 --pw 337 --blocks 4 \ |
| --vocoder-ckpt <vocos-mel-24khz state dict> --out decoder_init.pt |
| """ |
| import sys |
| from pathlib import Path |
|
|
| import torch |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
| from microtts.models import DecoderStudent |
|
|
|
|
| def main(): |
| import argparse |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--dim", type=int, default=130) |
| ap.add_argument("--pw", type=int, default=337) |
| ap.add_argument("--blocks", type=int, default=4) |
| ap.add_argument("--vocoder-ckpt", default=None, |
| help="state dict / checkpoint of a ConvNeXt-style neural vocoder") |
| ap.add_argument("--out", default="decoder_init.pt") |
| args = ap.parse_args() |
|
|
| if args.vocoder_ckpt: |
| obj = torch.load(args.vocoder_ckpt, map_location="cpu", weights_only=False) |
| sd = obj.state_dict() if hasattr(obj, "state_dict") else obj |
| else: |
| raise ImportError("pass --vocoder-ckpt pointing at a ConvNeXt-style vocoder state dict") |
|
|
| student = DecoderStudent(dim=args.dim, pw_hidden=args.pw, blocks=args.blocks) |
| ssd = student.state_dict() |
|
|
| |
| ssd["embed.weight"] = sd["backbone.embed.weight"][:args.dim].contiguous() |
| ssd["embed.bias"] = sd["backbone.embed.bias"][:args.dim].contiguous() |
| ssd["norm.weight"] = sd["backbone.norm.weight"][:args.dim].contiguous() |
| ssd["norm.bias"] = sd["backbone.norm.bias"][:args.dim].contiguous() |
| for i in range(args.blocks): |
| src = f"backbone.convnext.{i}" |
| dst = f"blocks.{i}" |
| ssd[f"{dst}.dw.weight"] = sd[f"{src}.dwconv.weight"][:args.dim].contiguous() |
| ssd[f"{dst}.dw.bias"] = sd[f"{src}.dwconv.bias"][:args.dim].contiguous() |
| ssd[f"{dst}.norm.weight"] = sd[f"{src}.norm.weight"][:args.dim].contiguous() |
| ssd[f"{dst}.norm.bias"] = sd[f"{src}.norm.bias"][:args.dim].contiguous() |
| |
| ssd[f"{dst}.pw0.weight"] = sd[f"{src}.pwconv1.weight"][:args.pw, :args.dim].contiguous() |
| ssd[f"{dst}.pw0.bias"] = sd[f"{src}.pwconv1.bias"][:args.pw].contiguous() |
| ssd[f"{dst}.pw1.weight"] = sd[f"{src}.pwconv2.weight"][:args.dim, :args.pw].contiguous() |
| ssd[f"{dst}.pw1.bias"] = sd[f"{src}.pwconv2.bias"][:args.dim].contiguous() |
| ssd[f"{dst}.gamma"] = sd[f"{src}.gamma"][:args.dim].contiguous() |
| |
| ssd["fnorm.weight"] = sd["backbone.final_layer_norm.weight"][:args.dim].contiguous() |
| ssd["fnorm.bias"] = sd["backbone.final_layer_norm.bias"][:args.dim].contiguous() |
| ssd["head.weight"] = sd["head.out.weight"][:, :args.dim].contiguous() |
| ssd["head.bias"] = sd["head.out.bias"].contiguous() |
|
|
| student.load_state_dict(ssd) |
| torch.save(student.state_dict(), args.out) |
| n = sum(p.numel() for p in student.parameters()) |
| print(f"saved {args.out} | student params: {n:,}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
