Files
scraed-LanPaint/tests/test_reshape_mask.py
T
scraed 6948b2c766 Retire sampler hyperparameters, add MinStepFrac, and rework the video mask union
Retire Beta/Friction/EarlyStop/InnerThreshold/InnerPatience/MinStepFrac
from all sampler nodes: the widgets are gone, the values are fixed
constants, and old prompts still validate via hidden inputs. A value
sanitizer falls back to defaults on invalid widget values, and the
frontend migrates old workflows by key-based widgets_values mapping
(LAYOUTS table) instead of positional arrays, so kept parameters
(PromptMode, Inpainting_mode) survive on old workflows.

MinStepFrac defaults to 1.0: the inner step size is pinned and the
inner-step count ramps linearly with the remaining noise fraction.

Video mask union: resample nearest-exact first, then a slice-level
sliding max-pool (experimental order) so the mask follows the painted
region's motion through the video.

Update Example_29 with the new output and re-export the workflow JSON.
2026-08-12 17:24:03 +08:00

122 lines
4.6 KiB
Python

import importlib
import sys
import types
import pytest
import torch
def _repeat_to_batch_size(tensor: torch.Tensor, batch_size: int) -> torch.Tensor:
if tensor.shape[0] == batch_size:
return tensor
if tensor.shape[0] == 1:
return tensor.repeat((batch_size,) + (1,) * (tensor.ndim - 1))
repeats = (batch_size + tensor.shape[0] - 1) // tensor.shape[0]
return tensor.repeat((repeats,) + (1,) * (tensor.ndim - 1))[:batch_size]
def _import_nodes(monkeypatch, comfyui_version: str):
comfy_mod = types.ModuleType("comfy")
comfy_mod.__path__ = []
comfy_utils_mod = types.ModuleType("comfy.utils")
comfy_utils_mod.repeat_to_batch_size = _repeat_to_batch_size
comfy_samplers_mod = types.ModuleType("comfy.samplers")
class DummyKSAMPLER: ...
comfy_samplers_mod.KSAMPLER = DummyKSAMPLER
comfy_model_base_mod = types.ModuleType("comfy.model_base")
class ModelType:
FLUX = "FLUX"
FLOW = "FLOW"
class WAN22: ...
comfy_model_base_mod.ModelType = ModelType
comfy_model_base_mod.WAN22 = WAN22
comfyui_version_mod = types.ModuleType("comfyui_version")
comfyui_version_mod.__version__ = comfyui_version
comfy_mod.utils = comfy_utils_mod
comfy_mod.samplers = comfy_samplers_mod
comfy_mod.model_base = comfy_model_base_mod
monkeypatch.setitem(sys.modules, "comfy", comfy_mod)
monkeypatch.setitem(sys.modules, "comfy.utils", comfy_utils_mod)
monkeypatch.setitem(sys.modules, "comfy.samplers", comfy_samplers_mod)
monkeypatch.setitem(sys.modules, "comfy.model_base", comfy_model_base_mod)
monkeypatch.setitem(sys.modules, "nodes", types.ModuleType("nodes"))
monkeypatch.setitem(sys.modules, "latent_preview", types.ModuleType("latent_preview"))
monkeypatch.setitem(sys.modules, "comfyui_version", comfyui_version_mod)
sys.modules.pop("src.LanPaint.nodes", None)
return importlib.import_module("src.LanPaint.nodes")
@pytest.mark.parametrize("comfyui_version", ["0.5.0", "0.6.0"])
def test_reshape_mask_accepts_bhw_and_5d_output_shape(monkeypatch, comfyui_version: str) -> None:
lanpaint_nodes = _import_nodes(monkeypatch, comfyui_version)
input_mask = torch.zeros((1, 4, 4))
output_shape = (1, 16, 1, 8, 8)
out = lanpaint_nodes.reshape_mask(input_mask, output_shape, video_inpainting=False)
assert tuple(out.shape) == output_shape
def test_prepare_mask_accepts_hw_and_moves_device(monkeypatch) -> None:
lanpaint_nodes = _import_nodes(monkeypatch, "0.5.0")
input_mask = torch.zeros((4, 4))
output_shape = (2, 3, 8, 8)
out = lanpaint_nodes.prepare_mask(input_mask, output_shape, device=torch.device("cpu"), video_inpainting=False)
assert tuple(out.shape) == output_shape
assert out.device.type == "cpu"
def test_video_union_covers_picked_frames(monkeypatch) -> None:
"""EXPERIMENTAL order (nearest-exact first, slice-level union): a stroke
on a frame that the resample picks survives, and the sliding max spreads
it to neighboring slices."""
lanpaint_nodes = _import_nodes(monkeypatch, "0.6.0")
# 8 frames -> 2 slices; nearest-exact picks frames {2, 6}. Strokes on
# picked frames 2 and 6, on rows/cols the 8->4 grid samples ({1,3,5,7}).
mask = torch.zeros(8, 8, 8)
mask[2, 5, 5] = 1.0
mask[6, 7, 7] = 1.0
out = lanpaint_nodes.reshape_mask(mask, (1, 16, 2, 4, 4), video_inpainting=True)
assert tuple(out.shape) == (1, 16, 2, 4, 4)
# both slices are marked (slice-level union spreads the picked values)
assert out[0, 0, 0].max().item() == 1.0
assert out[0, 0, 1].max().item() == 1.0
# the stroke positions land on the token containing the stroke pixel
assert out[0, 0, 0, 2, 2].item() == 1.0
assert out[0, 0, 1, 3, 3].item() == 1.0
def test_video_union_drops_unpicked_frames(monkeypatch) -> None:
"""EXPERIMENTAL order: a stroke on a frame that no slice picks (e.g.
frame 3, since nearest-exact picks {2, 6}) is silently lost."""
lanpaint_nodes = _import_nodes(monkeypatch, "0.6.0")
mask = torch.zeros(8, 8, 8)
mask[3, 5, 5] = 1.0
out = lanpaint_nodes.reshape_mask(mask, (1, 16, 2, 4, 4), video_inpainting=True)
assert out.max().item() == 0.0
def test_video_union_short_sequence(monkeypatch) -> None:
"""A 3-frame sequence downsamples to one slice; the stroke must survive
if it lands on the frame the resample picks (floor(0.5*3) = frame 1)."""
lanpaint_nodes = _import_nodes(monkeypatch, "0.6.0")
mask = torch.zeros(3, 8, 8)
mask[1, 5, 5] = 1.0
out = lanpaint_nodes.reshape_mask(mask, (1, 16, 1, 4, 4), video_inpainting=True)
assert tuple(out.shape) == (1, 16, 1, 4, 4)
assert out.max().item() == 1.0