Update app.py

app.py

@@ -2,11 +2,12 @@ from fastapi import FastAPI
 from fastapi.responses import StreamingResponse
 from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech
 import torch
-import soundfile as sf
 import io
 import os
 import logging
+import numpy as np
 from pydub import AudioSegment
+import soundfile as sf
 
 logging.basicConfig(level=logging.INFO)
 
@@ -14,24 +15,31 @@ os.environ["HF_HOME"] = "/home/user/.cache/huggingface"
 
 app = FastAPI()
 
-# Load processor & model
+# Load model & processor
 processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
 model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts")
 
-# Dummy speaker embedding (flat voice)
+# Dummy speaker embedding
 speaker_embeddings = torch.zeros((1, 512))
 
 @app.get("/")
 def home():
-    return {"message": "Welcome to the SpeechT5 TTS API. Use /speak?text=Hello"}
+    return {"message": "Welcome to SpeechT5 TTS API. Use /speak_wav or /speak_mp3"}
+
+def float32_to_int16(waveform: np.ndarray):
+    """Convert float32 [-1,1] waveform to int16 PCM"""
+    waveform = np.clip(waveform, -1.0, 1.0)
+    waveform_int16 = (waveform * 32767).astype(np.int16)
+    return waveform_int16
 
 @app.get("/speak_wav")
 def speak_wav(text: str):
     inputs = processor(text=text, return_tensors="pt")
     speech = model.generate_speech(inputs["input_ids"], speaker_embeddings)
+
     buf = io.BytesIO()
-    # PCM16 WAV
-    sf.write(buf, speech.numpy(), 16000, format="WAV", subtype="PCM_16")
+    waveform_int16 = float32_to_int16(speech.numpy())
+    sf.write(buf, waveform_int16, 16000, format="WAV")
     buf.seek(0)
     return StreamingResponse(buf, media_type="audio/wav")
 
@@ -39,11 +47,16 @@ def speak_wav(text: str):
 def speak_mp3(text: str):
     inputs = processor(text=text, return_tensors="pt")
     speech = model.generate_speech(inputs["input_ids"], speaker_embeddings)
-    # Convert to WAV buffer first
+
+    # Convert float32 -> int16
+    waveform_int16 = float32_to_int16(speech.numpy())
+
+    # WAV buffer
     wav_buf = io.BytesIO()
-    sf.write(wav_buf, speech.numpy(), 16000, format="WAV", subtype="PCM_16")
+    sf.write(wav_buf, waveform_int16, 16000, format="WAV")
     wav_buf.seek(0)
-    # Convert WAV → MP3
+
+    # Convert to MP3
     audio = AudioSegment.from_wav(wav_buf)
     mp3_buf = io.BytesIO()
     audio.export(mp3_buf, format="mp3")