From 012dc673bc17f4f6dd8599ed2245d2bcc0b1489f Mon Sep 17 00:00:00 2001 From: comfyanonymous <121283862+comfyanonymous@users.noreply.github.com> Date: Tue, 17 Mar 2026 20:12:55 -0400 Subject: [PATCH] Lower memory usage. --- __init__.py | 21 +++++---------------- 1 file changed, 5 insertions(+), 16 deletions(-) diff --git a/__init__.py b/__init__.py index 3023ebf..9fb9495 100644 --- a/__init__.py +++ b/__init__.py @@ -50,7 +50,7 @@ class RTXVideoSuperResolution(io.ComfyNode): @classmethod def execute(cls, images: torch.Tensor, resize_type: UpscaleTypedDict, quality: str) -> io.NodeOutput: - _, h, w, _ = images.shape + b, h, w, c = images.shape selected_type = resize_type["resize_type"] if selected_type == UpscaleType.SCALE_BY: @@ -79,34 +79,23 @@ class RTXVideoSuperResolution(io.ComfyNode): } selected_quality = quality_mapping.get(quality, nvvfx.effects.QualityLevel.HIGH) - upscaled_batches = [] - with nvvfx.VideoSuperRes(selected_quality) as sr: sr.output_width = output_width sr.output_height = output_height sr.load() + out_tensor = torch.empty((images.shape[0], output_height, output_width, c), device=images.device, dtype=images.dtype) for i in range(0, images.shape[0], batch_size): batch = images[i:i + batch_size] - batch_cuda = batch.cuda().permute(0, 3, 1, 2).contiguous() - - batch_outputs = [] + batch_cuda = batch.cuda().permute(0, 3, 1, 2).float().contiguous() for j in range(batch_cuda.shape[0]): input_frame = batch_cuda[j] dlpack_out = sr.run(input_frame).image - output = torch.from_dlpack(dlpack_out).clone() - batch_outputs.append(output) + out_tensor[i + j: i + j + 1] = torch.from_dlpack(dlpack_out).movedim(0, -1).unsqueeze(0) - batch_out_tensor = torch.stack(batch_outputs, dim=0) - - batch_out_tensor = batch_out_tensor.permute(0, 2, 3, 1).cpu() - upscaled_batches.append(batch_out_tensor) - - final_images = torch.cat(upscaled_batches, dim=0) - - return io.NodeOutput(final_images) + return io.NodeOutput(out_tensor) class NVVFXVideoExtension(ComfyExtension):