WildAi
e2fdbbf453
The round-4 fused crop loops re-derived each frame's detection image with one single-row converter call per frame - N host sync points where the round-3 crop did one batch. The grouped form converts the chunk's gathered rows in one elementwise pass (bitwise the per-row rows) and is bounded by the caller's chunk; the paste-fallback row form stays. Pins guard the conversion count. The paste stage itself is proven byte-equivalent to 32e6628 (dispatch, cores, tails); timed CPU stubs show no superlinear term in either tree - the real-box burn needs the owner's log to attribute (the device-mode verdict under real VRAM pressure is the only paste-stage quantity this diff can move).
880 lines
33 KiB
Python
880 lines
33 KiB
Python
"""P3.2 - the face processor of ``src/keep_processor.py``.
|
|
|
|
The tests drive the processor with **stubs** (a stub face helper and a stub
|
|
KEEP net) and assert the things the processor is responsible for:
|
|
|
|
* the chunk math of the sequence path (``max_clip_length``),
|
|
* the paste-back (the output frame count and shape),
|
|
* the ``tiled_scale`` wrapper of the optional ComfyUI upscalers,
|
|
* the **console routing**: the four steps are logged and the console writes
|
|
of the vendored facelib reach the logger of the pack.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import importlib
|
|
import logging
|
|
import types
|
|
import weakref
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def processor_module():
|
|
return importlib.import_module("ComfyUI_KEEP.src.keep_processor")
|
|
|
|
|
|
class _StubKeepNet:
|
|
"""A KEEP net that records the clips it is called with."""
|
|
|
|
def __init__(self) -> None:
|
|
self.calls: list[tuple] = []
|
|
|
|
def __call__(self, frames, need_upscale=False):
|
|
self.calls.append((tuple(frames.shape), need_upscale))
|
|
return torch.zeros_like(frames)
|
|
|
|
|
|
#: The vendored 5-point 512 template (``face_restoration_helper.py:52-53``).
|
|
FACE_TEMPLATE_512 = np.array(
|
|
[[192.98138, 239.94708], [318.90277, 240.1936], [256.63416, 314.01935],
|
|
[201.26117, 371.41043], [313.08905, 371.15118]]
|
|
)
|
|
|
|
|
|
class _StubFaceHelper:
|
|
"""The facelib helper, reduced to the calls the processor makes."""
|
|
|
|
def __init__(self) -> None:
|
|
self.device = torch.device("cpu")
|
|
self.upscale_factor = 1.0
|
|
self.is_gray = False
|
|
self.cropped_faces: list = []
|
|
self.restored_faces: list = []
|
|
self.all_landmarks_5: list = []
|
|
self.affine_matrices: list = []
|
|
self.input_img = None
|
|
# the template the vendored helper estimates the crop affines with
|
|
self.face_template = FACE_TEMPLATE_512
|
|
self.calls: list[str] = []
|
|
|
|
def clean_all(self):
|
|
self.calls.append("clean_all")
|
|
self.cropped_faces = []
|
|
self.restored_faces = []
|
|
|
|
def read_image(self, image):
|
|
self.calls.append("read_image")
|
|
# the vendored helper keeps the image it read and upscales it when
|
|
# min(h, w) < 512 (``face_restoration_helper.py:179-184``) - the stub
|
|
# reproduces the upscale, or the paste canvas of the small frames
|
|
# would stay at the raw size (the round-1 defect).
|
|
self.input_img = image
|
|
if min(image.shape[:2]) < 512:
|
|
factor = 512.0 / min(image.shape[:2])
|
|
self.input_img = cv2.resize(
|
|
image, (0, 0), fx=factor, fy=factor, interpolation=cv2.INTER_LINEAR)
|
|
|
|
def align_warp_face(self):
|
|
self.calls.append("align_warp_face")
|
|
|
|
def add_restored_face(self, face, **kwargs):
|
|
self.calls.append("add_restored_face")
|
|
self.restored_faces = [face]
|
|
|
|
def get_inverse_affine(self, *args, **kwargs):
|
|
self.calls.append("get_inverse_affine")
|
|
|
|
def paste_faces_to_input_image(self, *args, **kwargs):
|
|
self.calls.append("paste_faces_to_input_image")
|
|
return np.zeros((512, 512, 3), dtype=np.uint8)
|
|
|
|
|
|
def _make_pack(keep_net, face_helper):
|
|
"""A model pack the processor can consume."""
|
|
return types.SimpleNamespace(
|
|
keep_net=keep_net,
|
|
face_helper=face_helper,
|
|
bg_upscale_model=None,
|
|
face_upscale_model=None,
|
|
device=torch.device("cpu"),
|
|
model_type_str="KEEP",
|
|
)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_the_sequence_is_split_into_clips_of_max_clip_length(processor_module):
|
|
"""The chunk math of the sequence path (``max_clip_length``)."""
|
|
keep_net = _StubKeepNet()
|
|
face_helper = _StubFaceHelper()
|
|
processor = processor_module.KEEPFaceProcessor(_make_pack(keep_net, face_helper))
|
|
|
|
frames = torch.rand((4, 16, 16, 3), dtype=torch.float32)
|
|
output = processor.process_image_sequence(
|
|
image_sequence_tensor=frames,
|
|
final_upscale_factor=1.0,
|
|
has_aligned_frames=True,
|
|
only_center_face=True,
|
|
draw_box=False,
|
|
max_clip_length=2,
|
|
)
|
|
|
|
# 4 frames, clips of 2 -> two calls of two frames each
|
|
assert len(keep_net.calls) == 2, keep_net.calls
|
|
for shape, need_upscale in keep_net.calls:
|
|
assert shape == (1, 2, 3, 512, 512), shape
|
|
assert need_upscale is False
|
|
# the paste-back produced one frame per input frame
|
|
assert tuple(output.shape) == (4, 512, 512, 3), output.shape
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_an_empty_sequence_is_returned_unchanged(processor_module):
|
|
"""No frames -> no work, the input is returned."""
|
|
keep_net = _StubKeepNet()
|
|
processor = processor_module.KEEPFaceProcessor(
|
|
_make_pack(keep_net, _StubFaceHelper())
|
|
)
|
|
frames = torch.zeros((0, 16, 16, 3), dtype=torch.float32)
|
|
output = processor.process_image_sequence(
|
|
image_sequence_tensor=frames,
|
|
final_upscale_factor=1.0,
|
|
has_aligned_frames=True,
|
|
only_center_face=True,
|
|
draw_box=False,
|
|
max_clip_length=20,
|
|
)
|
|
assert output is frames
|
|
assert keep_net.calls == []
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_the_upscaler_goes_through_tiled_scale(processor_module):
|
|
"""The optional ComfyUI upscaler is called through ``tiled_scale``."""
|
|
calls: list[dict] = []
|
|
|
|
class _StubUpscaleModel:
|
|
scale = 2.0
|
|
|
|
def __init__(self):
|
|
self.model = self._model
|
|
|
|
def _model(self, samples):
|
|
calls.append({"shape": tuple(samples.shape)})
|
|
return samples * 2.0
|
|
|
|
processor = processor_module.KEEPFaceProcessor(
|
|
_make_pack(_StubKeepNet(), _StubFaceHelper())
|
|
)
|
|
image = np.zeros((8, 8, 3), dtype=np.uint8)
|
|
result = processor._run_upscaler(_StubUpscaleModel(), image)
|
|
|
|
assert len(calls) == 1, calls
|
|
assert calls[0]["shape"] == (1, 3, 8, 8), calls
|
|
assert result.shape == (8, 8, 3), result.shape
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_without_an_upscaler_the_image_is_untouched(processor_module):
|
|
"""No upscaler -> the image is returned unchanged."""
|
|
processor = processor_module.KEEPFaceProcessor(
|
|
_make_pack(_StubKeepNet(), _StubFaceHelper())
|
|
)
|
|
image = np.full((4, 4, 3), 7, dtype=np.uint8)
|
|
assert processor._run_upscaler(None, image) is image
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_thirty_frames_are_split_into_clips_of_twenty(processor_module):
|
|
"""The chunk math of the sequence path for 30 frames / ``max_clip_length=20``."""
|
|
keep_net = _StubKeepNet()
|
|
face_helper = _StubFaceHelper()
|
|
processor = processor_module.KEEPFaceProcessor(_make_pack(keep_net, face_helper))
|
|
|
|
frames = torch.rand((30, 16, 16, 3), dtype=torch.float32)
|
|
output = processor.process_image_sequence(
|
|
image_sequence_tensor=frames,
|
|
final_upscale_factor=1.0,
|
|
has_aligned_frames=True,
|
|
only_center_face=True,
|
|
draw_box=False,
|
|
max_clip_length=20,
|
|
)
|
|
|
|
# 30 frames, clips of 20 -> one call of 20 frames and one of 10
|
|
assert len(keep_net.calls) == 2, keep_net.calls
|
|
shapes = [shape for shape, _ in keep_net.calls]
|
|
assert shapes == [(1, 20, 3, 512, 512), (1, 10, 3, 512, 512)], shapes
|
|
assert tuple(output.shape) == (30, 512, 512, 3), output.shape
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# P4.6 - the VRAM is released through ComfyUI
|
|
# ---------------------------------------------------------------------------
|
|
@pytest.mark.unit
|
|
def test_the_vram_is_released_through_comfyui(processor_module):
|
|
"""``_soft_empty_cache`` goes through ``model_management`` (P4.6)."""
|
|
model_management = importlib.import_module("comfy.model_management")
|
|
model_management.calls.clear()
|
|
|
|
assert processor_module._soft_empty_cache() is True
|
|
assert ("soft_empty_cache",) in model_management.calls, model_management.calls
|
|
|
|
|
|
@pytest.mark.mock_integration
|
|
def test_the_sequence_releases_the_vram_at_every_chunk_boundary(
|
|
processor_module, monkeypatch
|
|
):
|
|
"""P4.6: every chunk boundary releases the cached VRAM.
|
|
|
|
30 frames with ``max_clip_length=20`` are two chunks - the release must
|
|
happen after each of them (the next chunk is allocated right after).
|
|
"""
|
|
events: list[str] = []
|
|
|
|
class _RecordingKeepNet(_StubKeepNet):
|
|
def __call__(self, frames, need_upscale=False):
|
|
events.append("net")
|
|
return super().__call__(frames, need_upscale=need_upscale)
|
|
|
|
monkeypatch.setattr(
|
|
processor_module, "_soft_empty_cache", lambda: events.append("empty_cache")
|
|
)
|
|
|
|
processor = processor_module.KEEPFaceProcessor(
|
|
_make_pack(_RecordingKeepNet(), _StubFaceHelper())
|
|
)
|
|
frames = torch.rand((30, 16, 16, 3), dtype=torch.float32)
|
|
processor.process_image_sequence(
|
|
image_sequence_tensor=frames,
|
|
final_upscale_factor=1.0,
|
|
has_aligned_frames=True,
|
|
only_center_face=True,
|
|
draw_box=False,
|
|
max_clip_length=20,
|
|
)
|
|
|
|
net_positions = [index for index, event in enumerate(events) if event == "net"]
|
|
assert len(net_positions) == 2, events
|
|
for position in net_positions:
|
|
assert events[position + 1] == "empty_cache", (
|
|
f"the chunk boundary did not release the VRAM: {events}"
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# D4 - the console routing of the processor
|
|
# ---------------------------------------------------------------------------
|
|
class _DetectingFaceHelper(_StubFaceHelper):
|
|
"""A face helper that **detects** a face - and prints like the facelib.
|
|
|
|
The vendored helper writes to the console (``print()``) while it reads
|
|
the image, while it detects and while it pastes a face back: the
|
|
processor must route those writes to the logger of the pack.
|
|
|
|
The zero landmarks of the stub are a degenerate set - the crop of the
|
|
redesigned path estimates the affine with the real vendored call, and
|
|
the ``LMEDS`` failure of that call is what prints the warning of the
|
|
facelib.
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.landmarks = np.zeros(10, dtype=np.float32)
|
|
|
|
def get_face_landmarks_5(self, **kwargs):
|
|
self.calls.append("get_face_landmarks_5")
|
|
print("model_rootpath is None. Using default relative path")
|
|
self.all_landmarks_5 = [self.landmarks]
|
|
return 1
|
|
|
|
def paste_faces_to_input_image(self, *args, **kwargs):
|
|
super().paste_faces_to_input_image(*args, **kwargs)
|
|
print("Warning: Mismatch between restored_faces and inverse_affine_matrices.")
|
|
return np.zeros((512, 512, 3), dtype=np.uint8)
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def _pack_records():
|
|
"""The records the logger of the pack receives while the block runs."""
|
|
pack_logger = importlib.import_module("ComfyUI_KEEP.src.keep_logger").logger
|
|
records: list[logging.LogRecord] = []
|
|
|
|
class _Capture(logging.Handler):
|
|
def emit(self, record):
|
|
records.append(record)
|
|
|
|
handler = _Capture()
|
|
pack_logger.addHandler(handler)
|
|
try:
|
|
yield records
|
|
finally:
|
|
pack_logger.removeHandler(handler)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_the_four_steps_of_the_sequence_are_logged(processor_module):
|
|
"""The sequence path logs its four steps - tagged, once each."""
|
|
processor = processor_module.KEEPFaceProcessor(
|
|
_make_pack(_StubKeepNet(), _DetectingFaceHelper())
|
|
)
|
|
frames = torch.rand((2, 16, 16, 3), dtype=torch.float32)
|
|
|
|
with _pack_records() as records:
|
|
processor.process_image_sequence(
|
|
image_sequence_tensor=frames,
|
|
final_upscale_factor=1.0,
|
|
has_aligned_frames=False,
|
|
only_center_face=True,
|
|
draw_box=False,
|
|
)
|
|
|
|
steps = [
|
|
"[KEEP] Detecting face landmarks",
|
|
"[KEEP] Cropping and aligning faces",
|
|
"[KEEP] Restoring faces with KEEP",
|
|
"[KEEP] Pasting faces and finalizing frames",
|
|
]
|
|
messages = [record.getMessage() for record in records]
|
|
# the records of the console routing are in between - the four steps are
|
|
# logged once each and in this order
|
|
assert [message for message in messages if message in steps] == steps
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_the_console_writes_of_the_detection_reach_the_logger(processor_module):
|
|
"""A print of the vendored detector becomes a tagged record."""
|
|
processor = processor_module.KEEPFaceProcessor(
|
|
_make_pack(_StubKeepNet(), _DetectingFaceHelper())
|
|
)
|
|
frames = torch.rand((1, 16, 16, 3), dtype=torch.float32)
|
|
|
|
with _pack_records() as records:
|
|
processor.process_image_sequence(
|
|
image_sequence_tensor=frames,
|
|
final_upscale_factor=1.0,
|
|
has_aligned_frames=False,
|
|
only_center_face=True,
|
|
draw_box=False,
|
|
)
|
|
|
|
messages = [record.getMessage() for record in records]
|
|
assert (
|
|
"[KEEP] model_rootpath is None. Using default relative path" in messages
|
|
), messages
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_every_console_write_of_the_facelib_reaches_the_logger(processor_module):
|
|
"""The prints of the detection, of the alignment and of the paste-back."""
|
|
processor = processor_module.KEEPFaceProcessor(
|
|
_make_pack(_StubKeepNet(), _DetectingFaceHelper())
|
|
)
|
|
frames = torch.rand((1, 16, 16, 3), dtype=torch.float32)
|
|
|
|
with _pack_records() as records:
|
|
processor.process_image_sequence(
|
|
image_sequence_tensor=frames,
|
|
final_upscale_factor=1.0,
|
|
has_aligned_frames=False,
|
|
only_center_face=True,
|
|
draw_box=False,
|
|
)
|
|
|
|
messages = [record.getMessage() for record in records]
|
|
for text in (
|
|
"model_rootpath is None. Using default relative path",
|
|
"Warning: Could not estimate affine matrix for landmark set 0 (frame 0). Skipping this face.",
|
|
"Warning: Mismatch between restored_faces and inverse_affine_matrices.",
|
|
):
|
|
assert f"[KEEP] {text}" in messages, (text, messages)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_the_console_writes_of_the_image_path_reach_the_logger(processor_module):
|
|
"""The single image path routes the detection prints as well."""
|
|
processor = processor_module.KEEPFaceProcessor(
|
|
_make_pack(_StubKeepNet(), _DetectingFaceHelper())
|
|
)
|
|
image = np.zeros((16, 16, 3), dtype=np.uint8)
|
|
|
|
with _pack_records() as records:
|
|
processor.process_image(
|
|
cv2_image_orig=image,
|
|
final_upscale_factor=1.0,
|
|
has_aligned=False,
|
|
only_center_face=True,
|
|
draw_box=False,
|
|
)
|
|
|
|
messages = [record.getMessage() for record in records]
|
|
assert (
|
|
"[KEEP] model_rootpath is None. Using default relative path" in messages
|
|
), messages
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# The GPU video pipeline redesign - the batched sequence path
|
|
# ---------------------------------------------------------------------------
|
|
@pytest.mark.unit
|
|
def test_a_single_frame_chunk_is_duplicated_for_the_net(processor_module):
|
|
"""The 1-frame chunk rule of the KEEP contract.
|
|
|
|
A chunk of one frame is fed to the net as two copies and the first
|
|
output frame is kept - the chunk boundaries are a quality contract and
|
|
must not move.
|
|
"""
|
|
keep_net = _StubKeepNet()
|
|
processor = processor_module.KEEPFaceProcessor(
|
|
_make_pack(keep_net, _StubFaceHelper())
|
|
)
|
|
|
|
frames = torch.rand((1, 16, 16, 3), dtype=torch.float32)
|
|
output = processor.process_image_sequence(
|
|
image_sequence_tensor=frames,
|
|
final_upscale_factor=1.0,
|
|
has_aligned_frames=True,
|
|
only_center_face=True,
|
|
draw_box=False,
|
|
max_clip_length=20,
|
|
)
|
|
|
|
assert len(keep_net.calls) == 1, keep_net.calls
|
|
shape, need_upscale = keep_net.calls[0]
|
|
assert shape == (1, 2, 3, 512, 512), shape
|
|
assert need_upscale is False
|
|
assert tuple(output.shape) == (1, 512, 512, 3), output.shape
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_the_sequence_reads_and_measures_every_frame_once(processor_module):
|
|
"""Read once, affine once.
|
|
|
|
The redesigned path reads every frame a single time (the detection) and
|
|
estimates its affine a single time (the crop) - the paste loop reuses
|
|
the matrix instead of running a second ``align_warp_face``.
|
|
"""
|
|
keep_net = _StubKeepNet()
|
|
face_helper = _DetectingFaceHelper()
|
|
processor = processor_module.KEEPFaceProcessor(_make_pack(keep_net, face_helper))
|
|
|
|
# The frames are at the 512 floor: the paste fallback's deterministic
|
|
# re-derivation of the vendored upscale (D-R2.3.3) fires only below 512
|
|
# and would add one read per fallback frame - the read-once contract of
|
|
# the detection is pinned at the size where it holds for the whole run.
|
|
frames = torch.rand((3, 512, 640, 3), dtype=torch.float32)
|
|
processor.process_image_sequence(
|
|
image_sequence_tensor=frames,
|
|
final_upscale_factor=1.0,
|
|
has_aligned_frames=False,
|
|
only_center_face=True,
|
|
draw_box=False,
|
|
max_clip_length=20,
|
|
)
|
|
|
|
assert face_helper.calls.count("read_image") == 3, face_helper.calls
|
|
assert "align_warp_face" not in face_helper.calls, face_helper.calls
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_the_progress_reports_four_steps_per_frame(processor_module, monkeypatch):
|
|
"""The progress contract of the restructured path.
|
|
|
|
The total stays ``num_frames * 4`` and the four families of the updates
|
|
stay: the conversion (one batched update), the detection (one update),
|
|
the KEEP chunks (one update per chunk) and the paste (one update per
|
|
frame).
|
|
"""
|
|
events: list = []
|
|
|
|
class _RecordingProgress:
|
|
def __init__(self, total):
|
|
events.append(("total", total))
|
|
|
|
def update(self, value):
|
|
events.append(("update", value))
|
|
|
|
monkeypatch.setattr(processor_module, "ProgressReporter", _RecordingProgress)
|
|
|
|
processor = processor_module.KEEPFaceProcessor(
|
|
_make_pack(_StubKeepNet(), _DetectingFaceHelper())
|
|
)
|
|
frames = torch.rand((3, 16, 16, 3), dtype=torch.float32)
|
|
processor.process_image_sequence(
|
|
image_sequence_tensor=frames,
|
|
final_upscale_factor=1.0,
|
|
has_aligned_frames=False,
|
|
only_center_face=True,
|
|
draw_box=False,
|
|
max_clip_length=20,
|
|
)
|
|
|
|
assert events == [
|
|
("total", 15),
|
|
("update", 3), # the batched convert-in
|
|
("update", 3), # the detection
|
|
("update", 3), # the single KEEP chunk
|
|
("update", 3), # the paste stage's second convert-in (D-R2.3.1)
|
|
("update", 1), # the paste, frame by frame
|
|
("update", 1),
|
|
("update", 1),
|
|
], events
|
|
|
|
|
|
class _LandmarkFaceHelper(_StubFaceHelper):
|
|
"""A face helper that detects a face at real, non-degenerate landmarks."""
|
|
|
|
def __init__(self, landmarks) -> None:
|
|
super().__init__()
|
|
self.landmarks = np.asarray(landmarks, dtype=np.float64).reshape(10)
|
|
|
|
def get_face_landmarks_5(self, **kwargs):
|
|
self.calls.append("get_face_landmarks_5")
|
|
self.all_landmarks_5 = [self.landmarks]
|
|
return 1
|
|
|
|
|
|
class _RecordingKeepNet(_StubKeepNet):
|
|
"""A KEEP net that keeps the frames it was called with."""
|
|
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.frames: list = []
|
|
|
|
def __call__(self, frames, need_upscale=False):
|
|
self.frames.append(frames.detach().clone())
|
|
return super().__call__(frames, need_upscale=need_upscale)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_the_batched_crop_matches_the_vendored_warp(processor_module):
|
|
"""The crop of the redesigned path equals the vendored ``cv2.warpAffine``.
|
|
|
|
The reference is the warp of ``face_restoration_helper.py:326-328`` (the
|
|
512 template, the gray border) on the very frame the detector saw,
|
|
followed by the ``img2tensor`` + normalize tail of the processor. The
|
|
measured residual of the batched warp against cv2 is at most **1/255**
|
|
on this content (``tests/test_gpu_batch_equivalence.py``); the bound
|
|
below leaves the headroom of the float round trip.
|
|
|
|
The test also pins the **affine reuse**: the paste loop receives the
|
|
matrix the crop computed - the second ``align_warp_face`` of the old
|
|
path is gone.
|
|
"""
|
|
height, width = 512, 640
|
|
rng = np.random.default_rng(4)
|
|
small = rng.uniform(0.0, 255.0, size=(16, 20, 3)).astype(np.float32)
|
|
frame = np.clip(
|
|
cv2.resize(small, (width, height), interpolation=cv2.INTER_CUBIC), 0, 255
|
|
).astype(np.uint8)
|
|
landmarks = FACE_TEMPLATE_512 + np.array([100.0, 50.0])
|
|
|
|
face_helper = _LandmarkFaceHelper(landmarks)
|
|
keep_net = _RecordingKeepNet()
|
|
processor = processor_module.KEEPFaceProcessor(_make_pack(keep_net, face_helper))
|
|
|
|
frames = torch.from_numpy(frame.astype(np.float32) / 255.0).unsqueeze(0)
|
|
processor.process_image_sequence(
|
|
image_sequence_tensor=frames,
|
|
final_upscale_factor=1.0,
|
|
has_aligned_frames=False,
|
|
only_center_face=True,
|
|
draw_box=False,
|
|
max_clip_length=20,
|
|
)
|
|
|
|
# the crop the net was called with (the 1-frame chunk is duplicated)
|
|
assert len(keep_net.frames) == 1, keep_net.frames
|
|
cropped = keep_net.frames[0][0, 0]
|
|
assert tuple(cropped.shape) == (3, 512, 512), cropped.shape
|
|
|
|
# the reference: the vendored warp + the img2tensor tail of the processor
|
|
# the warp source is the BGR conversion of the input (the converter of
|
|
# the pack - the same array the crop warps)
|
|
image_utils = importlib.import_module("ComfyUI_KEEP.src.keep_image_utils")
|
|
frame_bgr = image_utils.comfy_batch_to_bgr_u8(frames)[0]
|
|
affine = cv2.estimateAffinePartial2D(
|
|
landmarks, FACE_TEMPLATE_512, method=cv2.LMEDS
|
|
)[0]
|
|
reference_u8 = cv2.warpAffine(
|
|
frame_bgr, affine, (512, 512),
|
|
borderMode=cv2.BORDER_CONSTANT, borderValue=(135, 133, 132),
|
|
)
|
|
reference = (
|
|
torch.from_numpy(np.ascontiguousarray(reference_u8[..., ::-1]) / 255.0)
|
|
.to(torch.float32)
|
|
.permute(2, 0, 1)
|
|
.sub_(0.5)
|
|
.div_(0.5)
|
|
)
|
|
|
|
# the crop is not the zero tensor of a no-face frame
|
|
assert cropped.abs().max() > 0.5, "the crop is empty"
|
|
|
|
difference = (cropped - reference).abs().max().item()
|
|
# a 1/255 difference of the uint8 crop is 1/32512.5 in the normalized scale
|
|
assert difference <= 2.0 / 255.0 / 127.5, (
|
|
f"the batched crop drifted {difference:.3e} from the vendored warp"
|
|
)
|
|
|
|
# the paste loop reused the matrix of the crop
|
|
assert len(face_helper.affine_matrices) == 1, face_helper.affine_matrices
|
|
assert np.allclose(face_helper.affine_matrices[0], affine), (
|
|
"the paste loop did not reuse the affine matrix of the crop"
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# P-M1 - the golden fixtures of the image paths (the pre-restructure outputs)
|
|
# ---------------------------------------------------------------------------
|
|
def _load_golden_capture():
|
|
"""The capture tool module - the single source of the golden cases."""
|
|
import importlib.util
|
|
|
|
source = (
|
|
Path(__file__).resolve().parent.parent / "tools" / "capture_round4_golden.py"
|
|
)
|
|
spec = importlib.util.spec_from_file_location("_round4_golden_capture", source)
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("case_name", ["image_multiface", "image_single"])
|
|
def test_the_image_golden_fixtures_are_reproduced_bit_for_bit(
|
|
processor_module, monkeypatch, case_name
|
|
):
|
|
"""P-M1: the image paths re-run their captured cases bitwise.
|
|
|
|
The fixtures come from ``tools/capture_round4_golden.py`` (captured
|
|
once from the pre-restructure code, D-R2.0): the input image, the
|
|
returned BGR image, the KEEP call shapes and the parse-call count (the
|
|
vendored paste parses every pasted face itself,
|
|
``face_restoration_helper.py:462``). The sub-512 input exercises the
|
|
vendored upscale in all three ``read_image`` stages of the image paths.
|
|
"""
|
|
helper_module = importlib.import_module(
|
|
"ComfyUI_KEEP.modules.deps.facelib.utils.face_restoration_helper")
|
|
monkeypatch.setattr(helper_module, "init_detection_model", lambda *a, **k: object())
|
|
monkeypatch.setattr(helper_module, "init_parsing_model", lambda *a, **k: object())
|
|
capture = _load_golden_capture()
|
|
stubs = capture._load_test_stubs()
|
|
case = capture.IMAGE_CASES[case_name]
|
|
fixture_dir = Path(__file__).resolve().parent / "fixtures"
|
|
golden = np.load(fixture_dir / f"round4_golden_{case_name}.npz")
|
|
|
|
image_bgr = capture.golden_frames(1, case["frame_seed"])[0]
|
|
assert np.array_equal(image_bgr, golden["input_image"]), (
|
|
"the re-run input drifted from the captured one"
|
|
)
|
|
faces = {key: capture.golden_face(stubs, key) for key in capture.FACE_OFFSETS}
|
|
face_helper = capture._make_helper(
|
|
stubs, helper_module,
|
|
[[faces[key] for key in case["faces"]]], capture.PARSE_SEEDS[case_name])
|
|
keep_net = stubs._KeepNet()
|
|
processor = processor_module.KEEPFaceProcessor(
|
|
_make_pack(keep_net, face_helper))
|
|
output = processor.process_image(
|
|
cv2_image_orig=image_bgr,
|
|
final_upscale_factor=1.0,
|
|
has_aligned=False,
|
|
only_center_face=case["only_center_face"],
|
|
draw_box=False,
|
|
)
|
|
|
|
assert np.array_equal(np.asarray(output), golden["output_image"]), (
|
|
"the returned image is not bitwise the captured one"
|
|
)
|
|
shapes, flags = capture._keep_records(keep_net)
|
|
assert np.array_equal(shapes, golden["keep_net_shapes"]), (
|
|
f"the KEEP call shapes moved: {shapes} vs {golden['keep_net_shapes']}"
|
|
)
|
|
assert np.array_equal(flags, golden["keep_net_need_upscale"])
|
|
assert face_helper.face_parse.calls == int(golden["parse_calls"]), (
|
|
f"the parse-call count moved: {face_helper.face_parse.calls} "
|
|
f"vs {int(golden['parse_calls'])}"
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# P-M3 / P-M4 / P-M5 - the memory bounds of the redesign (design section 3)
|
|
# ---------------------------------------------------------------------------
|
|
class _TinyParseNet:
|
|
"""A parse net with the vendored contract (``[0]`` are the logits)."""
|
|
|
|
def __init__(self, seed: int = 7) -> None:
|
|
generator = torch.Generator().manual_seed(seed)
|
|
self.weight = torch.randn(19, 3, 1, 1, generator=generator) * 0.35
|
|
|
|
def __call__(self, samples):
|
|
return (F.conv2d(samples, self.weight), None)
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_the_detection_pass_converts_at_most_a_group_at_a_time(
|
|
processor_module, monkeypatch
|
|
):
|
|
"""P-M3 (single-face leg): no whole-clip convert before the first KEEP call.
|
|
|
|
The streamed detection pass converts a group of at most
|
|
``max_clip_length`` rows at a time; the whole-clip uint8 stack of the
|
|
old detection pass is gone. Recording the converter's batch sizes:
|
|
everything before the first KEEP forward stays at or below the clip
|
|
length (3 frames, clips of 2 - a whole-clip convert would be a 3).
|
|
"""
|
|
conversions: list[int] = []
|
|
real_convert = processor_module.comfy_batch_to_bgr_u8
|
|
|
|
def recording_convert(tensor):
|
|
conversions.append(int(tensor.shape[0]))
|
|
return real_convert(tensor)
|
|
|
|
monkeypatch.setattr(processor_module, "comfy_batch_to_bgr_u8", recording_convert)
|
|
|
|
keep_at: list[int] = []
|
|
|
|
class _ProbeKeepNet(_StubKeepNet):
|
|
def __call__(self, frames, need_upscale=False):
|
|
keep_at.append(len(conversions))
|
|
return super().__call__(frames, need_upscale=need_upscale)
|
|
|
|
processor = processor_module.KEEPFaceProcessor(
|
|
_make_pack(_ProbeKeepNet(), _DetectingFaceHelper())
|
|
)
|
|
frames = torch.rand((3, 512, 640, 3), dtype=torch.float32)
|
|
processor.process_image_sequence(
|
|
image_sequence_tensor=frames,
|
|
final_upscale_factor=1.0,
|
|
has_aligned_frames=False,
|
|
only_center_face=True,
|
|
draw_box=False,
|
|
max_clip_length=2,
|
|
)
|
|
|
|
assert keep_at, "the probe net must be called"
|
|
before_the_first_keep = conversions[: keep_at[0]]
|
|
assert before_the_first_keep, (
|
|
"the detection pass must convert through the module function"
|
|
)
|
|
assert max(before_the_first_keep) <= 2, before_the_first_keep
|
|
# The fused crop loop re-derives the chunk's detection images in ONE
|
|
# group conversion per chunk (the grouped D-R2.2.1) - the per-row
|
|
# re-derivations of the first round-4 cut are gone: across the whole
|
|
# run only the detection groups, the crop groups and the paste stage's
|
|
# second convert-in may call the converter (2 * ceil(3 / 2) + 1 = 5).
|
|
assert len(conversions) <= 5, conversions
|
|
|
|
|
|
@pytest.mark.unit
|
|
def test_the_chunked_mode_keeps_one_tail_buffer(processor_module, monkeypatch):
|
|
"""P-M4: the forced chunked mode releases its staging tail (D-R2.3.5).
|
|
|
|
``device_output_mode`` is patched to "chunked" (on CPU the modes are
|
|
value-identical, the verdict just never selects a device mode by
|
|
itself): 3 frames at 64x64 with one black - gray - frame, so 2
|
|
eligible rows. Recording every ``torch.empty``: exactly one float32
|
|
allocation carries the whole-clip row count (the contractual
|
|
``final_batch``) and exactly one carries the eligible rows (the
|
|
chunked core's ``out``). After the return the eligible buffer is
|
|
gone - the processor drops the name AND the row views
|
|
(``batched_outputs`` becomes an index set) - while the returned IMAGE
|
|
is ``final_batch`` itself and keeps it alive.
|
|
"""
|
|
paste_mod = importlib.import_module("ComfyUI_KEEP.src.keep_gpu_paste")
|
|
monkeypatch.setattr(paste_mod, "device_output_mode", lambda *args, **kwargs: "chunked")
|
|
|
|
empties: list[tuple[tuple, torch.dtype, weakref.ref]] = []
|
|
real_empty = torch.empty
|
|
|
|
def recording_empty(*args, **kwargs):
|
|
tensor = real_empty(*args, **kwargs)
|
|
empties.append((tuple(tensor.shape), tensor.dtype, weakref.ref(tensor)))
|
|
return tensor
|
|
|
|
monkeypatch.setattr(torch, "empty", recording_empty)
|
|
|
|
landmarks = (FACE_TEMPLATE_512 - 256.0) * 0.1 + np.array([48.0, 32.0])
|
|
face_helper = _LandmarkFaceHelper(landmarks)
|
|
face_helper.face_parse = _TinyParseNet()
|
|
face_helper.use_parse = True
|
|
frames = torch.rand((3, 64, 64, 3), dtype=torch.float32)
|
|
frames[1] = 0.0 # the black frame is a gray frame - it leaves the batch
|
|
|
|
processor = processor_module.KEEPFaceProcessor(
|
|
_make_pack(_StubKeepNet(), face_helper)
|
|
)
|
|
output = processor.process_image_sequence(
|
|
image_sequence_tensor=frames,
|
|
final_upscale_factor=1.0,
|
|
has_aligned_frames=False,
|
|
only_center_face=True,
|
|
draw_box=False,
|
|
max_clip_length=20,
|
|
)
|
|
|
|
float_rows = [
|
|
(shape, ref) for shape, dtype, ref in empties
|
|
if dtype == torch.float32 and len(shape) == 4
|
|
]
|
|
final = [ref for shape, ref in float_rows if shape[0] == 3]
|
|
eligible = [ref for shape, ref in float_rows if shape[0] == 2]
|
|
assert len(final) == 1, float_rows
|
|
assert len(eligible) == 1, float_rows
|
|
# the returned IMAGE is the contractual whole-clip buffer
|
|
assert final[0]() is not None, "final_batch must stay alive after the return"
|
|
# the chunked core's eligible-rows buffer is released with its row views
|
|
assert eligible[0]() is None, (
|
|
"the chunked staging tail must be released after the indexed scatter"
|
|
)
|
|
assert tuple(output.shape) == (3, 512, 512, 3), output.shape
|
|
|
|
|
|
@pytest.mark.unit
|
|
@pytest.mark.parametrize("case_name", ["image_single", "image_multiface"])
|
|
def test_the_image_paths_release_the_helper_after_the_run(
|
|
processor_module, monkeypatch, case_name
|
|
):
|
|
"""P-M5 (image legs): no detection image survives the run (D-R2.4).
|
|
|
|
The image paths run the real vendored helper (the capture rig): the
|
|
``finally`` must clear the face lists and the ``input_img`` the same
|
|
way the sequence paths do - the sub-512 golden input leaves a full
|
|
upscaled frame in ``input_img`` unless it is released.
|
|
"""
|
|
helper_module = importlib.import_module(
|
|
"ComfyUI_KEEP.modules.deps.facelib.utils.face_restoration_helper")
|
|
monkeypatch.setattr(helper_module, "init_detection_model", lambda *a, **k: object())
|
|
monkeypatch.setattr(helper_module, "init_parsing_model", lambda *a, **k: object())
|
|
capture = _load_golden_capture()
|
|
stubs = capture._load_test_stubs()
|
|
case = capture.IMAGE_CASES[case_name]
|
|
image_bgr = capture.golden_frames(1, case["frame_seed"])[0]
|
|
faces = {key: capture.golden_face(stubs, key) for key in capture.FACE_OFFSETS}
|
|
face_helper = capture._make_helper(
|
|
stubs, helper_module,
|
|
[[faces[key] for key in case["faces"]]], capture.PARSE_SEEDS[case_name])
|
|
processor = processor_module.KEEPFaceProcessor(
|
|
_make_pack(stubs._KeepNet(), face_helper))
|
|
processor.process_image(
|
|
cv2_image_orig=image_bgr,
|
|
final_upscale_factor=1.0,
|
|
has_aligned=False,
|
|
only_center_face=case["only_center_face"],
|
|
draw_box=False,
|
|
)
|
|
|
|
assert face_helper.input_img is None, case_name
|
|
for name in (
|
|
"cropped_faces", "restored_faces", "affine_matrices",
|
|
"inverse_affine_matrices", "all_landmarks_5",
|
|
):
|
|
assert getattr(face_helper, name) == [], (case_name, name)
|