feat: add FLUX.3 Image Klein and Klein Base support (#2143)

README.md

@@ -15,6 +15,7 @@ API and command-line option may change frequently.***
 
 ## 🔥Important News
 
+* **2026/10/16** 🚀 stable-diffusion.cpp adds **Day-0 support for FLUX.3 Image Klein / Klein Base**
 * **2026/09/20** 🚀 stable-diffusion.cpp adds **Day-0 support for Qwen-Image-2.1**
 * **2026/08/20** 🚀 stable-diffusion.cpp now supports **LTX-2.5**
 * **2026/08/04** 🚀 stable-diffusion.cpp adds **Day-1 support for MiniMax-H3**
@@ -44,6 +45,7 @@ API and command-line option may change frequently.***
     - [SD3/SD3.5](./docs/sd3.md)
     - [FLUX.1-dev/FLUX.1-schnell](./docs/flux.md)
     - [FLUX.2-dev/FLUX.2-klein](./docs/flux2.md)
+    - [FLUX.3 Image Klein / Klein Base](./docs/flux3.md)
     - [Lens](./docs/lens.md)
     - [Chroma](./docs/chroma.md)
     - [Chroma1-Radiance](./docs/chroma_radiance.md)
@@ -190,6 +192,7 @@ These projects use `stable-diffusion.cpp` as a backend for their image generatio
 - [LocalAI](https://github.com/mudler/LocalAI)
 - [Neural-Pixel](https://github.com/Luiz-Alcantara/Neural-Pixel)
 - [KoboldCpp](https://github.com/LostRuins/koboldcpp)
+- [Kokon](https://github.com/GabrieleLozupone/kokon)
 
 ## Contributors
 

docs/backend.md

@@ -228,17 +228,20 @@ sd-cli -m model.safetensors -p "a cat" --conv-direct off,diffusion=on
 
 - `on` requests direct 2D and 3D convolution.
 - `off` uses the existing im2col-based convolution paths.
-- `auto` enables direct convolution on compute backends with implicit GEMM
-  implementations, without running a benchmark. A metadata-only probe first
-  checks support for a representative F16-weight/F32-input direct convolution
-  graph. Results and decision logs are cached per compute backend and dimension
-  in the context, as with flash attention. Vulkan enables both 2D and 3D.
+- `auto` first checks support for a representative F16-weight/F32-input direct
+  convolution graph. Vulkan then benchmarks direct convolution against im2col
+  using two small shapes per dimension, with warmup, output validation, and five
+  interleaved measurements per shape. Direct convolution must improve the sum of
+  median times by at least 5%; a failed probe defaults to `off`.
+  Other backends retain their implicit GEMM checks without benchmarking:
   CUDA 2D requires a compiled Turing-or-newer MMA architecture; HIP 2D requires
   the supported CDNA MFMA or RDNA WMMA architectures described below for flash
   attention. MUSA 2D has a SIMT implementation. CUDA, HIP, and MUSA 3D also have
   a SIMT implementation, so they do not require MMA. Other backends default to
-  `off`. Each dimension is selected independently, and all compute devices in
-  a module must qualify. Parameter-storage devices do not affect this choice.
+  `off`. Results and decision logs are cached per compute backend and dimension
+  in the context, as with flash attention. Each dimension is selected independently,
+  and all compute devices in a module must qualify. Parameter-storage devices do
+  not affect this choice.
 
 These are backend-level defaults. GGML still chooses the actual kernel based on
 weight type, shape, and hardware; enabling direct convolution does not guarantee

docs/backend_probe.md

@@ -5,7 +5,8 @@ policy stays separate: `src/core/probe/flash_attn.cpp`, for example, defines att
 shapes and its CPU/CUDA/HIP/Vulkan shortcuts, then calls `sd_run_backend_probe` for unknown
 configurations. `src/core/probe/conv_direct.cpp` uses metadata-only probe graphs
 to check direct 2D/3D convolution support, then enables known implicit GEMM
-paths. It does not allocate device tensors, execute kernels, or run measurements.
+paths on non-Vulkan backends. Vulkan compares direct convolution with im2col
+using the shared measurement machinery.
 
 For each representative shape, create a `BackendProbeCase`:
 

docs/edit.md

@@ -14,6 +14,7 @@ Depending on the architecture, different models handle reference images differen
 | [**Qwen Image Edit**](./qwen_image_edit.md) | `qwen` |
 | **Qwen Image LAYERED** | `qwen_layered` |
 | [**Flux.2 [Dev] / Flux.2 [Klein]**](./flux2.md) | `flux2` |
+| [**FLUX.3 Image Klein / Klein Base**](./flux3.md) | `flux2` |
 | [**Boogu Image Edit**](./boogu_image.md) | `z_image_omni` |
 | **Krea2 (Community Edit LoRAs)** | `krea2_ostris_edit` |
 | [**Mage-Flow-Edit**](./mage_flow.md#image-editing) | `mage_flow` |

docs/flux3.md

@@ -0,0 +1,48 @@
+# How to Use
+
+FLUX.3 Image Klein uses a FLUX.3 diffusion transformer, the FLUX.2 VAE, and Qwen3-VL 4B as the LLM text encoder.
+
+## Download weights
+
+- Download FLUX.3 Image Klein
+    - safetensors: https://huggingface.co/black-forest-labs/flux-3-image-klein/tree/main
+- Download FLUX.3 Image Klein Base
+    - safetensors: https://huggingface.co/black-forest-labs/flux-3-image-klein-base/tree/main
+- Download vae
+    - safetensors: https://huggingface.co/black-forest-labs/FLUX.2-dev/tree/main
+    - The `ae.safetensors` included with FLUX.3 can also be used.
+- Download Qwen3-VL 4B
+    - safetensors: https://huggingface.co/Comfy-Org/Krea-2/tree/main/text_encoders
+    - gguf: https://huggingface.co/Qwen/Qwen3-VL-4B-Instruct-GGUF/tree/main
+
+Use the root `model.safetensors` for the diffusion model.
+
+## Examples
+
+### FLUX.3 Image Klein
+
+The distilled model uses 16 steps and CFG 1.0.
+
+```
+.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\flux-3-image-klein.safetensors --llm ..\models\text_encoders\Qwen3-VL-4B-Instruct-Q4_K_M.gguf --vae ..\models\vae\flux2_ae.safetensors -p "a lovely cat holding a sign says 'flux3.cpp'" --cfg-scale 1.0 --steps 16 --sampling-method euler -v --offload-to-cpu
+```
+
+<img src="../assets/flux3/flux3-image-klein.png" />
+
+```
+.\bin\Release\sd-cli.exe --diffusion-model  ..\models\diffusion_models\flux-3-image-klein.safetensors --vae ..\models\vae\flux2_ae.safetensors --llm  ..\models\text_encoders\Qwen3-VL-4B-Instruct-Q4_K_M.gguf -r ..\assets\flux\flux1-dev-q8_0.png -p "change 'flux.cpp' to 'flux3-image-klein.cpp'" --cfg-scale 1.0 --sampling-method euler -v  --offload-to-cpu --steps 16
+```
+
+<img src="../assets/flux3/flux3-image-klein-edit.png" />
+
+Repeat `-r` to use multiple reference images, up to 10. See [Image Editing](./edit.md) for reference-image preprocessing options.
+
+### FLUX.3 Image Klein Base
+
+```
+.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\flux-3-image-klein-base.safetensors --llm ..\models\text_encoders\Qwen3-VL-4B-Instruct-Q4_K_M.gguf --vae ..\models\vae\flux2_ae.safetensors -p "a lovely cat holding a sign says 'flux3.cpp'" --cfg-scale 3.5 --steps 20 --sampling-method euler -v --offload-to-cpu
+```
+
+<img src="../assets/flux3/flux3-image-klein-base.png" />
+
+The default discrete scheduler uses the upstream time shift of 6.93 for both models. Use `--flow-shift` to override it.

docs/performance.md

@@ -5,8 +5,13 @@ eg.:
  - flux 768x768 ~600mb
  - SD2 768x768 ~1400mb
 
-For most backends, it slows things down, but for cuda it generally speeds it up too.
-At the moment, it is only supported for some models and some backends (like cpu, cuda/rocm, metal).
+`--flash-attn on|off|auto` controls flash attention. Its effect on speed depends
+on GPU architecture, backend, driver, model, and tensor shapes. Automatic
+selection uses backend rules and, for some paths, small probes; it does not
+guarantee the fastest end-to-end generation. Compare `on` and `off` with the
+same model and generation settings over repeated runs, separating startup/probe
+time from steady-state generation time. See [backend selection](backend.md)
+for automatic selection rules and per-module overrides.
 
 Watch for:
 ```
@@ -21,6 +26,19 @@ CPU fallback. It excludes weights and cache buffers. Within a runner lifecycle,
 the summary is printed only on the first graph or when backend capacities or the
 segment count change.
 
+## Tune direct convolution.
+
+`--conv-direct on|off|auto` selects direct convolution or the im2col-based
+path. Performance depends on GPU architecture, backend, driver, model, and
+tensor shapes, so enabling direct convolution can either improve or reduce speed.
+
+Automatic selection uses small benchmarks on Vulkan and implicit GEMM checks
+on other backends. It does not guarantee the fastest end-to-end generation.
+Compare `on` and `off` with the same model and generation settings over
+repeated runs, keeping flash attention unchanged and separating startup/probe
+time from steady-state generation time. See [backend selection](backend.md)
+for automatic selection rules and per-module overrides.
+
 ## Run conditional and unconditional CFG in one batched UNet forward.
 
 For UNet models, the conditional and unconditional guidance branches are

src/conditioning/conditioner.hpp

@@ -2006,6 +2006,7 @@ struct LLMEmbedder : public Conditioner {
         } else if (sd_version_is_pid(version)) {
             arch = LLM::LLMArch::GEMMA2_2B;
         } else if (version == VERSION_QWEN_IMAGE_2_1 ||
+                   sd_version_is_flux3(version) ||
                    sd_version_is_lingbot_video(version) ||
                    sd_version_is_ideogram4(version) ||
                    sd_version_is_boogu_image(version) ||
@@ -2201,6 +2202,9 @@ struct LLMEmbedder : public Conditioner {
             weights.insert(weights.end(), suffix_tokens.size(), 1.f);
         }
 
+        if (sd_version_is_flux3(version)) {
+            min_length = std::min(max_length, ((tokens.size() + 79) / 80) * 80);
+        }
         std::vector<float> mask;
         tokenizer->pad_tokens(tokens, &weights, &mask, min_length, max_length);
 
@@ -2254,6 +2258,15 @@ struct LLMEmbedder : public Conditioner {
             }
         }
 
+        sd::Tensor<int32_t> position_ids;
+        if (sd_version_is_flux3(version)) {
+            position_ids = sd::Tensor<int32_t>::zeros({static_cast<int64_t>(tokens.size() * 4)});
+            for (size_t axis = 0; axis < 4; ++axis) {
+                for (size_t i = 0; i < tokens.size(); ++i) {
+                    position_ids[axis * tokens.size() + i] = mask[i] == 0.f ? 1 : static_cast<int32_t>(i);
+                }
+            }
+        }
         auto hidden_states = llm->compute(n_threads,
                                           input_ids,
                                           attention_mask,
@@ -2262,7 +2275,8 @@ struct LLMEmbedder : public Conditioner {
                                           false,
                                           false,
                                           deepstack_image_embeds,
-                                          image_grids);
+                                          image_grids,
+                                          position_ids);
         if (hidden_states.empty()) {
             LOG_ERROR("LLM prompt encoding failed");
             return {};
@@ -3025,6 +3039,15 @@ struct LLMEmbedder : public Conditioner {
 
                 prompt += "<|im_end|>\n<|im_start|>assistant\n";
             }
+        } else if (sd_version_is_flux3(version)) {
+            prompt_template_encode_start_idx = 0;
+            max_length                       = 4096;
+            out_layers                       = {4, 8, 12, 16, 20, 24, 28, 32};
+            prompt                           = "<|im_start|>user\n";
+            prompt_attn_range.first          = static_cast<int>(prompt.size());
+            prompt += conditioner_params.text;
+            prompt_attn_range.second = static_cast<int>(prompt.size());
+            prompt += "<|im_end|>\n<|im_start|>assistant\n";
         } else if (version == VERSION_FLUX2_KLEIN) {
             prompt_template_encode_start_idx = 0;
             min_length                       = 512;

src/core/probe/conv_direct.cpp

@@ -1,26 +1,52 @@
 #include "core/probe/conv_direct.h"
 
+#include <cmath>
 #include <exception>
+#include <utility>
+#include <vector>
 
 #include "core/ggml_extend_backend.h"
 #include "core/probe/backend_probe.h"
 #include "core/util.h"
 
-static void build_conv_direct_probe(BackendProbeGraph& graph, bool conv3d) {
-    constexpr int channels = 64;
+static void build_conv_direct_probe(BackendProbeGraph& graph, bool conv3d, bool direct, int size = 32, int channels = 64) {
     if (conv3d) {
-        auto input   = ggml_new_tensor_4d(graph.ctx, GGML_TYPE_F32, 16, 16, 4, channels);
+        auto input   = ggml_new_tensor_4d(graph.ctx, GGML_TYPE_F32, size / 2, size / 2, 4, channels);
         auto weight  = ggml_new_tensor_4d(graph.ctx, GGML_TYPE_F16, 3, 3, 3, channels * channels);
         graph.inputs = {input, weight};
-        graph.output = ggml_conv_3d_direct(graph.ctx, weight, input, 1, 1, 1, 1, 1, 1, 1, 1, 1, channels, 1, channels);
+        graph.output = direct ? ggml_conv_3d_direct(graph.ctx, weight, input, 1, 1, 1, 1, 1, 1, 1, 1, 1, channels, 1, channels)
+                              : ggml_conv_3d(graph.ctx, weight, input, channels, 1, 1, 1, 1, 1, 1, 1, 1, 1);
     } else {
-        auto input   = ggml_new_tensor_4d(graph.ctx, GGML_TYPE_F32, 32, 32, channels, 1);
+        auto input   = ggml_new_tensor_4d(graph.ctx, GGML_TYPE_F32, size, size, channels, 1);
         auto weight  = ggml_new_tensor_4d(graph.ctx, GGML_TYPE_F16, 3, 3, channels, channels);
         graph.inputs = {input, weight};
-        graph.output = ggml_conv_2d_direct(graph.ctx, weight, input, 1, 1, 1, 1, 1, 1);
+        graph.output = direct ? ggml_conv_2d_direct(graph.ctx, weight, input, 1, 1, 1, 1, 1, 1)
+                              : ggml_conv_2d(graph.ctx, weight, input, 1, 1, 1, 1, 1, 1);
     }
 }
 
+static bool initialize_conv_probe(BackendProbeGraph& graph, bool conv3d) {
+    uint32_t state = 42;
+    for (auto* tensor : graph.inputs) {
+        const size_t count = static_cast<size_t>(ggml_nelements(tensor));
+        const float scale  = tensor->type == GGML_TYPE_F16 ? 1.f / std::sqrt(float(graph.inputs[0]->ne[conv3d ? 3 : 2] * (conv3d ? 27 : 9))) : 1.f;
+        std::vector<ggml_fp16_t> values(count);
+        for (auto& value : values) {
+            state = state * 1664525u + 1013904223u;
+            value = ggml_fp32_to_fp16((static_cast<float>(state >> 8) / 16777216.f - 0.5f) * scale);
+        }
+        if (tensor->type == GGML_TYPE_F16) {
+            ggml_backend_tensor_set(tensor, values.data(), 0, count * sizeof(ggml_fp16_t));
+        } else {
+            // im2col rounds inputs to F16; use the same input precision for both paths.
+            std::vector<float> input(count);
+            ggml_fp16_to_fp32_row(values.data(), input.data(), count);
+            ggml_backend_tensor_set(tensor, input.data(), 0, count * sizeof(float));
+        }
+    }
+    return true;
+}
+
 bool sd_probe_conv_direct(ggml_backend_t backend, bool conv3d) {
     if (backend == nullptr) {
         return false;
@@ -33,11 +59,30 @@ bool sd_probe_conv_direct(ggml_backend_t backend, bool conv3d) {
             LOG_INFO("conv%dd direct auto: %s -> off (probe graph unavailable)", dims, name);
             return false;
         }
-        build_conv_direct_probe(capability, conv3d);
+        build_conv_direct_probe(capability, conv3d, true);
         if (!capability.supported()) {
             LOG_INFO("conv%dd direct auto: %s -> off (unsupported probe graph)", dims, name);
             return false;
         }
+        if (sd_backend_is(backend, "Vulkan")) {
+            std::vector<BackendProbeCase> cases;
+            for (const auto& shape : {std::pair<int, int>{32, 64}, {16, 128}}) {
+                BackendProbeCase test;
+                test.build = [conv3d, shape](BackendProbeGraph& graph, bool direct) {
+                    build_conv_direct_probe(graph, conv3d, direct, shape.first, shape.second);
+                };
+                test.initialize = [conv3d](BackendProbeGraph& graph) { return initialize_conv_probe(graph, conv3d); };
+                cases.push_back(std::move(test));
+            }
+            const auto result = sd_run_backend_probe(backend, 0, cases);
+            if (!result.completed) {
+                LOG_WARN("conv%dd direct auto: %s -> off (%s)", dims, name, result.error.c_str());
+                return false;
+            }
+            LOG_INFO("conv%dd direct auto: %s -> %s (probe: im2col %.3f ms, direct %.3f ms)",
+                     dims, name, result.candidate_faster ? "on" : "off", result.reference_ms, result.candidate_ms);
+            return result.candidate_faster;
+        }
         if (sd_backend_supports_conv_implicit_gemm(backend, conv3d)) {
             LOG_INFO("conv%dd direct auto: %s -> on (implicit GEMM)", dims, name);
             return true;

src/model.h

@@ -66,6 +66,7 @@ enum SDVersion {
     VERSION_MING_IMAGE,
     VERSION_Z_IMAGE_L2P,
     VERSION_IRIS,
+    VERSION_FLUX3,
     VERSION_COUNT,
 };
 
@@ -118,6 +119,10 @@ static inline bool sd_version_is_flux(SDVersion version) {
     return false;
 }
 
+static inline bool sd_version_is_flux3(SDVersion version) {
+    return version == VERSION_FLUX3;
+}
+
 static inline bool sd_version_is_flux2(SDVersion version) {
     if (version == VERSION_FLUX2 || version == VERSION_FLUX2_KLEIN) {
         return true;
@@ -280,7 +285,13 @@ static inline bool sd_version_uses_flux_vae(SDVersion version) {
 }
 
 static inline bool sd_version_uses_flux2_vae(SDVersion version) {
-    if (sd_version_is_flux2(version) || sd_version_is_ernie_image(version) || sd_version_is_lens(version) || sd_version_is_ideogram4(version) || sd_version_is_sefi_image(version) || sd_version_is_llada_image(version)) {
+    if (sd_version_is_flux2(version) ||
+        sd_version_is_flux3(version) ||
+        sd_version_is_ernie_image(version) ||
+        sd_version_is_lens(version) ||
+        sd_version_is_ideogram4(version) ||
+        sd_version_is_sefi_image(version) ||
+        sd_version_is_llada_image(version)) {
         return true;
     }
     return false;
@@ -311,6 +322,7 @@ static inline bool sd_version_is_inpaint(SDVersion version) {
 static inline bool sd_version_is_dit(SDVersion version) {
     if (sd_version_is_flux(version) ||
         sd_version_is_flux2(version) ||
+        sd_version_is_flux3(version) ||
         sd_version_is_ltxav(version) ||
         sd_version_is_minimax_h3(version) ||
         sd_version_is_sd3(version) ||

src/model/common/ggml_block.hpp

@@ -14,6 +14,7 @@
 
 #include "core/ggml_extend.h"
 #include "core/ggml_runner.h"
+#include "core/ggml_extend_backend.h"
 #include "model.h"
 
 class GGMLBlock {
@@ -222,7 +223,7 @@ public:
                 if (force_prec_f32) {
                     ggml_prec_set_acc(fp8_matmul, GGML_PREC_F32);
                 }
-                supports_fp8_matmul = ggml_backend_supports_op(ctx->backend, fp8_matmul);
+                supports_fp8_matmul = ggml_backend_supports_op(ctx->backend, fp8_matmul) && !sd_backend_is_cpu(ctx->backend);
             }
             if (!supports_fp8_matmul) {
                 w = ggml_cast(ctx->ggml_ctx, w, GGML_TYPE_BF16);

src/model/diffusion/flux3.hpp

@@ -0,0 +1,256 @@
+#ifndef __SD_MODEL_DIFFUSION_FLUX3_HPP__
+#define __SD_MODEL_DIFFUSION_FLUX3_HPP__
+
+#include "model/diffusion/flux.hpp"
+
+namespace Flux3 {
+
+    struct Flux3Config {
+        int64_t hidden_size = 3072;
+        int64_t context_dim = 20480;
+        int64_t channels    = 128;
+        int64_t mlp_dim     = 9216;
+        int heads           = 24;
+        int depth           = 5;
+        int single_depth    = 28;
+        bool flow_map       = false;
+
+        static Flux3Config detect_from_weights(const String2TensorStorage& tensors, const std::string& prefix) {
+            Flux3Config config;
+            int depth = 0, single_depth = 0;
+            bool detected = false;
+            for (const auto& [name, tensor] : tensors) {
+                if (!starts_with(name, prefix + ".")) {
+                    continue;
+                }
+                std::string key = name.substr(prefix.size() + 1);
+                if (key == "txt_in.weight") {
+                    detected           = true;
+                    config.context_dim = tensor.ne[0];
+                    config.hidden_size = tensor.ne[1];
+                } else if (key == "emb_in.image.weight") {
+                    config.channels = tensor.ne[0];
+                } else if (key == "single_blocks.0.mlp_out.weight") {
+                    config.mlp_dim = tensor.ne[0];
+                } else if (key == "ref_time_in.in_layer.weight") {
+                    config.flow_map = true;
+                }
+                if (starts_with(key, "txt_mode_blocks.")) {
+                    depth = std::max(depth, atoi(key.substr(16).c_str()) + 1);
+                }
+                if (starts_with(key, "single_blocks.")) {
+                    single_depth = std::max(single_depth, atoi(key.substr(14).c_str()) + 1);
+                }
+            }
+            if (depth > 0)
+                config.depth = depth;
+            if (single_depth > 0)
+                config.single_depth = single_depth;
+            config.heads = static_cast<int>(config.hidden_size / 128);
+            if (detected) {
+                LOG_VERBOSE("flux3: hidden_size = %" PRId64 ", context_dim = %" PRId64 ", depth = %d, single_depth = %d, flow_map = %s",
+                            config.hidden_size, config.context_dim, config.depth, config.single_depth, config.flow_map ? "true" : "false");
+            }
+            return config;
+        }
+    };
+
+    static ggml_tensor* token_slice(ggml_context* ctx, ggml_tensor* x, int64_t start, int64_t length) {
+        return ggml_view_3d(ctx, x, x->ne[0], length, x->ne[2], x->nb[1], x->nb[2], start * x->nb[1]);
+    }
+
+    struct Block : public GGMLBlock {
+        int heads;
+
+        Block(const Flux3Config& config)
+            : heads(config.heads) {
+            for (const auto* name : {"q_proj", "k_proj", "v_proj", "attn_out"}) {
+                blocks[name] = std::make_shared<Linear>(config.hidden_size, config.hidden_size, false);
+            }
+            blocks["mlp_in"]  = std::make_shared<Linear>(config.hidden_size, 2 * config.mlp_dim, false);
+            blocks["mlp_out"] = std::make_shared<Linear>(config.mlp_dim, config.hidden_size, false);
+            blocks["norm"]    = std::make_shared<Flux::QKNorm>(config.hidden_size / heads);
+        }
+
+        ggml_tensor* linear(GGMLRunnerContext* ctx, const std::string& name, ggml_tensor* x) {
+            return std::dynamic_pointer_cast<Linear>(blocks[name])->forward(ctx, x);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe, const std::vector<Flux::ModulationOut>& mods, const std::vector<int64_t>& lengths) {
+            auto gc            = ctx->ggml_ctx;
+            ggml_tensor* x_mod = nullptr;
+            int64_t offset     = 0;
+            for (size_t i = 0; i < lengths.size(); ++i) {
+                auto part = token_slice(gc, x, offset, lengths[i]);
+                part      = Flux::modulate(gc, ggml_norm(gc, part, 1e-6f), mods[i].shift, mods[i].scale);
+                x_mod     = x_mod ? ggml_concat(gc, x_mod, part, 1) : part;
+                offset += lengths[i];
+            }
+            auto split = [&](ggml_tensor* v) {
+                return ggml_reshape_4d(gc, v, x->ne[0] / heads, heads, x->ne[1], x->ne[2]);
+            };
+            auto norm        = std::dynamic_pointer_cast<Flux::QKNorm>(blocks["norm"]);
+            auto q           = norm->query_norm(ctx, split(linear(ctx, "q_proj", x_mod)));
+            auto k           = norm->key_norm(ctx, split(linear(ctx, "k_proj", x_mod)));
+            auto v           = split(linear(ctx, "v_proj", x_mod));
+            auto attn        = Rope::attention(ctx, q, k, v, pe, nullptr);
+            auto mlp         = ggml_ext_silu_act(gc, linear(ctx, "mlp_in", x_mod));
+            auto update      = ggml_add(gc, linear(ctx, "attn_out", attn), linear(ctx, "mlp_out", mlp));
+            ggml_tensor* out = nullptr;
+            offset           = 0;
+            for (size_t i = 0; i < lengths.size(); ++i) {
+                auto part = ggml_mul(gc, token_slice(gc, update, offset, lengths[i]), mods[i].gate);
+                part      = ggml_add(gc, token_slice(gc, x, offset, lengths[i]), part);
+                out       = out ? ggml_concat(gc, out, part, 1) : part;
+                offset += lengths[i];
+            }
+            return out;
+        }
+    };
+
+    struct Model : public GGMLBlock {
+        Flux3Config config;
+
+        Model(const Flux3Config& config)
+            : config(config) {
+            blocks["txt_in"]  = std::make_shared<Linear>(config.context_dim, config.hidden_size, false);
+            blocks["time_in"] = std::make_shared<Flux::MLPEmbedder>(256, config.hidden_size, false);
+            if (config.flow_map) {
+                blocks["ref_time_in"] = std::make_shared<Flux::MLPEmbedder>(256, config.hidden_size, false);
+            }
+            for (const std::string mode : {"image", "image_cond", "txt"}) {
+                blocks["early_stream_modulations." + mode]  = std::make_shared<Flux::Modulation>(config.hidden_size, false, false);
+                blocks["single_stream_modulations." + mode] = std::make_shared<Flux::Modulation>(config.hidden_size, false, false);
+                const std::string base                      = mode == "txt" ? "txt_mode_blocks." : "content_mode_blocks." + mode + ".";
+                for (int i = 0; i < config.depth; ++i) {
+                    blocks[base + std::to_string(i)] = std::make_shared<Block>(config);
+                }
+                if (mode != "txt") {
+                    blocks["emb_in." + mode] = std::make_shared<Linear>(config.channels, config.hidden_size, false);
+                }
+            }
+            for (int i = 0; i < config.single_depth; ++i) {
+                blocks["single_blocks." + std::to_string(i)] = std::make_shared<Block>(config);
+            }
+            blocks["final_layer.image"] = std::make_shared<Flux::LastLayer>(config.hidden_size, 1, config.channels, false, false);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* image, ggml_tensor* text, ggml_tensor* reference, ggml_tensor* timestep, ggml_tensor* zero, ggml_tensor* target, ggml_tensor* pe) {
+            auto gc             = ctx->ggml_ctx;
+            const int64_t width = image->ne[0], height = image->ne[1];
+            auto time_in   = std::dynamic_pointer_cast<Flux::MLPEmbedder>(blocks["time_in"]);
+            auto image_vec = time_in->forward(ctx, ggml_ext_timestep_embedding(gc, timestep, 256, 10000, 1000.f));
+            auto clean_vec = time_in->forward(ctx, ggml_ext_timestep_embedding(gc, zero, 256, 10000, 1000.f));
+            if (config.flow_map) {
+                auto ref_time = std::dynamic_pointer_cast<Flux::MLPEmbedder>(blocks["ref_time_in"]);
+                auto ref_vec  = ref_time->forward(ctx, ggml_ext_timestep_embedding(gc, target, 256, 10000, 1000.f));
+                image_vec     = ggml_add(gc, image_vec, ref_vec);
+                clean_vec     = ggml_add(gc, clean_vec, ref_vec);
+            }
+            std::vector<std::string> modes{"image"};
+            std::vector<ggml_tensor*> inputs{DiT::patchify(gc, image, 1, 1)};
+            if (reference) {
+                modes.push_back("image_cond");
+                inputs.push_back(reference);
+            }
+            modes.push_back("txt");
+            inputs.push_back(text);
+            std::vector<Flux::ModulationOut> single_mods;
+            std::vector<int64_t> lengths;
+            ggml_tensor* joint = nullptr;
+            int64_t offset     = 0;
+            for (size_t m = 0; m < modes.size(); ++m) {
+                const auto& mode = modes[m];
+                auto vec         = m == 0 ? image_vec : clean_vec;
+                auto early       = std::dynamic_pointer_cast<Flux::Modulation>(blocks["early_stream_modulations." + mode])->forward(ctx, vec)[0];
+                single_mods.push_back(std::dynamic_pointer_cast<Flux::Modulation>(blocks["single_stream_modulations." + mode])->forward(ctx, vec)[0]);
+                lengths.push_back(inputs[m]->ne[1]);
+                auto h                 = std::dynamic_pointer_cast<Linear>(blocks[mode == "txt" ? "txt_in" : "emb_in." + mode])->forward(ctx, inputs[m]);
+                auto mode_pe           = ggml_view_4d(gc, pe, 2, 2, 64, lengths.back(), pe->nb[1], pe->nb[2], pe->nb[3], offset * pe->nb[3]);
+                const std::string base = mode == "txt" ? "txt_mode_blocks." : "content_mode_blocks." + mode + ".";
+                for (int i = 0; i < config.depth; ++i) {
+                    h = std::dynamic_pointer_cast<Block>(blocks[base + std::to_string(i)])->forward(ctx, h, mode_pe, {early}, {lengths.back()});
+                }
+                joint = joint ? ggml_concat(gc, joint, h, 1) : h;
+                offset += lengths.back();
+            }
+            for (int i = 0; i < config.single_depth; ++i) {
+                joint = std::dynamic_pointer_cast<Block>(blocks["single_blocks." + std::to_string(i)])->forward(ctx, joint, pe, single_mods, lengths);
+            }
+            auto out = std::dynamic_pointer_cast<Flux::LastLayer>(blocks["final_layer.image"])->forward(ctx, token_slice(gc, joint, 0, lengths[0]), image_vec);
+            return DiT::unpatchify(gc, out, height, width, 1, 1);
+        }
+    };
+
+    struct Flux3Runner : public DiffusionModelRunner {
+        Flux3Config config;
+        Model model;
+        std::vector<float> pe_values;
+        sd::Tensor<float> zero = sd::Tensor<float>::zeros({1});
+
+        Flux3Runner(ggml_backend_t backend, const String2TensorStorage& tensors, const std::string& prefix, std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : DiffusionModelRunner(backend, prefix, weight_manager),
+              config(Flux3Config::detect_from_weights(tensors, prefix)),
+              model(config) {
+            model.init(params_ctx, tensors, prefix);
+        }
+
+        std::string get_desc() override { return "flux3"; }
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
+            model.get_param_tensors(tensors, prefix);
+        }
+
+        sd::Tensor<float> compute(int n_threads, const DiffusionParams& params) override {
+            if (params.ref_latents && params.ref_image_params.pass_to_dit && params.ref_latents->size() > 10) {
+                LOG_ERROR("FLUX.3 supports at most 10 reference images");
+                return {};
+            }
+            const auto* extra = diffusion_extra_as<Flux3DiffusionExtra>(params);
+            GGML_ASSERT(params.x && params.context && params.timesteps);
+            GGML_ASSERT(!config.flow_map || extra->target_timestep);
+            auto get_graph = [&]() {
+                auto gf   = new_graph_custom(FLUX_GRAPH_SIZE);
+                auto x    = make_input(*params.x);
+                auto text = make_input(*params.context);
+                GGML_ASSERT(x->ne[3] == 1 && text->ne[2] == 1);
+                Rope::Embedding embedding;
+                auto append_image = [&](int64_t w, int64_t h, float index) {
+                    embedding.positions.append_image(static_cast<int>(h), static_cast<int>(w), 1, 0.5f, 0.5f);
+                    for (int64_t y = 0; y < h; ++y) {
+                        for (int64_t x = 0; x < w; ++x) {
+                            embedding.ids.push_back({0.f, y * 0.5f, x * 0.5f, index});
+                        }
+                    }
+                };
+                append_image(x->ne[0], x->ne[1], 0.f);
+                ggml_tensor* references = nullptr;
+                if (params.ref_latents && params.ref_image_params.pass_to_dit) {
+                    int index = 0;
+                    for (const auto& ref : *params.ref_latents) {
+                        auto r = make_input(ref);
+                        append_image(r->ne[0], r->ne[1], -float(++index));
+                        r          = DiT::patchify(compute_ctx, r, 1, 1);
+                        references = references ? ggml_concat(compute_ctx, references, r, 1) : r;
+                    }
+                }
+                for (int64_t i = 0; i < text->ne[1]; ++i) {
+                    embedding.ids.push_back({0.f, 0.f, 0.f, static_cast<float>(i)});
+                }
+                embedding.positions.append_tokens(text->ne[1]);
+                embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, {32, 32, 32, 32}, embedding.layout, &embedding.frequencies);
+                pe_values        = finish_rope_pe(std::move(embedding));
+                auto pe          = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, 64, pe_values.size() / 256);
+                set_backend_tensor_data(pe, pe_values.data());
+                auto runner_ctx = get_context();
+                auto out        = model.forward(&runner_ctx, x, text, references, make_input(*params.timesteps), make_input(zero),
+                                         config.flow_map ? make_input(*extra->target_timestep) : nullptr, pe);
+                ggml_build_forward_expand(gf, out);
+                return gf;
+            };
+            return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), params.x->dim());
+        }
+    };
+
+}  // namespace Flux3
+
+#endif  // __SD_MODEL_DIFFUSION_FLUX3_HPP__

src/model/diffusion/model.hpp

@@ -64,6 +64,10 @@ struct FluxDiffusionExtra {
     float pulid_id_weight               = 1.0f;
 };
 
+struct Flux3DiffusionExtra {
+    const sd::Tensor<float>* target_timestep = nullptr;
+};
+
 struct AnimaDiffusionExtra {
     const sd::Tensor<int32_t>* t5_ids   = nullptr;
     const sd::Tensor<float>* t5_weights = nullptr;
@@ -149,6 +153,7 @@ using DiffusionExtraParams = std::variant<std::monostate,
                                           UNetDiffusionExtra,
                                           SkipLayerDiffusionExtra,
                                           FluxDiffusionExtra,
+                                          Flux3DiffusionExtra,
                                           AnimaDiffusionExtra,
                                           QwenImage21DiffusionExtra,
                                           WanDiffusionExtra,

src/model/te/llm.hpp

@@ -2338,7 +2338,8 @@ namespace LLM {
                                  const DeepStackImageEmbeds& deepstack_image_embeds_tensor,
                                  const std::vector<ImageGrid>& image_grids,
                                  std::set<int> out_layers,
-                                 bool return_all_hidden_states = false) {
+                                 bool return_all_hidden_states           = false,
+                                 const sd::Tensor<int32_t>& position_ids = {}) {
             ggml_cgraph* gf        = new_graph_custom(LLM_GRAPH_SIZE);
             ggml_tensor* input_ids = make_input(input_ids_tensor);
             std::vector<std::pair<int, ggml_tensor*>> image_embeds;
@@ -2402,6 +2403,10 @@ namespace LLM {
                 }
             }
 
+            if (!position_ids.empty()) {
+                GGML_ASSERT(position_ids.numel() == static_cast<int64_t>(input_pos_vec.size()));
+                input_pos_vec = position_ids.values();
+            }
             auto input_pos = ggml_new_tensor_1d(compute_ctx,
                                                 GGML_TYPE_I32,
                                                 input_pos_vec.size());
@@ -2475,7 +2480,8 @@ namespace LLM {
                                   bool return_all_hidden_states                      = false,
                                   bool auto_runner_end                               = true,
                                   const DeepStackImageEmbeds& deepstack_image_embeds = {},
-                                  const std::vector<ImageGrid>& image_grids          = {}) {
+                                  const std::vector<ImageGrid>& image_grids          = {},
+                                  const sd::Tensor<int32_t>& position_ids            = {}) {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(input_ids,
                                    attention_mask,
@@ -2483,7 +2489,8 @@ namespace LLM {
                                    deepstack_image_embeds,
                                    image_grids,
                                    out_layers,
-                                   return_all_hidden_states);
+                                   return_all_hidden_states,
+                                   position_ids);
             };
             return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, auto_runner_end),
                                                    input_ids.dim() + 1);

src/model_loader.cpp

@@ -438,6 +438,11 @@ bool ModelLoader::init_from_diffusers_file(const std::string& file_path, const s
 SDVersion ModelLoader::get_sd_version() const {
     TensorStorage token_embedding_weight, input_block_weight, context_ebedding_weight;
 
+    if (tensor_storage_map.find("model.diffusion_model.emb_in.image.weight") != tensor_storage_map.end() &&
+        tensor_storage_map.find("model.diffusion_model.txt_mode_blocks.0.q_proj.weight") != tensor_storage_map.end()) {
+        return VERSION_FLUX3;
+    }
+
     bool has_multiple_encoders = false;
     bool is_unet               = false;
 

src/pipeline/diffusion_engine.cpp

@@ -111,6 +111,7 @@ const char* model_version_to_str[] = {
     "Ming-Image",
     "Z-Image L2P",
     "Iris-3B",
+    "Flux.3 Image Klein",
 };
 
 static_assert(VERSION_COUNT == sizeof(model_version_to_str) / sizeof(model_version_to_str[0]),
@@ -1514,6 +1515,9 @@ bool StableDiffusionGGML::build_denoiser() {
             } else {
                 default_flow_shift = 3.f;
             }
+        } else if (sd_version_is_flux3(version)) {
+            pred_type          = FLUX_FLOW_PRED;
+            default_flow_shift = 6.93f;
         } else if (sd_version_is_flux(version) ||
                    sd_version_is_flux2(version) ||
                    sd_version_is_longcat(version) ||
@@ -1579,7 +1583,11 @@ bool StableDiffusionGGML::build_denoiser() {
         }
         case FLUX_FLOW_PRED: {
             LOG_INFO("running in Flux FLOW mode");
-            denoiser = std::make_shared<FluxFlowDenoiser>();
+            if (sd_version_is_flux3(version)) {
+                denoiser = std::make_shared<Flux3FlowDenoiser>();
+            } else {
+                denoiser = std::make_shared<FluxFlowDenoiser>();
+            }
             break;
         }
         case SEFI_FLOW_PRED: {
@@ -2558,6 +2566,13 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
         if (sd_version_is_hunyuan_video(version) && step + 1 < sigmas.size()) {
             hunyuan_timestep_r_tensor = sd::Tensor<float>::from_vector({sigmas[step + 1]});
         }
+        sd::Tensor<float> flux3_target_timestep;
+        if (version == VERSION_FLUX3) {
+            // The sign marks intermediate evaluations; the magnitude is the one-based step.
+            const size_t target_index = static_cast<size_t>(std::abs(step));
+            GGML_ASSERT(target_index > 0 && target_index < sigmas.size());
+            flux3_target_timestep = sd::Tensor<float>::from_vector({sigmas[target_index]});
+        }
         sd::Tensor<float> noised_input = x * c_in;
         if (apply_denoise_mask && (version == VERSION_WAN2_2_TI2V || sd_version_is_ltxav(version) || sd_version_is_lingbot_video(version))) {
             noised_input = noised_input * denoise_mask + sampling_init_latent * (1.0f - denoise_mask);
@@ -2634,6 +2649,8 @@ sd::Tensor<float> StableDiffusionGGML::sample(const std::shared_ptr<DiffusionMod
             } else if (sd_version_is_flux(version) || sd_version_is_flux2(version) || sd_version_is_longcat(version) || sd_version_is_sefi_image(version)) {
                 diffusion_params.extra = FluxDiffusionExtra{&guidance_tensor,
                                                             local_skip_layers};
+            } else if (sd_version_is_flux3(version)) {
+                diffusion_params.extra = Flux3DiffusionExtra{flux3_target_timestep.empty() ? nullptr : &flux3_target_timestep};
             } else if (version == VERSION_QWEN_IMAGE_2_1) {
                 diffusion_params.extra = QwenImage21DiffusionExtra{&condition.c_token_types};
             } else if (sd_version_is_anima(version)) {
@@ -3180,6 +3197,8 @@ std::string StableDiffusionGGML::get_default_ref_image_preset(SDVersion version)
         return "longcat";
     } else if (sd_version_is_flux(version)) {
         return "flux_kontext";
+    } else if (sd_version_is_flux3(version)) {
+        return "flux2";
     } else if (sd_version_is_flux2(version) || sd_version_is_sefi_image(version)) {
         return "flux2";
     } else if (version == VERSION_QWEN_IMAGE_LAYERED) {

src/pipeline/model_builders.cpp

@@ -15,6 +15,7 @@
 #include "model/diffusion/control.hpp"
 #include "model/diffusion/ernie_image.hpp"
 #include "model/diffusion/flux.hpp"
+#include "model/diffusion/flux3.hpp"
 #include "model/diffusion/hidream_o1.hpp"
 #include "model/diffusion/hunyuan.hpp"
 #include "model/diffusion/ideogram4.hpp"
@@ -173,6 +174,11 @@ namespace sd::model_builders {
                                                                   version,
                                                                   weight_manager,
                                                                   sd_ctx_params->model_args);
+        } else if (sd_version_is_flux3(version)) {
+            result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),
+                                                               tensor_storage_map, version, "", false, weight_manager, tokenizers);
+            result.diffusion   = std::make_shared<Flux3::Flux3Runner>(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION),
+                                                                    tensor_storage_map, "model.diffusion_model", weight_manager);
         } else if (sd_version_is_flux2(version) || sd_version_is_sefi_image(version)) {
             bool is_chroma     = false;
             result.conditioner = std::make_shared<LLMEmbedder>(ctx.backends.runtime_backend(SDBackendModule::TE),

src/runtime/denoiser.hpp

@@ -1452,6 +1452,29 @@ struct FluxFlowDenoiser : public DiscreteFlowDenoiser {
     }
 };
 
+struct Flux3FlowDenoiser : public FluxFlowDenoiser {
+    float t_to_sigma(float t) override {
+        float u = (t + 1.f) / TIMESTEPS;
+        return shift * u / (1.f + (shift - 1.f) * u);
+    }
+
+    std::vector<float> get_sigmas(uint32_t n, int image_seq_len, scheduler_t scheduler_type, SDVersion version, const char* extra_sample_args = nullptr) override {
+        if (scheduler_type != DISCRETE_SCHEDULER) {
+            return Denoiser::get_sigmas(n, image_seq_len, scheduler_type, version, extra_sample_args);
+        }
+        if (n == 0) {
+            return {};
+        }
+        std::vector<float> sigmas;
+        sigmas.reserve(n + 1);
+        for (uint32_t i = 0; i <= n; ++i) {
+            float u = 1.f - static_cast<float>(i) / n;
+            sigmas.push_back(shift * u / (1.f + (shift - 1.f) * u));
+        }
+        return sigmas;
+    }
+};
+
 struct SefiFlowDenoiser;
 
 struct SefiFlowDenoiser : public FluxFlowDenoiser {