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:
@@ -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
|
||||||
|
|||||||
@@ -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) |
|
||||||
|
|
||||||
## Лицензия
|
## Лицензия
|
||||||
|
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
@@ -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",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -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__
|
||||||
|
|||||||
@@ -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)
|
||||||
@@ -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"
|
||||||
|
|||||||
@@ -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))
|
||||||
@@ -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()
|
||||||
@@ -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)
|
||||||
|
|
||||||
|
|||||||
@@ -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
@@ -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()
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
@@ -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())
|
||||||
Reference in New Issue
Block a user