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