104 lines
3.7 KiB
Python
104 lines
3.7 KiB
Python
"""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)
|