# Video Detail Refiner # v0.2.1 # SOLRICKS import logging import os import threading import torch import torch.nn.functional as F try: from safetensors.torch import load_file except ImportError: # ComfyUI normally includes safetensors, but keep import safe. load_file = None try: from .model import DLAANet from .utils import rgb_luma, clamp01 except ImportError: from model import DLAANet from utils import rgb_luma, clamp01 try: import comfy.model_management as mm except ImportError: mm = None try: from comfy.utils import ProgressBar except ImportError: ProgressBar = None log = logging.getLogger("VideoDetailRefiner") MIN_EFFECT_STRENGTH = 1e-5 OOM_TILE_STEPS = (1536, 1024, 512, 256) def load_torch_weights(path, device): """Load legacy .pth weights safely when the installed torch supports weights_only.""" try: return torch.load(path, map_location=device, weights_only=True) except TypeError: return torch.load(path, map_location=device) # The presets are intentionally separate from VideoTAADLAA presets. # This node is a detail/texture refiner, not a full anti-aliasing pipeline. RESTORE_PRESETS = { "Photo": { "model_weight": 0.48, "detail_gain": 0.44, "texture_gain": 0.34, "temporal_gain": 0.00, "detail_limit": 0.018, "texture_limit": 0.014, "edge_threshold": 0.056, "highlight_protect": 0.76, "chroma_guard": 0.42, "fine_line_strength": 0.18, "dehalo_strength": 0.050, "micro_contrast_strength": 0.035, "motion_suppression": 0.00, }, "Video Balanced": { "model_weight": 0.55, "detail_gain": 0.55, "texture_gain": 0.45, "temporal_gain": 0.25, "detail_limit": 0.020, "texture_limit": 0.018, "edge_threshold": 0.055, "highlight_protect": 0.66, "chroma_guard": 0.36, "fine_line_strength": 0.22, "dehalo_strength": 0.060, "micro_contrast_strength": 0.040, "motion_suppression": 0.85, }, "Video High Detail": { "model_weight": 0.72, "detail_gain": 0.78, "texture_gain": 0.65, "temporal_gain": 0.20, "detail_limit": 0.026, "texture_limit": 0.022, "edge_threshold": 0.050, "highlight_protect": 0.60, "chroma_guard": 0.44, "fine_line_strength": 0.32, "dehalo_strength": 0.080, "micro_contrast_strength": 0.052, "motion_suppression": 0.75, }, "Soft": { "model_weight": 0.36, "detail_gain": 0.34, "texture_gain": 0.24, "temporal_gain": 0.38, "detail_limit": 0.014, "texture_limit": 0.011, "edge_threshold": 0.070, "highlight_protect": 0.78, "chroma_guard": 0.56, "fine_line_strength": 0.16, "dehalo_strength": 0.045, "micro_contrast_strength": 0.026, "motion_suppression": 0.95, }, "Performance": { "model_weight": 0.00, "detail_gain": 0.42, "texture_gain": 0.00, "temporal_gain": 0.18, "detail_limit": 0.016, "texture_limit": 0.000, "edge_threshold": 0.064, "highlight_protect": 0.72, "chroma_guard": 0.30, "fine_line_strength": 0.14, "dehalo_strength": 0.040, "micro_contrast_strength": 0.022, "motion_suppression": 0.80, }, } PRESET_ALIASES = { # Step 1 / older workflow names "Balanced": "Video Balanced", "High Detail": "Video High Detail", "Soft Restore": "Soft", # extra forgiving aliases "Video Detail": "Video High Detail", "Detail": "Video High Detail", "Photo Restore": "Photo", } class VideoDetailRefiner: def __init__(self): self.model_cache = {} self._model_lock = threading.Lock() self._warned = set() self._warning_lock = threading.Lock() self._sobel_cache = {} @classmethod def INPUT_TYPES(cls): return { "required": { "images": ("IMAGE",), "preset": ([ "Photo", "Video Balanced", "Video High Detail", "Soft", "Performance", ],), "detail_strength": ( "FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.05}, ), "texture_strength": ( "FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.05}, ), "temporal_stability": ( "FLOAT", {"default": 1.00, "min": 0.00, "max": 2.00, "step": 0.05}, ), }, } RETURN_TYPES = ("IMAGE",) FUNCTION = "execute" CATEGORY = "CustomPostProcess" def _warn_once(self, key, message): with self._warning_lock: if key in self._warned: return self._warned.add(key) log.warning(message) def _normalize_preset(self, preset): preset = PRESET_ALIASES.get(preset, preset) return preset if preset in RESTORE_PRESETS else "Video Balanced" def _resolve_run_mode(self, preset, frame_count): """Make image/video behavior explicit and safe for old workflows.""" is_single_image = frame_count == 1 if is_single_image: # Single images should default to the Photo tuning unless the user explicitly # picked Performance/Soft. Old workflows with Video presets stay safe. if preset in ("Video Balanced", "Video High Detail"): return "image", "Photo" return "image", preset # Multi-frame batches should use video tuning. If Photo is used on a video, # fall back to Video Balanced so temporal stabilization is enabled. if preset == "Photo": return "video", "Video Balanced" return "video", preset 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 _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): # Match the TAA node style but keep CPU/VRAM conservative. if vram_mb <= 0: return 512 if vram_mb <= 8192: return 512 if vram_mb <= 16384: return 1024 return 1536 def _tile_retry_sequence(self, initial_tile_size): initial_tile_size = int(initial_tile_size) sequence = [] if initial_tile_size not in OOM_TILE_STEPS: sequence.append(initial_tile_size) for tile in OOM_TILE_STEPS: if tile <= initial_tile_size and tile not in sequence: sequence.append(tile) if 256 not in sequence: sequence.append(256) return sequence def _load_texture_model(self, device): key = str(device) with self._model_lock: if key in self.model_cache: return self.model_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") model_path = None if os.path.exists(safetensors_path): model_path = safetensors_path elif os.path.exists(pth_path): model_path = pth_path if model_path is None: self._warn_once( "missing_texture_model", "[Video Detail Refiner] DLAATexture model not found. Running classical detail pass only.", ) self.model_cache[key] = None return None try: net = DLAANet().to(device) if model_path.endswith(".safetensors"): if load_file is None: raise ImportError("safetensors is not available") state_dict = load_file(model_path, device=str(device)) else: state_dict = load_torch_weights(model_path, 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"] state_dict = { name: tensor for name, tensor in state_dict.items() if name != "jitter_offsets" } load_result = net.load_state_dict(state_dict, strict=False) missing = [name for name in load_result.missing_keys if name != "jitter_offsets"] if missing: log.warning("[Video Detail Refiner] Texture model missing keys: %s", missing) if load_result.unexpected_keys: log.warning( "[Video Detail Refiner] Texture model unexpected keys: %s", load_result.unexpected_keys, ) net = net.float().eval() self.model_cache[key] = net n_params = sum(p.numel() for p in net.parameters()) log.info("[Video Detail Refiner] Loaded %s", os.path.basename(model_path)) log.info("[Video Detail Refiner] Texture model parameters: %.2fM", n_params / 1e6) return net except Exception as e: self._warn_once( "texture_model_load_failed", f"[Video Detail Refiner] Texture model could not be loaded: {type(e).__name__}: {e}", ) self.model_cache[key] = None return None def _get_sobel_kernels(self, device, dtype): cache_key = (str(device), dtype) if cache_key not in self._sobel_cache: sobel_x = torch.tensor( [[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], device=device, dtype=dtype, ).view(1, 1, 3, 3) sobel_y = torch.tensor( [[-1, -2, -1], [0, 0, 0], [1, 2, 1]], device=device, dtype=dtype, ).view(1, 1, 3, 3) self._sobel_cache[cache_key] = (sobel_x, sobel_y) return self._sobel_cache[cache_key] def _luma_edge(self, image, net=None): luma = rgb_luma(image) if net is not None: sobel_x = net.sobel_x.to(device=image.device, dtype=image.dtype) sobel_y = net.sobel_y.to(device=image.device, dtype=image.dtype) else: sobel_x, sobel_y = self._get_sobel_kernels(image.device, image.dtype) sx = F.conv2d(luma, sobel_x, padding=1) sy = F.conv2d(luma, sobel_y, padding=1) edge = torch.sqrt(sx * sx + sy * sy + 1e-6) return luma, edge 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, image, tile_size=1024, overlap=32): if net is None: return image tile_size = int(tile_size) overlap = int(overlap) if tile_size <= 0: raise ValueError(f"Invalid tile_size: {tile_size}") if overlap < 0: raise ValueError(f"Invalid overlap: {overlap}") if overlap * 2 >= tile_size: raise ValueError(f"Invalid tiling settings: tile_size={tile_size}, overlap={overlap}") b, c, h, w = image.shape if h <= tile_size and w <= tile_size: return torch.clamp(net(image), 0.0, 1.0) step = tile_size - overlap * 2 if step <= 0: raise ValueError(f"Invalid tiling settings: tile_size={tile_size}, overlap={overlap}") out = torch.zeros_like(image) weight = torch.zeros(b, 1, h, w, device=image.device, dtype=image.dtype) weight_cache = {} y0 = 0 while y0 < h: y1 = min(y0 + tile_size, h) y0c = max(0, y1 - tile_size) x0 = 0 while x0 < w: x1 = min(x0 + tile_size, w) x0c = max(0, x1 - tile_size) tile = image[:, :, y0c:y1, x0c:x1] tile_out = net(tile) th, tw = tile.shape[2], tile.shape[3] cache_key = (th, tw, overlap, image.device, image.dtype) if cache_key not in weight_cache: weight_cache[cache_key] = self._tile_weight_map( th, tw, overlap, image.device, image.dtype, ) w_map = weight_cache[cache_key] out[:, :, y0c:y1, x0c:x1] += tile_out * w_map weight[:, :, y0c:y1, x0c: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 _run_model_with_retry(self, net, image, initial_tile_size, frame_index=None): if net is None: return None, initial_tile_size last_error = None for tile_size in self._tile_retry_sequence(initial_tile_size): try: model_out = self._tiled_forward(net, image, tile_size=tile_size, overlap=32) return model_out, tile_size except torch.OutOfMemoryError as e: last_error = e if torch.cuda.is_available(): torch.cuda.empty_cache() log.warning( "[Video Detail Refiner] OOM at tile size %d%s, retrying smaller tile.", tile_size, f" on frame {frame_index}" if frame_index is not None else "", ) except RuntimeError as e: if "out of memory" not in str(e).lower(): raise last_error = e if torch.cuda.is_available(): torch.cuda.empty_cache() log.warning( "[Video Detail Refiner] OOM at tile size %d%s, retrying smaller tile.", tile_size, f" on frame {frame_index}" if frame_index is not None else "", ) self._warn_once( "texture_model_oom_skip", "[Video Detail Refiner] Texture model pass failed even at tile size 256. Skipping model pass.", ) if last_error is not None: log.debug("[Video Detail Refiner] Final OOM error: %r", last_error) return None, 256 def _motion_protect(self, cfg, temporal_motion, run_mode): if run_mode != "video": return 1.0 return 1.0 - min(max(float(temporal_motion) * cfg["motion_suppression"], 0.0), 0.95) def _classical_detail_pass(self, image, cfg, detail_strength, temporal_motion, run_mode): detail_strength = max(0.0, min(float(detail_strength), 2.0)) if detail_strength <= MIN_EFFECT_STRENGTH: return image luma, edge = self._luma_edge(image) local_3 = F.avg_pool2d( F.pad(image, [1, 1, 1, 1], mode="reflect"), 3, stride=1, ) local_7 = F.avg_pool2d( F.pad(image, [3, 3, 3, 3], mode="reflect"), 7, stride=1, ) fine_residual = image - local_3 soft_residual = image - local_7 residual = fine_residual * 0.72 + soft_residual * 0.28 detail_energy = rgb_luma(fine_residual.abs()) detail_mask = torch.sigmoid((detail_energy - 0.0045) * 80.0) edge_mask = torch.sigmoid((edge - cfg["edge_threshold"]) * 16.0) highlight_mask = torch.sigmoid((luma - 0.82) * 12.0) highlight_protect = 1.0 - highlight_mask * cfg["highlight_protect"] motion_protect = self._motion_protect(cfg, temporal_motion, run_mode) mask = ( detail_mask * (0.62 + edge_mask * 0.38) * highlight_protect * motion_protect ).clamp(0.0, 1.0) residual = residual.clamp(-cfg["detail_limit"], cfg["detail_limit"]) out = image + residual * mask * cfg["detail_gain"] * detail_strength return torch.clamp(out, 0.0, 1.0) def _micro_contrast_lite(self, image, cfg, detail_strength, temporal_motion, run_mode): strength = cfg["micro_contrast_strength"] * max(0.0, min(float(detail_strength), 2.0)) if strength <= MIN_EFFECT_STRENGTH: return image radius = 5 pad = radius // 2 local_avg = F.avg_pool2d( F.pad(image, [pad, pad, pad, pad], mode="reflect"), radius, stride=1, ) residual = (image - local_avg).clamp(-0.014, 0.014) luma = rgb_luma(image) detail_energy = rgb_luma(residual.abs()) detail_mask = torch.sigmoid((detail_energy - 0.0055) * 80.0) highlight_mask = torch.sigmoid((luma - 0.84) * 12.0) protect = 1.0 - highlight_mask * cfg["highlight_protect"] motion_protect = self._motion_protect(cfg, temporal_motion, run_mode) out = image + residual * detail_mask * protect * motion_protect * strength return torch.clamp(out, 0.0, 1.0) def _fine_line_recovery(self, image, cfg, detail_strength, temporal_motion, run_mode): strength = cfg["fine_line_strength"] * max(0.0, min(float(detail_strength), 2.0)) if strength <= MIN_EFFECT_STRENGTH: return image luma, edge = self._luma_edge(image) local_avg = F.avg_pool2d(image, 3, stride=1, padding=1) detail = image - local_avg detail_energy = rgb_luma(detail.abs()) edge_mask = torch.sigmoid((edge - cfg["edge_threshold"] * 0.92) * 18.0) detail_mask = torch.sigmoid((detail_energy - 0.0060) * 90.0) dark_line_bias = torch.sigmoid((0.48 - luma) * 8.0) highlight_mask = torch.sigmoid((luma - 0.82) * 12.0) line_mask = ( edge_mask * detail_mask * (0.45 + dark_line_bias * 0.55) * (1.0 - highlight_mask * cfg["highlight_protect"]) ).clamp(0.0, 1.0) motion_protect = self._motion_protect(cfg, temporal_motion, run_mode) recovery = detail.clamp(-cfg["detail_limit"] * 0.75, cfg["detail_limit"] * 0.75) out = image + recovery * line_mask * strength * motion_protect return torch.clamp(out, 0.0, 1.0) def _dehalo_lite(self, image, cfg, detail_strength): strength = cfg["dehalo_strength"] * max(0.0, min(float(detail_strength), 2.0)) if strength <= MIN_EFFECT_STRENGTH: return image luma, edge = self._luma_edge(image) edge_mask = torch.sigmoid((edge - cfg["edge_threshold"] * 1.35) * 14.0) local_avg = F.avg_pool2d( F.pad(image, [2, 2, 2, 2], mode="reflect"), 5, stride=1, ) halo_residual = image - local_avg bright_halo = halo_residual.clamp(min=0.0) dark_halo = (-halo_residual).clamp(min=0.0) dark_protect = torch.sigmoid((0.22 - luma) * 12.0) light_protect = torch.sigmoid((luma - 0.78) * 12.0) bright_reduce = bright_halo * edge_mask * strength * (1.0 - light_protect * 0.55) dark_restore = dark_halo * edge_mask * strength * 0.35 * (1.0 - dark_protect * 0.35) out = image - bright_reduce + dark_restore return torch.clamp(out, 0.0, 1.0) def _texture_model_pass(self, net, image, cfg, texture_strength, tile_size, temporal_motion, run_mode, frame_index): texture_strength = max(0.0, min(float(texture_strength), 2.0)) if ( net is None or texture_strength <= MIN_EFFECT_STRENGTH or cfg["model_weight"] <= MIN_EFFECT_STRENGTH or cfg["texture_gain"] <= MIN_EFFECT_STRENGTH ): return image, tile_size model_out, used_tile = self._run_model_with_retry( net, image, initial_tile_size=tile_size, frame_index=frame_index, ) if model_out is None: return image, used_tile if model_out.shape != image.shape: self._warn_once( "texture_model_shape_mismatch", "[Video Detail Refiner] Texture model output shape mismatch. Skipping model pass.", ) return image, used_tile # Keep global brightness stable before extracting the generated texture. model_out = model_out - model_out.mean(dim=(1, 2, 3), keepdim=True) + image.mean(dim=(1, 2, 3), keepdim=True) blur_kernel = 5 if run_mode == "image" else 7 model_blur = F.avg_pool2d( F.pad(model_out, [blur_kernel // 2] * 4, mode="reflect"), blur_kernel, stride=1, ) texture_delta = model_out - model_blur luma, edge = self._luma_edge(image, net) highlight_mask = torch.sigmoid((luma - 0.82) * 12.0) edge_mask = torch.sigmoid((edge - cfg["edge_threshold"]) * 16.0) motion_protect = self._motion_protect(cfg, temporal_motion, run_mode) texture_delta = texture_delta.clamp(-cfg["texture_limit"], cfg["texture_limit"]) mask = ( (0.34 + edge_mask * 0.66) * (1.0 - highlight_mask * cfg["highlight_protect"]) * motion_protect ).clamp(0.0, 1.0) out = image + texture_delta * mask * cfg["model_weight"] * cfg["texture_gain"] * texture_strength return torch.clamp(out, 0.0, 1.0), used_tile def _chroma_guard(self, image, source, strength): strength = clamp01(strength) if strength <= MIN_EFFECT_STRENGTH: return image source_luma = rgb_luma(source) out_luma = rgb_luma(image) source_chroma = source - source_luma out_chroma = image - out_luma chroma_delta = out_chroma - source_chroma chroma_energy = rgb_luma(chroma_delta.abs()) chroma_mask = torch.sigmoid((chroma_energy - 0.010) * 80.0) guarded_chroma = torch.lerp(out_chroma, source_chroma, chroma_mask * strength) out = out_luma + guarded_chroma # Preserve restored luma target after chroma. luma_error = rgb_luma(out) - out_luma out = out - luma_error return torch.clamp(out, 0.0, 1.0) def _temporal_stabilize_detail(self, current, previous, cfg, temporal_stability): temporal_stability = max(0.0, min(float(temporal_stability), 2.0)) if previous is None or previous.shape != current.shape or temporal_stability <= MIN_EFFECT_STRENGTH: 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 base_motion = torch.abs(rgb_luma(current_base) - rgb_luma(previous_base)) stable_mask = (1.0 - torch.clamp(base_motion / 0.080, 0.0, 1.0)).clamp(0.0, 1.0) detail_energy = rgb_luma(current_detail.abs()) detail_mask = torch.sigmoid((detail_energy - 0.006) * 80.0) blend = ( stable_mask * detail_mask * cfg["temporal_gain"] * temporal_stability ).clamp(0.0, 0.45) stabilized_detail = torch.lerp(current_detail, previous_detail, blend) return torch.clamp(current_base + stabilized_detail, 0.0, 1.0) def _estimate_motion(self, current, previous): if previous is None or previous.shape != current.shape: return 0.0 return min(max(torch.abs(rgb_luma(current) - rgb_luma(previous)).mean().item() * 10.0, 0.0), 1.0) def _process_frame( self, rgb, previous_out, net, cfg, detail_strength, texture_strength, temporal_stability, tile_size, run_mode, frame_index, ): source = rgb temporal_motion = self._estimate_motion(rgb, previous_out) if run_mode == "video" else 0.0 out = self._classical_detail_pass(rgb, cfg, detail_strength, temporal_motion, run_mode) out = self._micro_contrast_lite(out, cfg, detail_strength, temporal_motion, run_mode) out = self._fine_line_recovery(out, cfg, detail_strength, temporal_motion, run_mode) out = self._dehalo_lite(out, cfg, detail_strength) out, used_tile = self._texture_model_pass( net, out, cfg, texture_strength, tile_size, temporal_motion, run_mode, frame_index, ) out = self._chroma_guard(out, source, cfg["chroma_guard"]) if run_mode == "video": out = self._temporal_stabilize_detail(out, previous_out, cfg, temporal_stability) return torch.clamp(out, 0.0, 1.0), used_tile def execute(self, images, preset, detail_strength=1.0, texture_strength=1.0, temporal_stability=1.0): detail_strength = max(0.0, min(float(detail_strength), 2.0)) texture_strength = max(0.0, min(float(texture_strength), 2.0)) temporal_stability = max(0.0, min(float(temporal_stability), 2.0)) preset = self._normalize_preset(preset) if len(images.shape) != 4: raise ValueError("VideoDetailRefiner expects an IMAGE tensor with shape [B, H, W, C].") frame_count, height, width, channels = images.shape if channels < 3: raise ValueError("VideoDetailRefiner expects RGB or RGBA images.") run_mode, effective_preset = self._resolve_run_mode(preset, frame_count) cfg = RESTORE_PRESETS.get(effective_preset, RESTORE_PRESETS["Video Balanced"]) device = self._get_device() vram_mb = self._vram_mb(device) tile_size = self._tile_size_for_vram(vram_mb) net = None if cfg["model_weight"] > MIN_EFFECT_STRENGTH and texture_strength > MIN_EFFECT_STRENGTH: net = self._load_texture_model(device) if log.isEnabledFor(logging.DEBUG): log.debug( "[Video Detail Refiner] mode=%s preset=%s effective=%s frames=%d size=%dx%d device=%s vram=%dMB tile=%d", run_mode, preset, effective_preset, frame_count, width, height, device, vram_mb, tile_size, ) # Process RGB, preserve alpha channels extra_channels = None if channels > 3: extra_channels = images[:, :, :, 3:].detach().cpu() out_channels = 3 if extra_channels is None else 3 + extra_channels.shape[-1] out_tensor = torch.empty((frame_count, height, width, out_channels), dtype=images.dtype, device="cpu") progress = ProgressBar(frame_count) if ProgressBar is not None else None previous_out = None current_tile_size = tile_size with torch.inference_mode(): for i in range(frame_count): if mm is not None and hasattr(mm, "throw_exception_if_processing_interrupted"): mm.throw_exception_if_processing_interrupted() rgb = images[i:i + 1].to(device).permute(0, 3, 1, 2).float()[:, :3] restored, used_tile = self._process_frame( rgb=rgb, previous_out=previous_out, net=net, cfg=cfg, detail_strength=detail_strength, texture_strength=texture_strength, temporal_stability=temporal_stability, tile_size=current_tile_size, run_mode=run_mode, frame_index=i, ) if used_tile < current_tile_size: # Keep the successful smaller tile for subsequent frames in the same run. current_tile_size = used_tile previous_out = restored.detach() if run_mode == "video" else None frame_out = restored.permute(0, 2, 3, 1).detach().cpu() if frame_out.dtype != out_tensor.dtype: frame_out = frame_out.to(out_tensor.dtype) if extra_channels is not None: frame_extra = extra_channels[i:i + 1] if frame_extra.dtype != frame_out.dtype: frame_extra = frame_extra.to(frame_out.dtype) frame_out = torch.cat((frame_out, frame_extra), dim=-1) out_tensor[i:i + 1].copy_(frame_out) if progress is not None: progress.update(1) if mm is not None and i > 0 and i % 50 == 0: mm.soft_empty_cache() return (out_tensor,) NODE_CLASS_MAPPINGS = { "VideoDetailRefiner": VideoDetailRefiner, } NODE_DISPLAY_NAME_MAPPINGS = { "VideoDetailRefiner": "✨ Video Detail Refiner", }