[PR #36] [CLOSED] fix: free dlpack tensor per-frame to prevent GPU memory leak in RTX VSR node #40

Closed
opened 2026-10-26 22:28:19 +00:00 by bsakaguchi · 0 comments
Owner

📋 Pull Request Information

Original PR: https://github.com/Comfy-Org/ComfyUI_Nvidia_VFX_Nodes/pull/36
Author: @CharlesMcquade
Created: 7/28/2026
Status: ❌ Closed

Base: main ← Head: fix/vsr-memory-leak


📝 Commits (2)

  • 77bc194 fix: free dlpack tensor per-frame to prevent GPU memory leak in RTX VSR node
  • 71a5f69 perf(vsr): avoid flushing CUDA cache per frame

📊 Changes

1 file changed (+5 additions, -2 deletions)

View changed files

📝 __init__.py (+5 -2)

📄 Description

The problem, from the user's seat

You drop a 200-frame video into ComfyUI, wire it through the RTX Video Super Resolution node, hit Queue, and walk away. A minute later nvidia-smi shows VRAM pegged at 100%, the ComfyUI prompt stops making progress, and the only way out is to kill the process. Short batches are fine — the node sings on 10 frames. It's only once you get past ~60 frames that it slowly suffocates.

This is exactly what andypotato reported in #16: "When upscaling batches of images it will slowly eat up all my system memory. At some point it will start swapping and then spiral the system into swapping death loop. Usually it happens around 60 - 70 images upscaled." It's been open and reproducible for a while. Tracked in the new issue #35.

What's actually leaking

In __init__.py, the execute() loop calls sr.run(input_frame) per frame and converts the returned dlpack tensor via torch.from_dlpack():

for j in range(batch_cuda.shape[0]):
    input_frame = batch_cuda[j]
    dlpack_out = sr.run(input_frame).image
    out_tensor[i + j: i + j + 1] = torch.from_dlpack(dlpack_out).movedim(0, -1).unsqueeze(0)

The problem: torch.from_dlpack(dlpack_out) produces a PyTorch GPU tensor that shares the dlpack's underlying storage. The dlpack_out variable gets overwritten on the next iteration, but neither the dlpack capsule nor the PyTorch view that aliases it is ever explicitly released. The .movedim() / .unsqueeze() calls produce views, not new allocations, so the original per-frame storage stays pinned for the life of the loop. The CUDA caching allocator holds onto all of it, VRAM climbs monotonically, and eventually the node OOMs or deadlocks on a blocking CUDA call.

The fix

Two changes to the per-frame loop, intended to release DLPack-backed storage each iteration without forcing a per-frame CUDA cache flush:

  1. .clone() the converted tensor — breaks the aliasing so the output owns its own storage and the dlpack-backed tensor becomes immediately releasable.
  2. del dlpack_out, result, out_frame — drops the references this iteration instead of waiting for the frame to go out of scope.
for j in range(batch_cuda.shape[0]):
    input_frame = batch_cuda[j]
    result = sr.run(input_frame)
    dlpack_out = result.image
    out_frame = torch.from_dlpack(dlpack_out).movedim(0, -1).unsqueeze(0).clone()
    out_tensor[i + j: i + j + 1] = out_frame
    del dlpack_out, result, out_frame

Verification

Tested locally with 100+ frame runs where the original code OOM'd around ~60 frames. With this patch VRAM stays flat across the whole run — nvidia-smi holds a steady plateau instead of the old monotonic climb. Output is pixel-identical (.clone() is a copy, not a transform).

Follow-up 71a5f69 removes the per-frame empty_cache() in response to the review: the allocator can reuse released blocks without a forced flush each frame. The 100+ frame observation above was made on the prior version with empty_cache(); GPU memory behavior after its removal has not been re-measured, so that specific result is not claimed for the current head.

Scope

  • One file changed: __init__.py
  • No API or behaviour change for existing graphs — same inputs, same outputs, just without the memory leak.

Closes #35. Related: #16.


🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.

## 📋 Pull Request Information **Original PR:** https://github.com/Comfy-Org/ComfyUI_Nvidia_VFX_Nodes/pull/36 **Author:** [@CharlesMcquade](https://github.com/CharlesMcquade) **Created:** 7/28/2026 **Status:** ❌ Closed **Base:** `main` ← **Head:** `fix/vsr-memory-leak` --- ### 📝 Commits (2) - [`77bc194`](https://github.com/Comfy-Org/ComfyUI_Nvidia_VFX_Nodes/commit/77bc194dc0da087fb842cc5b39d33a6926b44018) fix: free dlpack tensor per-frame to prevent GPU memory leak in RTX VSR node - [`71a5f69`](https://github.com/Comfy-Org/ComfyUI_Nvidia_VFX_Nodes/commit/71a5f69202ae0f5ffd38206b5b52ee10d6cb8fee) perf(vsr): avoid flushing CUDA cache per frame ### 📊 Changes **1 file changed** (+5 additions, -2 deletions) <details> <summary>View changed files</summary> 📝 `__init__.py` (+5 -2) </details> ### 📄 Description ## The problem, from the user's seat You drop a 200-frame video into ComfyUI, wire it through the `RTX Video Super Resolution` node, hit Queue, and walk away. A minute later `nvidia-smi` shows VRAM pegged at 100%, the ComfyUI prompt stops making progress, and the only way out is to kill the process. Short batches are fine — the node sings on 10 frames. It's only once you get past ~60 frames that it slowly suffocates. This is exactly what andypotato reported in #16: *"When upscaling batches of images it will slowly eat up all my system memory. At some point it will start swapping and then spiral the system into swapping death loop. Usually it happens around 60 - 70 images upscaled."* It's been open and reproducible for a while. Tracked in the new issue #35. ## What's actually leaking In `__init__.py`, the `execute()` loop calls `sr.run(input_frame)` per frame and converts the returned dlpack tensor via `torch.from_dlpack()`: ```python for j in range(batch_cuda.shape[0]): input_frame = batch_cuda[j] dlpack_out = sr.run(input_frame).image out_tensor[i + j: i + j + 1] = torch.from_dlpack(dlpack_out).movedim(0, -1).unsqueeze(0) ``` The problem: `torch.from_dlpack(dlpack_out)` produces a PyTorch GPU tensor that **shares** the dlpack's underlying storage. The `dlpack_out` variable gets overwritten on the next iteration, but neither the dlpack capsule nor the PyTorch view that aliases it is ever explicitly released. The `.movedim()` / `.unsqueeze()` calls produce *views*, not new allocations, so the original per-frame storage stays pinned for the life of the loop. The CUDA caching allocator holds onto all of it, VRAM climbs monotonically, and eventually the node OOMs or deadlocks on a blocking CUDA call. ## The fix Two changes to the per-frame loop, intended to release DLPack-backed storage each iteration without forcing a per-frame CUDA cache flush: 1. **`.clone()` the converted tensor** — breaks the aliasing so the output owns its own storage and the dlpack-backed tensor becomes immediately releasable. 2. **`del dlpack_out, result, out_frame`** — drops the references this iteration instead of waiting for the frame to go out of scope. ```python for j in range(batch_cuda.shape[0]): input_frame = batch_cuda[j] result = sr.run(input_frame) dlpack_out = result.image out_frame = torch.from_dlpack(dlpack_out).movedim(0, -1).unsqueeze(0).clone() out_tensor[i + j: i + j + 1] = out_frame del dlpack_out, result, out_frame ``` ## Verification Tested locally with 100+ frame runs where the original code OOM'd around ~60 frames. With this patch VRAM stays flat across the whole run — `nvidia-smi` holds a steady plateau instead of the old monotonic climb. Output is pixel-identical (`.clone()` is a copy, not a transform). Follow-up `71a5f69` removes the per-frame `empty_cache()` in response to the review: the allocator can reuse released blocks without a forced flush each frame. The 100+ frame observation above was made on the prior version *with* `empty_cache()`; GPU memory behavior after its removal has not been re-measured, so that specific result is not claimed for the current head. ## Scope - One file changed: `__init__.py` - No API or behaviour change for existing graphs — same inputs, same outputs, just without the memory leak. Closes #35. Related: #16. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
bsakaguchi 2026-10-26 22:28:19 +00:00
Sign in to join this conversation.
No labels
pull-request
No milestone
No project
No assignees
1 participant
Notifications
Due date
The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference
Comfy-Org/ComfyUI_Nvidia_VFX_Nodes#40
No description provided.