SpongeBobFan2002/openaudio-s1-mini

1import click
2import torch
3from loguru import logger
4
5
6@click.command()
7@click.argument("model_path")
8@click.argument("output_path")
9def main(model_path, output_path):
10    if model_path == output_path:
11        logger.error("Model path and output path are the same")
12        return
13
14    logger.info(f"Loading model from {model_path}")
15    state_dict = torch.load(model_path, map_location="cpu")["state_dict"]
16    torch.save(state_dict, output_path)
17    logger.info(f"Model saved to {output_path}")
18
19
20if __name__ == "__main__":
21    main()
22