VideoDetailRefiner.py

# 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",
}