Spaces:

zomehwh

/

vits-models

like 746

Running

You can remove the optimizer that is not needed during model inference. This will reduce the size of the model to about 160MB.

import torch
model_path = "/path/to/model.pth"
output_path = "/path/to/model-p.pth"
checkpoint_dict = torch.load(model_path, map_location="cpu")
checkpoint_dict_new = {
    "model": {}
}
for k, v in checkpoint_dict.items():
    if k == "optimizer": 
        print("remove optimizer")
        continue
    checkpoint_dict_new[k] = v
torch.save(checkpoint_dict_new, output_path)

You can also remove some weights irrelevant to inference and convert other weights to half-precision, it's OK. This will reduce the size of the model to about 60MB.

import torch
model_path = "/path/to/model-p.pth"
output_path = "/path/to/model-p2.pth"
checkpoint_dict = torch.load(model_path, map_location="cpu")
checkpoint_dict_new = {
    "model": {}
}
for k, v in checkpoint_dict.items():
    if k == "model":
        for k2, v2 in checkpoint_dict['model'].items():
            if "enc_q" in k2:
                print("remove enc_q")
                continue
            checkpoint_dict_new['model'][k2] = checkpoint_dict['model'][k2].half()
        continue
    checkpoint_dict_new[k] = v
torch.save(checkpoint_dict_new, output_path)

Thanks!