Update app.py

app.py

@@ -1,20 +1,24 @@
 import os
+
+# Set cache directories to /tmp to avoid permission issues in the container
+os.environ["HF_HOME"] = "/tmp"
+os.environ["TRANSFORMERS_CACHE"] = "/tmp"
+os.environ["TORCH_HOME"] = "/tmp"
+os.environ["XDG_CACHE_HOME"] = "/tmp"
+
 import io
 import re
 import numpy as np
 import scipy.io.wavfile
+import torch
 from fastapi import FastAPI
-from pydantic import BaseModel
 from fastapi.responses import StreamingResponse
-import torch
+from pydantic import BaseModel
 from transformers import VitsModel, AutoTokenizer
 
-# Use /tmp for cache to avoid permission errors
-os.environ["HF_HOME"] = "/tmp"
-
 app = FastAPI()
 
-# Load model and tokenizer once
+# Load model and tokenizer once at startup
 model = VitsModel.from_pretrained("Somali-tts/somali_tts_model")
 tokenizer = AutoTokenizer.from_pretrained("saleolow/somali-mms-tts")
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -94,7 +98,6 @@ async def synthesize(data: TextIn):
     with torch.no_grad():
         waveform = model(**inputs).waveform.squeeze().cpu().numpy()
 
-    # Convert waveform to WAV bytes
     buf = io.BytesIO()
     scipy.io.wavfile.write(buf, rate=model.config.sampling_rate, data=(waveform * 32767).astype(np.int16))
     buf.seek(0)