Update app.py

app.py

@@ -1,210 +1,161 @@
-from __future__ import annotations
-
+```python
+# app.py
 import os
-
-import gradio as gr
-import spaces
+import io
 import torch
 import torchaudio
-
+import numpy as np
+import gradio as gr
 from transformers import (
     SeamlessM4TFeatureExtractor,
     SeamlessM4TTokenizer,
     SeamlessM4Tv2ForSpeechToText,
+    AutoTokenizer,
+    AutoFeatureExtractor,
 )
-
-from lang_list import (
-    ASR_TARGET_LANGUAGE_NAMES,
-    LANGUAGE_NAME_TO_CODE,
-    S2ST_TARGET_LANGUAGE_NAMES,
-    S2TT_TARGET_LANGUAGE_NAMES,
-    T2ST_TARGET_LANGUAGE_NAMES,
-    TEXT_SOURCE_LANGUAGE_NAMES,
+from pydub import AudioSegment
+import nltk
+from parler_tts import ParlerTTSForConditionalGeneration
+from lang_list import LANGUAGE_NAME_TO_CODE, ASR_TARGET_LANGUAGE_NAMES, S2TT_TARGET_LANGUAGE_NAMES
+
+# Download punkt for sentence splitting
+nltk.download('punkt_tab')
+
+# Device and dtype
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+DTYPE = torch.bfloat16 if DEVICE != "cpu" else torch.float32
+SAMPLE_RATE = 16000
+
+# Load speech-to-text model
+stt_model = SeamlessM4Tv2ForSpeechToText.from_pretrained(
+    "ai4bharat/indic-seamless",
+    torch_dtype=DTYPE
+).to(DEVICE)
+feature_extractor = SeamlessM4TFeatureExtractor.from_pretrained(
+    "ai4bharat/indic-seamless"
+)
+tt_tokenizer = SeamlessM4TTokenizer.from_pretrained(
+    "ai4bharat/indic-seamless"
 )
 
-
-DESCRIPTION = """\
-### Translation Model for Indian Languages** 🎙️➡️📜 
-#### **How to Use:**  
-1. **Upload or record** an audio clip in any supported Indian language.  
-2. Click **"Translate"** to generate the corresponding text in the target language.  
-3. View or copy the output for further use.  
-🚀 Try it out and experience seamless speech to speech translation for Indian languages!
-"""
-
-hf_token = os.getenv("HF_TOKEN")
-device = "cuda:0" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
-torch_dtype = torch.bfloat16 if device != "cpu" else torch.float32
-
-model = SeamlessM4Tv2ForSpeechToText.from_pretrained("ai4bharat/indic-seamless", torch_dtype=torch_dtype, token=hf_token).to(device)
-processor = SeamlessM4TFeatureExtractor.from_pretrained("ai4bharat/indic-seamless", token=hf_token)
-tokenizer = SeamlessM4TTokenizer.from_pretrained("ai4bharat/indic-seamless", token=hf_token)
-
-CACHE_EXAMPLES = os.getenv("CACHE_EXAMPLES") == "1" and torch.cuda.is_available()
-
-AUDIO_SAMPLE_RATE = 16000
-MAX_INPUT_AUDIO_LENGTH = 60  # in seconds
-DEFAULT_TARGET_LANGUAGE = "Hindi"
-
-def preprocess_audio(input_audio: str) -> None:
-    arr, org_sr = torchaudio.load(input_audio)
-    new_arr = torchaudio.functional.resample(arr, orig_freq=org_sr, new_freq=AUDIO_SAMPLE_RATE)
-    max_length = int(MAX_INPUT_AUDIO_LENGTH * AUDIO_SAMPLE_RATE)
-    if new_arr.shape[1] > max_length:
-        new_arr = new_arr[:, :max_length]
-        gr.Warning(f"Input audio is too long. Only the first {MAX_INPUT_AUDIO_LENGTH} seconds is used.")
-    torchaudio.save(input_audio, new_arr, sample_rate=int(AUDIO_SAMPLE_RATE))
-
-def run_s2tt(input_audio: str, source_language: str, target_language: str) -> str:
-    # preprocess_audio(input_audio)
-    # source_language_code = LANGUAGE_NAME_TO_CODE[source_language]
-    target_language_code = LANGUAGE_NAME_TO_CODE[target_language]
-
-    input_audio, orig_freq = torchaudio.load(input_audio)
-    input_audio = torchaudio.functional.resample(input_audio, orig_freq=orig_freq, new_freq=16000)
-    audio_inputs= processor(input_audio, sampling_rate=16000, return_tensors="pt").to(device=device, dtype=torch_dtype)
-
-    text_out = model.generate(**audio_inputs, tgt_lang=target_language_code)[0].float().cpu().numpy().squeeze()
-
-    return tokenizer.decode(text_out, clean_up_tokenization_spaces=True, skip_special_tokens=True)
-
-def run_asr(input_audio: str, target_language: str) -> str:
-    # preprocess_audio(input_audio)
-    target_language_code = LANGUAGE_NAME_TO_CODE[target_language]
-
-    input_audio, orig_freq = torchaudio.load(input_audio)
-    input_audio = torchaudio.functional.resample(input_audio, orig_freq=orig_freq, new_freq=16000)
-    audio_inputs= processor(input_audio, sampling_rate=16000, return_tensors="pt").to(device=device, dtype=torch_dtype)
-
-    text_out = model.generate(**audio_inputs, tgt_lang=target_language_code)[0].float().cpu().numpy().squeeze()
-
-    return tokenizer.decode(text_out, clean_up_tokenization_spaces=True, skip_special_tokens=True)
-
-
-
-with gr.Blocks() as demo_s2st:
-    with gr.Row():
-        with gr.Column():
-            with gr.Group():
-                input_audio = gr.Audio(label="Input speech", type="filepath")
-                source_language = gr.Dropdown(
-                    label="Source language",
-                    choices=ASR_TARGET_LANGUAGE_NAMES,
-                    value="English",
-                )
-                target_language = gr.Dropdown(
-                    label="Target language",
-                    choices=S2ST_TARGET_LANGUAGE_NAMES,
-                    value=DEFAULT_TARGET_LANGUAGE,
-                )
-            btn = gr.Button("Translate")
-        with gr.Column():
-            with gr.Group():
-                output_audio = gr.Audio(
-                    label="Translated speech",
-                    autoplay=False,
-                    streaming=False,
-                    type="numpy",
-                )
-                output_text = gr.Textbox(label="Translated text")
-
-with gr.Blocks() as demo_s2tt:
-    with gr.Row():
-        with gr.Column():
-            with gr.Group():
-                input_audio = gr.Audio(label="Input speech", type="filepath")
-                source_language = gr.Dropdown(
-                    label="Source language",
-                    choices=ASR_TARGET_LANGUAGE_NAMES,
-                    value="English",
-                )
-                target_language = gr.Dropdown(
-                    label="Target language",
-                    choices=S2TT_TARGET_LANGUAGE_NAMES,
-                    value=DEFAULT_TARGET_LANGUAGE,
-                )
-            btn = gr.Button("Translate")
-        with gr.Column():
-            output_text = gr.Textbox(label="Translated text")
-
-    btn.click(
-        fn=run_s2tt,
-        inputs=[input_audio, source_language, target_language],
-        outputs=output_text,
-        api_name="s2tt",
+# Load TTS models
+repo_id = "ai4bharat/indic-parler-tts-pretrained"
+finetuned_repo_id = "ai4bharat/indic-parler-tts"
+
+tts_model = ParlerTTSForConditionalGeneration.from_pretrained(
+    repo_id,
+    attn_implementation="eager",
+    torch_dtype=DTYPE,
+).to(DEVICE)
+finetuned_tts = ParlerTTSForConditionalGeneration.from_pretrained(
+    finetuned_repo_id,
+    attn_implementation="eager",
+    torch_dtype=DTYPE,
+).to(DEVICE)
+
+tts_tokenizer = AutoTokenizer.from_pretrained(repo_id)
+description_tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-large")
+tts_feature_extractor = AutoFeatureExtractor.from_pretrained(repo_id)
+
+# Voice options - example speakers
+VOICES = [
+    "Sunita", "Suresh", "Aditi", "Prakash", "Rohit", "Anjali", "Jaya"
+]
+
+# Custom CSS for visual styling
+CSS = '''
+body { background-color: #f9fafb; }
+.gradio-container { max-width: 900px; margin: auto; padding: 20px; }
+.step-box { background: #ffffff; border-radius: 12px; padding: 16px; margin-bottom: 12px; box-shadow: 0 2px 8px rgba(0,0,0,0.1); }
+'''  
+
+# Helpers
+def numpy_to_mp3(audio_array, sampling_rate):
+    if np.issubdtype(audio_array.dtype, np.floating):
+        max_val = np.max(np.abs(audio_array))
+        audio_array = (audio_array / max_val) * 32767
+        audio_array = audio_array.astype(np.int16)
+    segment = AudioSegment(
+        audio_array.tobytes(),
+        frame_rate=sampling_rate,
+        sample_width=audio_array.dtype.itemsize,
+        channels=1
     )
-
-with gr.Blocks() as demo_t2st:
-    with gr.Row():
-        with gr.Column():
-            with gr.Group():
-                input_text = gr.Textbox(label="Input text")
-                with gr.Row():
-                    source_language = gr.Dropdown(
-                        label="Source language",
-                        choices=TEXT_SOURCE_LANGUAGE_NAMES,
-                        value="English",
-                    )
-                    target_language = gr.Dropdown(
-                        label="Target language",
-                        choices=T2ST_TARGET_LANGUAGE_NAMES,
-                        value=DEFAULT_TARGET_LANGUAGE,
-                    )
-            btn = gr.Button("Translate")
+    mp3_io = io.BytesIO()
+    segment.export(mp3_io, format="mp3", bitrate="320k")
+    return mp3_io.getvalue()
+
+# STT / Translation
+def transcribe_and_translate(audio_path, source_language, target_language):
+    wav, orig_sr = torchaudio.load(audio_path)
+    wav = torchaudio.functional.resample(wav, orig_freq=orig_sr, new_freq=SAMPLE_RATE)
+    inputs = feature_extractor(wav, sampling_rate=SAMPLE_RATE, return_tensors="pt").to(DEVICE, DTYPE)
+    tgt = LANGUAGE_NAME_TO_CODE[target_language]
+    gen = stt_model.generate(**inputs, tgt_lang=tgt)[0]
+    text = tt_tokenizer.decode(gen, skip_special_tokens=True, clean_up_tokenization_spaces=True)
+    return text
+
+# TTS generation
+def generate_tts(text, voice, finetuned=False):
+    description = f"{voice} speaks in a neutral tone with clear audio."
+    sentences = nltk.sent_tokenize(text)
+    all_audio = []
+    for sent in sentences:
+        desc_inputs = description_tokenizer(description, return_tensors="pt").to(DEVICE)
+        prompt_inputs = tts_tokenizer(sent, return_tensors="pt").to(DEVICE)
+        model = finetuned_tts if finetuned else tts_model
+        gen = model.generate(
+            input_ids=desc_inputs.input_ids,
+            attention_mask=desc_inputs.attention_mask,
+            prompt_input_ids=prompt_inputs.input_ids,
+            prompt_attention_mask=prompt_inputs.attention_mask,
+            do_sample=True,
+            return_dict_in_generate=True
+        )
+        if hasattr(gen, 'sequences') and hasattr(gen, 'audios_length'):
+            audio = gen.sequences[0, :gen.audios_length[0]]
+            audio_np = audio.to(torch.float32).cpu().numpy().flatten()
+            all_audio.append(audio_np)
+    combined = np.concatenate(all_audio)
+    return numpy_to_mp3(combined, tts_feature_extractor.sampling_rate)
+
+# Combined pipeline to reduce duplicate STT calls
+def pipeline(audio_path, source_language, target_language, voice, finetuned):
+    text = transcribe_and_translate(audio_path, source_language, target_language)
+    audio_bytes = generate_tts(text, voice, finetuned)
+    return text, audio_bytes
+
+# Gradio UI
+
+def build_ui():
+    with gr.Blocks(css=CSS) as demo:
+        gr.Markdown("🎙️ AUDIO TRANSLATION 🎙️")
+        # Usage Steps
         with gr.Column():
-            with gr.Group():
-                output_audio = gr.Audio(
-                    label="Translated speech",
-                    autoplay=False,
-                    streaming=False,
-                    type="numpy",
-                )
-                output_text = gr.Textbox(label="Translated text")
-
-
-
-with gr.Blocks() as demo_asr:
-    with gr.Row():
-        with gr.Column():
-            with gr.Group():
-                input_audio = gr.Audio(label="Input speech", type="filepath")
-                target_language = gr.Dropdown(
-                    label="Target language",
-                    choices=ASR_TARGET_LANGUAGE_NAMES,
-                    value=DEFAULT_TARGET_LANGUAGE,
-                )
-            btn = gr.Button("Transcribe")
-        with gr.Column():
-            output_text = gr.Textbox(label="Transcribed text")
-
-    btn.click(
-        fn=run_asr,
-        inputs=[input_audio, target_language],
-        outputs=output_text,
-        api_name="asr",
-    )
-
-
-with gr.Blocks(css="style.css") as demo:
-    gr.Markdown(DESCRIPTION)
-    gr.DuplicateButton(
-        value="Duplicate Space for private use",
-        elem_id="duplicate-button",
-        visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1",
-    )
-
-    with gr.Tabs():
-        # with gr.Tab(label="S2ST"):
-        #     demo_s2st.render()
-        with gr.Tab(label="S2TT"):
-            demo_s2tt.render()
-        # with gr.Tab(label="T2ST"):
-        #     demo_t2st.render()
-        # with gr.Tab(label="T2TT"):
-        #     demo_t2tt.render()
-        with gr.Tab(label="ASR"):
-            demo_asr.render()
-
-
-
-demo.launch(share=True)
+            gr.HTML("<div class='step-box'><strong>Step 1:</strong> Upload or record your audio clip.</div>")
+            gr.HTML("<div class='step-box'><strong>Step 2:</strong> Select the source and target languages.</div>")
+            gr.HTML("<div class='step-box'><strong>Step 3:</strong> Choose a voice persona.</div>")
+            gr.HTML("<div class='step-box'><strong>Step 4:</strong> (Optional) Toggle fine-tuned TTS for more natural speech.</div>")
+            gr.HTML("<div class='step-box'><strong>Step 5:</strong> Click <em>Run</em> and view your text & audio results on the right.</div>")
+        with gr.Row():
+            with gr.Column(scale=1):
+                audio_in = gr.Audio(label="Input Audio", type="filepath")
+                src = gr.Dropdown(ASR_TARGET_LANGUAGE_NAMES, label="Source Language", value="English")
+                tgt = gr.Dropdown(S2TT_TARGET_LANGUAGE_NAMES, label="Target Language", value="English")
+                voice = gr.Dropdown(VOICES, label="Voice Persona", value=VOICES[0])
+                finetune = gr.Checkbox(label="Use Fine-tuned TTS", value=False)
+                run_btn = gr.Button("Run", variant="primary")
+            with gr.Column(scale=1):
+                text_out = gr.Textbox(label="Translated Text")
+                audio_out = gr.Audio(label="Synthesized Speech", format="mp3")
+        run_btn.click(
+            fn=pipeline,
+            inputs=[audio_in, src, tgt, voice, finetune],
+            outputs=[text_out, audio_out]
+        )
+    return demo
+
+if __name__ == "__main__":
+    ui = build_ui()
+    ui.launch(share=True)
+```