### Run CLI Interactively

Source: https://github.com/black-forest-labs/flux2/blob/main/docs/flux2_with_prompt_upsampling.md

Start the CLI interactively to choose your prompt upsampling model. Ensure PYTHONPATH is set.

```bash
export PYTHONPATH=src
python scripts/cli.py
```

--------------------------------

### Load 4-bit Quantized FLUX.2 [dev] with Remote Text Encoder

Source: https://github.com/black-forest-labs/flux2/blob/main/model_cards/FLUX.2-dev.md

Example demonstrating how to load a 4-bit quantized FLUX.2 [dev] model with a remote text encoder for deployment on consumer GPUs. Ensure you have a Hugging Face token configured.

```python
import torch
from diffusers import Flux2Pipeline, Flux2Transformer2DModel
from diffusers.utils import load_image
from huggingface_hub import get_token
import requests
import io

repo_id = "diffusers/FLUX.2-dev-bnb-4bit"
device = "cuda:0"
torch_dtype = torch.bfloat16

def remote_text_encoder(prompts):
    response = requests.post(
        "https://remote-text-encoder-flux-2.huggingface.co/predict",
        json={"prompt": prompts},
        headers={
            "Authorization": f"Bearer {get_token()}",
            "Content-Type": "application/json"
        }
    )
    prompt_embeds = torch.load(io.BytesIO(response.content))

    return prompt_embeds.to(device)

pipe = Flux2Pipeline.from_pretrained(
    repo_id, transformer=transformer, text_encoder=None, torch_dtype=torch_dtype
).to(device)

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_embeds=remote_text_encoder(prompt),
    #image=load_image("https://huggingface.co/spaces/zerogpu-aoti/FLUX.1-Kontext-Dev-fp8-dynamic/resolve/main/cat.png") #optional image input
    generator=torch.Generator(device=device).manual_seed(42),
    num_inference_steps=50, #28 steps can be a good trade-off
    guidance_scale=4,
).images[0]

image.save("flux2_output.png")
```

--------------------------------

### Install FLUX.2 Dependencies

Source: https://github.com/black-forest-labs/flux2/blob/main/README.md

Installs the necessary dependencies for FLUX.2 using pip and a virtual environment. Ensure CUDA 12.9 and Python 3.12 are installed.

```bash
python3.12 -m venv .venv
source .venv/bin/activate
pip install -e . --extra-index-url https://download.pytorch.org/whl/cu129 --no-cache-dir
```

--------------------------------

### Install Diffusers and Dependencies

Source: https://github.com/black-forest-labs/flux2/blob/main/docs/flux2_dev_hf.md

Install the diffusers library from main and upgrade related dependencies. Ensure you have accepted the gating on the FLUX.2-dev repository.

```sh
pip install git+https://github.com/huggingface/diffusers.git
pip install --upgrade transformers accelerate bitsandbytes
```

--------------------------------

### Run FLUX.2 [dev] with 4-bit Quantization and Remote Text Encoder

Source: https://context7.com/black-forest-labs/flux2/llms.txt

This snippet demonstrates how to load and run the FLUX.2 [dev] model with 4-bit quantization for reduced VRAM usage. It utilizes a remote text encoder hosted by Hugging Face to avoid loading the large text encoder locally. Ensure you have installed the necessary libraries and logged in to Hugging Face Hub.

```python
# Install diffusers from main branch
# pip install git+https://github.com/huggingface/diffusers.git
# pip install --upgrade transformers accelerate bitsandbytes
# hf auth login   # accept gating at huggingface.co/black-forest-labs/FLUX.2-dev

import torch
import requests
import io
from diffusers import Flux2Pipeline
from huggingface_hub import get_token

# Option A: 4-bit quantized model + remote text encoder (~18 GB VRAM)
repo_id = "diffusers/FLUX.2-dev-bnb-4bit"
device = "cuda:0"

def remote_text_encoder(prompts):
    response = requests.post(
        "https://remote-text-encoder-flux-2.huggingface.co/predict",
        json={"prompt": prompts},
        headers={"Authorization": f"Bearer {get_token()}", "Content-Type": "application/json"},
    )
    return torch.load(io.BytesIO(response.content)).to(device)

pipe = Flux2Pipeline.from_pretrained(
    repo_id, text_encoder=None, torch_dtype=torch.bfloat16
).to(device)

image = pipe(
    prompt_embeds=remote_text_encoder("a hermit crab using a soda can as its shell"),
    generator=torch.Generator(device=device).manual_seed(42),
    num_inference_steps=50,
    guidance_scale=4,
).images[0]
image.save("flux2_output.png")

# Option B: Full model with CPU offload on H100 (80 GB VRAM)
# pipe = Flux2Pipeline.from_pretrained("black-forest-labs/FLUX.2-dev", torch_dtype=torch.bfloat16)
# pipe.enable_model_cpu_offload()
# image = pipe(prompt="...", num_inference_steps=50, guidance_scale=4).images[0]
```

--------------------------------

### CLI — Interactive generation and editing

Source: https://context7.com/black-forest-labs/flux2/llms.txt

Provides an interactive REPL for image generation and editing. Supports setting configuration keys, using text as prompts, and running generation by pressing Enter. Can be started with specific models, configurations, and options like CPU offloading or API-based upsampling.

```APIDOC
## CLI — Interactive generation and editing

The `scripts/cli.py` script provides a full interactive REPL for generation and editing. All configuration fields can be set with `key=value` syntax; plain text is treated as a prompt; pressing Enter runs generation.

### Usage Examples

```bash
# Install and set model paths (auto-downloads if unset)
export PYTHONPATH=src
export KLEIN_4B_MODEL_PATH="/models/flux-2-klein-4b.safetensors"
export AE_MODEL_PATH="/models/ae.safetensors"

# Start interactive session
python scripts/cli.py

# Or start with a specific model and skip model selection prompt
python scripts/cli.py --model_name flux.2-klein-4b

# Single evaluation (non-interactive, for scripting)
python scripts/cli.py \
  --model_name flux.2-klein-9b-kv \
  --single_eval \
  --prompt "a cat wearing sunglasses" \
  --width 768 \
  --height 768

# With CPU offloading for lower VRAM usage
python scripts/cli.py --model_name flux.2-dev --cpu_offloading

# With API-based prompt upsampling via OpenRouter
export OPENROUTER_API_KEY="sk-or-..."
python scripts/cli.py --upsample_prompt_mode openrouter --openrouter_model mistralai/pixtral-large-2411

# Inside the REPL:
# > prompt="a dragon in a forest"  width=1024  height=1024
# > seed=42  guidance=4.0  num_steps=50
# > input_images="ref1.jpg,ref2.jpg"   # enables image editing
# > match_image_size=0                  # match output to first input's dimensions
```

### Configuration Options

- **`--model_name`**: Specifies the model to use.
- **`--single_eval`**: Enables non-interactive, single evaluation mode.
- **`--prompt`**: Sets the text prompt for generation.
- **`--width`**: Sets the output image width.
- **`--height`**: Sets the output image height.
- **`--cpu_offloading`**: Enables CPU offloading for reduced VRAM usage.
- **`--upsample_prompt_mode`**: Sets the prompt upsampling mode (e.g., `openrouter`).
- **`--openrouter_model`**: Specifies the OpenRouter model to use when `upsample_prompt_mode` is `openrouter`.

**Interactive REPL Commands:**
- `key=value`: Sets configuration parameters (e.g., `prompt="text"`, `width=512`).
- `Enter`: Executes the generation or editing process with current settings.
```

--------------------------------

### Run CLI with flux.2-klein-9b-kv Model

Source: https://github.com/black-forest-labs/flux2/blob/main/docs/flux2_klein_kv_cache.md

Use the CLI to automatically select the KV cache model by specifying 'flux.2-klein-9b-kv'. Provide reference images and a prompt for editing.

```bash
PYTHONPATH=src python scripts/cli.py --model_name flux.2-klein-9b-kv
```

```bash
> input_images="ref1.jpg,ref2.jpg"
> prompt="a cat wearing sunglasses"
> run
```

--------------------------------

### Run FLUX.2 CLI

Source: https://github.com/black-forest-labs/flux2/blob/main/README.md

Launches the FLUX.2 command-line interface for interactive text-to-image generation and image editing. Ensure PYTHONPATH is set correctly.

```bash
PYTHONPATH=src python scripts/cli.py
```

--------------------------------

### Flux2 CLI for Interactive Generation

Source: https://context7.com/black-forest-labs/flux2/llms.txt

Provides an interactive REPL for generation and editing. Configuration fields can be set with key=value syntax. Plain text is treated as a prompt; Enter runs generation. Supports model path export, single evaluation, CPU offloading, and API-based upsampling.

```bash
# Install and set model paths (auto-downloads if unset)
export PYTHONPATH=src
export KLEIN_4B_MODEL_PATH="/models/flux-2-klein-4b.safetensors"
export AE_MODEL_PATH="/models/ae.safetensors"

# Start interactive session
python scripts/cli.py

# Or start with a specific model and skip model selection prompt
python scripts/cli.py --model_name flux.2-klein-4b

# Single evaluation (non-interactive, for scripting)
python scripts/cli.py \
  --model_name flux.2-klein-9b-kv \
  --single_eval \
  --prompt "a cat wearing sunglasses" \
  --width 768 \
  --height 768

# With CPU offloading for lower VRAM usage
python scripts/cli.py --model_name flux.2-dev --cpu_offloading

# With API-based prompt upsampling via OpenRouter
export OPENROUTER_API_KEY="sk-or-..."
python scripts/cli.py --upsample_prompt_mode openrouter --openrouter_model mistralai/pixtral-large-2411

# Inside the REPL:
# > prompt="a dragon in a forest"  width=1024  height=1024
# > seed=42  guidance=4.0  num_steps=50
# > input_images="ref1.jpg,ref2.jpg"   # enables image editing
# > match_image_size=0                  # match output to first input's dimensions
```

--------------------------------

### Enable Local Prompt Upsampling

Source: https://github.com/black-forest-labs/flux2/blob/main/docs/flux2_with_prompt_upsampling.md

Enable local prompt upsampling using the `--upsample_prompt_mode=local` flag. This method uses Mistral-Small-3.2-24B-Instruct-2506 and does not require API keys.

```bash
export PYTHONPATH=src
python scripts/cli.py --upsample_prompt_mode=local
```

--------------------------------

### Upsample Prompt with OpenRouterAPIClient

Source: https://context7.com/black-forest-labs/flux2/llms.txt

Uses multimodal models on OpenRouter for prompt upsampling. Sends images as base64-encoded PNG data URIs. Generally produces higher-quality expansions than local methods. API failures return the original prompt unchanged.

```python
import os
from PIL import Image
from flux2.openrouter_api_client import OpenRouterAPIClient

os.environ["OPENROUTER_API_KEY"] = "sk-or-..."   # or set in shell

client = OpenRouterAPIClient(
    model="mistralai/pixtral-large-2411",
    sampling_params={"temperature": 0.15},
    max_tokens=512,
)

# Text-only upsampling
result = client.upsample_prompt(
    txt=["a photo of a robot chef cooking pasta"],
    img=None,
)
print(result[0])
# "A hyperrealistic photograph of a sleek, chrome-plated humanoid robot chef ..."

# Image-conditioned upsampling
ref = Image.open("kitchen.jpg")
result_i2i = client.upsample_prompt(
    txt=["make it look futuristic"],
    img=[[ref]],   # list[list[PIL.Image]]
)
print(result_i2i[0])

# Error handling: on API failure, the original prompt is returned unchanged
result_safe = client.upsample_prompt(txt=["a landscape"], img=None)
assert len(result_safe) == 1  # always returns same-length list
```

--------------------------------

### Run CLI with OpenRouter Upsampling

Source: https://github.com/black-forest-labs/flux2/blob/main/docs/flux2_with_prompt_upsampling.md

Run the CLI with API-based prompt upsampling enabled via OpenRouter. You can specify a model using --openrouter_model.

```bash
export PYTHONPATH=src
python scripts/cli.py --upsample_prompt_mode=openrouter
```

--------------------------------

### Local prompt upsampling with `Mistral3SmallEmbedder.upsample_prompt`

Source: https://context7.com/black-forest-labs/flux2/llms.txt

Use `upsample_prompt` to expand prompts using `Mistral-Small-3.2-24B-Instruct-2506`. Supports text-only (T2I) and image-conditioned (I2I) modes. Reuses the existing text encoder, avoiding extra model downloads.

```python
from flux2.util import load_text_encoder
from PIL import Image

# Uses the same Mistral model already loaded for flux.2-dev
enc = load_text_encoder("flux.2-dev", device="cuda")
enc.eval()

# Text-to-image upsampling
short_prompt = "a meme about AI taking over"
upsampled = enc.upsample_prompt([short_prompt], img=None, temperature=0.15)
print(upsampled[0])
# "A split-panel meme: left side shows a frantic programmer surrounded by blinking..."

# Image-to-image upsampling (interprets what's in the reference image)
ref = Image.open("sketch.jpg")
upsampled_i2i = enc.upsample_prompt(
    ["add dramatic studio lighting"],
    img=[[ref]],         # list of per-prompt image lists
    temperature=0.15,
)
print(upsampled_i2i[0])
```

--------------------------------

### OpenRouterAPIClient.upsample_prompt

Source: https://context7.com/black-forest-labs/flux2/llms.txt

Upsamples text prompts using multimodal models available on OpenRouter. Supports text-only and image-conditioned upsampling. Returns higher-quality expansions compared to local methods. In case of API failure, the original prompt is returned.

```APIDOC
## `OpenRouterAPIClient.upsample_prompt` — API-based prompt upsampling

Uses any multimodal model available on [OpenRouter](https://openrouter.ai/) (e.g., `mistralai/pixtral-large-2411`, `qwen/qwen3-vl-235b-a22b-instruct`) to upsample prompts. Sends images as base64-encoded PNG data URIs. Generally produces higher-quality expansions than the local method.

### Parameters

- **txt** (list[str]) - Required - A list of text prompts to upsample.
- **img** (list[list[PIL.Image]] | None) - Optional - A list of lists of PIL Images for image-conditioned upsampling. If None, performs text-only upsampling.

### Request Example

```python
import os
from PIL import Image
from flux2.openrouter_api_client import OpenRouterAPIClient

os.environ["OPENROUTER_API_KEY"] = "sk-or-..."   # or set in shell

client = OpenRouterAPIClient(
    model="mistralai/pixtral-large-2411",
    sampling_params={"temperature": 0.15},
    max_tokens=512,
)

# Text-only upsampling
result = client.upsample_prompt(
    txt=["a photo of a robot chef cooking pasta"],
    img=None,
)
print(result[0])

# Image-conditioned upsampling
ref = Image.open("kitchen.jpg")
result_i2i = client.upsample_prompt(
    txt=["make it look futuristic"],
    img=[[ref]],   # list[list[PIL.Image]]
)
print(result_i2i[0])
```

### Response

- **list[str]** - A list of upsampled prompts. The length of the list is always the same as the input `txt` list. In case of API failure, the original prompts are returned.
```

--------------------------------

### Authenticate with Hugging Face CLI

Source: https://github.com/black-forest-labs/flux2/blob/main/docs/flux2_dev_hf.md

Log in to Hugging Face using the CLI for authentication. This is required before running inference.

```sh
hf auth login
```

--------------------------------

### Set OpenRouter API Key

Source: https://github.com/black-forest-labs/flux2/blob/main/docs/flux2_with_prompt_upsampling.md

Set your OpenRouter API key as an environment variable before using API-based prompt upsampling.

```bash
export OPENROUTER_API_KEY="<api_key>"
```

--------------------------------

### Set FLUX.2 Environment Variables

Source: https://github.com/black-forest-labs/flux2/blob/main/README.md

Configures environment variables for FLUX.2 model paths. If not set, weights will be downloaded automatically. Replace placeholders with actual paths.

```bash
export FLUX2_MODEL_PATH="<flux2_path>"
export AE_MODEL_PATH="<ae_path>"
export KLEIN_4B_MODEL_PATH="<klein_4b_path>"
export KLEIN_4B_BASE_MODEL_PATH="<klein_4b_base_path>"
export KLEIN_9B_MODEL_PATH="<klein_9b_path>"
export KLEIN_9B_KV_MODEL_PATH="<klein_9b_kv_path>"
export KLEIN_9B_BASE_MODEL_PATH="<klein_9b_base_path>"
```

--------------------------------

### Set KV Cache Model Path Environment Variable

Source: https://github.com/black-forest-labs/flux2/blob/main/docs/flux2_klein_kv_cache.md

Set the KLEIN_9B_KV_MODEL_PATH environment variable to point to the model file. The AE_MODEL_PATH is optional and will auto-download if not set.

```bash
export KLEIN_9B_KV_MODEL_PATH="/path/to/flux-2-klein-9b-kv.safetensors"
export AE_MODEL_PATH="/path/to/ae.safetensors"  # optional, auto-downloads if unset
```

--------------------------------

### Flux2 Pipeline with CPU Offload for High VRAM

Source: https://github.com/black-forest-labs/flux2/blob/main/docs/flux2_dev_hf.md

Use `pipe.enable_model_cpu_offload()` for GPUs with more than 80GB of VRAM. For H200, B200, or larger cards, direct `pipe.to(device)` is recommended instead of CPU offload.

```python
import torch
from diffusers import Flux2Pipeline
from diffusers.utils import load_image

repo_id = "black-forest-labs/FLUX.2-dev"
device = "cuda:0"
torch_dtype = torch.bfloat16

pipe = Flux2Pipeline.from_pretrained(
    repo_id, torch_dtype=torch_dtype
)
pipe.enable_model_cpu_offload() #no need to do cpu offload for >80G VRAM carts like H200, B200, etc. and do a `pipe.to(device)` instead

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."

#cat_image = load_image("https://huggingface.co/spaces/zerogpu-aoti/FLUX.1-Kontext-Dev-fp8-dynamic/resolve/main/cat.png")
image = pipe(
    prompt=prompt,
    #image=[cat_image] #multi-image input
    generator=torch.Generator(device=device).manual_seed(42),
    num_inference_steps=50,
    guidance_scale=4,
).images[0]

image.save("flux2_output.png")
```

--------------------------------

### KV-cache-accelerated denoising with `denoise_cached`

Source: https://context7.com/black-forest-labs/flux2/llms.txt

Use `denoise_cached` for accelerated denoising with KV-cache. Step 0 extracts the cache, and subsequent steps reuse it for efficiency. Requires `flux.2-klein-9b-kv` model.

```python
import torch
from flux2.util import load_flow_model, load_ae, load_text_encoder
from flux2.sampling import (
    batched_prc_img, batched_prc_txt, denoise_cached,
    encode_image_refs, get_schedule, scatter_ids,
)
from einops import rearrange
from PIL import Image

model = load_flow_model("flux.2-klein-9b-kv", device="cuda")
ae    = load_ae("flux.2-klein-9b-kv", device="cuda")
enc   = load_text_encoder("flux.2-klein-9b-kv", device="cuda")
model.eval(); ae.eval(); enc.eval()

ref_img = Image.open("cat.jpg")
width, height = 768, 768

with torch.no_grad():
    ctx = enc(["a cat wearing a party hat"]).to(torch.bfloat16)
    ctx, ctx_ids = batched_prc_txt(ctx)

    ref_tokens, ref_ids = encode_image_refs(ae, [ref_img])  # required for KV cache

    shape = (1, 128, height // 16, width // 16)
    gen = torch.Generator(device="cuda").manual_seed(7)
    randn = torch.randn(shape, generator=gen, dtype=torch.bfloat16, device="cuda")
    x, x_ids = batched_prc_img(randn)

    timesteps = get_schedule(num_steps=4, image_seq_len=x.shape[1])

    # KV-cached denoising: step 0 extracts cache, steps 1+ reuse it
    x = denoise_cached(
        model, x, x_ids, ctx, ctx_ids,
        timesteps=timesteps,
        guidance=1.0,
        img_cond_seq=ref_tokens,
        img_cond_seq_ids=ref_ids,
    )

    x = torch.cat(scatter_ids(x, x_ids)).squeeze(2)
    x = ae.decode(x).float().clamp(-1, 1)
    img_np = rearrange(x[0], "c h w -> h w c")
    Image.fromarray((127.5 * (img_np + 1.0)).cpu().byte().numpy()).save("edited_cat.png")
```

--------------------------------

### get_schedule

Source: https://context7.com/black-forest-labs/flux2/llms.txt

Generates an adaptive list of timesteps for the rectified flow denoiser. The schedule is shifted based on image sequence length and number of steps using an empirically-fitted `mu` parameter — longer sequences use a more aggressive shift to avoid over-denoising.

```APIDOC
## get_schedule — Compute denoising timestep schedule

Generates an adaptive list of timesteps for the rectified flow denoiser. The schedule is shifted based on image sequence length and number of steps using an empirically-fitted `mu` parameter — longer sequences use a more aggressive shift to avoid over-denoising.

```python
from flux2.sampling import get_schedule

# For a 768x1360 image (H//16 * W//16 = 48 * 85 = 4080 tokens), 50 steps
timesteps = get_schedule(num_steps=50, image_seq_len=4080)
print(f"Steps: {len(timesteps)}, first: {timesteps[0]:.4f}, last: {timesteps[-1]:.4f}")
# Steps: 51, first: ~0.9993, last: 0.0000

# For fast [klein] distilled models: 4 steps
timesteps_fast = get_schedule(num_steps=4, image_seq_len=1024)
print(timesteps_fast)
# [~0.9991, ~0.7437, ~0.4256, ~0.1569, 0.0]
```
```

--------------------------------

### Standard Denoising Loop with Flux2

Source: https://context7.com/black-forest-labs/flux2/llms.txt

Runs the iterative denoising loop for guidance-distilled models. Optionally appends reference image tokens to the noisy image sequence at each step for editing tasks. Requires loading the flow model, autoencoder, and text encoder.

```python
import torch
from flux2.util import load_flow_model, load_ae, load_text_encoder
from flux2.sampling import batched_prc_img, batched_prc_txt, denoise, encode_image_refs, get_schedule, scatter_ids
from einops import rearrange
from PIL import Image

model = load_flow_model("flux.2-klein-4b", device="cuda")
ae    = load_ae("flux.2-klein-4b", device="cuda")
enc   = load_text_encoder("flux.2-klein-4b", device="cuda")
model.eval(); ae.eval(); enc.eval()

width, height = 512, 512
seed = 42

with torch.no_grad():
    # Encode prompt
    ctx = enc(["a photo of a red apple on a wooden table"]).to(torch.bfloat16)
    ctx, ctx_ids = batched_prc_txt(ctx)

    # Encode reference image (omit for text-to-image)
    ref_tokens, ref_ids = encode_image_refs(ae, [])  # None, None for t2i

    # Initialize noise
    shape = (1, 128, height // 16, width // 16)
    gen = torch.Generator(device="cuda").manual_seed(seed)
    randn = torch.randn(shape, generator=gen, dtype=torch.bfloat16, device="cuda")
    x, x_ids = batched_prc_img(randn)

    timesteps = get_schedule(num_steps=4, image_seq_len=x.shape[1])

    x = denoise(
        model, x, x_ids, ctx, ctx_ids,
        timesteps=timesteps,
        guidance=1.0,
        img_cond_seq=ref_tokens,
        img_cond_seq_ids=ref_ids,
    )

    x = torch.cat(scatter_ids(x, x_ids)).squeeze(2)
    x = ae.decode(x).float().clamp(-1, 1)
    img_np = rearrange(x[0], "c h w -> h w c")
    Image.fromarray((127.5 * (img_np + 1.0)).cpu().byte().numpy()).save("output.png")
```

--------------------------------

### Compute Denoising Timestep Schedule

Source: https://context7.com/black-forest-labs/flux2/llms.txt

Generates an adaptive list of timesteps for the rectified flow denoiser. The schedule is shifted based on image sequence length and number of steps using an empirically-fitted `mu` parameter. Longer sequences use a more aggressive shift to avoid over-denoising.

```python
from flux2.sampling import get_schedule

# For a 768x1360 image (H//16 * W//16 = 48 * 85 = 4080 tokens), 50 steps
timesteps = get_schedule(num_steps=50, image_seq_len=4080)
print(f"Steps: {len(timesteps)}, first: {timesteps[0]:.4f}, last: {timesteps[-1]:.4f}")
# Steps: 51, first: ~0.9993, last: 0.0000

# For fast [klein] distilled models: 4 steps
timesteps_fast = get_schedule(num_steps=4, image_seq_len=1024)
print(timesteps_fast)
# [~0.9991, ~0.7437, ~0.4256, ~0.1569, 0.0]
```

--------------------------------

### denoise

Source: https://context7.com/black-forest-labs/flux2/llms.txt

Runs the iterative denoising loop for guidance-distilled models (all `[klein]` variants and `flux.2-dev`). Optionally appends reference image tokens to the noisy image sequence at each step for editing tasks.

```APIDOC
## denoise — Standard denoising loop

Runs the iterative denoising loop for guidance-distilled models (all `[klein]` variants and `flux.2-dev`). Optionally appends reference image tokens to the noisy image sequence at each step for editing tasks.

```python
import torch
from flux2.util import load_flow_model, load_ae, load_text_encoder
from flux2.sampling import batched_prc_img, batched_prc_txt, denoise, encode_image_refs, get_schedule

model = load_flow_model("flux.2-klein-4b", device="cuda")
ae    = load_ae("flux.2-klein-4b", device="cuda")
enc   = load_text_encoder("flux.2-klein-4b", device="cuda")
model.eval(); ae.eval(); enc.eval()

width, height = 512, 512
seed = 42

with torch.no_grad():
    # Encode prompt
    ctx = enc(["a photo of a red apple on a wooden table"]).to(torch.bfloat16)
    ctx, ctx_ids = batched_prc_txt(ctx)

    # Encode reference image (omit for text-to-image)
    ref_tokens, ref_ids = encode_image_refs(ae, [])  # None, None for t2i

    # Initialize noise
    shape = (1, 128, height // 16, width // 16)
    gen = torch.Generator(device="cuda").manual_seed(seed)
    randn = torch.randn(shape, generator=gen, dtype=torch.bfloat16, device="cuda")
    x, x_ids = batched_prc_img(randn)

    timesteps = get_schedule(num_steps=4, image_seq_len=x.shape[1])

    x = denoise(
        model, x, x_ids, ctx, ctx_ids,
        timesteps=timesteps,
        guidance=1.0,
        img_cond_seq=ref_tokens,
        img_cond_seq_ids=ref_ids,
    )

    from flux2.sampling import scatter_ids
    from einops import rearrange
    x = torch.cat(scatter_ids(x, x_ids)).squeeze(2)
    x = ae.decode(x).float().clamp(-1, 1)
    img_np = rearrange(x[0], "c h w -> h w c")
    from PIL import Image
    Image.fromarray((127.5 * (img_np + 1.0)).cpu().byte().numpy()).save("output.png")
```
```

--------------------------------

### Load FLUX.2 Transformer Model

Source: https://context7.com/black-forest-labs/flux2/llms.txt

Loads a `Flux2` model by name, downloading weights from Hugging Face if the corresponding environment variable is not set. Accepts an optional `debug_mode` flag for fast iteration and a `device` argument for CPU offloading.

```python
import torch
from flux2.util import load_flow_model

# Auto-download from HuggingFace (requires HF access to black-forest-labs/FLUX.2-klein-4B)
model = load_flow_model("flux.2-klein-4b", device="cuda")
model.eval()

# Or set env var to use a local weight file:
# export KLEIN_4B_MODEL_PATH="/path/to/flux-2-klein-4b.safetensors"
# model = load_flow_model("flux.2-klein-4b", device="cuda")

# Debug mode: reduced depth, no weight loading — fast for architecture testing
model_debug = load_flow_model("flux.2-dev", debug_mode=True, device="cuda")

# CPU offloading: load weights to CPU, move to GPU before inference
model_cpu = load_flow_model("flux.2-klein-9b", device="cpu")
model_cpu = model_cpu.to("cuda")  # move when ready to run
```

--------------------------------

### FLUX.2 Inference with 4-bit Text Encoder and Transformer (~20G VRAM)

Source: https://github.com/black-forest-labs/flux2/blob/main/docs/flux2_dev_hf.md

Load both the text-encoder and transformer in 4-bit, offloading the text-encoder to CPU to fit within ~20G VRAM. Requires `pipe.enable_model_cpu_offload()`.

```py
import torch
from diffusers import Flux2Pipeline, AutoModel
from transformers import Mistral3ForConditionalGeneration
from diffusers.utils import load_image

repo_id = "diffusers/FLUX.2-dev-bnb-4bit" #quantized text-encoder and DiT. VAE still in bf16
device = "cuda:0"
torch_dtype = torch.bfloat16

text_encoder = Mistral3ForConditionalGeneration.from_pretrained(
    repo_id, subfolder="text_encoder", torch_dtype=torch.bfloat16, device_map="cpu"
)
dit = AutoModel.from_pretrained(
    repo_id, subfolder="transformer", torch_dtype=torch.bfloat16, device_map="cpu"
)
pipe = Flux2Pipeline.from_pretrained(
    repo_id, text_encoder=text_encoder, transformer=dit, 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."
#cat_image = load_image("https://huggingface.co/spaces/zerogpu-aoti/FLUX.1-Kontext-Dev-fp8-dynamic/resolve/main/cat.png")
image = pipe(
    prompt=prompt,
    #image=[cat_image] #multi-image input
    generator=torch.Generator(device=device).manual_seed(42),
    num_inference_steps=50,
    guidance_scale=4,
).images[0]

image.save("flux2_output.png")
```

--------------------------------

### Load Text Encoder for FLUX.2

Source: https://context7.com/black-forest-labs/flux2/llms.txt

Loads the appropriate text encoder for the chosen model: `Mistral3SmallEmbedder` for `flux.2-dev`, or `Qwen3Embedder` for all `[klein]` variants. The encoder maps text prompts to dense hidden-state embeddings.

```python
from flux2.util import load_text_encoder
import torch

# Mistral-based encoder for FLUX.2 [dev]
text_enc = load_text_encoder("flux.2-dev", device="cuda")
text_enc.eval()

with torch.no_grad():
    ctx = text_enc(["a photo of a golden retriever at sunset"])
    # ctx: Tensor [1, 512, 15360] — batch x seq_len x (3 layers * hidden_dim)
    print(ctx.shape)  # torch.Size([1, 512, 15360])

# Qwen3-based encoder for [klein] models (no upsampling / content filtering support)
klein_enc = load_text_encoder("flux.2-klein-4b", device="cuda")
with torch.no_grad():
    ctx_klein = klein_enc(["a futuristic cityscape at night"])
    print(ctx_klein.shape)  # torch.Size([1, 512, 7680])
```

--------------------------------

### Classifier-Free Guidance denoising with `denoise_cfg`

Source: https://context7.com/black-forest-labs/flux2/llms.txt

Use `denoise_cfg` for denoising with base models. It runs paired unconditional and conditional forward passes and applies the CFG formula. Requires base models like `flux.2-klein-base-4b` or `flux.2-klein-base-9b`.

```python
import torch
from flux2.util import load_flow_model, load_ae, load_text_encoder
from flux2.sampling import batched_prc_img, batched_prc_txt, denoise_cfg, get_schedule, scatter_ids
from einops import rearrange
from PIL import Image

model = load_flow_model("flux.2-klein-base-4b", device="cuda")
ae    = load_ae("flux.2-klein-base-4b", device="cuda")
enc   = load_text_encoder("flux.2-klein-base-4b", device="cuda")
model.eval(); ae.eval(); enc.eval()

prompt = "a surreal painting of a clock melting over a desert landscape"

with torch.no_grad():
    # CFG requires [empty_ctx, prompt_ctx] concatenated on batch dim
    ctx_empty  = enc([""]).to(torch.bfloat16)
    ctx_prompt = enc([prompt]).to(torch.bfloat16)
    ctx = torch.cat([ctx_empty, ctx_prompt], dim=0)
    ctx, ctx_ids = batched_prc_txt(ctx)

    shape = (1, 128, 768 // 16, 768 // 16)
    gen = torch.Generator(device="cuda").manual_seed(0)
    randn = torch.randn(shape, generator=gen, dtype=torch.bfloat16, device="cuda")
    x, x_ids = batched_prc_img(randn)

    timesteps = get_schedule(num_steps=50, image_seq_len=x.shape[1])

    x = denoise_cfg(
        model, x, x_ids, ctx, ctx_ids,
        timesteps=timesteps,
        guidance=4.0,       # CFG scale
    )

    x = torch.cat(scatter_ids(x, x_ids)).squeeze(2)
    x = ae.decode(x).float().clamp(-1, 1)
    img_np = rearrange(x[0], "c h w -> h w c")
    Image.fromarray((127.5 * (img_np + 1.0)).cpu().byte().numpy()).save("base_output.png")
```

--------------------------------

### load_text_encoder

Source: https://context7.com/black-forest-labs/flux2/llms.txt

Loads the text encoder compatible with the specified FLUX.2 model. It uses Mistral-based encoders for `flux.2-dev` and Qwen3-based encoders for `[klein]` variants to convert text prompts into embeddings.

```APIDOC
## load_text_encoder — Load a text encoder

Loads the appropriate text encoder for the chosen model: `Mistral3SmallEmbedder` for `flux.2-dev`, or `Qwen3Embedder` for all `[klein]` variants. The encoder maps text prompts to dense hidden-state embeddings used as conditioning context.

```python
from flux2.util import load_text_encoder
import torch

# Mistral-based encoder for FLUX.2 [dev]
text_enc = load_text_encoder("flux.2-dev", device="cuda")
text_enc.eval()

with torch.no_grad():
    ctx = text_enc(["a photo of a golden retriever at sunset"])
    # ctx: Tensor [1, 512, 15360] — batch x seq_len x (3 layers * hidden_dim)
    print(ctx.shape)  # torch.Size([1, 512, 15360])

# Qwen3-based encoder for [klein] models (no upsampling / content filtering support)
klein_enc = load_text_encoder("flux.2-klein-4b", device="cuda")
with torch.no_grad():
    ctx_klein = klein_enc(["a futuristic cityscape at night"])
    print(ctx_klein.shape)  # torch.Size([1, 512, 7680])
```
```

--------------------------------

### Encode Reference Images for Editing

Source: https://context7.com/black-forest-labs/flux2/llms.txt

Converts a list of PIL images into tokenized latent reference sequences with positional IDs. Handles pixel-capping, RGB conversion, center cropping, and time-offset assignment.

```python
from flux2.util import load_ae
from flux2.sampling import encode_image_refs
from PIL import Image
import torch

ae = load_ae("flux.2-klein-4b", device="cuda")
ae.eval()

```

--------------------------------

### FLUX.2 Inference with Remote Text Encoder (4-bit)

Source: https://github.com/black-forest-labs/flux2/blob/main/docs/flux2_dev_hf.md

Run FLUX.2 with a remote text encoder and a quantized transformer, suitable for ~18G VRAM. Text embeddings are calculated in the cloud.

```py
import torch
from diffusers import Flux2Pipeline
from diffusers.utils import load_image
from huggingface_hub import get_token
import requests
import io

repo_id = "diffusers/FLUX.2-dev-bnb-4bit" #quantized text-encoder and DiT. VAE still in bf16
device = "cuda:0"
torch_dtype = torch.bfloat16

def remote_text_encoder(prompts):
    response = requests.post(
        "https://remote-text-encoder-flux-2.huggingface.co/predict",
        json={"prompt": prompts},
        headers={
            "Authorization": f"Bearer {get_token()}",
            "Content-Type": "application/json"
        }
    )
    prompt_embeds = torch.load(io.BytesIO(response.content))

    return prompt_embeds.to(device)

pipe = Flux2Pipeline.from_pretrained(
    repo_id, text_encoder=None, torch_dtype=torch_dtype
).to(device)

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."

#cat_image = load_image("https://huggingface.co/spaces/zerogpu-aoti/FLUX.1-Kontext-Dev-fp8-dynamic/resolve/main/cat.png")
image = pipe(
    prompt_embeds=remote_text_encoder(prompt),
    #image=[cat_image] #optional multi-image input
    generator=torch.Generator(device=device).manual_seed(42),
    num_inference_steps=50, #28 steps can be a good trade-off
    guidance_scale=4,
).images[0]

image.save("flux2_output.png")
```

--------------------------------

### Content Safety Filtering with Mistral3SmallEmbedder

Source: https://context7.com/black-forest-labs/flux2/llms.txt

Filters text or images for unsafe content using Mistral-Small-3.2-24B-Instruct-2506. Images are pre-filtered by Falconsai/nsfw_image_detection. Returns True if content is flagged as unsafe. Accepts PIL images, file paths, or torch.Tensors.

```python
from flux2.util import load_text_encoder
from PIL import Image

enc = load_text_encoder("flux.2-dev", device="cuda")
enc.eval()

# Text content filtering
is_flagged = enc.test_txt("a photo of a cat")
print(is_flagged)  # False

is_flagged_bad = enc.test_txt("<harmful prompt here>")
print(is_flagged_bad)  # True

# Image content filtering (accepts PIL, path, or torch.Tensor)
img = Image.open("output.png")
is_img_flagged = enc.test_image(img)
print(is_img_flagged)  # False — safe output

# Also accepts a file path
is_img_flagged_path = enc.test_image("suspicious_output.png")

# Integrate into generation pipeline:
if enc.test_txt(prompt):
    print("Prompt flagged — skipping generation")
elif not enc.test_image(generated_image):
    generated_image.save("output.png")
else:
    print("Output flagged — discarding")
```

--------------------------------

### Flux2 Pipeline with Remote Text-Encoder for H100

Source: https://github.com/black-forest-labs/flux2/blob/main/docs/flux2_dev_hf.md

Utilize a remote text-encoder for maximum speed on H100 GPUs. This approach moves the text-encoding process to a separate service, reducing the load on the local GPU.

```python
import torch
from diffusers import Flux2Pipeline
from diffusers.utils import load_image
from huggingface_hub import get_token
import requests
import io

repo_id = "black-forest-labs/FLUX.2-dev"
device = "cuda:0"
torch_dtype = torch.bfloat16

def remote_text_encoder(prompts):
    response = requests.post(
        "https://remote-text-encoder-flux-2.huggingface.co/predict",
        json={"prompt": prompts},
        headers={
            "Authorization": f"Bearer {get_token()}",
            "Content-Type": "application/json"
        }
    )
    assert response.status_code == 200, f"{response.status_code=}"
    prompt_embeds = torch.load(io.BytesIO(response.content))

    return prompt_embeds.to(device)

pipe = Flux2Pipeline.from_pretrained(
    repo_id, text_encoder=None, torch_dtype=torch_dtype
).to(device)

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."

#cat_image = load_image("https://huggingface.co/spaces/zerogpu-aoti/FLUX.1-Kontext-Dev-fp8-dynamic/resolve/main/cat.png")
image = pipe(
    prompt_embeds=remote_text_encoder(prompt),
    #image=[cat_image] #optional multi-image input
    generator=torch.Generator(device=device).manual_seed(42),
    num_inference_steps=50,
    guidance_scale=4,
).images[0]

image.save("flux2_output.png")
```

--------------------------------

### load_ae

Source: https://context7.com/black-forest-labs/flux2/llms.txt

Loads the FLUX.2 AutoEncoder, which is used for encoding images into latent tokens and decoding latent representations back into pixel tensors. Weights are shared across models and downloaded from a default Hugging Face repository unless overridden by an environment variable.

```APIDOC
## load_ae — Load the FLUX.2 AutoEncoder

Loads the `AutoEncoder` used to encode reference images into latent tokens and decode denoised latents back to pixel tensors. Weights are shared across all models and downloaded from `black-forest-labs/FLUX.2-dev` unless `AE_MODEL_PATH` is set.

```python
from flux2.util import load_ae

ae = load_ae("flux.2-klein-4b", device="cuda")
ae.eval()

# Encode a PIL image to latent space
from PIL import Image
import torch, torchvision

img = Image.open("input.jpg").convert("RGB").resize((512, 512))
x = torchvision.transforms.ToTensor()(img)  # [3, H, W] in [0, 1]
x = (2 * x - 1).unsqueeze(0).to(torch.bfloat16).cuda()  # normalize to [-1, 1]

with torch.no_grad():
    z = ae.encode(x)       # -> Tensor [1, 128, H//16, W//16]
    x_rec = ae.decode(z)   # -> Tensor [1, 3, H, W] in [-1, 1]
    # Convert back to PIL
    pixel = x_rec.clamp(-1, 1)[0]
    pixel = (127.5 * (pixel + 1.0)).byte().permute(1, 2, 0).cpu().numpy()
    Image.fromarray(pixel).save("reconstructed.png")
```
```