From 1cc367271ae498db8c4b82b66fd9a8a2d31bdd3c Mon Sep 17 00:00:00 2001 From: MSXYZ-GenAI Date: Thu, 30 Apr 2026 16:10:45 +0300 Subject: [PATCH] Add subpixel refinement --- VideoTAADLAA.py | 151 +++++++++++++++++++++++++++++++++++++++++++++++- config.py | 20 ++++++- 2 files changed, 167 insertions(+), 4 deletions(-) diff --git a/VideoTAADLAA.py b/VideoTAADLAA.py index 7ac3757..e4ade84 100644 --- a/VideoTAADLAA.py +++ b/VideoTAADLAA.py @@ -544,6 +544,140 @@ class VideoTAADLAA: return torch.clamp(out, 0.0, 1.0) + def _chroma_edge_cleanup(self, x, net, strength): + strength = max(0.0, min(float(strength), 1.0)) + + if strength <= 1e-5: + return x + + luma = rgb_luma(x) + + sx = F.conv2d(luma, net.sobel_x, padding=1) + sy = F.conv2d(luma, net.sobel_y, padding=1) + edge = torch.sqrt(sx * sx + sy * sy + 1e-6) + + edge_mask = torch.sigmoid( + (edge - self.detail_chroma_edge_threshold) * self.edge_aa_slope + ) + + chroma = x - luma + + chroma_blur = F.avg_pool2d( + F.pad(chroma, [1, 1, 1, 1], mode="reflect"), + 3, + stride=1, + ) + + chroma_delta = chroma_blur - chroma + chroma_residual = rgb_luma(chroma_delta.abs()) + chroma_amount = rgb_luma(chroma.abs()) + + chroma_mask = torch.sigmoid( + (chroma_amount - self.detail_chroma_saturation_threshold) * 20.0 + ) + + fringe_mask = torch.sigmoid( + (chroma_residual - self.detail_chroma_saturation_threshold * 0.25) * 80.0 + ) + + dark_protect = torch.sigmoid( + (luma - self.detail_chroma_dark_protect) * 12.0 + ) + + mask = (edge_mask * chroma_mask * fringe_mask * dark_protect).clamp(0.0, 1.0) + + chroma_delta = chroma_delta.clamp( + -self.detail_chroma_cleanup_limit, + self.detail_chroma_cleanup_limit, + ) + + clean_chroma = chroma + chroma_delta * mask * strength + out = luma + clean_chroma + + # keep luma stable while cleaning chroma + luma_error = rgb_luma(out) - luma + out = out - luma_error + + return torch.clamp(out, 0.0, 1.0) + + def _subpixel_edge_reconstruction(self, x, net, motion_gate, strength): + strength = max(0.0, min(float(strength), 1.0)) + + if strength <= 1e-5: + return x + + luma = rgb_luma(x) + + sx = F.conv2d(luma, net.sobel_x, padding=1) + sy = F.conv2d(luma, net.sobel_y, padding=1) + edge = torch.sqrt(sx * sx + sy * sy + 1e-6) + + edge_mask = torch.sigmoid( + (edge - self.detail_subpixel_edge_threshold) * + self.detail_subpixel_edge_slope + ) + + B, C, H, W = x.shape + + nx = sx / edge.clamp(min=1e-6) + ny = sy / edge.clamp(min=1e-6) + + sample_scale = float(self.detail_subpixel_sample_scale) + + offset_x = nx.squeeze(1) * sample_scale * (2.0 / max(W - 1, 1)) + offset_y = ny.squeeze(1) * sample_scale * (2.0 / max(H - 1, 1)) + + yy, xx = torch.meshgrid( + torch.linspace(-1.0, 1.0, H, device=x.device, dtype=x.dtype), + torch.linspace(-1.0, 1.0, W, device=x.device, dtype=x.dtype), + indexing="ij", + ) + + base_grid = torch.stack((xx, yy), dim=-1).unsqueeze(0).repeat(B, 1, 1, 1) + + offset = torch.stack((offset_x, offset_y), dim=-1) + + grid_pos = base_grid + offset + grid_neg = base_grid - offset + + sample_pos = F.grid_sample( + x, + grid_pos, + mode="bilinear", + padding_mode="reflection", + align_corners=False, + ) + + sample_neg = F.grid_sample( + x, + grid_neg, + mode="bilinear", + padding_mode="reflection", + align_corners=False, + ) + + reconstructed = (sample_pos + sample_neg) * 0.5 + + delta = reconstructed - x + delta = delta.clamp( + -self.detail_subpixel_delta_limit, + self.detail_subpixel_delta_limit, + ) + + stable_motion = 1.0 - min( + max(float(motion_gate) * self.detail_subpixel_motion_protect, 0.0), + 1.0, + ) + + blend = (edge_mask * strength * stable_motion).clamp( + 0.0, + self.detail_subpixel_blend_limit, + ) + + out = x + delta * blend + + return torch.clamp(out, 0.0, 1.0) + def _fur_hair_stabilizer(self, x, previous, net, motion_gate, strength): strength = max(0.0, min(float(strength), 1.0)) @@ -554,8 +688,6 @@ class VideoTAADLAA: return x luma = rgb_luma(x) - prev_luma = rgb_luma(previous) - sx = F.conv2d(luma, net.sobel_x, padding=1) sy = F.conv2d(luma, net.sobel_y, padding=1) edge = torch.sqrt(sx * sx + sy * sy + 1e-6) @@ -1366,6 +1498,21 @@ class VideoTAADLAA: self.detail_dehalo_strength, ) + if preset == "Detail": + dlaa_out = self._subpixel_edge_reconstruction( + dlaa_out, + net, + motion_gate, + self.detail_subpixel_edge_strength, + ) + + if preset == "Detail": + dlaa_out = self._chroma_edge_cleanup( + dlaa_out, + net, + self.detail_chroma_cleanup_strength, + ) + dlaa_out = self._apply_tone_and_color_pass( dlaa_out, rgb, diff --git a/config.py b/config.py index 34e0231..b98999f 100644 --- a/config.py +++ b/config.py @@ -93,13 +93,29 @@ NODE_DEFAULTS = { "detail_micro_contrast_radius": 5, "detail_micro_contrast_limit": 0.016, "detail_micro_contrast_highlight_protect": 0.70, - + # Edge dehalo "detail_dehalo_strength": 0.14, "detail_dehalo_threshold": 0.095, "detail_dehalo_dark_protect": 0.35, "detail_dehalo_light_protect": 0.55, - + + # Chroma edge cleanup + "detail_chroma_cleanup_strength": 0.18, # best range 0.10 - 0.25 + "detail_chroma_edge_threshold": 0.055, + "detail_chroma_saturation_threshold": 0.070, + "detail_chroma_cleanup_limit": 0.014, + "detail_chroma_dark_protect": 0.30, + + # Subpixel edge reconstruction + "detail_subpixel_edge_strength": 0.14, # best range 0.10 - 0.20 + "detail_subpixel_edge_threshold": 0.070, + "detail_subpixel_edge_slope": 16.0, + "detail_subpixel_sample_scale": 0.35, + "detail_subpixel_blend_limit": 0.24, + "detail_subpixel_delta_limit": 0.016, + "detail_subpixel_motion_protect": 0.75, + # Fur / hair directional stabilizer "detail_fur_stabilizer_strength": 0.08, "detail_fur_edge_threshold": 0.05,