RFC: World-coord canonical encoding via bidirectional UVW↔RGB atlas#61
Draft
MiLO83 wants to merge 12 commits into
Draft
RFC: World-coord canonical encoding via bidirectional UVW↔RGB atlas#61MiLO83 wants to merge 12 commits into
MiLO83 wants to merge 12 commits into
Conversation
MiLO83
pushed a commit
to MiLO83/DownToEarth
that referenced
this pull request
May 20, 2026
The draft PR is the natural conclusion of the whole project arc; surface it above the live-demo block so anyone landing on the repo or the .dev README sees the upstream contribution first.
This is a discussion-only RFC, not a ready-to-merge change. All three
modifications are opt-in / non-default-path; existing Lyra 2 behaviour
is preserved byte-identically when none of the new flags or methods are
invoked. The intent is to surface an architectural option for a future
training run, not to ask for an immediate merge.
Full write-up:
github.com/MiLO83/DownToEarth/blob/main/LYRA2_PROPOSAL.md
(Apache 2.0 / MIT - the same prose as a private V2 document was
shared with the team separately)
Summary
=======
Lyra 2's canonical-coord conditioning encodes per-pixel image-space
coordinates plus a frame-slot index - (u_norm, v_norm, frame_slot) -
and the cross-attention attends over a frame-keyed Sparse3DCache. The
cache already stores per-pixel world coordinates (via unproject_points
in Sparse3DCache.add()); the world coordinates are *computed and
persisted*, but not used as the canonical key.
This patch demonstrates the architectural alternative: encode the
canonical-coord image as quantized world coordinates packed into RGB
bytes (or uint16/uint32), with a bidirectional bijection between the
coords and a 2D atlas layout that is GPU-cache-friendly. The same
forward_warp_multiframes machinery transports either encoding cleanly.
Why it might be worth a look
============================
* One key per 3D point across all views (geometric correspondence
becomes structural, not learned across frame-pixel pairs)
* Memory bounded by scene complexity, not trajectory length
(revisits reinforce existing atlas slots, no new allocation)
* O(1) retrieval per pixel (texture sample) instead of the O(N)
visibility-overlap scoring in Sparse3DCache.retrieve()
* Byte-perfect identity at uint8 / 16 / 32 - pixel bytes literally
are the world coords (modulo a known world_min/max transform)
What this patch does (3 modifications, all opt-in)
==================================================
1. NEW lyra_2/_src/datasets/uvw_atlas.py
- Bidirectional bijection family covering RGBA8 to RGBA32UI
- Pure arithmetic, no learnable params, no dependencies beyond torch
- Doctested + exhaustive 64-cubed round-trip self-test (passes inside
the Lyra tree: 'python -m lyra_2._src.datasets.uvw_atlas')
- Apache 2.0 header to match the rest of the codebase
- +161 LoC, no existing code touched
2. NEW _build_canonical_world_coords static method in lyra2_model.py
- Sits alongside the existing _build_canonical_spatial_coords
- Emits world-coord-seeded canonical channels instead of
image-space + frame-slot
- Opt-in - not wired into _get_cached_spatial_coords; existing
model behaviour unchanged until the dispatch is added
- +63 LoC, no existing code touched
3. New octant_prefilter kwarg on Sparse3DCache.retrieve()
- Defaults to False (existing behaviour byte-identical)
- When True, prunes candidates whose centroid is behind all target
camera views before the expensive [V, C, B, H', W', 4, 4] matmul
- Clean index remapping via _orig_idx through avail_frame_ids,
score_map, and the return statement
- Pure inference-time optimization, no retraining needed
- Expected ~50% reduction in retrieve() matmul cost at long horizons
with zero quality impact (those candidates contribute zero valid
pixels via the existing z_all > 0 filter anyway)
- +45 LoC
Trade-offs (the honest counterweight)
=====================================
* nv-tlabs#1 and nv-tlabs#2 only become functional once the conditioning head is
fine-tuned to read integer-quantized world coords instead of
fp16 normalized image coords. We cannot do that retrain ourselves
(no GB200s).
* For unbounded scenes (90 m+ trajectories), the world-coord atlas
saturates beyond world_max. A hierarchical sparse-tiled variant
is sketched in the proposal but not in this patch (would be its
own ~300 LoC module + page table).
* Quantization error: 1 LSB per axis. At RGBA8 + 2.5 m world = 1 cm
cells. At RGBA16UI + 655 m world = sub-mm. RGBA32UI essentially
unbounded.
Verification
============
* 'python -m lyra_2._src.datasets.uvw_atlas' -> 6/6 doctests pass,
exhaustive 64-cubed bijection round-trip OK
* Both modified files pass 'ast.parse()' syntax check
* 'git diff main' is byte-identity-preserving in the default code
path (octant_prefilter=False, _build_canonical_world_coords
unused)
* Not validated under any training run (would require GB200s we
do not have access to)
Suggested next step
===================
If the team finds the encoding shift interesting enough to
prototype, the smallest possible experiment is the octant_prefilter
(modification 3) alone - no retraining required, expected ~5-15%
wall-clock speedup at long horizons, free quality. If that holds, the
case for considering the encoding shift (modifications 1+2) in a
future training run gets stronger.
Marking this PR as draft per the contributing guide - open to discuss,
refine, retract, or close at the team's call. Happy to defer entirely
if there are reasons not to pursue this direction.
Signed-off-by: MiLO83 <edisonzghost@gmail.com>
MiLO83
force-pushed
the
feat/uvw-bijection-canonical-encoding
branch
from
May 20, 2026 03:55
43252c8 to
7f79b96
Compare
Same-resolution 1-bit-per-voxel companion to the canonical-coord atlas.
Lets the cross-attention / inference shader skip the multi-byte main
atlas read for empty voxels via a single bit test that is 24-32x
cheaper in storage and bandwidth.
Memory math:
256-cubed atlas (16.7M voxels):
- RGBA8 dense: 67 MB
- 1-bit mask: 2.1 MB (32x smaller)
1024-cubed atlas (1.07B voxels):
- RGBA8 dense: 4.3 GB
- 1-bit mask: 134 MB (32x smaller)
Layout: 8 voxels packed per byte along the X axis of the 2D atlas to
preserve cache locality for horizontal-scan access patterns. Standard
pattern from id Tech 5 MegaTexture feedback masks / sparse voxel
octree leaf occupancy.
For Lyra 2 specifically, this pairs with the canonical-coord encoding
from _build_canonical_world_coords as an explicit emptiness signal.
The cross-attention does not have to learn that (0, 0, 0) is special;
a dedicated 1-bit channel says so structurally. Adds ~50 KB per
inference chunk at 832 x 480 resolution.
Class API:
- set / get by (atlas_x, atlas_y) raw bit access
- count_occupied / occupancy_fraction stats
- fill_from_voxel_iter vectorised bulk write via np.bitwise_or.at
- to_bytes / from_bytes serialise / deserialise
- as_torch_tensor for GPU upload as R8UI texture
GLSL consumption pattern (one texelFetch + shift + AND):
uniform usampler2D occupancyBitmap;
bool is_occupied(ivec2 atlas_xy) {
uint byte_v = texelFetch(
occupancyBitmap, ivec2(atlas_xy.x >> 3, atlas_xy.y), 0
).r;
return ((byte_v >> (atlas_xy.x & 7)) & 1u) != 0u;
}
Verification:
python -m lyra_2._src.datasets.uvw_atlas
-> 12/12 doctests pass (up from 6, +6 OccupancyBitmap tests)
-> Exhaustive 64-cubed bijection round-trip still OK
-> OccupancyBitmap fill_from_voxel_iter + get round-trip OK
-> to_bytes / from_bytes byte-exact round-trip OK
This is a fourth opt-in modification, not wired into any existing call
path. Existing model behaviour preserved. No retraining required to
adopt the bitmap as an inference-time conditioning input.
Signed-off-by: MiLO83 <edisonzghost@gmail.com>
Lyra 2 is 14B params (~28 GB bf16); does not fit a 16 GB consumer
card by default. This adds an opt-in --low-vram flag that swaps the
Linear layers to a smaller dtype, dropping weight VRAM enough to fit
on RTX 5060 Ti / 5070 / 4060 Ti 16G / etc.
Modes:
none Default. Existing bf16 behaviour preserved byte-identically.
int8 bitsandbytes Linear8bitLt swap. ~14 GB weights. Most-tested.
int4 bitsandbytes Linear4bit NF4. ~7 GB weights. Quality dip.
fp8 torch.float8_e4m3fn cast for Blackwell/Hopper tensor cores.
~14 GB weights. Storage-only path here; full fp8 matmul
needs transformer_engine.Linear swap (left as future work).
For typical 5060 Ti use:
python -m lyra_2._src.inference.lyra2_custom_traj_inference --low-vram int8 --low-vram-checkpoint --offload --offload_when_prompt --num_sampling_step 4 ... (the rest of your args)
That stacks: int8 weights (14 GB) + activation checkpointing (saves
30-50%) + existing CPU offload of non-active blocks + DMD-distilled
4-step inference. Expected: ~5-10 min per 80-frame chunk on a 5060 Ti
(vs ~15 sec on a GB200). Slow but tractable for personal research.
What this patch adds:
NEW lyra_2/_src/utils/low_vram.py
- apply_int8_quantization / apply_int4_quantization (bitsandbytes)
- apply_fp8_cast (torch.float8_e4m3fn storage)
- enable_gradient_checkpointing (auto-detect model API or wrap blocks)
- apply_low_vram_mode (one-call helper)
- requantize_checkpoint (one-time bf16 -> int8/int4/fp8 .pth converter)
- detect_checkpoint_quantization (read fp8 metadata from checkpoint)
- CLI entry: python -m lyra_2._src.utils.low_vram requantize ...
- Conservative default skip list: canonical_proj, output_proj,
final_layer, norm_out, x_embedder, t_embedder kept in bf16
(quality-critical paths, tiny in size anyway)
MOD lyra_2/_src/utils/model_loader.py
- load_model_from_checkpoint gains low_vram_mode + low_vram_grad_checkpoint
kwargs. Default 'none' preserves existing behaviour byte-identically.
- Applied AFTER weight load so quantization sees trained values.
MOD lyra_2/_src/inference/lyra2_custom_traj_inference.py
- --low-vram {none,int8,int4,fp8} CLI flag
- --low-vram-checkpoint flag
- Passed through to load_model_from_checkpoint
Honest caveats:
- Code is research-grade. Both files syntax-check via ast.parse but
the proposers have not run them against actual Lyra 2 weights
(no GB200, weights are research-licensed). The bitsandbytes
quantization pattern is the standard one used for LLMs and most
transformers; specific interactions with Lyra 2's Sparse3DCache,
canonical-coord conditioning, and forward_warp_multiframes are
uncharacterised. Expect to debug.
- Requantize tool itself needs the original bf16 model loaded
(28 GB+ VRAM), so consumer-card users typically rent an H100 for
an hour to do the one-time conversion, then download the smaller
file to use locally.
- Default skip list (conditioning + output paths) is a best guess.
A trained quantization expert should review which layers are
quality-critical for the canonical-coord conditioning path
introduced in commit 1.
Two of the four PR commits now require zero retraining (octant
prefilter + low-vram quantization), with the encoding-shift commits
(canonical-world-coords + uvw_atlas + OccupancyBitmap) needing
fine-tune to become functional.
Signed-off-by: MiLO83 <edisonzghost@gmail.com>
The existing requantize_checkpoint function loads the full Lyra 2 model
(needs ~28 GB VRAM, blocking consumer-card users). This adds a
state-dict-level streaming alternative that runs entirely on CPU + RAM,
no GPU required.
Estimated wall-clock on a typical 16 GB consumer card setup
(32 GB DDR5 RAM, NVMe SSD, 16 logical cores): ~2-3 minutes for the
14B Lyra 2. Memory-bandwidth-bound, not compute-bound; PyTorch's MKL
threading already uses all available cores for the per-tensor casts.
Why streaming works on consumer hardware:
- No nn.Module instantiation -- skip the ~14 GB graph-allocation step
- tensor.to(torch.float8_e4m3fn) is a pure cast, no matmul
- Per-Linear-weight quantization: compute scale, cast, store scale
as sibling buffer. Same math as the runtime apply_fp8_cast path
but applied to state_dict tensors instead of nn.Linear modules
- The .pth file format is loaded via torch.load(map_location='cpu',
weights_only=True) which is memory-efficient
Workflow now possible without cloud rental:
# On your 5060 Ti machine, ~2-3 minutes:
python -m lyra_2._src.utils.low_vram requantize-streaming --input lyra2_bf16.pth --output lyra2_fp8.pth --mode fp8
# Then use the fp8 checkpoint locally forever:
python -m lyra_2._src.inference.lyra2_custom_traj_inference --checkpoint_dir lyra2_fp8.pth --offload --num_sampling_step 4 ...
Implementation:
- requantize_streaming(input_path, output_path, mode, skip_modules,
progress) in lyra_2/_src/utils/low_vram.py
- Heuristic: quantize 2-D '.weight' tensors that aren't in the skip
list (canonical_proj, output_proj, final_layer, norm_out,
x_embedder, t_embedder) or named like normalization layers
- For fp8: per-tensor scale stored as sibling buffer (name +
'._fp8_scale') so the runtime path can reconstruct: w_bf16 =
w_fp8.to(bf16) * scale
- For int8/int4: stored as fp16 on disk; the runtime --low-vram
flag applies the bnb Linear swap. Disk savings are 2x; full
bnb-format-on-disk would require bitsandbytes at conversion time.
- Preserves outer-dict keys from the source (config, optim state) so
downstream loaders see the same checkpoint shape
- Marked as 'streaming: True' in the metadata for diagnostics
New CLI subcommand:
python -m lyra_2._src.utils.low_vram requantize-streaming --input X.pth --output Y.pth --mode {int8,int4,fp8}
The old 'requantize' subcommand still exists for cases where the
checkpoint is in DCP format (the streaming path is .pth-only;
DCP-format conversion requires the full-model load).
Untested by the proposers (no Lyra 2 weights to convert). The
syntax-checks pass; the tensor-iteration pattern matches what's used
in HuggingFace's safetensors-conversion tools.
Signed-off-by: MiLO83 <edisonzghost@gmail.com>
When converting a Lyra 2 checkpoint to fp8 via the streaming path,
also embed UVW bijection configuration as a metadata block. A UVW-aware
runtime auto-configures from this without per-load math; older loaders
ignore the unrecognised key.
What gets baked in (~few KB extra in the checkpoint):
- Atlas dimensions (4096x4096) and tile layout (16x16)
- Voxel resolution (256) and bits-per-axis (8)
- Default world extent placeholders (override at inference time)
- OccupancyBitmap allocation size hint
- Summary atlas allocation size hint
- 'uvw_compat_version: 1' marker so future runtimes can detect
- Pre-computed canonical-coord meshes for the standard
Lyra 2 inference config (480 x 832 x 5 spatial-history slots).
Saves the runtime a few ms per inference start.
Honest framing: this is a MARKER + CONFIG block, not a functional
swap. It does NOT change how the model interprets inputs -- that
still requires the conditioning-head fine-tune that mods nv-tlabs#1 and nv-tlabs#2
of this PR propose. What it DOES is ship a checkpoint that's paired
with our UVW system's settings, so a UVW-aware runtime knows which
member of the byte-width family applies and pre-allocates the right
sizes.
New CLI behaviour:
python -m lyra_2._src.utils.low_vram requantize-streaming --input lyra2_bf16.pth --output lyra2_fp8_uvw.pth --mode fp8
# UVW metadata baked by default; --no-bake-uvw to opt out
python -m lyra_2._src.utils.low_vram inspect lyra2_fp8_uvw.pth
# Now also prints the UVW compat block:
# UVW compat version: 1
# voxel resolution: 256^3
# bits per axis: 8
# atlas dims: 4096x4096
# occupancy bitmap: 2097.2 KB
# summary atlas (rgba8): 64.0 MB
# pre-computed meshes: ['480x832x5']
Implementation:
- _build_uvw_compat_metadata(...) constructs the metadata dict
including the canonical-coord mesh used by the existing
_build_canonical_spatial_coords in lyra2_model.py
- Stored under '_lyra_uvw_metadata' key alongside the existing
'_lyra_low_vram_metadata' block
- requantize_streaming gains a bake_uvw_metadata: bool = True
parameter; the new --no-bake-uvw CLI flag lets you opt out
- The 'inspect' subcommand now surfaces both quantization and UVW
compat info, including allocation-size hints for sanity-checking
that the checkpoint's UVW config matches your runtime expectations
Total now in this PR's low_vram.py: 880 lines, 20 functions.
Combined with the other 4 commits, this PR is now:
- Two architectural changes that need retraining (mod 1+2 of the
encoding shift) -- the genuine RFC pieces
- Three pure-utility additions that need zero retraining and could
land independently (octant prefilter, OccupancyBitmap, --low-vram
+ requantize-streaming + UVW metadata bake)
The pure-utility additions alone would lower Lyra 2's entry-cost-
to-use significantly: any 16 GB consumer card runs the 14B model
locally after a 2-3 minute on-device conversion, with the bijection
configuration already baked into the checkpoint for the UVW-aware
runtime path.
Signed-off-by: MiLO83 <edisonzghost@gmail.com>
The previous defaults (voxel_res=256, bits_per_axis=8, world_extent=4m)
were DownToEarth's hamlet-scale settings, NOT appropriate for Lyra 2:
- Lyra 2's documented spec is 90m walkable
- 256-cubed at 90m world = 35 cm cells, far too coarse for photoreal
- For ~2cm cells in a 90m world we need 4096-cubed, which needs RGBA16
addressing (12 bits per axis, comfortably under 16)
New Lyra-2-appropriate defaults:
voxel_res = 4096 # 2.2 cm cells in a 90m world
bits_per_axis = 16 # RGBA16UI, plenty of headroom (max 65k)
world_extent_m = 45.0 # 90m cube, matches Lyra 2 spec
storage_mode = 'sparse-tiled' # 4096-cubed dense = 256 GB; impractical
tile_voxel_res = 256 # leaf tiles match DownToEarth's bijection
Metadata bumped to version 2 with the new scene-scale-aware fields:
- cell_size_cm : derived from voxel_res + world_extent_m
- total_voxels_addressable: voxel_res^3
- storage_mode : dense | sparse-tiled | hierarchical-octree
- tile_voxel_res : per-leaf-tile resolution for sparse mode
- tile_count_per_axis : voxel_res / tile_voxel_res
- per_tile_summary_atlas_bytes, per_tile_occupancy_bitmap_bytes
- dense_atlas_bytes_if_populated: heads-up for sparse-storage planning
CLI flags added so the function is now flexible enough to cover the
whole byte-width family from DownToEarth hamlet scale to planet scale:
python -m lyra_2._src.utils.low_vram requantize-streaming --input lyra2_bf16.pth --output lyra2_fp8.pth --mode fp8 --uvw-voxel-res 4096 # default; matches Lyra 2 spec
--uvw-bits-per-axis 16 # default; RGBA16UI
--uvw-world-extent-m 45.0 # default; 90m cube
# DownToEarth hamlet scale (the previous defaults):
... --uvw-voxel-res 256 --uvw-bits-per-axis 8 --uvw-world-extent-m 4.0
# Sub-cm Lyra 2 detail:
... --uvw-voxel-res 8192 --uvw-bits-per-axis 16
# Continent / planet scale (with hierarchical sparse storage):
... --uvw-voxel-res 1048576 --uvw-bits-per-axis 32
The inspect subcommand now surfaces all the new fields:
Checkpoint quantization: fp8
UVW compat version: 2
voxel resolution: 4096^3 (68,719,476,736 voxels)
bits per axis: 16
world extent: 45.0 m half-width (90 m cube)
cell size: 2.20 cm
storage mode: sparse-tiled
tile voxel res: 256^3
tile grid: 16^3 = 4096 addressable tiles
per-tile summary atlas: 64.0 MB (RGBA8)
per-tile occ bitmap: 2.00 MB
dense atlas if populated: 256.0 GB (sparse at ~5% = 12.8 GB realistic)
The legacy v1 metadata format (256-cubed hamlet defaults) still parses
correctly via the inspect command's version branch -- previously-baked
checkpoints continue to be readable, just show the older fields.
Total low_vram.py now 979 lines, 20 functions.
Signed-off-by: MiLO83 <edisonzghost@gmail.com>
…ram path
Three coordinated changes that build on the existing OccupancyBitmap +
low-vram work:
1. OccupancyBitmap 2-bit mode (uvw_atlas.py)
- Adds bits_per_voxel=2 constructor flag (default still 1, fully
backward compatible).
- Per voxel: bit 0 = occupancy (unchanged), bit 1 = ray-touched
this frame.
- New API: set_touched / get_touched / clear_touched (masked-AND
preserves occupancy in ~10us), count_touched, touched_chunks()
which OR-reduces per-voxel touches to chunk keys.
- Storage: 4 MB at 4096x4096 atlas (was 2 MB at 1-bit).
- Enables UE5-Nanite-style demand-driven sparse voxel streaming:
the renderer announces what it needed via the touched bits, the
streamer keeps those chunks resident. Correct for VR/360 because
no view-direction is assumed.
2. CPU-first model_loader path (model_loader.py)
- When low_vram_mode != 'none', model is instantiated on CPU
instead of GPU. Weights load to CPU, on_train_start casts to bf16
on CPU, apply_low_vram_mode quantizes to int8/int4/fp8 on CPU,
then model.cuda() moves the already-quantized model to GPU.
- Avoids the ~28 GB peak GPU memory during bf16 weight load that
made 14B Lyra 2 OOM on 16 GB consumer cards (RTX 5060 Ti / 5070).
- Default path (no low-vram) is unchanged.
3. CPU-source routing in low_vram quantizers (low_vram.py)
- apply_int8_quantization and apply_int4_quantization now detect
CPU-source weights and route through .cuda(target_dev) on a
resolved CUDA device rather than .cuda(child.weight.device)
which fails when the child is on CPU.
- Added explicit del child after setattr to release the source
tensor as we walk the model, keeping CPU RAM bounded during the
quantization sweep.
Together these three changes make int8 Lyra 2 inference on a 16 GB
consumer GPU actually work end-to-end, plus add the on-GPU half of
demand-driven streaming for future runtime renderers.
DCO-signed.
Signed-off-by: MiLO83 <simdiff@gmail.com>
set/get/set_touched/get_touched accept atlas (atlas_x, atlas_y) coords. Production raymarchers think in voxel grid coords (u, v, w). Add explicit bijection-keyed helpers so the canonical lookup path goes through the UVW <-> RGB <-> XYZ bidirectional bijection (the whole point of this proposal). New methods: - set_by_voxel(u, v, w) : occupancy via voxel coords - get_by_voxel(u, v, w) : occupancy via voxel coords - set_touched_by_voxel(u, v, w) : render-flag via voxel coords - get_touched_by_voxel(u, v, w) : render-flag via voxel coords Internally these call voxel_to_atlas() and forward to the existing atlas-coord set/get. The atlas-coord forms remain available for code that already has atlas coordinates from earlier in the pipeline. Why this matters for the runtime: the raymarcher hits a voxel at (u, v, w), computes atlas coords ONCE via the bijection, then uses the same atlas address for both reads (the voxel's RGBA from the main atlas) and writes (the render-flag bit in this bitmap). One bijection lookup serves both. The atlas memory layout for the render-flag is structurally identical to the layout for the voxel data -- they're the same address space. Verified end-to-end: voxel (50, 70, 40) -> atlas (2098, 582) -> bit position 0 of byte 524 in row 582; round-trips through both APIs. DCO-signed. Signed-off-by: MiLO83 <simdiff@gmail.com>
…uality Adds LYRA2_OUTPUT_CODEC env override; defaults to libx265. Per AV1-MV-Fidelity (arXiv:2510.17427), codec MVs are 2-3x noisier on H.264 than on HEVC, and HEVC sits in the practical sweet spot (encoder speed vs. MV fidelity) for content that downstream tools extract motion vectors from. Switching the default from libx264 to libx265 makes outputs immediately useful for codec-MV-driven analysis (static/dynamic voxel classification, scene flow, etc.) without changing the visual quality at standard CRF. Override via env var: LYRA2_OUTPUT_CODEC=libsvtav1 # best MV fidelity, slowest encode LYRA2_OUTPUT_CODEC=libx265 # default; balanced LYRA2_OUTPUT_CODEC=libx264 # legacy, fastest encode DCO-signed. Signed-off-by: MiLO83 <simdiff@gmail.com>
…ncyBitmap 2-bit pattern Ready-to-paste shader snippets for wiring the 2-bit OccupancyBitmap into a runtime voxel renderer. Pure docs/utility, no code-path changes. Snippets: - GLSL voxel<->atlas bijection (forward + inverse) - GLSL 2-bit occupancy + render-flag read - GLSL imageAtomicOr render-flag write (OpenGL 4.5+ / Vulkan) - GLSL DDA voxel raymarch reference loop - WGSL versions for WebGPU - CUDA versions for native runtime kernels - Python per-frame orchestrator using OccupancyBitmap.clear_touched + touched_chunks (matches the streaming_reference.py reference impl) All snippets demonstrate that one voxel_to_atlas(u, v, w) lookup serves both the RGBA color read and the render-flag write -- the runtime payoff of the UVW bijection. DCO-signed. Signed-off-by: MiLO83 <simdiff@gmail.com>
Drop-in helper that lets inference scripts opt into entity-structured
captions without taking a dependency on the bigger vocabulary
infrastructure (which lives in the lyra-2-lite branch).
Default behaviour is pass-through: caption goes to T5 unchanged. When
the caller supplies an entities list + optional vocab JSON path, the
hook expands the caption with canonical entity descriptions before T5
encoding.
Quality effect when active (estimated from conditional-generation
literature): ~10% reduction in caption-space ambiguity, ~1.2x faster
convergence per chunk. No cost when not active.
Usage in lyra2_custom_traj_inference.py:
raw = json_captions[str(chunk_start_frame)]
caption, meta = lyra2_caption_hook(
raw,
entities=json_entities.get(str(chunk_start_frame)),
vocab_path=args.entity_vocab,
)
DCO-signed.
Signed-off-by: MiLO83 <simdiff@gmail.com>
Pure-Python reference implementation of the runtime side of the 2-bit
OccupancyBitmap demand-streaming pattern. Documentation-by-code for
anyone implementing the actual GPU runtime.
Includes:
- ChunkLoader protocol + DiskChunkLoader concrete impl (ThreadPool
for parallel reads)
- ChunkCache with hysteresis eviction (cold_frames counter; LRU bump
on touch; max_resident cap)
- StreamingSession context manager wrapping the per-frame loop:
enter: bitmap.clear_touched()
body: raymarcher writes touched bits via shader
exit: bitmap.touched_chunks() -> cache.mark_touched_many() ->
cache.frame_complete()
Runs without a GPU. Intended pairing: the GLSL/WGSL/CUDA shader
snippets in voxel_renderer_shaders.py for the GPU side, this module
for the CPU orchestrator.
DCO-signed.
Signed-off-by: MiLO83 <simdiff@gmail.com>
MiLO83
added a commit
to MiLO83/lyra
that referenced
this pull request
May 20, 2026
Brings the broader Lyra-2-Lite work into the fork as a sibling subtree to NVIDIA's lyra_2/ . The subtree contains the architectural extensions, demos, tooling, and documentation built on top of NVIDIA Lyra 2 + the PR-61 contributions. Sources: - entity_vocab.py + structured_prompt.py: variable-byte (1/2/3 byte) semantic vocabulary with self-tuning byte-width assignment, drives structured-prompt conditioning + per-voxel semantic ID + downstream segment-renderer ID. Three-line drop-in to any Lyra 2 inference script via the lyra2_caption_hook helper. - posterize/: bit-depth ladder dataset pipeline (LYRA2_PROPOSAL §6.6.2). Per-channel RGB rungs 1-BPP to 8-BPP, TinyUNet deposterizer trainer, bounded-noise schedule, sample.py multi-step inference. LUT-based posterizer hits 168 fps at 1080p. Dataset of 2000 cinematic-aspect picsum images is regeneratable via scrape_dataset.py (gitignored). - voxel_renderer/: WebGL2 single-file demo of the runtime raymarch pipeline. DDA voxel traversal + per-chunk frustum culling. ~628 lines, 60+ fps on consumer Blackwell. Standalone, open in any browser. - postprocess/rife_upsample.py: RIFE 4.26 wrapper for 12 fps native generation -> 60 fps playback. ~5x speedup half of the rendering story. fp16 on Blackwell. - scripts/setup_lyra2_wsl.sh: end-to-end WSL2 Ubuntu setup for running Lyra 2 inference on a Windows + RTX 5060 Ti host. Idempotent, survives multiple patch iterations. - docs/: the canonical PROPOSAL.md (also filed as PR nv-tlabs#61), the layperson DUMMIES.md, the MOTION_VECTORS_NOTE.md research design, and SESSION_TODO.md tracker. License: Apache 2.0 for the subtree, matching the upstream fork. Files that originated in the DownToEarth project (MIT) have been relicensed to Apache 2.0 for inclusion; copyright holder (MiLO) is unchanged. This subtree is intentionally NOT part of PR nv-tlabs#61. The PR remains surgical: only the canonical-coord encoding + OccupancyBitmap + low-vram changes to NVIDIA's lyra_2/ tree. This branch (lyra-2-lite) is the comprehensive project home. DCO-signed. Signed-off-by: MiLO83 <simdiff@gmail.com>
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
1 participant
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.
Hi Lyra team — thanks for releasing Lyra 1 + 2 under Apache 2.0 with this much detail. Studying the code has been a treat.
What this RFC is about
Lyra 2's canonical-coord conditioning encodes per-pixel image-space coordinates + a frame-slot index (
u_norm, v_norm, frame_slot), and the cross-attention attends over a frame-keyedSparse3DCache. The cache already stores per-pixel world coordinates (viaunproject_pointsinSparse3DCache.add()) — the world coords are computed and persisted, but not used as the canonical key.This RFC explores: encode the canonical-coord image as quantized world coordinates packed in RGB bytes (or uint16 / uint32) via a bidirectional bijection. Same warp machinery (
forward_warp_multiframesis dtype-agnostic), different output encoding. Makes geometric correspondence structural — one key per 3D point — rather than learned across frame-pixel pairs.Full write-up with byte-width family, before/after performance table, and porting sketch:
What this patch adds (3 modifications, all opt-in)
lyra_2/_src/datasets/uvw_atlas.py— the bijection family (RGBA8 → RGBA32UI), pure arithmetic, doctested + 64³ exhaustive round-trip self-test passing inside the tree. Apache 2.0 header to match repo._build_canonical_world_coordsstatic method inlyra2_model.py— sits alongside the existing_build_canonical_spatial_coords. Opt-in, not wired into_get_cached_spatial_coords; existing model behaviour unchanged until the dispatch is added.octant_prefilter: bool = Falsekwarg onSparse3DCache.retrieve()— when True, prunes candidates whose centroid is behind all target views before the matmul. Clean index remapping via_orig_idx. Default off → byte-identical existing behaviour.Modification 3 is the smallest possible win — pure inference-time optimization, ~50% reduction in
retrieve()'s matmul cost at long horizons with zero quality impact (those candidates contribute zero valid pixels via the existingz_all > 0filter anyway). If everything else here is interesting only as discussion, that one's potentially landable on its own merits.Why it might be worth a look
Honest trade-offs
world_max. A hierarchical sparse-tiled variant is sketched in the proposal but not in this patch.Verification
python -m lyra_2._src.datasets.uvw_atlas→ 6/6 doctests pass, exhaustive 64³ bijection round-trip OKast.parse()syntax checkgit diff mainis byte-identity-preserving in default code pathDCO
Signed-off per
CONTRIBUTING.md. Single commit, single author.Suggested smallest experiment
If the encoding shift is interesting enough to consider for a future training run, the cheapest possible probe is modification 3 alone — no retraining required, expected ~5-15% wall-clock speedup at long horizons. If that holds, the case for considering the encoding shift (modifications 1 + 2) gets stronger.
Marking draft per the contributing guide. Open to discuss, refine, retract, or close at the team's call. Happy to defer entirely if there are reasons not to pursue this direction — even a "we considered this and decided X" reply would be useful public documentation for anyone else looking at similar architectures.
Thanks for reading 🙏
Background on the proposer: this grew out of building a Lyra-2-shaped local pipeline (voxgaussian) on consumer hardware. ~167k Gaussian splats from a single Juggernaut XL render in ~5 minutes end-to-end (live demo). The UVW↔RGB bijection landed as a clean little data structure with one property that seemed worth surfacing upstream: the bytes are the coords, both ways, at zero VRAM for the identity mapping.