Update app.py

app.py

@@ -1,41 +1,33 @@
+import os
 import gradio as gr
 import torch
 import numpy as np
-import os
-from transformers import pipeline, AutoProcessor, CsmForConditionalGeneration
-from pyannote.audio import Model, Inference
+from transformers import pipeline
 from dia.model import Dia
 from dac.utils import load_model as load_dac_model
 from accelerate import init_empty_weights, load_checkpoint_and_dispatch
 
-# Access HF_TOKEN from environment variables (Secrets)
-HF_TOKEN = os.environ.get("HF_TOKEN")
+# Retrieve your HF token from the Space secrets
+HF_TOKEN = os.environ["HF_TOKEN"]
 
-# Device mapping for 4× L4 GPU distribution
+# Automatically shard across 4× L4 GPUs
 device_map = "auto"
 
-print("Loading models...")
-
-# Load Descript Audio Codec (RVQ) at startup
-print("Loading RVQ Codec...")
+# 1. Load RVQ codec
 rvq = load_dac_model(tag="latest", model_type="44khz")
 rvq.eval()
 if torch.cuda.is_available():
     rvq = rvq.to("cuda")
 
-# Load segmentation model with authentication
-print("Loading Segmentation Model...")
-seg_model = Model.from_pretrained(
-    "pyannote/segmentation",
-    use_auth_token=HF_TOKEN
+# 2. Load VAD via Hugging Face pipeline (no segmentation mismatch)
+vad_pipe = pipeline(
+    "voice-activity-detection",
+    model="pyannote/voice-activity-detection",
+    use_auth_token=HF_TOKEN,
+    device=0 if torch.cuda.is_available() else -1
 )
-seg_inference = Inference(seg_model, device=0 if torch.cuda.is_available() else -1)
 
-# Use segmentation model for VAD
-vad = seg_inference
-
-# Load Ultravox via generic pipeline (without specifying task)
-print("Loading Ultravox...")
+# 3. Load Ultravox (speech-in → text+LLM)
 ultravox_pipe = pipeline(
     model="fixie-ai/ultravox-v0_4",
     trust_remote_code=True,
@@ -43,8 +35,7 @@ ultravox_pipe = pipeline(
     torch_dtype=torch.float16
 )
 
-# Load Diffusion model
-print("Loading Diffusion Model...")
+# 4. Load diffusion prosody model
 diff_pipe = pipeline(
     "audio-to-audio",
     model="teticio/audio-diffusion-instrumental-hiphop-256",
@@ -53,8 +44,7 @@ diff_pipe = pipeline(
     torch_dtype=torch.float16
 )
 
-# Load Dia TTS with multi-GPU dispatch
-print("Loading Dia TTS...")
+# 5. Load Dia TTS with multi-GPU dispatch
 with init_empty_weights():
     dia = Dia.from_pretrained(
         "nari-labs/Dia-1.6B",
@@ -68,66 +58,40 @@ dia = load_checkpoint_and_dispatch(
     dtype=torch.float16
 )
 
-print("All models loaded successfully!")
-
-# Gradio inference function
+# 6. Inference function
 def process_audio(audio):
-    try:
-        if audio is None:
-            return None, "No audio input provided"
-        
-        sr, array = audio
-        
-        # Ensure audio is numpy array
-        if torch.is_tensor(array):
-            array = array.numpy()
-        
-        # VAD segmentation
-        segments = vad({"waveform": torch.tensor(array).unsqueeze(0), "sample_rate": sr})
-        
-        # RVQ encode/decode
-        audio_tensor = torch.tensor(array).unsqueeze(0)
-        if torch.cuda.is_available():
-            audio_tensor = audio_tensor.to("cuda")
-        codes = rvq.encode(audio_tensor)
-        decoded = rvq.decode(codes)
-        array = decoded.squeeze().cpu().numpy()
-        
-        # Ultravox ASR→LLM
-        ultra_out = ultravox_pipe({"array": array, "sampling_rate": sr})
-        text = ultra_out.get("text", "I understand your audio input.")
-        
-        # Diffusion-based prosody enhancement
-        prosody_audio = diff_pipe({"array": decoded.cpu().numpy(), "sampling_rate": sr})["array"][0]
-        
-        # Dia TTS
-        tts_audio = dia.generate(f"[emotion:neutral] {text}")
-        tts_np = tts_audio.squeeze().cpu().numpy()
-        
-        # Normalize
-        tts_np = tts_np / np.max(np.abs(tts_np)) * 0.95
-        
-        return (sr, tts_np), text
-        
-    except Exception as e:
-        print(f"Error in process_audio: {e}")
-        return None, f"Processing error: {str(e)}"
-
-# Gradio UI
-with gr.Blocks(title="Maya-AI: Supernatural Speech Agent") as demo:
-    gr.Markdown("# Maya-AI: Supernatural Speech Agent")
-    gr.Markdown("Record audio to interact with the AI agent that understands emotions and responds naturally.")
-    
-    with gr.Row():
-        with gr.Column():
-            audio_in = gr.Audio(source="microphone", type="numpy", label="Record Your Voice")
-            btn = gr.Button("Send", variant="primary")
-        
-        with gr.Column():
-            audio_out = gr.Audio(label="AI Response")
-            txt_out = gr.Textbox(label="Transcribed & Generated Text", lines=3)
-    
-    btn.click(fn=process_audio, inputs=audio_in, outputs=[audio_out, txt_out])
+    sr, array = audio
+    array = array.numpy() if torch.is_tensor(array) else array
+
+    # Voice activity detection
+    speech = vad_pipe(array, sampling_rate=sr)[0]["chunks"]
+
+    # RVQ encode/decode
+    x = torch.tensor(array).unsqueeze(0).to("cuda")
+    codes = rvq.encode(x)
+    decoded = rvq.decode(codes).squeeze().cpu().numpy()
+
+    # Ultravox ASR → LLM
+    out = ultravox_pipe({"array": decoded, "sampling_rate": sr})
+    text = out.get("text", "")
+
+    # Diffusion-based prosody
+    pros = diff_pipe({"array": decoded, "sampling_rate": sr})["array"][0]
+
+    # Dia TTS synth
+    tts = dia.generate(f"[emotion:neutral] {text}").squeeze().cpu().numpy()
+    tts = tts / np.max(np.abs(tts)) * 0.95
+
+    return (sr, tts), text
+
+# 7. Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("## Maya-AI: Supernatural Conversational Agent")
+    audio_in = gr.Audio(source="microphone", type="numpy", label="Your Voice")
+    send = gr.Button("Send")
+    audio_out = gr.Audio(label="AI’s Response")
+    text_out = gr.Textbox(label="Generated Text")
+    send.click(process_audio, inputs=audio_in, outputs=[audio_out, text_out])
 
 if __name__ == "__main__":
     demo.launch()