diff --git a/README.md b/README.md index f6d1bceed..1414dfe52 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 000000000..411502269 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 000000000..f392eb86d 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 000000000..7d45f3472 Binary files /dev/null and b/assets/flux3/flux3-image-klein.png differ diff --git a/docs/edit.md b/docs/edit.md index 0bd3b5620..62ecff2af 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 000000000..dd11e4044 --- /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 +``` + + + +``` +.\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 +``` + + + +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 +``` + + + +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 0062c4e3e..d5fa6750d 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 mask; tokenizer->pad_tokens(tokens, &weights, &mask, min_length, max_length); @@ -2254,6 +2258,15 @@ struct LLMEmbedder : public Conditioner { } } + sd::Tensor position_ids; + if (sd_version_is_flux3(version)) { + position_ids = sd::Tensor::zeros({static_cast(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(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(prompt.size()); + prompt += conditioner_params.text; + prompt_attn_range.second = static_cast(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 0d291ce43..baf3a2a34 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 000000000..952587b01 --- /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(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(config.hidden_size, config.hidden_size, false); + } + blocks["mlp_in"] = std::make_shared(config.hidden_size, 2 * config.mlp_dim, false); + blocks["mlp_out"] = std::make_shared(config.mlp_dim, config.hidden_size, false); + blocks["norm"] = std::make_shared(config.hidden_size / heads); + } + + ggml_tensor* linear(GGMLRunnerContext* ctx, const std::string& name, ggml_tensor* x) { + return std::dynamic_pointer_cast(blocks[name])->forward(ctx, x); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe, const std::vector& mods, const std::vector& 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(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(config.context_dim, config.hidden_size, false); + blocks["time_in"] = std::make_shared(256, config.hidden_size, false); + if (config.flow_map) { + blocks["ref_time_in"] = std::make_shared(256, config.hidden_size, false); + } + for (const std::string mode : {"image", "image_cond", "txt"}) { + blocks["early_stream_modulations." + mode] = std::make_shared(config.hidden_size, false, false); + blocks["single_stream_modulations." + mode] = std::make_shared(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(config); + } + if (mode != "txt") { + blocks["emb_in." + mode] = std::make_shared(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(config); + } + blocks["final_layer.image"] = std::make_shared(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(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(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 modes{"image"}; + std::vector 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 single_mods; + std::vector 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(blocks["early_stream_modulations." + mode])->forward(ctx, vec)[0]; + single_mods.push_back(std::dynamic_pointer_cast(blocks["single_stream_modulations." + mode])->forward(ctx, vec)[0]); + lengths.push_back(inputs[m]->ne[1]); + auto h = std::dynamic_pointer_cast(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(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(blocks["single_blocks." + std::to_string(i)])->forward(ctx, joint, pe, single_mods, lengths); + } + auto out = std::dynamic_pointer_cast(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 pe_values; + sd::Tensor zero = sd::Tensor::zeros({1}); + + Flux3Runner(ggml_backend_t backend, const String2TensorStorage& tensors, const std::string& prefix, std::shared_ptr 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& tensors, const std::string& prefix) override { + model.get_param_tensors(tensors, prefix); + } + + sd::Tensor 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(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(h), static_cast(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(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 957d575c1..955c5277a 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* target_timestep = nullptr; +}; + struct AnimaDiffusionExtra { const sd::Tensor* t5_ids = nullptr; const sd::Tensor* t5_weights = nullptr; @@ -149,6 +153,7 @@ using DiffusionExtraParams = std::variant& image_grids, std::set out_layers, - bool return_all_hidden_states = false) { + bool return_all_hidden_states = false, + const sd::Tensor& position_ids = {}) { ggml_cgraph* gf = new_graph_custom(LLM_GRAPH_SIZE); ggml_tensor* input_ids = make_input(input_ids_tensor); std::vector> image_embeds; @@ -2402,6 +2403,10 @@ namespace LLM { } } + if (!position_ids.empty()) { + GGML_ASSERT(position_ids.numel() == static_cast(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& image_grids = {}) { + const std::vector& image_grids = {}, + const sd::Tensor& 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 c3096ebb5..c24351fb3 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 cdb6e907d..daa15175e 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(); + if (sd_version_is_flux3(version)) { + denoiser = std::make_shared(); + } else { + denoiser = std::make_shared(); + } break; } case SEFI_FLOW_PRED: { @@ -2558,6 +2566,13 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptr::from_vector({sigmas[step + 1]}); } + sd::Tensor 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(std::abs(step)); + GGML_ASSERT(target_index > 0 && target_index < sigmas.size()); + flux3_target_timestep = sd::Tensor::from_vector({sigmas[target_index]}); + } sd::Tensor 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 StableDiffusionGGML::sample(const std::shared_ptrmodel_args); + } else if (sd_version_is_flux3(version)) { + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, version, "", false, weight_manager, tokenizers); + result.diffusion = std::make_shared(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(ctx.backends.runtime_backend(SDBackendModule::TE), diff --git a/src/runtime/denoiser.hpp b/src/runtime/denoiser.hpp index 02ebe57e0..6fbcb43df 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 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 sigmas; + sigmas.reserve(n + 1); + for (uint32_t i = 0; i <= n; ++i) { + float u = 1.f - static_cast(i) / n; + sigmas.push_back(shift * u / (1.f + (shift - 1.f) * u)); + } + return sigmas; + } +}; + struct SefiFlowDenoiser; struct SefiFlowDenoiser : public FluxFlowDenoiser {