diff --git a/CLAUDE.md b/CLAUDE.md index 00babf5..c672afa 100644 --- a/CLAUDE.md +++ b/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 diff --git a/DESCRIPTION b/DESCRIPTION index 221bc0c..f85ffa1 100644 --- a/DESCRIPTION +++ b/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")), diff --git a/NAMESPACE b/NAMESPACE index 2736e3c..de07e6d 100644 --- a/NAMESPACE +++ b/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) diff --git a/NEWS.md b/NEWS.md index f27aed1..1c29c74 100644 --- a/NEWS.md +++ b/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)` diff --git a/R/checkpoint_flux.R b/R/checkpoint_flux.R index 6f99524..fc38ca0 100644 --- a/R/checkpoint_flux.R +++ b/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 { diff --git a/R/dit_flux2_modules.R b/R/dit_flux2_modules.R index 8a7d5a6..c7c21a2 100644 --- a/R/dit_flux2_modules.R +++ b/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) diff --git a/R/dit_flux3.R b/R/dit_flux3.R new file mode 100644 index 0000000..0d95cb4 --- /dev/null +++ b/R/dit_flux3.R @@ -0,0 +1,234 @@ +#' FLUX 3 Transformer +#' +#' Fresh R port of Flux3Transformer2DModel from the diffusers reference +#' implementation (Apache-2.0, +#' src/diffusers/models/transformers/transformer_flux3.py). Defaults are +#' the FLUX 3 Image Klein configuration: 5 per-mode blocks for each of +#' the image, reference-image and text modes, then 28 joint blocks. +#' Conditioning is a per-token timestep embedding (the image tokens at +#' the flow timestep, text and reference tokens at 0); the step-distilled +#' (flow-map) checkpoint adds the embedding of the timestep each step +#' jumps to. +#' +#' @name dit_flux3 +NULL + +# Linear(dim, 3 dim) on an already SiLU-activated embedding. Reference: +# Flux3Modulation. +flux3_modulation <- torch::nn_module( + "flux3_modulation", + initialize = function(dim) { + self$linear <- linear_noinit(dim, dim * 3L, bias = FALSE) +}, + forward = function(temb_act) { + self$linear(temb_act) +} +) + +#' FLUX 3 transformer model +#' +#' The forward mirrors the reference: per-mode blocks for the text, +#' image and (optional) reference tokens, each with its own rotary +#' embedding, then joint blocks over [text; image; reference] and an +#' output projection of the image tokens only. Modulation is computed +#' once per mode on the [B, 1, dim] embedding and broadcast over the +#' mode's tokens (the reference expands to [B, S, dim] first; every row +#' is identical). +#' +#' @param in_channels Integer. Packed latent channels (128). +#' @param cond_in_channels Integer. Reference-image token channels (128). +#' @param num_layers Integer. Per-mode blocks per mode (5). +#' @param num_single_layers Integer. Joint blocks (28). +#' @param attention_head_dim Integer. Per-head dimension (128). +#' @param num_attention_heads Integer. Attention heads (24). +#' @param joint_attention_dim Integer. Text embedding dim (20480). +#' @param timestep_channels Integer. Sinusoid width (256). +#' @param mlp_ratio Numeric. Feed-forward multiplier (3.0). +#' @param axes_dims_rope Integer vector. Per-axis rotary dims. +#' @param rope_theta Numeric. Rotary base frequency (10000). +#' @param eps Numeric. Norm epsilon. +#' @param flow_map Logical. Step-distilled checkpoint: adds +#' \code{target_timestep_embedder} and requires \code{target_timestep}. +#' +#' @return Module whose forward(hidden_states, encoder_hidden_states, +#' timestep, img_ids, txt_ids, image_cond_hidden_states, image_cond_ids, +#' target_timestep, chunk_size) returns the predicted velocity of the +#' image tokens [B, S_img, in_channels]. Timesteps are flow times in +#' [0, 1]; they are scaled by 1000 internally. +#' +#' @export +flux3_transformer <- torch::nn_module( + "flux3_transformer", + initialize = function(in_channels = 128L, + cond_in_channels = 128L, + num_layers = 5L, + num_single_layers = 28L, + attention_head_dim = 128L, + num_attention_heads = 24L, + joint_attention_dim = 20480L, + timestep_channels = 256L, + mlp_ratio = 3.0, + axes_dims_rope = c(32L, 32L, 32L, 32L), + rope_theta = 10000, + eps = 1e-6, + flow_map = FALSE) { + inner_dim <- num_attention_heads * attention_head_dim + self$inner_dim <- inner_dim + self$timestep_channels <- as.integer(timestep_channels) + self$axes_dims_rope <- as.integer(axes_dims_rope) + self$rope_theta <- rope_theta + self$flow_map <- isTRUE(flow_map) + + self$timestep_embedder <- ltx23_timestep_embedding(timestep_channels, + inner_dim, bias = FALSE) + if (self$flow_map) { + self$target_timestep_embedder <- ltx23_timestep_embedding( + timestep_channels, inner_dim, bias = FALSE) + } + + modes <- c("image", "image_cond", "txt") + mods <- function() { + torch::nn_module_dict(stats::setNames( + lapply(modes, function(m) flux3_modulation(inner_dim)), modes)) + } + self$early_stream_modulations <- mods() + self$single_stream_modulations <- mods() + + self$x_embedders <- torch::nn_module_dict(list( + image = linear_noinit(in_channels, inner_dim, bias = FALSE), + image_cond = linear_noinit(cond_in_channels, inner_dim, + bias = FALSE))) + self$context_embedder <- linear_noinit(joint_attention_dim, + inner_dim, bias = FALSE) + + blocks <- function(n, ctor) { + torch::nn_module_list(lapply(seq_len(n), function(i) { + ctor(inner_dim, num_attention_heads, attention_head_dim, + mlp_ratio = mlp_ratio, eps = eps) + })) + } + self$mode_blocks <- torch::nn_module_dict(list( + image = blocks(num_layers, flux3_mode_block), + image_cond = blocks(num_layers, flux3_mode_block))) + self$txt_mode_blocks <- blocks(num_layers, flux3_mode_block) + self$single_transformer_blocks <- blocks(num_single_layers, + flux3_single_block) + + self$norm_out <- flux_ada_layer_norm_continuous(inner_dim, inner_dim, + bias = FALSE) + self$proj_out <- linear_noinit(inner_dim, in_channels, bias = FALSE) +}, + embed_timestep = function(embedder, t, dtype) { + proj <- ltx23_get_timestep_embedding(t$to(dtype = dtype)$mul(1000), + self$timestep_channels, flip_sin_to_cos = TRUE, + downscale_freq_shift = 0) + embedder(proj$to(dtype = dtype)) +}, + forward = function(hidden_states, encoder_hidden_states, timestep, + img_ids, txt_ids, image_cond_hidden_states = NULL, + image_cond_ids = NULL, target_timestep = NULL, + chunk_size = NULL) { + if (is.null(image_cond_hidden_states) != is.null(image_cond_ids)) { + stop("`image_cond_hidden_states` and `image_cond_ids` go together.") + } + if (is.null(target_timestep) == self$flow_map) { + stop("`target_timestep` is required by, and only by, a flow-map ", + "transformer (flow_map = TRUE).") + } + dtype <- hidden_states$dtype + batch_size <- hidden_states$shape[1] + image_seq_len <- hidden_states$shape[2] + has_cond <- !is.null(image_cond_hidden_states) + + # 1. Timestep embeddings [B, dim]: image tokens at `timestep`, text + # and reference tokens at 0; the flow-map target is added to both + image_temb <- self$embed_timestep(self$timestep_embedder, timestep, dtype) + zero_temb <- self$embed_timestep(self$timestep_embedder, + torch::torch_zeros(batch_size, device = hidden_states$device), dtype) + if (self$flow_map) { + target_temb <- self$embed_timestep(self$target_timestep_embedder, + target_timestep, dtype) + image_temb <- image_temb + target_temb + zero_temb <- zero_temb + target_temb + } + zero_act <- torch::nnf_silu(zero_temb)$unsqueeze(2L) # [B, 1, dim] + image_act <- torch::nnf_silu(image_temb)$unsqueeze(2L) + + # 2. Text context + txt_early <- self$early_stream_modulations$txt(zero_act) + txt_single <- self$single_stream_modulations$txt(zero_act) + encoder_hidden_states <- self$context_embedder(encoder_hidden_states) + txt_rope <- flux_pos_embed(txt_ids, self$axes_dims_rope, self$rope_theta) + + # 3. Image tokens + image_early <- self$early_stream_modulations$image(image_act) + image_single <- self$single_stream_modulations$image(image_act) + hidden_states <- self$x_embedders$image(hidden_states) + image_rope <- flux_pos_embed(img_ids, self$axes_dims_rope, self$rope_theta) + + # 4. Reference tokens (optional) + if (has_cond) { + cond_early <- self$early_stream_modulations$image_cond(zero_act) + cond_single <- self$single_stream_modulations$image_cond(zero_act) + image_cond_hidden_states <- self$x_embedders$image_cond( + image_cond_hidden_states) + cond_rope <- flux_pos_embed(image_cond_ids, self$axes_dims_rope, + self$rope_theta) + } + + # 5. Per-mode blocks (each mode attends within itself only) + for (i in seq_along(self$txt_mode_blocks)) { + encoder_hidden_states <- self$txt_mode_blocks[[i]]( + encoder_hidden_states, txt_early, txt_rope, chunk_size) + } + for (i in seq_along(self$mode_blocks$image)) { + hidden_states <- self$mode_blocks$image[[i]]( + hidden_states, image_early, image_rope, chunk_size) + } + if (has_cond) { + for (i in seq_along(self$mode_blocks$image_cond)) { + image_cond_hidden_states <- self$mode_blocks$image_cond[[i]]( + image_cond_hidden_states, cond_early, cond_rope, chunk_size) + } + } + + # 6. Joint blocks over [txt, image, image_cond] + if (has_cond) { + cond_len <- image_cond_hidden_states$shape[2] + content <- torch::torch_cat(list(hidden_states, + image_cond_hidden_states), dim = 2L) + content_mod <- torch::torch_cat(list( + image_single$expand(c(-1L, image_seq_len, -1L)), + cond_single$expand(c(-1L, cond_len, -1L))), dim = 2L) + ropes <- list(txt_rope, image_rope, cond_rope) + } else { + content <- hidden_states + content_mod <- image_single + ropes <- list(txt_rope, image_rope) + } + joint_rope <- list( + torch::torch_cat(lapply(ropes, `[[`, 1L), dim = 1L), + torch::torch_cat(lapply(ropes, `[[`, 2L), dim = 1L)) + rm(hidden_states, image_cond_hidden_states) + + block_gc <- isTRUE(getOption("diffuseR.block_gc")) + for (i in seq_along(self$single_transformer_blocks)) { + res <- self$single_transformer_blocks[[i]]( + hidden_states = content, + encoder_hidden_states = encoder_hidden_states, + temb_mod = content_mod, + temb_mod_context = txt_single, + image_rotary_emb = joint_rope, + chunk_size = chunk_size) + encoder_hidden_states <- res[[1]] + content <- res[[2]] + if (block_gc) { + gc(verbose = FALSE) + } + } + + # 7. Output: the image tokens sit first in the content sequence + out <- content$narrow(2L, 1L, image_seq_len) + self$proj_out(self$norm_out(out, image_temb)) +} +) diff --git a/R/dit_flux3_modules.R b/R/dit_flux3_modules.R new file mode 100644 index 0000000..f15defb --- /dev/null +++ b/R/dit_flux3_modules.R @@ -0,0 +1,167 @@ +#' FLUX 3 Transformer Building Blocks +#' +#' Fresh R port of the FLUX 3 blocks from the diffusers reference +#' implementation (Apache-2.0, +#' src/diffusers/models/transformers/transformer_flux3.py). Every block +#' is a parallel attention + SwiGLU feed-forward computed from one +#' modulated input: per-mode blocks self-attend within one token mode +#' (image, reference image, or text), joint blocks attend over the +#' concatenated [text; content] sequence. Modulation is per token, +#' computed at model level. Module field names mirror the diffusers +#' state-dict keys 1:1; every linear is bias-free. Reuses +#' \code{flux2_feed_forward}, \code{flux_apply_rotary_emb} and +#' \code{.ltx23_sdpa}. +#' +#' @name dit_flux3_modules +NULL + +#' FLUX 3 self-attention +#' +#' Separate q/k/v projections, per-head RMS norms and rotary embeddings +#' applied in float32 (cast back to the value dtype before attention), +#' and one output projection. Reference: Flux3Attention + +#' Flux3AttnProcessor. +#' +#' @param query_dim Integer. Model dimension. +#' @param heads Integer. Attention heads. +#' @param dim_head Integer. Per-head dimension. +#' @param eps Numeric. RMS norm epsilon. +#' +#' @return Module whose forward(hidden_states, image_rotary_emb, +#' chunk_size) returns the attended states [B, S, query_dim]. +#' +#' @export +flux3_attention <- torch::nn_module( + "flux3_attention", + initialize = function(query_dim, heads, dim_head, eps = 1e-6) { + inner_dim <- heads * dim_head + self$heads <- heads + self$eps <- eps + self$to_q <- linear_noinit(query_dim, inner_dim, bias = FALSE) + self$to_k <- linear_noinit(query_dim, inner_dim, bias = FALSE) + self$to_v <- linear_noinit(query_dim, inner_dim, bias = FALSE) + self$norm_q <- ltx23_rms_norm(dim_head, eps = eps) + self$norm_k <- ltx23_rms_norm(dim_head, eps = eps) + self$to_out <- torch::nn_module_list(list( + linear_noinit(inner_dim, query_dim, bias = FALSE))) +}, + forward = function(hidden_states, image_rotary_emb = NULL, + chunk_size = NULL) { + value <- self$to_v(hidden_states)$unflatten(3L, c(self$heads, -1L)) + query <- self$to_q(hidden_states)$unflatten(3L, c(self$heads, -1L)) + key <- self$to_k(hidden_states)$unflatten(3L, c(self$heads, -1L)) + + # [B, S, H, D] -> [B, H, S, D]; norms and rotary in float32 + query <- .flux3_rms_norm_f32(query, self$norm_q)$transpose(2L, 3L) + key <- .flux3_rms_norm_f32(key, self$norm_k)$transpose(2L, 3L) + value <- value$transpose(2L, 3L) + if (!is.null(image_rotary_emb)) { + query <- flux_apply_rotary_emb(query, image_rotary_emb) + key <- flux_apply_rotary_emb(key, image_rotary_emb) + } + query <- query$to(dtype = value$dtype) + key <- key$to(dtype = value$dtype) + + attn <- .ltx23_sdpa(query, key, value, chunk_size = chunk_size) + attn <- attn$transpose(2L, 3L)$flatten(start_dim = 3L) + self$to_out[[1]](attn$to(dtype = hidden_states$dtype)) +} +) + +# torch.nn.RMSNorm applied to a float32 copy of x: the reference calls +# norm_q(query.float()), so the weight multiplies in float32 too. +.flux3_rms_norm_f32 <- function(x, norm) { + x <- x$to(dtype = torch::torch_float32()) + variance <- x$pow(2)$mean(dim = -1L, keepdim = TRUE) + x * torch::torch_rsqrt(variance + norm$eps) * + norm$weight$to(dtype = torch::torch_float32()) +} + +# Split a per-token modulation tensor [B, S or 1, 3 * dim] into +# (shift, scale, gate). Reference: torch.chunk(temb_mod, 3, dim=-1). +.flux3_mod_split <- function(mod) { + mod$chunk(3L, dim = -1L) +} + +#' FLUX 3 per-mode transformer block +#' +#' Parameterless LayerNorm, per-token (shift, scale) modulation, then +#' parallel attention and feed-forward from the same normed input, added +#' back under one gate. Self-attention stays within a single token mode. +#' Reference: Flux3ModeTransformerBlock. +#' +#' @param dim Integer. Model dimension. +#' @param num_attention_heads Integer. Attention heads. +#' @param attention_head_dim Integer. Per-head dimension. +#' @param mlp_ratio Numeric. Feed-forward multiplier (3.0). +#' @param eps Numeric. Norm epsilon. +#' +#' @return Module whose forward(hidden_states, temb_mod, +#' image_rotary_emb, chunk_size) returns the updated states. +#' +#' @export +flux3_mode_block <- torch::nn_module( + "flux3_mode_block", + initialize = function(dim, num_attention_heads, attention_head_dim, + mlp_ratio = 3.0, eps = 1e-6) { + self$norm <- torch::nn_layer_norm(dim, eps = eps, + elementwise_affine = FALSE) + self$attn <- flux3_attention(dim, num_attention_heads, + attention_head_dim, eps = eps) + self$ff <- flux2_feed_forward(dim, dim, mult = mlp_ratio) +}, + forward = function(hidden_states, temb_mod, image_rotary_emb = NULL, + chunk_size = NULL) { + mod <- .flux3_mod_split(temb_mod) + norm_h <- self$norm(hidden_states) * mod[[2]]$add(1) + mod[[1]] + attn_out <- self$attn(norm_h, image_rotary_emb = image_rotary_emb, + chunk_size = chunk_size) + hidden_states + mod[[3]] * (attn_out + self$ff(norm_h)) +} +) + +#' FLUX 3 joint (single-stream) block +#' +#' Text and content tokens are normed and modulated separately, then +#' concatenated (text first) for one parallel attention + feed-forward; +#' the output is split back and added under per-stream gates. +#' Reference: Flux3SingleTransformerBlock. +#' +#' @inheritParams flux3_mode_block +#' +#' @return Module whose forward(hidden_states, encoder_hidden_states, +#' temb_mod, temb_mod_context, image_rotary_emb, chunk_size) returns +#' \code{list(encoder_hidden_states, hidden_states)}. +#' +#' @export +flux3_single_block <- torch::nn_module( + "flux3_single_block", + initialize = function(dim, num_attention_heads, attention_head_dim, + mlp_ratio = 3.0, eps = 1e-6) { + self$norm <- torch::nn_layer_norm(dim, eps = eps, + elementwise_affine = FALSE) + self$attn <- flux3_attention(dim, num_attention_heads, + attention_head_dim, eps = eps) + self$ff <- flux2_feed_forward(dim, dim, mult = mlp_ratio) +}, + forward = function(hidden_states, encoder_hidden_states, temb_mod, + temb_mod_context, image_rotary_emb = NULL, + chunk_size = NULL) { + mod <- .flux3_mod_split(temb_mod) + c_mod <- .flux3_mod_split(temb_mod_context) + + norm_h <- self$norm(hidden_states) * mod[[2]]$add(1) + mod[[1]] + norm_c <- self$norm(encoder_hidden_states) * c_mod[[2]]$add(1) + + c_mod[[1]] + combined <- torch::torch_cat(list(norm_c, norm_h), dim = 2L) + + output <- self$attn(combined, image_rotary_emb = image_rotary_emb, + chunk_size = chunk_size) + self$ff(combined) + txt_len <- encoder_hidden_states$shape[2] + context_out <- output$narrow(2L, 1L, txt_len) + content_out <- output$narrow(2L, txt_len + 1L, hidden_states$shape[2]) + + list(encoder_hidden_states + c_mod[[3]] * context_out, + hidden_states + mod[[3]] * content_out) +} +) diff --git a/R/download_flux3.R b/R/download_flux3.R new file mode 100644 index 0000000..0950f70 --- /dev/null +++ b/R/download_flux3.R @@ -0,0 +1,157 @@ +#' Download and Prepare FLUX 3 Image Klein Weights +#' +#' Downloads FLUX 3 Image Klein from HuggingFace (gated; FLUX Kommunity +#' License) and quantizes the 5.6B transformer to a local fp8 (~6 GB) or +#' NF4 artifact. Only the diffusers-format files are fetched: the bf16 +#' transformer shards, the Qwen3-VL text encoder, the VAE and the +#' tokenizer. BFL's own root \code{model.safetensors} and the +#' \code{variants/} quantizations are not used. +#' +#' @name download_flux3 +NULL + +.flux3_transformer_files <- c( + "transformer/config.json", + "transformer/diffusion_pytorch_model.safetensors.index.json", + sprintf("transformer/diffusion_pytorch_model-%05d-of-00002.safetensors", + 1:2) +) + +.flux3_support_files <- c( + "LICENSE.md", + "NOTICE.md", + "vae/config.json", + "vae/diffusion_pytorch_model.safetensors", + "text_encoder/config.json", + "text_encoder/model.safetensors.index.json", + sprintf("text_encoder/model-%05d-of-00002.safetensors", 1:2), + "text_encoder/tokenizer.json", + "text_encoder/tokenizer_config.json", + "text_encoder/chat_template.json" +) + +# hub_download pinned to the validated snapshot, with the gated-repo +# 401/403 turned into an actionable error +.flux3_download <- function(file, revision = .flux3_revision) { + tryCatch( + hfhub::hub_download(.flux3_repo, file, revision = revision), + error = function(e) { + msg <- conditionMessage(e) + if (grepl("401|403|[Uu]nauthorized|[Ff]orbidden", msg)) { + stop( + "FLUX 3 Image Klein is a gated HuggingFace repo under the ", + "FLUX Kommunity License. To download:\n", + " 1. Log in at https://huggingface.co/", .flux3_repo, + " and accept the license.\n", + " 2. Create a read token at https://huggingface.co/settings/tokens\n", + " 3. Sys.setenv(HF_TOKEN = \"hf_...\") and retry.", + call. = FALSE + ) + } + stop(e) + } + ) +} + +#' Download FLUX 3 Image Klein and build the quantized artifact +#' +#' Skips work already done: a valid quantized manifest short-circuits +#' the transformer download; cached files are not re-fetched. Needs +#' \code{HF_TOKEN} for an account that accepted the license. The bf16 +#' transformer source (~11 GB in the HuggingFace cache) may be deleted +#' after quantization. +#' +#' The FLUX Kommunity License permits non-commercial use, and commercial +#' use of outputs by Qualifying Users (under US$5M gross annualized +#' revenue) who filter or review outputs before distributing them and +#' disclose AI generation where the law requires it. Read LICENSE.md in +#' the downloaded snapshot for the full terms. The quantized artifact is +#' a Derivative; keep it local. +#' +#' @param quantize Logical. Build the quantized artifact. +#' @param precision "auto" (default: fp8 when safetensors supports +#' float8, else nf4), "fp8" (~6 GB), or "nf4". +#' @param output_dir Directory for the quantized artifact. +#' @param text_encoders Logical. Also fetch the Qwen3-VL text encoder, +#' tokenizer and VAE (~9.2 GB). +#' @param verbose Logical. +#' +#' @return Invisibly, a list with \code{transformer_dir}, +#' \code{artifact_dir}, and \code{support} (named file paths). +#' +#' @export +download_flux3_klein <- function(quantize = TRUE, + precision = c("auto", "fp8", "nf4"), + output_dir = NULL, text_encoders = TRUE, + verbose = TRUE) { + precision <- match.arg(precision) + precision <- .flux_resolve_precision(precision, .flux3_artifact_prefix()) + precision <- .st_graceful_precision(precision, mode = "write") + output_dir <- output_dir %||% paste0(.flux3_artifact_prefix(), precision) + if (!requireNamespace("hfhub", quietly = TRUE)) { + stop("The hfhub package is required to download model weights.") + } + result <- list(transformer_dir = NULL, artifact_dir = output_dir, + support = character(0)) + + manifest_path <- file.path(output_dir, "manifest.json") + have_artifact <- file.exists(manifest_path) && { + m <- jsonlite::fromJSON(manifest_path) + all(file.exists(file.path(output_dir, m$shards))) + } + + if (!have_artifact || !quantize) { + if (!.hub_all_cached(.flux3_repo, .flux3_transformer_files)) { + free <- .ltx23_disk_free_gb(path.expand("~")) + if (!is.na(free) && free < 40) { + warning(sprintf( + "Only %.0f GB free; the downloads + %s artifact need ~30 GB.", + free, precision)) + } + ok <- .ltx23_consent(paste0( + "FLUX 3 Image Klein: the 11 GB bf16 transformer plus a local ", + precision, " artifact (FLUX Kommunity License, gated)")) + if (!ok) { + stop("Download cancelled.", call. = FALSE) + } + if (verbose) { + message("Downloading the FLUX 3 Image Klein transformer (11 GB)...") + } + } + paths <- vapply(.flux3_transformer_files, .flux3_download, character(1)) + result$transformer_dir <- dirname(paths[[1]]) + + if (quantize && !have_artifact) { + if (verbose) { + message("Quantizing transformer linears to ", precision, + " (one-time)...") + } + flux_quantize(result$transformer_dir, output_dir, + format = precision, verbose = verbose) + if (verbose) { + message(toupper(precision), " artifact ready: ", output_dir, + "\nThe 11 GB source in the HuggingFace cache may be ", + "deleted if you do not need bf16 weights.") + } + } + } else if (verbose) { + message(toupper(precision), " artifact already present: ", output_dir) + } + + if (text_encoders) { + if (!.hub_all_cached(.flux3_repo, .flux3_support_files)) { + ok <- .ltx23_consent( + "the Qwen3-VL-4B text encoder, tokenizer and VAE (~9.2 GB)") + if (!ok) { + stop("Download cancelled.", call. = FALSE) + } + if (verbose) { + message("Downloading text encoder + VAE...") + } + } + result$support <- vapply(.flux3_support_files, .flux3_download, + character(1)) + } + + invisible(result) +} diff --git a/R/encode_flux3.R b/R/encode_flux3.R new file mode 100644 index 0000000..1687562 --- /dev/null +++ b/R/encode_flux3.R @@ -0,0 +1,70 @@ +#' Encode a prompt for FLUX 3 +#' +#' The Qwen3-VL prompt path of the reference pipeline (diffusers +#' Flux3TextEncoderStep._get_qwen3_vl_prompt_embeds, Apache-2.0): the +#' prompt goes through the chat template (one user turn plus the +#' generation prompt, no thinking block), is right-padded to a multiple +#' of \code{pad_multiple} tokens (capped at \code{max_sequence_length}), +#' and the hidden states of \code{out_layers} are concatenated per +#' token. The padding tokens get position id 1, as in the training +#' environment, and the transformer attends to them, so their states are +#' part of the embedding. +#' +#' FLUX 3 expects a full scene caption: a short prompt gives generic +#' output. Expand short requests with an LLM first (BFL publishes the +#' system prompts with flux-3-image-klein-pup). +#' +#' @param prompt Character scalar. +#' @param model A \code{\link{qwen3_encoder}} loaded from FLUX 3's +#' \code{text_encoder} with \code{\link{load_qwen3_text_encoder}}. +#' @param tokenizer A \code{\link{qwen_bpe_tokenizer}} from the same +#' directory. +#' @param max_sequence_length Integer. Token cap (8192). +#' @param pad_multiple Integer. Window rounding (80). +#' @param out_layers Integer vector. Hidden-state layers to stack +#' (4, 8, ..., 32). +#' @param device Character or NULL. Defaults to the model's device. +#' +#' @return \code{list(prompt_embeds, txt_ids)}: embeddings [1, L, +#' length(out_layers) * hidden] and position ids [L, 4] from +#' \code{\link{flux3_prepare_text_ids}}. +#' +#' @export +flux3_encode_prompt <- function(prompt, model, tokenizer, + max_sequence_length = 8192L, + pad_multiple = 80L, + out_layers = seq(4L, 32L, by = 4L), + device = NULL) { + if (!is.character(prompt) || length(prompt) != 1L) { + stop("`prompt` must be a single string.") + } + ids <- encode_qwen(tokenizer, prompt, max_length = NULL, + chat_template = TRUE, enable_thinking = TRUE)$input_ids[[1]] + win <- .flux3_text_window(ids, tokenizer$pad_id, max_sequence_length, + pad_multiple) + + device <- device %||% model$model$embed_tokens$weight$device + long <- torch::torch_long() + input_ids <- torch::torch_tensor(matrix(win$ids + 1L, nrow = 1L), + dtype = long, device = device) + mask <- torch::torch_tensor(matrix(win$mask, nrow = 1L), dtype = long, + device = device) + states <- torch::with_no_grad(model(input_ids, attention_mask = mask, + out_layers = out_layers, + position_ids = win$positions)) + list(prompt_embeds = torch::torch_cat(states, dim = -1L), + txt_ids = flux3_prepare_text_ids(length(win$ids), device = device)) +} + +# Truncate to max_len, right-pad to a multiple of pad_multiple (capped at +# max_len). Padding positions are 1, matching the reference's +# arange(...).masked_fill(attention_mask == 0, 1). +.flux3_text_window <- function(ids, pad_id, max_len, pad_multiple) { + ids <- utils::head(as.integer(ids), max_len) + n <- length(ids) + padded <- min(ceiling(n / pad_multiple) * pad_multiple, max_len) + pad <- padded - n + list(ids = c(ids, rep(as.integer(pad_id), pad)), + mask = c(rep(1L, n), rep(0L, pad)), + positions = c(seq_len(n) - 1L, rep(1L, pad))) +} diff --git a/R/img2img_flux2.R b/R/img2img_flux2.R index e57d798..0da4133 100644 --- a/R/img2img_flux2.R +++ b/R/img2img_flux2.R @@ -69,33 +69,7 @@ img2img_flux2 <- function(prompt, image, pipeline = NULL, width = NULL, #' #' @export flux2_preprocess_image <- function(image, max_area = 1024 * 1024) { - if (is.character(image)) { - if (!file.exists(image)) { - stop("image file not found: ", image, call. = FALSE) - } - image <- if (grepl("\\.png$", image, ignore.case = TRUE)) { - png::readPNG(image) - } else if (grepl("\\.jpe?g$", image, ignore.case = TRUE)) { - jpeg::readJPEG(image) - } else { - stop("image must be a .png, .jpg or .jpeg file", call. = FALSE) - } - } - if (!is.array(image) && !is.matrix(image)) { - stop("image must be a file path or an [H, W, 3] array", call. = FALSE) - } - if (length(dim(image)) == 2L) { - image <- array(rep(image, 3L), dim = c(dim(image), 3L)) - } - if (length(dim(image)) != 3L) { - stop("image must be an [H, W, 3] array", call. = FALSE) - } - if (dim(image)[3] == 4L) { - image <- image[,, 1:3, drop = FALSE] - } - if (dim(image)[3] != 3L) { - stop("image must have 3 colour channels", call. = FALSE) - } + image <- .read_rgb_image(image) h <- dim(image)[1] w <- dim(image)[2] @@ -126,6 +100,40 @@ flux2_preprocess_image <- function(image, max_area = 1024 * 1024) { x$unsqueeze(1L)$mul(2)$sub(1) } +# Read a PNG/JPEG path or take an array; return an [H, W, 3] array in +# [0, 1]. Grayscale is replicated to 3 channels; alpha is dropped (as +# PIL's convert("RGB") does). +.read_rgb_image <- function(image) { + if (is.character(image)) { + if (!file.exists(image)) { + stop("image file not found: ", image, call. = FALSE) + } + image <- if (grepl("\\.png$", image, ignore.case = TRUE)) { + png::readPNG(image) + } else if (grepl("\\.jpe?g$", image, ignore.case = TRUE)) { + jpeg::readJPEG(image) + } else { + stop("image must be a .png, .jpg or .jpeg file", call. = FALSE) + } + } + if (!is.array(image) && !is.matrix(image)) { + stop("image must be a file path or an [H, W, 3] array", call. = FALSE) + } + if (length(dim(image)) == 2L) { + image <- array(rep(image, 3L), dim = c(dim(image), 3L)) + } + if (length(dim(image)) != 3L) { + stop("image must be an [H, W, 3] array", call. = FALSE) + } + if (dim(image)[3] == 4L) { + image <- image[,, 1:3, drop = FALSE] + } + if (dim(image)[3] != 3L) { + stop("image must have 3 colour channels", call. = FALSE) + } + image +} + # PIL's Lanczos-3 resampling weights for one axis (ImagingResample # precompute_coeffs): the kernel support widens by the downscale factor, # which is what antialiases, and each row is normalized to sum to 1. diff --git a/R/img2img_flux3.R b/R/img2img_flux3.R new file mode 100644 index 0000000..52745d6 --- /dev/null +++ b/R/img2img_flux3.R @@ -0,0 +1,67 @@ +#' Edit an image with FLUX 3 Image Klein +#' +#' Image + caption to image with the same weights as +#' \code{\link{txt2img_flux3}}. The reference is kept whole (see +#' \code{\link{flux3_preprocess_image}}): scaled so its canvas fits +#' \code{max_area}, reflect-padded to the 16 px grid, VAE-encoded once, +#' and passed to every transformer forward as its own token mode +#' (position ids at L = -1). The output defaults to the canvas size and +#' is cropped back to the reference's box. Reference: diffusers +#' Flux3DistilledModularPipeline with \code{image =} (Apache-2.0). +#' +#' Write the caption as a full description of the edited result (the +#' reference with the edit folded in); a one-line instruction reframes +#' the shot. BFL's prompt upsampler (flux-3-image-klein-pup) writes such +#' captions from an instruction plus the reference image. +#' +#' @param prompt Character. The full caption of the edited image. +#' @param image Path to a PNG/JPEG, or an [H, W, 3] array in [0, 1]. +#' @param pipeline A \code{flux3_pipeline} loaded with +#' \code{encoder = TRUE}; NULL loads one (passing \code{...} through). +#' @param width,height Integers divisible by 16, or NULL (default) for +#' the reference canvas, cropped back to the reference afterwards. +#' @param max_area Numeric. Reference canvas area cap (reference: 4 +#' megapixels). The reference adds \code{area / 256} tokens to every +#' forward, so lower it on small cards. +#' @param num_inference_steps Integer. Denoising steps (16). +#' @param max_sequence_length Integer. Qwen3-VL token cap (8192). +#' @param seed Integer or NULL. See \code{\link{txt2img_flux3}}. +#' @param prompt_embeds Optional precomputed [1, S, 20480] embeddings. +#' @param save_file Logical. Write a PNG. +#' @param filename Output path (default derived from the prompt). +#' @param verbose Logical, or one of "silent", "progress", "steps". +#' @param ... Passed to \code{\link{flux3_load_pipeline}} when +#' \code{pipeline} is NULL. +#' +#' @return Invisibly, \code{list(image, metadata)} where \code{image} is +#' an [H, W, 3] array in [0, 1]. +#' +#' @export +img2img_flux3 <- function(prompt, image, pipeline = NULL, width = NULL, + height = NULL, max_area = 4 * 1024 * 1024, + num_inference_steps = 16L, + max_sequence_length = 8192L, seed = NULL, + prompt_embeds = NULL, save_file = TRUE, + filename = NULL, verbose = TRUE, ...) { + level <- .verbosity(verbose) + pre <- flux3_preprocess_image(image, max_area = max_area) + canvas_w <- as.integer(pre$image$shape[4]) + canvas_h <- as.integer(pre$image$shape[3]) + # Crop back only when the output is the canvas (the reference's + # default); an explicit size is returned whole + crop <- if (is.null(width) && is.null(height)) pre$crop else NULL + width <- as.integer(width %||% canvas_w) + height <- as.integer(height %||% canvas_h) + if (level != "silent") { + message(sprintf("FLUX 3 klein edit: %dx%d from a %dx%d reference canvas, %d steps", + width, height, canvas_w, canvas_h, + as.integer(num_inference_steps))) + } + if (is.null(pipeline)) { + pipeline <- flux3_load_pipeline(..., encoder = TRUE, + verbose = level == "steps") + } + .flux3_generate(prompt, pipeline, width, height, num_inference_steps, + max_sequence_length, seed, prompt_embeds, save_file, + filename, level, reference = pre$image, crop = crop) +} diff --git a/R/preprocess_flux3.R b/R/preprocess_flux3.R new file mode 100644 index 0000000..09c6027 --- /dev/null +++ b/R/preprocess_flux3.R @@ -0,0 +1,88 @@ +#' Preprocess a reference image for FLUX 3 +#' +#' The single-reference ("fill") path of the reference pipeline +#' (diffusers Flux3ImageProcessor, Apache-2.0, following BFL's +#' \code{reference_canvas}): the image is kept whole. It is downscaled, +#' keeping its aspect ratio, until its canvas (its size rounded up to the +#' 16 px token grid, sides clamped to [64, \code{max_side}]) fits +#' \code{max_area}; backing the scale off by 0.5\% per try when the +#' rounding overshoots. The resize is PIL's antialiased Lanczos-3 in +#' floating point. The image is then centered on the canvas and the +#' border filled with its reflection (numpy \code{"reflect"}, or +#' \code{"symmetric"} when a side is at most 16 px). +#' +#' @param image Path to a PNG/JPEG, or an [H, W, C] array in [0, 1]. +#' @param max_area Numeric. Canvas area cap (reference: 4 megapixels). +#' Lower it to bound memory; the reference tokens are +#' \code{canvas area / 256}. +#' @param max_side Integer. Canvas side cap. +#' +#' @return \code{list(image, crop)}: a float32 tensor [1, 3, H, W] in +#' [-1, 1] at the canvas size, and the (left, top, right, bottom) box +#' of the image inside the canvas (0-based, right/bottom exclusive), +#' used to crop the generated image back. +#' +#' @export +flux3_preprocess_image <- function(image, max_area = 4 * 1024 * 1024, + max_side = 4096L) { + image <- .read_rgb_image(image) + h <- dim(image)[1] + w <- dim(image)[2] + + scale <- min(1, sqrt(max_area / (w * h)), max_side / max(w, h)) + repeat { + nh <- max(1, round(h * scale)) + nw <- max(1, round(w * scale)) + canvas <- .flux3_canvas(nh, nw, max_side) + if (prod(canvas) <= max_area) { + break + } + scale <- scale * 0.995 + } + + x <- torch::torch_tensor(image, dtype = torch::torch_float32()) + x <- x$permute(c(3L, 1L, 2L)) # [3, H, W] + if (nh != h || nw != w) { + x <- .lanczos_resize(x, nh, nw) + } + + ch <- canvas[1] + cw <- canvas[2] + top <- (ch - nh) %/% 2 + left <- (cw - nw) %/% 2 + mode <- if (min(nh, nw) > 16) "reflect" else "symmetric" + rows <- .flux3_pad_index(nh, top, ch - nh - top, mode) + cols <- .flux3_pad_index(nw, left, cw - nw - left, mode) + x <- x[, rows, cols, drop = FALSE] + + list(image = x$unsqueeze(1L)$mul(2)$sub(1), + crop = c(left = left, top = top, right = left + nw, bottom = top + nh)) +} + +# Canvas (height, width): each side rounded up to a multiple of 16 and +# clamped to [min_side, max_side]. Reference: _get_canvas_height_width. +.flux3_canvas <- function(h, w, max_side = 4096L, min_side = 64L) { + vapply(c(h, w), function(side) { + min(max_side, max(min_side, ceiling(side / 16) * 16)) + }, numeric(1)) +} + +# 1-based source indices for numpy.pad(mode = "reflect" | "symmetric") +# along one axis of length n, padded by `before` and `after`. Both modes +# are periodic extensions, so any pad width works. +.flux3_pad_index <- function(n, before, after, mode) { + j <- seq(-before, n + after - 1) + if (n == 1) { + return(rep(1L, length(j))) + } + if (mode == "reflect") { + period <- 2 * (n - 1) + i <- j %% period + i <- ifelse(i >= n, period - i, i) + } else { + period <- 2 * n + i <- j %% period + i <- ifelse(i >= n, period - 1 - i, i) + } + as.integer(i) + 1L +} diff --git a/R/quantize_flux.R b/R/quantize_flux.R index 5c8cc83..0491792 100644 --- a/R/quantize_flux.R +++ b/R/quantize_flux.R @@ -179,6 +179,35 @@ NULL args } +.flux3_transformer_args <- function(config) { + if (is.null(config)) { + return(list()) + } + args <- list( + in_channels = config$in_channels, + cond_in_channels = config$cond_in_channels, + num_layers = config$num_layers, + num_single_layers = config$num_single_layers, + attention_head_dim = config$attention_head_dim, + num_attention_heads = config$num_attention_heads, + joint_attention_dim = config$joint_attention_dim, + timestep_channels = config$timestep_channels, + axes_dims_rope = config$axes_dims_rope + ) + args <- Filter(function(x) !is.null(x) && length(x) > 0L, args) + args <- lapply(args, function(x) if (is.numeric(x)) as.integer(x) else x) + for (field in c("mlp_ratio", "rope_theta", "eps")) { + v <- config[[field]] + if (!is.null(v) && length(v) == 1L) { + args[[field]] <- as.numeric(v) + } + } + if (!is.null(config$flow_map)) { + args$flow_map <- isTRUE(config$flow_map) + } + args +} + .zimage_transformer_args <- function(config) { if (is.null(config)) { return(list()) @@ -282,6 +311,12 @@ NULL if (family == "flux2") { list(model_fn = flux2_transformer, args_fn = .flux2_transformer_args, is_quant_key = flux2_is_quant_key, shard_prefix = "flux2-klein") + } else if (family == "flux3") { + # skeleton: built uninitialized in the compute dtype (5.6B params + # would otherwise init in float32, ~22 GB, before the cast) + list(model_fn = flux3_transformer, args_fn = .flux3_transformer_args, + is_quant_key = flux3_is_quant_key, shard_prefix = "flux3-klein", + skeleton = TRUE) } else if (family == "zimage") { list(model_fn = zimage_transformer, args_fn = .zimage_transformer_args, @@ -500,7 +535,13 @@ flux_load_transformer <- function(ckpt, device = "cuda", dtype = "bfloat16", hooks <- .flux_family_hooks(ckpt$config) args <- utils::modifyList(hooks$args_fn(ckpt$config), list(...)) - model <- do.call(hooks$model_fn, args) + model <- if (isTRUE(hooks$skeleton)) { + # Every parameter is filled below (unfilled ones are an error) + do.call(.construct_skeleton, c(list(hooks$model_fn), args, + list(dtype = .flux_dtype(dtype)))) + } else { + do.call(hooks$model_fn, args) + } if (format == "full") { model$to(dtype = .flux_dtype(dtype)) diff --git a/R/qwen3_text_encoder.R b/R/qwen3_text_encoder.R index c68390a..19a59aa 100644 --- a/R/qwen3_text_encoder.R +++ b/R/qwen3_text_encoder.R @@ -16,15 +16,22 @@ NULL # Llama-convention rotary tables: cos/sin [S, head_dim/2] at the given # theta, applied with the split-half kernel (x1*cos - x2*sin | x2*cos + # x1*sin), which matches rotate_half exactly. -.qwen3_rope_tables <- function(seq_len, head_dim, theta, device) { +# `positions` (0-based integers, length seq_len) overrides 0..seq_len-1. +.qwen3_rope_tables <- function(seq_len, head_dim, theta, device, + positions = NULL) { f32 <- torch::torch_float32() inv_freq <- 1.0 / torch::torch_pow( theta, torch::torch_arange(start = 0, end = head_dim - 2, step = 2, dtype = f32, device = device) / head_dim ) - pos <- torch::torch_arange(start = 0, end = seq_len - 1, dtype = f32, - device = device) + pos <- if (is.null(positions)) { + torch::torch_arange(start = 0, end = seq_len - 1, dtype = f32, + device = device) + } else { + torch::torch_tensor(as.numeric(positions), dtype = f32, + device = device) + } freqs <- pos$unsqueeze(2L) * inv_freq$unsqueeze(1L) # [S, D/2] # [1, 1, S, D/2] for broadcasting against [B, H, S, D/2] list( @@ -40,14 +47,14 @@ NULL self$num_heads <- num_heads self$num_kv_heads <- num_kv_heads self$head_dim <- head_dim - self$q_proj <- torch::nn_linear(hidden_size, num_heads * head_dim, - bias = FALSE) - self$k_proj <- torch::nn_linear(hidden_size, num_kv_heads * head_dim, - bias = FALSE) - self$v_proj <- torch::nn_linear(hidden_size, num_kv_heads * head_dim, - bias = FALSE) - self$o_proj <- torch::nn_linear(num_heads * head_dim, hidden_size, - bias = FALSE) + self$q_proj <- linear_noinit(hidden_size, num_heads * head_dim, + bias = FALSE) + self$k_proj <- linear_noinit(hidden_size, num_kv_heads * head_dim, + bias = FALSE) + self$v_proj <- linear_noinit(hidden_size, num_kv_heads * head_dim, + bias = FALSE) + self$o_proj <- linear_noinit(num_heads * head_dim, hidden_size, + bias = FALSE) self$q_norm <- ltx23_rms_norm(head_dim, eps = eps) self$k_norm <- ltx23_rms_norm(head_dim, eps = eps) }, @@ -83,12 +90,12 @@ NULL .qwen3_mlp <- torch::nn_module( "qwen3_mlp", initialize = function(hidden_size, intermediate_size) { - self$gate_proj <- torch::nn_linear(hidden_size, intermediate_size, - bias = FALSE) - self$up_proj <- torch::nn_linear(hidden_size, intermediate_size, - bias = FALSE) - self$down_proj <- torch::nn_linear(intermediate_size, hidden_size, - bias = FALSE) + self$gate_proj <- linear_noinit(hidden_size, intermediate_size, + bias = FALSE) + self$up_proj <- linear_noinit(hidden_size, intermediate_size, + bias = FALSE) + self$down_proj <- linear_noinit(intermediate_size, hidden_size, + bias = FALSE) }, forward = function(x) { self$down_proj(torch::nnf_silu(self$gate_proj(x)) * self$up_proj(x)) @@ -125,10 +132,13 @@ NULL #' @param rms_norm_eps Numeric. #' #' @return Module whose forward(input_ids, attention_mask = NULL, -#' out_layers) returns a list of hidden-state tensors [B, S, hidden], -#' one per requested layer (a value of k means the state after k -#' layers, matching HF \code{output.hidden_states[k]}). Runs only to -#' \code{max(out_layers)}. \code{input_ids} are 1-based. +#' out_layers, position_ids = NULL) returns a list of hidden-state +#' tensors [B, S, hidden], one per requested layer (a value of k means +#' the state after k layers, matching HF \code{output.hidden_states[k]}, +#' which carries the final norm when k is the last layer). +#' Runs only to \code{max(out_layers)}. \code{input_ids} are 1-based; +#' \code{position_ids} (0-based integers, length S, shared by the batch) +#' override the default 0..S-1 rotary positions. #' #' @export qwen3_encoder <- torch::nn_module( @@ -159,14 +169,15 @@ qwen3_encoder <- torch::nn_module( self$model <- inner() }, forward = function(input_ids, attention_mask = NULL, - out_layers = c(9L, 18L, 27L)) { + out_layers = c(9L, 18L, 27L), position_ids = NULL) { x <- self$model$embed_tokens(input_ids) b <- input_ids$shape[1] s <- input_ids$shape[2] device <- x$device dt <- x$dtype - rope <- .qwen3_rope_tables(s, self$head_dim, self$rope_theta, device) + rope <- .qwen3_rope_tables(s, self$head_dim, self$rope_theta, device, + positions = position_ids) # Additive causal (+ padding) mask [B, 1, S, S] in the compute dtype: # the fused SDPA kernel requires the bias dtype to match the query. @@ -185,11 +196,13 @@ qwen3_encoder <- torch::nn_module( out_layers <- sort(as.integer(out_layers)) states <- vector("list", length(out_layers)) + n_layers <- length(self$model$layers) for (i in seq_len(max(out_layers))) { x <- self$model$layers[[i]](x, rope, mask) hit <- which(out_layers == i) if (length(hit)) { - states[[hit]] <- x + # HF's last hidden_states entry carries the final norm + states[[hit]] <- if (i == n_layers) self$model$norm(x) else x } } states @@ -200,10 +213,15 @@ qwen3_encoder <- torch::nn_module( #' #' Streams the (possibly sharded) safetensors weights into #' \code{\link{qwen3_encoder}}. The LM head is tied to the embeddings -#' and skipped. +#' and skipped. A Qwen3-VL checkpoint (FLUX 3's text encoder: config +#' under \code{text_config}, weights under \code{model.language_model.*}) +#' loads its text model only; the vision tower is skipped. For text-only +#' input with 1D positions, Qwen3-VL's interleaved MRoPE reduces to the +#' plain Qwen3 rotary (all three axes carry the same position). #' #' @param model_path Directory with \code{config.json} and -#' \code{model*.safetensors} (FLUX.2-klein's \code{text_encoder}). +#' \code{model*.safetensors} (FLUX.2-klein's or FLUX 3's +#' \code{text_encoder}). #' @param device Character. Target device. #' @param dtype Character. "bfloat16" (GPU) or "float32" (CPU). #' @param verbose Logical. @@ -220,6 +238,11 @@ load_qwen3_text_encoder <- function(model_path, device = "cpu", if (file.exists(config_path)) { config <- jsonlite::fromJSON(config_path, simplifyVector = TRUE) } + # Qwen3-VL: text model config nested, weights under language_model + vl <- !is.null(config$text_config) + if (vl) { + config <- config$text_config + } args <- list(vocab_size = config$vocab_size, hidden_size = config$hidden_size, @@ -230,15 +253,20 @@ load_qwen3_text_encoder <- function(model_path, device = "cpu", head_dim = config$head_dim) args <- Filter(function(x) !is.null(x) && length(x) > 0L, args) args <- lapply(args, as.integer) - if (!is.null(config$rope_theta)) { - args$rope_theta <- as.numeric(config$rope_theta) + # transformers >= 5 writes rope_theta under rope_parameters + rope_theta <- config$rope_theta %||% config$rope_parameters$rope_theta + if (!is.null(rope_theta)) { + args$rope_theta <- as.numeric(rope_theta) } if (!is.null(config$rms_norm_eps)) { args$rms_norm_eps <- as.numeric(config$rms_norm_eps) } args <- utils::modifyList(args, list(...)) - model <- do.call(qwen3_encoder, args) - model$to(dtype = .flux_dtype(dtype)) + # Build at the target dtype: a 4B fp32 build before the bf16 cast + # needs 16 GB of host RAM + model <- do.call(.construct_skeleton, + c(list(qwen3_encoder), args, + list(dtype = .flux_dtype(dtype)))) opened <- .flux_open_sharded_dir(model_path, "model") ckpt <- structure( @@ -250,6 +278,12 @@ load_qwen3_text_encoder <- function(model_path, device = "cpu", if (identical(key, "lm_head.weight")) { return(NA_character_) # tied to embed_tokens } + if (vl) { + if (startsWith(key, "model.visual.")) { + return(NA_character_) + } + key <- sub("^model\\.language_model\\.", "model.", key) + } key } res <- ltx23_load_group(ckpt, ckpt$keys, model, map_key = map_key, diff --git a/R/recommend.R b/R/recommend.R index 5005067..7cc4209 100644 --- a/R/recommend.R +++ b/R/recommend.R @@ -199,6 +199,7 @@ recommend <- function(model = c("sd21", "sdxl", "flux1", "flux2", "zimage", flux1 = c(bf16 = 34, fp8 = 22, nf4 = 17), # DiT + T5 bf16 host copy flux2 = c(bf16 = 17, fp8 = 12, nf4 = 10), # DiT + Qwen3 bf16 zimage = c(bf16 = 21, fp8 = 14, nf4 = 12), # DiT + Qwen3 bf16 + flux3 = c(bf16 = 20, fp8 = 15, nf4 = 12), # DiT + Qwen3-VL text bf16 ltx = c(fp8 = 34, nf4 = 28) # checkpoint + Gemma3 NF4 ) v <- unname(sets[[model]][precision]) diff --git a/R/rope_flux3.R b/R/rope_flux3.R new file mode 100644 index 0000000..e4658df --- /dev/null +++ b/R/rope_flux3.R @@ -0,0 +1,68 @@ +#' FLUX 3 position ids +#' +#' 4-axis (T, H, W, L) position ids for the FLUX 3 transformer, from the +#' diffusers reference (Apache-2.0, +#' src/diffusers/modular_pipelines/flux3/before_denoise.py). Image tokens +#' sit at T = 0 with the H/W axes scaled by the image RoPE scale (0.5); +#' the generated image has L = 0 and the i-th reference L = -i. Text +#' advances along L. Rotary frequencies come from +#' \code{\link{flux_pos_embed}} with \code{axes_dim = c(32, 32, 32, 32)}, +#' \code{theta = 10000}. +#' +#' @name rope_flux3 +NULL + +#' FLUX 3 image token position ids +#' +#' @param height Integer. Token grid rows (pixels / 16). +#' @param width Integer. Token grid columns (pixels / 16). +#' @param l_coord Numeric. L coordinate: 0 for the generated image, +#' \code{-i} for the i-th reference image. +#' @param rope_scale Numeric. Scale of the H/W axes (0.5). +#' @param device Character. Device. +#' +#' @return Float tensor [height * width, 4], row-major over (h, w). +#' +#' @export +flux3_prepare_image_ids <- function(height, width, l_coord = 0, + rope_scale = 0.5, device = "cpu") { + h <- rep(seq_len(height) - 1, each = width) + w <- rep(seq_len(width) - 1, times = height) + ids <- cbind(0, h * rope_scale, w * rope_scale, l_coord) + torch::torch_tensor(ids, dtype = torch::torch_float32(), device = device) +} + +#' FLUX 3 text token position ids +#' +#' @param len Integer. Text sequence length (padded). +#' @param device Character. Device. +#' +#' @return Float tensor [len, 4]: (0, 0, 0, 0..len-1). +#' +#' @export +flux3_prepare_text_ids <- function(len, device = "cpu") { + ids <- cbind(0, 0, 0, seq_len(len) - 1) + torch::torch_tensor(ids, dtype = torch::torch_float32(), device = device) +} + +#' FLUX 3 distilled sampling schedule +#' +#' FlowMatch Euler with the static shift of the FLUX 3 recipe (6.93) +#' over \code{sigmas = linspace(1, 1 / n, n)}. Reference: +#' Flux3SetTimestepsStep. +#' +#' @param num_inference_steps Integer (16 for the distilled checkpoint). +#' @param shift Numeric. Static timestep shift. +#' @param device Character. Device. +#' +#' @return A FlowMatch schedule from \code{\link{flowmatch_set_timesteps}}. +#' +#' @export +flux3_schedule <- function(num_inference_steps = 16L, shift = 6.93, + device = "cpu") { + n <- as.integer(num_inference_steps) + schedule <- flowmatch_scheduler_create(shift = shift) + flowmatch_set_timesteps(schedule, num_inference_steps = n, + device = device, + sigmas = seq(1, 1 / n, length.out = n)) +} diff --git a/R/txt2img.R b/R/txt2img.R index 0b4f2c9..9ef2f8a 100644 --- a/R/txt2img.R +++ b/R/txt2img.R @@ -12,7 +12,8 @@ #' img <- txt2img("a cat wearing sunglasses in space", device = "cuda") #' } txt2img <- function(prompt, - model_name = c("sd21", "sdxl", "flux1", "flux2", "zimage"), + model_name = c("sd21", "sdxl", "flux1", "flux2", "zimage", + "flux3"), ...) { model_name <- match.arg(model_name) switch(model_name, @@ -22,6 +23,7 @@ txt2img <- function(prompt, "flux1" = txt2img_flux(prompt, ...), "flux2" = txt2img_flux2(prompt, ...), "zimage" = txt2img_zimage(prompt, ...), + "flux3" = txt2img_flux3(prompt, ...), # "sd3" = txt2img_sd3(prompt, ...), stop("Unsupported model: ", model_name) ) diff --git a/R/txt2img_flux3.R b/R/txt2img_flux3.R new file mode 100644 index 0000000..e631ae0 --- /dev/null +++ b/R/txt2img_flux3.R @@ -0,0 +1,470 @@ +#' FLUX 3 Image Klein Pipeline +#' +#' FLUX 3 Image Klein text-to-image and single-reference editing, ported +#' from the diffusers reference (Apache-2.0, +#' src/diffusers/modular_pipelines/flux3/). The distilled checkpoint is a +#' flow map: 16 steps, no classifier-free guidance, and every step also +#' tells the transformer the timestep it jumps to. The FlowMatch +#' schedule uses the static shift 6.93. Latent noise is drawn directly in +#' the packed 128-channel space. +#' +#' The weights are under the FLUX Kommunity License (non-commercial use, +#' plus commercial use of outputs by Qualifying Users with under US$5M +#' revenue, who must filter or review outputs; see the license in the +#' downloaded repository). diffuseR never redistributes them. +#' +#' @name txt2img_flux3 +NULL + +.flux3_repo <- "black-forest-labs/flux-3-image-klein" + +# The snapshot the port was validated against; support files resolve +# straight out of snapshots// (see .flux2_rev). +.flux3_revision <- "337337621b8526e120f2275d250242f4a6bedd75" + +.flux3_artifact_prefix <- function() { + file.path(tools::R_user_dir("diffuseR", "data"), "flux3-klein-") +} + +# Resolve a FLUX 3 file from the HuggingFace cache +.flux3_cached <- function(file, revision = .flux3_revision) { + if (!requireNamespace("hfhub", quietly = TRUE)) { + stop("The hfhub package is required to locate model files.") + } + rev <- .flux2_rev(revision) + tryCatch( + do.call(hfhub::hub_download, + c(list(.flux3_repo, file, local_files_only = TRUE), rev)), + error = function(e) { + stop("Missing ", file, " in the HuggingFace cache; ", + "run download_flux3_klein() first.", call. = FALSE) + } + ) +} + +#' Load the FLUX 3 Image Klein pipeline +#' +#' Loads the quantized transformer artifact plus the VAE (encoder and +#' decoder), the Qwen3-VL text encoder (text model only) and the +#' tokenizer from the HuggingFace cache populated by +#' \code{\link{download_flux3_klein}}. With phase offload each component +#' visits the GPU only for its own phase, so the GPU peak is the largest +#' phase rather than the sum: the ~8 GB Qwen3-VL encode, or the fp8 +#' transformer (~6 GB) plus activations. +#' +#' @param model_dir Quantized artifact directory (default: the +#' \code{download_flux3_klein} location for \code{precision}), or a raw +#' diffusers transformer directory. +#' @param device Character. Compute device. +#' @param precision "auto" (default: reuse an existing artifact, else +#' fp8 when safetensors supports float8, else nf4), "fp8", "nf4", or +#' "bf16" (the raw diffusers transformer). +#' @param text_device Device for the Qwen3-VL encoder (default: +#' \code{device}; it encodes in its own phase and offloads). "cpu" +#' keeps it in float32 on the host. +#' @param attn_chunk Integer or NULL. Attention query-chunk override. +#' @param phase_offload Logical. One GPU tenant per phase. +#' @param pin Logical or NULL. Page-lock the phase-swapped weights (see +#' \code{\link{staging}}). +#' @param encoder Logical. Also load the VAE encoder, which +#' \code{\link{img2img_flux3}} needs. +#' @param revision Exact 40-hex commit for the support files (default: +#' the validated snapshot). +#' @param dtype Compute dtype, "bfloat16" or "float32". NULL (default): +#' bfloat16 on CUDA, float32 on the CPU. "bfloat16" on the CPU halves +#' host memory (the raw transformer is ~11 GB instead of ~22 GB). +#' @param verbose Logical. +#' +#' @return A \code{flux3_pipeline} list. +#' +#' @export +flux3_load_pipeline <- function(model_dir = NULL, device = "cuda", + precision = c("auto", "fp8", "nf4", "bf16"), + text_device = NULL, attn_chunk = NULL, + phase_offload = TRUE, pin = NULL, + encoder = TRUE, revision = .flux3_revision, + dtype = NULL, verbose = TRUE) { + .flux2_rev(revision) + precision <- match.arg(precision) + if (precision != "bf16") { + precision <- .flux_resolve_precision(precision, + .flux3_artifact_prefix()) + } + text_device <- text_device %||% device + if (is.null(model_dir)) { + model_dir <- if (precision == "bf16") { + dirname(.flux3_cached("transformer/config.json", revision)) + } else { + paste0(.flux3_artifact_prefix(), precision) + } + } + + ckpt <- if (file.exists(file.path(model_dir, "manifest.json"))) { + flux_open_quantized(model_dir) + } else { + flux_open_checkpoint(model_dir) + } + pin <- .resolve_pin(pin, "flux3", ckpt$format %||% "full") + + if (!nzchar(Sys.getenv("PYTORCH_CUDA_ALLOC_CONF"))) { + Sys.setenv(PYTORCH_CUDA_ALLOC_CONF = "backend:native") + } + if (device == "cuda") { + # Gate footprint sized to the largest phase (the ~8 GB Qwen3-VL + # encode), as for klein + .flux_gc_gates(footprint_gb = 12) + } + component_device <- if (phase_offload) "cpu" else device + compute_dtype <- dtype %||% if (device == "cpu") "float32" else "bfloat16" + compute_dtype <- match.arg(compute_dtype, c("bfloat16", "float32")) + + pipe <- list( + format = ckpt$format %||% "full", + device = device, + compute_dtype = compute_dtype, + text_device = text_device, + phase_offload = phase_offload, + attn_chunk = if (is.null(attn_chunk)) NULL else as.integer(attn_chunk), + config = ckpt$config + ) + + if (verbose) { + message("Loading FLUX 3 transformer (", pipe$format, ")...") + } + pipe$transformer <- flux_load_transformer( + ckpt, device = component_device, dtype = compute_dtype, + pin = FALSE, fp8_resident = FALSE, verbose = verbose + ) + pipe$fp8_resident <- identical(pipe$format, "fp8") && device == "cuda" + + if (verbose) { + message("Loading FLUX 3 VAE...") + } + vae_config <- jsonlite::fromJSON(.flux3_cached("vae/config.json", revision)) + vae_path <- .flux3_cached("vae/diffusion_pytorch_model.safetensors", + revision) + latent_channels <- as.integer(vae_config$latent_channels %||% 32L) + pipe$vae_bn_eps <- vae_config$batch_norm_eps %||% 1e-4 + pipe$decoder <- load_flux2_vae_decoder(vae_path, + latent_channels = latent_channels, + verbose = verbose) + if (encoder) { + pipe$encoder <- load_flux2_vae_encoder(vae_path, + latent_channels = latent_channels, + verbose = verbose) + } + for (nm in intersect(c("decoder", "encoder"), names(pipe))) { + pipe[[nm]]$to(device = component_device) + # The checkpoint VAE is float32; cast once so the per-phase + # onload is device-only + if (compute_dtype == "bfloat16") { + pipe[[nm]]$to(dtype = torch::torch_bfloat16()) + } + } + + if (verbose) { + message("Loading Qwen3-VL text encoder...") + } + te_dir <- dirname(.flux3_cached("text_encoder/config.json", revision)) + for (f in c("text_encoder/model.safetensors.index.json", + sprintf("text_encoder/model-%05d-of-00002.safetensors", 1:2))) { + .flux3_cached(f, revision) + } + pipe$text_encoder <- load_qwen3_text_encoder( + te_dir, device = if (phase_offload) "cpu" else text_device, + dtype = if (text_device == "cpu") "float32" else "bfloat16", + verbose = verbose + ) + pipe$tokenizer <- qwen_bpe_tokenizer( + .flux3_cached("text_encoder/tokenizer.json", revision)) + + components <- intersect(c("transformer", "decoder", "encoder"), names(pipe)) + if (!identical(text_device, "cpu")) { + components <- c(components, "text_encoder") + } + pipe$staging <- .flux_build_staging(pipe, pin, phase_offload, device, + components, verbose = verbose) + + structure(pipe, class = "flux3_pipeline") +} + +# Flow-map Euler loop: one CFG-free forward per step, given the current +# flow time and the time it jumps to (the next scheduler timestep, 0 +# after the last). Reference-image tokens (img2img) are a separate +# transformer input; only the image tokens come back. +.flux3_denoise <- function(transformer, latents, schedule, prompt_embeds, + txt_ids, img_ids, compute_dtype, + cond_tokens = NULL, cond_ids = NULL, + chunk_size = NULL, verbose = TRUE) { + timesteps <- as.numeric(schedule$timesteps$cpu()) + n <- length(timesteps) + pb <- .denoise_progress(n, NULL, verbose) + f32 <- torch::torch_float32() + device <- latents$device + as_t <- function(v) { + torch::torch_tensor(v, dtype = f32, device = device)$div(1000)$to(dtype = compute_dtype)$reshape(1L) + } + + torch::with_no_grad({ + for (i in seq_len(n)) { + t <- timesteps[[i]] + next_t <- if (i < n) timesteps[[i + 1L]] else 0 + noise_pred <- transformer(hidden_states = latents, + encoder_hidden_states = prompt_embeds, + timestep = as_t(t), img_ids = img_ids, + txt_ids = txt_ids, + image_cond_hidden_states = cond_tokens, + image_cond_ids = cond_ids, + target_timestep = as_t(next_t), + chunk_size = chunk_size) + # Euler step in float32, result back in the compute dtype + # (diffusers FlowMatchEulerDiscreteScheduler.step) + step <- flowmatch_scheduler_step( + noise_pred$to(dtype = f32), t, latents$to(dtype = f32), + schedule + ) + latents <- step$prev_sample$to(dtype = compute_dtype) + schedule <- step$schedule + rm(noise_pred, step) + pb$tick(i) + } + }) + pb$done() + latents +} + +#' Generate an image with FLUX 3 Image Klein +#' +#' Step-distilled text-to-image (16 steps, no guidance): Qwen3-VL prompt +#' encoding, flow-map FlowMatch denoising with the static shift 6.93, and +#' 32-channel VAE decode through the BatchNorm latent statistics. +#' +#' FLUX 3 takes full scene captions: a short prompt gives generic +#' output. Expand short requests into one detailed caption (framing, +#' subjects with colors, materials and positions, background, lighting) +#' with an LLM first. Layout boxes go in the caption as BFL documents +#' (a JSON list of elements with \code{bbox = [top, left, bottom, +#' right]} on a 0 to 1000 grid). +#' +#' @param prompt Character. The (full) caption. +#' @param pipeline A \code{flux3_pipeline} from +#' \code{\link{flux3_load_pipeline}}; NULL loads one (passing +#' \code{...} through). +#' @param width,height Integers, divisible by 16. +#' @param num_inference_steps Integer. Denoising steps (16). +#' @param max_sequence_length Integer. Qwen3-VL token cap (8192). +#' @param seed Integer or NULL. Latents are drawn on the CPU in float32 +#' in the packed shape, so a seed matches a diffusers run with a CPU +#' generator. +#' @param prompt_embeds Optional precomputed [1, S, 20480] embeddings. +#' @param save_file Logical. Write a PNG. +#' @param filename Output path (default derived from the prompt). +#' @param verbose Logical, or one of "silent", "progress", "steps". +#' @param ... Passed to \code{\link{flux3_load_pipeline}} when +#' \code{pipeline} is NULL. +#' +#' @return Invisibly, \code{list(image, metadata)} where \code{image} is +#' an [H, W, 3] array in [0, 1]. +#' +#' @export +txt2img_flux3 <- function(prompt, pipeline = NULL, width = 1024L, + height = 1024L, num_inference_steps = 16L, + max_sequence_length = 8192L, seed = NULL, + prompt_embeds = NULL, save_file = TRUE, + filename = NULL, verbose = TRUE, ...) { + level <- .verbosity(verbose) + if (level != "silent") { + message(sprintf("FLUX 3 klein: %dx%d, %d steps", as.integer(width), + as.integer(height), as.integer(num_inference_steps))) + } + if (is.null(pipeline)) { + pipeline <- flux3_load_pipeline(..., verbose = level == "steps") + } + .flux3_generate(prompt, pipeline, width, height, num_inference_steps, + max_sequence_length, seed, prompt_embeds, save_file, + filename, level) +} + +# The shared FLUX 3 generation body. `reference` is NULL (txt2img) or the +# list from flux3_preprocess_image (img2img); `crop` is the box to crop +# the decoded image back to, or NULL. +.flux3_generate <- function(prompt, pipeline, width, height, + num_inference_steps, max_sequence_length, seed, + prompt_embeds, save_file, filename, level, + reference = NULL, crop = NULL) { + verbose <- level == "steps" + device <- pipeline$device + width <- as.integer(width) + height <- as.integer(height) + if (width %% 16L != 0L || height %% 16L != 0L) { + stop("width and height must be divisible by 16") + } + f32 <- torch::torch_float32() + compute_dtype <- .flux_dtype(pipeline$compute_dtype %||% + if (device == "cpu") "float32" else "bfloat16") + + phase_offload <- isTRUE(pipeline$phase_offload) && device != "cpu" + staging <- pipeline$staging %||% list() + onload <- function(what) { + module <- pipeline[[what]] + if (phase_offload) { + st <- staging[[what]] + if (is.null(st)) { + module$to(device = device) + } else { + .staged_onload(st, device) + } + } + module + } + offload <- function(what) { + module <- pipeline[[what]] + if (phase_offload) { + st <- staging[[what]] + if (is.null(st)) { + module$to(device = "cpu") + } else { + .staged_offload(st) + } + clear_vram() + } + invisible(module) + } + + t0 <- Sys.time() + + # --- Phase 1: text encoding -------------------------------------------------- + if (is.null(prompt_embeds)) { + if (verbose) { + message("Encoding prompt (Qwen3-VL)...") + } + gpu_encode <- !identical(pipeline$text_device, "cpu") + if (gpu_encode) { + onload("text_encoder") + } + te_device <- pipeline$text_encoder$model$embed_tokens$weight$device + enc <- flux3_encode_prompt(prompt, pipeline$text_encoder, + pipeline$tokenizer, + max_sequence_length = max_sequence_length, + device = te_device) + prompt_embeds <- enc$prompt_embeds + rm(enc) + if (gpu_encode) { + offload("text_encoder") + } + } + prompt_embeds <- prompt_embeds$to(device = device, dtype = compute_dtype) + txt_ids <- flux3_prepare_text_ids(prompt_embeds$shape[2], device = device) + + # --- Phase 2: latents, ids, schedule ------------------------------------------ + h2 <- height %/% 16L + w2 <- width %/% 16L + if (!is.null(seed)) { + torch::torch_manual_seed(seed) + } + # Drawn in float32 (the reference's RNG stream), then held in the + # compute dtype as the reference does + latents <- torch::torch_randn(c(1L, 128L, h2, w2), dtype = f32) + latents <- flux2_pack_latents(latents)$to(device = device, + dtype = compute_dtype) + img_ids <- flux3_prepare_image_ids(h2, w2, device = device) + + # --- Reference image (img2img): encode once, reuse every step ------------------- + cond_tokens <- NULL + cond_ids <- NULL + if (!is.null(reference)) { + if (is.null(pipeline$encoder)) { + stop("this pipeline has no VAE encoder; load it with ", + "flux3_load_pipeline(encoder = TRUE)", call. = FALSE) + } + if (verbose) { + message("Encoding reference image...") + } + encoder <- onload("encoder") + ref <- flux2_encode_image(encoder, reference, + bn_eps = pipeline$vae_bn_eps %||% 1e-4) + offload("encoder") + cond_tokens <- ref$tokens$to(device = device, dtype = compute_dtype) + cond_ids <- flux3_prepare_image_ids(ref$height, ref$width, + l_coord = -1, device = device) + rm(ref) + } + + n_steps <- as.integer(num_inference_steps) + sched <- flux3_schedule(n_steps) + + # --- Phase 3: denoise ------------------------------------------------------------ + transformer <- onload("transformer") + fp8_manual <- isTRUE(pipeline$fp8_resident) && + is.null(staging[["transformer"]]) + if (fp8_manual) { + .flux_fp8_to_device(transformer, device) + } + if (verbose) { + message(sprintf("Denoising: %d steps at %dx%d...", n_steps, width, + height)) + } + latents <- .flux3_denoise(transformer, latents, sched, prompt_embeds, + txt_ids, img_ids, compute_dtype, + cond_tokens = cond_tokens, cond_ids = cond_ids, + chunk_size = pipeline$attn_chunk, + verbose = level) + rm(cond_tokens) + if (fp8_manual && phase_offload) { + .flux_fp8_to_device(pipeline$transformer, "cpu") + } + offload("transformer") + ltx23_release_dequant_buffers() + + # --- Phase 4: decode --------------------------------------------------------------- + if (verbose) { + message("Decoding...") + } + # [1, h*w, 128] -> [1, 128, h, w] + latents <- latents$reshape(c(1L, h2, w2, 128L))$permute(c(1L, 4L, 2L, 3L)) + latents <- flux2_bn_normalize(latents, pipeline$decoder$bn$running_mean, + pipeline$decoder$bn$running_var, + eps = pipeline$vae_bn_eps %||% 1e-4, + inverse = TRUE) + latents <- flux2_unpatchify_latents(latents) + + decoder <- onload("decoder") + torch::with_no_grad({ + dec_param <- decoder$post_quant_conv$weight + img <- decoder(latents$to(device = dec_param$device, + dtype = dec_param$dtype)) + img <- img$to(dtype = f32)$cpu() + }) + offload("decoder") + + img <- img$squeeze(1)$permute(c(2L, 3L, 1L)) + img <- img$add(1)$div(2)$clamp(0, 1) + img_array <- as.array(img) + if (!is.null(crop)) { + img_array <- img_array[(crop[["top"]] + 1):crop[["bottom"]], + (crop[["left"]] + 1):crop[["right"]],, + drop = FALSE] + } + + gen_seconds <- as.numeric(difftime(Sys.time(), t0, units = "secs")) + if (verbose) { + message(sprintf("Generated in %.1f s", gen_seconds)) + } + if (save_file) { + if (is.null(filename)) { + filename <- filename_from_prompt(prompt) + } + save_image(img_array, filename) + if (verbose) { + message("Saved to ", filename) + } + } + + metadata <- list( + prompt = prompt, width = dim(img_array)[2], + height = dim(img_array)[1], steps = n_steps, seed = seed, + model = "flux3-image-klein", precision = pipeline$format, + seconds = gen_seconds + ) + invisible(list(image = img_array, metadata = metadata)) +} diff --git a/R/vae_decoder.R b/R/vae_decoder.R index 17731c8..2cea400 100644 --- a/R/vae_decoder.R +++ b/R/vae_decoder.R @@ -79,11 +79,10 @@ VAEAttentionBlock <- torch::nn_module( k <- self$to_k(x) v <- self$to_v(x) - # Scaled dot-product attention - scale <- 1.0 / sqrt(channels) - attn <- torch::torch_bmm(q, k$transpose(2, 3)) * scale - attn <- torch::nnf_softmax(attn, dim = -1) - out <- torch::torch_bmm(attn, v) + # Single-head attention; fused SDPA avoids the (h*w)^2 score matrix + # (8 GiB in bf16 for a 2048x2048 image) + out <- .ltx23_sdpa(q$unsqueeze(2), k$unsqueeze(2), v$unsqueeze(2)) + out <- out$squeeze(2) # Project out out <- self$to_out[[1]](out) diff --git a/inst/COPYRIGHTS b/inst/COPYRIGHTS index 06608d1..1b24f6f 100644 --- a/inst/COPYRIGHTS +++ b/inst/COPYRIGHTS @@ -35,6 +35,13 @@ package are R ports of, or derived from, third-party work: licenses (CreativeML Open RAIL++-M and similar) and are not part of, nor redistributed with, this package. + FLUX 3 Image Klein weights (Black Forest Labs Inc.) are downloaded by + the user under the FLUX Kommunity License after accepting it on + HuggingFace, and are not part of, nor redistributed with, this + package. The FLUX 3 code in diffuseR is ported from diffusers + (Apache License 2.0); no code from Black Forest Labs' repositories is + included. + 5. The R Foundation https://www.r-project.org/logo/ The diffuseR logo (man/figures/logo.png and its dark-mode variant diff --git a/inst/REFERENCES.md b/inst/REFERENCES.md index b316ee0..7133fd7 100644 --- a/inst/REFERENCES.md +++ b/inst/REFERENCES.md @@ -48,6 +48,18 @@ module; this file documents the actual lineage, idea by idea. | Qwen2 byte-level BPE tokenization (GPT-2 byte table, GPT-4-style split regex, rank-based merges) | GPT-2 BPE (Radford et al. 2019; openai/gpt-2 `encoder.py`, MIT); format facts from HuggingFace tokenizers documentation | | Weights | black-forest-labs/FLUX.2-klein-4B (Apache-2.0, ungated; downloaded by the user, never redistributed) | +## FLUX 3 Image Klein + +| What | Source | +|---|---| +| Transformer: per-mode blocks (image, reference image, text each self-attend within the mode), joint blocks over [text; content], parallel attention + SwiGLU from one modulated input, per-token timestep modulation per mode, QK RMS norm and RoPE in float32, flow-map target-timestep embedder, AdaLayerNorm output on image tokens | Ported from HuggingFace **diffusers** (Apache-2.0): `models/transformers/transformer_flux3.py` | +| Pipeline flow: position ids (image H/W scaled by 0.5, references at L = -1 - i, text along L), static-shift (6.93) linspace sigma schedule, flow-map step (timestep and the next one, 0 last), latents held in the compute dtype with float32 Euler steps, packed-space noise, BN de-normalize + unpatchify decode | diffusers `modular_pipelines/flux3/` (`before_denoise.py`, `denoise.py`, `decoders.py`) (Apache-2.0) | +| Prompt encoding contract: chat template with generation prompt (no thinking block), right-padding to a multiple of 80 (cap 8192), padding at position id 1, hidden states of layers 4, 8, ..., 32 stacked per token | diffusers `modular_pipelines/flux3/encoders.py` (Apache-2.0) | +| Qwen3-VL text model (same layers as Qwen3; interleaved MRoPE reduces to plain RoPE for text-only 1D positions; nested `text_config`, `model.language_model.*` keys) | HuggingFace **transformers** (Apache-2.0): `models/qwen3_vl/modeling_qwen3_vl.py`, loaded with the Qwen3 port above | +| Reference preprocessing ("fill"): canvas = size rounded up to the 16 px grid and clamped to [64, 4096], downscale until the canvas fits 4 MP (0.5% back-off), numpy reflect / symmetric padding, crop box for the output | diffusers `modular_pipelines/flux3/image_processor.py` (Apache-2.0), following BFL's `reference_canvas` (described there); numpy `pad` mode semantics | +| VAE (`AutoencoderKLFlux2`) encode and decode | The FLUX.2 klein ports above, unchanged | +| Weights | black-forest-labs/flux-3-image-klein (FLUX Kommunity License, gated; downloaded by the user, never redistributed) | + ## Z-Image Turbo | What | Source | diff --git a/inst/tinytest/fixtures/flux3_model.safetensors b/inst/tinytest/fixtures/flux3_model.safetensors new file mode 100644 index 0000000..c746f55 Binary files /dev/null and b/inst/tinytest/fixtures/flux3_model.safetensors differ diff --git a/inst/tinytest/fixtures/flux3_preprocess.safetensors b/inst/tinytest/fixtures/flux3_preprocess.safetensors new file mode 100644 index 0000000..57c6be0 Binary files /dev/null and b/inst/tinytest/fixtures/flux3_preprocess.safetensors differ diff --git a/inst/tinytest/fixtures/flux3_qwen.safetensors b/inst/tinytest/fixtures/flux3_qwen.safetensors new file mode 100644 index 0000000..0799e30 Binary files /dev/null and b/inst/tinytest/fixtures/flux3_qwen.safetensors differ diff --git a/inst/tinytest/fixtures/flux3_qwen_tiny_ckpt/config.json b/inst/tinytest/fixtures/flux3_qwen_tiny_ckpt/config.json new file mode 100644 index 0000000..67c1581 --- /dev/null +++ b/inst/tinytest/fixtures/flux3_qwen_tiny_ckpt/config.json @@ -0,0 +1,64 @@ +{ + "architectures": [ + "Qwen3VLForConditionalGeneration" + ], + "dtype": "float32", + "image_token_id": 97, + "model_type": "qwen3_vl", + "text_config": { + "attention_bias": false, + "attention_dropout": 0.0, + "head_dim": 8, + "hidden_act": "silu", + "hidden_size": 32, + "initializer_range": 0.02, + "intermediate_size": 64, + "max_position_embeddings": 512, + "model_type": "qwen3_vl_text", + "num_attention_heads": 4, + "num_hidden_layers": 4, + "num_key_value_heads": 2, + "pad_token_id": null, + "rms_norm_eps": 1e-06, + "rope_parameters": { + "mrope_interleaved": true, + "mrope_section": [ + 2, + 1, + 1 + ], + "rope_theta": 5000000.0, + "rope_type": "default" + }, + "tie_word_embeddings": true, + "use_cache": true, + "vocab_size": 100 + }, + "tie_word_embeddings": true, + "transformers_version": "5.21.0", + "video_token_id": 98, + "vision_config": { + "deepstack_visual_indexes": [ + 0 + ], + "depth": 1, + "hidden_act": "gelu_pytorch_tanh", + "hidden_size": 16, + "in_channels": 3, + "initializer_range": 0.02, + "intermediate_size": 32, + "model_type": "qwen3_vl_vision", + "num_heads": 2, + "num_position_embeddings": 16, + "out_hidden_size": 32, + "patch_size": 4, + "rope_parameters": { + "rope_theta": 10000.0, + "rope_type": "axial" + }, + "spatial_merge_size": 2, + "temporal_patch_size": 2 + }, + "vision_end_token_id": 96, + "vision_start_token_id": 95 +} diff --git a/inst/tinytest/fixtures/flux3_qwen_tiny_ckpt/generation_config.json b/inst/tinytest/fixtures/flux3_qwen_tiny_ckpt/generation_config.json new file mode 100644 index 0000000..5f8c8d6 --- /dev/null +++ b/inst/tinytest/fixtures/flux3_qwen_tiny_ckpt/generation_config.json @@ -0,0 +1,7 @@ +{ + "_from_model_config": true, + "output_attentions": false, + "output_hidden_states": false, + "transformers_version": "5.21.0", + "use_cache": true +} diff --git a/inst/tinytest/fixtures/flux3_qwen_tiny_ckpt/model.safetensors b/inst/tinytest/fixtures/flux3_qwen_tiny_ckpt/model.safetensors new file mode 100644 index 0000000..4f618ec Binary files /dev/null and b/inst/tinytest/fixtures/flux3_qwen_tiny_ckpt/model.safetensors differ diff --git a/inst/tinytest/fixtures/flux3_tiny_ckpt/config.json b/inst/tinytest/fixtures/flux3_tiny_ckpt/config.json new file mode 100644 index 0000000..7be3c4e --- /dev/null +++ b/inst/tinytest/fixtures/flux3_tiny_ckpt/config.json @@ -0,0 +1,22 @@ +{ + "_class_name": "Flux3Transformer2DModel", + "_diffusers_version": "0.42.0.dev0", + "attention_head_dim": 16, + "axes_dims_rope": [ + 4, + 4, + 4, + 4 + ], + "cond_in_channels": 8, + "eps": 1e-06, + "flow_map": true, + "in_channels": 8, + "joint_attention_dim": 32, + "mlp_ratio": 3.0, + "num_attention_heads": 2, + "num_layers": 1, + "num_single_layers": 2, + "rope_theta": 10000, + "timestep_channels": 256 +} diff --git a/inst/tinytest/fixtures/flux3_tiny_ckpt/diffusion_pytorch_model-00001-of-00002.safetensors b/inst/tinytest/fixtures/flux3_tiny_ckpt/diffusion_pytorch_model-00001-of-00002.safetensors new file mode 100644 index 0000000..2900c17 Binary files /dev/null and b/inst/tinytest/fixtures/flux3_tiny_ckpt/diffusion_pytorch_model-00001-of-00002.safetensors differ diff --git a/inst/tinytest/fixtures/flux3_tiny_ckpt/diffusion_pytorch_model-00002-of-00002.safetensors b/inst/tinytest/fixtures/flux3_tiny_ckpt/diffusion_pytorch_model-00002-of-00002.safetensors new file mode 100644 index 0000000..520e49c Binary files /dev/null and b/inst/tinytest/fixtures/flux3_tiny_ckpt/diffusion_pytorch_model-00002-of-00002.safetensors differ diff --git a/inst/tinytest/fixtures/flux3_tiny_ckpt/diffusion_pytorch_model.safetensors.index.json b/inst/tinytest/fixtures/flux3_tiny_ckpt/diffusion_pytorch_model.safetensors.index.json new file mode 100644 index 0000000..a50c3ee --- /dev/null +++ b/inst/tinytest/fixtures/flux3_tiny_ckpt/diffusion_pytorch_model.safetensors.index.json @@ -0,0 +1,62 @@ +{ + "metadata": { + "total_size": 429696 + }, + "weight_map": { + "context_embedder.weight": "diffusion_pytorch_model-00001-of-00002.safetensors", + "early_stream_modulations.image.linear.weight": "diffusion_pytorch_model-00001-of-00002.safetensors", + "early_stream_modulations.image_cond.linear.weight": "diffusion_pytorch_model-00001-of-00002.safetensors", + "early_stream_modulations.txt.linear.weight": "diffusion_pytorch_model-00001-of-00002.safetensors", + "mode_blocks.image.0.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors", + "mode_blocks.image.0.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors", + "mode_blocks.image.0.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors", + "mode_blocks.image.0.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors", + "mode_blocks.image.0.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors", + "mode_blocks.image.0.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors", + "mode_blocks.image.0.ff.linear_in.weight": "diffusion_pytorch_model-00001-of-00002.safetensors", + "mode_blocks.image.0.ff.linear_out.weight": "diffusion_pytorch_model-00001-of-00002.safetensors", + "mode_blocks.image_cond.0.attn.norm_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors", + "mode_blocks.image_cond.0.attn.norm_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors", + "mode_blocks.image_cond.0.attn.to_k.weight": "diffusion_pytorch_model-00001-of-00002.safetensors", + "mode_blocks.image_cond.0.attn.to_out.0.weight": "diffusion_pytorch_model-00001-of-00002.safetensors", + "mode_blocks.image_cond.0.attn.to_q.weight": "diffusion_pytorch_model-00001-of-00002.safetensors", + "mode_blocks.image_cond.0.attn.to_v.weight": "diffusion_pytorch_model-00001-of-00002.safetensors", + "mode_blocks.image_cond.0.ff.linear_in.weight": "diffusion_pytorch_model-00001-of-00002.safetensors", + "mode_blocks.image_cond.0.ff.linear_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors", + "norm_out.linear.weight": "diffusion_pytorch_model-00002-of-00002.safetensors", + "proj_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors", + "single_stream_modulations.image.linear.weight": "diffusion_pytorch_model-00001-of-00002.safetensors", + "single_stream_modulations.image_cond.linear.weight": "diffusion_pytorch_model-00001-of-00002.safetensors", + "single_stream_modulations.txt.linear.weight": "diffusion_pytorch_model-00001-of-00002.safetensors", + "single_transformer_blocks.0.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors", + "single_transformer_blocks.0.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors", + "single_transformer_blocks.0.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors", + "single_transformer_blocks.0.attn.to_out.0.weight": "diffusion_pytorch_model-00002-of-00002.safetensors", + "single_transformer_blocks.0.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors", + "single_transformer_blocks.0.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors", + "single_transformer_blocks.0.ff.linear_in.weight": "diffusion_pytorch_model-00002-of-00002.safetensors", + "single_transformer_blocks.0.ff.linear_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors", + "single_transformer_blocks.1.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors", + "single_transformer_blocks.1.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors", + "single_transformer_blocks.1.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors", + "single_transformer_blocks.1.attn.to_out.0.weight": "diffusion_pytorch_model-00002-of-00002.safetensors", + "single_transformer_blocks.1.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors", + "single_transformer_blocks.1.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors", + "single_transformer_blocks.1.ff.linear_in.weight": "diffusion_pytorch_model-00002-of-00002.safetensors", + "single_transformer_blocks.1.ff.linear_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors", + "target_timestep_embedder.linear_1.weight": "diffusion_pytorch_model-00001-of-00002.safetensors", + "target_timestep_embedder.linear_2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors", + "timestep_embedder.linear_1.weight": "diffusion_pytorch_model-00001-of-00002.safetensors", + "timestep_embedder.linear_2.weight": "diffusion_pytorch_model-00001-of-00002.safetensors", + "txt_mode_blocks.0.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors", + "txt_mode_blocks.0.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors", + "txt_mode_blocks.0.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00002.safetensors", + "txt_mode_blocks.0.attn.to_out.0.weight": "diffusion_pytorch_model-00002-of-00002.safetensors", + "txt_mode_blocks.0.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00002.safetensors", + "txt_mode_blocks.0.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00002.safetensors", + "txt_mode_blocks.0.ff.linear_in.weight": "diffusion_pytorch_model-00002-of-00002.safetensors", + "txt_mode_blocks.0.ff.linear_out.weight": "diffusion_pytorch_model-00002-of-00002.safetensors", + "x_embedders.image.weight": "diffusion_pytorch_model-00001-of-00002.safetensors", + "x_embedders.image_cond.weight": "diffusion_pytorch_model-00001-of-00002.safetensors" + } +} diff --git a/inst/tinytest/test_flux3_preprocess.R b/inst/tinytest/test_flux3_preprocess.R new file mode 100644 index 0000000..dc3be6a --- /dev/null +++ b/inst/tinytest/test_flux3_preprocess.R @@ -0,0 +1,53 @@ +# FLUX 3 reference-image preprocessing against diffusers +# Flux3ImageProcessor (fixture generated by +# tools/gen_fixtures_flux3_preprocess.py, checked in). + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} +if (!requireNamespace("safetensors", quietly = TRUE)) { + exit_file("safetensors not installed") +} + +library(diffuseR) + +fixture_path <- system.file("tinytest", "fixtures", + "flux3_preprocess.safetensors", package = "diffuseR") +if (fixture_path == "") fixture_path <- "fixtures/flux3_preprocess.safetensors" +if (!file.exists(fixture_path)) exit_file("flux3 preprocess fixtures missing") + +fx <- safetensors::safe_load_file(fixture_path, framework = "torch") + +max_abs_diff <- function(a, b) { + as.numeric(torch::torch_max(torch::torch_abs( + a$to(dtype = torch::torch_float32()) - b$to(dtype = torch::torch_float32()) + ))) +} + +# a: 100x70 -> 112x80 canvas (reflect); b: 20x12 -> 64x64 canvas +# (symmetric, side <= 16, pad wider than the image); c: already on grid +for (name in c("a", "b", "c")) { + arr <- as.array(fx[[paste0(name, "_input")]]$to(dtype = torch::torch_float64())) / 255 + res <- flux3_preprocess_image(arr) + ref <- fx[[paste0(name, "_output")]] + expect_equal(as.integer(res$image$shape), as.integer(ref$shape), + info = name) + expect_true(max_abs_diff(res$image, ref) < 1e-6, info = name) + expect_equal(unname(res$crop), + as.numeric(as.array(fx[[paste0(name, "_crop")]])), info = name) +} + +# Downscale path: the canvas fits the cap (here 64x64 = 4096 px area) +arr <- array(runif(300 * 200 * 3), dim = c(200L, 300L, 3L)) +res <- flux3_preprocess_image(arr, max_area = 96 * 64) +expect_true(prod(res$image$shape[3:4]) <= 96 * 64) +expect_equal(res$image$shape[3] %% 16, 0) +expect_equal(res$image$shape[4] %% 16, 0) +crop <- res$crop +expect_true(crop[["right"]] - crop[["left"]] <= res$image$shape[4]) + +# numpy.pad index maps +expect_equal(diffuseR:::.flux3_pad_index(4, 2, 2, "reflect"), + c(3L, 2L, 1L, 2L, 3L, 4L, 3L, 2L)) +expect_equal(diffuseR:::.flux3_pad_index(3, 4, 0, "symmetric"), + c(3L, 3L, 2L, 1L, 1L, 2L, 3L)) diff --git a/inst/tinytest/test_flux3_qwen.R b/inst/tinytest/test_flux3_qwen.R new file mode 100644 index 0000000..d2f269f --- /dev/null +++ b/inst/tinytest/test_flux3_qwen.R @@ -0,0 +1,65 @@ +# Parity test for the FLUX 3 text-encoder path against a tiny random-init +# transformers Qwen3VLForConditionalGeneration (fixture generated by +# tools/gen_fixtures_flux3_qwen.py, checked in). + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} +if (!requireNamespace("safetensors", quietly = TRUE)) { + exit_file("safetensors not installed") +} + +library(diffuseR) + +fixture_path <- system.file("tinytest", "fixtures", "flux3_qwen.safetensors", + package = "diffuseR") +if (fixture_path == "") fixture_path <- "fixtures/flux3_qwen.safetensors" +ckpt_dir <- system.file("tinytest", "fixtures", "flux3_qwen_tiny_ckpt", + package = "diffuseR") +if (ckpt_dir == "") ckpt_dir <- "fixtures/flux3_qwen_tiny_ckpt" +if (!file.exists(fixture_path) || !dir.exists(ckpt_dir)) { + exit_file("flux3 qwen fixtures missing") +} + +fx <- safetensors::safe_load_file(fixture_path, framework = "torch") + +max_abs_diff <- function(a, b) { + as.numeric(torch::torch_max(torch::torch_abs( + a$to(dtype = torch::torch_float32()) - b$to(dtype = torch::torch_float32()) + ))) +} + +# The Qwen3-VL layout loads (text model only, vision tower skipped) with +# the nested text config, including rope_theta under rope_parameters +model <- load_qwen3_text_encoder(ckpt_dir, device = "cpu", + dtype = "float32", verbose = FALSE) +expect_equal(model$rope_theta, 5e6) +expect_equal(length(model$model$layers), 4L) + +ids <- fx$input_ids$to(dtype = torch::torch_long()) + 1L +mask <- fx$attention_mask$to(dtype = torch::torch_long()) +positions <- as.integer(as.array(fx$position_ids)) +expect_equal(positions, c(0:8, 1L, 1L, 1L)) + +states <- torch::with_no_grad(model(ids, attention_mask = mask, + out_layers = c(1L, 2L, 4L), position_ids = positions)) +stacked <- torch::torch_cat(states, dim = -1L) +expect_equal(as.integer(stacked$shape), as.integer(fx$stacked$shape)) +# Padding states included: the transformer attends to them +expect_true(max_abs_diff(stacked, fx$stacked) < 1e-5) + +# Default positions (0..S-1) change the padding states only +states_default <- torch::with_no_grad(model(ids, attention_mask = mask, + out_layers = c(1L, 2L, 4L))) +default <- torch::torch_cat(states_default, dim = -1L) +expect_true(max_abs_diff(default[, 1:9, ], fx$stacked[, 1:9, ]) < 1e-5) +expect_true(max_abs_diff(default[, 10:12, ], fx$stacked[, 10:12, ]) > 1e-4) + +# Window: rounded up to the pad multiple, capped, pads at position 1 +win <- diffuseR:::.flux3_text_window(1:9, 0L, 8192L, 4L) +expect_equal(length(win$ids), 12L) +expect_equal(win$mask, c(rep(1L, 9), 0L, 0L, 0L)) +expect_equal(win$positions, c(0:8, 1L, 1L, 1L)) +win <- diffuseR:::.flux3_text_window(1:100, 0L, 90L, 80L) +expect_equal(length(win$ids), 90L) +expect_equal(sum(win$mask), 90L) diff --git a/inst/tinytest/test_flux3_transformer.R b/inst/tinytest/test_flux3_transformer.R new file mode 100644 index 0000000..ca84a7e --- /dev/null +++ b/inst/tinytest/test_flux3_transformer.R @@ -0,0 +1,96 @@ +# Parity test for the FLUX 3 transformer against a tiny random-init +# diffusers Flux3Transformer2DModel (fixture generated by +# tools/gen_fixtures_flux3_model.py, checked in). + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} +if (!requireNamespace("safetensors", quietly = TRUE)) { + exit_file("safetensors not installed") +} + +library(diffuseR) + +fixture_path <- system.file("tinytest", "fixtures", "flux3_model.safetensors", + package = "diffuseR") +if (fixture_path == "") fixture_path <- "fixtures/flux3_model.safetensors" +if (!file.exists(fixture_path)) exit_file("flux3 model fixtures missing") + +fx <- safetensors::safe_load_file(fixture_path, framework = "torch") + +max_abs_diff <- function(a, b) { + as.numeric(torch::torch_max(torch::torch_abs( + a$to(dtype = torch::torch_float32()) - b$to(dtype = torch::torch_float32()) + ))) +} + +model <- flux3_transformer( + in_channels = 8L, + cond_in_channels = 8L, + num_layers = 1L, + num_single_layers = 2L, + attention_head_dim = 16L, + num_attention_heads = 2L, + joint_attention_dim = 32L, + axes_dims_rope = c(4L, 4L, 4L, 4L), + flow_map = TRUE +) +model$eval() + +# Strict both ways: the R module tree matches the diffusers state dict +weights <- fx[grep("^model\\.", names(fx))] +names(weights) <- sub("^model\\.", "", names(weights)) +dests <- c(model$named_parameters(), model$named_buffers()) +expect_equal(character(0), setdiff(names(weights), names(dests))) +expect_equal(character(0), setdiff(names(dests), names(weights))) +torch::with_no_grad({ + for (name in names(weights)) dests[[name]]$copy_(weights[[name]]) +}) + +# Position-id builders match the reference exactly +expect_true(max_abs_diff(flux3_prepare_text_ids(7L), fx$txt_ids) == 0) +expect_true(max_abs_diff(flux3_prepare_image_ids(4L, 6L), fx$img_ids) == 0) +expect_true(max_abs_diff(flux3_prepare_image_ids(3L, 5L, l_coord = -1), + fx$cond_ids) == 0) + +# Text-to-image forward +out <- torch::with_no_grad(model( + hidden_states = fx$hidden, + encoder_hidden_states = fx$encoder, + timestep = fx$timestep, + img_ids = fx$img_ids, + txt_ids = fx$txt_ids, + target_timestep = fx$target_timestep +)) +expect_equal(as.integer(out$shape), as.integer(fx$out_t2i$shape)) +expect_true(max_abs_diff(out, fx$out_t2i) < 1e-5) + +# With one reference image +out <- torch::with_no_grad(model( + hidden_states = fx$hidden, + encoder_hidden_states = fx$encoder, + timestep = fx$timestep, + img_ids = fx$img_ids, + txt_ids = fx$txt_ids, + image_cond_hidden_states = fx$cond, + image_cond_ids = fx$cond_ids, + target_timestep = fx$target_timestep +)) +expect_equal(as.integer(out$shape), as.integer(fx$out_ref$shape)) +expect_true(max_abs_diff(out, fx$out_ref) < 1e-5) + +# The flow-map checkpoint requires the target timestep +expect_error(model( + hidden_states = fx$hidden, + encoder_hidden_states = fx$encoder, + timestep = fx$timestep, + img_ids = fx$img_ids, + txt_ids = fx$txt_ids +), "target_timestep") + +# Distilled schedule: static shift 6.93 over linspace(1, 1/16, 16) +sched <- flux3_schedule(16L) +expect_true(max(abs(as.numeric(sched$timesteps) - + as.numeric(fx$timesteps))) < 1e-3) +expect_true(max(abs(as.numeric(sched$sigmas) - + as.numeric(fx$sigmas))) < 1e-6) diff --git a/inst/tinytest/test_txt2img_flux3.R b/inst/tinytest/test_txt2img_flux3.R new file mode 100644 index 0000000..9af2fa6 --- /dev/null +++ b/inst/tinytest/test_txt2img_flux3.R @@ -0,0 +1,93 @@ +# FLUX 3 pipeline wiring on the CPU with tiny random-init components: +# the flow-map denoise loop (timestep / target timestep pairs, reference +# tokens as their own input) and end-to-end txt2img / img2img runs. + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} + +library(diffuseR) + +# --- Denoise loop with a recording fake transformer ---------------------------------- +seen <- new.env() +seen$t <- numeric(0) +seen$target <- numeric(0) +fake <- torch::nn_module( + forward = function(hidden_states, encoder_hidden_states, timestep, img_ids, + txt_ids, image_cond_hidden_states = NULL, + image_cond_ids = NULL, target_timestep = NULL, + chunk_size = NULL) { + seen$t <- c(seen$t, as.numeric(timestep)) + seen$target <- c(seen$target, as.numeric(target_timestep)) + seen$cond <- if (is.null(image_cond_hidden_states)) 0L else + image_cond_hidden_states$shape[2] + torch::torch_zeros_like(hidden_states) + } +)() +sched <- flux3_schedule(4L) +lat <- torch::torch_randn(1L, 6L, 128L) +out <- diffuseR:::.flux3_denoise(fake, lat, sched, + prompt_embeds = torch::torch_randn(1L, 3L, 24L), + txt_ids = flux3_prepare_text_ids(3L), + img_ids = flux3_prepare_image_ids(2L, 3L), + compute_dtype = torch::torch_float32(), + cond_tokens = torch::torch_randn(1L, 4L, 128L), + cond_ids = flux3_prepare_image_ids(2L, 2L, l_coord = -1), + verbose = FALSE) +ts <- as.numeric(sched$timesteps) / 1000 +expect_equal(seen$t, ts, tolerance = 1e-6) +expect_equal(seen$target, c(ts[-1], 0), tolerance = 1e-6) # jumps to the next, 0 last +expect_equal(seen$cond, 4L) +expect_equal(as.integer(out$shape), c(1L, 6L, 128L)) +expect_true(as.numeric((out - lat)$abs()$max()) < 1e-6) # zero velocity + +# --- End to end with tiny components -------------------------------------------------- +torch::torch_manual_seed(8) +transformer <- flux3_transformer(in_channels = 128L, cond_in_channels = 128L, + num_layers = 1L, num_single_layers = 1L, attention_head_dim = 8L, + num_attention_heads = 2L, joint_attention_dim = 24L, + axes_dims_rope = c(2L, 2L, 2L, 2L), flow_map = TRUE) +transformer$eval() +decoder <- flux2_vae_decoder(latent_channels = 32L, + block_channels = c(32L, 32L, 16L, 8L), norm_groups = 8L) +decoder$eval() +pipeline <- structure(list(transformer = transformer, decoder = decoder, + device = "cpu", text_device = "cpu", phase_offload = FALSE, + fp8_resident = FALSE, format = "full", attn_chunk = NULL, + vae_bn_eps = 1e-4), class = "flux3_pipeline") +emb <- torch::torch_randn(1L, 7L, 24L) + +res <- txt2img_flux3("tiny smoke test", pipeline = pipeline, width = 64L, + height = 48L, num_inference_steps = 2L, seed = 42L, prompt_embeds = emb, + save_file = FALSE, verbose = FALSE) +expect_equal(dim(res$image), c(48L, 64L, 3L)) +expect_true(all(is.finite(res$image))) +expect_true(all(res$image >= 0) && all(res$image <= 1)) +expect_equal(res$metadata$model, "flux3-image-klein") + +# Same seed, same image +res2 <- txt2img_flux3("tiny smoke test", pipeline = pipeline, width = 64L, + height = 48L, num_inference_steps = 2L, seed = 42L, prompt_embeds = emb, + save_file = FALSE, verbose = FALSE) +expect_equal(res$image, res2$image) + +# Without an encoder the edit says how to load one +ref <- array(0.5, c(50L, 70L, 3L)) +expect_error(img2img_flux3("x", ref, pipeline = pipeline, + num_inference_steps = 1L, prompt_embeds = emb, save_file = FALSE, + verbose = FALSE), "encoder = TRUE") + +pipeline$encoder <- flux2_vae_encoder(latent_channels = 32L, + block_channels = c(8L, 16L, 32L, 32L), norm_groups = 8L) +pipeline$encoder$eval() +# 70x50 reference -> 80x64 canvas; the output is cropped back to 70x50 +res <- img2img_flux3("tiny edit", ref, pipeline = pipeline, + num_inference_steps = 2L, seed = 42L, prompt_embeds = emb, + save_file = FALSE, verbose = FALSE) +expect_equal(dim(res$image), c(50L, 70L, 3L)) +expect_true(all(is.finite(res$image))) +# An explicit size is returned whole +res <- img2img_flux3("tiny edit", ref, pipeline = pipeline, width = 64L, + height = 64L, num_inference_steps = 1L, seed = 42L, prompt_embeds = emb, + save_file = FALSE, verbose = FALSE) +expect_equal(dim(res$image), c(64L, 64L, 3L)) diff --git a/man/dit_flux3.Rd b/man/dit_flux3.Rd new file mode 100644 index 0000000..dc27954 --- /dev/null +++ b/man/dit_flux3.Rd @@ -0,0 +1,15 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{dit_flux3} +\alias{dit_flux3} +\title{FLUX 3 Transformer} +\description{ +Fresh R port of Flux3Transformer2DModel from the diffusers reference +implementation (Apache-2.0, +src/diffusers/models/transformers/transformer_flux3.py). Defaults are +the FLUX 3 Image Klein configuration: 5 per-mode blocks for each of +the image, reference-image and text modes, then 28 joint blocks. +Conditioning is a per-token timestep embedding (the image tokens at +the flow timestep, text and reference tokens at 0); the step-distilled +(flow-map) checkpoint adds the embedding of the timestep each step +jumps to. +} diff --git a/man/dit_flux3_modules.Rd b/man/dit_flux3_modules.Rd new file mode 100644 index 0000000..8895ce1 --- /dev/null +++ b/man/dit_flux3_modules.Rd @@ -0,0 +1,17 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{dit_flux3_modules} +\alias{dit_flux3_modules} +\title{FLUX 3 Transformer Building Blocks} +\description{ +Fresh R port of the FLUX 3 blocks from the diffusers reference +implementation (Apache-2.0, +src/diffusers/models/transformers/transformer_flux3.py). Every block +is a parallel attention + SwiGLU feed-forward computed from one +modulated input: per-mode blocks self-attend within one token mode +(image, reference image, or text), joint blocks attend over the +concatenated [text; content] sequence. Modulation is per token, +computed at model level. Module field names mirror the diffusers +state-dict keys 1:1; every linear is bias-free. Reuses +\code{flux2_feed_forward}, \code{flux_apply_rotary_emb} and +\code{.ltx23_sdpa}. +} diff --git a/man/download_flux3.Rd b/man/download_flux3.Rd new file mode 100644 index 0000000..d0386d1 --- /dev/null +++ b/man/download_flux3.Rd @@ -0,0 +1,12 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{download_flux3} +\alias{download_flux3} +\title{Download and Prepare FLUX 3 Image Klein Weights} +\description{ +Downloads FLUX 3 Image Klein from HuggingFace (gated; FLUX Kommunity +License) and quantizes the 5.6B transformer to a local fp8 (~6 GB) or +NF4 artifact. Only the diffusers-format files are fetched: the bf16 +transformer shards, the Qwen3-VL text encoder, the VAE and the +tokenizer. BFL's own root \code{model.safetensors} and the +\code{variants/} quantizations are not used. +} diff --git a/man/download_flux3_klein.Rd b/man/download_flux3_klein.Rd new file mode 100644 index 0000000..2503a64 --- /dev/null +++ b/man/download_flux3_klein.Rd @@ -0,0 +1,46 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{download_flux3_klein} +\alias{download_flux3_klein} +\title{Download FLUX 3 Image Klein and build the quantized artifact} +\usage{ +download_flux3_klein( + quantize = TRUE, + precision = c("auto", "fp8", "nf4"), + output_dir = NULL, + text_encoders = TRUE, + verbose = TRUE +) +} +\arguments{ +\item{quantize}{Logical. Build the quantized artifact.} + +\item{precision}{"auto" (default: fp8 when safetensors supports +float8, else nf4), "fp8" (~6 GB), or "nf4".} + +\item{output_dir}{Directory for the quantized artifact.} + +\item{text_encoders}{Logical. Also fetch the Qwen3-VL text encoder, +tokenizer and VAE (~9.2 GB).} + +\item{verbose}{Logical.} +} +\value{ +Invisibly, a list with \code{transformer_dir}, + \code{artifact_dir}, and \code{support} (named file paths). +} +\description{ +Skips work already done: a valid quantized manifest short-circuits +the transformer download; cached files are not re-fetched. Needs +\code{HF_TOKEN} for an account that accepted the license. The bf16 +transformer source (~11 GB in the HuggingFace cache) may be deleted +after quantization. +} +\details{ +The FLUX Kommunity License permits non-commercial use, and commercial +use of outputs by Qualifying Users (under US$5M gross annualized +revenue) who filter or review outputs before distributing them and +disclose AI generation where the law requires it. Read LICENSE.md in +the downloaded snapshot for the full terms. The quantized artifact is +a Derivative; keep it local. + +} diff --git a/man/flux3_attention.Rd b/man/flux3_attention.Rd new file mode 100644 index 0000000..ceadf27 --- /dev/null +++ b/man/flux3_attention.Rd @@ -0,0 +1,26 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{flux3_attention} +\alias{flux3_attention} +\title{FLUX 3 self-attention} +\usage{ +flux3_attention(query_dim, heads, dim_head, eps = 1e-06) +} +\arguments{ +\item{query_dim}{Integer. Model dimension.} + +\item{heads}{Integer. Attention heads.} + +\item{dim_head}{Integer. Per-head dimension.} + +\item{eps}{Numeric. RMS norm epsilon.} +} +\value{ +Module whose forward(hidden_states, image_rotary_emb, + chunk_size) returns the attended states [B, S, query_dim]. +} +\description{ +Separate q/k/v projections, per-head RMS norms and rotary embeddings +applied in float32 (cast back to the value dtype before attention), +and one output projection. Reference: Flux3Attention + +Flux3AttnProcessor. +} diff --git a/man/flux3_encode_prompt.Rd b/man/flux3_encode_prompt.Rd new file mode 100644 index 0000000..29a6204 --- /dev/null +++ b/man/flux3_encode_prompt.Rd @@ -0,0 +1,55 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{flux3_encode_prompt} +\alias{flux3_encode_prompt} +\title{Encode a prompt for FLUX 3} +\usage{ +flux3_encode_prompt( + prompt, + model, + tokenizer, + max_sequence_length = 8192L, + pad_multiple = 80L, + out_layers = seq(4L, 32L, by = 4L), + device = NULL +) +} +\arguments{ +\item{prompt}{Character scalar.} + +\item{model}{A \code{\link{qwen3_encoder}} loaded from FLUX 3's +\code{text_encoder} with \code{\link{load_qwen3_text_encoder}}.} + +\item{tokenizer}{A \code{\link{qwen_bpe_tokenizer}} from the same +directory.} + +\item{max_sequence_length}{Integer. Token cap (8192).} + +\item{pad_multiple}{Integer. Window rounding (80).} + +\item{out_layers}{Integer vector. Hidden-state layers to stack +(4, 8, ..., 32).} + +\item{device}{Character or NULL. Defaults to the model's device.} +} +\value{ +\code{list(prompt_embeds, txt_ids)}: embeddings [1, L, + length(out_layers) * hidden] and position ids [L, 4] from + \code{\link{flux3_prepare_text_ids}}. +} +\description{ +The Qwen3-VL prompt path of the reference pipeline (diffusers +Flux3TextEncoderStep._get_qwen3_vl_prompt_embeds, Apache-2.0): the +prompt goes through the chat template (one user turn plus the +generation prompt, no thinking block), is right-padded to a multiple +of \code{pad_multiple} tokens (capped at \code{max_sequence_length}), +and the hidden states of \code{out_layers} are concatenated per +token. The padding tokens get position id 1, as in the training +environment, and the transformer attends to them, so their states are +part of the embedding. +} +\details{ +FLUX 3 expects a full scene caption: a short prompt gives generic +output. Expand short requests with an LLM first (BFL publishes the +system prompts with flux-3-image-klein-pup). + +} diff --git a/man/flux3_is_quant_key.Rd b/man/flux3_is_quant_key.Rd new file mode 100644 index 0000000..08f84aa --- /dev/null +++ b/man/flux3_is_quant_key.Rd @@ -0,0 +1,16 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{flux3_is_quant_key} +\alias{flux3_is_quant_key} +\title{Test whether a FLUX 3 key is in the quantization cast set} +\usage{ +flux3_is_quant_key(key) +} +\arguments{ +\item{key}{Character vector of parameter names (diffusers-style).} +} +\value{ +Logical vector. +} +\description{ +Test whether a FLUX 3 key is in the quantization cast set +} diff --git a/man/flux3_load_pipeline.Rd b/man/flux3_load_pipeline.Rd new file mode 100644 index 0000000..57aa0ce --- /dev/null +++ b/man/flux3_load_pipeline.Rd @@ -0,0 +1,65 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{flux3_load_pipeline} +\alias{flux3_load_pipeline} +\title{Load the FLUX 3 Image Klein pipeline} +\usage{ +flux3_load_pipeline( + model_dir = NULL, + device = "cuda", + precision = c("auto", "fp8", "nf4", "bf16"), + text_device = NULL, + attn_chunk = NULL, + phase_offload = TRUE, + pin = NULL, + encoder = TRUE, + revision = .flux3_revision, + dtype = NULL, + verbose = TRUE +) +} +\arguments{ +\item{model_dir}{Quantized artifact directory (default: the +\code{download_flux3_klein} location for \code{precision}), or a raw +diffusers transformer directory.} + +\item{device}{Character. Compute device.} + +\item{precision}{"auto" (default: reuse an existing artifact, else +fp8 when safetensors supports float8, else nf4), "fp8", "nf4", or +"bf16" (the raw diffusers transformer).} + +\item{text_device}{Device for the Qwen3-VL encoder (default: +\code{device}; it encodes in its own phase and offloads). "cpu" +keeps it in float32 on the host.} + +\item{attn_chunk}{Integer or NULL. Attention query-chunk override.} + +\item{phase_offload}{Logical. One GPU tenant per phase.} + +\item{pin}{Logical or NULL. Page-lock the phase-swapped weights (see +\code{\link{staging}}).} + +\item{encoder}{Logical. Also load the VAE encoder, which +\code{\link{img2img_flux3}} needs.} + +\item{revision}{Exact 40-hex commit for the support files (default: +the validated snapshot).} + +\item{dtype}{Compute dtype, "bfloat16" or "float32". NULL (default): +bfloat16 on CUDA, float32 on the CPU. "bfloat16" on the CPU halves +host memory (the raw transformer is ~11 GB instead of ~22 GB).} + +\item{verbose}{Logical.} +} +\value{ +A \code{flux3_pipeline} list. +} +\description{ +Loads the quantized transformer artifact plus the VAE (encoder and +decoder), the Qwen3-VL text encoder (text model only) and the +tokenizer from the HuggingFace cache populated by +\code{\link{download_flux3_klein}}. With phase offload each component +visits the GPU only for its own phase, so the GPU peak is the largest +phase rather than the sum: the ~8 GB Qwen3-VL encode, or the fp8 +transformer (~6 GB) plus activations. +} diff --git a/man/flux3_mode_block.Rd b/man/flux3_mode_block.Rd new file mode 100644 index 0000000..5acff00 --- /dev/null +++ b/man/flux3_mode_block.Rd @@ -0,0 +1,34 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{flux3_mode_block} +\alias{flux3_mode_block} +\title{FLUX 3 per-mode transformer block} +\usage{ +flux3_mode_block( + dim, + num_attention_heads, + attention_head_dim, + mlp_ratio = 3, + eps = 1e-06 +) +} +\arguments{ +\item{dim}{Integer. Model dimension.} + +\item{num_attention_heads}{Integer. Attention heads.} + +\item{attention_head_dim}{Integer. Per-head dimension.} + +\item{mlp_ratio}{Numeric. Feed-forward multiplier (3.0).} + +\item{eps}{Numeric. Norm epsilon.} +} +\value{ +Module whose forward(hidden_states, temb_mod, + image_rotary_emb, chunk_size) returns the updated states. +} +\description{ +Parameterless LayerNorm, per-token (shift, scale) modulation, then +parallel attention and feed-forward from the same normed input, added +back under one gate. Self-attention stays within a single token mode. +Reference: Flux3ModeTransformerBlock. +} diff --git a/man/flux3_prepare_image_ids.Rd b/man/flux3_prepare_image_ids.Rd new file mode 100644 index 0000000..eb50405 --- /dev/null +++ b/man/flux3_prepare_image_ids.Rd @@ -0,0 +1,31 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{flux3_prepare_image_ids} +\alias{flux3_prepare_image_ids} +\title{FLUX 3 image token position ids} +\usage{ +flux3_prepare_image_ids( + height, + width, + l_coord = 0, + rope_scale = 0.5, + device = "cpu" +) +} +\arguments{ +\item{height}{Integer. Token grid rows (pixels / 16).} + +\item{width}{Integer. Token grid columns (pixels / 16).} + +\item{l_coord}{Numeric. L coordinate: 0 for the generated image, +\code{-i} for the i-th reference image.} + +\item{rope_scale}{Numeric. Scale of the H/W axes (0.5).} + +\item{device}{Character. Device.} +} +\value{ +Float tensor [height * width, 4], row-major over (h, w). +} +\description{ +FLUX 3 image token position ids +} diff --git a/man/flux3_prepare_text_ids.Rd b/man/flux3_prepare_text_ids.Rd new file mode 100644 index 0000000..b8daae2 --- /dev/null +++ b/man/flux3_prepare_text_ids.Rd @@ -0,0 +1,18 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{flux3_prepare_text_ids} +\alias{flux3_prepare_text_ids} +\title{FLUX 3 text token position ids} +\usage{ +flux3_prepare_text_ids(len, device = "cpu") +} +\arguments{ +\item{len}{Integer. Text sequence length (padded).} + +\item{device}{Character. Device.} +} +\value{ +Float tensor [len, 4]: (0, 0, 0, 0..len-1). +} +\description{ +FLUX 3 text token position ids +} diff --git a/man/flux3_preprocess_image.Rd b/man/flux3_preprocess_image.Rd new file mode 100644 index 0000000..890a9ba --- /dev/null +++ b/man/flux3_preprocess_image.Rd @@ -0,0 +1,34 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{flux3_preprocess_image} +\alias{flux3_preprocess_image} +\title{Preprocess a reference image for FLUX 3} +\usage{ +flux3_preprocess_image(image, max_area = 4 * 1024 * 1024, max_side = 4096L) +} +\arguments{ +\item{image}{Path to a PNG/JPEG, or an [H, W, C] array in [0, 1].} + +\item{max_area}{Numeric. Canvas area cap (reference: 4 megapixels). +Lower it to bound memory; the reference tokens are +\code{canvas area / 256}.} + +\item{max_side}{Integer. Canvas side cap.} +} +\value{ +\code{list(image, crop)}: a float32 tensor [1, 3, H, W] in + [-1, 1] at the canvas size, and the (left, top, right, bottom) box + of the image inside the canvas (0-based, right/bottom exclusive), + used to crop the generated image back. +} +\description{ +The single-reference ("fill") path of the reference pipeline +(diffusers Flux3ImageProcessor, Apache-2.0, following BFL's +\code{reference_canvas}): the image is kept whole. It is downscaled, +keeping its aspect ratio, until its canvas (its size rounded up to the +16 px token grid, sides clamped to [64, \code{max_side}]) fits +\code{max_area}; backing the scale off by 0.5\% per try when the +rounding overshoots. The resize is PIL's antialiased Lanczos-3 in +floating point. The image is then centered on the canvas and the +border filled with its reflection (numpy \code{"reflect"}, or +\code{"symmetric"} when a side is at most 16 px). +} diff --git a/man/flux3_schedule.Rd b/man/flux3_schedule.Rd new file mode 100644 index 0000000..f084608 --- /dev/null +++ b/man/flux3_schedule.Rd @@ -0,0 +1,22 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{flux3_schedule} +\alias{flux3_schedule} +\title{FLUX 3 distilled sampling schedule} +\usage{ +flux3_schedule(num_inference_steps = 16L, shift = 6.93, device = "cpu") +} +\arguments{ +\item{num_inference_steps}{Integer (16 for the distilled checkpoint).} + +\item{shift}{Numeric. Static timestep shift.} + +\item{device}{Character. Device.} +} +\value{ +A FlowMatch schedule from \code{\link{flowmatch_set_timesteps}}. +} +\description{ +FlowMatch Euler with the static shift of the FLUX 3 recipe (6.93) +over \code{sigmas = linspace(1, 1 / n, n)}. Reference: +Flux3SetTimestepsStep. +} diff --git a/man/flux3_single_block.Rd b/man/flux3_single_block.Rd new file mode 100644 index 0000000..7c7c77a --- /dev/null +++ b/man/flux3_single_block.Rd @@ -0,0 +1,35 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{flux3_single_block} +\alias{flux3_single_block} +\title{FLUX 3 joint (single-stream) block} +\usage{ +flux3_single_block( + dim, + num_attention_heads, + attention_head_dim, + mlp_ratio = 3, + eps = 1e-06 +) +} +\arguments{ +\item{dim}{Integer. Model dimension.} + +\item{num_attention_heads}{Integer. Attention heads.} + +\item{attention_head_dim}{Integer. Per-head dimension.} + +\item{mlp_ratio}{Numeric. Feed-forward multiplier (3.0).} + +\item{eps}{Numeric. Norm epsilon.} +} +\value{ +Module whose forward(hidden_states, encoder_hidden_states, + temb_mod, temb_mod_context, image_rotary_emb, chunk_size) returns + \code{list(encoder_hidden_states, hidden_states)}. +} +\description{ +Text and content tokens are normed and modulated separately, then +concatenated (text first) for one parallel attention + feed-forward; +the output is split back and added under per-stream gates. +Reference: Flux3SingleTransformerBlock. +} diff --git a/man/flux3_transformer.Rd b/man/flux3_transformer.Rd new file mode 100644 index 0000000..7c62796 --- /dev/null +++ b/man/flux3_transformer.Rd @@ -0,0 +1,65 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{flux3_transformer} +\alias{flux3_transformer} +\title{FLUX 3 transformer model} +\usage{ +flux3_transformer( + in_channels = 128L, + cond_in_channels = 128L, + num_layers = 5L, + num_single_layers = 28L, + attention_head_dim = 128L, + num_attention_heads = 24L, + joint_attention_dim = 20480L, + timestep_channels = 256L, + mlp_ratio = 3, + axes_dims_rope = c(32L, 32L, 32L, 32L), + rope_theta = 10000, + eps = 1e-06, + flow_map = FALSE +) +} +\arguments{ +\item{in_channels}{Integer. Packed latent channels (128).} + +\item{cond_in_channels}{Integer. Reference-image token channels (128).} + +\item{num_layers}{Integer. Per-mode blocks per mode (5).} + +\item{num_single_layers}{Integer. Joint blocks (28).} + +\item{attention_head_dim}{Integer. Per-head dimension (128).} + +\item{num_attention_heads}{Integer. Attention heads (24).} + +\item{joint_attention_dim}{Integer. Text embedding dim (20480).} + +\item{timestep_channels}{Integer. Sinusoid width (256).} + +\item{mlp_ratio}{Numeric. Feed-forward multiplier (3.0).} + +\item{axes_dims_rope}{Integer vector. Per-axis rotary dims.} + +\item{rope_theta}{Numeric. Rotary base frequency (10000).} + +\item{eps}{Numeric. Norm epsilon.} + +\item{flow_map}{Logical. Step-distilled checkpoint: adds +\code{target_timestep_embedder} and requires \code{target_timestep}.} +} +\value{ +Module whose forward(hidden_states, encoder_hidden_states, + timestep, img_ids, txt_ids, image_cond_hidden_states, image_cond_ids, + target_timestep, chunk_size) returns the predicted velocity of the + image tokens [B, S_img, in_channels]. Timesteps are flow times in + [0, 1]; they are scaled by 1000 internally. +} +\description{ +The forward mirrors the reference: per-mode blocks for the text, +image and (optional) reference tokens, each with its own rotary +embedding, then joint blocks over [text; image; reference] and an +output projection of the image tokens only. Modulation is computed +once per mode on the [B, 1, dim] embedding and broadcast over the +mode's tokens (the reference expands to [B, S, dim] first; every row +is identical). +} diff --git a/man/img2img_flux3.Rd b/man/img2img_flux3.Rd new file mode 100644 index 0000000..3430fda --- /dev/null +++ b/man/img2img_flux3.Rd @@ -0,0 +1,75 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{img2img_flux3} +\alias{img2img_flux3} +\title{Edit an image with FLUX 3 Image Klein} +\usage{ +img2img_flux3( + prompt, + image, + pipeline = NULL, + width = NULL, + height = NULL, + max_area = 4 * 1024 * 1024, + num_inference_steps = 16L, + max_sequence_length = 8192L, + seed = NULL, + prompt_embeds = NULL, + save_file = TRUE, + filename = NULL, + verbose = TRUE, + ... +) +} +\arguments{ +\item{prompt}{Character. The full caption of the edited image.} + +\item{image}{Path to a PNG/JPEG, or an [H, W, 3] array in [0, 1].} + +\item{pipeline}{A \code{flux3_pipeline} loaded with +\code{encoder = TRUE}; NULL loads one (passing \code{...} through).} + +\item{max_area}{Numeric. Reference canvas area cap (reference: 4 +megapixels). The reference adds \code{area / 256} tokens to every +forward, so lower it on small cards.} + +\item{num_inference_steps}{Integer. Denoising steps (16).} + +\item{max_sequence_length}{Integer. Qwen3-VL token cap (8192).} + +\item{seed}{Integer or NULL. See \code{\link{txt2img_flux3}}.} + +\item{prompt_embeds}{Optional precomputed [1, S, 20480] embeddings.} + +\item{save_file}{Logical. Write a PNG.} + +\item{filename}{Output path (default derived from the prompt).} + +\item{verbose}{Logical, or one of "silent", "progress", "steps".} + +\item{...}{Passed to \code{\link{flux3_load_pipeline}} when +\code{pipeline} is NULL.} + +\item{width,height}{Integers divisible by 16, or NULL (default) for +the reference canvas, cropped back to the reference afterwards.} +} +\value{ +Invisibly, \code{list(image, metadata)} where \code{image} is + an [H, W, 3] array in [0, 1]. +} +\description{ +Image + caption to image with the same weights as +\code{\link{txt2img_flux3}}. The reference is kept whole (see +\code{\link{flux3_preprocess_image}}): scaled so its canvas fits +\code{max_area}, reflect-padded to the 16 px grid, VAE-encoded once, +and passed to every transformer forward as its own token mode +(position ids at L = -1). The output defaults to the canvas size and +is cropped back to the reference's box. Reference: diffusers +Flux3DistilledModularPipeline with \code{image =} (Apache-2.0). +} +\details{ +Write the caption as a full description of the edited result (the +reference with the edit folded in); a one-line instruction reframes +the shot. BFL's prompt upsampler (flux-3-image-klein-pup) writes such +captions from an instruction plus the reference image. + +} diff --git a/man/load_qwen3_text_encoder.Rd b/man/load_qwen3_text_encoder.Rd index 00308f2..a688a49 100644 --- a/man/load_qwen3_text_encoder.Rd +++ b/man/load_qwen3_text_encoder.Rd @@ -13,7 +13,8 @@ load_qwen3_text_encoder( } \arguments{ \item{model_path}{Directory with \code{config.json} and -\code{model*.safetensors} (FLUX.2-klein's \code{text_encoder}).} +\code{model*.safetensors} (FLUX.2-klein's or FLUX 3's +\code{text_encoder}).} \item{device}{Character. Target device.} @@ -29,5 +30,9 @@ The loaded \code{qwen3_encoder} in eval mode. \description{ Streams the (possibly sharded) safetensors weights into \code{\link{qwen3_encoder}}. The LM head is tied to the embeddings -and skipped. +and skipped. A Qwen3-VL checkpoint (FLUX 3's text encoder: config +under \code{text_config}, weights under \code{model.language_model.*}) +loads its text model only; the vision tower is skipped. For text-only +input with 1D positions, Qwen3-VL's interleaved MRoPE reduces to the +plain Qwen3 rotary (all three axes carry the same position). } diff --git a/man/qwen3_encoder.Rd b/man/qwen3_encoder.Rd index 1e3dafc..3c0a5cc 100644 --- a/man/qwen3_encoder.Rd +++ b/man/qwen3_encoder.Rd @@ -26,10 +26,13 @@ qwen3_encoder( } \value{ Module whose forward(input_ids, attention_mask = NULL, - out_layers) returns a list of hidden-state tensors [B, S, hidden], - one per requested layer (a value of k means the state after k - layers, matching HF \code{output.hidden_states[k]}). Runs only to - \code{max(out_layers)}. \code{input_ids} are 1-based. + out_layers, position_ids = NULL) returns a list of hidden-state + tensors [B, S, hidden], one per requested layer (a value of k means + the state after k layers, matching HF \code{output.hidden_states[k]}, + which carries the final norm when k is the last layer). + Runs only to \code{max(out_layers)}. \code{input_ids} are 1-based; + \code{position_ids} (0-based integers, length S, shared by the batch) + override the default 0..S-1 rotary positions. } \description{ Defaults are the Qwen3-4B configuration used by FLUX.2 klein. The diff --git a/man/rope_flux3.Rd b/man/rope_flux3.Rd new file mode 100644 index 0000000..e016af0 --- /dev/null +++ b/man/rope_flux3.Rd @@ -0,0 +1,14 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{rope_flux3} +\alias{rope_flux3} +\title{FLUX 3 position ids} +\description{ +4-axis (T, H, W, L) position ids for the FLUX 3 transformer, from the +diffusers reference (Apache-2.0, +src/diffusers/modular_pipelines/flux3/before_denoise.py). Image tokens +sit at T = 0 with the H/W axes scaled by the image RoPE scale (0.5); +the generated image has L = 0 and the i-th reference L = -i. Text +advances along L. Rotary frequencies come from +\code{\link{flux_pos_embed}} with \code{axes_dim = c(32, 32, 32, 32)}, +\code{theta = 10000}. +} diff --git a/man/txt2img.Rd b/man/txt2img.Rd index 779013d..0abcb5a 100644 --- a/man/txt2img.Rd +++ b/man/txt2img.Rd @@ -5,7 +5,7 @@ \usage{ txt2img( prompt, - model_name = c("sd21", "sdxl", "flux1", "flux2", "zimage"), + model_name = c("sd21", "sdxl", "flux1", "flux2", "zimage", "flux3"), ... ) } diff --git a/man/txt2img_flux3.Rd b/man/txt2img_flux3.Rd new file mode 100644 index 0000000..f7c639e --- /dev/null +++ b/man/txt2img_flux3.Rd @@ -0,0 +1,66 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{txt2img_flux3} +\alias{txt2img_flux3} +\title{Generate an image with FLUX 3 Image Klein} +\usage{ +txt2img_flux3( + prompt, + pipeline = NULL, + width = 1024L, + height = 1024L, + num_inference_steps = 16L, + max_sequence_length = 8192L, + seed = NULL, + prompt_embeds = NULL, + save_file = TRUE, + filename = NULL, + verbose = TRUE, + ... +) +} +\arguments{ +\item{prompt}{Character. The (full) caption.} + +\item{pipeline}{A \code{flux3_pipeline} from +\code{\link{flux3_load_pipeline}}; NULL loads one (passing +\code{...} through).} + +\item{num_inference_steps}{Integer. Denoising steps (16).} + +\item{max_sequence_length}{Integer. Qwen3-VL token cap (8192).} + +\item{seed}{Integer or NULL. Latents are drawn on the CPU in float32 +in the packed shape, so a seed matches a diffusers run with a CPU +generator.} + +\item{prompt_embeds}{Optional precomputed [1, S, 20480] embeddings.} + +\item{save_file}{Logical. Write a PNG.} + +\item{filename}{Output path (default derived from the prompt).} + +\item{verbose}{Logical, or one of "silent", "progress", "steps".} + +\item{...}{Passed to \code{\link{flux3_load_pipeline}} when +\code{pipeline} is NULL.} + +\item{width,height}{Integers, divisible by 16.} +} +\value{ +Invisibly, \code{list(image, metadata)} where \code{image} is + an [H, W, 3] array in [0, 1]. +} +\description{ +Step-distilled text-to-image (16 steps, no guidance): Qwen3-VL prompt +encoding, flow-map FlowMatch denoising with the static shift 6.93, and +32-channel VAE decode through the BatchNorm latent statistics. +} +\details{ +FLUX 3 takes full scene captions: a short prompt gives generic +output. Expand short requests into one detailed caption (framing, +subjects with colors, materials and positions, background, lighting) +with an LLM first. Layout boxes go in the caption as BFL documents +(a JSON list of elements with \code{bbox = [top, left, bottom, +right]} on a 0 to 1000 grid). + +} diff --git a/tools/gen_fixtures_flux3_model.py b/tools/gen_fixtures_flux3_model.py new file mode 100644 index 0000000..f184872 --- /dev/null +++ b/tools/gen_fixtures_flux3_model.py @@ -0,0 +1,112 @@ +# Generate a tiny random-init Flux3Transformer2DModel parity fixture for +# the R port (text-to-image and one-reference forwards, position ids, +# the distilled sigma schedule), plus a save_pretrained checkpoint for +# loader tests. +# +# Uses the diffusers reference (Apache-2.0). Run via +# tools/gen_fixtures.sh; never executed at package test/run time. + +import os +import shutil +import sys + +import numpy as np +import torch +from safetensors.torch import save_file + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "ref", "upstream", "diffusers", "src")) + +from diffusers import FlowMatchEulerDiscreteScheduler, Flux3Transformer2DModel # noqa: E402 +from diffusers.modular_pipelines.flux3.before_denoise import ( # noqa: E402 + Flux3TextRoPEInputsStep, + _prepare_image_ids, +) + +OUT_DIR = os.path.join(os.path.dirname(__file__), "..", "inst", "tinytest", "fixtures") +CKPT_DIR = os.path.join(OUT_DIR, "flux3_tiny_ckpt") +os.makedirs(OUT_DIR, exist_ok=True) + +torch.manual_seed(53) + +model = Flux3Transformer2DModel( + in_channels=8, + cond_in_channels=8, + num_layers=1, + num_single_layers=2, + attention_head_dim=16, + num_attention_heads=2, + joint_attention_dim=32, + timestep_channels=256, + mlp_ratio=3.0, + axes_dims_rope=(4, 4, 4, 4), + rope_theta=10000, + flow_map=True, +) +model.eval() + +B, S_TXT = 2, 7 +GRID_H, GRID_W = 4, 6 +COND_H, COND_W = 3, 5 + +hidden = torch.randn(B, GRID_H * GRID_W, 8) +encoder = torch.randn(B, S_TXT, 32) +cond = torch.randn(B, COND_H * COND_W, 8) +timestep = torch.tensor([1.0, 0.4]) +target_timestep = torch.tensor([0.75, 0.0]) + +txt_ids = Flux3TextRoPEInputsStep._prepare_text_ids(encoder) +img_ids = _prepare_image_ids(GRID_H, GRID_W, 0.5) +cond_ids = _prepare_image_ids(COND_H, COND_W, 0.5, l_coord=-1) + +with torch.no_grad(): + out_t2i = model( + hidden_states=hidden, + encoder_hidden_states=encoder, + timestep=timestep, + img_ids=img_ids, + txt_ids=txt_ids, + target_timestep=target_timestep, + return_dict=False, + )[0] + out_ref = model( + hidden_states=hidden, + encoder_hidden_states=encoder, + timestep=timestep, + img_ids=img_ids, + txt_ids=txt_ids, + image_cond_hidden_states=cond, + image_cond_ids=cond_ids, + target_timestep=target_timestep, + return_dict=False, + )[0] + +# Distilled schedule: static shift 6.93, sigmas linspace(1, 1/16, 16) +scheduler = FlowMatchEulerDiscreteScheduler(shift=6.93) +scheduler.set_timesteps(sigmas=np.linspace(1.0, 1 / 16, 16)) + +fx = {f"model.{k}": v for k, v in model.state_dict().items()} +fx.update( + { + "hidden": hidden, + "encoder": encoder, + "cond": cond, + "timestep": timestep, + "target_timestep": target_timestep, + "txt_ids": txt_ids, + "img_ids": img_ids, + "cond_ids": cond_ids, + "out_t2i": out_t2i, + "out_ref": out_ref, + "timesteps": scheduler.timesteps.float(), + "sigmas": scheduler.sigmas.float(), + } +) +fx = {k: v.contiguous() for k, v in fx.items()} +save_file(fx, os.path.join(OUT_DIR, "flux3_model.safetensors"), + metadata={"purpose": "diffuseR FLUX 3 test fixture"}) +print(f"wrote {len(fx)} tensors to {OUT_DIR}/flux3_model.safetensors") + +if os.path.isdir(CKPT_DIR): + shutil.rmtree(CKPT_DIR) +model.save_pretrained(CKPT_DIR, max_shard_size="250KB") +print(f"wrote checkpoint to {CKPT_DIR}: {sorted(os.listdir(CKPT_DIR))}") diff --git a/tools/gen_fixtures_flux3_preprocess.py b/tools/gen_fixtures_flux3_preprocess.py new file mode 100644 index 0000000..ae307f2 --- /dev/null +++ b/tools/gen_fixtures_flux3_preprocess.py @@ -0,0 +1,38 @@ +# Generate FLUX 3 reference-image preprocessing fixtures: the +# single-reference "fill" path (canvas rounding, reflect / symmetric +# padding, crop box) on small synthetic images. +# +# Uses diffusers Flux3ImageProcessor (Apache-2.0). Run via +# tools/gen_fixtures.sh. + +import os +import sys + +import numpy as np +import PIL.Image +import torch +from safetensors.torch import save_file + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "ref", "upstream", "diffusers", "src")) + +from diffusers.modular_pipelines.flux3.image_processor import Flux3ImageProcessor # noqa: E402 + +OUT_DIR = os.path.join(os.path.dirname(__file__), "..", "inst", "tinytest", "fixtures") +os.makedirs(OUT_DIR, exist_ok=True) + +rng = np.random.default_rng(55) +proc = Flux3ImageProcessor() + +fx = {} +# (height, width): padded with reflect; a side <= 16 switches to symmetric +for name, (h, w) in {"a": (70, 100), "b": (12, 20), "c": (64, 128)}.items(): + arr = rng.integers(0, 256, size=(h, w, 3), dtype=np.uint8) + image, kwargs = proc.preprocess(PIL.Image.fromarray(arr), resize_mode="fill") + fx[f"{name}_input"] = torch.from_numpy(arr.copy()) + fx[f"{name}_output"] = image.float() + fx[f"{name}_crop"] = torch.tensor(kwargs["reference_crop"]) + +fx = {k: v.contiguous() for k, v in fx.items()} +save_file(fx, os.path.join(OUT_DIR, "flux3_preprocess.safetensors"), + metadata={"purpose": "diffuseR FLUX 3 test fixture"}) +print(f"wrote {len(fx)} tensors: " + ", ".join(f"{k} {tuple(v.shape)}" for k, v in fx.items())) diff --git a/tools/gen_fixtures_flux3_qwen.py b/tools/gen_fixtures_flux3_qwen.py new file mode 100644 index 0000000..b625098 --- /dev/null +++ b/tools/gen_fixtures_flux3_qwen.py @@ -0,0 +1,91 @@ +# Generate a tiny random-init Qwen3-VL text-encoder fixture for the FLUX 3 +# prompt path: a right-padded prompt window with the pipeline's explicit +# position ids (padding at position 1), hidden states stacked from the +# configured layers. Saves a save_pretrained checkpoint (vision tower +# included, so the loader's skip is exercised). +# +# Uses HF transformers' Qwen3VLForConditionalGeneration (Apache-2.0) and +# mirrors diffusers Flux3TextEncoderStep._get_qwen3_vl_prompt_embeds +# (Apache-2.0). Run via tools/gen_fixtures.sh. + +import os +import shutil + +import torch +from safetensors.torch import save_file +from transformers import Qwen3VLConfig, Qwen3VLForConditionalGeneration + +OUT_DIR = os.path.join(os.path.dirname(__file__), "..", "inst", "tinytest", "fixtures") +CKPT_DIR = os.path.join(OUT_DIR, "flux3_qwen_tiny_ckpt") +os.makedirs(OUT_DIR, exist_ok=True) + +torch.manual_seed(54) + +config = Qwen3VLConfig( + text_config=dict( + vocab_size=100, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=4, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + rope_theta=5e6, + rms_norm_eps=1e-6, + max_position_embeddings=512, + rope_scaling={"rope_type": "default", "mrope_section": [2, 1, 1], + "mrope_interleaved": True}, + tie_word_embeddings=True, + ), + vision_config=dict( + depth=1, + hidden_size=16, + intermediate_size=32, + num_heads=2, + out_hidden_size=32, + patch_size=4, + spatial_merge_size=2, + temporal_patch_size=2, + num_position_embeddings=16, + deepstack_visual_indexes=[0], + ), + image_token_id=97, + video_token_id=98, + vision_start_token_id=95, + vision_end_token_id=96, + tie_word_embeddings=True, +) +model = Qwen3VLForConditionalGeneration(config) +model.eval() + +# One prompt, right-padded to a window of 12 (pad id 0) +real = [5, 23, 61, 7, 19, 88, 42, 3, 9] +window = 12 +input_ids = torch.tensor([real + [0] * (window - len(real))]) +attention_mask = torch.tensor([[1] * len(real) + [0] * (window - len(real))]) +position_ids = torch.arange(window)[None].masked_fill(attention_mask == 0, 1) + +layers = (1, 2, 4) +with torch.no_grad(): + out = model(input_ids=input_ids, attention_mask=attention_mask, + position_ids=position_ids, output_hidden_states=True, + use_cache=False) +stacked = torch.stack([out.hidden_states[k] for k in layers], dim=1) +b, c, s, h = stacked.shape +stacked = stacked.permute(0, 2, 1, 3).reshape(b, s, c * h) + +fx = { + "input_ids": input_ids, + "attention_mask": attention_mask, + "position_ids": position_ids, + "stacked": stacked, +} +fx = {k: v.contiguous() for k, v in fx.items()} +save_file(fx, os.path.join(OUT_DIR, "flux3_qwen.safetensors"), + metadata={"purpose": "diffuseR FLUX 3 test fixture"}) +print(f"wrote {len(fx)} tensors to {OUT_DIR}/flux3_qwen.safetensors") + +if os.path.isdir(CKPT_DIR): + shutil.rmtree(CKPT_DIR) +model.save_pretrained(CKPT_DIR) +print(f"wrote checkpoint to {CKPT_DIR}: {sorted(os.listdir(CKPT_DIR))}")