| """Load a training checkpoint and save a lean inference-only model file. |
| |
| Keeps only the fp16 model weights (plus the config needed to rebuild the |
| architecture). Drops optimizer state, scaler state, and everything else. |
| |
| Usage: |
| python quantize.py --input ckpt.pt --output evotalk_fp16.pt |
| """ |
|
|
| import os |
| import argparse |
|
|
| import torch |
|
|
| from model import EvoTalk |
|
|
|
|
| def quantize(input_path, output_path): |
| ckpt = torch.load(input_path, map_location="cpu", weights_only=False) |
|
|
| config = ckpt.get("config", None) |
| if config is None: |
| raise ValueError("checkpoint has no config; cannot rebuild the model") |
|
|
| model = EvoTalk(config) |
| model.load_state_dict(ckpt["model"]) |
| model.eval() |
|
|
| state = { |
| "model": {k: v.half() for k, v in model.state_dict().items()}, |
| "config": config, |
| "step": ckpt.get("step", 0), |
| } |
|
|
| os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True) |
| torch.save(state, output_path) |
| print(f"saved fp16 weights -> {output_path}") |
| print(f"weights: {sum(v.numel() for v in state['model'].values()) / 1e6:.2f}M " |
| f"({os.path.getsize(output_path) / 1e6:.1f} MB)") |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--input", type=str, default="ckpt.pt") |
| parser.add_argument("--output", type=str, default="evotalk_fp16.pt") |
| args = parser.parse_args() |
|
|
| quantize(args.input, args.output) |
|
|
