Update app.py

app.py

@@ -1,28 +1,64 @@
 import os
+import io
 import torch
 import whisper
 import streamlit as st
 from groq import Groq
-from TTS.api import TTS
 from dotenv import load_dotenv
 from tempfile import NamedTemporaryFile
 from streamlit_webrtc import webrtc_streamer, WebRtcMode, ClientSettings
 import av
 import numpy as np
-import scipy.io.wavfile
-import scipy.sparse
-
-from huggingface_hub import HfApi
-
-# will use api to restart space on a unrecoverable error
-api = HfApi(token=HF_TOKEN)
+import uuid
+import time
 
-# Load API key from Hugging Face
+# Load environment variables
 load_dotenv()
 API_KEY = os.getenv("GROQ_API_KEY")
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+# By using XTTS you agree to CPML license
+os.environ["COQUI_TOS_AGREED"] = "1"
+
+# For proper language detection
+import langid
+
+# Import TTS components
+from TTS.api import TTS
+from TTS.tts.configs.xtts_config import XttsConfig
+from TTS.tts.models.xtts import Xtts
+from TTS.utils.generic_utils import get_user_data_dir
+
+# Download and configure XTTS model
+print("Downloading Coqui XTTS V2 if not already downloaded")
+from TTS.utils.manage import ModelManager
+
+model_name = "tts_models/multilingual/multi-dataset/xtts_v2"
+ModelManager().download_model(model_name)
+model_path = os.path.join(get_user_data_dir("tts"), model_name.replace("/", "--"))
+print("XTTS downloaded")
+
+config = XttsConfig()
+config.load_json(os.path.join(model_path, "config.json"))
+
+model = Xtts.init_from_config(config)
+model.load_checkpoint(
+    config,
+    checkpoint_path=os.path.join(model_path, "model.pth"),
+    vocab_path=os.path.join(model_path, "vocab.json"),
+    eval=True,
+    use_deepspeed=True,
+)
+if torch.cuda.is_available():
+    model.cuda()
+
+supported_languages = config.languages
 
 # LLM Response Function
 def get_llm_response(api_key, user_input):
+    if not api_key:
+        return "API key not found. Please set the GROQ_API_KEY environment variable."
+    
     client = Groq(api_key=api_key)
     prompt = (
         "IMPORTANT: You are an AI assistant that MUST provide responses in 25 words or less.\n"
@@ -37,141 +73,230 @@ def get_llm_response(api_key, user_input):
         "Your response will be converted to speech. Maximum 25 words."
     )
     
-    chat_completion = client.chat.completions.create(
-        messages=[
-            {"role": "system", "content": prompt},
-            {"role": "user", "content": user_input}
-        ],
-        model="llama3-8b-8192",
-        temperature=0.5,
-        top_p=1,
-        stream=False,
-    )
-    return chat_completion.choices[0].message.content
+    try:
+        chat_completion = client.chat.completions.create(
+            messages=[
+                {"role": "system", "content": prompt},
+                {"role": "user", "content": user_input}
+            ],
+            model="llama3-8b-8192",
+            temperature=0.5,
+            top_p=1,
+            stream=False,
+        )
+        return chat_completion.choices[0].message.content
+    except Exception as e:
+        return f"Error with LLM: {str(e)}"
 
 # Transcribe Audio
 def transcribe_audio(audio_path, model_size="base"):
-    model = whisper.load_model(model_size)
-    result = model.transcribe(audio_path)
-    return result["text"]
+    try:
+        model = whisper.load_model(model_size)
+        result = model.transcribe(audio_path)
+        return result["text"]
+    except Exception as e:
+        return f"Error transcribing audio: {str(e)}"
 
-# Generate Speech
-def generate_speech(text, output_file, speaker_wav, language="en", use_gpu=True):
+# Generate Speech using the configured XTTS model
+def generate_speech(text, output_file, speaker_wav, language="en"):
     if not os.path.exists(speaker_wav):
         raise FileNotFoundError("Reference audio file not found. Please upload or record a valid audio.")
     
-    tts = TTS("tts_models/multilingual/multi-dataset/xtts_v2", gpu=use_gpu)
-    tts.tts_to_file(
-        text=text,
-        file_path=output_file,
-        speaker_wav=speaker_wav,
-        language=language,
-    )
+    if language not in supported_languages:
+        st.warning(f"Language {language} is not supported. Defaulting to English.")
+        language = "en"
+    
+    # Detect language if text is long enough
+    detected_lang = langid.classify(text)[0]
+    if detected_lang == "zh":
+        detected_lang = "zh-cn"
+    
+    # Use the configured model directly
+    try:
+        t_latent = time.time()
+        gpt_cond_latent, speaker_embedding = model.get_conditioning_latents(
+            audio_path=speaker_wav, 
+            gpt_cond_len=30, 
+            gpt_cond_chunk_len=4,
+            max_ref_length=60
+        )
+        
+        out = model.inference(
+            text,
+            language,
+            gpt_cond_latent,
+            speaker_embedding,
+            repetition_penalty=5.0,
+            temperature=0.75,
+        )
+        
+        # Save the audio to file
+        torch.tensor(out["wav"]).unsqueeze(0).cpu().numpy()
+        import soundfile as sf
+        sf.write(output_file, out["wav"], 24000, 'PCM_24')
+        
+        return True, "Speech generated successfully"
+    except Exception as e:
+        return False, f"Error generating speech: {str(e)}"
 
-# Audio Frame Processing
+# Audio Frame Processing for WebRTC
 class AudioProcessor:
     def __init__(self):
         self.audio_frames = []
+        self.sample_rate = 24000  # XTTS expects 24kHz
 
     def recv(self, frame):
-        self.audio_frames.append(frame.to_ndarray().tobytes())
+        sound = frame.to_ndarray()
+        self.audio_frames.append(sound)
         return frame
 
     def save_audio(self, file_path):
-        with open(file_path, "wb") as f:
-            for frame in self.audio_frames:
-                f.write(frame)
+        if not self.audio_frames:
+            return None
+            
+        # Concatenate audio frames
+        concat_audio = np.concatenate(self.audio_frames, axis=0)
+        
+        # Save as WAV file
+        import soundfile as sf
+        sf.write(file_path, concat_audio, self.sample_rate)
+        
         return file_path
 
 # Streamlit App
 def main():
     st.set_page_config(page_title="Vocal AI", layout="wide")
-    st.sidebar.title("Vocal-AI Settings")
-
+    
+    st.title("VocaL AI - Voice Cloning Assistant")
+    st.write("Clone your voice and interact with an AI assistant that responds in your voice!")
+    
+    st.sidebar.title("Settings")
+    
+    # Language selection
+    language = st.sidebar.selectbox(
+        "Output Language", 
+        supported_languages,
+        index=supported_languages.index("en") if "en" in supported_languages else 0
+    )
+    
+    # TOS agreement
+    agree_tos = st.sidebar.checkbox("I agree to the Coqui Public Model License (CPML)", value=False)
+    
     # User option for reference audio (Record or Upload)
     ref_audio_choice = st.sidebar.radio("Reference Audio", ("Upload", "Record"))
 
     ref_audio_path = None
     reference_audio_processor = None
 
-    if ref_audio_choice == "Upload":
-        reference_audio = st.sidebar.file_uploader("Upload Reference Audio", type=["wav", "mp3", "ogg"])
-        if reference_audio:
-            with NamedTemporaryFile(delete=False, suffix=".wav") as temp_ref_audio:
-                temp_ref_audio.write(reference_audio.read())
-                ref_audio_path = temp_ref_audio.name
-    else:
-        st.sidebar.write("Record your reference audio:")
-        reference_audio_processor = AudioProcessor()
-        webrtc_streamer(
-            key="ref_audio",
-            mode=WebRtcMode.SENDRECV,
-            client_settings=ClientSettings(rtc_configuration={"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]}),
-            audio_receiver_size=1024,
-            video_processor_factory=None,
-            audio_processor_factory=lambda: reference_audio_processor,
-        )
+    col1, col2 = st.columns(2)
+    
+    with col1:
+        st.header("Step 1: Provide Reference Voice")
+        if ref_audio_choice == "Upload":
+            reference_audio = st.file_uploader("Upload Reference Audio", type=["wav", "mp3", "ogg"])
+            if reference_audio:
+                with NamedTemporaryFile(delete=False, suffix=".wav") as temp_ref_audio:
+                    temp_ref_audio.write(reference_audio.read())
+                    ref_audio_path = temp_ref_audio.name
+                st.audio(ref_audio_path)
+        else:
+            st.write("Record your reference voice:")
+            reference_audio_processor = AudioProcessor()
+            webrtc_ctx = webrtc_streamer(
+                key="ref_audio",
+                mode=WebRtcMode.SENDRECV,
+                client_settings=ClientSettings(
+                    rtc_configuration={"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]},
+                    media_stream_constraints={"audio": True, "video": False},
+                ),
+                audio_receiver_size=1024,
+                video_processor_factory=None,
+                audio_processor_factory=lambda: reference_audio_processor,
+            )
+            
+            if webrtc_ctx.state.playing and reference_audio_processor is not None:
+                st.info("Recording... Speak into your microphone.")
+            
+            if st.button("Save Reference Audio"):
+                if reference_audio_processor and reference_audio_processor.audio_frames:
+                    with NamedTemporaryFile(delete=False, suffix=".wav") as temp_ref_audio:
+                        reference_audio_processor.save_audio(temp_ref_audio.name)
+                        ref_audio_path = temp_ref_audio.name
+                    st.success("Reference audio saved!")
+                    st.audio(ref_audio_path)
+                else:
+                    st.error("No audio recorded. Please speak into your microphone.")
 
-    st.title("Welcome to VocaL AI")
-    st.write("### How to Use")
-    st.write("1. Upload or record a reference audio file.")
-    st.write("2. Choose between text or audio input.")
-    st.write("3. If audio input is selected, record and submit your audio.")
-    st.write("4. Click 'Generate Speech' to hear the AI response in your cloned voice.")
-
-    # User Input (Text or Audio)
-    input_type = st.radio("Choose Input Type", ("Text", "Audio"))
-    user_input = None
-    user_audio_processor = None
-
-    if input_type == "Text":
-        user_input = st.text_area("Enter your text here")
-    else:
-        st.write("Record your voice:")
-        user_audio_processor = AudioProcessor()
-        webrtc_streamer(
-            key="user_audio",
-            mode=WebRtcMode.SENDRECV,
-            client_settings=ClientSettings(rtc_configuration={"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]}),
-            audio_receiver_size=1024,
-            video_processor_factory=None,
-            audio_processor_factory=lambda: user_audio_processor,
-        )
+    with col2:
+        st.header("Step 2: Ask Something")
+        # User Input (Text or Audio)
+        input_type = st.radio("Choose Input Type", ("Text", "Audio"))
+        user_input = None
+        user_audio_processor = None
 
-    if st.button("Generate Speech"):
-        # Handle Reference Audio
-        if reference_audio_processor:
-            with NamedTemporaryFile(delete=False, suffix=".wav") as temp_ref_audio:
-                reference_audio_processor.save_audio(temp_ref_audio.name)
-                ref_audio_path = temp_ref_audio.name
+        if input_type == "Text":
+            user_input = st.text_area("Enter your question or prompt here")
+        else:
+            st.write("Record your question:")
+            user_audio_processor = AudioProcessor()
+            webrtc_ctx_user = webrtc_streamer(
+                key="user_audio",
+                mode=WebRtcMode.SENDRECV,
+                client_settings=ClientSettings(
+                    rtc_configuration={"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]},
+                    media_stream_constraints={"audio": True, "video": False},
+                ),
+                audio_receiver_size=1024,
+                video_processor_factory=None,
+                audio_processor_factory=lambda: user_audio_processor,
+            )
+            
+            if webrtc_ctx_user.state.playing and user_audio_processor is not None:
+                st.info("Recording... Ask your question")
+            
+            if st.button("Process Recording"):
+                if user_audio_processor and user_audio_processor.audio_frames:
+                    with NamedTemporaryFile(delete=False, suffix=".wav") as temp_user_audio:
+                        user_audio_processor.save_audio(temp_user_audio.name)
+                        user_input = transcribe_audio(temp_user_audio.name)
+                        st.write(f"Transcribed: {user_input}")
+                else:
+                    st.error("No audio recorded. Please speak into your microphone.")
 
+    # Process and generate response
+    if st.button("Generate AI Response in My Voice"):
+        if not agree_tos:
+            st.error("Please agree to the Coqui Public Model License to continue.")
+            return
+            
         if not ref_audio_path:
-            st.error("Please upload or record reference audio.")
+            st.error("Please provide reference audio (upload or record).")
             return
 
-        # Handle User Input
-        if input_type == "Audio":
-            if user_audio_processor:
-                with NamedTemporaryFile(delete=False, suffix=".wav") as temp_user_audio:
-                    user_audio_processor.save_audio(temp_user_audio.name)
-                    user_input = transcribe_audio(temp_user_audio.name)
-                    os.unlink(temp_user_audio.name)
-
         if not user_input:
-            st.error("Please enter text or record audio.")
+            st.error("Please enter text or record a question.")
             return
 
-        # Get AI Response
-        response_text = get_llm_response(API_KEY, user_input)
-
-        # Generate Speech
-        output_audio_path = "output_speech.wav"
-        try:
-            generate_speech(response_text, output_audio_path, ref_audio_path)
-            os.unlink(ref_audio_path)
-            st.audio(output_audio_path, format="audio/wav")
-        except FileNotFoundError as e:
-            st.error(str(e))
+        with st.spinner("Processing..."):
+            # Get AI Response
+            llm_response = get_llm_response(API_KEY, user_input)
+            st.subheader("AI Response:")
+            st.write(llm_response)
+            
+            # Generate Speech
+            output_audio_path = f"output_speech_{uuid.uuid4()}.wav"
+            success, message = generate_speech(
+                llm_response, 
+                output_audio_path, 
+                ref_audio_path, 
+                language
+            )
+            
+            if success:
+                st.subheader("Listen to the response in your voice:")
+                st.audio(output_audio_path, format="audio/wav")
+            else:
+                st.error(message)
 
 if __name__ == "__main__":
-    main()
+    main()