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.

2026-06-08 06:21:44 +03:00
parent 8a366ed43d
commit 15f89b395d
14 changed files with 903 additions and 70 deletions
@@ -118,6 +118,87 @@ def test_nms() -> None:
    check(len(no_nms) == 3, "without nms_iou, all detections are concatenated")


+def test_diffusion_restorer() -> None:
+    print("restore: diffusion inpaint (mask + fake backend, no network)")
+    from hvideotool.core.restore.diffusion import (
+        DiffusionBackend,
+        DiffusionRestorer,
+        InpaintParams,
+    )
+    from hvideotool.core.restore.factory import restorer_needs_detections
+    from hvideotool.core.restore.mask import detections_to_mask, mask_is_empty
+
+    check(restorer_needs_detections("diffusion") is True, "diffusion needs detections")
+    check(
+        restorer_needs_detections("deepmosaics") is False,
+        "deepmosaics doesn't need detections",
+    )
+
+    img = np.zeros((40, 40, 3), dtype=np.uint8)
+    dets = [_det(0.9, (10, 10, 12, 12), model="a")]
+    mask = detections_to_mask(dets, img.shape, dilate=0, blur=0)
+    check(mask[16, 16] == 255 and mask[2, 2] == 0, "mask filled inside bbox, empty outside")
+    check(mask_is_empty(detections_to_mask([], img.shape)), "no detections => empty mask")
+
+    class _RecordingBackend(DiffusionBackend):
+        def __init__(self):
+            self.calls = []
+
+        def inpaint(self, image_bgr, mask, params, should_cancel=None):
+            self.calls.append(mask.copy())
+            return image_bgr.copy()
+
+    backend = _RecordingBackend()
+    r = DiffusionRestorer(backend, InpaintParams(), mask_dilate=0, mask_blur=0)
+    check(r.needs_detections is True and r.temporal is False, "DiffusionRestorer flags")
+    same = r.restore(img, [])
+    check(
+        len(backend.calls) == 0 and np.array_equal(same, img),
+        "no dets => backend skipped, original copy returned",
+    )
+    r.restore(img, dets)
+    check(len(backend.calls) == 1, "backend called once when detections present")
+    check(np.any(backend.calls[0] > 0), "backend received a non-empty mask")
+
+
+def test_restore_dialog_diffusion() -> None:
+    print("restore dialog: diffusion engine fields + connection probe")
+    from PySide6.QtWidgets import QApplication
+
+    from hvideotool.core.restore.swarmui import SwarmUIBackend
+    from hvideotool.ui.restore_dialog import RestoreDialog
+
+    _ensure_app(QApplication)
+    cfg = AppConfig()
+    cfg.restorer = "diffusion"
+    cfg.diff_url = "http://localhost:7801"
+    dlg = RestoreDialog(cfg)
+    check(hasattr(dlg, "diff_test_btn"), "connection-test button exists")
+    check(
+        dlg.diff_group.isVisibleTo(dlg) and not dlg.dm_group.isVisibleTo(dlg),
+        "diffusion engine shows diffusion group, hides DeepMosaics group",
+    )
+    dlg.engine.setCurrentIndex(dlg.engine.findData("deepmosaics"))
+    check(
+        dlg.dm_group.isVisibleTo(dlg) and not dlg.diff_group.isVisibleTo(dlg),
+        "switching to DeepMosaics swaps the visible group",
+    )
+    # apply_to_config writes the diffusion fields back.
+    dlg.engine.setCurrentIndex(dlg.engine.findData("diffusion"))
+    dlg.diff_prompt.setText("clean skin")
+    dlg.diff_steps.setValue(33)
+    dlg.apply_to_config()
+    check(cfg.restorer == "diffusion" and cfg.diff_prompt == "clean skin" and cfg.diff_steps == 33,
+          "apply_to_config persists diffusion fields")
+
+    # ping() against a dead port raises a clear, actionable RuntimeError (no GUI/modal).
+    try:
+        SwarmUIBackend("http://127.0.0.1:1", timeout=1.0).ping()
+        check(False, "ping should raise when no server is listening")
+    except RuntimeError as e:
+        check("SwarmUI" in str(e), "ping raises actionable RuntimeError when server is down")
+
+
 def test_extract_dialog_options() -> None:
    print("extract dialog: options() includes JPEG quality")
    from PySide6.QtWidgets import QApplication
@@ -141,6 +222,10 @@ def test_settings_roundtrip() -> None:
        cfg.cross_model_nms = True
        cfg.nms_iou = 0.55
        cfg.default_threshold = 0.3
+        cfg.diff_url = "http://localhost:9999"
+        cfg.diff_prompt = "test prompt"
+        cfg.diff_steps = 42
+        cfg.diff_denoise = 0.7
        proj = Project.create(Path(d) / "P", name="P")
        proj.update_from_config(cfg)
        proj.save()
@@ -150,6 +235,9 @@ def test_settings_roundtrip() -> None:
        check(cfg2.model_thresholds == {"penis": 0.42}, "model_thresholds persisted")
        check(cfg2.cross_model_nms is True, "cross_model_nms persisted")
        check(abs(cfg2.nms_iou - 0.55) < 1e-9, "nms_iou persisted")
+        check(cfg2.diff_url == "http://localhost:9999", "diff_url persisted")
+        check(cfg2.diff_prompt == "test prompt", "diff_prompt persisted")
+        check(cfg2.diff_steps == 42 and abs(cfg2.diff_denoise - 0.7) < 1e-9, "diff params persisted")
        check(not (proj.root / "project.json.tmp").exists(), "project.json.tmp cleaned up")


@@ -227,6 +315,8 @@ def main() -> int:
    tests = [
        test_cache_atomic_roundtrip,
        test_nms,
+        test_diffusion_restorer,
+        test_restore_dialog_diffusion,
        test_extract_dialog_options,
        test_settings_roundtrip,
        test_mainwindow_filter_and_jump,
@@ -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())