diff --git a/README.md b/README.md index 55e52c6..8f25908 100644 --- a/README.md +++ b/README.md @@ -68,7 +68,7 @@ State-of-the-art image and video enhancement using NVIDIA RTX Video SDK. It exec - **Upscale:** AI-powered VSR and High Bitrate upscaling. - **Smart Sizing:** Multiple resize modes including Constant Megapixel targets. - **Efficiency:** Frame-by-frame processing for minimal VRAM usage. -- **Memory Control:** Disk-backed (mmap) fallback (`use_mmap`, on by default) is only taken as the last tier of the VRAM -> RAM -> disk chain, after automatic model unloading when memory runs short (`auto_unload_models`, on by default). `use_mmap` off = never use disk (uses RAM when possible, else a clear error). +- **Memory Control:** The output batch is allocated lazily (like the reference NVIDIA node — the kernel decides, no up-front memory pressure, no temp file). A disk-backed (mmap) fallback (`use_mmap`, off by default) is opt-in for very long video batches: when enabled it is the last tier of the VRAM -> RAM -> disk chain, taken only when available RAM is still short after automatic model unloading (`auto_unload_models`, on by default). **Warning:** enabling `use_mmap` writes a multi-giB `.mmap` temp file to your temp drive for the whole run. ![RTX_UpscalerRefiner.png](assets/RTX_UpscalerRefiner.png) diff --git a/docs/rtx_upscaler_refiner.md b/docs/rtx_upscaler_refiner.md index 6339ccc..3b53432 100644 --- a/docs/rtx_upscaler_refiner.md +++ b/docs/rtx_upscaler_refiner.md @@ -47,20 +47,14 @@ Check these settings first when users report degraded quality: | **divisible_by** | Use **8** for RTX-style image/video upscaling. Use **32** only when the next model in the workflow requires it. | | **resize_method** | When the source and target aspect ratios don't match: `Center Crop` fills the whole target frame by cutting edges; `Letterbox` fits the whole image inside and adds black bars. | | **device_id** | If you have multiple GPUs, set this to the index of your RTX card (usually `0`). | -| **use_mmap** | **On** (default): permits the disk-backed (mmap) fallback. Preference is always VRAM -> RAM -> disk; disk is used only when free VRAM/RAM is still short after `auto_unload_models`. **Off**: never use disk — if memory stays short the output uses RAM (CPU) when possible, otherwise a clear error is raised. | +| **use_mmap** | **Off** (default): disk is never used — the output is allocated lazily in memory (VRAM when it fits, otherwise RAM) and the kernel decides. **On**: opt-in disk-backed (mmap) fallback as the last tier of the VRAM -> RAM -> disk chain; it writes a multi-giB `.mmap` temp file to your temp drive for the whole run. Only enable it for very long video batches that genuinely exceed available RAM. | | **auto_unload_models** | **On** (default): when free VRAM/RAM is insufficient for the output, unloads ComfyUI-managed models (a full `unload_all_models`, like the manual empty-cache path but more thorough) and re-checks before falling back to disk. **Off**: fall back immediately without unloading. | ## 💾 Disk-Backed (mmap) Output — `use_mmap` -`use_mmap` (on by default) *permits* the disk-backed (mmap) fallback; it is the last tier of the VRAM -> RAM -> disk chain and is only used when free VRAM/RAM is still insufficient after `auto_unload_models`. RAM/VRAM is always preferred when it fits; the file is deleted as soon as the output tensor is released. When `use_mmap` is **off**, disk is never used — the output uses RAM when possible, and otherwise the node raises a clear error telling you to enable `use_mmap` or reduce the batch/resolution. +`use_mmap` is **off by default**: the output is allocated lazily in memory (VRAM when it fits, otherwise RAM) and the kernel decides — matching the reference NVIDIA RTX node, so reserving a large video batch never forces the full footprint up front and no temp file is ever created. When `use_mmap` is **enabled** (opt-in, only for very long video batches that genuinely exceed available RAM), it *permits* the disk-backed (mmap) fallback; it is the last tier of the VRAM -> RAM -> disk chain and is only used when free VRAM/RAM is still insufficient after `auto_unload_models`. The file is deleted as soon as the output tensor is released. -> **Disk space (issue #29):** disk-backed files can be several GiB for long 4K -> batches (8.99 GB was observed with Frame Interpolation). The node checks free -> disk space before creating the file and raises a clear error when the temp -> drive cannot hold output + 1 GiB reserve. On Windows, an unlinked mmap file -> only releases its disk space when the last handle to it is closed — while a -> downstream node still holds the output, the space stays charged. Process -> batches in smaller chunks to cap simultaneous disk usage. +> **Warning:** enabling `use_mmap` writes a temp file to your ComfyUI temp drive that can grow to the **full output size** (several GiB for long 4K batches; 8.99 GB was observed with Frame Interpolation) and stays on disk for the whole run. The node checks free disk space before creating the file and raises a clear error when the temp drive cannot hold output + 1 GiB reserve. On Windows, an unlinked mmap file only releases its disk space when the last handle to it is closed — while a downstream node still holds the output, the space stays charged. ## 💡 Pro-Tips @@ -70,9 +64,9 @@ Check these settings first when users report degraded quality: ## CPU Batch Storage -CPU output batches use the currently available RAM reported by `psutil`, not a fixed RAM cap. The node reserves 25% of total RAM on systems below 32 GiB (minimum 1 GiB), rising smoothly to an 8 GiB maximum reserve on larger systems. It uses disk-backed temporary storage only when allocating the complete output batch would cross that reserve. The check runs at allocation time rather than at ComfyUI startup, so it accounts for models and other workloads loaded after startup. GPU outputs retain their separate 8 GiB VRAM safety limit. +CPU output batches use the currently available RAM reported by `psutil`, not a fixed RAM cap. The node reserves 25% of total RAM on systems below 32 GiB (minimum 1 GiB), rising smoothly to an 8 GiB maximum reserve on larger systems. The reserve is used to decide whether the opt-in disk-backed (`use_mmap`) fallback is needed; when `use_mmap` is off (default), the output is allocated lazily in RAM regardless and the kernel decides. The check runs at allocation time rather than at ComfyUI startup, so it accounts for models and other workloads loaded after startup. GPU outputs retain their separate 24 GiB VRAM safety limit. -- **`use_mmap`** (on by default) *permits* the disk-backed path as the final tier of the VRAM -> RAM -> disk chain, taken only when free VRAM/RAM is still insufficient after auto-unload. **Off** disables disk entirely (the node uses RAM when possible and raises a clear error otherwise). +- **`use_mmap`** (off by default) *permits* the disk-backed path as the final tier of the VRAM -> RAM -> disk chain, taken only when available RAM cannot back the batch and `auto_unload_models` did not make room. **Off** (default) disables disk entirely: the output is allocated lazily in memory and a genuine OOM surfaces naturally instead of being pre-empted by a temp file. - **`auto_unload_models`** (on by default) is the first remediation when the active tier (VRAM, or RAM on CPU) is short: it fully unloads ComfyUI-managed models, re-checks that tier, and only then drops to a lower tier. --- diff --git a/js/minimax_h3_director_v2.js b/js/minimax_h3_director_v2.js index a177931..c440224 100644 --- a/js/minimax_h3_director_v2.js +++ b/js/minimax_h3_director_v2.js @@ -72,7 +72,7 @@ const PP_STAGE_FIELDS = { { key: "resize_method", type: "combo", options: ["Center Crop (Fill)", "Letterbox (Fit)"], default: "Center Crop (Fill)" }, { key: "device_id", type: "number", min: 0, max: 8, step: 1, default: 0 }, { key: "empty_cache", type: "boolean", default: false }, - { key: "use_mmap", type: "boolean", default: true }, + { key: "use_mmap", type: "boolean", default: false }, { key: "auto_unload_models", type: "boolean", default: true }, ], watermark: [ diff --git a/nodes/helper_minimax_h3_director_execute_v2.py b/nodes/helper_minimax_h3_director_execute_v2.py index 7efc83c..f506649 100644 --- a/nodes/helper_minimax_h3_director_execute_v2.py +++ b/nodes/helper_minimax_h3_director_execute_v2.py @@ -5,7 +5,7 @@ DEFAULT_POSTPROCESS_RECIPE = ( {"id": "frame_interpolation", "enabled": False, "factor": 2, "model": "rife_v4.26.safetensors"}, {"id": "torch_resize", "enabled": False, "size_mode": "Multiplier", "scale_multiplier": 2, "interpolation": "Lanczos"}, {"id": "model_upscale", "enabled": False, "model_name": "2x-AnimeSharpV4_RCAN.safetensors"}, - {"id": "rtx_refiner", "enabled": False, "denoise": False, "denoise_quality": "Ultra", "deblur": False, "deblur_quality": "Ultra", "upscale": "VSR", "upscale_quality": "Ultra", "resize_type": "Scale", "scale": 2.0, "megapixels": 2.0, "width": 1920, "height": 1080, "divisible_by": "8", "ratio_preset": "16:9", "resize_method": "Center Crop (Fill)", "device_id": 0, "empty_cache": False, "use_mmap": True, "auto_unload_models": True}, + {"id": "rtx_refiner", "enabled": False, "denoise": False, "denoise_quality": "Ultra", "deblur": False, "deblur_quality": "Ultra", "upscale": "VSR", "upscale_quality": "Ultra", "resize_type": "Scale", "scale": 2.0, "megapixels": 2.0, "width": 1920, "height": 1080, "divisible_by": "8", "ratio_preset": "16:9", "resize_method": "Center Crop (Fill)", "device_id": 0, "empty_cache": False, "use_mmap": False, "auto_unload_models": True}, {"id": "watermark", "enabled": False, "watermark_path": "", "position": "bottom-right"}, ) diff --git a/nodes/nodes_rtx_upscaler_refiner.py b/nodes/nodes_rtx_upscaler_refiner.py index 0f02abe..66a1b14 100644 --- a/nodes/nodes_rtx_upscaler_refiner.py +++ b/nodes/nodes_rtx_upscaler_refiner.py @@ -4,7 +4,7 @@ import os import torch import torch.nn.functional as F import contextlib -from typing import Tuple +from typing import Optional, Tuple import folder_paths try: @@ -37,7 +37,10 @@ DIVISIBLE_BY_VALUES = ["8", "16", "32", "64", "128"] COMMON_RATIOS = ["1:1", "4:3", "3:2", "16:9", "21:9"] RESIZE_METHODS = ["Center Crop (Fill)", "Letterbox (Fit)"] MAX_CHUNK_OUTPUT_PIXELS = 1024 * 1024 * 16 -MAX_GPU_OUTPUT_BYTES = 8 * 1024 * 1024 * 1024 +# Ceiling for output kept on the GPU. Sized for a 32 GB card: a large batch can +# stay in VRAM when free VRAM allows (the 15% headroom check still self-limits), +# instead of being bounced to RAM just because the old 8 GiB cap was hit. +MAX_GPU_OUTPUT_BYTES = 24 * 1024 * 1024 * 1024 # --- Helpers --- @@ -86,18 +89,25 @@ def _allocate_output_tensor( dtype: torch.dtype, device: torch.device, soft_empty_cache: bool = False, - allow_mmap: bool = True, + allow_mmap: bool = False, auto_unload_models: bool = True, ): """Allocate the output tensor. - Preference order (always): VRAM -> RAM -> disk. + Allocation is LAZY (``torch.empty``, no zero-fill) to match the reference + NVIDIA RTX node: physical pages are faulted in only as frames are written, + so reserving a large video batch never forces the full footprint up front. + This is what keeps the node from OOMing at allocation time on long videos. + Preference order: VRAM -> RAM -> disk. + - VRAM tier self-limits to a 24 GiB ceiling plus a 15% free-VRAM headroom. - ``auto_unload_models``: when the active tier is short, fully unload - ComfyUI-managed models and re-check that tier before dropping to a lower - tier. This is the first remediation, not a last resort. - - ``allow_mmap`` (the ``use_mmap`` switch): when False, disk is never used; - if every non-disk tier is short, raise a clear error. + ComfyUI-managed models and re-check before dropping a tier. + - ``allow_mmap`` (the ``use_mmap`` switch, OFF by default): disk is used + only when enabled AND the batch will not fit in available RAM. When off, + the node always allocates lazily in memory and lets the kernel decide + (matching the reference); a genuine OOM surfaces naturally rather than + being pre-empted by a temp file. Returns ``(tensor, mmap_path)`` where ``mmap_path`` is None for in-memory tensors and the temp file path for the disk-backed fallback. @@ -114,7 +124,8 @@ def _allocate_output_tensor( and _can_fit_in_vram(required_bytes, device) ) - def ram_fits() -> bool: + def ram_will_fit() -> bool: + # Used only to decide whether the opt-in disk tier is actually needed. return can_allocate_in_ram(required_bytes) def try_mmap(): @@ -128,18 +139,17 @@ def _allocate_output_tensor( ), ) - def no_disk_allowed(short_of: str): - raise RuntimeError( - f"RTX output requires {required_bytes / 1024 ** 3:.2f} GiB and " - f"{short_of} is still insufficient, but 'use_mmap' is off (never use " - "disk-backed output). Enable 'use_mmap', free system memory, or reduce " - "the batch size / target resolution, then retry." - ) + def disk_needed_and_allowed() -> bool: + """Disk is the fallback only when enabled and RAM will not back the batch.""" + return allow_mmap and not ram_will_fit() + + def alloc_vram(): + return torch.empty(shape, device=device, dtype=dtype), None if device.type == "cuda": # Tier 1: VRAM. if gpu_fits(): - return torch.zeros(shape, device=device, dtype=dtype), None + return alloc_vram() # First remediation: auto-unload, re-check VRAM. if auto_unload_models and unload_all_comfy_models(): log_dasiwa("RTX Upscaler & Refiner", @@ -147,36 +157,32 @@ def _allocate_output_tensor( if gpu_fits(): log_dasiwa("RTX Upscaler & Refiner", "VRAM sufficient after auto-unload; keeping GPU output.") - return torch.zeros(shape, device=device, dtype=dtype), None - # Tier 2: RAM (always preferred over disk). - if ram_fits(): + return alloc_vram() + # Tier 2: RAM (lazy; preferred over disk). + if disk_needed_and_allowed(): log_dasiwa("RTX Upscaler & Refiner", - f"VRAM short for {required_bytes / 1024 ** 3:.2f} GiB; using RAM (CPU) instead of disk.") - return torch.zeros(shape, dtype=dtype), None - # Tier 3: disk (only if permitted). - if not allow_mmap: - no_disk_allowed("VRAM and RAM") + f"VRAM and RAM short for {required_bytes / 1024 ** 3:.2f} GiB; " + "falling back to disk-backed (mmap) output.") + return try_mmap() log_dasiwa("RTX Upscaler & Refiner", - f"VRAM and RAM short for {required_bytes / 1024 ** 3:.2f} GiB; " + f"VRAM short for {required_bytes / 1024 ** 3:.2f} GiB; " + "allocating lazily in RAM (CPU).") + return torch.empty(shape, dtype=dtype), None + + # CPU device: Tier 1 is RAM (lazy; no up-front pre-check, the kernel decides). + if auto_unload_models and not ram_will_fit(): + if unload_all_comfy_models(): + log_dasiwa("RTX Upscaler & Refiner", + "Auto-unload: RAM short; unloaded ComfyUI-managed models to make room.") + if disk_needed_and_allowed(): + log_dasiwa("RTX Upscaler & Refiner", + f"RAM short for {required_bytes / 1024 ** 3:.2f} GiB; " "falling back to disk-backed (mmap) output.") return try_mmap() - - # CPU device: Tier 1 is RAM. - if ram_fits(): - return torch.zeros(shape, dtype=dtype), None - if auto_unload_models and unload_all_comfy_models(): - log_dasiwa("RTX Upscaler & Refiner", - "Auto-unload: RAM insufficient; unloaded ComfyUI-managed models and re-checking.") - if ram_fits(): - log_dasiwa("RTX Upscaler & Refiner", - "RAM sufficient after auto-unload; using RAM instead of disk.") - return torch.zeros(shape, dtype=dtype), None - if not allow_mmap: - no_disk_allowed("RAM") log_dasiwa("RTX Upscaler & Refiner", - f"RAM short for {required_bytes / 1024 ** 3:.2f} GiB; " - "falling back to disk-backed (mmap) output.") - return try_mmap() + f"Output is {required_bytes / 1024 ** 3:.2f} GiB; " + "allocating lazily in RAM (CPU, kernel decides).") + return torch.empty(shape, dtype=dtype), None def _aligned_aspect_size( target_width: float, @@ -513,7 +519,7 @@ class DaSiWa_RTX_UpscalerRefiner: }, "optional": { "empty_cache": ("BOOLEAN", {"default": False, "description": "Run a lightweight GC + empty_cache before allocation; also calls model_management.soft_empty_cache() to ask ComfyUI models to unload, which can free VRAM but may slow subsequent nodes."}), - "use_mmap": ("BOOLEAN", {"default": True, "description": "Allow the disk-backed (mmap) output fallback. Preference is always VRAM -> RAM -> disk. Only when free VRAM/RAM is still insufficient after 'auto_unload_models' does the output spill to a temp .mmap file. Off: never use disk — if memory stays short the output uses RAM when possible, otherwise a clear error is raised (free RAM or reduce the batch/resolution)."}), + "use_mmap": ("BOOLEAN", {"default": False, "description": "OFF (default): never use disk — the output is allocated lazily in memory (VRAM when it fits, otherwise RAM) and the kernel decides, like the reference NVIDIA node. **WARN:** enable this ONLY for very long video batches that genuinely exceed available RAM; when on, disk is used as the last tier of the VRAM -> RAM -> disk chain and a multi-giB .mmap temp file is written to your temp drive for the whole run."}), "auto_unload_models": ("BOOLEAN", {"default": True, "description": "When VRAM/RAM is insufficient for the output, automatically unload ComfyUI-managed models (like the manual 'empty cache' but a full unload) and re-check before falling back to disk. Off: skip the unload and fall back directly."}), }, } @@ -555,7 +561,7 @@ class DaSiWa_RTX_UpscalerRefiner: resize_method, device_id, empty_cache=False, - use_mmap=True, + use_mmap=False, auto_unload_models=True, ): if not torch.cuda.is_available(): @@ -662,13 +668,22 @@ class DaSiWa_RTX_UpscalerRefiner: if upscale_effect: frame = _run_vfx_effect(upscale_effect, frame, cuda_device) - output_frame = ( - frame.permute(1, 2, 0) - .contiguous() - .clamp(0.0, 1.0) - .to(device=out_device, dtype=out_dtype, non_blocking=False) - ) - out[global_index].copy_(output_frame, non_blocking=False) + row = out[global_index] # (H,W,3) view into the output, on out_device + if out_device == cuda_device and out_dtype == frame.dtype: + # In-VRAM path: zero extra buffers. A zero-copy permute + # view feeds an in-place copy into the row, then clamp + # in place (the row aliases the output, never the input). + # Matches the reference node's single in-place write. + row.copy_(frame.permute(1, 2, 0), non_blocking=True) + row.clamp_(0.0, 1.0) + else: + # Cross-device/dtype path (e.g. CPU output): one + # materialized buffer for the transpose + clamp; the + # copy_ below transfers and casts into the row. + row.copy_( + frame.permute(1, 2, 0).clamp(0.0, 1.0), + non_blocking=(out_device.type == "cuda"), + ) torch.cuda.synchronize(cuda_device) if out_device.type == "cuda": diff --git a/test/test_minimax_h3_director_execute.py b/test/test_minimax_h3_director_execute.py index 12d1474..314b985 100644 --- a/test/test_minimax_h3_director_execute.py +++ b/test/test_minimax_h3_director_execute.py @@ -45,5 +45,5 @@ def test_rtx_refiner_recipe_exposes_full_node_options(): assert rtx["resize_type"] == "Scale" assert rtx["scale"] == 2.0 assert rtx["divisible_by"] == "8" - assert rtx["use_mmap"] is True + assert rtx["use_mmap"] is False assert rtx["auto_unload_models"] is True