|
|
| """
|
| Utility to prune training checkpoints down to the tensors required for
|
| IndexTTS2 inference.
|
|
|
| Typical usage:
|
|
|
| uv run python tools/prune_gpt_checkpoint.py \
|
| --input trained_ckpts_multilingual/model_step33000.pth \
|
| --output checkpoints/gpt_multilingual_pruned.pth
|
|
|
| The resulting file mirrors the structure of the shipped inference
|
| checkpoints (an OrderedDict of model weights) so it can be loaded by
|
| `indextts/infer_v2_modded.py` and the WebUI.
|
| """
|
|
|
| from __future__ import annotations
|
|
|
| import argparse
|
| from collections import OrderedDict
|
| from pathlib import Path
|
| from typing import Iterable, Tuple
|
|
|
| import torch
|
|
|
|
|
| def parse_args() -> argparse.Namespace:
|
| parser = argparse.ArgumentParser(
|
| description="Strip training artefacts (optimizer, scheduler, scaler, etc.) "
|
| "from a checkpoint and emit an inference-ready weight file."
|
| )
|
| parser.add_argument(
|
| "--input",
|
| type=Path,
|
| required=True,
|
| help="Path to the training checkpoint (.pth) produced by train_gpt_v2.py.",
|
| )
|
| parser.add_argument(
|
| "--output",
|
| type=Path,
|
| required=True,
|
| help="Destination path for the pruned checkpoint.",
|
| )
|
| parser.add_argument(
|
| "--dtype",
|
| choices=("keep", "float32", "float16"),
|
| default="keep",
|
| help="Optional tensor dtype override for the saved weights (default: keep original).",
|
| )
|
| parser.add_argument(
|
| "--dry-run",
|
| action="store_true",
|
| help="Parse and report summary without writing the output file.",
|
| )
|
| return parser.parse_args()
|
|
|
|
|
| def _coerce_dtype(tensor: torch.Tensor, dtype: str) -> torch.Tensor:
|
| if dtype == "keep":
|
| return tensor
|
| target = torch.float16 if dtype == "float16" else torch.float32
|
| if tensor.dtype == target:
|
| return tensor
|
| return tensor.to(dtype=target)
|
|
|
|
|
| def summarise_state(state: OrderedDict) -> Tuple[int, int]:
|
| tensor_count = 0
|
| total_params = 0
|
| for value in state.values():
|
| if isinstance(value, torch.Tensor):
|
| tensor_count += 1
|
| total_params += value.numel()
|
| return tensor_count, total_params
|
|
|
|
|
| def main() -> None:
|
| args = parse_args()
|
|
|
| if not args.input.exists():
|
| raise FileNotFoundError(f"Input checkpoint not found: {args.input}")
|
|
|
| checkpoint = torch.load(args.input, map_location="cpu")
|
| if isinstance(checkpoint, dict) and "model" in checkpoint:
|
| state_dict = checkpoint["model"]
|
| elif isinstance(checkpoint, OrderedDict):
|
| state_dict = checkpoint
|
| else:
|
| raise RuntimeError(
|
| "Unrecognised checkpoint structure. Expected either a plain state_dict "
|
| "or a dict containing the 'model' key produced by train_gpt_v2.py."
|
| )
|
|
|
| if not isinstance(state_dict, dict):
|
| raise RuntimeError("Model state should be a dict-like object.")
|
|
|
| dtype_choice = args.dtype
|
|
|
| pruned = OrderedDict()
|
| for key, value in state_dict.items():
|
| if isinstance(value, torch.Tensor):
|
| tensor = value.detach().cpu()
|
| tensor = _coerce_dtype(tensor, dtype_choice)
|
| pruned[key] = tensor
|
| else:
|
|
|
| pruned[key] = value
|
|
|
| tensor_count, param_total = summarise_state(pruned)
|
| print(f"[Prune] Retained {tensor_count} tensors / {param_total:,} parameters.")
|
|
|
| if args.dry_run:
|
| print("[Prune] Dry run enabled; no file was written.")
|
| return
|
|
|
| args.output.parent.mkdir(parents=True, exist_ok=True)
|
| torch.save(pruned, args.output)
|
| size_mb = args.output.stat().st_size / (1024 * 1024)
|
| print(f"[Prune] Wrote {args.output} ({size_mb:.2f} MiB).")
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
