Dinamush commited on
Commit
ee6a4bd
·
1 Parent(s): 7419da4

Enhance README and backend API for GPU inference support. Updated README to include GPU visibility verification, runtime controls, and examples for inference settings. Modified backend API to add new endpoints for provider health checks and improved inference handling with batch processing capabilities. Updated database schema to include scan_recursive setting and enhanced telemetry tracking for inference runs. Improved tests to validate new API responses and functionality.

Browse files
README.md CHANGED
@@ -105,11 +105,44 @@ The frontend automatically falls back to local mock mode when backend requests f
105
 
106
  ## GPU Acceleration
107
 
108
- Inference attempts to use ONNX Runtime CUDA provider when available. If CUDA/cuDNN dependencies are missing, runtime falls back and logs provider errors. Ensure CUDA 12.x + cuDNN 9.x compatibility for GPU inference.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
109
 
110
  ## API Surface (High Level)
111
 
112
  - `GET /health`
 
 
113
  - `GET/PUT /api/settings`
114
  - `GET /api/tags`
115
  - `POST /api/runs/start`
 
105
 
106
  ## GPU Acceleration
107
 
108
+ Inference attempts to use ONNX Runtime CUDA provider when available. If CUDA/cuDNN dependencies are missing, runtime falls back to CPU and logs provider errors.
109
+
110
+ ### Verify GPU visibility
111
+
112
+ ```bash
113
+ curl http://127.0.0.1:8000/health/providers
114
+ curl http://127.0.0.1:8000/api/providers
115
+ ```
116
+
117
+ `likely_device: "gpu"` means CUDA provider is visible and CPU is not being force-disabled.
118
+
119
+ ### Runtime controls
120
+
121
+ - `MAX_INFERENCE_WORKERS` (default `2`, clamped to `1..16`): controls conservative thread-pool parallelism for per-image inference.
122
+ - `FORCE_CPU_INFERENCE=true`: force reported/expected CPU path even if CUDA provider is available.
123
+ - `INFERENCE_MODE=batch|single` (default `batch`): prefer GPU-first batched inference or legacy single-image inference.
124
+ - `INFERENCE_BATCH_SIZE` (default `8`, clamped to `1..64`): request batch size for batched inference. If batch inference fails, runtime auto-falls back by splitting batches down to single-image.
125
+ - `QUEUE_SHUFFLE_ENABLED=true|false` (default `true`): stochastic queue ordering toggle.
126
+ - `QUEUE_SHUFFLE_SEED=<int>` (default `run_id`): deterministic seed for reproducible queue shuffling.
127
+
128
+ Examples:
129
+
130
+ ```bash
131
+ # CPU-safe baseline
132
+ FORCE_CPU_INFERENCE=true MAX_INFERENCE_WORKERS=1 ../.venv/bin/uvicorn app.main:app --host 127.0.0.1 --port 8000 --reload
133
+
134
+ # Throughput mode (tune workers to your CPU/GPU memory limits)
135
+ MAX_INFERENCE_WORKERS=4 ../.venv/bin/uvicorn app.main:app --host 127.0.0.1 --port 8000 --reload
136
+
137
+ # GPU-first batching + seeded stochastic queue
138
+ INFERENCE_MODE=batch INFERENCE_BATCH_SIZE=8 QUEUE_SHUFFLE_ENABLED=true QUEUE_SHUFFLE_SEED=1337 MAX_INFERENCE_WORKERS=4 ../.venv/bin/uvicorn app.main:app --host 127.0.0.1 --port 8000 --reload
139
+ ```
140
 
141
  ## API Surface (High Level)
142
 
143
  - `GET /health`
144
+ - `GET /health/providers`
145
+ - `GET /api/providers`
146
  - `GET/PUT /api/settings`
147
  - `GET /api/tags`
148
  - `POST /api/runs/start`
backend/app.db-journal ADDED
Binary file (16.9 kB). View file
 
backend/app/api.py CHANGED
@@ -1,9 +1,14 @@
1
  from __future__ import annotations
2
 
 
3
  import logging
 
4
  import threading
 
 
5
  from datetime import datetime, timezone
6
  from pathlib import Path
 
7
 
8
  from fastapi import APIRouter, HTTPException, Query
9
  from fastapi.responses import FileResponse
@@ -24,6 +29,7 @@ from .services import (
24
  choose_best_tags,
25
  discover_tag_folders,
26
  extract_scores,
 
27
  load_known_tags,
28
  migrate_file,
29
  resolve_settings,
@@ -42,12 +48,188 @@ SUPPORTED_PREVIEW_SUFFIXES = {
42
  ".webp": "image/webp",
43
  ".bmp": "image/bmp",
44
  }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
45
 
46
 
47
  def _now_iso() -> str:
48
  return datetime.now(timezone.utc).isoformat()
49
 
50
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
51
  def _settings_from_db() -> AppSettings:
52
  row = fetch_one("SELECT * FROM settings WHERE id = 1")
53
  if not row:
@@ -57,6 +239,7 @@ def _settings_from_db() -> AppSettings:
57
  categories_root=row["categories_root"],
58
  confidence_threshold=float(row["confidence_threshold"]),
59
  default_migrate_mode=row["default_migrate_mode"],
 
60
  )
61
 
62
 
@@ -90,6 +273,7 @@ def _run_status_from_row(row: dict) -> RunStatusResponse:
90
  pct = 0.0 if total <= 0 else min(100.0, (processed / total) * 100.0)
91
  item_count_row = fetch_one("SELECT COUNT(*) AS cnt FROM items WHERE run_id = ?", (row["id"],))
92
  has_items = bool(item_count_row and int(item_count_row["cnt"]) > 0)
 
93
  return RunStatusResponse(
94
  run_id=row["id"],
95
  status=row.get("status") or "pending",
@@ -102,6 +286,10 @@ def _run_status_from_row(row: dict) -> RunStatusResponse:
102
  last_error=row.get("last_error"),
103
  cancel_requested=bool(row.get("cancel_requested") or 0),
104
  has_items=has_items,
 
 
 
 
105
  )
106
 
107
 
@@ -123,80 +311,175 @@ def _execute_run(
123
  categories_root: Path,
124
  confidence_threshold: float,
125
  matched_tags: set[str],
 
126
  ) -> None:
127
  try:
128
  execute(
129
  "UPDATE runs SET status = 'running', started_at = ?, last_error = NULL WHERE id = ?",
130
  (_now_iso(), run_id),
131
  )
132
- scan_output = scan_images(root_repo)
133
  execute("UPDATE runs SET total_images = ? WHERE id = ?", (scan_output.stats.eligible_images, run_id))
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
134
 
135
  processed = 0
136
  failed = 0
137
- for image_path in scan_output.image_paths:
138
- if _is_cancel_requested(run_id):
139
- execute(
140
- "UPDATE runs SET status = 'cancelled', finished_at = ? WHERE id = ?",
141
- (_now_iso(), run_id),
142
- )
143
- logger.info("run_cancelled run_id=%d processed=%d", run_id, processed)
144
- return
145
-
146
- relative_path = str(image_path.relative_to(root_repo))
147
- try:
148
- scores = extract_scores(image_path)
149
- primary_tag, primary_score, secondary = choose_best_tags(scores, matched_tags)
150
- needs_review = False
151
- reason = None
152
- except Exception:
153
- logger.exception("inference_failed run_id=%d image=%s", run_id, image_path)
154
- scores, primary_tag, primary_score, secondary = {}, None, None, []
155
- needs_review = True
156
- reason = "Inference failed for this image; requires manual review."
157
- failed += 1
158
 
159
- if primary_tag is None and reason is None:
160
- needs_review = True
161
- reason = "No matching tags found in selected folders."
162
- elif (
163
- primary_score is not None
164
- and primary_score < confidence_threshold
165
- and reason is None
166
- ):
167
- needs_review = True
168
- reason = f"Below threshold ({primary_score:.3f} < {confidence_threshold:.3f})."
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
169
 
170
- suggested_destination = (
171
- str(categories_root / primary_tag) if primary_tag is not None else None
172
- )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
173
  execute(
174
  """
175
- INSERT INTO items (
176
- run_id, file_path, relative_path, primary_tag, primary_score, secondary_json,
177
- full_scores_json, suggested_destination, final_tag, final_destination, status, needs_review, review_reason
178
- ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
 
 
 
179
  """,
180
- (
181
- run_id,
182
- str(image_path),
183
- relative_path,
184
- primary_tag,
185
- primary_score,
186
- to_json(secondary),
187
- to_json(scores),
188
- suggested_destination,
189
- primary_tag,
190
- suggested_destination,
191
- "approved" if not needs_review else "proposed",
192
- 1 if needs_review else 0,
193
- reason,
194
- ),
195
  )
196
- processed += 1
197
- _update_run_progress(run_id, processed, failed)
198
- if processed % 10 == 0:
199
- logger.info("run_progress run_id=%d processed=%d total=%d", run_id, processed, scan_output.stats.eligible_images)
200
 
201
  execute(
202
  "UPDATE runs SET status = 'completed', finished_at = ? WHERE id = ?",
@@ -222,7 +505,8 @@ def save_settings(payload: SaveSettingsRequest) -> AppSettings:
222
  execute(
223
  """
224
  UPDATE settings
225
- SET root_repo = ?, categories_root = ?, confidence_threshold = ?, default_migrate_mode = ?
 
226
  WHERE id = 1
227
  """,
228
  (
@@ -230,6 +514,7 @@ def save_settings(payload: SaveSettingsRequest) -> AppSettings:
230
  payload.categories_root,
231
  payload.confidence_threshold,
232
  payload.default_migrate_mode,
 
233
  ),
234
  )
235
  except Exception:
@@ -246,6 +531,59 @@ def search_tags(query: str = Query("", min_length=0), limit: int = 50) -> dict:
246
  return {"items": known[:limit], "count": len(known)}
247
 
248
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
249
  @router.post("/runs/start", response_model=StartRunResponse)
250
  def start_run(payload: StartRunRequest) -> StartRunResponse:
251
  logger.info("run_start_requested")
@@ -287,7 +625,8 @@ def start_run(payload: StartRunRequest) -> StartRunResponse:
287
 
288
  worker = threading.Thread(
289
  target=_execute_run,
290
- args=(run_id, root_repo, categories_root, resolved.confidence_threshold, matched_tags),
 
291
  daemon=True,
292
  )
293
  worker.start()
 
1
  from __future__ import annotations
2
 
3
+ import os
4
  import logging
5
+ import random
6
  import threading
7
+ from concurrent.futures import FIRST_COMPLETED, Future, ThreadPoolExecutor, wait
8
+ from dataclasses import dataclass
9
  from datetime import datetime, timezone
10
  from pathlib import Path
11
+ import time
12
 
13
  from fastapi import APIRouter, HTTPException, Query
14
  from fastapi.responses import FileResponse
 
29
  choose_best_tags,
30
  discover_tag_folders,
31
  extract_scores,
32
+ extract_scores_batch,
33
  load_known_tags,
34
  migrate_file,
35
  resolve_settings,
 
48
  ".webp": "image/webp",
49
  ".bmp": "image/bmp",
50
  }
51
+ _RUN_TELEMETRY: dict[int, dict[str, object]] = {}
52
+ _RUN_TELEMETRY_LOCK = threading.Lock()
53
+
54
+
55
+ @dataclass
56
+ class _ImageInferenceResult:
57
+ image_path: Path
58
+ scores: dict[str, float]
59
+ primary_tag: str | None
60
+ primary_score: float | None
61
+ secondary: list[dict[str, float]]
62
+ needs_review: bool
63
+ reason: str | None
64
+ inference_failed: bool
65
 
66
 
67
  def _now_iso() -> str:
68
  return datetime.now(timezone.utc).isoformat()
69
 
70
 
71
+ def _get_max_inference_workers() -> int:
72
+ raw = os.getenv("MAX_INFERENCE_WORKERS", "2").strip()
73
+ try:
74
+ value = int(raw)
75
+ except ValueError:
76
+ value = 2
77
+ return max(1, min(value, 16))
78
+
79
+
80
+ def _get_inference_mode() -> str:
81
+ raw = os.getenv("INFERENCE_MODE", "batch").strip().lower()
82
+ return "single" if raw == "single" else "batch"
83
+
84
+
85
+ def _get_inference_batch_size() -> int:
86
+ raw = os.getenv("INFERENCE_BATCH_SIZE", "8").strip()
87
+ try:
88
+ value = int(raw)
89
+ except ValueError:
90
+ value = 8
91
+ return max(1, min(value, 64))
92
+
93
+
94
+ def _get_queue_shuffle_enabled() -> bool:
95
+ raw = os.getenv("QUEUE_SHUFFLE_ENABLED", "true").strip().lower()
96
+ return raw not in {"0", "false", "no", "off"}
97
+
98
+
99
+ def _get_queue_shuffle_seed(run_id: int) -> int:
100
+ raw = os.getenv("QUEUE_SHUFFLE_SEED", "").strip()
101
+ if raw:
102
+ try:
103
+ return int(raw)
104
+ except ValueError:
105
+ pass
106
+ return int(run_id)
107
+
108
+
109
+ def _set_run_telemetry(run_id: int, **kwargs) -> None:
110
+ with _RUN_TELEMETRY_LOCK:
111
+ telemetry = _RUN_TELEMETRY.get(run_id, {})
112
+ telemetry.update(kwargs)
113
+ _RUN_TELEMETRY[run_id] = telemetry
114
+
115
+
116
+ def _get_run_telemetry(run_id: int) -> dict[str, object]:
117
+ with _RUN_TELEMETRY_LOCK:
118
+ return dict(_RUN_TELEMETRY.get(run_id, {}))
119
+
120
+
121
+ def _is_provider_related_error(err: Exception) -> bool:
122
+ msg = str(err).lower()
123
+ keywords = ("cuda", "cudnn", "executionprovider", "provider", "onnxruntime", "gpu")
124
+ return any(k in msg for k in keywords)
125
+
126
+
127
+ def _infer_one_image(
128
+ image_path: Path, matched_tags: set[str], confidence_threshold: float
129
+ ) -> _ImageInferenceResult:
130
+ try:
131
+ scores = extract_scores(image_path)
132
+ primary_tag, primary_score, secondary = choose_best_tags(scores, matched_tags)
133
+ needs_review = False
134
+ reason = None
135
+ if primary_tag is None:
136
+ needs_review = True
137
+ reason = "No matching tags found in selected folders."
138
+ elif primary_score is not None and primary_score < confidence_threshold:
139
+ needs_review = True
140
+ reason = f"Below threshold ({primary_score:.3f} < {confidence_threshold:.3f})."
141
+ return _ImageInferenceResult(
142
+ image_path=image_path,
143
+ scores=scores,
144
+ primary_tag=primary_tag,
145
+ primary_score=primary_score,
146
+ secondary=secondary,
147
+ needs_review=needs_review,
148
+ reason=reason,
149
+ inference_failed=False,
150
+ )
151
+ except Exception as err:
152
+ if _is_provider_related_error(err):
153
+ logger.exception("inference_provider_failure image=%s", image_path)
154
+ else:
155
+ logger.exception("inference_failed image=%s", image_path)
156
+ return _ImageInferenceResult(
157
+ image_path=image_path,
158
+ scores={},
159
+ primary_tag=None,
160
+ primary_score=None,
161
+ secondary=[],
162
+ needs_review=True,
163
+ reason="Inference failed for this image; requires manual review.",
164
+ inference_failed=True,
165
+ )
166
+
167
+
168
+ def _infer_batch_with_fallback(
169
+ image_paths: list[Path],
170
+ matched_tags: set[str],
171
+ confidence_threshold: float,
172
+ requested_mode: str,
173
+ ) -> tuple[list[_ImageInferenceResult], float, str]:
174
+ if not image_paths:
175
+ return [], 0.0, "none"
176
+
177
+ if requested_mode != "batch" or len(image_paths) == 1:
178
+ start = time.perf_counter()
179
+ rows = [_infer_one_image(image_paths[0], matched_tags, confidence_threshold)]
180
+ elapsed_ms = (time.perf_counter() - start) * 1000.0
181
+ mode = "single" if requested_mode == "single" else "single_fallback"
182
+ return rows, elapsed_ms, mode
183
+
184
+ start = time.perf_counter()
185
+ try:
186
+ scores_by_image = extract_scores_batch(image_paths)
187
+ if len(scores_by_image) != len(image_paths):
188
+ raise RuntimeError("Batch inference result count mismatch")
189
+ rows: list[_ImageInferenceResult] = []
190
+ for image_path, scores in zip(image_paths, scores_by_image):
191
+ primary_tag, primary_score, secondary = choose_best_tags(scores, matched_tags)
192
+ needs_review = False
193
+ reason = None
194
+ if primary_tag is None:
195
+ needs_review = True
196
+ reason = "No matching tags found in selected folders."
197
+ elif primary_score is not None and primary_score < confidence_threshold:
198
+ needs_review = True
199
+ reason = f"Below threshold ({primary_score:.3f} < {confidence_threshold:.3f})."
200
+ rows.append(
201
+ _ImageInferenceResult(
202
+ image_path=image_path,
203
+ scores=scores,
204
+ primary_tag=primary_tag,
205
+ primary_score=primary_score,
206
+ secondary=secondary,
207
+ needs_review=needs_review,
208
+ reason=reason,
209
+ inference_failed=False,
210
+ )
211
+ )
212
+ elapsed_ms = (time.perf_counter() - start) * 1000.0
213
+ return rows, elapsed_ms, "batch"
214
+ except Exception as err:
215
+ if _is_provider_related_error(err):
216
+ logger.exception("batch_inference_provider_failure batch_size=%d", len(image_paths))
217
+ else:
218
+ logger.exception("batch_inference_failed batch_size=%d", len(image_paths))
219
+ if len(image_paths) > 1:
220
+ mid = len(image_paths) // 2
221
+ left_rows, left_ms, _ = _infer_batch_with_fallback(
222
+ image_paths[:mid], matched_tags, confidence_threshold, "batch"
223
+ )
224
+ right_rows, right_ms, _ = _infer_batch_with_fallback(
225
+ image_paths[mid:], matched_tags, confidence_threshold, "batch"
226
+ )
227
+ return left_rows + right_rows, left_ms + right_ms, "batch_fallback"
228
+ row = _infer_one_image(image_paths[0], matched_tags, confidence_threshold)
229
+ elapsed_ms = (time.perf_counter() - start) * 1000.0
230
+ return [row], elapsed_ms, "single_fallback"
231
+
232
+
233
  def _settings_from_db() -> AppSettings:
234
  row = fetch_one("SELECT * FROM settings WHERE id = 1")
235
  if not row:
 
239
  categories_root=row["categories_root"],
240
  confidence_threshold=float(row["confidence_threshold"]),
241
  default_migrate_mode=row["default_migrate_mode"],
242
+ scan_recursive=bool(row.get("scan_recursive", 1)),
243
  )
244
 
245
 
 
273
  pct = 0.0 if total <= 0 else min(100.0, (processed / total) * 100.0)
274
  item_count_row = fetch_one("SELECT COUNT(*) AS cnt FROM items WHERE run_id = ?", (row["id"],))
275
  has_items = bool(item_count_row and int(item_count_row["cnt"]) > 0)
276
+ telemetry = _get_run_telemetry(row["id"])
277
  return RunStatusResponse(
278
  run_id=row["id"],
279
  status=row.get("status") or "pending",
 
286
  last_error=row.get("last_error"),
287
  cancel_requested=bool(row.get("cancel_requested") or 0),
288
  has_items=has_items,
289
+ inference_mode=telemetry.get("inference_mode"),
290
+ batch_size=telemetry.get("batch_size"),
291
+ avg_infer_ms_per_image=telemetry.get("avg_infer_ms_per_image"),
292
+ queue_seed=telemetry.get("queue_seed"),
293
  )
294
 
295
 
 
311
  categories_root: Path,
312
  confidence_threshold: float,
313
  matched_tags: set[str],
314
+ scan_recursive: bool = True,
315
  ) -> None:
316
  try:
317
  execute(
318
  "UPDATE runs SET status = 'running', started_at = ?, last_error = NULL WHERE id = ?",
319
  (_now_iso(), run_id),
320
  )
321
+ scan_output = scan_images(root_repo, exclude_dirs={categories_root}, recursive=scan_recursive)
322
  execute("UPDATE runs SET total_images = ? WHERE id = ?", (scan_output.stats.eligible_images, run_id))
323
+ logger.info(
324
+ "run_scan_complete run_id=%d total_files=%d eligible=%d "
325
+ "ignored_gif=%d ignored_unsupported=%d failed_to_read=%d",
326
+ run_id,
327
+ scan_output.stats.total_files,
328
+ scan_output.stats.eligible_images,
329
+ scan_output.stats.ignored_gif,
330
+ scan_output.stats.ignored_unsupported,
331
+ scan_output.stats.failed_to_read,
332
+ )
333
+
334
+ queue_shuffle_enabled = _get_queue_shuffle_enabled()
335
+ queue_seed = _get_queue_shuffle_seed(run_id)
336
+ ordered_paths = list(scan_output.image_paths)
337
+ if queue_shuffle_enabled:
338
+ rng = random.Random(queue_seed)
339
+ rng.shuffle(ordered_paths)
340
+ inference_mode = _get_inference_mode()
341
+ configured_batch_size = _get_inference_batch_size()
342
+ batch_size = 1 if inference_mode == "single" else configured_batch_size
343
+ _set_run_telemetry(
344
+ run_id,
345
+ queue_seed=queue_seed,
346
+ inference_mode=inference_mode,
347
+ batch_size=batch_size,
348
+ avg_infer_ms_per_image=0.0,
349
+ )
350
+ logger.info(
351
+ "run_queue_config run_id=%d shuffle=%s seed=%d mode=%s batch_size=%d",
352
+ run_id,
353
+ queue_shuffle_enabled,
354
+ queue_seed,
355
+ inference_mode,
356
+ batch_size,
357
+ )
358
 
359
  processed = 0
360
  failed = 0
361
+ cancelled = False
362
+ infer_elapsed_ms_total = 0.0
363
+ infer_sample_count = 0
364
+ max_workers = _get_max_inference_workers()
365
+ logger.info(
366
+ "run_inference_workers run_id=%d workers=%d total_images=%d",
367
+ run_id,
368
+ max_workers,
369
+ scan_output.stats.eligible_images,
370
+ )
 
 
 
 
 
 
 
 
 
 
 
371
 
372
+ with ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix="infer") as executor:
373
+ pending: dict[Future[tuple[list[_ImageInferenceResult], float, str]], list[Path]] = {}
374
+ batches = [
375
+ ordered_paths[idx : idx + batch_size]
376
+ for idx in range(0, len(ordered_paths), batch_size)
377
+ ]
378
+ iterator = iter(batches)
379
+
380
+ def _submit_until_capacity() -> None:
381
+ while len(pending) < max_workers:
382
+ try:
383
+ next_batch = next(iterator)
384
+ except StopIteration:
385
+ return
386
+ future = executor.submit(
387
+ _infer_batch_with_fallback,
388
+ next_batch,
389
+ matched_tags,
390
+ confidence_threshold,
391
+ inference_mode,
392
+ )
393
+ pending[future] = next_batch
394
+
395
+ _submit_until_capacity()
396
+ while pending:
397
+ if _is_cancel_requested(run_id):
398
+ cancelled = True
399
+ for future in pending:
400
+ future.cancel()
401
+ break
402
+
403
+ done, _ = wait(set(pending.keys()), return_when=FIRST_COMPLETED)
404
+ for future in done:
405
+ pending.pop(future)
406
+ if future.cancelled():
407
+ continue
408
+ batch_results, elapsed_ms, used_mode = future.result()
409
+ infer_elapsed_ms_total += elapsed_ms
410
+ infer_sample_count += len(batch_results)
411
+ avg_ms = (
412
+ infer_elapsed_ms_total / infer_sample_count if infer_sample_count else 0.0
413
+ )
414
+ _set_run_telemetry(
415
+ run_id,
416
+ inference_mode=used_mode if used_mode != "single_fallback" else "single",
417
+ batch_size=batch_size,
418
+ avg_infer_ms_per_image=avg_ms,
419
+ )
420
 
421
+ for result in batch_results:
422
+ relative_path = str(result.image_path.relative_to(root_repo))
423
+ suggested_destination = (
424
+ str(categories_root / result.primary_tag)
425
+ if result.primary_tag is not None
426
+ else None
427
+ )
428
+ execute(
429
+ """
430
+ INSERT INTO items (
431
+ run_id, file_path, relative_path, primary_tag, primary_score, secondary_json,
432
+ full_scores_json, suggested_destination, final_tag, final_destination, status, needs_review, review_reason
433
+ ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
434
+ """,
435
+ (
436
+ run_id,
437
+ str(result.image_path),
438
+ relative_path,
439
+ result.primary_tag,
440
+ result.primary_score,
441
+ to_json(result.secondary),
442
+ to_json(result.scores),
443
+ suggested_destination,
444
+ result.primary_tag,
445
+ suggested_destination,
446
+ "approved" if not result.needs_review else "proposed",
447
+ 1 if result.needs_review else 0,
448
+ result.reason,
449
+ ),
450
+ )
451
+ processed += 1
452
+ if result.inference_failed:
453
+ failed += 1
454
+ _update_run_progress(run_id, processed, failed)
455
+ if processed % 10 == 0:
456
+ logger.info(
457
+ "run_progress run_id=%d processed=%d total=%d avg_infer_ms=%.2f",
458
+ run_id,
459
+ processed,
460
+ scan_output.stats.eligible_images,
461
+ avg_ms,
462
+ )
463
+ _submit_until_capacity()
464
+
465
+ if cancelled:
466
+ # Reset partial run artifacts so cancelled runs do not look like
467
+ # "missing file" runs with incomplete queues.
468
+ execute("DELETE FROM items WHERE run_id = ?", (run_id,))
469
  execute(
470
  """
471
+ UPDATE runs
472
+ SET status = 'cancelled',
473
+ finished_at = ?,
474
+ total_images = 0,
475
+ processed_images = 0,
476
+ failed_images = 0
477
+ WHERE id = ?
478
  """,
479
+ (_now_iso(), run_id),
 
 
 
 
 
 
 
 
 
 
 
 
 
 
480
  )
481
+ logger.info("run_cancelled run_id=%d processed=%d", run_id, processed)
482
+ return
 
 
483
 
484
  execute(
485
  "UPDATE runs SET status = 'completed', finished_at = ? WHERE id = ?",
 
505
  execute(
506
  """
507
  UPDATE settings
508
+ SET root_repo = ?, categories_root = ?, confidence_threshold = ?,
509
+ default_migrate_mode = ?, scan_recursive = ?
510
  WHERE id = 1
511
  """,
512
  (
 
514
  payload.categories_root,
515
  payload.confidence_threshold,
516
  payload.default_migrate_mode,
517
+ 1 if payload.scan_recursive else 0,
518
  ),
519
  )
520
  except Exception:
 
531
  return {"items": known[:limit], "count": len(known)}
532
 
533
 
534
+ @router.get("/scan/preview")
535
+ def scan_preview() -> dict:
536
+ """Run image discovery on the configured root_repo and return stats without inference."""
537
+ settings = _settings_from_db()
538
+ if not settings.root_repo:
539
+ raise HTTPException(status_code=400, detail="root_repo is not configured in settings")
540
+ root_repo = Path(settings.root_repo).expanduser()
541
+ exclude_dirs: set[Path] = set()
542
+ if settings.categories_root:
543
+ exclude_dirs.add(Path(settings.categories_root).expanduser())
544
+ try:
545
+ output = scan_images(
546
+ root_repo,
547
+ exclude_dirs=exclude_dirs or None,
548
+ recursive=settings.scan_recursive,
549
+ )
550
+ except ValueError as err:
551
+ raise HTTPException(status_code=400, detail=str(err))
552
+ return {
553
+ "root_repo": str(root_repo),
554
+ "recursive": settings.scan_recursive,
555
+ "excluded_dirs": [str(d) for d in exclude_dirs],
556
+ "stats": output.stats.model_dump(),
557
+ "sample_paths": [str(p) for p in output.image_paths[:20]],
558
+ }
559
+
560
+
561
+ @router.get("/providers")
562
+ def get_providers() -> dict[str, object]:
563
+ force_cpu = os.getenv("FORCE_CPU_INFERENCE", "").strip().lower() in {"1", "true", "yes", "on"}
564
+ try:
565
+ import onnxruntime as ort
566
+
567
+ providers = list(ort.get_available_providers())
568
+ cuda_available = "CUDAExecutionProvider" in providers
569
+ return {
570
+ "available_providers": providers,
571
+ "cuda_available": cuda_available,
572
+ "cpu_available": "CPUExecutionProvider" in providers,
573
+ "forced_cpu": force_cpu,
574
+ "likely_device": "cpu" if force_cpu else ("gpu" if cuda_available else "cpu"),
575
+ }
576
+ except Exception as err:
577
+ return {
578
+ "available_providers": [],
579
+ "cuda_available": False,
580
+ "cpu_available": True,
581
+ "forced_cpu": force_cpu,
582
+ "likely_device": "cpu",
583
+ "error": str(err),
584
+ }
585
+
586
+
587
  @router.post("/runs/start", response_model=StartRunResponse)
588
  def start_run(payload: StartRunRequest) -> StartRunResponse:
589
  logger.info("run_start_requested")
 
625
 
626
  worker = threading.Thread(
627
  target=_execute_run,
628
+ args=(run_id, root_repo, categories_root, resolved.confidence_threshold, matched_tags,
629
+ resolved.scan_recursive),
630
  daemon=True,
631
  )
632
  worker.start()
backend/app/main.py CHANGED
@@ -1,6 +1,7 @@
1
  from __future__ import annotations
2
 
3
  import logging
 
4
  import time
5
 
6
  from fastapi import FastAPI
@@ -29,6 +30,31 @@ def _configure_logging() -> None:
29
  )
30
 
31
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
32
  @app.middleware("http")
33
  async def log_requests(request: Request, call_next):
34
  start = time.perf_counter()
@@ -57,6 +83,8 @@ async def log_requests(request: Request, call_next):
57
  @app.on_event("startup")
58
  def startup() -> None:
59
  _configure_logging()
 
 
60
  logger.info("initializing database at startup")
61
  init_db()
62
  logger.info("startup complete")
@@ -67,4 +95,9 @@ def health() -> dict[str, str]:
67
  return {"status": "ok"}
68
 
69
 
 
 
 
 
 
70
  app.include_router(router)
 
1
  from __future__ import annotations
2
 
3
  import logging
4
+ import os
5
  import time
6
 
7
  from fastapi import FastAPI
 
30
  )
31
 
32
 
33
+ def _get_provider_snapshot() -> dict[str, object]:
34
+ force_cpu = os.getenv("FORCE_CPU_INFERENCE", "").strip().lower() in {"1", "true", "yes", "on"}
35
+ try:
36
+ import onnxruntime as ort
37
+
38
+ providers = list(ort.get_available_providers())
39
+ cuda_available = "CUDAExecutionProvider" in providers
40
+ return {
41
+ "available_providers": providers,
42
+ "cuda_available": cuda_available,
43
+ "cpu_available": "CPUExecutionProvider" in providers,
44
+ "forced_cpu": force_cpu,
45
+ "likely_device": "cpu" if force_cpu else ("gpu" if cuda_available else "cpu"),
46
+ }
47
+ except Exception as err:
48
+ return {
49
+ "available_providers": [],
50
+ "cuda_available": False,
51
+ "cpu_available": True,
52
+ "forced_cpu": force_cpu,
53
+ "likely_device": "cpu",
54
+ "error": str(err),
55
+ }
56
+
57
+
58
  @app.middleware("http")
59
  async def log_requests(request: Request, call_next):
60
  start = time.perf_counter()
 
83
  @app.on_event("startup")
84
  def startup() -> None:
85
  _configure_logging()
86
+ provider_state = _get_provider_snapshot()
87
+ logger.info("onnx_provider_state state=%s", provider_state)
88
  logger.info("initializing database at startup")
89
  init_db()
90
  logger.info("startup complete")
 
95
  return {"status": "ok"}
96
 
97
 
98
+ @app.get("/health/providers")
99
+ def health_providers() -> dict[str, object]:
100
+ return {"status": "ok", **_get_provider_snapshot()}
101
+
102
+
103
  app.include_router(router)
backend/app/schemas.py CHANGED
@@ -15,6 +15,7 @@ class AppSettings(BaseModel):
15
  categories_root: str = ""
16
  confidence_threshold: float = 0.6
17
  default_migrate_mode: MigrateMode = "copy"
 
18
 
19
 
20
  class SaveSettingsRequest(AppSettings):
@@ -88,6 +89,10 @@ class RunStatusResponse(BaseModel):
88
  last_error: str | None = None
89
  cancel_requested: bool = False
90
  has_items: bool = False
 
 
 
 
91
 
92
 
93
  class UpdateItemRequest(BaseModel):
 
15
  categories_root: str = ""
16
  confidence_threshold: float = 0.6
17
  default_migrate_mode: MigrateMode = "copy"
18
+ scan_recursive: bool = True
19
 
20
 
21
  class SaveSettingsRequest(AppSettings):
 
89
  last_error: str | None = None
90
  cancel_requested: bool = False
91
  has_items: bool = False
92
+ inference_mode: str | None = None
93
+ batch_size: int | None = None
94
+ avg_infer_ms_per_image: float | None = None
95
+ queue_seed: int | None = None
96
 
97
 
98
  class UpdateItemRequest(BaseModel):
backend/app/services.py CHANGED
@@ -10,7 +10,7 @@ from PIL import Image
10
 
11
  from .schemas import AppSettings, FolderMapping, MigrationResult, ScanStats
12
 
13
- SUPPORTED_IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tiff"}
14
  VIDEO_EXTENSIONS = {
15
  ".mp4",
16
  ".mov",
@@ -22,6 +22,7 @@ VIDEO_EXTENSIONS = {
22
  ".m4v",
23
  }
24
  logger = logging.getLogger(__name__)
 
25
 
26
 
27
  def normalize_tag_name(value: str) -> str:
@@ -78,19 +79,50 @@ class ScanOutput:
78
  stats: ScanStats
79
 
80
 
81
- def scan_images(root_repo: Path) -> ScanOutput:
 
 
 
 
82
  if not root_repo.exists() or not root_repo.is_dir():
83
  raise ValueError(f"root_repo does not exist or is not a directory: {root_repo}")
84
 
 
 
 
 
 
 
 
 
 
85
  total_files = 0
86
  eligible: list[Path] = []
87
  ignored_unsupported = 0
88
  ignored_gif = 0
89
  failed_to_read = 0
90
 
91
- for path in root_repo.rglob("*"):
92
- if not path.is_file():
 
 
 
 
 
93
  continue
 
 
 
 
 
 
 
 
 
 
 
 
 
94
  total_files += 1
95
  ext = path.suffix.lower()
96
  if ext == ".gif":
@@ -99,17 +131,37 @@ def scan_images(root_repo: Path) -> ScanOutput:
99
  if ext in VIDEO_EXTENSIONS:
100
  ignored_unsupported += 1
101
  continue
102
- if ext not in SUPPORTED_IMAGE_EXTENSIONS:
103
- ignored_unsupported += 1
 
 
 
104
  continue
105
 
 
106
  try:
107
  with Image.open(path) as img:
108
- img.verify()
 
109
  eligible.append(path)
110
  except Exception:
111
- failed_to_read += 1
112
-
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
113
  return ScanOutput(
114
  image_paths=eligible,
115
  stats=ScanStats(
@@ -147,6 +199,57 @@ def extract_scores(image_path: Path) -> dict[str, float]:
147
  return scores
148
 
149
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
150
  def choose_best_tags(
151
  scores: dict[str, float], allowed_tags: set[str], max_secondary: int = 3
152
  ) -> tuple[str | None, float | None, list[dict[str, float]]]:
@@ -226,4 +329,5 @@ def resolve_settings(
226
  confidence_threshold if confidence_threshold is not None else current.confidence_threshold
227
  ),
228
  default_migrate_mode=current.default_migrate_mode,
 
229
  )
 
10
 
11
  from .schemas import AppSettings, FolderMapping, MigrationResult, ScanStats
12
 
13
+ SUPPORTED_IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".jfif", ".png", ".bmp", ".webp", ".tiff"}
14
  VIDEO_EXTENSIONS = {
15
  ".mp4",
16
  ".mov",
 
22
  ".m4v",
23
  }
24
  logger = logging.getLogger(__name__)
25
+ _BATCH_INFERENCE_SUPPORTED: bool | None = None
26
 
27
 
28
  def normalize_tag_name(value: str) -> str:
 
79
  stats: ScanStats
80
 
81
 
82
+ def scan_images(
83
+ root_repo: Path,
84
+ exclude_dirs: set[Path] | None = None,
85
+ recursive: bool = True,
86
+ ) -> ScanOutput:
87
  if not root_repo.exists() or not root_repo.is_dir():
88
  raise ValueError(f"root_repo does not exist or is not a directory: {root_repo}")
89
 
90
+ # Resolve exclude_dirs to absolute paths so prefix matching is reliable.
91
+ resolved_excludes: set[Path] = set()
92
+ if exclude_dirs:
93
+ for d in exclude_dirs:
94
+ try:
95
+ resolved_excludes.add(d.resolve())
96
+ except OSError:
97
+ pass
98
+
99
  total_files = 0
100
  eligible: list[Path] = []
101
  ignored_unsupported = 0
102
  ignored_gif = 0
103
  failed_to_read = 0
104
 
105
+ iterator = root_repo.rglob("*") if recursive else root_repo.iterdir()
106
+ for path in iterator:
107
+ try:
108
+ if not path.is_file():
109
+ continue
110
+ except OSError:
111
+ failed_to_read += 1
112
  continue
113
+
114
+ # Skip any path that lives inside an excluded directory.
115
+ if resolved_excludes:
116
+ try:
117
+ resolved_path = path.resolve()
118
+ if any(
119
+ resolved_path == exc or resolved_path.is_relative_to(exc)
120
+ for exc in resolved_excludes
121
+ ):
122
+ continue
123
+ except OSError:
124
+ pass
125
+
126
  total_files += 1
127
  ext = path.suffix.lower()
128
  if ext == ".gif":
 
131
  if ext in VIDEO_EXTENSIONS:
132
  ignored_unsupported += 1
133
  continue
134
+
135
+ if ext in SUPPORTED_IMAGE_EXTENSIONS:
136
+ # Known image extension: accept without Pillow verification.
137
+ # Corrupt or unreadable files are handled gracefully during inference.
138
+ eligible.append(path)
139
  continue
140
 
141
+ # Unknown or missing extension: try content-based detection.
142
  try:
143
  with Image.open(path) as img:
144
+ if not img.format:
145
+ raise ValueError("Not a recognizable image format")
146
  eligible.append(path)
147
  except Exception:
148
+ if ext == "":
149
+ failed_to_read += 1
150
+ else:
151
+ ignored_unsupported += 1
152
+
153
+ logger.info(
154
+ "scan_complete root=%s recursive=%s total_files=%d eligible=%d "
155
+ "ignored_gif=%d ignored_unsupported=%d failed_to_read=%d excluded_dirs=%d",
156
+ root_repo,
157
+ recursive,
158
+ total_files,
159
+ len(eligible),
160
+ ignored_gif,
161
+ ignored_unsupported,
162
+ failed_to_read,
163
+ len(resolved_excludes),
164
+ )
165
  return ScanOutput(
166
  image_paths=eligible,
167
  stats=ScanStats(
 
199
  return scores
200
 
201
 
202
+ def extract_scores_batch(image_paths: list[Path]) -> list[dict[str, float]]:
203
+ if not image_paths:
204
+ return []
205
+ global _BATCH_INFERENCE_SUPPORTED
206
+
207
+ if _BATCH_INFERENCE_SUPPORTED is False:
208
+ return [extract_scores(p) for p in image_paths]
209
+
210
+ from imgutils.tagging import get_mldanbooru_tags
211
+
212
+ try:
213
+ raw = get_mldanbooru_tags(
214
+ [str(p) for p in image_paths],
215
+ threshold=0.0,
216
+ size=448,
217
+ keep_ratio=True,
218
+ drop_overlap=False,
219
+ use_real_name=False,
220
+ )
221
+ except TypeError as err:
222
+ # Current imgutils build treats list input as invalid image type.
223
+ if "Unknown image type" in str(err):
224
+ if _BATCH_INFERENCE_SUPPORTED is not False:
225
+ logger.warning("batch inference not supported by imgutils; using per-image fallback")
226
+ _BATCH_INFERENCE_SUPPORTED = False
227
+ return [extract_scores(p) for p in image_paths]
228
+ raise
229
+
230
+ parsed: list[dict[str, float]] = []
231
+ if isinstance(raw, list) and len(raw) == len(image_paths):
232
+ _BATCH_INFERENCE_SUPPORTED = True
233
+ for entry in raw:
234
+ if isinstance(entry, dict):
235
+ parsed.append({str(k): float(v) for k, v in entry.items()})
236
+ continue
237
+ if isinstance(entry, list):
238
+ scores: dict[str, float] = {}
239
+ for item in entry:
240
+ if isinstance(item, (list, tuple)) and len(item) >= 2:
241
+ scores[str(item[0])] = float(item[1])
242
+ parsed.append(scores)
243
+ continue
244
+ raise TypeError("Unexpected batch inference entry format")
245
+ return parsed
246
+
247
+ # If the backend or library returns an unexpected shape, fall back to per-image inference.
248
+ logger.warning("batch inference unsupported format; falling back to per-image path")
249
+ _BATCH_INFERENCE_SUPPORTED = False
250
+ return [extract_scores(p) for p in image_paths]
251
+
252
+
253
  def choose_best_tags(
254
  scores: dict[str, float], allowed_tags: set[str], max_secondary: int = 3
255
  ) -> tuple[str | None, float | None, list[dict[str, float]]]:
 
329
  confidence_threshold if confidence_threshold is not None else current.confidence_threshold
330
  ),
331
  default_migrate_mode=current.default_migrate_mode,
332
+ scan_recursive=current.scan_recursive,
333
  )
backend/app/storage.py CHANGED
@@ -27,7 +27,8 @@ def init_db() -> None:
27
  root_repo TEXT NOT NULL DEFAULT '',
28
  categories_root TEXT NOT NULL DEFAULT '',
29
  confidence_threshold REAL NOT NULL DEFAULT 0.6,
30
- default_migrate_mode TEXT NOT NULL DEFAULT 'copy'
 
31
  );
32
 
33
  INSERT OR IGNORE INTO settings (id) VALUES (1);
@@ -67,6 +68,7 @@ def init_db() -> None:
67
  );
68
  """
69
  )
 
70
  _ensure_runs_columns(conn)
71
  _ensure_items_columns(conn)
72
  except sqlite3.DatabaseError:
@@ -74,6 +76,17 @@ def init_db() -> None:
74
  raise
75
 
76
 
 
 
 
 
 
 
 
 
 
 
 
77
  def _ensure_runs_columns(conn: sqlite3.Connection) -> None:
78
  expected_columns = {
79
  "status": "TEXT NOT NULL DEFAULT 'pending'",
 
27
  root_repo TEXT NOT NULL DEFAULT '',
28
  categories_root TEXT NOT NULL DEFAULT '',
29
  confidence_threshold REAL NOT NULL DEFAULT 0.6,
30
+ default_migrate_mode TEXT NOT NULL DEFAULT 'copy',
31
+ scan_recursive INTEGER NOT NULL DEFAULT 1
32
  );
33
 
34
  INSERT OR IGNORE INTO settings (id) VALUES (1);
 
68
  );
69
  """
70
  )
71
+ _ensure_settings_columns(conn)
72
  _ensure_runs_columns(conn)
73
  _ensure_items_columns(conn)
74
  except sqlite3.DatabaseError:
 
76
  raise
77
 
78
 
79
+ def _ensure_settings_columns(conn: sqlite3.Connection) -> None:
80
+ expected_columns = {
81
+ "scan_recursive": "INTEGER NOT NULL DEFAULT 1",
82
+ }
83
+ rows = conn.execute("PRAGMA table_info(settings)").fetchall()
84
+ existing = {row[1] for row in rows}
85
+ for name, definition in expected_columns.items():
86
+ if name not in existing:
87
+ conn.execute(f"ALTER TABLE settings ADD COLUMN {name} {definition}")
88
+
89
+
90
  def _ensure_runs_columns(conn: sqlite3.Connection) -> None:
91
  expected_columns = {
92
  "status": "TEXT NOT NULL DEFAULT 'pending'",
backend/tests/test_api_run_progress.py CHANGED
@@ -1,4 +1,6 @@
1
  import time
 
 
2
  from pathlib import Path
3
 
4
  from fastapi.testclient import TestClient
@@ -9,6 +11,17 @@ from app.services import ScanOutput
9
  from app.storage import execute
10
 
11
 
 
 
 
 
 
 
 
 
 
 
 
12
  def _wait_for_status(client: TestClient, run_id: int, terminal: set[str], timeout_s: float = 5.0):
13
  start = time.time()
14
  last = None
@@ -32,7 +45,7 @@ def test_run_progress_reaches_completed(monkeypatch, tmp_path: Path):
32
  file_path = root / "a.jpg"
33
  file_path.write_text("fake", encoding="utf-8")
34
 
35
- def fake_scan_images(_root):
36
  return ScanOutput(
37
  image_paths=[file_path],
38
  stats=ScanStats(
@@ -90,7 +103,7 @@ def test_item_preview_returns_image(monkeypatch, tmp_path: Path):
90
  file_path = root / "preview.jpg"
91
  file_path.write_bytes(b"fake-image-bytes")
92
 
93
- def fake_scan_images(_root):
94
  return ScanOutput(
95
  image_paths=[file_path],
96
  stats=ScanStats(
@@ -149,7 +162,7 @@ def test_selected_tag_wins_when_global_top_not_selected(monkeypatch, tmp_path: P
149
  file_path = root / "s.png"
150
  file_path.write_text("fake", encoding="utf-8")
151
 
152
- def fake_scan_images(_root):
153
  return ScanOutput(
154
  image_paths=[file_path],
155
  stats=ScanStats(
@@ -217,7 +230,7 @@ def test_item_scores_debug_endpoint(monkeypatch, tmp_path: Path):
217
  file_path = root / "z.jpg"
218
  file_path.write_text("fake", encoding="utf-8")
219
 
220
- def fake_scan_images(_root):
221
  return ScanOutput(
222
  image_paths=[file_path],
223
  stats=ScanStats(
@@ -280,7 +293,7 @@ def test_run_cancel_sets_cancelled(monkeypatch, tmp_path: Path):
280
  p.write_text("fake", encoding="utf-8")
281
  file_paths.append(p)
282
 
283
- def fake_scan_images(_root):
284
  return ScanOutput(
285
  image_paths=file_paths,
286
  stats=ScanStats(
@@ -326,7 +339,179 @@ def test_run_cancel_sets_cancelled(monkeypatch, tmp_path: Path):
326
  assert final is not None
327
  assert final["cancel_requested"] is True
328
  assert final["status"] in {"cancelled", "completed"}
 
 
 
 
 
 
329
 
330
  # Keep test environment clean of inserted rows.
331
  execute("DELETE FROM items WHERE run_id = ?", (run_id,))
332
  execute("DELETE FROM runs WHERE id = ?", (run_id,))
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  import time
2
+ import os
3
+ import random
4
  from pathlib import Path
5
 
6
  from fastapi.testclient import TestClient
 
11
  from app.storage import execute
12
 
13
 
14
+ def test_providers_endpoint_shape():
15
+ with TestClient(app) as client:
16
+ resp = client.get("/api/providers")
17
+ resp.raise_for_status()
18
+ payload = resp.json()
19
+ assert "available_providers" in payload
20
+ assert "likely_device" in payload
21
+ assert "cuda_available" in payload
22
+ assert "forced_cpu" in payload
23
+
24
+
25
  def _wait_for_status(client: TestClient, run_id: int, terminal: set[str], timeout_s: float = 5.0):
26
  start = time.time()
27
  last = None
 
45
  file_path = root / "a.jpg"
46
  file_path.write_text("fake", encoding="utf-8")
47
 
48
+ def fake_scan_images(_root, **kwargs):
49
  return ScanOutput(
50
  image_paths=[file_path],
51
  stats=ScanStats(
 
103
  file_path = root / "preview.jpg"
104
  file_path.write_bytes(b"fake-image-bytes")
105
 
106
+ def fake_scan_images(_root, **kwargs):
107
  return ScanOutput(
108
  image_paths=[file_path],
109
  stats=ScanStats(
 
162
  file_path = root / "s.png"
163
  file_path.write_text("fake", encoding="utf-8")
164
 
165
+ def fake_scan_images(_root, **kwargs):
166
  return ScanOutput(
167
  image_paths=[file_path],
168
  stats=ScanStats(
 
230
  file_path = root / "z.jpg"
231
  file_path.write_text("fake", encoding="utf-8")
232
 
233
+ def fake_scan_images(_root, **kwargs):
234
  return ScanOutput(
235
  image_paths=[file_path],
236
  stats=ScanStats(
 
293
  p.write_text("fake", encoding="utf-8")
294
  file_paths.append(p)
295
 
296
+ def fake_scan_images(_root, **kwargs):
297
  return ScanOutput(
298
  image_paths=file_paths,
299
  stats=ScanStats(
 
339
  assert final is not None
340
  assert final["cancel_requested"] is True
341
  assert final["status"] in {"cancelled", "completed"}
342
+ if final["status"] == "cancelled":
343
+ items_resp = client.get(f"/api/runs/{run_id}/items")
344
+ items_resp.raise_for_status()
345
+ assert items_resp.json() == []
346
+ assert final["processed_images"] == 0
347
+ assert final["total_images"] == 0
348
 
349
  # Keep test environment clean of inserted rows.
350
  execute("DELETE FROM items WHERE run_id = ?", (run_id,))
351
  execute("DELETE FROM runs WHERE id = ?", (run_id,))
352
+
353
+
354
+ def test_seeded_queue_shuffle_is_deterministic(monkeypatch, tmp_path: Path):
355
+ root = tmp_path / "root_seed"
356
+ cats = tmp_path / "cats_seed"
357
+ root.mkdir()
358
+ cats.mkdir()
359
+ (cats / "1girl").mkdir()
360
+ file_paths = []
361
+ for idx in range(6):
362
+ p = root / f"{idx}.jpg"
363
+ p.write_text("fake", encoding="utf-8")
364
+ file_paths.append(p)
365
+
366
+ observed_order: list[str] = []
367
+
368
+ def fake_scan_images(_root, **kwargs):
369
+ return ScanOutput(
370
+ image_paths=file_paths,
371
+ stats=ScanStats(
372
+ total_files=len(file_paths),
373
+ eligible_images=len(file_paths),
374
+ ignored_unsupported=0,
375
+ ignored_gif=0,
376
+ failed_to_read=0,
377
+ ),
378
+ )
379
+
380
+ def record_scores(path: Path):
381
+ observed_order.append(path.name)
382
+ return {"1girl": 0.88}
383
+
384
+ monkeypatch.setattr("app.api.scan_images", fake_scan_images)
385
+ monkeypatch.setattr("app.api.extract_scores", record_scores)
386
+ monkeypatch.setattr("app.api.load_known_tags", lambda _p: {"1girl"})
387
+ monkeypatch.setattr(
388
+ "app.api.discover_tag_folders",
389
+ lambda _root, _tags, _selected: [
390
+ FolderMapping(
391
+ folder_name="1girl", normalized_name="1girl", matched_tag="1girl", matched=True
392
+ )
393
+ ],
394
+ )
395
+
396
+ prev_shuffle = os.environ.get("QUEUE_SHUFFLE_ENABLED")
397
+ prev_seed = os.environ.get("QUEUE_SHUFFLE_SEED")
398
+ prev_workers = os.environ.get("MAX_INFERENCE_WORKERS")
399
+ prev_mode = os.environ.get("INFERENCE_MODE")
400
+ os.environ["QUEUE_SHUFFLE_ENABLED"] = "true"
401
+ os.environ["QUEUE_SHUFFLE_SEED"] = "1337"
402
+ os.environ["MAX_INFERENCE_WORKERS"] = "1"
403
+ os.environ["INFERENCE_MODE"] = "single"
404
+ try:
405
+ with TestClient(app) as client:
406
+ start_resp = client.post(
407
+ "/api/runs/start",
408
+ json={
409
+ "root_repo": str(root),
410
+ "categories_root": str(cats),
411
+ "confidence_threshold": 0.6,
412
+ "selected_folders": ["1girl"],
413
+ },
414
+ )
415
+ start_resp.raise_for_status()
416
+ run_id = start_resp.json()["run_id"]
417
+ final = _wait_for_status(client, run_id, {"completed", "failed", "cancelled"}, timeout_s=8.0)
418
+ assert final is not None
419
+ assert final["status"] == "completed"
420
+ expected = [p.name for p in file_paths]
421
+ random.Random(1337).shuffle(expected)
422
+ assert observed_order == expected
423
+ execute("DELETE FROM items WHERE run_id = ?", (run_id,))
424
+ execute("DELETE FROM runs WHERE id = ?", (run_id,))
425
+ finally:
426
+ if prev_shuffle is None:
427
+ os.environ.pop("QUEUE_SHUFFLE_ENABLED", None)
428
+ else:
429
+ os.environ["QUEUE_SHUFFLE_ENABLED"] = prev_shuffle
430
+ if prev_seed is None:
431
+ os.environ.pop("QUEUE_SHUFFLE_SEED", None)
432
+ else:
433
+ os.environ["QUEUE_SHUFFLE_SEED"] = prev_seed
434
+ if prev_workers is None:
435
+ os.environ.pop("MAX_INFERENCE_WORKERS", None)
436
+ else:
437
+ os.environ["MAX_INFERENCE_WORKERS"] = prev_workers
438
+ if prev_mode is None:
439
+ os.environ.pop("INFERENCE_MODE", None)
440
+ else:
441
+ os.environ["INFERENCE_MODE"] = prev_mode
442
+
443
+
444
+ def test_batch_mode_falls_back_to_single(monkeypatch, tmp_path: Path):
445
+ root = tmp_path / "root_batch_fallback"
446
+ cats = tmp_path / "cats_batch_fallback"
447
+ root.mkdir()
448
+ cats.mkdir()
449
+ (cats / "1girl").mkdir()
450
+ file_paths = []
451
+ for idx in range(8):
452
+ p = root / f"{idx}.jpg"
453
+ p.write_text("fake", encoding="utf-8")
454
+ file_paths.append(p)
455
+
456
+ def fake_scan_images(_root, **kwargs):
457
+ return ScanOutput(
458
+ image_paths=file_paths,
459
+ stats=ScanStats(
460
+ total_files=len(file_paths),
461
+ eligible_images=len(file_paths),
462
+ ignored_unsupported=0,
463
+ ignored_gif=0,
464
+ failed_to_read=0,
465
+ ),
466
+ )
467
+
468
+ def fail_batch(_paths):
469
+ raise RuntimeError("synthetic batch failure")
470
+
471
+ monkeypatch.setattr("app.api.scan_images", fake_scan_images)
472
+ monkeypatch.setattr("app.api.extract_scores_batch", fail_batch)
473
+ monkeypatch.setattr("app.api.extract_scores", lambda _p: {"1girl": 0.9})
474
+ monkeypatch.setattr("app.api.load_known_tags", lambda _p: {"1girl"})
475
+ monkeypatch.setattr(
476
+ "app.api.discover_tag_folders",
477
+ lambda _root, _tags, _selected: [
478
+ FolderMapping(
479
+ folder_name="1girl", normalized_name="1girl", matched_tag="1girl", matched=True
480
+ )
481
+ ],
482
+ )
483
+
484
+ prev_mode = os.environ.get("INFERENCE_MODE")
485
+ prev_batch = os.environ.get("INFERENCE_BATCH_SIZE")
486
+ os.environ["INFERENCE_MODE"] = "batch"
487
+ os.environ["INFERENCE_BATCH_SIZE"] = "4"
488
+ try:
489
+ with TestClient(app) as client:
490
+ start_resp = client.post(
491
+ "/api/runs/start",
492
+ json={
493
+ "root_repo": str(root),
494
+ "categories_root": str(cats),
495
+ "confidence_threshold": 0.6,
496
+ "selected_folders": ["1girl"],
497
+ },
498
+ )
499
+ start_resp.raise_for_status()
500
+ run_id = start_resp.json()["run_id"]
501
+ final = _wait_for_status(client, run_id, {"completed", "failed", "cancelled"}, timeout_s=8.0)
502
+ assert final is not None
503
+ assert final["status"] == "completed"
504
+ assert final["failed_images"] == 0
505
+ assert final["processed_images"] == len(file_paths)
506
+ assert final["inference_mode"] in {"single", "batch_fallback", "single_fallback"}
507
+ execute("DELETE FROM items WHERE run_id = ?", (run_id,))
508
+ execute("DELETE FROM runs WHERE id = ?", (run_id,))
509
+ finally:
510
+ if prev_mode is None:
511
+ os.environ.pop("INFERENCE_MODE", None)
512
+ else:
513
+ os.environ["INFERENCE_MODE"] = prev_mode
514
+ if prev_batch is None:
515
+ os.environ.pop("INFERENCE_BATCH_SIZE", None)
516
+ else:
517
+ os.environ["INFERENCE_BATCH_SIZE"] = prev_batch
backend/tests/test_services.py CHANGED
@@ -38,14 +38,80 @@ def test_scan_images_filters_types(tmp_path: Path) -> None:
38
  (tmp_path / "video.mp4").write_text("not-video", encoding="utf-8")
39
  (tmp_path / "anim.gif").write_text("gif", encoding="utf-8")
40
  (tmp_path / "doc.txt").write_text("txt", encoding="utf-8")
 
 
41
  (tmp_path / "broken.png").write_text("broken", encoding="utf-8")
42
 
43
  result = scan_images(tmp_path)
44
- assert len(result.image_paths) == 1
45
- assert result.stats.eligible_images == 1
 
 
46
  assert result.stats.ignored_gif == 1
47
- assert result.stats.ignored_unsupported >= 2
48
- assert result.stats.failed_to_read == 1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
49
 
50
 
51
  def test_choose_best_tags_single_primary_and_secondary() -> None:
 
38
  (tmp_path / "video.mp4").write_text("not-video", encoding="utf-8")
39
  (tmp_path / "anim.gif").write_text("gif", encoding="utf-8")
40
  (tmp_path / "doc.txt").write_text("txt", encoding="utf-8")
41
+ # broken.png has a known image extension so it is accepted at scan time;
42
+ # corrupt files are caught later during inference (not at discovery).
43
  (tmp_path / "broken.png").write_text("broken", encoding="utf-8")
44
 
45
  result = scan_images(tmp_path)
46
+ names = {p.name for p in result.image_paths}
47
+ assert "ok.jpg" in names
48
+ assert "broken.png" in names # accepted at scan; inference will handle corruption
49
+ assert result.stats.eligible_images == 2
50
  assert result.stats.ignored_gif == 1
51
+ assert result.stats.ignored_unsupported >= 1 # doc.txt
52
+ assert result.stats.failed_to_read == 0 # no OSError-level failures
53
+
54
+
55
+ def test_scan_images_includes_extensionless_valid_image(tmp_path: Path) -> None:
56
+ extless = tmp_path / "no_extension_image"
57
+ Image.new("RGB", (16, 16), color="purple").save(extless, format="PNG")
58
+
59
+ result = scan_images(tmp_path)
60
+ names = {p.name for p in result.image_paths}
61
+ assert "no_extension_image" in names
62
+ assert result.stats.eligible_images == 1
63
+
64
+
65
+ def test_scan_images_includes_unknown_extension_if_decodable(tmp_path: Path) -> None:
66
+ odd_ext = tmp_path / "odd_format.weird"
67
+ Image.new("RGB", (16, 16), color="yellow").save(odd_ext, format="PNG")
68
+
69
+ result = scan_images(tmp_path)
70
+ names = {p.name for p in result.image_paths}
71
+ assert "odd_format.weird" in names
72
+ assert result.stats.eligible_images == 1
73
+ assert result.stats.ignored_unsupported == 0
74
+
75
+
76
+ def test_scan_images_rejects_unknown_extension_when_not_image(tmp_path: Path) -> None:
77
+ non_image = tmp_path / "not_image.weird"
78
+ non_image.write_text("plain text", encoding="utf-8")
79
+
80
+ result = scan_images(tmp_path)
81
+ names = {p.name for p in result.image_paths}
82
+ assert "not_image.weird" not in names
83
+ assert result.stats.eligible_images == 0
84
+ assert result.stats.ignored_unsupported == 1
85
+
86
+
87
+ def test_scan_images_descends_into_subdirectories(tmp_path: Path) -> None:
88
+ nested = tmp_path / "nested"
89
+ nested.mkdir()
90
+ nested_img = nested / "nested.jpg"
91
+ Image.new("RGB", (16, 16), color="blue").save(nested_img)
92
+
93
+ top_img = tmp_path / "top.jpg"
94
+ Image.new("RGB", (16, 16), color="green").save(top_img)
95
+
96
+ result = scan_images(tmp_path)
97
+ names = {p.name for p in result.image_paths}
98
+ assert "top.jpg" in names
99
+ assert "nested.jpg" in names
100
+
101
+
102
+ def test_scan_images_excludes_specified_directory(tmp_path: Path) -> None:
103
+ included_dir = tmp_path / "included"
104
+ included_dir.mkdir()
105
+ Image.new("RGB", (16, 16), color="red").save(included_dir / "in.jpg")
106
+
107
+ excluded_dir = tmp_path / "excluded"
108
+ excluded_dir.mkdir()
109
+ Image.new("RGB", (16, 16), color="blue").save(excluded_dir / "out.jpg")
110
+
111
+ result = scan_images(tmp_path, exclude_dirs={excluded_dir})
112
+ names = {p.name for p in result.image_paths}
113
+ assert "in.jpg" in names
114
+ assert "out.jpg" not in names
115
 
116
 
117
  def test_choose_best_tags_single_primary_and_secondary() -> None:
frontend/src/App.jsx CHANGED
@@ -6,6 +6,7 @@ const DEFAULT_SETTINGS = {
6
  categories_root: "",
7
  confidence_threshold: 0.6,
8
  default_migrate_mode: "copy",
 
9
  };
10
  const ACTIVE_RUN_STORAGE_KEY = "imageClassifierActiveRunId";
11
 
@@ -18,6 +19,7 @@ function App() {
18
  const [items, setItems] = useState([]);
19
  const [loading, setLoading] = useState(false);
20
  const [error, setError] = useState("");
 
21
  const [migrateMode, setMigrateMode] = useState("copy");
22
  const [selectedIds, setSelectedIds] = useState([]);
23
  const [tagQuery, setTagQuery] = useState("");
@@ -65,6 +67,9 @@ function App() {
65
  if (storedRunId > 0) {
66
  setRunId(storedRunId);
67
  }
 
 
 
68
  }, []);
69
 
70
  useEffect(() => {
@@ -353,6 +358,13 @@ function App() {
353
  Backend is offline. Running in browser-only demo mode with local mock data.
354
  </div>
355
  )}
 
 
 
 
 
 
 
356
 
357
  {error && <div className="error">{error}</div>}
358
 
@@ -396,6 +408,14 @@ function App() {
396
  <option value="move">move</option>
397
  </select>
398
  </label>
 
 
 
 
 
 
 
 
399
  <button disabled={loading || opsLoading.saving}>Save Settings</button>
400
  </form>
401
  </section>
@@ -477,6 +497,17 @@ function App() {
477
  {Number(runStatus.progress_pct || 0).toFixed(1)}%)
478
  </span>
479
  <span>Failed: {runStatus.failed_images}</span>
 
 
 
 
 
 
 
 
 
 
 
480
  {runStatus.cancel_requested && <span>Cancel requested</span>}
481
  </div>
482
  )}
 
6
  categories_root: "",
7
  confidence_threshold: 0.6,
8
  default_migrate_mode: "copy",
9
+ scan_recursive: true,
10
  };
11
  const ACTIVE_RUN_STORAGE_KEY = "imageClassifierActiveRunId";
12
 
 
19
  const [items, setItems] = useState([]);
20
  const [loading, setLoading] = useState(false);
21
  const [error, setError] = useState("");
22
+ const [providerInfo, setProviderInfo] = useState(null);
23
  const [migrateMode, setMigrateMode] = useState("copy");
24
  const [selectedIds, setSelectedIds] = useState([]);
25
  const [tagQuery, setTagQuery] = useState("");
 
67
  if (storedRunId > 0) {
68
  setRunId(storedRunId);
69
  }
70
+ api.getProviders()
71
+ .then((info) => setProviderInfo(info))
72
+ .catch(() => setProviderInfo(null));
73
  }, []);
74
 
75
  useEffect(() => {
 
358
  Backend is offline. Running in browser-only demo mode with local mock data.
359
  </div>
360
  )}
361
+ {providerInfo && !offlineMode && (
362
+ <div className="stats">
363
+ <span>Inference device: {providerInfo.likely_device || "unknown"}</span>
364
+ <span>CUDA available: {String(Boolean(providerInfo.cuda_available))}</span>
365
+ <span>Forced CPU: {String(Boolean(providerInfo.forced_cpu))}</span>
366
+ </div>
367
+ )}
368
 
369
  {error && <div className="error">{error}</div>}
370
 
 
408
  <option value="move">move</option>
409
  </select>
410
  </label>
411
+ <label style={{ flexDirection: "row", alignItems: "center", gap: "0.5rem" }}>
412
+ <input
413
+ type="checkbox"
414
+ checked={Boolean(settings.scan_recursive)}
415
+ onChange={(e) => setSettings({ ...settings, scan_recursive: e.target.checked })}
416
+ />
417
+ Scan subfolders recursively
418
+ </label>
419
  <button disabled={loading || opsLoading.saving}>Save Settings</button>
420
  </form>
421
  </section>
 
497
  {Number(runStatus.progress_pct || 0).toFixed(1)}%)
498
  </span>
499
  <span>Failed: {runStatus.failed_images}</span>
500
+ {runStatus.inference_mode && <span>Inference mode: {runStatus.inference_mode}</span>}
501
+ {runStatus.batch_size ? <span>Batch size: {runStatus.batch_size}</span> : null}
502
+ {runStatus.avg_infer_ms_per_image !== null &&
503
+ runStatus.avg_infer_ms_per_image !== undefined ? (
504
+ <span>
505
+ Avg infer ms/image: {Number(runStatus.avg_infer_ms_per_image).toFixed(1)}
506
+ </span>
507
+ ) : null}
508
+ {runStatus.queue_seed !== null && runStatus.queue_seed !== undefined ? (
509
+ <span>Queue seed: {runStatus.queue_seed}</span>
510
+ ) : null}
511
  {runStatus.cancel_requested && <span>Cancel requested</span>}
512
  </div>
513
  )}
frontend/src/api.js CHANGED
@@ -6,6 +6,7 @@ const defaultSettings = {
6
  categories_root: "",
7
  confidence_threshold: 0.6,
8
  default_migrate_mode: "copy",
 
9
  };
10
  const mockTags = [
11
  "1girl",
@@ -167,6 +168,10 @@ function computeMockStatus(run) {
167
  last_error: run.last_error || null,
168
  cancel_requested: Boolean(run.cancel_requested),
169
  has_items: (run.items || []).length > 0,
 
 
 
 
170
  };
171
  }
172
 
@@ -177,6 +182,16 @@ function mockRequest(path, options = {}) {
177
  if (path === "/settings" && method === "GET") {
178
  return Promise.resolve(mockState.settings);
179
  }
 
 
 
 
 
 
 
 
 
 
180
  if (path.startsWith("/tags") && method === "GET") {
181
  const queryString = path.includes("?") ? path.split("?")[1] : "";
182
  const params = new URLSearchParams(queryString);
@@ -357,6 +372,7 @@ export const api = {
357
  return request(`/tags?${params.toString()}`);
358
  },
359
  getSettings: () => request("/settings"),
 
360
  saveSettings: (payload) =>
361
  request("/settings", {
362
  method: "PUT",
 
6
  categories_root: "",
7
  confidence_threshold: 0.6,
8
  default_migrate_mode: "copy",
9
+ scan_recursive: true,
10
  };
11
  const mockTags = [
12
  "1girl",
 
168
  last_error: run.last_error || null,
169
  cancel_requested: Boolean(run.cancel_requested),
170
  has_items: (run.items || []).length > 0,
171
+ inference_mode: "single",
172
+ batch_size: 1,
173
+ avg_infer_ms_per_image: null,
174
+ queue_seed: null,
175
  };
176
  }
177
 
 
182
  if (path === "/settings" && method === "GET") {
183
  return Promise.resolve(mockState.settings);
184
  }
185
+ if (path === "/providers" && method === "GET") {
186
+ return Promise.resolve({
187
+ available_providers: ["CPUExecutionProvider"],
188
+ cuda_available: false,
189
+ cpu_available: true,
190
+ forced_cpu: false,
191
+ likely_device: "cpu",
192
+ mock_mode: true,
193
+ });
194
+ }
195
  if (path.startsWith("/tags") && method === "GET") {
196
  const queryString = path.includes("?") ? path.split("?")[1] : "";
197
  const params = new URLSearchParams(queryString);
 
372
  return request(`/tags?${params.toString()}`);
373
  },
374
  getSettings: () => request("/settings"),
375
+ getProviders: () => request("/providers"),
376
  saveSettings: (payload) =>
377
  request("/settings", {
378
  method: "PUT",