WildAi
69b14b8a0e
New src/keep_gpu_ops.py: affine capture with the vendored cv2 calls, batched crop/inverse warps (grid_sample, in-bounds compositing of the gray border, uint8 requantization), a batched form of the vendored tensor2img chain, the ParseNet colormap/blur/postprocess, and the batched ComfyUI converters in keep_image_utils.py (one host transfer per batch; the /255 reproduces img2tensor's double rounding bit for bit). Every op carries a measured-tolerance pin in tests/test_gpu_batch_equivalence.py; the uint8-bound warps measure up to 8/255 against cv2's fixed-point sampling on adversarial content - the numbers that drive the tolerance decision D2.
577 lines
24 KiB
Python
577 lines
24 KiB
Python
"""The measured old-vs-new equivalence of the batched GPU ops.
|
|
|
|
The video pipeline redesign (design section 7 of
|
|
``.dev/docs/2026-10-25-keep-gpu-video-pipeline.md``) replaces the per-frame
|
|
cv2 / numpy loops of ``process_image_sequence`` with one batched torch op
|
|
each. Every op is pinned here against the vendored reference it replaces,
|
|
with the tolerance the runs measured - measured, not assumed (the pattern of
|
|
``tests/test_image_op_equivalence.py``).
|
|
|
|
Bound provenance:
|
|
|
|
- the converters, ``tensor2img_batch`` and the parse colormap are the same
|
|
elementwise chains as their per-frame references - pinned **bit-exact**;
|
|
- the warps replace ``cv2.warpAffine``, which samples in 1/32 px fixed point
|
|
(``INTER_TAB_SIZE``) and weights, while torch samples in float32: the
|
|
residual per pixel is bounded by ``local gradient / 64`` (the coordinate
|
|
quantization) plus the fixed-point weight rounding - the tests measure it
|
|
on smooth content and on adversarial noise and pin the measured bound;
|
|
- the parse blur keeps the exact ``cv2.getGaussianKernel`` kernel and the
|
|
``BORDER_REFLECT101`` extension, so its residual against the two cv2 calls
|
|
is float summation order only;
|
|
- the ParseNet forward is pinned batch-invariant (a batched forward must
|
|
return the per-frame pixels).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
|
|
#: The vendored 5-point 512 template (``face_restoration_helper.py:52-53``).
|
|
FACE_TEMPLATE = np.array(
|
|
[[192.98138, 239.94708], [318.90277, 240.1936], [256.63416, 314.01935],
|
|
[201.26117, 371.41043], [313.08905, 371.15118]]
|
|
)
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def gpu_ops():
|
|
return importlib.import_module("ComfyUI_KEEP.src.keep_gpu_ops")
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def image_utils():
|
|
return importlib.import_module("ComfyUI_KEEP.src.keep_image_utils")
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def img_util():
|
|
return importlib.import_module("ComfyUI_KEEP.modules.deps.basicsr.utils.img_util")
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def helper_module():
|
|
return importlib.import_module("ComfyUI_KEEP.modules.deps.facelib.utils.face_restoration_helper")
|
|
|
|
|
|
@pytest.fixture()
|
|
def face_helper(helper_module, monkeypatch):
|
|
"""A vendored ``FaceRestoreHelper`` without its weight-backed models.
|
|
|
|
The detection / parsing initializations are stubbed (the weights are not
|
|
on this machine and the affine / paste math under test does not touch
|
|
them), so the reference methods run the real vendored code.
|
|
"""
|
|
monkeypatch.setattr(helper_module, "init_detection_model", lambda *a, **k: object())
|
|
monkeypatch.setattr(helper_module, "init_parsing_model", lambda *a, **k: object())
|
|
return helper_module.FaceRestoreHelper(
|
|
upscale_factor=1, face_size=512, crop_ratio=(1, 1),
|
|
det_model="retinaface_resnet50", save_ext="png", use_parse=False,
|
|
device="cpu", model_rootpath="unused",
|
|
)
|
|
|
|
|
|
def synthetic_landmarks(num_frames: int, frame_hw: tuple[int, int], seed: int) -> np.ndarray:
|
|
"""One face per frame: the 512 template at a jittered in-frame position."""
|
|
rng = np.random.default_rng(seed)
|
|
height, width = frame_hw
|
|
base_scale = min(height, width) / 512.0 * 0.72
|
|
landmarks = np.zeros((num_frames, 5, 2), dtype=np.float64)
|
|
for i in range(num_frames):
|
|
scale = base_scale + rng.uniform(-0.005, 0.005)
|
|
center = np.array(
|
|
[width * 0.5 + rng.uniform(-8, 8), height * 0.46 + rng.uniform(-6, 6)]
|
|
)
|
|
landmarks[i] = (FACE_TEMPLATE - 256.0) * scale + center
|
|
return landmarks
|
|
|
|
|
|
def smooth_frames_u8(num_frames: int, frame_hw: tuple[int, int], seed: int) -> np.ndarray:
|
|
"""Low-frequency content: the realistic worst case for the warp residual."""
|
|
rng = np.random.default_rng(seed)
|
|
height, width = frame_hw
|
|
small = rng.uniform(0.0, 255.0, size=(num_frames, 12, 22, 3)).astype(np.float32)
|
|
frames = np.stack(
|
|
[cv2.resize(small[i], (width, height), interpolation=cv2.INTER_CUBIC) for i in range(num_frames)]
|
|
)
|
|
return np.clip(frames, 0, 255).astype(np.uint8)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# The batched converters (design sections 7.1, 7.4, 7.8)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_comfy_batch_to_bgr_u8_is_bit_exact_against_the_per_frame_loop(image_utils):
|
|
"""The batch conversion equals the per-frame ``comfy_image_to_cv2`` loop."""
|
|
tensor = torch.rand((4, 16, 24, 3), generator=torch.Generator().manual_seed(7))
|
|
|
|
batched = image_utils.comfy_batch_to_bgr_u8(tensor)
|
|
per_frame = np.stack(
|
|
[image_utils.comfy_image_to_cv2(tensor[i : i + 1]) for i in range(tensor.shape[0])]
|
|
)
|
|
|
|
assert batched.shape == (4, 16, 24, 3)
|
|
assert batched.dtype == np.uint8
|
|
assert np.array_equal(batched, per_frame)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_comfy_batch_to_bgr_u8_matches_the_former_numpy_implementation(image_utils):
|
|
"""The scale / cast / flip reproduce the former numpy + cv2 path exactly."""
|
|
tensor = torch.rand((3, 9, 11, 3), generator=torch.Generator().manual_seed(13))
|
|
|
|
batched = image_utils.comfy_batch_to_bgr_u8(tensor)
|
|
|
|
reference = np.stack(
|
|
[
|
|
cv2.cvtColor((tensor[i].numpy() * 255).astype(np.uint8), cv2.COLOR_RGB2BGR)
|
|
for i in range(tensor.shape[0])
|
|
]
|
|
)
|
|
assert np.array_equal(batched, reference)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_bgr_u8_batch_to_comfy_is_bit_exact_against_the_per_frame_loop(image_utils):
|
|
"""The batch conversion equals the per-frame ``cv2_to_comfy_image`` loop.
|
|
|
|
Both accepted inputs are covered: the numpy stack of the per-frame paste
|
|
outputs and the torch uint8 tensor the batched blend produces on the
|
|
device.
|
|
"""
|
|
rng = np.random.default_rng(11)
|
|
frames = rng.integers(0, 256, size=(5, 12, 18, 3), dtype=np.uint8)
|
|
|
|
from_numpy = image_utils.bgr_u8_batch_to_comfy(frames)
|
|
from_torch = image_utils.bgr_u8_batch_to_comfy(torch.from_numpy(frames.copy()))
|
|
per_frame = torch.cat([image_utils.cv2_to_comfy_image(frames[i]) for i in range(5)], dim=0)
|
|
|
|
assert from_numpy.shape == (5, 12, 18, 3)
|
|
assert from_numpy.dtype == torch.float32
|
|
assert torch.equal(from_numpy, from_torch)
|
|
assert torch.equal(from_numpy, per_frame)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_bgr_u8_batch_to_normalized_is_bit_exact_against_img2tensor(image_utils, img_util):
|
|
"""The batched crop tail equals ``img2tensor`` + the normalize replacement.
|
|
|
|
The reference is the per-frame chain of ``keep_processor.py:311-313``:
|
|
``img2tensor(face / 255., bgr2rgb=True, float32=True)`` (the numpy
|
|
float64 division, then the float32 cast) and the in-place
|
|
``(t - 0.5) / 0.5``. The float64 division of the batch form reproduces
|
|
the double rounding bit for bit.
|
|
"""
|
|
rng = np.random.default_rng(23)
|
|
faces = rng.integers(0, 256, size=(3, 3, 48, 64), dtype=np.uint8)
|
|
|
|
batched = image_utils.bgr_u8_batch_to_normalized(torch.from_numpy(faces))
|
|
reference = torch.stack(
|
|
[
|
|
img_util.img2tensor(
|
|
faces[i].transpose(1, 2, 0) / 255.0, bgr2rgb=True, float32=True
|
|
).sub_(0.5).div_(0.5)
|
|
for i in range(3)
|
|
]
|
|
)
|
|
|
|
assert batched.dtype == torch.float32
|
|
assert torch.equal(batched, reference)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Affine capture (design section 7.3) - identity by construction
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_estimate_affines_match_the_vendored_align_warp_face(gpu_ops, face_helper):
|
|
"""The captured matrices are the ones ``align_warp_face`` appends."""
|
|
landmarks = synthetic_landmarks(3, (720, 1280), seed=5)
|
|
# the vendored method reads ``input_img`` for the crop warp (any image
|
|
# works - only the matrices are under test here)
|
|
face_helper.input_img = np.zeros((128, 128, 3), dtype=np.uint8)
|
|
face_helper.all_landmarks_5 = [landmarks[i] for i in range(3)]
|
|
face_helper.align_warp_face()
|
|
captured = gpu_ops.estimate_affines(landmarks, face_helper.face_template)
|
|
|
|
assert len(captured) == 3
|
|
for reference, matrix in zip(face_helper.affine_matrices, captured):
|
|
assert matrix is not None
|
|
assert matrix.dtype == np.float64
|
|
assert np.array_equal(reference, matrix)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_estimate_affines_preserves_the_lmeds_none(gpu_ops, face_helper):
|
|
"""Degenerate landmarks keep the ``None`` of the vendored ``:315-319`` branch."""
|
|
degenerate = np.zeros((5, 2), dtype=np.float64)
|
|
|
|
face_helper.input_img = np.zeros((128, 128, 3), dtype=np.uint8)
|
|
face_helper.all_landmarks_5 = [degenerate]
|
|
face_helper.align_warp_face()
|
|
captured = gpu_ops.estimate_affines([degenerate], face_helper.face_template)
|
|
|
|
assert captured == [None]
|
|
assert face_helper.affine_matrices == [None]
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_inverse_affines_match_the_vendored_get_inverse_affine(gpu_ops, face_helper):
|
|
"""The inverse warps equal ``get_inverse_affine`` for upscale 1 and 2."""
|
|
landmarks = synthetic_landmarks(3, (360, 640), seed=9)
|
|
affines = gpu_ops.estimate_affines(landmarks, face_helper.face_template)
|
|
|
|
for upscale in (1, 2):
|
|
face_helper.upscale_factor = upscale
|
|
face_helper.affine_matrices = list(affines)
|
|
face_helper.get_inverse_affine(None)
|
|
captured = gpu_ops.inverse_affines(affines, upscale_factor=upscale)
|
|
for reference, matrix in zip(face_helper.inverse_affine_matrices, captured):
|
|
assert np.array_equal(reference, matrix)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# The batched warps (design sections 7.4 and 7.7)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_warp_affine_batch_is_exact_for_identity_and_integer_shift(gpu_ops):
|
|
"""The normalized theta math lands cv2-exactly on aligned geometries.
|
|
|
|
An identity matrix and a one-pixel integer translation sample exact pixel
|
|
centers, where cv2's 1/32 px fixed point and the torch float grid agree
|
|
to the last bit - so the quantized outputs must be pixel-identical.
|
|
"""
|
|
rng = np.random.default_rng(31)
|
|
frames = rng.integers(0, 256, size=(2, 3, 48, 64), dtype=np.uint8)
|
|
source = torch.from_numpy(frames).permute(0, 3, 1, 2).float()
|
|
identity = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])
|
|
shift = np.array([[1.0, 0.0, 1.0], [0.0, 1.0, 0.0]])
|
|
|
|
warped = gpu_ops.warp_affine_batch(source, [identity, shift], (64, 48))
|
|
warped_u8 = warped.round().clamp_(0, 255).to(torch.uint8).permute(0, 2, 3, 1).numpy()
|
|
reference = np.stack(
|
|
[cv2.warpAffine(frames[i], matrix, (64, 48)) for i, matrix in enumerate((identity, shift))]
|
|
)
|
|
|
|
assert np.array_equal(warped_u8, reference)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_warp_affine_batch_rejects_mismatched_and_none_matrices(gpu_ops):
|
|
"""A None matrix or a wrong count fails loudly (the no-face frames are
|
|
filtered by the caller, like today's zero-tensor branch)."""
|
|
source = torch.rand((2, 3, 16, 16))
|
|
with pytest.raises(ValueError):
|
|
gpu_ops.warp_affine_batch(source, [np.eye(2, 3)], (16, 16))
|
|
with pytest.raises(ValueError):
|
|
gpu_ops.warp_affine_batch(source, [np.eye(2, 3), None], (16, 16))
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_crop_faces_batch_matches_cv2_on_smooth_content(gpu_ops):
|
|
"""The batched crop vs the vendored warp: smooth content.
|
|
|
|
Measured over 4 frames of 1280x720 low-frequency content with jittered
|
|
landmark geometries: the max residual is **1/255** (2.1% of the pixels
|
|
differ at all) - cv2 samples in 1/32 px fixed point, torch in float32.
|
|
"""
|
|
frames = smooth_frames_u8(4, (720, 1280), seed=3)
|
|
landmarks = synthetic_landmarks(4, (720, 1280), seed=4)
|
|
affines = gpu_ops.estimate_affines(landmarks, FACE_TEMPLATE)
|
|
source = torch.from_numpy(frames).permute(0, 3, 1, 2).float()
|
|
|
|
cropped = gpu_ops.crop_faces_batch(source, affines).permute(0, 2, 3, 1).numpy()
|
|
reference = np.stack(
|
|
[
|
|
cv2.warpAffine(
|
|
frames[i], affines[i], (512, 512),
|
|
borderMode=cv2.BORDER_CONSTANT, borderValue=(135, 133, 132),
|
|
)
|
|
for i in range(4)
|
|
]
|
|
)
|
|
|
|
assert cropped.dtype == np.uint8
|
|
diff = np.abs(reference.astype(int) - cropped.astype(int))
|
|
assert diff.max() <= 1, f"the batched crop drifted {diff.max()}/255 from cv2"
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_crop_faces_batch_matches_cv2_on_noise_content(gpu_ops):
|
|
"""The batched crop vs the vendored warp: adversarial content.
|
|
|
|
Pure per-pixel noise is the worst case for cv2's fixed-point sampling
|
|
(gradients up to 255 per pixel). Measured: max **6/255** (p99.9 = 4,
|
|
62% of the pixels differ by at least 1). The pinned bound is the
|
|
quantization analysis ceiling of ~2 x 255 / 64.
|
|
"""
|
|
rng = np.random.default_rng(6)
|
|
frames = rng.integers(0, 256, size=(4, 720, 1280, 3), dtype=np.uint8)
|
|
landmarks = synthetic_landmarks(4, (720, 1280), seed=7)
|
|
affines = gpu_ops.estimate_affines(landmarks, FACE_TEMPLATE)
|
|
source = torch.from_numpy(frames).permute(0, 3, 1, 2).float()
|
|
|
|
cropped = gpu_ops.crop_faces_batch(source, affines).permute(0, 2, 3, 1).numpy()
|
|
reference = np.stack(
|
|
[
|
|
cv2.warpAffine(
|
|
frames[i], affines[i], (512, 512),
|
|
borderMode=cv2.BORDER_CONSTANT, borderValue=(135, 133, 132),
|
|
)
|
|
for i in range(4)
|
|
]
|
|
)
|
|
|
|
diff = np.abs(reference.astype(int) - cropped.astype(int))
|
|
assert diff.max() <= 8, f"the batched crop drifted {diff.max()}/255 from cv2"
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_crop_faces_batch_composites_the_vendored_gray_border(gpu_ops):
|
|
"""A fully out-of-frame face returns the exact gray constant everywhere.
|
|
|
|
The border composite must reproduce cv2's ``BORDER_CONSTANT`` value
|
|
(135, 133, 132) - not the black a zero-padded ``grid_sample`` gives.
|
|
"""
|
|
outside = FACE_TEMPLATE + np.array([2000.0, 2000.0]) # the face sits far outside
|
|
affines = gpu_ops.estimate_affines([outside], FACE_TEMPLATE)
|
|
frame = torch.full((1, 3, 64, 64), 200.0)
|
|
|
|
cropped = gpu_ops.crop_faces_batch(frame, affines).permute(0, 2, 3, 1).numpy()
|
|
reference = cv2.warpAffine(
|
|
np.full((64, 64, 3), 200, dtype=np.uint8), affines[0], (512, 512),
|
|
borderMode=cv2.BORDER_CONSTANT, borderValue=(135, 133, 132),
|
|
)
|
|
|
|
assert np.array_equal(cropped[0], reference)
|
|
assert np.all(cropped[0][:, :, 0] == 135)
|
|
assert np.all(cropped[0][:, :, 1] == 133)
|
|
assert np.all(cropped[0][:, :, 2] == 132)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_warp_affine_batch_matches_cv2_for_the_inverse_paste_warp(gpu_ops, face_helper):
|
|
"""The batched inverse warp of the faces vs the vendored paste warp.
|
|
|
|
The paste warps the uint8 ``tensor2img`` faces back onto the (upscaled)
|
|
frame with ``cv2.warpAffine(face, inverse_affine, (w_up, h_up))``
|
|
(``face_restoration_helper.py:406``, the default black border). Both
|
|
upscale factors are measured at the production 720p canvas: the max
|
|
residual is **7/255** for upscale 1 and 2 (18% of the pixels differ).
|
|
"""
|
|
rng = np.random.default_rng(8)
|
|
faces = rng.integers(0, 256, size=(3, 3, 512, 512), dtype=np.uint8)
|
|
landmarks = synthetic_landmarks(3, (720, 1280), seed=10)
|
|
affines = gpu_ops.estimate_affines(landmarks, face_helper.face_template)
|
|
source = torch.from_numpy(faces).float()
|
|
|
|
for upscale in (1, 2):
|
|
inverses = gpu_ops.inverse_affines(affines, upscale_factor=upscale)
|
|
out_size = (1280 * upscale, 720 * upscale) # the cv2 (w_up, h_up) order
|
|
warped = (
|
|
gpu_ops.warp_affine_batch(source, inverses, out_size)
|
|
.round().clamp_(0, 255).to(torch.uint8).permute(0, 2, 3, 1).numpy()
|
|
)
|
|
faces_hwc = faces.transpose(0, 2, 3, 1) # cv2 wants (H, W, C)
|
|
reference = np.stack(
|
|
[cv2.warpAffine(faces_hwc[i], inverses[i], out_size) for i in range(3)]
|
|
)
|
|
diff = np.abs(reference.astype(int) - warped.astype(int))
|
|
assert diff.max() <= 8, (
|
|
f"upscale {upscale}: the batched inverse warp drifted {diff.max()}/255 from cv2"
|
|
)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_warp_affine_batch_matches_cv2_for_the_float_mask_warp(gpu_ops):
|
|
"""The float warp (the parse mask path) matches the cv2 float warp.
|
|
|
|
The parse mask is warped as float32 with ``INTER_LINEAR`` semantics
|
|
(``face_restoration_helper.py:477``). cv2's fixed-point path quantizes
|
|
the affine coefficients to 1/1024, so the position error grows with the
|
|
output coordinate - measured at the production 720p canvas:
|
|
**0.0257** mask units on per-pixel-independent content (the worst case)
|
|
and **0.0075** on the smoothed binary mask the pipeline actually warps.
|
|
"""
|
|
rng = np.random.default_rng(12)
|
|
uniform = rng.uniform(0.0, 1.0, size=(3, 1, 512, 512)).astype(np.float32)
|
|
binary = (rng.uniform(size=(3, 1, 512, 512)) > 0.5).astype(np.float32) * 255.0
|
|
for i in range(3): # the production shape: the colormap output, blurred
|
|
blurred = cv2.GaussianBlur(binary[i, 0], (101, 101), 11)
|
|
binary[i, 0] = cv2.GaussianBlur(blurred, (101, 101), 11)
|
|
landmarks = synthetic_landmarks(3, (720, 1280), seed=13)
|
|
affines = gpu_ops.estimate_affines(landmarks, FACE_TEMPLATE)
|
|
inverses = gpu_ops.inverse_affines(affines, upscale_factor=1.0)
|
|
|
|
reference = np.stack(
|
|
[cv2.warpAffine(uniform[i, 0], inverses[i], (1280, 720)) for i in range(3)]
|
|
)[:, None] # (N, 1, H, W) like the batched form
|
|
warped = gpu_ops.warp_affine_batch(torch.from_numpy(uniform), inverses, (1280, 720)).numpy()
|
|
diff = np.abs(reference - warped).max()
|
|
assert diff <= 0.03, f"the float mask warp drifted {diff:.2e} mask units from cv2"
|
|
|
|
reference = np.stack(
|
|
[cv2.warpAffine(binary[i, 0] / 255.0, inverses[i], (1280, 720)) for i in range(3)]
|
|
)[:, None]
|
|
warped = gpu_ops.warp_affine_batch(torch.from_numpy(binary / 255.0), inverses, (1280, 720)).numpy()
|
|
diff = np.abs(reference - warped).max()
|
|
assert diff <= 0.01, f"the smoothed mask warp drifted {diff:.2e} mask units from cv2"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# tensor2img batch (design section 7.6) - bit-exact
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_tensor2img_batch_is_bit_exact_against_the_vendored_loop(gpu_ops, img_util):
|
|
"""The batched ``tensor2img`` equals the per-frame vendored call.
|
|
|
|
The values overshoot (-1, 1) on purpose so the clamp branch is covered;
|
|
the torch uint8 cast of the rounded values is exact, like the numpy
|
|
``astype`` of the vendored code.
|
|
"""
|
|
tensor = torch.randn((5, 3, 32, 32), generator=torch.Generator().manual_seed(17)) * 1.5
|
|
|
|
batched = gpu_ops.tensor2img_batch(tensor, min_max=(-1, 1))
|
|
reference = np.stack(
|
|
[img_util.tensor2img(tensor[i], rgb2bgr=True, min_max=(-1, 1)) for i in range(5)]
|
|
)
|
|
|
|
assert batched.shape == (5, 32, 32, 3)
|
|
assert batched.dtype == np.uint8
|
|
assert np.array_equal(batched, reference)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_tensor2img_u8_batch_stays_on_the_device_and_matches_the_numpy_form(gpu_ops):
|
|
"""The paste-stage face stack is the same quantization, without the transfer."""
|
|
tensor = torch.tanh(torch.randn((3, 3, 16, 16), generator=torch.Generator().manual_seed(19)))
|
|
|
|
u8 = gpu_ops.tensor2img_u8_batch(tensor, min_max=(-1, 1))
|
|
as_numpy = gpu_ops.tensor2img_batch(tensor, min_max=(-1, 1))
|
|
|
|
assert u8.dtype == torch.uint8
|
|
assert torch.equal(u8, torch.from_numpy(as_numpy))
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# The parse postprocess (design section 7.7, step 4)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_parse_colormap_batch_matches_the_vendored_assignment_loop(gpu_ops, helper_module):
|
|
"""The argmax + gather writes the values of the vendored per-class loop."""
|
|
logits = torch.randn((2, 19, 32, 32), generator=torch.Generator().manual_seed(21))
|
|
|
|
batched = gpu_ops.parse_colormap_batch(logits).squeeze(1).numpy()
|
|
reference = []
|
|
for i in range(2):
|
|
out_parse = logits[i].argmax(dim=0).cpu().numpy()
|
|
raw = np.zeros(out_parse.shape, dtype=np.float32)
|
|
for map_idx, color_val in enumerate(gpu_ops.PARSE_MASK_COLORMAP):
|
|
raw[out_parse == map_idx] = color_val
|
|
reference.append(raw)
|
|
reference = np.stack(reference)
|
|
|
|
assert np.array_equal(batched, reference)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_gaussian_blur_matches_two_cv2_gaussian_blurs(gpu_ops):
|
|
"""The batched two-pass blur vs the vendored ``GaussianBlur`` pair.
|
|
|
|
The binary 0/255 masks are the realistic parse output, the uniform
|
|
random masks the adversarial one; both keep the exact kernel and the
|
|
reflect101 extension, so the residual is float summation order only -
|
|
measured **1.7e-4** mask units (0..255 scale) at worst. The pinned
|
|
bound still trips on any structural change (a wrong border mode would
|
|
show up in the tens).
|
|
"""
|
|
rng = np.random.default_rng(3)
|
|
binary = (rng.uniform(size=(2, 1, 512, 512)) > 0.5).astype(np.float32) * 255.0
|
|
uniform = rng.uniform(0.0, 255.0, size=(2, 1, 512, 512)).astype(np.float32)
|
|
|
|
for masks_np in (binary, uniform):
|
|
masks = torch.from_numpy(masks_np)
|
|
blurred = gpu_ops.gaussian_blur_2d_reflect101(masks, 101, 11.0)
|
|
blurred = gpu_ops.gaussian_blur_2d_reflect101(blurred, 101, 11.0)
|
|
reference = np.stack(
|
|
[
|
|
cv2.GaussianBlur(
|
|
cv2.GaussianBlur(masks_np[i, 0], (101, 101), 11), (101, 101), 11
|
|
)
|
|
for i in range(2)
|
|
]
|
|
)
|
|
diff = np.abs(reference - blurred[:, 0].numpy()).max()
|
|
assert diff <= 0.01, f"the batched blur drifted {diff:.4f} mask units from cv2"
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_parse_soft_masks_match_the_vendored_postprocess(gpu_ops):
|
|
"""The full batched postprocess vs the vendored ``:463-472`` chain.
|
|
|
|
Measured drift: **8.9e-7** in [0, 1] mask units (~2e-4 on the 0..255
|
|
scale, the float summation order of the two blur passes).
|
|
"""
|
|
logits = torch.randn((2, 19, 512, 512), generator=torch.Generator().manual_seed(27))
|
|
|
|
soft = gpu_ops.parse_soft_masks(logits).numpy()[:, 0] # (N, H, W) like the reference
|
|
reference = []
|
|
for i in range(2):
|
|
out_parse = logits[i].argmax(dim=0).cpu().numpy()
|
|
raw = np.zeros(out_parse.shape, dtype=np.float32)
|
|
for map_idx, color_val in enumerate(gpu_ops.PARSE_MASK_COLORMAP):
|
|
raw[out_parse == map_idx] = color_val
|
|
blurred = cv2.GaussianBlur(raw, (101, 101), 11)
|
|
blurred = cv2.GaussianBlur(blurred, (101, 101), 11)
|
|
thres = 10
|
|
blurred[:thres, :] = 0
|
|
blurred[-thres:, :] = 0
|
|
blurred[:, :thres] = 0
|
|
blurred[:, -thres:] = 0
|
|
reference.append(blurred / 255.0)
|
|
reference = np.stack(reference)
|
|
|
|
diff = np.abs(reference - soft).max()
|
|
assert diff <= 1.0 / 255.0, f"the soft masks drifted {diff:.2e} from the vendored chain"
|
|
# the border zeroing is exact in both
|
|
assert np.all(soft[0, :10, :] == 0.0) and np.all(reference[0, :10, :] == 0.0)
|
|
assert np.all(soft[0, :, -10:] == 0.0) and np.all(reference[0, :, -10:] == 0.0)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_parse_forward_chunks_are_batch_invariant():
|
|
"""A chunked ParseNet forward returns the per-frame pixels.
|
|
|
|
The batched paste forwards the parse net in chunks (the VRAM-aware cap
|
|
of the design); batch independence in eval mode is the property that
|
|
makes the chunk size a pure performance knob. ``in_size`` 256 keeps the
|
|
probe cheap - the claim is resolution-independent (BatchNorm in eval
|
|
mode uses running stats, so no cross-sample path exists).
|
|
"""
|
|
parsing = importlib.import_module("ComfyUI_KEEP.modules.deps.facelib.parsing.parsenet")
|
|
net = parsing.ParseNet(in_size=256, out_size=256).eval()
|
|
torch.manual_seed(29)
|
|
inputs = torch.randn(8, 3, 256, 256)
|
|
|
|
with torch.no_grad():
|
|
per_frame = torch.cat([net(inputs[i : i + 1])[0] for i in range(8)], dim=0)
|
|
chunked = torch.cat([net(inputs[j : j + 4])[0] for j in range(0, 8, 4)], dim=0)
|
|
whole = net(inputs)[0]
|
|
|
|
assert torch.equal(per_frame, chunked)
|
|
assert torch.equal(per_frame, whole)
|