Update app.py

app.py

@@ -28,26 +28,30 @@ def load_model():
         processor = ColQwen2_5OmniProcessor.from_pretrained("manu/colqwen-omni-v0.1")
     return model, processor
 
-def chunk_audio(audio_file, chunk_length=30):
+def chunk_audio(audio_file_path, chunk_length=30):
     """Split audio into chunks"""
-    audio = AudioSegment.from_file(audio_file.name)
-    
-    audios = []
-    target_rate = 16000
-    chunk_length_ms = chunk_length * 1000
-    
-    for i in range(0, len(audio), chunk_length_ms):
-        chunk = audio[i:i + chunk_length_ms]
-        chunk = chunk.set_channels(1).set_frame_rate(target_rate)
+    try:
+        # audio_file_path is already a string path when type="filepath"
+        audio = AudioSegment.from_file(audio_file_path)
         
-        buf = io.BytesIO()
-        chunk.export(buf, format="wav")
-        buf.seek(0)
+        audios = []
+        target_rate = 16000
+        chunk_length_ms = chunk_length * 1000
         
-        rate, data = wavfile.read(buf)
-        audios.append(data)
-    
-    return audios
+        for i in range(0, len(audio), chunk_length_ms):
+            chunk = audio[i:i + chunk_length_ms]
+            chunk = chunk.set_channels(1).set_frame_rate(target_rate)
+            
+            buf = io.BytesIO()
+            chunk.export(buf, format="wav")
+            buf.seek(0)
+            
+            rate, data = wavfile.read(buf)
+            audios.append(data)
+        
+        return audios
+    except Exception as e:
+        raise gr.Error(f"Error processing audio file: {str(e)}. Make sure ffmpeg is installed.")
 
 @spaces.GPU(duration=120)
 def embed_audio_chunks(audios):
@@ -108,66 +112,73 @@ def audio_to_base64(data, rate=16000):
     encoded_string = base64.b64encode(buf.read()).decode("utf-8")
     return encoded_string
 
-def process_audio_rag(audio_file, query, chunk_length=30, use_openai=False, openai_key=None):
+def process_audio_rag(audio_file_path, query, chunk_length=30, use_openai=False, openai_key=None):
     """Main processing function"""
-    if not audio_file:
+    if not audio_file_path:
         return "Please upload an audio file", None, None
     
-    # Chunk audio
-    audios = chunk_audio(audio_file, chunk_length)
-    
-    # Embed chunks
-    embeddings = embed_audio_chunks(audios)
-    
-    # Search for relevant chunks
-    top_indices = search_audio(query, embeddings, audios)
-    
-    # Prepare results
-    result_text = f"Found {len(top_indices)} relevant audio chunks:\n"
-    result_text += f"Chunk indices: {top_indices}\n\n"
-    
-    # Save first result as audio file
-    first_chunk_path = "result_chunk.wav"
-    wavfile.write(first_chunk_path, 16000, audios[top_indices[0]])
+    if not query:
+        return "Please enter a search query", None, None
     
-    # Optional: Use OpenAI for answer generation
-    if use_openai and openai_key:
-        from openai import OpenAI
-        client = OpenAI(api_key=openai_key)
+    try:
+        # Chunk audio
+        audios = chunk_audio(audio_file_path, chunk_length)
         
-        content = [{"type": "text", "text": f"Answer the query using the audio files. Query: {query}"}]
+        # Embed chunks
+        embeddings = embed_audio_chunks(audios)
         
-        for idx in top_indices[:3]:  # Use top 3 chunks
-            content.extend([
-                {"type": "text", "text": f"Audio chunk #{idx}:"},
-                {
-                    "type": "input_audio",
-                    "input_audio": {
-                        "data": audio_to_base64(audios[idx]),
-                        "format": "wav"
+        # Search for relevant chunks
+        top_indices = search_audio(query, embeddings, audios)
+        
+        # Prepare results
+        result_text = f"Found {len(top_indices)} relevant audio chunks:\n"
+        result_text += f"Chunk indices: {top_indices}\n\n"
+        
+        # Save first result as audio file
+        first_chunk_path = "result_chunk.wav"
+        wavfile.write(first_chunk_path, 16000, audios[top_indices[0]])
+        
+        # Optional: Use OpenAI for answer generation
+        if use_openai and openai_key:
+            from openai import OpenAI
+            client = OpenAI(api_key=openai_key)
+            
+            content = [{"type": "text", "text": f"Answer the query using the audio files. Query: {query}"}]
+            
+            for idx in top_indices[:3]:  # Use top 3 chunks
+                content.extend([
+                    {"type": "text", "text": f"Audio chunk #{idx}:"},
+                    {
+                        "type": "input_audio",
+                        "input_audio": {
+                            "data": audio_to_base64(audios[idx]),
+                            "format": "wav"
+                        }
                     }
-                }
-            ])
+                ])
+            
+            try:
+                completion = client.chat.completions.create(
+                    model="gpt-4o-audio-preview",
+                    messages=[{"role": "user", "content": content}]
+                )
+                result_text += f"\nOpenAI Answer: {completion.choices[0].message.content}"
+            except Exception as e:
+                result_text += f"\nOpenAI Error: {str(e)}"
         
-        try:
-            completion = client.chat.completions.create(
-                model="gpt-4o-audio-preview",
-                messages=[{"role": "user", "content": content}]
-            )
-            result_text += f"\nOpenAI Answer: {completion.choices[0].message.content}"
-        except Exception as e:
-            result_text += f"\nOpenAI Error: {str(e)}"
-    
-    # Create audio visualization
-    import matplotlib.pyplot as plt
-    fig, ax = plt.subplots(figsize=(10, 4))
-    ax.plot(audios[top_indices[0]])
-    ax.set_title(f"Waveform of top matching chunk (#{top_indices[0]})")
-    ax.set_xlabel("Samples")
-    ax.set_ylabel("Amplitude")
-    plt.tight_layout()
-    
-    return result_text, first_chunk_path, fig
+        # Create audio visualization
+        import matplotlib.pyplot as plt
+        fig, ax = plt.subplots(figsize=(10, 4))
+        ax.plot(audios[top_indices[0]])
+        ax.set_title(f"Waveform of top matching chunk (#{top_indices[0]})")
+        ax.set_xlabel("Samples")
+        ax.set_ylabel("Amplitude")
+        plt.tight_layout()
+        
+        return result_text, first_chunk_path, fig
+        
+    except Exception as e:
+        return f"Error: {str(e)}", None, None
 
 # Create Gradio interface
 with gr.Blocks(title="AudioRAG Demo") as demo: