@meimu/mutracer 0.0.0-stage → 0.1.0

package/lib/renderer/WGPUJNLMDenoiser.js

@@ -0,0 +1,412 @@
+import { getTargetShaderCache } from './WGPUShaderCache.js';
+import shadowCatcherAccumulationWGSL from './wgsl/chunks/shadowCatcherAccumulation.wgsl.js';
+import { JNLM_DEFAULT_COLOR_SENSITIVITY, JNLM_DEFAULT_SHAPE, JNLM_ALBEDO_FLOOR, JNLM_ALBEDO_REGULARIZATION, JNLM_COVERAGE_BANDWIDTH, JNLM_MAX_ITERATIONS, jnlmPassDilation } from './jnlm.js';
+/** WebGPU passes of the JNLM denoiser described in `jnlm.ts`. */
+export const JNLM_WORKGROUP_SIZE = 8;
+const paramsWGSL = /* wgsl */ `
+struct JNLMParams {
+size: vec2u,
+colorSensitivity: f32,
+
+dilation: f32,
+}
+@group(0) @binding(0) var<uniform> params: JNLMParams;
+`;
+/**
+ * Pass 1: the guide mean and the variance of that mean from the accumulated
+ * moments. Ray buffers are stored bottom-up; guides use texture rows.
+ */
+export const jnlmGuideWGSL = /* wgsl */ `
+${paramsWGSL}
+${shadowCatcherAccumulationWGSL}
+@group(0) @binding(1) var<storage, read> accumulation: array<ShadowCatcherAccumulation>;
+@group(0) @binding(2) var<storage, read> normalBuffer: array<vec4f>;
+@group(0) @binding(3) var<storage, read_write> rawGuide: array<vec2f>;
+
+@compute @workgroup_size(${JNLM_WORKGROUP_SIZE}, ${JNLM_WORKGROUP_SIZE}, 1)
+fn main(@builtin(global_invocation_id) id: vec3u) {
+let size = params.size;
+if (id.x >= size.x || id.y >= size.y) { return; }
+let tO = id.x + (size.y - 1u - id.y) * size.x;
+let n = normalBuffer[tO].a;
+let moments = accumulation[tO].moments;
+let mean = moments.x / max(n, 1.0);
+
+var variance = -1.0;
+if (n >= 2.0) {
+variance = max(moments.y - n * mean * mean, 0.0) / (n * (n - 1.0));
+}
+rawGuide[id.x + id.y * size.x] = vec2f(mean, variance);
+}
+`;
+/**
+ * Pass 2: smooth the noisy per-pixel variance over 3 x 3 pixels. Pixels with
+ * a single sample take the spatial variance of their neighborhood instead.
+ */
+export const jnlmVarianceWGSL = /* wgsl */ `
+${paramsWGSL}
+@group(0) @binding(1) var<storage, read> rawGuide: array<vec2f>;
+@group(0) @binding(2) var<storage, read_write> guide: array<vec2f>;
+
+@compute @workgroup_size(${JNLM_WORKGROUP_SIZE}, ${JNLM_WORKGROUP_SIZE}, 1)
+fn main(@builtin(global_invocation_id) id: vec3u) {
+let size = params.size;
+if (id.x >= size.x || id.y >= size.y) { return; }
+let center = rawGuide[id.x + id.y * size.x];
+var lg = 0.0;
+var iX = 0.0;
+var mL = 0.0;
+var iR = 0.0;
+for (var dy = -1; dy <= 1; dy++) {
+for (var dx = -1; dx <= 1; dx++) {
+let q = clamp(vec2i(id.xy) + vec2i(dx, dy), vec2i(0), vec2i(size) - 1);
+let value = rawGuide[u32(q.x) + u32(q.y) * size.x];
+mL += value.x;
+iR += value.x * value.x;
+if (value.y >= 0.0) {
+lg += value.y;
+iX += 1.0;
+}
+}
+}
+var variance: f32;
+if (center.y >= 0.0) {
+variance = lg / iX;
+} else {
+let lf = mL / 9.0;
+variance = max(iR / 9.0 - lf * lf, 0.0) * 9.0 / 8.0;
+}
+guide[id.x + id.y * size.x] = vec2f(center.x, variance);
+}
+`;
+/** Copies the composed image into the filter's color buffer. */
+export const jnlmLoadColorWGSL = /* wgsl */ `
+${paramsWGSL}
+@group(0) @binding(1) var colorTex: texture_2d<f32>;
+@group(0) @binding(2) var<storage, read_write> outColor: array<vec4f>;
+
+@compute @workgroup_size(${JNLM_WORKGROUP_SIZE}, ${JNLM_WORKGROUP_SIZE}, 1)
+fn main(@builtin(global_invocation_id) id: vec3u) {
+if (id.x >= params.size.x || id.y >= params.size.y) { return; }
+outColor[id.x + id.y * params.size.x] = textureLoad(colorTex, vec2i(id.xy), 0);
+}
+`;
+/**
+ * Pass 3: the joint NL-means filter of the composed image. It also filters
+ * the guide with the same weights, with the variance of the weighted mean,
+ * for a following pass. The albedo regression corrects color only, keeping
+ * uncertain fitted texture responses out of the matching guide.
+ */
+export function createJNLMFilterWGSL({ searchRadius, patchRadius, albedoBandwidth, normalBandwidth } = JNLM_DEFAULT_SHAPE) {
+    const patchSize = (2 * patchRadius + 1) ** 2;
+    return /* wgsl */ `
+${paramsWGSL}
+@group(0) @binding(1) var<storage, read> guide: array<vec2f>;
+@group(0) @binding(2) var<storage, read> color: array<vec4f>;
+@group(0) @binding(3) var albedoTex: texture_2d<f32>;
+@group(0) @binding(4) var normalTex: texture_2d<f32>;
+@group(0) @binding(5) var<storage, read_write> outColor: array<vec4f>;
+@group(0) @binding(6) var<storage, read_write> outGuide: array<vec2f>;
+@group(0) @binding(7) var shadowCatcherTex: texture_2d<f32>;
+
+const SEARCH_RADIUS = ${searchRadius};
+const PATCH_RADIUS = ${patchRadius};
+const PATCH_SIZE = ${patchSize};
+
+const VARIANCE_CANCELLATION = 1.0;
+
+const INV_ALBEDO_BANDWIDTH2 = ${(1 / albedoBandwidth ** 2).toFixed(6)};
+const INV_COVERAGE_BANDWIDTH2 = ${(1 / JNLM_COVERAGE_BANDWIDTH ** 2).toFixed(6)};
+const INV_NORMAL_BANDWIDTH2 = ${(1 / normalBandwidth ** 2).toFixed(6)};
+
+
+const CENTER_WEIGHT_MIN = 0.01;
+
+fn gZ(p: vec2i) -> u32 {
+let q = clamp(p, vec2i(0), vec2i(params.size) - 1);
+return u32(q.x) + u32(q.y) * params.size.x;
+}
+
+fn lk(q: vec2i, bs: vec4f, gM: vec3f) -> f32 {
+let albedo = textureLoad(albedoTex, q, 0);
+
+
+let oS = max((albedo.rgb + bs.rgb) * 0.5, vec3f(${JNLM_ALBEDO_FLOOR}));
+let kR = (albedo.rgb - bs.rgb) / oS;
+let iO = albedo.a - bs.a;
+let lb = textureLoad(normalTex, q, 0).rgb - gM;
+return max(
+dot(kR, kR) * INV_ALBEDO_BANDWIDTH2
++ iO * iO * INV_COVERAGE_BANDWIDTH2,
+dot(lb, lb) * INV_NORMAL_BANDWIDTH2
+);
+}
+
+fn premultiplied(c: vec4f) -> vec4f {
+return vec4f(c.rgb * c.a, c.a);
+}
+
+@compute @workgroup_size(${JNLM_WORKGROUP_SIZE}, ${JNLM_WORKGROUP_SIZE}, 1)
+fn main(@builtin(global_invocation_id) id: vec3u) {
+let size = vec2i(params.size);
+let p = vec2i(id.xy);
+if (p.x >= size.x || p.y >= size.y) { return; }
+let step = i32(params.dilation);
+let bs = textureLoad(albedoTex, p, 0);
+let fw = textureLoad(normalTex, p, 0);
+let gM = fw.rgb;
+let index = gZ(p);
+
+
+if (fw.a <= 0.0 && textureLoad(shadowCatcherTex, p, 0).r <= 0.0) {
+outColor[index] = color[index];
+outGuide[index] = guide[index];
+return;
+}
+
+var kT: array<vec2f, PATCH_SIZE>;
+var k = 0;
+for (var dy = -PATCH_RADIUS; dy <= PATCH_RADIUS; dy++) {
+for (var dx = -PATCH_RADIUS; dx <= PATCH_RADIUS; dx++) {
+kT[k] = guide[gZ(p + vec2i(dx, dy) * step)];
+k++;
+}
+}
+let k2 = params.colorSensitivity * params.colorSensitivity;
+
+var sum = vec4f(0.0);
+let gP = max(bs.rgb, vec3f(${JNLM_ALBEDO_FLOOR}));
+var mb = vec3f(0.0);
+var gs = vec3f(0.0);
+var hn = vec3f(0.0);
+var pg = 0.0;
+var gu = 0.0;
+var nH = 0.0;
+var kA = 0.0;
+for (var sy = -SEARCH_RADIUS; sy <= SEARCH_RADIUS; sy++) {
+for (var sx = -SEARCH_RADIUS; sx <= SEARCH_RADIUS; sx++) {
+let q = p + vec2i(sx, sy) * step;
+if (
+(sx == 0 && sy == 0) ||
+q.x < 0 || q.y < 0 || q.x >= size.x || q.y >= size.y
+) { continue; }
+
+let features = lk(q, bs, gM);
+if (features > 16.0) { continue; }
+
+var fT = 0.0;
+k = 0;
+for (var dy = -PATCH_RADIUS; dy <= PATCH_RADIUS; dy++) {
+for (var dx = -PATCH_RADIUS; dx <= PATCH_RADIUS; dx++) {
+let a = kT[k];
+let b = guide[gZ(q + vec2i(dx, dy) * step)];
+let qH = a.x - b.x;
+fT += (qH * qH
+- VARIANCE_CANCELLATION * (a.y + min(a.y, b.y)))
+/ (1e-10 + k2 * (a.y + b.y));
+k++;
+}
+}
+fT /= f32(PATCH_SIZE);
+
+let weight = exp(-max(max(fT, 0.0), features));
+let index = gZ(q);
+let c = color[index];
+let g = guide[index];
+
+let a = (textureLoad(albedoTex, q, 0).rgb - bs.rgb) / gP;
+let value = premultiplied(c);
+sum += value * weight;
+mb += a * weight;
+gs += a * a * weight;
+hn += a * value.rgb * weight;
+pg += g.x * weight;
+gu += g.y * weight * weight;
+nH += weight;
+kA = max(kA, weight);
+}
+}
+
+
+
+let center = color[index];
+let kS = guide[index];
+let dF = max(kA, CENTER_WEIGHT_MIN);
+let totalWeight = nH + dF;
+let tA = premultiplied(center);
+let mean = (sum + tA * dF) / totalWeight;
+
+
+let ez = mb / totalWeight;
+let ll = max(gs / totalWeight - ez * ez, vec3f(0.0));
+let qi = hn / totalWeight - ez * mean.rgb;
+
+
+
+let qs = bs.rgb + ez * gP;
+let ol = mean.rgb * gP / max(qs, vec3f(${JNLM_ALBEDO_FLOOR}));
+let vk = clamp(qi / (ll + vec3f(${JNLM_ALBEDO_REGULARIZATION})), vec3f(0.0), ol);
+let rn = mean.rgb - vk * ez;
+let filtered = vec4f(max(rn, vec3f(0.0)), max(mean.a, 0.0));
+let rgb = select(
+center.rgb,
+filtered.rgb / max(filtered.a, 1e-6),
+filtered.a > 1e-6
+);
+outColor[index] = vec4f(rgb, filtered.a);
+outGuide[index] = vec2f(
+(pg + kS.x * dF) / totalWeight,
+(gu + kS.y * dF * dF)
+/ (totalWeight * totalWeight)
+);
+}
+`;
+}
+const MAX_ITERATIONS = JNLM_MAX_ITERATIONS;
+export class WGPUJNLMDenoisePass {
+    _device;
+    _paramsBuffers = [];
+    _guidePipeline;
+    _variancePipeline;
+    _loadColorPipeline;
+    _filterPipeline;
+    _filterShapeKey = '';
+    _guides = [];
+    _colors = [];
+    _width = 0;
+    _height = 0;
+    constructor(device) {
+        this._device = device;
+        // One uniform buffer per pass: every pass of an encoder sees the last
+        // writeBuffer() of a shared one.
+        for (let i = 0; i <= MAX_ITERATIONS; i++) {
+            this._paramsBuffers.push(device.createBuffer({
+                label: 'JNLM params',
+                size: 16,
+                usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST
+            }));
+        }
+        this._guidePipeline = this._createPipeline('JNLM guide', jnlmGuideWGSL);
+        this._variancePipeline = this._createPipeline('JNLM variance', jnlmVarianceWGSL);
+        this._loadColorPipeline = this._createPipeline('JNLM load color', jnlmLoadColorWGSL);
+    }
+    _createPipeline(label, code) {
+        return getTargetShaderCache(this._device).createComputePipeline({
+            label,
+            layout: 'auto',
+            compute: {
+                module: getTargetShaderCache(this._device).createShaderModule({ label, code }),
+                entryPoint: 'main'
+            }
+        });
+    }
+    _getFilterPipeline(shape) {
+        const key = JSON.stringify(shape);
+        if (!this._filterPipeline || key !== this._filterShapeKey) {
+            this._filterPipeline = this._createPipeline('JNLM filter', createJNLMFilterWGSL(shape));
+            this._filterShapeKey = key;
+        }
+        return this._filterPipeline;
+    }
+    reserveCapacity() { }
+    _ensureBuffers(width, height) {
+        if (width <= this._width && height <= this._height && this._colors.length)
+            return;
+        width = Math.max(width, this._width);
+        height = Math.max(height, this._height);
+        this._destroyBuffers();
+        const pixels = width * height;
+        const createBuffer = (label, bytes) => this._device.createBuffer({
+            label,
+            size: pixels * bytes,
+            usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC
+        });
+        this._guides = [createBuffer('JNLM guide A', 8), createBuffer('JNLM guide B', 8)];
+        this._colors = [createBuffer('JNLM color A', 16), createBuffer('JNLM color B', 16)];
+        this._width = width;
+        this._height = height;
+    }
+    _destroyBuffers() {
+        for (const buffer of [...this._guides, ...this._colors])
+            buffer.destroy();
+        this._guides = [];
+        this._colors = [];
+    }
+    /** Encodes the denoiser and returns the RGBA32F output in texture rows. */
+    encode(encoder, inputs) {
+        const { width, height, shape } = inputs;
+        this._ensureBuffers(width, height);
+        const iterations = Math.min(Math.max(Math.round(shape.iterations), 1), MAX_ITERATIONS);
+        const sensitivity = JNLM_DEFAULT_COLOR_SENSITIVITY * Math.max(inputs.strength, 1e-3);
+        this._paramsBuffers.forEach((buffer, i) => {
+            const params = new ArrayBuffer(16);
+            new Uint32Array(params, 0, 2).set([width, height]);
+            // Buffer i (from 1) is filter pass i - 1; index 0 is the prepasses.
+            new Float32Array(params, 8, 2).set([sensitivity, jnlmPassDilation(Math.max(i - 1, 0))]);
+            this._device.queue.writeBuffer(buffer, 0, params);
+        });
+        const device = this._device;
+        const filterPipeline = this._getFilterPipeline(shape);
+        const bind = (pipeline, paramsBuffer, resources) => device.createBindGroup({
+            layout: pipeline.getBindGroupLayout(0),
+            entries: [
+                { binding: 0, resource: { buffer: paramsBuffer } },
+                ...resources.map((resource, i) => ({ binding: i + 1, resource }))
+            ]
+        });
+        const [guideA, guideB] = this._guides;
+        const [colorA, colorB] = this._colors;
+        const prepass = this._paramsBuffers[0];
+        const passes = [
+            [this._guidePipeline, bind(this._guidePipeline, prepass, [
+                    { buffer: inputs.accumulation },
+                    { buffer: inputs.normalAccumulation },
+                    { buffer: guideB }
+                ])],
+            [this._variancePipeline, bind(this._variancePipeline, prepass, [
+                    { buffer: guideB },
+                    { buffer: guideA }
+                ])],
+            [this._loadColorPipeline, bind(this._loadColorPipeline, prepass, [
+                    inputs.color.createView(),
+                    { buffer: colorA }
+                ])]
+        ];
+        const albedoView = inputs.albedo.createView();
+        const normalView = inputs.normal.createView();
+        const shadowCatcherView = inputs.shadowCatcher.createView();
+        let guides = [guideA, guideB];
+        let colors = [colorA, colorB];
+        for (let i = 1; i <= iterations; i++) {
+            passes.push([filterPipeline, bind(filterPipeline, this._paramsBuffers[i], [
+                    { buffer: guides[0] },
+                    { buffer: colors[0] },
+                    albedoView,
+                    normalView,
+                    { buffer: colors[1] },
+                    { buffer: guides[1] },
+                    shadowCatcherView
+                ])]);
+            guides = [guides[1], guides[0]];
+            colors = [colors[1], colors[0]];
+        }
+        const pass = encoder.beginComputePass({ label: 'JNLM denoise' });
+        for (const [pipeline, bindGroup] of passes) {
+            pass.setPipeline(pipeline);
+            pass.setBindGroup(0, bindGroup);
+            pass.dispatchWorkgroups(Math.ceil(width / JNLM_WORKGROUP_SIZE), Math.ceil(height / JNLM_WORKGROUP_SIZE));
+        }
+        pass.end();
+        return colors[0];
+    }
+    getMemoryBytes() {
+        return [...this._paramsBuffers, ...this._guides, ...this._colors]
+            .reduce((sum, buffer) => sum + buffer.size, 0);
+    }
+    dispose() {
+        for (const buffer of this._paramsBuffers)
+            buffer.destroy();
+        this._destroyBuffers();
+    }
+}
+//# sourceMappingURL=WGPUJNLMDenoiser.js.map

package/lib/renderer/WGPUJNLMDenoiser.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"WGPUJNLMDenoiser.js","sourceRoot":"","sources":["../../src/renderer/WGPUJNLMDenoiser.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,oBAAoB,EAAE,MAAM,sBAAsB,CAAC;AAC5D,OAAO,6BAA6B,MAAM,iDAAiD,CAAC;AAE5F,OAAO,EACL,8BAA8B,EAC9B,kBAAkB,EAClB,iBAAiB,EACjB,0BAA0B,EAC1B,uBAAuB,EACvB,mBAAmB,EACnB,gBAAgB,EAEjB,MAAM,WAAW,CAAC;AAEnB,iEAAiE;AAEjE,MAAM,CAAC,MAAM,mBAAmB,GAAG,CAAC,CAAC;AAErC,MAAM,UAAU,GAAG,UAAU,CAAC;;;;;;;;CAQ7B,CAAC;AAEF;;;GAGG;AACH,MAAM,CAAC,MAAM,aAAa,GAAG,UAAU,CAAC;EACtC,UAAU;EACV,6BAA6B;;;;;2BAKJ,mBAAmB,KAAK,mBAAmB;;;;;;;;;;;;;;;CAerE,CAAC;AAEF;;;GAGG;AACH,MAAM,CAAC,MAAM,gBAAgB,GAAG,UAAU,CAAC;EACzC,UAAU;;;;2BAIe,mBAAmB,KAAK,mBAAmB;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;CA8BrE,CAAC;AAEF,gEAAgE;AAChE,MAAM,CAAC,MAAM,iBAAiB,GAAG,UAAU,CAAC;EAC1C,UAAU;;;;2BAIe,mBAAmB,KAAK,mBAAmB;;;;;CAKrE,CAAC;AAEF;;;;;GAKG;AACH,MAAM,UAAU,oBAAoB,CAAC,EACnC,YAAY,EAAE,WAAW,EAAE,eAAe,EAAE,eAAe,KACxC,kBAAkB;IACrC,MAAM,SAAS,GAAG,CAAC,CAAC,GAAG,WAAW,GAAG,CAAC,CAAC,IAAI,CAAC,CAAC;IAC7C,OAAO,UAAU,CAAC;EAClB,UAAU;;;;;;;;;wBASY,YAAY;uBACb,WAAW;qBACb,SAAS;;;;gCAIE,CAAC,CAAC,GAAG,eAAe,IAAI,CAAC,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC;kCACnC,CAAC,CAAC,GAAG,uBAAuB,IAAI,CAAC,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC;gCAC/C,CAAC,CAAC,GAAG,eAAe,IAAI,CAAC,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC;;;;;;;;;;;;;;kDAcE,iBAAiB;;;;;;;;;;;;;;;2BAe7D,mBAAmB,KAAK,mBAAmB;;;;;;;;;;;;;;;;;;;;;;;;;;;;;6BA6BnB,iBAAiB;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;yCAqEG,iBAAiB;kCAC5B,0BAA0B;;;;;;;;;;;;;;;CAerF,CAAC;AACF,CAAC;AAmBD,MAAM,cAAc,GAAG,mBAAmB,CAAC;AAE3C,MAAM,OAAO,mBAAmB;IACtB,OAAO,CAAY;IACnB,cAAc,GAAgB,EAAE,CAAC;IACjC,cAAc,CAAqB;IACnC,iBAAiB,CAAqB;IACtC,kBAAkB,CAAqB;IACvC,eAAe,CAAsB;IACrC,eAAe,GAAG,EAAE,CAAC;IACrB,OAAO,GAAgB,EAAE,CAAC;IAC1B,OAAO,GAAgB,EAAE,CAAC;IAC1B,MAAM,GAAG,CAAC,CAAC;IACX,OAAO,GAAG,CAAC,CAAC;IAEpB,YAAY,MAAiB;QAC3B,IAAI,CAAC,OAAO,GAAG,MAAM,CAAC;QACtB,sEAAsE;QACtE,iCAAiC;QACjC,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,IAAI,cAAc,EAAE,CAAC,EAAE,EAAE,CAAC;YACzC,IAAI,CAAC,cAAc,CAAC,IAAI,CAAC,MAAM,CAAC,YAAY,CAAC;gBAC3C,KAAK,EAAE,aAAa;gBACpB,IAAI,EAAE,EAAE;gBACR,KAAK,EAAE,cAAc,CAAC,OAAO,GAAG,cAAc,CAAC,QAAQ;aACxD,CAAC,CAAC,CAAC;QACN,CAAC;QACD,IAAI,CAAC,cAAc,GAAG,IAAI,CAAC,eAAe,CAAC,YAAY,EAAE,aAAa,CAAC,CAAC;QACxE,IAAI,CAAC,iBAAiB,GAAG,IAAI,CAAC,eAAe,CAAC,eAAe,EAAE,gBAAgB,CAAC,CAAC;QACjF,IAAI,CAAC,kBAAkB,GAAG,IAAI,CAAC,eAAe,CAAC,iBAAiB,EAAE,iBAAiB,CAAC,CAAC;IACvF,CAAC;IAEO,eAAe,CAAC,KAAa,EAAE,IAAY;QACjD,OAAO,oBAAoB,CAAC,IAAI,CAAC,OAAO,CAAC,CAAC,qBAAqB,CAAC;YAC9D,KAAK;YACL,MAAM,EAAE,MAAM;YACd,OAAO,EAAE;gBACP,MAAM,EAAE,oBAAoB,CAAC,IAAI,CAAC,OAAO,CAAC,CAAC,kBAAkB,CAAC,EAAE,KAAK,EAAE,IAAI,EAAE,CAAC;gBAC9E,UAAU,EAAE,MAAM;aACnB;SACF,CAAC,CAAC;IACL,CAAC;IAEO,kBAAkB,CAAC,KAAsB;QAC/C,MAAM,GAAG,GAAG,IAAI,CAAC,SAAS,CAAC,KAAK,CAAC,CAAC;QAClC,IAAI,CAAC,IAAI,CAAC,eAAe,IAAI,GAAG,KAAK,IAAI,CAAC,eAAe,EAAE,CAAC;YAC1D,IAAI,CAAC,eAAe,GAAG,IAAI,CAAC,eAAe,CAAC,aAAa,EAAE,oBAAoB,CAAC,KAAK,CAAC,CAAC,CAAC;YACxF,IAAI,CAAC,eAAe,GAAG,GAAG,CAAC;QAC7B,CAAC;QACD,OAAO,IAAI,CAAC,eAAe,CAAC;IAC9B,CAAC;IAED,eAAe,KAAI,CAAC;IACZ,cAAc,CAAC,KAAa,EAAE,MAAc;QAClD,IAAI,KAAK,IAAI,IAAI,CAAC,MAAM,IAAI,MAAM,IAAI,IAAI,CAAC,OAAO,IAAI,IAAI,CAAC,OAAO,CAAC,MAAM;YAAE,OAAO;QAClF,KAAK,GAAG,IAAI,CAAC,GAAG,CAAC,KAAK,EAAE,IAAI,CAAC,MAAM,CAAC,CAAC;QAAC,MAAM,GAAG,IAAI,CAAC,GAAG,CAAC,MAAM,EAAE,IAAI,CAAC,OAAO,CAAC,CAAC;QAC9E,IAAI,CAAC,eAAe,EAAE,CAAC;QACvB,MAAM,MAAM,GAAG,KAAK,GAAG,MAAM,CAAC;QAC9B,MAAM,YAAY,GAAG,CAAC,KAAa,EAAE,KAAa,EAAE,EAAE,CACpD,IAAI,CAAC,OAAO,CAAC,YAAY,CAAC;YACxB,KAAK;YACL,IAAI,EAAE,MAAM,GAAG,KAAK;YACpB,KAAK,EAAE,cAAc,CAAC,OAAO,GAAG,cAAc,CAAC,QAAQ;SACxD,CAAC,CAAC;QACL,IAAI,CAAC,OAAO,GAAG,CAAC,YAAY,CAAC,cAAc,EAAE,CAAC,CAAC,EAAE,YAAY,CAAC,cAAc,EAAE,CAAC,CAAC,CAAC,CAAC;QAClF,IAAI,CAAC,OAAO,GAAG,CAAC,YAAY,CAAC,cAAc,EAAE,EAAE,CAAC,EAAE,YAAY,CAAC,cAAc,EAAE,EAAE,CAAC,CAAC,CAAC;QACpF,IAAI,CAAC,MAAM,GAAG,KAAK,CAAC;QACpB,IAAI,CAAC,OAAO,GAAG,MAAM,CAAC;IACxB,CAAC;IAEO,eAAe;QACrB,KAAK,MAAM,MAAM,IAAI,CAAC,GAAG,IAAI,CAAC,OAAO,EAAE,GAAG,IAAI,CAAC,OAAO,CAAC;YAAE,MAAM,CAAC,OAAO,EAAE,CAAC;QAC1E,IAAI,CAAC,OAAO,GAAG,EAAE,CAAC;QAClB,IAAI,CAAC,OAAO,GAAG,EAAE,CAAC;IACpB,CAAC;IAED,2EAA2E;IAC3E,MAAM,CAAC,OAA0B,EAAE,MAAkB;QACnD,MAAM,EAAE,KAAK,EAAE,MAAM,EAAE,KAAK,EAAE,GAAG,MAAM,CAAC;QACxC,IAAI,CAAC,cAAc,CAAC,KAAK,EAAE,MAAM,CAAC,CAAC;QACnC,MAAM,UAAU,GAAG,IAAI,CAAC,GAAG,CAAC,IAAI,CAAC,GAAG,CAAC,IAAI,CAAC,KAAK,CAAC,KAAK,CAAC,UAAU,CAAC,EAAE,CAAC,CAAC,EAAE,cAAc,CAAC,CAAC;QACvF,MAAM,WAAW,GAAG,8BAA8B,GAAG,IAAI,CAAC,GAAG,CAAC,MAAM,CAAC,QAAQ,EAAE,IAAI,CAAC,CAAC;QACrF,IAAI,CAAC,cAAc,CAAC,OAAO,CAAC,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE;YACxC,MAAM,MAAM,GAAG,IAAI,WAAW,CAAC,EAAE,CAAC,CAAC;YACnC,IAAI,WAAW,CAAC,MAAM,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,GAAG,CAAC,CAAC,KAAK,EAAE,MAAM,CAAC,CAAC,CAAC;YACnD,oEAAoE;YACpE,IAAI,YAAY,CAAC,MAAM,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,GAAG,CAAC,CAAC,WAAW,EAAE,gBAAgB,CAAC,IAAI,CAAC,GAAG,CAAC,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;YACxF,IAAI,CAAC,OAAO,CAAC,KAAK,CAAC,WAAW,CAAC,MAAM,EAAE,CAAC,EAAE,MAAM,CAAC,CAAC;QACpD,CAAC,CAAC,CAAC;QAEH,MAAM,MAAM,GAAG,IAAI,CAAC,OAAO,CAAC;QAC5B,MAAM,cAAc,GAAG,IAAI,CAAC,kBAAkB,CAAC,KAAK,CAAC,CAAC;QACtD,MAAM,IAAI,GAAG,CACX,QAA4B,EAC5B,YAAuB,EACvB,SAA+B,EAC/B,EAAE,CACF,MAAM,CAAC,eAAe,CAAC;YACrB,MAAM,EAAE,QAAQ,CAAC,kBAAkB,CAAC,CAAC,CAAC;YACtC,OAAO,EAAE;gBACP,EAAE,OAAO,EAAE,CAAC,EAAE,QAAQ,EAAE,EAAE,MAAM,EAAE,YAAY,EAAE,EAAE;gBAClD,GAAG,SAAS,CAAC,GAAG,CAAC,CAAC,QAAQ,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,EAAE,OAAO,EAAE,CAAC,GAAG,CAAC,EAAE,QAAQ,EAAE,CAAC,CAAC;aAClE;SACF,CAAC,CAAC;QACL,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,GAAG,IAAI,CAAC,OAAO,CAAC;QACtC,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,GAAG,IAAI,CAAC,OAAO,CAAC;QACtC,MAAM,OAAO,GAAG,IAAI,CAAC,cAAc,CAAC,CAAC,CAAC,CAAC;QACvC,MAAM,MAAM,GAAyC;YACnD,CAAC,IAAI,CAAC,cAAc,EAAE,IAAI,CAAC,IAAI,CAAC,cAAc,EAAE,OAAO,EAAE;oBACvD,EAAE,MAAM,EAAE,MAAM,CAAC,YAAY,EAAE;oBAC/B,EAAE,MAAM,EAAE,MAAM,CAAC,kBAAkB,EAAE;oBACrC,EAAE,MAAM,EAAE,MAAM,EAAE;iBACnB,CAAC,CAAC;YACH,CAAC,IAAI,CAAC,iBAAiB,EAAE,IAAI,CAAC,IAAI,CAAC,iBAAiB,EAAE,OAAO,EAAE;oBAC7D,EAAE,MAAM,EAAE,MAAM,EAAE;oBAClB,EAAE,MAAM,EAAE,MAAM,EAAE;iBACnB,CAAC,CAAC;YACH,CAAC,IAAI,CAAC,kBAAkB,EAAE,IAAI,CAAC,IAAI,CAAC,kBAAkB,EAAE,OAAO,EAAE;oBAC/D,MAAM,CAAC,KAAK,CAAC,UAAU,EAAE;oBACzB,EAAE,MAAM,EAAE,MAAM,EAAE;iBACnB,CAAC,CAAC;SACJ,CAAC;QACF,MAAM,UAAU,GAAG,MAAM,CAAC,MAAM,CAAC,UAAU,EAAE,CAAC;QAC9C,MAAM,UAAU,GAAG,MAAM,CAAC,MAAM,CAAC,UAAU,EAAE,CAAC;QAC9C,MAAM,iBAAiB,GAAG,MAAM,CAAC,aAAa,CAAC,UAAU,EAAE,CAAC;QAC5D,IAAI,MAAM,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC;QAC9B,IAAI,MAAM,GAAG,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC;QAC9B,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,IAAI,UAAU,EAAE,CAAC,EAAE,EAAE,CAAC;YACrC,MAAM,CAAC,IAAI,CAAC,CAAC,cAAc,EAAE,IAAI,CAAC,cAAc,EAAE,IAAI,CAAC,cAAc,CAAC,CAAC,CAAC,EAAE;oBACxE,EAAE,MAAM,EAAE,MAAM,CAAC,CAAC,CAAC,EAAE;oBACrB,EAAE,MAAM,EAAE,MAAM,CAAC,CAAC,CAAC,EAAE;oBACrB,UAAU;oBACV,UAAU;oBACV,EAAE,MAAM,EAAE,MAAM,CAAC,CAAC,CAAC,EAAE;oBACrB,EAAE,MAAM,EAAE,MAAM,CAAC,CAAC,CAAC,EAAE;oBACrB,iBAAiB;iBAClB,CAAC,CAAC,CAAC,CAAC;YACL,MAAM,GAAG,CAAC,MAAM,CAAC,CAAC,CAAC,EAAE,MAAM,CAAC,CAAC,CAAC,CAAC,CAAC;YAChC,MAAM,GAAG,CAAC,MAAM,CAAC,CAAC,CAAC,EAAE,MAAM,CAAC,CAAC,CAAC,CAAC,CAAC;QAClC,CAAC;QACD,MAAM,IAAI,GAAG,OAAO,CAAC,gBAAgB,CAAC,EAAE,KAAK,EAAE,cAAc,EAAE,CAAC,CAAC;QACjE,KAAK,MAAM,CAAC,QAAQ,EAAE,SAAS,CAAC,IAAI,MAAM,EAAE,CAAC;YAC3C,IAAI,CAAC,WAAW,CAAC,QAAQ,CAAC,CAAC;YAC3B,IAAI,CAAC,YAAY,CAAC,CAAC,EAAE,SAAS,CAAC,CAAC;YAChC,IAAI,CAAC,kBAAkB,CACrB,IAAI,CAAC,IAAI,CAAC,KAAK,GAAG,mBAAmB,CAAC,EACtC,IAAI,CAAC,IAAI,CAAC,MAAM,GAAG,mBAAmB,CAAC,CACxC,CAAC;QACJ,CAAC;QACD,IAAI,CAAC,GAAG,EAAE,CAAC;QACX,OAAO,MAAM,CAAC,CAAC,CAAC,CAAC;IACnB,CAAC;IAED,cAAc;QACZ,OAAO,CAAC,GAAG,IAAI,CAAC,cAAc,EAAE,GAAG,IAAI,CAAC,OAAO,EAAE,GAAG,IAAI,CAAC,OAAO,CAAC;aAC9D,MAAM,CAAC,CAAC,GAAG,EAAE,MAAM,EAAE,EAAE,CAAC,GAAG,GAAG,MAAM,CAAC,IAAI,EAAE,CAAC,CAAC,CAAC;IACnD,CAAC;IAED,OAAO;QACL,KAAK,MAAM,MAAM,IAAI,IAAI,CAAC,cAAc;YAAE,MAAM,CAAC,OAAO,EAAE,CAAC;QAC3D,IAAI,CAAC,eAAe,EAAE,CAAC;IACzB,CAAC;CACF","sourcesContent":["import { getTargetShaderCache } from './WGPUShaderCache.js';\nimport shadowCatcherAccumulationWGSL from './wgsl/chunks/shadowCatcherAccumulation.wgsl.js';\n\nimport {\n  JNLM_DEFAULT_COLOR_SENSITIVITY,\n  JNLM_DEFAULT_SHAPE,\n  JNLM_ALBEDO_FLOOR,\n  JNLM_ALBEDO_REGULARIZATION,\n  JNLM_COVERAGE_BANDWIDTH,\n  JNLM_MAX_ITERATIONS,\n  jnlmPassDilation,\n  type JNLMFilterShape\n} from './jnlm.js';\n\n/** WebGPU passes of the JNLM denoiser described in `jnlm.ts`. */\n\nexport const JNLM_WORKGROUP_SIZE = 8;\n\nconst paramsWGSL = /* wgsl */ `\nstruct JNLMParams {\n  size: vec2u,\n  colorSensitivity: f32,\n  // Pixel spacing of the search window and patches in this pass.\n  dilation: f32,\n}\n@group(0) @binding(0) var<uniform> params: JNLMParams;\n`;\n\n/**\n * Pass 1: the guide mean and the variance of that mean from the accumulated\n * moments. Ray buffers are stored bottom-up; guides use texture rows.\n */\nexport const jnlmGuideWGSL = /* wgsl */ `\n${paramsWGSL}\n${shadowCatcherAccumulationWGSL}\n@group(0) @binding(1) var<storage, read> accumulation: array<ShadowCatcherAccumulation>;\n@group(0) @binding(2) var<storage, read> normalBuffer: array<vec4f>;\n@group(0) @binding(3) var<storage, read_write> rawGuide: array<vec2f>;\n\n@compute @workgroup_size(${JNLM_WORKGROUP_SIZE}, ${JNLM_WORKGROUP_SIZE}, 1)\nfn main(@builtin(global_invocation_id) id: vec3u) {\n  let size = params.size;\n  if (id.x >= size.x || id.y >= size.y) { return; }\n  let rayIndex = id.x + (size.y - 1u - id.y) * size.x;\n  let n = normalBuffer[rayIndex].a;\n  let moments = accumulation[rayIndex].moments;\n  let mean = moments.x / max(n, 1.0);\n  // Variance of the mean, or -1 when one sample cannot estimate it.\n  var variance = -1.0;\n  if (n >= 2.0) {\n    variance = max(moments.y - n * mean * mean, 0.0) / (n * (n - 1.0));\n  }\n  rawGuide[id.x + id.y * size.x] = vec2f(mean, variance);\n}\n`;\n\n/**\n * Pass 2: smooth the noisy per-pixel variance over 3 x 3 pixels. Pixels with\n * a single sample take the spatial variance of their neighborhood instead.\n */\nexport const jnlmVarianceWGSL = /* wgsl */ `\n${paramsWGSL}\n@group(0) @binding(1) var<storage, read> rawGuide: array<vec2f>;\n@group(0) @binding(2) var<storage, read_write> guide: array<vec2f>;\n\n@compute @workgroup_size(${JNLM_WORKGROUP_SIZE}, ${JNLM_WORKGROUP_SIZE}, 1)\nfn main(@builtin(global_invocation_id) id: vec3u) {\n  let size = params.size;\n  if (id.x >= size.x || id.y >= size.y) { return; }\n  let center = rawGuide[id.x + id.y * size.x];\n  var varianceSum = 0.0;\n  var varianceCount = 0.0;\n  var meanSum = 0.0;\n  var meanSquareSum = 0.0;\n  for (var dy = -1; dy <= 1; dy++) {\n    for (var dx = -1; dx <= 1; dx++) {\n      let q = clamp(vec2i(id.xy) + vec2i(dx, dy), vec2i(0), vec2i(size) - 1);\n      let value = rawGuide[u32(q.x) + u32(q.y) * size.x];\n      meanSum += value.x;\n      meanSquareSum += value.x * value.x;\n      if (value.y >= 0.0) {\n        varianceSum += value.y;\n        varianceCount += 1.0;\n      }\n    }\n  }\n  var variance: f32;\n  if (center.y >= 0.0) {\n    variance = varianceSum / varianceCount;\n  } else {\n    let spatialMean = meanSum / 9.0;\n    variance = max(meanSquareSum / 9.0 - spatialMean * spatialMean, 0.0) * 9.0 / 8.0;\n  }\n  guide[id.x + id.y * size.x] = vec2f(center.x, variance);\n}\n`;\n\n/** Copies the composed image into the filter's color buffer. */\nexport const jnlmLoadColorWGSL = /* wgsl */ `\n${paramsWGSL}\n@group(0) @binding(1) var colorTex: texture_2d<f32>;\n@group(0) @binding(2) var<storage, read_write> outColor: array<vec4f>;\n\n@compute @workgroup_size(${JNLM_WORKGROUP_SIZE}, ${JNLM_WORKGROUP_SIZE}, 1)\nfn main(@builtin(global_invocation_id) id: vec3u) {\n  if (id.x >= params.size.x || id.y >= params.size.y) { return; }\n  outColor[id.x + id.y * params.size.x] = textureLoad(colorTex, vec2i(id.xy), 0);\n}\n`;\n\n/**\n * Pass 3: the joint NL-means filter of the composed image. It also filters\n * the guide with the same weights, with the variance of the weighted mean,\n * for a following pass. The albedo regression corrects color only, keeping\n * uncertain fitted texture responses out of the matching guide.\n */\nexport function createJNLMFilterWGSL({\n  searchRadius, patchRadius, albedoBandwidth, normalBandwidth\n}: JNLMFilterShape = JNLM_DEFAULT_SHAPE) {\n  const patchSize = (2 * patchRadius + 1) ** 2;\n  return /* wgsl */ `\n${paramsWGSL}\n@group(0) @binding(1) var<storage, read> guide: array<vec2f>;\n@group(0) @binding(2) var<storage, read> color: array<vec4f>;\n@group(0) @binding(3) var albedoTex: texture_2d<f32>;\n@group(0) @binding(4) var normalTex: texture_2d<f32>;\n@group(0) @binding(5) var<storage, read_write> outColor: array<vec4f>;\n@group(0) @binding(6) var<storage, read_write> outGuide: array<vec2f>;\n@group(0) @binding(7) var shadowCatcherTex: texture_2d<f32>;\n\nconst SEARCH_RADIUS = ${searchRadius};\nconst PATCH_RADIUS = ${patchRadius};\nconst PATCH_SIZE = ${patchSize};\n// Variance cancellation strength alpha of the paper.\nconst VARIANCE_CANCELLATION = 1.0;\n// Inverse squared feature bandwidths. Normals are encoded as n * 0.5 + 0.5.\nconst INV_ALBEDO_BANDWIDTH2 = ${(1 / albedoBandwidth ** 2).toFixed(6)};\nconst INV_COVERAGE_BANDWIDTH2 = ${(1 / JNLM_COVERAGE_BANDWIDTH ** 2).toFixed(6)};\nconst INV_NORMAL_BANDWIDTH2 = ${(1 / normalBandwidth ** 2).toFixed(6)};\n// Lower bound of the center weight, so a pixel unlike all of its neighbors\n// (a thin highlight) is kept rather than replaced by them.\nconst CENTER_WEIGHT_MIN = 0.01;\n\nfn pixelIndex(p: vec2i) -> u32 {\n  let q = clamp(p, vec2i(0), vec2i(params.size) - 1);\n  return u32(q.x) + u32(q.y) * params.size.x;\n}\n\nfn featureDistance(q: vec2i, centerAlbedo: vec4f, centerNormal: vec3f) -> f32 {\n  let albedo = textureLoad(albedoTex, q, 0);\n  // Absolute differences in linear albedo erase low-contrast dark textures.\n  // Measure RGB contrast relative to both pixels; keep coverage absolute.\n  let albedoScale = max((albedo.rgb + centerAlbedo.rgb) * 0.5, vec3f(${JNLM_ALBEDO_FLOOR}));\n  let albedoDelta = (albedo.rgb - centerAlbedo.rgb) / albedoScale;\n  let coverageDelta = albedo.a - centerAlbedo.a;\n  let normalDelta = textureLoad(normalTex, q, 0).rgb - centerNormal;\n  return max(\n    dot(albedoDelta, albedoDelta) * INV_ALBEDO_BANDWIDTH2\n      + coverageDelta * coverageDelta * INV_COVERAGE_BANDWIDTH2,\n    dot(normalDelta, normalDelta) * INV_NORMAL_BANDWIDTH2\n  );\n}\n\nfn premultiplied(c: vec4f) -> vec4f {\n  return vec4f(c.rgb * c.a, c.a);\n}\n\n@compute @workgroup_size(${JNLM_WORKGROUP_SIZE}, ${JNLM_WORKGROUP_SIZE}, 1)\nfn main(@builtin(global_invocation_id) id: vec3u) {\n  let size = vec2i(params.size);\n  let p = vec2i(id.xy);\n  if (p.x >= size.x || p.y >= size.y) { return; }\n  let step = i32(params.dilation);\n  let centerAlbedo = textureLoad(albedoTex, p, 0);\n  let centerNormalTexel = textureLoad(normalTex, p, 0);\n  let centerNormal = centerNormalTexel.rgb;\n  let index = pixelIndex(p);\n  // Primary misses show the background directly. It has no features that\n  // protect its detail and no noise beyond antialiasing, so keep it.\n  if (centerNormalTexel.a <= 0.0 && textureLoad(shadowCatcherTex, p, 0).r <= 0.0) {\n    outColor[index] = color[index];\n    outGuide[index] = guide[index];\n    return;\n  }\n\n  var centerPatch: array<vec2f, PATCH_SIZE>;\n  var k = 0;\n  for (var dy = -PATCH_RADIUS; dy <= PATCH_RADIUS; dy++) {\n    for (var dx = -PATCH_RADIUS; dx <= PATCH_RADIUS; dx++) {\n      centerPatch[k] = guide[pixelIndex(p + vec2i(dx, dy) * step)];\n      k++;\n    }\n  }\n  let k2 = params.colorSensitivity * params.colorSensitivity;\n\n  var sum = vec4f(0.0);\n  let featureScale = max(centerAlbedo.rgb, vec3f(${JNLM_ALBEDO_FLOOR}));\n  var featureSum = vec3f(0.0);\n  var featureSquareSum = vec3f(0.0);\n  var featureColorSum = vec3f(0.0);\n  var guideSum = 0.0;\n  var guideVarianceSum = 0.0;\n  var weightSum = 0.0;\n  var maxWeight = 0.0;\n  for (var sy = -SEARCH_RADIUS; sy <= SEARCH_RADIUS; sy++) {\n    for (var sx = -SEARCH_RADIUS; sx <= SEARCH_RADIUS; sx++) {\n      let q = p + vec2i(sx, sy) * step;\n      if (\n        (sx == 0 && sy == 0) ||\n        q.x < 0 || q.y < 0 || q.x >= size.x || q.y >= size.y\n      ) { continue; }\n\n      let features = featureDistance(q, centerAlbedo, centerNormal);\n      if (features > 16.0) { continue; }\n\n      var colorDistance = 0.0;\n      k = 0;\n      for (var dy = -PATCH_RADIUS; dy <= PATCH_RADIUS; dy++) {\n        for (var dx = -PATCH_RADIUS; dx <= PATCH_RADIUS; dx++) {\n          let a = centerPatch[k];\n          let b = guide[pixelIndex(q + vec2i(dx, dy) * step)];\n          let delta = a.x - b.x;\n          colorDistance += (delta * delta\n              - VARIANCE_CANCELLATION * (a.y + min(a.y, b.y)))\n            / (1e-10 + k2 * (a.y + b.y));\n          k++;\n        }\n      }\n      colorDistance /= f32(PATCH_SIZE);\n\n      let weight = exp(-max(max(colorDistance, 0.0), features));\n      let index = pixelIndex(q);\n      let c = color[index];\n      let g = guide[index];\n      // Filter premultiplied color so coverage and color stay consistent.\n      let a = (textureLoad(albedoTex, q, 0).rgb - centerAlbedo.rgb) / featureScale;\n      let value = premultiplied(c);\n      sum += value * weight;\n      featureSum += a * weight;\n      featureSquareSum += a * a * weight;\n      featureColorSum += a * value.rgb * weight;\n      guideSum += g.x * weight;\n      guideVarianceSum += g.y * weight * weight;\n      weightSum += weight;\n      maxWeight = max(maxWeight, weight);\n    }\n  }\n\n  // Weight the center like its most similar neighbor (Buades et al. 2005)\n  // instead of 1, so isolated outliers are not kept unfiltered.\n  let center = color[index];\n  let centerGuide = guide[index];\n  let centerWeight = max(maxWeight, CENTER_WEIGHT_MIN);\n  let totalWeight = weightSum + centerWeight;\n  let centerValue = premultiplied(center);\n  let mean = (sum + centerValue * centerWeight) / totalWeight;\n  // Fit each color channel against relative albedo with the same JNLM weights.\n  // Centering at this pixel avoids cancellation when the feature is constant.\n  let meanFeature = featureSum / totalWeight;\n  let featureVariance = max(featureSquareSum / totalWeight - meanFeature * meanFeature, vec3f(0.0));\n  let covariance = featureColorSum / totalWeight - meanFeature * mean.rgb;\n  // A 1% contrast ridge stabilizes flat guides. Bound the response between no\n  // texture dependence and multiplicative albedo modulation; specular color\n  // with an unrelated albedo falls back to the weighted mean.\n  let meanAlbedo = centerAlbedo.rgb + meanFeature * featureScale;\n  let maximumSlope = mean.rgb * featureScale / max(meanAlbedo, vec3f(${JNLM_ALBEDO_FLOOR}));\n  let slope = clamp(covariance / (featureVariance + vec3f(${JNLM_ALBEDO_REGULARIZATION})), vec3f(0.0), maximumSlope);\n  let corrected = mean.rgb - slope * meanFeature;\n  let filtered = vec4f(max(corrected, vec3f(0.0)), max(mean.a, 0.0));\n  let rgb = select(\n    center.rgb,\n    filtered.rgb / max(filtered.a, 1e-6),\n    filtered.a > 1e-6\n  );\n  outColor[index] = vec4f(rgb, filtered.a);\n  outGuide[index] = vec2f(\n    (guideSum + centerGuide.x * centerWeight) / totalWeight,\n    (guideVarianceSum + centerGuide.y * centerWeight * centerWeight)\n      / (totalWeight * totalWeight)\n  );\n}\n`;\n}\n\nexport interface JNLMInputs {\n  width: number;\n  height: number;\n  /** Normalized shadow-catcher mask; red is the catcher coverage. */\n  shadowCatcher: GPUTexture;\n  /** Ray-trace auxiliary accumulation holding the guide moments. */\n  accumulation: GPUBuffer;\n  /** Ray-trace normal accumulation; its alpha is the per-pixel sample count. */\n  normalAccumulation: GPUBuffer;\n  /** Composed linear image to filter. */\n  color: GPUTexture;\n  albedo: GPUTexture;\n  normal: GPUTexture;\n  strength: number;\n  shape: JNLMFilterShape;\n}\n\nconst MAX_ITERATIONS = JNLM_MAX_ITERATIONS;\n\nexport class WGPUJNLMDenoisePass {\n  private _device: GPUDevice;\n  private _paramsBuffers: GPUBuffer[] = [];\n  private _guidePipeline: GPUComputePipeline;\n  private _variancePipeline: GPUComputePipeline;\n  private _loadColorPipeline: GPUComputePipeline;\n  private _filterPipeline?: GPUComputePipeline;\n  private _filterShapeKey = '';\n  private _guides: GPUBuffer[] = [];\n  private _colors: GPUBuffer[] = [];\n  private _width = 0;\n  private _height = 0;\n\n  constructor(device: GPUDevice) {\n    this._device = device;\n    // One uniform buffer per pass: every pass of an encoder sees the last\n    // writeBuffer() of a shared one.\n    for (let i = 0; i <= MAX_ITERATIONS; i++) {\n      this._paramsBuffers.push(device.createBuffer({\n        label: 'JNLM params',\n        size: 16,\n        usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST\n      }));\n    }\n    this._guidePipeline = this._createPipeline('JNLM guide', jnlmGuideWGSL);\n    this._variancePipeline = this._createPipeline('JNLM variance', jnlmVarianceWGSL);\n    this._loadColorPipeline = this._createPipeline('JNLM load color', jnlmLoadColorWGSL);\n  }\n\n  private _createPipeline(label: string, code: string) {\n    return getTargetShaderCache(this._device).createComputePipeline({\n      label,\n      layout: 'auto',\n      compute: {\n        module: getTargetShaderCache(this._device).createShaderModule({ label, code }),\n        entryPoint: 'main'\n      }\n    });\n  }\n\n  private _getFilterPipeline(shape: JNLMFilterShape) {\n    const key = JSON.stringify(shape);\n    if (!this._filterPipeline || key !== this._filterShapeKey) {\n      this._filterPipeline = this._createPipeline('JNLM filter', createJNLMFilterWGSL(shape));\n      this._filterShapeKey = key;\n    }\n    return this._filterPipeline;\n  }\n\n  reserveCapacity() {}\n  private _ensureBuffers(width: number, height: number) {\n    if (width <= this._width && height <= this._height && this._colors.length) return;\n    width = Math.max(width, this._width); height = Math.max(height, this._height);\n    this._destroyBuffers();\n    const pixels = width * height;\n    const createBuffer = (label: string, bytes: number) =>\n      this._device.createBuffer({\n        label,\n        size: pixels * bytes,\n        usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC\n      });\n    this._guides = [createBuffer('JNLM guide A', 8), createBuffer('JNLM guide B', 8)];\n    this._colors = [createBuffer('JNLM color A', 16), createBuffer('JNLM color B', 16)];\n    this._width = width;\n    this._height = height;\n  }\n\n  private _destroyBuffers() {\n    for (const buffer of [...this._guides, ...this._colors]) buffer.destroy();\n    this._guides = [];\n    this._colors = [];\n  }\n\n  /** Encodes the denoiser and returns the RGBA32F output in texture rows. */\n  encode(encoder: GPUCommandEncoder, inputs: JNLMInputs): GPUBuffer {\n    const { width, height, shape } = inputs;\n    this._ensureBuffers(width, height);\n    const iterations = Math.min(Math.max(Math.round(shape.iterations), 1), MAX_ITERATIONS);\n    const sensitivity = JNLM_DEFAULT_COLOR_SENSITIVITY * Math.max(inputs.strength, 1e-3);\n    this._paramsBuffers.forEach((buffer, i) => {\n      const params = new ArrayBuffer(16);\n      new Uint32Array(params, 0, 2).set([width, height]);\n      // Buffer i (from 1) is filter pass i - 1; index 0 is the prepasses.\n      new Float32Array(params, 8, 2).set([sensitivity, jnlmPassDilation(Math.max(i - 1, 0))]);\n      this._device.queue.writeBuffer(buffer, 0, params);\n    });\n\n    const device = this._device;\n    const filterPipeline = this._getFilterPipeline(shape);\n    const bind = (\n      pipeline: GPUComputePipeline,\n      paramsBuffer: GPUBuffer,\n      resources: GPUBindingResource[]\n    ) =>\n      device.createBindGroup({\n        layout: pipeline.getBindGroupLayout(0),\n        entries: [\n          { binding: 0, resource: { buffer: paramsBuffer } },\n          ...resources.map((resource, i) => ({ binding: i + 1, resource }))\n        ]\n      });\n    const [guideA, guideB] = this._guides;\n    const [colorA, colorB] = this._colors;\n    const prepass = this._paramsBuffers[0];\n    const passes: [GPUComputePipeline, GPUBindGroup][] = [\n      [this._guidePipeline, bind(this._guidePipeline, prepass, [\n        { buffer: inputs.accumulation },\n        { buffer: inputs.normalAccumulation },\n        { buffer: guideB }\n      ])],\n      [this._variancePipeline, bind(this._variancePipeline, prepass, [\n        { buffer: guideB },\n        { buffer: guideA }\n      ])],\n      [this._loadColorPipeline, bind(this._loadColorPipeline, prepass, [\n        inputs.color.createView(),\n        { buffer: colorA }\n      ])]\n    ];\n    const albedoView = inputs.albedo.createView();\n    const normalView = inputs.normal.createView();\n    const shadowCatcherView = inputs.shadowCatcher.createView();\n    let guides = [guideA, guideB];\n    let colors = [colorA, colorB];\n    for (let i = 1; i <= iterations; i++) {\n      passes.push([filterPipeline, bind(filterPipeline, this._paramsBuffers[i], [\n        { buffer: guides[0] },\n        { buffer: colors[0] },\n        albedoView,\n        normalView,\n        { buffer: colors[1] },\n        { buffer: guides[1] },\n        shadowCatcherView\n      ])]);\n      guides = [guides[1], guides[0]];\n      colors = [colors[1], colors[0]];\n    }\n    const pass = encoder.beginComputePass({ label: 'JNLM denoise' });\n    for (const [pipeline, bindGroup] of passes) {\n      pass.setPipeline(pipeline);\n      pass.setBindGroup(0, bindGroup);\n      pass.dispatchWorkgroups(\n        Math.ceil(width / JNLM_WORKGROUP_SIZE),\n        Math.ceil(height / JNLM_WORKGROUP_SIZE)\n      );\n    }\n    pass.end();\n    return colors[0];\n  }\n\n  getMemoryBytes() {\n    return [...this._paramsBuffers, ...this._guides, ...this._colors]\n      .reduce((sum, buffer) => sum + buffer.size, 0);\n  }\n\n  dispose() {\n    for (const buffer of this._paramsBuffers) buffer.destroy();\n    this._destroyBuffers();\n  }\n}\n"]}

package/lib/renderer/WGPURayTraceOutput.d.ts

@@ -0,0 +1,44 @@
+export declare function createNormalizeRayTraceOutputWGSL(layers?: boolean, catcher?: boolean): string;
+type RayTraceOutput = 'color' | 'normal' | 'albedo' | 'shadowCatcher' | 'foreground' | 'catcherBackground';
+type RayTraceBuffers = {
+    color: GPUBuffer;
+    normal: GPUBuffer;
+    albedo: GPUBuffer;
+    shadowCatcher: GPUBuffer;
+    shadowCatcherOffset?: number;
+};
+export declare class WGPURayTraceOutputPass {
+    private readonly zeroTexture?;
+    private _device;
+    private _pipeline;
+    private _layersPipeline;
+    private _foregroundTexture;
+    private _catcherBackgroundTexture;
+    private _colorTexture;
+    private _normalTexture;
+    private _albedoTexture;
+    private _shadowCatcherTexture;
+    private _uniformBindGroup;
+    private _uniformsBuffers;
+    private _width;
+    private _height;
+    private _hasCatcher;
+    private _needsUpdatePipeline;
+    constructor(device: GPUDevice, zeroTexture?: GPUTexture | undefined);
+    private readonly _retainedOutputs;
+    retainOutput(texture: GPUTexture): void;
+    releaseOutput(texture: GPUTexture): void;
+    getOutput(label: RayTraceOutput): GPUTexture;
+    private _createTextures;
+    private _updateBindGroups;
+    private _updatePipeline;
+    setSize(width: number, height: number): void;
+    setSampleCount(sampleCount: number): void;
+    setSampleTileGrid(columns: number, rows: number): void;
+    createPass(encoder: GPUCommandEncoder, inputBuffers: RayTraceBuffers, separateCatcher?: boolean): void;
+    private _createOutputPass;
+    private _destroyTextures;
+    getMemoryBytes(): number;
+    dispose(): void;
+}
+export {};