FLUX.2 Klein 4B text-to-image support

CLAUDE.md

@@ -317,6 +317,7 @@ See cornyverse CLAUDE.md for safetensors package setup (use cornball-ai fork unt
 
 ### Model Support
 - [x] Add FLUX model support (FLUX.1-schnell, see below)
+- [x] Add FLUX.2 support (klein-4B, see below)
 - [ ] Add SD3 model support
 - [ ] ControlNet integration
 
@@ -347,6 +348,31 @@ the shipped tokenizer_2 uses Metaspace prepend "always" (a spiece
 conversion gives "never" - different ids for every prompt), and quantized
 residents must follow the compute dtype (bf16 GPU / fp32 CPU).
 
+### FLUX.2 Klein 4B (Complete)
+
+Step-distilled FLUX.2 (4 steps, no CFG): 4B MMDiT (shared modulation,
+SwiGLU, parallel single blocks) + Qwen3-4B text encoder (chat template,
+mid-stack hidden states) + 32-channel `AutoencoderKLFlux2` (BatchNorm
+latent stats) + FlowMatch with the BFL empirical dynamic shift.
+
+```r
+download_flux2_klein()   # ungated, Apache-2.0; ~16 GB download,
+                         # one-time fp8 quantize to a 3.9 GB artifact
+txt2img_flux2("An astronaut riding a horse on Mars, photorealistic",
+              seed = 7)  # or txt2img("...", model_name = "flux2")
+```
+
+Measured on the RTX 5060 Ti 16 GB (fp8 GPU-resident, Qwen3 bf16
+phase-onloaded): 1024x1024 in ~48 s (peak 8.2 GB), 512x512 in ~40 s;
+pipeline load 31 s. Cast census is exactly 104 weights.
+
+Perf lesson that cost an afternoon: generation was 93.8% R garbage
+collection until the tokenizer stopped holding 151k-binding
+environments — R gc cost scales with live object count (12 vs 203 ms),
+and torch's allocator callbacks run gc hundreds of times per
+generation. Keep big lookup tables as atomic vectors (integer-id BPE
+via findInterval), never as environments.
+
 ### LTX-2.3 Video Generation (clean-room rewrite in progress)
 
 The original LTX-2.0 port was removed and is being replaced by a ground-up

DESCRIPTION

@@ -1,6 +1,6 @@
 Package: diffuseR
 Title: Functional Interface to Diffusion Models in R
-Version: 0.1.0.1
+Version: 0.1.0.2
 Authors@R: c(
     person("Troy", "Hernandez", email = "troy@cornball.ai", role = c("aut", "cre"),
            comment = c(ORCID = "0009-0005-4248-604X")),

NAMESPACE

@@ -10,11 +10,14 @@ export(ddim_scheduler_step)
 export(decode_bpe)
 export(download_component)
 export(download_flux1)
+export(download_flux2_klein)
 export(download_ltx2)
 export(download_model)
 export(encode_bpe)
+export(encode_qwen)
 export(encode_unigram)
 export(encode_with_gemma3)
+export(encode_with_qwen3)
 export(encode_with_t5)
 export(filename_from_prompt)
 export(flowmatch_calculate_shift)
@@ -41,6 +44,23 @@ export(flux_quantize)
 export(flux_single_block)
 export(flux_transformer)
 export(flux_unpack_latents)
+export(flux2_bn_normalize)
+export(flux2_double_block)
+export(flux2_empirical_mu)
+export(flux2_feed_forward)
+export(flux2_is_quant_key)
+export(flux2_load_pipeline)
+export(flux2_modulation)
+export(flux2_pack_latents)
+export(flux2_parallel_self_attention)
+export(flux2_patchify_latents)
+export(flux2_prepare_latent_ids)
+export(flux2_prepare_text_ids)
+export(flux2_single_block)
+export(flux2_transformer)
+export(flux2_unpack_latents_with_ids)
+export(flux2_unpatchify_latents)
+export(flux2_vae_decoder)
 export(gemma3_config_ltx2)
 export(gemma3_text_model)
 export(gemma3_tokenizer)
@@ -49,9 +69,11 @@ export(is_blackwell_gpu)
 export(latents_to_video)
 export(load_decoder_safetensors)
 export(load_decoder_weights)
+export(load_flux2_vae_decoder)
 export(load_gemma3_text_encoder)
 export(load_model_component)
 export(load_pipeline)
+export(load_qwen3_text_encoder)
 export(load_t5_text_encoder)
 export(load_text_encoder_safetensors)
 export(load_text_encoder_weights)
@@ -144,6 +166,8 @@ export(offload_to_cpu)
 export(post_quant_conv)
 export(preprocess_image)
 export(quant_conv)
+export(qwen_bpe_tokenizer)
+export(qwen3_encoder)
 export(save_frames)
 export(save_image)
 export(save_video)
@@ -156,6 +180,7 @@ export(text_encoder2_native)
 export(tokenize_gemma3)
 export(txt2img)
 export(txt2img_flux)
+export(txt2img_flux2)
 export(txt2img_sd21)
 export(txt2img_sdxl)
 export(txt2vid_ltx2)
@@ -171,6 +196,7 @@ export(write_wav)
 
 S3method(print,bpe_tokenizer)
 S3method(print,ltx23_checkpoint)
+S3method(print,qwen_tokenizer)
 S3method(print,unigram_tokenizer)
 
 importFrom(utils,head)

R/checkpoint_flux.R

@@ -42,6 +42,46 @@ flux_is_quant_key <- function(key) {
     grepl(.flux_quant_cast_pattern, key)
 }
 
+# FLUX.2 cast set: block linears, the three shared modulation
+# projections, and the context embedder (7680 x 3072). Everything else
+# (x_embedder, timestep MLP, norm_out, q/k norms) stays in the resident
+# dtype. Full klein-4B census: 5 double x 12 + 20 single x 2 + 3
+# modulations + context_embedder = 104 cast weights, ~3.9B of 4B params.
+.flux2_quant_cast_pattern <- paste0(
+                                    "^(",
+                                    "transformer_blocks\\.[0-9]+\\.(",
+                                    "attn\\.(to_q|to_k|to_v|add_q_proj|add_k_proj|add_v_proj|to_out\\.0|to_add_out)",
+                                    "|ff\\.(linear_in|linear_out)|ff_context\\.(linear_in|linear_out)",
+                                    ")",
+                                    "|single_transformer_blocks\\.[0-9]+\\.attn\\.(to_qkv_mlp_proj|to_out)",
+                                    "|double_stream_modulation_img\\.linear",
+                                    "|double_stream_modulation_txt\\.linear",
+                                    "|single_stream_modulation\\.linear",
+                                    "|context_embedder",
+                                    ")\\.weight$"
+)
+
+#' Test whether a FLUX.2 key is in the quantization cast set
+#'
+#' @param key Character vector of parameter names (diffusers-style).
+#'
+#' @return Logical vector.
+#'
+#' @export
+flux2_is_quant_key <- function(key) {
+    grepl(.flux2_quant_cast_pattern, key)
+}
+
+# Model family from a diffusers transformer config
+.flux_family <- function(config) {
+    cls <- config$`_class_name` %||% "FluxTransformer2DModel"
+    if (identical(cls, "Flux2Transformer2DModel")) {
+        "flux2"
+    } else {
+        "flux1"
+    }
+}
+
 #' Open a FLUX transformer checkpoint directory
 #'
 #' Opens a diffusers-layout transformer directory lazily (headers only).

R/dit_flux2.R

@@ -0,0 +1,163 @@
+#' FLUX.2 Transformer (MMDiT)
+#'
+#' Fresh R port of Flux2Transformer2DModel from the diffusers reference
+#' implementation (Apache-2.0,
+#' src/diffusers/models/transformers/transformer_flux2.py). Defaults are
+#' the klein-4B configuration (5 double + 20 single blocks). Guidance
+#' embeddings (FLUX.2-dev) are not implemented; klein is step-distilled
+#' with \code{guidance_embeds = false}. Timestep conditioning has no
+#' pooled-text component, and the three modulation projections are
+#' shared across all blocks.
+#'
+#' @name dit_flux2
+NULL
+
+# Timestep-only conditioning: sinusoid(256) -> MLP, bias-free. Matches
+# diffusers Flux2TimestepGuidanceEmbeddings state-dict names (klein has
+# no guidance_embedder).
+flux2_time_guidance_embed <- torch::nn_module(
+    "flux2_time_guidance_embed",
+    initialize = function(embedding_dim, in_channels = 256L) {
+    self$in_channels <- in_channels
+    self$timestep_embedder <- ltx23_timestep_embedding(in_channels,
+        embedding_dim, bias = FALSE)
+},
+    forward = function(timestep) {
+    proj <- ltx23_get_timestep_embedding(timestep, self$in_channels,
+        flip_sin_to_cos = TRUE, downscale_freq_shift = 0)
+    self$timestep_embedder(proj$to(dtype = self$timestep_embedder$linear_1$weight$dtype))
+}
+)
+
+#' FLUX.2 transformer model
+#'
+#' Shared modulation computed once per forward; double blocks over
+#' separate text/image streams, then single (parallel) blocks over the
+#' concatenated [text; image] sequence. Rotary embeddings are
+#' precomputed by the caller with \code{\link{flux_pos_embed}}
+#' (\code{axes_dim = c(32, 32, 32, 32)}, \code{theta = 2000}) over the
+#' concatenated [text; image] 4-axis position ids.
+#'
+#' @param in_channels Integer. Packed latent channels (128).
+#' @param num_layers Integer. Double-stream block count (klein-4B: 5).
+#' @param num_single_layers Integer. Single-stream block count (20).
+#' @param attention_head_dim Integer. Per-head dimension.
+#' @param num_attention_heads Integer. Attention heads.
+#' @param joint_attention_dim Integer. Text embedding dim (7680).
+#' @param mlp_ratio Numeric. Feed-forward multiplier (3.0).
+#' @param timestep_guidance_channels Integer. Sinusoid width (256).
+#' @param axes_dims_rope Integer vector. Per-axis rotary dims.
+#' @param rope_theta Numeric. Rotary base frequency (2000).
+#' @param eps Numeric. Norm epsilon.
+#' @param out_channels Integer or NULL. Defaults to \code{in_channels}.
+#'
+#' @return Module whose forward(hidden_states, encoder_hidden_states,
+#'   timestep, image_rotary_emb) returns the predicted velocity for the
+#'   image tokens [B, S_img, out_channels]. \code{timestep} is in sigma
+#'   space (0-1); it is scaled by 1000 internally.
+#'
+#' @export
+flux2_transformer <- torch::nn_module(
+                                      "flux2_transformer",
+                                      initialize = function(in_channels = 128L,
+        num_layers = 5L,
+        num_single_layers = 20L,
+        attention_head_dim = 128L,
+        num_attention_heads = 24L,
+        joint_attention_dim = 7680L,
+        mlp_ratio = 3.0,
+        timestep_guidance_channels = 256L,
+        axes_dims_rope = c(32L, 32L, 32L, 32L),
+        rope_theta = 2000,
+        eps = 1e-6,
+        out_channels = NULL) {
+    inner_dim <- num_attention_heads * attention_head_dim
+    self$inner_dim <- inner_dim
+    self$axes_dims_rope <- as.integer(axes_dims_rope)
+    self$rope_theta <- rope_theta
+    self$out_channels <- as.integer(out_channels %||% in_channels)
+
+    self$time_guidance_embed <- flux2_time_guidance_embed(inner_dim,
+        as.integer(timestep_guidance_channels))
+    self$double_stream_modulation_img <- flux2_modulation(inner_dim, 2L)
+    self$double_stream_modulation_txt <- flux2_modulation(inner_dim, 2L)
+    self$single_stream_modulation <- flux2_modulation(inner_dim, 1L)
+
+    self$x_embedder <- torch::nn_linear(in_channels, inner_dim, bias = FALSE)
+    self$context_embedder <- torch::nn_linear(joint_attention_dim,
+        inner_dim, bias = FALSE)
+
+    self$transformer_blocks <- torch::nn_module_list(
+        lapply(seq_len(num_layers), function(i) {
+        flux2_double_block(inner_dim, num_attention_heads,
+                           attention_head_dim, mlp_ratio = mlp_ratio,
+                           eps = eps)
+    })
+    )
+    self$single_transformer_blocks <- torch::nn_module_list(
+        lapply(seq_len(num_single_layers), function(i) {
+        flux2_single_block(inner_dim, num_attention_heads,
+                           attention_head_dim, mlp_ratio = mlp_ratio,
+                           eps = eps)
+    })
+    )
+
+    self$norm_out <- flux_ada_layer_norm_continuous(inner_dim, inner_dim,
+        bias = FALSE)
+    self$proj_out <- torch::nn_linear(inner_dim, self$out_channels,
+                                      bias = FALSE)
+},
+                                      forward = function(hidden_states, encoder_hidden_states, timestep,
+        image_rotary_emb, chunk_size = NULL) {
+    hidden_states <- self$x_embedder(hidden_states)
+    timestep <- timestep$to(dtype = hidden_states$dtype)$mul(1000)
+    temb <- self$time_guidance_embed(timestep)
+
+    mod_img <- self$double_stream_modulation_img(temb)
+    mod_txt <- self$double_stream_modulation_txt(temb)
+    mod_single <- self$single_stream_modulation(temb)
+
+    encoder_hidden_states <- self$context_embedder(encoder_hidden_states)
+
+    block_gc <- isTRUE(getOption("diffuseR.block_gc"))
+    for (i in seq_along(self$transformer_blocks)) {
+        res <- self$transformer_blocks[[i]](
+                                      hidden_states = hidden_states,
+                                      encoder_hidden_states = encoder_hidden_states,
+                                      temb_mod_img = mod_img,
+                                      temb_mod_txt = mod_txt,
+                                      image_rotary_emb = image_rotary_emb,
+                                      chunk_size = chunk_size
+        )
+        encoder_hidden_states <- res[[1]]
+        hidden_states <- res[[2]]
+        if (block_gc) {
+            gc(verbose = FALSE)
+        }
+    }
+
+    txt_len <- encoder_hidden_states$shape[2]
+    hidden_states <- torch::torch_cat(
+                                      list(encoder_hidden_states, hidden_states),
+                                      dim = 2L
+    )
+    for (i in seq_along(self$single_transformer_blocks)) {
+        hidden_states <- self$single_transformer_blocks[[i]](
+                                      hidden_states = hidden_states,
+                                      temb_mod = mod_single,
+                                      image_rotary_emb = image_rotary_emb,
+                                      chunk_size = chunk_size
+        )
+        if (block_gc) {
+            gc(verbose = FALSE)
+        }
+    }
+    hidden_states <- hidden_states$narrow(
+        2L, txt_len + 1L,
+        hidden_states$shape[2] - txt_len
+    )
+
+    hidden_states <- self$norm_out(hidden_states, temb)
+    self$proj_out(hidden_states)
+}
+)