initial commit

app.py

@@ -393,42 +393,51 @@ TTS_MODEL_HUB_ID = "MoHamdyy/transformer-tts-ljspeech"
 ASR_HUB_ID       = "MoHamdyy/whisper-stt-model"
 MARIAN_HUB_ID    = "MoHamdyy/marian-ar-en-translation"
 
-TTS_MODEL = None
-stt_processor = None
-stt_model = None
-mt_tokenizer = None
-mt_model = None
+
 
 # Wrap model loading in a function to clearly see when it happens or to potentially delay it.
 # For Spaces, global loading is fine and preferred as it happens once.
 print("--- Starting Model Loading ---")
 try:
-    print(f"Loading TTS model ({TTS_MODEL_HUB_ID}) to {DEVICE}...")
+    print("Loading TTS model...")
+    # Download the .pt file from its repo
     tts_model_path = hf_hub_download(repo_id=TTS_MODEL_HUB_ID, filename="train_SimpleTransfromerTTS.pt")
-    state = torch.load(tts_model_path, map_location=DEVICE) # Load to target device directly
-    TTS_MODEL = TransformerTTS(device=DEVICE).to(DEVICE)
-    model_state_dict = state.get("model", state.get("state_dict", state))
-    TTS_MODEL.load_state_dict(model_state_dict)
+    state = torch.load(tts_model_path, map_location=DEVICE)
+    TTS_MODEL = TransformerTTS().to(DEVICE)
+    # Check for the correct key in the state dictionary
+    if "model" in state:
+        TTS_MODEL.load_state_dict(state["model"])
+    elif "state_dict" in state:
+        TTS_MODEL.load_state_dict(state["state_dict"])
+    else:
+        TTS_MODEL.load_state_dict(state) # Assume the whole file is the state_dict
     TTS_MODEL.eval()
     print("TTS model loaded successfully.")
 except Exception as e:
     print(f"Error loading TTS model: {e}")
+    TTS_MODEL = None
 
+# Load STT (Whisper) Model from Hub
 try:
-    print(f"Loading STT (Whisper) model ({ASR_HUB_ID}) to {DEVICE}...")
+    print("Loading STT (Whisper) model...")
     stt_processor = WhisperProcessor.from_pretrained(ASR_HUB_ID)
     stt_model = WhisperForConditionalGeneration.from_pretrained(ASR_HUB_ID).to(DEVICE).eval()
     print("STT model loaded successfully.")
 except Exception as e:
     print(f"Error loading STT model: {e}")
+    stt_processor = None
+    stt_model = None
 
+# Load TTT (MarianMT) Model from Hub
 try:
-    print(f"Loading TTT (MarianMT) model ({MARIAN_HUB_ID}) to {DEVICE}...")
+    print("Loading TTT (MarianMT) model...")
     mt_tokenizer = MarianTokenizer.from_pretrained(MARIAN_HUB_ID)
     mt_model = MarianMTModel.from_pretrained(MARIAN_HUB_ID).to(DEVICE).eval()
     print("TTT model loaded successfully.")
 except Exception as e:
     print(f"Error loading TTT model: {e}")
+    mt_tokenizer = None
+    mt_model = None
 print("--- Model Loading Complete ---")