FLUX.1-schnell text-to-image support

.gitignore

@@ -51,5 +51,7 @@ rsconnect/
 
 # Upstream reference checkouts and working notes (never committed)
 ref/
+/ref
 tasks/
 .venv/
+tools/cache/

CLAUDE.md

@@ -316,10 +316,37 @@ See cornyverse CLAUDE.md for safetensors package setup (use cornball-ai fork unt
   - Handles Blackwell workaround automatically
 
 ### Model Support
-- [ ] Add FLUX model support
+- [x] Add FLUX model support (FLUX.1-schnell, see below)
 - [ ] Add SD3 model support
 - [ ] ControlNet integration
 
+### FLUX.1-schnell (Complete)
+
+4-step distilled text-to-image, CFG-free. 12B MMDiT + T5-XXL + CLIP-L +
+16-channel VAE, ported from diffusers/transformers (provenance in
+`inst/REFERENCES.md`).
+
+```r
+library(diffuseR)
+download_flux1()   # gated HF repo: accept license + set HF_TOKEN; ~34 GB
+                   # download, one-time NF4 quantize to a 6.8 GB artifact
+txt2img_flux("An astronaut riding a horse on Mars, photorealistic",
+             seed = 7)   # or txt2img("...", model_name = "flux1")
+```
+
+Measured on the RTX 5060 Ti 16 GB (NF4 resident, T5 float32 on CPU):
+1024x1024 in ~2 min wall (peak 8.7 GB alloc / 9.0 GB reserved),
+512x512 in ~1.5 min (peak 8.0 GB). CPU-only works too (~9 min at 256px).
+
+Key components: `flux_transformer` (19 double + 38 single blocks),
+`t5_encoder` + `unigram_tokenizer` (pure R SentencePiece Viterbi),
+`flux_quantize`/`flux_load_transformer` (NF4 ~6.8 GB resident or fp8
+~12 GB streamed; cast census is exactly 494 weights), `flux_load_pipeline`
++ `txt2img_flux` with per-phase GPU offloading. Gotchas that bit once:
+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).
+
 ### LTX-2.3 Video Generation (clean-room rewrite in progress)
 
 The original LTX-2.0 port was removed and is being replaced by a ground-up

DESCRIPTION

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

NAMESPACE

@@ -3,31 +3,57 @@
 export(auto_devices)
 export(bpe_tokenizer)
 export(clear_vram)
+export(clip_pooled_output)
 export(CLIPTokenizer)
 export(ddim_scheduler_create)
 export(ddim_scheduler_step)
 export(decode_bpe)
 export(download_component)
+export(download_flux1)
 export(download_ltx2)
 export(download_model)
 export(encode_bpe)
+export(encode_unigram)
 export(encode_with_gemma3)
+export(encode_with_t5)
 export(filename_from_prompt)
 export(flowmatch_calculate_shift)
 export(flowmatch_scale_noise)
 export(flowmatch_scheduler_create)
 export(flowmatch_scheduler_step)
 export(flowmatch_set_timesteps)
+export(flux_ada_layer_norm_continuous)
+export(flux_ada_layer_norm_zero)
+export(flux_ada_layer_norm_zero_single)
+export(flux_apply_rotary_emb)
+export(flux_attention)
+export(flux_double_block)
+export(flux_is_quant_key)
+export(flux_load_pipeline)
+export(flux_load_transformer)
+export(flux_memory_profile)
+export(flux_open_checkpoint)
+export(flux_open_quantized)
+export(flux_pack_latents)
+export(flux_pos_embed)
+export(flux_prepare_latent_image_ids)
+export(flux_quantize)
+export(flux_single_block)
+export(flux_transformer)
+export(flux_unpack_latents)
 export(gemma3_config_ltx2)
 export(gemma3_text_model)
 export(gemma3_tokenizer)
 export(img2img)
 export(is_blackwell_gpu)
 export(latents_to_video)
+export(load_decoder_safetensors)
 export(load_decoder_weights)
 export(load_gemma3_text_encoder)
 export(load_model_component)
 export(load_pipeline)
+export(load_t5_text_encoder)
+export(load_text_encoder_safetensors)
 export(load_text_encoder_weights)
 export(load_text_encoder2_weights)
 export(load_to_gpu)
@@ -124,23 +150,27 @@ export(save_video)
 export(save_video_ltx23)
 export(scheduler_add_noise)
 export(sdxl_memory_profile)
+export(t5_encoder)
 export(text_encoder_native)
 export(text_encoder2_native)
 export(tokenize_gemma3)
 export(txt2img)
+export(txt2img_flux)
 export(txt2img_sd21)
 export(txt2img_sdxl)
 export(txt2vid_ltx2)
 export(unet_native)
 export(unet_native_from_torchscript)
 export(unet_sdxl_native)
 export(unet_sdxl_native_from_torchscript)
+export(unigram_tokenizer)
 export(vae_decoder_native)
 export(vocab_size)
 export(vram_report)
 export(write_wav)
 
 S3method(print,bpe_tokenizer)
 S3method(print,ltx23_checkpoint)
+S3method(print,unigram_tokenizer)
 
 importFrom(utils,head)

R/checkpoint_flux.R

@@ -0,0 +1,141 @@
+#' FLUX Checkpoint Readers
+#'
+#' FLUX transformers ship in the diffusers layout: a directory with
+#' \code{config.json}, one or more \code{diffusion_pytorch_model*.safetensors}
+#' shards, and (when sharded) a
+#' \code{diffusion_pytorch_model.safetensors.index.json} weight map.
+#' These helpers open that layout behind the same checkpoint interface as
+#' \code{\link{ltx23_open_checkpoint}}, so the LTX group loaders and
+#' quantization machinery work unchanged. FLUX module names mirror the
+#' checkpoint keys 1:1 - no key mapping is needed.
+#'
+#' @name checkpoint_flux
+NULL
+
+# Quantization cast set: every large linear weight in the FLUX blocks,
+# including the adaLN modulation linears (norm*.linear are 3.2B params
+# across the model; leaving them bf16 would not fit resident on 16 GB).
+# Everything else (embedders, q/k norms, final norm, biases) stays bf16.
+# Full-size census: 19 double blocks x 14 + 38 single blocks x 6 = 494
+# cast weights, ~11.8B of the 12B parameters.
+.flux_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\\.net\\.(0\\.proj|2)|ff_context\\.net\\.(0\\.proj|2)",
+                                   "|norm1\\.linear|norm1_context\\.linear",
+                                   ")",
+                                   "|single_transformer_blocks\\.[0-9]+\\.(",
+                                   "attn\\.(to_q|to_k|to_v)|proj_mlp|proj_out|norm\\.linear",
+                                   ")",
+                                   ")\\.weight$"
+)
+
+#' Test whether a FLUX key is in the quantization cast set
+#'
+#' @param key Character vector of parameter names (diffusers-style).
+#'
+#' @return Logical vector.
+#'
+#' @export
+flux_is_quant_key <- function(key) {
+    grepl(.flux_quant_cast_pattern, key)
+}
+
+#' Open a FLUX transformer checkpoint directory
+#'
+#' Opens a diffusers-layout transformer directory lazily (headers only).
+#' Sharded checkpoints are resolved through the index.json weight map;
+#' single-file checkpoints are opened directly. The transformer
+#' \code{config.json} is attached as \code{$config}.
+#'
+#' @param transformer_dir Directory containing \code{config.json} and the
+#'   \code{diffusion_pytorch_model*.safetensors} file(s).
+#'
+#' @return An object of class \code{ltx23_checkpoint} (shared checkpoint
+#'   interface): list with \code{handle$get_tensor}, \code{keys},
+#'   \code{config}, and \code{path}.
+#'
+#' @export
+flux_open_checkpoint <- function(transformer_dir) {
+    if (!requireNamespace("safetensors", quietly = TRUE)) {
+        stop("The safetensors package is required to read FLUX checkpoints.")
+    }
+    transformer_dir <- path.expand(transformer_dir)
+    if (!dir.exists(transformer_dir)) {
+        stop("Checkpoint directory not found: ", transformer_dir)
+    }
+
+    config <- NULL
+    config_path <- file.path(transformer_dir, "config.json")
+    if (file.exists(config_path)) {
+        config <- jsonlite::fromJSON(config_path, simplifyVector = TRUE)
+    }
+
+    opened <- .flux_open_sharded_dir(transformer_dir, "diffusion_pytorch_model")
+
+    structure(
+              list(handle = opened$handle, keys = opened$keys, version = NULL,
+                   config = config, path = transformer_dir),
+              class = "ltx23_checkpoint"
+    )
+}
+
+# Open a HF-layout safetensors directory (sharded via <base>.safetensors
+# .index.json, or a single <base>.safetensors) lazily; returns the
+# handle/keys pair shared by all checkpoint objects.
+.flux_open_sharded_dir <- function(dir, base) {
+    index_path <- file.path(dir, paste0(base, ".safetensors.index.json"))
+    if (file.exists(index_path)) {
+        index <- jsonlite::fromJSON(index_path, simplifyVector = TRUE)
+        weight_map <- unlist(index$weight_map)
+        shard_files <- unique(weight_map)
+        missing <- shard_files[!file.exists(file.path(dir, shard_files))]
+        if (length(missing)) {
+            stop("Missing checkpoint shards: ", paste(missing, collapse = ", "))
+        }
+        handles <- lapply(file.path(dir, shard_files), function(p) {
+            safetensors::safetensors$new(p, framework = "torch")
+        })
+        names(handles) <- shard_files
+        keys <- names(weight_map)
+        handle <- list(
+                       get_tensor = function(key) {
+            shard <- weight_map[[key]]
+            if (is.null(shard) || is.na(shard)) {
+                stop("Key not found in checkpoint index: ", key)
+            }
+            handles[[shard]]$get_tensor(key)
+        }
+        )
+    } else {
+        single_path <- file.path(dir, paste0(base, ".safetensors"))
+        if (!file.exists(single_path)) {
+            stop("No ", base, " safetensors (or index) in ", dir)
+        }
+        h <- safetensors::safetensors$new(single_path, framework = "torch")
+        keys <- setdiff(h$keys(), "__metadata__")
+        handle <- list(get_tensor = function(key) h$get_tensor(key))
+    }
+    list(handle = handle, keys = keys)
+}
+
+#' Open a quantized FLUX artifact directory
+#'
+#' Opens the sharded NF4/fp8 artifact written by
+#' \code{\link{flux_quantize}} through the shared checkpoint interface.
+#' The manifest's embedded transformer config and \code{format} ride
+#' along, so \code{\link{flux_load_transformer}} needs nothing else.
+#'
+#' @param dir The quantized artifact directory (with manifest.json).
+#'
+#' @return An \code{ltx23_checkpoint} with \code{$format} set.
+#'
+#' @export
+flux_open_quantized <- function(dir) {
+    manifest_path <- file.path(dir, "manifest.json")
+    if (!file.exists(manifest_path)) {
+        stop("No manifest.json in ", dir, "; run flux_quantize() first.")
+    }
+    ltx23_open_fp8_checkpoint(dir)
+}