Update app.py

app.py

@@ -2,138 +2,86 @@ import os
 import gradio as gr
 import torch
 import numpy as np
-from transformers import pipeline
+from transformers import pipeline, AutoModel
 from diffusers import DiffusionPipeline
 from pyannote.audio import Pipeline as PyannotePipeline
 from dia.model import Dia
 from dac.utils import load_model as load_dac_model
 
-# Environment token from HF Secrets
-HF_TOKEN = os.environ["HF_TOKEN"]
+# 1. Retrieve HF token and set device mapping
+HF_TOKEN   = os.environ["HF_TOKEN"]
+device_map = "auto"  # auto-shard models across 4×L4 GPUs
 
-print("Loading models...")
-
-# 1. Load RVQ Codec (Descript Audio Codec)
 print("Loading RVQ Codec...")
 rvq = load_dac_model(tag="latest", model_type="44khz")
 rvq.eval()
 if torch.cuda.is_available():
     rvq = rvq.to("cuda")
 
-# 2. Load Voice Activity Detection
-print("Loading VAD...")
+print("Loading VAD pipeline...")
 vad_pipe = PyannotePipeline.from_pretrained(
     "pyannote/voice-activity-detection",
     use_auth_token=HF_TOKEN
 )
 
-# 3. Load Ultravox ASR+LLM
-print("Loading Ultravox...")
+print("Loading Ultravox pipeline...")
 ultravox_pipe = pipeline(
     model="fixie-ai/ultravox-v0_4",
     trust_remote_code=True,
-    device_map="auto",
+    device_map=device_map,
     torch_dtype=torch.float16
 )
 
-# 4. Load Audio Diffusion Model
-print("Loading Audio Diffusion...")
+print("Loading Audio Diffusion model...")
 diff_pipe = DiffusionPipeline.from_pretrained(
     "teticio/audio-diffusion-instrumental-hiphop-256",
     torch_dtype=torch.float16
 ).to("cuda")
 
-# 5. Load Dia TTS Model (WITHOUT meta tensor approach)
-print("Loading Dia TTS...")
-# Direct loading without init_empty_weights to avoid meta tensor issues
-dia = Dia.from_pretrained("nari-labs/Dia-1.6B")
+print("Loading Dia TTS (sharded across GPUs)...")
+dia = Dia.from_pretrained(
+    "nari-labs/Dia-1.6B",
+    device_map=device_map,
+    torch_dtype=torch.float16,
+    trust_remote_code=True
+)
 
 print("All models loaded successfully!")
 
-# Audio processing function
 def process_audio(audio):
-    try:
-        if audio is None:
-            return None, "No audio input provided"
-        
-        sr, array = audio
-        
-        # Ensure numpy array
-        if torch.is_tensor(array):
-            array = array.numpy()
-        
-        # Voice Activity Detection
-        vad_result = vad_pipe({"waveform": torch.tensor(array).unsqueeze(0), "sample_rate": sr})
-        
-        # RVQ encode/decode for audio compression
-        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).squeeze().cpu().numpy()
-        
-        # Ultravox ASR + LLM processing
-        ultra_out = ultravox_pipe({"array": decoded, "sampling_rate": sr})
-        text = ultra_out.get("text", "I understand your audio input.")
-        
-        # Audio diffusion for prosody enhancement
-        try:
-            prosody_result = diff_pipe(raw_audio=decoded)
-            if "audios" in prosody_result:
-                prosody_audio = prosody_result["audios"][0]
-            else:
-                prosody_audio = decoded
-        except Exception as e:
-            print(f"Diffusion processing error: {e}")
-            prosody_audio = decoded
-        
-        # Dia TTS generation
-        tts_output = dia.generate(f"[emotion:neutral] {text}")
-        
-        # Convert to numpy and normalize
-        if torch.is_tensor(tts_output):
-            tts_np = tts_output.squeeze().cpu().numpy()
-        else:
-            tts_np = tts_output
-        
-        # Normalize audio output
-        if len(tts_np) > 0:
-            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 Interface
+    sr, array = audio
+    array = array.numpy() if torch.is_tensor(array) else array
+
+    # 1. Voice activity detection
+    vad_pipe({"waveform": torch.tensor(array).unsqueeze(0), "sample_rate": sr})
+
+    # 2. RVQ encode/decode
+    x = torch.tensor(array).unsqueeze(0).to("cuda")
+    codes   = rvq.encode(x)
+    decoded = rvq.decode(codes).squeeze().cpu().numpy()
+
+    # 3. Ultravox ASR → text
+    out = ultravox_pipe({"array": decoded, "sampling_rate": sr})
+    text = out.get("text", "")
+
+    # 4. Prosody diffusion
+    pros = diff_pipe(raw_audio=decoded)["audios"][0]
+
+    # 5. Dia TTS synthesis
+    tts = dia.generate(f"[emotion:neutral] {text}")
+    tts_np = tts.squeeze().cpu().numpy()
+    tts_np = tts_np / np.max(np.abs(tts_np)) * 0.95 if tts_np.size else tts_np
+
+    return (sr, tts_np), text
+
+# Gradio UI
 with gr.Blocks(title="Maya AI 📈") as demo:
-    gr.Markdown("# Maya-AI: Supernatural Conversational 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(
-                sources=["microphone"], 
-                type="numpy", 
-                label="Record Your Voice"
-            )
-            send_btn = gr.Button("Send", variant="primary")
-        
-        with gr.Column():
-            audio_out = gr.Audio(label="AI Response")
-            text_out = gr.Textbox(
-                label="Generated Text", 
-                lines=3,
-                placeholder="AI response will appear here..."
-            )
-    
-    # Event handler
-    send_btn.click(
-        fn=process_audio,
-        inputs=audio_in,
-        outputs=[audio_out, text_out]
-    )
+    gr.Markdown("## Maya-AI Supernatural Conversational Agent")
+    audio_in  = gr.Audio(source="microphone", type="numpy", label="Your Voice")
+    send_btn  = gr.Button("Send")
+    audio_out = gr.Audio(label="AI Response")
+    text_out  = gr.Textbox(label="Generated Text")
+    send_btn.click(process_audio, inputs=audio_in, outputs=[audio_out, text_out])
 
 if __name__ == "__main__":
     demo.launch()