Optimize TTS for speed - add fp16, batch processing, remove sentence splitting

api.py

@@ -40,10 +40,20 @@ if DEVICE == "cuda":
 
 print("Loading Indic Parler-TTS model...")
 model = ParlerTTSForConditionalGeneration.from_pretrained(MODEL_ID).to(DEVICE)
+
+# Optimize model for inference
+if DEVICE == "cuda":
+    model = model.half()  # Use half precision (fp16) for faster inference
+    model.eval()
+
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 description_tokenizer = AutoTokenizer.from_pretrained(model.config.text_encoder._name_or_path)
 SAMPLE_RATE = model.config.sampling_rate
-print("Model loaded!")
+
+# Disable gradients for inference
+torch.set_grad_enabled(False)
+
+print("Model loaded and optimized!")
 
 # Named speakers
 SPEAKERS = {
@@ -92,10 +102,8 @@ def split_sentences(text):
 
 
 def clean_urdu_text(text):
-    """Clean and normalize Urdu text."""
-    text = re.sub(r'[a-zA-Z]+', '', text)
-    text = re.sub(r'\s+', ' ', text)
-    text = text.strip()
+    """Minimal text cleaning - preserve content."""
+    text = re.sub(r'\s+', ' ', text).strip()
     if text and text[-1] not in '۔.!?،':
         text += '۔'
     return text
@@ -103,7 +111,7 @@ def clean_urdu_text(text):
 
 @spaces.GPU()
 def generate_speech_internal(text, speaker, pitch, rate, temperature, do_sample):
-    """Internal function for speech generation."""
+    """Internal function for speech generation - optimized for speed."""
     if not text.strip():
         return None
 
@@ -113,47 +121,38 @@ def generate_speech_internal(text, speaker, pitch, rate, temperature, do_sample)
         gender = "female" if "Female" in speaker or speaker in ["Divya", "Rani"] else "male"
         description = build_description(speaker_name, gender, pitch, rate)
 
-        sentences = split_sentences(text)
-        if not sentences:
-            sentences = [text.strip()]
+        # Tokenize in batch for efficiency
+        desc_tokens = description_tokenizer(description, return_tensors="pt").to(DEVICE)
+        prompt_tokens = tokenizer(text, return_tensors="pt").to(DEVICE)
 
-        all_audio = []
+        # Set seed for reproducibility
         seed = torch.randint(0, 2**32, (1,)).item()
-
-        for sentence in sentences:
-            desc_tokens = description_tokenizer(description, return_tensors="pt")
-            prompt_tokens = tokenizer(sentence, return_tensors="pt")
-
-            torch.manual_seed(seed)
-            if torch.cuda.is_available():
-                torch.cuda.manual_seed(seed)
-
-            with torch.no_grad():
-                generation = model.generate(
-                    input_ids=desc_tokens.input_ids.to(DEVICE),
-                    attention_mask=desc_tokens.attention_mask.to(DEVICE),
-                    prompt_input_ids=prompt_tokens.input_ids.to(DEVICE),
-                    prompt_attention_mask=prompt_tokens.attention_mask.to(DEVICE),
-                    do_sample=do_sample,
-                    temperature=temperature if do_sample else 1.0,
-                    min_new_tokens=10,
-                )
-
-            audio_chunk = generation.cpu().numpy().squeeze()
-            audio_chunk = (audio_chunk * 32767).astype(np.int16)
-            all_audio.append(audio_chunk)
-
-            silence = np.zeros(int(SAMPLE_RATE * 0.3), dtype=np.int16)
-            all_audio.append(silence)
-
-        if not all_audio:
-            return None
-
-        audio = np.concatenate(all_audio)
+        torch.manual_seed(seed)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed(seed)
+
+        # Generate audio in one pass
+        with torch.no_grad():
+            generation = model.generate(
+                input_ids=desc_tokens.input_ids,
+                attention_mask=desc_tokens.attention_mask,
+                prompt_input_ids=prompt_tokens.input_ids,
+                prompt_attention_mask=prompt_tokens.attention_mask,
+                do_sample=do_sample,
+                temperature=temperature if do_sample else 1.0,
+                min_new_tokens=10,
+                max_new_tokens=1024,
+            )
+
+        # Convert to audio format
+        audio = generation.cpu().numpy().squeeze()
+        audio = (audio * 32767).astype(np.int16)
         return audio
 
     except Exception as e:
         print(f"Error generating speech: {e}")
+        import traceback
+        traceback.print_exc()
         return None