fix: base dir

apps/audio_cloning/vallex/g2p/__init__.py

@@ -1,5 +1,6 @@
 """from https://github.com/keithito/tacotron"""
 
+import logging
 import os
 
 # import utils.g2p.cleaners
@@ -9,19 +10,29 @@ import apps.audio_cloning.vallex.g2p.cleaners as cleaners
 
 from .symbols import symbols
 
+logger = logging.getLogger(__name__)
+
 # Mappings from symbol to numeric ID and vice versa:
 _symbol_to_id = {s: i for i, s in enumerate(symbols)}
 _id_to_symbol = {i: s for i, s in enumerate(symbols)}
 
 
-BASE_DIR = os.getenv("HF_HOME", ".")
-TOKENIZER_PATH = os.path.join(BASE_DIR, "apps/audio_cloning/vallex/g2p/bpe_1024.json")
+PREPARED_BASE_DIR = "."
+TOKENIZER_PATH = os.path.join(
+    PREPARED_BASE_DIR, "apps/audio_cloning/vallex/g2p/bpe_1024.json"
+)
 
 
 class PhonemeBpeTokenizer:
     def __init__(self, tokenizer_path=TOKENIZER_PATH):
         print(f"Initializing PhonemeBpeTokenizer with tokenizer path: {tokenizer_path}")
-        self.tokenizer = Tokenizer.from_file(tokenizer_path)
+        try:
+            self.tokenizer = Tokenizer.from_file(tokenizer_path)
+        except Exception as e:
+            logger.error(
+                f"Error initializing PhonemeBpeTokenizer when reading file: {tokenizer_path}: {e}"
+            )
+            raise e
 
     def tokenize(self, text):
         # 1. convert text to phoneme

apps/audio_cloning/vallex/main.py

@@ -41,8 +41,10 @@ from .models.vallex import VALLE
 logger = logging.getLogger(__name__)
 
 # set base directory
-BASE_DIR = os.getenv("HF_HOME", ".")
-logger.info("Base directory: %s", BASE_DIR)
+OUTPUT_BASE_DIR = os.getenv("HF_HOME", ".")
+PREPARED_BASE_DIR = "."
+logger.info("Base directory: %s", OUTPUT_BASE_DIR)
+logger.info("Prepared base directory: %s", PREPARED_BASE_DIR)
 
 # set languages
 langid.set_languages(["en", "zh", "ja"])
@@ -90,7 +92,9 @@ else:
 
 # set text tokenizer and collater
 logger.info("Setting text tokenizer and collater...")
-tokenizer_path = os.path.join(BASE_DIR, "apps/audio_cloning/vallex/g2p/bpe_69.json")
+tokenizer_path = os.path.join(
+    PREPARED_BASE_DIR, "apps/audio_cloning/vallex/g2p/bpe_69.json"
+)
 text_tokenizer = PhonemeBpeTokenizer(tokenizer_path=tokenizer_path)
 text_collater = get_text_token_collater()
 
@@ -104,7 +108,7 @@ if torch.cuda.is_available():
 logger.info("Device set to %s", device)
 
 # Download VALL-E-X model weights if not exists
-OUTPUT_DIR_CHECKPOINTS = os.path.join(BASE_DIR, "models/checkpoints")
+OUTPUT_DIR_CHECKPOINTS = os.path.join(OUTPUT_BASE_DIR, "models/checkpoints")
 OUTPUT_FILENAME_CHECKPOINTS = "vallex-checkpoint.pt"
 OUTPUT_PATH_CHECKPOINTS = os.path.join(
     OUTPUT_DIR_CHECKPOINTS, OUTPUT_FILENAME_CHECKPOINTS
@@ -142,7 +146,9 @@ model = VALLE(
     prepend_bos=True,
     num_quantizers=NUM_QUANTIZERS,
 )
-checkpoint = torch.load(OUTPUT_PATH_CHECKPOINTS, map_location="cpu", weights_only=False)
+checkpoint = torch.load(
+    OUTPUT_PATH_CHECKPOINTS, map_location=device, weights_only=False
+)
 missing_keys, unexpected_keys = model.load_state_dict(checkpoint["model"], strict=True)
 assert not missing_keys
 model.eval()
@@ -155,7 +161,7 @@ audio_tokenizer = AudioTokenizer(device)
 vocos = Vocos.from_pretrained("charactr/vocos-encodec-24khz").to(device)
 
 # initialize ASR model
-OUTPUT_DIR_WHISPER = os.path.join(BASE_DIR, "models/whisper")
+OUTPUT_DIR_WHISPER = os.path.join(OUTPUT_BASE_DIR, "models/whisper")
 if not os.path.exists(OUTPUT_DIR_WHISPER):
     os.makedirs(OUTPUT_DIR_WHISPER, exist_ok=True)
 
@@ -176,7 +182,7 @@ except Exception as e:
 
 # Initialize Voice Presets
 logger.info("Initializing Voice Presets...")
-PRESETS_DIR = os.path.join(BASE_DIR, "apps/audio_cloning/vallex/presets")
+PRESETS_DIR = os.path.join(PREPARED_BASE_DIR, "apps/audio_cloning/vallex/presets")
 preset_list = os.walk(PRESETS_DIR).__next__()[2]
 preset_list = [preset[:-4] for preset in preset_list if preset.endswith(".npz")]