Update app.py

app.py

@@ -2,6 +2,7 @@ import os
 import shutil
 import subprocess
 import traceback
+import wave
 from uuid import uuid4
 from fastapi import FastAPI, UploadFile, HTTPException
 from markitdown import MarkItDown
@@ -11,14 +12,14 @@ import pytesseract
 from pdf2image import convert_from_path
 import riva.client as rcli
 
-app = FastAPI(title="Free MarkItDown Service with Official NVIDIA Riva ASR")
+app = FastAPI(title="Free MarkItDown & Real-Time Voice Processing")
 md = MarkItDown()
 
 IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".bmp", ".tiff", ".webp"}
 AUDIO_EXTENSIONS = {".mp3", ".wav", ".m4a", ".ogg", ".flac", ".opus"}
 LEGACY_OFFICE_EXTENSIONS = {".doc", ".xls", ".ppt", ".rtf"}
 
-# Official NVIDIA ASR Configuration
+# NVIDIA Riva gRPC Configuration Endpoint
 NVIDIA_STAGE_SERVER = "grpc.nvcf.nvidia.com:443"
 NVIDIA_PARAKEET_FUNCTION_ID = "b0e8b4a5-217c-40b7-9b96-17d84e666317"
 
@@ -26,7 +27,7 @@ def get_nvidia_api_keys():
     """Dynamically collects NVIDIA API keys from environment variables."""
     keys = []
     
-    # Check for numbered keys: NVIDIA_API_KEY_1, NVIDIA_API_KEY_2, etc.
+    # Check for numbered environment keys: NVIDIA_API_KEY_1, NVIDIA_API_KEY_2, etc.
     index = 1
     while True:
         key = os.getenv(f"NVIDIA_API_KEY_{index}")
@@ -36,7 +37,7 @@ def get_nvidia_api_keys():
         else:
             break
 
-    # Fallback to single NVIDIA_API_KEY
+    # Fallback to single NVIDIA_API_KEY if numbered keys aren't set
     if not keys:
         single_key = os.getenv("NVIDIA_API_KEY")
         if single_key and single_key.strip():
@@ -45,77 +46,83 @@ def get_nvidia_api_keys():
     return keys
 
 def transcribe_with_nvidia_riva(audio_path: str) -> str:
-    """Transcribes audio using NVIDIA's standard gRPC client with multi-key rotation."""
+    """Transcribes audio using NVIDIA's gRPC client with multi-key failover."""
     api_keys = get_nvidia_api_keys()
     
     if not api_keys:
         return "Audio Error: No NVIDIA API keys (NVIDIA_API_KEY_1, NVIDIA_API_KEY_2...) set in environment variables."
 
-    wav_path = audio_path
-    converted_wav = None
+    converted_wav = f"/tmp/{uuid4()}.wav"
 
-    # Convert audio to mono 16kHz WAV format (required by Riva)
-    file_ext = os.path.splitext(audio_path)[1].lower()
-    if file_ext != ".wav":
-        converted_wav = f"/tmp/{uuid4()}.wav"
-        try:
-            subprocess.run(
-                ["ffmpeg", "-y", "-i", audio_path, "-ac", "1", "-ar", "16000", converted_wav],
-                check=True,
-                stdout=subprocess.DEVNULL,
-                stderr=subprocess.DEVNULL
-            )
-            wav_path = converted_wav
-        except Exception as e:
-            return f"Audio Conversion Error (FFmpeg): {str(e)}"
+    # 1. Convert any input format (MP3, OPUS, uploaded WAV, recorded WAV) 
+    # into strictly 16kHz, Mono, 16-bit PCM WAV using FFmpeg
+    try:
+        subprocess.run(
+            [
+                "ffmpeg", "-y", "-i", audio_path, 
+                "-ac", "1", 
+                "-ar", "16000", 
+                "-sample_fmt", "s16", 
+                converted_wav
+            ],
+            check=True,
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL
+        )
+    except Exception as e:
+        return f"Audio Pre-processing Error (FFmpeg): {str(e)}"
 
     last_error = ""
+    audio_data = None
 
-    # Rotate through keys if an error/rate limit occurs
-    for idx, key in enumerate(api_keys, start=1):
-        try:
-            # Construct standard NVIDIA gRPC metadata headers
-            auth = rcli.Auth(
-                use_ssl=True,
-                uri=NVIDIA_STAGE_SERVER,
-                metadata_args=[
-                    ("authorization", f"Bearer {key}"),
-                    ("function-id", NVIDIA_PARAKEET_FUNCTION_ID)
-                ]
-            )
-
-            asr_service = rcli.ASRService(auth)
-            config = rcli.RecognitionConfig(
-                encoding=rcli.AudioEncoding.LINEAR_PCM,
-                sample_rate_hertz=16000,
-                language_code="en-US",
-                max_alternatives=1
-            )
-
-            with open(wav_path, "rb") as fh:
-                audio_data = fh.read()
-
-            response = asr_service.offline_recognize(audio_data, config)
-            
-            # Process returned transcript
-            results = response.results
-            transcripts = [res.alternatives[0].transcript for res in results if res.alternatives]
-            text = " ".join(transcripts).strip()
-            
-            if converted_wav and os.path.exists(converted_wav):
-                os.remove(converted_wav)
-
-            return f"## Audio Transcription\n\n{text}" if text else "Audio processed, but no speech was detected."
-
-        except Exception as e:
-            last_error = str(e)
-            print(f"Key #{idx} failed with error: {last_error}. Retrying with next key...")
-            continue
+    # 2. Extract ONLY raw PCM frames, stripping the 44-byte WAV header
+    try:
+        with wave.open(converted_wav, "rb") as wf:
+            audio_data = wf.readframes(wf.getnframes())
+    except Exception as e:
+        last_error = f"WAV Parse Error: {str(e)}"
 
-    if converted_wav and os.path.exists(converted_wav):
-        os.remove(converted_wav)
+    try:
+        if audio_data:
+            # 3. Rotate through API keys upon rate-limiting or errors
+            for idx, key in enumerate(api_keys, start=1):
+                try:
+                    auth = rcli.Auth(
+                        use_ssl=True,
+                        uri=NVIDIA_STAGE_SERVER,
+                        metadata_args=[
+                            ("authorization", f"Bearer {key}"),
+                            ("function-id", NVIDIA_PARAKEET_FUNCTION_ID)
+                        ]
+                    )
+
+                    asr_service = rcli.ASRService(auth)
+                    config = rcli.RecognitionConfig(
+                        encoding=rcli.AudioEncoding.LINEAR_PCM,
+                        sample_rate_hertz=16000,
+                        language_code="en-US",
+                        max_alternatives=1,
+                        audio_channel_count=1
+                    )
+
+                    response = asr_service.offline_recognize(audio_data, config)
+                    
+                    results = response.results
+                    transcripts = [res.alternatives[0].transcript for res in results if res.alternatives]
+                    text = " ".join(transcripts).strip()
+
+                    return f"## Audio Transcription\n\n{text}" if text else "Audio processed, but no speech was detected."
+
+                except Exception as e:
+                    last_error = str(e)
+                    print(f"NVIDIA API Key #{idx} failed: {last_error}. Trying next key...")
+                    continue
+
+        return f"Audio Transcription Error: All provided NVIDIA API keys failed. Last error: {last_error}"
 
-    return f"Audio Transcription Error: All provided NVIDIA API keys failed. Last error: {last_error}"
+    finally:
+        if os.path.exists(converted_wav):
+            os.remove(converted_wav)
 
 def run_free_ocr(file_path: str) -> str:
     """Fallback OCR pipeline using Tesseract (images and PDFs only)."""
@@ -174,7 +181,7 @@ def process_input(source_path_or_url: str, original_filename: str = None) -> str
     temp_converted_office = None
 
     try:
-        # Route 1: Audio files -> Standard NVIDIA gRPC ASR
+        # Route 1: Audio files (Uploaded or Recorded) -> NVIDIA Riva gRPC
         if ext in AUDIO_EXTENSIONS:
             return transcribe_with_nvidia_riva(working_path)
 
@@ -244,58 +251,47 @@ async def convert_file_api(file: UploadFile = None, url: str = None):
             os.remove(temp_path)
 
 # -------------------------------------------------------------
-# Updated Gradio Web UI
+# Gradio UI with Real-time Recording & File Handling
 # -------------------------------------------------------------
-def gradio_handler(file_obj, mic_file_path, url_input):
+def gradio_handler(mic_file_path, uploaded_file_obj, url_input):
     """
-    Priority Order:
-    1. Recorded Audio (Microphone)
-    2. Uploaded File
-    3. URL Input
+    Evaluates inputs with priority given to live microphone recordings.
     """
-    # Case 1: Real-time recording from the microphone
     if mic_file_path is not None:
         return process_input(mic_file_path, original_filename="recorded_audio.wav")
-        
-    # Case 2: Standard file upload (.pdf, .doc, .mp3, .opus, etc.)
-    elif file_obj is not None:
-        orig_name = getattr(file_obj, "orig_name", file_obj.name)
-        return process_input(file_obj.name, original_filename=orig_name)
-        
-    # Case 3: URL string
+    elif uploaded_file_obj is not None:
+        orig_name = getattr(uploaded_file_obj, "orig_name", uploaded_file_obj.name)
+        return process_input(uploaded_file_obj.name, original_filename=orig_name)
     elif url_input and url_input.strip():
         return process_input(url_input)
-        
-    return "Please record audio, upload a file, or enter a valid URL."
-
+    
+    return "Please record voice, upload a file, or provide a URL."
 
-with gr.Blocks(title="Free MarkItDown & Audio Converter") as demo:
-    gr.Markdown("# 📄 Document & Real-Time Audio Converter")
+with gr.Blocks(title="Free Converter & Voice Transcriber") as demo:
+    gr.Markdown("# 📄 Document & Real-Time Voice Transcriber")
     
     with gr.Row():
         with gr.Column():
-            gr.Markdown("### Option 1: Record Audio Live")
-            # Microphone component (saves recorded audio to temporary .wav file automatically)
+            gr.Markdown("### Option 1: Record Live Voice")
             mic_input = gr.Audio(
                 sources=["microphone"], 
                 type="filepath", 
-                label="Click red icon to record voice"
+                label="Click microphone to start recording"
             )
             
             gr.Markdown("---")
             gr.Markdown("### Option 2: Upload File or URL")
-            file_input = gr.File(label="Upload Document or Audio (.pdf, .mp3, .opus, etc.)")
+            file_input = gr.File(label="Upload File (.pdf, .doc, .mp3, .wav, .png, etc.)")
             url_input = gr.Textbox(label="Or enter URL", placeholder="https://example.com/document.pdf")
             
-            submit_btn = gr.Button("Process & Convert", variant="primary")
+            submit_btn = gr.Button("Process & Transcribe", variant="primary")
         
         with gr.Column():
-            output_text = gr.Textbox(label="Converted Output / Transcription", lines=22)
+            output_text = gr.Textbox(label="Result Output / Markdown", lines=24)
 
-    # Trigger processing on button click
     submit_btn.click(
         fn=gradio_handler,
-        inputs=[file_input, mic_input, url_input],
+        inputs=[mic_input, file_input, url_input],
         outputs=output_text
     )