Introduce diffusion-inpaint restoration engine in HVideoTool: added support for a new restoration method that regenerates masked regions via an external SwarmUI server, requiring YOLO detections for mask creation. Updated configuration management to include diffusion parameters, enhanced the UI for engine selection, and improved documentation in README and CLAUDE.md to guide users on the new functionality.

2026-06-08 06:21:44 +03:00
parent 8a366ed43d
commit 15f89b395d
14 changed files with 903 additions and 70 deletions
@@ -48,10 +48,12 @@ Keep this scope sharp:

 - Primary job is **detection + overlay/inspection**. A **restoration** ("расцензурить")
  step was added later (user-requested): on-demand (single frame **or** whole-project
-  batch into `restored/`), behind a `Restorer` interface. **Detection is YOLO-only and
-  restoration is DeepMosaics-only** — the noisy classic-CV detector (+ the `combined`
-  composite) and the cv2 `inpaint` baseline (filled but didn't reconstruct) were
-  **removed** as "works poorly". Detection is **multi-model** (ADetailer-style): drop YOLO
+  batch into `restored/`), behind a `Restorer` interface. **Detection is YOLO-only.**
+  Restoration has **two engine families**: **DeepMosaics** (default — *reconstructs*
+  mosaic, locates it itself) and **diffusion-inpaint** (SwarmUI — *regenerates* the masked
+  region; opt-in, see below). The noisy classic-CV detector (+ the `combined` composite)
+  and the cv2 `inpaint` baseline (filled but didn't reconstruct) were **removed** as
+  "works poorly". Detection is **multi-model** (ADetailer-style): drop YOLO
  weights under `models/yolo/<category>/`, tick which ones are active in the toolbar
  "Модели" menu, and a detect runs **all** ticked models and merges results (each tagged
  with its category → its own overlay colour). DeepMosaics has two engines: **image** (per-frame) and
@@ -59,10 +61,17 @@ Keep this scope sharp:
  **vendored** under `core/restore/_deepmosaics/` and run in-process (user supplies only
  the weights). Because of that vendoring the **whole project is GPL-3.0**. LADA
  (BasicVSR++) is a possible future engine, not wired.
- Still **no diffusion / ControlNet / SDXL**. (`xinsir/controlnet-union-sdxl-1.0` was
-  rejected early — a generative *conditioning* model, not a censorship restorer. Don't
-  reintroduce it.) Restoration, if upgraded, uses a mosaic-removal model (DeepMosaics/
-  LADA), not a general text-to-image diffusion pipeline.
+- **Diffusion-inpaint is now allowed** (the earlier ban was lifted by the user). It is a
+  *second* restoration engine (`restorer="diffusion"`), NOT a replacement for DeepMosaics
+  and NOT the default. It **regenerates** the censored region with an external diffusion
+  server (SwarmUI) over HTTP — it does *not* reconstruct the original, it draws plausible
+  new content from a prompt + the YOLO mask. So it's best where DeepMosaics is helpless
+  (black bars / solid fill), per-frame only (a video sequence flickers), and **needs
+  detections** (the mask). The diffusion model runs in SwarmUI's process, so this path
+  adds **no torch dependency** to the app. The backend is abstract (`DiffusionBackend`);
+  SwarmUI is the first impl — ComfyUI/A1111 could be added later as another backend.
+  (`xinsir/controlnet-union-sdxl-1.0` is still not used — it's a *conditioning* model, a
+  poor fit; "diffusion-inpaint" here means a standard SD/SDXL inpaint via SwarmUI.)
 - It detects **already-censored** regions, not "content that should be censored"
  (i.e. not an NSFW classifier).
 - Video is only a one-shot frame-extraction convenience (see below): "Создать из
@@ -121,11 +130,14 @@ hvideotool/
    ├── video/
    │   ├── extract.py   # extract_frames(): ffmpeg CLI (cv2 fallback) -> JPGs; keyframe/step modes + downscale
    │   └── frame.py     # Frame dataclass (image BGR, index, pts) — the detector input type
-    ├── restore/         # "un-censor" (DeepMosaics only; per-frame OR temporal)
-    │   ├── base.py      #   Restorer ABC: restore(image, dets, should_cancel) + restore_sequence (batch/temporal) + .temporal flag; Cancelled exc
-    │   ├── factory.py   #   build_restorer(name, config) -> deepmosaics | deepmosaics_video (lada = TODO)
+    ├── restore/         # "un-censor": DeepMosaics (reconstruct) OR diffusion-inpaint (regenerate)
+    │   ├── base.py      #   Restorer ABC: restore(image, dets, should_cancel) + restore_sequence (batch/temporal) + .temporal/.needs_detections flags; Cancelled exc
+    │   ├── factory.py   #   build_restorer(name, config) -> deepmosaics | deepmosaics_video | diffusion (lada = TODO); restorer_needs_detections(name)
    │   ├── deepmosaics.py #  DeepMosaicsRestorer (image, per-frame) + DeepMosaicsVideoRestorer (BVDNet, temporal); in-process, load once; uses _deepmosaics/
-    │   └── _deepmosaics/ #  VENDORED DeepMosaics models/+util/ (GPL-3.0) — added to sys.path at import
+    │   ├── _deepmosaics/ #  VENDORED DeepMosaics models/+util/ (GPL-3.0) — added to sys.path at import
+    │   ├── mask.py      #   detections_to_mask() — rasterise dets → uint8 mask (dilate/blur) for diffusion inpaint
+    │   ├── diffusion.py #   DiffusionRestorer (needs_detections=True, per-frame) + DiffusionBackend ABC + InpaintParams
+    │   └── swarmui.py   #   SwarmUIBackend — HTTP to a SwarmUI server (stdlib urllib, no torch dep); GetNewSession + GenerateText2Image
    └── detection/        # YOLO only, multi-model
        ├── base.py      # Detector ABC: detect(frame) -> list[Detection]
        ├── factory.py   # build_detector(config) -> MultiYoloDetector over config.detector_models
@@ -260,9 +272,11 @@ hvideotool/
  created lazily on first move), removing them from list/`_files`/cache. `_unique_dest`
  avoids clobbering (`foo.jpg` → `foo (1).jpg`). Use case: flag good frames into a
  training/example set while inspecting detections.
- **Restoration ("Расцензурить кадр" / "Расцензурить все").** DeepMosaics-only, and it
-  **locates the mosaic itself** — so restoration is **fully decoupled from detection**: no
-  detector runs in either path (detections are passed as `[]`). Single-frame: a toolbar
+- **Restoration ("Расцензурить кадр" / "Расцензурить все").** The **default DeepMosaics**
+  engine **locates the mosaic itself** — so for it restoration is **decoupled from
+  detection**: no detector runs, detections are passed as `[]`. (The **diffusion** engine
+  is the exception — it masks the detections, so the UI feeds it `_results`; see the
+  diffusion bullet above and `restorer.needs_detections`.) Single-frame: a toolbar
  action reads the current frame + runs `self._restorer` (built via `build_restorer`) **on
  a background job** (`_restore_current`; `done` stores `_restored[path]`, **auto-saves it
  to the project's `restored/`** (so a single restore persists like the batch, not just in
@@ -320,6 +334,26 @@ hvideotool/
  passed to it. To add another engine (e.g. LADA), implement `core/restore/base.Restorer`
  (set `.temporal` + override `restore_sequence` if it needs neighbours) and register it in
  `restore/factory.build_restorer`.
+  **Diffusion-inpaint engine (`restorer="diffusion"`, `core/restore/diffusion.py`).** A
+  second engine *family* that **regenerates** the censored region instead of reconstructing
+  it. `DiffusionRestorer` (`needs_detections=True`, per-frame): builds an inpaint mask from
+  the frame's YOLO detections (`mask.detections_to_mask`, with `diff_mask_dilate`/
+  `diff_mask_blur`) and hands `(image, mask, InpaintParams)` to a pluggable
+  `DiffusionBackend`. First backend is `SwarmUIBackend` (`swarmui.py`): stdlib-`urllib`
+  HTTP to a running SwarmUI server (`GetNewSession` → `GenerateText2Image` with base64
+  init+mask images, `diff_prompt`/`diff_negative`/`diff_steps`/`diff_cfg`/`diff_denoise`/
+  `diff_seed`/`diff_model`) → decode the returned image. The diffusion model runs in
+  SwarmUI's process, so **no torch dep is added here**. Because it needs a mask, the UI
+  feeds it real detections (`_restore_current` captures `_results[path]`; `_restore_all`
+  builds `dets_by_index` for the hit frames) and gates it like "Расцензурить найденное"
+  (requires detection computed + at least one hit) via `restorer_needs_detections`. Frames
+  with no detection come back unchanged. Per-frame only → flickers on video; best for
+  black bars / solid fill where DeepMosaics can't help. Config fields `diff_*` persist in
+  `project.json` + `settings.json`; engine chosen in `RestoreDialog` (its diffusion field
+  group shows when the engine is selected). The dialog has a **"Проверить соединение"**
+  button (`RestoreDialog._test_connection` → `SwarmUIBackend.ping()`, a fresh
+  `GetNewSession` with a short 15s timeout) that reports ✓/✗ inline — lets the user verify
+  SwarmUI is reachable without running a restore.
 - **Navigation bar** under the image (`_build_nav_bar`): prev/next frame (◀ ▶, keys
  `,`/`.`), a scrubber `frame_slider` across the whole sequence, a clickable `pos_label`
  (a flat `QPushButton` "row / n" → `_jump_to_frame`, a "go to frame N" `QInputDialog` —
@@ -367,12 +401,15 @@ hvideotool/
 - `ui/` must not import `torch` / `ultralytics` directly. It builds detectors only via
  `core/detection/factory.build_detector` and talks to `core/` through the `Detector`
  interface and the `Detection`/`CensorType` types.
- Detection is YOLO-only (multi-model) and restoration is DeepMosaics-only. New detection
-  kinds plug in by adding more `.pt` under `models/yolo/<category>/` — no code change. The
-  toolbar shows a "Модели" menu of checkable models (no detector dropdown); restoration
-  engines are chosen in `RestoreDialog`. If you re-add a different engine *kind*, implement
-  `core/detection/base.Detector` / `core/restore/base.Restorer` and register it in the
-  respective `factory`. (No legacy-settings migration is kept while in active development —
+- Detection is YOLO-only (multi-model); restoration is DeepMosaics (default) **or**
+  diffusion-inpaint (SwarmUI). New detection kinds plug in by adding more `.pt` under
+  `models/yolo/<category>/` — no code change. The toolbar shows a "Модели" menu of checkable
+  models (no detector dropdown); restoration engines are chosen in `RestoreDialog`. A new
+  restoration *engine kind* = implement `core/restore/base.Restorer` and register it in
+  `restore/factory.build_restorer`; a new *diffusion backend* = implement
+  `core/restore/diffusion.DiffusionBackend` (keep it out-of-process — no torch dep in the
+  app). If you re-add a different detector kind, implement `core/detection/base.Detector`
+  and register it in its `factory`. (No legacy-settings migration is kept while in active development —
  old `settings.json`/`project.json` keys are simply ignored, not coerced.)

 ## Commands
@@ -416,7 +453,9 @@ frame directly.
 - **classic-CV / inpaint were removed (worked poorly).** The classic-CV detector was
  noisy/approximate on real footage (false positives on skin/hair/fabric/JPEG; missed
  real mosaic after downscale) and the `combined` mode + cv2 `inpaint` baseline went with
-  it. Detection is YOLO-only, restoration is DeepMosaics-only. No backward-compat shims
+  it. Detection is YOLO-only; restoration is DeepMosaics (default) or diffusion-inpaint.
+  (The removed cv2 `inpaint` was a *classic* fill; the new diffusion-inpaint is a different
+  thing — a real generative SD/SDXL inpaint via SwarmUI.) No backward-compat shims
  while in active development — stale keys in old `settings.json`/`project.json` are just
  ignored (a project with no valid model selection default-ticks all discovered models).
 - **Domain matters.** LADA is trained on REAL video (JAV). It detects some anime mosaic
@@ -443,8 +482,17 @@ frame directly.
  otherwise it aliases a buffer that gets freed → garbage/crash.
 - **Always use `core/imageio.py`** (`imread_unicode`/`imwrite_unicode`) for images —
  `cv2.imread`/`imwrite` silently fail on non-ASCII Windows paths.
- Don't reintroduce any generative / ControlNet dependency, nor the removed video
-  *playback pipeline* (PyAV, producer/consumer worker threads, player, project session).
+- **Diffusion-inpaint runs out-of-process (SwarmUI), so `ui/` and the diffusion path add
+  no torch/diffusers dependency** — `swarmui.py` uses only stdlib `urllib`. Keep it that
+  way: the diffusion model lives in the SwarmUI server, we just POST image+mask+prompt.
+  Don't add `diffusers`/in-process SD to the app. The diffusion engine **needs detections**
+  (it masks them) — the UI feeds them via `dets_by_index` / captured `_results` and gates
+  it like "Расцензурить найденное" (`restorer_needs_detections` + `Restorer.needs_detections`);
+  DeepMosaics still gets `[]` (it self-locates). A new diffusion backend = another
+  `DiffusionBackend` impl, not new app deps.
+- Don't reintroduce the removed video *playback pipeline* (PyAV, producer/consumer worker
+  threads, player, project session). (Generative/diffusion inpaint via an external server
+  IS now allowed — see the diffusion engine; the old blanket "no generative" ban is lifted.)
  (The new `core/project.py` is an on-disk layout, not that thread-based "project
  model".) NOTE: a **single** background `Job` thread for detect/restore (`ui/workers.py`)
  IS in scope now (keeps the GUI responsive) — that's different from the rejected