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
+97
View File
@@ -0,0 +1,97 @@
"""Probe the diffusion-inpaint path against a LIVE SwarmUI server (no GUI).
Exercises exactly what the app does — builds an inpaint mask from a detection and runs
``DiffusionRestorer`` over a real SwarmUI backend — so you can verify the HTTP/API
plumbing (field names drift between SwarmUI versions) independently of the GUI.
Prereqs: a running SwarmUI server with an inpaint-capable checkpoint loaded.
Run (PowerShell)::
.venv\\Scripts\\python.exe scripts\\swarmui_probe.py `
--url http://localhost:7801 `
--image "D:\\path\\to\\frame.jpg" `
--out "D:\\path\\to\\out.png" `
--prompt "clean skin" --steps 25 --denoise 1.0
With no --image a 512x512 synthetic frame is used (a grey box censored in the centre).
Exits non-zero on any error and prints the server's reply on failure.
"""
from __future__ import annotations
import argparse
import sys
from pathlib import Path
import numpy as np
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from hvideotool.core.detection.types import CensorType, Detection # noqa: E402
from hvideotool.core.imageio import imread_unicode, imwrite_unicode # noqa: E402
from hvideotool.core.restore.diffusion import DiffusionRestorer, InpaintParams # noqa: E402
from hvideotool.core.restore.swarmui import SwarmUIBackend # noqa: E402
def _synthetic() -> tuple[np.ndarray, Detection]:
"""A 512x512 frame with a 'censored' grey block in the middle + a matching detection."""
img = np.full((512, 512, 3), 60, np.uint8)
img[40:472, 40:472] = (120, 90, 70) # some background
img[180:332, 180:332] = 128 # the "censored" block
det = Detection(type=CensorType.MOSAIC, score=0.99, bbox=(180, 180, 152, 152), label="mosaic")
return img, det
def main() -> int:
ap = argparse.ArgumentParser(description="Probe diffusion-inpaint against a live SwarmUI")
ap.add_argument("--url", default="http://localhost:7801")
ap.add_argument("--image", default=None, help="frame to inpaint (default: synthetic)")
ap.add_argument("--out", default="swarmui_probe_out.png")
ap.add_argument("--model", default=None, help="checkpoint name as SwarmUI knows it")
ap.add_argument("--prompt", default="")
ap.add_argument("--negative", default="")
ap.add_argument("--steps", type=int, default=25)
ap.add_argument("--cfg", type=float, default=7.0)
ap.add_argument("--denoise", type=float, default=1.0)
ap.add_argument("--seed", type=int, default=-1)
ap.add_argument("--dilate", type=int, default=4)
ap.add_argument("--blur", type=int, default=8)
args = ap.parse_args()
if args.image:
img = imread_unicode(args.image)
if img is None:
print(f"Не удалось прочитать {args.image}", file=sys.stderr)
return 2
h, w = img.shape[:2]
# No detector here — mask the central third so there's something to inpaint.
bx, by = w // 3, h // 3
det = Detection(type=CensorType.MOSAIC, score=0.99, bbox=(bx, by, w // 3, h // 3), label="mosaic")
else:
img, det = _synthetic()
backend = SwarmUIBackend(args.url)
params = InpaintParams(
prompt=args.prompt, negative=args.negative, model=args.model,
steps=args.steps, cfg=args.cfg, denoise=args.denoise, seed=args.seed,
)
restorer = DiffusionRestorer(backend, params, mask_dilate=args.dilate, mask_blur=args.blur)
print(f"→ SwarmUI {args.url} frame={img.shape[1]}x{img.shape[0]} steps={args.steps} denoise={args.denoise}")
try:
out = restorer.restore(img, [det])
except Exception as e: # noqa: BLE001 — surface the server error verbatim
print(f"ОШИБКА: {e}", file=sys.stderr)
return 1
if not imwrite_unicode(args.out, out):
print(f"Не удалось записать {args.out}", file=sys.stderr)
return 3
changed = int(np.count_nonzero(np.any(out.astype(int) - img.astype(int) != 0, axis=2)))
print(f"✓ Готово → {args.out} (изменено пикселей: {changed})")
return 0
if __name__ == "__main__":
raise SystemExit(main())