Update app.py

app.py

@@ -1,131 +1,222 @@
 import gradio as gr
-import librosa
 import numpy as np
 import torch
-
-from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan
-
-
-checkpoint = "Chithekitale/PhD_tts_updated"
-processor = SpeechT5Processor.from_pretrained(checkpoint)
+import tempfile
+import os
+from scipy.io.wavfile import write
+from transformers import (
+    SpeechT5Processor,
+    SpeechT5ForTextToSpeech,
+    SpeechT5HifiGan
+)
+
+# =========================
+# Model loading
+# =========================
+checkpoint = "Chithekitale/Chichewa_tts_v2"
+
+processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
 model = SpeechT5ForTextToSpeech.from_pretrained(checkpoint)
 vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")
 
-
+# Make all keys consistent
 speaker_embeddings = {
-    "BDL": "spkemb/cmu_us_bdl_arctic-wav-arctic_a0009.npy",
-    "CLB": "spkemb/cmu_us_clb_arctic-wav-arctic_a0144.npy",
-    "KSP": "spkemb/cmu_us_ksp_arctic-wav-arctic_b0087.npy",
-    "RMS": "spkemb/cmu_us_rms_arctic-wav-arctic_b0353.npy",
-    "SLT": "spkemb/cmu_us_slt_arctic-wav-arctic_a0508.npy",
+    "SPK1": "spkemb/speaker_2.npy",
+    "SPK2": "spkemb/speaker_1.npy",
+    "SPK3": "spkemb/cmu_us_ksp_arctic-wav-arctic_b0087.npy",
+    "SPK4": "spkemb/cmu_us_rms_arctic-wav-arctic_b0353.npy",
+    "SPK5": "spkemb/cmu_us_slt_arctic-wav-arctic_a0508.npy",
 }
 
+SPEAKER_CHOICES = [
+    "SPK1 (female)",
+    "SPK2 (male)",
+    "SPK3 (male)",
+    "SPK4 (male)",
+    "SPK5 (female)"
+]
 
-def predict(text, speaker):
-    if len(text.strip()) == 0:
-        return (16000, np.zeros(0).astype(np.int16))
+EXAMPLES = [
+    ["Ndapita, koma ndibweranso pompano.", "SPK1 (female)"],
+    ["Koma apapa zikuoneka kuti ziyenda bwino.", "SPK2 (male)"],
+    ["Ineyo ndikuona kuti sizizasithanso.", "SPK3 (male)"],
+    ["Mwina kusogolo kuno anthu ena azalimba mtima, koma panopana ndakaika.", "SPK4 (male)"],
+    ["Simungasankhe munthu oti bola linamukana.", "SPK5 (female)"],
+    ["Kodi chimanga panopa chikugulisidwa zingati, kapena nanunso simukudziwa?", "SPK5 (female)"],
+]
 
-    inputs = processor(text=text, return_tensors="pt")
+SAMPLE_RATE = 16000
 
-    # limit input length
-    input_ids = inputs["input_ids"]
-    input_ids = input_ids[..., :model.config.max_text_positions]
+# =========================
+# Helpers
+# =========================
+def get_speaker_key(speaker_label: str) -> str:
+    # "SPK1 (female)" -> "SPK1"
+    return speaker_label.split()[0]
 
-    if speaker == "Surprise Me!":
-        # load one of the provided speaker embeddings at random
-        idx = np.random.randint(len(speaker_embeddings))
-        key = list(speaker_embeddings.keys())[idx]
-        speaker_embedding = np.load(speaker_embeddings[key])
+def load_speaker_embedding(speaker: str) -> np.ndarray:
+    speaker_key = get_speaker_key(speaker)
 
-        # randomly shuffle the elements
-        np.random.shuffle(speaker_embedding)
+    if speaker_key not in speaker_embeddings:
+        raise ValueError(f"Unknown speaker key: {speaker_key}")
 
-        # randomly flip half the values
-        x = (np.random.rand(512) >= 0.5) * 1.0
-        x[x == 0] = -1.0
-        speaker_embedding *= x
+    path = speaker_embeddings[speaker_key]
 
-        #speaker_embedding = np.random.rand(512).astype(np.float32) * 0.3 - 0.15
-    else:
-        speaker_embedding = np.load(speaker_embeddings[speaker[:3]])
+    try:
+        speaker_embedding = np.load(path).astype(np.float32)
+    except Exception as e:
+        raise FileNotFoundError(
+            f"Could not load speaker embedding file: {path}. Error: {e}"
+        )
 
-    speaker_embedding = torch.tensor(speaker_embedding).unsqueeze(0)
+    if speaker_embedding.ndim == 2:
+        speaker_embedding = speaker_embedding.mean(axis=0)
 
-    speech = model.generate_speech(input_ids, speaker_embedding, vocoder=vocoder)
+    speaker_embedding = np.squeeze(speaker_embedding)
 
-    speech = (speech.numpy() * 32767).astype(np.int16)
-    return (16000, speech)
+    if speaker_embedding.shape != (512,):
+        raise ValueError(
+            f"Unexpected speaker embedding shape after processing: "
+            f"{speaker_embedding.shape}. Expected (512,)"
+        )
 
+    return speaker_embedding
 
-title = "SpeechT5: Speech Synthesis"
+def save_audio_to_wav(audio: np.ndarray, sample_rate: int = SAMPLE_RATE) -> str:
+    """
+    Save generated int16 audio to a temporary WAV file and return its path.
+    """
+    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
+    temp_file.close()
+    write(temp_file.name, sample_rate, audio)
+    return temp_file.name
 
-description = """
-The <b>SpeechT5</b> model is pre-trained on text as well as speech inputs, with targets that are also a mix of text and speech.
-By pre-training on text and speech at the same time, it learns unified representations for both, resulting in improved modeling capabilities.
+# =========================
+# Inference
+# =========================
+def predict(text, speaker):
+    try:
+        if not text or len(text.strip()) == 0:
+            return None, None, "Please enter some Chichewa text."
 
-SpeechT5 can be fine-tuned for different speech tasks. This space demonstrates the <b>text-to-speech</b> (TTS) checkpoint for the English language.
+        inputs = processor(text=text, return_tensors="pt")
+        input_ids = inputs["input_ids"][..., :model.config.max_text_positions]
 
-See also the <a href="https://huggingface.co/spaces/Matthijs/speecht5-asr-demo">speech recognition (ASR) demo</a>
-and the <a href="https://huggingface.co/spaces/Matthijs/speecht5-vc-demo">voice conversion demo</a>.
+        speaker_embedding = load_speaker_embedding(speaker)
+        speaker_embedding = torch.tensor(
+            speaker_embedding, dtype=torch.float32
+        ).unsqueeze(0)
 
-Refer to <a href="https://colab.research.google.com/drive/1i7I5pzBcU3WDFarDnzweIj4-sVVoIUFJ">this Colab notebook</a> to learn how to fine-tune the SpeechT5 TTS model on your own dataset or language.
+        with torch.no_grad():
+            speech = model.generate_speech(
+                input_ids,
+                speaker_embedding,
+                vocoder=vocoder
+            )
 
-<b>How to use:</b> Enter some English text and choose a speaker. The output is a mel spectrogram, which is converted to a mono 16 kHz waveform by the
-HiFi-GAN vocoder. Because the model always applies random dropout, each attempt will give slightly different results.
-The <em>Surprise Me!</em> option creates a completely randomized speaker.
-"""
+        speech = speech.cpu().numpy()
 
-article = """
-<div style='margin:20px auto;'>
-
-<p>References: <a href="https://arxiv.org/abs/2110.07205">SpeechT5 paper</a> |
-<a href="https://github.com/microsoft/SpeechT5/">original GitHub</a> |
-<a href="https://huggingface.co/mechanicalsea/speecht5-tts">original weights</a></p>
-
-<pre>
-@article{Ao2021SpeechT5,
-  title   = {SpeechT5: Unified-Modal Encoder-Decoder Pre-training for Spoken Language Processing},
-  author  = {Junyi Ao and Rui Wang and Long Zhou and Chengyi Wang and Shuo Ren and Yu Wu and Shujie Liu and Tom Ko and Qing Li and Yu Zhang and Zhihua Wei and Yao Qian and Jinyu Li and Furu Wei},
-  eprint={2110.07205},
-  archivePrefix={arXiv},
-  primaryClass={eess.AS},
-  year={2021}
-}
-</pre>
+        # Normalize safely before int16 conversion
+        max_val = np.max(np.abs(speech))
+        if max_val > 0:
+            speech = speech / max_val
 
-<p>Speaker embeddings were generated from <a href="http://www.festvox.org/cmu_arctic/">CMU ARCTIC</a> using <a href="https://huggingface.co/mechanicalsea/speecht5-vc/blob/main/manifest/utils/prep_cmu_arctic_spkemb.py">this script</a>.</p>
+        speech = (speech * 32767).astype(np.int16)
 
-</div>
-"""
+        # Save WAV file for downloading
+        wav_path = save_audio_to_wav(speech, SAMPLE_RATE)
 
-examples = [
-    ["It is not in the stars to hold our destiny but in ourselves.", "BDL (male)"],
-    ["The octopus and Oliver went to the opera in October.", "CLB (female)"],
-    ["She sells seashells by the seashore. I saw a kitten eating chicken in the kitchen.", "RMS (male)"],
-    ["Brisk brave brigadiers brandished broad bright blades, blunderbusses, and bludgeons—balancing them badly.", "SLT (female)"],
-    ["A synonym for cinnamon is a cinnamon synonym.", "BDL (male)"],
-    ["How much wood would a woodchuck chuck if a woodchuck could chuck wood? He would chuck, he would, as much as he could, and chuck as much wood as a woodchuck would if a woodchuck could chuck wood.", "CLB (female)"],
-]
+        status = f"Generated speech successfully using speaker: {speaker}"
+        return (SAMPLE_RATE, speech), wav_path, status
+
+    except Exception as e:
+        return None, None, f"Error during generation: {str(e)}"
+
+def clear_all():
+    return "", "SPK1 (female)", None, None, "Ready."
+
+# =========================
+# UI
+# =========================
+custom_css = """
+.gradio-container {
+    max-width: 1100px !important;
+    margin: 0 auto;
+}
+.hero {
+    text-align: center;
+    padding: 10px 0 0 0;
+}
+.section-note {
+    font-size: 0.95rem;
+    opacity: 0.9;
+}
+"""
 
-gr.Interface(
-    fn=predict,
-    inputs=[
-        gr.Text(label="Input Text"),
-        gr.Radio(label="Speaker", choices=[
-            "BDL (male)",
-            "CLB (female)",
-            "KSP (male)",
-            "RMS (male)",
-            "SLT (female)",
-            "Surprise Me!"
-        ],
-        value="BDL (male)"),
-    ],
-    outputs=[
-        gr.Audio(label="Generated Speech", type="numpy"),
-    ],
-    title=title,
-    description=description,
-    article=article,
-    examples=examples,
-).launch()
+with gr.Blocks(css=custom_css, title="Chichewa Speech Synthesis Demo") as demo:
+    gr.HTML(
+        """
+        <div class="hero">
+            <h1>Rule-Intergrated Chichewa Speech Synthesis</h1>
+            <p class="section-note">
+                Enter Chichewa text, choose a speaker voice, and generate speech audio.
+            </p>
+        </div>
+        """
+    )
+
+    with gr.Row():
+        with gr.Column(scale=5):
+            text_input = gr.Textbox(
+                label="Input Text",
+                placeholder="Type Chichewa text here...",
+                lines=6
+            )
+
+            speaker_input = gr.Radio(
+                label="Speaker Voice",
+                choices=SPEAKER_CHOICES,
+                value="SPK1 (female)"
+            )
+
+            with gr.Row():
+                generate_btn = gr.Button("Generate Speech", variant="primary")
+                clear_btn = gr.Button("Clear")
+
+            status_box = gr.Textbox(
+                label="System Status",
+                value="Ready.",
+                interactive=False
+            )
+
+        with gr.Column(scale=5):
+            audio_output = gr.Audio(
+                label="Generated Speech",
+                type="numpy",
+                autoplay=False
+            )
+
+            download_file = gr.File(
+                label="Download Audio File"
+            )
+
+    gr.Markdown("### Example Inputs")
+    gr.Examples(
+        examples=EXAMPLES,
+        inputs=[text_input, speaker_input]
+    )
+
+    generate_btn.click(
+        fn=predict,
+        inputs=[text_input, speaker_input],
+        outputs=[audio_output, download_file, status_box],
+        show_progress="full"
+    )
+
+    clear_btn.click(
+        fn=clear_all,
+        inputs=[],
+        outputs=[text_input, speaker_input, audio_output, download_file, status_box]
+    )
+
+demo.launch()