Kernels:

---

phanerozoic

/

pathtracer-diff

like 0

"""Hero video: THE DRAGON'S HALL - one scene, every gradient, live.

A single continuous shot in the style of the old 3D benchmarks: a dark
marble hall, a pedestal with gold trim and a glass slab, glowing runes on
the back wall, warm braziers, atmospheric fog, and a faintly luminous
crystal dragon. From a flat gray start, Adam through render() recovers
simultaneously:
  - the 128x128 per-texel floor albedo (49,152 unknowns),
  - the gold trim's conductor tint and the pedestal's plastic albedo,
  - the dragon's emission glow and the braziers' warm emission,
  - a 16x64 per-texel emissive rune band (3,072 unknowns),
  - the participating medium (sigma_a, sigma_s) filling the hall,
and, once the appearance has settled, geometry_grad() alone slides the
displaced dragon (5,205 vertices) onto its pedestal by its shadow and
silhouette. The optimization runs at spp 8; every displayed frame is
re-rendered clean from the current parameters. Layout: one widescreen
live view with a small fixed TARGET inset and a progress bar.

Dragon mesh: the Stanford dragon, courtesy of the Stanford Computer
Graphics Laboratory (dragon_vrip_res4).

Output: media/inverse.mp4 (1024x620, 30 fps) + media/inverse.gif.
"""
import math
import os
import sys

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)

import imageio
import numpy as np
import torch
from PIL import Image, ImageDraw, ImageFont

import load_local

ptd = load_local.load()

DEV = "cuda"
RW, RH = 1024, 576
HEADER = 44
CW, CH = RW, RH + HEADER
FPS = 30
INSET_W, INSET_H = 256, 144

_FONT_CANDIDATES = (
    "DejaVuSansMono-Bold.ttf",
    "/usr/share/fonts/truetype/dejavu/DejaVuSansMono-Bold.ttf",
    "consolab.ttf",
)


def _font(size):
    """First available bold monospace face, else PIL's default."""
    for path in _FONT_CANDIDATES:
        try:
            return ImageFont.truetype(path, size)
        except OSError:
            continue
    return ImageFont.load_default()


FONT, FONT_BIG, FONT_MID, FONT_SM = (_font(20), _font(44), _font(25),
                                     _font(16))

MP4 = os.path.join(ROOT, "media", "inverse.mp4")
GIF = os.path.join(ROOT, "media", "inverse.gif")
vw = imageio.get_writer(MP4, fps=FPS, codec="libx264", quality=8,
                        pixelformat="yuv420p", macro_block_size=1)
gif_entries = []  # (PIL frame, seconds, static?)


def emit(fr, seconds, static=False):
    arr = np.asarray(fr, dtype=np.uint8)
    for _ in range(max(1, round(seconds * FPS))):
        vw.append_data(arr)
    gif_entries.append((fr, seconds, static))


def quad(a, b, c, d):
    return [[a, b, c], [a, c, d]]


def add_box(verts, faces, mat, uv, lo, hi, mid):
    x0, y0, z0 = lo
    x1, y1, z1 = hi
    b = len(verts)
    verts.extend([(x0, y0, z0), (x1, y0, z0), (x1, y0, z1), (x0, y0, z1),
                  (x0, y1, z0), (x1, y1, z0), (x1, y1, z1), (x0, y1, z1)])
    for q in [quad(b, b + 1, b + 2, b + 3), quad(b + 4, b + 7, b + 6, b + 5),
              quad(b, b + 4, b + 5, b + 1), quad(b + 3, b + 2, b + 6, b + 7),
              quad(b, b + 3, b + 7, b + 4), quad(b + 1, b + 5, b + 6, b + 2)]:
        faces.extend(q)
        mat.extend([mid, mid])
        uv.extend([[(0, 0)] * 3] * 2)


def load_dragon(path, height=1.5, yaw_deg=205.0):
    """Parse the ASCII PLY, center on xz, base at y=0, scale to `height`,
    yaw around y."""
    verts, faces = [], []
    with open(path, "r", encoding="ascii") as f:
        n_v = n_f = 0
        props = 0
        for line in f:
            line = line.strip()
            if line.startswith("element vertex"):
                n_v = int(line.split()[-1])
            elif line.startswith("element face"):
                n_f = int(line.split()[-1])
            elif line.startswith("property float") or \
                    line.startswith("property double"):
                props += 1
            elif line == "end_header":
                break
        for _ in range(n_v):
            xs = f.readline().split()
            verts.append([float(xs[0]), float(xs[1]), float(xs[2])])
        for _ in range(n_f):
            xs = f.readline().split()
            assert xs[0] == "3"
            faces.append([int(xs[1]), int(xs[2]), int(xs[3])])
    v = torch.tensor(verts, dtype=torch.float32)
    fc = torch.tensor(faces, dtype=torch.int64)
    lo, hi = v.amin(0), v.amax(0)
    v = v - torch.tensor([(lo[0] + hi[0]) / 2, lo[1], (lo[2] + hi[2]) / 2])
    v = v * (height / float(hi[1] - lo[1]))
    a = math.radians(yaw_deg)
    ca, sa = math.cos(a), math.sin(a)
    x, y, z = v[:, 0].clone(), v[:, 1], v[:, 2].clone()
    v[:, 0] = ca * x + sa * z
    v[:, 2] = -sa * x + ca * z
    return v, fc


DRAGON_V, DRAGON_F = load_dragon(os.path.join(ROOT, "media",
                                              "dragon_res4.ply"))
print(f"dragon: {DRAGON_V.shape[0]} verts, {DRAGON_F.shape[0]} faces")

DRAGON_BASE_Y = 1.78           # levitating just above the glass slab
DRAGON_HOME = (0.0, -1.5)      # pedestal center (x, z)

# materials: 0 floor (diffuse, 128^2 tex) | 1 walls | 2 pedestal (plastic)
# 3 gold trim | 4 glass slab | 5 dragon (rough dielectric + glow)
# 6 braziers (warm emitters) | 7 rune band (emissive texels) | 8 key light


def build_scene(floor_tex, gold, plast, glow, brazier, runes, sa, ss,
                occ_t):
    verts, faces, mat, uv = [], [], [], []
    b = len(verts)
    verts.extend([(-6, 0, -6), (6, 0, -6), (6, 0, 6), (-6, 0, 6)])
    faces.extend(quad(b, b + 1, b + 2, b + 3))
    mat.extend([0, 0])
    uv.extend([[(0, 0), (1, 0), (1, 1)], [(0, 0), (1, 1), (0, 1)]])
    b = len(verts)  # back wall
    verts.extend([(-6, 0, -6), (6, 0, -6), (6, 7, -6), (-6, 7, -6)])
    faces.extend(quad(b, b + 1, b + 2, b + 3))
    mat.extend([1, 1])
    uv.extend([[(0, 0)] * 3] * 2)
    b = len(verts)  # left wall
    verts.extend([(-6, 0, 6), (-6, 0, -6), (-6, 7, -6), (-6, 7, 6)])
    faces.extend(quad(b, b + 1, b + 2, b + 3))
    mat.extend([1, 1])
    uv.extend([[(0, 0)] * 3] * 2)
    b = len(verts)  # right wall
    verts.extend([(6, 0, -6), (6, 0, 6), (6, 7, 6), (6, 7, -6)])
    faces.extend(quad(b, b + 1, b + 2, b + 3))
    mat.extend([1, 1])
    uv.extend([[(0, 0)] * 3] * 2)
    b = len(verts)  # ceiling
    verts.extend([(-6, 7, -6), (6, 7, -6), (6, 7, 6), (-6, 7, 6)])
    faces.extend(quad(b, b + 1, b + 2, b + 3))
    mat.extend([1, 1])
    uv.extend([[(0, 0)] * 3] * 2)

    add_box(verts, faces, mat, uv, (-1.3, 0.0, -2.8), (1.3, 1.0, -0.2), 2)
    add_box(verts, faces, mat, uv, (-1.45, 1.0, -2.95), (1.45, 1.25, -0.05),
            3)
    add_box(verts, faces, mat, uv, (-1.05, 1.25, -2.55), (1.05, 1.55, -0.45),
            4)

    b = len(verts)  # braziers: warm sconce panels on the side walls
    verts.extend([(-5.98, 1.4, -2.9), (-5.98, 1.4, -2.0),
                  (-5.98, 3.2, -2.0), (-5.98, 3.2, -2.9)])
    faces.extend(quad(b, b + 1, b + 2, b + 3))
    mat.extend([6, 6])
    uv.extend([[(0, 0)] * 3] * 2)
    b = len(verts)
    verts.extend([(5.98, 1.4, -2.0), (5.98, 1.4, -2.9),
                  (5.98, 3.2, -2.9), (5.98, 3.2, -2.0)])
    faces.extend(quad(b, b + 1, b + 2, b + 3))
    mat.extend([6, 6])
    uv.extend([[(0, 0)] * 3] * 2)

    b = len(verts)  # rune band on the back wall
    verts.extend([(-4.5, 3.2, -5.98), (4.5, 3.2, -5.98),
                  (4.5, 4.2, -5.98), (-4.5, 4.2, -5.98)])
    faces.extend(quad(b, b + 1, b + 2, b + 3))
    mat.extend([7, 7])
    uv.extend([[(0, 0), (1, 0), (1, 1)], [(0, 0), (1, 1), (0, 1)]])

    b = len(verts)  # key light
    verts.extend([(-1.6, 6.98, -2.9), (1.6, 6.98, -2.9),
                  (1.6, 6.98, -0.1), (-1.6, 6.98, -0.1)])
    faces.extend(quad(b, b + 1, b + 2, b + 3))
    mat.extend([8, 8])
    uv.extend([[(0, 0)] * 3] * 2)

    occ0 = len(verts)
    dx = DRAGON_HOME[0] + float(occ_t[0])
    dz = DRAGON_HOME[1] + float(occ_t[1])
    dv = DRAGON_V + torch.tensor([dx, DRAGON_BASE_Y, dz])
    verts = torch.cat([torch.tensor(verts, dtype=torch.float32), dv])
    faces = torch.cat([torch.tensor(faces, dtype=torch.int64),
                       DRAGON_F + occ0])
    mat.extend([5] * DRAGON_F.shape[0])
    uv.extend([[(0, 0)] * 3] * DRAGON_F.shape[0])

    scene = ptd.Scene(
        verts, faces, mat,
        albedo=[floor_tex,
                torch.tensor([0.24, 0.22, 0.21], device=DEV),
                plast, gold,
                torch.tensor([1.0, 1.0, 1.0], device=DEV),
                torch.tensor([1.0, 1.0, 1.0], device=DEV),
                torch.tensor([0.55, 0.35, 0.20], device=DEV),
                torch.tensor([0.05, 0.06, 0.09], device=DEV),
                torch.tensor([0.8, 0.8, 0.8], device=DEV)],
        emission=[torch.zeros(3), torch.zeros(3), torch.zeros(3),
                  torch.zeros(3), torch.zeros(3), glow, brazier, runes,
                  torch.tensor([20.0, 17.5, 14.0])],
        material_types=[ptd.DIFFUSE, ptd.DIFFUSE, ptd.PLASTIC,
                        ptd.CONDUCTOR, ptd.DIELECTRIC, ptd.ROUGH_DIELECTRIC,
                        ptd.DIFFUSE, ptd.DIFFUSE, ptd.DIFFUSE],
        roughness=[0.3, 0.3, 0.25, 0.12, 0.0, 0.08, 0.3, 0.3, 0.3],
        ior=[1.5] * 9, uvs=uv, medium=(sa, ss))
    scene.med_sbar = 0.05  # fixed detached rate across the optimization
    return scene, occ0


def marble(n=128):
    y, x = torch.meshgrid(torch.linspace(0, 1, n), torch.linspace(0, 1, n),
                          indexing="ij")
    v1 = torch.sin(9.0 * x + 3.5 * torch.sin(4.0 * y + 1.0)) ** 6
    v2 = torch.sin(7.0 * y + 3.0 * torch.sin(5.0 * x + 2.2) + 1.3) ** 8
    base = torch.tensor([0.055, 0.055, 0.075])
    teal = torch.tensor([0.10, 0.45, 0.50])
    goldv = torch.tensor([0.55, 0.38, 0.12])
    t = base + v1.unsqueeze(-1) * teal + v2.unsqueeze(-1) * goldv
    return t.clamp(0.02, 0.98).to(DEV)


def rune_band(h=16, w=64):
    torch.manual_seed(4)
    e = torch.full((h, w, 3), 0.06)
    e[..., 2] = 0.16
    for k in range(10):
        cx = 4 + k * 6
        g = torch.rand(5)
        if g[0] > 0.3:
            e[3:13, cx:cx + 2] = 0
            e[3:13, cx:cx + 2, 1] = 3.2
            e[3:13, cx:cx + 2, 2] = 4.0
        if g[1] > 0.4:
            e[3:5, cx:cx + 4, 1] = 3.2
            e[3:5, cx:cx + 4, 2] = 4.0
        if g[2] > 0.4:
            e[11:13, cx - 2:cx + 3, 1] = 3.2
            e[11:13, cx - 2:cx + 3, 2] = 4.0
        if g[3] > 0.5:
            e[7:9, cx:cx + 4, 1] = 3.2
            e[7:9, cx:cx + 4, 2] = 4.0
    return e.to(DEV)


cam = ptd.Camera(position=(0.0, 3.5, 9.6), look_at=(0.0, 1.9, -1.6),
                 vfov_deg=33.0)


def tonemap(img):
    x = (img * 1.35).clamp(0, 1) ** (1 / 2.2)
    return (x * 255).byte().cpu().numpy()


target_np_inset = None


def frame_of(live_img, label, progress):
    fr = Image.new("RGB", (CW, CH), (10, 10, 12))
    fr.paste(Image.fromarray(tonemap(live_img)), (0, HEADER))
    d = ImageDraw.Draw(fr)
    d.text((12, 10), label, font=FONT, fill=(120, 235, 160))
    # target inset
    x0 = CW - INSET_W - 14
    y0 = HEADER + 12
    d.rectangle([x0 - 2, y0 - 2, x0 + INSET_W + 1, y0 + INSET_H + 1],
                outline=(230, 230, 235))
    fr.paste(target_np_inset, (x0, y0))
    d.text((x0 + 6, y0 + INSET_H - 22), "TARGET", font=FONT_SM,
           fill=(235, 235, 240))
    # progress bar
    d.rectangle([0, CH - 5, int(CW * progress), CH - 1],
                fill=(90, 220, 140))
    return fr


def title_card(title, lines, seconds=2.2):
    fr = Image.new("RGB", (CW, CH), (10, 10, 12))
    d = ImageDraw.Draw(fr)
    y = CH // 2 - 40 - 20 * len(lines)
    d.text((CW // 2, y), title, font=FONT_BIG, fill=(235, 235, 240),
           anchor="mm")
    for i, ln in enumerate(lines):
        d.text((CW // 2, y + 64 + 34 * i), ln, font=FONT_MID,
               fill=(150, 200, 165), anchor="mm")
    emit(fr, seconds, static=True)


# ---- targets
tex_t = marble()
gold_t = torch.tensor([1.0, 0.71, 0.29], device=DEV)
plast_t = torch.tensor([0.55, 0.06, 0.08], device=DEV)
glow_t = torch.tensor([0.30, 0.55, 0.90], device=DEV)
brazier_t = torch.tensor([14.0, 8.0, 3.0], device=DEV)
runes_t = rune_band()
sa_t = torch.tensor([0.010, 0.014, 0.022], device=DEV)
ss_t = torch.tensor([0.020, 0.017, 0.012], device=DEV)
t_home = torch.zeros(2)

target_scene, _ = build_scene(tex_t, gold_t, plast_t, glow_t, brazier_t,
                              runes_t, sa_t, ss_t, t_home)
target_hi = ptd.render(target_scene, cam, RH, RW, spp=1024, max_bounces=6,
                       seed=9).detach()
target_lo = ptd.render(target_scene, cam, RH, RW, spp=8, max_bounces=6,
                       seed=3).detach()
target_np_inset = Image.fromarray(tonemap(target_hi)).resize(
    (INSET_W, INSET_H), Image.LANCZOS)

title_card("THE DRAGON'S HALL",
           ["pathtracer-diff: one scene, every gradient, live",
            "floor texels - gold - plastic - glass - crystal - emission",
            "runes - fog - and the dragon's position",
            "all optimized through one CUDA kernel"], 3.0)

# ---- parameters (flat gray start; dragon displaced; hall clear of fog).
# The dragon's glow is a STATIC prop: optimizing the emission of an object
# that starts in the wrong place teaches it to imitate the wall behind it.
glow = glow_t
tex = torch.full((128, 128, 3), 0.18, device=DEV, requires_grad=True)
gold = torch.full((3,), 0.35, device=DEV, requires_grad=True)
plast = torch.full((3,), 0.30, device=DEV, requires_grad=True)
brazier = torch.full((3,), 2.0, device=DEV, requires_grad=True)
runes = torch.full((16, 64, 3), 0.30, device=DEV, requires_grad=True)
sa = torch.full((3,), 0.002, device=DEV, requires_grad=True)
ss = torch.full((3,), 0.002, device=DEV, requires_grad=True)
occ_t = torch.tensor([-3.1, 3.4], requires_grad=True)  # front-left floor

opt = torch.optim.Adam([
    {"params": [tex], "lr": 0.08},
    {"params": [gold], "lr": 0.08},
    {"params": [plast], "lr": 0.05},
    {"params": [brazier], "lr": 0.15},
    {"params": [runes], "lr": 0.09},
])
opt_f = torch.optim.Adam([sa, ss], lr=0.0025)
opt_g = torch.optim.Adam([occ_t], lr=0.09)

# one continuous shot, phased couplings: appearance runs throughout,
# geometry joins at 350 in the still-clear hall (a wrong-appearance
# mismatch field repels it), fog at 700 once the textures are trustworthy
# (else extinction becomes the cheap global-darkness knob); the tail
# repairs the texels scarred while the dragon was displaced
ITERS = 1400
FOG_START = 700
GEO_START = 350
APP_FREEZE = 1400
losses = []
for it in range(ITERS):
    scene, occ0 = build_scene(tex, gold, plast, glow, brazier, runes, sa,
                              ss, occ_t.detach())
    opt.zero_grad()
    opt_f.zero_grad()
    img = ptd.render(scene, cam, RH, RW, spp=12, max_bounces=6, seed=3)
    loss = (img - target_lo).square().mean()
    loss.backward()
    if it < APP_FREEZE:
        opt.step()
        if it >= FOG_START:
            opt_f.step()
    if it >= GEO_START:
        dldi = (2.0 / img.numel()) * (img.detach() - target_lo)
        gv = ptd.geometry_grad(scene, cam, dldi, spp=12,
                               edge_samples=1 << 18, seed=17 + it)
        opt_g.zero_grad()
        g = gv[occ0:occ0 + DRAGON_V.shape[0]][:, [0, 2]].sum(dim=0).cpu()
        gn = g.norm()
        if gn > 0.5:
            g = g * (0.5 / gn)
        occ_t.grad = g
        if float(occ_t.detach().norm()) < 0.8:
            for grp in opt_g.param_groups:
                grp["lr"] = 0.04  # settle gently once nearly home
        opt_g.step()
        with torch.no_grad():
            occ_t.clamp_(-4.0, 4.0)
    with torch.no_grad():
        tex.clamp_(0.02, 0.98)
        gold.clamp_(0.02, 1.0)
        plast.clamp_(0.02, 0.98)
        brazier.clamp_(0.0, 40.0)
        runes.clamp_(0.0, 10.0)
        sa.clamp_(1e-4, 0.5)
        ss.clamp_(1e-4, 0.5)
    losses.append(loss.item())
    if it % 8 == 0:
        # duty-cycle breather: softens sustained power transients on the
        # host (the machine hard-reset once under continuous max load)
        import time as _time
        torch.cuda.synchronize()
        _time.sleep(0.25)
    if it % 4 == 0:
        st_err = ((sa + ss).detach() - (sa_t + ss_t)).abs().max().item()
        disp = ptd.render(scene, cam, RH, RW, spp=48, max_bounces=6,
                          seed=9).detach()
        lab = (f"LIVE  iter {it:04d}  loss {loss.item():8.2e}  "
               f"fog err {st_err:6.4f}  dragon off {occ_t.norm():5.2f}")
        emit(frame_of(disp, lab, (it + 1) / ITERS), 1.0 / FPS)

scene, _ = build_scene(tex, gold, plast, glow, brazier, runes,
                       sa.detach(), ss.detach(), occ_t.detach())
final_hi = ptd.render(scene, cam, RH, RW, spp=1024, max_bounces=6,
                      seed=9).detach()
st_err = ((sa + ss).detach() - (sa_t + ss_t)).abs().max().item()
emit(frame_of(final_hi,
              f"LIVE  iter {ITERS}  loss {losses[-1]:8.2e}  "
              f"fog err {st_err:6.4f}  dragon off {occ_t.norm():5.2f}"
              f"  (spp 1024)", 1.0), 3.2)

title_card("pathtracer-diff",
           [f"loss {losses[0]:.1e} -> {losses[-1]:.1e}   "
            f"dragon offset 4.60 -> {occ_t.norm():.3f}",
            "49,152 floor texels + 3,072 rune texels + materials + fog",
            "Stanford dragon: Stanford Computer Graphics Laboratory",
            "hf.co/kernels/phanerozoic/pathtracer-diff"], 3.0)

vw.close()

# gif fallback: 3/4 scale, statics one frame, animation ~7.5 fps
GW, GH = CW * 3 // 4, CH * 3 // 4
frames, durs = [], []
carry = 0.0
idx = 0
for fr, sec, static in gif_entries:
    if static:
        if carry > 0 and frames:
            durs[-1] += int(carry * 1000)
            carry = 0.0
        frames.append(fr.resize((GW, GH), Image.LANCZOS))
        durs.append(int(sec * 1000))
        idx = 0
    else:
        carry += sec
        if idx % 4 == 0:
            frames.append(fr.resize((GW, GH), Image.LANCZOS))
            durs.append(max(60, int(carry * 1000)))
            carry = 0.0
        idx += 1
frames[0].save(GIF, save_all=True, append_images=frames[1:], duration=durs,
               loop=0, optimize=True)

mp4_mb = os.path.getsize(MP4) / 1e6
gif_mb = os.path.getsize(GIF) / 1e6
print(f"loss {losses[0]:.2e} -> {losses[-1]:.2e} in {ITERS} steps")
print(f"floor tex mean err {(tex.detach() - tex_t).abs().mean():.3f}; "
      f"gold max err {(gold.detach() - gold_t).abs().max():.3f}; "
      f"plastic max err {(plast.detach() - plast_t).abs().max():.3f}")
print(f"brazier max err {(brazier.detach() - brazier_t).abs().max():.2f}; "
      f"runes mean err {(runes.detach() - runes_t).abs().mean():.3f}")
print(f"fog sigma_t err {st_err:.4f}; dragon offset residual "
      f"{occ_t.detach().norm():.3f} (start 4.60)")
print(f"wrote {MP4} ({mp4_mb:.1f} MB) and {GIF} ({gif_mb:.1f} MB, "
      f"{len(frames)} frames)")