Add NVIDIA DLSS Ray Reconstruction denoiser

[fknfilewalker]

The implementation is following Mitsuba's OptiX denoiser as well as Blender's DLSS integration (PR #153077).
Proof of Concept, prompted collaboratively with @WeiPhil via Claude.
Feel free to close this if not needed.

Using DLSS with a pre-built Mitsuba (Vibed)

Two independent gates: the first is decided when Mitsuba is compiled, the second on the machine that runs it.

Gate 1: compile

Enable via MI_ENABLE_DLSS, with the entire class body sitting behind #if defined(MI_ENABLE_CUDA) && defined(MI_ENABLE_DLSS). Without it the bindings fall back to a stub whose is_available() is a hardcoded return false. For the build only the headers are needed.

Gate 2: the runtime library

A wheel built with the flag on still couldn't bundle nvngx_dlssd.dll; the same licensing reason Blender declined, and why CMakeLists.txt keeps it out of install. Therefore it needs to be manually provided by the user:

1. pip install mitsuba
2. Download NVIDIA/DLSS and take
   lib/Windows_x86_64/rel/nvngx_dlssd.dll (or
   lib/Linux_x86_64/rel/libnvidia-ngx-dlssd.so.*)
3. Either drop it into site-packages/mitsuba/, or point MI_DLSS_LIBRARY_PATH
   at the directory holding it
4. RTX GPU, driver 590+, and mi.set_variant('cuda_ad_rgb')
5. Confirm with mi.DLSSDenoiser.is_available()

Run

Show Real-Time Example Code (requires: 'imgui-bundle' and 'pyopengl') (DLSS part is vibed)

import drjit as dr
from drjit.auto import TensorXf, Quaternion4f, Array3f, Float
import mitsuba as mi
from imgui_bundle import imgui, immapp
mi.set_variant('cuda_ad_rgb' if dr.has_backend(dr.JitBackend.CUDA) else 'llvm_ad_rgb')

# DLSS Ray Reconstruction is only present in CUDA builds that were configured
# with `MI_ENABLE_DLSS`, and only usable on a supported driver/GPU combination.
# `is_available()` checks all of that, and returns false under a non-CUDA
# variant, so it subsumes a backend test.
DLSS_AVAILABLE = mi.DLSSDenoiser.is_available()

class GlTexture:
    def __init__(self, width, height):
        import OpenGL.GL as gl
        self.width, self.height, self.id = width, height, gl.glGenTextures(1)
        gl.glBindTexture(gl.GL_TEXTURE_2D, self.id)
        gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MIN_FILTER, gl.GL_LINEAR)
        gl.glTexParameteri(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MAG_FILTER, gl.GL_LINEAR)
        gl.glTexImage2D(
            gl.GL_TEXTURE_2D, 0, gl.GL_RGBA32F, width, height, 0,
            gl.GL_RGBA, gl.GL_FLOAT, None
        )
        try:
            self.interop = dr.cuda.GLInterop.from_texture(self.id)
        except Exception as e:
            print("Warning: Could not create CUDA-OpenGL interop:", e)
            self.interop = None

    def upload(self, img):
        if dr.backend_v(img) == dr.JitBackend.CUDA:
            self.interop.map().upload(img).unmap()
        elif dr.backend_v(img) == dr.JitBackend.LLVM:
            import OpenGL.GL as gl
            gl.glBindTexture(gl.GL_TEXTURE_2D, self.id)
            gl.glTexSubImage2D(
                gl.GL_TEXTURE_2D, 0, 0, 0, self.width, self.height,
                gl.GL_RGBA, gl.GL_FLOAT, img.to_numpy()
            )

class Camera:
    def __init__(self):
        self.q, self.c, self.d = Quaternion4f(0, 0, 0, 1), Array3f(0, 0, 0), Float(14.0)
        self.resolution, self.scale = dr.scalar.Array2u(256), 1.0
        self.gl_texture: GlTexture | None = None
        self.start_pos = None
        self.sensor = mi.load_dict({
            'type': 'perspective',
            'fov': 39.3077,
            'to_world': mi.ScalarTransform4f(),
            'sampler': {
                'type': 'independent',
                'sample_count': 1
            },
            'film': {
                'type': 'hdrfilm',
                'width': self.resolution[0],
                'height': self.resolution[1],
                # Both the accumulation loop and the denoisers want unfiltered
                # samples: a filter spanning several pixels is indistinguishable
                # from scene detail to a denoiser.
                'rfilter': {
                    'type': 'box',
                },
                'pixel_format': 'rgb',
            },
        })
        self.params = mi.traverse(self.sensor)
        self.update()
        # Camera of the frame that was rendered last, used to derive the
        # screen-space motion vectors that DLSS needs
        self.prev_to_world = self.to_world

    def map_to_sphere(self, px, py, size):
        x = -(1.0 - (px / size[0]) * 2.0)
        y = (1.0 - (py / size[1]) * 2.0)
        length2 = x*x + y*y
        if length2 > 1.0:
            return Array3f(x, y, 0.0) / dr.sqrt(length2)
        else:
            return Array3f(x, y, dr.sqrt(max(0.0, 1.0 - length2)))

    def rotate(self, a, b, size):
        a = self.map_to_sphere(*a, size)
        b = self.map_to_sphere(*b, size)
        perp = dr.cross(b, a)
        if dr.norm(perp) > 1e-5:
            q = Quaternion4f(perp.x, perp.y, perp.z, dr.dot(a, b))
        else:
            q = Quaternion4f(0, 0, 0, 1)
        return dr.normalize(self.q * q)

    def update(self, q=Quaternion4f(0, 0, 0, 1)):
        self.to_world = mi.Transform4f().look_at(
                origin=dr.quat_apply(q, Array3f(0, 0, self.d)) + self.c,
                target=self.c,
                up=dr.quat_apply(q, Array3f(0, 1, 0))
            )
        self.params['to_world'] = self.to_world
        self.params['film.size'] = self.resolution
        self.film = dr.zeros(TensorXf, shape=(self.resolution[1], self.resolution[0], 4))
        self.params.update()

    def set_resolution(self, resolution):
        """
        Resize the film, returning whether anything changed.

        This is the only place that may touch `self.resolution`: the film size
        of the sensor and the accumulation buffer have to be updated together,
        and going through `update()` guarantees that.
        """
        resolution = dr.scalar.Array2u(resolution)
        if dr.all(resolution == self.resolution):
            return False
        self.resolution = resolution
        self.update(self.q)
        return True

    def handle_gl_texture(self, resolution):
        ' Only call if gl texture is needed and with a valid opengl context '
        if not self.gl_texture or self.gl_texture.width != resolution[0] or self.gl_texture.height != resolution[1]:
            self.gl_texture = GlTexture(resolution[0], resolution[1])

    def process_imgui_inputs(self, pos, size):
        """
        Handle the mouse inputs of a viewport of `size` pixels at `pos`.

        The arcball is parameterized by the extent that the image covers *on
        screen*, which is not the render resolution: `scale` magnifies it, and
        DLSS upscaling makes the rendered frame smaller than the one shown.
        """
        from imgui_bundle import imgui
        io = imgui.get_io()
        if imgui.is_mouse_clicked(0):
            self.start_pos = io.mouse_pos - pos
        if imgui.is_mouse_released(0) and self.start_pos:
            self.q = self.rotate(self.start_pos, io.mouse_pos - pos, size)
            self.update(self.q)
            self.start_pos = None
        if imgui.is_mouse_dragging(0) and self.start_pos:
            self.update(self.rotate(self.start_pos, io.mouse_pos - pos, size))
        elif imgui.is_mouse_dragging(1):
            # Pan by the world-space distance that the cursor travelled on the
            # plane through the target, so that the scene keeps up with the
            # cursor at any viewport extent
            fov = dr.deg2rad(self.params['x_fov'][0])
            per_pixel = 2.0 * dr.tan(0.5 * fov) * self.d / size[0]
            self.c += dr.quat_apply(
                self.q, Array3f(-io.mouse_delta.x, io.mouse_delta.y, 0.0) * per_pixel)
            self.update(self.q)
        if io.mouse_wheel != 0:
            self.d = self.d - io.mouse_wheel * 0.3
            if self.d < 0.1:
                self.d = 0.1
            self.update(self.q)


class Denoiser:
    """
    Drives `mi.DLSSDenoiser` from the frames produced by this viewport.

    DLSS Ray Reconstruction consumes one independently rendered frame at a
    time together with a G-buffer (diffuse albedo, shading normals, linear
    camera-space depth) and screen-space motion vectors, and accumulates
    detail across frames. The G-buffer is obtained by wrapping the scene
    integrator in an `aov` integrator, the motion vectors are derived by
    reprojecting the shading points of the current frame into the camera of
    the previous one.
    """

    AOVS = 'albedo:albedo,sh_normal:sh_normal,depth:depth'
    QUALITIES = ['high', 'balanced', 'fast']
    UPSCALE_FACTORS = [1, 2, 3]

    def __init__(self, scene):
        self.integrator = mi.load_dict({
            'type': 'aov',
            'aovs': Denoiser.AOVS,
            'img': scene.integrator(),
        })
        self.quality = 'high'
        self.upscale = 1
        self.active = False
        self.invalidate()

    def output_resolution(self, resolution):
        return dr.scalar.Array2u(resolution * self.upscale)

    def invalidate(self):
        """Force the denoiser (and its temporal history) to be rebuilt"""
        self.denoiser = None
        self.rays = None

    def prepare(self, camera):
        out_res = self.output_resolution(camera.resolution)
        if self.denoiser is None:
            self.denoiser = mi.DLSSDenoiser(mi.ScalarVector2u(camera.resolution),
                                            mi.ScalarVector2u(out_res),
                                            self.quality)
            self.rays = self.pixel_rays(camera)
        return out_res

    def pixel_rays(self, camera):
        """Per-pixel camera-space near plane point and direction (cached)"""
        width, height = int(camera.resolution[0]), int(camera.resolution[1])
        index = dr.arange(mi.UInt32, width * height)
        px = mi.Float(index % width) + 0.5
        py = mi.Float(index // width) + 0.5

        projection = camera.sensor.projection_transform()
        near_p = mi.Point3f(projection.inverse() @
                            mi.Point3f(px / width, py / height, 0.0))

        return projection, near_p, dr.normalize(mi.Vector3f(near_p)), px, py

    def __call__(self, camera, img):
        """Denoise one rendered frame, given as the `aov` integrator output"""
        width, height = int(camera.resolution[0]), int(camera.resolution[1])
        projection, near_p, direction, px, py = self.rays

        # Channel layout of the `aov` integrator: the nested image comes first,
        # followed by the AOVs in the order they were requested
        noisy = img[..., 0:3]
        albedo = img[..., 3:6]
        normals = img[..., 6:9]

        # The `depth` AOV is the distance travelled along the ray, and rays
        # start on the near plane rather than at the camera origin. Rays that
        # left the scene are pushed far away so that they still receive
        # plausible motion vectors.
        t = img[..., 9:10].array
        t = dr.select(t > 0.0, t, 1e7)

        p_camera = near_p + direction * t
        p_world = camera.to_world @ mi.Point3f(p_camera)

        # Reproject into the previous frame to obtain the motion vectors. Note
        # that Mitsuba's samplers always jitter randomly inside a pixel, so the
        # exact subpixel offset of a frame is not known to us and no jitter is
        # reported to DLSS. The random offsets still provide the sample
        # diversity that DLSS relies on.
        p_prev = mi.Point3f(camera.prev_to_world.inverse() @ p_world)
        uv_prev = projection @ p_prev
        visible = p_prev.z > 0.0

        flow = dr.concat((
            mi.TensorXf(dr.select(visible, uv_prev.x * width - px, 0.0),
                        shape=(height, width, 1)),
            mi.TensorXf(dr.select(visible, uv_prev.y * height - py, 0.0),
                        shape=(height, width, 1))), axis=2)

        # DLSS expects a linear camera-space depth
        depth = mi.TensorXf(p_camera.z, shape=(height, width, 1))

        return self.denoiser(noisy, albedo, normals, depth, flow=flow)


class State:
    def __init__(self, scene_file):
        self.scene = mi.load_file(scene_file)
        self.frame = dr.opaque(dr.auto.UInt, 0)
        self.accumulate = True
        self.camera = Camera()
        self.resolution_slider = int(self.camera.resolution[0])
        self.denoiser = Denoiser(self.scene) if DLSS_AVAILABLE else None

    def process_inputs(self, pos, size):
        if imgui.table_get_hovered_column() == 0:
            self.camera.process_imgui_inputs(pos, size)

# `dr.freeze` keys its recordings on the structure of the Dr.Jit variables it
# can discover, and the film resolution is not among them: it is a scalar C++
# member, and a change in the *width* of the film storage is explicitly allowed
# to replay an existing recording. Without the film size as additional state, a
# resolution change would therefore keep replaying a recording made for another
# resolution and hand back an image of the previous size.
#
# Each film size (and each integrator) thus gets its own recording, which means
# that re-tracing is by design here: the warning is disabled, and an LRU cache
# bounds the memory that the recordings occupy instead.
@dr.freeze(limit=8, warn_after=None,
           state_fn=lambda scene, seed=0, sensor=0, integrator=None:
               tuple(sensor.film().crop_size()))
def render(scene, seed=0, sensor=0, integrator=None):
    return mi.render(scene, spp=1, seed=seed, sensor=sensor,
                     integrator=integrator)

def show_image():
    camera = state.camera
    denoiser = state.denoiser

    if denoiser is not None and denoiser.active:
        denoiser.prepare(camera)
        img = render(state.scene, seed=state.frame, sensor=camera.sensor,
                     integrator=denoiser.integrator)
        # DLSS accumulates across frames on its own, so every frame is shown
        # as it comes out of the denoiser
        out = add_alpha(dr.linear_to_srgb(denoiser(camera, img)))
    else:
        img = add_alpha(render(state.scene, seed=state.frame, sensor=camera.sensor))
        # The accumulated samples are only meaningful for the film size they
        # were rendered at, so a resolution change restarts the accumulation
        if state.accumulate and camera.film.shape == img.shape:
            camera.film = camera.film + img
        else:
            camera.film = img
        # use alpha channel to store the sample count
        out = dr.linear_to_srgb(camera.film / camera.film[0][0][3])

    # Size the texture from the image that is about to be uploaded rather than
    # from the resolution that was requested, so that the two cannot disagree
    resolution = (out.shape[1], out.shape[0])
    camera.handle_gl_texture(resolution)
    camera.gl_texture.upload(out)

    camera.prev_to_world = camera.to_world
    state.frame += 1

    w = float(resolution[0]) * camera.scale
    h = float(resolution[1]) * camera.scale
    x = imgui.get_column_width()//2 - w//2
    y = imgui.get_window_viewport().size.y//2 - h//2
    imgui.set_cursor_pos_x(x)
    imgui.set_cursor_pos_y(y)
    imgui.image(imgui.ImTextureRef(camera.gl_texture.id), (w, h))
    state.process_inputs((x, y), (w, h))

def add_alpha(img):
    if img.shape[2] == 4:
        return TensorXf(img)
    alpha = dr.ones(TensorXf, shape=(img.shape[0], img.shape[1], 1))
    return dr.concat((TensorXf(img), alpha), axis=2)

def denoiser_settings():
    if state.denoiser is None:
        imgui.text_disabled("DLSS Ray Reconstruction is unavailable")
        imgui.text_disabled("(needs MI_ENABLE_DLSS, an RTX GPU, driver 590+)")
        return

    denoiser = state.denoiser
    changed, denoiser.active = imgui.checkbox("DLSS Ray Reconstruction",
                                              denoiser.active)
    if changed:
        denoiser.invalidate()
        state.camera.update(state.camera.q)

    if not denoiser.active:
        return

    changed, index = imgui.combo("Quality", Denoiser.QUALITIES.index(denoiser.quality),
                                 Denoiser.QUALITIES)
    if changed:
        denoiser.quality = Denoiser.QUALITIES[index]
        denoiser.invalidate()

    labels = [f"{f}x" for f in Denoiser.UPSCALE_FACTORS]
    changed, index = imgui.combo("Upscaling",
                                 Denoiser.UPSCALE_FACTORS.index(denoiser.upscale),
                                 labels)
    if changed:
        denoiser.upscale = Denoiser.UPSCALE_FACTORS[index]
        denoiser.invalidate()

def gui():
    imgui.set_next_window_size(imgui.get_io().display_size)
    imgui.set_next_window_pos((0, 0))
    if imgui.begin("##FullscreenWindow", None, imgui.WindowFlags_.no_resize | imgui.WindowFlags_.no_decoration | imgui.WindowFlags_.no_background | imgui.WindowFlags_.no_focus_on_appearing):
        if imgui.begin_table("Table", 2, imgui.TableFlags_.resizable | imgui.TableFlags_.sizing_stretch_prop, (-1, -1)):
            imgui.table_setup_column("Rendering", imgui.TableColumnFlags_.width_stretch)
            imgui.table_setup_column("Settings", imgui.TableColumnFlags_.width_fixed, 400)
            imgui.table_next_row()
            imgui.table_set_column_index(0)
            show_image()
            imgui.table_set_column_index(1)
            if imgui.begin_tab_bar("Settings"):
                if imgui.begin_tab_item("General")[0]:
                    if imgui.collapsing_header("Film", imgui.TreeNodeFlags_.default_open):
                        _, state.camera.scale = imgui.slider_float("Scale", state.camera.scale, 1.0, 10.0)
                        _, state.resolution_slider = imgui.slider_int("Resolution", state.resolution_slider, 64, 2048)
                        # Applying every intermediate value of a drag would
                        # re-trace `render()` dozens of times, so the film is
                        # only resized once the slider is let go of
                        if imgui.is_item_deactivated_after_edit():
                            changed = state.camera.set_resolution(state.resolution_slider)
                            if changed and state.denoiser is not None:
                                state.denoiser.invalidate()
                    denoising = state.denoiser is not None and state.denoiser.active
                    imgui.begin_disabled(denoising)
                    state.accumulate = imgui.checkbox("Accumulate", state.accumulate)[1]
                    imgui.end_disabled()
                    if imgui.collapsing_header("Denoiser", imgui.TreeNodeFlags_.default_open):
                        denoiser_settings()
                    imgui.end_tab_item()
                if imgui.begin_tab_item("Other")[0]:
                    imgui.text("More settings can be added here.")
                    imgui.end_tab_item()
                imgui.end_tab_bar()
            imgui.set_cursor_pos_y(imgui.get_cursor_pos_y() + imgui.get_content_region_avail().y - imgui.get_text_line_height_with_spacing())
            imgui.separator()
            imgui.text(f"{imgui.get_io().framerate :.2f} FPS | { 1000.0 / imgui.get_io().framerate :.1f} ms")
            imgui.end_table()
        imgui.end()

if __name__ == "__main__":
    state = State("scenes/cbox.xml")
    immapp.run(gui_function=gui, window_title="Real-Time Mitsuba", window_size=(1280, 720), fps_idle=0)

[wjakob]

I think it might be problematic for Dr.Jit to ship DLSS binaries from a legal point of view.

[wjakob]

My plan was to rely on https://github.com/mitsuba-renderer/mitsuba-oidn for denoising in Mitsuba, which is distributed as a separate package and supports CPU, Metal, and CUDA in a consistent way. It is also an open source project, which makes things legally unproblematic. What benefits would DLSS give?

[fknfilewalker]

I think it might be problematic for Dr.Jit to ship DLSS binaries from a legal point of view.

True, and this is why the user needs to set the environment variable MI_DLSS_LIBRARY_PATH that points to the dynamic lib.

OIDN

Many good points and it makes a lot of sense.

What benefits would DLSS give?

It is a temporal denoiser and upscaler. It basically allows mitsuba to be a real-time viewport (see video using code from above). It could also speed up rendering out long sequences. On the Metal side, metalfx would be similar. Now that I saw mitsuba-oidn, I kinda like the idea of making it a separate project.

Recording.2026-09-08.202151.mp4