"""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)