add gradio app

.gitignore

@@ -0,0 +1,223 @@
+# Project-specific
+data/
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[codz]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#   Usually these files are written by a python script from a template
+#   before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py.cover
+*.lcov
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+# Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+# uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+# poetry.lock
+# poetry.toml
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#   pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
+#   https://pdm-project.org/en/latest/usage/project/#working-with-version-control
+# pdm.lock
+# pdm.toml
+.pdm-python
+.pdm-build/
+
+# pixi
+#   Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
+# pixi.lock
+#   Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
+#   in the .venv directory. It is recommended not to include this directory in version control.
+.pixi/*
+!.pixi/config.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule*
+celerybeat.pid
+
+# Redis
+*.rdb
+*.aof
+*.pid
+
+# RabbitMQ
+mnesia/
+rabbitmq/
+rabbitmq-data/
+
+# ActiveMQ
+activemq-data/
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.envrc
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#   JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#   be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#   and can be added to the global gitignore or merged into this file.  For a more nuclear
+#   option (not recommended) you can uncomment the following to ignore the entire idea folder.
+# .idea/
+
+# Abstra
+#   Abstra is an AI-powered process automation framework.
+#   Ignore directories containing user credentials, local state, and settings.
+#   Learn more at https://abstra.io/docs
+.abstra/
+
+# Visual Studio Code
+#   Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore 
+#   that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
+#   and can be added to the global gitignore or merged into this file. However, if you prefer, 
+#   you could uncomment the following to ignore the entire vscode folder
+# .vscode/
+# Temporary file for partial code execution
+tempCodeRunnerFile.py
+
+# Ruff stuff:
+.ruff_cache/
+
+# PyPI configuration file
+.pypirc
+
+# Marimo
+marimo/_static/
+marimo/_lsp/
+__marimo__/
+
+# Streamlit
+.streamlit/secrets.toml

README.md

@@ -1,14 +1,30 @@
 ---
 title: DanbooruTagQuery
-emoji: 🦀
-colorFrom: gray
-colorTo: pink
+emoji: 🏷️
+colorFrom: indigo
+colorTo: purple
 sdk: gradio
-sdk_version: 6.20.0
-python_version: '3.12'
+sdk_version: 5.23.3
 app_file: app.py
 pinned: false
-license: other
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# 🏷️ DanbooruTagQuery
+
+ONNX-powered anime image tagger. Upload an image or paste a URL to get predicted Danbooru tags with scores, categories, and wiki lookups.
+
+## Local testing
+
+```bash
+# Use a local model file
+python app.py /path/to/DanbooruTagQuery_b16_448x448.onnx
+
+# Or point to a directory with .onnx + sidecar files
+export MODEL_DIR=/path/to/models/
+python app.py
+```
+
+## Credits
+
+- Model: [realphongha/danbooru-tag-query](https://huggingface.co/realphongha/danbooru-tag-query)
+- Dataset: [danbooru-tags](https://huggingface.co/datasets/qdlabs/danbooru-tags)

app.py

@@ -0,0 +1,718 @@
+#!/usr/bin/env python3
+"""
+DanbooruTagQuery — Hugging Face ZeroGPU Space
+
+Usage:
+  python app.py                          # download model from HF hub
+  python app.py /path/to/model.onnx      # use local model
+  MODEL_DIR=/path/to python app.py       # env var with model dir
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import sqlite3
+import sys
+import tempfile
+import threading
+import time
+from pathlib import Path
+
+import gradio as gr
+import numpy as np
+from PIL import Image
+
+# ── optional deps (loaded on demand) ────────────────────────────────────────
+
+_hf_hub = None
+
+def _import_hf_hub():
+    global _hf_hub
+    if _hf_hub is None:
+        import huggingface_hub as h
+        _hf_hub = h
+    return _hf_hub
+
+
+# ── constants ───────────────────────────────────────────────────────────────
+
+HF_REPO = "realphongha/danbooru-tag-query"
+MODELS_DIR = "models"
+
+IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
+IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
+
+CATEGORY_MAP = {
+    0: "general",
+    1: "artist",
+    3: "copyright",
+    4: "character",
+    5: "meta",
+}
+
+DEFAULT_TOP_K = None
+DEFAULT_MIN_SCORE = 0.2
+
+# ── TagCache (SQLite) ───────────────────────────────────────────────────────
+
+_DB_PATH = Path("data/tag_cache.db")
+
+
+class TagCache:
+    def __init__(self, db: str | Path = _DB_PATH):
+        self.db = Path(db)
+        self.db.parent.mkdir(parents=True, exist_ok=True)
+        self._conn = sqlite3.connect(str(self.db), check_same_thread=False)
+        self._lock = threading.Lock()
+        self._conn.execute(
+            """CREATE TABLE IF NOT EXISTS tag_cache (
+                name TEXT PRIMARY KEY,
+                category INTEGER,
+                wiki_body TEXT,
+                fetched_at REAL NOT NULL
+            )"""
+        )
+        self._conn.commit()
+
+    def get(self, name: str) -> tuple[int | None, str | None] | None:
+        with self._lock:
+            cur = self._conn.execute(
+                "SELECT category, wiki_body FROM tag_cache WHERE name = ?", (name,)
+            )
+            return cur.fetchone()
+
+    def set(self, name: str, category: int | None, wiki_body: str | None):
+        with self._lock:
+            self._conn.execute(
+                """INSERT OR REPLACE INTO tag_cache (name, category, wiki_body, fetched_at)
+                   VALUES (?, ?, ?, ?)""",
+                (name, category, wiki_body, time.time()),
+            )
+            self._conn.commit()
+
+    def get_many(self, names: list[str]) -> dict[str, tuple[int | None, str | None]]:
+        if not names:
+            return {}
+        with self._lock:
+            placeholders = ",".join("?" for _ in names)
+            cur = self._conn.execute(
+                f"SELECT name, category, wiki_body FROM tag_cache "
+                f"WHERE name IN ({placeholders})", names,
+            )
+            return {row[0]: (row[1], row[2]) for row in cur}
+
+    def bulk_set(self, items: list[tuple[str, int | None, str | None]]):
+        with self._lock:
+            now = time.time()
+            self._conn.executemany(
+                """INSERT OR REPLACE INTO tag_cache (name, category, wiki_body, fetched_at)
+                   VALUES (?, ?, ?, ?)""",
+                [(name, cat, body, now) for name, cat, body in items],
+            )
+            self._conn.commit()
+
+    def clear(self):
+        with self._lock:
+            self._conn.execute("DELETE FROM tag_cache")
+            self._conn.commit()
+
+    def size(self) -> int:
+        with self._lock:
+            cur = self._conn.execute("SELECT COUNT(*) FROM tag_cache")
+            return cur.fetchone()[0]
+
+
+# ── image preprocessing ────────────────────────────────────────────────────
+
+def preprocess(image: Image.Image, image_size: int = 448) -> np.ndarray:
+    w, h = image.size
+    scale = image_size / max(w, h)
+    new_w = int(w * scale)
+    new_h = int(h * scale)
+    image = image.resize((new_w, new_h), Image.BILINEAR)
+    canvas = Image.new("RGB", (image_size, image_size), (0, 0, 0))
+    left = (image_size - new_w) // 2
+    top = (image_size - new_h) // 2
+    canvas.paste(image, (left, top))
+    arr = np.asarray(canvas, dtype=np.float32).transpose(2, 0, 1) / 255.0
+    arr[0] = (arr[0] - IMAGENET_MEAN[0]) / IMAGENET_STD[0]
+    arr[1] = (arr[1] - IMAGENET_MEAN[1]) / IMAGENET_STD[1]
+    arr[2] = (arr[2] - IMAGENET_MEAN[2]) / IMAGENET_STD[2]
+    return arr[np.newaxis, ...]
+
+
+# ── sidecar loading ───────────────────────���────────────────────────────────
+
+def _sidecar_path(checkpoint: Path, suffix: str) -> Path:
+    if checkpoint.suffix == ".onnx":
+        return checkpoint.with_name(checkpoint.stem + suffix)
+    return checkpoint / suffix.lstrip(".")
+
+
+def load_tag_to_id(checkpoint: str | Path) -> dict[str, int]:
+    ckpt = Path(checkpoint)
+    path = _sidecar_path(ckpt, ".tag_to_id.json")
+    if not path.exists():
+        path = ckpt.parent / "tag_to_id.json"
+    if not path.exists():
+        raise FileNotFoundError(f"Missing tag map: {path}")
+    return json.loads(path.read_text())
+
+
+def load_config(checkpoint: str | Path) -> dict:
+    ckpt = Path(checkpoint)
+    path = _sidecar_path(ckpt, ".config.json")
+    if not path.exists():
+        path = ckpt.parent / "config.json"
+    if not path.exists():
+        return {"image_size": 448}
+    return json.loads(path.read_text())
+
+
+# ── Predictor (ONNX) ──────────────────────────────────────────────────────
+
+class Predictor:
+    def __init__(self, checkpoint: str | Path):
+        import onnxruntime as ort
+
+        self.checkpoint = str(checkpoint)
+        self.tag_to_id = load_tag_to_id(self.checkpoint)
+        cfg = load_config(self.checkpoint)
+        self.image_size = cfg.get("image_size", 448)
+
+        providers = [
+            ("CUDAExecutionProvider", {}),
+            "CPUExecutionProvider",
+        ]
+        try:
+            self._sess = ort.InferenceSession(self.checkpoint, providers=providers)
+        except Exception:
+            self._sess = ort.InferenceSession(
+                self.checkpoint, providers=["CPUExecutionProvider"]
+            )
+
+        self._input_name = self._sess.get_inputs()[0].name
+        self._output_name = self._sess.get_outputs()[0].name
+
+    def run(self, pixel_values: np.ndarray) -> np.ndarray:
+        raw = self._sess.run([self._output_name], {self._input_name: pixel_values})[0]
+        return 1.0 / (1.0 + np.exp(-raw))
+
+    @property
+    def num_classes(self) -> int:
+        return len(self.tag_to_id)
+
+
+# ── model discovery & loading (HF hub) ─────────────────────────────────────
+
+def discover_model_variants() -> list[str]:
+    try:
+        hf = _import_hf_hub()
+        api = hf.HfApi()
+        siblings = api.list_repo_files(HF_REPO, repo_type="model")
+        variants: set[str] = set()
+        for path in siblings:
+            if path.startswith(f"{MODELS_DIR}/") and "/" in path[len(MODELS_DIR) + 1:]:
+                variant = path.split("/")[1]
+                if variant:
+                    variants.add(variant)
+        return sorted(variants, reverse=True)
+    except Exception as exc:
+        print(f"Warning: could not discover models on hub: {exc}")
+        return []
+
+
+def resolve_model_path(variant: str) -> str:
+    onnx_path = f"{MODELS_DIR}/{variant}/model.onnx"
+    hf = _import_hf_hub()
+    return hf.hf_hub_download(
+        repo_id=HF_REPO,
+        filename=onnx_path,
+        repo_type="model",
+    )
+
+
+def resolve_sidecar_paths(variant: str) -> tuple[str, str]:
+    hf = _import_hf_hub()
+    config_path = hf.hf_hub_download(
+        repo_id=HF_REPO,
+        filename=f"{MODELS_DIR}/{variant}/config.json",
+        repo_type="model",
+    )
+    tagmap_path = hf.hf_hub_download(
+        repo_id=HF_REPO,
+        filename=f"{MODELS_DIR}/{variant}/tag_to_id.json",
+        repo_type="model",
+    )
+    return config_path, tagmap_path
+
+
+def download_model_variant(variant: str) -> Path:
+    onnx = Path(resolve_model_path(variant))
+    resolve_sidecar_paths(variant)
+    return onnx
+
+
+# ── tag enrichment (category lookup, no API) ───────────────────────────────
+
+_CACHE = TagCache()
+
+
+def prebuild_cache(tag_to_id: dict[str, int]):
+    """Pre-populate tag cache from HF dataset.
+
+    Every tag gets a category (defaults to 0 = general).
+    No lazy API calls needed.
+    """
+    try:
+        from datasets import load_dataset
+    except ImportError:
+        return 0, 0
+
+    tags = sorted(tag_to_id.keys(), key=lambda t: tag_to_id[t])
+
+    try:
+        ds = load_dataset("qdlabs/danbooru-tags", split="train")
+        cat_map = {row["name"]: row["category"] for row in ds}
+    except Exception:
+        return 0, 0
+
+    found = sum(1 for t in tags if t in cat_map)
+    todo = []
+    already = 0
+    for tag in tags:
+        cached = _CACHE.get(tag)
+        if cached is not None and cached[0] is not None:
+            already += 1
+            continue
+        cat = cat_map.get(tag, 0)   # 0 = general for unknown
+        existing_wiki = cached[1] if cached else None
+        todo.append((tag, cat, existing_wiki))
+
+    if todo:
+        _CACHE.bulk_set(todo)
+
+    return already + len(todo), found
+
+
+def enrich_tags(tags_scores: list[tuple[str, float]]) -> dict[str, dict]:
+    """Attach category info to each tag. No API calls."""
+    tags = [t for t, _ in tags_scores]
+    cached_map = _CACHE.get_many(tags)
+    result: dict[str, dict] = {}
+    for tag, score in tags_scores:
+        cached = cached_map.get(tag)
+        if cached is not None:
+            cat_id, _ = cached
+        else:
+            cat_id = None
+        result[tag] = {
+            "score": score,
+            "category": cat_id,
+            "category_name": CATEGORY_MAP.get(cat_id, "general"),
+        }
+    return result
+
+
+def format_tag(tag: str, use_underscore: bool) -> str:
+    return tag if use_underscore else tag.replace("_", " ")
+
+
+# ── Gradio UI ──────────────────────────────────────────────────────────────
+
+def build_app(predict_fn, model_choices: list[str]) -> gr.Blocks:
+    state = {
+        "all_logits": None,
+        "tag_metadata": None,
+        "current_image": None,
+        "predictor": None,
+        "predict_fn": predict_fn,
+    }
+
+    css = """
+    #csv-wrap { position: relative; }
+    #copy-csv-btn { position: absolute; top: 4px; right: 4px; z-index: 10;
+                    min-width: 0; padding: 0 6px; height: 24px;
+                    font-size: 13px; line-height: 24px; }
+    """
+
+    category_names = sorted(CATEGORY_MAP.values())
+
+    with gr.Blocks(title="DanbooruTagQuery", theme=gr.themes.Soft(), css=css) as app:
+        gr.Markdown("# 🏷️ DanbooruTagQuery")
+
+        with gr.Row():
+            # ── left: image + model ──
+            with gr.Column(scale=1):
+                image_input = gr.Image(
+                    label="Image",
+                    type="pil",
+                    sources=["upload", "clipboard"],
+                    height=300,
+                )
+                url_input = gr.Textbox(
+                    label="Image URL",
+                    placeholder="Paste image URL and press Enter",
+                )
+                with gr.Row():
+                    analyze_btn = gr.Button("🔍 Analyze", variant="primary", scale=2)
+                    clear_btn = gr.Button("🗑️ Clear", scale=1)
+
+                gr.Markdown("### 🤖 Model")
+                model_dropdown = gr.Dropdown(
+                    choices=model_choices,
+                    value=model_choices[0] if model_choices else None,
+                    label="Model variant",
+                    interactive=True,
+                )
+                model_status = gr.Markdown("Ready")
+
+            # ── right: params + categories ──
+            with gr.Column(scale=1):
+                top_k = gr.Number(
+                    label="Top-K", value=DEFAULT_TOP_K, minimum=0, step=1
+                )
+                min_score = gr.Slider(
+                    label="Min Score",
+                    value=DEFAULT_MIN_SCORE,
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.01,
+                )
+                sort_by = gr.Radio(
+                    label="Sort by", choices=["score", "name"], value="score"
+                )
+                use_underscore = gr.Checkbox(
+                    label="Use underscore (_)", value=False
+                )
+                categories = gr.CheckboxGroup(
+                    label="Categories",
+                    choices=category_names,
+                    value=["general"],
+                )
+
+        # ── outputs ──
+        with gr.Tabs():
+            with gr.TabItem("📋 Tag list"):
+                tag_table = gr.HTML(label="Tags")
+            with gr.TabItem("📝 Comma-separated"):
+                with gr.Column(elem_id="csv-wrap"):
+                    tag_string = gr.Textbox(label="Tags", lines=6, elem_id="csv-text")
+                    copy_btn = gr.Button("📋", elem_id="copy-csv-btn")
+
+        # ── status row ──
+        with gr.Row():
+            status = gr.Markdown("Ready. Load an image and click **Analyze**.")
+            clear_cache_btn = gr.Button("🧹 Clear cache", size="sm", elem_id="clear-cache-btn")
+
+        # ── tag score query ──
+        gr.Markdown("### 🔍 Tag score query")
+        with gr.Row():
+            tag_query = gr.Textbox(
+                label="Search tag",
+                placeholder="Type to search…",
+                scale=3,
+            )
+        tag_query_output = gr.HTML(label="Results")
+
+        # ── callbacks ──────────────────────────────────────────────────────
+
+        def refresh_results(
+            _top_k, _min_score, _sort_by, _use_underscore, _categories,
+        ):
+            if state["all_logits"] is None or state["tag_metadata"] is None:
+                return "<i>No results yet.</i>", ""
+
+            meta = state["tag_metadata"]
+            all_tags = list(meta.keys())
+            if _categories:
+                all_tags = [
+                    t for t in all_tags
+                    if meta[t]["category_name"] in _categories
+                ]
+
+            items = [(t, meta[t]["score"]) for t in all_tags]
+            if _sort_by == "name":
+                items.sort(key=lambda x: format_tag(x[0], _use_underscore))
+            else:
+                items.sort(key=lambda x: x[1], reverse=True)
+
+            items = [(t, s) for t, s in items if s >= _min_score]
+            if _top_k is not None and _top_k > 0:
+                items = items[:_top_k]
+
+            if not items:
+                return "<i>No tags pass the filters.</i>", ""
+
+            rows = []
+            for tag, score in items:
+                m = meta[tag]
+                link = (
+                    f'<a href="https://danbooru.donmai.us/posts?tags={tag}"'
+                    f' target="_blank">{tag}</a>'
+                )
+                display = format_tag(tag, _use_underscore)
+                rows.append(
+                    f"<tr>"
+                    f"<td>{link}</td>"
+                    f"<td>{display}</td>"
+                    f"<td style='text-align:right'>{score:.4f}</td>"
+                    f"<td><code>{m['category_name'] or '?'}</code></td>"
+                    f"</tr>"
+                )
+            table = (
+                '<table style="width:100%">'
+                '<thead><tr>'
+                '<th>Link</th><th>Tag</th>'
+                '<th style="text-align:right">Score</th>'
+                '<th>Category</th>'
+                '</tr></thead>'
+                '<tbody>' + "".join(rows) + '</tbody></table>'
+            )
+
+            csv = ", ".join(format_tag(t, _use_underscore) for t, _ in items)
+            return table, csv
+
+        def on_analyze(image, url):
+            if image is None and not url:
+                return *refresh_results(
+                    top_k.value, min_score.value,
+                    sort_by.value, use_underscore.value,
+                    categories.value,
+                ), "⚠️ No image loaded."
+
+            pil = image
+            if pil is None and url:
+                import requests as std_requests
+                try:
+                    resp = std_requests.get(url, timeout=30)
+                    resp.raise_for_status()
+                    tmp = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
+                    tmp.write(resp.content)
+                    tmp.close()
+                    pil = Image.open(tmp.name).convert("RGB")
+                except Exception as exc:
+                    return "<i>Error loading URL.</i>", "", f"❌ {exc}"
+
+            state["current_image"] = pil
+            fn = state["predict_fn"]
+            if fn is None:
+                return "<i>No model loaded.</i>", "", "❌ No model loaded."
+
+            t0 = time.time()
+            all_logits = fn(pil)
+            state["all_logits"] = all_logits
+            state["tag_metadata"] = enrich_tags(all_logits)
+
+            table, csv = refresh_results(
+                top_k.value, min_score.value,
+                sort_by.value, use_underscore.value,
+                categories.value,
+            )
+            elapsed = time.time() - t0
+            n = len(state["tag_metadata"])
+            cached = _CACHE.size()
+            return table, csv, f"✅ {n} tags · {cached} cached · {elapsed:.2f}s"
+
+        analyze_btn.click(
+            fn=on_analyze,
+            inputs=[image_input, url_input],
+            outputs=[tag_table, tag_string, status],
+        )
+
+        url_input.submit(
+            fn=on_analyze,
+            inputs=[image_input, url_input],
+            outputs=[tag_table, tag_string, status],
+        )
+
+        def on_clear():
+            state["all_logits"] = None
+            state["tag_metadata"] = None
+            state["current_image"] = None
+            return None, "", "<i>No results yet.</i>", "Cleared.", "", ""
+
+        clear_btn.click(
+            fn=on_clear,
+            inputs=[],
+            outputs=[image_input, url_input, tag_table, tag_string,
+                     status, tag_query, tag_query_output],
+        )
+
+        def on_clear_cache():
+            _CACHE.clear()
+            return "🧹 Cache cleared (0 entries)"
+
+        clear_cache_btn.click(fn=on_clear_cache, inputs=[], outputs=[status])
+
+        for widget in [top_k, min_score, sort_by, use_underscore, categories]:
+            widget.change(
+                fn=refresh_results,
+                inputs=[top_k, min_score, sort_by, use_underscore, categories],
+                outputs=[tag_table, tag_string],
+            )
+
+        def query_tag_score(query):
+            meta = state.get("tag_metadata")
+            if not meta or not query:
+                return ""
+            query_l = query.lower()
+            matches = sorted(
+                [(t, meta[t]["score"]) for t in meta if query_l in t.lower()],
+                key=lambda x: x[1], reverse=True,
+            )[:20]
+            if not matches:
+                return "<i>No matching tags.</i>"
+            rows = "".join(
+                f"<tr><td>{t}</td><td>{s:.4f}</td>"
+                f"<td><code>{meta[t]['category_name'] or '?'}</code></td></tr>"
+                for t, s in matches
+            )
+            return (f"<table style='width:100%'>"
+                    f"<tr><th>Tag</th><th>Score</th><th>Category</th></tr>"
+                    f"{rows}</table>")
+
+        tag_query.change(
+            fn=query_tag_score,
+            inputs=[tag_query],
+            outputs=[tag_query_output],
+        )
+
+        copy_btn.click(
+            fn=lambda: None,
+            inputs=[],
+            outputs=[],
+            js="""() => {
+                const tb = document.querySelector('#csv-text textarea');
+                if (tb) { navigator.clipboard.writeText(tb.value); }
+            }"""
+        )
+
+        # ── model switcher ────────────────────────────────────────────────
+
+        def on_model_change(variant):
+            if not variant:
+                return "⚠️ No model selected"
+            try:
+                onnx_path = download_model_variant(variant)
+                predictor = Predictor(onnx_path)
+                state["predictor"] = predictor
+
+                def new_predict_fn(image: Image.Image) -> list[tuple[str, float]]:
+                    tensor = preprocess(image, predictor.image_size)
+                    logits = predictor.run(tensor)[0]
+                    inv = {v: k for k, v in predictor.tag_to_id.items()}
+                    indices = np.argsort(logits)[::-1]
+                    return [(inv[int(i)], float(logits[i])) for i in indices]
+
+                state["predict_fn"] = new_predict_fn
+                state["all_logits"] = None
+                state["tag_metadata"] = None
+                _try_prebuild(predictor)
+                return f"✅ Switched to {variant} ({predictor.num_classes} tags)"
+            except Exception as exc:
+                return f"❌ Failed to load model: {exc}"
+
+        model_dropdown.change(
+            fn=on_model_change,
+            inputs=[model_dropdown],
+            outputs=[model_status],
+        )
+
+    return app
+
+
+# ── main ────────────────────────────────────────────────────────────────────
+
+def _try_prebuild(predictor: Predictor):
+    try:
+        total, found = prebuild_cache(predictor.tag_to_id)
+        if total:
+            print(f"Cache prebuilt: {total} tags ({found} with category from HF dataset)")
+    except Exception as exc:
+        print(f"Cache prebuild skipped: {exc}")
+
+
+def main():
+    model_arg = sys.argv[1] if len(sys.argv) > 1 else None
+    model_env = os.environ.get("MODEL_DIR")
+    model_variants: list[str] = []
+    initial_predict_fn = None
+
+    if model_arg:
+        onnx = Path(model_arg)
+        if not onnx.exists():
+            print(f"Error: {onnx} not found", file=sys.stderr)
+            sys.exit(1)
+        print(f"Loading local model: {onnx}")
+        predictor = Predictor(onnx)
+        _try_prebuild(predictor)
+
+        def _predict(image: Image.Image) -> list[tuple[str, float]]:
+            tensor = preprocess(image, predictor.image_size)
+            logits = predictor.run(tensor)[0]
+            inv = {v: k for k, v in predictor.tag_to_id.items()}
+            indices = np.argsort(logits)[::-1]
+            return [(inv[int(i)], float(logits[i])) for i in indices]
+
+        initial_predict_fn = _predict
+    elif model_env:
+        env_dir = Path(model_env)
+        if not env_dir.is_dir():
+            print(f"Error: MODEL_DIR {env_dir} is not a directory", file=sys.stderr)
+            sys.exit(1)
+        onnx_files = list(env_dir.glob("*.onnx"))
+        if not onnx_files:
+            print(f"Error: no .onnx files in {env_dir}", file=sys.stderr)
+            sys.exit(1)
+        onnx = onnx_files[0]
+        print(f"Loading local model from MODEL_DIR: {onnx}")
+        predictor = Predictor(onnx)
+        _try_prebuild(predictor)
+
+        def _predict(image: Image.Image) -> list[tuple[str, float]]:
+            tensor = preprocess(image, predictor.image_size)
+            logits = predictor.run(tensor)[0]
+            inv = {v: k for k, v in predictor.tag_to_id.items()}
+            indices = np.argsort(logits)[::-1]
+            return [(inv[int(i)], float(logits[i])) for i in indices]
+
+        initial_predict_fn = _predict
+    else:
+        print("Discovering model variants on HF hub …")
+        model_variants = discover_model_variants()
+        if not model_variants:
+            print("Warning: no models found on hub.")
+        else:
+            print(f"Found variants: {model_variants}")
+            default = model_variants[0]
+            print(f"Downloading default model: {default} …")
+            try:
+                onnx_path = download_model_variant(default)
+                predictor = Predictor(onnx_path)
+                _try_prebuild(predictor)
+
+                def _predict(image: Image.Image) -> list[tuple[str, float]]:
+                    tensor = preprocess(image, predictor.image_size)
+                    logits = predictor.run(tensor)[0]
+                    inv = {v: k for k, v in predictor.tag_to_id.items()}
+                    indices = np.argsort(logits)[::-1]
+                    return [(inv[int(i)], float(logits[i])) for i in indices]
+
+                initial_predict_fn = _predict
+                print(f"Loaded {default} ({predictor.num_classes} tags)")
+            except Exception as exc:
+                print(f"Error loading default model: {exc}")
+
+    app = build_app(initial_predict_fn, model_variants)
+
+    host = os.environ.get("GRADIO_SERVER_NAME") or os.environ.get("HOST")
+    port_str = os.environ.get("GRADIO_SERVER_PORT") or os.environ.get("PORT")
+    port = int(port_str) if port_str else None
+
+    app.launch(server_name=host, server_port=port)
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,7 @@
+onnxruntime-gpu==1.18.0
+gradio>=5
+numpy
+Pillow
+huggingface-hub
+datasets
+tqdm