Upload app.py with huggingface_hub

app.py

@@ -1,10 +1,10 @@
-import io
 import time
+import traceback
 import tempfile
 from pathlib import Path
 
 from fastapi import FastAPI, HTTPException
-from fastapi.responses import Response
+from fastapi.responses import Response, JSONResponse
 from pydantic import BaseModel, Field
 
 app = FastAPI(title="iamcodio Dia2 TTS")
@@ -16,10 +16,10 @@ def get_model():
     global model
     if model is None:
         from dia2 import Dia2
-        print("[dia2] Loading Dia2-2B model...")
+        print("[dia2] Loading Dia2-2B model...", flush=True)
         start = time.time()
         model = Dia2.from_repo("nari-labs/Dia2-2B", device="cuda", dtype="bfloat16")
-        print(f"[dia2] Model loaded in {time.time() - start:.1f}s")
+        print(f"[dia2] Model loaded in {time.time() - start:.1f}s", flush=True)
     return model
 
 
@@ -41,19 +41,19 @@ def generate(req: GenerateRequest):
     if not req.text or req.text.isspace():
         raise HTTPException(status_code=400, detail="Text input cannot be empty")
 
-    from dia2 import GenerationConfig, SamplingConfig
+    try:
+        from dia2 import GenerationConfig, SamplingConfig
 
-    dia = get_model()
-    config = GenerationConfig(
-        cfg_scale=req.cfg_scale,
-        audio=SamplingConfig(temperature=req.temperature, top_k=req.top_k),
-        use_cuda_graph=req.use_cuda_graph,
-    )
+        dia = get_model()
+        config = GenerationConfig(
+            cfg_scale=req.cfg_scale,
+            audio=SamplingConfig(temperature=req.temperature, top_k=req.top_k),
+            use_cuda_graph=req.use_cuda_graph,
+        )
 
-    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
-        tmp_path = f.name
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as f:
+            tmp_path = f.name
 
-    try:
         start = time.time()
         result = dia.generate(
             req.text,
@@ -62,9 +62,11 @@ def generate(req: GenerateRequest):
             verbose=True,
         )
         elapsed = time.time() - start
-        print(f"[dia2] Generated in {elapsed:.2f}s")
+        print(f"[dia2] Generated in {elapsed:.2f}s", flush=True)
 
         wav_bytes = Path(tmp_path).read_bytes()
+        Path(tmp_path).unlink(missing_ok=True)
+
         return Response(
             content=wav_bytes,
             media_type="audio/wav",
@@ -72,5 +74,9 @@ def generate(req: GenerateRequest):
                 "X-Generation-Time": f"{elapsed:.2f}",
             },
         )
-    finally:
-        Path(tmp_path).unlink(missing_ok=True)
+    except Exception as e:
+        traceback.print_exc()
+        return JSONResponse(
+            status_code=500,
+            content={"error": str(e), "type": type(e).__name__},
+        )