Update app.py

app.py

@@ -1,27 +1,27 @@
-import gradio as gr
+import os
 import torch
-import spaces
-from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline, AutoModelForSeq2SeqLM, AutoTokenizer
+from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
+from transformers import LlamaForCausalLM, LlamaTokenizer
 from datasets import load_dataset
-from openvoice.api import ToneColorConverter
 from openvoice import se_extractor
-from melo.api import TTS
-import pyaudio
-import wave
-import numpy as np
+from openvoice.api import BaseSpeakerTTS, ToneColorConverter
+import gradio as gr
+import spaces
 
-# Load ASR model and processor
+# Device setup
 torch_dtype = torch.float16
 
-asr_model_id = "openai/whisper-large-v3"
-asr_model = AutoModelForSpeechSeq2Seq.from_pretrained(asr_model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True)
-asr_processor = AutoProcessor.from_pretrained(asr_model_id)
-
-asr_pipeline = pipeline(
+# Whisper setup
+whisper_model_id = "openai/whisper-large-v3"
+whisper_model = AutoModelForSpeechSeq2Seq.from_pretrained(
+    whisper_model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
+)
+whisper_processor = AutoProcessor.from_pretrained(whisper_model_id)
+whisper_pipe = pipeline(
     "automatic-speech-recognition",
-    model=asr_model,
-    tokenizer=asr_processor.tokenizer,
-    feature_extractor=asr_processor.feature_extractor,
+    model=whisper_model,
+    tokenizer=whisper_processor.tokenizer,
+    feature_extractor=whisper_processor.feature_extractor,
     max_new_tokens=128,
     chunk_length_s=30,
     batch_size=16,
@@ -30,77 +30,64 @@ asr_pipeline = pipeline(
     device=device,
 )
 
-# Load text-to-text model and tokenizer
-text_model_id = "meta-llama/Meta-Llama-3-8B"
-text_model = AutoModelForSeq2SeqLM.from_pretrained(text_model_id)
-text_tokenizer = AutoTokenizer.from_pretrained(text_model_id)
-
-# Load TTS model and vocoder
-tts_converter_ckpt = 'checkpoints_v2/converter'
-tts_output_dir = 'outputs_v2'
-os.makedirs(tts_output_dir, exist_ok=True)
+# LLaMa3-8B setup
+llama_model_id = "meta-llama/Meta-Llama-3-8B"
+llama_tokenizer = LlamaTokenizer.from_pretrained(llama_model_id)
+llama_model = LlamaForCausalLM.from_pretrained(llama_model_id, torch_dtype=torch_dtype)
 
-tts_converter = ToneColorConverter(f'{tts_converter_ckpt}/config.json')
-tts_converter.load_ckpt(f'{tts_converter_ckpt}/checkpoint.pth')
+# OpenVoiceV2 setup
+ckpt_base = 'checkpoints/base_speakers/EN'
+ckpt_converter = 'checkpoints/converter'
+output_dir = 'outputs'
 
-reference_speaker = 'resources/example_reference.mp3' # This is the voice you want to clone
-target_se, _ = se_extractor.get_se(reference_speaker, tts_converter, vad=False)
+base_speaker_tts = BaseSpeakerTTS(f'{ckpt_base}/config.json',)
+base_speaker_tts.load_ckpt(f'{ckpt_base}/checkpoint.pth')
 
-def process_audio(input_audio):
-    # Perform ASR
-    asr_result = asr_pipeline(input_audio)["text"]
+tone_color_converter = ToneColorConverter(f'{ckpt_converter}/config.json',)
+tone_color_converter.load_ckpt(f'{ckpt_converter}/checkpoint.pth')
 
-    # Perform text-to-text processing
-    input_ids = text_tokenizer(asr_result, return_tensors="pt").input_ids.to(device)
-    generated_ids = text_model.generate(input_ids, max_length=512)
-    response_text = text_tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+os.makedirs(output_dir, exist_ok=True)
+source_se = torch.load(f'{ckpt_base}/en_default_se.pth').to(device)
 
-    # Perform TTS
-    tts_model = TTS(language='EN', device=device)
-    speaker_id = list(tts_model.hps.data.spk2id.values())[0]
-    tts_model.tts_to_file(response_text, speaker_id, f'{tts_output_dir}/tmp.wav')
-    save_path = f'{tts_output_dir}/output_v2.wav'
+def process_audio(input_audio):
+    # ASR with Whisper
+    whisper_result = whisper_pipe(input_audio)["text"]
+    
+    # Text generation with LLaMa
+    inputs = llama_tokenizer(whisper_result, return_tensors="pt").to(device)
+    outputs = llama_model.generate(**inputs, max_new_tokens=50)
+    generated_text = llama_tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    # TTS with OpenVoiceV2
+    reference_speaker = 'resources/example_reference.mp3'
+    target_se, _ = se_extractor.get_se(reference_speaker, tone_color_converter, target_dir='processed', vad=True)
+    save_path = f'{output_dir}/output_en_default.wav'
+    src_path = f'{output_dir}/tmp.wav'
+    
+    base_speaker_tts.tts(generated_text, src_path, speaker='default', language='English', speed=1.0)
+    tone_color_converter.convert(
+        audio_src_path=src_path, 
+        src_se=source_se, 
+        tgt_se=target_se, 
+        output_path=save_path,
+        message="@MyShell"
+    )
     
-    source_se = torch.load(f'checkpoints_v2/base_speakers/ses/english-american.pth', map_location=device)
-    tts_converter.convert(audio_src_path=f'{tts_output_dir}/tmp.wav', src_se=source_se, tgt_se=target_se, output_path=save_path, message="@MyShell")
-
     return save_path
-
-# Real-time audio processing
-
-def real_time_audio_processing():
-    p = pyaudio.PyAudio()
-    stream = p.open(format=pyaudio.paInt16, channels=1, rate=16000, input=True, frames_per_buffer=1024)
     
-    frames = []
-    print("Listening...")
-
-    while True:
-        data = stream.read(1024)
-        frames.append(data)
-        audio_data = np.frombuffer(data, dtype=np.int16)
-        if np.max(audio_data) > 3000:  # Simple VAD threshold
-            wf = wave.open("input_audio.wav", 'wb')
-            wf.setnchannels(1)
-            wf.setsampwidth(p.get_sample_size(pyaudio.paInt16))
-            wf.setframerate(16000)
-            wf.writeframes(b''.join(frames))
-            wf.close()
-            return "input_audio.wav"
-
-# Gradio Interface
-@spaces.GPU(duration=300)
-def main():
-    input_audio_path = real_time_audio_processing()
-    if input_audio_path:
-        output_audio_path = process_audio(input_audio_path)
-        return output_audio_path
+@spaces.GPU()
+def real_time_processing(input_audio):
+    return process_audio(input_audio)
 
+# Gradio interface
 iface = gr.Interface(
-    fn=main,
-    inputs=None,
-    outputs=gr.Audio(type="filepath"),
-    live=True
+    fn=real_time_processing,
+    inputs=gr.Audio(source="microphone", type="filepath"),
+    outputs=gr.Audio(type="file"),
+    live=True,
+    title="ASR + Text-to-Text + TTS with Whisper, LLaMa3-8B, and OpenVoiceV2",
+    description="Real-time processing using Whisper for ASR, LLaMa3-8B for text generation, and OpenVoiceV2 for TTS."
 )
 
-iface.launch()
+if __name__ == "__main__":
+    iface.launch()