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

This commit is contained in:
Leonid Pershin
2026-06-08 06:21:44 +03:00
parent 8a366ed43d
commit 15f89b395d
14 changed files with 903 additions and 70 deletions
+103
View File
@@ -0,0 +1,103 @@
"""Diffusion-inpaint restoration — redraw censored regions with a diffusion backend.
Unlike DeepMosaics (which *reconstructs* mosaic from its residual low-frequency data and
locates it itself), this engine *regenerates* the masked region with a diffusion inpaint
model: it builds a mask from the YOLO detections and hands ``(image, mask, params)`` to a
pluggable :class:`DiffusionBackend` (SwarmUI is the first, see ``swarmui.py``).
Consequences of that design:
- It **needs detections** (``needs_detections = True``) — a frame with none comes back
unchanged (no mask → nothing to regenerate). The caller feeds it the real detections.
- It's **per-frame** (``temporal = False``): each frame is generated independently, so a
video sequence will flicker. Best for stills / single frames, not coherent clips.
- The backend runs in a **separate process/server** (e.g. SwarmUI over HTTP), so this
path adds **no torch dependency** to the app and keeps the heavy model out-of-process.
The backend is abstract so other diffusion servers (ComfyUI/A1111) can be added later as
another :class:`DiffusionBackend`, without touching the restorer or the UI.
"""
from __future__ import annotations
from abc import ABC, abstractmethod
from dataclasses import dataclass
import numpy as np
from ..detection.types import Detection
from .base import CancelCheck, Cancelled, Restorer
from .mask import detections_to_mask, mask_is_empty
@dataclass(frozen=True)
class InpaintParams:
"""Generation knobs handed to a :class:`DiffusionBackend`."""
prompt: str = ""
negative: str = ""
model: str | None = None # checkpoint name as the backend knows it (None = current)
steps: int = 30
cfg: float = 7.0
denoise: float = 1.0 # 0..1 — how much to regenerate under the mask (1 = full)
seed: int = -1 # -1 = random each call
mask_blur: int = 8 # px feather applied by the backend at its mask edge
class DiffusionBackend(ABC):
"""A diffusion inpaint engine reachable from our process (typically over HTTP)."""
@property
def name(self) -> str:
return type(self).__name__
@abstractmethod
def inpaint(
self,
image_bgr: np.ndarray,
mask: np.ndarray,
params: InpaintParams,
should_cancel: CancelCheck | None = None,
) -> np.ndarray:
"""Regenerate the white area of ``mask`` in ``image_bgr``; return a new BGR image."""
raise NotImplementedError
class DiffusionRestorer(Restorer):
"""Restorer that masks the detected regions and inpaints them via a backend."""
temporal = False
needs_detections = True
def __init__(
self,
backend: DiffusionBackend,
params: InpaintParams,
*,
mask_dilate: int = 4,
mask_blur: int = 8,
) -> None:
self._backend = backend
self._params = params
self._dilate = mask_dilate
self._blur = mask_blur
@property
def name(self) -> str:
return f"Diffusion({self._backend.name})"
def restore(
self,
image: np.ndarray,
detections: list[Detection],
should_cancel: CancelCheck | None = None,
) -> np.ndarray:
if should_cancel is not None and should_cancel():
raise Cancelled("Восстановление отменено")
if not detections:
return image.copy() # no detections → no mask → nothing to regenerate
mask = detections_to_mask(
detections, image.shape, dilate=self._dilate, blur=self._blur
)
if mask_is_empty(mask):
return image.copy()
return self._backend.inpaint(image, mask, self._params, should_cancel)