From 573b53c729b0e9076020c94de61b995aa2f028e1 Mon Sep 17 00:00:00 2001 From: xmarre Date: Sun, 3 May 2026 00:50:10 +0200 Subject: [PATCH] Fix HDR LAB conversion crash with chunked processing --- superbeastsai_nodes.py | 123 ++++++++++++++++++++++++++++++++--------- 1 file changed, 98 insertions(+), 25 deletions(-) diff --git a/superbeastsai_nodes.py b/superbeastsai_nodes.py index 0ad64f1..e60dd78 100644 --- a/superbeastsai_nodes.py +++ b/superbeastsai_nodes.py @@ -650,39 +650,112 @@ def _f_lab_to_xyz(t): def _rgb_to_lab_components(img: Image.Image): - rgb = np.asarray(img.convert('RGB'), dtype=np.float32) / 255.0 - linear_rgb = _srgb_to_linear(rgb) - xyz = linear_rgb @ _SRGB_TO_XYZ.T - xyz = xyz / _D65_WHITE + rgb_u8 = np.ascontiguousarray(np.asarray(img.convert('RGB'), dtype=np.uint8)) + if rgb_u8.ndim != 3 or rgb_u8.shape[2] != 3: + raise ValueError(f"Expected RGB image array with shape HxWx3, got {rgb_u8.shape}") - f_xyz = _f_xyz_to_lab(xyz) - l = (116.0 * f_xyz[..., 1]) - 16.0 - a = 500.0 * (f_xyz[..., 0] - f_xyz[..., 1]) - b = 200.0 * (f_xyz[..., 1] - f_xyz[..., 2]) + height, width = rgb_u8.shape[:2] + l_uint8 = np.empty((height, width), dtype=np.uint8) + a_out = np.empty((height, width), dtype=np.float32) + b_out = np.empty((height, width), dtype=np.float32) - l_uint8 = np.clip(l * (255.0 / 100.0), 0, 255).astype(np.uint8) - return l_uint8, a.astype(np.float32, copy=False), b.astype(np.float32, copy=False) + max_chunk_pixels = 1_048_576 + rows_per_chunk = max(1, max_chunk_pixels // max(1, width)) + rgb_to_xyz = _SRGB_TO_XYZ + white_x = _D65_WHITE[0] + white_y = _D65_WHITE[1] + white_z = _D65_WHITE[2] + + for y0 in range(0, height, rows_per_chunk): + y1 = min(height, y0 + rows_per_chunk) + rgb = rgb_u8[y0:y1].astype(np.float32) / 255.0 + linear_rgb = _srgb_to_linear(rgb) + + r = linear_rgb[..., 0] + g = linear_rgb[..., 1] + bl = linear_rgb[..., 2] + + x = (rgb_to_xyz[0, 0] * r + rgb_to_xyz[0, 1] * g + rgb_to_xyz[0, 2] * bl) / white_x + y = (rgb_to_xyz[1, 0] * r + rgb_to_xyz[1, 1] * g + rgb_to_xyz[1, 2] * bl) / white_y + z = (rgb_to_xyz[2, 0] * r + rgb_to_xyz[2, 1] * g + rgb_to_xyz[2, 2] * bl) / white_z + + fx = _f_xyz_to_lab(x) + fy = _f_xyz_to_lab(y) + fz = _f_xyz_to_lab(z) + + l = (np.float32(116.0) * fy) - np.float32(16.0) + a = np.float32(500.0) * (fx - fy) + b = np.float32(200.0) * (fy - fz) + + l_uint8[y0:y1] = np.clip(l * (255.0 / 100.0), 0, 255).astype(np.uint8) + a_out[y0:y1] = a.astype(np.float32, copy=False) + b_out[y0:y1] = b.astype(np.float32, copy=False) + + return l_uint8, a_out, b_out + + +def _linear_srgb_channel_to_u8(channel: np.ndarray) -> np.ndarray: + """Convert one linear-light sRGB channel to uint8 sRGB. + + Kept channel-wise so HDR Effects does not allocate several full-frame + HxWx3 float32 temporaries while converting LAB back to RGB. Some NumPy + builds have crashed in those native ufunc/matmul paths under ComfyUI's + large-image memory pressure instead of raising MemoryError. + """ + srgb = np.clip(_linear_to_srgb(channel), 0.0, 1.0) + return (srgb * 255.0 + 0.5).astype(np.uint8) def _lab_components_to_rgb(l_uint8: np.ndarray, a: np.ndarray, b: np.ndarray) -> Image.Image: - lab = np.empty((l_uint8.shape[0], l_uint8.shape[1], 3), dtype=np.float32) - lab[..., 0] = np.asarray(l_uint8, dtype=np.float32) * (100.0 / 255.0) - lab[..., 1] = np.asarray(a, dtype=np.float32) - lab[..., 2] = np.asarray(b, dtype=np.float32) + l_uint8 = np.ascontiguousarray(l_uint8) + a = np.ascontiguousarray(a, dtype=np.float32) + b = np.ascontiguousarray(b, dtype=np.float32) - fy = (lab[..., 0] + 16.0) / 116.0 - fx = lab[..., 1] / 500.0 + fy - fz = fy - lab[..., 2] / 200.0 + if l_uint8.ndim != 2: + raise ValueError(f"Expected 2D LAB luminance channel, got shape {l_uint8.shape}") + if a.shape != l_uint8.shape or b.shape != l_uint8.shape: + raise ValueError( + "LAB channel shape mismatch: " + f"L={l_uint8.shape}, a={a.shape}, b={b.shape}" + ) - xyz = np.empty_like(lab) - xyz[..., 0] = _f_lab_to_xyz(fx) - xyz[..., 1] = _f_lab_to_xyz(fy) - xyz[..., 2] = _f_lab_to_xyz(fz) - xyz = xyz * _D65_WHITE + height, width = l_uint8.shape + rgb_u8 = np.empty((height, width, 3), dtype=np.uint8) - linear_rgb = xyz @ _XYZ_TO_SRGB.T - rgb = np.clip(_linear_to_srgb(linear_rgb), 0.0, 1.0) - return Image.fromarray((rgb * 255.0 + 0.5).astype(np.uint8), mode='RGB') + # Process rows in bounded chunks. This avoids the previous full-frame LAB, + # XYZ, linear-RGB and matrix-multiply temporaries while preserving the same + # D65/LAB math. Keep the chunk size large enough to avoid Python-loop cost + # on normal images but small enough for 4K/8K frames and batched Comfy runs. + max_chunk_pixels = 1_048_576 + rows_per_chunk = max(1, max_chunk_pixels // max(1, width)) + xyz_to_rgb = _XYZ_TO_SRGB + white_x = _D65_WHITE[0] + white_y = _D65_WHITE[1] + white_z = _D65_WHITE[2] + + for y0 in range(0, height, rows_per_chunk): + y1 = min(height, y0 + rows_per_chunk) + + l = l_uint8[y0:y1].astype(np.float32, copy=False) * (100.0 / 255.0) + fy = (l + 16.0) / 116.0 + fx = a[y0:y1] / 500.0 + fy + fz = fy - b[y0:y1] / 200.0 + + x = _f_lab_to_xyz(fx) * white_x + y = _f_lab_to_xyz(fy) * white_y + z = _f_lab_to_xyz(fz) * white_z + + # Inline the XYZ->linear sRGB matrix multiply channel-wise. This keeps + # execution inside simple 2D ufuncs instead of NumPy's 3D matmul path. + r = xyz_to_rgb[0, 0] * x + xyz_to_rgb[0, 1] * y + xyz_to_rgb[0, 2] * z + g = xyz_to_rgb[1, 0] * x + xyz_to_rgb[1, 1] * y + xyz_to_rgb[1, 2] * z + bl = xyz_to_rgb[2, 0] * x + xyz_to_rgb[2, 1] * y + xyz_to_rgb[2, 2] * z + + rgb_u8[y0:y1, :, 0] = _linear_srgb_channel_to_u8(r) + rgb_u8[y0:y1, :, 1] = _linear_srgb_channel_to_u8(g) + rgb_u8[y0:y1, :, 2] = _linear_srgb_channel_to_u8(bl) + + return Image.fromarray(rgb_u8, mode='RGB') def adjust_shadows(luminance_array, shadow_intensity, hdr_intensity):