Create dynamic ZeroGPU LLM playground

README.md

@@ -1,13 +1,41 @@
 ---
-title: Llm Zero Gpu Playground
-emoji: 🐨
-colorFrom: yellow
-colorTo: purple
+title: Dynamic LLM ZeroGPU Playground
+emoji: 🧪
+colorFrom: blue
+colorTo: indigo
 sdk: gradio
 sdk_version: 6.25.0
-python_version: '3.12'
 app_file: app.py
-pinned: false
+short_description: Load and test HF causal LLMs on ZeroGPU
+python_version: "3.10"
+startup_duration_timeout: 1h
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Dynamic LLM ZeroGPU Playground
+
+A small, general-purpose playground for testing standard Hugging Face causal
+language models by model ID.
+
+## How it works
+
+1. Enter a model ID such as `Qwen/Qwen2.5-0.5B-Instruct`.
+2. Click **Download** to fetch the snapshot into the Space's CPU-side cache.
+3. Click **Load** to put that model on the ZeroGPU worker.
+4. Chat with the model, or use **Unload** before switching models.
+5. Use **Delete from disk** to unload first, then remove all cached revisions for the selected ID.
+
+Only one model is kept active by the runtime. Switching models releases the
+previous model with `del`, `gc.collect()`, and `torch.cuda.empty_cache()` before
+the new model is loaded. Chat templates are used whenever the tokenizer
+provides `apply_chat_template()`.
+
+This MVP intentionally targets standard `transformers` +
+`AutoModelForCausalLM` checkpoints. The loader is isolated in
+`TransformersCausalLMRuntime`, so AWQ, GPTQ, or FP8 backends can be added later.
+
+## Notes
+
+- Model downloads happen on CPU and are never triggered by the chat handler.
+- A model must be downloaded before it can be loaded or used.
+- Large models may exceed ZeroGPU memory or take a long time to load.
+- Remote model code is disabled in this first version for safety and stability.

app.py

@@ -0,0 +1,233 @@
+"""Dynamic Hugging Face causal-LLM playground for 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 any library that may touch CUDA.
+import spaces
+import torch
+import gradio as gr
+
+from model_manager import ModelCache, TransformersCausalLMRuntime, validate_model_id
+
+
+logging.basicConfig(level=logging.INFO)
+LOGGER = logging.getLogger(__name__)
+
+DEFAULT_MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
+cache = ModelCache()
+runtime = TransformersCausalLMRuntime(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[:300]}"
+
+
+def _safe_generation_settings(
+    max_new_tokens: Any,
+    temperature: Any,
+    top_p: Any,
+) -> tuple[int, float, float]:
+    tokens = max(1, min(2048, 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 download_model(model_id: str) -> tuple[str, str]:
+    """Download a model snapshot on CPU into the playground cache."""
+
+    try:
+        model_id = validate_model_id(model_id)
+        snapshot_path = cache.download(model_id)
+        return (
+            f"Downloaded on CPU: `{model_id}`",
+            f"Disk cache: ready ({snapshot_path.name}).",
+        )
+    except Exception as exc:
+        message = _short_error("Could not download the model:", exc)
+        return message, cache.describe(model_id)
+
+
+@spaces.GPU(duration=420)
+def load_model_on_gpu(model_id: str) -> tuple[str, str, str]:
+    """Load one downloaded Transformers causal LM on the ZeroGPU worker."""
+
+    try:
+        model_id = validate_model_id(model_id)
+        cache.cached_snapshot(model_id)  # local-only check; never downloads here
+        active = runtime.ensure_loaded(model_id)
+        return (
+            f"Loaded on ZeroGPU: `{active}`",
+            active,
+            cache.describe(active),
+        )
+    except Exception as exc:
+        return _short_error("Could not load the model:", exc), "No model loaded", cache.describe(model_id)
+
+
+@spaces.GPU(duration=180)
+def chat_with_model(
+    message: str,
+    history: list[Any] | None,
+    model_id: str,
+    system_prompt: str,
+    max_new_tokens: Any,
+    temperature: Any,
+    top_p: Any,
+) -> str:
+    """Generate a reply using the selected cached Hugging Face model."""
+
+    try:
+        model_id = validate_model_id(model_id)
+        cache.cached_snapshot(model_id)  # a chat never performs a download
+        tokens, temp, nucleus = _safe_generation_settings(
+            max_new_tokens, temperature, top_p
+        )
+        return runtime.generate(
+            model_id=model_id,
+            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=30)
+def unload_model_on_gpu(model_id: str) -> tuple[str, str, str]:
+    """Unload the active model and release RAM/VRAM on the ZeroGPU worker."""
+
+    try:
+        runtime.unload()
+        return "Unloaded; RAM/VRAM cleanup requested.", "No model loaded", cache.describe(model_id)
+    except Exception as exc:
+        return _short_error("Could not unload the model:", exc), "Unknown", cache.describe(model_id)
+
+
+def delete_model_from_disk(model_id: str) -> tuple[str, str, str]:
+    """Delete every cached revision of the selected model from disk."""
+
+    try:
+        model_id = validate_model_id(model_id)
+        deleted = cache.delete(model_id)
+        if deleted:
+            status = f"Deleted from disk: `{model_id}`"
+        else:
+            status = f"No cached files found for `{model_id}`"
+        return status, "No model loaded", cache.describe(model_id)
+    except Exception as exc:
+        return _short_error("Could not delete the model cache:", exc), "Unknown", cache.describe(model_id)
+
+
+CSS = """
+#app-container { max-width: 1180px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+
+with gr.Blocks(title="Dynamic LLM ZeroGPU Playground", css=CSS) as demo:
+    gr.Markdown(
+        """
+        # Dynamic LLM ZeroGPU Playground
+
+        Download standard `transformers` causal language models, load one model
+        at a time on ZeroGPU, and test it through a Gradio chat interface.
+        Downloading and cache management stay on CPU; model loading and
+        inference use the GPU only when requested.
+        """
+    )
+
+    with gr.Row():
+        model_id = gr.Textbox(
+            label="Hugging Face model ID",
+            value=DEFAULT_MODEL_ID,
+            placeholder="namespace/model-name",
+            scale=4,
+        )
+        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",
+            value="No model loaded",
+            interactive=False,
+            scale=1,
+        )
+        cache_status = gr.Markdown("Disk cache: no model selected.", scale=1)
+
+    status = gr.Markdown("Status: enter a model ID, then click Download.")
+
+    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=2048, 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(height=520)
+    gr.ChatInterface(
+        fn=chat_with_model,
+        chatbot=chatbot,
+        additional_inputs=[model_id, system_prompt, max_new_tokens, temperature, top_p],
+        textbox=gr.Textbox(placeholder="Write a message…", container=False),
+        api_name="chat",
+    )
+
+    download_button.click(
+        fn=download_model,
+        inputs=[model_id],
+        outputs=[status, cache_status],
+        api_name="download",
+    )
+    load_button.click(
+        fn=load_model_on_gpu,
+        inputs=[model_id],
+        outputs=[status, current_model, cache_status],
+        api_name="load",
+    )
+    unload_button.click(
+        fn=unload_model_on_gpu,
+        inputs=[model_id],
+        outputs=[status, current_model, cache_status],
+        api_name="unload",
+    )
+    delete_event = delete_button.click(
+        # Release a possibly active GPU copy before removing its CPU cache.
+        # The actual deletion remains a CPU-only operation in the next step.
+        fn=unload_model_on_gpu,
+        inputs=[model_id],
+        outputs=[status, current_model, cache_status],
+    )
+    delete_event.then(
+        fn=delete_model_from_disk,
+        inputs=[model_id],
+        outputs=[status, current_model, cache_status],
+        api_name="delete_from_disk",
+    )
+
+
+if __name__ == "__main__":
+    demo.queue(default_concurrency_limit=1)
+    demo.launch(mcp_server=True)

model_manager.py

@@ -0,0 +1,264 @@
+"""CPU cache helpers and the single-model Transformers runtime.
+
+The runtime deliberately keeps model loading behind the GPU handlers in
+``app.py``.  That makes model IDs dynamic while still ensuring that no model
+weights are loaded during Space startup.
+"""
+
+from __future__ import annotations
+
+import gc
+import logging
+import os
+import re
+import threading
+from pathlib import Path
+from typing import Any
+
+import spaces
+import torch
+from huggingface_hub import scan_cache_dir, snapshot_download
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+LOGGER = logging.getLogger(__name__)
+MODEL_ID_PATTERN = re.compile(r"^[^/\s]+/[^/\s]+$")
+
+
+def validate_model_id(model_id: str) -> str:
+    """Validate and normalize a Hugging Face model repository ID."""
+
+    normalized = (model_id or "").strip()
+    if not MODEL_ID_PATTERN.fullmatch(normalized):
+        raise ValueError("Model ID must look like namespace/model-name.")
+    return normalized
+
+
+class ModelCache:
+    """A dedicated Hugging Face cache for downloaded model snapshots."""
+
+    def __init__(self, cache_dir: str | None = None) -> None:
+        default_dir = Path.home() / ".cache" / "huggingface" / "llm-playground"
+        self.root = Path(cache_dir or os.getenv("PLAYGROUND_CACHE_DIR", default_dir))
+        self.root.mkdir(parents=True, exist_ok=True)
+
+    def download(self, model_id: str) -> Path:
+        """Download a complete model snapshot without initializing a model."""
+
+        model_id = validate_model_id(model_id)
+        snapshot_path = snapshot_download(
+            repo_id=model_id,
+            repo_type="model",
+            cache_dir=str(self.root),
+        )
+        return Path(snapshot_path)
+
+    def cached_snapshot(self, model_id: str) -> Path:
+        """Return a locally cached snapshot, raising if it is not complete."""
+
+        model_id = validate_model_id(model_id)
+        try:
+            snapshot_path = snapshot_download(
+                repo_id=model_id,
+                repo_type="model",
+                cache_dir=str(self.root),
+                local_files_only=True,
+            )
+        except Exception as exc:  # hub versions expose different local-cache errors
+            raise FileNotFoundError(
+                f"{model_id} is not downloaded yet. Click Download first."
+            ) from exc
+        return Path(snapshot_path)
+
+    def describe(self, model_id: str) -> str:
+        """Return a short cache status for the UI."""
+
+        model_id = (model_id or "").strip()
+        if not model_id:
+            return "Disk cache: no model selected."
+        try:
+            path = self.cached_snapshot(model_id)
+        except (ValueError, FileNotFoundError):
+            return f"Disk cache: {model_id} is not downloaded."
+        return f"Disk cache: ready ({path.name})."
+
+    def delete(self, model_id: str) -> bool:
+        """Delete every cached revision of one model from this cache."""
+
+        model_id = validate_model_id(model_id)
+        cache_info = scan_cache_dir(cache_dir=str(self.root))
+        revisions = []
+        for repo in cache_info.repos:
+            if repo.repo_id == model_id:
+                revisions.extend(revision.commit_hash for revision in repo.revisions)
+
+        if not revisions:
+            return False
+
+        # The cache manager removes snapshots, refs, and blobs that are no
+        # longer shared by another cached revision.
+        cache_info.delete_revisions(*revisions).execute()
+        return True
+
+
+class TransformersCausalLMRuntime:
+    """Single active standard Transformers causal language model."""
+
+    def __init__(self, cache: ModelCache) -> None:
+        self.cache = cache
+        self._model: Any | None = None
+        self._tokenizer: Any | None = None
+        self._model_id: str | None = None
+        self._lock = threading.RLock()
+
+    @property
+    def active_model_id(self) -> str | None:
+        return self._model_id
+
+    def unload(self) -> None:
+        """Release the active model and clear CUDA's allocator cache."""
+
+        with self._lock:
+            old_model = self._model
+            self._model = None
+            self._tokenizer = None
+            self._model_id = None
+
+            if old_model is not None:
+                del old_model
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
+    def ensure_loaded(self, model_id: str) -> str:
+        """Load one cached model on CUDA, replacing any previously active model."""
+
+        model_id = validate_model_id(model_id)
+        with self._lock:
+            if self._model is not None and self._model_id == model_id:
+                return model_id
+
+            # Switching models always releases the old object before reading
+            # the new checkpoint, keeping the one-model invariant explicit.
+            self.unload()
+            snapshot_path = self.cache.cached_snapshot(model_id)
+
+            tokenizer = AutoTokenizer.from_pretrained(
+                str(snapshot_path),
+                local_files_only=True,
+                use_fast=True,
+                trust_remote_code=False,
+            )
+            model = AutoModelForCausalLM.from_pretrained(
+                str(snapshot_path),
+                local_files_only=True,
+                torch_dtype=torch.bfloat16,
+                low_cpu_mem_usage=True,
+                trust_remote_code=False,
+            )
+            model = model.to("cuda").eval()
+
+            if tokenizer.pad_token_id is None and tokenizer.eos_token_id is not None:
+                tokenizer.pad_token = tokenizer.eos_token
+            if getattr(model.config, "pad_token_id", None) is None:
+                model.config.pad_token_id = tokenizer.pad_token_id
+
+            self._tokenizer = tokenizer
+            self._model = model
+            self._model_id = model_id
+            return model_id
+
+    @staticmethod
+    def _history_to_messages(history: list[Any] | None) -> list[dict[str, str]]:
+        messages: list[dict[str, str]] = []
+        for item in history or []:
+            if isinstance(item, dict):
+                role = str(item.get("role", ""))
+                content = item.get("content", "")
+                if role in {"user", "assistant"} and isinstance(content, str):
+                    messages.append({"role": role, "content": content})
+            elif isinstance(item, (list, tuple)) and len(item) == 2:
+                user_text, assistant_text = item
+                if isinstance(user_text, str) and user_text:
+                    messages.append({"role": "user", "content": user_text})
+                if isinstance(assistant_text, str) and assistant_text:
+                    messages.append({"role": "assistant", "content": assistant_text})
+        return messages
+
+    @staticmethod
+    def _plain_prompt(messages: list[dict[str, str]]) -> str:
+        lines = [f"{message['role'].capitalize()}: {message['content']}" for message in messages]
+        return "\n".join(lines) + "\nAssistant:"
+
+    def _tokenize(self, messages: list[dict[str, str]]) -> Any:
+        assert self._tokenizer is not None
+        tokenizer = self._tokenizer
+
+        if hasattr(tokenizer, "apply_chat_template"):
+            try:
+                return tokenizer.apply_chat_template(
+                    messages,
+                    add_generation_prompt=True,
+                    tokenize=True,
+                    return_tensors="pt",
+                    return_dict=True,
+                )
+            except TypeError:
+                try:
+                    return tokenizer.apply_chat_template(
+                        messages,
+                        add_generation_prompt=True,
+                        tokenize=True,
+                        return_tensors="pt",
+                    )
+                except Exception:
+                    LOGGER.debug("Chat template without return_dict failed", exc_info=True)
+            except Exception:
+                LOGGER.debug("Chat template failed; using plain prompt", exc_info=True)
+
+        return tokenizer(self._plain_prompt(messages), return_tensors="pt")
+
+    def generate(
+        self,
+        model_id: str,
+        message: str,
+        history: list[Any] | None,
+        system_prompt: str,
+        max_new_tokens: int,
+        temperature: float,
+        top_p: float,
+    ) -> str:
+        """Generate one answer from the active cached Transformers model."""
+
+        with self._lock:
+            self.ensure_loaded(model_id)
+            assert self._model is not None
+            assert self._tokenizer is not None
+
+            messages: list[dict[str, str]] = []
+            if (system_prompt or "").strip():
+                messages.append({"role": "system", "content": system_prompt.strip()})
+            messages.extend(self._history_to_messages(history))
+            messages.append({"role": "user", "content": (message or "").strip()})
+
+            encoded = self._tokenize(messages)
+            encoded = {
+                key: value.to("cuda")
+                for key, value in encoded.items()
+                if torch.is_tensor(value)
+            }
+            input_length = int(encoded["input_ids"].shape[-1])
+
+            generation_kwargs: dict[str, Any] = {
+                "max_new_tokens": max_new_tokens,
+                "do_sample": temperature > 0,
+            }
+            if temperature > 0:
+                generation_kwargs.update({"temperature": temperature, "top_p": top_p})
+
+            with torch.inference_mode():
+                generated = self._model.generate(**encoded, **generation_kwargs)
+
+            new_tokens = generated[0, input_length:]
+            answer = self._tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
+            return answer or "The model returned an empty response."

requirements.txt

@@ -0,0 +1,4 @@
+transformers>=4.45.0,<5
+accelerate>=0.34.0
+safetensors>=0.4.3
+sentencepiece>=0.2.0