Update handler.py

handler.py

@@ -37,201 +37,196 @@ limitation = os.getenv("SYSTEM") == "spaces"
 
 config = Config()
 
-# Edge TTS
-tts_voice_list = edge_tts.list_voices()
-tts_voices = ["mn-MN-BataaNeural", "mn-MN-YesuiNeural"]  # Specific voices
+class EndpointHandler:
+    def __init__(self, model_dir=None):
+        self.model_dir = model_dir
+        self.hubert_model = self.load_hubert()
+        self.rmvpe_model = RMVPE("rmvpe.pt", config.is_half, config.device)
+        self.voice_mapping = {
+            "Mongolian Male": "mn-MN-BataaNeural",
+            "Mongolian Female": "mn-MN-YesuiNeural"
+        }
+        # Edge TTS
+        self.tts_voice_list = edge_tts.list_voices()
+        self.tts_voices = ["mn-MN-BataaNeural", "mn-MN-YesuiNeural"]  # Specific voices
 
-# RVC models
-model_root = "weights"
-models = [d for d in os.listdir(model_root) if os.path.isdir(f"{model_root}/{d}")]
-models.sort()
+        # RVC models
+        self.model_root = "weights"
+        self.models = [d for d in os.listdir(self.model_root) if os.path.isdir(f"{self.model_root}/{d}")]
+        self.models.sort()
 
-def get_unique_filename(extension):
-    return f"{uuid.uuid4()}.{extension}"
+    def get_unique_filename(self, extension):
+        return f"{uuid.uuid4()}.{extension}"
 
-def model_data(model_name):
-    # global n_spk, tgt_sr, net_g, vc, cpt, version, index_file
-    pth_path = [
-        f"{model_root}/{model_name}/{f}"
-        for f in os.listdir(f"{model_root}/{model_name}")
-        if f.endswith(".pth")
-    ][0]
-    print(f"Loading {pth_path}")
-    cpt = torch.load(pth_path, map_location="cpu")
-    tgt_sr = cpt["config"][-1]
-    cpt["config"][-3] = cpt["weight"]["emb_g.weight"].shape[0]  # n_spk
-    if_f0 = cpt.get("f0", 1)
-    version = cpt.get("version", "v1")
-    if version == "v1":
-        if if_f0 == 1:
-            net_g = SynthesizerTrnMs256NSFsid(*cpt["config"], is_half=config.is_half)
+    def model_data(self, model_name):
+        # global n_spk, tgt_sr, net_g, vc, cpt, version, index_file
+        pth_path = [
+            f"{self.model_root}/{model_name}/{f}"
+            for f in os.listdir(f"{self.model_root}/{model_name}")
+            if f.endswith(".pth")
+        ][0]
+        print(f"Loading {pth_path}")
+        cpt = torch.load(pth_path, map_location="cpu")
+        tgt_sr = cpt["config"][-1]
+        cpt["config"][-3] = cpt["weight"]["emb_g.weight"].shape[0]  # n_spk
+        if_f0 = cpt.get("f0", 1)
+        version = cpt.get("version", "v1")
+        if version == "v1":
+            if if_f0 == 1:
+                net_g = SynthesizerTrnMs256NSFsid(*cpt["config"], is_half=config.is_half)
+            else:
+                net_g = SynthesizerTrnMs256NSFsid_nono(*cpt["config"])
+        elif version == "v2":
+            if if_f0 == 1:
+                net_g = SynthesizerTrnMs768NSFsid(*cpt["config"], is_half=config.is_half)
+            else:
+                net_g = SynthesizerTrnMs768NSFsid_nono(*cpt["config"])
+        else:
+            raise ValueError("Unknown version")
+        del net_g.enc_q
+        net_g.load_state_dict(cpt["weight"], strict=False)
+        print("Model loaded")
+        net_g.eval().to(config.device)
+        if config.is_half:
+            net_g = net_g.half()
         else:
-            net_g = SynthesizerTrnMs256NSFsid_nono(*cpt["config"])
-    elif version == "v2":
-        if if_f0 == 1:
-            net_g = SynthesizerTrnMs768NSFsid(*cpt["config"], is_half=config.is_half)
+            net_g = net_g.float()
+        vc = VC(tgt_sr, config)
+        # n_spk = cpt["config"][-3]
+
+        index_files = [
+            f"{self.model_root}/{model_name}/{f}"
+            for f in os.listdir(f"{self.model_root}/{model_name}")
+            if f.endswith(".index")
+        ]
+        if len(index_files) == 0:
+            print("No index file found")
+            index_file = ""
         else:
-            net_g = SynthesizerTrnMs768NSFsid_nono(*cpt["config"])
-    else:
-        raise ValueError("Unknown version")
-    del net_g.enc_q
-    net_g.load_state_dict(cpt["weight"], strict=False)
-    print("Model loaded")
-    net_g.eval().to(config.device)
-    if config.is_half:
-        net_g = net_g.half()
-    else:
-        net_g = net_g.float()
-    vc = VC(tgt_sr, config)
-    # n_spk = cpt["config"][-3]
+            index_file = index_files[0]
+            print(f"Index file found: {index_file}")
 
-    index_files = [
-        f"{model_root}/{model_name}/{f}"
-        for f in os.listdir(f"{model_root}/{model_name}")
-        if f.endswith(".index")
-    ]
-    if len(index_files) == 0:
-        print("No index file found")
-        index_file = ""
-    else:
-        index_file = index_files[0]
-        print(f"Index file found: {index_file}")
+        return tgt_sr, net_g, vc, version, index_file, if_f0
 
-    return tgt_sr, net_g, vc, version, index_file, if_f0
+    def load_hubert(self):
+        # global hubert_model
+        models, _, _ = checkpoint_utils.load_model_ensemble_and_task(
+            ["hubert_base.pt"],
+            suffix="",
+        )
+        hubert_model = models[0]
+        hubert_model = hubert_model.to(config.device)
+        if config.is_half:
+            hubert_model = hubert_model.half()
+        else:
+            hubert_model = hubert_model.float()
+        return hubert_model.eval()
 
+    def get_model_names(self):
+        return [d for d in os.listdir(self.model_root) if os.path.isdir(f"{self.model_root}/{d}")]
 
-def load_hubert():
-    # global hubert_model
-    models, _, _ = checkpoint_utils.load_model_ensemble_and_task(
-        ["hubert_base.pt"],
-        suffix="",
-    )
-    hubert_model = models[0]
-    hubert_model = hubert_model.to(config.device)
-    if config.is_half:
-        hubert_model = hubert_model.half()
-    else:
-        hubert_model = hubert_model.float()
-    return hubert_model.eval()
+    def tts(self, model_name, tts_text, tts_voice, index_rate, use_uploaded_voice, uploaded_voice):
+        # Default values for parameters used in EdgeTTS
+        speed = 0  # Default speech speed
+        f0_up_key = 0  # Default pitch adjustment
+        f0_method = "rmvpe"  # Default pitch extraction method
+        protect = 0.33  # Default protect value
+        filter_radius = 3
+        resample_sr = 0
+        rms_mix_rate = 0.25
+        edge_time = 0  # Initialize edge_time
 
-def get_model_names():
-    model_root = "weights"  # Assuming this is where your models are stored
-    return [d for d in os.listdir(model_root) if os.path.isdir(f"{model_root}/{d}")]
+        edge_output_filename = self.get_unique_filename("mp3")
 
-def tts(model_name, tts_text, tts_voice, index_rate, use_uploaded_voice, uploaded_voice):
-    # Default values for parameters used in EdgeTTS
-    speed = 0  # Default speech speed
-    f0_up_key = 0  # Default pitch adjustment
-    f0_method = "rmvpe"  # Default pitch extraction method
-    protect = 0.33  # Default protect value
-    filter_radius = 3
-    resample_sr = 0
-    rms_mix_rate = 0.25
-    edge_time = 0  # Initialize edge_time
+        try:
+            if use_uploaded_voice:
+                if uploaded_voice is None:
+                    return "No voice file uploaded.", None, None
+                
+                # Process the uploaded voice file
+                with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_file:
+                    tmp_file.write(uploaded_voice)
+                    uploaded_file_path = tmp_file.name
 
-    edge_output_filename = get_unique_filename("mp3")
+                audio, sr = librosa.load(uploaded_file_path, sr=16000, mono=True)
+            else:
+                # EdgeTTS processing
+                if limitation and len(tts_text) > 4000:
+                    return (
+                        f"Text characters should be at most 280 in this huggingface space, but got {len(tts_text)} characters.",
+                        None,
+                        None,
+                    )
+                
+                # Invoke Edge TTS
+                t0 = time.time()
+                speed_str = f"+{speed}%" if speed >= 0 else f"{speed}%"
+                edge_tts.Communicate(tts_text, tts_voice, rate=speed_str).save(edge_output_filename)
+                t1 = time.time()
+                edge_time = t1 - t0
 
-    try:
-        if use_uploaded_voice:
-            if uploaded_voice is None:
-                return "No voice file uploaded.", None, None
-            
-            # Process the uploaded voice file
-            with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_file:
-                tmp_file.write(uploaded_voice)
-                uploaded_file_path = tmp_file.name
+                audio, sr = librosa.load(edge_output_filename, sr=16000, mono=True)
 
-            audio, sr = librosa.load(uploaded_file_path, sr=16000, mono=True)
-        else:
-            # EdgeTTS processing
-            if limitation and len(tts_text) > 4000:
+            # Common processing after loading the audio
+            duration = len(audio) / sr
+            print(f"Audio duration: {duration}s")
+            if limitation and duration >= 20:
                 return (
-                    f"Text characters should be at most 280 in this huggingface space, but got {len(tts_text)} characters.",
+                    f"Audio should be less than 20 seconds in this huggingface space, but got {duration}s.",
                     None,
                     None,
                 )
-            
-            # Invoke Edge TTS
-            t0 = time.time()
-            speed_str = f"+{speed}%" if speed >= 0 else f"{speed}%"
-            edge_tts.Communicate(tts_text, tts_voice, rate=speed_str).save(edge_output_filename)
-            t1 = time.time()
-            edge_time = t1 - t0
 
-            audio, sr = librosa.load(edge_output_filename, sr=16000, mono=True)
+            f0_up_key = int(f0_up_key)
+            tgt_sr, net_g, vc, version, index_file, if_f0 = self.model_data(model_name)
 
-        # Common processing after loading the audio
-        duration = len(audio) / sr
-        print(f"Audio duration: {duration}s")
-        if limitation and duration >= 20:
-            return (
-                f"Audio should be less than 20 seconds in this huggingface space, but got {duration}s.",
-                None,
+            # Setup for RMVPE or other pitch extraction methods
+            if f0_method == "rmvpe":
+                vc.model_rmvpe = self.rmvpe_model
+
+            # Perform voice conversion pipeline
+            times = [0, 0, 0]
+            audio_opt = vc.pipeline(
+                self.hubert_model,
+                net_g,
+                0,
+                audio,
+                edge_output_filename if not use_uploaded_voice else uploaded_file_path,
+                times,
+                f0_up_key,
+                f0_method,
+                index_file,
+                index_rate,
+                if_f0,
+                filter_radius,
+                tgt_sr,
+                resample_sr,
+                rms_mix_rate,
+                version,
+                protect,
                 None,
             )
 
-        f0_up_key = int(f0_up_key)
-        tgt_sr, net_g, vc, version, index_file, if_f0 = model_data(model_name)
-
-        # Setup for RMVPE or other pitch extraction methods
-        if f0_method == "rmvpe":
-            vc.model_rmvpe = rmvpe_model
-
-        # Perform voice conversion pipeline
-        times = [0, 0, 0]
-        audio_opt = vc.pipeline(
-            hubert_model,
-            net_g,
-            0,
-            audio,
-            edge_output_filename if not use_uploaded_voice else uploaded_file_path,
-            times,
-            f0_up_key,
-            f0_method,
-            index_file,
-            index_rate,
-            if_f0,
-            filter_radius,
-            tgt_sr,
-            resample_sr,
-            rms_mix_rate,
-            version,
-            protect,
-            None,
-        )
-
-        if tgt_sr != resample_sr and resample_sr >= 16000:
-            tgt_sr = resample_sr
-        
-        info = f"Success. Time: tts: {edge_time}s, npy: {times[0]}s, f0: {times[1]}s, infer: {times[2]}s"
-        print(info)
-        return (
-            info,
-            edge_output_filename if not use_uploaded_voice else None,   
-            (tgt_sr, audio_opt),
-            edge_output_filename  
-        )
-
-    except EOFError:
-        info = "Output not valid. This may occur when input text and speaker do not match."
-        print(info)
-        return info, None, None
-    except Exception as e:
-        traceback_info = traceback.format_exc()
-        print(traceback_info)
-        return str(e), None, None
-
-voice_mapping = {
-    "Mongolian Male": "mn-MN-BataaNeural",
-    "Mongolian Female": "mn-MN-YesuiNeural"
-}
-
-hubert_model = load_hubert()
-rmvpe_model = RMVPE("rmvpe.pt", config.is_half, config.device)
+            if tgt_sr != resample_sr and resample_sr >= 16000:
+                tgt_sr = resample_sr
+            
+            info = f"Success. Time: tts: {edge_time}s, npy: {times[0]}s, f0: {times[1]}s, infer: {times[2]}s"
+            print(info)
+            return (
+                info,
+                edge_output_filename if not use_uploaded_voice else None,   
+                (tgt_sr, audio_opt),
+                edge_output_filename  
+            )
 
-class EndpointHandler:
-    def __init__(self, model_dir=None):
-        self.model_dir = model_dir
+        except EOFError:
+            info = "Output not valid. This may occur when input text and speaker do not match."
+            print(info)
+            return info, None, None
+        except Exception as e:
+            traceback_info = traceback.format_exc()
+            print(traceback_info)
+            return str(e), None, None
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
         try:
@@ -265,11 +260,11 @@ class EndpointHandler:
             use_uploaded_voice = json_data["use_uploaded_voice"]
             voice_upload_file = json_data.get("voice_upload_file", None)
     
-            edge_tts_voice = voice_mapping.get(selected_voice)
+            edge_tts_voice = self.voice_mapping.get(selected_voice)
             if not edge_tts_voice:
                 raise ValueError(f"Invalid voice '{selected_voice}'.")
     
-            info, edge_tts_output_path, tts_output_data, edge_output_file = tts(
+            info, edge_tts_output_path, tts_output_data, edge_output_file = self.tts(
                 model_name,
                 tts_text,
                 edge_tts_voice,
@@ -299,6 +294,6 @@ class EndpointHandler:
             raise ValueError("Invalid JSON structure.")
     
     def save_audio_data_to_file(self, audio_data, sample_rate=40000):
-        file_path = get_unique_filename('wav')
+        file_path = self.get_unique_filename('wav')
         wavfile.write(file_path, sample_rate, audio_data)
         return file_path