DLPack and CUDA Stream Limitations

The Python wrappers accept CUDA arrays through DLPack and return DLPack output capsules. The guidance on raw capsules, input stream ordering, consumer stream handoff, and output-buffer lifetime applies to VideoSuperRes, VideoFrameGeneration, and TrueHDR.

The RGB8 output-conversion limitation applies only to effects configured for RGB8 output. TrueHDR returns RGB10A2 output and does not use this conversion path.

Raw Output Capsule

Each inference method returns an output dataclass whose image field is a raw DLPack capsule. PyTorch accepts this format directly with torch.from_dlpack(result.image). CuPy also accepts it through its legacy capsule compatibility path.

JAX 0.7.2 and later require an object that implements __dlpack__() and __dlpack_device__() instead of a raw capsule. An adapter wrapper is required until the wheel exposes the full Python DLPack protocol. Additionally, copying during DLPack import requires buffer semantics that the JAX PJRT C API backend did not support. As a workaround, wrap the raw capsule, import the DLPack buffer without copying, then create and synchronize a copy after the DLPack import. Compatibility with JAX versions earlier than 0.7.2 and other JAX backend configurations was not evaluated.

After completing the SDK producer work as described in CUDA Stream Ordering, use the following Python DLPack protocol adapter. result can be returned by VideoSuperRes.run(), VideoFrameGeneration.run(), VideoFrameGeneration.run_at_timestep(), or TrueHDR.run(). Set device_id to the CUDA device used by the effect.

import jax
import jax.numpy as jnp

class JaxDLPackAdapter:
    def __init__(self, capsule, device_id=0):
        self._capsule = capsule
        self._device_id = device_id

    def __dlpack__(self, stream=None, **kwargs):
        return self._capsule

    def __dlpack_device__(self):
        return (2, self._device_id)  # kDLCUDA

adapter = JaxDLPackAdapter(result.image, device_id=0)
output_view = jax.dlpack.from_dlpack(adapter, copy=False)
output = jnp.array(output_view, copy=True)
output.block_until_ready()

CUDA Stream Ordering

stream_ptr selects the CUDA stream used by the SDK effect and input transfer, but it is not a complete producer-to-consumer synchronization contract today.

The bindings acquire each input by calling the array producer’s __dlpack__() without passing the SDK stream.

When input writes and run(..., stream_ptr=...) use the same stream, CUDA orders the writes before the SDK input transfer, so no additional input synchronization is required. When they use different streams, synchronize the producer stream before calling run().

For output, ordering includes inference, RGB8 output conversion when applicable, and the framework import or copy. If all these operations use the same stream, CUDA orders them and device-wide synchronization is not required. If the consumer imports or copies the output on a different stream, synchronize the SDK producer stream before importing or copying the output.

The high-level wrappers currently call low-level get_output() without forwarding stream_ptr. Consequently, RGB8 inference may use the requested SDK stream while output conversion uses stream 0. With a nonzero stream_ptr, do not assume that inference and conversion share a stream. Synchronizing only the configured SDK stream may not cover the conversion work. In particular, non_blocking=True is unsafe because conversion may begin before inference finishes.

For the current high-level RGB8 path, use blocking mode. Device-wide synchronization before import is the conservative fallback because it waits for both the SDK stream and the RGB8 conversion stream. The following example assumes that sr is a loaded VideoSuperRes instance and frame is a compatible CUDA input:

stream = torch.cuda.current_stream()
result = sr.run(
    frame,
    stream_ptr=stream.cuda_stream,
    non_blocking=False,
)

# Wait for SDK work and RGB8 output conversion before import.
torch.cuda.synchronize()
output = torch.from_dlpack(result.image).clone()

# Wait for the clone before reusing or closing the effect or accessing CPU data.
stream.synchronize()

A synchronization only after clone() cannot order a clone that has already raced SDK producer work. If inference, conversion, and the copy are all ordered on the same stream, synchronizing that stream is sufficient and no pre-import device-wide synchronization is needed. Device-wide synchronization may reduce throughput.

Output Buffer Lifetime

The DLPack output references SDK-owned memory. It remains valid only while the effect and its current output buffer remain alive. The buffer may be reused or released by another run(), an output-size or encoding change, or close().

After SDK producer work has completed, import the capsule and immediately clone/copy the result. If the copy and the next run() use the same stream, CUDA orders the copy before reuse of the SDK output buffer. If they use different streams, synchronize the copy stream before reusing the buffer.

Always wait for the copy before close() or an output-size or encoding change because those host operations may release or replace the SDK buffer without CUDA stream ordering. After the copy has completed, the copied output is managed by the framework and can outlive the effect.

Do not retain a zero-copy tensor imported from result.image across another run(), an output-size or encoding change, or close().

For CuPy with the current high-level RGB8 path, call cupy.cuda.runtime.deviceSynchronize() before importing the output when the inference, conversion, and copy streams may differ. After the copy, synchronize the CuPy copy stream before close() or an output-size or encoding change. If the copy and next run() use the same stream, CUDA orders the copy before buffer reuse.