# Video TAA + DLAA # v0.1.2 # MSXYZ # Developed with AI-assisted tooling import os import logging import threading import torch import torch.nn as nn import torch.nn.functional as F from safetensors.torch import load_file try: import comfy.model_management as mm except ImportError: mm = None try: from comfy.utils import ProgressBar except ImportError: ProgressBar = None try: import comfy.utils as comfy_utils from spandrel import ModelLoader, MAIN_REGISTRY try: from spandrel_extra_arches import EXTRA_REGISTRY MAIN_REGISTRY.add(*EXTRA_REGISTRY) except Exception: pass except ImportError: comfy_utils = None ModelLoader = None log = logging.getLogger("VideoTAADLAA") LUMA_WEIGHTS = (0.2126, 0.7152, 0.0722) def rgb_luma(x: torch.Tensor) -> torch.Tensor: r, g, b = LUMA_WEIGHTS return r * x[:, 0:1] + g * x[:, 1:2] + b * x[:, 2:3] # Main tuning presets. PRESETS = { "Balanced": { "detail_boost": 1.00, "edge_boost": 1.00, "temporal_strength": 0.35, "micro_limit": 0.040, "luma_boost_mult": 1.00, "saturation_boost_mult": 1.00, "motion_threshold": 0.08, "taa_strength": 0.45, "dlaa_strength": 0.65, "tone_strength": 0.12, "edge_sharp_strength": 0.14, "motion_sensitivity": 0.08, "jitter_scale": 0.20, }, "Detail": { "detail_boost": 1.44, "edge_boost": 1.36, "temporal_strength": 0.00, "micro_limit": 0.050, "luma_boost_mult": 1.08, "saturation_boost_mult": 1.03, "motion_threshold": 0.050, "taa_strength": 0.00, "dlaa_strength": 1.25, "tone_strength": 0.06, "edge_sharp_strength": 0.14, "motion_sensitivity": 0.060, "jitter_scale": 0.0, }, "Smooth": { "detail_boost": 0.85, "edge_boost": 0.75, "temporal_strength": 0.45, "micro_limit": 0.025, "luma_boost_mult": 0.80, "saturation_boost_mult": 0.70, "motion_threshold": 0.10, "taa_strength": 0.65, "dlaa_strength": 0.65, "tone_strength": 0.16, "edge_sharp_strength": 0.08, "motion_sensitivity": 0.10, "jitter_scale": 0.25, }, } # Auto mode profiles picked from estimated frame motion. AUTO_STATIC = { "detail_boost": 1.12, "edge_boost": 1.15, "temporal_strength": 0.28, "micro_limit": 0.045, "luma_boost_mult": 1.15, "saturation_boost_mult": 1.05, "motion_threshold": 0.07, "taa_strength": 0.35, "dlaa_strength": 0.70, "tone_strength": 0.10, "edge_sharp_strength": 0.18, "motion_sensitivity": 0.07, "jitter_scale": 0.15, } AUTO_BALANCED = { "detail_boost": 1.00, "edge_boost": 1.00, "temporal_strength": 0.35, "micro_limit": 0.040, "luma_boost_mult": 1.00, "saturation_boost_mult": 1.00, "motion_threshold": 0.08, "taa_strength": 0.45, "dlaa_strength": 0.65, "tone_strength": 0.12, "edge_sharp_strength": 0.14, "motion_sensitivity": 0.08, "jitter_scale": 0.20, } AUTO_MOTION = { "detail_boost": 0.85, "edge_boost": 0.75, "temporal_strength": 0.42, "micro_limit": 0.025, "luma_boost_mult": 0.85, "saturation_boost_mult": 0.75, "motion_threshold": 0.10, "taa_strength": 0.60, "dlaa_strength": 0.60, "tone_strength": 0.16, "edge_sharp_strength": 0.08, "motion_sensitivity": 0.10, "jitter_scale": 0.25, } # Softer Detail settings for single images. SINGLE_IMAGE_DETAIL = { "detail_boost": 1.30, "edge_boost": 1.18, "temporal_strength": 0.00, "luma_boost_mult": 1.04, "saturation_boost_mult": 1.01, "taa_strength": 0.08, "dlaa_strength": 1.12, "edge_sharp_strength": 0.10, } # Reduces fine detail boost when frame-to-frame motion is detected. MOTION_SUPPRESSION = { "Detail": { "detail": 0.05, "edge": 0.08, "micro": 0.18, }, "Default": { "detail": 0.20, "edge": 0.15, "micro": 0.20, }, } class DLAANet(nn.Module): """DLAA network""" def __init__(self): super().__init__() base_channels = 192 self.register_buffer("sobel_x", torch.tensor([[-1,0,1],[-2,0,2],[-1,0,1]], dtype=torch.float32).view(1,1,3,3)) self.register_buffer("sobel_y", torch.tensor([[-1,-2,-1],[0,0,0],[1,2,1]], dtype=torch.float32).view(1,1,3,3)) self.register_buffer( "jitter_offsets", torch.tensor([ [ 0.0000, -0.1667], [-0.2500, 0.1667], [ 0.2500, -0.3889], [-0.3750, -0.0556], [ 0.1250, 0.2778], [-0.1250, -0.2778], [ 0.3750, 0.0556], [-0.4375, 0.3889], [ 0.0625, -0.4630], [-0.3125, 0.1296], [ 0.1875, -0.2963], [-0.4375, -0.0370], [ 0.3125, 0.2407], [-0.0625, -0.2037], [ 0.4375, 0.0185], [-0.4688, 0.3148], ], dtype=torch.float32) ) self.enc1 = nn.Sequential( nn.Conv2d(3, base_channels, 3, padding=1, bias=False), nn.GroupNorm(12, base_channels), nn.LeakyReLU(0.2, inplace=True), ) self.enc2 = nn.Sequential( nn.Conv2d(base_channels, base_channels, 3, padding=1, bias=False), nn.GroupNorm(12, base_channels), nn.LeakyReLU(0.2, inplace=True), ) self.bottleneck = nn.Sequential( nn.Conv2d(base_channels, base_channels, 3, padding=1, dilation=1, bias=False), nn.GroupNorm(12, base_channels), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(base_channels, base_channels, 3, padding=2, dilation=2, bias=False), nn.GroupNorm(12, base_channels), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(base_channels, base_channels, 3, padding=4, dilation=4, bias=False), nn.GroupNorm(12, base_channels), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(base_channels, base_channels, 3, padding=1, dilation=1, bias=False), nn.GroupNorm(12, base_channels), nn.LeakyReLU(0.2, inplace=True), ) self.dec = nn.Sequential( nn.Conv2d(base_channels * 2, base_channels, 3, padding=1, bias=False), nn.GroupNorm(12, base_channels), nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(base_channels, 64, 3, padding=1, bias=False), nn.GroupNorm(8, 64), nn.LeakyReLU(0.2, inplace=True), ) self.reconstructor = nn.Conv2d(64, 3, 3, padding=1, bias=False) self._init_weights() def _init_weights(self): for layer in self.modules(): if isinstance(layer, nn.Conv2d): nn.init.kaiming_normal_(layer.weight, nonlinearity="leaky_relu") elif isinstance(layer, nn.GroupNorm): nn.init.constant_(layer.weight, 1) nn.init.constant_(layer.bias, 0) nn.init.zeros_(self.reconstructor.weight) def forward(self, x): input_img = x feat_low = self.enc1(input_img) feat_high = self.enc2(feat_low) context = self.bottleneck(feat_high) decoded = self.dec(torch.cat([context, feat_low], dim=1)) residual = self.reconstructor(decoded) return torch.clamp(input_img + residual, 0.0, 1.0) class TAAState: def __init__(self, variance_gamma=1.5, edge_guard_strength=2.0, edge_guard_min=0.25, edge_guard_max=1.0 ): self.history = None self.frame_id = 0 self.variance_gamma = variance_gamma self.edge_guard_strength = edge_guard_strength self.edge_guard_min = edge_guard_min self.edge_guard_max = edge_guard_max def reset(self): self.history = None self.frame_id = 0 @staticmethod def smoothstep(x: torch.Tensor) -> torch.Tensor: return x * x * (3.0 - 2.0 * x) def update(self, frame, alpha, sensitivity): if self.history is None or self.history.shape != frame.shape: self.history = frame.detach().clone().to(frame.device) return frame local_mean = F.avg_pool2d(frame, kernel_size=3, stride=1, padding=1) local_sq_mean = F.avg_pool2d(frame * frame, kernel_size=3, stride=1, padding=1) local_var = (local_sq_mean - local_mean * local_mean).clamp(min=0.0) local_std = torch.sqrt(local_var + 1e-6) local_min = local_mean - local_std * self.variance_gamma local_max = local_mean + local_std * self.variance_gamma history_clipped = torch.maximum(torch.minimum(self.history, local_max), local_min) diff = torch.abs(frame - history_clipped).mean(dim=1, keepdim=True) raw_diff = torch.abs(frame - self.history).mean(dim=1, keepdim=True) disocclusion = ((raw_diff - sensitivity * 2.0) / (sensitivity + 1e-6)).clamp(0.0, 1.0) disocclusion = self.smoothstep(disocclusion) gray = rgb_luma(frame) edge_x = torch.abs(gray[:, :, :, 1:] - gray[:, :, :, :-1]) edge_y = torch.abs(gray[:, :, 1:, :] - gray[:, :, :-1, :]) edge_x = F.pad(edge_x, (0, 1, 0, 0)) edge_y = F.pad(edge_y, (0, 0, 0, 1)) edge_strength = (edge_x + edge_y).clamp(0.0, 1.0) edge_guard = ( self.edge_guard_max - edge_strength * self.edge_guard_strength ).clamp(self.edge_guard_min, self.edge_guard_max) motion_soft = ((diff - sensitivity) / (sensitivity + 1e-6)).clamp(0.0, 1.0) motion_soft = self.smoothstep(motion_soft) dynamic_alpha = alpha * (1.0 - motion_soft) * edge_guard confidence = torch.exp(-diff * 10.0) confidence = confidence * (1.0 - disocclusion) confidence = confidence.clamp(0.15, 1.0) dynamic_alpha = dynamic_alpha * confidence reject_strength = ((diff - sensitivity * 1.5) / (sensitivity + 1e-6)).clamp(0.0, 1.0) reject_strength = torch.maximum(reject_strength, disocclusion) reject_strength = self.smoothstep(reject_strength) history_clipped = torch.lerp(history_clipped, frame, reject_strength) out = torch.lerp(frame, history_clipped, dynamic_alpha) self.history = out.detach() return out class VideoTAADLAA: def __init__(self): self.net_cache = {} self._net_lock = threading.Lock() self.model_weight = 1.00 self.preset_model_weight = { "Balanced": 1.00, "Detail": 1.38, "Smooth": 0.90, "Auto": 1.10, } self.taa_alpha = 0.10 self.jitter_scale = 0.20 self.edge_threshold = 0.15 self.dlaa_blend_scale = 1.00 self.tone_curve_bias = 0.6 self.highlight_pre_blend = 0.15 self.highlight_post_blend = 0.08 self.highlight_threshold = 0.85 self.highlight_slope = 12.0 self.detail_base_scale = 9.0 self.detail_ref_scale = 0.02 self.detail_min_scale = 6.0 self.detail_max_scale = 12.0 self.detail_min_gain = 0.10 self.detail_max_gain = 0.26 self.detail_edge_boost = 0.35 self.detail_highlight_suppression = 0.5 self.edge_sharp_threshold = 0.08 self.edge_sharp_slope = 12.0 self.edge_aa_slope = 14.0 self.motion_gate_scale = 12.0 self.jitter_motion_damping = 8.0 self.detail_dark_luma_start = 0.20 self.detail_dark_luma_range = 0.30 self.detail_dark_mix_base = 0.6 self.detail_dark_mix_scale = 0.4 self.fine_detail_limit = 0.15 self.edge_detail_limit_scale = 0.7 self.detail_highlight_pre_scale = 0.45 self.detail_highlight_post_scale = 0.50 self.detail_blend_boost = 1.12 self.auto_default_scene_motion = 0.02 self.auto_static_motion_threshold = 0.015 self.auto_balanced_motion_threshold = 0.045 self.luma_boost_base = 0.03 self.saturation_boost_base = 0.06 self.luma_highlight_protect = 0.6 self.saturation_highlight_protect = 0.5 self.texture_pass_enabled = True self.texture_strength = 1.20 self.texture_limit = 0.05 self.texture_edge_threshold = 0.06 self.texture_edge_slope = 18.0 self.texture_line_suppression = 0.35 self.texture_motion_suppression = 0.70 self.texture_highlight_suppression = 0.50 self.texture_tile_overlap = 32 self.texture_log_interval = 30 self.texture_net_cache = {} self._texture_missing_warned = False @classmethod def INPUT_TYPES(cls): return { "required": { "images" : ("IMAGE",), "preset": (["Auto", "Balanced", "Detail", "Smooth"],), } } RETURN_TYPES = ("IMAGE",) FUNCTION = "execute" CATEGORY = "CustomPostProcess" def _net(self, device): key = str(device) if key in self.net_cache: return self.net_cache[key] # Model loading can be triggered from parallel executions. with self._net_lock: if key in self.net_cache: return self.net_cache[key] net = DLAANet().to(device) base_path = os.path.dirname(os.path.realpath(__file__)) safetensors_path = os.path.join(base_path, "DLAANet.safetensors") pth_path = os.path.join(base_path, "DLAANet.pth") if os.path.exists(safetensors_path): state_dict = load_file(safetensors_path, device=str(device)) log.info("[DLAA] Loaded DLAANet.safetensors") elif os.path.exists(pth_path): state_dict = torch.load(pth_path, map_location=device) log.info("[DLAA] Loaded DLAANet.pth") else: raise FileNotFoundError( f"[DLAA] Model not found. Expected: {safetensors_path} or {pth_path}" ) if "jitter_offsets" in state_dict: del state_dict["jitter_offsets"] net.load_state_dict(state_dict, strict=False) net = net.float() net.eval() n_params = sum(p.numel() for p in net.parameters()) log.info(f"[DLAA] Model parameters: {n_params/1e6:.2f}M") self.net_cache[key] = net return self.net_cache[key] def _texture_net(self, device): if not self.texture_pass_enabled: self.texture_net_cache = {} return None if ModelLoader is None: if not self._texture_missing_warned: log.warning("[DLAA] Texture model loader is not available, skipping texture pass.") self._texture_missing_warned = True return None key = str(device) if key in self.texture_net_cache: return self.texture_net_cache[key] with self._net_lock: if key in self.texture_net_cache: return self.texture_net_cache[key] base_path = os.path.dirname(os.path.realpath(__file__)) pth_path = os.path.join(base_path, "DLAATexture.pth") safetensors_path = os.path.join(base_path, "DLAATexture.safetensors") if os.path.exists(safetensors_path): model_path = safetensors_path elif os.path.exists(pth_path): model_path = pth_path else: if not self._texture_missing_warned: log.info("[DLAA] Texture pass model not found, skipping texture pass.") self._texture_missing_warned = True return None try: if model_path.endswith(".safetensors"): state_dict = load_file(model_path, device="cpu") else: state_dict = comfy_utils.load_torch_file(model_path, safe_load=False) if isinstance(state_dict, dict): if "params_ema" in state_dict: state_dict = state_dict["params_ema"] elif "params" in state_dict: state_dict = state_dict["params"] elif "state_dict" in state_dict: state_dict = state_dict["state_dict"] elif "model" in state_dict: state_dict = state_dict["model"] net = ModelLoader().load_from_state_dict(state_dict).eval() net = net.to(device) self.texture_net_cache[key] = net log.info( "[DLAA] Loaded texture model: %s", os.path.basename(model_path) ) except Exception as e: if not self._texture_missing_warned: import traceback log.error(f"[DLAA] Texture model exception:\n{traceback.format_exc()}") log.warning( f"[DLAA] Texture model could not be loaded, skipping texture pass: {type(e).__name__}: {e!r}" ) self._texture_missing_warned = True return None return self.texture_net_cache[key] def _tiled_forward(self, net, x: torch.Tensor, tile_size: int = 512, overlap: int = 32) -> torch.Tensor: B, C, H, W = x.shape if H <= tile_size and W <= tile_size: return torch.clamp(net(x), 0.0, 1.0) step = tile_size - overlap * 2 # Invalid overlap would stall the tile loops. if step <= 0: raise ValueError( f"Invalid tiling settings: tile_size={tile_size}, overlap={overlap}" ) out = torch.zeros_like(x) weight = torch.zeros(B, 1, H, W, device=x.device) y0 = 0 while y0 < H: y1 = min(y0 + tile_size, H) y0_c = max(0, y1 - tile_size) x0 = 0 while x0 < W: x1 = min(x0 + tile_size, W) x0_c = max(0, x1 - tile_size) tile = x[:, :, y0_c:y1, x0_c:x1] dlaa_out_tile = net(tile) th, tw = dlaa_out_tile.shape[2], dlaa_out_tile.shape[3] w_map = torch.ones(1, 1, th, tw, device=x.device) ramp = min(overlap, th // 4, tw // 4) if ramp > 0: for k in range(ramp): v = (k + 1) / (ramp + 1) w_map[:, :, k, :] = torch.clamp(w_map[:, :, k, :], max=v) w_map[:, :, th-1-k, :] = torch.clamp(w_map[:, :, th-1-k, :], max=v) w_map[:, :, :, k] = torch.clamp(w_map[:, :, :, k], max=v) w_map[:, :, :, tw-1-k] = torch.clamp(w_map[:, :, :, tw-1-k], max=v) out[:, :, y0_c:y1, x0_c:x1] += dlaa_out_tile * w_map weight[:, :, y0_c:y1, x0_c:x1] += w_map if x1 == W: break x0 += step if y1 == H: break y0 += step return torch.clamp(out / weight.clamp(min=1e-6), 0.0, 1.0) def _jitter(self, x, idx, scale, net): if scale < 1e-5: return x off = net.jitter_offsets[idx % net.jitter_offsets.shape[0]] B, C, H, W = x.shape theta = torch.eye(2, 3, device=x.device).unsqueeze(0).repeat(B, 1, 1) theta[:, 0, 2] = off[0] * scale / W theta[:, 1, 2] = off[1] * scale / H grid = F.affine_grid(theta, x.shape, align_corners=False) return F.grid_sample(x, grid, mode="bilinear", padding_mode="reflection", align_corners=False) def _edge_aa(self, x, thr, blur_radius, net): if blur_radius <= 0: return x gray = rgb_luma(x) sx = F.conv2d(gray, net.sobel_x, padding=1) sy = F.conv2d(gray, net.sobel_y, padding=1) edge = torch.sqrt(sx*sx + sy*sy + 1e-6) mask = torch.sigmoid((edge - thr) * self.edge_aa_slope) blurred = F.avg_pool2d( F.pad(x, [blur_radius]*4, mode="reflect"), blur_radius*2+1, stride=1 ) return x*(1.0 - mask) + blurred*mask def _temporal_refine(self, current, previous, strength=0.35, motion_threshold=0.08): if previous is None: return current if previous.shape != current.shape: return current curr_luma = rgb_luma(current) prev_luma = rgb_luma(previous) motion = torch.abs(curr_luma - prev_luma) motion = torch.clamp(motion / motion_threshold, 0.0, 1.0) blend_mask = (1.0 - motion) * strength refined = current * (1.0 - blend_mask) + previous * blend_mask return refined.clamp(0.0, 1.0) def _resolve_frame_config(self, preset, is_single_image, rgb, taa): if preset == "Auto": if taa.history is not None and taa.history.shape == rgb.shape: scene_motion = torch.abs(rgb - taa.history).mean().item() else: scene_motion = self.auto_default_scene_motion if scene_motion < self.auto_static_motion_threshold: frame_cfg = AUTO_STATIC elif scene_motion < self.auto_balanced_motion_threshold: frame_cfg = AUTO_BALANCED else: frame_cfg = AUTO_MOTION else: frame_cfg = PRESETS.get(preset, PRESETS["Balanced"]) cfg = frame_cfg.copy() if preset == "Detail" and is_single_image: cfg.update(SINGLE_IMAGE_DETAIL) return cfg def _apply_texture_pass( self, texture_net, dlaa_out, tile_size, motion_gate, dark_mask, highlight_mask, frame_index=None ): if texture_net is None: return dlaa_out gen_out = self._tiled_forward( texture_net, dlaa_out, tile_size=tile_size, overlap=self.texture_tile_overlap ) if gen_out.shape != dlaa_out.shape: log.warning( "[DLAA] Texture model output size does not match input, skipping texture pass." ) return dlaa_out gray = rgb_luma(dlaa_out) edge_x = torch.abs(gray[:, :, :, 1:] - gray[:, :, :, :-1]) edge_y = torch.abs(gray[:, :, 1:, :] - gray[:, :, :-1, :]) edge_x = F.pad(edge_x, (0, 1, 0, 0)) edge_y = F.pad(edge_y, (0, 0, 0, 1)) thin_edge_mask = (edge_x + edge_y).clamp(0.0, 1.0) thin_edge_mask = torch.sigmoid( (thin_edge_mask - self.texture_edge_threshold) * self.texture_edge_slope ) thin_edge_mask = F.avg_pool2d(thin_edge_mask, kernel_size=3, stride=1, padding=1) # Keep only the model texture gen_blur = F.avg_pool2d(gen_out, kernel_size=7, stride=1, padding=3) hallucinated_texture = gen_out - gen_blur raw_gen_delta = hallucinated_texture.abs().mean().item() texture_mask = dark_mask.expand_as(dlaa_out[:, 0:1, :, :]) texture_mask = texture_mask * ( 1.0 - highlight_mask * self.texture_highlight_suppression ) texture_mask = texture_mask * ( 1.0 - motion_gate * self.texture_motion_suppression ) texture_mask = texture_mask * ( 1.0 - thin_edge_mask * self.texture_line_suppression ) hallucinated_texture = hallucinated_texture * texture_mask hallucinated_texture = hallucinated_texture.clamp( -self.texture_limit, self.texture_limit ) out = dlaa_out + hallucinated_texture * self.texture_strength final_texture_delta = (out - dlaa_out).abs().mean().item() if ( frame_index is not None and self.texture_log_interval > 0 and frame_index % self.texture_log_interval == 0 ): log.info( "[DLAA] texture frame=%d raw_delta=%.6f final_delta=%.6f", frame_index, raw_gen_delta, final_texture_delta ) return torch.clamp(out, 0.0, 1.0) def execute(self, images, preset): if preset == "Sharp": preset = "Detail" elif preset == "Cinematic": preset = "Smooth" device = mm.get_torch_device() if mm else \ torch.device("cuda" if torch.cuda.is_available() else "cpu") B, H, W, C = images.shape is_single_image = (B == 1) if is_single_image: blur_radius = 0 else: blur_radius = 0 if preset == "Detail" else 1 # Temporal state belongs to this execution only. taa = TAAState() prev_dlaa_output = None net = self._net(device) texture_net = self._texture_net(device) if preset == "Detail" else None out_list = [] delta_sum = 0.0 delta_count = 0 if mm is not None: try: vram_mb = torch.cuda.get_device_properties(device).total_memory // (1024 * 1024) except: vram_mb = 8192 else: vram_mb = 8192 if vram_mb <= 8192: tile_size = 512 elif vram_mb <= 16384: tile_size = 1024 else: tile_size = 1024 if H > tile_size or W > tile_size: tile_step = tile_size - 64 tile_count = ((H + tile_step - 1) // tile_step) * ((W + tile_step - 1) // tile_step) log.debug(f"[DLAA] Tiled inference: {tile_count} tiles") progress = ProgressBar(B) if ProgressBar is not None else None with torch.inference_mode(): for i in range(B): img = images[i:i+1].to(device).permute(0, 3, 1, 2).float() rgb = img[:, :3] motion_gate = 0.0 # Auto is resolved per frame, after history is available. frame_cfg = self._resolve_frame_config( preset, is_single_image, rgb, taa ) detail_boost = frame_cfg["detail_boost"] edge_boost = frame_cfg["edge_boost"] temporal_strength = frame_cfg["temporal_strength"] micro_limit = frame_cfg["micro_limit"] luma_boost_mult = frame_cfg["luma_boost_mult"] saturation_boost_mult = frame_cfg["saturation_boost_mult"] motion_threshold = frame_cfg["motion_threshold"] taa_strength = frame_cfg["taa_strength"] dlaa_strength = frame_cfg["dlaa_strength"] tone_strength = frame_cfg["tone_strength"] edge_sharp_strength = frame_cfg["edge_sharp_strength"] motion_sensitivity = frame_cfg["motion_sensitivity"] preset_jitter_scale = frame_cfg["jitter_scale"] fid = taa.frame_id taa.frame_id = (fid + 1) % net.jitter_offsets.shape[0] # Jitter is damped on motion to avoid shimmer. adaptive_jitter_scale = preset_jitter_scale if taa.history is not None and taa.history.shape == rgb.shape: motion_estimate = torch.abs(rgb - taa.history).mean() motion_gate = torch.clamp( motion_estimate * self.motion_gate_scale, 0.0, 1.0 ) jitter_damping = ( 1.0 - motion_estimate * self.jitter_motion_damping ).clamp(0.45, 1.0) adaptive_jitter_scale = preset_jitter_scale * jitter_damping rgb = self._jitter(rgb, fid, adaptive_jitter_scale, net) rgb = self._edge_aa(rgb, self.edge_threshold, blur_radius, net) taa_out = taa.update(rgb, self.taa_alpha, motion_sensitivity) rgb = torch.lerp(rgb, taa_out, taa_strength) if dlaa_strength > 0: current_tile_size = tile_size min_tile_size = 128 dlaa_out = None last_oom_error = None while current_tile_size >= min_tile_size: try: dlaa_out = self._tiled_forward( net, rgb, tile_size=current_tile_size, overlap=32 ) if i == 0: raw_model_delta = (dlaa_out - rgb).abs().mean().item() log.info(f"[DLAA] raw_model_delta={raw_model_delta:.6f}") break except torch.OutOfMemoryError as e: last_oom_error = e if torch.cuda.is_available(): torch.cuda.empty_cache() next_tile_size = current_tile_size // 2 if next_tile_size < min_tile_size: log.error( "DLAA tiled inference failed at minimum tile size %d.", current_tile_size ) raise last_oom_error log.warning( "Out of VRAM at tile size %d, retrying with %d.", current_tile_size, next_tile_size ) current_tile_size = next_tile_size if dlaa_out is None: raise RuntimeError("DLAA tiled inference failed.") # Keep the model pass from shifting overall brightness. dlaa_mean = dlaa_out.mean(dim=(1,2,3), keepdim=True) rgb_mean = rgb.mean(dim=(1,2,3), keepdim=True) dlaa_out = dlaa_out - dlaa_mean + rgb_mean preset_model_weight = self.preset_model_weight.get(preset, 1.00) # The network is used as a residual refiner, not a full replacement. model_delta = dlaa_out - rgb model_delta_value = model_delta.abs().mean().item() delta_sum += model_delta_value delta_count += 1 dlaa_out = torch.clamp( rgb + model_delta * self.model_weight * preset_model_weight, 0.0, 1.0 ) if i == 0: log.info(f"[DLAA] model_delta_first={model_delta_value:.6f}") luma = rgb_luma(dlaa_out) highlight_mask = torch.sigmoid((luma - self.highlight_threshold) * self.highlight_slope) # Bright areas are easy to overcook, so pull them back early. highlight_pre_blend = self.highlight_pre_blend * ( self.detail_highlight_pre_scale if preset == "Detail" else 1.0 ) dlaa_out = torch.lerp(dlaa_out, rgb, highlight_mask * highlight_pre_blend) dlaa_out = torch.clamp(dlaa_out, 0.0, 1.0) motion_cfg = MOTION_SUPPRESSION["Detail"] if preset == "Detail" else MOTION_SUPPRESSION["Default"] detail_boost *= (1.0 - motion_gate * motion_cfg["detail"]) edge_boost *= (1.0 - motion_gate * motion_cfg["edge"]) micro_limit *= (1.0 - motion_gate * motion_cfg["micro"]) # Small residual detail pass after the model output. local_avg_rgb = F.avg_pool2d(dlaa_out, 3, stride=1, padding=1) fine_detail_rgb = dlaa_out - local_avg_rgb luma = rgb_luma(dlaa_out) fine_detail = rgb_luma(fine_detail_rgb) if preset in ["Detail"]: dark_mask = torch.clamp( (luma - self.detail_dark_luma_start) / self.detail_dark_luma_range, 0.0, 1.0 ) micro_limit = micro_limit * ( self.detail_dark_mix_base + self.detail_dark_mix_scale * dark_mask ) else: dark_mask = 1.0 detail_strength = fine_detail.abs().mean(dim=(1,2,3), keepdim=True) detail_scale = (self.detail_base_scale + (self.detail_ref_scale / (detail_strength + 1e-6))).clamp(self.detail_min_scale, self.detail_max_scale) fine_detail = fine_detail * torch.sigmoid(fine_detail * detail_scale) fine_detail = fine_detail.clamp(-self.fine_detail_limit, self.fine_detail_limit) local_detail = F.avg_pool2d(fine_detail.abs(), 7, stride=1, padding=3) global_detail = fine_detail.abs().mean(dim=(1,2,3), keepdim=True) edge_for_detail = torch.sqrt( F.conv2d(luma, net.sobel_x, padding=1) ** 2 + F.conv2d(luma, net.sobel_y, padding=1) ** 2 + 1e-6 ) edge_detail_weight = torch.sigmoid((edge_for_detail - self.edge_sharp_threshold) * self.edge_sharp_slope) detail_gain = (global_detail / (local_detail + 1e-6)).clamp(self.detail_min_gain, self.detail_max_gain) detail_gain = detail_gain * (1.0 - local_detail.clamp(0.0, 0.5)) detail_gain = detail_gain * (1.0 + edge_detail_weight * self.detail_edge_boost) detail_gain = detail_gain * (1.0 - highlight_mask * self.detail_highlight_suppression) micro_detail = fine_detail_rgb * detail_gain * detail_boost * dark_mask micro_detail = micro_detail.clamp(-micro_limit, micro_limit) dlaa_out = dlaa_out + micro_detail edge_mask = torch.sigmoid((edge_for_detail - self.edge_sharp_threshold) * self.edge_sharp_slope) edge_detail = fine_detail_rgb * edge_mask * (1.0 - highlight_mask) edge_boosted = edge_detail * edge_sharp_strength * edge_boost * dark_mask edge_boosted = edge_boosted.clamp( -micro_limit * self.edge_detail_limit_scale, micro_limit * self.edge_detail_limit_scale ) dlaa_out = dlaa_out + edge_boosted dlaa_out = self._apply_texture_pass( texture_net=texture_net, dlaa_out=dlaa_out, tile_size=current_tile_size, motion_gate=motion_gate, dark_mask=dark_mask, highlight_mask=highlight_mask, frame_index=i ) tone_mapped = dlaa_out / (dlaa_out + self.tone_curve_bias) dlaa_out = torch.lerp(dlaa_out, tone_mapped, highlight_mask * tone_strength) luma = rgb_luma(dlaa_out) luma_boost = self.luma_boost_base * luma_boost_mult * (1.0 - highlight_mask * self.luma_highlight_protect) dlaa_out = dlaa_out * (1.0 + luma_boost) mean_rgb = dlaa_out.mean(dim=1, keepdim=True) saturation_boost = ( self.saturation_boost_base * saturation_boost_mult * (1.0 - highlight_mask * self.saturation_highlight_protect) ) dlaa_out = mean_rgb + (dlaa_out - mean_rgb) * (1.0 + saturation_boost) highlight_post_blend = self.highlight_post_blend * ( self.detail_highlight_post_scale if preset == "Detail" else 1.0 ) dlaa_out = torch.lerp(dlaa_out, rgb, highlight_mask * highlight_post_blend) dlaa_out = torch.clamp(dlaa_out, 0.0, 1.0) if prev_dlaa_output is not None: if prev_dlaa_output.shape != dlaa_out.shape: prev_dlaa_output = None # Final temporal pass; Detail may keep this disabled. dlaa_out = self._temporal_refine( dlaa_out, prev_dlaa_output, strength=temporal_strength, motion_threshold=motion_threshold ) prev_dlaa_output = dlaa_out.detach() dlaa_out = torch.clamp(dlaa_out, 0.0, 1.0) blend_weight = dlaa_strength * self.dlaa_blend_scale if preset in ["Detail"]: blend_weight = min(blend_weight * self.detail_blend_boost, 1.0) rgb = torch.lerp(rgb, dlaa_out, blend_weight) rgb = torch.clamp(rgb, 0.0, 1.0) out_list.append(rgb.permute(0, 2, 3, 1).cpu()) if progress is not None: progress.update(1) if mm is not None and i > 0 and i % 50 == 0: mm.soft_empty_cache() if delta_count > 0: log.info(f"[DLAA] model_delta_avg={delta_sum / delta_count:.6f}") return (torch.cat(out_list, dim=0),) NODE_CLASS_MAPPINGS = {"VideoTAADLAA": VideoTAADLAA} NODE_DISPLAY_NAME_MAPPINGS = {"VideoTAADLAA": "🎮 Video Anti-Aliasing (TAA + DLAA)"}