欢迎 FLUX.2——Black Forest Labs 推出的全新开源图像生成模型 🤗
发布于 2025 年 11 月 25 日
FLUX.2 是 Black Forest Labs 继 Flux.1 系列之后推出的最新图像生成模型系列。这是一个采用全新架构、从头开始预训练的全新模型！本文将讨论 FLUX.2 引入的关键变化、在不同配置下执行推理的方法，以及 LoRA 微调。
🚨 FLUX.2 并非 FLUX.1 的直接替代品，而是一个全新的图像生成与编辑模型。
目录
- FLUX.2 简介
- 使用 Diffusers 进行推理
- 高级提示词
- LoRA 微调
FLUX.2 简要介绍
FLUX.2 既可用于图像引导的图像生成，也可用于文本引导的图像生成。此外，它还能将多张图像作为参考输入，同时生成最终输出图像。下面，我们简要讨论 FLUX.2 引入的关键变化。
文本编码器
首先，与 Flux.1 使用两个文本编码器不同，FLUX.2 仅使用一个文本编码器——Mistral Small 3.1。使用单一文本编码器大大简化了提示词嵌入向量的计算过程。该流水线支持最大序列长度（max_sequence_length）为 512。FLUX.2 并非使用单层输出作为提示词嵌入向量，而是堆叠中间层的输出，已知这种做法更为有益。
DiT
FLUX.2 沿用了与 Flux.1 相同的通用多模态扩散 Transformer（MM-DiT）+ 并行 DiT 架构。回顾一下，MM-DiT 块首先在独立的流中处理图像潜变量和条件文本，仅在注意力操作时将两者合并，因此被称为"双流"块。随后，并行块对拼接后的图像和文本流进行操作，可视为"单流"块。
从 Flux.1 到 FLUX.2，DiT 的关键变化如下：
时间信息和引导信息（以 AdaLayerNorm-Zero 调制参数的形式）分别在所有双流和单流 Transformer 块之间共享，而非像 Flux.1 那样为每个块设置独立的调制参数。
模型中的任何层都不使用偏置参数。具体来说，两个 Transformer 模块中的注意力子块和前馈（FF）子块，在其任何层中均不使用偏置参数。
在 Flux.1 中，单流 Transformer 模块将注意力输出投影与前馈输出投影融合在一起。FLUX.2 的单流模块还将注意力 QKV 投影与前馈输入投影融合，从而创建了一个完全并行的 Transformer 模块：
图片取自 ViT-22B 论文。
请注意，与上图所示的 ViT-22B 模块相比，FLUX.2 使用了 SwiGLU 风格的 MLP 激活函数，而非 GELU 激活函数（并且同样不使用偏置参数）。
- FLUX.2 中单流 Transformer 模块的比例更大（8 个双流模块对 48 个单流模块，而 Flux.1 的比例为 19/38）。这也意味着单流模块在 DiT 参数中占据了更大比例：Flux.1[dev]-12B 约 54% 的总参数位于双流模块中，而 FLUX.2[dev]-32B 约 24% 的参数位于双流模块中（约 73% 位于单流模块中）。
其他
- 新的自编码器，即 AutoencoderKLFlux2
- 整合分辨率相关时间步调度的更好方法
使用 Diffusers 进行推理
FLUX.2 使用了更大的 DiT 和 Mistral3 Small 作为其文本编码器。当两者一起使用且不进行任何卸载时，推理需要超过 80GB 的显存。在以下章节中，我们将展示如何在各种系统级限制下，以更易实现的方式对 FLUX.2 进行推理。
安装与身份验证
在尝试以下代码片段之前，请确保你已从 main 分支安装了 diffusers，并运行了 `hf auth login`。
pip uninstall diffusers -y && pip install git+https://github.com/huggingface/diffusers -U
常规推理
from diffusers import Flux2Pipeline
import torch
repo_id = "black-forest-labs/FLUX.2-dev"
pipe = Flux2Pipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
pipe.enable_model_cpu_offload()
image = pipe(
prompt="dog dancing near the sun",
num_inference_steps=50, # 28 is a good trade-off
guidance_scale=4,
height=1024,
width=1024
).images[0]
上述代码片段已在 H100 上测试过，如果不启用 CPU 卸载，它不足以在该 GPU 上运行推理。启用 CPU 卸载后，此配置运行大约需要 62GB 显存。
拥有 Hopper 系列 GPU 的用户可以利用 Flash Attention 3 来加速推理：
from diffusers import Flux2Pipeline
import torch
repo_id = "black-forest-labs/FLUX.2-dev"
pipe = Flux2Pipeline.from_pretrained(path, torch_dtype=torch.bfloat16)
+ pipe.transformer.set_attention_backend("_flash_3_hub")
pipe.enable_model_cpu_offload()
image = pipe(
prompt="dog dancing near the sun",
num_inference_steps=50,
guidance_scale=2.5,
height=1024,
width=1024
).images[0]
你可以在此处查看支持的注意力后端（我们有很多！）。
资源受限
使用 4 位量化
借助 bitsandbytes，我们可以将 Transformer 和文本编码器模型加载为 4-bit 精度，从而使拥有 24GB GPU 的用户能够在本地使用该模型。你可以在拥有约 20 GB 空闲显存的 GPU 上运行这段代码。
展开
import torch
from transformers import Mistral3ForConditionalGeneration
from diffusers import Flux2Pipeline, Flux2Transformer2DModel
repo_id = "diffusers/FLUX.2-dev-bnb-4bit"
device = "cuda:0"
torch_dtype = torch.bfloat16
transformer = Flux2Transformer2DModel.from_pretrained(
repo_id, subfolder="transformer", torch_dtype=torch_dtype, device_map="cpu"
)
text_encoder = Mistral3ForConditionalGeneration.from_pretrained(
repo_id, subfolder="text_encoder", dtype=torch_dtype, device_map="cpu"
)
pipe = Flux2Pipeline.from_pretrained(
repo_id, transformer=transformer, text_encoder=text_encoder, torch_dtype=torch_dtype
)
pipe.enable_model_cpu_offload()
prompt = "Realistic macro photograph of a hermit crab using a soda can as its shell, partially emerging from the can, captured with sharp detail and natural colors, on a sunlit beach with soft shadows and a shallow depth of field, with blurred ocean waves in the background. The can has the text `BFL Diffusers` on it and it has a color gradient that start with #FF5733 at the top and transitions to #33FF57 at the bottom."
image = pipe(
prompt=prompt,
generator=torch.Generator(device=device).manual_seed(42),
num_inference_steps=50, # 28 is a good trade-off
guidance_scale=4,
).images[0]
image.save("flux2_t2i_nf4.png")
请注意，我们使用的仓库包含了 FLUX.2 DiT 和 Mistral 文本编码器的 NF4 量化版本。
本地 + 远程
由于 Diffusers 管线的模块化设计，我们可以将各个模块分离并按顺序处理。我们将文本编码器解耦，并将其部署到推理端点（Inference Endpoint）上。这有助于我们释放显存，使其仅用于 DiT 和 VAE。
⚠️ 要使用远程文本编码器，你需要拥有一个有效的 token。如果你已经完成身份验证，则无需进一步操作。
下面的示例结合使用了本地推理和远程推理。此外，我们还通过 bitsandbytes 对 DiT 进行了 NF4 量化。
你可以在拥有 18 GB 显存的 GPU 上运行这段代码：
展开
from diffusers import Flux2Pipeline, Flux2Transformer2DModel
from diffusers import BitsAndBytesConfig as DiffBitsAndBytesConfig
from huggingface_hub import get_token
import requests
import torch
import io
def remote_text_encoder(prompts: str | list[str]):
def _encode_single(prompt: str):
response = requests.post(
"https://remote-text-encoder-flux-2.huggingface.co/predict",
json={"prompt": prompt},
headers={
"Authorization": f"Bearer {get_token()}",
"Content-Type": "application/json"
}
)
assert response.status_code == 200, f"{response.status_code=}"
return torch.load(io.BytesIO(response.content))
if isinstance(prompts, (list, tuple)):
embeds = [_encode_single(p) for p in prompts]
return torch.cat(embeds, dim=0)
return _encode_single(prompts).to("cuda")
repo_id = "black-forest-labs/FLUX.2-dev"
quantized_dit_id = "diffusers/FLUX.2-dev-bnb-4bit"
dit = Flux2Transformer2DModel.from_pretrained(
quantized_dit_id, subfolder="transformer", torch_dtype=torch_dtype, device_map="cpu"
)
pipe = Flux2Pipeline.from_pretrained(
repo_id,
text_encoder=None,
transformer=dit,
torch_dtype=torch.bfloat16,
)
pipe.enable_model_cpu_offload()
print("Running remote text encoder ☁️")
prompt1 = "a photo of a forest with mist swirling around the tree trunks. The word 'FLUX.2' is painted over it in big, red brush strokes with visible texture"
prompt2 = "a photo of a dense forest with rain. The word 'FLUX.2' is painted over it in big, red brush strokes with visible texture"
prompt_embeds = remote_text_encoder([prompt1, prompt2])
print("Done ✅")
out = pipe(
prompt_embeds=prompt_embeds,
generator=torch.Generator(device="cuda").manual_seed(42),
num_inference_steps=50, # 28 is a good trade-off
guidance_scale=4,
height=1024,
width=1024,
)
for idx, image in enumerate(out.images):
image.save(f"flux_out_{idx}.png")
对于显存更低的 GPU，我们提供了 group_offloading 功能，这使得显存低至 8GB 的 GPU 也能使用该模型。不过，你需要 32GB 的空闲内存。或者，如果你愿意牺牲一些速度，可以将 `low_cpu_mem_usage=True` 设置为 True，从而将内存需求降低到仅 10GB。
展开
import io
import os
import requests
import torch
from diffusers import Flux2Pipeline, Flux2Transformer2DModel
repo_id = "diffusers/FLUX.2-dev-bnb-4bit"
torch_dtype = torch.bfloat16
device = "cuda"
def remote_text_encoder(prompts: str | list[str]):
def _encode_single(prompt: str):
response = requests.post(
"https://remote-text-encoder-flux-2.huggingface.co/predict",
json={"prompt": prompt},
headers={"Authorization": f"Bearer {os.environ['HF_TOKEN']}", "Content-Type": "application/json"},
)
assert response.status_code == 200, f"{response.status_code=}"
return torch.load(io.BytesIO(response.content))
if isinstance(prompts, (list, tuple)):
embeds = [_encode_single(p) for p in prompts]
return torch.cat(embeds, dim=0)
return _encode_single(prompts).to("cuda")
transformer = Flux2Transformer2DModel.from_pretrained(
repo_id, subfolder="transformer", torch_dtype=torch_dtype, device_map="cpu"
)
pipe = Flux2Pipeline.from_pretrained(
repo_id,
text_encoder=None,
transformer=transformer,
torch_dtype=torch_dtype,
)
pipe.transformer.enable_group_offload(
onload_device=device,
offload_device="cpu",
offload_type="leaf_level",
use_stream=True,
# low_cpu_mem_usage=True # uncomment for lower RAM usage
)
pipe.to(device)
prompt = "a photo of a forest with mist swirling around the tree trunks. The word 'FLUX.2' is painted over it in big, red brush strokes with visible texture"
prompt_embeds = remote_text_encoder(prompt)
image = pipe(
prompt_embeds=prompt_embeds,
generator=torch.Generator(device=device).manual_seed(42),
num_inference_steps=50,
guidance_scale=4,
height=1024,
width=1024,
).images[0]
你可以在此处查看其他支持的量化后端，以及在此处查看其他节省内存的技术。
要查看不同量化方式对图像的影响，你可以在下面的交互式演示中尝试，或者直接在 FLUX.2 量化实验 Space 中独立访问。
多张图像作为参考
FLUX.2 支持使用多张图像作为输入，允许你最多使用 10 张图像。但请注意，每增加一张图像都需要更多的显存。你可以通过索引（例如，图像 1、图像 2）或自然语言（例如，袋鼠、乌龟）来引用这些图像。为了获得最佳效果，最好的方法是结合使用这两种方式。
展开
import torch
from transformers import Mistral3ForConditionalGeneration
from diffusers import Flux2Pipeline, Flux2Transformer2DModel
from diffusers.utils import load_image
repo_id = "diffusers-internal-dev/new-model-image-final-weights"
device = "cuda:0"
torch_dtype = torch.bfloat16
pipe = Flux2Pipeline.from_pretrained(
repo_id, torch_dtype=torch_dtype
)
pipe.enable_model_cpu_offload()
image_one = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/flux2_blog/kangaroo.png")
image_two = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/flux2_blog/turtle.png")
prompt = "the boxer kangaroo from image 1 and the martial artist turtle from image 2 are fighting in an epic battle scene at a beach of a tropical island, 35mm, depth of field, 50mm lens, f/3.5, cinematic lighting"
image = pipe(
prompt=prompt,
image=[image_one, image_two],
generator=torch.Generator(device=device).manual_seed(42),
num_inference_steps=50,
guidance_scale=2.5,
width=1024,
height=768,
).images[0]
image.save(f"./flux2_t2i.png")
多图像输入
高级提示词
FLUX.2 支持高级提示词技术，例如结构化 JSON 提示词、精确的十六进制颜色控制以及多参考图像编辑。除了提供更强的控制能力外，这还允许在保持其他属性整体不变的情况下，灵活地更改特定属性。
例如，我们先以这个 JSON 作为基础架构（取自官方 FLUX.2 提示词指南）：
{
"scene": "overall scene description",
"subjects": [
{
"description": "detailed subject description",
"position": "where in frame",
"action": "what they're doing"
}
],
"style": "artistic style",
"color_palette": ["#hex1", "#hex2", "#hex3"],
"lighting": "lighting description",
"mood": "emotional tone",
"background": "background details",
"composition": "framing and layout",
"camera": {
"angle": "camera angle",
"lens": "lens type",
"depth_of_field": "focus behavior"
}
}
在此基础上，我们将其转化为一个提示词，用于生成一张老式随身听放在地毯上的画面（只需将这段提示词传入你上面选定的 diffusers 推理示例即可）：
prompt = """
{
"scene": "Professional studio product photography setup with soft-textured carpet surface",
"subjects": [
{
"description": "Old silver Walkman placed on a carpet in the middle of an empty room",
"pose": "Stationary, lying flat",
"position": "Center foreground on carpeted surface",
"color_palette": ["brushed silver", "dark gray accents"]
}
],
"style": "Ultra-realistic product photography with commercial quality",
"color_palette": ["brushed silver", "neutral beige", "soft white highlights"],
"lighting": "Three-point softbox setup creating soft, diffused highlights with no harsh shadows",
"mood": "Clean, professional, minimalist",
"background": "Soft-textured carpet surface with subtle studio backdrop suggesting an empty room",
"composition": "rule of thirds",
"camera": {
"angle": "high angle",
"distance": "medium shot",
"focus": "Sharp focus on metallic Walkman textures and physical controls",
"lens-mm": 85,
"f-number": "f/5.6",
"ISO": 200
}
}
"""
现在，将地毯颜色改为特定的蓝绿色调（#367588），并添加连接到随身听的有线耳机：
prompt = """
{
"scene": "Professional studio product photography setup with soft-textured carpet surface",
"subjects": [
{
"description": "Old silver Walkman placed on a teal-blue carpet (#367588) in the middle of an empty room, with wired headphones plugged in",
"pose": "Stationary, lying flat",
"position": "Center foreground on carpeted surface",
"color_palette": ["brushed silver", "dark gray accents", "#367588"]
},
{
"description": "Wired headphones connected to the Walkman, cable loosely coiled on the carpet",
"pose": "Stationary",
"position": "Next to and partially in front of the Walkman on the carpet",
"color_palette": ["dark gray", "soft black", "#367588"]
}
],
"style": "Ultra-realistic product photography with commercial quality",
"color_palette": ["brushed silver", "#367588", "neutral beige", "soft white highlights"],
"lighting": "Three-point softbox setup creating soft, diffused highlights with no harsh shadows",
"mood": "Clean, professional, minimalist",
"background": "Soft-textured teal-blue carpet surface (#367588) with subtle studio backdrop suggesting an empty room",
"composition": "rule of thirds",
"camera": {
"angle": "high angle",
"distance": "medium shot",
"focus": "Sharp focus on metallic Walkman textures, wired headphones, and carpet fibers",
"lens-mm": 85,
"f-number": "f/5.6",
"ISO": 200
}
}
"""
地毯颜色现已与提供的十六进制色码匹配，耳机也已添加，整体场景有细微调整。
更多示例和详情请查阅官方提示词指南。
LoRA 微调
作为一款同时支持文生图和图生图的模型，FLUX.2 非常适合针对多种使用场景进行微调！然而，由于仅推理就需要超过 80GB 显存，LoRA 微调在消费级 GPU 上运行更具挑战。为了尽可能节省显存，我们将上述部分推理优化技术也应用于训练，并结合共享显存节省技术，大幅降低显存消耗。要训练该模型，你可以使用下方的 diffusers 代码或 Ostris 的 AI Toolkit。
我们提供了文生图和图生图两种训练脚本，本篇博客将重点介绍文生图训练示例。
微调显存优化
其中许多技术可以相互补充，并同时使用以进一步降低显存消耗。不过，某些技术可能互斥，因此在启动训练前务必检查确认。
展开查看所用显存节省技术的详情：
远程文本编码器：要利用远程文本编码进行训练，只需传入 `--remote_text_encoder` 参数。请注意，你必须已登录 Hugging Face 账户（`hf auth login`）或通过 `--hub_token` 传入一个 token。
CPU 卸载：通过传入 `--offload` 参数，VAE 和文本编码器将被卸载到 CPU 内存，仅在需要时移至 GPU。
潜在缓存：使用 VAE 对训练图像进行预编码，然后删除 VAE 以释放部分显存。要启用潜在缓存，只需传入 `--cache_latents` 参数。
QLoRA：基于量化的低精度训练——使用 8 位或 4 位量化。你可以使用以下标志：
- 基于 torchao 的 FP8 训练：通过传入 `--do_fp8_training` 启用 FP8 训练。由于我们使用的是 FP8 张量核心，因此需要计算能力至少为 8.9 或更高的 CUDA GPU。如果你希望在相对较旧的显卡上进行内存高效的训练，我们建议你查看其他训练器，例如 SimpleTuner、ai-toolkit 等。
- 基于 bitsandbytes 的 NF4 训练：或者，你也可以通过 bitsandbytes 使用 8 位或 4 位量化：传入 `--bnb_quantization_config_path`，并附带一个指向包含你配置的 json 文件的相应路径。详情请见下文。
梯度检查点与梯度累积：`--gradient_accumulation` 指的是在执行反向传播/更新步骤之前累积的更新步数。通过传入一个大于 1 的值，你可以减少反向传播/更新步骤的次数，从而也降低内存需求。* 使用 `--gradient_checkpointing`，我们可以通过不在前向传播过程中存储所有中间激活值来节省内存。相反，只存储这些激活值的一个子集（检查点），其余部分在反向传播过程中根据需要重新计算。请注意，这是以反向传播速度变慢为代价的。
8 位 Adam 优化器：在使用 AdamW 进行训练时（不适用于 prodigy），你可以传入 `--use_8bit_adam` 来降低训练的内存需求。如果这样做，请确保已安装 bitsandbytes。
在开始训练之前，请务必查看 README 以了解先决条件。
对于这个示例，我们将使用 `multimodalart/1920-raider-waite-tarot-public-domain` 数据集，并采用以下配置进行 FP8 训练。欢迎尝试更多超参数并分享你的结果 🤗
accelerate launch train_dreambooth_lora_flux2.py \
--pretrained_model_name_or_path="black-forest-labs/FLUX.2-dev" \
--mixed_precision="bf16" \
--gradient_checkpointing \
--remote_text_encoder \
--cache_latents \
--caption_column="caption"\
--do_fp8_training \
--dataset_name="multimodalart/1920-raider-waite-tarot-public-domain" \
--output_dir="tarot_card_Flux2_LoRA" \
--instance_prompt="trcrd tarot card" \
--resolution=1024 \
--train_batch_size=2 \
--guidance_scale=1 \
--gradient_accumulation_steps=1 \
--optimizer="adamW" \
--use_8bit_adam\
--learning_rate=1e-4 \
--report_to="wandb" \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=200 \
--checkpointing_steps=250\
--max_train_steps=1000 \
--rank=8\
--validation_prompt="a trtcrd of a person on a computer, on the computer you see a meme being made with an ancient looking trollface, 'the shitposter' arcana, in the style of TOK a trtcrd, tarot style" \
--validation_epochs=25 \
--seed="0"\
--push_to_hub
LoRA 微调
预训练的 FLUX.2
LoRA 微调后的 FLUX.2
左侧图像由预训练的 FLUX.2 模型生成，右侧图像由 LoRA 生成。
如果你的硬件与 FP8 训练不兼容，你可以使用 bitsandbytes 进行 QLoRA 训练。你首先需要定义一个如下的 `config.json` 文件：
{
"load_in_4bit": true,
"bnb_4bit_quant_type": "nf4"
}
然后将其路径传给 `--bnb_quantization_config_path`：
accelerate launch train_dreambooth_lora_flux2.py \
--pretrained_model_name_or_path="black-forest-labs/FLUX.2-dev" \
--mixed_precision="bf16" \
--gradient_checkpointing \
--remote_text_encoder \
--cache_latents \
--caption_column="caption"\
**--bnb_quantization_config_path="config.json" \**
--dataset_name="multimodalart/1920-raider-waite-tarot-public-domain" \
--output_dir="tarot_card_Flux2_LoRA" \
--instance_prompt="a tarot card" \
--resolution=1024 \
--train_batch_size=2 \
--guidance_scale=1 \
--gradient_accumulation_steps=1 \
--optimizer="adamW" \
--use_8bit_adam\
--learning_rate=1e-4 \
--report_to="wandb" \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=200 \
--max_train_steps=1000 \
--rank=8\
--validation_prompt="a trtcrd of a person on a computer, on the computer you see a meme being made with an ancient looking trollface, 'the shitposter' arcana, in the style of TOK a trtcrd, tarot style" \
--seed="0"
资源
- FLUX.2 发布公告
- Diffusers 文档
- FLUX.2 官方演示
- Hub 上的 FLUX.2
- FLUX.2 原始代码库
Welcome FLUX.2 - BFL’s new open image generation model 🤗
Published November 25, 2025
FLUX.2 is the recent series of image generation models from Black Forest Labs, preceded by the Flux.1 series. It is an entirely new model with a new architecture and pre-training done from scratch! 
In this post, we discuss the key changes introduced in FLUX.2, performing inference with it under various setups, and LoRA fine-tuning.
🚨 FLUX.2 is not meant to be a drop-in replacement of FLUX.1, but a new image generation and editing model.
Table of contents
FLUX.2: A Brief Introduction
FLUX.2 can be used for both image-guided and text-guided image generation. Furthermore, it can take multiple images as reference inputs, while producing the final output image. Below, we briefly discuss the key changes introduced in FLUX.2.
Text encoder
First, instead of two text encoders as in Flux.1, it uses a single text encoder — Mistral Small 3.1. Using a single text encoder greatly simplifies the process of computing prompt embeddings. The pipeline allows for a max_sequence_length of 512. Instead of using a single-layer output for the prompt embedding, FLUX.2 stacks outputs from intermediate layers, which have been known to be more beneficial.
DiT
FLUX.2 follows the same general multimodel diffusion transformer (MM-DiT) + parallel DiT architecture as Flux.1. As a refresher, MM-DiT blocks first process the image latents and conditioning text in separate streams, only joining the two together for the attention operation, and are thus referred to as “double-stream” blocks. The parallel blocks then operate on the concatenated image and text streams and can be regarded as “single-stream” blocks.
The key DiT changes from Flux.1 to FLUX.2 are as follows:
Time and guidance information (in the form of AdaLayerNorm-Zero modulation parameters) is shared across all double-stream and single-stream transformer blocks, respectively, rather than having individual modulation parameters for each block as in Flux.1.
None of the layers in the model use
biasparameters. In particular, neither the attention nor feedforward (FF) sub-blocks of either transformer block usebiasparameters in any of their layers.In Flux.1, the single-stream transformer blocks fused the attention output projection with the FF output projection. FLUX.2 single-stream blocks also fuse the attention QKV projections with the FF input projection, creating a fully parallel transformer block:
Figure taken from the ViT-22B paper.
Note that compared to the ViT-22B block depicted above, FLUX.2 uses a SwiGLU-style MLP activation rather than a GELU activation (and also doesn’t use bias parameters).
- A larger proportion of the transformer blocks in FLUX.2 are single-stream blocks (
8double-stream blocks to48single-stream blocks, compared to19/38for Flux.1). This also means that single-stream blocks make up a larger proportion of the DiT parameters:Flux.1[dev]-12Bhas ~54% of its total parameters in the double-stream blocks, whereasFLUX.2[dev]-32Bhas ~24% of its parameters in the double-stream blocks (and ~73% in the single-stream blocks).
Misc
- A new Autoencoder aka
AutoencoderKLFlux2 - Better way to incorporate resolution-dependent timestep schedules
Inference With Diffusers
FLUX.2 uses a larger DiT and Mistral3 Small as its text encoder. When used together without any kind of offloading, the inference takes more than 80GB VRAM. In the following sections, we show how to perform inference with FLUX.2 in more accessible ways, under various system-level constraints.
Installation and Authentication
Before you try out the following code snippets, make sure you have installed diffusers from main and have run hf auth login.
pip uninstall diffusers -y && pip install git+https://github.com/huggingface/diffusers -U
Regular Inference
from diffusers import Flux2Pipeline
import torch
repo_id = "black-forest-labs/FLUX.2-dev"
pipe = Flux2Pipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
pipe.enable_model_cpu_offload()
image = pipe(
prompt="dog dancing near the sun",
num_inference_steps=50, # 28 is a good trade-off
guidance_scale=4,
height=1024,
width=1024
).images[0]
The above code snippet was tested on an H100, and it isn’t sufficient to run inference on it without CPU offloading. With CPU offloading enabled, this setup takes ~62GB to run.
Users who have access to Hopper-series GPUs can take advantage of Flash Attention 3 to speed up inference:
from diffusers import Flux2Pipeline
import torch
repo_id = "black-forest-labs/FLUX.2-dev"
pipe = Flux2Pipeline.from_pretrained(path, torch_dtype=torch.bfloat16)
+ pipe.transformer.set_attention_backend("_flash_3_hub")
pipe.enable_model_cpu_offload()
image = pipe(
prompt="dog dancing near the sun",
num_inference_steps=50,
guidance_scale=2.5,
height=1024,
width=1024
).images[0]
You can check out the supported attention backends (we have many!) here.
Resource-constrained
Using 4-bit quantization
Using bitsandbytes, we can load the transformer and text encoder models in 4-bit, allowing owners of 24GB GPUs to use the model locally. You can run this snippet on a GPU with ~20 GB of free VRAM.
Unfold
import torch
from transformers import Mistral3ForConditionalGeneration
from diffusers import Flux2Pipeline, Flux2Transformer2DModel
repo_id = "diffusers/FLUX.2-dev-bnb-4bit"
device = "cuda:0"
torch_dtype = torch.bfloat16
transformer = Flux2Transformer2DModel.from_pretrained(
repo_id, subfolder="transformer", torch_dtype=torch_dtype, device_map="cpu"
)
text_encoder = Mistral3ForConditionalGeneration.from_pretrained(
repo_id, subfolder="text_encoder", dtype=torch_dtype, device_map="cpu"
)
pipe = Flux2Pipeline.from_pretrained(
repo_id, transformer=transformer, text_encoder=text_encoder, torch_dtype=torch_dtype
)
pipe.enable_model_cpu_offload()
prompt = "Realistic macro photograph of a hermit crab using a soda can as its shell, partially emerging from the can, captured with sharp detail and natural colors, on a sunlit beach with soft shadows and a shallow depth of field, with blurred ocean waves in the background. The can has the text `BFL Diffusers` on it and it has a color gradient that start with #FF5733 at the top and transitions to #33FF57 at the bottom."
image = pipe(
prompt=prompt,
generator=torch.Generator(device=device).manual_seed(42),
num_inference_steps=50, # 28 is a good trade-off
guidance_scale=4,
).images[0]
image.save("flux2_t2i_nf4.png")
Notice that we're using a repository that contains the NF4-quantized versions of the FLUX.2 DiT and the Mistral text encoder.
Local + remote
Due to the modular design of a Diffusers pipeline, we can isolate modules and work with them in sequence. We decouple the text encoder and deploy it to an Inference Endpoint. This helps us with freeing up the VRAM usage for the DiT and VAE only.
⚠️ To use the remote text encoder, you need to have a valid token. If you are already authenticated, no further action is needed.
The example below uses a combination of local and remote inference. Additionally, we quantize the DiT with NF4 quantization through bitsandbytes.
You can run this snippet on a GPU with 18 GB of VRAM:
Unfold
from diffusers import Flux2Pipeline, Flux2Transformer2DModel
from diffusers import BitsAndBytesConfig as DiffBitsAndBytesConfig
from huggingface_hub import get_token
import requests
import torch
import io
def remote_text_encoder(prompts: str | list[str]):
def _encode_single(prompt: str):
response = requests.post(
"https://remote-text-encoder-flux-2.huggingface.co/predict",
json={"prompt": prompt},
headers={
"Authorization": f"Bearer {get_token()}",
"Content-Type": "application/json"
}
)
assert response.status_code == 200, f"{response.status_code=}"
return torch.load(io.BytesIO(response.content))
if isinstance(prompts, (list, tuple)):
embeds = [_encode_single(p) for p in prompts]
return torch.cat(embeds, dim=0)
return _encode_single(prompts).to("cuda")
repo_id = "black-forest-labs/FLUX.2-dev"
quantized_dit_id = "diffusers/FLUX.2-dev-bnb-4bit"
dit = Flux2Transformer2DModel.from_pretrained(
quantized_dit_id, subfolder="transformer", torch_dtype=torch_dtype, device_map="cpu"
)
pipe = Flux2Pipeline.from_pretrained(
repo_id,
text_encoder=None,
transformer=dit,
torch_dtype=torch.bfloat16,
)
pipe.enable_model_cpu_offload()
print("Running remote text encoder ☁️")
prompt1 = "a photo of a forest with mist swirling around the tree trunks. The word 'FLUX.2' is painted over it in big, red brush strokes with visible texture"
prompt2 = "a photo of a dense forest with rain. The word 'FLUX.2' is painted over it in big, red brush strokes with visible texture"
prompt_embeds = remote_text_encoder([prompt1, prompt2])
print("Done ✅")
out = pipe(
prompt_embeds=prompt_embeds,
generator=torch.Generator(device="cuda").manual_seed(42),
num_inference_steps=50, # 28 is a good trade-off
guidance_scale=4,
height=1024,
width=1024,
)
for idx, image in enumerate(out.images):
image.save(f"flux_out_{idx}.png")
For GPUs with even lower VRAM, we have group_offloading, which allows GPUs with as little as 8GB of free VRAM to use this model. However, you'll need 32GB of free RAM. Alternatively, if you're willing to sacrifice some speed, you can set low_cpu_mem_usage=True to reduce the RAM requirement to just 10GB.
Unfold
import io
import os
import requests
import torch
from diffusers import Flux2Pipeline, Flux2Transformer2DModel
repo_id = "diffusers/FLUX.2-dev-bnb-4bit"
torch_dtype = torch.bfloat16
device = "cuda"
def remote_text_encoder(prompts: str | list[str]):
def _encode_single(prompt: str):
response = requests.post(
"https://remote-text-encoder-flux-2.huggingface.co/predict",
json={"prompt": prompt},
headers={"Authorization": f"Bearer {os.environ['HF_TOKEN']}", "Content-Type": "application/json"},
)
assert response.status_code == 200, f"{response.status_code=}"
return torch.load(io.BytesIO(response.content))
if isinstance(prompts, (list, tuple)):
embeds = [_encode_single(p) for p in prompts]
return torch.cat(embeds, dim=0)
return _encode_single(prompts).to("cuda")
transformer = Flux2Transformer2DModel.from_pretrained(
repo_id, subfolder="transformer", torch_dtype=torch_dtype, device_map="cpu"
)
pipe = Flux2Pipeline.from_pretrained(
repo_id,
text_encoder=None,
transformer=transformer,
torch_dtype=torch_dtype,
)
pipe.transformer.enable_group_offload(
onload_device=device,
offload_device="cpu",
offload_type="leaf_level",
use_stream=True,
# low_cpu_mem_usage=True # uncomment for lower RAM usage
)
pipe.to(device)
prompt = "a photo of a forest with mist swirling around the tree trunks. The word 'FLUX.2' is painted over it in big, red brush strokes with visible texture"
prompt_embeds = remote_text_encoder(prompt)
image = pipe(
prompt_embeds=prompt_embeds,
generator=torch.Generator(device=device).manual_seed(42),
num_inference_steps=50,
guidance_scale=4,
height=1024,
width=1024,
).images[0]
You can check out other supported quantization backends here and other memory-saving techniques here.
To check how different quantizations affect an image, you can play with the playground below or access it as standlone in the FLUX.2 Quantization experiments Space
Multiple images as reference
FLUX.2 supports using multiple images as inputs, allowing you to use up to 10 images. However, keep in mind that each additional image will require more VRAM. You can reference the images by index (e.g., image 1, image 2) or by natural language (e.g., the kangaroo, the turtle). For optimal results, the best approach is to use a combination of both methods.
Unfold
import torch
from transformers import Mistral3ForConditionalGeneration
from diffusers import Flux2Pipeline, Flux2Transformer2DModel
from diffusers.utils import load_image
repo_id = "diffusers-internal-dev/new-model-image-final-weights"
device = "cuda:0"
torch_dtype = torch.bfloat16
pipe = Flux2Pipeline.from_pretrained(
repo_id, torch_dtype=torch_dtype
)
pipe.enable_model_cpu_offload()
image_one = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/flux2_blog/kangaroo.png")
image_two = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/flux2_blog/turtle.png")
prompt = "the boxer kangaroo from image 1 and the martial artist turtle from image 2 are fighting in an epic battle scene at a beach of a tropical island, 35mm, depth of field, 50mm lens, f/3.5, cinematic lighting"
image = pipe(
prompt=prompt,
image=[image_one, image_two],
generator=torch.Generator(device=device).manual_seed(42),
num_inference_steps=50,
guidance_scale=2.5,
width=1024,
height=768,
).images[0]
image.save(f"./flux2_t2i.png")
Multi-image input
Advanced Prompting
FLUX.2 supports advanced prompting techniques like structured JSON prompting, precise hex color control, and multi-reference image editing. Aside for the added control, this also allows for flexibility in changing specific attributes while maintaining others overall the same.
For example, let's start with this json as the base schema (taken from the official FLUX.2 prompting guide):
{
"scene": "overall scene description",
"subjects": [
{
"description": "detailed subject description",
"position": "where in frame",
"action": "what they're doing"
}
],
"style": "artistic style",
"color_palette": ["#hex1", "#hex2", "#hex3"],
"lighting": "lighting description",
"mood": "emotional tone",
"background": "background details",
"composition": "framing and layout",
"camera": {
"angle": "camera angle",
"lens": "lens type",
"depth_of_field": "focus behavior"
}
}
Building up on that, let's turn it into a prompt for a shot of a good old fashion walkman on a carpet (simply pass this prompt to your chosen diffusers inference example from above):
prompt = """
{
"scene": "Professional studio product photography setup with soft-textured carpet surface",
"subjects": [
{
"description": "Old silver Walkman placed on a carpet in the middle of an empty room",
"pose": "Stationary, lying flat",
"position": "Center foreground on carpeted surface",
"color_palette": ["brushed silver", "dark gray accents"]
}
],
"style": "Ultra-realistic product photography with commercial quality",
"color_palette": ["brushed silver", "neutral beige", "soft white highlights"],
"lighting": "Three-point softbox setup creating soft, diffused highlights with no harsh shadows",
"mood": "Clean, professional, minimalist",
"background": "Soft-textured carpet surface with subtle studio backdrop suggesting an empty room",
"composition": "rule of thirds",
"camera": {
"angle": "high angle",
"distance": "medium shot",
"focus": "Sharp focus on metallic Walkman textures and physical controls",
"lens-mm": 85,
"f-number": "f/5.6",
"ISO": 200
}
}
"""
Now, let's change the color of the carpet to a specific teal-blue shade (#367588) and add wired headphones plugged into the walkman:
prompt = """
{
"scene": "Professional studio product photography setup with soft-textured carpet surface",
"subjects": [
{
"description": "Old silver Walkman placed on a teal-blue carpet (#367588) in the middle of an empty room, with wired headphones plugged in",
"pose": "Stationary, lying flat",
"position": "Center foreground on carpeted surface",
"color_palette": ["brushed silver", "dark gray accents", "#367588"]
},
{
"description": "Wired headphones connected to the Walkman, cable loosely coiled on the carpet",
"pose": "Stationary",
"position": "Next to and partially in front of the Walkman on the carpet",
"color_palette": ["dark gray", "soft black", "#367588"]
}
],
"style": "Ultra-realistic product photography with commercial quality",
"color_palette": ["brushed silver", "#367588", "neutral beige", "soft white highlights"],
"lighting": "Three-point softbox setup creating soft, diffused highlights with no harsh shadows",
"mood": "Clean, professional, minimalist",
"background": "Soft-textured teal-blue carpet surface (#367588) with subtle studio backdrop suggesting an empty room",
"composition": "rule of thirds",
"camera": {
"angle": "high angle",
"distance": "medium shot",
"focus": "Sharp focus on metallic Walkman textures, wired headphones, and carpet fibers",
"lens-mm": 85,
"f-number": "f/5.6",
"ISO": 200
}
}
"""
The carpet color now matches the hex code provided, and the headphones have been with small changes to the overall scene.
Check out the official prompting guide for more examples and details.
LoRA fine-tuning
Being both a text-to-image and an image-to-image model, FLUX.2 makes the perfect fine-tuning candidate for many use-cases! However, as inference alone takes more than 80GB of VRAM, LoRA fine-tuning is even more challenging to run on consumer GPUs. To squeeze in as much memory saving as we can, we utilize some of the inference optimizations described above for training as well, together with shared memory saving techniques, to substantially reduce memory consumption. To train it, you can use either the diffusers code below or Ostris' AI Toolkit.
We provide both text-to-image and image-to-image training scripts, for the purpose of this blog will focus on a text-to-image training example.
Memory optimizations for fine-tuning
Many of these techniques complement each other and can be used together to reduce memory consumption further. However, some techniques may be mutually exclusive, so be sure to check before launching a training run.
Unfold to check details on the memory-saving techniques used:
Remote Text Encoder: to leverage the remote text encoding for training, simply pass
--remote_text_encoder. Note that you must either be logged in to your Hugging Face account (hf auth login) OR pass a token with--hub_token.CPU Offloading: by passing
--offloadthe vae and text encoder to will be offloaded to CPU memory and only moved to GPU when needed.Latent Caching: Pre-encode the training images with the vae, and then delete it to free up some memory. To enable
latent_cachingsimply pass--cache_latents.QLoRA: Low Precision Training with Quantization - using 8-bit or 4-bit quantization. You can use the following flags:
- FP8 training with
torchao: enable FP8 training by passing--do_fp8_training. Since we are utilizing FP8 tensor cores, we need CUDA GPUs with compute capability at least 8.9 or greater. If you're looking for memory-efficient training on relatively older cards, we encourage you to check out other trainers likeSimpleTuner,ai-toolkit, etc. - NF4 training with
bitsandbytes: Alternatively, you can use 8-bit or 4-bit quantization withbitsandbytesby passing:---bnb_quantization_config_pathwith a corresponding path to a json file containing your config. see below for more details.
- FP8 training with
Gradient Checkpointing and Accumulation:
--gradient accumulationrefers to the number of updates steps to accumulate before performing a backward/update pass.by passing a value > 1 you can reduce the amount of backward/update passes and hence also memory reqs.* with--gradient checkpointingwe can save memory by not storing all intermediate activations during the forward pass.Instead, only a subset of these activations (the checkpoints) are stored and the rest is recomputed as needed during the backward pass. Note that this comes at the expanse of a slower backward pass.8-bit-Adam Optimizer: When training with
AdamW(doesn't apply toprodigy) You can pass--use_8bit_adamto reduce the memory requirements of training. Make sure to installbitsandbytesif you want to do so.
Please make sure to check out the README for prerequisites before starting training.
For this example, we’ll use multimodalart/1920-raider-waite-tarot-public-domain dataset with the following configuration using FP8 training. Feel free to experiment more with the hyper-parameters and share your results 🤗
accelerate launch train_dreambooth_lora_flux2.py \
--pretrained_model_name_or_path="black-forest-labs/FLUX.2-dev" \
--mixed_precision="bf16" \
--gradient_checkpointing \
--remote_text_encoder \
--cache_latents \
--caption_column="caption"\
--do_fp8_training \
--dataset_name="multimodalart/1920-raider-waite-tarot-public-domain" \
--output_dir="tarot_card_Flux2_LoRA" \
--instance_prompt="trcrd tarot card" \
--resolution=1024 \
--train_batch_size=2 \
--guidance_scale=1 \
--gradient_accumulation_steps=1 \
--optimizer="adamW" \
--use_8bit_adam\
--learning_rate=1e-4 \
--report_to="wandb" \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=200 \
--checkpointing_steps=250\
--max_train_steps=1000 \
--rank=8\
--validation_prompt="a trtcrd of a person on a computer, on the computer you see a meme being made with an ancient looking trollface, 'the shitposter' arcana, in the style of TOK a trtcrd, tarot style" \
--validation_epochs=25 \
--seed="0"\
--push_to_hub
LoRA finetuning
Pre-trained FLUX.2
LoRA fine-tuned FLUX.2
The left image was generated using the pre-trained FLUX.2 model, and the right image was produced the LoRA.
In case your hardware isn’t compatible with FP8 training, you can use QLoRA with bitsandbytes. You first need to define a config.json file like so:
{
"load_in_4bit": true,
"bnb_4bit_quant_type": "nf4"
}
And then pass its path to --bnb_quantization_config_path:
accelerate launch train_dreambooth_lora_flux2.py \
--pretrained_model_name_or_path="black-forest-labs/FLUX.2-dev" \
--mixed_precision="bf16" \
--gradient_checkpointing \
--remote_text_encoder \
--cache_latents \
--caption_column="caption"\
**--bnb_quantization_config_path="config.json" \**
--dataset_name="multimodalart/1920-raider-waite-tarot-public-domain" \
--output_dir="tarot_card_Flux2_LoRA" \
--instance_prompt="a tarot card" \
--resolution=1024 \
--train_batch_size=2 \
--guidance_scale=1 \
--gradient_accumulation_steps=1 \
--optimizer="adamW" \
--use_8bit_adam\
--learning_rate=1e-4 \
--report_to="wandb" \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=200 \
--max_train_steps=1000 \
--rank=8\
--validation_prompt="a trtcrd of a person on a computer, on the computer you see a meme being made with an ancient looking trollface, 'the shitposter' arcana, in the style of TOK a trtcrd, tarot style" \
--seed="0"
Resources
- FLUX.2 announcement post
- Diffusers documentation
- FLUX.2 official demo
- FLUX.2 on the Hub
- FLUX.2 original codebase