diff --git a/image-studio b/image-studio new file mode 160000 index 000000000..3d895c0ce --- /dev/null +++ b/image-studio @@ -0,0 +1 @@ +Subproject commit 3d895c0ceec95de1e094e22d7386ff4e72c8c912 diff --git a/src/auto_encoder_kl.hpp b/src/auto_encoder_kl.hpp index 13396e737..6dd369dd7 100644 --- a/src/auto_encoder_kl.hpp +++ b/src/auto_encoder_kl.hpp @@ -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 { diff --git a/src/flux.hpp b/src/flux.hpp index 0271ef6fb..b1f8101ef 100644 --- a/src/flux.hpp +++ b/src/flux.hpp @@ -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(new Linear(dim, dim * 3, qkv_bias)); + if (fused_qkv) { + blocks["qkv"] = std::shared_ptr(new Linear(dim, dim * 3, qkv_bias)); + } else { + blocks["to_q"] = std::shared_ptr(new Linear(dim, dim, qkv_bias)); + blocks["to_k"] = std::shared_ptr(new Linear(dim, dim, qkv_bias)); + blocks["to_v"] = std::shared_ptr(new Linear(dim, dim, qkv_bias)); + } blocks["norm"] = std::shared_ptr(new QKNorm(head_dim)); blocks["proj"] = std::shared_ptr(new Linear(dim, dim, proj_bias)); } std::vector pre_attention(GGMLRunnerContext* ctx, ggml_tensor* x) { - auto qkv_proj = std::dynamic_pointer_cast(blocks["qkv"]); - auto norm = std::dynamic_pointer_cast(blocks["norm"]); + auto norm = std::dynamic_pointer_cast(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(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(blocks["to_q"]); + auto k_proj = std::dynamic_pointer_cast(blocks["to_k"]); + auto v_proj = std::dynamic_pointer_cast(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) { @@ -272,7 +293,7 @@ namespace Flux { blocks["img_mod"] = std::shared_ptr(new Modulation(hidden_size, true)); } blocks["img_norm1"] = std::shared_ptr(new LayerNorm(hidden_size, 1e-6f, false)); - blocks["img_attn"] = std::shared_ptr(new SelfAttention(hidden_size, num_heads, qkv_bias, mlp_proj_bias)); + blocks["img_attn"] = std::shared_ptr(new SelfAttention(hidden_size, num_heads, qkv_bias, mlp_proj_bias, !share_modulation)); blocks["img_norm2"] = std::shared_ptr(new LayerNorm(hidden_size, 1e-6f, false)); if (use_yak_mlp) { @@ -285,7 +306,7 @@ namespace Flux { blocks["txt_mod"] = std::shared_ptr(new Modulation(hidden_size, true)); } blocks["txt_norm1"] = std::shared_ptr(new LayerNorm(hidden_size, 1e-6f, false)); - blocks["txt_attn"] = std::shared_ptr(new SelfAttention(hidden_size, num_heads, qkv_bias, mlp_proj_bias)); + blocks["txt_attn"] = std::shared_ptr(new SelfAttention(hidden_size, num_heads, qkv_bias, mlp_proj_bias, !share_modulation)); blocks["txt_norm2"] = std::shared_ptr(new LayerNorm(hidden_size, 1e-6f, false)); if (use_yak_mlp) { diff --git a/src/model.cpp b/src/model.cpp index 25d78b94e..66b28cf39 100644 --- a/src/model.cpp +++ b/src/model.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #include #include "model.h" @@ -28,6 +29,7 @@ #include "zip.h" #include "name_conversion.h" +#include "json.hpp" /*================================================= Preprocess ==================================================*/ @@ -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 shard_files; + for (auto& [tensor_name, shard_file] : weight_map->items()) { + (void)tensor_name; + shard_files.insert(shard_file.get()); + } + + 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 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)) { @@ -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) { diff --git a/src/model.h b/src/model.h index 8ecea16b5..a8bb17a95 100644 --- a/src/model.h +++ b/src/model.h @@ -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 = ""); diff --git a/src/name_conversion.cpp b/src/name_conversion.cpp index a9cae0a87..9f45d1718 100644 --- a/src/name_conversion.cpp +++ b/src/name_conversion.cpp @@ -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) + "."; @@ -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"; @@ -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"; @@ -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 ---