From 8ee81dfae0cfd5dad6400f0ab00123e209c05a31 Mon Sep 17 00:00:00 2001
From: XXXXRT666 <runtongzhong@gmail.com>
Date: Tue, 6 Aug 2024 18:00:42 +0800
Subject: [PATCH] check for models

---
 GPT_SoVITS/feature_extractor/cnhubert.py      | 8 +++++---
 GPT_SoVITS/prepare_datasets/1-get-text.py     | 2 ++
 GPT_SoVITS/prepare_datasets/3-get-semantic.py | 2 ++
 3 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/GPT_SoVITS/feature_extractor/cnhubert.py b/GPT_SoVITS/feature_extractor/cnhubert.py
index dc155bdd..ca08f9f9 100644
--- a/GPT_SoVITS/feature_extractor/cnhubert.py
+++ b/GPT_SoVITS/feature_extractor/cnhubert.py
@@ -4,7 +4,7 @@ import librosa
 import torch
 import torch.nn.functional as F
 import soundfile as sf
-import logging
+import logging, os
 
 logging.getLogger("numba").setLevel(logging.WARNING)
 
@@ -22,9 +22,11 @@ cnhubert_base_path = None
 class CNHubert(nn.Module):
     def __init__(self):
         super().__init__()
-        self.model = HubertModel.from_pretrained(cnhubert_base_path)
+        if os.path.exists(cnhubert_base_path):...
+        else:raise FileNotFoundError(cnhubert_base_path)
+        self.model = HubertModel.from_pretrained(cnhubert_base_path, local_files_only=True)
         self.feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(
-            cnhubert_base_path
+            cnhubert_base_path, local_files_only=True
         )
 
     def forward(self, x):
diff --git a/GPT_SoVITS/prepare_datasets/1-get-text.py b/GPT_SoVITS/prepare_datasets/1-get-text.py
index 7e183103..a1049b31 100644
--- a/GPT_SoVITS/prepare_datasets/1-get-text.py
+++ b/GPT_SoVITS/prepare_datasets/1-get-text.py
@@ -54,6 +54,8 @@ if os.path.exists(txt_path) == False:
     #     device = "mps"
     else:
         device = "cpu"
+    if os.path.exists(bert_pretrained_dir):...
+    else:raise FileNotFoundError(bert_pretrained_dir)
     tokenizer = AutoTokenizer.from_pretrained(bert_pretrained_dir)
     bert_model = AutoModelForMaskedLM.from_pretrained(bert_pretrained_dir)
     if is_half == True:
diff --git a/GPT_SoVITS/prepare_datasets/3-get-semantic.py b/GPT_SoVITS/prepare_datasets/3-get-semantic.py
index 6adece61..842a4dce 100644
--- a/GPT_SoVITS/prepare_datasets/3-get-semantic.py
+++ b/GPT_SoVITS/prepare_datasets/3-get-semantic.py
@@ -34,6 +34,8 @@ logging.getLogger("numba").setLevel(logging.WARNING)
 # os.environ["CUDA_VISIBLE_DEVICES"]=sys.argv[5]
 # opt_dir="/data/docker/liujing04/gpt-vits/fine_tune_dataset/%s"%exp_name
 
+if os.path.exists(pretrained_s2G):...
+else:raise FileNotFoundError(pretrained_s2G)
 
 hubert_dir = "%s/4-cnhubert" % (opt_dir)
 semantic_path = "%s/6-name2semantic-%s.tsv" % (opt_dir, i_part)