Version2

app.py

@@ -40,17 +40,18 @@ import torch
 from parler_tts import ParlerTTSForConditionalGeneration
 from transformers import AutoTokenizer
 import soundfile as sf
-import uuid
+import io
 
 from fastapi import FastAPI
 from pydantic import BaseModel
-from fastapi.responses import FileResponse
+from fastapi.responses import StreamingResponse
 
 app = FastAPI()
 
-device = "cuda:0" if torch.cuda.is_available() else "cpu"
+# Select device
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load model ONCE when server starts
+# Load model ONCE when the server starts
 model = ParlerTTSForConditionalGeneration.from_pretrained(
     "Surya132/tts-custom"
 ).to(device)
@@ -59,22 +60,33 @@ tokenizer = AutoTokenizer.from_pretrained(
     "Surya132/tts-custom"
 )
 
+# Request schema
 class TTSRequest(BaseModel):
     prompt: str
     description: str
 
 
+@app.get("/")
+def root():
+    return {"message": "TTS API running. Use POST /generate"}
+
+
 @app.post("/generate")
 def generate_audio(req: TTSRequest):
 
-    description_inputs = description_tokenizer(
-        req.description, return_tensors="pt"
+    # Tokenize description
+    description_inputs = tokenizer(
+        req.description,
+        return_tensors="pt"
     ).to(device)
 
+    # Tokenize prompt
     prompt_inputs = tokenizer(
-        req.prompt, return_tensors="pt"
+        req.prompt,
+        return_tensors="pt"
     ).to(device)
 
+    # Generate audio
     generation = model.generate(
         input_ids=description_inputs.input_ids,
         attention_mask=description_inputs.attention_mask,
@@ -84,13 +96,14 @@ def generate_audio(req: TTSRequest):
 
     audio = generation.cpu().numpy().squeeze()
 
-    filename = f"audio_{uuid.uuid4()}.wav"
-
-    sf.write(filename, audio, model.config.sampling_rate)
+    # Write audio to memory buffer instead of file
+    buffer = io.BytesIO()
+    sf.write(buffer, audio, model.config.sampling_rate, format="WAV")
+    buffer.seek(0)
 
-    return FileResponse(filename, media_type="audio/wav")
+    return StreamingResponse(buffer, media_type="audio/wav")
 
-import uvicorn
 
 if __name__ == "__main__":
+    import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)