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.

This commit is contained in:
Leonid Pershin
2026-06-08 06:21:44 +03:00
parent 8a366ed43d
commit 15f89b395d
14 changed files with 903 additions and 70 deletions
+72 -24
View File
@@ -48,10 +48,12 @@ Keep this scope sharp:
- Primary job is **detection + overlay/inspection**. A **restoration** ("расцензурить") - Primary job is **detection + overlay/inspection**. A **restoration** ("расцензурить")
step was added later (user-requested): on-demand (single frame **or** whole-project 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 batch into `restored/`), behind a `Restorer` interface. **Detection is YOLO-only.**
restoration is DeepMosaics-only** — the noisy classic-CV detector (+ the `combined` Restoration has **two engine families**: **DeepMosaics** (default — *reconstructs*
composite) and the cv2 `inpaint` baseline (filled but didn't reconstruct) were mosaic, locates it itself) and **diffusion-inpaint** (SwarmUI — *regenerates* the masked
**removed** as "works poorly". Detection is **multi-model** (ADetailer-style): drop YOLO 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 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 "Модели" 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 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 **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 the weights). Because of that vendoring the **whole project is GPL-3.0**. LADA
(BasicVSR++) is a possible future engine, not wired. (BasicVSR++) is a possible future engine, not wired.
- Still **no diffusion / ControlNet / SDXL**. (`xinsir/controlnet-union-sdxl-1.0` was - **Diffusion-inpaint is now allowed** (the earlier ban was lifted by the user). It is a
rejected early — a generative *conditioning* model, not a censorship restorer. Don't *second* restoration engine (`restorer="diffusion"`), NOT a replacement for DeepMosaics
reintroduce it.) Restoration, if upgraded, uses a mosaic-removal model (DeepMosaics/ and NOT the default. It **regenerates** the censored region with an external diffusion
LADA), not a general text-to-image diffusion pipeline. 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" - It detects **already-censored** regions, not "content that should be censored"
(i.e. not an NSFW classifier). (i.e. not an NSFW classifier).
- Video is only a one-shot frame-extraction convenience (see below): "Создать из - Video is only a one-shot frame-extraction convenience (see below): "Создать из
@@ -121,11 +130,14 @@ hvideotool/
├── video/ ├── video/
│ ├── extract.py # extract_frames(): ffmpeg CLI (cv2 fallback) -> JPGs; keyframe/step modes + downscale │ ├── 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 │ └── frame.py # Frame dataclass (image BGR, index, pts) — the detector input type
├── restore/ # "un-censor" (DeepMosaics only; per-frame OR temporal) ├── restore/ # "un-censor": DeepMosaics (reconstruct) OR diffusion-inpaint (regenerate)
│ ├── base.py # Restorer ABC: restore(image, dets, should_cancel) + restore_sequence (batch/temporal) + .temporal flag; Cancelled exc │ ├── 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 (lada = TODO) │ ├── 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.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 └── detection/ # YOLO only, multi-model
├── base.py # Detector ABC: detect(frame) -> list[Detection] ├── base.py # Detector ABC: detect(frame) -> list[Detection]
├── factory.py # build_detector(config) -> MultiYoloDetector over config.detector_models ├── 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` 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 avoids clobbering (`foo.jpg``foo (1).jpg`). Use case: flag good frames into a
training/example set while inspecting detections. training/example set while inspecting detections.
- **Restoration ("Расцензурить кадр" / "Расцензурить все").** DeepMosaics-only, and it - **Restoration ("Расцензурить кадр" / "Расцензурить все").** The **default DeepMosaics**
**locates the mosaic itself** — so restoration is **fully decoupled from detection**: no engine **locates the mosaic itself** — so for it restoration is **decoupled from
detector runs in either path (detections are passed as `[]`). Single-frame: a toolbar 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 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 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 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` 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 (set `.temporal` + override `restore_sequence` if it needs neighbours) and register it in
`restore/factory.build_restorer`. `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 - **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 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` (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 - `ui/` must not import `torch` / `ultralytics` directly. It builds detectors only via
`core/detection/factory.build_detector` and talks to `core/` through the `Detector` `core/detection/factory.build_detector` and talks to `core/` through the `Detector`
interface and the `Detection`/`CensorType` types. interface and the `Detection`/`CensorType` types.
- Detection is YOLO-only (multi-model) and restoration is DeepMosaics-only. New detection - Detection is YOLO-only (multi-model); restoration is DeepMosaics (default) **or**
kinds plug in by adding more `.pt` under `models/yolo/<category>/` — no code change. The diffusion-inpaint (SwarmUI). New detection kinds plug in by adding more `.pt` under
toolbar shows a "Модели" menu of checkable models (no detector dropdown); restoration `models/yolo/<category>/` — no code change. The toolbar shows a "Модели" menu of checkable
engines are chosen in `RestoreDialog`. If you re-add a different engine *kind*, implement models (no detector dropdown); restoration engines are chosen in `RestoreDialog`. A new
`core/detection/base.Detector` / `core/restore/base.Restorer` and register it in the restoration *engine kind* = implement `core/restore/base.Restorer` and register it in
respective `factory`. (No legacy-settings migration is kept while in active development `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.) old `settings.json`/`project.json` keys are simply ignored, not coerced.)
## Commands ## Commands
@@ -416,7 +453,9 @@ frame directly.
- **classic-CV / inpaint were removed (worked poorly).** The classic-CV detector was - **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 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 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 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). 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 - **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. otherwise it aliases a buffer that gets freed → garbage/crash.
- **Always use `core/imageio.py`** (`imread_unicode`/`imwrite_unicode`) for images — - **Always use `core/imageio.py`** (`imread_unicode`/`imwrite_unicode`) for images —
`cv2.imread`/`imwrite` silently fail on non-ASCII Windows paths. `cv2.imread`/`imwrite` silently fail on non-ASCII Windows paths.
- Don't reintroduce any generative / ControlNet dependency, nor the removed video - **Diffusion-inpaint runs out-of-process (SwarmUI), so `ui/` and the diffusion path add
*playback pipeline* (PyAV, producer/consumer worker threads, player, project session). 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 (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`) 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 IS in scope now (keeps the GUI responsive) — that's different from the rejected
+45 -9
View File
@@ -104,9 +104,10 @@
отдельно от `frames/`, чтобы результаты не попадали обратно в список кадров). Прогресс, отдельно от `frames/`, чтобы результаты не попадали обратно в список кадров). Прогресс,
отмена (**«■ Стоп»**) и предпросмотр текущего кадра работают как при детекции. отмена (**«■ Стоп»**) и предпросмотр текущего кадра работают как при детекции.
Расцензуривание — только через **DeepMosaics** (код **встроен** в приложение, vendored, Доступны два **семейства** движков: **DeepMosaics** (восстанавливает мозаику) и
GPL-3.0; ставить отдельно не нужно — требуются только **веса** и желательно **GPU **diffusion-inpaint (SwarmUI)** (перерисовывает область заново). DeepMosaics — по
NVIDIA/CUDA**). Два движка: умолчанию; код **встроен** (vendored, GPL-3.0; ставить отдельно не нужно — нужны только
**веса** и желательно **GPU NVIDIA/CUDA**). Три движка в списке:
- **DeepMosaics — картинка** — реальное генеративное удаление мозаики **покадрово**. - **DeepMosaics — картинка** — реальное генеративное удаление мозаики **покадрово**.
- **DeepMosaics — видео (BVDNet)** — **временно́й** движок: использует **соседние кадры** - **DeepMosaics — видео (BVDNet)** — **временно́й** движок: использует **соседние кадры**
@@ -114,8 +115,12 @@ NVIDIA/CUDA**). Два движка:
роликах. Из-за рекуррентности обрабатывает **непрерывный диапазон по порядку** — т.е. роликах. Из-за рекуррентности обрабатывает **непрерывный диапазон по порядку** — т.е.
запускайте его через **«Расцензурить все»** (одиночный «Расцензурить кадр» сведётся к запускайте его через **«Расцензурить все»** (одиночный «Расцензурить кадр» сведётся к
окну из одного кадра). Нужна **видеомодель** `clean_youknow_video.pth`. окну из одного кадра). Нужна **видеомодель** `clean_youknow_video.pth`.
- **Diffusion-inpaint (SwarmUI)** — **перерисовывает** область цензуры заново диффузионной
inpaint-моделью по **маске из YOLO-детекций** и промпту (не восстанавливает оригинал!).
Лучше всего для **чёрных плашек / сплошной заливки**, где DeepMosaics бессилен. Покадрово
→ на роликах будет **мерцание**. См. «Настройка SwarmUI» ниже.
На аниме качество ограничено (модели обучены на реальном видео). На аниме качество DeepMosaics ограничено (модели обучены на реальном видео).
### Настройка DeepMosaics ### Настройка DeepMosaics
@@ -144,6 +149,34 @@ Baidu код `1x0a`):
> `numpy 2.x` (проверено). [LADA](https://github.com/ladaapp/lada) (видеомодель, > `numpy 2.x` (проверено). [LADA](https://github.com/ladaapp/lada) (видеомодель,
> лучшее качество на реальном видео) пока не подключён. > лучшее качество на реальном видео) пока не подключён.
### Настройка Diffusion-inpaint (SwarmUI)
Этот движок перерисовывает область **по маске из детекций YOLO**, поэтому **сначала
посчитайте детекцию** («Детектировать все» или «Рассчитать кадр»), а затем запускайте
расцензуривание — кадры без детекций остаются без изменений. Диффузионная модель крутится
в **отдельном сервере SwarmUI**, приложение лишь шлёт ему по HTTP картинку + маску + промпт
(никаких `torch`/`diffusers` в самом приложении на этом пути).
1. Установите и запустите [SwarmUI](https://github.com/mcmonkeyprojects/SwarmUI), загрузите
в нём inpaint-чекпойнт (SD/SDXL). По умолчанию сервер слушает `http://localhost:7801`.
2. В приложении: **Файл → Движок восстановления…** → выберите **Diffusion-inpaint
(SwarmUI)** и задайте:
- **SwarmUI URL** (по умолчанию `http://localhost:7801`);
- **Чекпойнт** — имя модели как её знает SwarmUI (пусто = текущая в сервере);
- **Промпт / Negative** — что нарисовать в области под цензурой / чего избегать;
- **Шаги / CFG / Denoise / Seed** — параметры генерации (`Denoise` 0..1, 1 = полностью
перерисовать; `Seed` `-1` = случайный);
- **Маска: расширить / размытие** (px) — расширение и мягкость края маски.
- Кнопка **«Проверить соединение»** дёргает SwarmUI и сразу показывает ✓ (сервер
отвечает) или ✗ с текстом ошибки — удобно убедиться в адресе до расцензуривания.
3. Запустите **«Расцензурить кадр»** или **«Расцензурить все»** — движок прогонит только
кадры с детекциями.
> Diffusion **выдумывает** правдоподобное содержимое, а не восстанавливает оригинал. Это
> сознательно второй движок (не замена DeepMosaics) под случаи, где под цензурой не
> осталось данных (чёрные плашки). Бэкенд абстрактный — позже можно добавить ComfyUI/A1111
> как ещё одну реализацию `DiffusionBackend`.
## Что НЕ делает (осознанно вне области задачи) ## Что НЕ делает (осознанно вне области задачи)
- Не **генерирует** изображения через диффузию (никакого ControlNet/SDXL). - Не **генерирует** изображения через диффузию (никакого ControlNet/SDXL).
@@ -270,11 +303,14 @@ hvideotool/
│ ├── types.py # Detection (+ label/категория, .display), CensorType │ ├── types.py # Detection (+ label/категория, .display), CensorType
│ ├── cache.py # кэш детекций (ключ = набор моделей + conf/imgsz) │ ├── cache.py # кэш детекций (ключ = набор моделей + conf/imgsz)
│ └── yolo.py # YOLO-детектор (Ultralytics, маски→полигоны, ярлык категории) │ └── yolo.py # YOLO-детектор (Ultralytics, маски→полигоны, ярлык категории)
└── restore/ # только DeepMosaics └── restore/ # DeepMosaics (восстановление) или diffusion-inpaint (перерисовка)
├── base.py # Restorer (ABC): restore() + restore_sequence() + .temporal ├── base.py # Restorer (ABC): restore() + restore_sequence() + .temporal/.needs_detections
├── factory.py # build_restorer -> deepmosaics | deepmosaics_video ├── factory.py # build_restorer -> deepmosaics | deepmosaics_video | diffusion
├── deepmosaics.py # движки «картинка» (покадрово) и «видео» (BVDNet) ├── deepmosaics.py # движки «картинка» (покадрово) и «видео» (BVDNet)
── _deepmosaics/ # встроенный код DeepMosaics (GPL-3.0) ── _deepmosaics/ # встроенный код DeepMosaics (GPL-3.0)
├── mask.py # маска из детекций (для diffusion inpaint)
├── diffusion.py # DiffusionRestorer + DiffusionBackend (ABC) + InpaintParams
└── swarmui.py # бэкенд SwarmUI (HTTP, stdlib urllib — без torch)
``` ```
Детекция синхронная (по клику/по кнопке «Детектировать все»); тяжёлый YOLO на CPU Детекция синхронная (по клику/по кнопке «Детектировать все»); тяжёлый YOLO на CPU
@@ -289,7 +325,7 @@ hvideotool/
| GUI | PySide6 (Qt 6) | | GUI | PySide6 (Qt 6) |
| Обработка картинок | OpenCV / NumPy | | Обработка картинок | OpenCV / NumPy |
| Детектор | Ultralytics YOLO, мульти-модель (models/yolo/<кат>) | | Детектор | Ultralytics YOLO, мульти-модель (models/yolo/<кат>) |
| Расцензуривание | DeepMosaics (встроен) + PyTorch/CUDA | | Расцензуривание | DeepMosaics (встроен) + PyTorch/CUDA, или diffusion-inpaint через SwarmUI (HTTP) |
## Лицензия ## Лицензия
+16
View File
@@ -75,3 +75,19 @@ class AppConfig:
# (instead of the original mosaic frames) for stronger temporal coherence. Slightly # (instead of the original mosaic frames) for stronger temporal coherence. Slightly
# out-of-distribution for the net (trained on mosaic windows) — toggle in the dialog. # out-of-distribution for the net (trained on mosaic windows) — toggle in the dialog.
dm_feed_restored: bool = True dm_feed_restored: bool = True
# --- diffusion-inpaint restoration ("diffusion" restorer) ---
# Regenerates the masked (detected) regions via an external diffusion server. Needs
# YOLO detections for the mask; the model runs out-of-process (no torch dep here).
# The backend is pluggable; only SwarmUI is wired so far.
diff_backend: str = "swarmui" # only "swarmui" implemented
diff_url: str = "http://localhost:7801" # SwarmUI server base URL
diff_model: str | None = None # checkpoint name as the server knows it
diff_prompt: str = ""
diff_negative: str = ""
diff_steps: int = 30
diff_cfg: float = 7.0
diff_denoise: float = 1.0 # 0..1, 1 = fully regenerate under the mask
diff_seed: int = -1 # -1 = random each call
diff_mask_dilate: int = 4 # px to grow the mask before inpaint
diff_mask_blur: int = 8 # px feather of the mask edge
+11
View File
@@ -47,6 +47,17 @@ _SETTING_KEYS = (
"dm_model", "dm_model",
"dm_gpu", "dm_gpu",
"dm_feed_restored", "dm_feed_restored",
"diff_backend",
"diff_url",
"diff_model",
"diff_prompt",
"diff_negative",
"diff_steps",
"diff_cfg",
"diff_denoise",
"diff_seed",
"diff_mask_dilate",
"diff_mask_blur",
) )
+6
View File
@@ -37,6 +37,12 @@ class Restorer(ABC):
#: leave this False; the temporal DeepMosaics (BVDNet) sets it True. #: leave this False; the temporal DeepMosaics (BVDNet) sets it True.
temporal: bool = False temporal: bool = False
#: Whether this engine needs the frame's detections (it builds an inpaint mask from
#: them). DeepMosaics locates the mosaic itself, so it leaves this False and the
#: caller passes ``[]``; the diffusion engine sets it True and must be fed the real
#: detections (a frame with none comes back unchanged).
needs_detections: bool = False
@property @property
def name(self) -> str: def name(self) -> str:
return type(self).__name__ return type(self).__name__
+103
View File
@@ -0,0 +1,103 @@
"""Diffusion-inpaint restoration — redraw censored regions with a diffusion backend.
Unlike DeepMosaics (which *reconstructs* mosaic from its residual low-frequency data and
locates it itself), this engine *regenerates* the masked region with a diffusion inpaint
model: it builds a mask from the YOLO detections and hands ``(image, mask, params)`` to a
pluggable :class:`DiffusionBackend` (SwarmUI is the first, see ``swarmui.py``).
Consequences of that design:
- It **needs detections** (``needs_detections = True``) — a frame with none comes back
unchanged (no mask → nothing to regenerate). The caller feeds it the real detections.
- It's **per-frame** (``temporal = False``): each frame is generated independently, so a
video sequence will flicker. Best for stills / single frames, not coherent clips.
- The backend runs in a **separate process/server** (e.g. SwarmUI over HTTP), so this
path adds **no torch dependency** to the app and keeps the heavy model out-of-process.
The backend is abstract so other diffusion servers (ComfyUI/A1111) can be added later as
another :class:`DiffusionBackend`, without touching the restorer or the UI.
"""
from __future__ import annotations
from abc import ABC, abstractmethod
from dataclasses import dataclass
import numpy as np
from ..detection.types import Detection
from .base import CancelCheck, Cancelled, Restorer
from .mask import detections_to_mask, mask_is_empty
@dataclass(frozen=True)
class InpaintParams:
"""Generation knobs handed to a :class:`DiffusionBackend`."""
prompt: str = ""
negative: str = ""
model: str | None = None # checkpoint name as the backend knows it (None = current)
steps: int = 30
cfg: float = 7.0
denoise: float = 1.0 # 0..1 — how much to regenerate under the mask (1 = full)
seed: int = -1 # -1 = random each call
mask_blur: int = 8 # px feather applied by the backend at its mask edge
class DiffusionBackend(ABC):
"""A diffusion inpaint engine reachable from our process (typically over HTTP)."""
@property
def name(self) -> str:
return type(self).__name__
@abstractmethod
def inpaint(
self,
image_bgr: np.ndarray,
mask: np.ndarray,
params: InpaintParams,
should_cancel: CancelCheck | None = None,
) -> np.ndarray:
"""Regenerate the white area of ``mask`` in ``image_bgr``; return a new BGR image."""
raise NotImplementedError
class DiffusionRestorer(Restorer):
"""Restorer that masks the detected regions and inpaints them via a backend."""
temporal = False
needs_detections = True
def __init__(
self,
backend: DiffusionBackend,
params: InpaintParams,
*,
mask_dilate: int = 4,
mask_blur: int = 8,
) -> None:
self._backend = backend
self._params = params
self._dilate = mask_dilate
self._blur = mask_blur
@property
def name(self) -> str:
return f"Diffusion({self._backend.name})"
def restore(
self,
image: np.ndarray,
detections: list[Detection],
should_cancel: CancelCheck | None = None,
) -> np.ndarray:
if should_cancel is not None and should_cancel():
raise Cancelled("Восстановление отменено")
if not detections:
return image.copy() # no detections → no mask → nothing to regenerate
mask = detections_to_mask(
detections, image.shape, dilate=self._dilate, blur=self._blur
)
if mask_is_empty(mask):
return image.copy()
return self._backend.inpaint(image, mask, self._params, should_cancel)
+44 -7
View File
@@ -1,13 +1,16 @@
"""Restorer factory: build a Restorer from the app config. """Restorer factory: build a Restorer from the app config.
Only DeepMosaics is supported (the cv2 inpaint baseline was removed — it filled but Engines:
did not reconstruct). The DeepMosaics network code is vendored (``_deepmosaics/``, - ``deepmosaics``: per-frame generative mosaic removal (image model). Vendored network
GPL-3.0) and run in-process; the user supplies only the weights (+ ``mosaic_position.pth`` code (``_deepmosaics/``, GPL-3.0) run in-process; user supplies only the weights
alongside). A CUDA GPU is recommended. (+ ``mosaic_position.pth`` alongside). Locates the mosaic itself — needs no detections.
- ``deepmosaics_video``: temporal variant (BVDNet) using neighbouring frames for coherence
— needs the ``clean_*_video.pth`` weights and a contiguous frame sequence.
- ``diffusion``: diffusion-inpaint that *regenerates* masked regions via an external
diffusion server (SwarmUI). Builds the mask from the YOLO detections, so it **needs
detections** (see ``restorer_needs_detections``); adds no torch dep (runs out-of-process).
- ``deepmosaics``: per-frame generative mosaic removal (image model). A CUDA GPU is recommended (for DeepMosaics in-process; for diffusion it's the server's GPU).
- ``deepmosaics_video``: temporal variant (BVDNet) that uses neighbouring frames for
coherence — needs the ``clean_*_video.pth`` weights and a contiguous frame sequence.
""" """
from __future__ import annotations from __future__ import annotations
@@ -36,8 +39,42 @@ def build_restorer(name: str = "deepmosaics", config: AppConfig | None = None) -
config.dm_dir, config.dm_model, config.dm_gpu, config.dm_dir, config.dm_model, config.dm_gpu,
feed_restored=getattr(config, "dm_feed_restored", True), feed_restored=getattr(config, "dm_feed_restored", True),
) )
if name == "diffusion":
backend_name = (getattr(config, "diff_backend", "swarmui") or "swarmui")
if backend_name != "swarmui":
raise ValueError(
f"Diffusion-бэкенд не поддержан: {backend_name!r} (доступен только swarmui)."
)
from .diffusion import DiffusionRestorer, InpaintParams
from .swarmui import SwarmUIBackend
backend = SwarmUIBackend(config.diff_url)
params = InpaintParams(
prompt=config.diff_prompt,
negative=config.diff_negative,
model=config.diff_model,
steps=config.diff_steps,
cfg=config.diff_cfg,
denoise=config.diff_denoise,
seed=config.diff_seed,
mask_blur=config.diff_mask_blur,
)
return DiffusionRestorer(
backend, params,
mask_dilate=config.diff_mask_dilate, mask_blur=config.diff_mask_blur,
)
if name == "lada": if name == "lada":
raise ValueError( raise ValueError(
"Движок LADA пока не подключён. Используйте DeepMosaics. См. README." "Движок LADA пока не подключён. Используйте DeepMosaics. См. README."
) )
raise ValueError(f"Неизвестный режим восстановления: {name!r}") raise ValueError(f"Неизвестный режим восстановления: {name!r}")
def restorer_needs_detections(name: str) -> bool:
"""Whether engine ``name`` needs the frame's detections (to build an inpaint mask).
Lets the UI decide — *without* building the engine — whether to feed real detections
and whether to require that detection has been computed. Mirrors
``Restorer.needs_detections`` for the engines that build lazily on a worker thread.
"""
return name == "diffusion"
+57
View File
@@ -0,0 +1,57 @@
"""Build an inpaint mask (255 = regenerate) from detections.
Used by the diffusion restorer. Unlike DeepMosaics — which locates the mosaic itself —
a diffusion-inpaint backend needs an explicit mask of the region to redraw. We rasterise
each detection's polygon (or its bbox when there's no polygon) onto a single-channel
uint8 mask, optionally growing (dilate) and feathering (blur) the edges so the inpaint
blends into the surrounding pixels.
Pure NumPy/OpenCV — no torch, no Qt.
"""
from __future__ import annotations
from collections.abc import Sequence
import cv2
import numpy as np
from ..detection.types import Detection
def detections_to_mask(
detections: Sequence[Detection],
shape: tuple[int, ...],
*,
dilate: int = 0,
blur: int = 0,
) -> np.ndarray:
"""Rasterise ``detections`` onto a single-channel uint8 mask (255 = regenerate).
``shape`` is the image shape (``(h, w)`` or ``(h, w, c)``). ``dilate`` grows the mask
by that many pixels (ellipse kernel) so the inpaint covers the censored edge; ``blur``
feathers the edge with a Gaussian so the boundary blends. Both are no-ops at 0.
"""
h, w = int(shape[0]), int(shape[1])
mask = np.zeros((h, w), dtype=np.uint8)
for d in detections:
if len(d.polygon) >= 3:
poly = np.array(
[[int(round(x)), int(round(y))] for x, y in d.polygon], dtype=np.int32
)
cv2.fillPoly(mask, [poly], 255)
else:
x, y, bw, bh = (int(round(v)) for v in d.bbox)
cv2.rectangle(mask, (x, y), (x + bw, y + bh), 255, thickness=-1)
if dilate > 0:
k = 2 * int(dilate) + 1
mask = cv2.dilate(mask, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k)))
if blur > 0:
k = 2 * int(blur) + 1
mask = cv2.GaussianBlur(mask, (k, k), 0)
return mask
def mask_is_empty(mask: np.ndarray) -> bool:
"""True if nothing is masked (so there's nothing to inpaint)."""
return not bool(np.any(mask))
+133
View File
@@ -0,0 +1,133 @@
"""SwarmUI diffusion backend — talk to a running SwarmUI server over HTTP.
SwarmUI (a REST wrapper over ComfyUI) exposes ``/API/GetNewSession`` to obtain a session
id, then ``/API/GenerateText2Image`` to run a generation. For inpaint we send the frame
and the mask as base64 PNG plus the prompt/params, and read the produced image back.
Implementation notes:
- Uses only stdlib ``urllib`` — **no new dependency**; the diffusion model runs in
SwarmUI's own process (so our app never imports torch on this path).
- Exact API field names drift between SwarmUI versions, so the request body is built in
one place (:meth:`_build_payload`) for easy tuning; errors surface the URL + a hint.
- The response may carry image data inline (``data:`` URI) or as a server-relative path
— :meth:`_fetch_image_bytes` handles both.
"""
from __future__ import annotations
import base64
import json
import urllib.error
import urllib.request
import cv2
import numpy as np
from .base import Cancelled
from .diffusion import DiffusionBackend, InpaintParams
class SwarmUIBackend(DiffusionBackend):
def __init__(self, url: str | None, timeout: float = 600.0) -> None:
self._url = (url or "http://localhost:7801").rstrip("/")
self._timeout = timeout
self._session: str | None = None
@property
def name(self) -> str:
return "SwarmUI"
# ------------------------------------------------------------------ HTTP
def _post(self, route: str, payload: dict) -> dict:
data = json.dumps(payload).encode("utf-8")
req = urllib.request.Request(
self._url + route, data=data, headers={"Content-Type": "application/json"}
)
try:
with urllib.request.urlopen(req, timeout=self._timeout) as resp:
return json.loads(resp.read().decode("utf-8"))
except urllib.error.URLError as e:
raise RuntimeError(
f"Не удалось связаться со SwarmUI ({self._url}{route}): {e}.\n"
"Проверьте, что сервер SwarmUI запущен и адрес верный "
"(Файл → Движок восстановления…)."
) from e
def ping(self) -> str:
"""Open a fresh session to verify the server is reachable; return the session id.
Used by the settings dialog's "Проверить соединение" — forces a new
``GetNewSession`` (ignores any cached id) so repeated checks really re-test, and
raises a clear RuntimeError (URL + hint) if the server is down/unreachable.
"""
self._session = None
return self._session_id()
def _session_id(self) -> str:
if self._session is None:
r = self._post("/API/GetNewSession", {})
self._session = r.get("session_id") or r.get("sessionId")
if not self._session:
raise RuntimeError(f"SwarmUI не вернул session_id: {r}")
return self._session
@staticmethod
def _b64_png(img: np.ndarray) -> str:
ok, buf = cv2.imencode(".png", img)
if not ok:
raise RuntimeError("Не удалось закодировать изображение в PNG для SwarmUI")
return base64.b64encode(buf.tobytes()).decode("ascii")
def _build_payload(
self, session: str, image_b64: str, mask_b64: str, params: InpaintParams, h: int, w: int
) -> dict:
"""Map our params onto SwarmUI's GenerateText2Image body (centralised for tuning)."""
payload = {
"session_id": session,
"images": 1,
"prompt": params.prompt,
"negativeprompt": params.negative,
"width": w,
"height": h,
"steps": int(params.steps),
"cfgscale": float(params.cfg),
"seed": int(params.seed),
"initimage": image_b64,
"maskimage": mask_b64, # white = regenerate
"initimagecreativity": float(params.denoise), # 0..1 inpaint denoise
"maskblur": int(params.mask_blur),
}
if params.model:
payload["model"] = params.model
return payload
# --------------------------------------------------------------- backend
def inpaint(self, image_bgr, mask, params, should_cancel=None):
if should_cancel is not None and should_cancel():
raise Cancelled("Восстановление отменено")
session = self._session_id()
h, w = image_bgr.shape[:2]
payload = self._build_payload(
session, self._b64_png(image_bgr), self._b64_png(mask), params, h, w
)
if should_cancel is not None and should_cancel():
raise Cancelled("Восстановление отменено")
resp = self._post("/API/GenerateText2Image", payload)
return self._decode_result(resp)
def _decode_result(self, resp: dict) -> np.ndarray:
images = resp.get("images") or []
if not images:
raise RuntimeError(f"SwarmUI не вернул изображений (ответ: {resp})")
raw = self._fetch_image_bytes(images[0])
arr = cv2.imdecode(np.frombuffer(raw, np.uint8), cv2.IMREAD_COLOR)
if arr is None:
raise RuntimeError("Не удалось декодировать результат SwarmUI")
return arr
def _fetch_image_bytes(self, ref: str) -> bytes:
if ref.startswith("data:"): # inline base64 data URI
return base64.b64decode(ref.split(",", 1)[1])
url = ref if ref.startswith("http") else f"{self._url}/{ref.lstrip('/')}"
with urllib.request.urlopen(url, timeout=self._timeout) as r:
return r.read()
+27
View File
@@ -52,6 +52,22 @@ def apply(config: AppConfig) -> None:
setattr(config, key, data[key]) setattr(config, key, data[key])
if "dm_feed_restored" in data: if "dm_feed_restored" in data:
config.dm_feed_restored = bool(data["dm_feed_restored"]) config.dm_feed_restored = bool(data["dm_feed_restored"])
# diffusion-inpaint engine
for key in ("diff_backend", "diff_url", "diff_model", "diff_prompt", "diff_negative"):
if key in data:
setattr(config, key, data[key])
if "diff_steps" in data:
config.diff_steps = int(data["diff_steps"])
if "diff_cfg" in data:
config.diff_cfg = float(data["diff_cfg"])
if "diff_denoise" in data:
config.diff_denoise = float(data["diff_denoise"])
if "diff_seed" in data:
config.diff_seed = int(data["diff_seed"])
if "diff_mask_dilate" in data:
config.diff_mask_dilate = int(data["diff_mask_dilate"])
if "diff_mask_blur" in data:
config.diff_mask_blur = int(data["diff_mask_blur"])
def save(config: AppConfig) -> None: def save(config: AppConfig) -> None:
@@ -69,6 +85,17 @@ def save(config: AppConfig) -> None:
dm_model=config.dm_model, dm_model=config.dm_model,
dm_gpu=config.dm_gpu, dm_gpu=config.dm_gpu,
dm_feed_restored=config.dm_feed_restored, dm_feed_restored=config.dm_feed_restored,
diff_backend=config.diff_backend,
diff_url=config.diff_url,
diff_model=config.diff_model,
diff_prompt=config.diff_prompt,
diff_negative=config.diff_negative,
diff_steps=config.diff_steps,
diff_cfg=config.diff_cfg,
diff_denoise=config.diff_denoise,
diff_seed=config.diff_seed,
diff_mask_dilate=config.diff_mask_dilate,
diff_mask_blur=config.diff_mask_blur,
) )
_write(data) _write(data)
+34 -6
View File
@@ -67,7 +67,7 @@ from ..core.detection.factory import build_detector
from ..core.detection.types import Detection from ..core.detection.types import Detection
from ..core.imageio import imread_unicode, imwrite_unicode from ..core.imageio import imread_unicode, imwrite_unicode
from ..core.project import PROJECT_FILE, Project from ..core.project import PROJECT_FILE, Project
from ..core.restore.factory import build_restorer from ..core.restore.factory import build_restorer, restorer_needs_detections
from ..core.video.extract import extract_frames from ..core.video.extract import extract_frames
from ..core.video.frame import Frame from ..core.video.frame import Frame
from .extract_dialog import ExtractDialog from .extract_dialog import ExtractDialog
@@ -1220,13 +1220,24 @@ class MainWindow(QMainWindow):
return return
path = self._current path = self._current
key = str(path) key = str(path)
# The diffusion engine builds its mask from detections — capture them on the GUI
# thread (DeepMosaics ignores them). Warn if it needs them but none are computed.
needs_dets = restorer_needs_detections(self._cfg.restorer)
dets_for_restore = list(self._results.get(key, []))
if needs_dets and not dets_for_restore:
self.statusBar().showMessage(
"Diffusion перерисовывает по детекции — на этом кадре цензура не найдена "
"(сначала «Рассчитать кадр»)"
)
return
def fn(job): def fn(job):
img = imread_unicode(key) img = imread_unicode(key)
if img is None: if img is None:
raise RuntimeError(f"Не удалось прочитать: {path.name}") raise RuntimeError(f"Не удалось прочитать: {path.name}")
restorer = self._make_restorer() restorer = self._make_restorer()
restored = restorer.restore(img, [], should_cancel=lambda: job.cancelled) dets = dets_for_restore if restorer.needs_detections else []
restored = restorer.restore(img, dets, should_cancel=lambda: job.cancelled)
return ("restored", key, restored, restorer.name) return ("restored", key, restored, restorer.name)
def done(result, cancelled): def done(result, cancelled):
@@ -1275,9 +1286,17 @@ class MainWindow(QMainWindow):
return return
files = list(self._files) # snapshot — favorites/move mutate self._files files = list(self._files) # snapshot — favorites/move mutate self._files
is_temporal = self._cfg.restorer == "deepmosaics_video" is_temporal = self._cfg.restorer == "deepmosaics_video"
# The diffusion engine masks the detections, so it inherently runs only on hits and
# requires detection to be computed — same gating as the explicit "найденное" mode.
needs_dets = restorer_needs_detections(self._cfg.restorer)
hits = [i for i, p in enumerate(files) if self._results.get(str(p))] hits = [i for i, p in enumerate(files) if self._results.get(str(p))]
# Per-index detection snapshot (GUI-thread read) — fed to engines that need a mask.
dets_by_index = (
{i: list(self._results.get(str(files[i]), [])) for i in hits} if needs_dets else {}
)
need_hits = only_detected or needs_dets
if only_detected: if need_hits:
if not self._results: if not self._results:
self.statusBar().showMessage( self.statusBar().showMessage(
"Детекция не посчитана — сначала «Детектировать все» (или «Расцензурить все»)" "Детекция не посчитана — сначала «Детектировать все» (или «Расцензурить все»)"
@@ -1331,7 +1350,8 @@ class MainWindow(QMainWindow):
should_cancel=lambda: job.cancelled, should_cancel=lambda: job.cancelled,
) )
else: else:
indices = hits if only_detected else range(len(files)) indices = hits if need_hits else range(len(files))
use_dets = restorer.needs_detections
for i in indices: for i in indices:
if job.cancelled: if job.cancelled:
break break
@@ -1339,7 +1359,8 @@ class MainWindow(QMainWindow):
if not force and out_path(p).is_file(): if not force and out_path(p).is_file():
emit(i, None, verb="Пропуск") # already restored — count, don't rewrite emit(i, None, verb="Пропуск") # already restored — count, don't rewrite
continue continue
emit(i, restorer.restore(get_frame(i), [], should_cancel=lambda: job.cancelled)) dets = dets_by_index.get(i, []) if use_dets else []
emit(i, restorer.restore(get_frame(i), dets, should_cancel=lambda: job.cancelled))
frame_cache.pop(i, None) # per-frame: don't accumulate frame_cache.pop(i, None) # per-frame: don't accumulate
return None return None
@@ -1359,7 +1380,14 @@ class MainWindow(QMainWindow):
self._start_job(fn, total, on_done=done) self._start_job(fn, total, on_done=done)
def _make_restorer(self): def _make_restorer(self):
key = (self._cfg.restorer, self._cfg.dm_dir, self._cfg.dm_model, self._cfg.dm_gpu) c = self._cfg
key = (
c.restorer, c.dm_dir, c.dm_model, c.dm_gpu,
# diffusion identity — changing any of these must rebuild the engine
c.diff_backend, c.diff_url, c.diff_model, c.diff_prompt, c.diff_negative,
c.diff_steps, c.diff_cfg, c.diff_denoise, c.diff_seed,
c.diff_mask_dilate, c.diff_mask_blur,
)
if key != self._restorer_key: if key != self._restorer_key:
self._restorer = build_restorer(self._cfg.restorer, self._cfg) # may raise self._restorer = build_restorer(self._cfg.restorer, self._cfg) # may raise
self._restorer_key = key self._restorer_key = key
+168 -24
View File
@@ -1,28 +1,38 @@
"""Configure the restoration ("расцензурить") engine. """Configure the restoration ("расцензурить") engine.
Restoration is DeepMosaics-only — pick the per-frame ("картинка") or temporal ("видео") Two kinds of engine:
engine. The network code is vendored (built-in); the user picks a clean model from a • **DeepMosaics** (картинка/видео) — vendored, in-process; *reconstructs* mosaic and
dropdown of the bundled ``models/deepmosaics`` weights (or browses to another locates it itself. Pick a clean model from the bundled ``models/deepmosaics`` weights
``clean_*.pth``) — for the video engine only ``clean_*_video.pth`` is offered. (or browse). ``mosaic_position.pth`` must sit beside it. CUDA GPU strongly recommended.
``mosaic_position.pth`` must sit beside the chosen model. A CUDA GPU is strongly • **Diffusion-inpaint (SwarmUI)** — *regenerates* the detected (masked) regions via an
recommended (GPU id, -1 = CPU/slow). external SwarmUI server over HTTP. Needs YOLO detections for the mask and a running
SwarmUI; the model runs in SwarmUI's process (no torch here). Per-frame (best for
stills — a video sequence will flicker).
The per-engine fields live in two group widgets that are shown/hidden by the engine combo.
""" """
from __future__ import annotations from __future__ import annotations
from pathlib import Path from pathlib import Path
from PySide6.QtCore import Qt
from PySide6.QtWidgets import ( from PySide6.QtWidgets import (
QApplication,
QCheckBox, QCheckBox,
QComboBox, QComboBox,
QDialog, QDialog,
QDialogButtonBox, QDialogButtonBox,
QDoubleSpinBox,
QFileDialog, QFileDialog,
QFormLayout, QFormLayout,
QHBoxLayout, QHBoxLayout,
QLabel, QLabel,
QLineEdit, QLineEdit,
QMessageBox,
QPushButton, QPushButton,
QSpinBox,
QVBoxLayout,
QWidget, QWidget,
) )
@@ -35,22 +45,40 @@ class RestoreDialog(QDialog):
super().__init__(parent) super().__init__(parent)
self._cfg = config self._cfg = config
self.setWindowTitle("Движок восстановления") self.setWindowTitle("Движок восстановления")
self.setMinimumWidth(560) self.setMinimumWidth(580)
self.engine = QComboBox() self.engine = QComboBox()
self.engine.addItem("DeepMosaics — картинка (покадрово, нужна модель+GPU)", "deepmosaics") self.engine.addItem("DeepMosaics — картинка (покадрово, нужна модель+GPU)", "deepmosaics")
self.engine.addItem("DeepMosaics — видео (соседние кадры, лучше для роликов)", "deepmosaics_video") self.engine.addItem("DeepMosaics — видео (соседние кадры, лучше для роликов)", "deepmosaics_video")
self.engine.addItem("Diffusion-inpaint (SwarmUI) — перерисовка по маске YOLO", "diffusion")
self.engine.setCurrentIndex(max(0, self.engine.findData(config.restorer))) self.engine.setCurrentIndex(max(0, self.engine.findData(config.restorer)))
self.engine.currentIndexChanged.connect(self._sync) self.engine.currentIndexChanged.connect(self._sync)
# Model dropdown — bundled clean models, plus the configured one if external. root = QVBoxLayout(self)
top = QFormLayout()
top.addRow("Движок:", self.engine)
root.addLayout(top)
root.addWidget(self._build_deepmosaics_group(config))
root.addWidget(self._build_diffusion_group(config))
buttons = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel)
buttons.accepted.connect(self.accept)
buttons.rejected.connect(self.reject)
root.addWidget(buttons)
self._sync()
# ----------------------------------------------------------- DeepMosaics UI
def _build_deepmosaics_group(self, config: AppConfig) -> QWidget:
self.dm_group = QWidget()
form = QFormLayout(self.dm_group)
form.setContentsMargins(0, 0, 0, 0)
self.model_combo = QComboBox() self.model_combo = QComboBox()
self._populate_models(config.dm_model) self._populate_models(config.dm_model)
self.dm_gpu = QLineEdit(config.dm_gpu or "0") self.dm_gpu = QLineEdit(config.dm_gpu or "0")
self.dm_gpu.setPlaceholderText("0 = первая CUDA-карта, -1 = CPU (медленно)") self.dm_gpu.setPlaceholderText("0 = первая CUDA-карта, -1 = CPU (медленно)")
# Video engine only: feed already-restored past frames into the temporal window.
self.feed_restored = QCheckBox( self.feed_restored = QCheckBox(
"Подавать уже расцензуренные прошлые кадры в окно (эксперим.)" "Подавать уже расцензуренные прошлые кадры в окно (эксперим.)"
) )
@@ -62,8 +90,6 @@ class RestoreDialog(QDialog):
"не гарантирован; выключите для точной реализации DeepMosaics." "не гарантирован; выключите для точной реализации DeepMosaics."
) )
form = QFormLayout(self)
form.addRow("Движок:", self.engine)
form.addRow("Модель:", self._with_browse(self.model_combo, self._browse_model)) form.addRow("Модель:", self._with_browse(self.model_combo, self._browse_model))
form.addRow("GPU id:", self.dm_gpu) form.addRow("GPU id:", self.dm_gpu)
form.addRow("", self.feed_restored) form.addRow("", self.feed_restored)
@@ -76,13 +102,94 @@ class RestoreDialog(QDialog):
) )
hint.setWordWrap(True) hint.setWordWrap(True)
form.addRow(hint) form.addRow(hint)
return self.dm_group
buttons = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel) # ------------------------------------------------------------ Diffusion UI
buttons.accepted.connect(self.accept) def _build_diffusion_group(self, config: AppConfig) -> QWidget:
buttons.rejected.connect(self.reject) self.diff_group = QWidget()
form.addRow(buttons) form = QFormLayout(self.diff_group)
self._sync() form.setContentsMargins(0, 0, 0, 0)
self.diff_url = QLineEdit(config.diff_url or "http://localhost:7801")
self.diff_url.setPlaceholderText("http://localhost:7801")
self.diff_test_btn = QPushButton("Проверить соединение")
self.diff_test_btn.clicked.connect(self._test_connection)
self.diff_test_status = QLabel("")
self.diff_test_status.setWordWrap(True)
self.diff_model = QLineEdit(config.diff_model or "")
self.diff_model.setPlaceholderText("имя чекпойнта в SwarmUI (пусто = текущий)")
self.diff_prompt = QLineEdit(config.diff_prompt or "")
self.diff_prompt.setPlaceholderText("что нарисовать в области под цензурой")
self.diff_negative = QLineEdit(config.diff_negative or "")
self.diff_negative.setPlaceholderText("чего избегать (negative prompt)")
self.diff_steps = QSpinBox()
self.diff_steps.setRange(1, 150)
self.diff_steps.setValue(int(config.diff_steps))
self.diff_cfg = QDoubleSpinBox()
self.diff_cfg.setRange(0.0, 30.0)
self.diff_cfg.setSingleStep(0.5)
self.diff_cfg.setValue(float(config.diff_cfg))
self.diff_denoise = QDoubleSpinBox()
self.diff_denoise.setRange(0.0, 1.0)
self.diff_denoise.setSingleStep(0.05)
self.diff_denoise.setValue(float(config.diff_denoise))
self.diff_denoise.setToolTip("0..1 — насколько перерисовать область (1 = полностью)")
self.diff_seed = QSpinBox()
self.diff_seed.setRange(-1, 2_147_483_647)
self.diff_seed.setValue(int(config.diff_seed))
self.diff_seed.setSpecialValueText("случайный") # at -1
self.diff_dilate = QSpinBox()
self.diff_dilate.setRange(0, 200)
self.diff_dilate.setValue(int(config.diff_mask_dilate))
self.diff_dilate.setToolTip("Расширить маску на N px (закрыть край цензуры)")
self.diff_blur = QSpinBox()
self.diff_blur.setRange(0, 200)
self.diff_blur.setValue(int(config.diff_mask_blur))
self.diff_blur.setToolTip("Размытие края маски, px (мягкий стык)")
form.addRow("SwarmUI URL:", self.diff_url)
form.addRow("", self.diff_test_btn)
form.addRow("", self.diff_test_status)
form.addRow("Чекпойнт:", self.diff_model)
form.addRow("Промпт:", self.diff_prompt)
form.addRow("Negative:", self.diff_negative)
steps_row = QWidget()
h = QHBoxLayout(steps_row)
h.setContentsMargins(0, 0, 0, 0)
h.addWidget(QLabel("Шаги:"))
h.addWidget(self.diff_steps)
h.addWidget(QLabel("CFG:"))
h.addWidget(self.diff_cfg)
h.addWidget(QLabel("Denoise:"))
h.addWidget(self.diff_denoise)
h.addWidget(QLabel("Seed:"))
h.addWidget(self.diff_seed)
form.addRow("", steps_row)
mask_row = QWidget()
hm = QHBoxLayout(mask_row)
hm.setContentsMargins(0, 0, 0, 0)
hm.addWidget(QLabel("Маска: расширить, px:"))
hm.addWidget(self.diff_dilate)
hm.addWidget(QLabel("размытие, px:"))
hm.addWidget(self.diff_blur)
form.addRow("", mask_row)
hint = QLabel(
"Diffusion перерисовывает область ЦЕНЗУРЫ заново (не восстанавливает оригинал),\n"
"опираясь на маску из детекций YOLO и промпт. Нужен запущенный сервер SwarmUI\n"
"и посчитанная детекция. Покадрово — на роликах будет мерцание. Лучше для\n"
"чёрных плашек/заливки, где DeepMosaics бессилен."
)
hint.setWordWrap(True)
form.addRow(hint)
return self.diff_group
# --------------------------------------------------------------- helpers
def _populate_models(self, current: str | None) -> None: def _populate_models(self, current: str | None) -> None:
self.model_combo.clear() self.model_combo.clear()
is_video = self.engine.currentData() == "deepmosaics_video" is_video = self.engine.currentData() == "deepmosaics_video"
@@ -92,7 +199,6 @@ class RestoreDialog(QDialog):
models = discover_models() # per-frame engine: image clean models only models = discover_models() # per-frame engine: image clean models only
for name, path in models: for name, path in models:
self.model_combo.addItem(name, path) self.model_combo.addItem(name, path)
# Keep an externally-configured model selectable even if it's outside the folder.
if current and self.model_combo.findData(current) < 0: if current and self.model_combo.findData(current) < 0:
self.model_combo.addItem(Path(current).stem + " (внешняя)", current) self.model_combo.addItem(Path(current).stem + " (внешняя)", current)
if self.model_combo.count() == 0: if self.model_combo.count() == 0:
@@ -111,13 +217,16 @@ class RestoreDialog(QDialog):
return w return w
def _sync(self) -> None: def _sync(self) -> None:
is_dm = self.engine.currentData() in ("deepmosaics", "deepmosaics_video") engine = self.engine.currentData()
is_video = self.engine.currentData() == "deepmosaics_video" is_dm = engine in ("deepmosaics", "deepmosaics_video")
# The model list differs per engine (image vs video weights) — repopulate. is_video = engine == "deepmosaics_video"
self._populate_models(self._cfg.dm_model) is_diff = engine == "diffusion"
self.model_combo.setEnabled(is_dm) if is_dm: # repopulate model list (image vs video weights differ)
self.dm_gpu.setEnabled(is_dm) self._populate_models(self._cfg.dm_model)
self.feed_restored.setEnabled(is_video) # only the temporal engine has a window self.feed_restored.setEnabled(is_video)
self.dm_group.setVisible(is_dm)
self.diff_group.setVisible(is_diff)
self.adjustSize()
def _browse_model(self) -> None: def _browse_model(self) -> None:
p, _ = QFileDialog.getOpenFileName(self, "Веса DeepMosaics", "", "Веса (*.pth);;Все файлы (*.*)") p, _ = QFileDialog.getOpenFileName(self, "Веса DeepMosaics", "", "Веса (*.pth);;Все файлы (*.*)")
@@ -126,8 +235,43 @@ class RestoreDialog(QDialog):
self.model_combo.addItem(Path(p).stem, p) self.model_combo.addItem(Path(p).stem, p)
self.model_combo.setCurrentIndex(self.model_combo.findData(p)) self.model_combo.setCurrentIndex(self.model_combo.findData(p))
def _test_connection(self) -> None:
"""Ping SwarmUI (GetNewSession) with the current URL and report OK / the error."""
from ..core.restore.swarmui import SwarmUIBackend
url = self.diff_url.text().strip() or "http://localhost:7801"
self.diff_test_status.setText("Проверка…")
self.diff_test_btn.setEnabled(False)
QApplication.setOverrideCursor(Qt.WaitCursor)
QApplication.processEvents()
try:
session = SwarmUIBackend(url, timeout=15.0).ping() # short timeout for the probe
except Exception as e: # noqa: BLE001 — show the server/connection error verbatim
self.diff_test_status.setText(f"<span style='color:#c0392b'>✗ {e}</span>")
QMessageBox.warning(self, "SwarmUI: соединение", str(e))
else:
self.diff_test_status.setText(
f"<span style='color:#27ae60'>✓ Соединение OK (session: {session})</span>"
)
finally:
QApplication.restoreOverrideCursor()
self.diff_test_btn.setEnabled(True)
def apply_to_config(self) -> None: def apply_to_config(self) -> None:
self._cfg.restorer = self.engine.currentData() self._cfg.restorer = self.engine.currentData()
# DeepMosaics
self._cfg.dm_model = self.model_combo.currentData() self._cfg.dm_model = self.model_combo.currentData()
self._cfg.dm_gpu = self.dm_gpu.text().strip() or "0" self._cfg.dm_gpu = self.dm_gpu.text().strip() or "0"
self._cfg.dm_feed_restored = self.feed_restored.isChecked() self._cfg.dm_feed_restored = self.feed_restored.isChecked()
# Diffusion (SwarmUI)
self._cfg.diff_backend = "swarmui"
self._cfg.diff_url = self.diff_url.text().strip() or "http://localhost:7801"
self._cfg.diff_model = self.diff_model.text().strip() or None
self._cfg.diff_prompt = self.diff_prompt.text()
self._cfg.diff_negative = self.diff_negative.text()
self._cfg.diff_steps = self.diff_steps.value()
self._cfg.diff_cfg = self.diff_cfg.value()
self._cfg.diff_denoise = self.diff_denoise.value()
self._cfg.diff_seed = self.diff_seed.value()
self._cfg.diff_mask_dilate = self.diff_dilate.value()
self._cfg.diff_mask_blur = self.diff_blur.value()
+90
View File
@@ -118,6 +118,87 @@ def test_nms() -> None:
check(len(no_nms) == 3, "without nms_iou, all detections are concatenated") check(len(no_nms) == 3, "without nms_iou, all detections are concatenated")
def test_diffusion_restorer() -> None:
print("restore: diffusion inpaint (mask + fake backend, no network)")
from hvideotool.core.restore.diffusion import (
DiffusionBackend,
DiffusionRestorer,
InpaintParams,
)
from hvideotool.core.restore.factory import restorer_needs_detections
from hvideotool.core.restore.mask import detections_to_mask, mask_is_empty
check(restorer_needs_detections("diffusion") is True, "diffusion needs detections")
check(
restorer_needs_detections("deepmosaics") is False,
"deepmosaics doesn't need detections",
)
img = np.zeros((40, 40, 3), dtype=np.uint8)
dets = [_det(0.9, (10, 10, 12, 12), model="a")]
mask = detections_to_mask(dets, img.shape, dilate=0, blur=0)
check(mask[16, 16] == 255 and mask[2, 2] == 0, "mask filled inside bbox, empty outside")
check(mask_is_empty(detections_to_mask([], img.shape)), "no detections => empty mask")
class _RecordingBackend(DiffusionBackend):
def __init__(self):
self.calls = []
def inpaint(self, image_bgr, mask, params, should_cancel=None):
self.calls.append(mask.copy())
return image_bgr.copy()
backend = _RecordingBackend()
r = DiffusionRestorer(backend, InpaintParams(), mask_dilate=0, mask_blur=0)
check(r.needs_detections is True and r.temporal is False, "DiffusionRestorer flags")
same = r.restore(img, [])
check(
len(backend.calls) == 0 and np.array_equal(same, img),
"no dets => backend skipped, original copy returned",
)
r.restore(img, dets)
check(len(backend.calls) == 1, "backend called once when detections present")
check(np.any(backend.calls[0] > 0), "backend received a non-empty mask")
def test_restore_dialog_diffusion() -> None:
print("restore dialog: diffusion engine fields + connection probe")
from PySide6.QtWidgets import QApplication
from hvideotool.core.restore.swarmui import SwarmUIBackend
from hvideotool.ui.restore_dialog import RestoreDialog
_ensure_app(QApplication)
cfg = AppConfig()
cfg.restorer = "diffusion"
cfg.diff_url = "http://localhost:7801"
dlg = RestoreDialog(cfg)
check(hasattr(dlg, "diff_test_btn"), "connection-test button exists")
check(
dlg.diff_group.isVisibleTo(dlg) and not dlg.dm_group.isVisibleTo(dlg),
"diffusion engine shows diffusion group, hides DeepMosaics group",
)
dlg.engine.setCurrentIndex(dlg.engine.findData("deepmosaics"))
check(
dlg.dm_group.isVisibleTo(dlg) and not dlg.diff_group.isVisibleTo(dlg),
"switching to DeepMosaics swaps the visible group",
)
# apply_to_config writes the diffusion fields back.
dlg.engine.setCurrentIndex(dlg.engine.findData("diffusion"))
dlg.diff_prompt.setText("clean skin")
dlg.diff_steps.setValue(33)
dlg.apply_to_config()
check(cfg.restorer == "diffusion" and cfg.diff_prompt == "clean skin" and cfg.diff_steps == 33,
"apply_to_config persists diffusion fields")
# ping() against a dead port raises a clear, actionable RuntimeError (no GUI/modal).
try:
SwarmUIBackend("http://127.0.0.1:1", timeout=1.0).ping()
check(False, "ping should raise when no server is listening")
except RuntimeError as e:
check("SwarmUI" in str(e), "ping raises actionable RuntimeError when server is down")
def test_extract_dialog_options() -> None: def test_extract_dialog_options() -> None:
print("extract dialog: options() includes JPEG quality") print("extract dialog: options() includes JPEG quality")
from PySide6.QtWidgets import QApplication from PySide6.QtWidgets import QApplication
@@ -141,6 +222,10 @@ def test_settings_roundtrip() -> None:
cfg.cross_model_nms = True cfg.cross_model_nms = True
cfg.nms_iou = 0.55 cfg.nms_iou = 0.55
cfg.default_threshold = 0.3 cfg.default_threshold = 0.3
cfg.diff_url = "http://localhost:9999"
cfg.diff_prompt = "test prompt"
cfg.diff_steps = 42
cfg.diff_denoise = 0.7
proj = Project.create(Path(d) / "P", name="P") proj = Project.create(Path(d) / "P", name="P")
proj.update_from_config(cfg) proj.update_from_config(cfg)
proj.save() proj.save()
@@ -150,6 +235,9 @@ def test_settings_roundtrip() -> None:
check(cfg2.model_thresholds == {"penis": 0.42}, "model_thresholds persisted") check(cfg2.model_thresholds == {"penis": 0.42}, "model_thresholds persisted")
check(cfg2.cross_model_nms is True, "cross_model_nms persisted") check(cfg2.cross_model_nms is True, "cross_model_nms persisted")
check(abs(cfg2.nms_iou - 0.55) < 1e-9, "nms_iou persisted") check(abs(cfg2.nms_iou - 0.55) < 1e-9, "nms_iou persisted")
check(cfg2.diff_url == "http://localhost:9999", "diff_url persisted")
check(cfg2.diff_prompt == "test prompt", "diff_prompt persisted")
check(cfg2.diff_steps == 42 and abs(cfg2.diff_denoise - 0.7) < 1e-9, "diff params persisted")
check(not (proj.root / "project.json.tmp").exists(), "project.json.tmp cleaned up") check(not (proj.root / "project.json.tmp").exists(), "project.json.tmp cleaned up")
@@ -227,6 +315,8 @@ def main() -> int:
tests = [ tests = [
test_cache_atomic_roundtrip, test_cache_atomic_roundtrip,
test_nms, test_nms,
test_diffusion_restorer,
test_restore_dialog_diffusion,
test_extract_dialog_options, test_extract_dialog_options,
test_settings_roundtrip, test_settings_roundtrip,
test_mainwindow_filter_and_jump, test_mainwindow_filter_and_jump,
+97
View File
@@ -0,0 +1,97 @@
"""Probe the diffusion-inpaint path against a LIVE SwarmUI server (no GUI).
Exercises exactly what the app does — builds an inpaint mask from a detection and runs
``DiffusionRestorer`` over a real SwarmUI backend — so you can verify the HTTP/API
plumbing (field names drift between SwarmUI versions) independently of the GUI.
Prereqs: a running SwarmUI server with an inpaint-capable checkpoint loaded.
Run (PowerShell)::
.venv\\Scripts\\python.exe scripts\\swarmui_probe.py `
--url http://localhost:7801 `
--image "D:\\path\\to\\frame.jpg" `
--out "D:\\path\\to\\out.png" `
--prompt "clean skin" --steps 25 --denoise 1.0
With no --image a 512x512 synthetic frame is used (a grey box censored in the centre).
Exits non-zero on any error and prints the server's reply on failure.
"""
from __future__ import annotations
import argparse
import sys
from pathlib import Path
import numpy as np
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from hvideotool.core.detection.types import CensorType, Detection # noqa: E402
from hvideotool.core.imageio import imread_unicode, imwrite_unicode # noqa: E402
from hvideotool.core.restore.diffusion import DiffusionRestorer, InpaintParams # noqa: E402
from hvideotool.core.restore.swarmui import SwarmUIBackend # noqa: E402
def _synthetic() -> tuple[np.ndarray, Detection]:
"""A 512x512 frame with a 'censored' grey block in the middle + a matching detection."""
img = np.full((512, 512, 3), 60, np.uint8)
img[40:472, 40:472] = (120, 90, 70) # some background
img[180:332, 180:332] = 128 # the "censored" block
det = Detection(type=CensorType.MOSAIC, score=0.99, bbox=(180, 180, 152, 152), label="mosaic")
return img, det
def main() -> int:
ap = argparse.ArgumentParser(description="Probe diffusion-inpaint against a live SwarmUI")
ap.add_argument("--url", default="http://localhost:7801")
ap.add_argument("--image", default=None, help="frame to inpaint (default: synthetic)")
ap.add_argument("--out", default="swarmui_probe_out.png")
ap.add_argument("--model", default=None, help="checkpoint name as SwarmUI knows it")
ap.add_argument("--prompt", default="")
ap.add_argument("--negative", default="")
ap.add_argument("--steps", type=int, default=25)
ap.add_argument("--cfg", type=float, default=7.0)
ap.add_argument("--denoise", type=float, default=1.0)
ap.add_argument("--seed", type=int, default=-1)
ap.add_argument("--dilate", type=int, default=4)
ap.add_argument("--blur", type=int, default=8)
args = ap.parse_args()
if args.image:
img = imread_unicode(args.image)
if img is None:
print(f"Не удалось прочитать {args.image}", file=sys.stderr)
return 2
h, w = img.shape[:2]
# No detector here — mask the central third so there's something to inpaint.
bx, by = w // 3, h // 3
det = Detection(type=CensorType.MOSAIC, score=0.99, bbox=(bx, by, w // 3, h // 3), label="mosaic")
else:
img, det = _synthetic()
backend = SwarmUIBackend(args.url)
params = InpaintParams(
prompt=args.prompt, negative=args.negative, model=args.model,
steps=args.steps, cfg=args.cfg, denoise=args.denoise, seed=args.seed,
)
restorer = DiffusionRestorer(backend, params, mask_dilate=args.dilate, mask_blur=args.blur)
print(f"→ SwarmUI {args.url} frame={img.shape[1]}x{img.shape[0]} steps={args.steps} denoise={args.denoise}")
try:
out = restorer.restore(img, [det])
except Exception as e: # noqa: BLE001 — surface the server error verbatim
print(f"ОШИБКА: {e}", file=sys.stderr)
return 1
if not imwrite_unicode(args.out, out):
print(f"Не удалось записать {args.out}", file=sys.stderr)
return 3
changed = int(np.count_nonzero(np.any(out.astype(int) - img.astype(int) != 0, axis=2)))
print(f"✓ Готово → {args.out} (изменено пикселей: {changed})")
return 0
if __name__ == "__main__":
raise SystemExit(main())