Update app.py

app.py

@@ -1,121 +1,106 @@
-import torch
-import torchaudio
 import gradio as gr
-import soundfile as sf
+import torch
+import spaces
+from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline, AutoModelForSeq2SeqLM, AutoTokenizer
+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 transformers import WhisperForCTC, WhisperProcessor, AutoModelForSeq2SeqLM, AutoTokenizer
-from transformers import OpenVoiceV2Processor, OpenVoiceV2
 
 # Load ASR model and processor
-processor_asr = WhisperProcessor.from_pretrained("openai/whisper-large-v3")
-model_asr = WhisperForCTC.from_pretrained("openai/whisper-large-v3")
+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(
+    "automatic-speech-recognition",
+    model=asr_model,
+    tokenizer=asr_processor.tokenizer,
+    feature_extractor=asr_processor.feature_extractor,
+    max_new_tokens=128,
+    chunk_length_s=30,
+    batch_size=16,
+    return_timestamps=True,
+    torch_dtype=torch_dtype,
+    device=device,
+)
 
 # Load text-to-text model and tokenizer
-text_model = AutoModelForSeq2SeqLM.from_pretrained("meta-llama/Meta-Llama-3-8B")
-tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B")
-
-# Load TTS model
-tts_processor = OpenVoiceV2Processor.from_pretrained("myshell-ai/OpenVoiceV2")
-tts_model = OpenVoiceV2.from_pretrained("myshell-ai/OpenVoiceV2")
-
-@spaces.GPU()
-# ASR function
-def transcribe(audio):
-    waveform, sample_rate = torchaudio.load(audio)
-    inputs = processor_asr(waveform, sampling_rate=sample_rate, return_tensors="pt", padding=True)
-    with torch.no_grad():
-        logits = model_asr(inputs.input_values).logits
-    predicted_ids = torch.argmax(logits, dim=-1)
-    transcription = processor_asr.batch_decode(predicted_ids)
-    return transcription[0]
-    
-@spaces.GPU()
-# Text-to-text function
-def generate_response(text):
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
-    outputs = text_model.generate(**inputs)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response
+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)
+
+tts_converter = ToneColorConverter(f'{tts_converter_ckpt}/config.json')
+tts_converter.load_ckpt(f'{tts_converter_ckpt}/checkpoint.pth')
+
+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)
+
+def process_audio(input_audio):
+    # Perform ASR
+    asr_result = asr_pipeline(input_audio)["text"]
+
+    # 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)
+
+    # 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'
     
-@spaces.GPU()
-# TTS function
-def synthesize_speech(text):
-    inputs = tts_processor(text, return_tensors="pt")
-    with torch.no_grad():
-        mel_outputs, mel_outputs_postnet, _, alignments = tts_model.inference(inputs.input_ids)
-        audio = tts_model.infer(mel_outputs_postnet)
-    return audio
+    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)
     
-@spaces.GPU()
-# Real-time processing function
-def real_time_pipeline():
-    # Adjust this part to handle live recording using soundfile and play back using simpleaudio
-    import simpleaudio as sa
-    import tempfile
-    import time
-
-    wake_word = "hello mate"
-    wake_word_detected = False
-
-    print("Listening for wake word...")
-
-    with tempfile.NamedTemporaryFile(delete=False) as tmp_wav_file:
-        tmp_wav_path = tmp_wav_file.name
-
-    try:
-        while True:
-            # Capture audio here (this is a simplified example, you need actual audio capture logic)
-            time.sleep(2)  # Simulate 2 seconds of audio capture
-
-            # Save the captured audio to the temp file for ASR
-            data, sample_rate = sf.read(tmp_wav_path)
-            sf.write(tmp_wav_path, data, sample_rate)
-
-            # Step 1: Transcribe audio to text
-            transcription = transcribe(tmp_wav_path).lower()
-
-            if wake_word in transcription:
-                wake_word_detected = True
-                print("Wake word detected. Processing audio...")
-
-                while wake_word_detected:
-                    # Capture audio here (this is a simplified example, you need actual audio capture logic)
-                    time.sleep(2)  # Simulate 2 seconds of audio capture
-
-                    # Save the captured audio to the temp file for ASR
-                    data, sample_rate = sf.read(tmp_wav_path)
-                    sf.write(tmp_wav_path, data, sample_rate)
-
-                    # Step 1: Transcribe audio to text
-                    transcription = transcribe(tmp_wav_path)
-
-                    # Step 2: Generate response using text-to-text model
-                    response = generate_response(transcription)
-
-                    # Step 3: Synthesize speech from text
-                    synthesized_audio = synthesize_speech(response)
-
-                    # Save the synthesized audio to a temporary file
-                    output_path = "output.wav"
-                    torchaudio.save(output_path, synthesized_audio.squeeze(1), 22050)
-
-                    # Play the synthesized audio using simpleaudio
-                    wave_obj = sa.WaveObject.from_wave_file(output_path)
-                    play_obj = wave_obj.play()
-                    play_obj.wait_done()
-    except KeyboardInterrupt:
-        print("Stopping...")
-
-# Gradio interface
-gr_interface = gr.Interface(
-    fn=real_time_pipeline, 
-    inputs=None, 
-    outputs=None,
-    live=True,
-    title="Real-Time Audio-to-Audio Model",
-    description="ASR + Text-to-Text Model + TTS with Human-like Voice and Emotions"
+    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
+
+iface = gr.Interface(
+    fn=main,
+    inputs=None,
+    outputs=gr.Audio(type="filepath"),
+    live=True
 )
 
-
-iface.launch(inline=False)
+iface.launch()