diff --git a/VideoTAADLAA.py b/VideoTAADLAA.py index ef92c24..7a07fa0 100644 --- a/VideoTAADLAA.py +++ b/VideoTAADLAA.py @@ -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),)