Temporal stability improvements + Optical Flow)
This commit is contained in:
MSXYZ-GenAI
+158
-124
@@ -3,110 +3,159 @@
|
||||
# v0.1.1 - Temporal stability improvements, ghosting reduction
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import logging
|
||||
import comfy.model_management as mm
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
logger = logging.getLogger("VideoTAADLAA")
|
||||
|
||||
# =========================
|
||||
# DLAA CORE
|
||||
# =========================
|
||||
class _DLAANet(nn.Module):
|
||||
"""
|
||||
Lightweight CNN-based sharpening module used as a DL-inspired post-process.
|
||||
|
||||
Not:
|
||||
Bu yapı gerçek bir eğitilmiş DLAA modeli değildir.
|
||||
Sadece inference-time çalışan küçük bir CNN ile
|
||||
detay keskinleştirme ve aliasing azaltma hedeflenmiştir.
|
||||
"""
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
# R, G, B için 3 kanal (groups=3 ile ultra hızlı)
|
||||
self.conv = nn.Conv2d(3, 3, 3, padding=1, bias=False, groups=3)
|
||||
|
||||
# Keskinleştirme çekirdeğini (Sharpening kernel)
|
||||
kernel = torch.tensor([
|
||||
[0.0, -1.0, 0.0],
|
||||
[-1.0, 5.0, -1.0],
|
||||
[0.0, -1.0, 0.0]
|
||||
], dtype=torch.float32) * 0.1
|
||||
|
||||
# Ağırlıkları tek seferde sabitliyoruz
|
||||
with torch.no_grad():
|
||||
for i in range(3):
|
||||
self.conv.weight[i, 0] = kernel
|
||||
|
||||
# Diğer fonksiyonlar (edge_aa ve sobel)
|
||||
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))
|
||||
|
||||
# 8-Noktalı Jitter Sekansı (Daha iyi TAA sonucu sağlar)
|
||||
self.register_buffer("jitter_offsets", torch.tensor([
|
||||
[0.125, 0.375], [-0.375, -0.125], [0.375, -0.375], [-0.125, 0.125],
|
||||
[0.250, 0.250], [-0.250, -0.250], [-0.250, 0.250], [0.250, -0.250]
|
||||
], dtype=torch.float32))
|
||||
|
||||
def forward(self, x):
|
||||
return torch.clamp(x + self.conv(x) * 0.12, 0.0, 1.0)
|
||||
|
||||
# =========================
|
||||
# TEMPORAL STATE
|
||||
# =========================
|
||||
class _TAAState:
|
||||
"""
|
||||
TAA (Zamansal Kenar Yumuşatma) için önceki karelerin geçmişini tutar.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.history = None
|
||||
self.frame_id = 0
|
||||
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
|
||||
|
||||
def update(self, frame, alpha, sensitivity):
|
||||
if self.history is None or self.history.shape != frame.shape:
|
||||
# Internal helpers
|
||||
@staticmethod
|
||||
def _optical_flow(prev: torch.Tensor, curr: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
CPU üzerinde akış (hareket) hesaplar ve düzeltme tensörü döndürür.
|
||||
Boyut: (1, H, W, 2)
|
||||
"""
|
||||
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)
|
||||
|
||||
# Calculate flow
|
||||
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, 2) — raw pixel offsets
|
||||
|
||||
H, W = flow_np.shape[:2]
|
||||
flow = torch.from_numpy(flow_np).to(curr.device).float() # (H, W, 2)
|
||||
|
||||
# Since the grid ranges from -1 to 1 (a total distance of 2),
|
||||
flow[..., 0] /= (W / 2.0)
|
||||
flow[..., 1] /= (H / 2.0)
|
||||
|
||||
return flow.unsqueeze(0) # (1, H, W, 2)
|
||||
|
||||
def update(
|
||||
self,
|
||||
frame: torch.Tensor, # (1, 3, H, W) float32
|
||||
alpha: float, # base temporal blend weight
|
||||
sensitivity: float, # motion sensitivity
|
||||
) -> torch.Tensor:
|
||||
|
||||
if self.history is None:
|
||||
self.history = frame.clone()
|
||||
self.prev_frame = frame.clone()
|
||||
self.frame_id += 1
|
||||
return frame
|
||||
|
||||
# Neighborhood clamping: geçmiş kareyi mevcut frame'in 3x3 lokal min/max sınırına çeker
|
||||
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)
|
||||
B, C, H, W = frame.shape
|
||||
|
||||
# Hareketli bölgelerde alpha'yı düşür: ghosting azalır, keskin geçişler korunur
|
||||
diff = torch.abs(frame - history_clipped).mean(dim=1, keepdim=True)
|
||||
motion_mask = torch.sigmoid((diff - sensitivity) * 20.0)
|
||||
dynamic_alpha = alpha * (1.0 - motion_mask)
|
||||
# highlight mask
|
||||
luma = (
|
||||
frame[:, 0:1] * 0.299
|
||||
+ frame[:, 1:2] * 0.587
|
||||
+ frame[:, 2:3] * 0.114
|
||||
)
|
||||
highlight_mask = torch.sigmoid((luma - 0.75) * 15.0) # (1,1,H,W)
|
||||
|
||||
out = torch.lerp(frame, history_clipped, dynamic_alpha)
|
||||
self.history = out.detach() # detach(): gradient graph birikimini önler
|
||||
# Neighbourhood clamping
|
||||
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)
|
||||
|
||||
# Optical-flow warp of history
|
||||
flow = self._optical_flow(self.prev_frame, frame) # (1, H, W, 2)
|
||||
|
||||
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) # (H, W, 2)
|
||||
sample_grid = (base_grid.unsqueeze(0) + flow).clamp(-1, 1) # (1, H, W, 2)
|
||||
|
||||
warped_history = F.grid_sample(
|
||||
self.history,
|
||||
sample_grid,
|
||||
mode="bilinear",
|
||||
padding_mode="border",
|
||||
align_corners=False,
|
||||
) # (1, 3, H, W)
|
||||
|
||||
# Occlusion
|
||||
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-reactive
|
||||
motion = (frame - self.prev_frame).abs().mean(dim=1, keepdim=True)
|
||||
reactive = torch.sigmoid(((motion / max(sensitivity, 0.001)) - 0.3) * 10.0)
|
||||
|
||||
temporal_weight = (
|
||||
alpha
|
||||
* (1.0 - reactive * 0.9)
|
||||
* (1.0 - highlight_mask * 0.85)
|
||||
) # per-pixel [0, alpha]
|
||||
|
||||
out = torch.lerp(frame, history_mix, temporal_weight)
|
||||
|
||||
self.history = out.detach()
|
||||
self.prev_frame = frame.detach()
|
||||
self.frame_id += 1
|
||||
|
||||
return out
|
||||
|
||||
# =========================
|
||||
# MAIN NODE
|
||||
# =========================
|
||||
|
||||
# DLAA-style sharpening
|
||||
def _dlaa_sharpen(rgb: torch.Tensor, strength: float) -> torch.Tensor:
|
||||
"""
|
||||
DLAA detay geri kazanımına benzer sonuçlar veren, sistemi yormayan
|
||||
keskinleştirme işlemi. Herhangi bir çözünürlükte, yapay zeka modeli
|
||||
(öğrenilmiş ağırlıklar) kullanmadan çalışır.
|
||||
"""
|
||||
blurred = F.avg_pool2d(rgb, kernel_size=3, stride=1, padding=1)
|
||||
detail = rgb - blurred # high-frequency residual
|
||||
sharpened = rgb + detail * 0.5 # amplify edges
|
||||
return torch.lerp(rgb, sharpened, strength).clamp(0.0, 1.0)
|
||||
|
||||
|
||||
# COMFYUI NODE
|
||||
class VideoTAADLAA:
|
||||
def __init__(self):
|
||||
self.net_cache = {} # Device başına ayrı model (CPU/GPU desteği)
|
||||
self.taa = _TAAState()
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"taa_strength": ("FLOAT", {"default": 0.7, "min": 0, "max": 1, "step": 0.05}),
|
||||
"taa_alpha": ("FLOAT", {"default": 0.6, "min": 0, "max": 0.95, "step": 0.01}),
|
||||
"motion_sensitivity": ("FLOAT", {"default": 0.05, "min": 0.01, "max": 0.5, "step": 0.01}),
|
||||
"jitter_scale": ("FLOAT", {"default": 0.0, "min": 0, "max": 1, "step": 0.01}),
|
||||
"dlaa_strength": ("FLOAT", {"default": 0.6, "min": 0, "max": 1, "step": 0.05}),
|
||||
"edge_threshold": ("FLOAT", {"default": 0.25, "min": 0, "max": 1, "step": 0.01}),
|
||||
"blur_radius": ("INT", {"default": 1, "min": 0, "max": 5, "step": 1}),
|
||||
"reset_history": ("BOOLEAN", {"default": True}),
|
||||
"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}),
|
||||
"dlaa_strength": ("FLOAT", {"default": 0.60, "min": 0.00, "max": 1.00, "step": 0.05}),
|
||||
"reset_history": ("BOOLEAN", {"default": True}),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -114,71 +163,56 @@ class VideoTAADLAA:
|
||||
FUNCTION = "execute"
|
||||
CATEGORY = "CustomPostProcess"
|
||||
|
||||
def _net(self, device):
|
||||
if device not in self.net_cache:
|
||||
self.net_cache[device] = _DLAANet().to(device).eval()
|
||||
return self.net_cache[device]
|
||||
def execute(
|
||||
self,
|
||||
images: torch.Tensor,
|
||||
taa_alpha: float,
|
||||
motion_sensitivity: float,
|
||||
dlaa_strength: float,
|
||||
reset_history: bool = True,
|
||||
) -> tuple[torch.Tensor]:
|
||||
|
||||
def jitter(self, x, idx, scale, net):
|
||||
if scale < 1e-5: return x
|
||||
off = net.jitter_offsets[idx % 8] # 8 noktaya güncellendi
|
||||
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) # epsilon: NaN önlemi
|
||||
mask = torch.sigmoid((edge - thr) * 12.0)
|
||||
blurred = F.avg_pool2d(F.pad(x, [blur]*4, mode="reflect"), blur*2+1, stride=1)
|
||||
return torch.lerp(x, blurred, mask) # Optimizasyon: x*(1-mask) + blurred*mask yerine lerp kullanıldı.
|
||||
|
||||
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 reset_history: self.taa.reset()
|
||||
net = self._net(device)
|
||||
|
||||
if reset_history:
|
||||
self.taa.reset()
|
||||
|
||||
B, H, W, C = images.shape
|
||||
has_alpha = (C == 4) # RGBA kontrolü
|
||||
out_list = []
|
||||
|
||||
out_frames: list[torch.Tensor] = []
|
||||
|
||||
with torch.no_grad():
|
||||
for i in range(B):
|
||||
img = images[i:i+1].to(device).permute(0, 3, 1, 2).float()
|
||||
|
||||
rgb = img[:, :3]
|
||||
if has_alpha:
|
||||
alpha_channel = img[:, 3:] # Saydamlığı güvenli bir yere al
|
||||
# images are (B, H, W, C) — convert to (1, C, H, W)
|
||||
rgb = images[i : i + 1].to(device).permute(0, 3, 1, 2).float()
|
||||
|
||||
# Stage 1: Spatial preprocessing
|
||||
rgb = self.jitter(rgb, self.taa.frame_id, jitter_scale, net)
|
||||
rgb = self.edge_aa(rgb, edge_threshold, blur_radius, net)
|
||||
# process the first 3 channels
|
||||
rgb_3 = rgb[:, :3]
|
||||
|
||||
# Stage 2: Temporal blending
|
||||
taa_out = self.taa.update(rgb, taa_alpha, motion_sensitivity)
|
||||
rgb = torch.lerp(rgb, taa_out, taa_strength)
|
||||
# 1. TAA
|
||||
taa_out = self.taa.update(rgb_3, taa_alpha, motion_sensitivity)
|
||||
rgb_3 = taa_out
|
||||
|
||||
# Stage 3: CNN sharpening
|
||||
if dlaa_strength > 0:
|
||||
rgb = torch.lerp(rgb, net(rgb), dlaa_strength)
|
||||
|
||||
# Alpha kanalını geri ekle (Maskeli videoların bozulmasını engeller)
|
||||
if has_alpha:
|
||||
final_tensor = torch.cat([rgb, alpha_channel], dim=1)
|
||||
# 2. DLAA sharpening
|
||||
if dlaa_strength > 0.0:
|
||||
rgb_3 = _dlaa_sharpen(rgb_3, dlaa_strength)
|
||||
|
||||
# Alpha channel RGBA input
|
||||
if C == 4:
|
||||
alpha_ch = rgb[:, 3:4]
|
||||
result = torch.cat([rgb_3, alpha_ch], dim=1)
|
||||
else:
|
||||
final_tensor = rgb
|
||||
result = rgb_3
|
||||
|
||||
# Back to (1, H, W, C)
|
||||
out_frames.append(result.permute(0, 2, 3, 1).cpu())
|
||||
|
||||
return (torch.cat(out_frames, dim=0),)
|
||||
|
||||
out_list.append(final_tensor.permute(0, 2, 3, 1).cpu())
|
||||
|
||||
return (torch.cat(out_list, dim=0),)
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"VideoTAADLAA": VideoTAADLAA
|
||||
"VideoTAADLAA": VideoTAADLAA,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"VideoTAADLAA": "🎮 Video TAA + DLAA Anti-Aliasing"
|
||||
"VideoTAADLAA": "🎮 Video TAA + DLAA Anti-Aliasing",
|
||||
}
|
||||
Reference in New Issue
Block a user