Update app.py

app.py

@@ -1,52 +1,90 @@
 import gradio as gr
+import torch
+import os
 from TTS.api import TTS
 from huggingface_hub import hf_hub_download
-import os
 
 # --- IMPORTING YOUR SEPARATE ROMANIZER ---
-from romanizer import sinhala_to_roman
+# Ensure romanizer.py is in the same directory
+try:
+    from romanizer import sinhala_to_roman
+except ImportError:
+    print("Warning: romanizer.py not found. Sinhala might not process correctly.")
+    def sinhala_to_roman(text): return text
+
+# --- MODEL LOADING LOGIC ---
 
-def load_my_model(repo_id):
+def load_standard_model(repo_id):
     print(f"Downloading {repo_id}...")
     model_path = hf_hub_download(repo_id=repo_id, filename="best_model.pth")
     config_path = hf_hub_download(repo_id=repo_id, filename="config.json")
-    # Initialize TTS
     return TTS(model_path=model_path, config_path=config_path, gpu=False)
 
-# Load Models
-print("Initializing Sinhala...")
-sin_tts = load_my_model("E-motionAssistant/text-to-speech-VITS-sinhala")
+def load_eng_model_with_surgery():
+    repo_id = "E-motionAssistant/text-to-speech-VITS-english"
+    print("--- Starting Weights Surgery for English ---")
+    
+    model_path = hf_hub_download(repo_id=repo_id, filename="best_model.pth")
+    config_path = hf_hub_download(repo_id=repo_id, filename="config.json")
+    
+    # Load and fix the tensor mismatch (137 -> 131)
+    checkpoint = torch.load(model_path, map_location="cpu")
+    raw_weights = checkpoint['model']['text_encoder.emb.weight']
+    
+    if raw_weights.shape[0] == 137:
+        checkpoint['model']['text_encoder.emb.weight'] = raw_weights[:131, :]
+        fixed_model_path = "fixed_eng_model.pth"
+        torch.save(checkpoint, fixed_model_path)
+        print("Surgery complete.")
+        return TTS(model_path=fixed_model_path, config_path=config_path, gpu=False)
+    
+    return TTS(model_path=model_path, config_path=config_path, gpu=False)
 
-print("Initializing Tamil...")
-tam_tts = load_my_model("E-motionAssistant/text-to-speech-VITS-tamil")
+# --- INITIALIZE ALL MODELS ---
+
+print("Initializing Models...")
+models = {
+    "Sinhala": load_standard_model("E-motionAssistant/text-to-speech-VITS-sinhala"),
+    "Tamil": load_standard_model("E-motionAssistant/text-to-speech-VITS-tamil"),
+    "English": load_eng_model_with_surgery()
+}
+print("All systems online.")
+
+# --- INFERENCE FUNCTION ---
 
 def generate_voice(text, language):
     try:
+        engine = models.get(language)
+        if not engine:
+            return None
+
+        # Apply specific preprocessing
         if language == "Sinhala":
-            # Use your separate function
             processed_text = sinhala_to_roman(text)
-            print(f"Input: {text} -> Romanized: {processed_text}")
-            engine = sin_tts
+            print(f"Sinhala Romanized: {processed_text}")
         else:
-            processed_text = text 
-            engine = tam_tts
+            processed_text = text
 
-        output_path = "output.wav"
-        engine.tts_to_file(text=processed_text, file_path=output_path)
+        output_path = f"output_{language.lower()}.wav"
+        engine.tts_to_file(text=str(processed_text), file_path=output_path)
         return output_path
+
     except Exception as e:
-        print(f"Error: {e}")
+        print(f"Synthesis Error ({language}): {e}")
         return None
 
-# Gradio Interface
+# --- GRADIO INTERFACE ---
+
 demo = gr.Interface(
     fn=generate_voice,
     inputs=[
-        gr.Textbox(label="Input Text"),
-        gr.Dropdown(["Sinhala", "Tamil"], label="Select Language")
+        gr.Textbox(label="Input Text", placeholder="Enter text here..."),
+        gr.Dropdown(["English", "Sinhala", "Tamil"], label="Select Language", value="English")
     ],
     outputs=gr.Audio(label="Synthesized Speech", type="filepath"),
-    title="Multilingual VITS TTS"
+    title="Trilingual VITS TTS System",
+    description="A unified interface for English (with weight surgery), Sinhala (romanized), and Tamil speech synthesis."
 )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()