Skip to content

[draft] Attempt to implement the Flux Klein bonsai model #1580

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Juste-Leo2 wants to merge 2 commits into from
Closed

[draft] Attempt to implement the Flux Klein bonsai model #1580

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d57d6d9
backup: FLUX.2 Klein-4B (bonsai) inference support
Juste-Leo2 May 29, 2026
ce0fc14
Force VAE weights to fp32 at init (force_upcast workaround for Flux2/…
Juste-Leo2 May 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions image-studio
Open in desktop
Submodule image-studio added at 3d895c
10 changes: 9 additions & 1 deletion src/auto_encoder_kl.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -700,7 +700,15 @@ struct AutoEncoderKL : public VAE {
}
}
ae = AutoEncoderKLModel(version, decode_only, use_linear_projection, use_video_decoder, tensor_storage_map, prefix);
ae.init(params_ctx, tensor_storage_map, prefix);
{
auto f32_map = tensor_storage_map;
for (auto& [name, storage] : f32_map) {
if (starts_with(name, prefix)) {
storage.expected_type = GGML_TYPE_F32;
}
}
ae.init(params_ctx, f32_map, prefix);
}
}

void set_conv2d_scale(float scale) override {
Expand Down
61 changes: 41 additions & 20 deletions src/flux.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,34 +86,55 @@ namespace Flux {
struct SelfAttention : public GGMLBlock {
public:
int64_t num_heads;
int64_t hidden_size;
bool fused_qkv;

public:
SelfAttention(int64_t dim,
int64_t num_heads = 8,
bool qkv_bias = false,
bool proj_bias = true)
: num_heads(num_heads) {
int64_t num_heads = 8,
bool qkv_bias = false,
bool proj_bias = true,
bool fused_qkv = true)
: num_heads(num_heads), hidden_size(dim), fused_qkv(fused_qkv) {
int64_t head_dim = dim / num_heads;
blocks["qkv"] = std::shared_ptr<GGMLBlock>(new Linear(dim, dim * 3, qkv_bias));
if (fused_qkv) {
blocks["qkv"] = std::shared_ptr<GGMLBlock>(new Linear(dim, dim * 3, qkv_bias));
} else {
blocks["to_q"] = std::shared_ptr<GGMLBlock>(new Linear(dim, dim, qkv_bias));
blocks["to_k"] = std::shared_ptr<GGMLBlock>(new Linear(dim, dim, qkv_bias));
blocks["to_v"] = std::shared_ptr<GGMLBlock>(new Linear(dim, dim, qkv_bias));
}
blocks["norm"] = std::shared_ptr<GGMLBlock>(new QKNorm(head_dim));
blocks["proj"] = std::shared_ptr<GGMLBlock>(new Linear(dim, dim, proj_bias));
}

std::vector<ggml_tensor*> pre_attention(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
auto norm = std::dynamic_pointer_cast<QKNorm>(blocks["norm"]);
auto norm = std::dynamic_pointer_cast<QKNorm>(blocks["norm"]);
int64_t head_dim = hidden_size / num_heads;

auto qkv = qkv_proj->forward(ctx, x);
int64_t head_dim = qkv->ne[0] / 3 / num_heads;
auto q = ggml_view_4d(ctx->ggml_ctx, qkv, head_dim, num_heads, qkv->ne[1], qkv->ne[2],
qkv->nb[0] * head_dim, qkv->nb[1], qkv->nb[2], 0);
auto k = ggml_view_4d(ctx->ggml_ctx, qkv, head_dim, num_heads, qkv->ne[1], qkv->ne[2],
qkv->nb[0] * head_dim, qkv->nb[1], qkv->nb[2], (qkv->nb[0]) * qkv->ne[0] / 3);
auto v = ggml_view_4d(ctx->ggml_ctx, qkv, head_dim, num_heads, qkv->ne[1], qkv->ne[2],
qkv->nb[0] * head_dim, qkv->nb[1], qkv->nb[2], (qkv->nb[0]) * 2 * qkv->ne[0] / 3);
q = norm->query_norm(ctx, q);
k = norm->key_norm(ctx, k);
return {q, k, v};
if (fused_qkv) {
auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
auto qkv = qkv_proj->forward(ctx, x);
auto q = ggml_view_4d(ctx->ggml_ctx, qkv, head_dim, num_heads, qkv->ne[1], qkv->ne[2],
qkv->nb[0] * head_dim, qkv->nb[1], qkv->nb[2], 0);
auto k = ggml_view_4d(ctx->ggml_ctx, qkv, head_dim, num_heads, qkv->ne[1], qkv->ne[2],
qkv->nb[0] * head_dim, qkv->nb[1], qkv->nb[2], (qkv->nb[0]) * qkv->ne[0] / 3);
auto v = ggml_view_4d(ctx->ggml_ctx, qkv, head_dim, num_heads, qkv->ne[1], qkv->ne[2],
qkv->nb[0] * head_dim, qkv->nb[1], qkv->nb[2], (qkv->nb[0]) * 2 * qkv->ne[0] / 3);
q = norm->query_norm(ctx, q);
k = norm->key_norm(ctx, k);
return {q, k, v};
} else {
auto q_proj = std::dynamic_pointer_cast<Linear>(blocks["to_q"]);
auto k_proj = std::dynamic_pointer_cast<Linear>(blocks["to_k"]);
auto v_proj = std::dynamic_pointer_cast<Linear>(blocks["to_v"]);
auto q = ggml_reshape_4d(ctx->ggml_ctx, q_proj->forward(ctx, x), head_dim, num_heads, x->ne[1], x->ne[2]);
auto k = ggml_reshape_4d(ctx->ggml_ctx, k_proj->forward(ctx, x), head_dim, num_heads, x->ne[1], x->ne[2]);
auto v = ggml_reshape_4d(ctx->ggml_ctx, v_proj->forward(ctx, x), head_dim, num_heads, x->ne[1], x->ne[2]);
q = norm->query_norm(ctx, q);
k = norm->key_norm(ctx, k);
return {q, k, v};
}
}

ggml_tensor* post_attention(GGMLRunnerContext* ctx, ggml_tensor* x) {
Expand Down Expand Up @@ -272,7 +293,7 @@ namespace Flux {
blocks["img_mod"] = std::shared_ptr<GGMLBlock>(new Modulation(hidden_size, true));
}
blocks["img_norm1"] = std::shared_ptr<GGMLBlock>(new LayerNorm(hidden_size, 1e-6f, false));
blocks["img_attn"] = std::shared_ptr<GGMLBlock>(new SelfAttention(hidden_size, num_heads, qkv_bias, mlp_proj_bias));
blocks["img_attn"] = std::shared_ptr<GGMLBlock>(new SelfAttention(hidden_size, num_heads, qkv_bias, mlp_proj_bias, !share_modulation));

blocks["img_norm2"] = std::shared_ptr<GGMLBlock>(new LayerNorm(hidden_size, 1e-6f, false));
if (use_yak_mlp) {
Expand All @@ -285,7 +306,7 @@ namespace Flux {
blocks["txt_mod"] = std::shared_ptr<GGMLBlock>(new Modulation(hidden_size, true));
}
blocks["txt_norm1"] = std::shared_ptr<GGMLBlock>(new LayerNorm(hidden_size, 1e-6f, false));
blocks["txt_attn"] = std::shared_ptr<GGMLBlock>(new SelfAttention(hidden_size, num_heads, qkv_bias, mlp_proj_bias));
blocks["txt_attn"] = std::shared_ptr<GGMLBlock>(new SelfAttention(hidden_size, num_heads, qkv_bias, mlp_proj_bias, !share_modulation));

blocks["txt_norm2"] = std::shared_ptr<GGMLBlock>(new LayerNorm(hidden_size, 1e-6f, false));
if (use_yak_mlp) {
Expand Down
74 changes: 73 additions & 1 deletion src/model.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
#include <string>
#include <thread>
#include <unordered_map>
#include <unordered_set>
#include <vector>

#include "model.h"
Expand All @@ -28,6 +29,7 @@
#include "zip.h"

#include "name_conversion.h"
#include "json.hpp"

/*================================================= Preprocess ==================================================*/

Expand Down Expand Up @@ -208,8 +210,77 @@ void ModelLoader::add_tensor_storage(const TensorStorage& tensor_storage) {
tensor_storage_map[tensor_storage.name] = tensor_storage;
}

bool is_sharded_safetensors_directory(const std::string& file_path) {
if (!is_directory(file_path)) {
return false;
}
std::string index_path = path_join(file_path, "model.safetensors.index.json");
return file_exists(index_path);
}

bool ModelLoader::init_from_sharded_safetensors_file(const std::string& file_path, const std::string& prefix) {
std::string index_path = path_join(file_path, "model.safetensors.index.json");

nlohmann::json index;
{
std::ifstream f(index_path);
if (!f.is_open()) {
LOG_ERROR("failed to open %s", index_path.c_str());
return false;
}
f >> index;
}

auto weight_map = index.find("weight_map");
if (weight_map == index.end()) {
LOG_ERROR("no weight_map in %s", index_path.c_str());
return false;
}

// Collect unique shard files
std::unordered_set<std::string> shard_files;
for (auto& [tensor_name, shard_file] : weight_map->items()) {
(void)tensor_name;
shard_files.insert(shard_file.get<std::string>());
}

size_t file_index_offset = file_paths_.size();

for (const auto& shard_file_name : shard_files) {
std::string shard_path = path_join(file_path, shard_file_name);
LOG_DEBUG("init from shard '%s'", shard_path.c_str());

std::vector<TensorStorage> tensor_storages;
std::string error;
if (!read_safetensors_file(shard_path, tensor_storages, &error)) {
LOG_ERROR("failed to read shard %s: %s", shard_path.c_str(), error.c_str());
return false;
}

file_paths_.push_back(shard_path);
size_t file_index = file_index_offset;

for (auto& tensor_storage : tensor_storages) {
if (!starts_with(tensor_storage.name, prefix)) {
tensor_storage.name = prefix + tensor_storage.name;
}
tensor_storage.file_index = file_index;
add_tensor_storage(tensor_storage);
}
file_index_offset++;
}

LOG_INFO("loaded %zu tensors from %zu shards in %s",
tensor_storage_map.size(), shard_files.size(), file_path.c_str());
return true;
}

bool ModelLoader::init_from_file(const std::string& file_path, const std::string& prefix) {
if (is_directory(file_path)) {
if (is_sharded_safetensors_directory(file_path)) {
LOG_INFO("load %s using sharded safetensors format", file_path.c_str());
return init_from_sharded_safetensors_file(file_path, prefix);
}
LOG_INFO("load %s using diffusers format", file_path.c_str());
return init_from_diffusers_file(file_path, prefix);
} else if (is_gguf_file(file_path)) {
Expand Down Expand Up @@ -452,7 +523,8 @@ SDVersion ModelLoader::get_sd_version() {
if (tensor_storage.name.find("llm_adapter.blocks.0.cross_attn.q_proj.weight") != std::string::npos) {
return VERSION_ANIMA;
}
if (tensor_storage.name.find("model.diffusion_model.double_stream_modulation_img.lin.weight") != std::string::npos) {
if (tensor_storage.name.find("model.diffusion_model.double_stream_modulation_img.lin.weight") != std::string::npos ||
tensor_storage.name.find("model.diffusion_model.double_stream_modulation_img.linear.weight") != std::string::npos) {
is_flux2 = true;
}
if (tensor_storage.name.find("single_blocks.47.linear1.weight") != std::string::npos) {
Expand Down
1 change: 1 addition & 0 deletions src/model.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -252,6 +252,7 @@ class ModelLoader {
bool init_from_torch_zip_file(const std::string& file_path, const std::string& prefix = "");
bool init_from_torch_legacy_file(const std::string& file_path, const std::string& prefix = "");
bool init_from_diffusers_file(const std::string& file_path, const std::string& prefix = "");
bool init_from_sharded_safetensors_file(const std::string& file_path, const std::string& prefix = "");

public:
bool init_from_file(const std::string& file_path, const std::string& prefix = "");
Expand Down
58 changes: 44 additions & 14 deletions src/name_conversion.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -538,12 +538,26 @@ std::string convert_diffusers_dit_to_original_flux(std::string name) {
flux_name_map["time_text_embed.guidance_embedder.linear_2.weight"] = "guidance_in.out_layer.weight";
flux_name_map["time_text_embed.guidance_embedder.linear_2.bias"] = "guidance_in.out_layer.bias";

// --- time_guidance_embed (flux.2 klein) ---
flux_name_map["time_guidance_embed.timestep_embedder.linear_1.weight"] = "time_in.in_layer.weight";
flux_name_map["time_guidance_embed.timestep_embedder.linear_1.bias"] = "time_in.in_layer.bias";
flux_name_map["time_guidance_embed.timestep_embedder.linear_2.weight"] = "time_in.out_layer.weight";
flux_name_map["time_guidance_embed.timestep_embedder.linear_2.bias"] = "time_in.out_layer.bias";

// --- context_embedder / x_embedder ---
flux_name_map["context_embedder.weight"] = "txt_in.weight";
flux_name_map["context_embedder.bias"] = "txt_in.bias";
flux_name_map["x_embedder.weight"] = "img_in.weight";
flux_name_map["x_embedder.bias"] = "img_in.bias";

// --- shared modulation (flux.2 klein) ---
flux_name_map["double_stream_modulation_img.linear.weight"] = "double_stream_modulation_img.lin.weight";
flux_name_map["double_stream_modulation_img.linear.bias"] = "double_stream_modulation_img.lin.bias";
flux_name_map["double_stream_modulation_txt.linear.weight"] = "double_stream_modulation_txt.lin.weight";
flux_name_map["double_stream_modulation_txt.linear.bias"] = "double_stream_modulation_txt.lin.bias";
flux_name_map["single_stream_modulation.linear.weight"] = "single_stream_modulation.lin.weight";
flux_name_map["single_stream_modulation.linear.bias"] = "single_stream_modulation.lin.bias";

// --- double transformer blocks ---
for (int i = 0; i < num_layers; ++i) {
std::string block_prefix = "transformer_blocks." + std::to_string(i) + ".";
Expand All @@ -554,20 +568,20 @@ std::string convert_diffusers_dit_to_original_flux(std::string name) {
flux_name_map[block_prefix + "norm1_context.linear.weight"] = dst_prefix + "txt_mod.lin.weight";
flux_name_map[block_prefix + "norm1_context.linear.bias"] = dst_prefix + "txt_mod.lin.bias";

// attn
flux_name_map[block_prefix + "attn.to_q.weight"] = dst_prefix + "img_attn.qkv.weight";
flux_name_map[block_prefix + "attn.to_q.bias"] = dst_prefix + "img_attn.qkv.bias";
flux_name_map[block_prefix + "attn.to_k.weight"] = dst_prefix + "img_attn.qkv.weight.1";
flux_name_map[block_prefix + "attn.to_k.bias"] = dst_prefix + "img_attn.qkv.bias.1";
flux_name_map[block_prefix + "attn.to_v.weight"] = dst_prefix + "img_attn.qkv.weight.2";
flux_name_map[block_prefix + "attn.to_v.bias"] = dst_prefix + "img_attn.qkv.bias.2";

flux_name_map[block_prefix + "attn.add_q_proj.weight"] = dst_prefix + "txt_attn.qkv.weight";
flux_name_map[block_prefix + "attn.add_q_proj.bias"] = dst_prefix + "txt_attn.qkv.bias";
flux_name_map[block_prefix + "attn.add_k_proj.weight"] = dst_prefix + "txt_attn.qkv.weight.1";
flux_name_map[block_prefix + "attn.add_k_proj.bias"] = dst_prefix + "txt_attn.qkv.bias.1";
flux_name_map[block_prefix + "attn.add_v_proj.weight"] = dst_prefix + "txt_attn.qkv.weight.2";
flux_name_map[block_prefix + "attn.add_v_proj.bias"] = dst_prefix + "txt_attn.qkv.bias.2";
// attn (separate q/k/v)
flux_name_map[block_prefix + "attn.to_q.weight"] = dst_prefix + "img_attn.to_q.weight";
flux_name_map[block_prefix + "attn.to_q.bias"] = dst_prefix + "img_attn.to_q.bias";
flux_name_map[block_prefix + "attn.to_k.weight"] = dst_prefix + "img_attn.to_k.weight";
flux_name_map[block_prefix + "attn.to_k.bias"] = dst_prefix + "img_attn.to_k.bias";
flux_name_map[block_prefix + "attn.to_v.weight"] = dst_prefix + "img_attn.to_v.weight";
flux_name_map[block_prefix + "attn.to_v.bias"] = dst_prefix + "img_attn.to_v.bias";

flux_name_map[block_prefix + "attn.add_q_proj.weight"] = dst_prefix + "txt_attn.to_q.weight";
flux_name_map[block_prefix + "attn.add_q_proj.bias"] = dst_prefix + "txt_attn.to_q.bias";
flux_name_map[block_prefix + "attn.add_k_proj.weight"] = dst_prefix + "txt_attn.to_k.weight";
flux_name_map[block_prefix + "attn.add_k_proj.bias"] = dst_prefix + "txt_attn.to_k.bias";
flux_name_map[block_prefix + "attn.add_v_proj.weight"] = dst_prefix + "txt_attn.to_v.weight";
flux_name_map[block_prefix + "attn.add_v_proj.bias"] = dst_prefix + "txt_attn.to_v.bias";

// norm
flux_name_map[block_prefix + "attn.norm_q.weight"] = dst_prefix + "img_attn.norm.query_norm.scale";
Expand All @@ -591,6 +605,16 @@ std::string convert_diffusers_dit_to_original_flux(std::string name) {
flux_name_map[block_prefix + "ff_context.net.2.weight"] = dst_prefix + "txt_mlp.2.weight";
flux_name_map[block_prefix + "ff_context.net.2.bias"] = dst_prefix + "txt_mlp.2.bias";

// ff.linear_in / linear_out (flux.2 klein)
flux_name_map[block_prefix + "ff.linear_in.weight"] = dst_prefix + "img_mlp.0.weight";
flux_name_map[block_prefix + "ff.linear_in.bias"] = dst_prefix + "img_mlp.0.bias";
flux_name_map[block_prefix + "ff.linear_out.weight"] = dst_prefix + "img_mlp.2.weight";
flux_name_map[block_prefix + "ff.linear_out.bias"] = dst_prefix + "img_mlp.2.bias";
flux_name_map[block_prefix + "ff_context.linear_in.weight"] = dst_prefix + "txt_mlp.0.weight";
flux_name_map[block_prefix + "ff_context.linear_in.bias"] = dst_prefix + "txt_mlp.0.bias";
flux_name_map[block_prefix + "ff_context.linear_out.weight"] = dst_prefix + "txt_mlp.2.weight";
flux_name_map[block_prefix + "ff_context.linear_out.bias"] = dst_prefix + "txt_mlp.2.bias";

// output projections
flux_name_map[block_prefix + "attn.to_out.0.weight"] = dst_prefix + "img_attn.proj.weight";
flux_name_map[block_prefix + "attn.to_out.0.bias"] = dst_prefix + "img_attn.proj.bias";
Expand Down Expand Up @@ -622,6 +646,12 @@ std::string convert_diffusers_dit_to_original_flux(std::string name) {
flux_name_map[dst_prefix + "norm.key_norm.weight"] = dst_prefix + "norm.key_norm.scale";
flux_name_map[block_prefix + "proj_out.weight"] = dst_prefix + "linear2.weight";
flux_name_map[block_prefix + "proj_out.bias"] = dst_prefix + "linear2.bias";

// fused to_qkv_mlp_proj (flux.2 klein)
flux_name_map[block_prefix + "attn.to_qkv_mlp_proj.weight"] = dst_prefix + "linear1.weight";
flux_name_map[block_prefix + "attn.to_qkv_mlp_proj.bias"] = dst_prefix + "linear1.bias";
flux_name_map[block_prefix + "attn.to_out.weight"] = dst_prefix + "linear2.weight";
flux_name_map[block_prefix + "attn.to_out.bias"] = dst_prefix + "linear2.bias";
}

// --- final layers ---
Expand Down
Loading