FLUX.2 Klein 4B text-to-image support#21
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.
…r tuning Generation was 93.8% R gc (Rprof): torch's allocator callbacks fire hundreds of gcs per generation, and the tokenizer's 151k-binding vocab and merge environments made each one 17x slower. The tokenizer now keeps three atomic vectors and a byte table (BPE on integer ids via findInterval); fp8 linears dequantize into persistent buffers; klein runs backend:native with the footprint sized to the largest phase. 512px: 133.5s -> 39.3s; 1024px: 193s -> 48.1s (peak 8.22GB).
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.
FLUX.2 klein-4B: step-distilled text-to-image (4 steps, no CFG) on consumer GPUs. Stacked on the FLUX.1 branch (#17); the diff here is FLUX.2-only.
What's new
flux2_transformer): 5 double + 20 single blocks with model-level shared modulation, SwiGLU feed-forwards (gate fused intolinear_in), ViT-22B-style parallel single blocks (fused QKV+MLP projections), 4-axis RoPE (theta 2000), all linears bias-free.qwen3_encoder): GQA with per-head q/k norms, causal + padding mask, mid-stack hidden states (9/18/27) concatenated per token; runs only as deep as the last requested layer.qwen_bpe_tokenizer): GPT-2 byte table + GPT-4-style split regex + rank-based merges over integer ids (three atomic vectors, findInterval lookups — see perf note), with the Qwen3 chat template.flux2_vae_decoder): post_quant_conv + reused decoder body at 32 latent channels + BatchNorm latent statistics (FLUX.2 has no scalar scaling factor).txt2img_flux2,flux2_load_pipeline,download_flux2_klein): empirical-mu dynamic shifting through the existing FlowMatch scheduler, GPU-resident fp8 (~3.9 GB artifact from the ungated repo), per-phase offloading;flux_quantize/flux_load_transformernow auto-dispatch FLUX.1 vs FLUX.2 on the config class.Verification
R CMD check: 0 errors, 0 warnings.Performance note
Initial generations ran 130-200 s with the GPU at 2% utilization;
Rprof(gc.profiling=TRUE)showed 93.8% of wall time in R gc. torch's allocator callbacks legitimately run gc hundreds of times per generation, and the tokenizer's 151k-binding vocab/merge environments made each gc 17x slower (12 -> 203 ms). The integer-id tokenizer rewrite plus buffered fp8 dequant and allocator tuning brought 1024px from 193 s to 48 s. Lesson recorded: big lookup tables as atomic vectors, never environments.Provenance
All ported code traces to Apache-2.0 sources (diffusers
transformer_flux2.py/pipeline_flux2_klein.py/autoencoder_kl_flux2.py, transformersmodels/qwen3/), recorded per-technique ininst/REFERENCES.md. Weights are user-downloaded from the ungated, Apache-2.0black-forest-labs/FLUX.2-klein-4Brepo and never redistributed.Notes
guidance_scalehas no effect (distilled,guidance_embeds=false).