modal_deploy/deploy_omni.py

"""
PregoPal x MiniCPM-o-4_5 - Modal deploy (llama.cpp-omni full-duplex voice upgrade)

Architecture:
  FastAPI (ASGI) <-> llama-server (OpenBMB/llama.cpp-omni subprocess)
                         |
                  Modal Volume: GGUF models (vision + audio + TTS)

Usage:
    pip install modal
    modal token new
    modal deploy modal_deploy.deploy_omni

Test:
    modal run -m modal_deploy.deploy_omni::test_inference
    modal run -m modal_deploy.deploy_omni::diagnose_volume

API:
    POST /v1/chat/completions   - OpenAI compatible (text + multimodal, streaming)
    POST /v1/audio/speech       - TTS: text -> voice WAV
    POST /v1/audio/transcriptions - STT: voice -> text
    POST /v1/embeddings         - Embeddings
    GET  /health                - Health check (audio/vision/TTS status)
    GET  /v1/models             - Model list
"""

import os
import modal
from modal import Image, App, Volume, asgi_app

# ============================================================================
# 1. IMAGE - Build OpenBMB/llama.cpp-omni from source
#    Source is copied from local llamacpp_omni/ (repo no longer public on GitHub)
# ============================================================================

_omni_image = (
    Image.debian_slim(python_version="3.11")
    .apt_install(
        "curl",
    )
    # Install CUDA Toolkit for compiling llama.cpp CUDA kernels
    .run_commands(
        "curl -L -o /tmp/cuda-keyring.deb https://developer.download.nvidia.com/compute/cuda/repos/debian12/x86_64/cuda-keyring_1.1-1_all.deb",
        "dpkg -i /tmp/cuda-keyring.deb",
        "apt-get update",
        "apt-get install -y cuda-toolkit-12-4 cuda-compiler-12-4 cuda-driver-dev-12-4",
    )
    .apt_install(
        "curl",
        "git",
        "build-essential",
        "cmake",
        "libcurl4-openssl-dev",
        "libsndfile1",
        "libasound2-dev",
        "pkg-config",
    )
    .pip_install(
        "fastapi",
        "uvicorn[standard]",
        "httpx",
        "numpy",
        "Pillow",
        "soundfile",
    )
    # Copy local llamacpp_omni source into image (repo no longer public)
    .add_local_dir(
        os.path.join(os.path.dirname(os.path.abspath(__file__)), "llamacpp_omni"),
        "/llama.cpp-omni",
        copy=True,
    )
.run_commands(
        "cd /llama.cpp-omni && cmake -B build "
        "-DGGML_CUDA=ON "
        "-DLLAMA_BUILD_SERVER=ON "
        "-DLLAMA_BUILD_TESTS=OFF "
        "-DLLAMA_BUILD_EXAMPLES=OFF "
        "-DLLAMA_CUDA_FORCE_MMQ=ON "
        "-DGGML_CUDA_NO_VMM=ON "
        "-DCMAKE_CUDA_ARCHITECTURES='75;89' "
        "-DCMAKE_BUILD_TYPE=Release "
        "-DCMAKE_CUDA_COMPILER=/usr/local/cuda-12/bin/nvcc",
        # Build llama-server first (includes all CUDA kernels + main libs)
        "cd /llama.cpp-omni && cmake --build build --config Release -j $(nproc) --target llama-server",
        # Then build llama-omni-server (links against already compiled libs)
        "cd /llama.cpp-omni && touch tools/server/server-omni.cpp",
        "cd /llama.cpp-omni && cmake --build build --config Release -j $(nproc) --target llama-omni-server",
        "ls -lh /llama.cpp-omni/build/bin/llama-server /llama.cpp-omni/build/bin/llama-omni-server",
    )
)

# ============================================================================
# 2. CONSTANTS
# ============================================================================

MODEL_DIR = "/models"
MODEL_SUBDIR = f"{MODEL_DIR}/MiniCPM-o-4_5-gguf"

MAIN_GGUF = "MiniCPM-o-4_5-Q4_K_M.gguf"
VISION_MMPROJ = "vision/MiniCPM-o-4_5-vision-F16.gguf"
AUDIO_MMPROJ = "audio/MiniCPM-o-4_5-audio-F16.gguf"
TTS_BASE_LM = "tts/MiniCPM-o-4_5-tts-F16.gguf"
TTS_ACOUSTIC = "tts/MiniCPM-o-4_5-projector-F16.gguf"
TOKEN2WAV_DIR = "token2wav-gguf"

LLAMA_SERVER_PORT = 8081
OMNI_SERVER_PORT = 8082
OMNI_TTS_WAV_DIR = "/tmp/omni_output"

model_volume = Volume.from_name("minicpm-o-4_5-models", create_if_missing=True)
app = App("prego-pal-minicpm-omni")


def get_model_paths(base_dir: str) -> dict:
    paths = {
        "main": os.path.join(base_dir, MAIN_GGUF),
        "vision": os.path.join(base_dir, VISION_MMPROJ),
        "audio": os.path.join(base_dir, AUDIO_MMPROJ),
        "tts_base_lm": os.path.join(base_dir, TTS_BASE_LM),
        "tts_acoustic": os.path.join(base_dir, TTS_ACOUSTIC),
        "token2wav_dir": os.path.join(base_dir, TOKEN2WAV_DIR),
    }
    for key, path in paths.items():
        if key == "token2wav_dir":
            exists = os.path.isdir(path)
        else:
            exists = os.path.isfile(path)
        print(f"[PregoPal] {key}: {path} (exists={exists})")
    return paths


# ============================================================================
# 3. ASGI APP - FastAPI lifespan + llama-server + llama-omni-server
# ============================================================================

@app.function(
    image=_omni_image,
    volumes={MODEL_DIR: model_volume},
    gpu="T4",
    timeout=1200,
    scaledown_window=300,
)
@modal.concurrent(max_inputs=10)
@asgi_app()
def serve():
    """
    FastAPI ASGI app. Launches llama-server + llama-omni-server subprocesses.
    serve() is sync; async logic lives in lifespan context manager.
    """
    import asyncio
    import json
    import logging
    import subprocess
    from contextlib import asynccontextmanager
    from fastapi import FastAPI, Request
    from fastapi.responses import StreamingResponse, JSONResponse
    from fastapi.middleware.cors import CORSMiddleware
    import httpx

    logging.basicConfig(level=logging.INFO)
    logger = logging.getLogger("prego-pal-omni")

    paths = get_model_paths(MODEL_SUBDIR)

    # Build llama-server command
    llama_server_bin = "/llama.cpp-omni/build/bin/llama-server"
    if not os.path.isfile(llama_server_bin):
        llama_server_bin = "/llama.cpp-omni/build/bin/Release/llama-server"

    cmd = [
        llama_server_bin,
        "-m", paths["main"],
        "--mmproj", paths["vision"],
        "--host", "127.0.0.1",
        "--port", str(LLAMA_SERVER_PORT),
        "-ngl", "99",
        "-c", "8192",
        "--no-mmap",
        "--jinja",
    ]

    # llama-omni-server binary (standalone omni HTTP API for full-duplex)
    omni_server_bin = "/llama.cpp-omni/build/bin/llama-omni-server"
    if not os.path.isfile(omni_server_bin):
        omni_server_bin = "/llama.cpp-omni/build/bin/Release/llama-omni-server"

    omni_cmd = [
        omni_server_bin,
        "--port", str(OMNI_SERVER_PORT),
        "--host", "127.0.0.1",
        "-ngl", "99",
    ]

    # Check token2wav and TTS model directories
    # llama-omni-server auto-detects model files from model_dir
    # via omni_init API body. No need to pass them as CLI args.

    @asynccontextmanager
    async def lifespan(web_app: FastAPI):
        """Async lifecycle: start both subprocesses, cleanup on shutdown."""
        processes = []

        # Start llama-server (standard chat/tts/stt)
        logger.info("[PregoPal] Starting llama-server...")
        ls_proc = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )
        processes.append(("llama-server", ls_proc))

        # Start llama-omni-server (full-duplex omni API)
        logger.info(f"[PregoPal] Starting llama-omni-server on port {OMNI_SERVER_PORT}...")
        os.makedirs(OMNI_TTS_WAV_DIR, exist_ok=True)
        omni_proc = subprocess.Popen(
            omni_cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )
        processes.append(("llama-omni-server", omni_proc))

        base_url = f"http://127.0.0.1:{LLAMA_SERVER_PORT}"
        omni_url = f"http://127.0.0.1:{OMNI_SERVER_PORT}"
        ls_ready = False
        omni_ready = False
        for i in range(45):
            await asyncio.sleep(2)
            try:
                async with httpx.AsyncClient(timeout=5.0) as client:
                    if not ls_ready:
                        r = await client.get(f"{base_url}/health")
                        if r.status_code == 200:
                            ls_ready = True
                            logger.info(f"[PregoPal] llama-server ready (attempt {i+1})")
                    if not omni_ready:
                        try:
                            r2 = await client.get(f"{omni_url}/health")
                            if r2.status_code == 200:
                                omni_ready = True
                                logger.info(f"[PregoPal] llama-omni-server ready (attempt {i+1})")
                        except Exception:
                            pass
            except Exception:
                if i > 0 and i % 5 == 0:
                    logger.info(f"[PregoPal] Waiting (llama-server={ls_ready}, omni-server={omni_ready})...")

        if not ls_ready:
            for name, proc in processes:
                stderr_lines = []
                try:
                    for _ in range(20):
                        line = proc.stderr.readline()
                        if line:
                            stderr_lines.append(line.strip())
                except Exception:
                    pass
                logger.error(f"[PregoPal] {name} stderr:\n" + "\n".join(stderr_lines[-10:]))
                proc.terminate()
            raise RuntimeError("llama-server failed to start within 90s")

        if not omni_ready:
            logger.warning("[PregoPal] llama-omni-server not ready - omni endpoints will be unavailable")

        web_app.state.llama_base_url = base_url
        web_app.state.llama_base_url = omni_url
        web_app.state.llama_client = httpx.AsyncClient(base_url=base_url, timeout=120.0)
        web_app.state.llama_client = httpx.AsyncClient(base_url=omni_url, timeout=600.0)
        web_app.state.omni_temp_dir = OMNI_TTS_WAV_DIR

        yield

        logger.info("[PregoPal] Shutting down...")
        for name, proc in processes:
            logger.info(f"[PregoPal] Terminating {name}...")
            proc.terminate()
            proc.wait(timeout=30)
        await web_app.state.llama_client.aclose()
        logger.info("[PregoPal] Shutdown complete")

    web_app = FastAPI(
        title="PregoPal MiniCPM-o-4_5 Omni API",
        lifespan=lifespan,
    )
    web_app.add_middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )

    base_url = f"http://127.0.0.1:{LLAMA_SERVER_PORT}"

    # ---- Proxy Endpoints ----

    @web_app.post("/v1/chat/completions")
    async def chat_completions(request: Request):
        body = await request.json()
        stream = body.get("stream", False)
        client = web_app.state.llama_client

        if stream:
            async def event_stream():
                async with httpx.AsyncClient(timeout=120.0) as sclient:
                    async with sclient.stream(
                        "POST", f"{base_url}/v1/chat/completions", json=body
                    ) as resp:
                        async for chunk in resp.aiter_lines():
                            if chunk:
                                yield chunk + "\n"
            return StreamingResponse(event_stream(), media_type="text/event-stream")

        try:
            resp = await client.post("/v1/chat/completions", json=body)
            return JSONResponse(resp.json(), status_code=resp.status_code)
        except Exception as e:
            logger.error(f"[PregoPal] Chat completion proxy error: {e}")
            return JSONResponse({"error": str(e)}, status_code=502)

    @web_app.post("/v1/audio/speech")
    async def audio_speech(request: Request):
        """TTS: text -> speech WAV"""
        body = await request.json()
        client = web_app.state.llama_client
        try:
            resp = await client.post("/v1/audio/speech", json=body)
            return StreamingResponse(
                resp.aiter_bytes(),
                media_type=resp.headers.get("content-type", "audio/wav"),
            )
        except Exception as e:
            logger.error(f"[PregoPal] TTS error: {e}")
            return JSONResponse({"error": str(e)}, status_code=502)

    @web_app.post("/v1/audio/speech/stream")
    async def audio_speech_stream(request: Request):
        """Streaming TTS"""
        body = await request.json()
        try:
            async with httpx.AsyncClient(timeout=120.0) as sclient:
                async with sclient.stream(
                    "POST", f"{base_url}/v1/audio/speech/stream", json=body
                ) as resp:
                    async def audio_stream():
                        async for chunk in resp.aiter_bytes():
                            yield chunk
                    return StreamingResponse(
                        audio_stream(),
                        media_type=resp.headers.get("content-type", "audio/wav"),
                    )
        except Exception as e:
            logger.error(f"[PregoPal] Stream TTS error: {e}")
            return JSONResponse({"error": str(e)}, status_code=502)

    @web_app.post("/v1/audio/transcriptions")
    async def audio_transcriptions(request: Request):
        """STT: speech -> text"""
        body = await request.json()
        client = web_app.state.llama_client
        try:
            resp = await client.post("/v1/audio/transcriptions", json=body)
            return JSONResponse(resp.json(), status_code=resp.status_code)
        except Exception as e:
            logger.error(f"[PregoPal] STT error: {e}")
            return JSONResponse({"error": str(e)}, status_code=502)

    @web_app.post("/v1/embeddings")
    async def embeddings(request: Request):
        body = await request.json()
        client = web_app.state.llama_client
        try:
            resp = await client.post("/v1/embeddings", json=body)
            return JSONResponse(resp.json(), status_code=resp.status_code)
        except Exception as e:
            logger.error(f"[PregoPal] Embeddings proxy error: {e}")
            return JSONResponse({"error": str(e)}, status_code=502)

    @web_app.get("/health")
    async def health():
        try:
            client = web_app.state.llama_client
            ls_resp = await client.get("/health")
            ls_status = ls_resp.json()
        except Exception as e:
            ls_status = {"error": str(e)}
        omni_status = {}
        try:
            oc = web_app.state.llama_client
            omni_r = await oc.get("/health")
            omni_status = omni_r.json()
        except Exception as e:
            omni_status = {"error": str(e)}
        return {
            "status": "ok",
            "model": "MiniCPM-o-4_5",
            "engine": "llama.cpp-omni",
            "cuda": True,
            "vision": os.path.isfile(paths["vision"]),
            "audio": os.path.isfile(paths["audio"]),
            "tts_base_lm": os.path.isfile(paths["tts_base_lm"]),
            "tts_acoustic": os.path.isfile(paths["tts_acoustic"]),
            "token2wav_dir": os.path.isdir(paths["token2wav_dir"]),
            "llama_server_status": ls_status,
            "omni_server_status": omni_status,
        }

    @web_app.get("/v1/models")
    async def list_models():
        try:
            client = web_app.state.llama_client
            resp = await client.get("/v1/models")
            return JSONResponse(resp.json(), status_code=resp.status_code)
        except Exception:
            return JSONResponse({
                "object": "list",
                "data": [{
                    "id": "MiniCPM-o-4_5",
                    "object": "model",
                    "created": 1,
                    "owned_by": "prego-pal",
                }],
            })

    # ===========================================================================
    # Omni Full-Duplex Voice Endpoints
    # Proxy/compatibility layer between frontend and llama-omni-server
    # ===========================================================================

    @web_app.post("/v1/omni/init")
    async def omni_init(request: Request):
        """
        Initialize omni context on llama-omni-server.
        Body: {
            "session_id": str,
            "media_type": int (0=text, 1=image, 2=audio+vision),
            "use_tts": bool,
            "duplex_mode": bool,
            "model_dir": str (omit to use default /models/MiniCPM-o-4_5-gguf),
            "voice_audio": str (base64 audio for voice cloning, optional),
            "voice_clone_prompt": str (optional),
            "assistant_prompt": str (optional),
        }
        """
        body = await request.json()
        body.setdefault("media_type", 2)
        body.setdefault("use_tts", True)
        body.setdefault("duplex_mode", True)
        body.setdefault("model_dir", os.path.join(MODEL_DIR, "MiniCPM-o-4_5-gguf"))
        body.setdefault("tts_bin_dir", body["model_dir"] + "/tts")
        body.setdefault("output_dir", OMNI_TTS_WAV_DIR)
        body.setdefault("token2wav_device", "gpu:0")

        oc = web_app.state.llama_client
        try:
            resp = await oc.post("/v1/stream/omni_init", json=body)
            return JSONResponse(resp.json(), status_code=resp.status_code)
        except Exception as e:
            logger.error(f"[Omni] Init error: {e}")
            return JSONResponse({"error": str(e)}, status_code=502)

    @web_app.post("/v1/omni/prefill")
    async def omni_prefill(request: Request):
        """
        Send audio/image to omni for streaming prefill.
        Body: {
            "cnt": int,  # chunk counter
            "audio": str (base64-encoded PCM 16kHz 16-bit, optional),
            "image": str (base64-encoded JPEG, optional),
            "text": str (optional, text input),
            "last_chunk": bool (True if this is the final chunk),
        }
        The audio/image base64 is decoded to a temp file, then proxied.
        """
        body = await request.json()
        cnt = body.get("cnt", 0)
        temp_dir = web_app.state.omni_temp_dir

        # Decode audio base64 -> temp WAV file
        audio_path = ""
        audio_b64 = body.get("audio", "")
        if audio_b64:
            import base64, io, soundfile as sf, numpy as np
            audio_bytes = base64.b64decode(audio_b64)
            audio_np = np.frombuffer(audio_bytes, dtype=np.int16).astype(np.float32) / 32768.0
            temp_audio = os.path.join(temp_dir, f"prefill_{cnt}.wav")
            sf.write(temp_audio, audio_np, 16000, format='WAV', subtype='PCM_16')
            audio_path = temp_audio

        # Decode image base64 -> temp PNG file
        img_path = ""
        img_b64 = body.get("image", "")
        if img_b64:
            import base64
            img_bytes = base64.b64decode(img_b64)
            temp_img = os.path.join(temp_dir, f"prefill_{cnt}.png")
            with open(temp_img, "wb") as f:
                f.write(img_bytes)
            img_path = temp_img

        oc = web_app.state.llama_client
        cpp_req = {
            "audio_path_prefix": audio_path,
            "img_path_prefix": img_path,
            "cnt": cnt,
        }
        text = body.get("text", "")
        if text:
            cpp_req["text"] = text
        if "max_slice_nums" in body:
            cpp_req["max_slice_nums"] = body["max_slice_nums"]

        try:
            resp = await oc.post("/v1/stream/prefill", json=cpp_req)
            return JSONResponse(resp.json(), status_code=resp.status_code)
        except Exception as e:
            logger.error(f"[Omni] Prefill error: {e}")
            return JSONResponse({"error": str(e)}, status_code=502)

    @web_app.post("/v1/omni/generate")
    async def omni_generate(request: Request):
        """
        Start omni streaming generation (SSE).
        Proxies llama-omni-server /v1/stream/decode SSE.
        The SSE output contains {"content": str, "stop": bool, ...}
        where content is mixed text + audio tokens.
        Audio tokens are decoded to WAV by llama-omni-server internally
        (token2wav) and written to OMNI_TTS_WAV_DIR as wav_NNN.wav.
        """
        body = await request.json()
        debug_dir = body.get("debug_dir", OMNI_TTS_WAV_DIR)
        stream = body.get("stream", True)
        round_idx = body.get("round_idx", -1)
        length_penalty = body.get("length_penalty", None)

        oc = web_app.state.llama_client
        cpp_req = {
            "debug_dir": debug_dir,
            "stream": stream,
            "round_idx": round_idx,
        }
        if length_penalty is not None:
            cpp_req["length_penalty"] = length_penalty

        async def sse_proxy():
            async with httpx.AsyncClient(base_url=web_app.state.llama_base_url, timeout=600.0) as sclient:
                async with sclient.stream("POST", "/v1/stream/decode", json=cpp_req) as resp:
                    async for chunk in resp.aiter_lines():
                        if chunk:
                            yield chunk + "\n"

        return StreamingResponse(sse_proxy(), media_type="text/event-stream")

    @web_app.post("/v1/omni/break")
    async def omni_break():
        """Break: re-initialize omni context (current gen is discarded).
        llama-omni-server has no standalone break endpoint, so we
        reinit via /v1/stream/omni_init which frees old context.
        """
        oc = web_app.state.llama_client
        # Re-init with same defaults — this frees any active omni context
        init_body = {
            "media_type": 2,
            "use_tts": True,
            "duplex_mode": True,
            "model_dir": os.path.join(MODEL_DIR, "MiniCPM-o-4_5-gguf"),
        }
        try:
            resp = await oc.post("/v1/stream/omni_init", json=init_body)
            return JSONResponse({"status": "reinitialized", "response": resp.json()}, status_code=resp.status_code)
        except Exception as e:
            logger.error(f"[Omni] Break error: {e}")
            return JSONResponse({"error": str(e)}, status_code=502)

    @web_app.post("/v1/omni/stop")
    async def omni_stop(request: Request):
        """Stop and free omni context by reinit with empty state."""
        oc = web_app.state.llama_client
        init_body = {
            "media_type": 0,
            "use_tts": False,
            "duplex_mode": False,
            "model_dir": os.path.join(MODEL_DIR, "MiniCPM-o-4_5-gguf"),
        }
        try:
            resp = await oc.post("/v1/stream/omni_init", json=init_body)
            return JSONResponse({"status": "stopped", "response": resp.json()}, status_code=resp.status_code)
        except Exception as e:
            logger.error(f"[Omni] Stop error: {e}")
            return JSONResponse({"error": str(e)}, status_code=502)

    @web_app.get("/v1/omni/tts_wav/{round_dir:path}/{filename:path}")
    async def omni_tts_wav(round_dir: str, filename: str):
        """Serve TTS WAV files generated by llama-omni-server."""
        import os
        wav_path = os.path.join(OMNI_TTS_WAV_DIR, round_dir, filename)
        if not os.path.isfile(wav_path):
            return JSONResponse({"error": f"WAV not found: {wav_path}"}, status_code=404)
        from fastapi.responses import FileResponse
        return FileResponse(wav_path, media_type="audio/wav")

    @web_app.get("/")
    async def root():
        return {
            "service": "PregoPal MiniCPM-o-4_5 Omni API",
            "version": "3.0.0",
            "model": MAIN_GGUF,
            "engine": "llama.cpp-omni (OpenBMB)",
            "endpoints": {
                "chat": "POST /v1/chat/completions (text+multimodal, streaming)",
                "tts": "POST /v1/audio/speech (text->speech)",
                "tts_stream": "POST /v1/audio/speech/stream (streaming TTS)",
                "stt": "POST /v1/audio/transcriptions (speech->text)",
                "embeddings": "POST /v1/embeddings",
                "models": "GET /v1/models",
                "health": "GET /health",
                "omni_init": "POST /v1/omni/init",
                "omni_prefill": "POST /v1/omni/prefill (base64 audio)",
                "omni_generate": "POST /v1/omni/generate (SSE)",
                "omni_break": "POST /v1/omni/break",
                "omni_stop": "POST /v1/omni/stop",
                "omni_tts_wav": "GET /v1/omni/tts_wav/{round_dir}/{filename}",
            },
        }

    return web_app


# ============================================================================
# 4. DIAGNOSE VOLUME
# ============================================================================

@app.function(
    image=_omni_image,
    volumes={MODEL_DIR: model_volume},
    timeout=120,
)
def diagnose_volume():
    """Check model file integrity in Modal Volume."""
    print(f"\n{'='*60}")
    print(f"[Diagnose] {MODEL_SUBDIR}")
    print(f"{'='*60}")
    for root, dirs, files in os.walk(MODEL_SUBDIR):
        level = root.replace(MODEL_SUBDIR, "").count(os.sep)
        indent = " " * 2 * level
        print(f"{indent}{os.path.basename(root)}/")
        subindent = " " * 2 * (level + 1)
        for file in sorted(files):
            fpath = os.path.join(root, file)
            size = os.path.getsize(fpath)
            print(f"{subindent}{file}  ({size:,} bytes = {size/1024**3:.2f} GB)")

    paths = get_model_paths(MODEL_SUBDIR)
    all_ok = True
    for key, path in paths.items():
        if key == "token2wav_dir":
            ok = os.path.isdir(path)
        else:
            ok = os.path.isfile(path)
        status = "OK" if ok else "MISSING"
        if not ok:
            all_ok = False
        print(f"  [{status}] {key}: {path}")

    if all_ok:
        print(f"\n[OK] All model files found! Ready to deploy.")
    else:
        print(f"\n[FAIL] Some files missing. Check uploads.")

    main_path = paths["main"]
    if os.path.isfile(main_path):
        with open(main_path, "rb") as f:
            magic = f.read(4)
        if magic == b"GGUF":
            print("[OK] Main model is valid GGUF")
        else:
            print(f"[WARN] Main model NOT valid GGUF (magic={magic.hex()})")

    print(f"\n{'='*60}")
    print(f"[Diagnose] CUDA library files in container")
    print(f"{'='*60}")
    import subprocess
    result = subprocess.run(
        "find / -name 'libcuda*' -type f,l 2>/dev/null | head -30",
        shell=True, capture_output=True, text=True
    )
    cuda_files = result.stdout.strip().split("\\n")
    for f in cuda_files:
        if f:
            size = os.path.getsize(f) if os.path.exists(f) else 0
            print(f"  {f}  ({size:,} bytes)")


# ============================================================================
# 5. TEST INFERENCE (standalone - not via ASGI)
# ============================================================================

@app.function(
    image=_omni_image,
    volumes={MODEL_DIR: model_volume},
    gpu="T4",
    timeout=900,
)
def test_inference():
    """Test llama-server text+multimodal inference on Modal T4."""
    import subprocess
    import time
    import httpx

    print("[PregoPal] ========== TEST INFERENCE (llama-server) ==========")

    vision_path = os.path.join(MODEL_SUBDIR, VISION_MMPROJ)
    cmd = [
        "/llama.cpp-omni/build/bin/llama-server",
        "-m", os.path.join(MODEL_SUBDIR, MAIN_GGUF),
        "--mmproj", vision_path,
        "--host", "127.0.0.1",
        "--port", "8081",
        "-ngl", "99",
        "-c", "4096",
        "--no-mmap",
        "--jinja",
    ]

    print("[PregoPal] Starting llama-server...")
    print(f"[PregoPal] cmd: {' '.join(cmd)}")
    server_proc = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True, bufsize=1)

    # Collect stderr in a background thread
    stderr_lines = []
    import threading, queue
    q = queue.Queue()
    def _reader():
        for line in iter(server_proc.stderr.readline, ''):
            q.put(line.rstrip())
        q.put(None)
    thr = threading.Thread(target=_reader, daemon=True)
    thr.start()

    base_url = "http://127.0.0.1:8081"
    ready = False
    start = time.time()
    last_log = time.time()
    for i in range(180):
        time.sleep(1)
        now = time.time()
        # Drain stderr from queue (non-blocking)
        while True:
            try:
                line = q.get_nowait()
            except queue.Empty:
                break
            if line is None:
                break
            print(f"[llama-server] {line}")
            last_log = now
        
        # Check if process died
        ret = server_proc.poll()
        if ret is not None:
            print(f"[PregoPal] PROCESS EXITED with code {ret}")
            # Drain remaining
            while True:
                try:
                    line = q.get_nowait()
                except queue.Empty:
                    break
                if line is None:
                    break
                print(f"[llama-server] {line}")
            break
        
        # Only print waiting msg every 10s
        elapsed = now - start
        if elapsed - (i // 10) * 10 < 2:
            try:
                r = httpx.get(f"{base_url}/health", timeout=3.0)
                if r.status_code == 200:
                    ready = True
                    print(f"[PregoPal] llama-server ready ({elapsed:.0f}s)")
                    break
            except Exception:
                pass
        
        if i % 10 == 0:
            print(f"[PregoPal] Waiting ({elapsed:.0f}s)...")

    if not ready:
        # Drain remaining stderr
        time.sleep(0.5)
        while True:
            try:
                line = q.get_nowait()
            except queue.Empty:
                break
            if line is None:
                break
            print(f"[llama-server] {line}")
        print(f"[PregoPal] Timed out ({time.time()-start:.0f}s). Check above for [llama-server] lines.")
        server_proc.terminate()
        return

    client = httpx.Client(base_url=base_url, timeout=120.0)

    try:
        # Test 1: Simple text (Chinese)
        print("\n[Test 1] Chinese...")
        t0 = time.time()
        resp = client.post("/v1/chat/completions", json={
            "messages": [{"role": "user", "content": "你好，请简单介绍一下你自己"}],
            "max_tokens": 100, "temperature": 0.3,
        })
        t1 = time.time()
        data = resp.json()
        content = data.get("choices", [{}])[0].get("message", {}).get("content", "")
        print(f"Response ({t1-t0:.1f}s): status={resp.status_code}")
        print(f"  content: {content}")
        print(f"  finish_reason: {data.get('choices',[{}])[0].get('finish_reason', 'N/A')}")
        print(f"  usage: {data.get('usage', {})}")

        # Test 2: Simple text (English) - need more tokens + higher temp
        print("\n[Test 2] English...")
        t0 = time.time()
        resp = client.post("/v1/chat/completions", json={
            "messages": [{"role": "user", "content": "What is the capital of France? Answer in one short sentence."}],
            "max_tokens": 100, "temperature": 0.5,
        })
        t1 = time.time()
        data = resp.json()
        content = data.get("choices", [{}])[0].get("message", {}).get("content", "")
        print(f"Response ({t1-t0:.1f}s): status={resp.status_code}")
        print(f"  content: {content}")
        print(f"  finish_reason: {data.get('choices',[{}])[0].get('finish_reason', 'N/A')}")
        print(f"  usage: {data.get('usage', {})}")

        # Test 3: Health
        print("\n[Test 3] Health...")
        resp = client.get("/health")

        # Test 3: Health
        print("\n[Test 3] Health...")
        resp = client.get("/health")
        info = resp.json()
        print(f"Health: model={info.get('model')}, cuda={info.get('cuda')}, "
              f"vision={info.get('vision')}, audio={info.get('audio')}")

        print(f"\n{'='*50}")
        print("[OK] All tests passed!")
        print(f"{'='*50}")

    except Exception as e:
        print(f"[PregoPal] Test error: {e}")
        raise
    finally:
        server_proc.terminate()
        server_proc.wait(timeout=10)


# ============================================================================
# 6. LOCAL ENTRY POINT
# ============================================================================

# ============================================================================
# 7. OMNI TEST (runs inside deployed Modal container)
# ============================================================================

@app.function(
    image=_omni_image,
    volumes={MODEL_DIR: model_volume},
    gpu="T4",
    timeout=600,
)
def test_omni():
    """
    Test llama-server built-in omni endpoints (/v1/stream/*).
    Launches llama-server with omni models and tests omni_init.
    Uses the same llama-server launch pattern as test_inference.
    """
    import subprocess
    import time
    import httpx
    import json

    print("[PregoPal] ========== TEST OMNI (llama-server built-in endpoints) ==========")

    paths = get_model_paths(MODEL_SUBDIR)

    # Build llama-server command with omni model paths
    ls_bin = "/llama.cpp-omni/build/bin/llama-server"
    llm_cmd = [
        ls_bin,
        "-m", paths["main"],
        "--mmproj", paths["vision"],
        "--host", "127.0.0.1",
        "--port", str(LLAMA_SERVER_PORT),
        "-ngl", "99",
        "-c", "8192",
        "--no-mmap",
        "--jinja",
    ]

    print(f"  Starting llama-server with omni support...")
    ls_proc = subprocess.Popen(llm_cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True, bufsize=1)

    base_url = f"http://127.0.0.1:{LLAMA_SERVER_PORT}"

    # Wait for server
    start = time.time()
    ready = False
    for _ in range(90):
        time.sleep(1)
        try:
            r = httpx.get(f"{base_url}/health", timeout=3)
            if r.status_code == 200:
                ready = True
                print(f"[PregoPal] llama-server ready ({time.time()-start:.0f}s)")
                break
        except:
            pass

    if not ready:
        print(f"[FAIL] llama-server not ready after {time.time()-start:.0f}s")
        # Read stderr
        import select
        stderr_lines = []
        for _ in range(50):
            line = ls_proc.stderr.readline()
            if not line:
                break
            stderr_lines.append(line.rstrip())
        print(f"[llama-server stderr] ({len(stderr_lines)} lines):")
        for line in stderr_lines[-30:]:
            print(f"  {line}")
        ls_proc.terminate()
        return

    # ====== Test omni_init ======
    print("\n=== 1. Omni Init ===")
    init_body = {
        "media_type": 2,      # audio + vision
        "use_tts": True,
        "duplex_mode": True,
        "model_dir": MODEL_SUBDIR,
        "tts_bin_dir": f"{MODEL_SUBDIR}/tts",
        "output_dir": OMNI_TTS_WAV_DIR,
        "token2wav_device": "gpu:0",
    }
    try:
        # llama-server has built-in /v1/stream/omni_init
        r = httpx.post(f"{base_url}/v1/stream/omni_init", json=init_body, timeout=120)
        print(f"Status: {r.status_code}")
        data = r.json()
        print(f"Body: {json.dumps(data, indent=2, ensure_ascii=False)}")
        if not data.get("success"):
            print("[FAIL] Omni Init failed")
            # Read stderr
            import select
            stderr_lines = []
            for _ in range(30):
                line = ls_proc.stderr.readline()
                if not line:
                    break
                stderr_lines.append(line.rstrip())
            print(f"[llama-server stderr] ({len(stderr_lines)} lines):")
            for line in stderr_lines[-20:]:
                print(f"  {line}")
            ls_proc.terminate()
            return
        print("\n=== [OK] Omni Init passed! ===")
        ls_proc.terminate()
    except Exception as e:
        print(f"[FAIL] Error: {e}")
        # Read stderr
        stderr_lines = []
        for _ in range(30):
            line = ls_proc.stderr.readline()
            if not line:
                break
            stderr_lines.append(line.rstrip())
        print(f"[llama-server stderr] ({len(stderr_lines)} lines):")
        for line in stderr_lines[-20:]:
            print(f"  {line}")
        ls_proc.terminate()
        return


if __name__ == "__main__":
    import sys
    if len(sys.argv) > 1:
        if sys.argv[1] == "test_inference":
            test_inference.local()
        elif sys.argv[1] == "diagnose_volume":
            diagnose_volume.local()
        elif sys.argv[1] == "test_omni":
            test_omni.local()