refactor: decouple ASR processing from HTTP endpoint and update Modal environment to CUDA 12.3 registry image

backend/asr.py

@@ -6,27 +6,21 @@ import tempfile
 from pathlib import Path
 
 import modal
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import JSONResponse
 
-# ---------------------------------------------------------------------------
-# Constants
-# ---------------------------------------------------------------------------
 MINUTES = 60
 MODEL_NAME = "large-v3-turbo"
 COMPUTE_TYPE = "float16"
-ASR_PORT = 8000
 MODELS_DIR = Path("/models")
 
-# ---------------------------------------------------------------------------
-# Volume — persists downloaded model weights across deploys
-# ---------------------------------------------------------------------------
 volume = modal.Volume.from_name("aiko-asr-models", create_if_missing=True)
 
-# ---------------------------------------------------------------------------
-# Image — debian slim + CUDA cublas runtime + Python deps
-# ---------------------------------------------------------------------------
 image = (
-    modal.Image.debian_slim(python_version="3.12")
-    .apt_install("libcublas12")
+    modal.Image.from_registry(
+        "nvidia/cuda:12.3.2-runtime-ubuntu22.04",
+        add_python="3.12",
+    )
     .pip_install(
         "faster-whisper",
         "fastapi",
@@ -35,15 +29,10 @@ image = (
     )
 )
 
-# ---------------------------------------------------------------------------
-# App
-# ---------------------------------------------------------------------------
 app = modal.App("aiko-asr", image=image)
+web_app = FastAPI()
 
 
-# ---------------------------------------------------------------------------
-# ASR Server class
-# ---------------------------------------------------------------------------
 @app.cls(
     gpu="T4",
     timeout=10 * MINUTES,
@@ -56,9 +45,7 @@ class ASRServer:
 
     @modal.enter()
     def startup(self):
-        """Load the Whisper model once — reused across all requests."""
         from faster_whisper import WhisperModel
-
         print(f"Loading faster-whisper {MODEL_NAME} ...")
         self.model = WhisperModel(
             MODEL_NAME,
@@ -67,23 +54,14 @@ class ASRServer:
             download_root=str(MODELS_DIR),
         )
         print("Model ready.")
-        volume.commit()  # persist downloaded weights
-
-    @modal.web_endpoint(method="GET")
-    def health(self):
-        return {"status": "ok", "model": MODEL_NAME}
-
-    @modal.web_endpoint(method="POST")
-    async def transcribe(self, audio: "UploadFile"):  # type: ignore[name-defined]
-        from fastapi import File, UploadFile
-        from fastapi.responses import JSONResponse
-
-        suffix = Path(audio.filename or "audio.wav").suffix or ".wav"
+        volume.commit()
 
+    @modal.method()
+    def _transcribe(self, audio_bytes: bytes, suffix: str) -> dict:
+        import tempfile
         with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
-            tmp.write(await audio.read())
+            tmp.write(audio_bytes)
             tmp_path = tmp.name
-
         try:
             segments, info = self.model.transcribe(
                 tmp_path,
@@ -92,53 +70,30 @@ class ASRServer:
                 condition_on_previous_text=False,
             )
             text = " ".join(s.text.strip() for s in segments).strip()
-            return JSONResponse({
+            return {
                 "text": text,
                 "language": info.language,
                 "language_probability": round(info.language_probability, 3),
-            })
+            }
         finally:
             os.unlink(tmp_path)
 
 
-# ---------------------------------------------------------------------------
-# Local test entrypoint:  modal run asr.py [path/to/audio.wav]
-# ---------------------------------------------------------------------------
-@app.local_entrypoint()
-def main():
-    import sys
-    import wave
-    import struct
-    import math
-    import httpx
-
-    test_audio = sys.argv[1] if len(sys.argv) > 1 else None
+@web_app.get("/health")
+def health():
+    return {"status": "ok", "model": MODEL_NAME}
 
-    # Generate a 1-second 440 Hz sine-wave WAV if no file provided
-    if test_audio is None:
-        test_audio = "/tmp/asr_test_tone.wav"
-        with wave.open(test_audio, "w") as wf:
-            wf.setnchannels(1)
-            wf.setsampwidth(2)
-            wf.setframerate(16000)
-            frames = [
-                struct.pack("<h", int(32767 * math.sin(2 * math.pi * 440 * i / 16000)))
-                for i in range(16000)
-            ]
-            wf.writeframes(b"".join(frames))
-        print(f"No audio file given — generated test tone at {test_audio}")
-
-    print(f"Testing with {test_audio} ...")
 
+@web_app.post("/transcribe")
+async def transcribe(audio: UploadFile = File(...)):
+    suffix = Path(audio.filename or "audio.wav").suffix or ".wav"
+    audio_bytes = await audio.read()
     server = ASRServer()
-    url = server.transcribe.web_url
+    result = server._transcribe.remote(audio_bytes, suffix)
+    return JSONResponse(result)
 
-    with open(test_audio, "rb") as f:
-        resp = httpx.post(
-            url,
-            files={"audio": (Path(test_audio).name, f, "audio/wav")},
-            timeout=60,
-        )
 
-    resp.raise_for_status()
-    print("✓", resp.json())
+@app.function()
+@modal.asgi_app()
+def fastapi_app():
+    return web_app