nodes/wildminder-ComfyUI-KEEP
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

perf(keep): derive the chunk's detection frames in one grouped pass

Browse Source 
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).
This commit is contained in:
WildAi
2026-10-29 00:55:54 +03:00
parent 3af08f4357
commit e2fdbbf453
3 changed files with 79 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/keep_processor.py
+64 -25
View File
@@ -695,27 +695,54 @@ class KEEPFaceProcessor:
self.face_helper.clean_all()
self.face_helper.input_img = None
def _detection_frames(self,
image_sequence_tensor: torch.Tensor,
indices,
needs_upscale: list[bool]) -> dict[int, np.ndarray]:
"""The images the detector saw for ``indices``, in one group pass.
The round-4 re-derivation (D-R2.2.1) of the released
``detection_frames`` list, in the **grouped** form: one batched
conversion of the gathered rows instead of one single-row
conversion per frame - the converter is elementwise, so every row
is bitwise the row the per-row form produced (and the round-3 crop
consumed), while the chunk pays one host transfer instead of
``len(indices)``. For the sub-512 frames the vendored
``read_image`` re-runs to reproduce its INTER_LINEAR upscale - the
deterministic image the crop warps consume and the paste sizes its
canvas from (``face_restoration_helper.py:182-184``). The caller
keeps the helper's console writes captured; the rows are bounded by
the caller's chunk (at most ``max_clip_length`` of them).
Returns:
``{index: the detection image}`` - the raw converted row, or
the vendored upscale of it.
"""
unique = sorted({int(index) for index in indices})
rows = comfy_batch_to_bgr_u8(image_sequence_tensor[unique])
frames: dict[int, np.ndarray] = {}
for position, index in enumerate(unique):
row = rows[position]
if needs_upscale[index]:
# the same vendored re-derivation the per-row form runs -
# the one ``read_image`` call site of the sequence path.
self.face_helper.read_image(row)
frames[index] = self.face_helper.input_img
else:
frames[index] = row
return frames
def _detection_frame(self,
image_sequence_tensor: torch.Tensor,
index: int,
needs_upscale: list[bool]) -> np.ndarray:
"""The image the detector saw for ``index`` (round 4, D-R2.2.1).
"""The image the detector saw for ``index`` - the group form, one row.
The row of the batched conversion, re-derived from the input tensor
- the converter is elementwise, so the row is bitwise the row the
detection pass converted. For the sub-512 frames the vendored
``read_image`` re-runs to reproduce its INTER_LINEAR upscale - the
deterministic image today's crop warps consume and the paste sizes
its canvas from (``face_restoration_helper.py:182-184``). The same
derivation serves the crop source and the paste fallback's
``input_img`` (D-R2.3.3); the call sites keep the helper's console
writes captured.
The paste fallback's per-frame derivation (D-R2.3.3): the same
values :meth:`_detection_frames` returns, for a single index.
"""
row = comfy_batch_to_bgr_u8(image_sequence_tensor[index:index + 1])[0]
if not needs_upscale[index]:
return row
self.face_helper.read_image(row)
return self.face_helper.input_img
return self._detection_frames(
image_sequence_tensor, [index], needs_upscale)[index]
@torch.no_grad()
def process_image_sequence(self,
@@ -886,14 +913,24 @@ class KEEPFaceProcessor:
# console writes are routed to the logger of the pack
# (D4) - the derivation's ``read_image`` included.
with capture_console_to_logger():
# The chunk's detection images in ONE group pass
# (D-R2.2.1, grouped): the conversion is elementwise,
# so the rows are bitwise the rows the per-row form
# produced - one host transfer per chunk instead of
# one per row.
chunk_detection_frames = (
self._detection_frames(
image_sequence_tensor, croppable_rows,
needs_upscale)
if croppable_rows else {})
for row, i in enumerate(croppable_rows):
staged_u8[row] = cv2.warpAffine(
self._detection_frame(
image_sequence_tensor, i, needs_upscale),
chunk_detection_frames[i],
affine_matrices[i], (512, 512),
borderMode=cv2.BORDER_CONSTANT,
borderValue=(135, 133, 132),
)
del chunk_detection_frames
# The frames without a (croppable) face carry the base's
# neutral crop (``torch.zeros(3, 512, 512) * 0.5`` of the
# base - 0.5 in the normalized scale, not 0): the
@@ -1266,23 +1303,25 @@ class KEEPFaceProcessor:
keep_interrupt.check_interrupted()
end_idx = min(start_idx + max_clip_length, total_faces)
staged_u8 = np.empty((end_idx - start_idx, 512, 512, 3), np.uint8)
chunk_frames: dict[int, np.ndarray] = {}
# The crop step calls into the vendored facelib: its console
# writes are routed to the logger of the pack (D4) - the
# derivation's ``read_image`` included.
with capture_console_to_logger():
# The chunk's involved frames in ONE group pass
# (D-R2.2.1, grouped): the conversion is elementwise, so
# the rows are bitwise the rows the per-row form produced
# - one host transfer per chunk instead of one per
# involved frame (a frame's faces are contiguous, so the
# involved frames stay ``max_clip_length``-bounded).
chunk_frames = self._detection_frames(
image_sequence_tensor,
face_frames[start_idx:end_idx], needs_upscale)
for row, k in enumerate(range(start_idx, end_idx)):
# The per-face poll of today's crop loop (the cadence
# pin counts it; the chunk top polled above).
keep_interrupt.check_interrupted()
frame = face_frames[k]
detection_frame = chunk_frames.get(frame)
if detection_frame is None:
detection_frame = self._detection_frame(
image_sequence_tensor, frame, needs_upscale)
chunk_frames[frame] = detection_frame
staged_u8[row] = cv2.warpAffine(
detection_frame, face_affines[k], (512, 512),
chunk_frames[face_frames[k]], face_affines[k], (512, 512),
borderMode=cv2.BORDER_CONSTANT,
borderValue=(135, 133, 132),
)
tests/test_keep_processor.py
+6
View File
@@ -765,6 +765,12 @@ def test_the_detection_pass_converts_at_most_a_group_at_a_time(
"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
tests/test_multiface_sequence.py
+9 -2
View File
@@ -960,8 +960,9 @@ def test_the_detection_pass_converts_at_most_a_group_at_a_time(
"""P-M3 (multi-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 and the fused loop re-derives
single rows; the whole-clip uint8 stack of the old detection pass is
``max_clip_length`` rows at a time and the fused loop re-derives the
chunk's involved frames in one group conversion per chunk (the grouped
D-R2.2.1); 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. The frames sit
at the 512 floor, so the pin measures exactly the group conversions
@@ -996,6 +997,12 @@ def test_the_detection_pass_converts_at_most_a_group_at_a_time(
"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 involved frames in ONE
# group conversion per chunk (the grouped D-R2.2.1) - the per-frame
# 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(4 / 2) + 1 = 5).
assert len(conversions) <= 5, conversions
# ---------------------------------------------------------------------------