app.py

"""Dynamic quantized LLM playground for Hugging Face ZeroGPU Spaces."""

from __future__ import annotations

import logging
import os
from typing import Any

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

# ZeroGPU must be imported before torch or a library that may initialize CUDA.
import spaces
import torch  # noqa: F401  # imported after spaces by design
import gradio as gr

from backend_router import (
    BACKEND_AUTO,
    BACKEND_CHOICES,
    BackendRouterError,
    ModelInspection,
)
from model_manager import ModelCache, ModelRuntime, validate_model_id


logging.basicConfig(level=logging.INFO)
LOGGER = logging.getLogger(__name__)

DEFAULT_MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
cache = ModelCache()
runtime = ModelRuntime(cache)


def _short_error(prefix: str, exc: Exception) -> str:
    LOGGER.exception("%s", prefix)
    detail = str(exc).strip().splitlines()[0] if str(exc).strip() else exc.__class__.__name__
    return f"Error: {prefix} {detail[:320]}"


def _safe_generation_settings(
    max_new_tokens: Any,
    temperature: Any,
    top_p: Any,
) -> tuple[int, float, float]:
    tokens = max(1, min(8192, int(max_new_tokens or 256)))
    temp = max(0.0, min(2.0, float(temperature or 0.7)))
    nucleus = max(0.05, min(1.0, float(top_p or 0.95)))
    return tokens, temp, nucleus


def _dropdown_update(inspection: ModelInspection | None) -> Any:
    choices = inspection.gguf_files if inspection else []
    value = inspection.default_gguf if inspection else None
    return gr.update(choices=choices, value=value)


def _inspection_markdown(
    inspection: ModelInspection | None,
    backend: str | None = None,
    selected_file: str | None = None,
) -> str:
    if inspection is None:
        return "**Detected format:** not inspected yet  \n**Backend:** Auto"
    return inspection.markdown(backend, selected_file)


def inspect_model(model_id: str) -> tuple[str, Any, str, str]:
    """Inspect repository metadata and list GGUF choices without downloading weights."""

    try:
        model_id = validate_model_id(model_id)
        inspection = cache.inspect_remote(model_id)
        selected = inspection.default_gguf
        status = f"Inspected `{model_id}` on CPU."
        if inspection.gguf_files:
            status += " Select a GGUF file, then click Download."
        else:
            status += " No GGUF file was found; Auto will use Transformers."
        return (
            status,
            _dropdown_update(inspection),
            _inspection_markdown(inspection, selected_file=selected),
            cache.describe(model_id, selected),
        )
    except Exception as exc:
        return (
            _short_error("Could not inspect the repository:", exc),
            _dropdown_update(None),
            _inspection_markdown(None),
            cache.describe(model_id),
        )


def download_model(
    model_id: str,
    backend_choice: str,
    gguf_file: str | None,
) -> tuple[str, Any, str, str]:
    """Download a standard snapshot or exactly one selected GGUF on CPU."""

    try:
        model_id = validate_model_id(model_id)
        inspection = cache.inspect_remote(model_id)
        selected = (gguf_file or inspection.default_gguf or "").strip()
        backend = cache.router.resolve_backend(inspection, backend_choice, selected or None)

        if backend == "llama.cpp":
            path = cache.download_gguf(model_id, selected)
            status = f"Downloaded one GGUF file on CPU: `{selected}`"
            cache_status = cache.describe(model_id, selected)
            return (
                status,
                _dropdown_update(inspection),
                _inspection_markdown(inspection, backend, selected),
                cache_status,
            )

        path = cache.download(model_id)
        return (
            f"Downloaded Transformers files on CPU: `{model_id}`",
            _dropdown_update(inspection),
            _inspection_markdown(inspection, backend),
            f"Disk cache: snapshot ready (`{path.name}`).",
        )
    except Exception as exc:
        return (
            _short_error("Could not download the model:", exc),
            _dropdown_update(None),
            _inspection_markdown(None),
            cache.describe(model_id, gguf_file),
        )



def _chat_gpu_duration(
    message: str,
    history: list[Any] | None,
    model_id: str,
    backend_choice: str,
    gguf_file: str | None,
    system_prompt: str,
    max_new_tokens: Any,
    temperature: Any,
    top_p: Any,
) -> int:
    """Reserve only the GPU time a text request is likely to need.

    ZeroGPU checks the declared duration against the visitor's remaining quota
    before the call starts, so a large fixed reservation wastes Free-tier quota.
    Keep ordinary short chats cheap while allowing longer 2K-token generations.
    """
    try:
        tokens = max(1, min(2048, int(max_new_tokens or 256)))
    except (TypeError, ValueError):
        tokens = 256

    # 256 tokens -> 38s, 1024 -> 62s, 2048 -> 94s.
    # Clamp below 120s so Free-tier calls stay well inside the daily 300s budget.
    return max(30, min(120, 30 + (tokens + 31) // 32))


@spaces.GPU(duration=120)
def load_model_on_gpu(
    model_id: str,
    backend_choice: str,
    gguf_file: str | None,
) -> tuple[str, str, str, str]:
    """Load one selected model on ZeroGPU using the routed backend."""

    try:
        model_id = validate_model_id(model_id)
        target = runtime.ensure_loaded(model_id, backend_choice, gguf_file)
        selected = target.selected_file
        return (
            f"Loaded on ZeroGPU: `{model_id}` via `{target.backend}`",
            runtime.active_label(),
            _inspection_markdown(target.inspection, target.backend, selected),
            cache.describe(model_id, selected),
        )
    except Exception as exc:
        return (
            _short_error("Could not load the model:", exc),
            "No model loaded",
            _inspection_markdown(None),
            cache.describe(model_id, gguf_file),
        )


@spaces.GPU(duration=_chat_gpu_duration)
def chat_with_model(
    message: str,
    history: list[Any] | None,
    model_id: str,
    backend_choice: str,
    gguf_file: str | None,
    system_prompt: str,
    max_new_tokens: Any,
    temperature: Any,
    top_p: Any,
) -> str:
    """Generate a reply through the active Transformers or llama.cpp runtime."""

    try:
        model_id = validate_model_id(model_id)
        tokens, temp, nucleus = _safe_generation_settings(
            max_new_tokens, temperature, top_p
        )
        return runtime.generate(
            model_id=model_id,
            requested_backend=backend_choice,
            selected_file=gguf_file,
            message=message,
            history=history,
            system_prompt=system_prompt,
            max_new_tokens=tokens,
            temperature=temp,
            top_p=nucleus,
        )
    except Exception as exc:
        return _short_error("Could not generate a reply:", exc)


@spaces.GPU(duration=15)
def unload_model_on_gpu(model_id: str) -> tuple[str, str, str, str]:
    """Unload the active runtime and release RAM/VRAM."""

    try:
        active_file = runtime.active_selected_file
        runtime.unload()
        return (
            "Unloaded; RAM/VRAM cleanup requested.",
            "No model loaded",
            "**Detected format:** none active  \n**Backend:** none",
            cache.describe(model_id, active_file),
        )
    except Exception as exc:
        return (
            _short_error("Could not unload the model:", exc),
            "Unknown",
            _inspection_markdown(None),
            cache.describe(model_id),
        )


def delete_model_from_disk(
    model_id: str,
    gguf_file: str | None,
) -> tuple[str, str, str, str, Any]:
    """Remove all cached revisions/files for the selected model on CPU."""

    try:
        model_id = validate_model_id(model_id)
        deleted = cache.delete(model_id)
        status = (
            f"Deleted from disk: `{model_id}`"
            if deleted
            else f"No cached files found for `{model_id}`"
        )
        return (
            status,
            "No model loaded",
            "**Detected format:** none active  \n**Backend:** none",
            cache.describe(model_id),
            _dropdown_update(None),
        )
    except Exception as exc:
        return (
            _short_error("Could not delete the model cache:", exc),
            "Unknown",
            _inspection_markdown(None),
            cache.describe(model_id, gguf_file),
            _dropdown_update(None),
        )


CSS = """
#app-container { max-width: 1180px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""


with gr.Blocks(title="Quantized LLM ZeroGPU Playground", css=CSS) as demo:
    gr.Markdown(
        """
        # Quantized LLM ZeroGPU Playground

        Download and test one Hugging Face LLM at a time. **Auto** routes
        standard Transformers checkpoints to Transformers and GGUF files to
        the CUDA-enabled llama.cpp backend. GGUF repositories are inspected
        first so you can select only the Q4/Q5/Q8 (or newer) quant you want.
        """
    )

    with gr.Row():
        model_id = gr.Textbox(
            label="Hugging Face model ID",
            value=DEFAULT_MODEL_ID,
            placeholder="namespace/model-name",
            scale=4,
        )
        backend_choice = gr.Radio(
            label="Backend",
            choices=BACKEND_CHOICES,
            value=BACKEND_AUTO,
            scale=2,
        )

    with gr.Row():
        inspect_button = gr.Button("Inspect / list GGUF", variant="secondary")
        gguf_file = gr.Dropdown(
            label="GGUF quant file (choose one)",
            choices=[],
            value=None,
            allow_custom_value=False,
            interactive=True,
            scale=4,
        )

    with gr.Row():
        download_button = gr.Button("Download", variant="secondary", scale=1)
        load_button = gr.Button("Load", variant="primary", scale=1)
        unload_button = gr.Button("Unload", variant="secondary", scale=1)
        delete_button = gr.Button("Delete from disk", variant="stop", scale=1)

    with gr.Row():
        current_model = gr.Textbox(
            label="Active model / backend",
            value="No model loaded",
            interactive=False,
            scale=1,
        )
        cache_status = gr.Markdown("Disk cache: no model selected.")

    status = gr.Markdown("Status: inspect a repository, then Download and Load it.")
    format_backend = gr.Markdown(
        "**Detected format:** not inspected yet  \n**Backend:** Auto"
    )

    with gr.Accordion("Generation settings", open=True):
        system_prompt = gr.Textbox(
            label="System prompt",
            value="You are a helpful assistant.",
            lines=2,
        )
        with gr.Row():
            max_new_tokens = gr.Slider(
                label="Max new tokens", minimum=1, maximum=8192, value=256, step=1
            )
            temperature = gr.Slider(
                label="Temperature", minimum=0, maximum=2, value=0.7, step=0.05
            )
            top_p = gr.Slider(label="Top-p", minimum=0.05, maximum=1, value=0.95, step=0.05)

    chatbot = gr.Chatbot(type="messages", height=520, allow_tags=False)
    gr.ChatInterface(
        fn=chat_with_model,
        chatbot=chatbot,
        type="messages",
        additional_inputs=[
            model_id,
            backend_choice,
            gguf_file,
            system_prompt,
            max_new_tokens,
            temperature,
            top_p,
        ],
        textbox=gr.Textbox(placeholder="Write a message…", container=False),
        api_name="chat",
    )

    inspect_button.click(
        fn=inspect_model,
        inputs=[model_id],
        outputs=[status, gguf_file, format_backend, cache_status],
        api_name="inspect",
    )
    download_button.click(
        fn=download_model,
        inputs=[model_id, backend_choice, gguf_file],
        outputs=[status, gguf_file, format_backend, cache_status],
        api_name="download",
    )
    load_button.click(
        fn=load_model_on_gpu,
        inputs=[model_id, backend_choice, gguf_file],
        outputs=[status, current_model, format_backend, cache_status],
        api_name="load",
    )
    unload_button.click(
        fn=unload_model_on_gpu,
        inputs=[model_id],
        outputs=[status, current_model, format_backend, cache_status],
        api_name="unload",
    )
    delete_event = delete_button.click(
        # Release a possibly active GPU copy before deleting its CPU cache.
        fn=unload_model_on_gpu,
        inputs=[model_id],
        outputs=[status, current_model, format_backend, cache_status],
    )
    delete_event.then(
        fn=delete_model_from_disk,
        inputs=[model_id, gguf_file],
        outputs=[status, current_model, format_backend, cache_status, gguf_file],
        api_name="delete_from_disk",
    )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1)
    demo.launch(mcp_server=True)