feat: add Flux2 Klein runtime support

README.md

@@ -18,9 +18,10 @@
 runtime for image generation, image editing, and video generation across local,
 edge, and resource-constrained deployment environments.
 
-It supports major DiT model families including **FLUX.1, Stable Diffusion 3/3.5,
-Qwen-Image, and Wan**, with explicit control over model loading, memory usage,
-graph execution, quantization, device placement, and backend selection.
+It supports major DiT model families including **FLUX.1, FLUX.2 [klein] 4B,
+Stable Diffusion 3/3.5, Qwen-Image, and Wan**, with explicit control over model
+loading, memory usage, graph execution, quantization, device placement, and
+backend selection.
 
 ## Features
 
@@ -32,7 +33,8 @@ graph execution, quantization, device placement, and backend selection.
 
 - **Unified across tasks and model families**
   - **Text-to-image**, **image editing**, and **video generation** in one runtime
-  - SD3/SD3.5, FLUX.1, FLUX.1-Kontext, Qwen-Image, Qwen-Image-Edit, and Wan 2.1
+  - SD3/SD3.5, FLUX.1, FLUX.1-Kontext, FLUX.2 [klein] 4B, Qwen-Image,
+    Qwen-Image-Edit, and Wan 2.1
   - **Few-step distilled models** auto-detected — Turbo / Lightning / schnell default to a **4–8 step** schedule
   - Shared **C API, CLI, HTTP server, and Python** interfaces across every family
 
@@ -63,17 +65,18 @@ graph execution, quantization, device placement, and backend selection.
 
 ## Supported Models
 
-This release focuses on the model families below, each with a base
-checkpoint and a **few-step distilled variant**. Some source files contain
-experimental model scaffolding beyond this table; those are not part of the
-current support commitment unless documented in
+This release focuses on the model families below. Most ship as a base
+checkpoint plus a **few-step distilled variant**; some source files contain
+experimental model scaffolding beyond this table, and those are not part of
+the current support commitment unless documented in
 [Supported Models](docs/models.md).
 
 | Model family | Task | Base checkpoint | Distilled variant (few-step) | Status |
 |---|---|---|---|---|
 | **SD3 / SD3.5** | Text-to-image | `stabilityai/stable-diffusion-3-medium` | SD3.5-medium-turbo | Supported |
 | **FLUX.1** | Text-to-image | `black-forest-labs/FLUX.1-dev` | FLUX.1-schnell | Supported |
 | **FLUX.1-Kontext** | Image editing / reference-guided | `black-forest-labs/FLUX.1-Kontext-dev` | Kontext Lightning | Supported |
+| **FLUX.2 [klein] 4B** | Text-to-image / image editing | `black-forest-labs/FLUX.2-klein-4B` | Native FLUX.2 [klein] 4B checkpoint | Supported |
 | **Qwen-Image** | Text-to-image | `Qwen/Qwen-Image` | Qwen-Image Lightning *(LoRA)* | Supported |
 | **Qwen-Image-Edit** | Image editing | `Qwen/Qwen-Image-Edit` | Qwen-Image-Edit Lightning *(LoRA)* | Supported |
 | **Wan 2.1** | Video generation | `Wan-AI/Wan2.1-T2V-1.3B` (and 14B) | Wan2.1-T2V-1.3B Distill | Supported (Vulkan still optimizing) |

docs/models.md

@@ -15,6 +15,7 @@ support commitment.
 | SD3 / SD3.5 | Text-to-image | [`stabilityai/stable-diffusion-3-medium-diffusers`](https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers) (and SD3.5 siblings) | Diffusers-style directory or component weights | CUDA first, CPU/Vulkan functional, Metal experimental | Supported |
 | FLUX.1 | Text-to-image | [`black-forest-labs/FLUX.1-dev`](https://huggingface.co/black-forest-labs/FLUX.1-dev) | Diffusers-style directory, top-level FLUX safetensors, or components | CUDA first, CPU/Vulkan functional, Metal experimental | Supported |
 | FLUX.1-Kontext | Image editing / reference-guided generation | [`black-forest-labs/FLUX.1-Kontext-dev`](https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev) | Diffusers-style directory or components | CUDA first, CPU/Vulkan functional, Metal experimental | Supported |
+| FLUX.2 [klein] 4B | Text-to-image / image editing | [`black-forest-labs/FLUX.2-klein-4B`](https://huggingface.co/black-forest-labs/FLUX.2-klein-4B) | Diffusers-style directory or components | CUDA first, CPU/Vulkan functional, Metal experimental | Supported |
 | Qwen-Image | Text-to-image | [`Qwen/Qwen-Image`](https://huggingface.co/Qwen/Qwen-Image) | Diffusers-style directory or components | CUDA first, CPU/Vulkan functional, Metal experimental | Supported |
 | Qwen-Image-Edit | Image editing | [`Qwen/Qwen-Image-Edit`](https://huggingface.co/Qwen/Qwen-Image-Edit) | Diffusers-style directory or components | CUDA first, CPU/Vulkan functional, Metal experimental | Supported |
 | Wan 2.1 | Video generation | [`Wan-AI/Wan2.1-T2V-1.3B-Diffusers`](https://huggingface.co/Wan-AI/Wan2.1-T2V-1.3B-Diffusers) (and Wan2.1-T2V-14B-Diffusers) | Diffusers-style directory or components | CUDA first, CPU functional for validation, Metal/Vulkan experimental | Supported (Vulkan still optimizing) |

src/core/optimization/cache/cache_types.cpp

@@ -34,6 +34,11 @@ CacheModelSpec cache_model_spec_for_version(SDVersion version) {
     CacheModelSpec spec;
     spec.version = version;
 
+    if (version == VERSION_FLUX2_KLEIN) {
+        spec.model_name = "flux2-klein";
+        spec.block_count = 5 + 20;  // double_blocks + single_blocks
+        return spec;
+    }
     if (ed_version_is_flux(version) || ed_version_is_flux2(version)) {
         spec.model_name = "flux";
         spec.block_count = 19 + 38;  // double_blocks + single_blocks

src/core/optimization/cache/model/dit_model_cache_contract.cpp

@@ -18,6 +18,12 @@ struct FamilyInfo {
 
 FamilyInfo family_info_for(SDVersion version) {
     FamilyInfo info;
+    if (version == VERSION_FLUX2_KLEIN) {
+        info.family = ModelFamily::Flux;
+        info.num_blocks = 5 + 20;  // double_blocks + single_blocks
+        info.dual_stream = true;
+        return info;
+    }
     if (ed_version_is_flux(version) || ed_version_is_flux2(version)) {
         info.family = ModelFamily::Flux;
         info.num_blocks = 19 + 38;  // double_blocks + single_blocks

src/core/runtime/model_loader.cpp

@@ -359,6 +359,9 @@ static SDVersion infer_transformer_file_version(const std::string& file_path) {
         if (contains(klass, "Kontext")) {
             return VERSION_FLUX_KONTEXT;
         }
+        if (contains(klass, "Flux2")) {
+            return contains(klass, "Klein") ? VERSION_FLUX2_KLEIN : VERSION_FLUX2;
+        }
         if (contains(klass, "Flux")) {
             return VERSION_FLUX;
         }
@@ -377,6 +380,10 @@ static bool is_flux1_family_version(SDVersion version) {
     return version == VERSION_FLUX || version == VERSION_FLUX_KONTEXT;
 }
 
+static bool is_flux2_family_version(SDVersion version) {
+    return ed_version_is_flux2(version);
+}
+
 static std::string resolve_flux_transformer_component_path(const std::string& file_path) {
     std::string normalized = file_path;
     std::replace(normalized.begin(), normalized.end(), '\\', '/');
@@ -450,6 +457,9 @@ static SDVersion infer_diffusers_version(const std::string& dir_path) {
             if (contains(klass, "Kontext")) {
                 return VERSION_FLUX_KONTEXT;
             }
+            if (contains(klass, "Flux2")) {
+                return contains(klass, "Klein") ? VERSION_FLUX2_KLEIN : VERSION_FLUX2;
+            }
             if (contains(klass, "Flux")) {
                 return VERSION_FLUX;
             }
@@ -478,6 +488,9 @@ static SDVersion infer_diffusers_version(const std::string& dir_path) {
             if (contains(klass, "Kontext")) {
                 return VERSION_FLUX_KONTEXT;
             }
+            if (contains(klass, "Flux2")) {
+                return contains(klass, "Klein") ? VERSION_FLUX2_KLEIN : VERSION_FLUX2;
+            }
             if (contains(klass, "Flux")) {
                 return VERSION_FLUX;
             }
@@ -539,6 +552,29 @@ static std::vector<std::string> component_weight_candidates(const std::string& c
     return candidates;
 }
 
+static std::string find_top_level_safetensors_file(const std::string& dir_path) {
+    std::error_code ec;
+    if (!fs::is_directory(dir_path, ec)) {
+        return {};
+    }
+
+    for (const auto& entry : fs::directory_iterator(dir_path, ec)) {
+        if (ec) {
+            break;
+        }
+        if (!entry.is_regular_file()) {
+            continue;
+        }
+
+        const std::string path = entry.path().string();
+        if (has_suffix(path, ".safetensors")) {
+            return path;
+        }
+    }
+
+    return {};
+}
+
 static uint16_t f8_e4m3_to_f16(uint8_t f8) {
     const uint32_t exponent_bias = 7;
     if (f8 == 0xff) {
@@ -1281,7 +1317,7 @@ bool ModelLoader::init_from_diffusers_directory(const std::string& dir_path, con
         }
         const std::string component_dir = path_join(dir_path, component.dir);
         if (!is_directory(component_dir)) {
-            if (component.required_for_flux && is_flux1_family_version(version_)) {
+            if (component.required_for_flux && (is_flux1_family_version(version_) || is_flux2_family_version(version_))) {
                 LOG_WARN("diffusers component '%s' not found", component.dir);
             }
             continue;
@@ -1297,6 +1333,15 @@ bool ModelLoader::init_from_diffusers_directory(const std::string& dir_path, con
                 continue;
             }
         }
+        if (is_flux2_family_version(version_)) {
+            if (std::strcmp(component.dir, "text_encoder") == 0) {
+                component_prefix = "text_encoders.llm.";
+            } else if (std::strcmp(component.dir, "text_encoder_2") == 0 ||
+                       std::strcmp(component.dir, "text_encoder_3") == 0 ||
+                       std::strcmp(component.dir, "unet") == 0) {
+                continue;
+            }
+        }
         if (ed_version_is_qwen_image(version_) || ed_version_is_qwen_image_edit(version_)) {
             if (std::strcmp(component.dir, "text_encoder") == 0) {
                 component_prefix = "text_encoders.llm.";
@@ -1308,6 +1353,20 @@ bool ModelLoader::init_from_diffusers_directory(const std::string& dir_path, con
         }
         bool loaded = false;
         std::set<std::string> tried;
+
+        if (is_flux2_family_version(version_) && std::strcmp(component.dir, "transformer") == 0) {
+            const std::string top_level_flux = find_top_level_safetensors_file(dir_path);
+            if (!top_level_flux.empty()) {
+                const size_t before = tensor_storage_map_.size();
+                loaded = init_from_safetensors_file(top_level_flux, component_prefix);
+                if (loaded) {
+                    LOG_INFO("loaded diffusers component '%s' from top-level Flux2 weights '%s' (%zu tensors)",
+                             component.dir,
+                             top_level_flux.c_str(),
+                             tensor_storage_map_.size() - before);
+                }
+            }
+        }
         if (is_flux1_family_version(version_) && std::strcmp(component.dir, "transformer") == 0) {
             const std::vector<std::string> top_level_flux_weights = {
                 path_join(dir_path, "flux1-kontext-dev.safetensors"),
@@ -1390,6 +1449,8 @@ SDVersion ModelLoader::get_ld_version() {
     bool has_transformer_blocks = false;
     bool has_unet = false;
     bool has_second_text_encoder = false;
+    bool has_flux2 = false;
+    bool has_single_block_47 = false;
 
     TensorStorage input_block_weight;
     TensorStorage token_embedding_weight;
@@ -1417,9 +1478,16 @@ SDVersion ModelLoader::get_ld_version() {
         if (contains(name, "model.diffusion_model.double_blocks.") || contains(name, "transformer.double_blocks.")) {
             has_flux_double = true;
         }
+        if (contains(name, "model.diffusion_model.double_stream_modulation_img.lin.weight") ||
+            contains(name, "model.diffusion_model.double_stream_modulation_img.linear.weight")) {
+            has_flux2 = true;
+        }
         if (contains(name, "single_transformer_blocks.") || contains(name, "single_blocks.")) {
             has_flux_single = true;
         }
+        if (contains(name, "single_blocks.47.linear1.weight")) {
+            has_single_block_47 = true;
+        }
         if (contains(name, "transformer.transformer_blocks.") || contains(name, "transformer_blocks.")) {
             has_transformer_blocks = true;
         }
@@ -1437,6 +1505,12 @@ SDVersion ModelLoader::get_ld_version() {
         }
     }
 
+    if (has_flux2) {
+        if (has_single_block_47) {
+            return VERSION_FLUX2;
+        }
+        return VERSION_FLUX2_KLEIN;
+    }
     if (has_flux_double || (has_transformer_blocks && has_flux_single)) {
         if (input_block_weight.ne[0] == 384) {
             return VERSION_FLUX_FILL;

src/dit_models/components/common/rope.hpp

@@ -785,6 +785,7 @@ namespace Rope {
                                              ggml_tensor* mask,
                                              float kv_scale        = 1.0f,
                                              bool rope_interleaved = true,
+                                             bool use_fused_rope   = true,
                                              bool k_rope_f16_for_flash = false,
                                              bool qk_rope_bf16_roundtrip = false) {
         // q,k,v: [N, L, n_head, d_head]
@@ -796,16 +797,17 @@ namespace Rope {
             ctx->flash_attn_enabled &&
             k->ne[2] % 256 != 0;
 
-        q = apply_rope(ctx->ggml_ctx, q, pe, rope_interleaved, ctx->backend);  // [N*n_head, L, d_head]
-        if (k_rope_f16_for_flash && ctx->flash_attn_enabled && !will_pad_kv_for_flash_attn) {
+        ggml_backend_t rope_backend = use_fused_rope ? ctx->backend : nullptr;
+        q = apply_rope(ctx->ggml_ctx, q, pe, rope_interleaved, rope_backend);  // [N*n_head, L, d_head]
+        if (k_rope_f16_for_flash && use_fused_rope && ctx->flash_attn_enabled && !will_pad_kv_for_flash_attn) {
             ggml_tensor* k_f16 = edgedit::ggml_ext::apply_rope_f16(ctx->ggml_ctx, k, pe, rope_interleaved);
             if (k_f16 != nullptr && ggml_backend_supports_op(ctx->backend, k_f16)) {
                 k = k_f16;  // flash attention consumes K as F16, so avoid a separate cast node
             } else {
-                k = apply_rope(ctx->ggml_ctx, k, pe, rope_interleaved, ctx->backend);
+                k = apply_rope(ctx->ggml_ctx, k, pe, rope_interleaved, rope_backend);
             }
         } else {
-            k = apply_rope(ctx->ggml_ctx, k, pe, rope_interleaved, ctx->backend);  // [N*n_head, L, d_head]
+            k = apply_rope(ctx->ggml_ctx, k, pe, rope_interleaved, rope_backend);  // [N*n_head, L, d_head]
         }
         if (qk_rope_bf16_roundtrip) {
             q = ggml_cast(ctx->ggml_ctx, ggml_cast(ctx->ggml_ctx, q, GGML_TYPE_BF16), GGML_TYPE_F32);

src/dit_models/components/text_encoders/llm.hpp

@@ -1612,7 +1612,7 @@ namespace LLM {
                 params.rms_norm_eps = 1e-5f;
             } else if (arch == LLMArch::QWEN3 || arch == LLMArch::QWEN3_VL) {
                 params.head_dim     = 128;
-                params.num_heads    = 64;
+                params.num_heads    = 32;
                 params.num_kv_heads = 8;
                 params.qkv_bias     = false;
                 params.qk_norm      = true;