perf(rtx upscaler): lazy output allocation + mmap off by default

Align memory behaviour with the reference Comfy-Org NVIDIA RTX node:
- torch.empty instead of torch.zeros: physical pages fault in as frames
  are written, so reserving large video batches no longer forces the
  full footprint up front (the OOM source).
- 24 GiB VRAM output ceiling (32 GB card): large batches stay in VRAM
  instead of bouncing to the CPU tier.
- Drop the conservative up-front RAM pre-check from the default path;
  the kernel decides. 'use_mmap' (now off by default, with a warning in
  the help text) is the opt-in last tier of VRAM -> RAM -> disk.
- Collapse per-frame output write to 0 extra buffers in the in-VRAM case
  (strided view + in-place copy_); cross-device/dtype still uses one.
- Director recipe + JS UI default follow use_mmap: false.

Tests: 26/26 green; allocator exercised in isolation for both tiers.
This commit is contained in:
Darksidewalker
2026-08-23 13:55:05 +02:00
parent 24ea823e75
commit 6785f0dcd6
6 changed files with 75 additions and 66 deletions
+1 -1
View File
@@ -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)
+5 -11
View File
@@ -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.
---
+1 -1
View File
@@ -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: [
@@ -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"},
)
+66 -51
View File
@@ -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":
+1 -1
View File
@@ -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