SOLRICKS-comfyui-solricks/VideoTAADLAA.py

# Created by MSXYZ (AI-assisted)
# Temporal Anti-Aliasing (TAA) + Lightweight DLAA-style Sharpening
# v0.1.1 - DLAANet with residual learning and orthogonal weights

import torch
import torch.nn as nn
import torch.nn.functional as F
import logging
import comfy.model_management as mm

logger = logging.getLogger("VideoTAADLAA")

class _DLAANet(nn.Module):
    def __init__(self):
        super().__init__()
        
        # Sobel ve Jitter tanımları
        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))

        # Detay ayrıştırıcı
        self.extract_feature = nn.Conv2d(3, 16, 3, padding=1, bias=False)
        
        # Doku koruyucu
        self.refiner = nn.Conv2d(16, 16, 3, padding=1, bias=False)
        
        # Artifactleri temizle
        self.reconstructor = nn.Conv2d(16, 3, 3, padding=1, bias=False)
        
        # DLAANet with residual learning and orthogonal weights
        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 for stability
        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)

        # Motion masking
        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.30, "min": 0, "max": 0.9, "step": 0.01}),
                "motion_sensitivity": ("FLOAT", {"default": 0.05, "min": 0.0, "max": 0.3, "step": 0.01}),
                "jitter_scale": ("FLOAT", {"default": 0.02, "min": 0, "max": 0.08, "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.35, "step": 0.01}),
                "blur_radius": ("INT", {"default": 0, "min": 0, "max": 5, "step": 1}),
                "reset_history": ("BOOLEAN", {"default": False}),
            }
        }

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

    def _net(self, device):
        key = str(device)
        if key not in self.net_cache:
            self.net_cache[key] = _DLAANet().to(device).eval()
        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
        
        # Grayscale for edge detection
        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=False):

        device = mm.get_torch_device()

        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()
                # split RGB channels
                rgb = img[:, :3]

                # 1. spatial preprocessing step
                fid = self.taa.frame_id
                self.taa.frame_id = (fid + 1) % 10000
                
                rgb = self.jitter(rgb, fid, jitter_scale, net)
                rgb = self.edge_aa(rgb, edge_threshold, blur_radius, net)
                
                # 2. temporal update
                taa_out = self.taa.update(rgb, taa_alpha, motion_sensitivity)
                rgb = torch.lerp(rgb, taa_out, taa_strength)
                
                # 3. sharpening pass
                if dlaa_strength > 0:
                    # compute luma from original image
                    luma_orig = 0.299*rgb[:,0:1] + 0.587*rgb[:,1:2] + 0.114*rgb[:,2:3]
                    refined_output = net(rgb)
                    residual = refined_output - rgb
                    
                    # apply residual mostly to luminance to avoid color shifts
                    luma_res = 0.299*residual[:,0:1] + 0.587*residual[:,1:2] + 0.114*residual[:,2:3]
                    rgb = rgb + (luma_res * dlaa_strength * 1.25)
                    
                    # slight contrast adjustment
                    mean_luma = torch.mean(luma_orig, dim=(1,2,3), keepdim=True)
                    rgb = (rgb - mean_luma) * 1.04 + mean_luma
                    
                    rgb = torch.lerp(luma_orig, rgb, 1.15)
                    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"}