"""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)