diff --git a/VideoTAADLAA.py b/VideoTAADLAA.py index 94aae9b..c2fb987 100644 --- a/VideoTAADLAA.py +++ b/VideoTAADLAA.py @@ -142,14 +142,20 @@ class VideoTAADLAA: ) state_dict = { - key: tensor - for key, tensor in state_dict.items() - if key != "jitter_offsets" + param_name: tensor + for param_name, tensor in state_dict.items() + if param_name != "jitter_offsets" } load_result = net.load_state_dict(state_dict, strict=False) - if load_result.missing_keys: - log.warning("[DLAA] Main model missing keys: %s", load_result.missing_keys) + + missing_keys = [ + name for name in load_result.missing_keys + if name != "jitter_offsets" + ] + + if missing_keys: + log.warning("[DLAA] Main model missing keys: %s", missing_keys) if load_result.unexpected_keys: log.warning("[DLAA] Main model unexpected keys: %s", load_result.unexpected_keys) @@ -165,7 +171,13 @@ class VideoTAADLAA: return net def _texture_net(self, device): - if not self.texture_pass_enabled: + texture_enabled = getattr( + self, + "texture_pass_enabled", + NODE_DEFAULTS.get("texture_pass_enabled", True) + ) + + if not texture_enabled: return None key = str(device) @@ -208,14 +220,20 @@ class VideoTAADLAA: state_dict = state_dict["model"] state_dict = { - key: tensor - for key, tensor in state_dict.items() - if key != "jitter_offsets" + param_name: tensor + for param_name, tensor in state_dict.items() + if param_name != "jitter_offsets" } load_result = net.load_state_dict(state_dict, strict=False) - if load_result.missing_keys: - log.warning("[DLAA] Texture model missing keys: %s", load_result.missing_keys) + + missing_keys = [ + name for name in load_result.missing_keys + if name != "jitter_offsets" + ] + + if missing_keys: + log.warning("[DLAA] Main model missing keys: %s", missing_keys) if load_result.unexpected_keys: log.warning("[DLAA] Texture model unexpected keys: %s", load_result.unexpected_keys) @@ -287,6 +305,7 @@ class VideoTAADLAA: out = torch.zeros_like(x) weight = torch.zeros(B, 1, H, W, device=x.device, dtype=x.dtype) + weight_cache = {} y0 = 0 while y0 < H: @@ -304,13 +323,18 @@ class VideoTAADLAA: tile_h = tile.shape[2] tile_w = tile.shape[3] - w_map = self._tile_weight_map( - tile_h, - tile_w, - overlap, - x.device, - x.dtype, - ) + cache_key = (tile_h, tile_w, overlap, x.device, x.dtype) + + if cache_key not in weight_cache: + weight_cache[cache_key] = self._tile_weight_map( + tile_h, + tile_w, + overlap, + x.device, + x.dtype, + ) + + w_map = weight_cache[cache_key] out[:, :, y0_c:y1, x0_c:x1] += dlaa_out_tile * w_map weight[:, :, y0_c:y1, x0_c:x1] += w_map @@ -1143,23 +1167,24 @@ class VideoTAADLAA: texture_intensity, motion_stability, ): + # Old preset names from earlier workflows if preset == "Sharp": preset = "Detail" elif preset == "Cinematic": preset = "Smooth" - + dlaa_intensity = max(0.0, min(float(dlaa_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, dlaa_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: if preset == "Photo": @@ -1167,19 +1192,19 @@ class VideoTAADLAA: return 0 return 1 - + def _frame_edge_aa_strength(self, preset, is_single_image): if is_single_image: if preset == "Photo": return self.photo_edge_aa_strength return 0.0 - + if preset == "Detail": return self.detail_edge_aa_strength - + if preset == "Photo": return self.photo_edge_aa_strength - + return 1.0 def _vram_mb(self, device): @@ -1187,30 +1212,30 @@ class VideoTAADLAA: 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") @@ -1218,11 +1243,17 @@ class VideoTAADLAA: 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"]) + + texture_enabled = getattr( + self, + "texture_pass_enabled", + NODE_DEFAULTS.get("texture_pass_enabled", True) + ) - if not self.texture_pass_enabled: + if not texture_enabled: return None if not texture_cfg.get("enabled", True): @@ -1232,7 +1263,7 @@ class VideoTAADLAA: return None return self._texture_net(device) - + def _frame_params(self, frame_cfg, dlaa_intensity): detail_boost = frame_cfg["detail_boost"] edge_boost = frame_cfg["edge_boost"] @@ -1247,12 +1278,12 @@ class VideoTAADLAA: edge_sharp_strength = frame_cfg["edge_sharp_strength"] motion_sensitivity = frame_cfg["motion_sensitivity"] jitter_scale = frame_cfg["jitter_scale"] - + detail_boost *= dlaa_intensity edge_boost *= dlaa_intensity edge_sharp_strength *= dlaa_intensity micro_limit *= dlaa_intensity - + return { "detail_boost": detail_boost, "edge_boost": edge_boost, @@ -1268,7 +1299,7 @@ class VideoTAADLAA: "motion_sensitivity": motion_sensitivity, "jitter_scale": jitter_scale, } - + def _apply_jitter_and_taa( self, rgb, @@ -1291,23 +1322,23 @@ class VideoTAADLAA: # 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_value = float(motion_estimate.item()) - + motion_gate = min( max(motion_value * self.motion_gate_scale, 0.0), 1.0, ) - + jitter_damping = min( max(1.0 - motion_value * self.jitter_motion_damping, 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, @@ -1316,18 +1347,18 @@ class VideoTAADLAA: net, edge_aa_strength, ) - + 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( @@ -1336,41 +1367,41 @@ class VideoTAADLAA: 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, @@ -1383,12 +1414,12 @@ class VideoTAADLAA: 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() @@ -1398,12 +1429,12 @@ class VideoTAADLAA: 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) @@ -1414,9 +1445,9 @@ class VideoTAADLAA: ) 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, @@ -1427,13 +1458,13 @@ class VideoTAADLAA: 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, @@ -1449,16 +1480,16 @@ class VideoTAADLAA: # 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 * ( @@ -1466,28 +1497,28 @@ class VideoTAADLAA: ) 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, @@ -1495,25 +1526,25 @@ class VideoTAADLAA: 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, @@ -1526,7 +1557,7 @@ class VideoTAADLAA: ): 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 * @@ -1534,7 +1565,7 @@ class VideoTAADLAA: (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 * @@ -1547,9 +1578,9 @@ class VideoTAADLAA: 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, @@ -1579,19 +1610,19 @@ class VideoTAADLAA: 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, @@ -1628,7 +1659,7 @@ class VideoTAADLAA: debug_stats, frame_index, ) - + dlaa_out, model_delta_value = self._apply_model_residual( rgb, dlaa_out, @@ -1636,13 +1667,13 @@ class VideoTAADLAA: 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, @@ -1651,7 +1682,7 @@ class VideoTAADLAA: motion_gate, motion_stability, ) - + dlaa_out, texture_dark_mask = self._apply_detail_pass( dlaa_out, net, @@ -1691,21 +1722,18 @@ class VideoTAADLAA: self.detail_specular_strength, ) - if preset == "Detail": dlaa_out = self._micro_contrast( dlaa_out, highlight_mask, self.detail_micro_contrast_strength, ) - if preset == "Detail": dlaa_out = self._edge_dehalo( dlaa_out, net, self.detail_dehalo_strength, ) - if preset == "Detail": dlaa_out = self._subpixel_edge_reconstruction( dlaa_out, net, @@ -1713,14 +1741,12 @@ class VideoTAADLAA: self.detail_subpixel_edge_strength, ) - if preset == "Detail": dlaa_out = self._chroma_edge_cleanup( dlaa_out, net, self.detail_chroma_cleanup_strength, ) - if preset == "Detail": dlaa_out = self._local_tone_mapping( dlaa_out, highlight_mask, @@ -1746,7 +1772,6 @@ class VideoTAADLAA: self.detail_specular_temporal_strength, ) - if preset == "Detail": dlaa_out = self._fur_hair_stabilizer( dlaa_out, prev_dlaa_output, @@ -1754,7 +1779,7 @@ class VideoTAADLAA: motion_gate, self.detail_fur_stabilizer_strength, ) - + rgb, prev_dlaa_output = self._apply_final_temporal_and_blend( rgb, dlaa_out, @@ -1784,7 +1809,7 @@ class VideoTAADLAA: B, H, W, C = images.shape is_single_image = (B == 1) - + blur_radius = self._frame_blur_radius(preset, is_single_image) edge_aa_strength = self._frame_edge_aa_strength( preset, diff --git a/config.py b/config.py index c3f300f..fc650df 100644 --- a/config.py +++ b/config.py @@ -6,19 +6,19 @@ NODE_DEFAULTS = { # Model blend "model_weight": 1.00, "dlaa_blend_scale": 1.00, - + # TAA / jitter "taa_alpha": 0.10, "jitter_scale": 0.20, "edge_threshold": 0.15, - + # Highlight handling "tone_curve_bias": 0.6, "highlight_pre_blend": 0.15, "highlight_post_blend": 0.08, "highlight_threshold": 0.85, "highlight_slope": 12.0, - + # Detail shaping "detail_base_scale": 9.0, "detail_ref_scale": 0.02, @@ -28,40 +28,40 @@ NODE_DEFAULTS = { "detail_max_gain": 0.26, "detail_edge_boost": 0.35, "detail_highlight_suppression": 0.5, - + # Edge detail "edge_sharp_threshold": 0.08, "edge_sharp_slope": 12.0, "edge_aa_slope": 14.0, "edge_detail_limit_scale": 0.7, - + # Motion handling "motion_gate_scale": 12.0, "jitter_motion_damping": 8.0, - + # Dark-area detail behavior "detail_dark_luma_start": 0.20, "detail_dark_luma_range": 0.30, "detail_dark_mix_base": 0.6, "detail_dark_mix_scale": 0.4, - + # Detail limits "fine_detail_limit": 0.15, "detail_highlight_pre_scale": 0.45, "detail_highlight_post_scale": 0.50, "detail_blend_boost": 1.12, - + # Auto preset motion thresholds "auto_default_scene_motion": 0.02, "auto_static_motion_threshold": 0.015, "auto_balanced_motion_threshold": 0.045, - + # Luma / saturation boost "luma_boost_base": 0.03, "saturation_boost_base": 0.06, "luma_highlight_protect": 0.6, "saturation_highlight_protect": 0.5, - + # Texture pass "texture_pass_enabled": True, "texture_tile_overlap": 32, @@ -101,14 +101,14 @@ NODE_DEFAULTS = { "detail_dehalo_light_protect": 0.55, # Chroma edge cleanup - "detail_chroma_cleanup_strength": 0.18, # best range 0.10 - 0.25 + "detail_chroma_cleanup_strength": 0.18, # best range 0.10 - 0.25 "detail_chroma_edge_threshold": 0.055, "detail_chroma_saturation_threshold": 0.070, "detail_chroma_cleanup_limit": 0.014, "detail_chroma_dark_protect": 0.30, # Subpixel edge reconstruction - "detail_subpixel_edge_strength": 0.14, # best range 0.10 - 0.20 + "detail_subpixel_edge_strength": 0.14, # best range 0.10 - 0.20 "detail_subpixel_edge_threshold": 0.070, "detail_subpixel_edge_slope": 16.0, "detail_subpixel_sample_scale": 0.35, @@ -117,7 +117,7 @@ NODE_DEFAULTS = { "detail_subpixel_motion_protect": 0.75, # Temporal specular stabilizer - "detail_specular_temporal_strength": 0.05, # best range 0.05 - 0.15 + "detail_specular_temporal_strength": 0.05, # best range 0.05 - 0.15 "detail_specular_temporal_threshold": 0.56, "detail_specular_temporal_slope": 10.0, "detail_specular_temporal_detail_threshold": 0.006, @@ -126,7 +126,7 @@ NODE_DEFAULTS = { "detail_specular_temporal_motion_protect": 0.85, # Local tone mapping - "detail_local_tonemap_strength": 0.07, # best range 0.04 - 0.10 + "detail_local_tonemap_strength": 0.07, # best range 0.04 - 0.10 "detail_local_tonemap_radius": 11, "detail_local_tonemap_limit": 0.018, "detail_local_tonemap_shadow_lift": 0.012,