LTX-2.3 video generation: clean-room implementation with GPU-poor support#16
Merged
Conversation
The LTX-2.0 port is superseded by a ground-up LTX-2.3 implementation written from the Apache-2.0 diffusers reference. Generic VRAM utilities and the SDXL memory profile move to R/vram.R and R/memory_sdxl.R; encode_with_gemma3() now returns raw stacked hidden states for the downstream connector modules. Kept: Gemma3 encoder, BPE tokenizer, FlowMatch scheduler.
- DESCRIPTION: cph entries for cornball.ai, HuggingFace (diffusers ports), and Lightricks (constants/format facts); ORCID for maintainer - inst/COPYRIGHTS: enumerates derivation sources and weight licensing - ref/ and tasks/ gitignored + Rbuildignored; upstream reference checkouts live in ref/upstream/ and are never committed - CLAUDE.md: LTX section rewritten for the 2.3 clean-room rewrite; stale 2.0 API docs removed, generic implementation facts kept - Drop tracked pyrotechnics conversions in ref/ (superseded by the real diffusers checkout in ref/upstream/)
ltx23_open_checkpoint() validates model_version metadata and parses the embedded component config; ltx23_split_keys() routes the flat key space to component groups (dit, connectors, vae, audio_vae, vocoder); ltx23_load_group() streams tensors one at a time into module params with two-sided coverage reporting. Tested against a tiny fake checkpoint with official-style keys.
Split and interleaved RoPE application plus the audio/video coordinate and frequency embedder (float64 frequencies per checkpoint config). Parity-tested against fixtures generated from the diffusers reference via tools/gen_fixtures.sh (uv-run, CPU torch; not used at test time).
Dual-stream transformer with gated attention, 9-param cross-attention modulation, prompt AdaLN, split RoPE, a2v/v2a cross attention with global+per-block modulation, STG perturbation, and optional query chunking for the manual SDPA. Includes the official-checkpoint key mapper. Parity-tested against a tiny reference model at 1e-4.
Per-token RMS norm over raw stacked Gemma3 states, per-modality sqrt-ratio scaling and projections, learnable-register padding replacement with stable front-alignment, per-connector 1D split RoPE, and 8-layer gated connector transformers. Parity-tested at 1e-4.
Per-channel RMS norm, causal 3D convs, pixel-shuffle up/downsamplers, 2.3 block structure (encoder last block 1024; 4-up-block decoder with mixed spatiotemporal/temporal/spatial upsampling, no residuals, zeros padding), latent statistics buffers, and the flat-to-nested official key mapper. Parity-tested at 1e-5.
Audio VAE: causal conv2d (height axis), pixel norm, nin_shortcut resnets, nearest-upsample stages, 4F-3 frame crop. Vocoder: BigVGAN style with snakebeta + anti-aliased kaiser-sinc resampling (base-R besselI window), 3-branch-mean resnets, checkpoint-loaded STFT/mel bases, BWE residual over a hann-resampled skip to 48kHz. Encoder deliberately not ported (t2v never encodes audio).
txt2vid_ltx2(): Gemma3 + connectors text encoding, joint audio/video denoising over the official 8-step distilled schedule (CFG-free, Euler velocity steps in float32, hoisted RoPE coords), packed-latent audio statistics handling, video/audio decoding, base-R WAV writer, and av muxing. ltx23_load_pipeline() streams all components from the single checkpoint with strict key coverage. Smoke-tested end-to-end with tiny random-weight components.
ltx23_quantize_fp8() streams the 46GB checkpoint into fp8 shards using the official cast policy (DiT attention/FFN linears, absmax/448 scales); ltx23_load_transformer_fp8() swaps those linears for CPU-resident (optionally pinned) fp8 modules that dequantize on the compute device. ltx23_memory_profile() + ltx23_tune_gc() port the whisper allocator tuning; download_ltx2() adds consent-gated hfhub downloads. Also fixes R-scalar dtype promotion throughout the LTX modules (infix scalar arithmetic promotes bf16 to fp32; switched to tensor methods), which broke bf16 inference.
ltx23_latent_upsampler ports the 2.3 upscaler (Conv3d ResBlock stages around a per-frame 2x pixel shuffle; no rational resampler) with AdaIN and tone-map latent filters from the Apache diffusers reference. txt2vid_ltx2(two_stage = TRUE) denoises at half resolution, upsamples the unnormalized latents, re-noises BOTH modalities at the stage-2 entry sigma, and refines over the stage-2 schedule. ltx23_load_pipeline() now accepts the fp8 artifact directory and loads the transformer with CPU-resident streaming fp8 weights.
Parses the diffusers reference and the R port with tree-sitter and checks every top-level reference entity against a curated mapping that also documents deliberate omissions. Fix the Gemma3 text encoder shard names in download_ltx2().
Repair the corrupted roxygen blocks in CLIPTokenizer.R and filename_from_prompt.R (stray backtick comment lines), switch Blackwell detection to torch::cuda_get_device_capability (0-based device index), and complete missing @param entries. Check is now clean apart from the CRAN-incoming and Rd-width NOTEs.
Connectors, VAEs, vocoder, and upsampler now move to the compute device only for their phase and return to the CPU afterwards, leaving the denoise phase as the sole GPU tenant (frees several GB of resident weights during the transformer steps).
QLoRA-style 4-bit NormalFloat: per-64-block absmax against the 16-level quantile code, packed nibble pairs, pure-torch bucketize/index_select quantize and dequantize (no custom kernels). At ~4.5 bits/param the 22B transformer stays resident on a 16GB card, removing the ~21GB/step PCIe streaming of the fp8 path (~9% weight round-trip error vs fp8's ~1%). ltx23_load_pipeline dispatches on the artifact manifest format.
high = NF4-resident transformer + phase offloading (measured baseline: fp8 streaming peaked 11.6GB at 512x320x49 without offloading); medium/low = fp8 CPU-streaming with pinned memory and tighter caps.
With ~11.4GB of resident NF4 weights, the per-linear dequantized bf16 weights (up to 128MB each, 28 per block) piled up faster than R's GC collected them and OOM'd the first step. Each distinct weight shape now dequantizes into one long-lived buffer, released at the denoise/decode boundary.
The nibble/index/value temporaries (up to ~300MB per linear, 28 linears per block) were R garbage between block-level gc calls and still OOM'd with 11.4GB resident. Unpack, index lookup (torch_index_select_out with an allocating fallback), scaling, and the output write now all run in-place against per-device scratch buffers.
The text connectors are 3.3B params (6.35GB bf16) and cannot share a 16GB card with the resident NF4 transformer (11.4GB): the earlier OOMs happened during connector onload, not denoising (a bare NF4 forward runs flat at 12.0GB across all 48 blocks). Each phase is now truly the sole GPU tenant: text encode alone, then the transformer onloads for denoising and offloads before decode (~3s of transfers per generation).
The 768x512x121 target hit the unchunked [1,32,6144,6144] attention matrix (2.25GB). Query chunks are now sized from a memory budget (diffuseR.attn_budget, default 1GB) against the actual key length, with any explicit attn_chunk acting as an upper bound.
At 768x512x121 (S=6144) the chunked attention temporaries were multi-GB R garbage per block and fragmented the allocator (1.7GB reserved but unallocated at OOM). The score matrix and context output now compute in-place into reusable per-shape buffers via torch's internal matmul_out/softmax_out (with an allocating fallback), released together with the dequant buffers at the denoise/decode boundary.
ltx23_load_pipeline defaults PYTORCH_CUDA_ALLOC_CONF to expandable_segments before the first CUDA allocation; the high profile caps at the hardware-validated 768x512x121 (peak 15.7GB on a 16GB card) with adaptive attention chunking handling the budget.
Port the diffusers spatial + temporal tiled decode (overlapping tiles, crossfaded seams; reference defaults 512px/448 stride, 16f/8 stride) -- required at 1280x704x121 where full-decode activations run ~7GB per tensor. ltx23_tune_gc now follows the documented allocator semantics: raising cuda_allocator_reserved_rate disables collection for large-resident workloads, so it instead lowers allocated_rate and defaults expandable segments.
Idea-by-idea references for the LTX-2.3 implementation, all independent of Wan2GP: diffusers/transformers module ports, QLoRA NF4, official Lightricks fp8 policy facts, diffusers offloading/tiling/slicing, the mlverse allocator docs, and our own measured engineering. Also: tune_gc follows the storm-stopper recipe (A/B measured inert here: 91s vs 106s per step with R-gc ~86-89% either way -- the cost is our explicit per-block gc syncs, not the allocator callback; JIT-compiling the block per the torch skill is the roadmap), and per-block gc is now the opt-out diffuseR.block_gc option (required for both quantized paths at high resolution).
nnf_gelu allocates a second copy of its input, which at 1280x704 (14080 video tokens) doubles the 440MB FFN intermediate -- that allocation was the OOM point on the 16GB card. Above a size threshold (diffuseR.gelu_inplace_min, default 1e8 elements) the proj output is now mutated in place via the internal torch_gelu_ kernel, chunked so internal temporaries stay small, with the allocating path as fallback. Parity vs nnf_gelu is exact (tested).
Validated on the RTX 5060 Ti 16GB: 1280x704x121f with audio in 1354s at a 15.7GB peak. The 5e8 default put ~1GB of score scratch on the card at 14080 video tokens (one ~500MB buffer per attention shape); 1.5e8 keeps the pair under ~400MB and is a no-op below ~2000 tokens where the sequence fits in one chunk anyway.
The forward collected the state after every layer AND the post-norm output -- 50 entries for the 48-layer model. HF transformers (the reference and what the LTX connectors were trained against) records the state BEFORE each layer plus the post-final-norm output: 49 entries, with the un-normed last-layer output never included. The extra state made the connector projection fail its first matmul (1024x192000 vs 188160x4096) on the first real-prompt encode. Regression test with a tiny config pins the count and the exact identity of the last collected state.
The venv rode in with the original LTX-2.0 commit and its lib64 symlink dangles once uv relinks the interpreter, which breaks the R CMD build copy step. Venvs are regenerable; ignore them everywhere.
The describe-block '9\%' double escape truncated ltx23_memory_profile.Rd mid-item (reworded); the index_select_out resolver now returns the function into a local before calling so codetools sees a callable binding. Check: 0 errors, 0 warnings, 2 benign NOTEs.
The 48-block step now runs as one torch::jit_compile'd TorchScript call (the JIT-decode pattern from the torch skill, proven in whisper and chatterbox): weights passed per call as a flat List[Tensor] with a fixed 114-slot per-block layout, NF4 dequant chunked in-script, attention through fused scaled_dot_product_attention instead of a materialized score matrix, and no R-side intermediates -- which removes the per-block gc() syncs that profiling showed were ~86% of eager step time. Gated on diffuseR.jit_blocks (default TRUE) and NF4-quantized 2.3 blocks with no STG/self-mask/isolation extras; anything else falls back to the eager loop. Parity tests: single block and 2-block stack vs eager at 1e-4 (fp32 CPU), mask effectiveness, NULL masks.
Two integration bugs the tiny-config test missed: the split-rope freqs from the model are per-head 4D [B, H, T, r] (the script now mirrors both branches of the eager apply; tests cover 4D and 3D), and nn_module_list children are named, which made the packed weight list marshal as Dict[str, Tensor] instead of List[Tensor] (unnamed now). Measured on the RTX 5060 Ti at the 768x512x121 shape (S=6144): eager 99.9s/step at 88% R-gc; compiled 8.8s/step at 1% R-gc.
Contributor
Author
|
Roadmap follow-ups pushed to this branch: TorchScript-compiled block stack (
NF4 vs fp8 quality A/B (same seed, 1280x704, files in the render dir): fp8 is crisper on fine geometry, NF4 slightly softer with minor composition drift, both prompt-faithful and artifact-free. fp8 streaming was also faster than eager NF4 at this resolution (882s vs 1331s), but NF4+JIT now leads at 389s. Full suite: 346 tests, 0 failures. |
The video decoder, audio decoder, and vocoder are static feed-forward graphs, so torch::jit_trace converts them wholesale: one crossing per tile/forward, libtorch frees intermediates eagerly. Traces are shape-specialized (runtime sizes bake in as constants; wrong shapes error), so they cache per instance/shape/dtype/device/tag and re-trace on miss -- at 1280x704 the tiled decode has ~4 distinct tile shapes. A trace pins the weight tensors it captured, so the pipeline drops all traces on every phase offload. diffuseR.jit_vae (default TRUE) with the eager modules as fallback; parity exact on tiny configs (10 tests: direct, tiled, retrace-on-new-shape, cache reuse, release, audio decoder, BWE vocoder).
Profiling the 1280x704 decode phase: 32.5s eager with 86% of it in R gc -- and the storm survives every allocator option because it is freeing work, not trigger frequency (each gc releases a tile's worth of CUDA blocks; expandable_segments makes each free a page-unmap). PYTORCH_CUDA_ALLOC_CONF=backend:native alone cuts decode to 21.0s / 51% gc. jit_trace of the decoders works and is guarded three ways after root-causing a lantern tracer hazard: R gc firing mid-recording (via the allocator callback under memory pressure) corrupts captured argument values -- observed as garbage/swapped narrow starts on the full-size decoder, 5/5 clean re-traces once gc cannot fire mid-trace. The traced path now gcs before every trace, tryCatches, and validates each fresh trace against eager output (mismatch = permanent per-shape blacklist), so silent corruption is impossible. It still re-pays trace+validation every render (traces release on phase offload, they pin captured GPU weights), which measured slower than eager for a single decode (39.8s vs 32.6s), so it ships opt-in via options(diffuseR.jit_vae = TRUE). Also: ltx23_load_pipeline now applies ltx23_tune_gc; .detect_vram falls back to nvidia-smi instead of guessing 8GB when gpuctl is absent (the guess silently skewed the tuned rate); per-tile gc in the tiled decode is gated to the eager path; verbose renders print per-phase timings (video decode / array conversion / audio chain / encode+mux). 356 tests, 0 failures.
A/B on the 1280x704 tiled decode (fresh process per condition, per the mlverse memory-management article): the two callback gates we had never set -- torch.cuda_allocator_allocated_rate and torch.cuda_allocator_allocated_reserved_rate (defaults 0.8) -- take the decode from 32.7s to 21.5s under expandable segments and cut the R-gc share from 50% to 32% on the native backend, with no wall-time downside anywhere. ltx23_tune_gc now sets both to 0.95 (only-if-unset). The loader's PYTORCH_CUDA_ALLOC_CONF default is now per-format: nf4 gets backend:native (the compiled block stack keeps intermediates out of R, native frees are cheap, and the 1280x704 fit is validated -- full corgi render 327s vs 389s under expandable); fp8/eager keep expandable_segments, the only configuration their max-resolution fit was validated under.
The ~110s of per-render 'phase traffic' was never PCIe transfer: it was allocator pool regrowth. clear_vram()'s cuda_empty_cache between phases returned every block to the driver, so each phase re-grew the pool one cudaMalloc per tensor (~15ms x 5530 = 83.5s measured for the NF4 transformer's onload alone; 4.6s after a pre-warm; back to 83.3s after one empty_cache). The loader now pre-warms the pool with a single footprint-sized allocation freed into the cache, and the pipeline's phase offloads gc without emptying the cache so the next phase carves from cached blocks. Corgi 1280x704x121: 327s -> 221s (video decode phase 35.6s -> 24.9s from the same fix). Pinned staging is implemented and parity-tested (pin once at load, non-blocking DMA onload, offload as a pointer swap since inference never mutates weights) but ships opt-in (diffuseR.pin_staging): page-locking ~20GB adds ~30s to load and saves only ~2.7s per render once transfers are no longer the bottleneck -- breaks even around 11 renders per pipeline.
Port of the frame-conditioning mechanics from the Apache-2.0 diffusers reference (pipeline_ltx2_image2video.py, pipeline_ltx2_condition.py), restricted to prefix conditioning at latent index 0 with strength 1: txt2vid_ltx2(image=) seeds the first frame from a PNG/JPEG or array; txt2vid_ltx2(condition_video=, conditioning_frames=) freezes the tail of a previous clip as the new clip's opening frames (the continuation primitive behind chained talking-head chunks). Conditioned tokens are VAE-encoded (argmax + per-channel normalization), see a per-token video timestep of zero, and are frozen through the Euler updates; audio_timestep stays scalar. The JIT block stack handles the per-token modulation unchanged (parity-tested). Real-weight validation: encoder round trip on a 1280x704 frame at 38.8 dB PSNR. 20 helper unit tests (preprocess geometry, mask placement, partial-noise blend, argmax encode), 6 pipeline smoke tests, 2 JIT per-token parity tests.
Lip-sync support: txt2vid_ltx2(audio=) encodes user audio (MP3/WAV via av, or a [2, samples] matrix) into clean audio latents that the video attends to while denoising -- the audio stream sees timestep/sigma zero and skips its Euler updates (the audio-side mirror of the Apache i2v conditioning pattern), and the original samples are muxed into the output. New pieces: the LTX2AudioEncoder port (Apache diffusers reference; the checkpoint's 44 encoder tensors now load, audio_vae census 102/102, 0 unmapped / 0 unfilled) and a 16 kHz log-mel frontend built from the checkpoint's preprocessing spec (STFT 1024/160 causal, 64 slaney-normed mel bins to 8 kHz). Both basis constructors were verified against the checkpoint's stored vocoder bases (identical up to bf16 rounding), and the full chain validated on real speech: 4.5s TTS through encode -> decode -> vocoder returns 48 kHz audio with 0.91 envelope correlation to the input. Tests: 13 encoder structure/shape tests, 11 frontend/geometry tests, 4 pipeline smoke tests (audio-conditioned run, original audio carried through).
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.
What
Ground-up LTX-2.3 audio+video generation, written from the Apache-2.0 HuggingFace diffusers reference. Replaces the previous LTX-2.0 port entirely (removed in the first commits of this branch).
txt2vid_ltx2()— 22B dual-stream DiT (48 blocks, split RoPE, gated attention, 2.3 prompt/cross-attn modulation), text connectors with learnable registers, video VAE with spatial+temporal tiled decode, audio VAE, BigVGAN-style vocoder with bandwidth extension to 48kHz, Gemma3-12B text encoder, single-stage distilled (8 steps, CFG=1) and two-stage HQ with the x2 latent upsampler.inst/REFERENCES.md; attribution ininst/COPYRIGHTSand DESCRIPTION. Weights are user-downloaded from Lightricks under the LTX-2 Community License, never redistributed.Tests / checks
R CMD check: 0 errors, 0 warnings, 2 benign NOTEs (test-fixture size, one long usage line)tools/ltx23_compare.R): every diffusers LTX-2.3 entity ported or a documented skipVersion bumped to 0.1.0.