app.py

"""
Dedicated DiffusionGemma GGUF ZeroGPU API Space.

This is intentionally not a universal GGUF template. It targets only:
  unsloth/diffusiongemma-26B-A4B-it-GGUF

Runtime shape:
  gradio.Server -> @app.api queue -> @spaces.GPU -> llama-diffusion-cli subprocess

The OpenAI-compatible /v1/chat/completions route is provided for convenience, but
/gradio_api/call/chat remains the most native ZeroGPU path.
"""

from __future__ import annotations

# Keep this import before any torch import. This repo is GGUF/subprocess based and
# does not import torch, but the ordering is retained for ZeroGPU compatibility.
try:
    import spaces  # type: ignore
except Exception:  # pragma: no cover - local development fallback only
    class _LocalSpaces:
        def GPU(self, *args, **kwargs):
            def decorator(fn):
                return fn
            # Support both @spaces.GPU and @spaces.GPU(...)
            if args and callable(args[0]) and len(args) == 1 and not kwargs:
                return args[0]
            return decorator
    spaces = _LocalSpaces()  # type: ignore

import json
import os
import shlex
import subprocess
import threading
import time
import uuid
from typing import Any

from fastapi import Request
from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse
from gradio import Server

from src.config import settings
from src.errors import ApiError, public_error_payload
from src.handlers import run_final_from_payload
from src.models import MODEL_REGISTRY, default_model_id, public_model_list
from src.openai import (
    extract_openai_payload,
    make_chat_completion,
    make_sse_chunk,
    make_sse_done,
)
from src.prepare import prepare_runtime_if_requested, runtime_events, runtime_status

prepare_runtime_if_requested()

app = Server(title="DiffusionGemma GGUF ZeroGPU API")
_api_lock = threading.Lock()


def estimate_gpu_duration(model_path: str, prompt: str, params: dict[str, Any]) -> int:
    """Return requested ZeroGPU duration for one llama-diffusion-cli call."""
    max_tokens = int(params.get("max_tokens") or settings.default_max_tokens)
    max_steps = int(params.get("max_denoising_steps") or settings.diffusion_max_steps)

    # DiffusionGemma uses 256-token canvases. This is intentionally conservative
    # because GGUF + subprocess runtime differs from the official Transformers demo.
    blocks = max(1, (max_tokens + 255) // 256)
    seconds = settings.gpu_base_seconds + blocks * max_steps * settings.seconds_per_block_step
    seconds = max(settings.min_gpu_duration_seconds, int(seconds))
    return min(settings.max_gpu_duration_seconds, seconds)


@spaces.GPU(size=settings.zero_gpu_size, duration=estimate_gpu_duration)
def _gpu_generate(model_path: str, prompt: str, params: dict[str, Any]) -> dict[str, Any]:
    """The only function that actually invokes GPU compute."""
    from src.runner import run_diffusion_cli

    return run_diffusion_cli(model_path=model_path, prompt=prompt, params=params)


@app.api(
    name="chat",
    concurrency_limit=1,
    concurrency_id="diffusiongemma_gpu_queue",
    time_limit=settings.api_time_limit_seconds,
)
def chat(payload: dict[str, Any]) -> dict[str, Any]:
    """Primary Gradio queue endpoint.

    Expected payload:
      {
        "model": "unsloth/diffusiongemma-26B-A4B-it-GGUF:Q4_K_M",
        "messages": [{"role": "user", "content": "..."}],
        "max_tokens": 512,
        "thinking": false
      }
    """
    try:
        return run_final_from_payload(payload, gpu_generate=_gpu_generate)
    except ApiError as exc:
        return public_error_payload(exc)
    except Exception as exc:  # keep Gradio response JSON-serializable
        return public_error_payload(ApiError("internal_error", str(exc), status_code=500))


@app.get("/health")
async def health() -> dict[str, Any]:
    return {
        "ok": True,
        "service": "diffusiongemma-gguf-zerogpu-api",
        "default_model": default_model_id(),
        "zero_gpu_size": settings.zero_gpu_size,
        "runner": "llama-diffusion-cli",
        "runtime": runtime_status(),
    }


@app.get("/v1/models")
async def v1_models() -> dict[str, Any]:
    return {"object": "list", "data": public_model_list()}


def _probe_command(cmd: list[str], timeout: int = 12) -> dict[str, Any]:
    try:
        proc = subprocess.run(
            cmd,
            text=True,
            capture_output=True,
            timeout=timeout,
            check=False,
        )
        return {
            "cmd": shlex.join(cmd),
            "returncode": proc.returncode,
            "stdout_tail": (proc.stdout or "")[-2000:],
            "stderr_tail": (proc.stderr or "")[-2000:],
        }
    except Exception as exc:  # noqa: BLE001 - diagnostic endpoint
        return {"cmd": shlex.join(cmd), "error": str(exc)}


@spaces.GPU(size=settings.zero_gpu_size, duration=60)
def _gpu_probe(include_cli: bool = False) -> dict[str, Any]:
    diagnostics = [
        _probe_command(["nvidia-smi", "-L"]),
        _probe_command(["nvidia-smi"]),
        _probe_command(["bash", "-lc", "printf 'CUDA_VISIBLE_DEVICES=%s\\nLD_LIBRARY_PATH=%s\\n' \"$CUDA_VISIBLE_DEVICES\" \"$LD_LIBRARY_PATH\""]),
        _probe_command([
            "bash",
            "-lc",
            "ls -l ${CUDA_HOME:-/usr/local/cuda}/lib64/libcudart.so* "
            "${CUDA_HOME:-/usr/local/cuda}/lib64/libcublas.so* "
            "/usr/local/lib/python*/site-packages/nvidia/*/lib/libcudart.so* "
            "/usr/local/lib/python*/site-packages/nvidia/*/lib/libcublas.so* 2>/dev/null || true",
        ]),
    ]
    cli_result = None
    if include_cli and settings.llama_diffusion_bin.exists():
        cli_result = _probe_command([str(settings.llama_diffusion_bin), "--help"])

    return {
        "ok": True,
        "cuda_visible_devices": os.getenv("CUDA_VISIBLE_DEVICES", ""),
        "diagnostics": diagnostics,
        "llama_diffusion_cli_probe": cli_result,
        "note": "This endpoint runs inside @spaces.GPU; a MIG value in CUDA_VISIBLE_DEVICES means ZeroGPU allocation is active.",
    }


@app.get("/v2")
async def v2_index() -> dict[str, Any]:
    return {
        "object": "api.catalog",
        "endpoints": [
            {"method": "GET", "path": "/health"},
            {"method": "GET", "path": "/v1/models"},
            {"method": "GET", "path": "/v2/runtime"},
            {"method": "GET", "path": "/v2/logs"},
            {"method": "GET", "path": "/v2/zerogpu/probe"},
            {"method": "POST", "path": "/v1/chat/completions"},
            {"method": "POST", "path": "/gradio_api/call/chat"},
        ],
        "default_model": default_model_id(),
    }


@app.get("/v2/runtime")
async def v2_runtime() -> dict[str, Any]:
    return {
        "ok": True,
        "service": "diffusiongemma-gguf-zerogpu-api",
        "default_model": default_model_id(),
        "zero_gpu_size": settings.zero_gpu_size,
        "runtime": runtime_status(),
        "limits": {
            "default_max_tokens": settings.default_max_tokens,
            "max_max_tokens": settings.max_max_tokens,
            "diffusion_max_steps": settings.diffusion_max_steps,
            "api_time_limit_seconds": settings.api_time_limit_seconds,
            "cli_timeout_seconds": settings.cli_timeout_seconds,
        },
        "build": {
            "build_llama_diffusion": settings.build_llama_diffusion,
            "llama_build_cuda": settings.llama_build_cuda,
            "llama_cmake_extra_args": settings.llama_cmake_extra_args,
            "llama_diffusion_bin_url": bool(settings.llama_diffusion_bin_url),
        },
        "paths": {
            "data_dir": str(settings.data_dir),
            "model_cache_dir": str(settings.model_cache_dir),
            "bin_dir": str(settings.bin_dir),
            "llama_src_dir": str(settings.llama_src_dir),
        },
    }


@app.get("/v2/logs")
async def v2_logs(limit: int = 100) -> dict[str, Any]:
    return {"object": "event.list", "data": runtime_events(limit)}


@app.get("/v2/zerogpu/probe")
async def v2_zerogpu_probe(request: Request) -> dict[str, Any]:
    include_cli = str(request.query_params.get("cli", "")).lower() in {"1", "true", "yes", "on"}
    return _gpu_probe(include_cli=include_cli)


@app.post("/v1/chat/completions")
async def v1_chat_completions(request: Request):
    """OpenAI-compatible convenience route.

    For ZeroGPU production use, the more robust design is an external gateway that
    calls /gradio_api/call/chat. This direct route is useful for clients that can
    point base_url directly at the Space.
    """
    body = await request.json()
    stream = bool(body.get("stream", False))
    request_model = body.get("model") or default_model_id()
    request_id = "chatcmpl-" + uuid.uuid4().hex
    created = int(time.time())

    try:
        internal_payload = extract_openai_payload(body)
    except ApiError as exc:
        return JSONResponse(public_error_payload(exc), status_code=exc.status_code)

    if stream:
        def event_stream():
            yield make_sse_chunk(
                request_id=request_id,
                created=created,
                model=request_model,
                delta={"role": "assistant"},
                finish_reason=None,
            )
            try:
                with _api_lock:
                    result = run_final_from_payload(internal_payload, gpu_generate=_gpu_generate)
                if "error" in result:
                    error_text = json.dumps(result["error"], ensure_ascii=False)
                    yield make_sse_chunk(
                        request_id=request_id,
                        created=created,
                        model=request_model,
                        delta={"content": error_text},
                        finish_reason="stop",
                    )
                    yield make_sse_done()
                    return

                # DiffusionGemma GGUF does not provide stable token-by-token deltas
                # in this first version. Send one final delta for SDK compatibility.
                content = result.get("content", "")
                if content:
                    yield make_sse_chunk(
                        request_id=request_id,
                        created=created,
                        model=request_model,
                        delta={"content": content},
                        finish_reason=None,
                    )
                yield make_sse_chunk(
                    request_id=request_id,
                    created=created,
                    model=request_model,
                    delta={},
                    finish_reason="stop",
                )
                yield make_sse_done()
            except ApiError as exc:
                yield make_sse_chunk(
                    request_id=request_id,
                    created=created,
                    model=request_model,
                    delta={"content": json.dumps(public_error_payload(exc), ensure_ascii=False)},
                    finish_reason="stop",
                )
                yield make_sse_done()
            except Exception as exc:
                yield make_sse_chunk(
                    request_id=request_id,
                    created=created,
                    model=request_model,
                    delta={"content": json.dumps(public_error_payload(ApiError("internal_error", str(exc), 500)), ensure_ascii=False)},
                    finish_reason="stop",
                )
                yield make_sse_done()

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

    try:
        with _api_lock:
            result = run_final_from_payload(internal_payload, gpu_generate=_gpu_generate)
        if "error" in result:
            return JSONResponse(result, status_code=500)
        return JSONResponse(make_chat_completion(request_id, created, request_model, result.get("content", "")))
    except ApiError as exc:
        return JSONResponse(public_error_payload(exc), status_code=exc.status_code)
    except Exception as exc:
        return JSONResponse(public_error_payload(ApiError("internal_error", str(exc), 500)), status_code=500)


@app.get("/", response_class=HTMLResponse)
async def homepage():
    index_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "static", "index.html")
    with open(index_path, "r", encoding="utf-8") as f:
        html = f.read()
    app_config = {
        "model": default_model_id(),
        "repoId": settings.gguf_repo_id,
        "filename": settings.gguf_filename,
        "modelSourceUrl": f"https://huggingface.co/{settings.gguf_repo_id}",
        "zeroGpuSize": settings.zero_gpu_size,
        "defaultMaxTokens": settings.default_max_tokens,
        "maxMaxTokens": settings.max_max_tokens,
        "diffusionMaxSteps": settings.diffusion_max_steps,
        "nGpuLayers": settings.n_gpu_layers,
        "thinkingDefault": settings.thinking_enabled_default,
        "diffusionVisualDefault": settings.diffusion_visual_default,
        "diffusionKvCache": settings.diffusion_kv_cache,
    }
    html = html.replace("{{APP_CONFIG_JSON}}", json.dumps(app_config, ensure_ascii=False))
    return HTMLResponse(html, headers={"Cache-Control": "no-store"})


# HF Spaces looks for a top-level demo/app object.
demo = app


if __name__ == "__main__":
    app.launch(
        server_name=os.environ.get("HOST", "0.0.0.0"),
        server_port=int(os.environ.get("PORT", "7860")),
        show_error=True,
    )