Update app.py

app.py

@@ -1,12 +1,10 @@
 import subprocess
 import sys
 import os
-import shutil
+import tempfile
 from huggingface_hub import hf_hub_download
 
-repo_id = "Praha-Labs/PrahaTTS-ML"
-
-# --- 1. PRE-FLIGHT: INSTALLS AND MEMORY PATCHING ---
+# --- 1. PRE-FLIGHT: BYPASS BUILD ISOLATION ---
 def pre_flight_setup():
     try:
         import chatterbox
@@ -16,49 +14,24 @@ def pre_flight_setup():
         subprocess.check_call([sys.executable, "-m", "pip", "install", "--no-build-isolation", "pkuseg==0.0.25"])
         subprocess.check_call([sys.executable, "-m", "pip", "install", "chatterbox-tts>=0.1.7"])
 
-    print("Downloading and applying Indic patch...")
-    config_path = hf_hub_download(repo_id=repo_id, filename="config_indic.py")
-    shutil.copy(config_path, "config_indic.py")
-    
-    # CRITICAL FIX: Purge chatterbox from Python's memory cache!
-    # This forces the library to completely reload *after* our patch is applied.
-    modules_to_remove = [mod for mod in sys.modules if mod.startswith("chatterbox")]
-    for mod in modules_to_remove:
-        del sys.modules[mod]
-    
-    import config_indic
-    if hasattr(config_indic, 'apply_config'):
-        try:
-            config_indic.apply_config()
-        except TypeError:
-            pass
-
-# Run the setup immediately
 pre_flight_setup()
 
-# --- 2. NOW WE IMPORT CHATTERBOX (CLEAN RELOAD) ---
+# --- 2. MAIN APPLICATION ---
 import gradio as gr
 import torch
+import torch.nn as nn
 import torchaudio as ta
 from peft import PeftModel
-import tempfile
 
 from chatterbox.tts import ChatterboxTTS
 from chatterbox.models.tokenizers import EnTokenizer
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
+repo_id = "Praha-Labs/PrahaTTS-ML"
 
 def load_model():
     print(f"Loading base Chatterbox model on {device}...")
     model = ChatterboxTTS.from_pretrained(device=device)
-    
-    import config_indic
-    # Sometimes the config requires the model object itself to resize layers
-    if hasattr(config_indic, 'apply_config'):
-        try:
-            config_indic.apply_config(model)
-        except TypeError:
-            pass
 
     print("Applying custom Indic tokenizer...")
     try:
@@ -67,10 +40,32 @@ def load_model():
     except Exception as e:
         print(f"Error during tokenizer inject: {e}")
 
+    # --- CRITICAL FIX: MANUALLY RESIZE PYTORCH EMBEDDINGS ---
+    # We must resize the base model's vocabulary layers to match the new 
+    # Malayalam vocab size (2573) before loading the adapter weights.
+    vocab_size = 2573
+    print(f"Resizing base model embeddings to handle vocab size of {vocab_size}...")
+    
+    target_layer = model.t3 if hasattr(model, 't3') else model
+    
+    if hasattr(target_layer, 'text_emb'):
+        embed_dim = target_layer.text_emb.embedding_dim
+        target_layer.text_emb = nn.Embedding(vocab_size, embed_dim)
+        
+    if hasattr(target_layer, 'text_head'):
+        in_features = target_layer.text_head.in_features
+        has_bias = target_layer.text_head.bias is not None
+        target_layer.text_head = nn.Linear(in_features, vocab_size, bias=has_bias)
+        
+    # Send resized layers to the correct device
+    target_layer.to(device)
+
     print("Loading LoRA adapter weights...")
     try:
-        # We wrap the ENTIRE model now, ensuring any saved embedding layers are swapped
-        model = PeftModel.from_pretrained(model, repo_id)
+        if hasattr(model, 't3'):
+            model.t3 = PeftModel.from_pretrained(model.t3, repo_id)
+        else:
+            model = PeftModel.from_pretrained(model, repo_id)
         print("LoRA adapter loaded successfully.")
     except Exception as e:
         print(f"Failed to load PEFT adapter: {e}")