Commit 8cd731e for stable-diffusion.cpp
commit 8cd731e7b4667eb8f40a153315b9a7125d09a58d
Author: leejet <leejet714@gmail.com>
Date: Fri Oct 16 02:42:07 2026 +0800
feat: add FLUX.3 Image Klein and Klein Base support (#2143)
diff --git a/README.md b/README.md
index 9020431..abb2dcb 100644
--- a/README.md
+++ b/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)
diff --git a/assets/flux3/flux3-image-klein-base.png b/assets/flux3/flux3-image-klein-base.png
new file mode 100644
index 0000000..4115022
Binary files /dev/null and b/assets/flux3/flux3-image-klein-base.png differ
diff --git a/assets/flux3/flux3-image-klein-edit.png b/assets/flux3/flux3-image-klein-edit.png
new file mode 100644
index 0000000..f392eb8
Binary files /dev/null and b/assets/flux3/flux3-image-klein-edit.png differ
diff --git a/assets/flux3/flux3-image-klein.png b/assets/flux3/flux3-image-klein.png
new file mode 100644
index 0000000..7d45f34
Binary files /dev/null and b/assets/flux3/flux3-image-klein.png differ
diff --git a/docs/edit.md b/docs/edit.md
index 0bd3b56..62ecff2 100644
--- a/docs/edit.md
+++ b/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` |
diff --git a/docs/flux3.md b/docs/flux3.md
new file mode 100644
index 0000000..dd11e40
--- /dev/null
+++ b/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.
diff --git a/src/conditioning/conditioner.hpp b/src/conditioning/conditioner.hpp
index 0062c4e..d5fa675 100644
--- a/src/conditioning/conditioner.hpp
+++ b/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;
diff --git a/src/model.h b/src/model.h
index 0d291ce..baf3a2a 100644
--- a/src/model.h
+++ b/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) ||
diff --git a/src/model/diffusion/flux3.hpp b/src/model/diffusion/flux3.hpp
new file mode 100644
index 0000000..952587b
--- /dev/null
+++ b/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__
diff --git a/src/model/diffusion/model.hpp b/src/model/diffusion/model.hpp
index 957d575..955c527 100644
--- a/src/model/diffusion/model.hpp
+++ b/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,
diff --git a/src/model/te/llm.hpp b/src/model/te/llm.hpp
index 0d00eed..a782ea7 100644
--- a/src/model/te/llm.hpp
+++ b/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);
diff --git a/src/model_loader.cpp b/src/model_loader.cpp
index c3096eb..c24351f 100644
--- a/src/model_loader.cpp
+++ b/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;
diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp
index cdb6e90..daa1517 100644
--- a/src/pipeline/diffusion_engine.cpp
+++ b/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) {
diff --git a/src/pipeline/model_builders.cpp b/src/pipeline/model_builders.cpp
index 10a8eef..4945679 100644
--- a/src/pipeline/model_builders.cpp
+++ b/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),
diff --git a/src/runtime/denoiser.hpp b/src/runtime/denoiser.hpp
index 02ebe57..6fbcb43 100644
--- a/src/runtime/denoiser.hpp
+++ b/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 {