Update app.py

app.py

@@ -3,21 +3,23 @@ import gradio as gr
 import torch
 import numpy as np
 from transformers import pipeline
+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
 from accelerate import init_empty_weights, load_checkpoint_and_dispatch
 
-# Environment token
+# Retrieve HF token from Secrets
 HF_TOKEN = os.environ["HF_TOKEN"]
 
-# Shard large models across 4× L4 GPUs
+# Automatic multi-GPU sharding across 4× L4 GPUs
 device_map = "auto"
 
-# 1. RVQ codec (Descript Audio Codec)
+# 1. Descript Audio Codec (RVQ)
 rvq = load_dac_model(tag="latest", model_type="44khz")
 rvq.eval()
-if torch.cuda.is_available(): rvq = rvq.to("cuda")
+if torch.cuda.is_available():
+    rvq = rvq.to("cuda")
 
 # 2. Voice Activity Detection via Pyannote
 vad_pipe = PyannotePipeline.from_pretrained(
@@ -25,7 +27,7 @@ vad_pipe = PyannotePipeline.from_pretrained(
     use_auth_token=HF_TOKEN
 )
 
-# 3. Ultravox pipeline (speech → text + LLM)
+# 3. Ultravox ASR+LLM
 ultravox_pipe = pipeline(
     model="fixie-ai/ultravox-v0_4",
     trust_remote_code=True,
@@ -33,16 +35,13 @@ ultravox_pipe = pipeline(
     torch_dtype=torch.float16
 )
 
-# 4. Diffusion-based prosody model
-diff_pipe = pipeline(
-    "audio-to-audio",
-    model="teticio/audio-diffusion-instrumental-hiphop-256",
-    trust_remote_code=True,
-    device_map=device_map,
+# 4. Audio Diffusion (direct load via Diffusers)
+diff_pipe = DiffusionPipeline.from_pretrained(
+    "teticio/audio-diffusion-instrumental-hiphop-256",
     torch_dtype=torch.float16
-)
+).to("cuda")
 
-# 5. Dia TTS loaded with multi-GPU dispatch
+# 5. Dia TTS (multi-GPU dispatch)
 with init_empty_weights():
     dia = Dia.from_pretrained(
         "nari-labs/Dia-1.6B",
@@ -56,42 +55,40 @@ dia = load_checkpoint_and_dispatch(
     dtype=torch.float16
 )
 
-# Inference function
+# 6. Inference Function
 def process_audio(audio):
     sr, array = audio
-    # Ensure numpy
-    if torch.is_tensor(array): array = array.numpy()
+    array = array.numpy() if torch.is_tensor(array) else array
 
-    # VAD: extract speech regions
-    chunks = vad_pipe(array, sampling_rate=sr)
+    # VAD
+    _ = vad_pipe(array, sampling_rate=sr)
 
     # RVQ encode/decode
-    x = torch.tensor(array).unsqueeze(0).to("cuda")
-    codes = rvq.encode(x)
+    tensor = torch.tensor(array).unsqueeze(0).to("cuda")
+    codes = rvq.encode(tensor)
     decoded = rvq.decode(codes).squeeze().cpu().numpy()
 
-    # Ultravox ASR + LLM
-    out = ultravox_pipe({"array": decoded, "sampling_rate": sr})
-    text = out.get("text", "")
+    # Ultravox inference
+    ultra_out = ultravox_pipe({"array": decoded, "sampling_rate": sr})
+    text = ultra_out.get("text", "")
 
-    # Diffusion prosody enhancement
-    pros_audio = diff_pipe({"array": decoded, "sampling_rate": sr})["array"][0]
+    # Diffusion enhancement
+    pros = diff_pipe(raw_audio=decoded)["audios"][0]
 
-    # 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
+    # Dia TTS
+    tts = dia.generate(f"[emotion:neutral] {text}").squeeze().cpu().numpy()
+    tts = tts / np.max(np.abs(tts)) * 0.95
 
-    return (sr, tts_np), text
+    return (sr, tts), text
 
-# Gradio UI
-with gr.Blocks(title="Maya AI 📈", theme=None) as demo:
+# 7. Gradio UI
+with gr.Blocks(title="Maya AI 📈") as demo:
     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’s Response")
-    text_out = gr.Textbox(label="Generated Text")
-    send_btn.click(process_audio, inputs=audio_in, outputs=[audio_out, text_out])
+    audio_input = gr.Audio(source="microphone", type="numpy", label="Your Voice")
+    send_button = gr.Button("Send")
+    audio_output = gr.Audio(label="AI’s Response")
+    text_output = gr.Textbox(label="Generated Text")
+    send_button.click(process_audio, inputs=audio_input, outputs=[audio_output, text_output])
 
 if __name__ == "__main__":
     demo.launch()