initial commit

app.py

@@ -442,112 +442,148 @@ print("--- Model Loading Complete ---")
 
 
 # --- Part 3: Full Pipeline Function for Gradio ---
-@spaces.GPU # <<< --- APPLY THE DECORATOR HERE --- <<<
-def full_speech_translation_pipeline_gradio(audio_input_path):
+@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
-    current_processing_device = next(stt_model.parameters()).device if stt_model else "CPU (STT model not loaded)"
-    print(f"--- @spaces.GPU function: Pipeline running on device: {current_processing_device} ---")
+    # 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)
-        # Raising gr.Error is better for UI feedback
-        raise gr.Error(error_msg)
-
+        raise gr.Error(error_msg) # Use Gradio's error for better UI feedback
 
-    if audio_input_path is None:
-        # This case should ideally be handled by Gradio's input validation or a check before calling.
-        # If it still occurs, provide a clear message.
+    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
+    # STT Stage (aligning with your original logic)
     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)
+        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.")
         target_sr_stt = stt_processor.feature_extractor.sampling_rate
-        if sr != target_sr_stt: wav = torchaudio.transforms.Resample(sr, target_sr_stt)(wav)
-        # Move wav to the STT model's device *before* converting to numpy if STT model is on GPU
-        audio_array_stt = wav.to(DEVICE).squeeze().cpu().numpy() # Process on DEVICE, then to CPU for 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 inputs are on the same device as the model
+        # 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)
         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_new_tokens=256)
+            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()
         print(f"STT Output: {arabic_transcript}")
     except Exception as e:
         print(f"STT Error: {e}")
-        raise gr.Error(f"STT processing failed: {e}")
-
+        raise gr.Error(f"STT processing failed: {str(e)}") # Show error in Gradio UI
 
-    # TTT Stage
+    # TTT Stage (aligning with your original logic)
     english_translation = "TTT Error: Processing failed."
     if arabic_transcript and not arabic_transcript.startswith("STT Error"):
         try:
             print("TTT: Translating to English...")
-            batch = mt_tokenizer(arabic_transcript, return_tensors="pt", padding=True, truncation=True).to(DEVICE)
+            # 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
             with torch.no_grad():
-                translated_ids = mt_model.generate(**batch, max_length=512) # max_new_tokens can also be used
+                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: {e}")
-
-    else:
-        if not arabic_transcript or arabic_transcript.startswith("STT Error"):
-            english_translation = "(Skipped TTT due to STT failure or empty STT output)"
+            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)"
         print(english_translation)
 
 
-    # TTS Stage
-    output_tts_audio_filepath = None
-    if english_translation and not english_translation.startswith("TTT Error") and TTS_MODEL:
-        try:
-            print("TTS: Synthesizing English speech...")
-            if not english_translation.strip():
-                 print("TTS Warning: Empty string for synthesis. Skipping TTS.")
-            else:
-                sequence = text_to_seq(english_translation).unsqueeze(0).to(DEVICE)
-                # max_length for TTS inference refers to max output mel frames
-                generated_mel = TTS_MODEL.inference(sequence, max_length=hp.max_mel_time - 50, stop_token_threshold=0.5)
-
-                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 :
-                    # TTS model's inverse_mel_spec_to_wav expects mel on DEVICE and returns wav on CPU
-                    # The mel from inference should be [N, mel_len, mel_bins]
-                    # inverse_mel_spec_to_wav might expect [mel_bins, mel_len]
-                    mel_for_vocoder = generated_mel.detach().squeeze(0).transpose(0, 1) # to [mel_len, mel_bins]
-                    audio_tensor = inverse_mel_spec_to_wav(mel_for_vocoder) # This function handles .to(DEVICE) internally
-                    synthesized_audio_np = audio_tensor.cpu().numpy() # Ensure output is on CPU for soundfile
-                    print(f"TTS: Synthesized audio shape: {synthesized_audio_np.shape}")
-
-                    timestamp = int(time.time()*1000) # more unique
-                    output_tts_audio_filepath = f"output_audio_{timestamp}.wav"
-                    sf.write(output_tts_audio_filepath, synthesized_audio_np, hp.sr)
-                    print(f"TTS: Synthesized audio saved to: {output_tts_audio_filepath}")
+    # 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:
-                    print("TTS Warning: Generated mel spectrogram was empty or invalid.")
-        except Exception as e:
-            print(f"TTS Error: {e}")
-            # Do not raise gr.Error here if a partial result (text) is still useful
-            # output_tts_audio_filepath will remain None
-            english_translation += f" (TTS Error: {e})" # Append error to text
-    else:
-        if not TTS_MODEL: print("TTS SKIPPED: Model not loaded.")
-        elif not (english_translation and not english_translation.startswith("TTT Error")):
-             print("TTS SKIPPED: (Due to TTT failure or empty TTT output)")
+                    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