Update app.py

app.py

@@ -1,142 +1,82 @@
 import gradio as gr
-from transformers import AutoProcessor, AutoModelForCausalLM, pipeline
 import torch
 import numpy as np
+from transformers import pipeline, AutoProcessor, CsmForConditionalGeneration
 from pyannote.audio import Pipeline as VAD
-import dac
+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
 
-# Load models with proper error handling
-def load_models():
-    try:
-        # Ultravox via transformers (no separate package needed)
-        ultra_proc = AutoProcessor.from_pretrained("fixie-ai/ultravox-v0_4", trust_remote_code=True)
-        ultra_model = AutoModelForCausalLM.from_pretrained("fixie-ai/ultravox-v0_4", device_map="auto", torch_dtype=torch.float16, trust_remote_code=True)
-        
-        # Speech emotion recognition via transformers pipeline
-        emotion_pipeline = pipeline("audio-classification", model="ehcalabres/wav2vec2-lg-xlsr-en-speech-emotion-recognition", device=0 if torch.cuda.is_available() else -1)
-        
-        # Audio diffusion (using transformers instead of torch.hub for HF compatibility)
-        from diffusers import DiffusionPipeline
-        diff_pipe = DiffusionPipeline.from_pretrained("teticio/audio-diffusion-instrumental-hiphop-256")
-        
-        # Descript Audio Codec
-        from dac.utils import load_model as load_dac_model
-        rvq = load_dac_model(tag="latest", model_type="44khz")
-        rvq.eval()
-        if torch.cuda.is_available():
-            rvq = rvq.to("cuda")
-        
-        # VAD
-        vad = VAD.from_pretrained("pyannote/voice-activity-detection")
-        
-        # Dia TTS
-        from dia.model import Dia
-        dia = Dia.from_pretrained("nari-labs/Dia-1.6B", compute_dtype="float16")
-        
-        return ultra_proc, ultra_model, emotion_pipeline, diff_pipe, rvq, vad, dia
-        
-    except Exception as e:
-        print(f"Error loading models: {e}")
-        return None, None, None, None, None, None, None
+# 2.1: Initialize Accelerator for 4×L4 GPU distribution
+from accelerate import init_empty_weights, load_checkpoint_and_dispatch
+device_map = "auto"  # accelerate automatically shards across 4 L4 GPUs[2]
 
-# Initialize models
-ultra_proc, ultra_model, emotion_pipeline, diff_pipe, rvq, vad, dia = load_models()
+# 2.2: Load Descript Audio Codec (RVQ) at startup
+rvq = load_dac_model(tag="latest", model_type="44khz")
+rvq.eval()
+if torch.cuda.is_available():
+    rvq = rvq.to("cuda")
 
+# 2.3: Load VAD pipeline
+vad = VAD.from_pretrained("pyannote/voice-activity-detection")
+
+# 2.4: Load Ultravox via audio-text-to-text pipeline
+ultravox_pipe = pipeline(
+    "audio-text-to-text",
+    model="fixie-ai/ultravox-v0_4",
+    trust_remote_code=True,
+    device_map=device_map,
+    torch_dtype=torch.float16
+)
+
+# 2.5: Load Diffusion model
+diff_pipe = pipeline(
+    "audio-to-audio",
+    model="teticio/audio-diffusion-instrumental-hiphop-256",
+    trust_remote_code=True,
+    device_map=device_map,
+    torch_dtype=torch.float16
+)
+
+# 2.6: Load Dia TTS with multi-GPU dispatch
+with init_empty_weights():
+    dia = Dia.from_pretrained("nari-labs/Dia-1.6B", torch_dtype=torch.float16, trust_remote_code=True)
+dia = load_checkpoint_and_dispatch(
+    dia, "nari-labs/Dia-1.6B", device_map=device_map, dtype=torch.float16
+)
+
+# 2.7: Gradio inference function
 def process_audio(audio):
-    try:
-        if audio is None:
-            return None, "No audio input provided"
-        
-        # Convert audio to proper format
-        audio_array = audio[1] if isinstance(audio, tuple) else audio["array"]
-        sample_rate = audio[0] if isinstance(audio, tuple) else audio["sampling_rate"]
-        
-        # Ensure audio is numpy array
-        if torch.is_tensor(audio_array):
-            audio_array = audio_array.numpy()
-        
-        # VAD processing
-        if vad is not None:
-            speech_segments = vad({"waveform": torch.from_numpy(audio_array).unsqueeze(0), "sample_rate": sample_rate})
-        
-        # Emotion recognition
-        emotion_result = "neutral"
-        if emotion_pipeline is not None:
-            try:
-                emotion_pred = emotion_pipeline(audio_array, sampling_rate=sample_rate)
-                emotion_result = emotion_pred[0]["label"] if emotion_pred else "neutral"
-            except:
-                emotion_result = "neutral"
-        
-        # RVQ encode/decode
-        if rvq is not None:
-            try:
-                audio_tensor = torch.from_numpy(audio_array).float().unsqueeze(0)
-                if torch.cuda.is_available():
-                    audio_tensor = audio_tensor.to("cuda")
-                encoded = rvq.encode(audio_tensor)
-                decoded_audio = rvq.decode(encoded)
-                if torch.cuda.is_available():
-                    decoded_audio = decoded_audio.cpu()
-                audio_array = decoded_audio.squeeze().numpy()
-            except Exception as e:
-                print(f"RVQ processing error: {e}")
-        
-        # Ultravox generation
-        response_text = "I understand your audio input."
-        if ultra_proc is not None and ultra_model is not None:
-            try:
-                inputs = ultra_proc(audio_array, sampling_rate=sample_rate, return_tensors="pt")
-                if torch.cuda.is_available():
-                    inputs = {k: v.to("cuda") for k, v in inputs.items()}
-                
-                with torch.no_grad():
-                    outputs = ultra_model.generate(**inputs, max_new_tokens=50)
-                    response_text = ultra_proc.decode(outputs[0], skip_special_tokens=True)
-            except Exception as e:
-                print(f"Ultravox generation error: {e}")
-                response_text = f"Detected emotion: {emotion_result}"
-        
-        # TTS generation
-        output_audio = None
-        if dia is not None:
-            try:
-                tts_text = f"[emotion:{emotion_result}] {response_text}"
-                output_audio = dia.generate(tts_text)
-                if torch.is_tensor(output_audio):
-                    output_audio = output_audio.cpu().numpy()
-                # Normalize audio
-                if output_audio is not None:
-                    output_audio = output_audio / np.max(np.abs(output_audio)) * 0.95
-            except Exception as e:
-                print(f"TTS generation error: {e}")
-        
-        return (sample_rate, output_audio) if output_audio is not None else None, response_text
-        
-    except Exception as e:
-        return None, f"Processing error: {str(e)}"
+    sr, array = audio["sampling_rate"], audio["array"]
+    # VAD segmentation
+    speech_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["text"]
+    # 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
 
-# Create Gradio interface
-with gr.Blocks(title="Supernatural Speech AI") as demo:
-    gr.Markdown("# Supernatural Speech AI Agent")
-    gr.Markdown("Record audio to interact with the AI agent that understands emotions and responds naturally.")
-    
-    with gr.Row():
-        with gr.Column():
-            audio_input = gr.Audio(sources=["microphone"], type="numpy", label="Record Audio")
-            process_btn = gr.Button("Process Audio", variant="primary")
-        
-        with gr.Column():
-            audio_output = gr.Audio(label="AI Response")
-            text_output = gr.Textbox(label="Response Text", lines=3)
-    
-    conversation_history = gr.State([])
-    
-    process_btn.click(
-        fn=process_audio,
-        inputs=[audio_input],
-        outputs=[audio_output, text_output]
-    )
+# 2.8: Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("## Supernatural Speech AI Agent")
+    audio_in = gr.Audio(source="microphone", type="numpy", label="Record Your Voice")
+    btn = gr.Button("Send")
+    audio_out = gr.Audio(label="AI Response")
+    txt_out = gr.Textbox(label="Transcribed & Generated Text")
+    btn.click(fn=process_audio, inputs=audio_in, outputs=[audio_out, txt_out])
 
 if __name__ == "__main__":
-    demo.queue(concurrency_limit=20, max_size=50).launch()
+    demo.launch()