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.

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

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",
 )
 
 

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__

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)

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"

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))

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()

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)
 

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

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"<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:
         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()