Update image

Dockerfile

@@ -18,5 +18,7 @@ COPY main.py .
 # Expose port 8080
 EXPOSE 8080
 
+ENV PYTHONUNBUFFERED=1
+
 # Start uvicorn
-CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8080"]
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8080", "--log-level", "debug"]

main.py

@@ -1,13 +1,18 @@
-import io
+from contextlib import asynccontextmanager
 from importlib.metadata import version
-import tempfile
-import os
+import asyncio
 import base64
-from typing import Dict, Any
+import io
 import logging
-from contextlib import asynccontextmanager
+import os
+import sys
+import tempfile
+import traceback
 
 from fastapi import FastAPI, Request, Response, HTTPException
+from fastapi.responses import JSONResponse
+
+from typing import Dict, Any
 
 # this must come before attempts to import from transformers or torch
 print(f'Versions: transformers: {version("transformers")}, torch: {version("torch")}')
@@ -99,11 +104,19 @@ class EndpointHandler():
                 # If already loaded (e.g. numpy array passed in some test environments)
                 audio_array = inputs
 
-            # 3. Prepare inputs using the recommended API
             processed_inputs = self.processor.apply_transcription_request(
                 audio=audio_array,
                 prompt=hotwords
-            ).to(self.model.device, self.model.dtype)
+            )
+
+            # Safely move to device and cast ONLY floating point tensors to the model's dtype.
+            # Casting integer tensors (like input_ids) to bfloat16/float16 will cause model errors.
+            for k, v in processed_inputs.items():
+                if isinstance(v, torch.Tensor):
+                    if torch.is_floating_point(v):
+                        processed_inputs[k] = v.to(device=self.model.device, dtype=self.model.dtype)
+                    else:
+                        processed_inputs[k] = v.to(device=self.model.device)
 
             # 4. Generate
             with torch.no_grad():
@@ -130,6 +143,7 @@ class EndpointHandler():
             return {"result": transcription}
 
         except Exception as e:
+            logger.exception("Inference failed due to exception", e)
             return {"error": f"Inference failed: {str(e)}"}
 
         finally:
@@ -175,9 +189,21 @@ async def predict(request: Request):
         # Fallback to raw bytes
         data = await request.body()
 
-    response = handler(data)
+    response = await asyncio.to_thread(handler, data)
 
     if "error" in response:
         raise HTTPException(status_code=500, detail=response["error"])
 
     return response
+
+@app.exception_handler(Exception)
+async def global_exception_handler(request: Request, exc: Exception):
+    # Force print to stderr
+    print(f"CRITICAL ERROR: {str(exc)}", file=sys.stderr)
+    print(traceback.format_exc(), file=sys.stderr)
+
+    # Return the error in the payload so you can see it in your HTTP client
+    return JSONResponse(
+        status_code=500,
+        content={"error": str(exc), "type": str(type(exc))}
+    )