v3: geometry gradients (dual-number interior + shadow and camera silhouette edge sampling)
1f9a369 verified | """pathtracer-diff: a differentiable Monte Carlo path tracer as one kernel. | |
| Forward renders a linear-radiance image with a megakernel path tracer: | |
| Lambertian diffuse, GGX conductor (VNDF-sampled, height-correlated Smith), | |
| smooth dielectric, rough plastic, and rough dielectric materials; area | |
| lights with per-texel emission and an importance-sampled equirectangular | |
| environment map; an optional homogeneous participating medium; multiple | |
| importance sampling by default; per-texel albedo with bilinear filtering; | |
| binned-SAH BVH. One thread owns one pixel, so the image is bitwise | |
| deterministic. | |
| Backward is exact path replay: sampling distributions depend only on | |
| geometry, frozen material parameters (roughness, ior), the environment CDF | |
| and medium rate (detached at Scene construction), and the counter-based | |
| Philox stream, never on the differentiable parameters. Every radiance term | |
| is a product of per-bounce factors that are affine in their vertex's | |
| albedo texels (Schlick Fresnel is affine in F0), times a linear emission | |
| or environment texel, with medium sigma-derivatives in closed form, so the | |
| replay differentiates exactly and scatters through the bilinear-footprint | |
| adjoint. render() gradients flow to albedo texels, emission texels, | |
| environment texels, and medium sigmas. geometry_grad() adds | |
| d(loss)/d(vertex positions) for direct lighting: a dual-number interior | |
| term plus edge-sampled shadow and camera silhouette boundary terms. | |
| from kernels import get_kernel | |
| ptd = get_kernel("phanerozoic/pathtracer-diff", version=1, trust_remote_code=True) | |
| scene = ptd.Scene(vertices, faces, material_ids, albedo, emission, | |
| uvs=uvs, material_types=types, roughness=rough, | |
| ior=ior, env=env_map) | |
| cam = ptd.Camera(position=(2.78, 2.73, -8.0), look_at=(2.78, 2.73, 2.8)) | |
| img = ptd.render(scene, cam, 512, 512, spp=64) | |
| img.sum().backward() # grads on albedo texels / emission / env texels | |
| gv = ptd.geometry_grad(scene, cam, dL_dimage) # grads on vertices | |
| """ | |
| import math | |
| import torch | |
| from ._ops import ops | |
| MAX_MATERIALS = 64 | |
| MAX_BOUNCES = 16 | |
| DIFFUSE, CONDUCTOR, DIELECTRIC, PLASTIC, ROUGH_DIELECTRIC = 0, 1, 2, 3, 4 | |
| _MODES = {"brdf": 0, "nee": 1, "mis": 2} | |
| __all__ = ["Scene", "Camera", "render", "geometry_grad", "ops", | |
| "MAX_MATERIALS", "MAX_BOUNCES", "DIFFUSE", "CONDUCTOR", | |
| "DIELECTRIC", "PLASTIC", "ROUGH_DIELECTRIC"] | |
| def _build_bvh(tris, leaf_size=4, sah_bins=16): | |
| """Binned-SAH BVH over [F, 9] triangles, built level-wise with fully | |
| vectorized torch ops (segmented reductions + stable-sort partitioning), | |
| so million-triangle scenes build in seconds. Returns (nodes_f [N, 6], | |
| nodes_i [N, 3], order [F]). nodes_i rows are (left, right, axis << 1) | |
| for internal nodes and (start, count, 1) for leaves; internal children | |
| are the (lower, upper) coordinate halves along the split axis, so | |
| traversal can visit the near child first.""" | |
| dev = "cuda" if torch.cuda.is_available() else "cpu" | |
| tris = tris.to(torch.float32).to(dev) | |
| F = tris.shape[0] | |
| v = tris.reshape(F, 3, 3) | |
| tlo = v.amin(dim=1) | |
| thi = v.amax(dim=1) | |
| cen = (tlo + thi) * 0.5 | |
| perm = torch.arange(F, device=dev) | |
| # active segments: contiguous [start, end) ranges of perm, one per node | |
| starts = torch.zeros(1, dtype=torch.long, device=dev) | |
| ends = torch.full((1,), F, dtype=torch.long, device=dev) | |
| ids = torch.zeros(1, dtype=torch.long, device=dev) | |
| INF = float("inf") | |
| nf_parts, ni_parts = [], [] # per-node rows appended in id order | |
| n_nodes = 1 | |
| def half_area(lo, hi): | |
| d = (hi - lo).clamp(min=0) | |
| return d[..., 0] * d[..., 1] + d[..., 1] * d[..., 2] + \ | |
| d[..., 2] * d[..., 0] | |
| while ids.numel() > 0: | |
| A = ids.numel() | |
| seg_len = ends - starts | |
| total = int(seg_len.sum()) | |
| # active positions and per-position segment index | |
| csum = torch.cumsum(seg_len, 0) - seg_len | |
| pos = starts.repeat_interleave(seg_len) + \ | |
| (torch.arange(total, device=dev) - | |
| csum.repeat_interleave(seg_len)) | |
| s = torch.arange(A, device=dev).repeat_interleave(seg_len) | |
| tri = perm[pos] | |
| # per-segment node bounds and centroid bounds | |
| nlo = torch.full((A, 3), INF, device=dev) | |
| nhi = torch.full((A, 3), -INF, device=dev) | |
| clo = torch.full((A, 3), INF, device=dev) | |
| chi = torch.full((A, 3), -INF, device=dev) | |
| s3 = s.unsqueeze(1).expand(-1, 3) | |
| nlo.scatter_reduce_(0, s3, tlo[tri], reduce="amin", include_self=True) | |
| nhi.scatter_reduce_(0, s3, thi[tri], reduce="amax", include_self=True) | |
| clo.scatter_reduce_(0, s3, cen[tri], reduce="amin", include_self=True) | |
| chi.scatter_reduce_(0, s3, cen[tri], reduce="amax", include_self=True) | |
| pad = 1e-5 * (1.0 + nlo.abs() + nhi.abs()) | |
| bounds = torch.cat([nlo - pad, nhi + pad], dim=1) # [A, 6] | |
| leaf_mask = seg_len <= leaf_size | |
| split_mask = ~leaf_mask | |
| # choose axis (largest centroid extent) and bin the centroids | |
| ext = (chi - clo).clamp(min=0) | |
| axis = ext.argmax(dim=1) # [A] | |
| ax_pos = axis[s] | |
| coord = cen[tri].gather(1, ax_pos.unsqueeze(1)).squeeze(1) | |
| seg_lo = clo.gather(1, axis.unsqueeze(1)).squeeze(1)[s] | |
| seg_ext = ext.gather(1, axis.unsqueeze(1)).squeeze(1)[s] | |
| t = (coord - seg_lo) / seg_ext.clamp(min=1e-30) | |
| bin_i = (t * sah_bins).long().clamp(0, sah_bins - 1) | |
| degen = seg_ext <= 1e-12 # all centroids equal on the axis | |
| # per-(segment, bin) counts and bounds | |
| key = s * sah_bins + bin_i | |
| counts = torch.bincount(key, minlength=A * sah_bins) \ | |
| .reshape(A, sah_bins) | |
| blo = torch.full((A * sah_bins, 3), INF, device=dev) | |
| bhi = torch.full((A * sah_bins, 3), -INF, device=dev) | |
| k3 = key.unsqueeze(1).expand(-1, 3) | |
| blo.scatter_reduce_(0, k3, tlo[tri], reduce="amin", include_self=True) | |
| bhi.scatter_reduce_(0, k3, thi[tri], reduce="amax", include_self=True) | |
| blo = blo.reshape(A, sah_bins, 3) | |
| bhi = bhi.reshape(A, sah_bins, 3) | |
| # prefix (left) and suffix (right) running bounds and counts | |
| plo = torch.cummin(blo, dim=1).values | |
| phi = torch.cummax(bhi, dim=1).values | |
| pn = torch.cumsum(counts, dim=1) | |
| slo = torch.flip(torch.cummin(torch.flip(blo, [1]), dim=1).values, [1]) | |
| shi = torch.flip(torch.cummax(torch.flip(bhi, [1]), dim=1).values, [1]) | |
| sn = torch.flip(torch.cumsum(torch.flip(counts, [1]), dim=1), [1]) | |
| # SAH cost of splitting after bin b (b = 0..bins-2) | |
| nl = pn[:, :-1].float() | |
| nr = sn[:, 1:].float() | |
| cost = half_area(plo[:, :-1], phi[:, :-1]) * nl + \ | |
| half_area(slo[:, 1:], shi[:, 1:]) * nr | |
| cost = torch.where((nl > 0) & (nr > 0), cost, | |
| torch.full_like(cost, INF)) | |
| best_cost, best_bin = cost.min(dim=1) | |
| sah_ok = torch.isfinite(best_cost) & ~degen[csum] # per segment | |
| # side: left if bin <= best_bin (SAH) else lower positional half | |
| # (median fallback for degenerate segments) | |
| side_left = bin_i <= best_bin[s] | |
| within = pos - starts[s] | |
| med_left = within < (seg_len[s] // 2) | |
| use_sah = sah_ok[s] | |
| side_left = torch.where(use_sah, side_left, med_left) | |
| side_left = side_left & split_mask[s] | |
| # per-segment left count; guard SAH splits that put everything on | |
| # one side (possible only via numerics) with the median fallback | |
| nl_seg = torch.zeros(A, dtype=torch.long, device=dev) | |
| nl_seg.scatter_add_(0, s, side_left.long()) | |
| bad = split_mask & ((nl_seg == 0) | (nl_seg == seg_len)) | |
| if bool(bad.any()): | |
| fix = bad[s] | |
| side_left = torch.where(fix, med_left & split_mask[s], side_left) | |
| nl_seg = torch.zeros(A, dtype=torch.long, device=dev) | |
| nl_seg.scatter_add_(0, s, side_left.long()) | |
| # partition each segment in place: stable sort by (segment, right?) | |
| k2 = s * 2 + (~side_left).long() | |
| order2 = torch.argsort(k2, stable=True) | |
| perm[pos] = tri[order2] | |
| # allocate children for splitting segments, emit node rows | |
| n_split = int(split_mask.sum()) | |
| child_rank = torch.cumsum(split_mask.long(), 0) - 1 | |
| left_ids = n_nodes + 2 * child_rank | |
| right_ids = left_ids + 1 | |
| ni = torch.empty(A, 3, dtype=torch.int64, device=dev) | |
| ni[:, 0] = torch.where(split_mask, left_ids, starts) | |
| ni[:, 1] = torch.where(split_mask, right_ids, seg_len) | |
| ni[:, 2] = torch.where(split_mask, axis * 2, | |
| torch.ones_like(axis)) | |
| nf_parts.append((ids, bounds)) | |
| ni_parts.append((ids, ni)) | |
| # next level | |
| mid = starts + nl_seg | |
| new_starts = torch.cat([starts[split_mask], mid[split_mask]]) | |
| new_ends = torch.cat([mid[split_mask], ends[split_mask]]) | |
| new_ids = torch.cat([left_ids[split_mask], right_ids[split_mask]]) | |
| n_nodes += 2 * n_split | |
| starts, ends, ids = new_starts, new_ends, new_ids | |
| nodes_f = torch.empty(n_nodes, 6, device=dev) | |
| nodes_i = torch.empty(n_nodes, 3, dtype=torch.int64, device=dev) | |
| for idv, rows in nf_parts: | |
| nodes_f[idv] = rows | |
| for idv, rows in ni_parts: | |
| nodes_i[idv] = rows | |
| return (nodes_f.cpu(), nodes_i.to(torch.int32).cpu(), perm.cpu()) | |
| def _build_bvh_reference(tris, leaf_size=4, force_split=8, sah_bins=16): | |
| """Recursive reference builder (kept for cross-checking the vectorized | |
| build in development).""" | |
| tris = tris.to(torch.float32) | |
| F = tris.shape[0] | |
| v = tris.reshape(F, 3, 3) | |
| lo = v.amin(dim=1) | |
| hi = v.amax(dim=1) | |
| cen = (lo + hi) * 0.5 | |
| nodes_f, nodes_i, order = [], [], [] | |
| def half_area(blo, bhi): | |
| d = (bhi - blo).clamp(min=0) | |
| return d[0] * d[1] + d[1] * d[2] + d[2] * d[0] | |
| stack = [] | |
| def alloc(idx): | |
| nid = len(nodes_f) | |
| nodes_f.append(None) | |
| nodes_i.append(None) | |
| stack.append((nid, idx)) | |
| return nid | |
| root = alloc(torch.arange(F)) | |
| while stack: | |
| nid, idx = stack.pop() | |
| n = idx.numel() | |
| blo = lo[idx].amin(dim=0) | |
| bhi = hi[idx].amax(dim=0) | |
| pad = 1e-5 * (1.0 + blo.abs() + bhi.abs()) | |
| nodes_f[nid] = torch.cat([blo - pad, bhi + pad]) | |
| if n <= leaf_size: | |
| start = len(order) | |
| order.extend(idx.tolist()) | |
| nodes_i[nid] = (start, n, 1) | |
| continue | |
| cb_lo = cen[idx].amin(dim=0) | |
| cb_hi = cen[idx].amax(dim=0) | |
| ext = cb_hi - cb_lo | |
| axis = int(ext.argmax()) | |
| split_pos = None | |
| if float(ext[axis]) > 1e-12: | |
| c = cen[idx, axis] | |
| edges = torch.linspace(float(cb_lo[axis]), float(cb_hi[axis]), | |
| sah_bins + 1) | |
| b = torch.bucketize(c, edges[1:-1]) | |
| counts = torch.bincount(b, minlength=sah_bins) | |
| # per-bin bounds over the node's triangles | |
| binlo = torch.full((sah_bins, 3), float("inf")) | |
| binhi = torch.full((sah_bins, 3), float("-inf")) | |
| binlo.scatter_reduce_(0, b.unsqueeze(1).expand(-1, 3), lo[idx], | |
| reduce="amin", include_self=True) | |
| binhi.scatter_reduce_(0, b.unsqueeze(1).expand(-1, 3), hi[idx], | |
| reduce="amax", include_self=True) | |
| best_cost, best_bin = None, None | |
| nl = 0 | |
| llo = torch.full((3,), float("inf")) | |
| lhi = torch.full((3,), float("-inf")) | |
| pre = [] | |
| for i in range(sah_bins - 1): | |
| if counts[i] > 0: | |
| llo = torch.minimum(llo, binlo[i]) | |
| lhi = torch.maximum(lhi, binhi[i]) | |
| nl += int(counts[i]) | |
| pre.append((nl, llo.clone(), lhi.clone())) | |
| nr = 0 | |
| rlo = torch.full((3,), float("inf")) | |
| rhi = torch.full((3,), float("-inf")) | |
| for i in range(sah_bins - 1, 0, -1): | |
| if counts[i] > 0: | |
| rlo = torch.minimum(rlo, binlo[i]) | |
| rhi = torch.maximum(rhi, binhi[i]) | |
| nr += int(counts[i]) | |
| nl_i, llo_i, lhi_i = pre[i - 1] | |
| if nl_i == 0 or nr == 0: | |
| continue | |
| cost = (half_area(llo_i, lhi_i) * nl_i + | |
| half_area(rlo, rhi) * nr) | |
| if best_cost is None or float(cost) < best_cost: | |
| best_cost = float(cost) | |
| best_bin = i - 1 | |
| if best_bin is not None: | |
| parent_area = half_area(blo, bhi) | |
| leaf_cost = float(n) * float(parent_area) | |
| split_cost = 0.125 * float(parent_area) + best_cost | |
| if split_cost < leaf_cost or n > force_split: | |
| mask = b <= best_bin | |
| left_idx = idx[mask] | |
| right_idx = idx[~mask] | |
| if left_idx.numel() > 0 and right_idx.numel() > 0: | |
| split_pos = (left_idx, right_idx) | |
| if split_pos is None: | |
| if n > force_split and float(ext[axis]) > 1e-12: | |
| srt = idx[torch.argsort(cen[idx, axis], stable=True)] | |
| mid = n // 2 | |
| split_pos = (srt[:mid], srt[mid:]) | |
| else: | |
| start = len(order) | |
| order.extend(idx.tolist()) | |
| nodes_i[nid] = (start, n, 1) | |
| continue | |
| left = alloc(split_pos[0]) | |
| right = alloc(split_pos[1]) | |
| nodes_i[nid] = (left, right, axis << 1) | |
| assert root == 0 | |
| return (torch.stack(nodes_f), | |
| torch.tensor(nodes_i, dtype=torch.int32), | |
| torch.tensor(order, dtype=torch.int64)) | |
| class Camera: | |
| """Pinhole camera. `tensor(H, W)` packs (pos, forward, right*tan(v/2)*aspect, | |
| up*tan(v/2)) as a float32 [12] tensor for the kernel.""" | |
| def __init__(self, position, look_at, up=(0.0, 1.0, 0.0), vfov_deg=40.0): | |
| p = torch.tensor(position, dtype=torch.float64) | |
| t = torch.tensor(look_at, dtype=torch.float64) | |
| u = torch.tensor(up, dtype=torch.float64) | |
| f = t - p | |
| f = f / f.norm() | |
| r = torch.linalg.cross(f, u) | |
| r = r / r.norm() | |
| uu = torch.linalg.cross(r, f) | |
| self.position, self.forward, self.right, self.up = p, f, r, uu | |
| self.vfov_deg = float(vfov_deg) | |
| def tensor(self, H, W, device="cuda"): | |
| th = math.tan(math.radians(self.vfov_deg) * 0.5) | |
| rs = self.right * (th * W / H) | |
| us = self.up * th | |
| return torch.cat([self.position, self.forward, rs, us]).to( | |
| device=device, dtype=torch.float32).contiguous() | |
| def _as_texture(t, device): | |
| t = torch.as_tensor(t, dtype=torch.float32) | |
| if t.dim() == 1 and t.numel() == 3: | |
| t = t.reshape(1, 1, 3) | |
| if t.dim() != 3 or t.shape[2] != 3 or t.shape[0] < 1 or t.shape[1] < 1: | |
| raise ValueError("each albedo texture must be [H, W, 3]") | |
| return t.to(device) | |
| class Scene: | |
| """Triangle-mesh scene. | |
| vertices [V, 3], faces [F, 3], material_ids [F]. albedo: [M, 3] | |
| constants or a list of M textures [Hm, Wm, 3] (constants are 1x1 | |
| textures; both forms receive gradients). emission [M, 3] (may require | |
| grad). uvs: None, [V, 2], or [F, 3, 2]; textures wrap. material_types | |
| [M] of DIFFUSE|CONDUCTOR|DIELECTRIC (default all DIFFUSE); roughness | |
| [M] GGX alpha for conductors (default 0.3, clamped >= 0.01); ior [M] | |
| for dielectrics (default 1.5). env: optional [Eh, Ew, 3] equirect | |
| radiance map (may require grad); its importance-sampling CDF is built | |
| detached at construction. Texture and env shapes are fixed at | |
| construction.""" | |
| def __init__(self, vertices, faces, material_ids, albedo, emission, | |
| uvs=None, material_types=None, roughness=None, ior=None, | |
| env=None, medium=None, device="cuda"): | |
| vertices = torch.as_tensor(vertices, dtype=torch.float32) | |
| faces = torch.as_tensor(faces, dtype=torch.int64) | |
| material_ids = torch.as_tensor(material_ids, dtype=torch.int64) | |
| is_texture_list = (isinstance(albedo, (list, tuple)) and | |
| any(torch.is_tensor(t) for t in albedo)) | |
| if is_texture_list: | |
| self.albedo = None | |
| self.albedo_textures = [_as_texture(t, device) for t in albedo] | |
| else: | |
| if not torch.is_tensor(albedo): | |
| albedo = torch.tensor(albedo, dtype=torch.float32, | |
| device=device) | |
| if albedo.dim() != 2 or albedo.shape[1] != 3: | |
| raise ValueError("albedo must be [M, 3] or a list of textures") | |
| self.albedo = albedo.to(device) | |
| self.albedo_textures = None | |
| M = (len(self.albedo_textures) if self.albedo is None | |
| else self.albedo.shape[0]) | |
| if M > MAX_MATERIALS: | |
| raise ValueError(f"at most {MAX_MATERIALS} materials, got {M}") | |
| emi_texture_list = (isinstance(emission, (list, tuple)) and | |
| any(torch.is_tensor(t) for t in emission)) | |
| if emi_texture_list: | |
| self.emission = None | |
| self.emission_textures = [_as_texture(t, device) for t in emission] | |
| if len(self.emission_textures) != M: | |
| raise ValueError("emission list must have M entries") | |
| else: | |
| if not torch.is_tensor(emission): | |
| emission = torch.tensor(emission, dtype=torch.float32, | |
| device=device) | |
| if emission.shape != (M, 3): | |
| raise ValueError("emission must be [M, 3]") | |
| self.emission = emission.to(device) | |
| self.emission_textures = None | |
| if int(material_ids.max()) >= M or int(material_ids.min()) < 0: | |
| raise ValueError("material_ids out of range") | |
| def mvec(x, default, dtype): | |
| if x is None: | |
| return torch.full((M,), default, dtype=dtype, device=device) | |
| t = torch.as_tensor(x, dtype=dtype).reshape(-1).to(device) | |
| if t.numel() != M: | |
| raise ValueError("per-material array must have M entries") | |
| return t | |
| self.mat_type = mvec(material_types, DIFFUSE, torch.int32).contiguous() | |
| if int(self.mat_type.max()) > 4 or int(self.mat_type.min()) < 0: | |
| raise ValueError("material_types must be 0..4") | |
| self.mat_rough = mvec(roughness, 0.3, torch.float32).contiguous() | |
| self.mat_ior = mvec(ior, 1.5, torch.float32).contiguous() | |
| def build_hdr(constants, textures): | |
| shapes = ([(1, 1)] * M if constants is not None else | |
| [(int(t.shape[0]), int(t.shape[1])) for t in textures]) | |
| hdr, off = [], 0 | |
| for (h, w) in shapes: | |
| hdr.append((off, w, h)) | |
| off += h * w | |
| return (torch.tensor(hdr, dtype=torch.int32, | |
| device=device).contiguous(), off) | |
| self.tex_hdr, self.n_texels = build_hdr(self.albedo, | |
| self.albedo_textures) | |
| self.emi_hdr, self.n_emi_texels = build_hdr(self.emission, | |
| self.emission_textures) | |
| # homogeneous medium: live sigmas (may require grad), frozen | |
| # (detached) sampling rate | |
| if medium is not None: | |
| if env is not None: | |
| raise ValueError("medium and env are mutually exclusive") | |
| sa, ss = medium | |
| sa = sa if torch.is_tensor(sa) else torch.tensor( | |
| sa, dtype=torch.float32, device=device) | |
| ss = ss if torch.is_tensor(ss) else torch.tensor( | |
| ss, dtype=torch.float32, device=device) | |
| if sa.numel() != 3 or ss.numel() != 3: | |
| raise ValueError("medium sigmas must be [3]") | |
| self.med_sa = sa.to(device) | |
| self.med_ss = ss.to(device) | |
| st = (sa.detach() + ss.detach()).reshape(-1) | |
| self.med_sbar = float(st.mean().clamp(min=1e-6)) | |
| else: | |
| self.med_sa = None | |
| self.med_ss = None | |
| self.med_sbar = 0.0 | |
| tris = vertices[faces].reshape(-1, 9) | |
| if uvs is None: | |
| uv_c = torch.zeros(tris.shape[0], 3, 2, dtype=torch.float32) | |
| else: | |
| uvs = torch.as_tensor(uvs, dtype=torch.float32) | |
| if uvs.dim() == 2 and uvs.shape[1] == 2: | |
| uv_c = uvs[faces] | |
| elif uvs.dim() == 3 and uvs.shape[1:] == (3, 2): | |
| uv_c = uvs | |
| else: | |
| raise ValueError("uvs must be [V, 2] or [F, 3, 2]") | |
| nodes_f, nodes_i, order = _build_bvh(tris) | |
| tris = tris[order] | |
| uv_c = uv_c[order] | |
| mat_ids = material_ids[order].to(torch.int32) | |
| self.face_verts = faces[order].to(torch.int32).to(device).contiguous() | |
| self.n_verts = int(vertices.shape[0]) | |
| self._vertices_cpu = vertices.cpu() | |
| self._faces_cpu = faces[order].cpu() | |
| self._edges = None | |
| if self.emission is not None: | |
| em_max = self.emission.detach().cpu().amax(dim=1) | |
| else: | |
| em_max = torch.stack([t.detach().max().cpu() | |
| for t in self.emission_textures]) | |
| emissive = (em_max > 0)[mat_ids.long()] | |
| lf = torch.nonzero(emissive, as_tuple=False).flatten().to(torch.int32) | |
| if lf.numel() > 0: | |
| t = tris[lf.long()].reshape(-1, 3, 3).double() | |
| e1 = t[:, 1] - t[:, 0] | |
| e2 = t[:, 2] - t[:, 0] | |
| areas = 0.5 * torch.linalg.cross(e1, e2).norm(dim=1) | |
| total = float(areas.sum()) | |
| cdf = (areas.cumsum(0) / areas.sum()).float() | |
| else: | |
| total = 0.0 | |
| cdf = torch.zeros(0, dtype=torch.float32) | |
| # environment map + detached sampling tables | |
| if env is not None: | |
| env = torch.as_tensor(env, dtype=torch.float32) | |
| if env.dim() != 3 or env.shape[2] != 3 or env.shape[0] < 2 or \ | |
| env.shape[1] < 2: | |
| raise ValueError("env must be [Eh, Ew, 3]") | |
| self.env = env.to(device) | |
| eh, ew = int(env.shape[0]), int(env.shape[1]) | |
| lum = self.env.detach().mean(dim=2).cpu().double() + 1e-8 | |
| sint = torch.sin((torch.arange(eh, dtype=torch.float64) + 0.5) | |
| * math.pi / eh).clamp(min=1e-4) | |
| roww = (lum.sum(dim=1) * sint) | |
| row_p = roww / roww.sum() | |
| cdf_m = row_p.cumsum(0).float() | |
| col_p = lum / lum.sum(dim=1, keepdim=True) | |
| cdf_c = col_p.cumsum(1).float() | |
| pdf_img = (row_p.unsqueeze(1) * col_p).float() # sums to 1 | |
| self.env_w, self.env_h = ew, eh | |
| self.env_cdf_m = cdf_m.to(device).contiguous() | |
| self.env_cdf_c = cdf_c.reshape(-1).to(device).contiguous() | |
| self.env_pdf = pdf_img.reshape(-1).to(device).contiguous() | |
| else: | |
| self.env = None | |
| self.env_w = self.env_h = 0 | |
| z = torch.zeros(0, dtype=torch.float32, device=device) | |
| self.env_cdf_m = z | |
| self.env_cdf_c = z | |
| self.env_pdf = z | |
| self.device = device | |
| self.tris = tris.to(device).contiguous() | |
| self.mat_ids = mat_ids.to(device).contiguous() | |
| self.uvs = uv_c.to(device).contiguous() | |
| self.nodes_f = nodes_f.to(device).contiguous() | |
| self.nodes_i = nodes_i.to(device).contiguous() | |
| self.light_faces = lf.to(device).contiguous() | |
| self.light_cdf = cdf.to(device).contiguous() | |
| self.total_light_area = total | |
| def n_faces(self): | |
| return self.tris.shape[0] | |
| def _flatten(constants, textures, hdr, what): | |
| if constants is not None: | |
| if constants.shape[0] != hdr.shape[0]: | |
| raise ValueError(f"{what}/material count changed") | |
| return constants.reshape(-1, 3) | |
| parts = [] | |
| for m, t in enumerate(textures): | |
| off, w, h = (int(x) for x in hdr[m]) | |
| if (int(t.shape[0]), int(t.shape[1])) != (h, w): | |
| raise ValueError(f"{what} shapes are fixed at construction") | |
| parts.append(t.reshape(-1, 3)) | |
| return torch.cat(parts, dim=0) | |
| def _flat_albedo(self): | |
| return self._flatten(self.albedo, self.albedo_textures, self.tex_hdr, | |
| "albedo") | |
| def _flat_emission(self): | |
| return self._flatten(self.emission, self.emission_textures, | |
| self.emi_hdr, "emission") | |
| def _flat_env(self): | |
| if self.env is None: | |
| return torch.zeros(0, 3, dtype=torch.float32, device=self.device) | |
| if (int(self.env.shape[0]), int(self.env.shape[1])) != \ | |
| (self.env_h, self.env_w): | |
| raise ValueError("env shape is fixed at construction") | |
| return self.env.reshape(-1, 3) | |
| class _RenderFn(torch.autograd.Function): | |
| def _args(scene, tex_flat, emi_flat, env_flat, med_sa, med_ss, cam_t, | |
| spp, max_bounces, mode, seed): | |
| return (scene.tris, scene.mat_ids, scene.uvs, scene.nodes_f, | |
| scene.nodes_i, scene.light_faces, scene.light_cdf, | |
| float(scene.total_light_area), | |
| tex_flat.detach().contiguous(), scene.tex_hdr, | |
| emi_flat.detach().contiguous(), scene.emi_hdr, | |
| scene.mat_type, scene.mat_rough, scene.mat_ior, | |
| med_sa.detach().contiguous(), med_ss.detach().contiguous(), | |
| float(scene.med_sbar), env_flat.detach().contiguous(), | |
| scene.env_cdf_m, scene.env_cdf_c, scene.env_pdf, scene.env_w, | |
| scene.env_h, cam_t, spp, max_bounces, mode, seed) | |
| def forward(ctx, tex_flat, emi_flat, env_flat, med_sa, med_ss, scene, | |
| cam_t, H, W, spp, max_bounces, mode, seed): | |
| image = torch.empty(H, W, 3, device=tex_flat.device, | |
| dtype=torch.float32) | |
| ops.pt_forward(*_RenderFn._args(scene, tex_flat, emi_flat, env_flat, | |
| med_sa, med_ss, cam_t, spp, | |
| max_bounces, mode, seed), image) | |
| ctx.scene = scene | |
| ctx.cam_t = cam_t | |
| ctx.params = (H, W, spp, max_bounces, mode, seed) | |
| ctx.save_for_backward(tex_flat, emi_flat, env_flat, med_sa, med_ss) | |
| return image | |
| def backward(ctx, grad_image): | |
| tex_flat, emi_flat, env_flat, med_sa, med_ss = ctx.saved_tensors | |
| scene, cam_t = ctx.scene, ctx.cam_t | |
| H, W, spp, max_bounces, mode, seed = ctx.params | |
| ga = torch.zeros_like(tex_flat) | |
| ge = torch.zeros_like(emi_flat) | |
| genv = torch.zeros_like(env_flat) | |
| gmed = torch.zeros(6 if med_sa.numel() else 0, device=tex_flat.device) | |
| ops.pt_backward(*_RenderFn._args(scene, tex_flat, emi_flat, env_flat, | |
| med_sa, med_ss, cam_t, spp, | |
| max_bounces, mode, seed), | |
| grad_image.contiguous(), ga, ge, genv, gmed) | |
| gsa = gmed[:3] if med_sa.numel() else None | |
| gss = gmed[3:] if med_sa.numel() else None | |
| return (ga, ge, genv, gsa, gss, None, None, None, None, None, None, | |
| None, None) | |
| def _scene_edges(scene): | |
| if scene._edges is None: | |
| fv = scene._faces_cpu | |
| vp = scene._vertices_cpu | |
| first = {} | |
| rows = [] | |
| for f in range(fv.shape[0]): | |
| ids = fv[f] | |
| for (ca, cb) in ((0, 1), (1, 2), (2, 0)): | |
| a, b = int(ids[ca]), int(ids[cb]) | |
| key = (min(a, b), max(a, b)) | |
| if key in first: | |
| rows[first[key]][3] = f | |
| else: | |
| first[key] = len(rows) | |
| rows.append([ca, cb, f, -1]) | |
| edges = torch.tensor(rows, dtype=torch.int32) | |
| p0 = vp[fv[edges[:, 2].long(), edges[:, 0].long()].long()] | |
| p1 = vp[fv[edges[:, 2].long(), edges[:, 1].long()].long()] | |
| lens = (p1 - p0).norm(dim=1).double() | |
| total = float(lens.sum()) | |
| cdf = (lens.cumsum(0) / lens.sum()).float() | |
| cdf = torch.cat([cdf, torch.tensor([total])]) | |
| scene._edges = (edges.to(scene.device).contiguous(), | |
| cdf.to(scene.device).contiguous()) | |
| return scene._edges | |
| def geometry_grad(scene, camera, grad_image, spp=16, edge_samples=1 << 16, | |
| seed=0): | |
| """d(loss)/d(vertex positions) for the DIRECT-lighting transport term: | |
| the detached-sampling interior derivative (dual-number shading + light | |
| area measure) plus edge-sampled shadow-silhouette and primary | |
| (camera-silhouette) boundary terms. Diffuse receivers/emitters carry | |
| the radiance jumps; indirect bounces are not differentiated with | |
| respect to geometry. grad_image is dLoss/dImage [H, W, 3]; returns | |
| grad_verts [V, 3].""" | |
| edges, ecdf = _scene_edges(scene) | |
| grad_verts = torch.zeros(scene.n_verts, 3, device=scene.device) | |
| cam_t = camera.tensor(grad_image.shape[0], grad_image.shape[1], | |
| device=scene.device) | |
| ops.pt_geometry_grad(scene.tris, scene.mat_ids, scene.uvs, scene.nodes_f, | |
| scene.nodes_i, scene.light_faces, scene.light_cdf, | |
| float(scene.total_light_area), | |
| scene._flat_albedo().detach().contiguous(), | |
| scene.tex_hdr, | |
| scene._flat_emission().detach().contiguous(), | |
| scene.emi_hdr, scene.mat_type, scene.mat_rough, | |
| scene.mat_ior, scene.face_verts, edges, ecdf, cam_t, | |
| int(spp), int(edge_samples), int(seed), | |
| grad_image.detach().contiguous(), grad_verts) | |
| return grad_verts | |
| def render(scene, camera, height, width, spp=64, max_bounces=4, | |
| estimator=None, nee=None, seed=0): | |
| """Render a linear-radiance image [H, W, 3] float32, differentiable with | |
| respect to the scene's albedo texels, emission, and environment texels. | |
| estimator: "mis" (default), "nee", or "brdf"; the legacy `nee` bool maps | |
| True->"nee", False->"brdf". A fixed seed renders the same paths every | |
| call, so the Monte Carlo objective is a deterministic function of the | |
| parameters.""" | |
| if not (1 <= max_bounces <= MAX_BOUNCES): | |
| raise ValueError(f"max_bounces must be in [1, {MAX_BOUNCES}]") | |
| if estimator is not None and nee is not None: | |
| raise ValueError("pass estimator or nee, not both") | |
| if estimator is None: | |
| estimator = "mis" if nee is None else ("nee" if nee else "brdf") | |
| if estimator not in _MODES: | |
| raise ValueError('estimator must be "mis", "nee", or "brdf"') | |
| if estimator == "nee" and scene.env is not None: | |
| raise ValueError('environment maps require the "mis" or "brdf" estimator') | |
| cam_t = camera.tensor(height, width, device=scene.device) | |
| if scene.med_sa is not None: | |
| med_sa, med_ss = scene.med_sa, scene.med_ss | |
| else: | |
| med_sa = torch.zeros(0, device=scene.device) | |
| med_ss = torch.zeros(0, device=scene.device) | |
| return _RenderFn.apply(scene._flat_albedo(), scene._flat_emission(), | |
| scene._flat_env(), med_sa, med_ss, scene, cam_t, | |
| height, width, int(spp), int(max_bounces), | |
| _MODES[estimator], int(seed)) | |