tuna2134/sbv2-v2
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Improve: Models can now be loaded directly onto NVIDIA GPUs

Browse Source
This commit is contained in:
tsukumi
2024-08-11 08:40:50 +09:00
parent 0e71d7ef7c
commit cfd6bd5bb6
4 changed files with 8 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
style_bert_vits2/models/infer.py
View File

@@ -86,10 +86,10 @@ def get_net_g(model_path: str, version: str, device: str, hps: HyperParameters):
_ = net_g.eval()
if model_path.endswith(".pth") or model_path.endswith(".pt"):
_ = utils.checkpoints.load_checkpoint(
model_path, net_g, None, skip_optimizer=True
model_path, net_g, None, skip_optimizer=True, device=device
)
elif model_path.endswith(".safetensors"):
_ = utils.safetensors.load_safetensors(model_path, net_g, True)
_ = utils.safetensors.load_safetensors(model_path, net_g, True, device=device)
else:
raise ValueError(f"Unknown model format: {model_path}")
return net_g