SOLRICKS-comfyui-solricks/VideoTAADLAA.py

# Created by MSXYZ (AI-assisted)
# Temporal Anti-Aliasing (TAA) + Lightweight DLAA-style Sharpening
# v0.1.1 - Active Weights, MSE Monitoring & Pixel Integrity

import torch
import torch.nn as nn
import torch.nn.functional as F
import os
import logging
try:
    import comfy.model_management as mm
except:
    mm = None

logger = logging.getLogger("VideoTAADLAA")


class _DLAANet(nn.Module):
    def __init__(self):
        super().__init__()
        
        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.25, 0.25], [-0.25, -0.25], [-0.25, 0.25], [0.25, -0.25]], dtype=torch.float32))

        self.extract_feature = nn.Conv2d(3, 16, 3, padding=1, bias=False)
        
        self.refiner = nn.Conv2d(16, 16, 3, padding=1, bias=False)
        self.reconstructor = nn.Conv2d(16, 3, 3, padding=1, bias=False)
        
        self._init_weights()

    def _init_weights(self):
        nn.init.orthogonal_(self.extract_feature.weight)
        nn.init.orthogonal_(self.refiner.weight)
        nn.init.dirac_(self.reconstructor.weight) 

    def forward(self, x):
        features = F.leaky_relu(self.extract_feature(x), 0.2)
        refined = F.leaky_relu(self.refiner(features), 0.2)
        residual = self.reconstructor(refined)
        return x + residual


class _TAAState:
    def __init__(self):
        self.history = None
        self.frame_id = 0

    def reset(self):
        self.history = None
        self.frame_id = 0

    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

        # Neighborhood clamping
        local_min = -F.max_pool2d(-frame, kernel_size=3, stride=1, padding=1)
        local_max = F.max_pool2d(frame, kernel_size=3, stride=1, padding=1)
        history_clipped = torch.clamp(self.history, local_min, local_max)

        diff = torch.abs(frame - history_clipped).mean(dim=1, keepdim=True)
        motion_mask = torch.sigmoid((diff - sensitivity) * 45.0)
        dynamic_alpha = alpha * (1.0 - motion_mask)

        out = torch.lerp(frame, history_clipped, dynamic_alpha)
        self.history = out.detach()
        return out

class VideoTAADLAA:
    def __init__(self):
        self.net_cache = {}
        self.taa = _TAAState()

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("IMAGE",),
                "taa_strength": ("FLOAT", {"default": 0.40, "min": 0, "max": 1, "step": 0.05}),
                "taa_alpha": ("FLOAT", {"default": 0.20, "min": 0, "max": 0.9, "step": 0.01}),
                "motion_sensitivity": ("FLOAT", {"default": 0.1, "min": 0, "max": 0.3, "step": 0.01}),
                "jitter_scale": ("FLOAT", {"default": 0.1, "min": 0, "max": 1, "step": 0.01}),
                "dlaa_strength": ("FLOAT", {"default": 0.40, "min": 0, "max": 1, "step": 0.05}),
                "edge_threshold": ("FLOAT", {"default": 0.25, "min": 0.05, "max": 0.3, "step": 0.01}),
                "blur_radius": ("INT", {"default": 0, "min": 0, "max": 3, "step": 1}),
                "reset_history": ("BOOLEAN", {"default": True}),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "execute"
    CATEGORY = "CustomPostProcess"

    def _net(self, device):
        key = str(device)
        if key not in self.net_cache:
            net = _DLAANet().to(device)
            
            base_path = os.path.dirname(os.path.realpath(__file__))
            path = os.path.join(base_path, "dlaanet.pth")
            
            try:
                state_dict = torch.load(path, map_location=device, weights_only=True)
                net.load_state_dict(state_dict)
            except Exception:
                pass # if an error, do nothing; let the program continue with random/dirac weights
            
            net.eval()
            self.net_cache[key] = net
            
        return self.net_cache[key]

    def jitter(self, x, idx, scale, net):
        if scale < 1e-5: return x
        off = net.jitter_offsets[idx % 4]
        B, C, H, W = x.shape
        
        theta = torch.eye(2, 3, device=x.device).unsqueeze(0).repeat(B, 1, 1)
        theta[:, 0, 2], theta[:, 1, 2] = off[0] * scale / W, 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, net):
        if blur <= 0: return x
        
        gray = 0.299*x[:, 0:1] + 0.587*x[:, 1:2] + 0.114*x[:, 2:3]
        sx, sy = F.conv2d(gray, net.sobel_x, padding=1), F.conv2d(gray, net.sobel_y, padding=1)
        edge = torch.sqrt(sx*sx + sy*sy + 1e-6)
        
        mask = torch.sigmoid((edge - thr) * 8.0)
        blurred = F.avg_pool2d(F.pad(x, [blur]*4, mode="reflect"), blur*2+1, stride=1)
        
        return x*(1-mask) + blurred*mask

    def execute(self, images, taa_strength, taa_alpha, motion_sensitivity,
                jitter_scale, dlaa_strength, edge_threshold,
                blur_radius, reset_history=True):

        device = mm.get_torch_device() if mm else torch.device("cuda" if torch.cuda.is_available() else "cpu")

        if reset_history:
            self.taa.reset()
            torch.cuda.empty_cache()

        net = self._net(device)

        B, H, W, C = images.shape
        out_list = []

        with torch.no_grad():
            for i in range(B):
                # convert to NCHW for PyTorch
                img = images[i:i+1].to(device).permute(0, 3, 1, 2).float()
                rgb = img[:, :3]

                fid = self.taa.frame_id
                self.taa.frame_id = (fid + 1) % 4
                
                rgb = self.jitter(rgb, fid, jitter_scale, net)
                rgb = self.edge_aa(rgb, edge_threshold, blur_radius, net)

                taa_out = self.taa.update(rgb, taa_alpha, motion_sensitivity)
                rgb = torch.lerp(rgb, taa_out, taa_strength)
                
                if dlaa_strength > 0:
                    luma_orig = 0.2126*rgb[:,0:1] + 0.7152*rgb[:,1:2] + 0.0722*rgb[:,2:3]

                    refined_output = torch.clamp(net(rgb), 0.0, 1.0)
                    diff_check = torch.mean((refined_output - rgb) ** 2).item()
                    if i == 0: 
                        print(f"\033[92m[DLAA] Model Delta (MSE): \033[93m{diff_check:.6f}\033[0m")
                    residual = refined_output - rgb

                    luma_res = 0.2126*residual[:,0:1] + 0.7152*residual[:,1:2] + 0.0722*residual[:,2:3]
                    
                    rgb = rgb + (luma_res * dlaa_strength * 1.4)
                    rgb = rgb * (1.0 + dlaa_strength * 0.4)
                    
                    rgb = torch.clamp(rgb, 0.0, 1.0)
                    rgb = torch.pow(rgb, 0.98)
                    rgb = torch.clamp(rgb, 0.0, 1.0)

                out_list.append(rgb.permute(0,2,3,1).cpu())

                if i % 10 == 0:
                    mm.soft_empty_cache()

        return (torch.cat(out_list, dim=0),)

NODE_CLASS_MAPPINGS = {"VideoTAADLAA": VideoTAADLAA}
NODE_DISPLAY_NAME_MAPPINGS = {"VideoTAADLAA": "🎮 Video TAA + DLAA Anti-Aliasing"}