Update app.py

app.py

@@ -3,7 +3,8 @@ import os
 import tempfile
 import uuid
 from pathlib import Path
-from typing import Optional
+from threading import Lock
+from typing import Dict, Optional
 
 import requests
 import torch
@@ -24,6 +25,8 @@ MAX_TEXT_LENGTH = 1000
 DEFAULT_LANGUAGE = "en"
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+JOBS: Dict[str, Dict[str, str]] = {}
+JOB_LOCK = Lock()
 
 # Set token in environment before importing
 if HF_TOKEN:
@@ -148,10 +151,20 @@ def _preprocess_audio_wav(path: str, target_sr: int = 24000, target_peak: float
     return path
 
 
-@app.post("/health")
-def health(x_api_key: Optional[str] = Header(default=None)):
-    _require_api_key(x_api_key)
-    return {"status": "ok", "model": "indextts2", "device": DEVICE}
+def _set_job(job_id: str, **kwargs):
+    with JOB_LOCK:
+        JOBS[job_id] = {**JOBS.get(job_id, {}), **kwargs}
+
+
+def _get_job(job_id: str) -> Optional[Dict[str, str]]:
+    with JOB_LOCK:
+        data = JOBS.get(job_id)
+        return dict(data) if data else None
+
+
+def _pop_job(job_id: str) -> Optional[Dict[str, str]]:
+    with JOB_LOCK:
+        return JOBS.pop(job_id, None)
 
 
 def _cleanup_files(*files: str):
@@ -164,59 +177,106 @@ def _cleanup_files(*files: str):
                 pass  # Ignore cleanup errors
 
 
-@app.post("/generate")
-def generate(
-    payload: GenerateRequest = Body(...),
-    background_tasks: BackgroundTasks = BackgroundTasks(),
-    x_api_key: Optional[str] = Header(default=None),
-):
-    _require_api_key(x_api_key)
-
+def _run_generate_job(job_id: str, payload: Dict[str, str]):
     speaker_file = None
     output_file = None
-
+    _set_job(job_id, status="processing")
     try:
-        speaker_file = _temp_speaker_file(payload.speaker_wav)
+        speaker_file = _temp_speaker_file(payload["speaker_wav"])
         speaker_file = _preprocess_audio_wav(speaker_file)
         output_file = os.path.join(tempfile.gettempdir(), f"indextts2-{uuid.uuid4()}.wav")
 
-        # IndexTTS2 inference
-        # Note: language parameter is kept for API compatibility but IndexTTS2
-        # handles multilingual automatically (supports English, Turkish, Chinese, etc.)
         tts_model.infer(
             spk_audio_prompt=speaker_file,
-            text=payload.text,
+            text=payload["text"],
             output_path=output_file,
-            use_random=False,  # Deterministic output
+            use_random=False,
             verbose=False,
         )
 
-        # Light post-process to avoid end-of-file artifacts
         output_file = _preprocess_audio_wav(output_file)
 
-        # Verify the output file was created
         if not Path(output_file).exists():
             raise RuntimeError(f"TTS generation failed: output file was not created at {output_file}")
 
-        # Schedule cleanup after response is sent
-        background_tasks.add_task(_cleanup_files, speaker_file, output_file)
+        _cleanup_files(speaker_file)
+        _set_job(job_id, status="completed", output_file=output_file)
+    except Exception as exc:
+        _cleanup_files(speaker_file, output_file)
+        _set_job(job_id, status="error", error=str(exc))
 
-        return FileResponse(output_file, media_type="audio/wav", filename="output.wav")
 
-    except HTTPException:
-        # Clean up on HTTPException
-        if speaker_file and Path(speaker_file).exists():
-            Path(speaker_file).unlink(missing_ok=True)
-        raise
-    except Exception as exc:  # pragma: no cover
-        # Clean up on error
-        if speaker_file and Path(speaker_file).exists():
-            Path(speaker_file).unlink(missing_ok=True)
-        if output_file and Path(output_file).exists():
-            Path(output_file).unlink(missing_ok=True)
-        return JSONResponse(status_code=500, content={"error": str(exc)})
+@app.post("/health")
+def health(x_api_key: Optional[str] = Header(default=None)):
+    _require_api_key(x_api_key)
+    return {"status": "ok", "model": "indextts2", "device": DEVICE}
+
+
+@app.post("/generate")
+def generate(
+    payload: GenerateRequest = Body(...),
+    background_tasks: BackgroundTasks = BackgroundTasks(),
+    x_api_key: Optional[str] = Header(default=None),
+):
+    _require_api_key(x_api_key)
+    job_id = str(uuid.uuid4())
+    _set_job(job_id, status="queued")
+
+    # Offload the long-running synthesis so the HTTP request stays fast (<100s)
+    background_tasks.add_task(_run_generate_job, job_id, payload.dict())
+
+    return JSONResponse(
+        status_code=202,
+        content={
+            "job_id": job_id,
+            "status": "queued",
+            "status_url": f"/status/{job_id}",
+            "result_url": f"/result/{job_id}",
+        },
+    )
+
+
+@app.get("/status/{job_id}")
+def job_status(job_id: str, x_api_key: Optional[str] = Header(default=None)):
+    _require_api_key(x_api_key)
+    job = _get_job(job_id)
+    if not job:
+        raise HTTPException(status_code=404, detail="Job not found")
+    payload: Dict[str, str] = {"job_id": job_id, "status": job.get("status", "unknown")}
+    if "error" in job:
+        payload["error"] = job["error"]
+    return payload
+
+
+@app.get("/result/{job_id}")
+def job_result(
+    job_id: str,
+    background_tasks: BackgroundTasks = BackgroundTasks(),
+    x_api_key: Optional[str] = Header(default=None),
+):
+    _require_api_key(x_api_key)
+    job = _get_job(job_id)
+    if not job:
+        raise HTTPException(status_code=404, detail="Job not found")
+    status = job.get("status")
+    if status != "completed":
+        raise HTTPException(status_code=409, detail=f"Job not ready (status={status})")
+
+    output_file = job.get("output_file")
+    if not output_file or not Path(output_file).exists():
+        _pop_job(job_id)
+        raise HTTPException(status_code=410, detail="Result expired or missing")
+
+    # Remove job from memory and cleanup output after sending
+    _pop_job(job_id)
+    background_tasks.add_task(_cleanup_files, output_file)
+
+    return FileResponse(output_file, media_type="audio/wav", filename="output.wav")
 
 
 @app.get("/")
 def root():
-    return {"name": "indextts2-api", "endpoints": ["/health", "/generate"]}
+    return {
+        "name": "indextts2-api",
+        "endpoints": ["/health", "/generate", "/status/{job_id}", "/result/{job_id}"],
+    }