Update handler.py

handler.py

@@ -1,79 +1,75 @@
-from typing import Dict, List, Any, Optional
-from voice_processing import tts, get_model_names, voice_mapping, get_unique_filename
+from pydantic import BaseModel
+from environs import Env
+from typing import List, Dict, Any
 import os
 import base64
 import numpy as np
+import librosa
 from scipy.io import wavfile
 import asyncio
-import time
+from voice_processing import tts, get_model_names, voice_mapping, get_unique_filename
 
 class EndpointHandler:
-    def __init__(self):
-        self.models = get_model_names()
-        self.voices = list(voice_mapping.keys())
-
-    async def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
-        model_name = data.get("model_name")
-        tts_text = data.get("tts_text", "Текстыг оруулна уу.")
-        selected_voice = data.get("selected_voice")
-        slang_rate = float(data.get("slang_rate", 0))
-        use_uploaded_voice = data.get("use_uploaded_voice", False)
-        voice_upload_file = data.get("voice_upload_file")
-
-        edge_tts_voice = 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 = await tts(
-            model_name, tts_text, edge_tts_voice, slang_rate, use_uploaded_voice, voice_upload_file
-        )
-
-        if edge_output_file and os.path.exists(edge_output_file):
-            os.remove(edge_output_file)
-
-        _, audio_output = tts_output_data
-
-        audio_file_path = self.save_audio_data_to_file(audio_output) if isinstance(audio_output, np.ndarray) else audio_output
-
-        audio_data_uri = self.get_audio_data_uri(audio_file_path)
-
-        if os.path.exists(audio_file_path):
-            os.remove(audio_file_path)
+    def __init__(self, model_dir=None):
+        self.model_dir = model_dir
 
-        return [{"info": info, "audio_data_uri": audio_data_uri}]
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        try:
+            if "inputs" in data:  # Check if data is in Hugging Face JSON format
+                return self.process_hf_input(data)
+            else:
+                return self.process_json_input(data)
+        except ValueError as e:
+            return {"error": str(e)}
+        except Exception as e:
+            return {"error": str(e)}
+
+    def process_json_input(self, json_data):
+        if all(key in json_data for key in ["model_name", "tts_text", "selected_voice", "slang_rate", "use_uploaded_voice"]):
+            model_name = json_data["model_name"]
+            tts_text = json_data["tts_text"]
+            selected_voice = json_data["selected_voice"]
+            slang_rate = json_data["slang_rate"]
+            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)
+            if not edge_tts_voice:
+                raise ValueError(f"Invalid voice '{selected_voice}'.")
+
+            info, edge_tts_output_path, tts_output_data, edge_output_file = asyncio.run(tts(
+                model_name, tts_text, edge_tts_voice, slang_rate, use_uploaded_voice, voice_upload_file
+            ))
+
+            if edge_output_file and os.path.exists(edge_output_file):
+                os.remove(edge_output_file)
+
+            _, audio_output = tts_output_data
+
+            audio_file_path = self.save_audio_data_to_file(audio_output) if isinstance(audio_output, np.ndarray) else audio_output
+
+            try:
+                with open(audio_file_path, 'rb') as file:
+                    audio_bytes = file.read()
+                audio_data_uri = f"data:audio/wav;base64,{base64.b64encode(audio_bytes).decode('utf-8')}"
+            except Exception as e:
+                raise Exception(f"Failed to read audio file: {e}")
+            finally:
+                if os.path.exists(audio_file_path):
+                    os.remove(audio_file_path)
+
+            return {"info": info, "audio_data_uri": audio_data_uri}
+        else:
+            raise ValueError("Invalid JSON structure.")
+
+    def process_hf_input(self, hf_data):
+        if "inputs" in hf_data:
+            actual_data = hf_data["inputs"]
+            return self.process_json_input(actual_data)
+        else:
+            return {"error": "Invalid Hugging Face JSON structure."}
 
     def save_audio_data_to_file(self, audio_data, sample_rate=40000):
         file_path = get_unique_filename('wav')
         wavfile.write(file_path, sample_rate, audio_data)
-        return file_path
-
-    def get_audio_data_uri(self, audio_file_path):
-        try:
-            with open(audio_file_path, 'rb') as file:
-                audio_bytes = file.read()
-            return f"data:audio/wav;base64,{base64.b64encode(audio_bytes).decode('utf-8')}"
-        except Exception as e:
-            raise ValueError(f"Failed to read audio file: {e}")
-
-async def periodic_cleanup(folder_path, interval_seconds, max_age_seconds):
-    while True:
-        now = time.time()
-        for filename in os.listdir(folder_path):
-            if filename.endswith(".mp3"):
-                file_path = os.path.join(folder_path, filename)
-                file_age = now - os.path.getmtime(file_path)
-                if file_age > max_age_seconds:
-                    os.remove(file_path)
-                    print(f"Deleted: {file_path}")
-        await asyncio.sleep(interval_seconds)
-
-async def start_periodic_cleanup():
-    folder_path = os.getcwd()  # Use current working directory
-    interval_seconds = 60 * 60 * 24  # Daily interval (24 hours)
-    max_age_seconds = 3600 * 3  # Maximum age of files (3 hours)
-    asyncio.create_task(periodic_cleanup(folder_path, interval_seconds, max_age_seconds))
-
-handler = EndpointHandler()
-
-if __name__ == "__main__":
-    asyncio.run(start_periodic_cleanup())
+        return file_path