FLUX.1-schnell text-to-image support#20
Merged
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 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.
…cks (parity-tested)
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.
Native FLUX.1-schnell support: 4-step distilled, CFG-free text-to-image on consumer GPUs.
What's new
flux_transformer): 19 double-stream + 38 single-stream blocks, joint attention with per-head RMS q/k norms, adaLN-Zero conditioning, interleaved RoPE over (txt, row, col) position ids. Module tree mirrors the diffusers state dict 1:1 (no key remapping).t5_encoder): relative position bias, unscaled/unbiased attention, gated-GELU FFN; runs float32 on CPU by default.unigram_tokenizer): Viterbi best-path segmentation, normalizer/Metaspace settings read from tokenizer.json.flux_quantize,flux_load_transformer): NF4 (~6.8 GB, GPU-resident) or fp8 (~12 GB, CPU-resident/streamed), reusing the LTX-2.3 machinery; cast census is exactly 494 weights.txt2img_flux,flux_load_pipeline,download_flux1): per-phase GPU offloading, VRAM-based profiles, gated-repo download with actionable auth errors, one-time local quantize.Verification
tools/gen_fixtures_flux*.py).tokenizer_2/tokenizer.json(25 prompts + padding/truncation cases).tools/compare_translation.R): no unintentional drift vs the Python reference.R CMD check: 0 errors, 0 warnings.Provenance
All ported code traces to Apache-2.0 sources (diffusers, HF transformers), recorded per-technique in
inst/REFERENCES.md. Weights are user-downloaded from the gatedblack-forest-labs/FLUX.1-schnellrepo (Apache-2.0) and never redistributed.Notes
feat/ltx23(reuses its quantization/offload/checkpoint machinery); the diff here is FLUX-only.max_sequence_lengthdefaults to 256 (schnell-documented) vs diffusers' 512; per-call override.