From 5337db991bc956b9e2d30bc6543b59f6f86e0e75 Mon Sep 17 00:00:00 2001 From: Tom Maiaroto Date: Sat, 15 Aug 2026 22:47:14 -0700 Subject: [PATCH] fix chromatic aberration edge wrapping by replacing torch.roll with edge replication padding --- stochastic_noise.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/stochastic_noise.py b/stochastic_noise.py index 33cc82c..4b6e0d4 100644 --- a/stochastic_noise.py +++ b/stochastic_noise.py @@ -621,13 +621,16 @@ class AfterDarkStochasticNoise: tensor_nchw = F.grid_sample(tensor_nchw, warped_grid, mode="bilinear", padding_mode="border", align_corners=False) out_image = tensor_nchw.permute(0, 2, 3, 1) - # 6. Subtle Chromatic Aberration (sub-pixel channel offset) + # 6. Subtle Chromatic Aberration (sub-pixel channel offset with edge replication) if eff_chromatic_aberration > 0.0 and C >= 3: shift_pixels = max(1, int(eff_chromatic_aberration * min(H, W))) - r_ch = torch.roll(out_image[..., 0], shifts=(shift_pixels, shift_pixels), dims=(1, 2)) - b_ch = torch.roll(out_image[..., 2], shifts=(-shift_pixels, -shift_pixels), dims=(1, 2)) - out_image[..., 0] = r_ch - out_image[..., 2] = b_ch + # Pad and slice Red channel (shift right & down) + r_padded = F.pad(out_image[..., 0].unsqueeze(1), (shift_pixels, 0, shift_pixels, 0), mode="replicate") + out_image[..., 0] = r_padded[:, 0, :H, :W] + + # Pad and slice Blue channel (shift left & up) + b_padded = F.pad(out_image[..., 2].unsqueeze(1), (0, shift_pixels, 0, shift_pixels), mode="replicate") + out_image[..., 2] = b_padded[:, 0, shift_pixels:shift_pixels+H, shift_pixels:shift_pixels+W] # 7. Film Grain / Stochastic Noise Generation if noise_level > 0.0: