diff --git a/CLAUDE.md b/CLAUDE.md index 4f3931f..1fd2eef 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -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//`, 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//` — 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//` — 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 diff --git a/README.md b/README.md index a37d3e0..ffe7cee 100644 --- a/README.md +++ b/README.md @@ -104,9 +104,10 @@ отдельно от `frames/`, чтобы результаты не попадали обратно в список кадров). Прогресс, отмена (**«■ Стоп»**) и предпросмотр текущего кадра работают как при детекции. -Расцензуривание — только через **DeepMosaics** (код **встроен** в приложение, vendored, -GPL-3.0; ставить отдельно не нужно — требуются только **веса** и желательно **GPU -NVIDIA/CUDA**). Два движка: +Доступны два **семейства** движков: **DeepMosaics** (восстанавливает мозаику) и +**diffusion-inpaint (SwarmUI)** (перерисовывает область заново). DeepMosaics — по +умолчанию; код **встроен** (vendored, GPL-3.0; ставить отдельно не нужно — нужны только +**веса** и желательно **GPU NVIDIA/CUDA**). Три движка в списке: - **DeepMosaics — картинка** — реальное генеративное удаление мозаики **покадрово**. - **DeepMosaics — видео (BVDNet)** — **временно́й** движок: использует **соседние кадры** @@ -114,8 +115,12 @@ NVIDIA/CUDA**). Два движка: роликах. Из-за рекуррентности обрабатывает **непрерывный диапазон по порядку** — т.е. запускайте его через **«Расцензурить все»** (одиночный «Расцензурить кадр» сведётся к окну из одного кадра). Нужна **видеомодель** `clean_youknow_video.pth`. +- **Diffusion-inpaint (SwarmUI)** — **перерисовывает** область цензуры заново диффузионной + inpaint-моделью по **маске из YOLO-детекций** и промпту (не восстанавливает оригинал!). + Лучше всего для **чёрных плашек / сплошной заливки**, где DeepMosaics бессилен. Покадрово + → на роликах будет **мерцание**. См. «Настройка SwarmUI» ниже. -На аниме качество ограничено (модели обучены на реальном видео). +На аниме качество DeepMosaics ограничено (модели обучены на реальном видео). ### Настройка DeepMosaics @@ -144,6 +149,34 @@ Baidu код `1x0a`): > `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). @@ -270,11 +303,14 @@ hvideotool/ │ ├── types.py # Detection (+ label/категория, .display), CensorType │ ├── cache.py # кэш детекций (ключ = набор моделей + conf/imgsz) │ └── yolo.py # YOLO-детектор (Ultralytics, маски→полигоны, ярлык категории) - └── restore/ # только DeepMosaics - ├── base.py # Restorer (ABC): restore() + restore_sequence() + .temporal - ├── factory.py # build_restorer -> deepmosaics | deepmosaics_video + └── restore/ # DeepMosaics (восстановление) или diffusion-inpaint (перерисовка) + ├── base.py # Restorer (ABC): restore() + restore_sequence() + .temporal/.needs_detections + ├── factory.py # build_restorer -> deepmosaics | deepmosaics_video | diffusion ├── 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 @@ -289,7 +325,7 @@ hvideotool/ | GUI | PySide6 (Qt 6) | | Обработка картинок | OpenCV / NumPy | | Детектор | Ultralytics YOLO, мульти-модель (models/yolo/<кат>) | -| Расцензуривание | DeepMosaics (встроен) + PyTorch/CUDA | +| Расцензуривание | DeepMosaics (встроен) + PyTorch/CUDA, или diffusion-inpaint через SwarmUI (HTTP) | ## Лицензия diff --git a/hvideotool/config.py b/hvideotool/config.py index 425847f..569bcbf 100644 --- a/hvideotool/config.py +++ b/hvideotool/config.py @@ -75,3 +75,19 @@ class AppConfig: # (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. 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 diff --git a/hvideotool/core/project.py b/hvideotool/core/project.py index 892a3bc..112e4ba 100644 --- a/hvideotool/core/project.py +++ b/hvideotool/core/project.py @@ -47,6 +47,17 @@ _SETTING_KEYS = ( "dm_model", "dm_gpu", "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", ) diff --git a/hvideotool/core/restore/base.py b/hvideotool/core/restore/base.py index ee4078e..7185a0f 100644 --- a/hvideotool/core/restore/base.py +++ b/hvideotool/core/restore/base.py @@ -37,6 +37,12 @@ class Restorer(ABC): #: leave this False; the temporal DeepMosaics (BVDNet) sets it True. 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 def name(self) -> str: return type(self).__name__ diff --git a/hvideotool/core/restore/diffusion.py b/hvideotool/core/restore/diffusion.py new file mode 100644 index 0000000..9e964d4 --- /dev/null +++ b/hvideotool/core/restore/diffusion.py @@ -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) diff --git a/hvideotool/core/restore/factory.py b/hvideotool/core/restore/factory.py index 636bd5c..02d4f34 100644 --- a/hvideotool/core/restore/factory.py +++ b/hvideotool/core/restore/factory.py @@ -1,13 +1,16 @@ """Restorer factory: build a Restorer from the app config. -Only DeepMosaics is supported (the cv2 inpaint baseline was removed — it filled but -did not reconstruct). The DeepMosaics network code is vendored (``_deepmosaics/``, -GPL-3.0) and run in-process; the user supplies only the weights (+ ``mosaic_position.pth`` -alongside). A CUDA GPU is recommended. +Engines: +- ``deepmosaics``: per-frame generative mosaic removal (image model). Vendored network + code (``_deepmosaics/``, GPL-3.0) run in-process; user supplies only the weights + (+ ``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). -- ``deepmosaics_video``: temporal variant (BVDNet) that uses neighbouring frames for - coherence — needs the ``clean_*_video.pth`` weights and a contiguous frame sequence. +A CUDA GPU is recommended (for DeepMosaics in-process; for diffusion it's the server's GPU). """ 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, 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": raise ValueError( "Движок LADA пока не подключён. Используйте DeepMosaics. См. README." ) 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" diff --git a/hvideotool/core/restore/mask.py b/hvideotool/core/restore/mask.py new file mode 100644 index 0000000..e7a3367 --- /dev/null +++ b/hvideotool/core/restore/mask.py @@ -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)) diff --git a/hvideotool/core/restore/swarmui.py b/hvideotool/core/restore/swarmui.py new file mode 100644 index 0000000..cec4970 --- /dev/null +++ b/hvideotool/core/restore/swarmui.py @@ -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() diff --git a/hvideotool/settings_store.py b/hvideotool/settings_store.py index 238b3a9..39f8a95 100644 --- a/hvideotool/settings_store.py +++ b/hvideotool/settings_store.py @@ -52,6 +52,22 @@ def apply(config: AppConfig) -> None: setattr(config, key, data[key]) if "dm_feed_restored" in data: 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: @@ -69,6 +85,17 @@ def save(config: AppConfig) -> None: dm_model=config.dm_model, dm_gpu=config.dm_gpu, 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) diff --git a/hvideotool/ui/main_window.py b/hvideotool/ui/main_window.py index c67c2b5..e44cde4 100644 --- a/hvideotool/ui/main_window.py +++ b/hvideotool/ui/main_window.py @@ -67,7 +67,7 @@ from ..core.detection.factory import build_detector from ..core.detection.types import Detection from ..core.imageio import imread_unicode, imwrite_unicode 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.frame import Frame from .extract_dialog import ExtractDialog @@ -1220,13 +1220,24 @@ class MainWindow(QMainWindow): return path = self._current 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): img = imread_unicode(key) if img is None: raise RuntimeError(f"Не удалось прочитать: {path.name}") 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) def done(result, cancelled): @@ -1275,9 +1286,17 @@ class MainWindow(QMainWindow): return files = list(self._files) # snapshot — favorites/move mutate self._files 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))] + # 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: self.statusBar().showMessage( "Детекция не посчитана — сначала «Детектировать все» (или «Расцензурить все»)" @@ -1331,7 +1350,8 @@ class MainWindow(QMainWindow): should_cancel=lambda: job.cancelled, ) 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: if job.cancelled: break @@ -1339,7 +1359,8 @@ class MainWindow(QMainWindow): if not force and out_path(p).is_file(): emit(i, None, verb="Пропуск") # already restored — count, don't rewrite 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 return None @@ -1359,7 +1380,14 @@ class MainWindow(QMainWindow): self._start_job(fn, total, on_done=done) 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: self._restorer = build_restorer(self._cfg.restorer, self._cfg) # may raise self._restorer_key = key diff --git a/hvideotool/ui/restore_dialog.py b/hvideotool/ui/restore_dialog.py index 9b6793f..d8d253f 100644 --- a/hvideotool/ui/restore_dialog.py +++ b/hvideotool/ui/restore_dialog.py @@ -1,28 +1,38 @@ """Configure the restoration ("расцензурить") engine. -Restoration is DeepMosaics-only — pick the per-frame ("картинка") or temporal ("видео") -engine. The network code is vendored (built-in); the user picks a clean model from a -dropdown of the bundled ``models/deepmosaics`` weights (or browses to another -``clean_*.pth``) — for the video engine only ``clean_*_video.pth`` is offered. -``mosaic_position.pth`` must sit beside the chosen model. A CUDA GPU is strongly -recommended (GPU id, -1 = CPU/slow). +Two kinds of engine: + • **DeepMosaics** (картинка/видео) — vendored, in-process; *reconstructs* mosaic and + locates it itself. Pick a clean model from the bundled ``models/deepmosaics`` weights + (or browse). ``mosaic_position.pth`` must sit beside it. CUDA GPU strongly recommended. + • **Diffusion-inpaint (SwarmUI)** — *regenerates* the detected (masked) regions via an + 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 pathlib import Path +from PySide6.QtCore import Qt from PySide6.QtWidgets import ( + QApplication, QCheckBox, QComboBox, QDialog, QDialogButtonBox, + QDoubleSpinBox, QFileDialog, QFormLayout, QHBoxLayout, QLabel, QLineEdit, + QMessageBox, QPushButton, + QSpinBox, + QVBoxLayout, QWidget, ) @@ -35,22 +45,40 @@ class RestoreDialog(QDialog): super().__init__(parent) self._cfg = config self.setWindowTitle("Движок восстановления") - self.setMinimumWidth(560) + self.setMinimumWidth(580) self.engine = QComboBox() self.engine.addItem("DeepMosaics — картинка (покадрово, нужна модель+GPU)", "deepmosaics") 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.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._populate_models(config.dm_model) self.dm_gpu = QLineEdit(config.dm_gpu or "0") self.dm_gpu.setPlaceholderText("0 = первая CUDA-карта, -1 = CPU (медленно)") - # Video engine only: feed already-restored past frames into the temporal window. self.feed_restored = QCheckBox( "Подавать уже расцензуренные прошлые кадры в окно (эксперим.)" ) @@ -62,8 +90,6 @@ class RestoreDialog(QDialog): "не гарантирован; выключите для точной реализации DeepMosaics." ) - form = QFormLayout(self) - form.addRow("Движок:", self.engine) form.addRow("Модель:", self._with_browse(self.model_combo, self._browse_model)) form.addRow("GPU id:", self.dm_gpu) form.addRow("", self.feed_restored) @@ -76,13 +102,94 @@ class RestoreDialog(QDialog): ) hint.setWordWrap(True) form.addRow(hint) + return self.dm_group - buttons = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel) - buttons.accepted.connect(self.accept) - buttons.rejected.connect(self.reject) - form.addRow(buttons) - self._sync() + # ------------------------------------------------------------ Diffusion UI + def _build_diffusion_group(self, config: AppConfig) -> QWidget: + self.diff_group = QWidget() + form = QFormLayout(self.diff_group) + 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: self.model_combo.clear() is_video = self.engine.currentData() == "deepmosaics_video" @@ -92,7 +199,6 @@ class RestoreDialog(QDialog): models = discover_models() # per-frame engine: image clean models only for name, path in models: 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: self.model_combo.addItem(Path(current).stem + " (внешняя)", current) if self.model_combo.count() == 0: @@ -111,13 +217,16 @@ class RestoreDialog(QDialog): return w def _sync(self) -> None: - is_dm = self.engine.currentData() in ("deepmosaics", "deepmosaics_video") - is_video = self.engine.currentData() == "deepmosaics_video" - # The model list differs per engine (image vs video weights) — repopulate. - self._populate_models(self._cfg.dm_model) - self.model_combo.setEnabled(is_dm) - self.dm_gpu.setEnabled(is_dm) - self.feed_restored.setEnabled(is_video) # only the temporal engine has a window + engine = self.engine.currentData() + is_dm = engine in ("deepmosaics", "deepmosaics_video") + is_video = engine == "deepmosaics_video" + is_diff = engine == "diffusion" + if is_dm: # repopulate model list (image vs video weights differ) + self._populate_models(self._cfg.dm_model) + 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: p, _ = QFileDialog.getOpenFileName(self, "Веса DeepMosaics", "", "Веса (*.pth);;Все файлы (*.*)") @@ -126,8 +235,43 @@ class RestoreDialog(QDialog): self.model_combo.addItem(Path(p).stem, 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"✗ {e}") + QMessageBox.warning(self, "SwarmUI: соединение", str(e)) + else: + self.diff_test_status.setText( + f"✓ Соединение OK (session: {session})" + ) + finally: + QApplication.restoreOverrideCursor() + self.diff_test_btn.setEnabled(True) + def apply_to_config(self) -> None: self._cfg.restorer = self.engine.currentData() + # DeepMosaics self._cfg.dm_model = self.model_combo.currentData() self._cfg.dm_gpu = self.dm_gpu.text().strip() or "0" 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() diff --git a/scripts/smoke_test.py b/scripts/smoke_test.py index 348b330..d3e6b5d 100644 --- a/scripts/smoke_test.py +++ b/scripts/smoke_test.py @@ -118,6 +118,87 @@ def test_nms() -> None: 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: print("extract dialog: options() includes JPEG quality") from PySide6.QtWidgets import QApplication @@ -141,6 +222,10 @@ def test_settings_roundtrip() -> None: cfg.cross_model_nms = True cfg.nms_iou = 0.55 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.update_from_config(cfg) proj.save() @@ -150,6 +235,9 @@ def test_settings_roundtrip() -> None: check(cfg2.model_thresholds == {"penis": 0.42}, "model_thresholds 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(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") @@ -227,6 +315,8 @@ def main() -> int: tests = [ test_cache_atomic_roundtrip, test_nms, + test_diffusion_restorer, + test_restore_dialog_diffusion, test_extract_dialog_options, test_settings_roundtrip, test_mainwindow_filter_and_jump, diff --git a/scripts/swarmui_probe.py b/scripts/swarmui_probe.py new file mode 100644 index 0000000..8e6a833 --- /dev/null +++ b/scripts/swarmui_probe.py @@ -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())