Grio43 commited on
Commit
b485bdf
·
verified ·
1 Parent(s): c7c5bd4

Add web interface (Gradio launcher + run.bat)

Browse files

Single-file Gradio launcher: bootstraps a venv, installs requirements, and serves a local UI on 127.0.0.1:7860. run.bat is a Windows convenience wrapper.

Files changed (2) hide show
  1. web_interface/app.py +656 -0
  2. web_interface/run.bat +25 -0
web_interface/app.py ADDED
@@ -0,0 +1,656 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Oppai ONNX tagger — single-file launcher.
3
+
4
+ First run: creates `.venv/`, installs requirements, re-execs inside the venv,
5
+ then starts a local Gradio web UI on http://127.0.0.1:7860 .
6
+
7
+ Subsequent runs: skip install (marker file) and start the UI immediately.
8
+
9
+ For most users:
10
+ Just run `run.bat` (or `py app.py`). The launcher will show a numbered
11
+ menu so you can pick an existing model or download one — no flags
12
+ required. Press Enter to accept the highlighted default at any prompt.
13
+
14
+ Advanced flags:
15
+ py app.py --reinstall # force re-install of requirements
16
+ py app.py --model-dir <folder> # skip the menu, load a specific folder
17
+
18
+ Models live in folders next to this script. Any folder containing
19
+ `model.onnx`, `selected_tags.csv`, and `preprocessing.json` is treated as a
20
+ model and will appear in the launcher menu and the UI's model picker. You
21
+ can also download variants from https://huggingface.co/Grio43/OppaiOracle
22
+ directly from the menu or from the UI.
23
+ """
24
+
25
+ from __future__ import annotations
26
+
27
+ import os
28
+ import subprocess
29
+ import sys
30
+ import venv
31
+ from pathlib import Path
32
+
33
+ ROOT = Path(__file__).resolve().parent
34
+ VENV_DIR = ROOT / ".venv"
35
+ MARKER = VENV_DIR / ".bootstrapped"
36
+
37
+ # Default folder if it exists; otherwise the first auto-discovered folder
38
+ # next to this script is used. Override with --model-dir or the UI picker.
39
+ DEFAULT_MODEL_DIR = ROOT / "V1.1_onnx"
40
+
41
+ # Variants published on HuggingFace that are usable with this ONNX runtime.
42
+ # First entry is the recommended default in interactive prompts.
43
+ HF_REPO_ID = "Grio43/OppaiOracle"
44
+ HF_VARIANTS = ["V1.1_onnx", "V1_onnx"]
45
+ HF_VARIANT_DESC = {
46
+ "V1.1_onnx": "448×448, higher accuracy",
47
+ "V1_onnx": "320×320, smaller and faster",
48
+ }
49
+
50
+ REQUIREMENTS = [
51
+ "onnxruntime>=1.20",
52
+ "pillow>=10.0",
53
+ "numpy>=1.26,<3",
54
+ "gradio>=4.44",
55
+ "huggingface_hub>=0.24",
56
+ ]
57
+
58
+
59
+ # ---------------------------------------------------------------------------
60
+ # Bootstrap
61
+ # ---------------------------------------------------------------------------
62
+
63
+ def _venv_python() -> Path:
64
+ if os.name == "nt":
65
+ return VENV_DIR / "Scripts" / "python.exe"
66
+ return VENV_DIR / "bin" / "python"
67
+
68
+
69
+ def _in_target_venv() -> bool:
70
+ # Belt-and-suspenders: compare both sys.executable and sys.prefix against
71
+ # the target venv. Windows Store Python uses reparse points that can make
72
+ # Path.resolve() on sys.executable return a path that differs from the
73
+ # venv's python.exe even when running inside it; sys.prefix is more
74
+ # reliable for that case. Either match counts as "in venv".
75
+ try:
76
+ target_py = _venv_python().resolve()
77
+ except OSError:
78
+ target_py = None
79
+ try:
80
+ target_dir = VENV_DIR.resolve()
81
+ except OSError:
82
+ target_dir = None
83
+ try:
84
+ if target_py is not None and Path(sys.executable).resolve() == target_py:
85
+ return True
86
+ except OSError:
87
+ pass
88
+ try:
89
+ if target_dir is not None and Path(sys.prefix).resolve() == target_dir:
90
+ return True
91
+ except OSError:
92
+ pass
93
+ return False
94
+
95
+
96
+ def _bootstrap(force_reinstall: bool) -> None:
97
+ if not VENV_DIR.exists():
98
+ print(f"[bootstrap] Creating virtualenv at {VENV_DIR} ...")
99
+ venv.EnvBuilder(with_pip=True, clear=False, upgrade_deps=False).create(VENV_DIR)
100
+
101
+ py = _venv_python()
102
+ needs_install = force_reinstall or not MARKER.exists()
103
+ if needs_install:
104
+ print("[bootstrap] Upgrading pip ...")
105
+ subprocess.check_call([str(py), "-m", "pip", "install", "--upgrade", "pip"])
106
+ print(f"[bootstrap] Installing: {', '.join(REQUIREMENTS)}")
107
+ subprocess.check_call([str(py), "-m", "pip", "install", *REQUIREMENTS])
108
+ MARKER.write_text("ok\n", encoding="utf-8")
109
+ else:
110
+ print("[bootstrap] Requirements already installed (delete .venv/.bootstrapped to redo).")
111
+
112
+ args = [a for a in sys.argv[1:] if a != "--reinstall"]
113
+ print("[bootstrap] Re-launching inside venv ...\n")
114
+ sys.exit(subprocess.call([str(py), str(Path(__file__).resolve()), *args]))
115
+
116
+
117
+ # ---------------------------------------------------------------------------
118
+ # App
119
+ # ---------------------------------------------------------------------------
120
+
121
+ REQUIRED_FILES = ("model.onnx", "selected_tags.csv", "preprocessing.json")
122
+
123
+
124
+ def _discover_model_dirs() -> list[Path]:
125
+ """Return every subdirectory of ROOT that looks like a usable model folder."""
126
+ out: list[Path] = []
127
+ if not ROOT.exists():
128
+ return out
129
+ for sub in sorted(ROOT.iterdir(), key=lambda p: p.name.lower()):
130
+ if not sub.is_dir():
131
+ continue
132
+ if all((sub / f).exists() for f in REQUIRED_FILES):
133
+ out.append(sub)
134
+ return out
135
+
136
+
137
+ def _variant_rank(name: str) -> int:
138
+ try:
139
+ return HF_VARIANTS.index(name)
140
+ except ValueError:
141
+ return len(HF_VARIANTS)
142
+
143
+
144
+ def _is_tty() -> bool:
145
+ try:
146
+ return sys.stdin.isatty() and sys.stdout.isatty()
147
+ except (AttributeError, OSError):
148
+ return False
149
+
150
+
151
+ def _prompt_choice(prompt: str, options: list[tuple[str, str]], default_idx: int = 0) -> str | None:
152
+ """Show a numbered terminal menu. Returns the chosen option's value, or None on EOF.
153
+
154
+ options: list of (display_text, value).
155
+ """
156
+ if not options:
157
+ return None
158
+ if not _is_tty():
159
+ return options[default_idx][1]
160
+
161
+ print()
162
+ print(prompt)
163
+ for i, (display, _) in enumerate(options, 1):
164
+ marker = " <- press Enter for this" if i - 1 == default_idx else ""
165
+ print(f" {i}) {display}{marker}")
166
+ while True:
167
+ try:
168
+ raw = input(f"Choice [1-{len(options)}, default {default_idx + 1}]: ").strip()
169
+ except EOFError:
170
+ return options[default_idx][1]
171
+ if not raw:
172
+ return options[default_idx][1]
173
+ try:
174
+ idx = int(raw) - 1
175
+ except ValueError:
176
+ print(f" Please enter a number 1-{len(options)}.")
177
+ continue
178
+ if 0 <= idx < len(options):
179
+ return options[idx][1]
180
+ print(f" Out of range. Pick 1-{len(options)}.")
181
+
182
+
183
+ def _download_variant(variant: str) -> Path | None:
184
+ """Download a HuggingFace variant into ROOT/<variant>. Returns the folder on success."""
185
+ try:
186
+ from huggingface_hub import snapshot_download
187
+ except ImportError:
188
+ print("[app] huggingface_hub is not installed — re-run with --reinstall.")
189
+ return None
190
+
191
+ print(f"[app] Downloading '{variant}' from huggingface.co/{HF_REPO_ID} ...")
192
+ try:
193
+ snapshot_download(
194
+ repo_id=HF_REPO_ID,
195
+ allow_patterns=[f"{variant}/*"],
196
+ local_dir=str(ROOT),
197
+ )
198
+ except Exception as e: # noqa: BLE001
199
+ print(f"[app] Download failed: {e}")
200
+ return None
201
+
202
+ target = ROOT / variant
203
+ missing = [f for f in REQUIRED_FILES if not (target / f).exists()]
204
+ if missing:
205
+ print(f"[app] Download finished but {target} is missing: {', '.join(missing)}")
206
+ return None
207
+ return target
208
+
209
+
210
+ def _interactive_pick_model() -> Path | None:
211
+ """Show a friendly menu so non-technical users can pick or download a model.
212
+
213
+ Returns the chosen model directory, or None if the user wants to start the
214
+ UI without loading anything (they can pick from the web UI then).
215
+ """
216
+ discovered = _discover_model_dirs()
217
+ discovered.sort(key=lambda p: (_variant_rank(p.name), p.name.lower()))
218
+ discovered_names = {p.name for p in discovered}
219
+
220
+ options: list[tuple[str, str]] = []
221
+ actions: list[tuple[str, str]] = [] # parallel list of (action, payload)
222
+
223
+ for p in discovered:
224
+ desc = HF_VARIANT_DESC.get(p.name, "model folder")
225
+ options.append((f"Use {p.name} ({desc})", str(p)))
226
+ actions.append(("load", str(p)))
227
+
228
+ for v in HF_VARIANTS:
229
+ if v in discovered_names:
230
+ continue
231
+ desc = HF_VARIANT_DESC.get(v, "")
232
+ suffix = f" ({desc})" if desc else ""
233
+ options.append((f"Download {v} from HuggingFace{suffix}", v))
234
+ actions.append(("download", v))
235
+
236
+ options.append(("Open the web UI without loading anything (pick later from the page)", "skip"))
237
+ actions.append(("skip", ""))
238
+
239
+ if not _is_tty() and discovered:
240
+ return discovered[0]
241
+ if not _is_tty():
242
+ return None
243
+
244
+ print()
245
+ print("=" * 50)
246
+ print(" Oppai ONNX Tagger")
247
+ print("=" * 50)
248
+ if discovered:
249
+ print(f"Found {len(discovered)} model folder(s) next to app.py.")
250
+ else:
251
+ print("No model folders found yet next to app.py.")
252
+ print(f"Pick a variant to download from huggingface.co/{HF_REPO_ID}.")
253
+
254
+ chosen = _prompt_choice("What would you like to do?", options, default_idx=0)
255
+ if chosen is None:
256
+ return None
257
+
258
+ idx = next(i for i, (_, v) in enumerate(options) if v == chosen)
259
+ action, payload = actions[idx]
260
+ if action == "load":
261
+ return Path(payload)
262
+ if action == "download":
263
+ return _download_variant(payload)
264
+ return None # skip
265
+
266
+
267
+ def _resolve_initial_model(cli_dir: str | None) -> Path | None:
268
+ if cli_dir:
269
+ p = Path(cli_dir).expanduser().resolve()
270
+ if not p.is_dir():
271
+ print(f"[app] --model-dir not a directory: {p}")
272
+ return None
273
+ missing = [f for f in REQUIRED_FILES if not (p / f).exists()]
274
+ if missing:
275
+ print(f"[app] --model-dir is missing required files: {', '.join(missing)}")
276
+ return None
277
+ return p
278
+ return _interactive_pick_model()
279
+
280
+
281
+ def _run_app() -> None:
282
+ import argparse
283
+ import csv
284
+ import json
285
+
286
+ import numpy as np
287
+ import onnxruntime as ort
288
+ import gradio as gr
289
+ from PIL import Image
290
+
291
+ parser = argparse.ArgumentParser(add_help=False)
292
+ parser.add_argument("--model-dir", type=str, default=None)
293
+ cli_args, _ = parser.parse_known_args()
294
+
295
+ cat_names = {0: "general", 1: "artist", 3: "copyright", 4: "character", 5: "meta"}
296
+ inv_cat_names = {v: k for k, v in cat_names.items()}
297
+
298
+ # Mutable holder so the UI can swap models without restarting the process.
299
+ state: dict = {
300
+ "session": None,
301
+ "tag_names": [],
302
+ "categories": [],
303
+ "skip_mask": None,
304
+ "image_size": 0,
305
+ "pad_color": (0, 0, 0),
306
+ "mean": None,
307
+ "std": None,
308
+ "breakeven_threshold": None,
309
+ "model_dir": None,
310
+ "providers": [],
311
+ }
312
+
313
+ def _ort_providers() -> list[str]:
314
+ available = ort.get_available_providers()
315
+ if "DmlExecutionProvider" in available:
316
+ return ["DmlExecutionProvider", "CPUExecutionProvider"]
317
+ if "CUDAExecutionProvider" in available:
318
+ return ["CUDAExecutionProvider", "CPUExecutionProvider"]
319
+ return ["CPUExecutionProvider"]
320
+
321
+ def load_model(model_dir: Path) -> str:
322
+ model_dir = Path(model_dir).expanduser().resolve()
323
+ if not model_dir.is_dir():
324
+ raise FileNotFoundError(f"not a directory: {model_dir}")
325
+ missing = [f for f in REQUIRED_FILES if not (model_dir / f).exists()]
326
+ if missing:
327
+ raise FileNotFoundError(
328
+ f"{model_dir} is missing required files: {', '.join(missing)}"
329
+ )
330
+
331
+ tag_names: list[str] = []
332
+ categories: list[int] = []
333
+ with (model_dir / "selected_tags.csv").open(encoding="utf-8") as f:
334
+ for row in csv.DictReader(f):
335
+ tag_names.append(row["name"])
336
+ categories.append(int(row["category"]))
337
+ n_tags = len(tag_names)
338
+
339
+ skip_mask = np.zeros(n_tags, dtype=bool)
340
+ for i, name in enumerate(tag_names):
341
+ if name in ("<PAD>", "<UNK>"):
342
+ skip_mask[i] = True
343
+
344
+ with (model_dir / "preprocessing.json").open(encoding="utf-8") as f:
345
+ preproc = json.load(f)
346
+ image_size = int(preproc["image_size"])
347
+ pad_color = tuple(int(c) for c in preproc["pad_color_rgb"])
348
+ mean = np.array(preproc["normalize_mean"], dtype=np.float32).reshape(3, 1, 1)
349
+ std = np.array(preproc["normalize_std"], dtype=np.float32).reshape(3, 1, 1)
350
+
351
+ # Calibrated breakeven (precision = recall) lives in pr_thresholds.json.
352
+ # It is tuned for whole-eval-set precision and is far too strict for
353
+ # interactive single-image tagging, so we surface it only as a hint.
354
+ breakeven_threshold = None
355
+ thr_path = model_dir / "pr_thresholds.json"
356
+ if thr_path.exists():
357
+ try:
358
+ with thr_path.open(encoding="utf-8") as f:
359
+ thr_data = json.load(f)
360
+ breakeven_threshold = float(thr_data["micro"]["pr_breakeven"]["threshold"])
361
+ except (OSError, KeyError, ValueError, json.JSONDecodeError):
362
+ pass
363
+
364
+ providers = _ort_providers()
365
+ sess_opts = ort.SessionOptions()
366
+ sess_opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
367
+ print(f"[app] Loading {model_dir / 'model.onnx'} ({image_size}×{image_size}) ...")
368
+ print(f"[app] Providers: {providers}")
369
+ session = ort.InferenceSession(
370
+ str(model_dir / "model.onnx"), sess_options=sess_opts, providers=providers
371
+ )
372
+
373
+ state.update(
374
+ session=session,
375
+ tag_names=tag_names,
376
+ categories=categories,
377
+ skip_mask=skip_mask,
378
+ image_size=image_size,
379
+ pad_color=pad_color,
380
+ mean=mean,
381
+ std=std,
382
+ breakeven_threshold=breakeven_threshold,
383
+ model_dir=model_dir,
384
+ providers=providers,
385
+ )
386
+ return _status_md()
387
+
388
+ def _status_md() -> str:
389
+ if state["session"] is None:
390
+ return (
391
+ "**No model loaded.** Drop an ONNX model folder next to "
392
+ "`app.py`, or use the **Download from HuggingFace** section below."
393
+ )
394
+ try:
395
+ display = state["model_dir"].relative_to(ROOT)
396
+ except ValueError:
397
+ display = state["model_dir"]
398
+ parts = [
399
+ f"**Loaded:** `{display}`",
400
+ f"{state['image_size']}×{state['image_size']}",
401
+ f"{len(state['tag_names'])} tags",
402
+ f"providers: {', '.join(state['providers'])}",
403
+ ]
404
+ if state["breakeven_threshold"] is not None:
405
+ parts.append(f"P=R breakeven: {state['breakeven_threshold']:.3f}")
406
+ return " — ".join(parts)
407
+
408
+ def _dropdown_choices() -> list[tuple[str, str]]:
409
+ out = []
410
+ for p in _discover_model_dirs():
411
+ try:
412
+ label = str(p.relative_to(ROOT))
413
+ except ValueError:
414
+ label = p.name
415
+ out.append((label, str(p)))
416
+ return out
417
+
418
+ def _current_value() -> str | None:
419
+ return str(state["model_dir"]) if state["model_dir"] else None
420
+
421
+ # Initial load (CLI override > default folder > first discovered)
422
+ initial = _resolve_initial_model(cli_args.model_dir)
423
+ if initial is not None:
424
+ try:
425
+ load_model(initial)
426
+ except Exception as e: # noqa: BLE001
427
+ print(f"[app] Initial model load failed: {e!r}")
428
+ else:
429
+ print("[app] No model folder found yet — pick or download one in the UI.")
430
+
431
+ def letterbox(img: Image.Image):
432
+ img = img.convert("RGB")
433
+ w, h = img.size
434
+ size = state["image_size"]
435
+ scale = min(size / w, size / h)
436
+ nw, nh = max(1, int(round(w * scale))), max(1, int(round(h * scale)))
437
+ resized = img.resize((nw, nh), Image.BICUBIC)
438
+ canvas = Image.new("RGB", (size, size), state["pad_color"])
439
+ x0 = (size - nw) // 2
440
+ y0 = (size - nh) // 2
441
+ canvas.paste(resized, (x0, y0))
442
+ mask = np.ones((size, size), dtype=bool) # True = padded
443
+ mask[y0:y0 + nh, x0:x0 + nw] = False
444
+ return canvas, mask
445
+
446
+ def preprocess(img: Image.Image):
447
+ canvas, mask = letterbox(img)
448
+ arr = np.asarray(canvas, dtype=np.float32) / 255.0
449
+ arr = arr.transpose(2, 0, 1) # CHW
450
+ arr = (arr - state["mean"]) / state["std"]
451
+ return arr.astype(np.float32), mask
452
+
453
+ def predict(image, threshold: float, max_tags, category_filter):
454
+ if state["session"] is None:
455
+ return "", "*no model loaded — pick or download one above*"
456
+ if image is None:
457
+ return "", "*upload an image to start*"
458
+ try:
459
+ max_tags_i = int(max_tags) if max_tags is not None else 0
460
+ if max_tags_i <= 0:
461
+ return "", "*no tags above threshold*"
462
+
463
+ # An empty list means "no categories selected" -> show nothing.
464
+ # `None` (event before component initialized) means "no filter".
465
+ if category_filter is None:
466
+ keep_cats = None
467
+ else:
468
+ keep_cats = {inv_cat_names[c] for c in category_filter if c in inv_cat_names}
469
+ if not keep_cats:
470
+ return "", "*no tags above threshold*"
471
+
472
+ pixel_values, padding_mask = preprocess(image)
473
+ outputs = state["session"].run(
474
+ ["probabilities"],
475
+ {
476
+ "pixel_values": pixel_values[None, ...],
477
+ "padding_mask": padding_mask[None, ...],
478
+ },
479
+ )
480
+ probs = outputs[0][0].astype(np.float32)
481
+ probs[state["skip_mask"]] = -1.0 # never surface PAD/UNK
482
+
483
+ order = np.argsort(-probs)
484
+ results = []
485
+ tag_names = state["tag_names"]
486
+ categories = state["categories"]
487
+ for idx in order:
488
+ p = float(probs[idx])
489
+ if p < threshold:
490
+ break
491
+ cat = categories[idx]
492
+ if keep_cats is not None and cat not in keep_cats:
493
+ continue
494
+ results.append((tag_names[idx], p, cat))
495
+ if len(results) >= max_tags_i:
496
+ break
497
+
498
+ if not results:
499
+ return "", "*no tags above threshold*"
500
+
501
+ comma = ", ".join(name.replace("_", " ") for name, _, _ in results)
502
+ lines = ["| # | Tag | Confidence | Category |", "|---|---|---|---|"]
503
+ for i, (name, p, cat) in enumerate(results, 1):
504
+ lines.append(f"| {i} | `{name}` | {p:.3f} | {cat_names.get(cat, str(cat))} |")
505
+ return comma, "\n".join(lines)
506
+ except Exception as e: # noqa: BLE001 — keep Gradio toast away
507
+ print(f"[app] predict() error: {e!r}")
508
+ return "", f"*error during inference: {e}*"
509
+
510
+ # --- UI callbacks ------------------------------------------------------
511
+
512
+ def on_refresh():
513
+ choices = _dropdown_choices()
514
+ return gr.update(choices=choices, value=_current_value()), _status_md()
515
+
516
+ def on_load(dropdown_value: str | None, custom_path: str):
517
+ target = (custom_path or "").strip() or dropdown_value
518
+ if not target:
519
+ return gr.update(), _status_md(), "Pick a model folder or paste a path first."
520
+ try:
521
+ load_model(Path(target))
522
+ except Exception as e: # noqa: BLE001
523
+ return gr.update(), _status_md(), f"Load failed: {e}"
524
+ choices = _dropdown_choices()
525
+ return (
526
+ gr.update(choices=choices, value=_current_value()),
527
+ _status_md(),
528
+ f"Loaded `{Path(target).name}`.",
529
+ )
530
+
531
+ def on_download(variant: str, progress=gr.Progress(track_tqdm=True)):
532
+ if not variant:
533
+ return gr.update(), _status_md(), "Pick a variant first."
534
+ try:
535
+ from huggingface_hub import snapshot_download
536
+ except ImportError:
537
+ return (
538
+ gr.update(),
539
+ _status_md(),
540
+ "huggingface_hub is not installed — re-run `app.py --reinstall`.",
541
+ )
542
+ progress(0, desc=f"Downloading {variant} from {HF_REPO_ID} ...")
543
+ try:
544
+ snapshot_download(
545
+ repo_id=HF_REPO_ID,
546
+ allow_patterns=[f"{variant}/*"],
547
+ local_dir=str(ROOT),
548
+ )
549
+ except Exception as e: # noqa: BLE001
550
+ return gr.update(), _status_md(), f"Download failed: {e}"
551
+
552
+ target = ROOT / variant
553
+ msg = f"Downloaded `{variant}`."
554
+ if all((target / f).exists() for f in REQUIRED_FILES):
555
+ try:
556
+ load_model(target)
557
+ msg += f" Loaded `{variant}`."
558
+ except Exception as e: # noqa: BLE001
559
+ msg += f" Load failed: {e}"
560
+ choices = _dropdown_choices()
561
+ return gr.update(choices=choices, value=_current_value()), _status_md(), msg
562
+
563
+ # --- UI layout ---------------------------------------------------------
564
+
565
+ with gr.Blocks(title="Oppai ONNX Tagger") as demo:
566
+ gr.Markdown(
567
+ "# Oppai ONNX Tagger\n"
568
+ "Upload an image and tweak the threshold / max tags. "
569
+ "Pick a model below or download one from "
570
+ "[Grio43/OppaiOracle](https://huggingface.co/Grio43/OppaiOracle)."
571
+ )
572
+
573
+ with gr.Accordion("Model", open=True):
574
+ with gr.Row():
575
+ model_dd = gr.Dropdown(
576
+ choices=_dropdown_choices(),
577
+ value=_current_value(),
578
+ label="Detected model folders (next to app.py)",
579
+ interactive=True,
580
+ scale=3,
581
+ )
582
+ refresh_btn = gr.Button("Refresh", scale=1)
583
+ with gr.Row():
584
+ custom_path = gr.Textbox(
585
+ label="…or paste a custom model folder path (overrides dropdown)",
586
+ placeholder=r"e.g. C:\models\my_onnx_folder",
587
+ scale=4,
588
+ )
589
+ load_btn = gr.Button("Load", variant="primary", scale=1)
590
+ with gr.Row():
591
+ hf_dd = gr.Dropdown(
592
+ choices=HF_VARIANTS,
593
+ value=HF_VARIANTS[0],
594
+ label=f"Download a variant from {HF_REPO_ID}",
595
+ scale=3,
596
+ )
597
+ download_btn = gr.Button("Download", scale=1)
598
+ status_md = gr.Markdown(_status_md())
599
+ action_msg = gr.Markdown("")
600
+
601
+ with gr.Row():
602
+ with gr.Column(scale=1):
603
+ inp = gr.Image(type="pil", label="Image", height=448)
604
+ threshold = gr.Slider(
605
+ 0.0, 1.0,
606
+ value=0.35,
607
+ step=0.005,
608
+ label="Threshold (interactive default 0.35; calibrated breakeven shown above)",
609
+ )
610
+ max_tags = gr.Slider(1, 200, value=50, step=1, label="Max tags")
611
+ cats = gr.CheckboxGroup(
612
+ choices=list(cat_names.values()),
613
+ value=list(cat_names.values()),
614
+ label="Categories to include",
615
+ )
616
+ btn = gr.Button("Tag image", variant="primary")
617
+ with gr.Column(scale=1):
618
+ tags_out = gr.Textbox(
619
+ label="Tags (comma-separated, underscores → spaces)",
620
+ lines=5,
621
+ )
622
+ table_out = gr.Markdown(label="Per-tag detail")
623
+
624
+ refresh_btn.click(on_refresh, outputs=[model_dd, status_md])
625
+ load_btn.click(on_load, inputs=[model_dd, custom_path], outputs=[model_dd, status_md, action_msg])
626
+ download_btn.click(on_download, inputs=[hf_dd], outputs=[model_dd, status_md, action_msg])
627
+
628
+ ev_inputs = [inp, threshold, max_tags, cats]
629
+ ev_outputs = [tags_out, table_out]
630
+ btn.click(predict, ev_inputs, ev_outputs)
631
+ inp.change(predict, ev_inputs, ev_outputs)
632
+ threshold.release(predict, ev_inputs, ev_outputs)
633
+ max_tags.release(predict, ev_inputs, ev_outputs)
634
+ cats.change(predict, ev_inputs, ev_outputs)
635
+
636
+ # CPU inference is ~1-3s per image; cap concurrency so spammed slider
637
+ # changes queue serially instead of fighting for the same model session.
638
+ demo.queue(default_concurrency_limit=1).launch(
639
+ server_name="127.0.0.1", server_port=7860, inbrowser=True
640
+ )
641
+
642
+
643
+ # ---------------------------------------------------------------------------
644
+ # Entrypoint
645
+ # ---------------------------------------------------------------------------
646
+
647
+ def main() -> None:
648
+ force = "--reinstall" in sys.argv[1:]
649
+ if not _in_target_venv():
650
+ _bootstrap(force_reinstall=force)
651
+ return # _bootstrap re-execs and exits
652
+ _run_app()
653
+
654
+
655
+ if __name__ == "__main__":
656
+ main()
web_interface/run.bat ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ @echo off
2
+ setlocal
3
+ cd /d "%~dp0"
4
+
5
+ rem Prefer the per-user py launcher; fall back to python on PATH.
6
+ where py >nul 2>nul
7
+ if %ERRORLEVEL%==0 (
8
+ py "%~dp0app.py" %*
9
+ ) else (
10
+ where python >nul 2>nul
11
+ if %ERRORLEVEL%==0 (
12
+ python "%~dp0app.py" %*
13
+ ) else (
14
+ echo [run.bat] No Python found. Install Python 3.10+ from python.org or the Microsoft Store.
15
+ pause
16
+ exit /b 1
17
+ )
18
+ )
19
+
20
+ if %ERRORLEVEL% NEQ 0 (
21
+ echo.
22
+ echo [run.bat] app.py exited with code %ERRORLEVEL%.
23
+ pause
24
+ )
25
+ endlocal