Update app.py

app.py

@@ -13,16 +13,57 @@ from helper import (
     ensure_wav_format,
 )
 import wave
+import shutil
 
 app = Flask(__name__)
 CORS(app)
 os.environ["COQUI_TOS_AGREED"] = "1"
 
 device = "cpu"
-tts = TTS("tts_models/multilingual/multi-dataset/xtts_v2").to(device)
-active_tasks = {}
 
-# Optional: set your max audio size (MB)
+# ============================================================
+# MODEL PATH CONFIGURATION
+# ============================================================
+
+# Mounted dataset path inside Hugging Face Space
+DATASET_MODEL_DIR = "/datasets/EllenBeta/Xtts_2/model"
+
+# Local cache fallback path
+LOCAL_MODEL_DIR = os.path.expanduser("~/.local/share/tts/xtts_v2_model")
+
+# If model already exists in dataset, use it directly (no download)
+if os.path.exists(DATASET_MODEL_DIR) and os.listdir(DATASET_MODEL_DIR):
+    print(f"✅ Loading model directly from dataset: {DATASET_MODEL_DIR}")
+    model_path = DATASET_MODEL_DIR
+else:
+    print("⬇️ Downloading XTTS model (first-time only)...")
+    tts = TTS("tts_models/multilingual/multi-dataset/xtts_v2").to(device)
+    
+    # Save the model locally for re-use
+    os.makedirs(LOCAL_MODEL_DIR, exist_ok=True)
+    src_model_dir = tts.model_dir
+
+    # Copy model to dataset path for future Space restarts
+    if os.path.exists(src_model_dir):
+        print(f"📦 Copying model files to dataset for reuse: {DATASET_MODEL_DIR}")
+        os.makedirs(DATASET_MODEL_DIR, exist_ok=True)
+        for item in os.listdir(src_model_dir):
+            s = os.path.join(src_model_dir, item)
+            d = os.path.join(DATASET_MODEL_DIR, item)
+            if os.path.isdir(s):
+                shutil.copytree(s, d, dirs_exist_ok=True)
+            else:
+                shutil.copy2(s, d)
+
+    model_path = src_model_dir
+    print("✅ Model cached successfully.")
+
+# Initialize TTS from cached or dataset path
+tts = TTS(model_path=model_path).to(device)
+
+# ============================================================
+
+active_tasks = {}
 MAX_AUDIO_SIZE_MB = 15
 
 
@@ -80,7 +121,6 @@ def generate_voice():
 def process_vox(user_id, text, video, audio_base64, task_id):
     temp_audio_path = None
     try:
-        # Handle audio or video input
         if audio_base64:
             if audio_base64.startswith("data:audio/"):
                 audio_base64 = audio_base64.split(",", 1)[1]
@@ -91,7 +131,6 @@ def process_vox(user_id, text, video, audio_base64, task_id):
             temp_audio_path = video_to_audio(video, output_path=temp_audio_path)
 
         temp_audio_path = ensure_wav_format(temp_audio_path)
-
         valid, msg = validate_audio_file(temp_audio_path, MAX_AUDIO_SIZE_MB)
         if not valid:
             raise Exception(f"Invalid audio file: {msg}")
@@ -111,7 +150,6 @@ def process_vox(user_id, text, video, audio_base64, task_id):
             duration = f"{dura:.2f}"
             title = text[:20]
 
-        # Upload + save metadata
         audio_url = save_to_dataset_repo(file_path, f"user/data/audios/{file_name}", file_name)
         active_tasks[task_id].update(
             {