Update app.py

app.py

@@ -1,46 +1,142 @@
 import gradio as gr
-from transformers import AutoProcessor, CsmForConditionalGeneration
-from dia.model import Dia
+from transformers import AutoProcessor, AutoModelForCausalLM, pipeline
+import torch
+import numpy as np
 from pyannote.audio import Pipeline as VAD
-import torch, numpy as np
+import dac
 
-# Load models
-ultra_proc = AutoProcessor.from_pretrained("fixie-ai/ultravox-v0_4")
-ultra_model = CsmForConditionalGeneration.from_pretrained("fixie-ai/ultravox-v0_4", device_map="auto", torch_dtype=torch.float16)
-ser = AutoProcessor.from_pretrained("r-f/wav2vec-english-speech-emotion-recognition")
-ser_model = torch.hub.load("jonatasgrosman/wav2vec2-large-xlsr-53-english", "wav2vec2_large_xlsr", pretrained=True).to("cuda")
-diff_pipe = torch.hub.load("teticio/audio-diffusion-instrumental-hiphop-256", "audio_diffusion").to("cuda")
-rvq = torch.hub.load("ibm/DAC.speech.v1.0", "DAC_speech_v1_0").to("cuda")
-vad = VAD.from_pretrained("pyannote/voice-activity-detection")
-dia = Dia.from_pretrained("nari-labs/Dia-1.6B", compute_dtype="float16")
+# 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
 
-def process(audio):
-    # VAD
-    speech = vad({"waveform": audio["array"], "sample_rate": audio["sampling_rate"]})
-    # RVQ encode/decode
-    codes = rvq.encode(audio["array"])
-    dec_audio = rvq.decode(codes)
-    # Emotion
-    emo_inputs = ser(audio["array"], sampling_rate=audio["sampling_rate"], return_tensors="pt")
-    emotion = ser_model(**emo_inputs).logits.argmax(-1).item()
-    # Ultravox generation
-    inputs = ultra_proc(audio["array"], sampling_rate=audio["sampling_rate"], return_tensors="pt").to("cuda")
-    speech_out = ultra_model.generate(**inputs, output_audio=True)
-    # Diffuse and clone voice
-    audio_diff = diff_pipe(speech_out.audio).audios[0]
-    # TTS
-    text = f"[S1][emotion={emotion}]" + " ".join(["..."]) # placeholder
-    dia_audio = dia.generate(text)
-    # Normalize
-    dia_audio = dia_audio / np.max(np.abs(dia_audio)) * 0.95
-    return 44100, dia_audio
+# Initialize models
+ultra_proc, ultra_model, emotion_pipeline, diff_pipe, rvq, vad, dia = load_models()
 
-with gr.Blocks() as demo:
-    state = gr.State([])
-    audio_in = gr.Audio(source="microphone", type="numpy")
-    chat = gr.Chatbot()
-    record = gr.Button("Record")
-    record.click(process, inputs=audio_in, outputs=[audio_in]).then(
-        lambda a: chat.update(value=[("User", ""), ("AI", "")]),
+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)}"
+
+# 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(source="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]
     )
+
+if __name__ == "__main__":
     demo.queue(concurrency_limit=20, max_size=50).launch()