initial commit

app.py

@@ -443,154 +443,75 @@ print("--- Model Loading Complete ---")
 
 # --- Part 3: Full Pipeline Function for Gradio ---
 @spaces.GPU # For ZeroGPU execution context
-def full_speech_translation_pipeline_gradio(audio_input_path: str):
-    # This print will show the device context *inside* the decorated function
-    # For ZeroGPU, this should ideally show 'cuda:X' when the function is executed
-    # Ensure models are loaded and on the correct device before this function is called.
-    # The global DEVICE variable should reflect the GPU if available.
-    if stt_model: # Check if at least one model is loaded to get device
-        current_pipeline_device = next(stt_model.parameters()).device
-        print(f"--- @spaces.GPU function: Pipeline attempting to run on device: {current_pipeline_device} ---")
-    else:
-        print(f"--- @spaces.GPU function: STT model not loaded, cannot determine precise pipeline device yet. Target: {DEVICE} ---")
-
-
-    if not all([TTS_MODEL, stt_processor, stt_model, mt_tokenizer, mt_model]):
-        error_msg = "Critical Error: One or more models failed to load during Space initialization. Cannot process."
-        print(error_msg)
-        raise gr.Error(error_msg) # Use Gradio's error for better UI feedback
-
-    if audio_input_path is None: # Gradio sends None if no file
-        raise gr.Error("No audio file provided. Please upload an audio file.")
-
-    # Check if the file path actually exists (Gradio should handle temp file creation)
-    # This check might be redundant if Gradio always provides a valid temp path for uploaded files.
-    if not os.path.exists(audio_input_path):
-        error_msg = f"Error: Audio file path provided by Gradio does not exist: {audio_input_path}"
-        print(error_msg)
-        raise gr.Error("Internal error: Audio file path not found.")
-
-
-    print(f"--- GRADIO PIPELINE START (GPU context): Processing {audio_input_path} ---")
-
-    # STT Stage (aligning with your original logic)
+def full_speech_translation_pipeline(audio_input_path: str):
+    print(f"--- PIPELINE START: Processing {audio_input_path} ---")
+    if audio_input_path is None or not os.path.exists(audio_input_path):
+        msg = "Error: Audio file not provided or not found."
+        print(msg)
+        # Return empty/default values
+        return "Error: No file", "", (hp.sr, np.array([]).astype(np.float32))
+
+    # STT Stage
     arabic_transcript = "STT Error: Processing failed."
     try:
         print("STT: Loading and resampling audio...")
         wav, sr = torchaudio.load(audio_input_path)
-        if wav.size(0) > 1: wav = wav.mean(dim=0, keepdim=True) # Ensure mono
-        
-        # Ensure stt_processor is available
-        if not stt_processor: raise gr.Error("STT Processor not loaded.")
+        if wav.size(0) > 1: wav = wav.mean(dim=0, keepdim=True)
         target_sr_stt = stt_processor.feature_extractor.sampling_rate
+        if sr != target_sr_stt: wav = torchaudio.transforms.Resample(sr, target_sr_stt)(wav)
+        audio_array_stt = wav.squeeze().cpu().numpy()
         
-        if sr != target_sr_stt:
-            # Resample op can be on CPU or GPU. If wav is already on GPU, it might be faster.
-            # However, for simplicity and consistency with original, let's keep it default (CPU).
-            resampler = torchaudio.transforms.Resample(sr, target_sr_stt)
-            wav = resampler(wav)
-            
-        audio_array_stt = wav.squeeze().cpu().numpy() # Whisper processor expects NumPy array
-
         print("STT: Extracting features and transcribing...")
-        # Ensure stt_model is available
-        if not stt_model: raise gr.Error("STT Model not loaded.")
-        inputs_stt = stt_processor(audio_array_stt, sampling_rate=target_sr_stt, return_tensors="pt").input_features.to(DEVICE)
+        inputs = stt_processor(audio_array_stt, sampling_rate=target_sr_stt, return_tensors="pt").input_features.to(DEVICE)
         forced_ids = stt_processor.get_decoder_prompt_ids(language="arabic", task="transcribe")
-        
         with torch.no_grad():
-            generated_ids = stt_model.generate(inputs_stt, forced_decoder_ids=forced_ids, max_length=448) # Using your original max_length
-        arabic_transcript = stt_processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
+            generated_ids = stt_model.generate(inputs, forced_decoder_ids=forced_ids, max_length=448)
+        arabic_transcript = stt_processor.decode(generated_ids[0], skip_special_tokens=True).strip()
         print(f"STT Output: {arabic_transcript}")
     except Exception as e:
         print(f"STT Error: {e}")
-        raise gr.Error(f"STT processing failed: {str(e)}") # Show error in Gradio UI
 
-    # TTT Stage (aligning with your original logic)
+    # TTT Stage
     english_translation = "TTT Error: Processing failed."
     if arabic_transcript and not arabic_transcript.startswith("STT Error"):
         try:
             print("TTT: Translating to English...")
-            # Ensure mt_tokenizer and mt_model are available
-            if not mt_tokenizer: raise gr.Error("Marian Tokenizer not loaded.")
-            if not mt_model: raise gr.Error("Marian Model not loaded.")
-                
-            batch = mt_tokenizer(arabic_transcript, return_tensors="pt", padding=True, truncation=True).to(DEVICE) # Added truncation
+            batch = mt_tokenizer(arabic_transcript, return_tensors="pt", padding=True).to(DEVICE)
             with torch.no_grad():
                 translated_ids = mt_model.generate(**batch, max_length=512)
             english_translation = mt_tokenizer.batch_decode(translated_ids, skip_special_tokens=True)[0].strip()
             print(f"TTT Output: {english_translation}")
         except Exception as e:
             print(f"TTT Error: {e}")
-            raise gr.Error(f"TTT processing failed: {str(e)}") # Show error in Gradio UI
-    elif arabic_transcript.startswith("STT Error"):
-        english_translation = "(Skipped TTT due to STT error)" # More specific message
-        print(english_translation)
-    else: # Handles empty arabic_transcript case
-        english_translation = "(Skipped TTT due to empty STT output)"
+    else:
+        english_translation = "(Skipped TTT due to STT failure)"
         print(english_translation)
 
-
-    # TTS Stage (aligning with your original logic for inference call)
-    output_tts_audio_filepath = None # Initialize for Gradio output
-    if english_translation and not english_translation.startswith("TTT Error") and not english_translation.startswith("(Skipped TTT"):
-        if TTS_MODEL:
-            try:
-                print("TTS: Synthesizing English speech...")
-                if not english_translation.strip():
-                    print("TTS Warning: Empty string provided for synthesis. Skipping TTS.")
-                    # No gr.Error here, just skip TTS
-                else:
-                    sequence = text_to_seq(english_translation).unsqueeze(0).to(DEVICE)
-                    
-                    # Call TTS_MODEL.inference exactly as in your original, working pipeline
-                    # This assumes your TTS_MODEL.inference returns (mel_spectrogram, other_data_like_stop_tokens)
-                    generated_mel, _ = TTS_MODEL.inference(
-                        sequence,
-                        max_length=hp.max_mel_time-20, # Your original max_length
-                        stop_token_threshold=0.5 # Your original threshold
-                        # with_tqdm=False # Not needed for Gradio backend
-                    )
-                    
-                    print(f"TTS: Generated mel shape: {generated_mel.shape if generated_mel is not None else 'None'}")
-                    if generated_mel is not None and generated_mel.numel() > 0 and generated_mel.shape[1] > 0: # Check time dimension
-                        # Your original processing of mel_for_vocoder
-                        mel_for_vocoder = generated_mel.detach().squeeze(0).transpose(0, 1)
-                        audio_tensor = inverse_mel_spec_to_wav(mel_for_vocoder) # This needs to be correct for your model
-                        synthesized_audio_np = audio_tensor.cpu().numpy()
-                        print(f"TTS: Synthesized audio shape: {synthesized_audio_np.shape}")
-
-                        # Save to a temporary file for Gradio Audio output
-                        timestamp = int(time.time()*1000) # For a somewhat unique filename
-                        output_tts_audio_filepath = f"synthesized_speech_{timestamp}.wav"
-                        sf.write(output_tts_audio_filepath, synthesized_audio_np, hp.sr)
-                        print(f"TTS: Synthesized audio saved to: {output_tts_audio_filepath}")
-                    else:
-                        print("TTS Warning: Generated mel spectrogram was empty or invalid.")
-                        # Do not raise gr.Error, let the text pass through.
-                        # output_tts_audio_filepath remains None.
-            except Exception as e:
-                print(f"TTS Error: {e}")
-                # Append error to text, don't stop the whole process if text is available
-                english_translation += f" (TTS synthesis failed: {str(e)})"
-                # output_tts_audio_filepath remains None.
-        else:
-            print("TTS SKIPPED: TTS_MODEL not loaded.")
-    elif not TTS_MODEL: # If TTS model didn't load but we reached here
-         print("TTS SKIPPED: Model not loaded.")
-    else: # If english_translation was invalid for TTS
-         print(f"TTS SKIPPED: English text was '{english_translation}'.")
-
-
-    print(f"--- GRADIO PIPELINE END (GPU context) ---")
-    # Return values in the order expected by Gradio's `outputs` components
-    return arabic_transcript, english_translation, output_tts_audio_filepath
+    # TTS Stage
+    synthesized_audio_np = np.array([]).astype(np.float32)
+    if english_translation and not english_translation.startswith("TTT Error"):
+        try:
+            print("TTS: Synthesizing English speech...")
+            sequence = text_to_seq(english_translation).unsqueeze(0).to(DEVICE)
+            generated_mel, _ = TTS_MODEL.inference(sequence, max_length=hp.max_mel_time-20, stop_token_threshold=0.5, with_tqdm=False)
+            
+            print(f"TTS: Generated mel shape: {generated_mel.shape if generated_mel is not None else 'None'}")
+            if generated_mel is not None and generated_mel.numel() > 0:
+                mel_for_vocoder = generated_mel.detach().squeeze(0).transpose(0, 1)
+                audio_tensor = inverse_mel_spec_to_wav(mel_for_vocoder)
+                synthesized_audio_np = audio_tensor.cpu().numpy()
+                print(f"TTS: Synthesized audio shape: {synthesized_audio_np.shape}")
+        except Exception as e:
+            print(f"TTS Error: {e}")
+    
+    print(f"--- PIPELINE END ---")
+    return arabic_transcript, english_translation, (hp.sr, synthesized_audio_np)
 
 
 # --- Part 4: Gradio Interface Definition ---
 # (Same as before)
 iface = gr.Interface(
-    fn=full_speech_translation_pipeline_gradio,
+    fn=full_speech_translation_pipeline,
     inputs=[
         gr.Audio(type="filepath", label="Upload Arabic Speech")
     ],