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