SOLRICKS-comfyui-solricks/VideoTAADLAA.py

# Created by MSXYZ (AI-assisted)
# Temporal Anti-Aliasing (TAA) + DLAA-style sharpening
# v0.1.1 - Temporal stability improvements, ghosting reduction

import torch
import torch.nn.functional as F
import logging
import comfy.model_management as mm
import cv2
import numpy as np

logger = logging.getLogger("VideoTAADLAA")


class _TAAState:
    def __init__(self):
        self.history: torch.Tensor | None = None
        self.prev_frame: torch.Tensor | None = None
        self.frame_id: int = 0

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

    @staticmethod
    def _optical_flow(prev: torch.Tensor, curr: torch.Tensor) -> torch.Tensor:
        # Move to CPU for OpenCV. Creates a bottleneck, but Farneback is far more 
        # stable for TAA motion vectors than native PyTorch approximations.
        def to_gray_u8(t: torch.Tensor) -> np.ndarray:
            arr = t[0].permute(1, 2, 0).detach().cpu().numpy()
            return cv2.cvtColor((arr * 255).astype(np.uint8), cv2.COLOR_RGB2GRAY)

        prev_gray = to_gray_u8(prev)
        curr_gray = to_gray_u8(curr)

        flow_np = cv2.calcOpticalFlowFarneback(
            prev_gray, curr_gray, None,
            pyr_scale=0.5, levels=3, winsize=15,
            iterations=3, poly_n=5, poly_sigma=1.2, flags=0,
        )

        H, W = flow_np.shape[:2]
        flow = torch.from_numpy(flow_np).to(curr.device).float()

        # Normalize to [-1, 1] for grid_sample
        flow[..., 0] /= (W / 2.0)
        flow[..., 1] /= (H / 2.0)

        return flow.unsqueeze(0)

    def update(self, frame: torch.Tensor, alpha: float, sensitivity: float) -> torch.Tensor:

        if self.history is None:
            self.history = frame.clone()
            self.prev_frame = frame.clone()
            self.frame_id += 1
            return frame

        B, C, H, W = frame.shape

        # Rec. 601 Luma
        luma = (
            frame[:, 0:1] * 0.299 +
            frame[:, 1:2] * 0.587 +
            frame[:, 2:3] * 0.114
        )
        
        # Highlight mask to prevent flickering in bright areas (Tuned: 0.75 threshold)
        highlight_mask = torch.sigmoid((luma - 0.75) * 15.0)

        # 3x3 Neighborhood AABB Clipping (Negative pooling trick for max efficiency)
        local_min = -F.max_pool2d(-frame, 3, 1, 1)
        local_max = F.max_pool2d(frame, 3, 1, 1)
        history_clamped = torch.minimum(torch.maximum(self.history, local_min), local_max)

        flow = self._optical_flow(self.prev_frame, frame)

        base_grid_y, base_grid_x = torch.meshgrid(
            torch.linspace(-1, 1, H, device=frame.device),
            torch.linspace(-1, 1, W, device=frame.device),
            indexing="ij",
        )

        base_grid = torch.stack((base_grid_x, base_grid_y), dim=-1)
        sample_grid = (base_grid.unsqueeze(0) + flow).clamp(-1, 1)

        warped_history = F.grid_sample(
            self.history,
            sample_grid,
            mode="bilinear",
            padding_mode="border",
            align_corners=False,
        )

        # Disocclusion mask: Drop history if flow diverges too much (threshold: 0.08)
        diff_flow = (frame - warped_history).abs().mean(dim=1, keepdim=True)
        occ_mask = torch.sigmoid((diff_flow - 0.08) * 20.0)

        history_mix = torch.lerp(warped_history, history_clamped, occ_mask)

        motion = (frame - self.prev_frame).abs().mean(dim=1, keepdim=True)
        reactive = torch.sigmoid(((motion / max(sensitivity, 0.001)) - 0.3) * 10.0)

        # Final blend weight calculation
        temporal_weight = alpha * (1.0 - reactive * 0.9) * (1.0 - highlight_mask * 0.85)

        out = torch.lerp(frame, history_mix, temporal_weight)

        self.history = out.detach()
        self.prev_frame = frame.detach()
        self.frame_id += 1

        return out


def _spatial_sharpen(rgb: torch.Tensor, strength: float) -> torch.Tensor:
    # High-pass filter (Unsharp Mask) to recover details lost during TAA
    blurred = F.avg_pool2d(rgb, 3, 1, 1)
    detail = rgb - blurred
    sharpened = rgb + detail * 0.5
    return torch.lerp(rgb, sharpened, strength).clamp(0.0, 1.0)


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

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "images": ("IMAGE",),
                "taa_alpha": ("FLOAT", {"default": 0.60, "min": 0.00, "max": 0.95, "step": 0.01}),
                "motion_sensitivity": ("FLOAT", {"default": 0.05, "min": 0.01, "max": 0.50, "step": 0.01}),
                "sharpen_strength": ("FLOAT", {"default": 0.60, "min": 0.00, "max": 1.00, "step": 0.05}),
                "reset_history": ("BOOLEAN", {"default": True}),
            }
        }

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

    def execute(self, images, taa_alpha, motion_sensitivity, sharpen_strength, reset_history=True):

        device = mm.get_torch_device()

        if reset_history:
            self.taa.reset()

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

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

                rgb_3 = self.taa.update(rgb_3, taa_alpha, motion_sensitivity)

                if sharpen_strength > 0:
                    rgb_3 = _spatial_sharpen(rgb_3, sharpen_strength)

                if C == 4:
                    rgb = torch.cat([rgb_3, rgb[:, 3:4]], dim=1)
                else:
                    rgb = rgb_3

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

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


NODE_CLASS_MAPPINGS = {
    "VideoTAADLAA": VideoTAADLAA,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "VideoTAADLAA": "🎮 Video TAA + Spatial Sharpen",
}