nodes/SOLRICKS-comfyui-solricks
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Advanced DLAA Network and Buffer Management

Browse Source
This commit is contained in:
MSXYZ-GenAI
2026-04-15 10:15:40 +03:00
parent e1f8c9b74b
commit f69261fafc
1 changed files with 22 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
VideoTAADLAA.py
+22 -48
View File
@@ -1,7 +1,7 @@
# Created By MSXYZ and Claude Opus 4.6
# GPU başında ciddi işler döndürüyoruz. 🚀
# TAA (Temporal Anti-Aliasing) + DLAA (Deep Learning Anti-Aliasing) Adaptation
# v0.1.1
# v0.1.2 - Fixed: Batch Size validation and Range Error
import torch
import torch.nn as nn
@@ -16,13 +16,9 @@ if not logger.handlers:
logger.addHandler(_handler)
logger.setLevel(logging.INFO)
# -----------------------------
# Gelişmiş DLAA Ağı ve Buffer Yönetimi
# -----------------------------
class _DLAANet(nn.Module):
def __init__(self):
super().__init__()
# Katmanlar
self.conv = nn.Sequential(
nn.Conv2d(3, 16, 3, padding=1, bias=False),
nn.ReLU(),
@@ -30,19 +26,14 @@ class _DLAANet(nn.Module):
nn.ReLU(),
nn.Conv2d(16, 3, 3, padding=1, bias=False),
)
# KERNEL OPTİMİZASYONU: Sabitleri GPU buffer'ına kaydediyoruz
# Böylece her karede CPU -> GPU transferi yapılmaz.
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._init_weights()
def _init_weights(self):
# Deterministik Sharpen Init
sharpen = torch.tensor([[0, -1, 0], [-1, 5, -1], [0, -1, 0]], dtype=torch.float32)
convs = [m for m in self.conv if isinstance(m, nn.Conv2d)]
nn.init.dirac_(convs[0].weight)
@@ -56,10 +47,6 @@ class _DLAANet(nn.Module):
def forward(self, x):
return torch.clamp(x + self.conv(x) * 0.15, 0.0, 1.0)
# -----------------------------
# Hareket Duyarlı TAA Yönetimi
# -----------------------------
class _TAAState:
def __init__(self):
self.history: torch.Tensor | None = None
@@ -75,12 +62,7 @@ class _TAAState:
self.frame_count += 1
return frame
# GHOSTING KONTROLÜ (Motion Masking):
# Kareler arası farkı hesapla (L1 norm)
diff = torch.abs(frame - self.history).mean(dim=1, keepdim=True)
# Hareketin olduğu yerlerde alpha'yı (geçmişin etkisini) düşür
# Eğer fark threshold'un üzerindeyse o bölgeyi yeni kareden al (ghosting'i kes)
motion_mask = torch.clamp(diff * motion_threshold * 10.0, 0.0, 1.0)
dynamic_alpha = alpha * (1.0 - motion_mask)
@@ -89,10 +71,6 @@ class _TAAState:
self.frame_count += 1
return out
# -----------------------------
# Ana ComfyUI Node
# -----------------------------
class VideoTAADLAA:
def __init__(self):
self._lock = threading.Lock()
@@ -110,7 +88,7 @@ class VideoTAADLAA:
"jitter_scale": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 2.0, "step": 0.1}),
"dlaa_strength": ("FLOAT", {"default": 0.6, "min": 0.0, "max": 1.0, "step": 0.05}),
"edge_threshold": ("FLOAT", {"default": 0.12, "min": 0.0, "max": 1.0, "step": 0.01}),
"blur_radius": ("INT", {"default": 1, "min": 1, "max": 5, "step": 1}),
"blur_radius": ("INT", {"default": 1, "min": 0, "max": 5, "step": 1}), # Min 0 yapıldı
"reset_history": ("BOOLEAN", {"default": False}),
"batch_size": ("INT", {"default": 0, "min": 0, "max": 64, "step": 1}),
}
@@ -127,27 +105,24 @@ class VideoTAADLAA:
return self._net_cache[key]
def _jitter(self, x: torch.Tensor, idx: int, scale: float, net: _DLAANet) -> torch.Tensor:
if scale <= 0:
return x
if scale <= 1e-4: 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] = off[0] * scale / W
theta[:, 1, 2] = off[1] * scale / H
grid = F.affine_grid(theta, x.shape, align_corners=False)
# BICUBIC JITTER: Bilinear yerine Bicubic kullanarak keskinliği koruyoruz
return F.grid_sample(x, grid, mode="bicubic", padding_mode="reflection", align_corners=False)
def _edge_aa(self, x: torch.Tensor, thr: float, blur: int, net: _DLAANet) -> torch.Tensor:
if blur <= 0: return x
gray = 0.299 * x[:, 0:1] + 0.587 * x[:, 1:2] + 0.114 * x[:, 2:3]
sx = F.conv2d(gray, net.sobel_x, padding=1)
sy = F.conv2d(gray, net.sobel_y, padding=1)
mask = (torch.sqrt(sx * sx + sy * sy + 1e-6) > thr).float()
if blur > 0:
k = blur * 2 + 1
blurred = F.avg_pool2d(F.pad(x, [blur] * 4, mode="reflect"), k, stride=1)
return x * (1.0 - mask) + blurred * mask
return x
k = blur * 2 + 1
blurred = F.avg_pool2d(F.pad(x, [blur] * 4, mode="reflect"), k, stride=1)
return x * (1.0 - mask) + blurred * mask
def execute(self, images, taa_strength, taa_alpha, motion_sensitivity, jitter_scale,
dlaa_strength, edge_threshold, blur_radius, reset_history, batch_size):
@@ -158,22 +133,26 @@ class VideoTAADLAA:
device = images.device
net = self._get_net(device)
if batch_size == 0:
# Maliyet hesabında C=4 (RGBA) ve Bicubic payı eklendi
B, H, W, C = images.shape
cost = H * W * C * 4 * 14 # Bicubic biraz daha fazla VRAM/İşlem ister
B, H, W, C = images.shape
# --- Kritik Hata Düzeltme: Batch Size Hesabı ---
working_batch_size = batch_size
if working_batch_size <= 0:
cost = H * W * C * 4 * 14
free, _ = torch.cuda.mem_get_info() if torch.cuda.is_available() else (0,0)
batch_size = max(1, min(16, int((free * 0.6) // cost))) if free > 0 else 1
working_batch_size = max(1, min(16, int((free * 0.6) // cost))) if free > 0 else 1
# Range hatasını önlemek için batch_size'ın en az 1 olduğundan emin oluyoruz
working_batch_size = max(1, working_batch_size)
# ----------------------------------------------
images = images.to(device)
B, H, W, C = images.shape
has_alpha = (C == 4)
outputs = []
with torch.no_grad():
for i in range(0, B, batch_size):
chunk = images[i:i + batch_size]
for i in range(0, B, working_batch_size):
chunk = images[i:i + working_batch_size]
img = chunk.permute(0, 3, 1, 2).float()
img_rgb = img[:, :3]
@@ -183,17 +162,12 @@ class VideoTAADLAA:
for f_idx in range(img_rgb.shape[0]):
f = img_rgb[f_idx:f_idx + 1]
# 1. Jitter (Bicubic Mode)
f = self._jitter(f, self._taa.frame_count, jitter_scale, net)
# 2. Edge AA (Optimized Kernels)
f = self._edge_aa(f, edge_threshold, blur_radius, net)
# 3. Motion-Adaptive TAA (Anti-Ghosting)
taa_out = self._taa.update(f, taa_alpha, motion_sensitivity)
f = torch.lerp(f, taa_out, taa_strength)
# 4. DLAA
if dlaa_strength > 0.0:
f = torch.lerp(f, net(f), dlaa_strength)
@@ -203,7 +177,7 @@ class VideoTAADLAA:
batch_out = torch.cat([batch_rgb, img_alpha], dim=1) if has_alpha else batch_rgb
outputs.append(batch_out.permute(0, 2, 3, 1).cpu())
logger.info(f"İşlendi: {min(i + batch_size, B)} / {B}")
logger.info(f"İşlendi: {min(i + working_batch_size, B)} / {B}")
return (torch.cat(outputs, dim=0),)