81 lines
3.3 KiB
Python
81 lines
3.3 KiB
Python
"""Restorer factory: build a Restorer from the app config.
|
|
|
|
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).
|
|
|
|
A CUDA GPU is recommended (for DeepMosaics in-process; for diffusion it's the server's GPU).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
from .base import Restorer
|
|
|
|
if TYPE_CHECKING: # avoid importing AppConfig at runtime here (not needed)
|
|
from ...config import AppConfig
|
|
|
|
|
|
def build_restorer(name: str = "deepmosaics", config: AppConfig | None = None) -> Restorer:
|
|
if config is None:
|
|
raise ValueError("Для DeepMosaics нужны настройки (config).")
|
|
if name == "deepmosaics":
|
|
from .deepmosaics import DeepMosaicsRestorer
|
|
|
|
return DeepMosaicsRestorer(
|
|
config.dm_dir, config.dm_model, config.dm_gpu
|
|
)
|
|
if name == "deepmosaics_video":
|
|
from .deepmosaics import DeepMosaicsVideoRestorer
|
|
|
|
return DeepMosaicsVideoRestorer(
|
|
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"
|