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