From a0796fb8d63ca35f5dd39f3df69504ed6686b7fc Mon Sep 17 00:00:00 2001 From: MSXYZ-GenAI Date: Wed, 15 Apr 2026 17:45:00 +0300 Subject: [PATCH] Batch Processing + Stability Improvements --- VideoTAADLAA.py | 293 ++++++++++++++++++++++++++++++++---------------- 1 file changed, 197 insertions(+), 96 deletions(-) diff --git a/VideoTAADLAA.py b/VideoTAADLAA.py index 56ae24d..eae1093 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 - ComfyUI Native Memory & Device Management +# v0.1.1 - Batch Processing + Stability Improvements import torch import torch.nn as nn @@ -11,87 +11,127 @@ import comfy.model_management as mm logger = logging.getLogger("VideoTAADLAA") if not logger.handlers: - _handler = logging.StreamHandler() - _handler.setFormatter(logging.Formatter("[%(name)s] %(message)s")) - logger.addHandler(_handler) + h = logging.StreamHandler() + h.setFormatter(logging.Formatter("[%(name)s] %(message)s")) + logger.addHandler(h) logger.setLevel(logging.INFO) -# --- (Sınıflarınız aynı kalıyor, sadece _DLAANet forward kısmında cihaz kontrolü gerekebilir) --- +# ========================= +# DLAA CORE (light CNN) +# ========================= class _DLAANet(nn.Module): def __init__(self): super().__init__() self.conv = nn.Sequential( nn.Conv2d(3, 16, 3, padding=1, bias=False), - nn.ReLU(), + nn.ReLU(inplace=True), nn.Conv2d(16, 16, 3, padding=1, bias=False), - nn.ReLU(), + nn.ReLU(inplace=True), nn.Conv2d(16, 3, 3, padding=1, bias=False), ) - 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): - sharpen = torch.tensor([[0, -1, 0], [-1, 5, -1], [0, -1, 0]], dtype=torch.float32) + 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() + + def _init(self): convs = [m for m in self.conv if isinstance(m, nn.Conv2d)] - nn.init.kaiming_normal_(convs[0].weight, mode='fan_out', nonlinearity='relu') - nn.init.dirac_(convs[1].weight) - nn.init.zeros_(convs[2].weight) + nn.init.kaiming_normal_(convs[0].weight) + + # FIX v0.2.0: dirac_ yerine identity-benzeri başlatma. + # dirac_ semantik olarak doğru ama in==out==16 durumunda + # bazı PyTorch versiyonlarında uyarı üretebilir. + # eye tabanlı init daha güvenli ve aynı etkiyi verir. with torch.no_grad(): - for i in range(3): - for j in range(3): - convs[2].weight[i, j] = sharpen * 0.1 + convs[1].weight.zero_() + for c in range(min(convs[1].weight.shape[0], convs[1].weight.shape[1])): + convs[1].weight[c, c, 1, 1] = 1.0 # merkez piksel = identity + + nn.init.zeros_(convs[2].weight) + + # Sharpen filtresini 16 kanala dağıt + sharpen = torch.tensor([ + [ 0.0, -1.0, 0.0], + [-1.0, 5.0, -1.0], + [ 0.0, -1.0, 0.0] + ], dtype=torch.float32) + + with torch.no_grad(): + for out_c in range(3): # RGB çıkış + for in_c in range(16): # 16 giriş kanalı + convs[2].weight[out_c, in_c] = (sharpen * 0.1) / 16.0 def forward(self, x): - return torch.clamp(x + self.conv(x) * 0.15, 0.0, 1.0) + return torch.clamp(x + self.conv(x) * 0.12, 0.0, 1.0) + +# ========================= +# TEMPORAL STATE +# ========================= class _TAAState: def __init__(self): - self.history: torch.Tensor | None = None - self.frame_count: int = 0 + self.history = None + self.frame_id = 0 def reset(self): self.history = None - self.frame_count = 0 + self.frame_id = 0 - def update(self, frame: torch.Tensor, alpha: float, motion_threshold: float) -> torch.Tensor: + def update(self, frame, alpha, motion_sensitivity): if self.history is None or self.history.shape != frame.shape: self.history = frame.clone() - self.frame_count += 1 + self.frame_id += 1 return frame - - diff = torch.abs(frame - self.history).mean(dim=1, keepdim=True) - motion_mask = torch.clamp(diff * motion_threshold * 10.0, 0.0, 1.0) - dynamic_alpha = alpha * (1.0 - motion_mask) - + + diff = torch.mean(torch.abs(frame - self.history), dim=1, keepdim=True) + + # soft sigmoid motion mask + motion = torch.sigmoid((diff - motion_sensitivity) * 8.0) + + dynamic_alpha = alpha * (1.0 - motion) + out = dynamic_alpha * self.history + (1.0 - dynamic_alpha) * frame + self.history = out.detach() - self.frame_count += 1 + self.frame_id += 1 return out + +# ========================= +# MAIN NODE +# ========================= class VideoTAADLAA: def __init__(self): - self._net_cache: dict[str, _DLAANet] = {} - self._taa = _TAAState() + self.net_cache = {} + self.taa = _TAAState() @classmethod def INPUT_TYPES(cls): return { "required": { "images": ("IMAGE",), - "taa_strength": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 1.0, "step": 0.05}), - "taa_alpha": ("FLOAT", {"default": 0.8, "min": 0.0, "max": 0.99, "step": 0.01}), - "motion_sensitivity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05}), - "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": 0, "max": 5, "step": 1}), - "reset_history": ("BOOLEAN", {"default": False}), - "batch_size": ("INT", {"default": 0, "min": 0, "max": 64, "step": 1}), + "taa_strength": ("FLOAT", {"default": 0.7, "min": 0, "max": 1, "step": 0.05}), + "taa_alpha": ("FLOAT", {"default": 0.8, "min": 0, "max": 0.99, "step": 0.01}), + "motion_sensitivity": ("FLOAT", {"default": 0.4, "min": 0, "max": 1, "step": 0.05}), + "jitter_scale": ("FLOAT", {"default": 0.1, "min": 0, "max": 2, "step": 0.05}), + "dlaa_strength": ("FLOAT", {"default": 0.6, "min": 0, "max": 1, "step": 0.05}), + "edge_threshold": ("FLOAT", {"default": 0.15, "min": 0, "max": 1, "step": 0.01}), + "blur_radius": ("INT", {"default": 1, "min": 0, "max": 5, "step": 1}), + "reset_history": ("BOOLEAN", {"default": True}), } } @@ -99,78 +139,139 @@ class VideoTAADLAA: FUNCTION = "execute" CATEGORY = "CustomPostProcess" - def _get_net(self, device: torch.device) -> _DLAANet: - key = str(device) - if key not in self._net_cache: - self._net_cache[key] = _DLAANet().to(device).eval() - return self._net_cache[key] + def _device_key(self, device: torch.device) -> str: + # FIX v0.2.0: "cuda" ve "cuda:0" aynı cache entry'i kullanır, + # gereksiz model kopyası ve bellek sızıntısı önlenir. + return f"{device.type}:{device.index if device.index is not None else 0}" + + def _net(self, device: torch.device) -> _DLAANet: + k = self._device_key(device) + if k not in self.net_cache: + self.net_cache[k] = _DLAANet().to(device).eval() + return self.net_cache[k] + + # ------------------------- + # jitter (stable bilinear) + # ------------------------- + def jitter(self, x: torch.Tensor, idx: int, scale: float, net: _DLAANet) -> torch.Tensor: + if scale < 1e-5: + return x - def _jitter(self, x: torch.Tensor, idx: int, scale: float, net: _DLAANet) -> torch.Tensor: - 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) - 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] + grid = F.affine_grid(theta, x.shape, align_corners=False) + return F.grid_sample(x, grid, mode="bilinear", padding_mode="reflection", align_corners=False) + + # ------------------------- + # edge AA (soft mask) + # ------------------------- + 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() + + edge = torch.sqrt(sx*sx + sy*sy + 1e-6) + + # soft mask + mask = torch.sigmoid((edge - thr) * 12.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 + blurred = F.avg_pool2d( + F.pad(x, [blur]*4, mode="reflect"), + k, stride=1 + ) - def execute(self, images, taa_strength, taa_alpha, motion_sensitivity, jitter_scale, - dlaa_strength, edge_threshold, blur_radius, reset_history, batch_size): - - # ComfyUI cihaz yönetimi + return x*(1-mask) + blurred*mask + + # ------------------------- + # DLAA batch forward + # ------------------------- + def _dlaa_batch(self, frames: torch.Tensor, strength: float, net: _DLAANet) -> torch.Tensor: + """ + FIX v0.2.0: DLAA artık tüm frame'leri tek seferde batch olarak işler. + Büyük videolarda belirgin hız artışı sağlar (özellikle GPU'da). + TAA history sıralı frame bağımlılığı gerektirdiği için döngü zorunlu, + ama DLAA bağımsız olduğundan batch'e alındı. + """ + if strength <= 0: + return frames + enhanced = net(frames) + return torch.lerp(frames, enhanced, strength) + + # ------------------------- + # EXECUTE + # ------------------------- + def execute( + self, + images: torch.Tensor, + taa_strength: float, + taa_alpha: float, + motion_sensitivity: float, + jitter_scale: float, + dlaa_strength: float, + edge_threshold: float, + blur_radius: int, + reset_history: bool = False, + ): device = mm.get_torch_device() - + if reset_history: - self._taa.reset() + self.taa.reset() + + net = self._net(device) - net = self._get_net(device) B, H, W, C = images.shape - - # ComfyUI bellek yönetimi (VRAM kullanımı) - free_mem = mm.get_free_memory() - working_batch_size = batch_size if batch_size > 0 else max(1, min(16, int(free_mem // (1024 * 1024 * 64)))) - images = images.to(device) + has_alpha = (C == 4) - outputs = [] + # --- TAA (sıralı, history bağımlı) --- + taa_frames = [] with torch.no_grad(): - 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] - img_alpha = img[:, 3:4] if has_alpha else None - results_rgb = [] + for i in range(B): + f = images[i:i+1].permute(0, 3, 1, 2).float() + rgb = f[:, :3] - for f_idx in range(img_rgb.shape[0]): - f = img_rgb[f_idx:f_idx + 1] - f = self._jitter(f, self._taa.frame_count, jitter_scale, net) - f = self._edge_aa(f, edge_threshold, blur_radius, net) - taa_out = self._taa.update(f, taa_alpha, motion_sensitivity) - f = torch.lerp(f, taa_out, taa_strength) - if dlaa_strength > 0.0: - f = torch.lerp(f, net(f), dlaa_strength) - results_rgb.append(f.clamp(0.0, 1.0)) + rgb = self.jitter(rgb, self.taa.frame_id, jitter_scale, net) + rgb = self.edge_aa(rgb, edge_threshold, blur_radius, net) - batch_rgb = torch.cat(results_rgb, dim=0) - 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).to("cpu")) - - # Belleği rahatlat + taa_out = self.taa.update(rgb, taa_alpha, motion_sensitivity) + rgb = torch.lerp(rgb, taa_out, taa_strength) + + taa_frames.append(rgb) + + # --- DLAA (batch, paralel) --- + with torch.no_grad(): + rgb_batch = torch.cat(taa_frames, dim=0) # (B, 3, H, W) + rgb_batch = self._dlaa_batch(rgb_batch, dlaa_strength, net) + rgb_batch = rgb_batch.clamp(0, 1) + + # --- Alpha kanalını geri ekle ve CPU'ya taşı --- + if has_alpha: + alpha_batch = images.permute(0, 3, 1, 2).float()[:, 3:4] # (B, 1, H, W) + rgb_batch = torch.cat([rgb_batch, alpha_batch], dim=1) # (B, 4, H, W) + + out = rgb_batch.permute(0, 2, 3, 1).cpu() # (B, H, W, C) + + # VRAM Temizliği mm.soft_empty_cache() - - return (torch.cat(outputs, dim=0),) -NODE_CLASS_MAPPINGS = {"VideoTAADLAA": VideoTAADLAA} -NODE_DISPLAY_NAME_MAPPINGS = {"VideoTAADLAA": "🎮 Video TAA + DLAA Anti-Aliasing"} \ No newline at end of file + return (out,) + + +NODE_CLASS_MAPPINGS = { + "VideoTAADLAA": VideoTAADLAA +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "VideoTAADLAA": "🎮 Video TAA + DLAA Anti-Aliasing" +} \ No newline at end of file