Update app.py

app.py

@@ -2,14 +2,24 @@ import os
 import sys
 import uuid
 import logging
+import io
+import base64
 from typing import Optional
 
-# CRITICAL: Set numba environment variables BEFORE any imports
+# CRITICAL: Set environment variables BEFORE any imports to fix PyTorch user issues
 os.environ['NUMBA_CACHE_DIR'] = '/tmp/numba_cache'
 os.environ['NUMBA_DISABLE_JIT'] = '0'  # Keep JIT enabled but control cache
+os.environ['TORCH_HOME'] = '/tmp/torch_cache'
+os.environ['TRANSFORMERS_CACHE'] = '/tmp/transformers_cache'
+os.environ['HF_HOME'] = '/tmp/huggingface_cache'
 
-# Create cache directory
-os.makedirs('/tmp/numba_cache', exist_ok=True)
+# Set user environment variables to avoid getpwuid errors
+os.environ['USER'] = 'appuser'
+os.environ['LOGNAME'] = 'appuser'
+
+# Create cache directories
+for cache_dir in ['/tmp/numba_cache', '/tmp/torch_cache', '/tmp/transformers_cache', '/tmp/huggingface_cache']:
+    os.makedirs(cache_dir, exist_ok=True)
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -28,9 +38,8 @@ try:
     from fastapi.responses import FileResponse, JSONResponse
     from fastapi.middleware.cors import CORSMiddleware
     import soundfile as sf
-    import io
     
-    # Now import NeuTTS - this should work with the numba cache fix
+    # Now import NeuTTS - this should work with the cache fixes
     from neuttsair.neutts import NeuTTSAir
     logger.info("✅ All imports successful")
     
@@ -114,7 +123,8 @@ async def test_endpoint():
         "status": "success",
         "message": "API is working",
         "model_loaded": tts is not None,
-        "numba_cache": os.environ.get('NUMBA_CACHE_DIR', 'not set')
+        "numba_cache": os.environ.get('NUMBA_CACHE_DIR', 'not set'),
+        "torch_cache": os.environ.get('TORCH_HOME', 'not set')
     }
 
 @app.post("/api/v1/synthesize")
@@ -209,7 +219,6 @@ async def synthesize_speech_base64(
         sf.write(buffer, wav, 24000, format='WAV')
         buffer.seek(0)
         
-        import base64
         audio_b64 = base64.b64encode(buffer.read()).decode('utf-8')
         
         # Clean up
@@ -231,6 +240,56 @@ async def synthesize_speech_base64(
         traceback.print_exc()
         raise HTTPException(500, f"Synthesis failed: {str(e)}")
 
+# Batch processing endpoint
+@app.post("/api/v1/batch-synthesize")
+async def batch_synthesize(
+    ref_text: str = Form(...),
+    ref_audio: UploadFile = File(...),
+    texts: str = Form(..., description="JSON array of texts to synthesize")
+):
+    """
+    Synthesize multiple texts with the same voice
+    """
+    try:
+        import json
+        text_list = json.loads(texts)
+        
+        # Initialize model if needed
+        tts_model = initialize_model()
+        
+        # Save reference audio
+        os.makedirs("uploads", exist_ok=True)
+        upload_path = f"uploads/{uuid.uuid4()}.wav"
+        with open(upload_path, "wb") as f:
+            content = await ref_audio.read()
+            f.write(content)
+        
+        # Encode reference once
+        ref_codes = tts_model.encode_reference(upload_path)
+        
+        results = []
+        for i, text in enumerate(text_list):
+            wav = tts_model.infer(text, ref_codes, ref_text)
+            output_path = f"outputs/{uuid.uuid4()}.wav"
+            sf.write(output_path, wav, 24000)
+            results.append({
+                "text": text,
+                "audio_file": output_path,
+                "index": i
+            })
+        
+        # Clean up upload file
+        try:
+            os.remove(upload_path)
+        except:
+            pass
+        
+        return {"generated_files": results}
+        
+    except Exception as e:
+        logger.error(f"Batch synthesis error: {str(e)}")
+        raise HTTPException(500, f"Batch synthesis failed: {str(e)}")
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)