Add PPISP USD export, SPG controller rendering, and SH bake validation#243
Closed
moennen wants to merge 42 commits into
Closed
Add PPISP USD export, SPG controller rendering, and SH bake validation#243moennen wants to merge 42 commits into
moennen wants to merge 42 commits into
Conversation
Decode NuRec USD/USDZ inputs into LightField export data so transcode can convert NuRec packages into standard ParticleField USD. Made-with: Cursor
Run OpenUSD validation after standard USD export so composition and stage metadata errors are caught before returning success. Made-with: Cursor
Add a no-parameter post-processing module for training and rendering when the renderer outputs linear RGB against sRGB targets. Made-with: Cursor
Expose export_usd.linear_srgb so LightField exports can author lin_rec709_scene instead of the default display-referred color space. Made-with: Cursor
Author per-camera RenderProducts and attach the PPISP SPG shader assets so standard USD exports can preserve the trained PPISP effect. Made-with: Cursor
Add a dedicated writer that maps PPISP exposure and fitted post-processing parameters to Omniverse USD attributes for deployments without reliable SPG support. Made-with: Cursor
Make export_ppisp the master PPISP export gate and use ov-post-processing only to select SPG, exposure fallback, fitted fallback, or combined SPG-plus-fallback output. Made-with: Cursor
Keep cameras, render products, and PPISP post-processing authored where the packaged root stage composes them, and align standalone USD export with PPISP configuration. Made-with: Cursor
Ensure checkpoint exports carry the PPISP SPG graph, camera timeline, and Gaussian render settings needed for Kit to render the post-processed LdrColor output. Made-with: Cursor
Add an explicit omni-usd opt-in so default USD exports remain neutral while Kit-specific PPISP SPG and ParticleField MDL material authoring are only emitted when requested. Made-with: Cursor
Default checkpoint exports to include loaded PPISP modules unless explicitly disabled, and scope neutral camera exposure to a hidden RenderProduct-local camera for PPISP SPG. Made-with: Cursor
Add a generic post-processing SH bake path with a PPISP adapter so standard USD exports can bake a fixed post-processing transform into Gaussian coefficients.
Remove the separate omni-usd and ov-post-processing controls while keeping PPISP native export behind the explicit omni-native post-processing mode.
Add fixed camera/frame controls for PPISP omni-native export and extend bake validation with L1 fitting plus reference baseline outputs.
Add a Viser-based image comparison tool for paired images and matched folders, with visual diff modes and metrics support.
Clamp the power branch away from zero so validation image conversion remains finite.
Display metrics as readable text, add folder-level aggregates, and compute FLIP by default while removing LPIPS from the viewer.
Add simple fixed-frame PPISP baking variants to the validation workflow and keep the fitted bake path aligned with the current L2 objective.
Expose and validate ParticleField sortingModeHint values so USD exports can target ray-hit sorting while preserving the existing cameraDistance default.
Adds the per-camera PPISP controller (CNN + adaptive avg pool + 3-layer
MLP) to the Omniverse-native USD export so the runtime SPG pipeline
predicts exposure/colour at frame time instead of relying on
time-sampled USD attributes baked at training.
Pipeline authored on each RenderProduct:
HdrColor -> PPISPController_<n> -> ControllerParams (1x9 float)
-> PPISP (dyn) -> PPISPColor -> LdrColor
The controller weights flatten into a single 241,961-element float[]
attribute on the Shader prim; the SPG slang reads them via a
StructuredBuffer<float>. The dynamic PPISP variant
(ppisp_usd_spg_dyn.slang) reads exposure/colour from the controller's
1x9 output texture instead of from USD inputs.
Added a slangpy-based headless harness under tools/render_ppisp_spg/
that compiles and dispatches the same shaders without booting Kit, plus
several validators (synthetic, real ppisp module, trained checkpoint)
that compare the slang result against the in-process PyTorch PPISP
forward pass.
End-to-end on a trained bonsai checkpoint: PSNR 63 dB, controller
9-float drift ~3e-7 vs torch — within fp32 rounding of the
rgba8_unorm output format.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Lets a downstream consumer that doesn't want runtime controller dispatch (e.g. for stricter portability or smaller assets) opt out of the controller export even when the checkpoint contains trained controllers. The exporter falls back to the static SPG path with time-sampled exposure / color attributes derived from ppisp.exposure_params and ppisp.color_params -- the same optimized per-frame parameters PPISP would use when the controller branch is bypassed. No effect on checkpoints trained without a controller. Available both as a CLI flag (--ignore-ppisp-controller) and a YAML key (export_usd.ignore-ppisp-controller). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
diagnose_controller.py runs three independent checks against a trained
checkpoint and pinpoints which layer is responsible when the controller-
driven render disagrees with the optimized-per-frame-params render:
H1 -- PyTorch controller(rgb) vs trained exposure_params/color_params.
Pure-python check; failure means the controller did not converge
to the per-frame state during distillation.
H2 -- slang controller vs PyTorch controller on the same HDR. Failure
means the slang shader / weight flatten / buffer upload diverges
from the trained module.
H3 -- ppisp_usd_spg_dyn.slang (reads ControllerParams texture) vs
ppisp_usd_spg.slang (USD attributes), both fed the same 9 floats.
Failure means the SPG plumbing between the two shaders is broken.
Each section prints per-frame numbers plus a clear pass/fail interpretation
threshold, and the summary names the failing hypothesis. Saves having
to instrument the export end-to-end again the next time something
disagrees in Kit.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Previously the PPISP-dyn shader's inputs:ControllerParams was connected
directly to PPISPController_n.outputs:ControllerParams. That's valid
UsdShade -- the slangpy harness happily resolves it -- but Kit's SPG
runtime walks RenderProduct.orderedVars and resolves connections through
omni:rtx:aov on RenderVar prims, never directly between Shader prims. A
direct Shader -> Shader hop leaves the consumer unbound at dispatch time
(or bound to whatever Kit falls back to), which manifests as a much-too-
exposed render in Omniverse even though slangpy looks correct.
Mirror the existing HdrColor / LdrColor idiom: insert an intermediate
"ControllerParams" RenderVar with `opaque omni:rtx:aov.connect =
PPISPController_n.outputs:ControllerParams`, add it to orderedVars, and
point PPISP.inputs:ControllerParams.connect at the RenderVar's
omni:rtx:aov attribute.
Wiring is now:
HdrColor (RenderVar+aov)
-> PPISPController_n.inputs:HdrColor.connect = HdrColor.omni:rtx:aov
PPISPController_n.outputs:ControllerParams
-> ControllerParams (RenderVar) omni:rtx:aov.connect
= PPISPController_n.outputs:ControllerParams
-> PPISP.inputs:ControllerParams.connect = ControllerParams.omni:rtx:aov
PPISP.outputs:PPISPColor
-> LdrColor (RenderVar) omni:rtx:aov.connect = PPISP.outputs:PPISPColor
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Kit's SPG lua sandbox does a textual deny-list match for the lua module-loader keyword 'require' and rejects any source containing it -- including inside string literals. Both new lua launchers had assert messages that read "... requires ...", which triggered the sandbox validator and caused the entire shader graph build to fail with: LuaSandbox: Lua source rejected - forbidden pattern 'require' found Output outputs:ControllerParams of node PPISPController_0 is missing shape spec. Can't compute output shape. (The shape error is a downstream effect: with the lua rejected, SPG never executes outputs["ControllerParams"] = slang.empty(...) so the shape can't be inferred.) Replace "requires" with "needs" in both messages. Also did a quick audit for other commonly denied tokens (dofile/loadfile/loadstring/ os./io./debug./setfenv/getfenv) -- both luas are clean. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Kit's SPG slang lua sandbox does not expose ``slang.StructuredBuffer`` (it's nil), so the previous launcher failed with "attempt to call a nil value (field 'StructuredBuffer')". Replace the hard-coded call with a small probe that tries every common HLSL/Slang buffer-resource name in turn (StructuredBuffer, RWStructuredBuffer, Buffer, RWBuffer, ByteAddressBuffer, RWByteAddressBuffer) and falls back to an explicit error message that lists every key currently in the ``slang.*`` table. If the probe doesn't find a match the user gets a precise log line naming the available helpers, which we can use to settle on the right one without further guessing. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The previous probe only listed pairs(slang), which under SPG's slang table iterates only the directly-stored keys (the truncated 'short, full, bo...' from the Kit log). The actual resource helpers like Texture2D/dispatch live behind a __index metatable, so pairs() misses them and the prior probe couldn't tell us anything useful. Expand the candidate list to cover everything plausible (cased and uncased forms: Buffer/buffer, Array/array, FloatArray/floatArray, image/Image, uniform/Uniform, list/List, etc.) and test each via direct field access (which goes through __index). Report both the direct pairs() keys and the candidate list in the error if nothing matches. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Kit's SPG was emitting one warning per dispatch: [Warning] [rtx.spg.slang] SpgSlangNode: Failed to find parameter 'params:weights' in shader reflection because the slang shader declared the buffer as ``g_Weights`` while the USD Shader prim exposes it as ``inputs:weights``. SPG's reflection-based auto-binding looks up parameters by name; the mismatch caused the spam even though the explicit lua bind in the dispatch worked. Rename the slang variable to ``weights`` so the names match. Validation still passes at PSNR 63 dB; max|Δ| in the controller's 9-float output remains ~3.58e-7. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…rning Kit's SPG resolves USD ``inputs:foo`` attributes against fields of the slang ParameterBlock -- its reflection lookup is ``params:foo``. With ``weights`` declared as a sibling top-level ``StructuredBuffer<float>`` the lookup failed every dispatch: [Warning] [rtx.spg.slang] SpgSlangNode: Failed to find parameter 'params:weights' in shader reflection The static PPISP shader follows the same convention -- every attribute (exposureOffset, vignetting*, crf*, ...) lives inside the PPISPParams struct that's wrapped by ParameterBlock<PPISPParams>. Move ``weights`` into PPISPControllerParams alongside priorExposure; update both the lua bind list (now passes the buffer as a positional argument to slang.ParameterBlock) and the slangpy harness's ShaderCursor (cur["g_Params"]["weights"] instead of cur["weights"]). Validation still passes at PSNR 63 dB; controller 9-float drift remains ~3.58e-7 vs torch. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The fit-by-bake export started Adam from the cloned checkpoint's SH
state, even though we have a closed-form one-shot bake (simple_bake)
that already lands close to the optimum. Use it as the initialization
so the fitting loop only refines the residual.
Adapter changes:
- PostProcessingBakeAdapter.initialize_fit(model, pp) -- new hook,
default no-op.
- bake_post_processing_into_sh now calls adapter.initialize_fit on
the cloned baked_model right before constructing the optimizer.
- PPISPPostProcessingBakeAdapter.initialize_fit calls simple_bake
with higher_order=True (so the spatial SH coefficients are
Jacobian-projected too) and apply_srgb_to_linear=True.
The new srgb_to_linear flag on simple_bake matters for the colour-
space round-trip: PPISP outputs display-referred values (its CRF
folds in gamma-like encoding). Storing those directly as linear SH
coefs leaves the asset double-encoded once a downstream consumer
applies linear_to_srgb (the validator does, Kit's tonemap does).
Applying srgb_to_linear before RGB2SH puts SH back in linear scene-
referred space; downstream linear_to_srgb then undoes the encoding
exactly. Verified srgb_to_linear∘linear_to_srgb round-trips to fp32
epsilon (max|Δ| ≈ 2.4e-7) on [0, 1].
Existing 19 export tests pass.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Adds two non-training view samplers to bake_post_processing_into_sh, on
top of the existing "iterate the train dataloader" mode:
* random-pair-slerp -- pick two distinct training views uniformly at
random and slerp between them at random s in [0, 1]. No global
structure but cheap.
* trajectory -- order the training views along an approximate
Hamiltonian path (nearest-neighbour seed + 2-opt refinement) using
a position+direction distance, arc-length-parameterise the path on
[0, 1], and per step sample a random t and slerp inside the
bracketing segment. Closer to the camera continuum a viewer would
fly through; better generalisation than discrete training poses.
New module ``post_processing_view_interpolation``:
- slerp_pose(pose_a, pose_b, s) -- quaternion slerp + translation lerp,
fp32-clean to within fp32 epsilon at s=0 / s=1.
- order_views_along_trajectory(poses, ...) -- NN + 2-opt with a
weighted (position L2 + 1 - cos(forward angle)) metric. Position
distances are mean-normalised so the rotation term lives on a
comparable scale across scenes.
- InterpolatedViewSampler -- wraps a template Batch and emits
steps_per_epoch synthetic Batches with only T_to_world replaced.
bake_post_processing_into_sh now accepts:
view_sampling_mode: "training" | "random-pair-slerp" | "trajectory"
interpolated_views_seed: optional int RNG seed
trajectory_weight_position / trajectory_weight_rotation: trajectory
metric weights.
Wired through USDExporter constructor + from_config (YAML keys
post-processing-bake-view-mode / -view-seed / -trajectory-weight-{position,rotation})
and the export_usd.py CLI (--post-processing-bake-view-mode etc.).
Default remains "training" so existing exports are unchanged.
Sanity checked on a 6-pose synthetic circle: shuffled inputs reorder to
a Hamiltonian path with uniform 0.2 arc-length steps; existing 19 export
tests pass.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
PPISP forward is display-referred, so the fitted SH side must encode to sRGB and clamp before MSE — without this the loss plateaued near 13 dB on real scenes. Mirrors post_processing_sh_bake_validation.py::_fitBakedSh. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Two bugs in simple_bake's higher_order=True path under apply_srgb_to_linear=True: 1. The Jacobian was dPPISP/dX while the DC bake landed in linear space via srgb_to_linear; the chain rule was missing srgb_to_linear'(PPISP). The composed Jacobian now matches the DC color space. 2. ~0.06% of Gaussians at PPISP-saturation extremes have pathological Jacobians (cond > 1e8, |J|_F > 1e4) that pump features_specular norms from O(1) to O(10^4). These outliers dominate Adam's adaptive variance and stall the fit at ~29 dB. Clipping |J|_F > 5 (and any non-finite J) to identity preserves the rotation for the well-behaved 99.94% while leaving the trained specular intact for outliers. On bonsai (1M Gaussians, 9-epoch fit), the higher-order PPISP warm-start went from 25.4 dB to 36.4 dB -- now within 0.2 dB of the DC-only sRGB warm-start (36.6 dB). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Empirical sweep showed random-pair sampling was always within noise of training views (sometimes -0.8 dB) and never improved over the trajectory sampler. The trajectory mode is retained for sparse-view datasets where interpolating between adjacent training poses can help generalisation. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Sweeps simple-bake and Adam-fit variants of the PPISP SH bake on a trained checkpoint, reports mean / median / min / max PSNR (optional SSIM, LPIPS) across the validation split, and writes per-frame numbers to metrics.json. Used to root-cause the higher-order warm-start regression (Jacobian chain rule + outlier clipping). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…ochs Update PPISP SH-bake export defaults to the configuration that wins on the bake-modes ablation: * PPISPPostProcessingBakeAdapter.initialize_fit now uses higher_order=False -- DC-only simple_bake leaves the trained features_specular intact, which Adam fine-tunes faster than recovering from a Jacobian-rotated specular. * post_processing_bake_view_mode default flipped from "training" to "trajectory" -- arc-length-sampled NN+2-opt path through training poses; +0.17 dB on bonsai at 15 epochs and expected to help more on sparse-view datasets. * post_processing_bake_epochs default raised from 1 to 13 -- the fit is far from converged at 1, and 13 sits in the knee of the PSNR-vs-epochs curve (~36.5-37 dB on bonsai). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Two bugs in the PPISP SH bake fit, each found by ablation on caterpillar: 1. The SH was fit in linear scene-referred space against a display-referred target with a linear_to_srgb step in the loss path. The chain rule through linear_to_srgb is wildly non-uniform on [0, 1] (~13x at darks, ~0.4x at brights). Adam over-weights dark-region updates and the high-degree SH bands oscillate to fit amplified noise -- visible as rainbow fringing around silhouettes. Standard 3DGS doesn't see this because it trains in gamma space directly with identity gradient through the loss. Fix: warm-start albedo with simple_bake(apply_srgb_to_linear=False) and strip linear_to_srgb from apply_fit_transform. SH eval is now display- referred; gradient is identity through the loss; the asset format aligns with no-PPISP exports (gamma SH + no post-processing layer). 2. Fixed-geometry SH refit on a smooth synthetic target (PPISP forward of the trained model's prediction) gives the high-degree bands too much license to align coherently into aliasing patterns -- nothing in the loss surface breaks the symmetry. Standard 3DGS training avoids this via density / split-clone-prune adapting the placement to the target. Fix: open features_albedo, features_specular and density to Adam, each at its 3DGS-standard learning rate (2.5e-3, 1.25e-4, 5e-2). On caterpillar (48 val frames, 13 epochs): before: mean 32.0 dB, worst 26.6 dB, strong rainbow on hard frames after: mean 43.4 dB, worst 39.0 dB, rainbow gone Default vignetting flipped from "achromatic-fit" to "none": the chromatic-vs-achromatic mismatch was contributing ~1 dB of error and the new asset format doesn't have a place to ship a runtime vignette anyway. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Match the new bake_post_processing_into_sh signature in the exporter and CLI: separate albedo / specular / density learning rates, with the 3DGS-standard ratio (specular = albedo / 20) preserved by default. * default --post-processing-bake-learning-rate flipped to 2.5e-3 * new --post-processing-bake-learning-rate-specular (default = albedo/20) * new --post-processing-bake-learning-rate-density (default 5e-2) * default --ppisp-bake-vignetting-mode flipped to "none" Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Drop the linear-SH + achromatic-vignette modes (fit-base, fit-base-srgb, fit-init, ...) -- they no longer match the production fit path and were only useful to debug the rainbow regression. Replace with a leaner catalogue aligned with the new defaults: simple one-shot DC-only bake (no fit, gamma SH) simple-higher-order one-shot DC + Jacobian-rotated specular (no fit) fit-color-only Adam on albedo + specular only (density ablation) fit Adam on albedo + specular + density (production) fit-trajectory fit + trajectory view sampling Eval likewise simplifies: SH is already display-referred so the baked side just clips to [0, 1]; reference uses no-vignette PPISP. On caterpillar (48 val frames, 13 epochs) the 'fit' mode lands at 43.4 dB mean / 39.0 dB worst-case, vs 32.0 / 26.6 with the old fit-base. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The well-conditioned gamma-SH fit converges fast: per-step loss curves on bonsai and caterpillar show diminishing returns past epoch 7. Cutting the default ~halves wall-clock time per export with marginal PSNR cost (rough estimate -0.5 dB mean, the asymptote is ~the same). Users wanting the last fraction of a dB can still set --post-processing-bake-epochs. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
A single multiplicative knob applied to the asset's SH-evaluated RGB. ``features_specular`` is scaled linearly; ``features_albedo`` picks up an extra ``(s - 1) * 0.5 / C0`` term to compensate for the constant offset baked into ``RGB2SH`` so a forward eval reproduces ``s * original_rgb``. Works uniformly across export modes (no-PPISP linear-SH, no-PPISP gamma, PPISP baked-sh) -- the offset compensation is structural, not colour- space-dependent. Default 1.0 keeps existing exports byte-identical. Useful for matching downstream tonemap exposure or compensating for runtime gain that the asset's consumer applies. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Adds a ``responsivityR/G/B`` float input on the PPISP SPG shader (both the static and controller-aware variants) that is premultiplied to the input HdrColor before exposure / vignetting / colour correction / CRF run. Default 1.0 per channel keeps exported assets visually identical; authors can scale per-channel sensitivity post-export by editing the USD ``inputs:responsivityR/G/B`` attributes on the per-camera shader prim. Touches: * both ``.slang`` shaders (struct field + premultiply at top of main) * both ``.slang.lua`` launchers (bind the new fields, default 1.0) * both ``.slang.usda`` schemas (declare inputs with default 1.0) * ``ppisp_writer.py`` writes the inputs explicitly via ``_set_responsivity_params`` so the attributes round-trip cleanly. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Remove local Cursor agent metadata from the branch and fix stale PPISP plan references so the PR content matches the exported tooling.
33e8f2a to
0e791a2
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR adds PPISP-aware USD export support, including SPG-based post-processing assets, per-camera render products, PPISP controller export, and post-processing bake paths for SH output.
It also adds validation and diagnostic tooling for PPISP USD/SPG output, image comparison workflows, bake-mode benchmarking, NuRec USD import/export refinements, and LightField USD validation.
Key Changes
linear-to-srgbpost-processing support and document/configure PPISP post-processing behavior.Test Plan
python -m pytest threedgrut/export/tests/test_export_import.py -vpytestis not installed in the current environment.