98 lines
4.1 KiB
Python
98 lines
4.1 KiB
Python
"""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())
|