Blame for native-eiffel/native_render.py

"""Render an exported three.js scene straight into a CUDA tensor with nvdiffrast.

Ports the three.js r161 shading the game uses (physically based lights, i.e.
``useLegacyLights = false``): MeshBasic / LineBasic unlit, MeshLambert and
MeshStandard (GGX direct specular, no environment map so no indirect
specular), hemisphere + directional lights, emissive, ACES filmic tone mapping
and the sRGB transfer. Front-face culling, double-sided normals and flat
shading follow WebGLRenderer. Depth-writing surfaces are resolved first;
surfaces that do not write depth (flash, tracer, impact) are composited over
them, with their opacity when transparent. Lines are one pixel wide, as WebGL
draws them, built as screen-aligned quads.

A frame never synchronizes with the host: culled triangles become degenerate
instead of being filtered out, the frame's scene data arrives in one copy, and
the per-pixel shading and compositing are one compiled function.

The output is an (H, W, 3) uint8 CUDA tensor, top row first — the same layout
the capture path hands to the model, without ever leaving the GPU.
"""
import math
import os

import numpy as np
import torch
import torch.nn.functional as F
import nvdiffrast.torch as dr

_BASIC, _LAMBERT, _STANDARD = 0, 1, 2
_TYPES = {'MeshBasicMaterial': _BASIC, 'LineBasicMaterial': _BASIC,
          'MeshLambertMaterial': _LAMBERT, 'MeshStandardMaterial': _STANDARD}
_FRONT, _BACK, _DOUBLE = 0, 1, 2
_ACES_TONE_MAPPING = 4

_ACES_IN = [[0.59719, 0.35458, 0.04823],
            [0.07600, 0.90834, 0.01566],
            [0.02840, 0.13383, 0.83777]]
_ACES_OUT = [[1.60475, -0.53108, -0.07367],
             [-0.10208, 1.10813, -0.00605],
             [-0.00327, -0.07276, 1.07602]]

# Per-triangle table columns: static material, per-frame material, geometry.
_KIND, _FLAT, _DOUBLE_SIDED, _TONE, _ROUGH, _METAL, _TRANSPARENT = range(7)
_COLOR, _OPACITY, _EMISSIVE, _EMISSIVE_INTENSITY = slice(7, 10), 10, slice(11, 14), 14
_FACE, _FRONT_FACING = slice(15, 18), 18


def _det3(m):
    """Determinant of (..., 3, 3) as a triple product; torch.linalg.det would sync with the host."""
    return (m[..., 0, :] * torch.cross(m[..., 1, :], m[..., 2, :], dim=-1)).sum(-1)


def _hex_to_srgb(value):
    value = value.lstrip('#')
    return [int(value[i:i + 2], 16) / 255 for i in (0, 2, 4)]


def _dot(a, b):
    return (a * b).sum(-1, keepdim=True)


def _ggx(l, v, n, f0, roughness):
    alpha = roughness * roughness
    h = F.normalize(l + v, dim=-1)
    dot_nl = _dot(n, l).clamp(0, 1)
    dot_nv = _dot(n, v).clamp(0, 1)
    dot_nh = _dot(n, h).clamp(0, 1)
    dot_vh = _dot(v, h).clamp(0, 1)
    fresnel = torch.exp2((-5.55473 * dot_vh - 6.98316) * dot_vh)
    f = f0 * (1 - fresnel) + fresnel
    a2 = alpha * alpha
    gv = dot_nl * torch.sqrt(a2 + (1 - a2) * dot_nv * dot_nv)
    gl = dot_nv * torch.sqrt(a2 + (1 - a2) * dot_nl * dot_nl)
    vis = 0.5 / (gv + gl).clamp(min=1e-6)
    d = a2 / (math.pi * (dot_nh * dot_nh * (a2 - 1) + 1) ** 2)
    return f * vis * d


def _to_output(linear, tone_mapped, exposure, aces_in, aces_out):
    x = linear * (exposure / 0.6)
    x = x @ aces_in.T
    a = x * (x + 0.0245786) - 0.000090537
    b = x * (0.983729 * x + 0.4329510) + 0.238081
    mapped = ((a / b) @ aces_out.T).clamp(0, 1)
    linear = torch.where(tone_mapped, mapped, linear).clamp(min=0)
    return torch.where(linear <= 0.0031308, linear * 12.92, linear.pow(0.41666) * 1.055 - 0.055).clamp(0, 1)


def _shade(rast, interp, table, eye, light_dirs, light_colors, hemi_sky, hemi_ground, hemi_up,
           exposure, tone_enabled, aces_in, aces_out):
    covered = rast[..., 3] > 0
    tid = (rast[..., 3].long() - 1).clamp(min=0)
    t = table[tid]
    kind = t[..., _KIND]
    flat = t[..., _FLAT:_FLAT + 1] > 0.5
    position, normal = interp[..., :3], interp[..., 3:6]
    normal = F.normalize(torch.where(flat, t[..., _FACE], normal), dim=-1)
    view = F.normalize(eye - position, dim=-1)
    back = (t[..., _DOUBLE_SIDED:_DOUBLE_SIDED + 1] > 0.5) & (t[..., _FRONT_FACING:_FRONT_FACING + 1] < 0.5) & ~flat
    normal = torch.where(back | (flat & (_dot(normal, view) < 0)), -normal, normal)

    color = t[..., _COLOR]
    metal = t[..., _METAL:_METAL + 1]
    standard = (kind == _STANDARD)[..., None]
    diffuse = torch.where(standard, color * (1 - metal), color)
    f0 = 0.04 + (color - 0.04) * metal

    # roughness = min(max(r, 0.0525) + geometryRoughness, 1); the normal
    # derivative is taken within each triangle, as a fragment quad would.
    def derivative(dim):
        forward = normal.roll(-1, dim) - normal
        backward = normal - normal.roll(1, dim)
        same_f = (tid.roll(-1, dim) == tid)[..., None]
        same_b = (tid.roll(1, dim) == tid)[..., None]
        return torch.where(same_f, forward, torch.where(same_b, backward, torch.zeros_like(forward))).abs()
    geometry = torch.maximum(derivative(0), derivative(1)).amax(-1, keepdim=True)
    roughness = (t[..., _ROUGH:_ROUGH + 1].clamp(min=0.0525) + geometry).clamp(max=1.0)

    irradiance = torch.zeros_like(position)
    for k in range(hemi_sky.shape[0]):
        w = 0.5 * _dot(normal, hemi_up[k]) + 0.5
        irradiance = irradiance + hemi_ground[k] + (hemi_sky[k] - hemi_ground[k]) * w
    specular = torch.zeros_like(position)
    for k in range(light_colors.shape[0]):
        l = F.normalize(light_dirs[k], dim=-1)
        direct = _dot(normal, l).clamp(0, 1) * light_colors[k]
        irradiance = irradiance + direct
        specular = specular + direct * _ggx(l, view, normal, f0, roughness)
    emissive = t[..., _EMISSIVE] * t[..., _EMISSIVE_INTENSITY:_EMISSIVE_INTENSITY + 1]
    lit = diffuse / math.pi * irradiance + emissive + torch.where(standard, specular, torch.zeros_like(specular))
    out = torch.where((kind == _BASIC)[..., None], color, lit)
    tone = (t[..., _TONE:_TONE + 1] > 0.5) & tone_enabled
    out = _to_output(out, tone, exposure, aces_in, aces_out)
    alpha = torch.where(t[..., _TRANSPARENT] > 0.5, t[..., _OPACITY], torch.ones_like(kind))
    return out, covered, rast[..., 2], alpha


def _compose(base_rast, base_interp, over_rast, over_interp, table, clear, eye, light_dirs, light_colors,
             hemi_sky, hemi_ground, hemi_up, exposure, tone_enabled, aces_in, aces_out):
    args = (table, eye, light_dirs, light_colors, hemi_sky, hemi_ground, hemi_up, exposure, tone_enabled,
            aces_in, aces_out)
    color, covered, z, _ = _shade(base_rast, base_interp, *args)
    color = torch.where(covered[..., None], color, clear)
    depth = torch.where(covered, z, torch.full_like(z, float('inf')))
    over, over_covered, over_z, alpha = _shade(over_rast, over_interp, *args)
    alpha = torch.where(over_covered & (over_z < depth), alpha, torch.zeros_like(alpha))[..., None]
    color = color * (1 - alpha) + over * alpha
    # WebGL rasterizes bottom row first; images are stored top row first.
    return color.flip(0).mul(255).round().clamp(0, 255).to(torch.uint8)


class NativeRenderer:
    def __init__(self, scene, arrays, device='cuda', compile=True, graphs=None):
        dev = self.device = torch.device(device)
        self.width, self.height = scene['width'], scene['height']
        self.ctx = dr.RasterizeCudaContext(device=dev)
        self.exposure = torch.tensor(float(scene['exposure']), device=dev)
        self.tone_enabled = torch.tensor(scene['toneMapping'] == _ACES_TONE_MAPPING, device=dev)
        self.clear = torch.tensor(_hex_to_srgb(scene['clearColor']), device=dev)
        self.aces_in = torch.tensor(_ACES_IN, device=dev)
        self.aces_out = torch.tensor(_ACES_OUT, device=dev)
        self._lights(scene['lights'])

        materials = scene['materials']
        side = np.array([m['side'] for m in materials])
        overlay = np.array([bool(m['transparent']) or not m['depthWrite'] for m in materials])
        static = np.array([[_TYPES[m['type']], m['flatShading'], m['side'] == _DOUBLE, bool(m['toneMapped']),
                            m['roughness'] if m['roughness'] is not None else 1.0,
                            m['metalness'] if m['metalness'] is not None else 0.0,
                            bool(m['transparent'])] for m in materials], np.float32)

        # One vertex copy per mesh draw, so each draw carries its own transform.
        positions, normals, vertex_draw, triangles, tri_draw = [], [], [], [], []
        base = 0
        self.lines = []   # (draw, geometry, material, segments)
        for d, draw in enumerate(scene['draws']):
            g = scene['geometries'][draw['geometry']]
            index = arrays[f'g{draw["geometry"]}_index']
            count = draw['count'] if draw['count'] >= 0 else len(index) - draw['start']
            index = index[draw['start']:draw['start'] + count]
            if draw['kind'] == 'line':
                if not g['dynamic']:
                    raise NotImplementedError('Static line geometry')
                self.lines.append((d, draw['geometry'], draw['material'], g['vertices'] - 1))
                continue
            used = np.unique(index)
            remap = np.full(g['vertices'], -1, np.int64)
            remap[used] = np.arange(len(used)) + base
            p = arrays[f'g{draw["geometry"]}_position'].reshape(-1, 3)[used]
            positions.append(p)
            normals.append(arrays[f'g{draw["geometry"]}_normal'].reshape(-1, 3)[used] if g['normal'] else np.zeros_like(p))
            vertex_draw.append(np.full(len(used), d))
            triangles.append(remap[index].reshape(-1, 3))
            tri_draw.append(np.full(len(index) // 3, d))
            base += len(used)
        self.mesh_vertices = base
        # Each line segment is a quad: four vertices, two triangles, after the meshes.
        for d, _, _, segments in self.lines:
            for s in range(segments):
                v = base + 4 * s
                triangles.append(np.array([[v, v + 1, v + 2], [v + 1, v + 3, v + 2]]))
                tri_draw.append(np.full(2, d))
            base += 4 * segments
        cat = lambda a, dtype: torch.from_numpy(np.ascontiguousarray(np.concatenate(a))).to(dev, dtype)
        self.positions = cat(positions, torch.float32)
        self.normals = cat(normals, torch.float32)
        self.vertex_draw = cat(vertex_draw, torch.long)
        self.triangles = cat(triangles, torch.int32)
        tri_draw = np.concatenate(tri_draw)
        self.tri_draw = torch.from_numpy(tri_draw).to(dev)
        self.draw_count = len(scene['draws'])
        draws = scene['draws']
        draw_object = np.array([d['object'] for d in draws])
        draw_material = np.array([d['material'] for d in draws])
        tri_material = draw_material[tri_draw]
        is_line = np.array([draws[d]['kind'] == 'line' for d in tri_draw])
        tri_side = np.where(is_line, _DOUBLE, side[tri_material])
        self.draw_object = torch.from_numpy(draw_object).to(dev)
        self.tri_object = torch.from_numpy(draw_object[tri_draw]).to(dev)
        self.tri_material = torch.from_numpy(tri_material).to(dev)
        self.tri_side = torch.from_numpy(tri_side).to(dev)
        self.tri_overlay = torch.from_numpy(overlay[tri_material]).to(dev)
        self.tri_static = torch.from_numpy(static[tri_material]).to(dev)
        self.zero_tri = torch.zeros_like(self.triangles)
        self._layout = None
        self._staging = None
        self._flat = None
        self._pixel = torch.tensor([2.0 / self.width, 2.0 / self.height], device=dev)
        self._xyw = torch.tensor([0, 1, 3], device=dev)
        self.use_graphs = graphs if graphs is not None else os.environ.get('NATIVE_RENDER_GRAPHS', '1') == '1'
        self._graph_prepare = self._graph_compose = None
        self._compose = torch.compile(_compose, dynamic=False, fullgraph=True) if compile else _compose

    def _lights(self, lights):
        sky, ground, up, colors = [], [], [], []
        for light in lights:
            if light.get('unsupported') or light['type'] == 'ambient':
                raise NotImplementedError(f'Light type {light["type"]}')
            if light['type'] == 'hemisphere':
                sky.append([c * light['intensity'] for c in light['sky']])
                ground.append([c * light['intensity'] for c in light['ground']])
                up.append(light.get('direction', [0, 1, 0]))
            elif light['type'] == 'directional':
                colors.append([c * light['intensity'] for c in light['color']])
        t = lambda v: torch.tensor(v, device=self.device, dtype=torch.float32).reshape(-1, 3)
        self.hemi_sky, self.hemi_ground, self.hemi_up, self.light_colors = t(sky), t(ground), t(up), t(colors)
        self.hemi_up = F.normalize(self.hemi_up, dim=-1)

    def _upload(self, frame):
        """All of a frame's scene data in one pinned copy into a fixed device buffer."""
        names = ['objects', 'materials', 'camera', 'lights'] + [f'g{g}_position' for _, g, _, _ in self.lines]
        if self._layout is None:
            offsets, total = {}, 0
            for name in names:
                offsets[name] = (total, frame[name].size)
                total += frame[name].size
            self._layout = offsets
            self._staging = torch.empty(total, dtype=torch.float32, pin_memory=True)
            self._flat = torch.empty(total, dtype=torch.float32, device=self.device)
            self._copied = torch.cuda.Event()
        else:
            # The previous frame's copy out of the staging buffer is long done
            # by now; this only guards against overwriting it mid-flight.
            self._copied.synchronize()
        host = self._staging.numpy()
        for name in names:
            start, size = self._layout[name]
            host[start:start + size] = frame[name]
        self._flat.copy_(self._staging, non_blocking=True)
        self._copied.record()

    @torch.no_grad()
    def render(self, frame):
        """Render one exported frame; returns a new (H, W, 3) uint8 CUDA tensor."""
        self._upload(frame)
        if not self.use_graphs:
            return self._raster_and_compose(*self._prepare(self._flat))
        if self._graph_prepare is None:
            self._capture()
        # Two graph replays and four nvdiffrast calls per frame; everything
        # else is recorded once, so the host only queues a handful of launches.
        self._graph_prepare.replay()
        base_rast, base_interp, over_rast, over_interp = self._raster(*self._prepared[:4])
        self._raster_in[0].copy_(base_rast[0])
        self._raster_in[1].copy_(base_interp[0])
        self._raster_in[2].copy_(over_rast[0])
        self._raster_in[3].copy_(over_interp[0])
        self._graph_compose.replay()
        return self._image.clone()

    @torch.no_grad()
    def coverage(self, frame):
        """Per draw: visible pixels and their mean x, y (0-1).

        The same culling and depth test as ``render``, so a part hidden behind
        a wall or outside the view counts zero; opaque and overlay layers are
        counted separately and summed. Not graphed; for data tools.
        """
        self._upload(frame)
        clip, _, base_tris, over_tris, _, _, _ = self._prepare(self._flat)
        n = self.draw_count
        pixels = torch.zeros(n, device=self.device)
        sum_x, sum_y = torch.zeros_like(pixels), torch.zeros_like(pixels)
        for tris in (base_tris, over_tris):
            rast, _ = dr.rasterize(self.ctx, clip, tris, (self.height, self.width))
            tid = rast[0, ..., 3].long()
            covered = tid > 0
            draw = self.tri_draw[tid[covered] - 1]
            ys, xs = torch.nonzero(covered, as_tuple=True)
            pixels += torch.bincount(draw, minlength=n).float()
            sum_x += torch.bincount(draw, weights=xs.float() / self.width, minlength=n)
            sum_y += torch.bincount(draw, weights=ys.float() / self.height, minlength=n)
        mean = lambda s: s / pixels.clamp(min=1)
        # nvdiffrast rows start at the bottom; report y from the top like the image.
        return torch.stack([pixels, mean(sum_x), 1 - mean(sum_y)], 1).cpu().numpy()

    def _capture(self):
        side = torch.cuda.Stream()
        side.wait_stream(torch.cuda.current_stream())
        with torch.cuda.stream(side):
            for _ in range(3):
                self._prepare(self._flat)
        torch.cuda.current_stream().wait_stream(side)
        self._graph_prepare = torch.cuda.CUDAGraph()
        with torch.cuda.graph(self._graph_prepare):
            self._prepared = self._prepare(self._flat)
        self._graph_prepare.replay()
        base_rast, base_interp, over_rast, over_interp = self._raster(*self._prepared[:4])
        self._raster_in = [base_rast[0].clone(), base_interp[0].clone(), over_rast[0].clone(), over_interp[0].clone()]
        compose = lambda: self._compose_from(*self._raster_in, *self._prepared[4:])
        side.wait_stream(torch.cuda.current_stream())
        with torch.cuda.stream(side):
            for _ in range(3):
                compose()
        torch.cuda.current_stream().wait_stream(side)
        self._graph_compose = torch.cuda.CUDAGraph()
        with torch.cuda.graph(self._graph_compose):
            self._image = compose()
        torch.cuda.synchronize()

    def _raster(self, clip, attrs, base_tris, over_tris):
        H, W = self.height, self.width
        base_rast, _ = dr.rasterize(self.ctx, clip, base_tris, (H, W))
        base_interp, _ = dr.interpolate(attrs, base_rast, base_tris)
        over_rast, _ = dr.rasterize(self.ctx, clip, over_tris, (H, W))
        over_interp, _ = dr.interpolate(attrs, over_rast, over_tris)
        return base_rast, base_interp, over_rast, over_interp

    def _raster_and_compose(self, clip, attrs, base_tris, over_tris, table, eye, light_dirs):
        base_rast, base_interp, over_rast, over_interp = self._raster(clip, attrs, base_tris, over_tris)
        return self._compose_from(base_rast[0], base_interp[0], over_rast[0], over_interp[0], table, eye, light_dirs)

    def _compose_from(self, base_rast, base_interp, over_rast, over_interp, table, eye, light_dirs):
        return self._compose(base_rast, base_interp, over_rast, over_interp, table, self.clear, eye, light_dirs,
                             self.light_colors, self.hemi_sky, self.hemi_ground, self.hemi_up, self.exposure,
                             self.tone_enabled, self.aces_in, self.aces_out)

    def _prepare(self, flat):
        """Device-only work before rasterization: transforms, culling, the triangle table."""
        H, W = self.height, self.width
        data = {name: flat[start:start + size] for name, (start, size) in self._layout.items()}
        objects = data['objects'].view(-1, 17)
        materials = data['materials'].view(-1, 8)
        camera = data['camera']
        visible = objects[:, 0] > 0.5
        model = objects[:, 1:].reshape(-1, 4, 4).transpose(1, 2)   # three stores column-major
        view_projection = camera[:16].reshape(4, 4).T @ camera[16:32].reshape(4, 4).T
        eye = camera[32:35]

        # Mesh vertices to world and clip space. The cofactor matrix is
        # det * inverse-transpose: the normal matrix up to scale (normals are
        # renormalized), and still defined for a zero-scaled object.
        draw_model = model[self.draw_object]
        vm = draw_model[self.vertex_draw]
        world = torch.einsum('vij,vj->vi', vm[:, :3, :3], self.positions) + vm[:, :3, 3]
        m3 = draw_model[:, :3, :3]
        c0, c1, c2 = m3[:, :, 0], m3[:, :, 1], m3[:, :, 2]
        cofactor = torch.stack([torch.cross(c1, c2, dim=-1), torch.cross(c2, c0, dim=-1), torch.cross(c0, c1, dim=-1)], 2)
        cofactor = cofactor * torch.where(_det3(m3) < 0, -1.0, 1.0)[:, None, None]
        world_normal = torch.einsum('vij,vj->vi', cofactor[self.vertex_draw], self.normals)
        clip = torch.cat([world, torch.ones_like(world[:, :1])], 1) @ view_projection.T

        # Lines: one-pixel quads around each projected segment.
        if self.lines:
            line_world, line_clip = [], []
            pixel = self._pixel
            for d, g, _, _ in self.lines:
                points = data[f'g{g}_position'].view(-1, 3)
                m = draw_model[d]
                pw = points @ m[:3, :3].T + m[:3, 3]
                pc = torch.cat([pw, torch.ones_like(pw[:, :1])], 1) @ view_projection.T
                a, b = pc[:-1], pc[1:]
                wa, wb = a[:, 3:].clamp(min=1e-6), b[:, 3:].clamp(min=1e-6)
                direction = (b[:, :2] / wb - a[:, :2] / wa) / pixel
                offset = F.normalize(torch.stack([-direction[:, 1], direction[:, 0]], 1), dim=-1) * 0.5 * pixel
                quad = [a.clone(), a.clone(), b.clone(), b.clone()]
                for q, sign, w in zip(quad, (1, -1, 1, -1), (wa, wa, wb, wb)):
                    q[:, :2] += sign * offset * w
                line_clip.append(torch.stack(quad, 1).reshape(-1, 4))
                line_world.append(torch.stack([pw[:-1], pw[:-1], pw[1:], pw[1:]], 1).reshape(-1, 3))
            line_world = torch.cat(line_world)
            world = torch.cat([world, line_world])
            world_normal = torch.cat([world_normal, torch.zeros_like(line_world)])
            clip = torch.cat([clip, torch.cat(line_clip)])

        # Culling as WebGL does it: hidden objects, then the face the material
        # does not draw. The homogeneous determinant keeps the orientation
        # right for triangles crossing the near plane. Culled triangles become
        # degenerate rather than being filtered, so nothing waits on the host.
        tri = self.triangles.long()
        facing = _det3(clip.index_select(1, self._xyw)[tri]) > 0
        keep = visible[self.tri_object]
        keep &= (self.tri_side == _DOUBLE) | ((self.tri_side == _FRONT) & facing) | ((self.tri_side == _BACK) & ~facing)
        corners = world[tri]
        face = torch.cross(corners[:, 1] - corners[:, 0], corners[:, 2] - corners[:, 0], dim=-1)
        table = torch.cat([self.tri_static, materials[self.tri_material], face, facing[:, None].float()], 1)

        clip = clip.unsqueeze(0).contiguous()
        attrs = torch.cat([world, world_normal], 1).unsqueeze(0).contiguous()
        base_tris = torch.where((keep & ~self.tri_overlay)[:, None], self.triangles, self.zero_tri).contiguous()
        over_tris = torch.where((keep & self.tri_overlay)[:, None], self.triangles, self.zero_tri).contiguous()
        return clip, attrs, base_tris, over_tris, table, eye, data['lights'].view(-1, 3)