How It Works

How It Works

Technical notes on what the extension does inside Forge Neo. Everything happens at runtime, in memory: no Forge Neo file, and no file of another extension, is changed.

The model

• Architecture. FLUX 3 Image Klein is a 5.6B rectified-flow transformer on the FLUX.2 autoencoder's latent (128 channels, 1/16 of the resolution). Five blocks per mode run the image, the reference images and the text each on its own, then 28 single-stream blocks run the joint sequence [image, references, text] under per-mode modulation. Each block is parallel attention and SwiGLU under one adaLN and one gate. Text and references sit at time 0; the image at the step's time.
• Port. The transformer follows Black Forest Labs' reference code (black-forest-labs/flux-3-image-klein, src/flux3/dit.py) with ComfyUI's batched layout. The CPU tests run the reference code itself, vendored unmodified, against the port with the same weights: base, references, flow map and fused projections give the same velocity.
• Loading. The detector reads the configuration from the tensor shapes (rows only, so 4-bit builds with packed columns detect correctly). The transformer is built through Forge's own Krea 2 loading path, which brings dtype, quantization (INT8 ConvRot, NVFP4) and offloading with it. The INT8 ConvRot files fuse q, k and v into one projection, and store a ComfyUI attention-kernel setting next to each block, which is dropped. The autoencoder and the text encoder load through Forge's own FLUX.2 and Krea 2 branches.

Text encoder

• The Qwen3-VL 4B language stack that Forge Neo already loads for Krea 2, under the plain instruct template (<|im_start|>user … <|im_end|>, no thinking block). The hidden states of layers 4, 8, … 32 are placed side by side for each token: 20480 features.
• The sequence is padded on the right with <|endoftext|> to a multiple of 80 tokens (at most 4096). The padding is masked out of the encoder's attention but its states are kept, since the transformer reads every text token; the padding sits at rope position 1, as in the reference code.
• Forge batches the positive and negative prompts by repeating the shorter one up to a common length. For FLUX 3 a repeated prompt is a different prompt, so prompts of different lengths run separately.
• Surrounding whitespace is dropped, so "", " " and a prompt holding only LoRA tags are all the empty prompt the base model's CFG uses.

Sampling

• FLUX.3 schedule: the reference grid, linspace(1, 0, n + 1) shifted by 6.93 t / (1 + 5.93 t), ending at 0. The model's own time shift is the same 6.93, so Forge's Simple schedule comes close, rounded to its 1000-step grid.
• Distilled model: a flow map, trained to jump from each sigma to the next one of the schedule. Forge only hands the model the current sigma, so the extension records every schedule Forge builds and looks the next sigma up at each step. With Euler and the FLUX.3 schedule this is exactly the reference sampler.
• Base model: plain flow matching with CFG; Euler with 50 steps and CFG 3.5 is the reference setting.
• The time is embedded at bfloat16 precision whatever the compute dtype, as the reference code does.

GPUs without bfloat16

On cards such as the RTX 20 series, Forge runs the model in float16. Qwen3-VL has a few massive activations - about 11000 in the text features - and three places outgrow float16's 65504, measured on an RTX 2070:

Where	Peak	Without the fix
Text projection (txt_in)	about 1.1 × 10⁵	infinite from the first step: a black picture with BF16 files
Output projections of the text tokens (attn_out + mlp_out) with long prompts	about 2.4 × 10⁵	infinite in the single blocks
Residual stream	about 1.1 × 10⁷	-

With the INT8 files the overflow did not turn into NaN, but into burnt-looking colours with the base model and crunchy textures with the distilled one. Under float16 the extension keeps the residual stream in float32, and runs those projections - linear and without bias - on inputs scaled down by a power of two (2⁻⁸ and 2⁻⁶), summing their outputs in float32. That is exact up to float16 rounding and costs no speed; the BF16 and INT8 files then give the same picture as the reference examples.

References

• img2img: with [Klein] Enable Reference on, the encoded input becomes the first reference, as Forge Neo does for FLUX.2 Klein.
• ImageStitch Integrated only encodes references for the models Forge Neo lists for it (Kontext, FLUX.2 Klein, Qwen-Image Edit, Wan, Anima, Krea 2). The extension adds FLUX 3 to that check when a FLUX 3 model is loaded, without borrowing FLUX.2's flag, which would also swap the Automatic schedule for FLUX.2's.
• References are told apart by their position on the fourth rope axis (-1, -2, …), as in the reference code. In box edits the first reference is ref_image_0, as in Black Forest Labs' cactus example.

Layouts

• The boxes are plain text in the prompt; the model was trained to read them.
• sd-dynamic-prompts reads {x} as a variant with one option and removes the braces, which turns the list of elements into a flat list. The text encoder rebuilds it - every element starts with its "id" - and uses the rebuilt list only when it is valid JSON. The infotext still shows the prompt as sd-dynamic-prompts left it.
• Forge's own prompt syntax leaves the list alone: [a:b:step] needs a number, and <id> tags are not extra networks.
• The Regional Boxes extension draws the boxes and only edits the prompt text for FLUX 3: it adds no generation setting. Write to prompt replaces the last list in the prompt (intact or brace-stripped) or appends one, in Black Forest Labs' formatting; Read from prompt parses it back. Coordinates are relative to the output's width and height.
• Other models. Qwen-Image 2.1 ignores these coordinates (tested: the teapot and vase prompt kept Qwen's own composition). Qwen edits by boxes drawn into the picture itself with a short note ("Inside the red box: …", then "Remove the red box outline"); FLUX 3 does the opposite - it reads the coordinates and treats a drawn box as part of the picture.

LoRAs

• Forge's own LoRA loader matches diffusion_model.<layer> and lora_unet_<layer> keys; the extension adds the transformer. prefix, bare layer names and LyCORIS names.
• On the INT8 ConvRot files, which fuse q, k and v, LoRAs for q_proj, k_proj and v_proj are applied to their third of the fused weight.
• No FLUX 3 LoRA existed when this was written, so the LoRA path is covered by tests but not by a real LoRA yet.

Files

File	Status
flux3_image_klein_sd_int8_convrot, flux3_image_klein_int8_convrot	Tested on an RTX 2070
flux3_image_klein_bf16	Tested on an RTX 2070 (float16 compute)
flux3_image_klein_sd_bf16	Detected; the same model as the INT8 file at full precision
flux3_image_klein_sd_nvfp4, flux3_image_klein_nvfp4	Detected; need an RTX 50 card, not tested
qwen3vl_4b_int8_convrot, Krea 2's qwen3vl_4b_fp8_scaled	Tested; practically the same picture
qwen3vl_4b_bf16	The same model at full precision, not tested
flux2-vae	Tested
qwen3.5_9b_flux3_prompt_upsampler_int8_convrot (prompt upsampler)	Not used: Forge Neo has no step for it

Measured on an RTX 2070 (8 GB)

Run	Time
Distilled, 1024×1024, 16 steps, INT8	about 25 s
Distilled, edit with one reference, 1024×1024	about 60 s
Distilled, three references (ImageStitch), 832×1248	about 85 s
Base, 832×1248, 50 steps, CFG 3.5, INT8	about 165 s (BF16: about 245 s)