diff --git a/CLAUDE.md b/CLAUDE.md index 54e8094..c9603d9 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -317,6 +317,7 @@ See cornyverse CLAUDE.md for safetensors package setup (use cornball-ai fork unt ### Model Support - [x] Add FLUX model support (FLUX.1-schnell, see below) +- [x] Add FLUX.2 support (klein-4B, see below) - [ ] Add SD3 model support - [ ] ControlNet integration @@ -347,6 +348,31 @@ the shipped tokenizer_2 uses Metaspace prepend "always" (a spiece conversion gives "never" - different ids for every prompt), and quantized residents must follow the compute dtype (bf16 GPU / fp32 CPU). +### FLUX.2 Klein 4B (Complete) + +Step-distilled FLUX.2 (4 steps, no CFG): 4B MMDiT (shared modulation, +SwiGLU, parallel single blocks) + Qwen3-4B text encoder (chat template, +mid-stack hidden states) + 32-channel `AutoencoderKLFlux2` (BatchNorm +latent stats) + FlowMatch with the BFL empirical dynamic shift. + +```r +download_flux2_klein() # ungated, Apache-2.0; ~16 GB download, + # one-time fp8 quantize to a 3.9 GB artifact +txt2img_flux2("An astronaut riding a horse on Mars, photorealistic", + seed = 7) # or txt2img("...", model_name = "flux2") +``` + +Measured on the RTX 5060 Ti 16 GB (fp8 GPU-resident, Qwen3 bf16 +phase-onloaded): 1024x1024 in ~48 s (peak 8.2 GB), 512x512 in ~40 s; +pipeline load 31 s. Cast census is exactly 104 weights. + +Perf lesson that cost an afternoon: generation was 93.8% R garbage +collection until the tokenizer stopped holding 151k-binding +environments — R gc cost scales with live object count (12 vs 203 ms), +and torch's allocator callbacks run gc hundreds of times per +generation. Keep big lookup tables as atomic vectors (integer-id BPE +via findInterval), never as environments. + ### LTX-2.3 Video Generation (clean-room rewrite in progress) The original LTX-2.0 port was removed and is being replaced by a ground-up diff --git a/DESCRIPTION b/DESCRIPTION index 9201130..5f6c71c 100644 --- a/DESCRIPTION +++ b/DESCRIPTION @@ -1,6 +1,6 @@ Package: diffuseR Title: Functional Interface to Diffusion Models in R -Version: 0.1.0.1 +Version: 0.1.0.2 Authors@R: c( person("Troy", "Hernandez", email = "troy@cornball.ai", role = c("aut", "cre"), comment = c(ORCID = "0009-0005-4248-604X")), diff --git a/NAMESPACE b/NAMESPACE index 3f00af7..aaeb02f 100644 --- a/NAMESPACE +++ b/NAMESPACE @@ -10,11 +10,14 @@ export(ddim_scheduler_step) export(decode_bpe) export(download_component) export(download_flux1) +export(download_flux2_klein) export(download_ltx2) export(download_model) export(encode_bpe) +export(encode_qwen) export(encode_unigram) export(encode_with_gemma3) +export(encode_with_qwen3) export(encode_with_t5) export(filename_from_prompt) export(flowmatch_calculate_shift) @@ -41,6 +44,23 @@ export(flux_quantize) export(flux_single_block) export(flux_transformer) export(flux_unpack_latents) +export(flux2_bn_normalize) +export(flux2_double_block) +export(flux2_empirical_mu) +export(flux2_feed_forward) +export(flux2_is_quant_key) +export(flux2_load_pipeline) +export(flux2_modulation) +export(flux2_pack_latents) +export(flux2_parallel_self_attention) +export(flux2_patchify_latents) +export(flux2_prepare_latent_ids) +export(flux2_prepare_text_ids) +export(flux2_single_block) +export(flux2_transformer) +export(flux2_unpack_latents_with_ids) +export(flux2_unpatchify_latents) +export(flux2_vae_decoder) export(gemma3_config_ltx2) export(gemma3_text_model) export(gemma3_tokenizer) @@ -49,9 +69,11 @@ export(is_blackwell_gpu) export(latents_to_video) export(load_decoder_safetensors) export(load_decoder_weights) +export(load_flux2_vae_decoder) export(load_gemma3_text_encoder) export(load_model_component) export(load_pipeline) +export(load_qwen3_text_encoder) export(load_t5_text_encoder) export(load_text_encoder_safetensors) export(load_text_encoder_weights) @@ -144,6 +166,8 @@ export(offload_to_cpu) export(post_quant_conv) export(preprocess_image) export(quant_conv) +export(qwen_bpe_tokenizer) +export(qwen3_encoder) export(save_frames) export(save_image) export(save_video) @@ -156,6 +180,7 @@ export(text_encoder2_native) export(tokenize_gemma3) export(txt2img) export(txt2img_flux) +export(txt2img_flux2) export(txt2img_sd21) export(txt2img_sdxl) export(txt2vid_ltx2) @@ -171,6 +196,7 @@ export(write_wav) S3method(print,bpe_tokenizer) S3method(print,ltx23_checkpoint) +S3method(print,qwen_tokenizer) S3method(print,unigram_tokenizer) importFrom(utils,head) diff --git a/R/checkpoint_flux.R b/R/checkpoint_flux.R index 841df5c..2012cb4 100644 --- a/R/checkpoint_flux.R +++ b/R/checkpoint_flux.R @@ -42,6 +42,46 @@ flux_is_quant_key <- function(key) { grepl(.flux_quant_cast_pattern, key) } +# FLUX.2 cast set: block linears, the three shared modulation +# projections, and the context embedder (7680 x 3072). Everything else +# (x_embedder, timestep MLP, norm_out, q/k norms) stays in the resident +# dtype. Full klein-4B census: 5 double x 12 + 20 single x 2 + 3 +# modulations + context_embedder = 104 cast weights, ~3.9B of 4B params. +.flux2_quant_cast_pattern <- paste0( + "^(", + "transformer_blocks\\.[0-9]+\\.(", + "attn\\.(to_q|to_k|to_v|add_q_proj|add_k_proj|add_v_proj|to_out\\.0|to_add_out)", + "|ff\\.(linear_in|linear_out)|ff_context\\.(linear_in|linear_out)", + ")", + "|single_transformer_blocks\\.[0-9]+\\.attn\\.(to_qkv_mlp_proj|to_out)", + "|double_stream_modulation_img\\.linear", + "|double_stream_modulation_txt\\.linear", + "|single_stream_modulation\\.linear", + "|context_embedder", + ")\\.weight$" +) + +#' Test whether a FLUX.2 key is in the quantization cast set +#' +#' @param key Character vector of parameter names (diffusers-style). +#' +#' @return Logical vector. +#' +#' @export +flux2_is_quant_key <- function(key) { + grepl(.flux2_quant_cast_pattern, key) +} + +# Model family from a diffusers transformer config +.flux_family <- function(config) { + cls <- config$`_class_name` %||% "FluxTransformer2DModel" + if (identical(cls, "Flux2Transformer2DModel")) { + "flux2" + } else { + "flux1" + } +} + #' Open a FLUX transformer checkpoint directory #' #' Opens a diffusers-layout transformer directory lazily (headers only). diff --git a/R/dit_flux2.R b/R/dit_flux2.R new file mode 100644 index 0000000..2322fd9 --- /dev/null +++ b/R/dit_flux2.R @@ -0,0 +1,163 @@ +#' FLUX.2 Transformer (MMDiT) +#' +#' Fresh R port of Flux2Transformer2DModel from the diffusers reference +#' implementation (Apache-2.0, +#' src/diffusers/models/transformers/transformer_flux2.py). Defaults are +#' the klein-4B configuration (5 double + 20 single blocks). Guidance +#' embeddings (FLUX.2-dev) are not implemented; klein is step-distilled +#' with \code{guidance_embeds = false}. Timestep conditioning has no +#' pooled-text component, and the three modulation projections are +#' shared across all blocks. +#' +#' @name dit_flux2 +NULL + +# Timestep-only conditioning: sinusoid(256) -> MLP, bias-free. Matches +# diffusers Flux2TimestepGuidanceEmbeddings state-dict names (klein has +# no guidance_embedder). +flux2_time_guidance_embed <- torch::nn_module( + "flux2_time_guidance_embed", + initialize = function(embedding_dim, in_channels = 256L) { + self$in_channels <- in_channels + self$timestep_embedder <- ltx23_timestep_embedding(in_channels, + embedding_dim, bias = FALSE) +}, + forward = function(timestep) { + proj <- ltx23_get_timestep_embedding(timestep, self$in_channels, + flip_sin_to_cos = TRUE, downscale_freq_shift = 0) + self$timestep_embedder(proj$to(dtype = self$timestep_embedder$linear_1$weight$dtype)) +} +) + +#' FLUX.2 transformer model +#' +#' Shared modulation computed once per forward; double blocks over +#' separate text/image streams, then single (parallel) blocks over the +#' concatenated [text; image] sequence. Rotary embeddings are +#' precomputed by the caller with \code{\link{flux_pos_embed}} +#' (\code{axes_dim = c(32, 32, 32, 32)}, \code{theta = 2000}) over the +#' concatenated [text; image] 4-axis position ids. +#' +#' @param in_channels Integer. Packed latent channels (128). +#' @param num_layers Integer. Double-stream block count (klein-4B: 5). +#' @param num_single_layers Integer. Single-stream block count (20). +#' @param attention_head_dim Integer. Per-head dimension. +#' @param num_attention_heads Integer. Attention heads. +#' @param joint_attention_dim Integer. Text embedding dim (7680). +#' @param mlp_ratio Numeric. Feed-forward multiplier (3.0). +#' @param timestep_guidance_channels Integer. Sinusoid width (256). +#' @param axes_dims_rope Integer vector. Per-axis rotary dims. +#' @param rope_theta Numeric. Rotary base frequency (2000). +#' @param eps Numeric. Norm epsilon. +#' @param out_channels Integer or NULL. Defaults to \code{in_channels}. +#' +#' @return Module whose forward(hidden_states, encoder_hidden_states, +#' timestep, image_rotary_emb) returns the predicted velocity for the +#' image tokens [B, S_img, out_channels]. \code{timestep} is in sigma +#' space (0-1); it is scaled by 1000 internally. +#' +#' @export +flux2_transformer <- torch::nn_module( + "flux2_transformer", + initialize = function(in_channels = 128L, + num_layers = 5L, + num_single_layers = 20L, + attention_head_dim = 128L, + num_attention_heads = 24L, + joint_attention_dim = 7680L, + mlp_ratio = 3.0, + timestep_guidance_channels = 256L, + axes_dims_rope = c(32L, 32L, 32L, 32L), + rope_theta = 2000, + eps = 1e-6, + out_channels = NULL) { + inner_dim <- num_attention_heads * attention_head_dim + self$inner_dim <- inner_dim + self$axes_dims_rope <- as.integer(axes_dims_rope) + self$rope_theta <- rope_theta + self$out_channels <- as.integer(out_channels %||% in_channels) + + self$time_guidance_embed <- flux2_time_guidance_embed(inner_dim, + as.integer(timestep_guidance_channels)) + self$double_stream_modulation_img <- flux2_modulation(inner_dim, 2L) + self$double_stream_modulation_txt <- flux2_modulation(inner_dim, 2L) + self$single_stream_modulation <- flux2_modulation(inner_dim, 1L) + + self$x_embedder <- torch::nn_linear(in_channels, inner_dim, bias = FALSE) + self$context_embedder <- torch::nn_linear(joint_attention_dim, + inner_dim, bias = FALSE) + + self$transformer_blocks <- torch::nn_module_list( + lapply(seq_len(num_layers), function(i) { + flux2_double_block(inner_dim, num_attention_heads, + attention_head_dim, mlp_ratio = mlp_ratio, + eps = eps) + }) + ) + self$single_transformer_blocks <- torch::nn_module_list( + lapply(seq_len(num_single_layers), function(i) { + flux2_single_block(inner_dim, num_attention_heads, + attention_head_dim, mlp_ratio = mlp_ratio, + eps = eps) + }) + ) + + self$norm_out <- flux_ada_layer_norm_continuous(inner_dim, inner_dim, + bias = FALSE) + self$proj_out <- torch::nn_linear(inner_dim, self$out_channels, + bias = FALSE) +}, + forward = function(hidden_states, encoder_hidden_states, timestep, + image_rotary_emb, chunk_size = NULL) { + hidden_states <- self$x_embedder(hidden_states) + timestep <- timestep$to(dtype = hidden_states$dtype)$mul(1000) + temb <- self$time_guidance_embed(timestep) + + mod_img <- self$double_stream_modulation_img(temb) + mod_txt <- self$double_stream_modulation_txt(temb) + mod_single <- self$single_stream_modulation(temb) + + encoder_hidden_states <- self$context_embedder(encoder_hidden_states) + + block_gc <- isTRUE(getOption("diffuseR.block_gc")) + for (i in seq_along(self$transformer_blocks)) { + res <- self$transformer_blocks[[i]]( + hidden_states = hidden_states, + encoder_hidden_states = encoder_hidden_states, + temb_mod_img = mod_img, + temb_mod_txt = mod_txt, + image_rotary_emb = image_rotary_emb, + chunk_size = chunk_size + ) + encoder_hidden_states <- res[[1]] + hidden_states <- res[[2]] + if (block_gc) { + gc(verbose = FALSE) + } + } + + txt_len <- encoder_hidden_states$shape[2] + hidden_states <- torch::torch_cat( + list(encoder_hidden_states, hidden_states), + dim = 2L + ) + for (i in seq_along(self$single_transformer_blocks)) { + hidden_states <- self$single_transformer_blocks[[i]]( + hidden_states = hidden_states, + temb_mod = mod_single, + image_rotary_emb = image_rotary_emb, + chunk_size = chunk_size + ) + if (block_gc) { + gc(verbose = FALSE) + } + } + hidden_states <- hidden_states$narrow( + 2L, txt_len + 1L, + hidden_states$shape[2] - txt_len + ) + + hidden_states <- self$norm_out(hidden_states, temb) + self$proj_out(hidden_states) +} +) diff --git a/R/dit_flux2_modules.R b/R/dit_flux2_modules.R new file mode 100644 index 0000000..21460b9 --- /dev/null +++ b/R/dit_flux2_modules.R @@ -0,0 +1,270 @@ +#' FLUX.2 Transformer Building Blocks +#' +#' Fresh R port of the FLUX.2 MMDiT blocks from the diffusers reference +#' implementation (Apache-2.0, +#' src/diffusers/models/transformers/transformer_flux2.py). Key +#' differences from FLUX.1: modulation is computed ONCE at model level +#' by shared \code{flux2_modulation} projections and passed into the +#' blocks (block norms are parameterless), feed-forwards use SwiGLU with +#' the gate fused into \code{linear_in}, the single-stream block is a +#' ViT-22B-style parallel block with fully fused projections, and every +#' linear is bias-free. Module field names mirror the diffusers +#' state-dict keys 1:1. Reuses \code{flux_attention} (bias = FALSE), +#' \code{ltx23_rms_norm}, \code{.ltx23_sdpa}, and +#' \code{flux_apply_rotary_emb}. +#' +#' @name dit_flux2_modules +NULL + +#' FLUX.2 shared modulation projection +#' +#' \code{linear(silu(temb))} producing \code{mod_param_sets} triples of +#' (shift, scale, gate). Computed once per forward at model level and +#' broadcast to every block. Reference: Flux2Modulation. +#' +#' @param dim Integer. Model dimension. +#' @param mod_param_sets Integer. Number of (shift, scale, gate) triples. +#' @param bias Logical. +#' +#' @export +flux2_modulation <- torch::nn_module( + "flux2_modulation", + initialize = function(dim, mod_param_sets = 2L, bias = FALSE) { + self$mod_param_sets <- mod_param_sets + self$linear <- torch::nn_linear(dim, dim * 3L * mod_param_sets, bias = bias) +}, + forward = function(temb) { + self$linear(torch::nnf_silu(temb)) +} +) + +# Split a modulation tensor into mod_param_sets triples of +# (shift, scale, gate), each [N, 1, dim]. Reference: Flux2Modulation.split. +.flux2_mod_split <- function(mod, mod_param_sets) { + if (mod$ndim == 2L) { + mod <- mod$unsqueeze(2L) + } + parts <- mod$chunk(3L * mod_param_sets, dim = -1L) + lapply(seq_len(mod_param_sets), function(i) { + parts[(3L * (i - 1L) + 1L):(3L * i)] + }) +} + +#' FLUX.2 feed-forward (fused SwiGLU) +#' +#' \code{linear_in} projects to twice the inner dim; SwiGLU gates the +#' first half with SiLU and multiplies by the second half; +#' \code{linear_out} projects back. Reference: Flux2FeedForward + +#' Flux2SwiGLU. +#' +#' @param dim Integer. Input dimension. +#' @param dim_out Integer. Output dimension (defaults to \code{dim}). +#' @param mult Numeric. Inner dim multiplier (FLUX.2: 3.0). +#' @param bias Logical. +#' +#' @export +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) +}, + forward = function(x) { + x <- self$linear_in(x) + half <- x$shape[length(x$shape)] %/% 2L + x <- torch::nnf_silu(x$narrow(-1L, 1L, half)) * + x$narrow(-1L, half + 1L, half) + self$linear_out(x) +} +) + +#' FLUX.2 parallel self-attention (single-stream) +#' +#' ViT-22B-style parallel block internals: one fused projection produces +#' QKV and the SwiGLU MLP input; one fused projection consumes +#' cat(attention output, MLP output). Reference: +#' Flux2ParallelSelfAttention + Flux2ParallelSelfAttnProcessor. +#' +#' @param query_dim Integer. Model dimension. +#' @param heads Integer. Attention heads. +#' @param dim_head Integer. Per-head dimension. +#' @param mlp_ratio Numeric. MLP hidden multiplier (FLUX.2: 3.0). +#' @param eps Numeric. RMS norm epsilon. +#' @param bias Logical. +#' +#' @export +flux2_parallel_self_attention <- torch::nn_module( + "flux2_parallel_self_attention", + initialize = function(query_dim, heads, dim_head, mlp_ratio = 3.0, + eps = 1e-6, bias = FALSE) { + inner_dim <- heads * dim_head + self$heads <- heads + self$inner_dim <- inner_dim + self$mlp_hidden_dim <- as.integer(query_dim * mlp_ratio) + + self$to_qkv_mlp_proj <- torch::nn_linear(query_dim, + inner_dim * 3L + self$mlp_hidden_dim * 2L, bias = bias) + 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_linear(inner_dim + self$mlp_hidden_dim, + query_dim, bias = bias) +}, + forward = function(hidden_states, image_rotary_emb = NULL, + chunk_size = NULL) { + proj <- self$to_qkv_mlp_proj(hidden_states) + qkv <- proj$narrow(-1L, 1L, 3L * self$inner_dim) + mlp <- proj$narrow(-1L, 3L * self$inner_dim + 1L, self$mlp_hidden_dim * 2L) + + parts <- qkv$chunk(3L, dim = -1L) + query <- parts[[1]]$unflatten(3L, c(self$heads, -1L)) + key <- parts[[2]]$unflatten(3L, c(self$heads, -1L)) + value <- parts[[3]]$unflatten(3L, c(self$heads, -1L)) + + query <- self$norm_q(query) + key <- self$norm_k(key) + + # [B, S, H, D] -> [B, H, S, D] + query <- query$transpose(2L, 3L) + key <- key$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) + } + attn <- .ltx23_sdpa(query, key, value, chunk_size = chunk_size) + attn <- attn$transpose(2L, 3L)$flatten(start_dim = 3L) + attn <- attn$to(dtype = hidden_states$dtype) + + # SwiGLU on the fused MLP half + half <- self$mlp_hidden_dim + mlp <- torch::nnf_silu(mlp$narrow(-1L, 1L, half)) * + mlp$narrow(-1L, half + 1L, half) + + # Attention half first, then the MLP half + self$to_out(torch::torch_cat(list(attn, mlp), dim = -1L)) +} +) + +#' FLUX.2 double-stream (MMDiT) block +#' +#' Image and text streams with externally supplied (shift, scale, gate) +#' modulation triples, joint attention (txt first), and SwiGLU +#' feed-forwards. Reference: Flux2TransformerBlock. +#' +#' @param dim Integer. Model dimension. +#' @param num_attention_heads Integer. Attention heads. +#' @param attention_head_dim Integer. Per-head dimension. +#' @param mlp_ratio Numeric. FF multiplier (FLUX.2: 3.0). +#' @param eps Numeric. Norm epsilon. +#' @param bias Logical. +#' +#' @return Module whose forward(hidden_states, encoder_hidden_states, +#' temb_mod_img, temb_mod_txt, image_rotary_emb) returns +#' \code{list(encoder_hidden_states, hidden_states)}. +#' +#' @export +flux2_double_block <- torch::nn_module( + "flux2_double_block", + initialize = function(dim, num_attention_heads, attention_head_dim, + mlp_ratio = 3.0, eps = 1e-6, bias = FALSE) { + self$norm1 <- torch::nn_layer_norm(dim, eps = eps, + elementwise_affine = FALSE) + self$norm1_context <- torch::nn_layer_norm(dim, eps = eps, + elementwise_affine = FALSE) + self$attn <- flux_attention(dim, num_attention_heads, + attention_head_dim, added_kv = TRUE, + eps = eps, bias = bias) + self$norm2 <- torch::nn_layer_norm(dim, eps = eps, + elementwise_affine = FALSE) + self$ff <- flux2_feed_forward(dim, dim, mult = mlp_ratio, bias = bias) + self$norm2_context <- torch::nn_layer_norm(dim, eps = eps, + elementwise_affine = FALSE) + self$ff_context <- flux2_feed_forward(dim, dim, mult = mlp_ratio, + bias = bias) +}, + forward = function(hidden_states, encoder_hidden_states, temb_mod_img, + temb_mod_txt, image_rotary_emb = NULL, + chunk_size = NULL) { + mi <- .flux2_mod_split(temb_mod_img, 2L) + mt <- .flux2_mod_split(temb_mod_txt, 2L) + # Each triple is (shift, scale, gate) + msa <- mi[[1]] + mlp <- mi[[2]] + c_msa <- mt[[1]] + c_mlp <- mt[[2]] + + norm_h <- self$norm1(hidden_states) * msa[[2]]$add(1) + msa[[1]] + norm_c <- self$norm1_context(encoder_hidden_states) * + c_msa[[2]]$add(1) + c_msa[[1]] + + attn_out <- self$attn( + hidden_states = norm_h, + encoder_hidden_states = norm_c, + image_rotary_emb = image_rotary_emb, + chunk_size = chunk_size + ) + + hidden_states <- hidden_states + msa[[3]] * attn_out[[1]] + norm_h <- self$norm2(hidden_states) * mlp[[2]]$add(1) + mlp[[1]] + hidden_states <- hidden_states + mlp[[3]] * self$ff(norm_h) + + encoder_hidden_states <- encoder_hidden_states + c_msa[[3]] * attn_out[[2]] + norm_c <- self$norm2_context(encoder_hidden_states) * + c_mlp[[2]]$add(1) + c_mlp[[1]] + encoder_hidden_states <- encoder_hidden_states + + c_mlp[[3]] * self$ff_context(norm_c) + + if (encoder_hidden_states$dtype == torch::torch_float16()) { + encoder_hidden_states <- encoder_hidden_states$clamp(-65504, 65504) + } + list(encoder_hidden_states, hidden_states) +} +) + +#' FLUX.2 single-stream block (parallel) +#' +#' Parameterless LayerNorm with external modulation, then the fused +#' parallel attention+MLP. Operates on the pre-concatenated [text; image] +#' sequence (the reference model concatenates once before the stack). +#' Reference: Flux2SingleTransformerBlock. +#' +#' @param dim Integer. Model dimension. +#' @param num_attention_heads Integer. Attention heads. +#' @param attention_head_dim Integer. Per-head dimension. +#' @param mlp_ratio Numeric. MLP multiplier (FLUX.2: 3.0). +#' @param eps Numeric. Norm epsilon. +#' @param bias Logical. +#' +#' @return Module whose forward(hidden_states, temb_mod, +#' image_rotary_emb) returns the joint hidden states. +#' +#' @export +flux2_single_block <- torch::nn_module( + "flux2_single_block", + initialize = function(dim, num_attention_heads, attention_head_dim, + mlp_ratio = 3.0, eps = 1e-6, bias = FALSE) { + self$norm <- torch::nn_layer_norm(dim, eps = eps, + elementwise_affine = FALSE) + self$attn <- flux2_parallel_self_attention( + dim, num_attention_heads, attention_head_dim, + mlp_ratio = mlp_ratio, eps = eps, bias = bias + ) +}, + forward = function(hidden_states, temb_mod, image_rotary_emb = NULL, + chunk_size = NULL) { + mod <- .flux2_mod_split(temb_mod, 1L)[[1]] + norm_h <- self$norm(hidden_states) * mod[[2]]$add(1) + mod[[1]] + attn_out <- self$attn( + hidden_states = norm_h, + image_rotary_emb = image_rotary_emb, + chunk_size = chunk_size + ) + hidden_states <- hidden_states + mod[[3]] * attn_out + if (hidden_states$dtype == torch::torch_float16()) { + hidden_states <- hidden_states$clamp(-65504, 65504) + } + hidden_states +} +) diff --git a/R/dit_flux_modules.R b/R/dit_flux_modules.R index 959a17e..afe04b0 100644 --- a/R/dit_flux_modules.R +++ b/R/dit_flux_modules.R @@ -74,12 +74,14 @@ flux_ada_layer_norm_zero_single <- torch::nn_module( #' #' @param dim Integer. Model dimension. #' @param cond_dim Integer. Conditioning embedding dimension. +#' @param bias Logical. Bias on the projection (TRUE for FLUX.1, FALSE +#' for FLUX.2). #' #' @export flux_ada_layer_norm_continuous <- torch::nn_module( "flux_ada_layer_norm_continuous", - initialize = function(dim, cond_dim = dim) { - self$linear <- torch::nn_linear(cond_dim, 2L * dim, bias = TRUE) + initialize = function(dim, cond_dim = dim, bias = TRUE) { + self$linear <- torch::nn_linear(cond_dim, 2L * dim, bias = bias) self$norm <- torch::nn_layer_norm(dim, eps = 1e-6, elementwise_affine = FALSE) }, @@ -107,12 +109,14 @@ flux_ada_layer_norm_continuous <- torch::nn_module( #' @param added_kv Logical. Add text-stream projections (double blocks). #' @param pre_only Logical. Skip the output projection (single blocks). #' @param eps Numeric. RMS norm epsilon. +#' @param bias Logical. Bias on the linear projections (TRUE for FLUX.1, +#' FALSE for FLUX.2). #' #' @export flux_attention <- torch::nn_module( "flux_attention", initialize = function(query_dim, heads, dim_head, added_kv = FALSE, - pre_only = FALSE, eps = 1e-6) { + pre_only = FALSE, eps = 1e-6, bias = TRUE) { inner_dim <- heads * dim_head self$heads <- heads self$dim_head <- dim_head @@ -121,22 +125,22 @@ flux_attention <- torch::nn_module( self$norm_q <- ltx23_rms_norm(dim_head, eps = eps) self$norm_k <- ltx23_rms_norm(dim_head, eps = eps) - self$to_q <- torch::nn_linear(query_dim, inner_dim, bias = TRUE) - self$to_k <- torch::nn_linear(query_dim, inner_dim, bias = TRUE) - self$to_v <- torch::nn_linear(query_dim, inner_dim, bias = TRUE) + self$to_q <- torch::nn_linear(query_dim, inner_dim, bias = bias) + self$to_k <- torch::nn_linear(query_dim, inner_dim, bias = bias) + self$to_v <- torch::nn_linear(query_dim, inner_dim, bias = bias) if (!pre_only) { self$to_out <- torch::nn_module_list(list( - torch::nn_linear(inner_dim, query_dim, bias = TRUE) + torch::nn_linear(inner_dim, query_dim, bias = bias) )) } if (added_kv) { self$norm_added_q <- ltx23_rms_norm(dim_head, eps = eps) self$norm_added_k <- ltx23_rms_norm(dim_head, eps = eps) - self$add_q_proj <- torch::nn_linear(query_dim, inner_dim, bias = TRUE) - self$add_k_proj <- torch::nn_linear(query_dim, inner_dim, bias = TRUE) - self$add_v_proj <- torch::nn_linear(query_dim, inner_dim, bias = TRUE) - self$to_add_out <- torch::nn_linear(inner_dim, query_dim, bias = TRUE) + self$add_q_proj <- torch::nn_linear(query_dim, inner_dim, bias = bias) + self$add_k_proj <- torch::nn_linear(query_dim, inner_dim, bias = bias) + self$add_v_proj <- torch::nn_linear(query_dim, inner_dim, bias = bias) + self$to_add_out <- torch::nn_linear(inner_dim, query_dim, bias = bias) } }, forward = function(hidden_states, encoder_hidden_states = NULL, diff --git a/R/dit_ltx23_modules.R b/R/dit_ltx23_modules.R index 25e71df..9752fde 100644 --- a/R/dit_ltx23_modules.R +++ b/R/dit_ltx23_modules.R @@ -86,9 +86,10 @@ ltx23_get_timestep_embedding <- function(timesteps, embedding_dim, # diffusers TimestepEmbedding state-dict names. ltx23_timestep_embedding <- torch::nn_module( "ltx23_timestep_embedding", - initialize = function(in_channels, time_embed_dim) { - self$linear_1 <- torch::nn_linear(in_channels, time_embed_dim) - self$linear_2 <- torch::nn_linear(time_embed_dim, time_embed_dim) + initialize = function(in_channels, time_embed_dim, bias = TRUE) { + self$linear_1 <- torch::nn_linear(in_channels, time_embed_dim, bias = bias) + self$linear_2 <- torch::nn_linear(time_embed_dim, time_embed_dim, + bias = bias) }, forward = function(sample) { self$linear_2(torch::nnf_silu(self$linear_1(sample))) diff --git a/R/download_flux2.R b/R/download_flux2.R new file mode 100644 index 0000000..05171bd --- /dev/null +++ b/R/download_flux2.R @@ -0,0 +1,139 @@ +#' Download and Prepare FLUX.2 Klein 4B Weights +#' +#' Downloads FLUX.2-klein-4B from HuggingFace (Apache-2.0, ungated) and +#' quantizes the 4B transformer to a local fp8 (~4 GB) or NF4 (~2.3 GB) +#' artifact. +#' +#' @name download_flux2 +NULL + +.flux2_repo <- "black-forest-labs/FLUX.2-klein-4B" + +.flux2_transformer_files <- c("transformer/config.json", + "transformer/diffusion_pytorch_model.safetensors") + +.flux2_support_files <- c( + "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), + "tokenizer/tokenizer.json", + "tokenizer/tokenizer_config.json", + "tokenizer/special_tokens_map.json", + "scheduler/scheduler_config.json" +) + +#' Download FLUX.2-klein-4B and build the quantized artifact +#' +#' Skips work already done: a valid quantized manifest short-circuits +#' the transformer download; cached files are not re-fetched. No token +#' is needed (the repo is ungated). The bf16 transformer source +#' (~7.8 GB in the HuggingFace cache) may be deleted after quantization. +#' +#' @param quantize Logical. Build the quantized artifact. +#' @param precision "fp8" (~4 GB, GPU-resident; near-bf16 quality) or +#' "nf4" (~2.3 GB). +#' @param output_dir Directory for the quantized artifact. +#' @param text_encoders Logical. Also fetch the Qwen3 text encoder, +#' tokenizer, VAE, and scheduler config (~8.3 GB). +#' @param verbose Logical. +#' +#' @return Invisibly, a list with \code{transformer_dir}, +#' \code{artifact_dir}, and \code{support} (named file paths). +#' +#' @export +download_flux2_klein <- function(quantize = TRUE, + precision = c("fp8", "nf4"), + output_dir = NULL, text_encoders = TRUE, + verbose = TRUE) { + precision <- match.arg(precision) + if (is.null(output_dir)) { + output_dir <- file.path(tools::R_user_dir("diffuseR", "data"), + paste0("flux2-klein-4b-", 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) { + cached <- tryCatch( + hfhub::hub_download(.flux2_repo, .flux2_transformer_files[[2]], + local_files_only = TRUE), + error = function(e) NULL + ) + if (is.null(cached) && !have_artifact) { + free <- .ltx23_disk_free_gb(path.expand("~")) + if (!is.na(free) && free < 25) { + warning(sprintf( + "Only %.0f GB free; the download + %s artifact need ~25 GB.", + free, precision + )) + } + ok <- .ltx23_consent(paste0( + "FLUX.2-klein-4B: the 7.8 GB bf16 transformer plus a local ", + precision, " artifact (Apache-2.0, ungated)" + )) + if (!ok) { + stop("Download cancelled.", call. = FALSE) + } + if (verbose) { + message("Downloading the FLUX.2-klein-4B transformer (7.8 GB)...") + } + } + paths <- vapply(.flux2_transformer_files, function(f) { + hfhub::hub_download(.flux2_repo, f) + }, 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, "\n", + "The 7.8 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) { + have_te <- !is.null(tryCatch( + hfhub::hub_download(.flux2_repo, .flux2_support_files[[5]], + local_files_only = TRUE), + error = function(e) NULL + )) + if (!have_te) { + ok <- .ltx23_consent( + "the Qwen3-4B text encoder, tokenizer, and VAE (~8.3 GB)" + ) + if (!ok) { + stop("Download cancelled.", call. = FALSE) + } + if (verbose) { + message("Downloading text encoder + VAE...") + } + } + result$support <- vapply(.flux2_support_files, function(f) { + hfhub::hub_download(.flux2_repo, f) + }, character(1)) + } + + invisible(result) +} diff --git a/R/fp8_ltx23.R b/R/fp8_ltx23.R index 361be65..88233ba 100644 --- a/R/fp8_ltx23.R +++ b/R/fp8_ltx23.R @@ -62,12 +62,18 @@ ltx23_fp8_linear <- torch::nn_module( invisible(self) }, forward = function(x) { - # Transfer fp8 bytes first, cast on the compute device second - w <- self$weight_fp8$to(device = x$device, non_blocking = TRUE) - w <- w$to(dtype = x$dtype) * self$weight_scale$to(device = x$device, dtype = x$dtype) - out <- torch::nnf_linear(x, w, self$bias) - rm(w) - out + # Transfer fp8 bytes first (no-op when resident), then cast into a + # persistent per-shape buffer and scale in place: zero fresh + # allocations per call, same math as the allocating upcast + w8 <- self$weight_fp8$to(device = x$device, non_blocking = TRUE) + w <- .ltx23_get_dequant_buffer( + c(self$out_features, self$in_features), x$dtype, x$device + ) + torch::with_no_grad({ + w$copy_(w8) + w$mul_(self$weight_scale$to(device = x$device, dtype = x$dtype)) + }) + torch::nnf_linear(x, w, self$bias) } ) diff --git a/R/models2devices.R b/R/models2devices.R index c0e8638..7b5f3f9 100644 --- a/R/models2devices.R +++ b/R/models2devices.R @@ -58,7 +58,8 @@ get_required_components <- function(model_name) { "sd21" = c("unet", "decoder", "text_encoder", "encoder"), "sdxl" = c("unet", "decoder", "text_encoder", "text_encoder2", "encoder"), - "flux1" = c("transformer", "decoder", "text_encoder", "text_encoder2") + "flux1" = c("transformer", "decoder", "text_encoder", "text_encoder2"), + "flux2" = c("transformer", "decoder", "text_encoder") # "sd3" = c("transformer", "decoder", "text_encoder", "text_encoder2", "text_encoder3", "encoder"), # "cascade" = c("prior", "decoder", "text_encoder", "vqgan") ) diff --git a/R/quantize_flux.R b/R/quantize_flux.R index 9021d6d..6f4376e 100644 --- a/R/quantize_flux.R +++ b/R/quantize_flux.R @@ -41,6 +41,63 @@ NULL lapply(args, function(x) if (is.numeric(x)) as.integer(x) else x) } +.flux2_transformer_args <- function(config) { + if (is.null(config)) { + return(list()) + } + if (isTRUE(config$guidance_embeds)) { + stop("This checkpoint uses guidance embeddings (FLUX.2-dev); ", + "only the klein variants (guidance_embeds = false) are supported.") + } + args <- list( + in_channels = config$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_guidance_channels = config$timestep_guidance_channels, + axes_dims_rope = config$axes_dims_rope, + out_channels = config$out_channels + ) + 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) + } + } + args +} + +# Move plain-field fp8 weights (and their scales) onto a device; used +# for resident fp8 where the whole quantized model fits on the GPU +.flux_fp8_to_device <- function(module, device) { + for (name in names(module$children)) { + child <- module$children[[name]] + if (!is.null(child$weight_fp8)) { + child$weight_fp8 <- child$weight_fp8$to(device = device) + child$weight_scale <- child$weight_scale$to(device = device) + } + .flux_fp8_to_device(child, device) + } + invisible(module) +} + +# Family-specific hooks for quantization and loading +.flux_family_hooks <- function(config) { + if (.flux_family(config) == "flux2") { + list(model_fn = flux2_transformer, args_fn = .flux2_transformer_args, + is_quant_key = flux2_is_quant_key, shard_prefix = "flux2-klein") + } else { + list(model_fn = flux_transformer, + args_fn = .flux_transformer_args, + is_quant_key = flux_is_quant_key, + shard_prefix = "flux1") + } +} + #' Quantize a FLUX transformer to NF4 or fp8 shards #' #' Streams the bf16 diffusers checkpoint tensor by tensor. Cast-set @@ -86,6 +143,7 @@ flux_quantize <- function(transformer_dir, output_dir = NULL, dir.create(output_dir, recursive = TRUE, showWarnings = FALSE) ckpt <- flux_open_checkpoint(transformer_dir) + hooks <- .flux_family_hooks(ckpt$config) if (format == "fp8") { fp8 <- torch::torch_float8_e4m3fn() } @@ -99,8 +157,8 @@ flux_quantize <- function(transformer_dir, output_dir = NULL, if (!length(shard)) { return() } - fname <- sprintf("flux1-%s-%05d.safetensors", format, - length(shard_files) + 1L) + fname <- sprintf("%s-%s-%05d.safetensors", hooks$shard_prefix, + format, length(shard_files) + 1L) safetensors::safe_save_file(shard, file.path(output_dir, fname)) shard_files[[length(shard_files) + 1L]] <<- fname if (verbose) { @@ -117,7 +175,7 @@ flux_quantize <- function(transformer_dir, output_dir = NULL, key <- keys[[i]] tensor <- ckpt$handle$get_tensor(key) - if (flux_is_quant_key(key)) { + if (hooks$is_quant_key(key)) { torch::with_no_grad({ if (format == "nf4") { q <- ltx23_nf4_quantize(tensor) @@ -192,7 +250,11 @@ flux_quantize <- function(transformer_dir, output_dir = NULL, #' quantized formats this sets the resident (non-quantized) tensors #' and must match the compute dtype: bfloat16 for GPU compute, #' float32 for CPU compute. -#' @param pin Logical. Pin fp8 host memory for faster transfers. +#' @param pin Logical. Pin fp8 host memory for faster transfers +#' (streamed fp8 only). +#' @param fp8_resident Logical. Keep the fp8 weights on \code{device} +#' instead of streaming from the CPU - right for models whose whole +#' quantized footprint fits in VRAM (FLUX.2 klein-4B: ~4 GB). #' @param verbose Logical. #' @param ... Overrides for \code{\link{flux_transformer}} arguments #' (tiny test configs). @@ -201,12 +263,14 @@ flux_quantize <- function(transformer_dir, output_dir = NULL, #' #' @export flux_load_transformer <- function(ckpt, device = "cuda", dtype = "bfloat16", - pin = TRUE, verbose = TRUE, ...) { + pin = TRUE, fp8_resident = FALSE, + verbose = TRUE, ...) { stopifnot(inherits(ckpt, "ltx23_checkpoint")) format <- ckpt$format %||% "full" + hooks <- .flux_family_hooks(ckpt$config) - args <- utils::modifyList(.flux_transformer_args(ckpt$config), list(...)) - model <- do.call(flux_transformer, args) + args <- utils::modifyList(hooks$args_fn(ckpt$config), list(...)) + model <- do.call(hooks$model_fn, args) if (format == "full") { model$to(dtype = .flux_dtype(dtype)) @@ -244,7 +308,7 @@ flux_load_transformer <- function(ckpt, device = "cuda", dtype = "bfloat16", for (i in seq_along(main_keys)) { key <- main_keys[[i]] - if (flux_is_quant_key(key) && + if (hooks$is_quant_key(key) && paste0(key, sib_suffix) %in% sib_keys) { segments <- strsplit(key, ".", fixed = TRUE)[[1]] parent <- .ltx23_walk_module(model, utils::head(segments, -2L)) @@ -307,7 +371,7 @@ flux_load_transformer <- function(ckpt, device = "cuda", dtype = "bfloat16", paste(utils::head(unmapped, 3), collapse = ", ")) } # Weight params replaced by quantized modules won't be "filled" - expected_missing <- flux_is_quant_key(names(dests)) + expected_missing <- hooks$is_quant_key(names(dests)) unfilled <- setdiff(names(dests)[!expected_missing], filled) if (length(unfilled)) { stop("FLUX ", format, " load: ", length(unfilled), @@ -316,8 +380,12 @@ flux_load_transformer <- function(ckpt, device = "cuda", dtype = "bfloat16", } # NF4: everything (packed buffers included) onto the device. - # FP8: residents move; fp8 weights are plain fields and stay CPU-side. + # FP8: residents move; the plain-field fp8 weights stay CPU-side + # unless fp8_resident moves them (small models that fit on the GPU). model$to(device = device) + if (format == "fp8" && fp8_resident) { + .flux_fp8_to_device(model, device) + } model$eval() # Block intermediates are large at image resolutions; per-block gc # keeps the quantized-linear temporaries bounded diff --git a/R/qwen3_text_encoder.R b/R/qwen3_text_encoder.R new file mode 100644 index 0000000..7d09e32 --- /dev/null +++ b/R/qwen3_text_encoder.R @@ -0,0 +1,296 @@ +#' Qwen3 Text Encoder +#' +#' Fresh R port of the Qwen3 decoder stack from HuggingFace transformers +#' (Apache-2.0, src/transformers/models/qwen3/), used by FLUX.2 klein as +#' its text encoder (Qwen3-4B: 36 layers, hidden 2560, 32 query / 8 KV +#' heads, head_dim 128, SwiGLU 9728, RoPE theta 1e6). The pipeline +#' consumes mid-stack hidden states (layers 9, 18, 27 for klein-4B) +#' concatenated per token, so the forward runs only as deep as the last +#' requested layer; the LM head is never needed (embeddings are tied). +#' Causal attention with the tokenizer's padding mask, matching the +#' reference exactly. +#' +#' @name qwen3_text_encoder +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) { + 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) + 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( + freqs$cos()$unsqueeze(1L)$unsqueeze(1L), + freqs$sin()$unsqueeze(1L)$unsqueeze(1L) + ) +} + +.qwen3_attention <- torch::nn_module( + "qwen3_attention", + initialize = function(hidden_size, num_heads, num_kv_heads, head_dim, + eps = 1e-6) { + 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_norm <- ltx23_rms_norm(head_dim, eps = eps) + self$k_norm <- ltx23_rms_norm(head_dim, eps = eps) +}, + forward = function(x, rope, mask = NULL) { + shape <- x$shape + b <- shape[1] + s <- shape[2] + + q <- self$q_proj(x)$view(c(b, s, self$num_heads, self$head_dim)) + k <- self$k_proj(x)$view(c(b, s, self$num_kv_heads, self$head_dim)) + v <- self$v_proj(x)$view(c(b, s, self$num_kv_heads, self$head_dim)) + + q <- self$q_norm(q)$transpose(2L, 3L) # [B, H, S, D] + k <- self$k_norm(k)$transpose(2L, 3L) + v <- v$transpose(2L, 3L) + + q <- ltx23_apply_split_rotary_emb(q, rope) + k <- ltx23_apply_split_rotary_emb(k, rope) + + # GQA: repeat KV heads to match the query heads + groups <- self$num_heads %/% self$num_kv_heads + if (groups > 1L) { + k <- k$repeat_interleave(groups, dim = 2L) + v <- v$repeat_interleave(groups, dim = 2L) + } + + out <- .ltx23_sdpa(q, k, v, attention_mask = mask) + out <- out$transpose(2L, 3L)$reshape(c(b, s, -1L)) + self$o_proj(out) +} +) + +.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) +}, + forward = function(x) { + self$down_proj(torch::nnf_silu(self$gate_proj(x)) * self$up_proj(x)) +} +) + +.qwen3_layer <- torch::nn_module( + "qwen3_layer", + initialize = function(hidden_size, num_heads, num_kv_heads, head_dim, + intermediate_size, eps = 1e-6) { + self$self_attn <- .qwen3_attention(hidden_size, num_heads, num_kv_heads, + head_dim, eps = eps) + self$mlp <- .qwen3_mlp(hidden_size, intermediate_size) + self$input_layernorm <- ltx23_rms_norm(hidden_size, eps = eps) + self$post_attention_layernorm <- ltx23_rms_norm(hidden_size, eps = eps) +}, + forward = function(x, rope, mask = NULL) { + x <- x + self$self_attn(self$input_layernorm(x), rope, mask) + x + self$mlp(self$post_attention_layernorm(x)) +} +) + +#' Qwen3 encoder stack +#' +#' Defaults are the Qwen3-4B configuration used by FLUX.2 klein. The +#' module tree mirrors the checkpoint keys (\code{model.embed_tokens}, +#' \code{model.layers.*}, \code{model.norm}); the tied LM head carries +#' no weights of its own and is not implemented. +#' +#' @param vocab_size,hidden_size,intermediate_size,num_hidden_layers +#' Integers. +#' @param num_attention_heads,num_key_value_heads,head_dim Integers. +#' @param rope_theta Numeric. +#' @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. +#' +#' @export +qwen3_encoder <- torch::nn_module( + "qwen3_encoder", + initialize = function(vocab_size = 151936L, hidden_size = 2560L, + intermediate_size = 9728L, + num_hidden_layers = 36L, + num_attention_heads = 32L, + num_key_value_heads = 8L, head_dim = 128L, + rope_theta = 1e6, rms_norm_eps = 1e-6) { + self$head_dim <- head_dim + self$rope_theta <- rope_theta + + inner <- torch::nn_module( + "qwen3_model", + initialize = function() { + self$embed_tokens <- torch::nn_embedding(vocab_size, hidden_size) + self$layers <- torch::nn_module_list( + lapply(seq_len(num_hidden_layers), function(i) { + .qwen3_layer(hidden_size, num_attention_heads, + num_key_value_heads, head_dim, intermediate_size, + eps = rms_norm_eps) + }) + ) + self$norm <- ltx23_rms_norm(hidden_size, eps = rms_norm_eps) + } + ) + self$model <- inner() +}, + forward = function(input_ids, attention_mask = NULL, + out_layers = c(9L, 18L, 27L)) { + x <- self$model$embed_tokens(input_ids) + b <- input_ids$shape[1] + s <- input_ids$shape[2] + device <- x$device + f32 <- torch::torch_float32() + + rope <- .qwen3_rope_tables(s, self$head_dim, self$rope_theta, device) + + # Additive causal (+ padding) mask [B, 1, S, S] in float32 + neg <- -3.4e38 + causal <- torch::torch_full(c(s, s), neg, dtype = f32, + device = device)$triu(diagonal = 1L) + mask <- causal$unsqueeze(1L)$unsqueeze(1L)$expand(c(b, 1L, s, s)) + if (!is.null(attention_mask)) { + pad <- (1 - attention_mask$to(dtype = f32, device = device))$mul(neg) + mask <- mask + pad$unsqueeze(2L)$unsqueeze(2L) + } + + out_layers <- sort(as.integer(out_layers)) + states <- vector("list", length(out_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 + } + } + states +} +) + +#' Load a Qwen3 encoder from a transformers directory +#' +#' Streams the (possibly sharded) safetensors weights into +#' \code{\link{qwen3_encoder}}. The LM head is tied to the embeddings +#' and skipped. +#' +#' @param model_path Directory with \code{config.json} and +#' \code{model*.safetensors} (FLUX.2-klein's \code{text_encoder}). +#' @param device Character. Target device. +#' @param dtype Character. "bfloat16" (GPU) or "float32" (CPU). +#' @param verbose Logical. +#' @param ... Overrides for \code{\link{qwen3_encoder}} arguments. +#' +#' @return The loaded \code{qwen3_encoder} in eval mode. +#' +#' @export +load_qwen3_text_encoder <- function(model_path, device = "cpu", + dtype = "float32", verbose = TRUE, ...) { + model_path <- path.expand(model_path) + config <- NULL + config_path <- file.path(model_path, "config.json") + if (file.exists(config_path)) { + config <- jsonlite::fromJSON(config_path, simplifyVector = TRUE) + } + + args <- list(vocab_size = config$vocab_size, + hidden_size = config$hidden_size, + intermediate_size = config$intermediate_size, + num_hidden_layers = config$num_hidden_layers, + num_attention_heads = config$num_attention_heads, + num_key_value_heads = config$num_key_value_heads, + 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) + } + 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)) + + opened <- .flux_open_sharded_dir(model_path, "model") + ckpt <- structure( + list(handle = opened$handle, keys = opened$keys, + version = NULL, config = config, path = model_path), + class = "ltx23_checkpoint" + ) + map_key <- function(key) { + if (identical(key, "lm_head.weight")) { + return(NA_character_) # tied to embed_tokens + } + key + } + res <- ltx23_load_group(ckpt, ckpt$keys, model, map_key = map_key, + verbose = verbose) + if (length(res$unmapped) || length(res$unfilled)) { + stop("Qwen3 encoder load: ", length(res$unmapped), + " unmapped keys, ", length(res$unfilled), " unfilled params") + } + + model$to(device = device) + model$eval() + model +} + +#' Encode prompts with the Qwen3 encoder for FLUX.2 +#' +#' Tokenizes with the chat template, runs the encoder with the padding +#' mask, and concatenates the requested mid-stack hidden states per +#' token, matching Flux2KleinPipeline._get_qwen3_prompt_embeds. +#' +#' @param prompts Character vector. +#' @param model A \code{\link{qwen3_encoder}}. +#' @param tokenizer A \code{\link{qwen_bpe_tokenizer}}. +#' @param max_sequence_length Integer. Fixed token length (klein: 512). +#' @param out_layers Integer vector. Hidden-state layers (klein-4B: +#' 9, 18, 27). +#' @param device Device for the input ids (defaults to the model's). +#' +#' @return Tensor [length(prompts), max_sequence_length, +#' 3 * hidden_size]. +#' +#' @export +encode_with_qwen3 <- function(prompts, model, tokenizer, + max_sequence_length = 512L, + out_layers = c(9L, 18L, 27L), device = NULL) { + enc <- encode_qwen(tokenizer, prompts, max_length = max_sequence_length, + chat_template = TRUE) + device <- device %||% model$model$embed_tokens$weight$device + long <- torch::torch_long() + ids <- torch::torch_tensor(enc$input_ids + 1L, dtype = long, + device = device) + mask <- torch::torch_tensor(enc$attention_mask, dtype = long, + device = device) + + states <- torch::with_no_grad(model(ids, attention_mask = mask, + out_layers = out_layers)) + # stack(dim=1) + permute + reshape == per-token concatenation + torch::torch_cat(states, dim = -1L) +} diff --git a/R/rope_flux2.R b/R/rope_flux2.R new file mode 100644 index 0000000..db13f15 --- /dev/null +++ b/R/rope_flux2.R @@ -0,0 +1,86 @@ +#' FLUX.2 Position Ids and Empirical Shift +#' +#' Fresh R port of the FLUX.2 position-id builders and the empirical +#' timestep-shift formula from the diffusers reference (Apache-2.0, +#' src/diffusers/pipelines/flux2/pipeline_flux2_klein.py). FLUX.2 uses +#' 4-axis rotary position ids (T, H, W, L): text tokens carry only the +#' L axis (sequence position), image latents carry H and W, and the T +#' axis distinguishes reference images (unused for txt2img). Frequencies +#' come from \code{\link{flux_pos_embed}} with +#' \code{axes_dim = c(32, 32, 32, 32)} and \code{theta = 2000}. +#' +#' @name rope_flux2 +NULL + +#' Build FLUX.2 text position ids +#' +#' Columns (T, H, W, L) with only L varying: 0..len-1. Reference: +#' Flux2KleinPipeline._prepare_text_ids. +#' +#' @param len Integer. Text sequence length. +#' @param device Device for the resulting tensor. +#' +#' @return Float tensor [len, 4]. +#' +#' @export +flux2_prepare_text_ids <- function(len, device = "cpu") { + f32 <- torch::torch_float32() + ids <- torch::torch_zeros(len, 4L, dtype = f32, device = device) + ids[, 4] <- torch::torch_arange(start = 0, end = len - 1, dtype = f32, + device = device) + ids +} + +#' Build FLUX.2 latent position ids +#' +#' Columns (T, H, W, L) with H and W carrying the packed-grid position +#' (row-major: H varies slowest), T = L = 0. Reference: +#' Flux2KleinPipeline._prepare_latent_ids. +#' +#' @param height Integer. Packed grid height (pixel height / 16). +#' @param width Integer. Packed grid width (pixel width / 16). +#' @param device Device for the resulting tensor. +#' +#' @return Float tensor [height * width, 4]. +#' +#' @export +flux2_prepare_latent_ids <- function(height, width, device = "cpu") { + f32 <- torch::torch_float32() + ids <- torch::torch_zeros(height, width, 4L, dtype = f32, device = device) + rows <- torch::torch_arange(start = 0, end = height - 1, dtype = f32, + device = device) + cols <- torch::torch_arange(start = 0, end = width - 1, dtype = f32, + device = device) + ids[,, 2] <- ids[,, 2] + rows$unsqueeze(2L) + ids[,, 3] <- ids[,, 3] + cols$unsqueeze(1L) + ids$reshape(c(height * width, 4L)) +} + +#' Empirical timestep shift for FLUX.2 +#' +#' BFL's piecewise-linear fit of the dynamic-shifting mu as a function of +#' image sequence length and step count; replaces FLUX.1's +#' calculate_shift. Reference: compute_empirical_mu (adapted from BFL +#' sampling.py). +#' +#' @param image_seq_len Integer. Packed image token count. +#' @param num_steps Integer. Inference steps. +#' +#' @return Numeric mu for \code{\link{flowmatch_set_timesteps}}. +#' +#' @export +flux2_empirical_mu <- function(image_seq_len, num_steps) { + a1 <- 8.73809524e-05 + b1 <- 1.89833333 + a2 <- 0.00016927 + b2 <- 0.45666666 + + if (image_seq_len > 4300) { + return(a2 * image_seq_len + b2) + } + m_200 <- a2 * image_seq_len + b2 + m_10 <- a1 * image_seq_len + b1 + a <- (m_200 - m_10) / 190.0 + b <- m_200 - 200.0 * a + a * num_steps + b +} diff --git a/R/tokenizer_qwen.R b/R/tokenizer_qwen.R new file mode 100644 index 0000000..0ab7a9d --- /dev/null +++ b/R/tokenizer_qwen.R @@ -0,0 +1,265 @@ +#' Qwen2 Byte-Level BPE Tokenizer +#' +#' Pure R implementation of the Qwen2 tokenizer (HuggingFace +#' tokenizer.json, BPE model with ByteLevel pre-tokenization), as used by +#' FLUX.2 klein's Qwen3 text encoder. Text is split with the GPT-4-style +#' regex, each pre-token's UTF-8 bytes are mapped through the GPT-2 +#' byte-to-unicode table, and rank-based BPE merges produce the ids. +#' Added tokens (\code{<|im_start|>}, \code{}, ...) are split out +#' literally before byte-level encoding. +#' +#' Limitation: the NFC normalizer is not applied (base R has no NFC); +#' input is assumed to already be NFC, which holds for ordinary text. +#' +#' @name tokenizer_qwen +NULL + +# GPT-2 byte-to-unicode table: printable bytes map to themselves, +# everything else to 256+n. Returns a character vector indexed by +# byte value + 1. +.qwen_byte_table <- function() { + bs <- c(33:126, 161:172, 174:255) + cs <- bs + n <- 0L + for (b in 0:255) { + if (!(b %in% bs)) { + bs <- c(bs, b) + cs <- c(cs, 256L + n) + n <- n + 1L + } + } + out <- character(256) + out[bs + 1L] <- vapply(cs, intToUtf8, character(1)) + out +} + +#' Load a Qwen2 byte-level BPE tokenizer +#' +#' @param tokenizer_path Path to a tokenizer.json (or a directory +#' containing one). +#' +#' @return A \code{qwen_tokenizer} object. +#' +#' @export +qwen_bpe_tokenizer <- function(tokenizer_path) { + path <- path.expand(tokenizer_path) + if (dir.exists(path)) { + path <- file.path(path, "tokenizer.json") + } + if (!file.exists(path)) { + stop("tokenizer.json not found: ", path) + } + + tj <- jsonlite::fromJSON(path, simplifyVector = TRUE) + model <- tj$model + if (!identical(model$type, "BPE")) { + stop("Expected a BPE tokenizer, got: ", model$type %||% "none") + } + + # Integer-id BPE with no persistent string maps: R gc cost scales + # with live object count, and a 151k-binding vocab environment made + # every gc ~17x slower (measured 12 -> 203 ms), which multiplied + # into minutes across the allocator-triggered gcs of a generation. + # All lookups compile to three atomic vectors + a 256-int byte table. + vocab <- unlist(model$vocab) # named int: piece -> id (0-based) + piece_names <- names(vocab) + + merges <- model$merges + if (is.matrix(merges)) { + a <- merges[, 1] + b <- merges[, 2] + } else { + sp <- regexpr(" ", merges, fixed = TRUE) + a <- substr(merges, 1L, sp - 1L) + b <- substr(merges, sp + 1L, nchar(merges)) + } + id_a <- unname(vocab[match(a, piece_names)]) + id_b <- unname(vocab[match(b, piece_names)]) + id_r <- unname(vocab[match(paste0(a, b), piece_names)]) + ok <- !is.na(id_a) & !is.na(id_b) & !is.na(id_r) + key <- id_a[ok] * 2097152 + id_b[ok] # ids < 2^21: exact in doubles + rank <- seq_along(id_a)[ok] + result <- id_r[ok] + ord <- order(key) + + # Initial ids for the 256 byte-unicode characters + byte_chars <- .qwen_byte_table() + byte_ids <- unname(vocab[match(byte_chars, piece_names)]) + + # Pre-tokenization split regex (GPT-4 style) + split_regex <- NULL + pres <- tj$pre_tokenizer$pretokenizers + if (!is.null(pres) && "pattern" %in% names(pres)) { + split_regex <- pres$pattern$Regex[!is.na(pres$pattern$Regex)][1] + } + if (is.null(split_regex) || !nzchar(split_regex)) { + stop("No Split pre-tokenizer regex found in ", path) + } + + added <- tj$added_tokens + added_env <- list2env(stats::setNames(as.list(added$id), added$content), + parent = emptyenv()) + + structure( + list( + merge_key = key[ord], + merge_rank = rank[ord], + merge_result = result[ord], + byte_ids = byte_ids, + n_pieces = length(vocab), + split_regex = split_regex, + added = added_env, + added_contents = added$content[order(-nchar(added$content))], + pad_id = get0("<|endoftext|>", envir = added_env, + ifnotfound = 151643L), + path = path + ), + class = "qwen_tokenizer" + ) +} + +#' @export +print.qwen_tokenizer <- function(x, ...) { + cat("\n") + cat(" vocab: ", x$n_pieces, "+", length(ls(x$added)), "added\n") + cat(" path: ", x$path, "\n") + invisible(x) +} + +# Rank-based BPE merge loop over token ids: adjacent-pair keys are +# looked up in the sorted merge table via findInterval (C binary +# search), vectorized across all pairs per iteration +.qwen_bpe_merge_ids <- function(ids, tokenizer) { + mk <- tokenizer$merge_key + while (length(ids) > 1L) { + keys <- ids[-length(ids)] * 2097152 + ids[-1L] + pos <- findInterval(keys, mk) + hit <- pos > 0L + hit[hit] <- mk[pos[hit]] == keys[hit] + if (!any(hit)) { + break + } + ranks <- rep(Inf, length(keys)) + ranks[hit] <- tokenizer$merge_rank[pos[hit]] + i <- which.min(ranks) + ids[i] <- tokenizer$merge_result[pos[i]] + ids <- ids[-(i + 1L)] + } + ids +} + +# Encode one plain-text segment (no added tokens inside) +.qwen_encode_segment <- function(text, tokenizer) { + if (!nzchar(text)) { + return(integer(0)) + } + text <- enc2utf8(text) + m <- gregexpr(tokenizer$split_regex, text, perl = TRUE)[[1]] + pre_tokens <- regmatches(text, list(m))[[1]] + + out <- lapply(pre_tokens, function(tok) { + bytes <- as.integer(charToRaw(tok)) + .qwen_bpe_merge_ids(tokenizer$byte_ids[bytes + 1L], tokenizer) + }) + as.integer(unlist(out)) +} + +# Split text on added-token literals (longest first), returning a list +# of list(text=, id=) chunks +.qwen_split_added <- function(text, tokenizer) { + chunks <- list(list(text = text, id = NA_integer_)) + for (content in tokenizer$added_contents) { + out <- list() + for (chunk in chunks) { + if (!is.na(chunk$id) || !grepl(content, chunk$text, fixed = TRUE)) { + out[[length(out) + 1L]] <- chunk + next + } + parts <- strsplit(chunk$text, content, fixed = TRUE)[[1]] + # strsplit drops trailing separators; recover the layout + n_seps <- lengths(regmatches(chunk$text, + gregexpr(content, chunk$text, fixed = TRUE))) + if (length(parts) == 0L) { + parts <- "" + } + for (i in seq_along(parts)) { + if (nzchar(parts[i])) { + out[[length(out) + 1L]] <- list(text = parts[i], + id = NA_integer_) + } + if (i <= n_seps) { + out[[length(out) + 1L]] <- list( + text = content, + id = get0(content, envir = tokenizer$added)) + } + } + } + chunks <- out + } + chunks +} + +#' Encode prompts with the Qwen tokenizer +#' +#' With \code{chat_template = TRUE} (the FLUX.2 klein pipeline behavior) +#' each prompt is wrapped as a single user turn with the generation +#' prompt and a disabled thinking block, matching +#' \code{apply_chat_template(..., add_generation_prompt = TRUE, +#' enable_thinking = FALSE)}. Right-pads with \code{<|endoftext|>}. +#' +#' @param tokenizer A \code{\link{qwen_bpe_tokenizer}}. +#' @param texts Character vector of prompts. +#' @param max_length Integer. Fixed sequence length (klein: 512). NULL +#' for no truncation/padding. +#' @param chat_template Logical. Wrap in the Qwen3 chat template. +#' +#' @return List with \code{input_ids} and \code{attention_mask} integer +#' matrices [length(texts), max_length] (ragged lists when +#' \code{max_length} is NULL). Ids are 0-based. +#' +#' @export +encode_qwen <- function(tokenizer, texts, max_length = 512L, + chat_template = TRUE) { + stopifnot(inherits(tokenizer, "qwen_tokenizer")) + + encode_one <- function(text) { + if (chat_template) { + text <- paste0("<|im_start|>user\n", text, "<|im_end|>\n", + "<|im_start|>assistant\n\n\n\n\n") + } + ids <- integer(0) + for (chunk in .qwen_split_added(text, tokenizer)) { + if (!is.na(chunk$id)) { + ids <- c(ids, as.integer(chunk$id)) + } else { + ids <- c(ids, .qwen_encode_segment(chunk$text, tokenizer)) + } + } + ids + } + + all_ids <- lapply(as.character(texts), encode_one) + + if (is.null(max_length)) { + return(list( + input_ids = all_ids, + attention_mask = lapply(all_ids, function(x) rep(1L, length(x))) + )) + } + + all_ids <- lapply(all_ids, function(ids) { + if (length(ids) > max_length) ids[seq_len(max_length)] else ids + }) + n <- length(all_ids) + input_ids <- matrix(as.integer(tokenizer$pad_id), nrow = n, + ncol = max_length) + attention_mask <- matrix(0L, nrow = n, ncol = max_length) + for (i in seq_len(n)) { + len <- length(all_ids[[i]]) + if (len > 0L) { + input_ids[i, seq_len(len)] <- all_ids[[i]] + attention_mask[i, seq_len(len)] <- 1L + } + } + list(input_ids = input_ids, attention_mask = attention_mask) +} diff --git a/R/txt2img.R b/R/txt2img.R index 476e282..2d1742b 100644 --- a/R/txt2img.R +++ b/R/txt2img.R @@ -11,12 +11,14 @@ #' \dontrun{ #' img <- txt2img("a cat wearing sunglasses in space", device = "cuda") #' } -txt2img <- function(prompt, model_name = c("sd21", "sdxl", "flux1"), ...) { +txt2img <- function(prompt, model_name = c("sd21", "sdxl", "flux1", "flux2"), + ...) { switch(model_name, # "sd15" = txt2img_sd15(prompt, ...), "sd21" = txt2img_sd21(prompt, ...), "sdxl" = txt2img_sdxl(prompt, ...), "flux1" = txt2img_flux(prompt, ...), + "flux2" = txt2img_flux2(prompt, ...), # "sd3" = txt2img_sd3(prompt, ...), stop("Unsupported model: ", model_name) ) diff --git a/R/txt2img_flux2.R b/R/txt2img_flux2.R new file mode 100644 index 0000000..30241f1 --- /dev/null +++ b/R/txt2img_flux2.R @@ -0,0 +1,369 @@ +#' FLUX.2 Klein Text-to-Image Pipeline +#' +#' Klein text-to-image, ported from the diffusers reference +#' (Apache-2.0, src/diffusers/pipelines/flux2/pipeline_flux2_klein.py). +#' Klein-4B is step-distilled: no classifier-free guidance and no +#' guidance embedding; the FlowMatch schedule uses dynamic shifting with +#' the BFL empirical mu. Latent noise is drawn directly in the packed +#' 128-channel space; only VAE-encoded latents get the BatchNorm +#' normalization (txt2img never encodes). +#' +#' @name txt2img_flux2 +NULL + +# Resolve a FLUX.2-klein support file from the HuggingFace cache +.flux2_cached <- function(file) { + if (!requireNamespace("hfhub", quietly = TRUE)) { + stop("The hfhub package is required to locate model files.") + } + tryCatch( + hfhub::hub_download(.flux2_repo, file, local_files_only = TRUE), + error = function(e) { + stop("Missing ", file, " in the HuggingFace cache; ", + "run download_flux2_klein() first.", call. = FALSE) + } + ) +} + +#' Load the FLUX.2 klein pipeline +#' +#' Loads the quantized transformer artifact plus the FLUX.2 VAE decoder, +#' Qwen3 text encoder, and tokenizer from the HuggingFace cache +#' populated by \code{\link{download_flux2_klein}}. With fp8 precision +#' the ~4 GB transformer stays GPU-resident. +#' +#' @param model_dir Quantized artifact directory (default: the +#' \code{download_flux2_klein} location for \code{precision}), or a +#' raw diffusers transformer directory. +#' @param device Character. Compute device. +#' @param precision "fp8" (default) or "nf4". +#' @param text_device Device for the Qwen3 encoder (default: +#' \code{device}; it encodes in its own phase and offloads). +#' @param attn_chunk Integer or NULL. Attention query-chunk override. +#' @param phase_offload Logical. One GPU tenant per phase. +#' @param verbose Logical. +#' +#' @return A \code{flux2_pipeline} list. +#' +#' @export +flux2_load_pipeline <- function(model_dir = NULL, device = "cuda", + precision = c("fp8", "nf4"), + text_device = NULL, attn_chunk = NULL, + phase_offload = TRUE, verbose = TRUE) { + precision <- match.arg(precision) + if (is.null(text_device)) { + if (device == "cuda") { + text_device <- "cuda" + } else { + text_device <- "cpu" + } + } + if (is.null(model_dir)) { + model_dir <- file.path(tools::R_user_dir("diffuseR", "data"), + paste0("flux2-klein-4b-", precision)) + } + + ckpt <- if (file.exists(file.path(model_dir, "manifest.json"))) { + flux_open_quantized(model_dir) + } else { + flux_open_checkpoint(model_dir) + } + + if (!nzchar(Sys.getenv("PYTORCH_CUDA_ALLOC_CONF"))) { + # Resident fp8 has an NF4-like stable footprint: the native + # backend avoids expandable_segments' page-unmap cost on the + # per-step activation churn (see ltx23_load_pipeline) + Sys.setenv(PYTORCH_CUDA_ALLOC_CONF = "backend:native") + } + if (device == "cuda") { + # Sized to the LARGEST phase (the 8 GB Qwen3 encode), not the + # transformer: a low footprint puts the allocator's R-gc + # callback threshold under the working set, and every callback + # walks the ~300k-object tokenizer heap (measured: 13-20 s + # forwards at footprint 6 vs sub-second at 12) + ltx23_tune_gc(footprint_gb = 12) + } + + if (phase_offload) { + component_device <- "cpu" + } else { + component_device <- device + } + + pipe <- list( + format = ckpt$format %||% "full", + device = device, + 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 transformer (", pipe$format, ")...") + } + pipe$transformer <- flux_load_transformer( + ckpt, device = component_device, + dtype = if (device == "cpu") "float32" else "bfloat16", + pin = FALSE, + fp8_resident = FALSE, + verbose = verbose + ) + # Resident fp8 happens at onload time (the weights ride to the GPU + # with the phase and back off after) + pipe$fp8_resident <- identical(pipe$format, "fp8") && device == "cuda" + + if (verbose) { + message("Loading FLUX.2 VAE decoder...") + } + vae_config <- jsonlite::fromJSON(.flux2_cached("vae/config.json")) + pipe$vae_bn_eps <- vae_config$batch_norm_eps %||% 1e-4 + pipe$decoder <- load_flux2_vae_decoder( + .flux2_cached("vae/diffusion_pytorch_model.safetensors"), + latent_channels = as.integer(vae_config$latent_channels %||% 32L), + verbose = verbose + ) + pipe$decoder$to(device = component_device) + + if (verbose) { + message("Loading Qwen3 text encoder...") + } + te_dir <- dirname(.flux2_cached("text_encoder/config.json")) + 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(.flux2_cached("tokenizer/tokenizer.json")) + + structure(pipe, class = "flux2_pipeline") +} + +# Flow-matching Euler loop; klein is CFG-free (one forward per step) +.flux2_denoise <- function(transformer, latents, schedule, prompt_embeds, + image_rotary_emb, compute_dtype, + chunk_size = NULL, verbose = TRUE) { + timesteps <- as.numeric(schedule$timesteps$cpu()) + n <- length(timesteps) + pb <- if (verbose) { + utils::txtProgressBar(min = 0, max = n, style = 3) + } else { + NULL + } + f32 <- torch::torch_float32() + + torch::with_no_grad({ + for (i in seq_len(n)) { + t <- timesteps[[i]] + t_model <- torch::torch_tensor(t / 1000, dtype = compute_dtype, + device = latents$device)$reshape(1L) + + noise_pred <- transformer( + hidden_states = latents$to(dtype = compute_dtype), + encoder_hidden_states = prompt_embeds, + timestep = t_model, + image_rotary_emb = image_rotary_emb, + chunk_size = chunk_size + ) + + step <- flowmatch_scheduler_step( + noise_pred$to(dtype = f32), t, latents, schedule + ) + latents <- step$prev_sample + schedule <- step$schedule + rm(noise_pred, step) + if (!is.null(pb)) { + utils::setTxtProgressBar(pb, i) + } + } + }) + if (!is.null(pb)) { + close(pb) + } + latents +} + +#' Generate an image with FLUX.2 klein +#' +#' Step-distilled text-to-image (klein-4B: 4 steps, no guidance): Qwen3 +#' prompt encoding (chat template, mid-stack hidden states), FlowMatch +#' denoising with the empirical dynamic shift, and 32-channel VAE decode +#' through the BatchNorm latent statistics. +#' +#' @param prompt Character. The prompt. +#' @param pipeline A \code{flux2_pipeline} from +#' \code{\link{flux2_load_pipeline}}; NULL loads one (passing +#' \code{...} through). +#' @param width,height Integers, divisible by 16. +#' @param num_inference_steps Integer. Denoising steps (klein-4B: 4). +#' @param max_sequence_length Integer. Qwen3 token length (512). +#' @param seed Integer or NULL. Latents are drawn on the CPU in the +#' packed shape, so a seed matches a Python diffusers run with a CPU +#' generator. +#' @param prompt_embeds Optional precomputed [B, S, 7680] embeddings. +#' @param save_file Logical. Write a PNG. +#' @param filename Output path (default derived from the prompt). +#' @param verbose Logical. +#' @param ... Passed to \code{\link{flux2_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_flux2 <- function(prompt, pipeline = NULL, width = 1024L, + height = 1024L, num_inference_steps = 4L, + max_sequence_length = 512L, seed = NULL, + prompt_embeds = NULL, save_file = TRUE, + filename = NULL, verbose = TRUE, ...) { + if (is.null(pipeline)) { + pipeline <- flux2_load_pipeline(..., verbose = verbose) + } + 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 <- if (device == "cpu") { + f32 + } else { + torch::torch_bfloat16() + } + + phase_offload <- isTRUE(pipeline$phase_offload) && device != "cpu" + onload <- function(module) { + if (phase_offload) { + module$to(device = device) + } + module + } + offload <- function(module) { + if (phase_offload) { + module$to(device = "cpu") + clear_vram() + } + invisible(module) + } + + t0 <- Sys.time() + + # --- Phase 1: text encoding -------------------------------------------------- + if (is.null(prompt_embeds)) { + if (verbose) { + message("Encoding prompt (Qwen3)...") + } + onload(pipeline$text_encoder) + te_device <- pipeline$text_encoder$model$embed_tokens$weight$device + prompt_embeds <- encode_with_qwen3(prompt, pipeline$text_encoder, + pipeline$tokenizer, max_sequence_length = max_sequence_length, + device = te_device) + offload(pipeline$text_encoder) + } + prompt_embeds <- prompt_embeds$to(device = device, dtype = compute_dtype) + txt_len <- prompt_embeds$shape[2] + + # --- Phase 2: latents, rotary embeddings, schedule ----------------------------- + h2 <- height %/% 16L + w2 <- width %/% 16L + if (!is.null(seed)) { + torch::torch_manual_seed(seed) + } + # Noise directly in the packed/patchified space (no BN for txt2img) + latents <- torch::torch_randn(c(1L, 128L, h2, w2), dtype = f32) + latents <- flux2_pack_latents(latents)$to(device = device) + seq_img <- latents$shape[2] + + txt_ids <- flux2_prepare_text_ids(txt_len) + latent_ids <- flux2_prepare_latent_ids(h2, w2) + ids <- torch::torch_cat(list(txt_ids, latent_ids), dim = 1L) + rope <- flux_pos_embed( + ids, + axes_dim = pipeline$transformer$axes_dims_rope %||% c(32L, 32L, 32L, 32L), + theta = pipeline$transformer$rope_theta %||% 2000 + ) + rope <- list(rope[[1]]$to(device = device), rope[[2]]$to(device = device)) + + n_steps <- as.integer(num_inference_steps) + mu <- flux2_empirical_mu(seq_img, n_steps) + sched <- flowmatch_scheduler_create( + use_dynamic_shifting = TRUE, + time_shift_type = "exponential" + ) + sched <- flowmatch_set_timesteps( + sched, n_steps, mu = mu, + sigmas = seq(1, 1 / n_steps, length.out = n_steps) + ) + + # --- Phase 3: denoise ------------------------------------------------------------ + transformer <- onload(pipeline$transformer) + if (isTRUE(pipeline$fp8_resident)) { + .flux_fp8_to_device(transformer, device) + } + if (verbose) { + message(sprintf("Denoising: %d steps at %dx%d...", n_steps, width, + height)) + } + latents <- .flux2_denoise( + transformer, latents, sched, prompt_embeds, rope, + compute_dtype, chunk_size = pipeline$attn_chunk, + verbose = verbose + ) + if (isTRUE(pipeline$fp8_resident) && phase_offload) { + .flux_fp8_to_device(pipeline$transformer, "cpu") + } + offload(pipeline$transformer) + ltx23_release_dequant_buffers() + + # --- Phase 4: decode --------------------------------------------------------------- + if (verbose) { + message("Decoding...") + } + latents <- flux2_unpack_latents_with_ids(latents, latent_ids, h2, w2) + 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 <- pipeline$decoder + if (phase_offload) { + decoder$to(device = device, dtype = compute_dtype) + } + 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) + + 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 = width, height = height, + steps = n_steps, seed = seed, model = "flux2-klein-4b", + precision = pipeline$format, seconds = gen_seconds + ) + invisible(list(image = img_array, metadata = metadata)) +} diff --git a/R/vae_flux2.R b/R/vae_flux2.R new file mode 100644 index 0000000..9d18fed --- /dev/null +++ b/R/vae_flux2.R @@ -0,0 +1,242 @@ +#' FLUX.2 Latent Layout and VAE Helpers +#' +#' Fresh R port of the FLUX.2 latent packing chain from the diffusers +#' reference (Apache-2.0, src/diffusers/pipelines/flux2/ +#' pipeline_flux2_klein.py). The 32-channel VAE latent is patchified +#' 2x2 into 128 channels, normalized with the VAE's BatchNorm running +#' statistics (there is no scalar scaling/shift factor in FLUX.2), and +#' flattened to channels-last tokens for the transformer. +#' +#' @name vae_flux2 +NULL + +#' Patchify FLUX.2 latents (2x2 -> channels) +#' +#' [B, C, H, W] -> [B, 4C, H/2, W/2], channel order (C, ph, pw). +#' Reference: Flux2KleinPipeline._patchify_latents. +#' +#' @param latents Tensor [B, C, H, W]; H and W must be even. +#' +#' @return Tensor [B, C * 4, H / 2, W / 2]. +#' +#' @export +flux2_patchify_latents <- function(latents) { + shape <- latents$shape + b <- shape[1] + ch <- shape[2] + h <- shape[3] + w <- shape[4] + latents <- latents$view(c(b, ch, h %/% 2L, 2L, w %/% 2L, 2L)) + # Python permute (0, 1, 3, 5, 2, 4), 1-indexed here + latents <- latents$permute(c(1L, 2L, 4L, 6L, 3L, 5L)) + latents$reshape(c(b, ch * 4L, h %/% 2L, w %/% 2L)) +} + +#' Unpatchify FLUX.2 latents (channels -> 2x2) +#' +#' Inverse of \code{\link{flux2_patchify_latents}}. Reference: +#' Flux2KleinPipeline._unpatchify_latents. +#' +#' @param latents Tensor [B, 4C, H, W]. +#' +#' @return Tensor [B, C, H * 2, W * 2]. +#' +#' @export +flux2_unpatchify_latents <- function(latents) { + shape <- latents$shape + b <- shape[1] + ch <- shape[2] + h <- shape[3] + w <- shape[4] + latents <- latents$reshape(c(b, ch %/% 4L, 2L, 2L, h, w)) + # Python permute (0, 1, 4, 2, 5, 3), 1-indexed here + latents <- latents$permute(c(1L, 2L, 5L, 3L, 6L, 4L)) + latents$reshape(c(b, ch %/% 4L, h * 2L, w * 2L)) +} + +#' Pack patchified FLUX.2 latents into tokens +#' +#' [B, C, H, W] -> [B, H * W, C] (row-major spatial flatten, +#' channels-last). Reference: Flux2KleinPipeline._pack_latents. +#' +#' @param latents Tensor [B, C, H, W]. +#' +#' @return Tensor [B, H * W, C]. +#' +#' @export +flux2_pack_latents <- function(latents) { + shape <- latents$shape + b <- shape[1] + ch <- shape[2] + latents <- latents$reshape(c(b, ch, shape[3] * shape[4])) + latents$permute(c(1L, 3L, 2L)) +} + +#' Unpack FLUX.2 tokens back to a latent grid via position ids +#' +#' Scatters tokens to (H, W) positions taken from the id columns +#' (H = column 2, W = column 3, 0-based values). Reference: +#' Flux2KleinPipeline._unpack_latents_with_ids. +#' +#' @param x Tensor [B, S, C] of tokens. +#' @param ids Tensor [S, 4] (or [B, S, 4]) of position ids. +#' @param height,width Integers. Packed grid dimensions. +#' +#' @return Tensor [B, C, height, width]. +#' +#' @export +flux2_unpack_latents_with_ids <- function(x, ids, height, width) { + if (ids$ndim == 3L) { + ids <- ids[1,,] + } + ids <- ids$to(device = x$device) + long <- torch::torch_long() + h_ids <- ids[, 2]$to(dtype = long) + w_ids <- ids[, 3]$to(dtype = long) + flat <- h_ids$mul(width)$add(w_ids)$add(1L) # 1-based scatter index + + shape <- x$shape + b <- shape[1] + ch <- shape[3] + out <- torch::torch_zeros(b, height * width, ch, dtype = x$dtype, + device = x$device) + index <- flat$unsqueeze(1L)$unsqueeze(3L)$expand(c(b, -1L, ch)) + out$scatter_(2L, index, x) + out$view(c(b, height, width, ch))$permute(c(1L, 4L, 2L, 3L)) +} + +#' Normalize patchified latents with the VAE BatchNorm statistics +#' +#' FLUX.2 has no scalar scaling/shift factor; latents are standardized +#' per packed channel with the VAE's \code{bn.running_mean} / +#' \code{bn.running_var} (eps 1e-4). Reference: encode/decode paths of +#' Flux2KleinPipeline. +#' +#' @param latents Tensor [B, 128, H, W] (patchified). +#' @param bn_mean,bn_var Float tensors [128]. +#' @param eps Numeric. BatchNorm epsilon. +#' @param inverse Logical. De-normalize (decode path) instead. +#' +#' @return Tensor like \code{latents}. +#' +#' @export +flux2_bn_normalize <- function(latents, bn_mean, bn_var, eps = 1e-4, + inverse = FALSE) { + mean <- bn_mean$view(c(1L, -1L, 1L, 1L))$to(device = latents$device, + dtype = latents$dtype) + std <- bn_var$add(eps)$sqrt()$view(c(1L, -1L, 1L, 1L))$ + to(device = latents$device, dtype = latents$dtype) + if (inverse) { + latents$mul(std)$add(mean) + } else { + latents$sub(mean)$div(std) + } +} + +#' FLUX.2 VAE decoder +#' +#' The AutoencoderKLFlux2 decode path: post_quant_conv (1x1, 32 +#' channels) followed by the standard AutoencoderKL decoder body +#' (reused from \code{\link{vae_decoder_native}}), plus the BatchNorm +#' running statistics used for latent (de)normalization. Reference: +#' src/diffusers/models/autoencoders/autoencoder_kl_flux2.py. +#' +#' @param latent_channels Integer (32 for FLUX.2). +#' @param block_channels Decoder block channels (reversed encoder +#' block_out_channels). +#' @param norm_groups Integer. Group norm groups. +#' +#' @return Module whose forward(z) decodes [B, 32, H, W] latents to +#' [B, 3, 8H, 8W] images; \code{$bn$running_mean} / +#' \code{$bn$running_var} carry the normalization statistics. +#' +#' @export +flux2_vae_decoder <- torch::nn_module( + "flux2_vae_decoder", + initialize = function(latent_channels = 32L, + block_channels = c(512L, 512L, 256L, 128L), + norm_groups = 32L) { + self$post_quant_conv <- torch::nn_conv2d(latent_channels, latent_channels, + kernel_size = 1) + self$decoder <- vae_decoder_native( + latent_channels = latent_channels, + block_channels = block_channels, + norm_groups = norm_groups + ) + bn_stats <- torch::nn_module( + "flux2_bn_stats", + initialize = function(n) { + self$running_mean <- torch::nn_buffer(torch::torch_zeros(n)) + self$running_var <- torch::nn_buffer(torch::torch_ones(n)) + }, + forward = function(x) x + ) + self$bn <- bn_stats(latent_channels * 4L) +}, + forward = function(z) { + self$decoder(self$post_quant_conv(z)) +} +) + +#' Load the FLUX.2 VAE decoder from safetensors +#' +#' Loads the decoder half plus post_quant_conv and the BatchNorm running +#' statistics; encoder and quant_conv keys are skipped (txt2img needs no +#' encoder). +#' +#' @param path Path to the VAE .safetensors file (or a directory +#' containing diffusion_pytorch_model.safetensors). +#' @param latent_channels,block_channels,norm_groups Constructor +#' arguments for \code{\link{flux2_vae_decoder}}. +#' @param verbose Logical. +#' +#' @return The loaded \code{flux2_vae_decoder} in eval mode. +#' +#' @export +load_flux2_vae_decoder <- function(path, latent_channels = 32L, + block_channels = c(512L, 512L, 256L, 128L), + norm_groups = 32L, verbose = TRUE) { + path <- path.expand(path) + if (dir.exists(path)) { + path <- file.path(path, "diffusion_pytorch_model.safetensors") + } + dec <- flux2_vae_decoder(latent_channels = latent_channels, + block_channels = block_channels, + norm_groups = norm_groups) + + handle <- safetensors::safetensors$new(path, framework = "torch") + keys <- setdiff(handle$keys(), "__metadata__") + keep <- keys[!startsWith(keys, "encoder.") & + !startsWith(keys, "quant_conv.") & + keys != "bn.num_batches_tracked"] + + dests <- c(dec$named_parameters(), dec$named_buffers()) + filled <- character(0) + unmapped <- character(0) + torch::with_no_grad({ + for (key in keep) { + dest <- dests[[key]] + if (is.null(dest)) { + unmapped <- c(unmapped, key) + next + } + dest$copy_(handle$get_tensor(key)) + filled <- c(filled, key) + } + }) + unfilled <- setdiff(names(dests), filled) + if (length(unmapped)) { + stop("FLUX.2 VAE load: ", length(unmapped), " unmapped keys, e.g. ", + paste(utils::head(unmapped, 3), collapse = ", ")) + } + if (length(unfilled)) { + stop("FLUX.2 VAE load: ", length(unfilled), + " unfilled params, e.g. ", + paste(utils::head(unfilled, 3), collapse = ", ")) + } + if (verbose) { + message("Loaded ", length(filled), " FLUX.2 VAE tensors from ", path) + } + dec$eval() + dec +} diff --git a/inst/REFERENCES.md b/inst/REFERENCES.md index b0953cd..6c2c651 100644 --- a/inst/REFERENCES.md +++ b/inst/REFERENCES.md @@ -29,6 +29,17 @@ module; this file documents the actual lineage, idea by idea. | NF4/fp8 transformer quantization, cast-set policy, phase offloading, allocator tuning | Same sources as the LTX sections below, applied to the FLUX cast set | | Weights | black-forest-labs/FLUX.1-schnell (Apache-2.0; gated HuggingFace repo, downloaded by the user, never redistributed) | +## FLUX.2 klein + +| What | Source | +|---|---| +| MMDiT transformer (shared modulation, SwiGLU feed-forwards, ViT-22B parallel single blocks, 4-axis RoPE) | Ported from HuggingFace **diffusers** (Apache-2.0): `models/transformers/transformer_flux2.py` | +| Klein pipeline flow: Qwen3 prompt encoding contract (chat template, mid-stack hidden states 9/18/27), 4-channel position ids, latent patchify/pack chain, empirical dynamic-shift mu, BatchNorm latent statistics | diffusers `pipelines/flux2/pipeline_flux2_klein.py` (Apache-2.0); the empirical mu formula originates from BFL `sampling.py` (numeric facts) | +| 32-channel `AutoencoderKLFlux2` decode path (post_quant_conv + standard decoder + BN running stats) | diffusers `models/autoencoders/autoencoder_kl_flux2.py` (Apache-2.0); decoder body shared with the SD/FLUX.1 port | +| Qwen3 encoder (GQA, per-head q/k RMS norms, SwiGLU, RoPE theta 1e6, causal + padding mask) | Ported from HuggingFace **transformers** (Apache-2.0): `models/qwen3/` | +| 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) | + ## Quantization | What | Source | diff --git a/inst/tinytest/fixtures/dit_flux2.safetensors b/inst/tinytest/fixtures/dit_flux2.safetensors new file mode 100644 index 0000000..926c729 Binary files /dev/null and b/inst/tinytest/fixtures/dit_flux2.safetensors differ diff --git a/inst/tinytest/fixtures/flux2_model.safetensors b/inst/tinytest/fixtures/flux2_model.safetensors new file mode 100644 index 0000000..5f58291 Binary files /dev/null and b/inst/tinytest/fixtures/flux2_model.safetensors differ diff --git a/inst/tinytest/fixtures/flux2_tiny_ckpt/config.json b/inst/tinytest/fixtures/flux2_tiny_ckpt/config.json new file mode 100644 index 0000000..e0ed09c --- /dev/null +++ b/inst/tinytest/fixtures/flux2_tiny_ckpt/config.json @@ -0,0 +1,23 @@ +{ + "_class_name": "Flux2Transformer2DModel", + "_diffusers_version": "0.39.0.dev0", + "attention_head_dim": 8, + "axes_dims_rope": [ + 2, + 2, + 2, + 2 + ], + "eps": 1e-06, + "guidance_embeds": false, + "in_channels": 8, + "joint_attention_dim": 24, + "mlp_ratio": 3.0, + "num_attention_heads": 2, + "num_layers": 1, + "num_single_layers": 1, + "out_channels": null, + "patch_size": 1, + "rope_theta": 2000, + "timestep_guidance_channels": 256 +} diff --git a/inst/tinytest/fixtures/flux2_tiny_ckpt/diffusion_pytorch_model-00001-of-00003.safetensors b/inst/tinytest/fixtures/flux2_tiny_ckpt/diffusion_pytorch_model-00001-of-00003.safetensors new file mode 100644 index 0000000..1778c4d Binary files /dev/null and b/inst/tinytest/fixtures/flux2_tiny_ckpt/diffusion_pytorch_model-00001-of-00003.safetensors differ diff --git a/inst/tinytest/fixtures/flux2_tiny_ckpt/diffusion_pytorch_model-00002-of-00003.safetensors b/inst/tinytest/fixtures/flux2_tiny_ckpt/diffusion_pytorch_model-00002-of-00003.safetensors new file mode 100644 index 0000000..38eb278 Binary files /dev/null and b/inst/tinytest/fixtures/flux2_tiny_ckpt/diffusion_pytorch_model-00002-of-00003.safetensors differ diff --git a/inst/tinytest/fixtures/flux2_tiny_ckpt/diffusion_pytorch_model-00003-of-00003.safetensors b/inst/tinytest/fixtures/flux2_tiny_ckpt/diffusion_pytorch_model-00003-of-00003.safetensors new file mode 100644 index 0000000..eed01d7 Binary files /dev/null and b/inst/tinytest/fixtures/flux2_tiny_ckpt/diffusion_pytorch_model-00003-of-00003.safetensors differ diff --git a/inst/tinytest/fixtures/flux2_tiny_ckpt/diffusion_pytorch_model.safetensors.index.json b/inst/tinytest/fixtures/flux2_tiny_ckpt/diffusion_pytorch_model.safetensors.index.json new file mode 100644 index 0000000..71b98b2 --- /dev/null +++ b/inst/tinytest/fixtures/flux2_tiny_ckpt/diffusion_pytorch_model.safetensors.index.json @@ -0,0 +1,36 @@ +{ + "metadata": { + "total_size": 77504 + }, + "weight_map": { + "context_embedder.weight": "diffusion_pytorch_model-00002-of-00003.safetensors", + "double_stream_modulation_img.linear.weight": "diffusion_pytorch_model-00001-of-00003.safetensors", + "double_stream_modulation_txt.linear.weight": "diffusion_pytorch_model-00001-of-00003.safetensors", + "norm_out.linear.weight": "diffusion_pytorch_model-00003-of-00003.safetensors", + "proj_out.weight": "diffusion_pytorch_model-00003-of-00003.safetensors", + "single_stream_modulation.linear.weight": "diffusion_pytorch_model-00002-of-00003.safetensors", + "single_transformer_blocks.0.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00003.safetensors", + "single_transformer_blocks.0.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00003.safetensors", + "single_transformer_blocks.0.attn.to_out.weight": "diffusion_pytorch_model-00003-of-00003.safetensors", + "single_transformer_blocks.0.attn.to_qkv_mlp_proj.weight": "diffusion_pytorch_model-00003-of-00003.safetensors", + "time_guidance_embed.timestep_embedder.linear_1.weight": "diffusion_pytorch_model-00001-of-00003.safetensors", + "time_guidance_embed.timestep_embedder.linear_2.weight": "diffusion_pytorch_model-00001-of-00003.safetensors", + "transformer_blocks.0.attn.add_k_proj.weight": "diffusion_pytorch_model-00002-of-00003.safetensors", + "transformer_blocks.0.attn.add_q_proj.weight": "diffusion_pytorch_model-00002-of-00003.safetensors", + "transformer_blocks.0.attn.add_v_proj.weight": "diffusion_pytorch_model-00002-of-00003.safetensors", + "transformer_blocks.0.attn.norm_added_k.weight": "diffusion_pytorch_model-00002-of-00003.safetensors", + "transformer_blocks.0.attn.norm_added_q.weight": "diffusion_pytorch_model-00002-of-00003.safetensors", + "transformer_blocks.0.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00003.safetensors", + "transformer_blocks.0.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00003.safetensors", + "transformer_blocks.0.attn.to_add_out.weight": "diffusion_pytorch_model-00002-of-00003.safetensors", + "transformer_blocks.0.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00003.safetensors", + "transformer_blocks.0.attn.to_out.0.weight": "diffusion_pytorch_model-00002-of-00003.safetensors", + "transformer_blocks.0.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00003.safetensors", + "transformer_blocks.0.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00003.safetensors", + "transformer_blocks.0.ff.linear_in.weight": "diffusion_pytorch_model-00002-of-00003.safetensors", + "transformer_blocks.0.ff.linear_out.weight": "diffusion_pytorch_model-00002-of-00003.safetensors", + "transformer_blocks.0.ff_context.linear_in.weight": "diffusion_pytorch_model-00002-of-00003.safetensors", + "transformer_blocks.0.ff_context.linear_out.weight": "diffusion_pytorch_model-00003-of-00003.safetensors", + "x_embedder.weight": "diffusion_pytorch_model-00002-of-00003.safetensors" + } +} diff --git a/inst/tinytest/fixtures/qwen3_flux2.safetensors b/inst/tinytest/fixtures/qwen3_flux2.safetensors new file mode 100644 index 0000000..fdbccbd Binary files /dev/null and b/inst/tinytest/fixtures/qwen3_flux2.safetensors differ diff --git a/inst/tinytest/fixtures/qwen3_tiny_ckpt/config.json b/inst/tinytest/fixtures/qwen3_tiny_ckpt/config.json new file mode 100644 index 0000000..9342884 --- /dev/null +++ b/inst/tinytest/fixtures/qwen3_tiny_ckpt/config.json @@ -0,0 +1,39 @@ +{ + "architectures": [ + "Qwen3ForCausalLM" + ], + "attention_bias": false, + "attention_dropout": 0.0, + "bos_token_id": null, + "dtype": "float32", + "eos_token_id": null, + "head_dim": 8, + "hidden_act": "silu", + "hidden_size": 32, + "initializer_range": 0.02, + "intermediate_size": 64, + "layer_types": [ + "full_attention", + "full_attention", + "full_attention", + "full_attention" + ], + "max_position_embeddings": 512, + "max_window_layers": 28, + "model_type": "qwen3", + "num_attention_heads": 4, + "num_hidden_layers": 4, + "num_key_value_heads": 2, + "pad_token_id": null, + "rms_norm_eps": 1e-06, + "rope_parameters": { + "rope_theta": 1000000.0, + "rope_type": "default" + }, + "sliding_window": null, + "tie_word_embeddings": true, + "transformers_version": "5.13.0", + "use_cache": true, + "use_sliding_window": false, + "vocab_size": 100 +} diff --git a/inst/tinytest/fixtures/qwen3_tiny_ckpt/generation_config.json b/inst/tinytest/fixtures/qwen3_tiny_ckpt/generation_config.json new file mode 100644 index 0000000..2299f67 --- /dev/null +++ b/inst/tinytest/fixtures/qwen3_tiny_ckpt/generation_config.json @@ -0,0 +1,7 @@ +{ + "_from_model_config": true, + "output_attentions": false, + "output_hidden_states": false, + "transformers_version": "5.13.0", + "use_cache": true +} diff --git a/inst/tinytest/fixtures/qwen3_tiny_ckpt/model.safetensors b/inst/tinytest/fixtures/qwen3_tiny_ckpt/model.safetensors new file mode 100644 index 0000000..1ddd99a Binary files /dev/null and b/inst/tinytest/fixtures/qwen3_tiny_ckpt/model.safetensors differ diff --git a/inst/tinytest/fixtures/qwen_tokenizer_cases.json b/inst/tinytest/fixtures/qwen_tokenizer_cases.json new file mode 100644 index 0000000..879377f --- /dev/null +++ b/inst/tinytest/fixtures/qwen_tokenizer_cases.json @@ -0,0 +1,777 @@ +{ + "cases": [ + { + "text": "a photo of a cat", + "ids": [ + 64, + 6548, + 315, + 264, + 8251 + ] + }, + { + "text": "A sunset over mountains, ultra detailed, 8k", + "ids": [ + 32, + 42984, + 916, + 23501, + 11, + 23998, + 11682, + 11, + 220, + 23, + 74 + ] + }, + { + "text": "Hello, world!", + "ids": [ + 9707, + 11, + 1879, + 0 + ] + }, + { + "text": "The quick brown fox jumps over the lazy dog.", + "ids": [ + 785, + 3974, + 13876, + 38835, + 34208, + 916, + 279, + 15678, + 5562, + 13 + ] + }, + { + "text": "it's a beautiful day; isn't it?", + "ids": [ + 275, + 594, + 264, + 6233, + 1899, + 26, + 4436, + 944, + 432, + 30 + ] + }, + { + "text": "3.14159 and 2,000,000 dollars", + "ids": [ + 18, + 13, + 16, + 19, + 16, + 20, + 24, + 323, + 220, + 17, + 11, + 15, + 15, + 15, + 11, + 15, + 15, + 15, + 11192 + ] + }, + { + "text": "state-of-the-art text-to-image generation", + "ids": [ + 2454, + 8668, + 10603, + 37821, + 1467, + 4686, + 13746, + 9471 + ] + }, + { + "text": " leading spaces", + "ids": [ + 220, + 6388, + 12621 + ] + }, + { + "text": "trailing spaces ", + "ids": [ + 376, + 14277, + 12621, + 262 + ] + }, + { + "text": "double and triple spaces", + "ids": [ + 4331, + 220, + 323, + 256, + 23725, + 256, + 12621 + ] + }, + { + "text": "UPPERCASE lowercase MiXeD", + "ids": [ + 3124, + 9654, + 40371, + 42047, + 20740, + 55, + 68, + 35 + ] + }, + { + "text": "email@example.com and https://example.org/path?q=1", + "ids": [ + 2332, + 35487, + 905, + 323, + 3703, + 1110, + 8687, + 2659, + 50976, + 43782, + 28, + 16 + ] + }, + { + "text": "quotes \"double\" and 'single'", + "ids": [ + 53282, + 330, + 4331, + 1, + 323, + 364, + 15338, + 6 + ] + }, + { + "text": "(parentheses) [brackets] {braces}", + "ids": [ + 49622, + 38322, + 8, + 508, + 1323, + 18382, + 60, + 314, + 1323, + 2434, + 92 + ] + }, + { + "text": "underscores_and_snake_case", + "ids": [ + 31009, + 7701, + 8378, + 28022, + 726, + 19096 + ] + }, + { + "text": "a", + "ids": [ + 64 + ] + }, + { + "text": "", + "ids": [] + }, + { + "text": "café résumé naïve", + "ids": [ + 924, + 58858, + 9333, + 1242, + 963, + 94880, + 586 + ] + }, + { + "text": "em—dash and en–dash", + "ids": [ + 336, + 2293, + 43519, + 323, + 662, + 4142, + 43519 + ] + }, + { + "text": "100% of $50 + €20", + "ids": [ + 16, + 15, + 15, + 4, + 315, + 400, + 20, + 15, + 488, + 12984, + 17, + 15 + ] + }, + { + "text": "emoji 🦊 and 中文字符 mixed in", + "ids": [ + 37523, + 11162, + 99, + 232, + 323, + 72858, + 16744, + 48391, + 9519, + 304 + ] + }, + { + "text": "newline\nand\ttab", + "ids": [ + 89202, + 198, + 437, + 58149 + ] + }, + { + "text": "An astronaut riding a horse on Mars, photorealistic", + "ids": [ + 2082, + 46633, + 19837, + 264, + 15223, + 389, + 21048, + 11, + 4503, + 89768, + 4532 + ] + }, + { + "text": "watercolor painting of a fox in a snowy forest", + "ids": [ + 12987, + 3423, + 18824, + 315, + 264, + 38835, + 304, + 264, + 89773, + 13638 + ] + }, + { + "text": "The transformer architecture uses self-attention mechanisms to model long-range dependencies. The transformer architecture uses self-attention mechanisms to model long-range dependencies. The transformer architecture uses self-attention mechanisms to model long-range dependencies. ", + "ids": [ + 785, + 42578, + 17646, + 5711, + 656, + 12, + 53103, + 23783, + 311, + 1614, + 1293, + 30508, + 19543, + 13, + 576, + 42578, + 17646, + 5711, + 656, + 12, + 53103, + 23783, + 311, + 1614, + 1293, + 30508, + 19543, + 13, + 576, + 42578, + 17646, + 5711, + 656, + 12, + 53103, + 23783, + 311, + 1614, + 1293, + 30508, + 19543, + 13, + 220 + ] + } + ], + "templated": [ + { + "text": "a photo of a cat", + "rendered": "<|im_start|>user\na photo of a cat<|im_end|>\n<|im_start|>assistant\n\n\n\n\n", + "max_length": 64, + "ids": [ + 151644, + 872, + 198, + 64, + 6548, + 315, + 264, + 8251, + 151645, + 198, + 151644, + 77091, + 198, + 151667, + 271, + 151668, + 271, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643 + ], + "mask": [ + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0 + ] + }, + { + "text": "emoji 🦊 and 中文字符 mixed in", + "rendered": "<|im_start|>user\nemoji 🦊 and 中文字符 mixed in<|im_end|>\n<|im_start|>assistant\n\n\n\n\n", + "max_length": 64, + "ids": [ + 151644, + 872, + 198, + 37523, + 11162, + 99, + 232, + 323, + 72858, + 16744, + 48391, + 9519, + 304, + 151645, + 198, + 151644, + 77091, + 198, + 151667, + 271, + 151668, + 271, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643 + ], + "mask": [ + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0 + ] + }, + { + "text": "The transformer architecture uses self-attention mechanisms to model long-range dependencies. The transformer architecture uses self-attention mechanisms to model long-range dependencies. The transformer architecture uses self-attention mechanisms to model long-range dependencies. ", + "rendered": "<|im_start|>user\nThe transformer architecture uses self-attention mechanisms to model long-range dependencies. The transformer architecture uses self-attention mechanisms to model long-range dependencies. The transformer architecture uses self-attention mechanisms to model long-range dependencies. <|im_end|>\n<|im_start|>assistant\n\n\n\n\n", + "max_length": 64, + "ids": [ + 151644, + 872, + 198, + 785, + 42578, + 17646, + 5711, + 656, + 12, + 53103, + 23783, + 311, + 1614, + 1293, + 30508, + 19543, + 13, + 576, + 42578, + 17646, + 5711, + 656, + 12, + 53103, + 23783, + 311, + 1614, + 1293, + 30508, + 19543, + 13, + 576, + 42578, + 17646, + 5711, + 656, + 12, + 53103, + 23783, + 311, + 1614, + 1293, + 30508, + 19543, + 13, + 220, + 151645, + 198, + 151644, + 77091, + 198, + 151667, + 271, + 151668, + 271, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643 + ], + "mask": [ + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0 + ] + } + ], + "meta": { + "pad_token": "<|endoftext|>", + "pad_token_id": 151643, + "padding_side": "right" + } +} \ No newline at end of file diff --git a/inst/tinytest/fixtures/rope_flux2.safetensors b/inst/tinytest/fixtures/rope_flux2.safetensors new file mode 100644 index 0000000..4f40633 Binary files /dev/null and b/inst/tinytest/fixtures/rope_flux2.safetensors differ diff --git a/inst/tinytest/fixtures/vae_flux2_io.safetensors b/inst/tinytest/fixtures/vae_flux2_io.safetensors new file mode 100644 index 0000000..eb533c9 Binary files /dev/null and b/inst/tinytest/fixtures/vae_flux2_io.safetensors differ diff --git a/inst/tinytest/fixtures/vae_flux2_tiny.safetensors b/inst/tinytest/fixtures/vae_flux2_tiny.safetensors new file mode 100644 index 0000000..7d23b30 Binary files /dev/null and b/inst/tinytest/fixtures/vae_flux2_tiny.safetensors differ diff --git a/inst/tinytest/test_dit_flux2.R b/inst/tinytest/test_dit_flux2.R new file mode 100644 index 0000000..b630267 --- /dev/null +++ b/inst/tinytest/test_dit_flux2.R @@ -0,0 +1,98 @@ +# Parity tests for the FLUX.2 transformer blocks against diffusers +# reference fixtures (generated by tools/gen_fixtures_flux2_dit.py). + +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", "dit_flux2.safetensors", + package = "diffuseR") +if (fixture_path == "") fixture_path <- "fixtures/dit_flux2.safetensors" +if (!file.exists(fixture_path)) exit_file("flux2 dit 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()) + ))) +} + +load_named_weights <- function(module, weights) { + dests <- c(module$named_parameters(), module$named_buffers()) + missing_dest <- setdiff(names(weights), names(dests)) + if (length(missing_dest)) { + stop("No destination for: ", paste(utils::head(missing_dest, 5), collapse = ", ")) + } + unfilled <- setdiff(names(dests), names(weights)) + if (length(unfilled)) { + stop("Unfilled params: ", paste(utils::head(unfilled, 5), collapse = ", ")) + } + torch::with_no_grad({ + for (name in names(weights)) dests[[name]]$copy_(weights[[name]]) + }) + invisible(module) +} + +fixture_group <- function(prefix) { + keys <- grep(paste0("^", prefix, "\\."), names(fx), value = TRUE) + w <- fx[keys] + names(w) <- sub(paste0("^", prefix, "\\."), "", keys) + w +} + +rope <- list(fx$rope_cos, fx$rope_sin) +DIM <- 16L +HEADS <- 2L +HEAD_DIM <- 8L + +# --- flux2_modulation -------------------------------------------------------------- + +mod2 <- flux2_modulation(DIM, mod_param_sets = 2L) +mod2$eval() +load_named_weights(mod2, fixture_group("mod2")) +expect_true(max_abs_diff(torch::with_no_grad(mod2(fx$temb)), fx$mod2_out) < 1e-5) + +mod1 <- flux2_modulation(DIM, mod_param_sets = 1L) +mod1$eval() +load_named_weights(mod1, fixture_group("mod1")) +expect_true(max_abs_diff(torch::with_no_grad(mod1(fx$temb)), fx$mod1_out) < 1e-5) + +# --- flux2_feed_forward (SwiGLU) ------------------------------------------------------ + +ff <- flux2_feed_forward(DIM, DIM, mult = 3.0) +ff$eval() +load_named_weights(ff, fixture_group("ff")) +expect_true(max_abs_diff(torch::with_no_grad(ff(fx$joint_x)), fx$ff_out) < 1e-5) + +# --- flux2_double_block ---------------------------------------------------------------- + +dbl <- flux2_double_block(DIM, HEADS, HEAD_DIM, mlp_ratio = 3.0) +dbl$eval() +load_named_weights(dbl, fixture_group("dbl")) +db <- torch::with_no_grad(dbl( + hidden_states = fx$img_x, + encoder_hidden_states = fx$txt_x, + temb_mod_img = fx$dbl_mod_img, + temb_mod_txt = fx$dbl_mod_txt, + image_rotary_emb = rope +)) +expect_true(max_abs_diff(db[[1]], fx$dbl_enc_out) < 1e-5) +expect_true(max_abs_diff(db[[2]], fx$dbl_hid_out) < 1e-5) + +# --- flux2_single_block ------------------------------------------------------------------ + +sgl <- flux2_single_block(DIM, HEADS, HEAD_DIM, mlp_ratio = 3.0) +sgl$eval() +load_named_weights(sgl, fixture_group("sgl")) +sg <- torch::with_no_grad(sgl( + hidden_states = fx$joint_x, + temb_mod = fx$sgl_mod, + image_rotary_emb = rope +)) +expect_true(max_abs_diff(sg, fx$sgl_out) < 1e-5) diff --git a/inst/tinytest/test_flux2_transformer.R b/inst/tinytest/test_flux2_transformer.R new file mode 100644 index 0000000..7f5f6db --- /dev/null +++ b/inst/tinytest/test_flux2_transformer.R @@ -0,0 +1,65 @@ +# Parity test for the full FLUX.2 transformer against a tiny random-init +# diffusers Flux2Transformer2DModel (fixture generated by +# tools/gen_fixtures_flux2_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", "flux2_model.safetensors", + package = "diffuseR") +if (fixture_path == "") fixture_path <- "fixtures/flux2_model.safetensors" +if (!file.exists(fixture_path)) exit_file("flux2 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 <- flux2_transformer( + in_channels = 8L, + num_layers = 1L, + num_single_layers = 1L, + attention_head_dim = 8L, + num_attention_heads = 2L, + joint_attention_dim = 24L, + mlp_ratio = 3.0, + axes_dims_rope = c(2L, 2L, 2L, 2L), + rope_theta = 2000 +) +model$eval() + +# Strict both ways: proves the R module tree matches the diffusers state +# dict exactly +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]]) +}) + +ids <- torch::torch_cat(list(fx$txt_ids, fx$img_ids), dim = 1L) +rope <- flux_pos_embed(ids, axes_dim = c(2L, 2L, 2L, 2L), theta = 2000) + +out <- torch::with_no_grad(model( + hidden_states = fx$hidden, + encoder_hidden_states = fx$encoder, + timestep = fx$timestep, + image_rotary_emb = rope +)) +expect_equal(as.integer(out$shape), as.integer(fx$out$shape)) +expect_true(max_abs_diff(out, fx$out) < 1e-4) + +# Position-id builders match the fixture ids exactly +expect_true(max_abs_diff(flux2_prepare_text_ids(7L), fx$txt_ids) == 0) +expect_true(max_abs_diff(flux2_prepare_latent_ids(6L, 10L), fx$img_ids) == 0) diff --git a/inst/tinytest/test_quantize_flux2.R b/inst/tinytest/test_quantize_flux2.R new file mode 100644 index 0000000..162f7e9 --- /dev/null +++ b/inst/tinytest/test_quantize_flux2.R @@ -0,0 +1,120 @@ +# FLUX.2 quantization round trip on the tiny sharded checkpoint: +# family auto-detection via config _class_name, cast census, NF4 and +# fp8 (streamed + resident) loads. Everything runs on CPU. + +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) + +ckpt_dir <- system.file("tinytest", "fixtures", "flux2_tiny_ckpt", + package = "diffuseR") +if (ckpt_dir == "") ckpt_dir <- "fixtures/flux2_tiny_ckpt" +if (!dir.exists(ckpt_dir)) exit_file("flux2 tiny checkpoint missing") + +fixture_path <- system.file("tinytest", "fixtures", "flux2_model.safetensors", + package = "diffuseR") +if (fixture_path == "") fixture_path <- "fixtures/flux2_model.safetensors" +if (!file.exists(fixture_path)) exit_file("flux2 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()) + ))) +} +cosine_sim <- function(a, b) { + a <- a$to(dtype = torch::torch_float32())$flatten() + b <- b$to(dtype = torch::torch_float32())$flatten() + as.numeric(torch::torch_dot(a, b) / (a$norm() * b$norm())) +} + +# --- family detection + cast census ------------------------------------------------ +# Tiny config: 1 double x 12 + 1 single x 2 + 3 modulations + +# context_embedder = 18 cast weights. (Full klein-4B: 104.) + +ckpt <- flux_open_checkpoint(ckpt_dir) +expect_equal(diffuseR:::.flux_family(ckpt$config), "flux2") +expect_equal(sum(flux2_is_quant_key(ckpt$keys)), 18L) + +# --- full-precision load through family dispatch ------------------------------------ + +model <- flux_load_transformer(ckpt, device = "cpu", dtype = "float32", + verbose = FALSE) +ids <- torch::torch_cat(list(fx$txt_ids, fx$img_ids), dim = 1L) +rope <- flux_pos_embed(ids, axes_dim = c(2L, 2L, 2L, 2L), theta = 2000) +out_full <- torch::with_no_grad(model( + hidden_states = fx$hidden, + encoder_hidden_states = fx$encoder, + timestep = fx$timestep, + image_rotary_emb = rope +)) +expect_true(max_abs_diff(out_full, fx$out) < 1e-4) + +# --- NF4 round trip ------------------------------------------------------------------ + +nf4_dir <- file.path(tempdir(), "flux2-tiny-nf4") +unlink(nf4_dir, recursive = TRUE) +manifest <- flux_quantize(ckpt_dir, output_dir = nf4_dir, format = "nf4", + verbose = FALSE) +expect_equal(manifest$cast, 18L) +expect_true(grepl("^flux2-klein-nf4", manifest$shards[[1]])) + +model_nf4 <- flux_load_transformer(flux_open_quantized(nf4_dir), + device = "cpu", verbose = FALSE) +out_nf4 <- torch::with_no_grad(model_nf4( + hidden_states = fx$hidden$to(dtype = torch::torch_bfloat16()), + encoder_hidden_states = fx$encoder$to(dtype = torch::torch_bfloat16()), + timestep = fx$timestep, + image_rotary_emb = rope +)) +expect_true(cosine_sim(out_nf4, out_full) > 0.98) +ltx23_release_dequant_buffers() + +# --- fp8 round trips (streamed and resident) ------------------------------------------ + +f8_ok <- tryCatch({ + x <- torch::torch_randn(2, 2)$to(dtype = torch::torch_float8_e4m3fn()) + tmp <- tempfile(fileext = ".safetensors") + safetensors::safe_save_file(list(w = x), tmp) + y <- safetensors::safe_load_file(tmp, framework = "torch") + unlink(tmp) + TRUE +}, error = function(e) FALSE) + +if (f8_ok) { + fp8_dir <- file.path(tempdir(), "flux2-tiny-fp8") + unlink(fp8_dir, recursive = TRUE) + manifest8 <- flux_quantize(ckpt_dir, output_dir = fp8_dir, format = "fp8", + verbose = FALSE) + expect_equal(manifest8$cast, 18L) + + model_fp8 <- flux_load_transformer(flux_open_quantized(fp8_dir), + device = "cpu", pin = FALSE, verbose = FALSE) + out_fp8 <- torch::with_no_grad(model_fp8( + hidden_states = fx$hidden$to(dtype = torch::torch_bfloat16()), + encoder_hidden_states = fx$encoder$to(dtype = torch::torch_bfloat16()), + timestep = fx$timestep, + image_rotary_emb = rope + )) + expect_true(cosine_sim(out_fp8, out_full) > 0.99) + + # Resident variant: weights moved to the compute device (CPU here just + # exercises the walker), forward unchanged + model_res <- flux_load_transformer(flux_open_quantized(fp8_dir), + device = "cpu", pin = FALSE, fp8_resident = TRUE, verbose = FALSE) + out_res <- torch::with_no_grad(model_res( + hidden_states = fx$hidden$to(dtype = torch::torch_bfloat16()), + encoder_hidden_states = fx$encoder$to(dtype = torch::torch_bfloat16()), + timestep = fx$timestep, + image_rotary_emb = rope + )) + expect_true(max_abs_diff(out_res, out_fp8) == 0) +} + +options(diffuseR.block_gc = NULL) diff --git a/inst/tinytest/test_qwen3_flux2.R b/inst/tinytest/test_qwen3_flux2.R new file mode 100644 index 0000000..c08352a --- /dev/null +++ b/inst/tinytest/test_qwen3_flux2.R @@ -0,0 +1,51 @@ +# Parity tests for the Qwen3 encoder port against HF transformers +# reference fixtures (generated by tools/gen_fixtures_qwen3.py). + +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", "qwen3_flux2.safetensors", + package = "diffuseR") +if (fixture_path == "") fixture_path <- "fixtures/qwen3_flux2.safetensors" +if (!file.exists(fixture_path)) exit_file("qwen3 fixtures missing") + +ckpt_dir <- system.file("tinytest", "fixtures", "qwen3_tiny_ckpt", + package = "diffuseR") +if (ckpt_dir == "") ckpt_dir <- "fixtures/qwen3_tiny_ckpt" +if (!dir.exists(ckpt_dir)) exit_file("qwen3 tiny checkpoint 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 <- load_qwen3_text_encoder(ckpt_dir, device = "cpu", + dtype = "float32", verbose = FALSE) + +ids <- fx$input_ids$to(dtype = torch::torch_long()) + 1L +mask <- fx$attention_mask$to(dtype = torch::torch_long()) + +states <- torch::with_no_grad(model(ids, attention_mask = mask, + out_layers = c(1L, 2L, 3L))) +expect_equal(as.integer(states[[1]]$shape), as.integer(fx$h1$shape)) +expect_true(max_abs_diff(states[[1]], fx$h1) < 1e-4) +expect_true(max_abs_diff(states[[2]], fx$h2) < 1e-4) +expect_true(max_abs_diff(states[[3]], fx$h3) < 1e-4) + +# Per-token concatenation matches the pipeline's stack/permute/reshape +stacked <- torch::torch_cat(states, dim = -1L) +expect_equal(as.integer(stacked$shape), as.integer(fx$stacked$shape)) +expect_true(max_abs_diff(stacked, fx$stacked) < 1e-4) + +# The mask matters: dropping it must change the padded batch's states +states_nomask <- torch::with_no_grad(model(ids, out_layers = c(3L))) +expect_true(max_abs_diff(states_nomask[[1]], fx$h3) > 1e-4) diff --git a/inst/tinytest/test_rope_flux2.R b/inst/tinytest/test_rope_flux2.R new file mode 100644 index 0000000..96c294d --- /dev/null +++ b/inst/tinytest/test_rope_flux2.R @@ -0,0 +1,75 @@ +# Parity tests for the FLUX.2 position ids, 4-axis RoPE, latent +# patchify/pack/unpack chain, and empirical mu against diffusers +# reference fixtures (generated by tools/gen_fixtures_flux2.py). + +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", "rope_flux2.safetensors", + package = "diffuseR") +if (fixture_path == "") fixture_path <- "fixtures/rope_flux2.safetensors" +if (!file.exists(fixture_path)) exit_file("flux2 rope 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()) + ))) +} + +# --- position ids --------------------------------------------------------------- + +txt_ids <- flux2_prepare_text_ids(7L) +expect_equal(as.integer(txt_ids$shape), as.integer(fx$txt_ids$shape)) +expect_true(max_abs_diff(txt_ids, fx$txt_ids) == 0) + +lat_ids <- flux2_prepare_latent_ids(6L, 10L) +expect_equal(as.integer(lat_ids$shape), as.integer(fx$latent_ids$shape)) +expect_true(max_abs_diff(lat_ids, fx$latent_ids) == 0) + +# --- 4-axis RoPE through flux_pos_embed (theta 2000) --------------------------- + +pf <- flux_pos_embed(fx$ids, axes_dim = c(32L, 32L, 32L, 32L), theta = 2000) +expect_equal(as.integer(pf[[1]]$shape), as.integer(fx$pos_full_cos$shape)) +expect_true(max_abs_diff(pf[[1]], fx$pos_full_cos) < 1e-6) +expect_true(max_abs_diff(pf[[2]], fx$pos_full_sin) < 1e-6) + +pt <- flux_pos_embed(fx$ids, axes_dim = c(2L, 2L, 2L, 2L), theta = 2000) +expect_true(max_abs_diff(pt[[1]], fx$pos_tiny_cos) < 1e-6) +expect_true(max_abs_diff(pt[[2]], fx$pos_tiny_sin) < 1e-6) + +# --- patchify / pack / unpack chain ---------------------------------------------- + +patched <- flux2_patchify_latents(fx$z32) +expect_equal(as.integer(patched$shape), as.integer(fx$patched$shape)) +expect_true(max_abs_diff(patched, fx$patched) == 0) + +expect_true(max_abs_diff(flux2_unpatchify_latents(fx$patched), fx$unpatched) == 0) +expect_true(max_abs_diff(flux2_unpatchify_latents(fx$patched), fx$z32) == 0) + +packed <- flux2_pack_latents(fx$patched) +expect_equal(as.integer(packed$shape), as.integer(fx$packed$shape)) +expect_true(max_abs_diff(packed, fx$packed) == 0) + +unpacked <- flux2_unpack_latents_with_ids(fx$packed, + flux2_prepare_latent_ids(6L, 10L), height = 6L, width = 10L) +expect_equal(as.integer(unpacked$shape), as.integer(fx$unpacked$shape)) +expect_true(max_abs_diff(unpacked, fx$unpacked) == 0) +expect_true(max_abs_diff(unpacked, fx$patched) == 0) + +# --- empirical mu ------------------------------------------------------------------ + +seqs <- as.numeric(fx$mu_seq) +steps <- as.numeric(fx$mu_steps) +vals <- as.numeric(fx$mu_vals) +for (i in seq_along(seqs)) { + expect_true(abs(flux2_empirical_mu(seqs[i], steps[i]) - vals[i]) < 1e-12, + info = sprintf("seq %d steps %d", seqs[i], steps[i])) +} diff --git a/inst/tinytest/test_tokenizer_qwen.R b/inst/tinytest/test_tokenizer_qwen.R new file mode 100644 index 0000000..e3aa5a7 --- /dev/null +++ b/inst/tinytest/test_tokenizer_qwen.R @@ -0,0 +1,71 @@ +# Exact token-id parity for the Qwen2 byte-level BPE tokenizer against +# HuggingFace reference cases (generated by +# tools/gen_qwen_tokenizer_cases.py, checked in). The 11 MB +# tokenizer.json is not shipped; the test locates one and skips if +# absent. + +library(diffuseR) + +find_qwen_tokenizer <- function() { + p <- Sys.getenv("DIFFUSER_QWEN_TOKENIZER", "") + if (nzchar(p) && file.exists(p)) { + return(p) + } + if (requireNamespace("hfhub", quietly = TRUE)) { + p <- tryCatch( + suppressMessages(hfhub::hub_download( + "black-forest-labs/FLUX.2-klein-4B", + "tokenizer/tokenizer.json", local_files_only = TRUE + )), + error = function(e) "" + ) + if (nzchar(p) && file.exists(p)) { + return(p) + } + } + p <- "../../tools/cache/tokenizer_qwen.json" + if (file.exists(p)) { + return(p) + } + "" +} + +tok_path <- find_qwen_tokenizer() +if (!nzchar(tok_path)) exit_file("no Qwen tokenizer.json available") + +cases_path <- system.file("tinytest", "fixtures", "qwen_tokenizer_cases.json", + package = "diffuseR") +if (cases_path == "") cases_path <- "fixtures/qwen_tokenizer_cases.json" +if (!file.exists(cases_path)) exit_file("qwen tokenizer cases missing") + +cases <- jsonlite::fromJSON(cases_path, simplifyVector = FALSE) + +tok <- qwen_bpe_tokenizer(tok_path) +expect_equal(tok$pad_id, 151643L) + +# --- raw encodings (no template, no padding) -------------------------------------- + +for (case in cases$cases) { + expected <- as.integer(unlist(case$ids)) + got <- encode_qwen(tok, case$text, max_length = NULL, + chat_template = FALSE)$input_ids[[1]] + expect_equal(got, expected, + info = sprintf("text: %s", substr(case$text, 1, 60))) +} + +# --- chat template + padding/truncation ---------------------------------------------- + +for (case in cases$templated) { + got <- encode_qwen(tok, case$text, max_length = case$max_length, + chat_template = TRUE) + expect_equal(as.integer(got$input_ids[1, ]), as.integer(unlist(case$ids)), + info = sprintf("templated: %s", substr(case$text, 1, 60))) + expect_equal(as.integer(got$attention_mask[1, ]), + as.integer(unlist(case$mask))) +} + +# --- batch shape ------------------------------------------------------------------------ + +batch <- encode_qwen(tok, c("a photo of a cat", "hello"), max_length = 32L) +expect_equal(dim(batch$input_ids), c(2L, 32L)) +expect_equal(dim(batch$attention_mask), c(2L, 32L)) diff --git a/inst/tinytest/test_txt2img_flux2.R b/inst/tinytest/test_txt2img_flux2.R new file mode 100644 index 0000000..aee5cb4 --- /dev/null +++ b/inst/tinytest/test_txt2img_flux2.R @@ -0,0 +1,64 @@ +# End-to-end smoke test for the FLUX.2 klein pipeline wiring on the CPU +# with tiny random-init components: pack -> denoise (2 steps, dynamic +# shifting) -> unpack-with-ids -> BN denorm -> unpatchify -> decode. + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} + +library(diffuseR) + +torch::torch_manual_seed(8) + +transformer <- flux2_transformer( + in_channels = 128L, + num_layers = 1L, + num_single_layers = 1L, + attention_head_dim = 8L, + num_attention_heads = 2L, + joint_attention_dim = 24L, + mlp_ratio = 3.0, + axes_dims_rope = c(2L, 2L, 2L, 2L), + rope_theta = 2000 +) +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, + config = list(in_channels = 128L), + vae_bn_eps = 1e-4 + ), + class = "flux2_pipeline" +) + +res <- txt2img_flux2( + "tiny smoke test", + pipeline = pipeline, + width = 64L, height = 64L, + num_inference_steps = 2L, + seed = 42L, + prompt_embeds = torch::torch_randn(1L, 7L, 24L), + save_file = FALSE, + verbose = FALSE +) + +expect_equal(dim(res$image), c(64L, 64L, 3L)) +expect_true(all(is.finite(res$image))) +expect_true(all(res$image >= 0) && all(res$image <= 1)) +expect_equal(res$metadata$steps, 2L) +expect_equal(res$metadata$model, "flux2-klein-4b") diff --git a/inst/tinytest/test_vae_flux2.R b/inst/tinytest/test_vae_flux2.R new file mode 100644 index 0000000..8b11556 --- /dev/null +++ b/inst/tinytest/test_vae_flux2.R @@ -0,0 +1,51 @@ +# Parity tests for the FLUX.2 VAE decoder (post_quant_conv + decoder +# body + BatchNorm latent statistics) against diffusers reference +# fixtures (generated by tools/gen_fixtures_flux2_vae.py). + +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 <- function(name) { + p <- system.file("tinytest", "fixtures", name, package = "diffuseR") + if (p == "") p <- file.path("fixtures", name) + p +} +vae_path <- fixture("vae_flux2_tiny.safetensors") +io_path <- fixture("vae_flux2_io.safetensors") +if (!file.exists(io_path)) exit_file("flux2 vae fixtures missing") + +fx <- safetensors::safe_load_file(io_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()) + ))) +} + +dec <- load_flux2_vae_decoder(vae_path, latent_channels = 32L, + block_channels = c(32L, 32L, 16L, 8L), norm_groups = 8L, + verbose = FALSE) + +# BN statistics loaded from the checkpoint +expect_true(max_abs_diff(dec$bn$running_mean, fx$bn_mean) == 0) +expect_true(max_abs_diff(dec$bn$running_var, fx$bn_var) == 0) + +# Decode parity (post_quant_conv + decoder body) +img <- torch::with_no_grad(dec(fx$latent)) +expect_equal(as.integer(img$shape), as.integer(fx$image$shape)) +expect_true(max_abs_diff(img, fx$image) < 1e-5) + +# BN normalize / denormalize parity + roundtrip +norm <- flux2_bn_normalize(fx$patched, dec$bn$running_mean, + dec$bn$running_var) +expect_true(max_abs_diff(norm, fx$normalized) < 1e-5) +denorm <- flux2_bn_normalize(fx$normalized, dec$bn$running_mean, + dec$bn$running_var, inverse = TRUE) +expect_true(max_abs_diff(denorm, fx$denormalized) < 1e-5) +expect_true(max_abs_diff(denorm, fx$patched) < 1e-5) diff --git a/man/dit_flux2.Rd b/man/dit_flux2.Rd new file mode 100644 index 0000000..1a3ede3 --- /dev/null +++ b/man/dit_flux2.Rd @@ -0,0 +1,14 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{dit_flux2} +\alias{dit_flux2} +\title{FLUX.2 Transformer (MMDiT)} +\description{ +Fresh R port of Flux2Transformer2DModel from the diffusers reference +implementation (Apache-2.0, +src/diffusers/models/transformers/transformer_flux2.py). Defaults are +the klein-4B configuration (5 double + 20 single blocks). Guidance +embeddings (FLUX.2-dev) are not implemented; klein is step-distilled +with \code{guidance_embeds = false}. Timestep conditioning has no +pooled-text component, and the three modulation projections are +shared across all blocks. +} diff --git a/man/dit_flux2_modules.Rd b/man/dit_flux2_modules.Rd new file mode 100644 index 0000000..43a8366 --- /dev/null +++ b/man/dit_flux2_modules.Rd @@ -0,0 +1,18 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{dit_flux2_modules} +\alias{dit_flux2_modules} +\title{FLUX.2 Transformer Building Blocks} +\description{ +Fresh R port of the FLUX.2 MMDiT blocks from the diffusers reference +implementation (Apache-2.0, +src/diffusers/models/transformers/transformer_flux2.py). Key +differences from FLUX.1: modulation is computed ONCE at model level +by shared \code{flux2_modulation} projections and passed into the +blocks (block norms are parameterless), feed-forwards use SwiGLU with +the gate fused into \code{linear_in}, the single-stream block is a +ViT-22B-style parallel block with fully fused projections, and every +linear is bias-free. Module field names mirror the diffusers +state-dict keys 1:1. Reuses \code{flux_attention} (bias = FALSE), +\code{ltx23_rms_norm}, \code{.ltx23_sdpa}, and +\code{flux_apply_rotary_emb}. +} diff --git a/man/download_flux2.Rd b/man/download_flux2.Rd new file mode 100644 index 0000000..6d84f95 --- /dev/null +++ b/man/download_flux2.Rd @@ -0,0 +1,9 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{download_flux2} +\alias{download_flux2} +\title{Download and Prepare FLUX.2 Klein 4B Weights} +\description{ +Downloads FLUX.2-klein-4B from HuggingFace (Apache-2.0, ungated) and +quantizes the 4B transformer to a local fp8 (~4 GB) or NF4 (~2.3 GB) +artifact. +} diff --git a/man/download_flux2_klein.Rd b/man/download_flux2_klein.Rd new file mode 100644 index 0000000..1c9e4ad --- /dev/null +++ b/man/download_flux2_klein.Rd @@ -0,0 +1,31 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{download_flux2_klein} +\alias{download_flux2_klein} +\title{Download FLUX.2-klein-4B and build the quantized artifact} +\usage{ +download_flux2_klein(quantize = TRUE, precision = c("fp8", "nf4"), + output_dir = NULL, text_encoders = TRUE, verbose = TRUE) +} +\arguments{ +\item{quantize}{Logical. Build the quantized artifact.} + +\item{precision}{"fp8" (~4 GB, GPU-resident; near-bf16 quality) or +"nf4" (~2.3 GB).} + +\item{output_dir}{Directory for the quantized artifact.} + +\item{text_encoders}{Logical. Also fetch the Qwen3 text encoder, +tokenizer, VAE, and scheduler config (~8.3 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. No token +is needed (the repo is ungated). The bf16 transformer source +(~7.8 GB in the HuggingFace cache) may be deleted after quantization. +} diff --git a/man/encode_qwen.Rd b/man/encode_qwen.Rd new file mode 100644 index 0000000..c4d784e --- /dev/null +++ b/man/encode_qwen.Rd @@ -0,0 +1,29 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{encode_qwen} +\alias{encode_qwen} +\title{Encode prompts with the Qwen tokenizer} +\usage{ +encode_qwen(tokenizer, texts, max_length = 512L, chat_template = TRUE) +} +\arguments{ +\item{tokenizer}{A \code{\link{qwen_bpe_tokenizer}}.} + +\item{texts}{Character vector of prompts.} + +\item{max_length}{Integer. Fixed sequence length (klein: 512). NULL +for no truncation/padding.} + +\item{chat_template}{Logical. Wrap in the Qwen3 chat template.} +} +\value{ +List with \code{input_ids} and \code{attention_mask} integer + matrices [length(texts), max_length] (ragged lists when + \code{max_length} is NULL). Ids are 0-based. +} +\description{ +With \code{chat_template = TRUE} (the FLUX.2 klein pipeline behavior) +each prompt is wrapped as a single user turn with the generation +prompt and a disabled thinking block, matching +\code{apply_chat_template(..., add_generation_prompt = TRUE, +enable_thinking = FALSE)}. Right-pads with \code{<|endoftext|>}. +} diff --git a/man/encode_with_qwen3.Rd b/man/encode_with_qwen3.Rd new file mode 100644 index 0000000..9cf9a03 --- /dev/null +++ b/man/encode_with_qwen3.Rd @@ -0,0 +1,31 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{encode_with_qwen3} +\alias{encode_with_qwen3} +\title{Encode prompts with the Qwen3 encoder for FLUX.2} +\usage{ +encode_with_qwen3(prompts, model, tokenizer, max_sequence_length = 512L, + out_layers = c(9L, 18L, 27L), device = NULL) +} +\arguments{ +\item{prompts}{Character vector.} + +\item{model}{A \code{\link{qwen3_encoder}}.} + +\item{tokenizer}{A \code{\link{qwen_bpe_tokenizer}}.} + +\item{max_sequence_length}{Integer. Fixed token length (klein: 512).} + +\item{out_layers}{Integer vector. Hidden-state layers (klein-4B: +9, 18, 27).} + +\item{device}{Device for the input ids (defaults to the model's).} +} +\value{ +Tensor [length(prompts), max_sequence_length, + 3 * hidden_size]. +} +\description{ +Tokenizes with the chat template, runs the encoder with the padding +mask, and concatenates the requested mid-stack hidden states per +token, matching Flux2KleinPipeline._get_qwen3_prompt_embeds. +} diff --git a/man/flux2_bn_normalize.Rd b/man/flux2_bn_normalize.Rd new file mode 100644 index 0000000..4146c4e --- /dev/null +++ b/man/flux2_bn_normalize.Rd @@ -0,0 +1,25 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{flux2_bn_normalize} +\alias{flux2_bn_normalize} +\title{Normalize patchified latents with the VAE BatchNorm statistics} +\usage{ +flux2_bn_normalize(latents, bn_mean, bn_var, eps = 1e-04, inverse = FALSE) +} +\arguments{ +\item{latents}{Tensor [B, 128, H, W] (patchified).} + +\item{eps}{Numeric. BatchNorm epsilon.} + +\item{inverse}{Logical. De-normalize (decode path) instead.} + +\item{bn_mean,bn_var}{Float tensors [128].} +} +\value{ +Tensor like \code{latents}. +} +\description{ +FLUX.2 has no scalar scaling/shift factor; latents are standardized +per packed channel with the VAE's \code{bn.running_mean} / +\code{bn.running_var} (eps 1e-4). Reference: encode/decode paths of +Flux2KleinPipeline. +} diff --git a/man/flux2_double_block.Rd b/man/flux2_double_block.Rd new file mode 100644 index 0000000..d6457c6 --- /dev/null +++ b/man/flux2_double_block.Rd @@ -0,0 +1,31 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{flux2_double_block} +\alias{flux2_double_block} +\title{FLUX.2 double-stream (MMDiT) block} +\usage{ +flux2_double_block(dim, num_attention_heads, attention_head_dim, mlp_ratio = 3, + eps = 1e-06, bias = FALSE) +} +\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. FF multiplier (FLUX.2: 3.0).} + +\item{eps}{Numeric. Norm epsilon.} + +\item{bias}{Logical.} +} +\value{ +Module whose forward(hidden_states, encoder_hidden_states, + temb_mod_img, temb_mod_txt, image_rotary_emb) returns + \code{list(encoder_hidden_states, hidden_states)}. +} +\description{ +Image and text streams with externally supplied (shift, scale, gate) +modulation triples, joint attention (txt first), and SwiGLU +feed-forwards. Reference: Flux2TransformerBlock. +} diff --git a/man/flux2_empirical_mu.Rd b/man/flux2_empirical_mu.Rd new file mode 100644 index 0000000..1a68113 --- /dev/null +++ b/man/flux2_empirical_mu.Rd @@ -0,0 +1,21 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{flux2_empirical_mu} +\alias{flux2_empirical_mu} +\title{Empirical timestep shift for FLUX.2} +\usage{ +flux2_empirical_mu(image_seq_len, num_steps) +} +\arguments{ +\item{image_seq_len}{Integer. Packed image token count.} + +\item{num_steps}{Integer. Inference steps.} +} +\value{ +Numeric mu for \code{\link{flowmatch_set_timesteps}}. +} +\description{ +BFL's piecewise-linear fit of the dynamic-shifting mu as a function of +image sequence length and step count; replaces FLUX.1's +calculate_shift. Reference: compute_empirical_mu (adapted from BFL +sampling.py). +} diff --git a/man/flux2_feed_forward.Rd b/man/flux2_feed_forward.Rd new file mode 100644 index 0000000..4be65a3 --- /dev/null +++ b/man/flux2_feed_forward.Rd @@ -0,0 +1,22 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{flux2_feed_forward} +\alias{flux2_feed_forward} +\title{FLUX.2 feed-forward (fused SwiGLU)} +\usage{ +flux2_feed_forward(dim, dim_out = NULL, mult = 3, bias = FALSE) +} +\arguments{ +\item{dim}{Integer. Input dimension.} + +\item{dim_out}{Integer. Output dimension (defaults to \code{dim}).} + +\item{mult}{Numeric. Inner dim multiplier (FLUX.2: 3.0).} + +\item{bias}{Logical.} +} +\description{ +\code{linear_in} projects to twice the inner dim; SwiGLU gates the +first half with SiLU and multiplies by the second half; +\code{linear_out} projects back. Reference: Flux2FeedForward + +Flux2SwiGLU. +} diff --git a/man/flux2_is_quant_key.Rd b/man/flux2_is_quant_key.Rd new file mode 100644 index 0000000..6738aa2 --- /dev/null +++ b/man/flux2_is_quant_key.Rd @@ -0,0 +1,16 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{flux2_is_quant_key} +\alias{flux2_is_quant_key} +\title{Test whether a FLUX.2 key is in the quantization cast set} +\usage{ +flux2_is_quant_key(key) +} +\arguments{ +\item{key}{Character vector of parameter names (diffusers-style).} +} +\value{ +Logical vector. +} +\description{ +Test whether a FLUX.2 key is in the quantization cast set +} diff --git a/man/flux2_load_pipeline.Rd b/man/flux2_load_pipeline.Rd new file mode 100644 index 0000000..65a1dcc --- /dev/null +++ b/man/flux2_load_pipeline.Rd @@ -0,0 +1,36 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{flux2_load_pipeline} +\alias{flux2_load_pipeline} +\title{Load the FLUX.2 klein pipeline} +\usage{ +flux2_load_pipeline(model_dir = NULL, device = "cuda", + precision = c("fp8", "nf4"), text_device = NULL, + attn_chunk = NULL, phase_offload = TRUE, verbose = TRUE) +} +\arguments{ +\item{model_dir}{Quantized artifact directory (default: the +\code{download_flux2_klein} location for \code{precision}), or a +raw diffusers transformer directory.} + +\item{device}{Character. Compute device.} + +\item{precision}{"fp8" (default) or "nf4".} + +\item{text_device}{Device for the Qwen3 encoder (default: +\code{device}; it encodes in its own phase and offloads).} + +\item{attn_chunk}{Integer or NULL. Attention query-chunk override.} + +\item{phase_offload}{Logical. One GPU tenant per phase.} + +\item{verbose}{Logical.} +} +\value{ +A \code{flux2_pipeline} list. +} +\description{ +Loads the quantized transformer artifact plus the FLUX.2 VAE decoder, +Qwen3 text encoder, and tokenizer from the HuggingFace cache +populated by \code{\link{download_flux2_klein}}. With fp8 precision +the ~4 GB transformer stays GPU-resident. +} diff --git a/man/flux2_modulation.Rd b/man/flux2_modulation.Rd new file mode 100644 index 0000000..620120f --- /dev/null +++ b/man/flux2_modulation.Rd @@ -0,0 +1,19 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{flux2_modulation} +\alias{flux2_modulation} +\title{FLUX.2 shared modulation projection} +\usage{ +flux2_modulation(dim, mod_param_sets = 2L, bias = FALSE) +} +\arguments{ +\item{dim}{Integer. Model dimension.} + +\item{mod_param_sets}{Integer. Number of (shift, scale, gate) triples.} + +\item{bias}{Logical.} +} +\description{ +\code{linear(silu(temb))} producing \code{mod_param_sets} triples of +(shift, scale, gate). Computed once per forward at model level and +broadcast to every block. Reference: Flux2Modulation. +} diff --git a/man/flux2_pack_latents.Rd b/man/flux2_pack_latents.Rd new file mode 100644 index 0000000..7c12a69 --- /dev/null +++ b/man/flux2_pack_latents.Rd @@ -0,0 +1,17 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{flux2_pack_latents} +\alias{flux2_pack_latents} +\title{Pack patchified FLUX.2 latents into tokens} +\usage{ +flux2_pack_latents(latents) +} +\arguments{ +\item{latents}{Tensor [B, C, H, W].} +} +\value{ +Tensor [B, H * W, C]. +} +\description{ +[B, C, H, W] -> [B, H * W, C] (row-major spatial flatten, +channels-last). Reference: Flux2KleinPipeline._pack_latents. +} diff --git a/man/flux2_parallel_self_attention.Rd b/man/flux2_parallel_self_attention.Rd new file mode 100644 index 0000000..ee36dcc --- /dev/null +++ b/man/flux2_parallel_self_attention.Rd @@ -0,0 +1,27 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{flux2_parallel_self_attention} +\alias{flux2_parallel_self_attention} +\title{FLUX.2 parallel self-attention (single-stream)} +\usage{ +flux2_parallel_self_attention(query_dim, heads, dim_head, mlp_ratio = 3, + eps = 1e-06, bias = FALSE) +} +\arguments{ +\item{query_dim}{Integer. Model dimension.} + +\item{heads}{Integer. Attention heads.} + +\item{dim_head}{Integer. Per-head dimension.} + +\item{mlp_ratio}{Numeric. MLP hidden multiplier (FLUX.2: 3.0).} + +\item{eps}{Numeric. RMS norm epsilon.} + +\item{bias}{Logical.} +} +\description{ +ViT-22B-style parallel block internals: one fused projection produces +QKV and the SwiGLU MLP input; one fused projection consumes +cat(attention output, MLP output). Reference: +Flux2ParallelSelfAttention + Flux2ParallelSelfAttnProcessor. +} diff --git a/man/flux2_patchify_latents.Rd b/man/flux2_patchify_latents.Rd new file mode 100644 index 0000000..1b2b73b --- /dev/null +++ b/man/flux2_patchify_latents.Rd @@ -0,0 +1,17 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{flux2_patchify_latents} +\alias{flux2_patchify_latents} +\title{Patchify FLUX.2 latents (2x2 -> channels)} +\usage{ +flux2_patchify_latents(latents) +} +\arguments{ +\item{latents}{Tensor [B, C, H, W]; H and W must be even.} +} +\value{ +Tensor [B, C * 4, H / 2, W / 2]. +} +\description{ +[B, C, H, W] -> [B, 4C, H/2, W/2], channel order (C, ph, pw). +Reference: Flux2KleinPipeline._patchify_latents. +} diff --git a/man/flux2_prepare_latent_ids.Rd b/man/flux2_prepare_latent_ids.Rd new file mode 100644 index 0000000..1ced386 --- /dev/null +++ b/man/flux2_prepare_latent_ids.Rd @@ -0,0 +1,22 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{flux2_prepare_latent_ids} +\alias{flux2_prepare_latent_ids} +\title{Build FLUX.2 latent position ids} +\usage{ +flux2_prepare_latent_ids(height, width, device = "cpu") +} +\arguments{ +\item{height}{Integer. Packed grid height (pixel height / 16).} + +\item{width}{Integer. Packed grid width (pixel width / 16).} + +\item{device}{Device for the resulting tensor.} +} +\value{ +Float tensor [height * width, 4]. +} +\description{ +Columns (T, H, W, L) with H and W carrying the packed-grid position +(row-major: H varies slowest), T = L = 0. Reference: +Flux2KleinPipeline._prepare_latent_ids. +} diff --git a/man/flux2_prepare_text_ids.Rd b/man/flux2_prepare_text_ids.Rd new file mode 100644 index 0000000..ba91920 --- /dev/null +++ b/man/flux2_prepare_text_ids.Rd @@ -0,0 +1,19 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{flux2_prepare_text_ids} +\alias{flux2_prepare_text_ids} +\title{Build FLUX.2 text position ids} +\usage{ +flux2_prepare_text_ids(len, device = "cpu") +} +\arguments{ +\item{len}{Integer. Text sequence length.} + +\item{device}{Device for the resulting tensor.} +} +\value{ +Float tensor [len, 4]. +} +\description{ +Columns (T, H, W, L) with only L varying: 0..len-1. Reference: +Flux2KleinPipeline._prepare_text_ids. +} diff --git a/man/flux2_single_block.Rd b/man/flux2_single_block.Rd new file mode 100644 index 0000000..0a72614 --- /dev/null +++ b/man/flux2_single_block.Rd @@ -0,0 +1,31 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{flux2_single_block} +\alias{flux2_single_block} +\title{FLUX.2 single-stream block (parallel)} +\usage{ +flux2_single_block(dim, num_attention_heads, attention_head_dim, mlp_ratio = 3, + eps = 1e-06, bias = FALSE) +} +\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. MLP multiplier (FLUX.2: 3.0).} + +\item{eps}{Numeric. Norm epsilon.} + +\item{bias}{Logical.} +} +\value{ +Module whose forward(hidden_states, temb_mod, + image_rotary_emb) returns the joint hidden states. +} +\description{ +Parameterless LayerNorm with external modulation, then the fused +parallel attention+MLP. Operates on the pre-concatenated [text; image] +sequence (the reference model concatenates once before the stack). +Reference: Flux2SingleTransformerBlock. +} diff --git a/man/flux2_transformer.Rd b/man/flux2_transformer.Rd new file mode 100644 index 0000000..dd13ae8 --- /dev/null +++ b/man/flux2_transformer.Rd @@ -0,0 +1,51 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{flux2_transformer} +\alias{flux2_transformer} +\title{FLUX.2 transformer model} +\usage{ +flux2_transformer(in_channels = 128L, num_layers = 5L, num_single_layers = 20L, + attention_head_dim = 128L, num_attention_heads = 24L, + joint_attention_dim = 7680L, mlp_ratio = 3, + timestep_guidance_channels = 256L, + axes_dims_rope = c(32L, 32L, 32L, 32L), rope_theta = 2000, + eps = 1e-06, out_channels = NULL) +} +\arguments{ +\item{in_channels}{Integer. Packed latent channels (128).} + +\item{num_layers}{Integer. Double-stream block count (klein-4B: 5).} + +\item{num_single_layers}{Integer. Single-stream block count (20).} + +\item{attention_head_dim}{Integer. Per-head dimension.} + +\item{num_attention_heads}{Integer. Attention heads.} + +\item{joint_attention_dim}{Integer. Text embedding dim (7680).} + +\item{mlp_ratio}{Numeric. Feed-forward multiplier (3.0).} + +\item{timestep_guidance_channels}{Integer. Sinusoid width (256).} + +\item{axes_dims_rope}{Integer vector. Per-axis rotary dims.} + +\item{rope_theta}{Numeric. Rotary base frequency (2000).} + +\item{eps}{Numeric. Norm epsilon.} + +\item{out_channels}{Integer or NULL. Defaults to \code{in_channels}.} +} +\value{ +Module whose forward(hidden_states, encoder_hidden_states, + timestep, image_rotary_emb) returns the predicted velocity for the + image tokens [B, S_img, out_channels]. \code{timestep} is in sigma + space (0-1); it is scaled by 1000 internally. +} +\description{ +Shared modulation computed once per forward; double blocks over +separate text/image streams, then single (parallel) blocks over the +concatenated [text; image] sequence. Rotary embeddings are +precomputed by the caller with \code{\link{flux_pos_embed}} +(\code{axes_dim = c(32, 32, 32, 32)}, \code{theta = 2000}) over the +concatenated [text; image] 4-axis position ids. +} diff --git a/man/flux2_unpack_latents_with_ids.Rd b/man/flux2_unpack_latents_with_ids.Rd new file mode 100644 index 0000000..99518aa --- /dev/null +++ b/man/flux2_unpack_latents_with_ids.Rd @@ -0,0 +1,22 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{flux2_unpack_latents_with_ids} +\alias{flux2_unpack_latents_with_ids} +\title{Unpack FLUX.2 tokens back to a latent grid via position ids} +\usage{ +flux2_unpack_latents_with_ids(x, ids, height, width) +} +\arguments{ +\item{x}{Tensor [B, S, C] of tokens.} + +\item{ids}{Tensor [S, 4] (or [B, S, 4]) of position ids.} + +\item{height,width}{Integers. Packed grid dimensions.} +} +\value{ +Tensor [B, C, height, width]. +} +\description{ +Scatters tokens to (H, W) positions taken from the id columns +(H = column 2, W = column 3, 0-based values). Reference: +Flux2KleinPipeline._unpack_latents_with_ids. +} diff --git a/man/flux2_unpatchify_latents.Rd b/man/flux2_unpatchify_latents.Rd new file mode 100644 index 0000000..a0eae60 --- /dev/null +++ b/man/flux2_unpatchify_latents.Rd @@ -0,0 +1,17 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{flux2_unpatchify_latents} +\alias{flux2_unpatchify_latents} +\title{Unpatchify FLUX.2 latents (channels -> 2x2)} +\usage{ +flux2_unpatchify_latents(latents) +} +\arguments{ +\item{latents}{Tensor [B, 4C, H, W].} +} +\value{ +Tensor [B, C, H * 2, W * 2]. +} +\description{ +Inverse of \code{\link{flux2_patchify_latents}}. Reference: +Flux2KleinPipeline._unpatchify_latents. +} diff --git a/man/flux2_vae_decoder.Rd b/man/flux2_vae_decoder.Rd new file mode 100644 index 0000000..0a3baa7 --- /dev/null +++ b/man/flux2_vae_decoder.Rd @@ -0,0 +1,28 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{flux2_vae_decoder} +\alias{flux2_vae_decoder} +\title{FLUX.2 VAE decoder} +\usage{ +flux2_vae_decoder(latent_channels = 32L, + block_channels = c(512L, 512L, 256L, 128L), norm_groups = 32L) +} +\arguments{ +\item{latent_channels}{Integer (32 for FLUX.2).} + +\item{block_channels}{Decoder block channels (reversed encoder +block_out_channels).} + +\item{norm_groups}{Integer. Group norm groups.} +} +\value{ +Module whose forward(z) decodes [B, 32, H, W] latents to + [B, 3, 8H, 8W] images; \code{$bn$running_mean} / + \code{$bn$running_var} carry the normalization statistics. +} +\description{ +The AutoencoderKLFlux2 decode path: post_quant_conv (1x1, 32 +channels) followed by the standard AutoencoderKL decoder body +(reused from \code{\link{vae_decoder_native}}), plus the BatchNorm +running statistics used for latent (de)normalization. Reference: +src/diffusers/models/autoencoders/autoencoder_kl_flux2.py. +} diff --git a/man/flux_ada_layer_norm_continuous.Rd b/man/flux_ada_layer_norm_continuous.Rd index 1504a4c..d9e5bfd 100644 --- a/man/flux_ada_layer_norm_continuous.Rd +++ b/man/flux_ada_layer_norm_continuous.Rd @@ -3,12 +3,15 @@ \alias{flux_ada_layer_norm_continuous} \title{FLUX continuous adaLN (final norm)} \usage{ -flux_ada_layer_norm_continuous(dim, cond_dim = dim) +flux_ada_layer_norm_continuous(dim, cond_dim = dim, bias = TRUE) } \arguments{ \item{dim}{Integer. Model dimension.} \item{cond_dim}{Integer. Conditioning embedding dimension.} + +\item{bias}{Logical. Bias on the projection (TRUE for FLUX.1, FALSE +for FLUX.2).} } \description{ Scale/shift conditioning of the final norm. Note the chunk order: diff --git a/man/flux_attention.Rd b/man/flux_attention.Rd index 9ba74f4..949bb69 100644 --- a/man/flux_attention.Rd +++ b/man/flux_attention.Rd @@ -4,7 +4,7 @@ \title{FLUX joint attention} \usage{ flux_attention(query_dim, heads, dim_head, added_kv = FALSE, pre_only = FALSE, - eps = 1e-06) + eps = 1e-06, bias = TRUE) } \arguments{ \item{query_dim}{Integer. Model dimension.} @@ -18,6 +18,9 @@ flux_attention(query_dim, heads, dim_head, added_kv = FALSE, pre_only = FALSE, \item{pre_only}{Logical. Skip the output projection (single blocks).} \item{eps}{Numeric. RMS norm epsilon.} + +\item{bias}{Logical. Bias on the linear projections (TRUE for FLUX.1, +FALSE for FLUX.2).} } \description{ Multi-head attention with per-head RMS q/k norms and rotary position diff --git a/man/flux_load_transformer.Rd b/man/flux_load_transformer.Rd index 112cb8f..75c4699 100644 --- a/man/flux_load_transformer.Rd +++ b/man/flux_load_transformer.Rd @@ -4,7 +4,7 @@ \title{Load a FLUX transformer from any checkpoint format} \usage{ flux_load_transformer(ckpt, device = "cuda", dtype = "bfloat16", pin = TRUE, - verbose = TRUE, ...) + fp8_resident = FALSE, verbose = TRUE, ...) } \arguments{ \item{ckpt}{A checkpoint from \code{\link{flux_open_checkpoint}} or @@ -17,7 +17,12 @@ quantized formats this sets the resident (non-quantized) tensors and must match the compute dtype: bfloat16 for GPU compute, float32 for CPU compute.} -\item{pin}{Logical. Pin fp8 host memory for faster transfers.} +\item{pin}{Logical. Pin fp8 host memory for faster transfers +(streamed fp8 only).} + +\item{fp8_resident}{Logical. Keep the fp8 weights on \code{device} +instead of streaming from the CPU - right for models whose whole +quantized footprint fits in VRAM (FLUX.2 klein-4B: ~4 GB).} \item{verbose}{Logical.} diff --git a/man/load_flux2_vae_decoder.Rd b/man/load_flux2_vae_decoder.Rd new file mode 100644 index 0000000..b183b0a --- /dev/null +++ b/man/load_flux2_vae_decoder.Rd @@ -0,0 +1,26 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{load_flux2_vae_decoder} +\alias{load_flux2_vae_decoder} +\title{Load the FLUX.2 VAE decoder from safetensors} +\usage{ +load_flux2_vae_decoder(path, latent_channels = 32L, + block_channels = c(512L, 512L, 256L, 128L), + norm_groups = 32L, verbose = TRUE) +} +\arguments{ +\item{path}{Path to the VAE .safetensors file (or a directory +containing diffusion_pytorch_model.safetensors).} + +\item{verbose}{Logical.} + +\item{latent_channels,block_channels,norm_groups}{Constructor +arguments for \code{\link{flux2_vae_decoder}}.} +} +\value{ +The loaded \code{flux2_vae_decoder} in eval mode. +} +\description{ +Loads the decoder half plus post_quant_conv and the BatchNorm running +statistics; encoder and quant_conv keys are skipped (txt2img needs no +encoder). +} diff --git a/man/load_qwen3_text_encoder.Rd b/man/load_qwen3_text_encoder.Rd new file mode 100644 index 0000000..f56e45c --- /dev/null +++ b/man/load_qwen3_text_encoder.Rd @@ -0,0 +1,28 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{load_qwen3_text_encoder} +\alias{load_qwen3_text_encoder} +\title{Load a Qwen3 encoder from a transformers directory} +\usage{ +load_qwen3_text_encoder(model_path, device = "cpu", dtype = "float32", + verbose = TRUE, ...) +} +\arguments{ +\item{model_path}{Directory with \code{config.json} and +\code{model*.safetensors} (FLUX.2-klein's \code{text_encoder}).} + +\item{device}{Character. Target device.} + +\item{dtype}{Character. "bfloat16" (GPU) or "float32" (CPU).} + +\item{verbose}{Logical.} + +\item{...}{Overrides for \code{\link{qwen3_encoder}} arguments.} +} +\value{ +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. +} diff --git a/man/qwen3_encoder.Rd b/man/qwen3_encoder.Rd new file mode 100644 index 0000000..68232d0 --- /dev/null +++ b/man/qwen3_encoder.Rd @@ -0,0 +1,32 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{qwen3_encoder} +\alias{qwen3_encoder} +\title{Qwen3 encoder stack} +\usage{ +qwen3_encoder(vocab_size = 151936L, hidden_size = 2560L, + intermediate_size = 9728L, num_hidden_layers = 36L, + num_attention_heads = 32L, num_key_value_heads = 8L, + head_dim = 128L, rope_theta = 1e+06, rms_norm_eps = 1e-06) +} +\arguments{ +\item{rope_theta}{Numeric.} + +\item{rms_norm_eps}{Numeric.} + +\item{vocab_size,hidden_size,intermediate_size,num_hidden_layers}{Integers.} + +\item{num_attention_heads,num_key_value_heads,head_dim}{Integers.} +} +\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. +} +\description{ +Defaults are the Qwen3-4B configuration used by FLUX.2 klein. The +module tree mirrors the checkpoint keys (\code{model.embed_tokens}, +\code{model.layers.*}, \code{model.norm}); the tied LM head carries +no weights of its own and is not implemented. +} diff --git a/man/qwen3_text_encoder.Rd b/man/qwen3_text_encoder.Rd new file mode 100644 index 0000000..4d5187e --- /dev/null +++ b/man/qwen3_text_encoder.Rd @@ -0,0 +1,15 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{qwen3_text_encoder} +\alias{qwen3_text_encoder} +\title{Qwen3 Text Encoder} +\description{ +Fresh R port of the Qwen3 decoder stack from HuggingFace transformers +(Apache-2.0, src/transformers/models/qwen3/), used by FLUX.2 klein as +its text encoder (Qwen3-4B: 36 layers, hidden 2560, 32 query / 8 KV +heads, head_dim 128, SwiGLU 9728, RoPE theta 1e6). The pipeline +consumes mid-stack hidden states (layers 9, 18, 27 for klein-4B) +concatenated per token, so the forward runs only as deep as the last +requested layer; the LM head is never needed (embeddings are tied). +Causal attention with the tokenizer's padding mask, matching the +reference exactly. +} diff --git a/man/qwen_bpe_tokenizer.Rd b/man/qwen_bpe_tokenizer.Rd new file mode 100644 index 0000000..a78349b --- /dev/null +++ b/man/qwen_bpe_tokenizer.Rd @@ -0,0 +1,17 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{qwen_bpe_tokenizer} +\alias{qwen_bpe_tokenizer} +\title{Load a Qwen2 byte-level BPE tokenizer} +\usage{ +qwen_bpe_tokenizer(tokenizer_path) +} +\arguments{ +\item{tokenizer_path}{Path to a tokenizer.json (or a directory +containing one).} +} +\value{ +A \code{qwen_tokenizer} object. +} +\description{ +Load a Qwen2 byte-level BPE tokenizer +} diff --git a/man/rope_flux2.Rd b/man/rope_flux2.Rd new file mode 100644 index 0000000..d65f231 --- /dev/null +++ b/man/rope_flux2.Rd @@ -0,0 +1,14 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{rope_flux2} +\alias{rope_flux2} +\title{FLUX.2 Position Ids and Empirical Shift} +\description{ +Fresh R port of the FLUX.2 position-id builders and the empirical +timestep-shift formula from the diffusers reference (Apache-2.0, +src/diffusers/pipelines/flux2/pipeline_flux2_klein.py). FLUX.2 uses +4-axis rotary position ids (T, H, W, L): text tokens carry only the +L axis (sequence position), image latents carry H and W, and the T +axis distinguishes reference images (unused for txt2img). Frequencies +come from \code{\link{flux_pos_embed}} with +\code{axes_dim = c(32, 32, 32, 32)} and \code{theta = 2000}. +} diff --git a/man/tokenizer_qwen.Rd b/man/tokenizer_qwen.Rd new file mode 100644 index 0000000..4667413 --- /dev/null +++ b/man/tokenizer_qwen.Rd @@ -0,0 +1,18 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{tokenizer_qwen} +\alias{tokenizer_qwen} +\title{Qwen2 Byte-Level BPE Tokenizer} +\description{ +Pure R implementation of the Qwen2 tokenizer (HuggingFace +tokenizer.json, BPE model with ByteLevel pre-tokenization), as used by +FLUX.2 klein's Qwen3 text encoder. Text is split with the GPT-4-style +regex, each pre-token's UTF-8 bytes are mapped through the GPT-2 +byte-to-unicode table, and rank-based BPE merges produce the ids. +Added tokens (\code{<|im_start|>}, \code{}, ...) are split out +literally before byte-level encoding. +} +\details{ +Limitation: the NFC normalizer is not applied (base R has no NFC); +input is assumed to already be NFC, which holds for ordinary text. + +} diff --git a/man/txt2img.Rd b/man/txt2img.Rd index 7eb0c03..b0f5628 100644 --- a/man/txt2img.Rd +++ b/man/txt2img.Rd @@ -3,7 +3,7 @@ \alias{txt2img} \title{Generate an image from a text prompt using a diffusion pipeline} \usage{ -txt2img(prompt, model_name = c("sd21", "sdxl", "flux1"), ...) +txt2img(prompt, model_name = c("sd21", "sdxl", "flux1", "flux2"), ...) } \arguments{ \item{prompt}{A character string prompt describing the image to generate.} diff --git a/man/txt2img_flux2.Rd b/man/txt2img_flux2.Rd new file mode 100644 index 0000000..ed48e3a --- /dev/null +++ b/man/txt2img_flux2.Rd @@ -0,0 +1,48 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{txt2img_flux2} +\alias{txt2img_flux2} +\title{Generate an image with FLUX.2 klein} +\usage{ +txt2img_flux2(prompt, pipeline = NULL, width = 1024L, height = 1024L, + num_inference_steps = 4L, max_sequence_length = 512L, + seed = NULL, prompt_embeds = NULL, save_file = TRUE, + filename = NULL, verbose = TRUE, ...) +} +\arguments{ +\item{prompt}{Character. The prompt.} + +\item{pipeline}{A \code{flux2_pipeline} from +\code{\link{flux2_load_pipeline}}; NULL loads one (passing +\code{...} through).} + +\item{num_inference_steps}{Integer. Denoising steps (klein-4B: 4).} + +\item{max_sequence_length}{Integer. Qwen3 token length (512).} + +\item{seed}{Integer or NULL. Latents are drawn on the CPU in the +packed shape, so a seed matches a Python diffusers run with a CPU +generator.} + +\item{prompt_embeds}{Optional precomputed [B, S, 7680] embeddings.} + +\item{save_file}{Logical. Write a PNG.} + +\item{filename}{Output path (default derived from the prompt).} + +\item{verbose}{Logical.} + +\item{...}{Passed to \code{\link{flux2_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 (klein-4B: 4 steps, no guidance): Qwen3 +prompt encoding (chat template, mid-stack hidden states), FlowMatch +denoising with the empirical dynamic shift, and 32-channel VAE decode +through the BatchNorm latent statistics. +} diff --git a/man/vae_flux2.Rd b/man/vae_flux2.Rd new file mode 100644 index 0000000..33170bc --- /dev/null +++ b/man/vae_flux2.Rd @@ -0,0 +1,12 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{vae_flux2} +\alias{vae_flux2} +\title{FLUX.2 Latent Layout and VAE Helpers} +\description{ +Fresh R port of the FLUX.2 latent packing chain from the diffusers +reference (Apache-2.0, src/diffusers/pipelines/flux2/ +pipeline_flux2_klein.py). The 32-channel VAE latent is patchified +2x2 into 128 channels, normalized with the VAE's BatchNorm running +statistics (there is no scalar scaling/shift factor in FLUX.2), and +flattened to channels-last tokens for the transformer. +} diff --git a/tools/compare_translation.R b/tools/compare_translation.R index 87feaad..ceab7a2 100644 --- a/tools/compare_translation.R +++ b/tools/compare_translation.R @@ -190,3 +190,32 @@ if (file.exists(t5_path)) { } else { cat("\n(t5 encoder pairing skipped: tools/cache/modeling_t5.py missing)\n") } + +# 4. FLUX.2 transformer stack + pipeline +tf2 <- parse_file(file.path(DIFFUSERS, "models/transformers/transformer_flux2.py"), + ts_language_python()) +report_pair("flux2 transformer", + py_scopes(tf2, c("Flux2SwiGLU", "Flux2FeedForward", "Flux2AttnProcessor", + "Flux2Attention", "Flux2ParallelSelfAttnProcessor", + "Flux2ParallelSelfAttention", "Flux2SingleTransformerBlock", + "Flux2TransformerBlock", "Flux2Modulation", "Flux2PosEmbed", + "Flux2TimestepGuidanceEmbeddings", "Flux2Transformer2DModel")), + c("R/dit_flux2_modules.R", "R/dit_flux2.R", "R/rope_flux2.R")) + +pipe2 <- parse_file(file.path(DIFFUSERS, "pipelines/flux2/pipeline_flux2_klein.py"), + ts_language_python()) +report_pair("flux2 klein pipeline", + py_scopes(pipe2, c("Flux2KleinPipeline", "compute_empirical_mu")), + c("R/txt2img_flux2.R", "R/vae_flux2.R", "R/rope_flux2.R")) + +# 5. Qwen3 encoder (modular reference; requires tools/cache/modular_qwen3.py) +q3_path <- "tools/cache/modular_qwen3.py" +if (file.exists(q3_path)) { + q3 <- parse_file(q3_path, ts_language_python()) + report_pair("qwen3 encoder", + py_scopes(q3, c("Qwen3MLP", "Qwen3Attention", "Qwen3DecoderLayer", + "Qwen3RMSNorm", "Qwen3Model")), + c("R/qwen3_text_encoder.R")) +} else { + cat("\n(qwen3 pairing skipped: tools/cache/modular_qwen3.py missing)\n") +} diff --git a/tools/gen_fixtures_flux2.py b/tools/gen_fixtures_flux2.py new file mode 100644 index 0000000..926ce47 --- /dev/null +++ b/tools/gen_fixtures_flux2.py @@ -0,0 +1,77 @@ +# Generate FLUX.2 phase-1 parity fixtures for the R port: 4-axis RoPE, +# position-id builders, patchify/pack/unpack chain, and empirical mu. +# +# Runs the diffusers reference (Apache-2.0) on small fixed inputs. Run +# via tools/gen_fixtures.sh; never executed at package test/run time. + +import os +import sys + +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.models.transformers.transformer_flux2 import Flux2PosEmbed # noqa: E402 +from diffusers.pipelines.flux2.pipeline_flux2_klein import ( # noqa: E402 + Flux2KleinPipeline, + compute_empirical_mu, +) + +OUT_DIR = os.path.join(os.path.dirname(__file__), "..", "inst", "tinytest", "fixtures") +os.makedirs(OUT_DIR, exist_ok=True) + +torch.manual_seed(47) +fx = {} + +# --- position ids --------------------------------------------------------------- +S_TXT = 7 +GRID_H, GRID_W = 6, 10 # asymmetric so an axis swap fails + +dummy_txt = torch.zeros(1, S_TXT, 4) +txt_ids = Flux2KleinPipeline._prepare_text_ids(dummy_txt)[0] # [7, 4] +fx["txt_ids"] = txt_ids.float() + +lat = torch.randn(1, 128, GRID_H, GRID_W) +latent_ids = Flux2KleinPipeline._prepare_latent_ids(lat)[0] # [60, 4] +fx["latent_ids"] = latent_ids.float() + +# --- Flux2PosEmbed: full axes (32,32,32,32) theta 2000 + tiny axes ------------- +ids = torch.cat([txt_ids.float(), latent_ids.float()], dim=0) # [67, 4] +fx["ids"] = ids + +cos_full, sin_full = Flux2PosEmbed(theta=2000, axes_dim=[32, 32, 32, 32])(ids) +fx["pos_full_cos"] = cos_full # [67, 128] +fx["pos_full_sin"] = sin_full + +cos_tiny, sin_tiny = Flux2PosEmbed(theta=2000, axes_dim=[2, 2, 2, 2])(ids) +fx["pos_tiny_cos"] = cos_tiny # [67, 8] +fx["pos_tiny_sin"] = sin_tiny + +# --- patchify / pack / unpack chain ---------------------------------------------- +z32 = torch.randn(2, 32, 12, 20) # unpatchified 32-channel latent +patched = Flux2KleinPipeline._patchify_latents(z32) # [2, 128, 6, 10] +packed = Flux2KleinPipeline._pack_latents(patched) # [2, 60, 128] +fx["z32"] = z32 +fx["patched"] = patched +fx["packed"] = packed +fx["unpatched"] = Flux2KleinPipeline._unpatchify_latents(patched) + +ids_b = Flux2KleinPipeline._prepare_latent_ids(patched) # [2, 60, 4] +unpacked = Flux2KleinPipeline._unpack_latents_with_ids( + packed, ids_b, height=6, width=10 +) # [2, 128, 6, 10] +fx["unpacked"] = unpacked + +# --- empirical mu ----------------------------------------------------------------- +cases = [(1024, 4), (4096, 4), (4096, 28), (256, 50), (4352, 4), (8192, 10)] +fx["mu_seq"] = torch.tensor([c[0] for c in cases], dtype=torch.float64) +fx["mu_steps"] = torch.tensor([c[1] for c in cases], dtype=torch.float64) +fx["mu_vals"] = torch.tensor( + [compute_empirical_mu(s, n) for s, n in cases], dtype=torch.float64 +) + +fx = {k: v.contiguous() for k, v in fx.items()} +save_file(fx, os.path.join(OUT_DIR, "rope_flux2.safetensors"), + metadata={"purpose": "diffuseR FLUX.2 test fixture"}) +print(f"wrote {len(fx)} tensors to {OUT_DIR}/rope_flux2.safetensors") diff --git a/tools/gen_fixtures_flux2_dit.py b/tools/gen_fixtures_flux2_dit.py new file mode 100644 index 0000000..3f7c9d9 --- /dev/null +++ b/tools/gen_fixtures_flux2_dit.py @@ -0,0 +1,116 @@ +# Generate FLUX.2 transformer-block parity fixtures for the R port. +# +# Runs the diffusers reference modules (Apache-2.0) with tiny random-init +# configs. Run via tools/gen_fixtures.sh; never executed at package +# test/run time. + +import os +import sys + +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.models.transformers.transformer_flux2 import ( # noqa: E402 + Flux2FeedForward, + Flux2Modulation, + Flux2PosEmbed, + Flux2SingleTransformerBlock, + Flux2TransformerBlock, +) + +OUT_DIR = os.path.join(os.path.dirname(__file__), "..", "inst", "tinytest", "fixtures") +os.makedirs(OUT_DIR, exist_ok=True) + +torch.manual_seed(48) +fx = {} + +# Tiny config: dim 16 = 2 heads x head_dim 8, rope axes (2,2,2,2), mlp 3.0 +DIM, HEADS, HEAD_DIM = 16, 2, 8 +B, S_TXT = 2, 7 +GRID_H, GRID_W = 6, 10 +S_IMG = GRID_H * GRID_W # 60 + +# 4-axis rope over [txt; img] ids +txt_ids = torch.zeros(S_TXT, 4) +txt_ids[:, 3] = torch.arange(S_TXT) +img_ids = torch.zeros(GRID_H, GRID_W, 4) +img_ids[..., 1] = torch.arange(GRID_H)[:, None] +img_ids[..., 2] = torch.arange(GRID_W)[None, :] +img_ids = img_ids.reshape(S_IMG, 4) +ids = torch.cat([txt_ids, img_ids], dim=0) +rope_cos, rope_sin = Flux2PosEmbed(theta=2000, axes_dim=[2, 2, 2, 2])(ids) +fx["rope_cos"] = rope_cos +fx["rope_sin"] = rope_sin + +img_x = torch.randn(B, S_IMG, DIM) +txt_x = torch.randn(B, S_TXT, DIM) +joint_x = torch.randn(B, S_TXT + S_IMG, DIM) +temb = torch.randn(B, DIM) +fx["img_x"] = img_x +fx["txt_x"] = txt_x +fx["joint_x"] = joint_x +fx["temb"] = temb + + +def add_state(prefix, module): + for k, v in module.state_dict().items(): + fx[f"{prefix}.{k}"] = v + + +# --- Flux2Modulation (2-set and 1-set) ------------------------------------------- +mod2 = Flux2Modulation(DIM, mod_param_sets=2) +add_state("mod2", mod2) +with torch.no_grad(): + fx["mod2_out"] = mod2(temb) + +mod1 = Flux2Modulation(DIM, mod_param_sets=1) +add_state("mod1", mod1) +with torch.no_grad(): + fx["mod1_out"] = mod1(temb) + +# --- Flux2FeedForward (SwiGLU) ------------------------------------------------------ +ff = Flux2FeedForward(DIM, DIM, mult=3.0) +add_state("ff", ff) +with torch.no_grad(): + fx["ff_out"] = ff(joint_x) + +# --- Flux2TransformerBlock (double) --------------------------------------------------- +dbl = Flux2TransformerBlock(dim=DIM, num_attention_heads=HEADS, + attention_head_dim=HEAD_DIM, mlp_ratio=3.0) +add_state("dbl", dbl) +with torch.no_grad(): + mod_img = mod2(temb) + mod_txt = mod2(temb) * 0.5 # distinct txt modulation + enc_out, hid_out = dbl( + hidden_states=img_x, + encoder_hidden_states=txt_x, + temb_mod_img=mod_img, + temb_mod_txt=mod_txt, + image_rotary_emb=(rope_cos, rope_sin), + ) +fx["dbl_mod_img"] = mod_img +fx["dbl_mod_txt"] = mod_txt +fx["dbl_enc_out"] = enc_out +fx["dbl_hid_out"] = hid_out + +# --- Flux2SingleTransformerBlock -------------------------------------------------------- +sgl = Flux2SingleTransformerBlock(dim=DIM, num_attention_heads=HEADS, + attention_head_dim=HEAD_DIM, mlp_ratio=3.0) +add_state("sgl", sgl) +with torch.no_grad(): + mod_single = mod1(temb) + joint_out = sgl( + hidden_states=joint_x, + encoder_hidden_states=None, + temb_mod=mod_single, + image_rotary_emb=(rope_cos, rope_sin), + ) +fx["sgl_mod"] = mod_single +fx["sgl_out"] = joint_out + +fx = {k: v.contiguous() for k, v in fx.items()} +save_file(fx, os.path.join(OUT_DIR, "dit_flux2.safetensors"), + metadata={"purpose": "diffuseR FLUX.2 test fixture"}) +print(f"wrote {len(fx)} tensors to {OUT_DIR}/dit_flux2.safetensors") diff --git a/tools/gen_fixtures_flux2_model.py b/tools/gen_fixtures_flux2_model.py new file mode 100644 index 0000000..f80892b --- /dev/null +++ b/tools/gen_fixtures_flux2_model.py @@ -0,0 +1,85 @@ +# Generate a tiny random-init Flux2Transformer2DModel parity fixture for +# the R port, 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 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 Flux2Transformer2DModel # noqa: E402 + +OUT_DIR = os.path.join(os.path.dirname(__file__), "..", "inst", "tinytest", "fixtures") +CKPT_DIR = os.path.join(OUT_DIR, "flux2_tiny_ckpt") +os.makedirs(OUT_DIR, exist_ok=True) + +torch.manual_seed(49) + +model = Flux2Transformer2DModel( + patch_size=1, + in_channels=8, + num_layers=1, + num_single_layers=1, + attention_head_dim=8, + num_attention_heads=2, + joint_attention_dim=24, + mlp_ratio=3.0, + timestep_guidance_channels=256, + axes_dims_rope=(2, 2, 2, 2), + rope_theta=2000, + guidance_embeds=False, +) +model.eval() + +B, S_TXT = 2, 7 +GRID_H, GRID_W = 6, 10 +S_IMG = GRID_H * GRID_W + +hidden = torch.randn(B, S_IMG, 8) +encoder = torch.randn(B, S_TXT, 24) +timestep = torch.tensor([1.0, 0.25]) # sigma space; model multiplies by 1000 + +txt_ids = torch.zeros(S_TXT, 4) +txt_ids[:, 3] = torch.arange(S_TXT) +img_ids = torch.zeros(GRID_H, GRID_W, 4) +img_ids[..., 1] = torch.arange(GRID_H)[:, None] +img_ids[..., 2] = torch.arange(GRID_W)[None, :] +img_ids = img_ids.reshape(S_IMG, 4) + +with torch.no_grad(): + out = model( + hidden_states=hidden, + encoder_hidden_states=encoder, + timestep=timestep, + img_ids=img_ids, + txt_ids=txt_ids, + guidance=None, + return_dict=False, + )[0] + +fx = {f"model.{k}": v for k, v in model.state_dict().items()} +fx.update( + { + "hidden": hidden, + "encoder": encoder, + "timestep": timestep, + "txt_ids": txt_ids, + "img_ids": img_ids, + "out": out, + } +) +fx = {k: v.contiguous() for k, v in fx.items()} +save_file(fx, os.path.join(OUT_DIR, "flux2_model.safetensors"), + metadata={"purpose": "diffuseR FLUX.2 test fixture"}) +print(f"wrote {len(fx)} tensors to {OUT_DIR}/flux2_model.safetensors") + +if os.path.isdir(CKPT_DIR): + shutil.rmtree(CKPT_DIR) +model.save_pretrained(CKPT_DIR, max_shard_size="30KB") +print(f"wrote checkpoint to {CKPT_DIR}: {sorted(os.listdir(CKPT_DIR))}") diff --git a/tools/gen_fixtures_flux2_vae.py b/tools/gen_fixtures_flux2_vae.py new file mode 100644 index 0000000..bc72275 --- /dev/null +++ b/tools/gen_fixtures_flux2_vae.py @@ -0,0 +1,70 @@ +# Generate FLUX.2 VAE decode parity fixtures for the R port. +# +# Uses the diffusers AutoencoderKLFlux2 (Apache-2.0) with a tiny +# random-init config and randomized BatchNorm running statistics. Run +# via tools/gen_fixtures.sh. + +import os +import sys + +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 AutoencoderKLFlux2 # noqa: E402 + +OUT_DIR = os.path.join(os.path.dirname(__file__), "..", "inst", "tinytest", "fixtures") +os.makedirs(OUT_DIR, exist_ok=True) + +torch.manual_seed(51) + +vae = AutoencoderKLFlux2( + in_channels=3, + out_channels=3, + down_block_types=("DownEncoderBlock2D",) * 4, + up_block_types=("UpDecoderBlock2D",) * 4, + block_out_channels=(8, 16, 32, 32), + layers_per_block=2, + latent_channels=32, + norm_num_groups=8, + use_quant_conv=True, + use_post_quant_conv=True, + sample_size=64, +) +vae.eval() + +# Randomize the BN running statistics so the (de)normalization is a real +# transform in the fixtures +with torch.no_grad(): + vae.bn.running_mean.normal_(0.0, 0.5) + vae.bn.running_var.uniform_(0.5, 2.0) + +latent = torch.randn(1, 32, 4, 4) +with torch.no_grad(): + image = vae.decode(latent, return_dict=False)[0] + +# Reference BN normalize/denormalize on a patchified map +patched = torch.randn(1, 128, 3, 5) +bn_mean = vae.bn.running_mean.view(1, -1, 1, 1) +bn_std = torch.sqrt(vae.bn.running_var + vae.config.batch_norm_eps).view(1, -1, 1, 1) +normalized = (patched - bn_mean) / bn_std +denormalized = normalized * bn_std + bn_mean + +sd = {k: v.contiguous() for k, v in vae.state_dict().items()} +save_file(sd, os.path.join(OUT_DIR, "vae_flux2_tiny.safetensors"), + metadata={"purpose": "diffuseR FLUX.2 test fixture"}) + +io = { + "latent": latent, + "image": image, + "patched": patched, + "normalized": normalized, + "denormalized": denormalized, + "bn_mean": vae.bn.running_mean.clone(), + "bn_var": vae.bn.running_var.clone(), +} +io = {k: v.contiguous() for k, v in io.items()} +save_file(io, os.path.join(OUT_DIR, "vae_flux2_io.safetensors"), + metadata={"purpose": "diffuseR FLUX.2 test fixture"}) +print(f"wrote vae_flux2_tiny + {len(io)} io tensors to {OUT_DIR}") diff --git a/tools/gen_fixtures_qwen3.py b/tools/gen_fixtures_qwen3.py new file mode 100644 index 0000000..b43ab9e --- /dev/null +++ b/tools/gen_fixtures_qwen3.py @@ -0,0 +1,75 @@ +# Generate Qwen3 encoder parity fixtures for the R port. +# +# Uses HF transformers' Qwen3ForCausalLM (Apache-2.0) with a tiny +# random-init config; captures the mid-stack hidden states the FLUX.2 +# klein pipeline consumes. Run via tools/gen_fixtures.sh. + +import os +import shutil + +import torch +from safetensors.torch import save_file +from transformers import Qwen3Config, Qwen3ForCausalLM + +OUT_DIR = os.path.join(os.path.dirname(__file__), "..", "inst", "tinytest", "fixtures") +CKPT_DIR = os.path.join(OUT_DIR, "qwen3_tiny_ckpt") +os.makedirs(OUT_DIR, exist_ok=True) + +torch.manual_seed(50) + +config = Qwen3Config( + 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=1e6, + rms_norm_eps=1e-6, + tie_word_embeddings=True, + attention_bias=False, + max_position_embeddings=512, +) +model = Qwen3ForCausalLM(config) +model.eval() + +# Right-padded batch with attention mask (the klein pipeline passes the +# mask to the text encoder) +input_ids = torch.tensor([ + [5, 23, 61, 7, 19, 88, 42, 3, 9, 0, 0, 0], + [33, 14, 2, 71, 55, 0, 0, 0, 0, 0, 0, 0], +]) +attention_mask = torch.tensor([ + [1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0], + [1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0], +]) + +with torch.no_grad(): + out = model(input_ids=input_ids, attention_mask=attention_mask, + output_hidden_states=True, use_cache=False) + +# Pipeline-style stack of mid-stack layers (1, 2, 3) -> [B, L, 3*hidden] +layers = (1, 2, 3) +stacked = torch.stack([out.hidden_states[k] for k in layers], dim=1) +stacked = stacked.permute(0, 2, 1, 3).reshape(input_ids.shape[0], + input_ids.shape[1], + 3 * config.hidden_size) + +fx = { + "input_ids": input_ids, + "attention_mask": attention_mask, + "h1": out.hidden_states[1], + "h2": out.hidden_states[2], + "h3": out.hidden_states[3], + "stacked": stacked, +} +fx = {k: v.contiguous() for k, v in fx.items()} +save_file(fx, os.path.join(OUT_DIR, "qwen3_flux2.safetensors"), + metadata={"purpose": "diffuseR FLUX.2 test fixture"}) +print(f"wrote {len(fx)} tensors to {OUT_DIR}/qwen3_flux2.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))}") diff --git a/tools/gen_qwen_tokenizer_cases.py b/tools/gen_qwen_tokenizer_cases.py new file mode 100644 index 0000000..22c19da --- /dev/null +++ b/tools/gen_qwen_tokenizer_cases.py @@ -0,0 +1,90 @@ +# Generate Qwen2 ByteLevel-BPE tokenizer parity cases for the R port +# (FLUX.2 klein's tokenizer), plus the rendered Qwen3 chat template. +# +# Downloads the tokenizer from the ungated black-forest-labs/ +# FLUX.2-klein-4B repo (a few MB). Writes: +# - tools/cache/tokenizer_qwen.json (dev copy, gitignored) +# - inst/tinytest/fixtures/qwen_tokenizer_cases.json (checked in) +# +# Run: +# uv run --no-project --with transformers --with torch \ +# --index https://download.pytorch.org/whl/cpu \ +# --index-strategy unsafe-best-match python tools/gen_qwen_tokenizer_cases.py + +import json +import os + +from transformers import AutoTokenizer + +ROOT = os.path.join(os.path.dirname(__file__), "..") +CACHE_DIR = os.path.join(ROOT, "tools", "cache") +FIXTURE = os.path.join(ROOT, "inst", "tinytest", "fixtures", "qwen_tokenizer_cases.json") +os.makedirs(CACHE_DIR, exist_ok=True) + +tok = AutoTokenizer.from_pretrained("black-forest-labs/FLUX.2-klein-4B", + subfolder="tokenizer") +tok.backend_tokenizer.save(os.path.join(CACHE_DIR, "tokenizer_qwen.json")) + +PROMPTS = [ + "a photo of a cat", + "A sunset over mountains, ultra detailed, 8k", + "Hello, world!", + "The quick brown fox jumps over the lazy dog.", + "it's a beautiful day; isn't it?", + "3.14159 and 2,000,000 dollars", + "state-of-the-art text-to-image generation", + " leading spaces", + "trailing spaces ", + "double and triple spaces", + "UPPERCASE lowercase MiXeD", + "email@example.com and https://example.org/path?q=1", + 'quotes "double" and \'single\'', + "(parentheses) [brackets] {braces}", + "underscores_and_snake_case", + "a", + "", + "café résumé naïve", + "em—dash and en–dash", + "100% of $50 + €20", + "emoji 🦊 and 中文字符 mixed in", + "newline\nand\ttab", + "An astronaut riding a horse on Mars, photorealistic", + "watercolor painting of a fox in a snowy forest", + ("The transformer architecture uses self-attention mechanisms " + "to model long-range dependencies. ") * 3, +] + +# Raw tokenizer parity (no special tokens, no template) +cases = [{"text": t, "ids": tok(t, add_special_tokens=False)["input_ids"]} + for t in PROMPTS] + +# Pipeline-style: chat template (user message, generation prompt, no +# thinking) rendered then padded/truncated exactly like the klein pipeline +templated = [] +for text in [PROMPTS[0], PROMPTS[20], PROMPTS[24]]: + messages = [{"role": "user", "content": text}] + rendered = tok.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True, + enable_thinking=False, + ) + enc = tok(rendered, padding="max_length", truncation=True, max_length=64) + templated.append({ + "text": text, + "rendered": rendered, + "max_length": 64, + "ids": enc["input_ids"], + "mask": enc["attention_mask"], + }) + +meta = { + "pad_token": tok.pad_token, + "pad_token_id": tok.pad_token_id, + "padding_side": tok.padding_side, +} + +with open(FIXTURE, "w", encoding="utf-8") as f: + json.dump({"cases": cases, "templated": templated, "meta": meta}, + f, ensure_ascii=False, indent=1) +print(f"wrote {len(cases)} cases + {len(templated)} templated to {FIXTURE}") +print("meta:", meta) +print("rendered example:", json.dumps(templated[0]["rendered"]))