FLUX 3 Image Klein: text-to-image and single-reference editing

CLAUDE.md

@@ -321,6 +321,7 @@ See cornyverse CLAUDE.md for safetensors package setup. CRAN safetensors
 - [x] Add FLUX model support (FLUX.1-schnell, see below)
 - [x] Add FLUX.2 support (klein-4B, see below)
 - [x] Add Z-Image-Turbo support (text rendering, see below)
+- [x] Add FLUX 3 Image Klein support (text-to-image + editing, see below)
 - [ ] Add SD3 model support
 - [ ] ControlNet integration
 
@@ -430,6 +431,52 @@ length (pads continue the ramp; the (0,0,0) pad ids in the reference
 are dead code), and the checkpoint scheduler uses static shift 3.0 —
 the pipeline's calculate_shift/mu path is dead for this model.
 
+### FLUX 3 Image Klein
+
+Step-distilled flow map (16 steps, no CFG): 5.6B transformer with
+per-mode blocks (image, reference image and text each self-attend within
+the mode, 5 blocks each) then 28 joint blocks, per-token timestep
+modulation, plus a target-timestep embedder (the step it jumps to) +
+Qwen3-VL-4B text model (layers 4, 8, ..., 32 stacked; prompt padded to
+a multiple of 80 with pad position ids 1) + the FLUX.2 VAE. FlowMatch
+with static shift 6.93. Gated, FLUX Kommunity License (non-commercial;
+commercial outputs for Qualifying Users under US$5M revenue who filter
+or review outputs). Weights never redistributed.
+
+```r
+download_flux3_klein()   # gated: accept the license + HF_TOKEN; ~20 GB
+                         # download, one-time fp8 quantize to 5.3 GB
+pipe <- flux3_load_pipeline()
+txt2img_flux3(caption, pipeline = pipe, seed = 7)
+img2img_flux3(caption, image = "photo.png", pipeline = pipe, seed = 7)
+```
+
+The model wants full scene captions (BFL: "do not use the model without
+a prompt upsampler"); diffuseR takes the caption as given, expansion is
+the caller's job. One reference is kept whole: scaled so its canvas fits
+4 MP (`max_area`), reflect-padded to the 16 px grid, output cropped back.
+
+Verified against diffusers on the real weights (CPU): Qwen3-VL embeds
+bit-identical (fp32); transformer forward max abs 2.9e-6 (fp32,
+depth-reduced to 4 joint blocks, all code paths); 256 px end to end in
+bf16: edit 46 dB PSNR, text-to-image same composition at 24.5 dB (bf16
+trajectory noise over 16 steps). fp8 census exactly 265 weights.
+Measured on the RTX 5060 Ti 16 GB (fp8 resident, phase offload):
+1024x1024 text-to-image ~44 s (peak 11.7 GB alloc / 12.0 GB reserved);
+1 MP edit ~50 s (13.8 / 14.1 GB); 4 MP edit (2048x2048 out) ~285 s
+(14.1 / 14.8 GB); pipeline load ~36 s. The 4 MP edit needed the VAE
+mid-block attention on fused SDPA (the manual form allocated an 8 GiB
+score matrix). GPU check against diffusers, 1 MP edit with each side's
+own text encoder: R bf16 52.2 dB PSNR, R fp8 36.7 dB. Edits that change
+little (no snow from a "covered in snow" caption) reproduce in
+diffusers; they come from the model without its prompt upsampler.
+
+Gotchas that bit once: HF's last `hidden_states` entry carries the final
+norm (the Qwen3 encoder now matches for k = num_layers); building 5.6B
+params in fp32 before the bf16 cast needs 22 GB of host RAM, so the
+loader constructs a `.construct_skeleton()` (FLUX 3 opts in via
+`skeleton = TRUE` in `.flux_family_hooks`).
+
 ### 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.2.2.15
+Version: 0.2.2.16
 Authors@R: c(
     person("Troy", "Hernandez", email = "troy@cornball.ai", role = c("aut", "cre"),
            comment = c(ORCID = "0009-0005-4248-604X")),

NAMESPACE

@@ -12,6 +12,7 @@ export(decode_bpe)
 export(download_component)
 export(download_flux1)
 export(download_flux2_klein)
+export(download_flux3_klein)
 export(download_ltx2)
 export(download_model)
 export(download_sd21)
@@ -71,6 +72,17 @@ export(flux2_unpack_latents_with_ids)
 export(flux2_unpatchify_latents)
 export(flux2_vae_decoder)
 export(flux2_vae_encoder)
+export(flux3_attention)
+export(flux3_encode_prompt)
+export(flux3_is_quant_key)
+export(flux3_load_pipeline)
+export(flux3_mode_block)
+export(flux3_prepare_image_ids)
+export(flux3_prepare_text_ids)
+export(flux3_preprocess_image)
+export(flux3_schedule)
+export(flux3_single_block)
+export(flux3_transformer)
 export(gemma3_config_ltx2)
 export(gemma3_encode_batch)
 export(gemma3_quantize_nf4)
@@ -81,6 +93,7 @@ export(gemma4_text_model)
 export(gemma4_tokenizer)
 export(img2img)
 export(img2img_flux2)
+export(img2img_flux3)
 export(is_blackwell_gpu)
 export(latents_to_video)
 export(load_decoder_safetensors)
@@ -222,6 +235,7 @@ export(tokenize_gemma3)
 export(txt2img)
 export(txt2img_flux)
 export(txt2img_flux2)
+export(txt2img_flux3)
 export(txt2img_sd21)
 export(txt2img_sdxl)
 export(txt2img_zimage)

NEWS.md

@@ -1,3 +1,30 @@
+# diffuseR 0.2.2.16
+
+* **FLUX 3 Image Klein.** `txt2img_flux3()` and `img2img_flux3()` run
+  Black Forest Labs' step-distilled FLUX 3 Image Klein (16 steps, no
+  guidance), ported from the diffusers pipeline. The weights are gated
+  under the FLUX Kommunity License; `download_flux3_klein()` fetches
+  them with `HF_TOKEN` and builds a 5.3 GB fp8 artifact.
+  - `flux3_load_pipeline()` swaps the transformer, Qwen3-VL text model
+    and VAE onto the GPU one phase at a time; `dtype = "bfloat16"` also
+    works on the CPU.
+  - The model expects full scene captions; expand short prompts first.
+  - `img2img_flux3()` keeps one reference whole (`max_area`, default
+    4 MP) and crops the output back to it; see
+    `flux3_preprocess_image()`.
+  - Building blocks: `flux3_transformer()`, `flux3_encode_prompt()`,
+    `flux3_prepare_image_ids()`, `flux3_prepare_text_ids()`,
+    `flux3_schedule()`, `flux3_is_quant_key()`.
+* `load_qwen3_text_encoder()` also loads Qwen3-VL text models, and the
+  encoder takes explicit `position_ids`. Requesting the last layer now
+  returns the post-norm state, as transformers does. The loader builds
+  the model at the target dtype, so a bf16 load no longer passes
+  through a 16 GB fp32 copy.
+* The VAE mid-block attention (shared by the SD, SDXL, FLUX.1 and
+  FLUX.2/FLUX 3 VAEs) uses fused scaled-dot-product attention when
+  torch provides it. A 2048x2048 image no longer allocates an 8 GiB
+  score matrix.
+
 # diffuseR 0.2.2.15
 
 * **FLUX.2 klein image + prompt editing.** `img2img_flux2(prompt, image)`

R/checkpoint_flux.R

@@ -84,6 +84,32 @@ zimage_is_quant_key <- function(key) {
     grepl(.zimage_quant_cast_pattern, key)
 }
 
+# FLUX 3 cast set: the six block linears (q/k/v/out, SwiGLU in/out) in
+# the per-mode and joint stacks, the six per-mode modulation projections
+# and the context embedder (20480 x 3072). Embedders, the timestep MLPs,
+# norm_out, proj_out and the q/k norms stay in the resident dtype. Full
+# Image Klein census: (5 x 3 mode + 28 joint blocks) x 6 + 6
+# modulations + context_embedder = 265 cast weights.
+.flux3_quant_cast_pattern <- paste0(
+                                    "^(",
+                                    "(mode_blocks\\.(image|image_cond)|txt_mode_blocks|single_transformer_blocks)",
+                                    "\\.[0-9]+\\.(attn\\.(to_q|to_k|to_v|to_out\\.0)|ff\\.(linear_in|linear_out))",
+                                    "|(early|single)_stream_modulations\\.(image|image_cond|txt)\\.linear",
+                                    "|context_embedder",
+                                    ")\\.weight$"
+)
+
+#' Test whether a FLUX 3 key is in the quantization cast set
+#'
+#' @param key Character vector of parameter names (diffusers-style).
+#'
+#' @return Logical vector.
+#'
+#' @export
+flux3_is_quant_key <- function(key) {
+    grepl(.flux3_quant_cast_pattern, key)
+}
+
 #' Test whether a FLUX.2 key is in the quantization cast set
 #'
 #' @param key Character vector of parameter names (diffusers-style).
@@ -100,6 +126,8 @@ flux2_is_quant_key <- function(key) {
     cls <- config$`_class_name` %||% "FluxTransformer2DModel"
     if (identical(cls, "Flux2Transformer2DModel")) {
         "flux2"
+    } else if (identical(cls, "Flux3Transformer2DModel")) {
+        "flux3"
     } else if (identical(cls, "ZImageTransformer2DModel")) {
         "zimage"
     } else {

R/dit_flux2_modules.R

@@ -74,9 +74,11 @@ flux2_feed_forward <- torch::nn_module(
                                        "flux2_feed_forward",
                                        initialize = function(dim, dim_out = NULL, mult = 3.0, bias = FALSE) {
     inner_dim <- as.integer(dim * mult)
-    self$linear_in <- torch::nn_linear(dim, inner_dim * 2L, bias = bias)
-    self$linear_out <- torch::nn_linear(inner_dim, dim_out %||% dim,
-                                        bias = bias)
+    # linear_noinit: skips the init writes inside .construct_skeleton()
+    # (FLUX 3 loads), initializes like nn_linear otherwise
+    self$linear_in <- linear_noinit(dim, inner_dim * 2L, bias = bias)
+    self$linear_out <- linear_noinit(inner_dim, dim_out %||% dim,
+                                     bias = bias)
 },
                                        forward = function(x) {
     x <- self$linear_in(x)