1210 lines
37 KiB
Python
1210 lines
37 KiB
Python
# Video TAA + DLAA
|
|
# v0.2.0
|
|
# MSXYZ
|
|
# AI-assisted
|
|
|
|
|
|
import os
|
|
import logging
|
|
import threading
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from safetensors.torch import load_file
|
|
|
|
|
|
try:
|
|
from .model import DLAANet
|
|
except ImportError:
|
|
from model import DLAANet
|
|
|
|
try:
|
|
from .utils import rgb_luma
|
|
except ImportError:
|
|
from utils import rgb_luma
|
|
|
|
try:
|
|
from .taa import TAAState
|
|
except ImportError:
|
|
from taa import TAAState
|
|
|
|
try:
|
|
from .config import NODE_DEFAULTS
|
|
except ImportError:
|
|
from config import NODE_DEFAULTS
|
|
|
|
try:
|
|
from .presets import (
|
|
PRESETS,
|
|
AUTO_STATIC,
|
|
AUTO_BALANCED,
|
|
AUTO_MOTION,
|
|
SINGLE_IMAGE_DETAIL,
|
|
MOTION_SUPPRESSION,
|
|
PRESET_MODEL_WEIGHT,
|
|
TEXTURE_PRESETS,
|
|
)
|
|
except ImportError:
|
|
from presets import (
|
|
PRESETS,
|
|
AUTO_STATIC,
|
|
AUTO_BALANCED,
|
|
AUTO_MOTION,
|
|
SINGLE_IMAGE_DETAIL,
|
|
MOTION_SUPPRESSION,
|
|
PRESET_MODEL_WEIGHT,
|
|
TEXTURE_PRESETS,
|
|
)
|
|
|
|
try:
|
|
import comfy.model_management as mm
|
|
except ImportError:
|
|
mm = None
|
|
|
|
try:
|
|
from comfy.utils import ProgressBar
|
|
except ImportError:
|
|
ProgressBar = None
|
|
|
|
log = logging.getLogger("VideoTAADLAA")
|
|
|
|
|
|
class VideoTAADLAA:
|
|
def __init__(self):
|
|
self.net_cache = {}
|
|
self._net_lock = threading.Lock()
|
|
|
|
self.texture_net_cache = {}
|
|
self._texture_lock = threading.Lock()
|
|
self._texture_missing_warned = False
|
|
|
|
for name, value in NODE_DEFAULTS.items():
|
|
setattr(self, name, value)
|
|
|
|
self.preset_model_weight = PRESET_MODEL_WEIGHT
|
|
self.texture_presets = TEXTURE_PRESETS
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"preset": (["Auto", "Balanced", "Detail", "Smooth"],),
|
|
},
|
|
"optional": {
|
|
"detail_intensity": ("FLOAT", {
|
|
"default": 1.00,
|
|
"min": 0.00,
|
|
"max": 2.00,
|
|
"step": 0.05,
|
|
}),
|
|
"texture_intensity": ("FLOAT", {
|
|
"default": 1.00,
|
|
"min": 0.00,
|
|
"max": 2.00,
|
|
"step": 0.05,
|
|
}),
|
|
"motion_stability": ("FLOAT", {
|
|
"default": 1.00,
|
|
"min": 0.50,
|
|
"max": 2.00,
|
|
"step": 0.05,
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "execute"
|
|
CATEGORY = "CustomPostProcess"
|
|
|
|
def _net(self, device):
|
|
key = str(device)
|
|
|
|
if key in self.net_cache:
|
|
return self.net_cache[key]
|
|
|
|
# main model
|
|
with self._net_lock:
|
|
if key in self.net_cache:
|
|
return self.net_cache[key]
|
|
|
|
net = DLAANet().to(device)
|
|
|
|
base_path = os.path.dirname(os.path.realpath(__file__))
|
|
|
|
safetensors_path = os.path.join(base_path, "DLAANet.safetensors")
|
|
pth_path = os.path.join(base_path, "DLAANet.pth")
|
|
|
|
if os.path.exists(safetensors_path):
|
|
state_dict = load_file(safetensors_path, device=str(device))
|
|
log.info("[DLAA] Loaded DLAANet.safetensors")
|
|
elif os.path.exists(pth_path):
|
|
state_dict = torch.load(pth_path, map_location=device)
|
|
log.info("[DLAA] Loaded DLAANet.pth")
|
|
else:
|
|
raise FileNotFoundError(
|
|
f"[DLAA] Model not found. Expected: {safetensors_path} or {pth_path}"
|
|
)
|
|
|
|
if "jitter_offsets" in state_dict:
|
|
del state_dict["jitter_offsets"]
|
|
|
|
net.load_state_dict(state_dict, strict=False)
|
|
net = net.float()
|
|
net.eval()
|
|
|
|
n_params = sum(p.numel() for p in net.parameters())
|
|
log.info(f"[DLAA] Main model parameters: {n_params / 1e6:.2f}M")
|
|
|
|
self.net_cache[key] = net
|
|
|
|
return self.net_cache[key]
|
|
|
|
|
|
def _texture_net(self, device):
|
|
if not self.texture_pass_enabled:
|
|
return None
|
|
|
|
key = str(device)
|
|
|
|
if key in self.texture_net_cache:
|
|
return self.texture_net_cache[key]
|
|
|
|
# DLAANet Texture Model Loading
|
|
with self._texture_lock:
|
|
if key in self.texture_net_cache:
|
|
return self.texture_net_cache[key]
|
|
|
|
base_path = os.path.dirname(os.path.realpath(__file__))
|
|
|
|
safetensors_path = os.path.join(base_path, "DLAATexture.safetensors")
|
|
pth_path = os.path.join(base_path, "DLAATexture.pth")
|
|
|
|
if os.path.exists(safetensors_path):
|
|
model_path = safetensors_path
|
|
elif os.path.exists(pth_path):
|
|
model_path = pth_path
|
|
else:
|
|
if not self._texture_missing_warned:
|
|
log.info("[DLAA] Texture model not found, skipping texture pass.")
|
|
self._texture_missing_warned = True
|
|
return None
|
|
|
|
try:
|
|
net = DLAANet().to(device)
|
|
|
|
if model_path.endswith(".safetensors"):
|
|
state_dict = load_file(model_path, device=str(device))
|
|
else:
|
|
state_dict = torch.load(model_path, map_location=device)
|
|
|
|
if isinstance(state_dict, dict):
|
|
if "state_dict" in state_dict:
|
|
state_dict = state_dict["state_dict"]
|
|
elif "params_ema" in state_dict:
|
|
state_dict = state_dict["params_ema"]
|
|
elif "params" in state_dict:
|
|
state_dict = state_dict["params"]
|
|
elif "model" in state_dict:
|
|
state_dict = state_dict["model"]
|
|
|
|
if "jitter_offsets" in state_dict:
|
|
del state_dict["jitter_offsets"]
|
|
|
|
net.load_state_dict(state_dict, strict=False)
|
|
net = net.float()
|
|
net.eval()
|
|
|
|
n_params = sum(p.numel() for p in net.parameters())
|
|
|
|
self.texture_net_cache[key] = net
|
|
|
|
log.info("[DLAA] Loaded %s", os.path.basename(model_path))
|
|
log.info(f"[DLAA] Texture model parameters: {n_params / 1e6:.2f}M")
|
|
|
|
except Exception as e:
|
|
if not self._texture_missing_warned:
|
|
import traceback
|
|
log.error(f"[DLAA] Texture model exception:\n{traceback.format_exc()}")
|
|
log.warning(
|
|
f"[DLAA] Texture model could not be loaded, skipping texture pass: {type(e).__name__}: {e!r}"
|
|
)
|
|
self._texture_missing_warned = True
|
|
return None
|
|
|
|
return self.texture_net_cache[key]
|
|
|
|
def _tile_weight_map(self, th, tw, overlap, device, dtype):
|
|
w_y = torch.ones(th, device=device, dtype=dtype)
|
|
w_x = torch.ones(tw, device=device, dtype=dtype)
|
|
|
|
ramp = min(overlap, th // 4, tw // 4)
|
|
|
|
if ramp > 0:
|
|
values = torch.linspace(
|
|
1.0 / (ramp + 1),
|
|
ramp / (ramp + 1),
|
|
ramp,
|
|
device=device,
|
|
dtype=dtype
|
|
)
|
|
|
|
w_y[:ramp] = values
|
|
w_y[-ramp:] = torch.flip(values, dims=[0])
|
|
|
|
w_x[:ramp] = values
|
|
w_x[-ramp:] = torch.flip(values, dims=[0])
|
|
|
|
return torch.minimum(
|
|
w_y.view(1, 1, th, 1),
|
|
w_x.view(1, 1, 1, tw)
|
|
)
|
|
|
|
def _tiled_forward(self, net, x: torch.Tensor, tile_size: int = 512, overlap: int = 32) -> torch.Tensor:
|
|
|
|
B, C, H, W = x.shape
|
|
|
|
if H <= tile_size and W <= tile_size:
|
|
return torch.clamp(net(x), 0.0, 1.0)
|
|
|
|
|
|
step = tile_size - overlap * 2
|
|
|
|
# avoid stuck tiling
|
|
if step <= 0:
|
|
raise ValueError(
|
|
f"Invalid tiling settings: tile_size={tile_size}, overlap={overlap}"
|
|
)
|
|
|
|
out = torch.zeros_like(x)
|
|
weight = torch.zeros(B, 1, H, W, device=x.device, dtype=x.dtype)
|
|
|
|
tile_h = min(tile_size, H)
|
|
tile_w = min(tile_size, W)
|
|
|
|
w_map = self._tile_weight_map(
|
|
tile_h,
|
|
tile_w,
|
|
overlap,
|
|
x.device,
|
|
x.dtype
|
|
)
|
|
|
|
y0 = 0
|
|
while y0 < H:
|
|
y1 = min(y0 + tile_size, H)
|
|
y0_c = max(0, y1 - tile_size)
|
|
|
|
x0 = 0
|
|
while x0 < W:
|
|
x1 = min(x0 + tile_size, W)
|
|
x0_c = max(0, x1 - tile_size)
|
|
|
|
tile = x[:, :, y0_c:y1, x0_c:x1]
|
|
dlaa_out_tile = net(tile)
|
|
|
|
out[:, :, y0_c:y1, x0_c:x1] += dlaa_out_tile * w_map
|
|
weight[:, :, y0_c:y1, x0_c:x1] += w_map
|
|
|
|
if x1 == W:
|
|
break
|
|
x0 += step
|
|
|
|
if y1 == H:
|
|
break
|
|
y0 += step
|
|
|
|
return torch.clamp(out / weight.clamp(min=1e-6), 0.0, 1.0)
|
|
|
|
def _jitter_count(self, net):
|
|
offsets = getattr(net, "jitter_offsets", None)
|
|
|
|
if offsets is None or offsets.shape[0] == 0:
|
|
return 0
|
|
|
|
return offsets.shape[0]
|
|
|
|
def _jitter(self, x, idx, scale, net):
|
|
if scale < 1e-5:
|
|
return x
|
|
|
|
offsets = getattr(net, "jitter_offsets", None)
|
|
|
|
if offsets is None or offsets.shape[0] == 0:
|
|
return x
|
|
|
|
off = offsets[idx % offsets.shape[0]].to(device=x.device, dtype=x.dtype)
|
|
|
|
B, C, H, W = x.shape
|
|
|
|
theta = torch.eye(
|
|
2,
|
|
3,
|
|
device=x.device,
|
|
dtype=x.dtype,
|
|
).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="bilinear",
|
|
padding_mode="reflection",
|
|
align_corners=False,
|
|
)
|
|
|
|
def _edge_aa(self, x, thr, blur_radius, net):
|
|
if blur_radius <= 0:
|
|
return x
|
|
|
|
gray = rgb_luma(x)
|
|
sx = F.conv2d(gray, net.sobel_x, padding=1)
|
|
sy = F.conv2d(gray, net.sobel_y, padding=1)
|
|
edge = torch.sqrt(sx*sx + sy*sy + 1e-6)
|
|
|
|
mask = torch.sigmoid((edge - thr) * self.edge_aa_slope)
|
|
|
|
blurred = F.avg_pool2d(
|
|
F.pad(x, [blur_radius]*4, mode="reflect"),
|
|
blur_radius*2+1,
|
|
stride=1
|
|
)
|
|
|
|
return x*(1.0 - mask) + blurred*mask
|
|
|
|
def _temporal_refine(self, current, previous, strength=0.35, motion_threshold=0.08):
|
|
if previous is None:
|
|
return current
|
|
|
|
if previous.shape != current.shape:
|
|
return current
|
|
|
|
curr_luma = rgb_luma(current)
|
|
prev_luma = rgb_luma(previous)
|
|
|
|
motion = torch.abs(curr_luma - prev_luma)
|
|
motion = torch.clamp(motion / motion_threshold, 0.0, 1.0)
|
|
|
|
blend_mask = (1.0 - motion) * strength
|
|
|
|
refined = current * (1.0 - blend_mask) + previous * blend_mask
|
|
return refined.clamp(0.0, 1.0)
|
|
|
|
def _stabilize_fine_detail(
|
|
self,
|
|
current,
|
|
previous,
|
|
strength,
|
|
motion_threshold,
|
|
):
|
|
if previous is None:
|
|
return current
|
|
|
|
if previous.shape != current.shape:
|
|
return current
|
|
|
|
if strength <= 1e-5:
|
|
return current
|
|
|
|
current_base = F.avg_pool2d(current, 3, stride=1, padding=1)
|
|
previous_base = F.avg_pool2d(previous, 3, stride=1, padding=1)
|
|
|
|
current_detail = current - current_base
|
|
previous_detail = previous - previous_base
|
|
|
|
current_luma = rgb_luma(current)
|
|
previous_luma = rgb_luma(previous)
|
|
|
|
motion = torch.abs(current_luma - previous_luma)
|
|
stable_area = 1.0 - torch.clamp(
|
|
motion / motion_threshold,
|
|
0.0,
|
|
1.0,
|
|
)
|
|
|
|
detail_energy = rgb_luma(current_detail.abs())
|
|
|
|
line_mask = torch.sigmoid(
|
|
(detail_energy - self.detail_shimmer_threshold) *
|
|
self.detail_shimmer_slope
|
|
)
|
|
|
|
blend_mask = (line_mask * stable_area * strength).clamp(
|
|
0.0,
|
|
self.detail_shimmer_max_blend,
|
|
)
|
|
|
|
stabilized_detail = torch.lerp(
|
|
current_detail,
|
|
previous_detail,
|
|
blend_mask,
|
|
)
|
|
|
|
return torch.clamp(current_base + stabilized_detail, 0.0, 1.0)
|
|
|
|
def _resolve_frame_config(self, preset, is_single_image, rgb, taa):
|
|
# Auto uses frame history
|
|
if preset == "Auto":
|
|
if taa.history is not None and taa.history.shape == rgb.shape:
|
|
scene_motion = torch.abs(rgb - taa.history).mean().item()
|
|
else:
|
|
scene_motion = self.auto_default_scene_motion
|
|
|
|
if scene_motion < self.auto_static_motion_threshold:
|
|
frame_cfg = AUTO_STATIC
|
|
elif scene_motion < self.auto_balanced_motion_threshold:
|
|
frame_cfg = AUTO_BALANCED
|
|
else:
|
|
frame_cfg = AUTO_MOTION
|
|
else:
|
|
frame_cfg = PRESETS.get(preset, PRESETS["Balanced"])
|
|
|
|
cfg = frame_cfg.copy()
|
|
|
|
if preset == "Detail" and is_single_image:
|
|
cfg.update(SINGLE_IMAGE_DETAIL)
|
|
|
|
return cfg
|
|
|
|
def _apply_texture_pass(
|
|
self,
|
|
texture_net,
|
|
dlaa_out,
|
|
tile_size,
|
|
motion_gate,
|
|
dark_mask,
|
|
highlight_mask,
|
|
preset,
|
|
texture_intensity=1.00,
|
|
motion_stability=1.00,
|
|
frame_index=None
|
|
):
|
|
if texture_net is None:
|
|
return dlaa_out
|
|
|
|
if texture_intensity <= 1e-5:
|
|
return dlaa_out
|
|
|
|
texture_cfg = self.texture_presets.get(preset, self.texture_presets["Balanced"])
|
|
if not texture_cfg.get("enabled", True):
|
|
return dlaa_out
|
|
|
|
gen_out = self._tiled_forward(
|
|
texture_net,
|
|
dlaa_out,
|
|
tile_size=tile_size,
|
|
overlap=self.texture_tile_overlap
|
|
)
|
|
|
|
if gen_out.shape != dlaa_out.shape:
|
|
log.warning(
|
|
"[DLAA] Texture model output size does not match input, skipping texture pass."
|
|
)
|
|
return dlaa_out
|
|
|
|
gray = rgb_luma(dlaa_out)
|
|
|
|
edge_x = torch.abs(gray[:, :, :, 1:] - gray[:, :, :, :-1])
|
|
edge_y = torch.abs(gray[:, :, 1:, :] - gray[:, :, :-1, :])
|
|
|
|
edge_x = F.pad(edge_x, (0, 1, 0, 0))
|
|
edge_y = F.pad(edge_y, (0, 0, 0, 1))
|
|
|
|
thin_edge_mask = (edge_x + edge_y).clamp(0.0, 1.0)
|
|
thin_edge_mask = torch.sigmoid(
|
|
(thin_edge_mask - texture_cfg["edge_threshold"]) * texture_cfg["edge_slope"]
|
|
)
|
|
thin_edge_mask = F.avg_pool2d(thin_edge_mask, kernel_size=3, stride=1, padding=1)
|
|
|
|
# isolate texture residual
|
|
blur_kernel = texture_cfg["blur_kernel"]
|
|
gen_blur = F.avg_pool2d(gen_out, kernel_size=blur_kernel, stride=1, padding=blur_kernel // 2)
|
|
hallucinated_texture = gen_out - gen_blur
|
|
|
|
debug_stats = log.isEnabledFor(logging.DEBUG)
|
|
raw_gen_delta = 0.0
|
|
if debug_stats:
|
|
raw_gen_delta = hallucinated_texture.abs().mean().item()
|
|
|
|
dark_base = texture_cfg["dark_base"]
|
|
texture_strength = texture_cfg["strength"] * texture_intensity
|
|
texture_limit = texture_cfg["limit"] * texture_intensity
|
|
motion_suppression = min(texture_cfg["motion_suppression"] * motion_stability, 0.98)
|
|
line_suppression = min(texture_cfg["line_suppression"] * motion_stability, 0.95)
|
|
|
|
texture_mask = dark_base + (1.0 - dark_base) * dark_mask
|
|
texture_mask = texture_mask * (
|
|
1.0 - highlight_mask * texture_cfg["highlight_suppression"]
|
|
)
|
|
texture_mask = texture_mask * (
|
|
1.0 - motion_gate * motion_suppression
|
|
)
|
|
texture_mask = texture_mask * (
|
|
1.0 - thin_edge_mask * line_suppression
|
|
)
|
|
|
|
hallucinated_texture = hallucinated_texture * texture_mask
|
|
hallucinated_texture = hallucinated_texture.clamp(
|
|
-texture_limit,
|
|
texture_limit
|
|
)
|
|
|
|
out = dlaa_out + hallucinated_texture * texture_strength
|
|
|
|
final_texture_delta = 0.0
|
|
|
|
if debug_stats:
|
|
final_texture_delta = (out - dlaa_out).abs().mean().item()
|
|
|
|
if (
|
|
debug_stats and
|
|
frame_index is not None and
|
|
self.texture_log_interval > 0 and
|
|
frame_index % self.texture_log_interval == 0
|
|
):
|
|
log.debug(
|
|
"[DLAA] texture frame=%d raw_delta=%.6f final_delta=%.6f",
|
|
frame_index,
|
|
raw_gen_delta,
|
|
final_texture_delta
|
|
)
|
|
|
|
return torch.clamp(out, 0.0, 1.0)
|
|
|
|
def _normalize_run_inputs(
|
|
self,
|
|
preset,
|
|
detail_intensity,
|
|
texture_intensity,
|
|
motion_stability,
|
|
):
|
|
if preset == "Sharp":
|
|
preset = "Detail"
|
|
elif preset == "Cinematic":
|
|
preset = "Smooth"
|
|
|
|
detail_intensity = max(0.0, min(float(detail_intensity), 2.0))
|
|
texture_intensity = max(0.0, min(float(texture_intensity), 2.0))
|
|
motion_stability = max(0.5, min(float(motion_stability), 2.0))
|
|
|
|
return preset, detail_intensity, texture_intensity, motion_stability
|
|
|
|
def _get_device(self):
|
|
if mm is not None:
|
|
return mm.get_torch_device()
|
|
|
|
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
def _frame_blur_radius(self, preset, is_single_image):
|
|
if is_single_image:
|
|
return 0
|
|
|
|
return 0 if preset == "Detail" else 1
|
|
|
|
def _vram_mb(self, device):
|
|
try:
|
|
torch_device = torch.device(device)
|
|
except Exception:
|
|
return 0
|
|
|
|
if torch_device.type != "cuda" or not torch.cuda.is_available():
|
|
return 0
|
|
|
|
try:
|
|
device_index = torch_device.index
|
|
|
|
if device_index is None:
|
|
device_index = torch.cuda.current_device()
|
|
|
|
return torch.cuda.get_device_properties(device_index).total_memory // (1024 * 1024)
|
|
except Exception:
|
|
return 0
|
|
|
|
def _tile_size_for_vram(self, vram_mb):
|
|
if vram_mb <= 8192:
|
|
return 512
|
|
|
|
return 1024
|
|
|
|
def _log_tiling(self, height, width, tile_size):
|
|
if height <= tile_size and width <= tile_size:
|
|
return
|
|
|
|
tile_step = tile_size - 64
|
|
tile_count = ((height + tile_step - 1) // tile_step) * ((width + tile_step - 1) // tile_step)
|
|
log.debug(f"[DLAA] Tiled inference: {tile_count} tiles")
|
|
|
|
def _load_frame(self, images, frame_index, device):
|
|
img = images[frame_index:frame_index + 1].to(device).permute(0, 3, 1, 2).float()
|
|
return img[:, :3]
|
|
|
|
def _texture_net_for_run(self, device, preset, texture_intensity):
|
|
texture_cfg = self.texture_presets.get(preset, self.texture_presets["Balanced"])
|
|
|
|
if not self.texture_pass_enabled:
|
|
return None
|
|
|
|
if not texture_cfg.get("enabled", True):
|
|
return None
|
|
|
|
if texture_intensity <= 1e-5:
|
|
return None
|
|
|
|
return self._texture_net(device)
|
|
|
|
def _frame_params(self, frame_cfg, detail_intensity):
|
|
detail_boost = frame_cfg["detail_boost"]
|
|
edge_boost = frame_cfg["edge_boost"]
|
|
temporal_strength = frame_cfg["temporal_strength"]
|
|
micro_limit = frame_cfg["micro_limit"]
|
|
luma_boost_mult = frame_cfg["luma_boost_mult"]
|
|
saturation_boost_mult = frame_cfg["saturation_boost_mult"]
|
|
motion_threshold = frame_cfg["motion_threshold"]
|
|
taa_strength = frame_cfg["taa_strength"]
|
|
dlaa_strength = frame_cfg["dlaa_strength"]
|
|
tone_strength = frame_cfg["tone_strength"]
|
|
edge_sharp_strength = frame_cfg["edge_sharp_strength"]
|
|
motion_sensitivity = frame_cfg["motion_sensitivity"]
|
|
jitter_scale = frame_cfg["jitter_scale"]
|
|
|
|
detail_boost *= detail_intensity
|
|
edge_boost *= detail_intensity
|
|
edge_sharp_strength *= detail_intensity
|
|
micro_limit *= detail_intensity
|
|
|
|
return {
|
|
"detail_boost": detail_boost,
|
|
"edge_boost": edge_boost,
|
|
"temporal_strength": temporal_strength,
|
|
"micro_limit": micro_limit,
|
|
"luma_boost_mult": luma_boost_mult,
|
|
"saturation_boost_mult": saturation_boost_mult,
|
|
"motion_threshold": motion_threshold,
|
|
"taa_strength": taa_strength,
|
|
"dlaa_strength": dlaa_strength,
|
|
"tone_strength": tone_strength,
|
|
"edge_sharp_strength": edge_sharp_strength,
|
|
"motion_sensitivity": motion_sensitivity,
|
|
"jitter_scale": jitter_scale,
|
|
}
|
|
|
|
def _apply_jitter_and_taa(
|
|
self,
|
|
rgb,
|
|
taa,
|
|
net,
|
|
motion_sensitivity,
|
|
preset_jitter_scale,
|
|
taa_strength,
|
|
blur_radius,
|
|
):
|
|
motion_gate = 0.0
|
|
fid = taa.frame_id
|
|
jitter_count = self._jitter_count(net)
|
|
|
|
if jitter_count > 0:
|
|
taa.frame_id = (fid + 1) % jitter_count
|
|
else:
|
|
taa.frame_id = fid + 1
|
|
|
|
# damp jitter on motion
|
|
adaptive_jitter_scale = preset_jitter_scale
|
|
|
|
if taa.history is not None and taa.history.shape == rgb.shape:
|
|
motion_estimate = torch.abs(rgb - taa.history).mean()
|
|
motion_gate = torch.clamp(
|
|
motion_estimate * self.motion_gate_scale,
|
|
0.0,
|
|
1.0,
|
|
)
|
|
|
|
jitter_damping = (
|
|
1.0 - motion_estimate * self.jitter_motion_damping
|
|
).clamp(0.45, 1.0)
|
|
adaptive_jitter_scale = preset_jitter_scale * jitter_damping
|
|
|
|
rgb = self._jitter(rgb, fid, adaptive_jitter_scale, net)
|
|
rgb = self._edge_aa(rgb, self.edge_threshold, blur_radius, net)
|
|
|
|
taa_out = taa.update(rgb, self.taa_alpha, motion_sensitivity)
|
|
rgb = torch.lerp(rgb, taa_out, taa_strength)
|
|
|
|
return rgb, motion_gate
|
|
|
|
def _run_dlaa_with_retry(self, net, rgb, tile_size, debug_stats, frame_index):
|
|
current_tile_size = tile_size
|
|
min_tile_size = 128
|
|
dlaa_out = None
|
|
last_oom_error = None
|
|
|
|
while current_tile_size >= min_tile_size:
|
|
try:
|
|
dlaa_out = self._tiled_forward(
|
|
net,
|
|
rgb,
|
|
tile_size=current_tile_size,
|
|
overlap=32,
|
|
)
|
|
|
|
if debug_stats and frame_index == 0:
|
|
raw_model_delta = (dlaa_out - rgb).abs().mean().item()
|
|
log.debug(f"[DLAA] raw_model_delta={raw_model_delta:.6f}")
|
|
|
|
break
|
|
|
|
except torch.OutOfMemoryError as e:
|
|
last_oom_error = e
|
|
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
next_tile_size = current_tile_size // 2
|
|
|
|
if next_tile_size < min_tile_size:
|
|
log.error(
|
|
"DLAA tiled inference failed at minimum tile size %d.",
|
|
current_tile_size,
|
|
)
|
|
raise last_oom_error
|
|
|
|
log.warning(
|
|
"Out of VRAM at tile size %d, retrying with %d.",
|
|
current_tile_size,
|
|
next_tile_size,
|
|
)
|
|
|
|
current_tile_size = next_tile_size
|
|
|
|
if dlaa_out is None:
|
|
raise RuntimeError("DLAA tiled inference failed.")
|
|
|
|
return dlaa_out, current_tile_size
|
|
|
|
def _apply_model_residual(
|
|
self,
|
|
rgb,
|
|
dlaa_out,
|
|
preset,
|
|
debug_stats,
|
|
frame_index,
|
|
):
|
|
# keep brightness stable
|
|
dlaa_mean = dlaa_out.mean(dim=(1, 2, 3), keepdim=True)
|
|
rgb_mean = rgb.mean(dim=(1, 2, 3), keepdim=True)
|
|
dlaa_out = dlaa_out - dlaa_mean + rgb_mean
|
|
|
|
preset_model_weight = self.preset_model_weight.get(preset, 1.00)
|
|
|
|
# residual model pass
|
|
model_delta = dlaa_out - rgb
|
|
|
|
model_delta_value = None
|
|
if debug_stats:
|
|
model_delta_value = model_delta.abs().mean().item()
|
|
|
|
dlaa_out = torch.clamp(
|
|
rgb + model_delta * self.model_weight * preset_model_weight,
|
|
0.0,
|
|
1.0,
|
|
)
|
|
|
|
if debug_stats and frame_index == 0:
|
|
log.debug(f"[DLAA] model_delta_first={model_delta_value:.6f}")
|
|
|
|
return dlaa_out, model_delta_value
|
|
|
|
def _apply_highlight_preblend(self, dlaa_out, rgb, preset):
|
|
luma = rgb_luma(dlaa_out)
|
|
highlight_mask = torch.sigmoid((luma - self.highlight_threshold) * self.highlight_slope)
|
|
|
|
# protect highlights early
|
|
highlight_pre_blend = self.highlight_pre_blend * (
|
|
self.detail_highlight_pre_scale if preset == "Detail" else 1.0
|
|
)
|
|
dlaa_out = torch.lerp(dlaa_out, rgb, highlight_mask * highlight_pre_blend)
|
|
dlaa_out = torch.clamp(dlaa_out, 0.0, 1.0)
|
|
|
|
return dlaa_out, highlight_mask
|
|
|
|
def _apply_motion_suppression(
|
|
self,
|
|
preset,
|
|
detail_boost,
|
|
edge_boost,
|
|
micro_limit,
|
|
motion_gate,
|
|
motion_stability,
|
|
):
|
|
motion_cfg = MOTION_SUPPRESSION["Detail"] if preset == "Detail" else MOTION_SUPPRESSION["Default"]
|
|
|
|
detail_boost *= (1.0 - motion_gate * motion_cfg["detail"] * motion_stability)
|
|
edge_boost *= (1.0 - motion_gate * motion_cfg["edge"] * motion_stability)
|
|
micro_limit *= (1.0 - motion_gate * motion_cfg["micro"] * motion_stability)
|
|
|
|
return detail_boost, edge_boost, micro_limit
|
|
|
|
def _apply_detail_pass(
|
|
self,
|
|
dlaa_out,
|
|
net,
|
|
preset,
|
|
detail_boost,
|
|
edge_boost,
|
|
micro_limit,
|
|
edge_sharp_strength,
|
|
highlight_mask,
|
|
):
|
|
|
|
# local detail pass
|
|
local_avg_rgb = F.avg_pool2d(dlaa_out, 3, stride=1, padding=1)
|
|
fine_detail_rgb = dlaa_out - local_avg_rgb
|
|
|
|
luma = rgb_luma(dlaa_out)
|
|
fine_detail = rgb_luma(fine_detail_rgb)
|
|
|
|
texture_dark_mask = torch.clamp(
|
|
(luma - self.detail_dark_luma_start) / self.detail_dark_luma_range,
|
|
0.0,
|
|
1.0,
|
|
)
|
|
|
|
if preset in ["Detail"]:
|
|
dark_mask = texture_dark_mask
|
|
micro_limit = micro_limit * (
|
|
self.detail_dark_mix_base + self.detail_dark_mix_scale * dark_mask
|
|
)
|
|
else:
|
|
dark_mask = 1.0
|
|
|
|
detail_strength = fine_detail.abs().mean(dim=(1, 2, 3), keepdim=True)
|
|
detail_scale = (
|
|
self.detail_base_scale + (self.detail_ref_scale / (detail_strength + 1e-6))
|
|
).clamp(self.detail_min_scale, self.detail_max_scale)
|
|
|
|
fine_detail = fine_detail * torch.sigmoid(fine_detail * detail_scale)
|
|
fine_detail = fine_detail.clamp(-self.fine_detail_limit, self.fine_detail_limit)
|
|
|
|
local_detail = F.avg_pool2d(fine_detail.abs(), 7, stride=1, padding=3)
|
|
global_detail = fine_detail.abs().mean(dim=(1, 2, 3), keepdim=True)
|
|
|
|
edge_for_detail = torch.sqrt(
|
|
F.conv2d(luma, net.sobel_x, padding=1) ** 2 +
|
|
F.conv2d(luma, net.sobel_y, padding=1) ** 2 +
|
|
1e-6
|
|
)
|
|
|
|
edge_detail_weight = torch.sigmoid(
|
|
(edge_for_detail - self.edge_sharp_threshold) * self.edge_sharp_slope
|
|
)
|
|
|
|
detail_gain = (global_detail / (local_detail + 1e-6)).clamp(
|
|
self.detail_min_gain,
|
|
self.detail_max_gain,
|
|
)
|
|
detail_gain = detail_gain * (1.0 - local_detail.clamp(0.0, 0.5))
|
|
detail_gain = detail_gain * (1.0 + edge_detail_weight * self.detail_edge_boost)
|
|
detail_gain = detail_gain * (1.0 - highlight_mask * self.detail_highlight_suppression)
|
|
|
|
micro_detail = fine_detail_rgb * detail_gain * detail_boost * dark_mask
|
|
micro_detail = micro_detail.clamp(-micro_limit, micro_limit)
|
|
dlaa_out = dlaa_out + micro_detail
|
|
|
|
edge_mask = torch.sigmoid(
|
|
(edge_for_detail - self.edge_sharp_threshold) * self.edge_sharp_slope
|
|
)
|
|
edge_detail = fine_detail_rgb * edge_mask * (1.0 - highlight_mask)
|
|
|
|
edge_boosted = edge_detail * edge_sharp_strength * edge_boost * dark_mask
|
|
edge_boosted = edge_boosted.clamp(
|
|
-micro_limit * self.edge_detail_limit_scale,
|
|
micro_limit * self.edge_detail_limit_scale,
|
|
)
|
|
dlaa_out = dlaa_out + edge_boosted
|
|
|
|
return dlaa_out, texture_dark_mask
|
|
|
|
def _apply_tone_and_color_pass(
|
|
self,
|
|
dlaa_out,
|
|
rgb,
|
|
preset,
|
|
highlight_mask,
|
|
tone_strength,
|
|
luma_boost_mult,
|
|
saturation_boost_mult,
|
|
):
|
|
tone_mapped = dlaa_out / (dlaa_out + self.tone_curve_bias)
|
|
dlaa_out = torch.lerp(dlaa_out, tone_mapped, highlight_mask * tone_strength)
|
|
|
|
luma = rgb_luma(dlaa_out)
|
|
luma_boost = (
|
|
self.luma_boost_base *
|
|
luma_boost_mult *
|
|
(1.0 - highlight_mask * self.luma_highlight_protect)
|
|
)
|
|
dlaa_out = dlaa_out * (1.0 + luma_boost)
|
|
|
|
mean_rgb = dlaa_out.mean(dim=1, keepdim=True)
|
|
saturation_boost = (
|
|
self.saturation_boost_base *
|
|
saturation_boost_mult *
|
|
(1.0 - highlight_mask * self.saturation_highlight_protect)
|
|
)
|
|
dlaa_out = mean_rgb + (dlaa_out - mean_rgb) * (1.0 + saturation_boost)
|
|
|
|
highlight_post_blend = self.highlight_post_blend * (
|
|
self.detail_highlight_post_scale if preset == "Detail" else 1.0
|
|
)
|
|
dlaa_out = torch.lerp(dlaa_out, rgb, highlight_mask * highlight_post_blend)
|
|
|
|
return torch.clamp(dlaa_out, 0.0, 1.0)
|
|
|
|
def _apply_final_temporal_and_blend(
|
|
self,
|
|
rgb,
|
|
dlaa_out,
|
|
prev_dlaa_output,
|
|
preset,
|
|
temporal_strength,
|
|
motion_threshold,
|
|
dlaa_strength,
|
|
):
|
|
if prev_dlaa_output is not None:
|
|
if prev_dlaa_output.shape != dlaa_out.shape:
|
|
prev_dlaa_output = None
|
|
|
|
if preset == "Detail":
|
|
dlaa_out = self._stabilize_fine_detail(
|
|
dlaa_out,
|
|
prev_dlaa_output,
|
|
strength=self.detail_shimmer_strength,
|
|
motion_threshold=motion_threshold,
|
|
)
|
|
|
|
# final temporal blend
|
|
dlaa_out = self._temporal_refine(
|
|
dlaa_out,
|
|
prev_dlaa_output,
|
|
strength=temporal_strength,
|
|
motion_threshold=motion_threshold,
|
|
)
|
|
|
|
prev_dlaa_output = dlaa_out.detach()
|
|
dlaa_out = torch.clamp(dlaa_out, 0.0, 1.0)
|
|
|
|
blend_weight = dlaa_strength * self.dlaa_blend_scale
|
|
|
|
if preset in ["Detail"]:
|
|
blend_weight = min(blend_weight * self.detail_blend_boost, 1.0)
|
|
|
|
rgb = torch.lerp(rgb, dlaa_out, blend_weight)
|
|
|
|
return rgb, prev_dlaa_output
|
|
|
|
def _apply_dlaa_pipeline(
|
|
self,
|
|
rgb,
|
|
net,
|
|
texture_net,
|
|
preset,
|
|
tile_size,
|
|
motion_gate,
|
|
detail_boost,
|
|
edge_boost,
|
|
temporal_strength,
|
|
micro_limit,
|
|
luma_boost_mult,
|
|
saturation_boost_mult,
|
|
motion_threshold,
|
|
dlaa_strength,
|
|
tone_strength,
|
|
edge_sharp_strength,
|
|
motion_stability,
|
|
texture_intensity,
|
|
prev_dlaa_output,
|
|
debug_stats,
|
|
frame_index,
|
|
):
|
|
if dlaa_strength <= 0:
|
|
return rgb, prev_dlaa_output, None
|
|
|
|
dlaa_out, current_tile_size = self._run_dlaa_with_retry(
|
|
net,
|
|
rgb,
|
|
tile_size,
|
|
debug_stats,
|
|
frame_index,
|
|
)
|
|
|
|
dlaa_out, model_delta_value = self._apply_model_residual(
|
|
rgb,
|
|
dlaa_out,
|
|
preset,
|
|
debug_stats,
|
|
frame_index,
|
|
)
|
|
|
|
dlaa_out, highlight_mask = self._apply_highlight_preblend(
|
|
dlaa_out,
|
|
rgb,
|
|
preset,
|
|
)
|
|
|
|
detail_boost, edge_boost, micro_limit = self._apply_motion_suppression(
|
|
preset,
|
|
detail_boost,
|
|
edge_boost,
|
|
micro_limit,
|
|
motion_gate,
|
|
motion_stability,
|
|
)
|
|
|
|
dlaa_out, texture_dark_mask = self._apply_detail_pass(
|
|
dlaa_out,
|
|
net,
|
|
preset,
|
|
detail_boost,
|
|
edge_boost,
|
|
micro_limit,
|
|
edge_sharp_strength,
|
|
highlight_mask,
|
|
)
|
|
|
|
dlaa_out = self._apply_texture_pass(
|
|
texture_net=texture_net,
|
|
dlaa_out=dlaa_out,
|
|
tile_size=current_tile_size,
|
|
motion_gate=motion_gate,
|
|
dark_mask=texture_dark_mask,
|
|
highlight_mask=highlight_mask,
|
|
preset=preset,
|
|
texture_intensity=texture_intensity,
|
|
motion_stability=motion_stability,
|
|
frame_index=frame_index,
|
|
)
|
|
|
|
dlaa_out = self._apply_tone_and_color_pass(
|
|
dlaa_out,
|
|
rgb,
|
|
preset,
|
|
highlight_mask,
|
|
tone_strength,
|
|
luma_boost_mult,
|
|
saturation_boost_mult,
|
|
)
|
|
|
|
rgb, prev_dlaa_output = self._apply_final_temporal_and_blend(
|
|
rgb,
|
|
dlaa_out,
|
|
prev_dlaa_output,
|
|
preset,
|
|
temporal_strength,
|
|
motion_threshold,
|
|
dlaa_strength,
|
|
)
|
|
|
|
return rgb, prev_dlaa_output, model_delta_value
|
|
|
|
def execute(self, images, preset, detail_intensity=1.00, texture_intensity=1.00, motion_stability=1.00):
|
|
preset, detail_intensity, texture_intensity, motion_stability = self._normalize_run_inputs(
|
|
preset,
|
|
detail_intensity,
|
|
texture_intensity,
|
|
motion_stability,
|
|
)
|
|
|
|
device = self._get_device()
|
|
|
|
B, H, W, C = images.shape
|
|
is_single_image = (B == 1)
|
|
blur_radius = self._frame_blur_radius(preset, is_single_image)
|
|
|
|
# per-run temporal state
|
|
taa = TAAState()
|
|
prev_dlaa_output = None
|
|
|
|
net = self._net(device)
|
|
texture_net = self._texture_net_for_run(device, preset, texture_intensity)
|
|
|
|
out_list = []
|
|
delta_sum = 0.0
|
|
delta_count = 0
|
|
debug_stats = log.isEnabledFor(logging.DEBUG)
|
|
|
|
vram_mb = self._vram_mb(device)
|
|
tile_size = self._tile_size_for_vram(vram_mb)
|
|
self._log_tiling(H, W, tile_size)
|
|
|
|
progress = ProgressBar(B) if ProgressBar is not None else None
|
|
|
|
with torch.inference_mode():
|
|
|
|
for i in range(B):
|
|
rgb = self._load_frame(images, i, device)
|
|
|
|
frame_cfg = self._resolve_frame_config(
|
|
preset,
|
|
is_single_image,
|
|
rgb,
|
|
taa,
|
|
)
|
|
params = self._frame_params(frame_cfg, detail_intensity)
|
|
|
|
rgb, motion_gate = self._apply_jitter_and_taa(
|
|
rgb,
|
|
taa,
|
|
net,
|
|
params["motion_sensitivity"],
|
|
params["jitter_scale"],
|
|
params["taa_strength"],
|
|
blur_radius,
|
|
)
|
|
|
|
rgb, prev_dlaa_output, model_delta_value = self._apply_dlaa_pipeline(
|
|
rgb=rgb,
|
|
net=net,
|
|
texture_net=texture_net,
|
|
preset=preset,
|
|
tile_size=tile_size,
|
|
motion_gate=motion_gate,
|
|
detail_boost=params["detail_boost"],
|
|
edge_boost=params["edge_boost"],
|
|
temporal_strength=params["temporal_strength"],
|
|
micro_limit=params["micro_limit"],
|
|
luma_boost_mult=params["luma_boost_mult"],
|
|
saturation_boost_mult=params["saturation_boost_mult"],
|
|
motion_threshold=params["motion_threshold"],
|
|
dlaa_strength=params["dlaa_strength"],
|
|
tone_strength=params["tone_strength"],
|
|
edge_sharp_strength=params["edge_sharp_strength"],
|
|
motion_stability=motion_stability,
|
|
texture_intensity=texture_intensity,
|
|
prev_dlaa_output=prev_dlaa_output,
|
|
debug_stats=debug_stats,
|
|
frame_index=i,
|
|
)
|
|
|
|
if debug_stats and model_delta_value is not None:
|
|
delta_sum += model_delta_value
|
|
delta_count += 1
|
|
|
|
rgb = torch.clamp(rgb, 0.0, 1.0)
|
|
out_list.append(rgb.permute(0, 2, 3, 1).cpu())
|
|
|
|
if progress is not None:
|
|
progress.update(1)
|
|
|
|
if mm is not None and i > 0 and i % 50 == 0:
|
|
mm.soft_empty_cache()
|
|
|
|
if debug_stats and delta_count > 0:
|
|
log.debug(f"[DLAA] model_delta_avg={delta_sum / delta_count:.6f}")
|
|
|
|
return (torch.cat(out_list, dim=0),)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {"VideoTAADLAA": VideoTAADLAA}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"VideoTAADLAA": "🎮 Video Anti-Aliasing (TAA + DLAA)"} |