| import os |
| import torch |
| import logging |
| import requests |
| import zipfile |
|
|
| from lib.models import ( |
| SynthesizerTrnMs256NSFsid, |
| SynthesizerTrnMs256NSFsid_nono, |
| SynthesizerTrnMs768NSFsid, |
| SynthesizerTrnMs768NSFsid_nono, |
| ) |
| from src.config import Config |
| from src.vc_infer_pipeline import VC |
|
|
| logging.getLogger("fairseq").setLevel(logging.WARNING) |
|
|
|
|
| class ModelLoader: |
| def __init__(self): |
| self.model_root = os.path.join(os.getcwd(), "weights") |
| self.config = Config() |
| self.model_list = [ |
| d |
| for d in os.listdir(self.model_root) |
| if os.path.isdir(os.path.join(self.model_root, d)) |
| ] |
| if len(self.model_list) == 0: |
| raise ValueError("No model found in `weights` folder") |
|
|
| self.model_name = "" |
| self.model_list.sort() |
|
|
| self.tgt_sr = None |
| self.net_g = None |
| self.vc = None |
| self.version = None |
| self.index_file = None |
| self.if_f0 = None |
|
|
| def _load_from_zip_url(self, url): |
| response = requests.get(url) |
| file_name = os.path.join( |
| self.model_root, os.path.basename(url[: url.index(".zip") + 4]) |
| ) |
| model_name = os.path.basename(file_name).replace(".zip", "") |
| print(f"Extraacting Model: {model_name}") |
|
|
| if response.status_code == 200: |
| with open(file_name, "wb") as file: |
| file.write(response.content) |
|
|
| with zipfile.ZipFile(file_name, "r") as zip_ref: |
| zip_ref.extractall(os.path.join(self.model_root, model_name)) |
| os.remove(file_name) |
| else: |
| print("Could not download model: {model_name}") |
| return model_name |
|
|
| def load(self, model_name): |
| if "http" in model_name: |
| model_name = self._load_from_zip_url(model_name) |
|
|
| pth_files = [ |
| os.path.join(self.model_root, model_name, f) |
| for f in os.listdir(os.path.join(self.model_root, model_name)) |
| if f.endswith(".pth") |
| ] |
| if len(pth_files) == 0: |
| raise ValueError(f"No pth file found in {self.model_root}/{model_name}") |
|
|
| self.model_name = model_name |
| pth_path = pth_files[0] |
| print(f"Loading {pth_path}, model: {model_name}") |
|
|
| cpt = torch.load(pth_path, map_location="cpu") |
| self.tgt_sr = cpt["config"][-1] |
| cpt["config"][-3] = cpt["weight"]["emb_g.weight"].shape[0] |
| self.if_f0 = cpt.get("f0", 1) |
| self.version = cpt.get("version", "v1") |
|
|
| if self.version == "v1": |
| if self.if_f0 == 1: |
| self.net_g = SynthesizerTrnMs256NSFsid( |
| *cpt["config"], is_half=self.config.is_half |
| ) |
| else: |
| self.net_g = SynthesizerTrnMs256NSFsid_nono(*cpt["config"]) |
| elif self.version == "v2": |
| if self.if_f0 == 1: |
| self.net_g = SynthesizerTrnMs768NSFsid( |
| *cpt["config"], is_half=self.config.is_half |
| ) |
| else: |
| self.net_g = SynthesizerTrnMs768NSFsid_nono(*cpt["config"]) |
| else: |
| raise ValueError("Unknown version") |
|
|
| del self.net_g.enc_q |
| self.net_g.load_state_dict(cpt["weight"], strict=False) |
| print("Model loaded") |
| self.net_g.eval().to(self.config.device) |
|
|
| if self.config.is_half: |
| self.net_g = self.net_g.half() |
| else: |
| self.net_g = self.net_g.float() |
|
|
| self.vc = VC(self.tgt_sr, self.config) |
|
|
| index_files = [ |
| os.path.join(self.model_root, model_name, f) |
| for f in os.listdir(os.path.join(self.model_root, model_name)) |
| if f.endswith(".index") |
| ] |
|
|
| if len(index_files) == 0: |
| print("No index file found") |
| self.index_file = "" |
| else: |
| self.index_file = index_files[0] |
| print(f"Index file found: {self.index_file}") |
|
|
| def load_hubert(self): |
| from fairseq import checkpoint_utils |
|
|
| models, _, _ = checkpoint_utils.load_model_ensemble_and_task( |
| ["weights/hubert_base.pt"], |
| suffix="", |
| ) |
| self.hubert_model = models[0] |
| self.hubert_model = self.hubert_model.to(self.config.device) |
|
|
| if self.config.is_half: |
| self.hubert_model = self.hubert_model.half() |
| else: |
| self.hubert_model = self.hubert_model.float() |
|
|
| return self.hubert_model.eval() |
|
|
