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:
@@ -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)
|
||||
Reference in New Issue
Block a user