Inference works again with abs pathes

tts/VALL-E-X/main.py

@@ -2,11 +2,8 @@ import argparse
 
 # python3 main.py -create_prompt_with_list 20.wav 21.wav -create_prompt_name alice
 # python3 main.py -use_prompt_name alice -output_file_name alice_says.wav -text "This is the text to be converted to speech."
-
 # python3 main.py -create_prompt_with_list 20.wav 21.wav 20.wav 21.wav -create_prompt_name alice -use_prompt_name alice -output_file_name alice_says.wav -text "This is my voice."
-
 # python3 main.py -create_prompt_with_list goku.wav -create_prompt_name goku -use_prompt_name goku -output_file_name goku_test.wav -text "oh my god krillin! where did you put my sensu beans!?"
-
 # python3 main.py -create_prompt_with_list
 
 from typing import List
@@ -16,203 +13,186 @@ import pathlib
 
 from utils.prompt_making import make_prompt
 import logging
-
+from utils.generation import SAMPLE_RATE, generate_audio,load_models_from_pathes
-from utils.generation import SAMPLE_RATE, generate_audio, preload_models,load_models_from_mount
 import soundfile as sf
 import numpy as np
 import torchaudio
 from typing import List
 import uuid
 import torchaudio.transforms as T
+class AudioOps():
     
-# class AudioOps():
+    @staticmethod
+    def stitch_wav_tensors_with_crossfade(wav_tensor1:np.array,wav_tensor2:np.array,sample_rate:int):
+        # Load two audio signals  as NumPy tensors
+        # Compute the length of the cross-fade window in samples
+        fade_len = int(0.01 * sample_rate)
 
-#     @staticmethod
+        # Create a cross-fade window
-#     def stitch_wav_tensors_with_crossfade(wav_tensor1:np.array,wav_tensor2:np.array,sample_rate:int):
+        fade_window = np.hanning(2*fade_len)
-#         # Load two audio signals  as NumPy tensors
-#         # Compute the length of the cross-fade window in samples
-#         fade_len = int(0.01 * sample_rate)
 
-#         # Create a cross-fade window
+        # Concatenate the signals with cross-fade
-#         fade_window = np.hanning(2*fade_len)
+        audio_out = np.concatenate((wav_tensor1[:-fade_len], 
+                                    wav_tensor1[-fade_len:] * fade_window[:fade_len] + wav_tensor2[:fade_len] * fade_window[fade_len:], 
+                                    wav_tensor2[fade_len:]))
+        return audio_out
     
-#         # Concatenate the signals with cross-fade
+    @staticmethod
-#         audio_out = np.concatenate((wav_tensor1[:-fade_len], 
+    def cut_audio_range(file_path, start_seconds, end_seconds, output_path="cut_audio.wav"):
-#                                     wav_tensor1[-fade_len:] * fade_window[:fade_len] + wav_tensor2[:fade_len] * fade_window[fade_len:], 
+        """
-#                                     wav_tensor2[fade_len:]))
+        Cuts out a specific segment from an audio file and saves the result.
-#         return audio_out
 
-#     @staticmethod
+        Args:
-#     def cut_audio_range(file_path, start_seconds, end_seconds, output_path="cut_audio.wav"):
+        - file_path (str): Path to the audio file from which the segment will be cut.
-#         """
+        - start_seconds (float): Start time of the segment to be cut out in seconds.
-#         Cuts out a specific segment from an audio file and saves the result.
+        - end_seconds (float): End time of the segment to be cut out in seconds.
+        - output_path (str): Path to save the resulting audio. Defaults to "cut_audio.wav".
         
-#         Args:
+        Returns:
-#         - file_path (str): Path to the audio file from which the segment will be cut.
+        - None
-#         - start_seconds (float): Start time of the segment to be cut out in seconds.
+        """
-#         - end_seconds (float): End time of the segment to be cut out in seconds.
+        # Load the audio file
-#         - output_path (str): Path to save the resulting audio. Defaults to "cut_audio.wav".
+        waveform, sr = torchaudio.load(file_path)
+        # Convert the range to sample indices
+        start_sample = int(start_seconds * sr)
+        end_sample = int(end_seconds * sr)
+        cut_tensor = waveform[:,start_sample:end_sample]
+        torchaudio.save(output_path, cut_tensor, sr)
         
-#         Returns:
+    @staticmethod
-#         - None
+    def stitch_wav_files_resample(files:List[str],target_sample_rate:int=SAMPLE_RATE):
-#         """
+        resampler = None 
-#         # Load the audio file
+        result = None
-#         waveform, sr = torchaudio.load(file_path)
+        for file in files:  
-#         # Convert the range to sample indices
+            if type(result) == type(None):
-#         start_sample = int(start_seconds * sr)
+                tensor, sr = torchaudio.load(pathlib.Path(file))
-#         end_sample = int(end_seconds * sr)
+                resampler = T.Resample(orig_freq=sr,new_freq=target_sample_rate) 
-#         cut_tensor = waveform[:,start_sample:end_sample]
+                resampled_tensor = resampler(tensor)
-#         torchaudio.save(output_path, cut_tensor, sr)
+                result = resampled_tensor.squeeze()
+                result = np.array(result)
+            else:  
+                tensor, _ = torchaudio.load(pathlib.Path(file))
+                resampled_tensor = resampler(tensor)
+                resampled_tensor = resampled_tensor.squeeze()
+                resampled_tensor = np.array(resampled_tensor)
+                result = AudioOps.stitch_wav_tensors_with_crossfade(result,resampled_tensor,sample_rate=target_sample_rate)
                 
-#     @staticmethod
+        return result
-#     def stitch_wav_files_resample(root_path:pathlib.Path, files:List[str],target_sample_rate:int=SAMPLE_RATE):
-#         resampler = None 
-#         result = None
-#         for file in files:  
-#             if type(result) == type(None):
-#                 tensor, sr = torchaudio.load(root_path / pathlib.Path(file))
-#                 resampler = T.Resample(orig_freq=sr,new_freq=target_sample_rate) 
-#                 resampled_tensor = resampler(tensor)
-#                 result = resampled_tensor.squeeze()
-#                 result = np.array(result)
-#             else:  
-#                 tensor, _ = torchaudio.load(root_path / pathlib.Path(file))
-#                 resampled_tensor = resampler(tensor)
-#                 resampled_tensor = resampled_tensor.squeeze()
-#                 resampled_tensor = np.array(resampled_tensor)
-#                 result = AudioOps.stitch_wav_tensors_with_crossfade(result,resampled_tensor,sample_rate=target_sample_rate)
     
-#         return result
+class VoiceDesigner():
+    def __init__(self) -> None:
+        # load models from host system.
+        self.logger = self.setup_logger()
+        #  host path for work
+        self.temp_path = None
         
-# class VoiceDesigner():
+    def tts_with_prompt(self, prompt_dir:pathlib.Path,
-#     def __init__(self) -> None:
+                        prompt_name:pathlib.Path,
-#         # load models from host system.
+                        audio_output_path:pathlib.Path, 
-#         self.logger = self.setup_logger()
+                        output_file_name:pathlib.Path,
+                        text:str):
         
-#         self.mount_folder = pathlib.Path.cwd() / pathlib.Path("Vall-E-mount") 
+        audio_array = generate_audio(text,prompt_name=prompt_name,prompt_dir=prompt_dir)
-#         # Path for the inputs bunch of audio files
+        # save audio to disk
-#         self.input_path = self.mount_folder / pathlib.Path("input")
+        sf.write(audio_output_path / output_file_name, audio_array, SAMPLE_RATE)
-#         # Path for the models to be used
-#         self.models_path = self.mount_folder / pathlib.Path("models") 
-#         # Path for the existing presets
-#         self.presets_path = self.mount_folder / pathlib.Path("presets")
-#         # Path for the new prompts created by the user
-#         self.prompts_path = self.mount_folder / pathlib.Path("prompts")
-#         # Path for the temporary files
-#         self.temp_path = self.mount_folder / pathlib.Path("temp")
-#         # Path for output
-#         self.output_path = self.mount_folder / pathlib.Path("output")
-#         # Will have to load whisper from here as well ... todo fix
-#         load_models_from_mount(model_path=self.models_path,
-#                            vocos_path=pathlib.Path("vocos-encodec-24khz"),
-#                            vall_e_x=pathlib.Path("vallex-checkpoint.pt"))
         
-#     def tts_with_prompt(self,prompt:pathlib.Path,text:str,output_file_name:str):
+    def initialize_voice_designer(self,vall_e_x_file_path:pathlib.Path,vocos_folder_path:pathlib.Path,whisper_folder_path:pathlib.Path):
-#         # TODO Filtering bad words
+        load_models_from_pathes(vocos_folder_path=vocos_folder_path,vall_e_x_file_path=vall_e_x_file_path,whisper_folder_path=whisper_folder_path)
-#         # TODO check if the files are actually wav ...
-#         audio_array = generate_audio(text,prompt=prompt)
-#         # save audio to disk
-#         sf.write(self.output_path / pathlib.Path(output_file_name), audio_array, SAMPLE_RATE)
         
-#     def create_prompt(self,audio_file_paths:List[pathlib.Path],prompt_file_name:str):
+    def create_prompt(self,audio_file_paths:List[pathlib.Path],prompt_file_name:str,prompt_output_path:str):
-#         # multiple audio files concat
+        # multiple audio files concat
-#         # concat audio and trim the audio if it is longer than 14-15 seconds# if prompt is too long it will error out anyways
+        # concat audio and trim the audio if it is longer than 14-15 seconds# if prompt is too long it will error out anyways
-#         if len(audio_file_paths) > 0:
+        if len(audio_file_paths) > 0:
-#             # TODO: going to assume brandon renames the files as they are uploaded
+            # TODO: going to assume brandon renames the files as they are uploaded
-#             try:# ./customs/ <-- folder it adds by default until we modified.
+            try:# ./customs/ <-- folder it adds by default until we modified.
-#                     # load the audio files and then concat them. 
+                    # load the audio files and then concat them. 
-#                 tensor_result = AudioOps.stitch_wav_files_resample(root_path=self.input_path,files=audio_file_paths)
+                tensor_result = AudioOps.stitch_wav_files_resample(files=audio_file_paths)
-#                 # save to disk
+                # save to disk
-#                 generated_name = f"{uuid.uuid4()}"
+                generated_name = f"{uuid.uuid4()}"
                 
-#                 concat_audio_file_path = self.temp_path / pathlib.Path(f"audio_concat_{generated_name}.wav")               
+                concat_audio_file_path = self.temp_path / pathlib.Path(f"audio_concat_{generated_name}.wav")               
-#                 sf.write(concat_audio_file_path, tensor_result, SAMPLE_RATE)
+                sf.write(concat_audio_file_path, tensor_result, SAMPLE_RATE)
-#                 # save result into processing...
+                # save result into processing...
-#                 # trim to 15 seconds.
+                # trim to 15 seconds.
-#                 audio_cut_path = self.temp_path / pathlib.Path(f"audio_cut_{generated_name}.wav")
+                audio_cut_path = self.temp_path / pathlib.Path(f"audio_cut_{generated_name}.wav")
-#                 AudioOps.cut_audio_range(file_path=concat_audio_file_path,start_seconds=0,end_seconds=15,output_path=audio_cut_path)
+                AudioOps.cut_audio_range(file_path=concat_audio_file_path,start_seconds=0,end_seconds=15,output_path=audio_cut_path)
-#                 # TODO get whisper x to work in the container and generate transcription that way or use a different container to do that work.
+                make_prompt(name=prompt_file_name,audio_prompt_output_path=prompt_output_path,audio_prompt_path=audio_cut_path)
-#                 make_prompt(name=prompt_file_name,audio_prompt_output_path=self.prompts_path,audio_prompt_path=audio_cut_path,transcript=None)
-#             except Exception as e:
-#                 self.logger.debug(f"Exception: {e.with_traceback()}")
-#                 raise
-#         else:
-#             self.logger.debug(f"No audio files to load")
                 
-#     def setup_logger(self):
+            except Exception as e:
-#         # Create a logger object
+                self.logger.debug(f"Exception: {e.with_traceback()}")
-#         logger = logging.getLogger(__name__)
+                raise
+        else:
+            self.logger.debug(f"No audio files to load")
             
-#         # Set the log level
+    def setup_logger(self):
-#         logger.setLevel(logging.DEBUG)
+        # Create a logger object
+        logger = logging.getLogger(__name__)
 
-#         # Create a file handler to log messages to a file
+        # Set the log level
-#         file_handler = logging.FileHandler('docker_valle-x.log')
+        logger.setLevel(logging.DEBUG)
-#         file_handler.setLevel(logging.DEBUG)
 
-#         # Create a console handler to log messages to the console
+        # Create a file handler to log messages to a file
-#         console_handler = logging.StreamHandler()
+        file_handler = logging.FileHandler('docker_valle-x.log')
-#         console_handler.setLevel(logging.DEBUG)
+        file_handler.setLevel(logging.DEBUG)
 
-#         # Create a formatter
+        # Create a console handler to log messages to the console
-#         formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+        console_handler = logging.StreamHandler()
+        console_handler.setLevel(logging.DEBUG)
 
-#         # Set the formatter for the handlers
+        # Create a formatter
-#         file_handler.setFormatter(formatter)
+        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
-#         console_handler.setFormatter(formatter)
         
-#         # Add the handlers to the logger
+        # Set the formatter for the handlers
-#         logger.addHandler(file_handler)
+        file_handler.setFormatter(formatter)
-#         logger.addHandler(console_handler)
+        console_handler.setFormatter(formatter)
 
-#         return logger
+        # Add the handlers to the logger
+        logger.addHandler(file_handler)
+        logger.addHandler(console_handler)
 
-#global voice_designer 
+        return logger
-#voice_designer = VoiceDesigner()
+
+global voice_designer 
+voice_designer = VoiceDesigner()
 
 def main(args):
-    # print(args.create_prompt_with_list)
-    # print(args.create_prompt_name)
-    
-    # print(args.text)
-    # print(args.use_prompt_name)
-    # print(args.output_file_name)
-    
-    # if args.create_prompt_name is not None and args.create_prompt_name is not None:
-    #     voice_designer.create_prompt(audio_file_paths=args.create_prompt_with_list,
-    #                                  prompt_file_name=args.create_prompt_name)
-    
-    # if args.text is not None and args.use_prompt_name is not None:
-    #     voice_designer.tts_with_prompt(prompt=args.use_prompt_name,
-    #                                    text=args.text,
-    #                                    output_file_name=args.output_file_name)    
-    
     
     print(args.mode)
     
     print(args.text)
     print(args.audio_wav_files)
-    print(args.whisper_file_path)
+    print(args.whisper_folder_path)
     print(args.vocos_folder_path)
     
     print(args.vallex_path)
     
     print(args.prompt_path)
-    print(args.audio_name)
     print(args.prompt_name)
     
+    print(args.audio_name)
+    print(args.audio_path)
+    
+    print(args.tmp_work_dir)
+    
+    voice_designer.temp_path = pathlib.Path(args.tmp_work_dir)
+
+    voice_designer.initialize_voice_designer(pathlib.Path(args.vallex_path),
+                                             pathlib.Path(args.vocos_folder_path),
+                                             pathlib.Path(args.whisper_folder_path))
+    
+    if args.mode == 0: # run inference
+        voice_designer.tts_with_prompt(prompt_dir=pathlib.Path(args.prompt_path),
+                                       prompt_name=pathlib.Path(args.prompt_name),
+                                       audio_output_path=pathlib.Path(args.audio_path),
+                                       output_file_name=pathlib.Path(args.audio_name),
+                                       text=args.text)
+    
+    elif args.mode == 1: # run create voice
+        voice_designer.create_prompt(audio_file_paths=args.audio_wav_files,
+                                     prompt_file_name=pathlib.Path(args.prompt_name))
+    
 if __name__ == "__main__":
-    
     parser = argparse.ArgumentParser(description="Process a list of speaker samples for zero shot - .wav file pathes, convert text to speech, and specify language.")
-
-    # parser.add_argument('-create_prompt_with_list', metavar='prompt wav files', nargs='+', 
-    #                     help='Takes in a list of .wav files for an audio prompt provide the absolute path. Must be less than 15 seconds otherwise will error out.')
-    # parser.add_argument('-create_prompt_name',metavar="prompt name to create",
-    #                     help='required with -create-prompt prompt output name this will be written to the mount. In Vall-E-mount/prompts')
-    # parser.add_argument("-use_prompt_name",metavar="prompt you want to use",type=str)
-    # parser.add_argument("-output_file_name",metavar="file name",type=str)
-    # parser.add_argument('-text', metavar='text to synthesize', type=str,help='Text to be converted to speech.')
-    
     # modes
     parser.add_argument('--mode', metavar='model mode',type=int,help='Mode 0 - (Inference) | Mode 1 - (Create)')
     parser.add_argument('--text', metavar='text to synthesize', type=str,help='Text to be converted to speech.')
@@ -220,7 +200,8 @@ if __name__ == "__main__":
     parser.add_argument('--audio-wav-files', metavar='wav to voice clone', type=str,nargs='+',help='Takes in a list of .wav files for an audio prompt provide the absolute path. Must be less than 15 seconds otherwise will error out. (Create Voice Only)')           
     
     # model pathes 
-    parser.add_argument("--whisper-file-path",metavar="abs whisper model path",type=str,help="Path to the Whisper Model .pt file")
+    parser.add_argument("--whisper-folder-path",metavar="abs whisper folder path",type=str,help="Path to the Whisper folder where the .pt file")
+    parser.add_argument("--whisper-model",metavar="model large or medium",type=str,help="medium or large")
     parser.add_argument("--vocos-folder-path",metavar="abs vocos model path",type=str,help="Path to vocos codec Model requires the (folder path)")
     parser.add_argument("--vallex-path",metavar="abs vall-e model path",type=str,help="Path to the Valle Model File")
     
@@ -230,13 +211,17 @@ if __name__ == "__main__":
     parser.add_argument("--prompt-path",metavar="abs path to prompt",type=str,help="Prompt Embedding Path (Inference and Create Voice)")
     parser.add_argument("--prompt-name",metavar="abs path to prompt",type=str,help="Prompt Name (Inference and Create Voice)")
     
+    parser.add_argument("--tmp-work-dir",metavar="abs tmp dir",type=str,help="Temporary work directory in host system.")
+    
     args = parser.parse_args()
 
     main(args)
 
 # all params command
+
 #/home/tensor/code/storyteller/storyteller-ml/tts/VALL-E-X/.venv/bin/python
 #/home/tensor/code/storyteller/storyteller-ml/tts/VALL-E-X/main.py --text "hello world" 
+
 # --mode 0 
 # --audio-wav-files "/home/tensor/code/TTSDockerContainer/Vall-E-mount/input/20.wav" "/home/tensor/code/TTSDockerContainer/Vall-E-mount/input/21.wav" 
 # --whisper-file-path "/home/tensor/code/storyteller/storyteller-ml/tts/VALL-E-X/whisper/medium.pt" 
@@ -246,6 +231,12 @@ if __name__ == "__main__":
 # --prompt-name "test_prompt.npz" 
 # --audio-path "/home/tensor/code/storyteller/VALL-E-X-TTS-Container/Vall-E-mount/output/"
 # --audio-name "test.wav"
+# --tmp-work-dir "/tmp"
 
 # all params command copy paste
-# /home/tensor/code/storyteller/storyteller-ml/tts/VALL-E-X/.venv/bin/python /home/tensor/code/storyteller/storyteller-ml/tts/VALL-E-X/main.py --text "hello world" --mode 0 --audio-wav-files "/home/tensor/code/TTSDockerContainer/Vall-E-mount/input/20.wav" "/home/tensor/code/TTSDockerContainer/Vall-E-mount/input/21.wav" --whisper-file-path "/home/tensor/code/storyteller/storyteller-ml/tts/VALL-E-X/whisper/medium.pt" --vocos-folder-path "/home/tensor/code/TTSDockerContainer/Vall-E-mount/models/vocos-encodec-24khz" --vallex-path "/home/tensor/code/TTSDockerContainer/Vall-E-mount/models/vallex-checkpoint.pt" --prompt-path "/home/tensor/code/storyteller/VALL-E-X-TTS-Container/Vall-E-mount/prompts" --prompt-name "test_prompt.npz" --audio-path "/home/tensor/code/storyteller/VALL-E-X-TTS-Container/Vall-E-mount/output/" --audio-name "test.wav"
+# /home/tensor/code/storyteller/storyteller-ml/tts/VALL-E-X/.venv/bin/python /home/tensor/code/storyteller/storyteller-ml/tts/VALL-E-X/main.py --text "hello world" --mode 0 --audio-wav-files "/home/tensor/code/TTSDockerContainer/Vall-E-mount/input/20.wav" "/home/tensor/code/TTSDockerContainer/Vall-E-mount/input/21.wav" --whisper-file-path "/home/tensor/code/storyteller/storyteller-ml/tts/VALL-E-X/whisper/medium.pt" --vocos-folder-path "/home/tensor/code/TTSDockerContainer/Vall-E-mount/models/vocos-encodec-24khz" --vallex-path "/home/tensor/code/TTSDockerContainer/Vall-E-mount/models/vallex-checkpoint.pt" --prompt-path "/home/tensor/code/storyteller/VALL-E-X-TTS-Container/Vall-E-mount/prompts" --prompt-name "test_prompt.npz" --audio-path "/home/tensor/code/storyteller/VALL-E-X-TTS-Container/Vall-E-mount/output/" --audio-name "test.wav" --tmp-work-dir "/tmp"
+
+# inference example
+
+#/home/tensor/code/storyteller/storyteller-ml/tts/VALL-E-X/.venv/bin/python /home/tensor/code/storyteller/storyteller-ml/tts/VALL-E-X/main.py --text "hello world" --mode 0 --audio-wav-files "/home/tensor/code/TTSDockerContainer/Vall-E-mount/input/20.wav" "/home/tensor/code/TTSDockerContainer/Vall-E-mount/input/21.wav" --whisper-folder-path "/home/tensor/code/storyteller/storyteller-ml/tts/VALL-E-X/whisper" --whisper-model medium --vocos-folder-path "/home/tensor/code/TTSDockerContainer/Vall-E-mount/models/vocos-encodec-24khz" --vallex-path "/home/tensor/code/TTSDockerContainer/Vall-E-mount/models/vallex-checkpoint.pt" --prompt-path "/home/tensor/code/storyteller/VALL-E-X-TTS-Container/Vall-E-mount/prompts" --prompt-name "goku.npz" --audio-path "/home/tensor/code/storyteller/VALL-E-X-TTS-Container/Vall-E-mount/output/" --audio-name "hello_world_test.wav" --tmp-work-dir "/tmp"
+

tts/VALL-E-X/utils/generation.py

@@ -4,6 +4,8 @@ import torch
 from vocos import Vocos
 from vocos.feature_extractors import FeatureExtractor, EncodecFeatures
 
+import whisper
+
 import logging
 import langid
 langid.set_languages(['en', 'zh', 'ja'])
@@ -40,14 +42,57 @@ checkpoints_dir = "./checkpoints/"
 model_checkpoint_name = "vallex-checkpoint.pt"
 
 model = None
-
 codec = None
-
 vocos = None
+whisper_model = None
 
 text_tokenizer = PhonemeBpeTokenizer(tokenizer_path="./utils/g2p/bpe_69.json")
 text_collater = get_text_token_collater()
 
+def load_models_from_pathes(vocos_folder_path:pathlib.Path,vall_e_x_file_path:pathlib.Path,whisper_folder_path:pathlib.Path):
+    global model, codec, vocos, whisper_model
+        
+    model = VALLE(
+        N_DIM,
+        NUM_HEAD,
+        NUM_LAYERS,
+        norm_first=True,
+        add_prenet=False,
+        prefix_mode=PREFIX_MODE,
+        share_embedding=True,
+        nar_scale_factor=1.0,
+        prepend_bos=True,
+        num_quantizers=NUM_QUANTIZERS,
+    ).to(device)
+    checkpoint = torch.load(pathlib.Path(vall_e_x_file_path), map_location='cpu')
+    missing_keys, unexpected_keys = model.load_state_dict(
+        checkpoint["model"], strict=True
+    )
+    
+    assert not missing_keys
+    model.eval()
+    # Encodec
+    codec = AudioTokenizer(device)
+
+    vocos = None 
+    vocos_model = Vocos.from_hparams(vocos_folder_path / pathlib.Path("config.yaml")) 
+    state_dict = torch.load(vocos_folder_path / pathlib.Path("pytorch_model.bin"), map_location="cpu")
+    if isinstance(vocos_model.feature_extractor, EncodecFeatures):
+        encodec_parameters = {
+            "feature_extractor.encodec." + key: value
+            for key, value in vocos_model.feature_extractor.encodec.state_dict().items()
+        }
+        state_dict.update(encodec_parameters)
+        
+    vocos_model.load_state_dict(state_dict)
+    vocos_model.eval()
+    vocos = vocos_model
+    vocos.to(device)
+    # TODO replace
+    whisper_model_name = "medium"
+    whisper_model = whisper.load_model(whisper_model_name,download_root=whisper_folder_path).cpu()
+
+# DEPRECATED
 def load_models_from_mount(model_path:pathlib.Path,vocos_path:pathlib.Path,vall_e_x:pathlib.Path):
     global model, codec, vocos
         
@@ -131,13 +176,11 @@ def preload_models():
     vocos = Vocos.from_pretrained('charactr/vocos-encodec-24khz').to(device)
 
 @torch.no_grad()
-def generate_audio(text, 
+def generate_audio(text:str, 
-                   prompt=None, 
+                   prompt_name:pathlib.Path=None, 
+                   prompt_dir:pathlib.Path=None,
                    language='auto', 
-                   accent='no-accent', 
+                   accent='no-accent'):
-                   mount_dir:pathlib.Path=pathlib.Path.cwd() / pathlib.Path("..") / pathlib.Path("Vall-E-mount"),
-                   presets_dir:pathlib.Path=pathlib.Path("presets"),
-                   prompt_dir:pathlib.Path=pathlib.Path("prompts")):
     
     global model, codec, vocos, text_tokenizer, text_collater
     text = text.replace("\n", "").strip(" ")
@@ -149,22 +192,12 @@ def generate_audio(text,
     text = lang_token + text + lang_token
 
     # load prompt
-    if prompt is not None:
+    if prompt_name is not None:
-        prompt_path = prompt
+        prompt_path = prompt_name
-        # if not os.path.exists(prompt_path):
+        path = prompt_dir / pathlib.Path(f"{prompt_name}")
-        #     prompt_path = "./presets/" + prompt + ".npz"
+        print(f"Loading Prompt From Path: {path}")
-        # if not os.path.exists(prompt_path):
+        if pathlib.Path.exists(path):
-        #     prompt_path = "./customs/" + prompt + ".npz"
+            prompt_path = path
-        # check preset dir
-        # mounted dirs
-        if pathlib.Path.exists(mount_dir / presets_dir / pathlib.Path(f"{prompt}.npz")):
-            prompt_path = mount_dir / presets_dir / pathlib.Path(f"{prompt}.npz")
-            
-        if pathlib.Path.exists(mount_dir / prompt_dir / pathlib.Path(f"{prompt}.npz")):
-            prompt_path = mount_dir / prompt_dir / pathlib.Path(f"{prompt}.npz")
-            
-        if not os.path.exists(prompt_path):
-            raise ValueError(f"Cannot find prompt {prompt}")
            
         prompt_data = np.load(prompt_path)
         audio_prompts = prompt_data['audio_tokens']

tts/VALL-E-X/utils/prompt_making.py

@@ -27,7 +27,7 @@ if torch.cuda.is_available():
 codec = AudioTokenizer(device)
 
 if not os.path.exists("./whisper/"): os.mkdir("./whisper/")
-whisper_model = None
+# whisper_model = None
 
 @torch.no_grad()
 def transcribe_one(model, audio_path):
@@ -54,7 +54,7 @@ def transcribe_one(model, audio_path):
         text_pr += "."
     return lang, text_pr
 
-def make_prompt(name, audio_prompt_path,audio_prompt_output_path:pathlib.Path=None, transcript=None):
+def make_prompt(name, audio_prompt_path,audio_prompt_output_path:pathlib.Path=None):
     global model, text_collater, text_tokenizer, codec
     wav_pr, sr = torchaudio.load(audio_prompt_path)
     # check length
@@ -62,7 +62,7 @@ def make_prompt(name, audio_prompt_path,audio_prompt_output_path:pathlib.Path=No
         raise ValueError(f"Prompt too long, expect length below 15 seconds, got {wav_pr / sr} seconds.")
     if wav_pr.size(0) == 2:
         wav_pr = wav_pr.mean(0, keepdim=True)
-    text_pr, lang_pr = make_transcript(name, wav_pr, sr, transcript)
+    text_pr, lang_pr = make_transcript(name, wav_pr, sr, None)
 
     # tokenize audio
     encoded_frames = tokenize_audio(codec, (wav_pr, sr))