server.py

#!/usr/bin/env python3
"""
voce — minimal STT server
Serves index.html, handles audio upload + SenseVoice transcription.
Port 3456.
"""

import os
import sys
import json
import uuid
import time
import mimetypes
import traceback
from pathlib import Path
from http.server import HTTPServer, BaseHTTPRequestHandler
from urllib.parse import urlparse

# ─── Paths ─────────────────────────────────────────────────────────────────────
BASE_DIR    = Path(__file__).parent.resolve()
UPLOAD_DIR  = BASE_DIR / "uploads"
UPLOAD_DIR.mkdir(exist_ok=True)

# ─── CUDA probe (defined before model load) ────────────────────────────────────
def _cuda_available():
    try:
        import torch
        return torch.cuda.is_available()
    except ImportError:
        return False


# ─── Model (loaded once at startup) ───────────────────────────────────────────
print("[voce] Loading SenseVoice model… (first run may download weights)", flush=True)
try:
    from funasr import AutoModel

    model = AutoModel(
        model="FunAudioLLM/SenseVoiceSmall",
        trust_remote_code=True,
        remote_code="./model.py",
        vad_model="fsmn-vad",
        vad_kwargs={"max_single_segment_time": 30000},
        device="cuda" if _cuda_available() else "cpu",
    )
    MODEL_LOADED = True
    print("[voce] Model ready.", flush=True)
except Exception as e:
    print(f"[voce] WARNING: Could not load SenseVoice model: {e}", flush=True)
    print("[voce]   STT endpoint will return a stub response.", flush=True)
    model = None
    MODEL_LOADED = False


# ─── Transcription ─────────────────────────────────────────────────────────────
EMOTION_MAP = {
    "<|HAPPY|>":    "happy",
    "<|SAD|>":      "sad",
    "<|ANGRY|>":    "angry",
    "<|NEUTRAL|>":  "neutral",
    "<|SURPRISED|>":"surprised",
    "<|FEARFUL|>":  "fearful",
    "<|DISGUSTED|>":"disgusted",
}

def parse_sense_voice_output(raw_text: str):
    """
    SenseVoice returns text like:
      <|zh|><|HAPPY|><|Speech|>你好世界
    Strip the tags and extract emotion.
    """
    import re
    emotion = "neutral"
    text = raw_text

    for tag, emo in EMOTION_MAP.items():
        if tag in raw_text:
            emotion = emo
            break

    text = re.sub(r"<\|[^|]+\|>", "", raw_text).strip()
    return text, emotion


def transcribe(audio_path: Path):
    if not MODEL_LOADED or model is None:
        # Stub: return placeholder so UI still works without GPU
        return {
            "text":    f"[SenseVoice not loaded] audio: {audio_path.name}",
            "emotion": "neutral",
        }

    try:
        res = model.generate(
            input=str(audio_path),
            cache={},
            language="auto",
            use_itn=True,
            batch_size_s=60,
            merge_vad=True,
            merge_length_s=15,
        )
        # res is a list of dicts; grab the first text
        raw = res[0]["text"] if res and "text" in res[0] else ""
        text, emotion = parse_sense_voice_output(raw)
        return {"text": text, "emotion": emotion}
    except Exception as e:
        traceback.print_exc()
        return {"text": "", "emotion": "neutral", "error": str(e)}


# ─── HTTP handler ──────────────────────────────────────────────────────────────
CORS_HEADERS = {
    "Access-Control-Allow-Origin":  "*",
    "Access-Control-Allow-Methods": "GET, POST, OPTIONS",
    "Access-Control-Allow-Headers": "Content-Type",
}

class VoceHandler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"
    server_version   = "voce/1.0"

    # ── Logging ──────────────────────────────────────────────────────────────
    def log_message(self, fmt, *args):
        print(f"[{self.address_string()}] {fmt % args}", flush=True)

    # ── Helpers ───────────────────────────────────────────────────────────────
    def _set_headers(self, code=200, content_type="application/json", extra=None):
        self.send_response(code)
        self.send_header("Content-Type", content_type)
        for k, v in CORS_HEADERS.items():
            self.send_header(k, v)
        if extra:
            for k, v in extra.items():
                self.send_header(k, v)
        self.end_headers()

    def _json(self, data, code=200):
        body = json.dumps(data, ensure_ascii=False).encode("utf-8")
        self._set_headers(code, extra={"Content-Length": str(len(body))})
        self.wfile.write(body)

    def _error(self, code, msg):
        self._json({"error": msg}, code)

    def _read_body(self, length=None):
        n = length or int(self.headers.get("Content-Length", 0))
        return self.rfile.read(n)

    # ── Routing ───────────────────────────────────────────────────────────────
    def do_OPTIONS(self):
        self._set_headers(204, "text/plain")

    def do_GET(self):
        path = urlparse(self.path).path
        if path in ("/", "/index.html"):
            self._serve_file(BASE_DIR / "index.html", "text/html; charset=utf-8")
        else:
            self._error(404, "Not found")

    def do_POST(self):
        path = urlparse(self.path).path
        if path == "/upload":
            self._handle_upload()
        elif path == "/stt":
            self._handle_stt()
        else:
            self._error(404, "Not found")

    # ── POST /upload ──────────────────────────────────────────────────────────
    def _handle_upload(self):
        content_type = self.headers.get("Content-Type", "")
        if "multipart/form-data" not in content_type:
            return self._error(400, "Expected multipart/form-data")

        try:
            body = self._read_body()
            filename, audio_bytes = _parse_multipart(body, content_type)
        except Exception as e:
            return self._error(400, f"Failed to parse multipart: {e}")

        if not audio_bytes:
            return self._error(400, "No audio data")

        # Sanitise extension
        ext = Path(filename).suffix.lower() if filename else ".webm"
        if ext not in {".webm", ".ogg", ".mp3", ".wav", ".m4a", ".flac", ".opus"}:
            ext = ".webm"

        uid = uuid.uuid4().hex[:12]
        safe_name = f"{uid}{ext}"
        dest = UPLOAD_DIR / safe_name

        dest.write_bytes(audio_bytes)
        print(f"[upload] Saved {safe_name} ({len(audio_bytes):,} bytes)", flush=True)

        self._json({"filename": safe_name, "size": len(audio_bytes)})

    # ── POST /stt ─────────────────────────────────────────────────────────────
    def _handle_stt(self):
        try:
            body = self._read_body()
            payload = json.loads(body)
        except Exception:
            return self._error(400, "Invalid JSON body")

        audio_path_str = payload.get("audio_path", "")
        if not audio_path_str:
            return self._error(400, "Missing audio_path")

        # Security: only allow filenames (no path traversal)
        audio_path = UPLOAD_DIR / Path(audio_path_str).name
        if not audio_path.exists():
            return self._error(404, f"Audio file not found: {audio_path_str}")

        t0 = time.time()
        result = transcribe(audio_path)
        elapsed = round(time.time() - t0, 2)
        result["elapsed_s"] = elapsed
        print(f"[stt] {audio_path.name} → emotion={result.get('emotion')} ({elapsed}s)", flush=True)

        self._json(result)

    # ── Static file serving ───────────────────────────────────────────────────
    def _serve_file(self, path: Path, content_type: str):
        if not path.exists():
            return self._error(404, "File not found")
        data = path.read_bytes()
        self._set_headers(200, content_type, {"Content-Length": str(len(data))})
        self.wfile.write(data)


# ─── Minimal multipart parser ──────────────────────────────────────────────────
def _parse_multipart(body: bytes, content_type: str):
    """
    Extract the first file part from a multipart/form-data body.
    Returns (filename, bytes).
    No external dependencies.
    """
    # Extract boundary
    boundary = None
    for part in content_type.split(";"):
        part = part.strip()
        if part.startswith("boundary="):
            boundary = part[len("boundary="):].strip('"')
            break
    if not boundary:
        raise ValueError("No boundary in Content-Type")

    sep = b"--" + boundary.encode()
    parts = body.split(sep)

    for raw in parts[1:]:  # skip preamble
        if raw.strip() in (b"", b"--", b"--\r\n"):
            continue
        if raw.startswith(b"--"):
            continue

        # Split headers / body
        if b"\r\n\r\n" in raw:
            headers_bytes, file_body = raw.split(b"\r\n\r\n", 1)
        elif b"\n\n" in raw:
            headers_bytes, file_body = raw.split(b"\n\n", 1)
        else:
            continue

        headers_str = headers_bytes.decode("utf-8", errors="replace")
        if "filename" not in headers_str:
            continue

        # Extract filename
        filename = "upload.webm"
        for line in headers_str.splitlines():
            if "Content-Disposition" in line and "filename=" in line:
                for seg in line.split(";"):
                    seg = seg.strip()
                    if seg.startswith("filename="):
                        filename = seg[9:].strip('"\'')
                        break

        # Strip trailing boundary marker
        file_body = file_body.rstrip(b"\r\n")
        if file_body.endswith(b"--"):
            file_body = file_body[:-2].rstrip(b"\r\n")

        return filename, file_body

    return None, b""


# ─── Entry point ───────────────────────────────────────────────────────────────
def main():
    host = os.environ.get("HOST", "0.0.0.0")
    port = int(os.environ.get("PORT", 3456))

    server = HTTPServer((host, port), VoceHandler)
    server.allow_reuse_address = True

    print(f"[voce] Listening on http://{host}:{port}", flush=True)
    print(f"[voce] Uploads → {UPLOAD_DIR}", flush=True)
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\n[voce] Shutting down.", flush=True)
        server.server_close()


if __name__ == "__main__":
    main()