Update app.py

app.py

@@ -1,51 +1,100 @@
-import gradio as gr
 import torch
-from parler_tts import ParlerTTSForConditionalGeneration
-from transformers import AutoTokenizer
 import soundfile as sf
 import uuid
+import gradio as gr
+import numpy as np
+import re
+from parler_tts import ParlerTTSForConditionalGeneration
+from transformers import AutoTokenizer
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# Load model and tokenizers
 model_name = "ai4bharat/indic-parler-tts"
-model = ParlerTTSForConditionalGeneration.from_pretrained(model_name).to(device)
+device = "cpu"
+
+print("Loading model...")
+model = ParlerTTSForConditionalGeneration.from_pretrained(model_name).to(device).eval()
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 desc_tokenizer = AutoTokenizer.from_pretrained(model.config.text_encoder._name_or_path)
 
+print("Applying dynamic quantization...")
+quantized_model = torch.quantization.quantize_dynamic(
+    model,
+    {torch.nn.Linear},
+    dtype=torch.qint8
+)
+
+# Sentence splitter (splits by full stop, exclamation, or question mark)
+def split_text(text, max_len=150):
+    # First, try to split by sentence punctuation
+    chunks = re.split(r'(?<=[.!?]) +', text)
+    
+    # If any chunk is still too long, split further
+    refined_chunks = []
+    for chunk in chunks:
+        if len(chunk) <= max_len:
+            refined_chunks.append(chunk)
+        else:
+            # Break on space while respecting max_len
+            words = chunk.split()
+            buffer = []
+            length = 0
+            for word in words:
+                buffer.append(word)
+                length += len(word) + 1
+                if length > max_len:
+                    refined_chunks.append(' '.join(buffer))
+                    buffer = []
+                    length = 0
+            if buffer:
+                refined_chunks.append(' '.join(buffer))
+    return refined_chunks
+
+# Main synthesis function
 def synthesize(language, text, gender, emotion, speed, pitch, quality):
-    desc = (
-        f"A native {language} {gender.lower()} speaker with a {emotion.lower()} and expressive tone, "
+    description = (
+        f"A native {language.lower()} {gender.lower()} speaker with a {emotion.lower()} and expressive tone, "
         f"speaking at a {speed.lower()} rate with {pitch.lower()} pitch and {quality.lower()} voice quality."
     )
-    desc_inputs = desc_tokenizer(desc, return_tensors="pt").to(device)
-    text_inputs = tokenizer(text, return_tensors="pt").to(device)
-
-    gen_audio = model.generate(
-        input_ids=desc_inputs.input_ids,
-        attention_mask=desc_inputs.attention_mask,
-        prompt_input_ids=text_inputs.input_ids,
-        prompt_attention_mask=torch.ones_like(text_inputs.input_ids).to(device)
-    )
 
-    audio_np = gen_audio.cpu().numpy().squeeze()
-    filename = f"{uuid.uuid4()}.wav"
-    sf.write(filename, audio_np, model.config.sampling_rate)
+    description_input = desc_tokenizer(description, return_tensors="pt").to(device)
+    
+    chunks = split_text(text)
+    audio_pieces = []
+
+    for chunk in chunks:
+        prompt_input = tokenizer(chunk, return_tensors="pt").to(device)
+        with torch.no_grad():
+            generation = quantized_model.generate(
+                input_ids=description_input.input_ids,
+                attention_mask=description_input.attention_mask,
+                prompt_input_ids=prompt_input.input_ids,
+                prompt_attention_mask=torch.ones_like(prompt_input.input_ids).to(device)
+            )
+        audio_chunk = generation.cpu().numpy().squeeze()
+        audio_pieces.append(audio_chunk)
+
+    # Concatenate all audio chunks
+    final_audio = np.concatenate(audio_pieces)
+
+    filename = f"{uuid.uuid4().hex}.wav"
+    sf.write(filename, final_audio, quantized_model.config.sampling_rate)
     return filename
 
+# Gradio Interface
 iface = gr.Interface(
     fn=synthesize,
     inputs=[
-        gr.Dropdown(["Malayalam", "English", "Hindi", "Tamil"], label="Language"),
-        gr.Textbox(label="Text to Synthesize", lines=4),
+        gr.Dropdown(["Malayalam", "Hindi", "Tamil", "English"], label="Language"),
+        gr.Textbox(label="Text to Synthesize", lines=6, placeholder="Enter your sentence here..."),
         gr.Radio(["Male", "Female"], label="Speaker Gender"),
         gr.Dropdown(["Neutral", "Happy", "Sad", "Angry"], label="Emotion"),
         gr.Dropdown(["Slow", "Moderate", "Fast"], label="Speaking Rate"),
         gr.Dropdown(["Low", "Normal", "High"], label="Pitch"),
         gr.Dropdown(["Basic", "Refined"], label="Voice Quality"),
     ],
-    outputs=gr.Audio(type="filepath", label="Synthesized Audio"),
-    allow_flagging="never",
-    title="Multilingual TTS using Indic Parler-TTS",
-    description="Type text, choose a speaker style, and get synthesized speech for Malayalam, Hindi, Tamil, or English."
+    outputs=gr.Audio(type="filepath", label="Synthesized Speech"),
+    title="Multilingual Indic TTS (Quantized + Chunked)",
+    description="Fast CPU-based TTS with quantized Parler-TTS and text chunking for Malayalam, Hindi, Tamil, and English.",
 )
 
 iface.launch()