Merge branch 'master' into hidream-o1

Co-authored-by: Piotr Wilkin <piotr.wilkin@syndatis.com>
Co-authored-by: Junmo Kim <me@junmo.kim>
Co-authored-by: leejet <leejet714@gmail.com>

CMakeLists.txt

@@ -106,7 +106,8 @@ if(SD_WEBP)
           "Or link against system library:\n  cmake (...) -DSD_USE_SYSTEM_WEBP=ON")
     endif()
     if(SD_USE_SYSTEM_WEBP)
-        find_package(WebP REQUIRED)
+        find_package(WebP)
+        if(WebP_FOUND)
             add_library(webp ALIAS WebP::webp)
             # libwebp CMake target naming is not consistent across versions/distros.
             # Some export WebP::libwebpmux, others export WebP::webpmux.
@@ -120,6 +121,14 @@ if(SD_WEBP)
                     "Expected WebP::libwebpmux or WebP::webpmux."
                 )
             endif()
+        else()
+            find_package(PkgConfig REQUIRED)
+            pkg_check_modules(WebP REQUIRED IMPORTED_TARGET GLOBAL libwebp)
+            pkg_check_modules(WebPMux REQUIRED IMPORTED_TARGET GLOBAL libwebpmux)
+            link_libraries(PkgConfig::WebP)
+            link_libraries(PkgConfig::WebPMux)
+            add_library(libwebpmux ALIAS PkgConfig::WebPMux)
+        endif()
     endif()
 endif()
 
@@ -133,6 +142,13 @@ if(SD_WEBM)
           "Or link against system library:\n  cmake (...) -DSD_USE_SYSTEM_WEBM=ON")
     endif()
     if(SD_USE_SYSTEM_WEBM)
+        find_package(PkgConfig)
+        if(PkgConfig_FOUND)
+            pkg_check_modules(WebM REQUIRED IMPORTED_TARGET GLOBAL libwebm)
+        endif()
+        if(PkgConfig_FOUND AND WebM_FOUND)
+            link_libraries(PkgConfig::WebM)
+        else()
             find_path(WEBM_INCLUDE_DIR
                 NAMES mkvmuxer/mkvmuxer.h mkvparser/mkvparser.h common/webmids.h
                 PATH_SUFFIXES webm
@@ -146,6 +162,7 @@ if(SD_WEBM)
                 IMPORTED_LOCATION "${WEBM_LIBRARY}"
                 INTERFACE_INCLUDE_DIRECTORIES "${WEBM_INCLUDE_DIR}")
         endif()
+    endif()
 endif()
 
 set(SD_LIB stable-diffusion)

README.md

@@ -58,6 +58,7 @@ API and command-line option may change frequently.***
     - [Ovis-Image](./docs/ovis_image.md)
     - [Anima](./docs/anima.md)
     - [ERNIE-Image](./docs/ernie_image.md)
+    - [HiDream-O1-Image](./docs/hidream_o1_image.md)
   - Image Edit Models
     - [FLUX.1-Kontext-dev](./docs/kontext.md)
     - [Qwen Image Edit series](./docs/qwen_image_edit.md)
@@ -148,6 +149,7 @@ If you want to improve performance or reduce VRAM/RAM usage, please refer to [pe
 - [Ovis-Image](./docs/ovis_image.md)
 - [Anima](./docs/anima.md)
 - [ERNIE-Image](./docs/ernie_image.md)
+- [HiDream-O1-Image](./docs/hidream_o1_image.md)
 - [LoRA](./docs/lora.md)
 - [LCM/LCM-LoRA](./docs/lcm.md)
 - [Using PhotoMaker to personalize image generation](./docs/photo_maker.md)
@@ -163,6 +165,7 @@ These projects wrap `stable-diffusion.cpp` for easier use in other languages/fra
 
 * Golang (non-cgo): [seasonjs/stable-diffusion](https://github.com/seasonjs/stable-diffusion)
 * Golang (cgo): [Binozo/GoStableDiffusion](https://github.com/Binozo/GoStableDiffusion)
+* Golang (non-cgo): [l8bloom/gosd](https://github.com/l8bloom/gosd)
 * C#: [DarthAffe/StableDiffusion.NET](https://github.com/DarthAffe/StableDiffusion.NET)
 * Python: [william-murray1204/stable-diffusion-cpp-python](https://github.com/william-murray1204/stable-diffusion-cpp-python)
 * Rust: [newfla/diffusion-rs](https://github.com/newfla/diffusion-rs)

docs/hidream_o1_image.md

@@ -0,0 +1,20 @@
+# How to Use
+
+## Download weights
+
+- Download HiDream-O1-Image-Dev
+    - safetensors: https://huggingface.co/Comfy-Org/HiDream-O1-Image/tree/main/checkpoints
+- Download HiDream-O1-Image
+    - safetensors: https://huggingface.co/Comfy-Org/HiDream-O1-Image/tree/main/checkpoints
+
+## Examples
+
+### HiDream-O1-Image-Dev
+
+```
+.\bin\Release\sd-cli.exe -m  ..\..\ComfyUI\models\diffusion_models\hidream_o1_image_dev_bf16.safetensors -p "a lovely cat holding a sign says 
+'hidream o1 cpp'" --cfg-scale 1.0  -v -H 1024 -W 1024
+```
+
+<img width="256" alt="HiDream-O1-Image-Dev example" src="../assets/hidream-o1/dev_example.png" />
+

examples/cli/README.md

@@ -103,6 +103,8 @@ Generation Options:
   --hires-upscaler <string>                highres fix upscaler, Lanczos, Nearest, Latent, Latent (nearest), Latent
                                            (nearest-exact), Latent (antialiased), Latent (bicubic), Latent (bicubic
                                            antialiased), or a model name under --hires-upscalers-dir (default: Latent)
+  --extra-sample-args <string>             extra sampler args, key=value list. Currently lcm supports noise_clip_std,
+                                           noise_scale_start, noise_scale_end
   -H, --height <int>                       image height, in pixel space (default: 512)
   -W, --width <int>                        image width, in pixel space (default: 512)
   --steps <int>                            number of sample steps (default: 20)
@@ -162,10 +164,10 @@ Generation Options:
   -s, --seed                               RNG seed (default: 42, use random seed for < 0)
   --sampling-method                        sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m,
                                            dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s,
-                                           er_sde] (default: euler for Flux/SD3/Wan, euler_a otherwise)
+                                           er_sde, euler_cfg_pp, euler_a_cfg_pp] (default: euler for Flux/SD3/Wan, euler_a otherwise)
   --high-noise-sampling-method             (high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a,
                                            dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep,
-                                           res_2s, er_sde] default: euler for Flux/SD3/Wan, euler_a otherwise
+                                           res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp] default: euler for Flux/SD3/Wan, euler_a otherwise
   --scheduler                              denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits,
                                            smoothstep, sgm_uniform, simple, kl_optimal, lcm, bong_tangent], default:
                                            discrete

examples/common/common.cpp

@@ -807,6 +807,10 @@ ArgOptions SDGenerationParams::get_options() {
          "Latent (antialiased), Latent (bicubic), Latent (bicubic antialiased), or a model name "
          "under --hires-upscalers-dir (default: Latent)",
          &hires_upscaler},
+        {"",
+         "--extra-sample-args",
+         "extra sampler args, key=value list. Currently lcm supports noise_clip_std, noise_scale_start, noise_scale_end",
+         &extra_sample_args},
     };
 
     options.int_options = {
@@ -1244,12 +1248,12 @@ ArgOptions SDGenerationParams::get_options() {
          on_seed_arg},
         {"",
          "--sampling-method",
-         "sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde] "
+         "sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp]"
          "(default: euler for Flux/SD3/Wan, euler_a otherwise)",
          on_sample_method_arg},
         {"",
          "--high-noise-sampling-method",
-         "(high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde]"
+         "(high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp]"
          " default: euler for Flux/SD3/Wan, euler_a otherwise",
          on_high_noise_sample_method_arg},
         {"",
@@ -1607,6 +1611,7 @@ bool SDGenerationParams::from_json_str(
 
     auto parse_sample_params_json = [&](const json& sample_json,
                                         sd_sample_params_t& target_params,
+                                        std::string& target_extra_sample_args,
                                         std::vector<int>& target_skip_layers,
                                         std::vector<float>* target_custom_sigmas) {
         if (sample_json.contains("sample_steps") && sample_json["sample_steps"].is_number_integer()) {
@@ -1621,6 +1626,9 @@ bool SDGenerationParams::from_json_str(
         if (sample_json.contains("flow_shift") && sample_json["flow_shift"].is_number()) {
             target_params.flow_shift = sample_json["flow_shift"];
         }
+        if (sample_json.contains("extra_sample_args") && sample_json["extra_sample_args"].is_string()) {
+            target_extra_sample_args = sample_json["extra_sample_args"].get<std::string>();
+        }
         if (target_custom_sigmas != nullptr &&
             sample_json.contains("custom_sigmas") &&
             sample_json["custom_sigmas"].is_array()) {
@@ -1668,11 +1676,12 @@ bool SDGenerationParams::from_json_str(
     };
 
     if (j.contains("sample_params") && j["sample_params"].is_object()) {
-        parse_sample_params_json(j["sample_params"], sample_params, skip_layers, &custom_sigmas);
+        parse_sample_params_json(j["sample_params"], sample_params, extra_sample_args, skip_layers, &custom_sigmas);
     }
     if (j.contains("high_noise_sample_params") && j["high_noise_sample_params"].is_object()) {
         parse_sample_params_json(j["high_noise_sample_params"],
                                  high_noise_sample_params,
+                                 high_noise_extra_sample_args,
                                  high_noise_skip_layers,
                                  nullptr);
     }
@@ -2099,6 +2108,8 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
     high_noise_sample_params.guidance.slg.layer_count = high_noise_skip_layers.size();
     sample_params.custom_sigmas                       = custom_sigmas.empty() ? nullptr : custom_sigmas.data();
     sample_params.custom_sigmas_count                 = static_cast<int>(custom_sigmas.size());
+    sample_params.extra_sample_args                   = extra_sample_args.empty() ? nullptr : extra_sample_args.c_str();
+    high_noise_sample_params.extra_sample_args        = high_noise_extra_sample_args.empty() ? nullptr : high_noise_extra_sample_args.c_str();
     cache_params.scm_mask                             = scm_mask.empty() ? nullptr : scm_mask.c_str();
 
     sd_pm_params_t pm_params = {
@@ -2168,6 +2179,8 @@ sd_vid_gen_params_t SDGenerationParams::to_sd_vid_gen_params_t() {
     high_noise_sample_params.guidance.slg.layer_count = high_noise_skip_layers.size();
     sample_params.custom_sigmas                       = custom_sigmas.empty() ? nullptr : custom_sigmas.data();
     sample_params.custom_sigmas_count                 = static_cast<int>(custom_sigmas.size());
+    sample_params.extra_sample_args                   = extra_sample_args.empty() ? nullptr : extra_sample_args.c_str();
+    high_noise_sample_params.extra_sample_args        = high_noise_extra_sample_args.empty() ? nullptr : high_noise_extra_sample_args.c_str();
     cache_params.scm_mask                             = scm_mask.empty() ? nullptr : scm_mask.c_str();
 
     params.loras                    = lora_vec.empty() ? nullptr : lora_vec.data();
@@ -2306,6 +2319,7 @@ static json build_sampling_metadata_json(const sd_sample_params_t& sample_params
         {"eta", sample_params.eta},
         {"shifted_timestep", sample_params.shifted_timestep},
         {"flow_shift", sample_params.flow_shift},
+        {"extra_sample_args", safe_json_string(sample_params.extra_sample_args)},
         {"guidance",
          {
              {"txt_cfg", sample_params.guidance.txt_cfg},
@@ -2497,6 +2511,9 @@ std::string get_image_params(const SDContextParams& ctx_params,
     }
     parameter_string += "Guidance: " + std::to_string(gen_params.sample_params.guidance.distilled_guidance) + ", ";
     parameter_string += "Eta: " + std::to_string(gen_params.sample_params.eta) + ", ";
+    if (!gen_params.extra_sample_args.empty()) {
+        parameter_string += "Extra sample args: " + gen_params.extra_sample_args + ", ";
+    }
     parameter_string += "Seed: " + std::to_string(seed) + ", ";
     parameter_string += "Size: " + std::to_string(gen_params.get_resolved_width()) + "x" + std::to_string(gen_params.get_resolved_height()) + ", ";
     parameter_string += "Model: " + sd_basename(ctx_params.model_path) + ", ";

examples/common/common.h

@@ -168,6 +168,8 @@ struct SDGenerationParams {
 
     sd_sample_params_t sample_params;
     sd_sample_params_t high_noise_sample_params;
+    std::string extra_sample_args;
+    std::string high_noise_extra_sample_args;
     std::vector<int> skip_layers            = {7, 8, 9};
     std::vector<int> high_noise_skip_layers = {7, 8, 9};
 

examples/server/README.md

@@ -205,6 +205,8 @@ Default Generation Options:
   --hires-upscaler <string>                highres fix upscaler, Lanczos, Nearest, Latent, Latent (nearest), Latent
                                            (nearest-exact), Latent (antialiased), Latent (bicubic), Latent (bicubic
                                            antialiased), or a model name under --hires-upscalers-dir (default: Latent)
+  --extra-sample-args <string>             extra sampler args, key=value list. Currently lcm supports noise_clip_std,
+                                           noise_scale_start, noise_scale_end
   -H, --height <int>                       image height, in pixel space (default: 512)
   -W, --width <int>                        image width, in pixel space (default: 512)
   --steps <int>                            number of sample steps (default: 20)
@@ -264,10 +266,10 @@ Default Generation Options:
   -s, --seed                               RNG seed (default: 42, use random seed for < 0)
   --sampling-method                        sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m,
                                            dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s,
-                                           er_sde] (default: euler for Flux/SD3/Wan, euler_a otherwise)
+                                           er_sde, euler_cfg_pp, euler_a_cfg_pp] (default: euler for Flux/SD3/Wan, euler_a otherwise)
   --high-noise-sampling-method             (high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a,
                                            dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep,
-                                           res_2s, er_sde] default: euler for Flux/SD3/Wan, euler_a otherwise
+                                           res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp] default: euler for Flux/SD3/Wan, euler_a otherwise
   --scheduler                              denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits,
                                            smoothstep, sgm_uniform, simple, kl_optimal, lcm, bong_tangent], default:
                                            discrete

examples/server/main.cpp

@@ -145,7 +145,7 @@ int main(int argc, const char** argv) {
     register_sdapi_endpoints(svr, runtime);
     register_sdcpp_api_endpoints(svr, runtime);
 
-    LOG_INFO("listening on: %s:%d\n", svr_params.listen_ip.c_str(), svr_params.listen_port);
+    LOG_INFO("listening on: http://%s:%d\n", svr_params.listen_ip.c_str(), svr_params.listen_port);
     svr.listen(svr_params.listen_ip, svr_params.listen_port);
 
     {

examples/server/routes_sdapi.cpp

@@ -67,6 +67,10 @@ static enum sample_method_t get_sdapi_sample_method(std::string name) {
         {"k_res_multistep", RES_MULTISTEP_SAMPLE_METHOD},
         {"res 2s", RES_2S_SAMPLE_METHOD},
         {"k_res_2s", RES_2S_SAMPLE_METHOD},
+        {"euler_cfg_pp", EULER_CFG_PP_SAMPLE_METHOD},
+        {"k_euler_cfg_pp", EULER_CFG_PP_SAMPLE_METHOD},
+        {"euler_a_cfg_pp", EULER_CFG_PP_SAMPLE_METHOD},
+        {"k_euler_a_cfg_pp", EULER_CFG_PP_SAMPLE_METHOD},
     };
     auto it = hardcoded.find(name);
     return it != hardcoded.end() ? it->second : SAMPLE_METHOD_COUNT;

include/stable-diffusion.h

@@ -37,7 +37,6 @@ enum rng_type_t {
 
 enum sample_method_t {
     EULER_SAMPLE_METHOD,
-    EULER_FLOW_FLASH_SAMPLE_METHOD,
     EULER_A_SAMPLE_METHOD,
     HEUN_SAMPLE_METHOD,
     DPM2_SAMPLE_METHOD,
@@ -52,6 +51,8 @@ enum sample_method_t {
     RES_MULTISTEP_SAMPLE_METHOD,
     RES_2S_SAMPLE_METHOD,
     ER_SDE_SAMPLE_METHOD,
+    EULER_CFG_PP_SAMPLE_METHOD,
+    EULER_A_CFG_PP_SAMPLE_METHOD,
     SAMPLE_METHOD_COUNT
 };
 
@@ -239,6 +240,7 @@ typedef struct {
     float* custom_sigmas;
     int custom_sigmas_count;
     float flow_shift;
+    const char* extra_sample_args;
 } sd_sample_params_t;
 
 typedef struct {

src/conditioner.hpp

@@ -17,9 +17,8 @@ struct SDCondition {
     sd::Tensor<int32_t> c_input_ids;
     sd::Tensor<int32_t> c_position_ids;
     sd::Tensor<int32_t> c_token_types;
-    sd::Tensor<int32_t> c_image_embed_ranges;
     sd::Tensor<int32_t> c_vinput_mask;
-    std::vector<sd::Tensor<float>> c_vlm_images;
+    std::vector<std::pair<int, sd::Tensor<float>>> c_image_embeds;
     std::vector<sd::Tensor<float>> c_ref_images;
 
     std::vector<sd::Tensor<float>> extra_c_crossattns;
@@ -35,13 +34,12 @@ struct SDCondition {
         if (!c_crossattn.empty() || !c_vector.empty() || !c_concat.empty() ||
             !c_t5_ids.empty() || !c_t5_weights.empty() ||
             !c_input_ids.empty() || !c_position_ids.empty() ||
-            !c_token_types.empty() || !c_image_embed_ranges.empty() ||
+            !c_token_types.empty() || !c_vinput_mask.empty()) {
-            !c_vinput_mask.empty()) {
             return false;
         }
 
-        for (const auto& tensor : c_vlm_images) {
+        for (const auto& image_embed : c_image_embeds) {
-            if (!tensor.empty()) {
+            if (!image_embed.second.empty()) {
                 return false;
             }
         }

src/denoiser.hpp

@@ -2,6 +2,7 @@
 #define __DENOISER_HPP__
 
 #include <cmath>
+#include <string>
 #include <utility>
 
 #include "ggml_extend.hpp"
@@ -752,7 +753,7 @@ struct Flux2FlowDenoiser : public FluxFlowDenoiser {
     }
 };
 
-typedef std::function<sd::Tensor<float>(const sd::Tensor<float>&, float, int)> denoise_cb_t;
+typedef std::function<sd::Tensor<float>(const sd::Tensor<float>&, float, int, sd::Tensor<float>*)> denoise_cb_t;
 
 static std::pair<float, float> get_ancestral_step(float sigma_from,
                                                   float sigma_to,
@@ -828,7 +829,7 @@ static sd::Tensor<float> sample_euler_ancestral(denoise_cb_t model,
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
         float sigma       = sigmas[i];
-        auto denoised_opt = model(x, sigma, i + 1);
+        auto denoised_opt = model(x, sigma, i + 1, nullptr);
         if (denoised_opt.empty()) {
             return {};
         }
@@ -851,7 +852,7 @@ static sd::Tensor<float> sample_euler_flow(denoise_cb_t model,
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
         float sigma       = sigmas[i];
-        auto denoised_opt = model(x, sigma, i + 1);
+        auto denoised_opt = model(x, sigma, i + 1, nullptr);
         if (denoised_opt.empty()) {
             return {};
         }
@@ -867,38 +868,13 @@ static sd::Tensor<float> sample_euler_flow(denoise_cb_t model,
     return x;
 }
 
-static sd::Tensor<float> sample_euler_flow_flash(denoise_cb_t model,
-                                                 sd::Tensor<float> x,
-                                                 const std::vector<float>& sigmas,
-                                                 std::shared_ptr<RNG> rng,
-                                                 float eta) {
-    float s_noise = eta;
-    int steps     = static_cast<int>(sigmas.size()) - 1;
-    for (int i = 0; i < steps; i++) {
-        float sigma       = sigmas[i];
-        float sigma_next  = sigmas[i + 1];
-        auto denoised_opt = model(x, sigma, i + 1);
-        if (denoised_opt.empty()) {
-            return {};
-        }
-        sd::Tensor<float> denoised = std::move(denoised_opt);
-        if (sigma_next == 0.0f) {
-            x = std::move(denoised);
-            continue;
-        }
-        auto noise = sd::Tensor<float>::randn_like(x, rng);
-        x          = sigma_next * noise * s_noise + (1.0f - sigma_next) * denoised;
-    }
-    return x;
-}
-
 static sd::Tensor<float> sample_euler(denoise_cb_t model,
                                       sd::Tensor<float> x,
                                       const std::vector<float>& sigmas) {
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
         float sigma       = sigmas[i];
-        auto denoised_opt = model(x, sigma, i + 1);
+        auto denoised_opt = model(x, sigma, i + 1, nullptr);
         if (denoised_opt.empty()) {
             return {};
         }
@@ -914,7 +890,7 @@ static sd::Tensor<float> sample_heun(denoise_cb_t model,
                                      const std::vector<float>& sigmas) {
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
-        auto denoised_opt = model(x, sigmas[i], -(i + 1));
+        auto denoised_opt = model(x, sigmas[i], -(i + 1), nullptr);
         if (denoised_opt.empty()) {
             return {};
         }
@@ -925,7 +901,7 @@ static sd::Tensor<float> sample_heun(denoise_cb_t model,
             x += d * dt;
         } else {
             sd::Tensor<float> x2 = x + d * dt;
-            auto denoised2_opt   = model(x2, sigmas[i + 1], i + 1);
+            auto denoised2_opt   = model(x2, sigmas[i + 1], i + 1, nullptr);
             if (denoised2_opt.empty()) {
                 return {};
             }
@@ -942,7 +918,7 @@ static sd::Tensor<float> sample_dpm2(denoise_cb_t model,
                                      const std::vector<float>& sigmas) {
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
-        auto denoised_opt = model(x, sigmas[i], -(i + 1));
+        auto denoised_opt = model(x, sigmas[i], -(i + 1), nullptr);
         if (denoised_opt.empty()) {
             return {};
         }
@@ -955,7 +931,7 @@ static sd::Tensor<float> sample_dpm2(denoise_cb_t model,
             float dt_1           = sigma_mid - sigmas[i];
             float dt_2           = sigmas[i + 1] - sigmas[i];
             sd::Tensor<float> x2 = x + d * dt_1;
-            auto denoised2_opt   = model(x2, sigma_mid, i + 1);
+            auto denoised2_opt   = model(x2, sigma_mid, i + 1, nullptr);
             if (denoised2_opt.empty()) {
                 return {};
             }
@@ -976,7 +952,7 @@ static sd::Tensor<float> sample_dpmpp_2s_ancestral(denoise_cb_t model,
 
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
-        auto denoised_opt = model(x, sigmas[i], -(i + 1));
+        auto denoised_opt = model(x, sigmas[i], -(i + 1), nullptr);
         if (denoised_opt.empty()) {
             return {};
         }
@@ -992,7 +968,7 @@ static sd::Tensor<float> sample_dpmpp_2s_ancestral(denoise_cb_t model,
             float s              = t + 0.5f * h;
             float sigma_s        = sigma_fn(s);
             sd::Tensor<float> x2 = (sigma_s / sigma_fn(t)) * x - (exp(-h * 0.5f) - 1) * denoised;
-            auto denoised2_opt   = model(x2, sigma_s, i + 1);
+            auto denoised2_opt   = model(x2, sigma_s, i + 1, nullptr);
             if (denoised2_opt.empty()) {
                 return {};
             }
@@ -1019,7 +995,7 @@ static sd::Tensor<float> sample_dpmpp_2s_ancestral_flow(denoise_cb_t model,
 
         bool opt_first_step = (1.0 - sigma < 1e-6);
 
-        auto denoised_opt = model(x, sigma, (opt_first_step ? 1 : -1) * (i + 1));
+        auto denoised_opt = model(x, sigma, (opt_first_step ? 1 : -1) * (i + 1), nullptr);
         if (denoised_opt.empty()) {
             return {};
         }
@@ -1048,8 +1024,8 @@ static sd::Tensor<float> sample_dpmpp_2s_ancestral_flow(denoise_cb_t model,
                 // so sigma_s = 1 = sigma, and sigma_s_i_ratio = sigma_s / sigma = 1
                 // u = (x*sigma_s_i_ratio)+(denoised*(1.0f-sigma_s_i_ratio))
                 //   = (x*1)+(denoised*0) = x
-                // so D_i = model(u, sigma_s, i + 1)
+                // so D_i = model(u, sigma_s, i + 1, nullptr)
-                //        = model(x, sigma,   i + 1)
+                //        = model(x, sigma,   i + 1, nullptr)
                 //        = denoised
                 D_i = denoised;
 
@@ -1082,7 +1058,7 @@ static sd::Tensor<float> sample_dpmpp_2s_ancestral_flow(denoise_cb_t model,
                 float sigma_s_i_ratio = sigma_s / sigma;
                 sd::Tensor<float> u   = (x * sigma_s_i_ratio) + (denoised * (1.0f - sigma_s_i_ratio));
 
-                auto denoised2_opt = model(u, sigma_s, i + 1);
+                auto denoised2_opt = model(u, sigma_s, i + 1, nullptr);
                 if (denoised2_opt.empty()) {
                     return {};
                 }
@@ -1109,7 +1085,7 @@ static sd::Tensor<float> sample_dpmpp_2m(denoise_cb_t model,
 
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
-        auto denoised_opt = model(x, sigmas[i], i + 1);
+        auto denoised_opt = model(x, sigmas[i], i + 1, nullptr);
         if (denoised_opt.empty()) {
             return {};
         }
@@ -1141,7 +1117,7 @@ static sd::Tensor<float> sample_dpmpp_2m_v2(denoise_cb_t model,
 
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
-        auto denoised_opt = model(x, sigmas[i], i + 1);
+        auto denoised_opt = model(x, sigmas[i], i + 1, nullptr);
         if (denoised_opt.empty()) {
             return {};
         }
@@ -1173,10 +1149,83 @@ static sd::Tensor<float> sample_lcm(denoise_cb_t model,
                                     sd::Tensor<float> x,
                                     const std::vector<float>& sigmas,
                                     std::shared_ptr<RNG> rng,
-                                    bool is_flow_denoiser) {
+                                    bool is_flow_denoiser,
+                                    const char* extra_sample_args = nullptr) {
+    struct LCMSampleArgs {
+        float noise_clip_std    = 0.0f;
+        float noise_scale_start = 1.0f;
+        float noise_scale_end   = 1.0f;
+    };
+
+    auto trim = [](std::string value) -> std::string {
+        const char* whitespace = " \t\r\n";
+        size_t begin           = value.find_first_not_of(whitespace);
+        if (begin == std::string::npos) {
+            return "";
+        }
+        size_t end = value.find_last_not_of(whitespace);
+        return value.substr(begin, end - begin + 1);
+    };
+
+    LCMSampleArgs args;
+    if (extra_sample_args != nullptr && extra_sample_args[0] != '\0') {
+        std::string raw(extra_sample_args);
+        size_t start                   = 0;
+        bool noise_scale_end_was_set   = false;
+        bool noise_scale_start_was_set = false;
+        auto parse_arg                 = [&](const std::string& item) {
+            std::string token = trim(item);
+            if (token.empty()) {
+                return;
+            }
+            size_t eq = token.find('=');
+            if (eq == std::string::npos) {
+                LOG_WARN("ignoring invalid lcm extra sample arg '%s'", token.c_str());
+                return;
+            }
+
+            std::string key   = trim(token.substr(0, eq));
+            std::string value = trim(token.substr(eq + 1));
+            float parsed      = 0.0f;
+            try {
+                size_t consumed = 0;
+                parsed          = std::stof(value, &consumed);
+                if (trim(value.substr(consumed)).size() != 0) {
+                    LOG_WARN("ignoring invalid lcm extra sample arg '%s'", token.c_str());
+                    return;
+                }
+            } catch (const std::exception&) {
+                LOG_WARN("ignoring invalid lcm extra sample arg '%s'", token.c_str());
+                return;
+            }
+
+            if (key == "noise_clip_std") {
+                args.noise_clip_std = parsed;
+            } else if (key == "noise_scale_start") {
+                args.noise_scale_start    = parsed;
+                noise_scale_start_was_set = true;
+            } else if (key == "noise_scale_end") {
+                args.noise_scale_end    = parsed;
+                noise_scale_end_was_set = true;
+            } else {
+                LOG_WARN("ignoring unknown lcm extra sample arg '%s'", key.c_str());
+            }
+        };
+
+        for (size_t pos = 0; pos <= raw.size(); ++pos) {
+            if (pos == raw.size() || raw[pos] == ',' || raw[pos] == ';') {
+                parse_arg(raw.substr(start, pos - start));
+                start = pos + 1;
+            }
+        }
+        if (noise_scale_start_was_set && !noise_scale_end_was_set) {
+            args.noise_scale_end = args.noise_scale_start;
+        }
+    }
+
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
-        auto denoised_opt = model(x, sigmas[i], i + 1);
+        auto denoised_opt = model(x, sigmas[i], i + 1, nullptr);
         if (denoised_opt.empty()) {
             return {};
         }
@@ -1185,7 +1234,27 @@ static sd::Tensor<float> sample_lcm(denoise_cb_t model,
             if (is_flow_denoiser) {
                 x *= (1 - sigmas[i + 1]);
             }
-            x += sd::Tensor<float>::randn_like(x, rng) * sigmas[i + 1];
+            auto noise = sd::Tensor<float>::randn_like(x, rng);
+            if (args.noise_clip_std > 0.0f && noise.numel() > 0) {
+                double mean = 0.0;
+                for (int64_t j = 0; j < noise.numel(); ++j) {
+                    mean += static_cast<double>(noise[j]);
+                }
+                mean /= static_cast<double>(noise.numel());
+
+                double variance = 0.0;
+                for (int64_t j = 0; j < noise.numel(); ++j) {
+                    double centered = static_cast<double>(noise[j]) - mean;
+                    variance += centered * centered;
+                }
+                variance /= static_cast<double>(noise.numel());
+
+                float clip_val = args.noise_clip_std * static_cast<float>(std::sqrt(variance));
+                noise          = sd::ops::clamp(noise, -clip_val, clip_val);
+            }
+            float t           = steps > 1 ? static_cast<float>(i) / static_cast<float>(steps - 1) : 0.0f;
+            float noise_scale = args.noise_scale_start + (args.noise_scale_end - args.noise_scale_start) * t;
+            x += noise * (sigmas[i + 1] * noise_scale);
         }
     }
     return x;
@@ -1202,7 +1271,7 @@ static sd::Tensor<float> sample_ipndm(denoise_cb_t model,
         float sigma      = sigmas[i];
         float sigma_next = sigmas[i + 1];
 
-        auto denoised_opt = model(x, sigma, i + 1);
+        auto denoised_opt = model(x, sigma, i + 1, nullptr);
         if (denoised_opt.empty()) {
             return {};
         }
@@ -1246,7 +1315,7 @@ static sd::Tensor<float> sample_ipndm_v(denoise_cb_t model,
         float sigma  = sigmas[i];
         float t_next = sigmas[i + 1];
 
-        auto denoised_opt = model(x, sigma, i + 1);
+        auto denoised_opt = model(x, sigma, i + 1, nullptr);
         if (denoised_opt.empty()) {
             return {};
         }
@@ -1308,7 +1377,7 @@ static sd::Tensor<float> sample_res_multistep(denoise_cb_t model,
 
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
-        auto denoised_opt = model(x, sigmas[i], i + 1);
+        auto denoised_opt = model(x, sigmas[i], i + 1, nullptr);
         if (denoised_opt.empty()) {
             return {};
         }
@@ -1385,7 +1454,7 @@ static sd::Tensor<float> sample_res_2s(denoise_cb_t model,
         float sigma_from = sigmas[i];
         float sigma_to   = sigmas[i + 1];
 
-        auto denoised_opt = model(x, sigma_from, -(i + 1));
+        auto denoised_opt = model(x, sigma_from, -(i + 1), nullptr);
         if (denoised_opt.empty()) {
             return {};
         }
@@ -1411,7 +1480,7 @@ static sd::Tensor<float> sample_res_2s(denoise_cb_t model,
             sd::Tensor<float> eps1 = denoised - x0;
             sd::Tensor<float> x2   = x0 + eps1 * (h * a21);
 
-            auto denoised2_opt = model(x2, sigma_c2, i + 1);
+            auto denoised2_opt = model(x2, sigma_c2, i + 1, nullptr);
             if (denoised2_opt.empty()) {
                 return {};
             }
@@ -1488,7 +1557,7 @@ static sd::Tensor<float> sample_er_sde(denoise_cb_t model,
 
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
-        sd::Tensor<float> denoised = model(x, sigmas[i], i + 1);
+        sd::Tensor<float> denoised = model(x, sigmas[i], i + 1, nullptr);
         if (denoised.empty()) {
             return {};
         }
@@ -1574,7 +1643,7 @@ static sd::Tensor<float> sample_ddim_trailing(denoise_cb_t model,
         float sigma    = sigmas[i];
         float sigma_to = sigmas[i + 1];
 
-        auto model_output_opt = model(x, sigma, i + 1);
+        auto model_output_opt = model(x, sigma, i + 1, nullptr);
         if (model_output_opt.empty()) {
             return {};
         }
@@ -1646,7 +1715,7 @@ static sd::Tensor<float> sample_tcd(denoise_cb_t model,
         int timestep_s    = (int)floor((1 - eta) * prev_timestep);
         float sigma       = sigmas[i];
 
-        auto model_output_opt = model(x, sigma, i + 1);
+        auto model_output_opt = model(x, sigma, i + 1, nullptr);
         if (model_output_opt.empty()) {
             return {};
         }
@@ -1674,6 +1743,56 @@ static sd::Tensor<float> sample_tcd(denoise_cb_t model,
     return x;
 }
 
+static sd::Tensor<float> sample_euler_cfg_pp(denoise_cb_t model,
+                                             sd::Tensor<float> x,
+                                             const std::vector<float>& sigmas) {
+    int steps = static_cast<int>(sigmas.size()) - 1;
+    for (int i = 0; i < steps; i++) {
+        float sigma = sigmas[i];
+        sd::Tensor<float> uncond_denoised;
+
+        auto denoised_opt = model(x, sigma, i + 1, &uncond_denoised);
+        if (denoised_opt.empty() || uncond_denoised.empty()) {
+            return {};
+        }
+
+        sd::Tensor<float> denoised = std::move(denoised_opt);
+        sd::Tensor<float> d        = (x - uncond_denoised) / sigma;
+
+        x = denoised + d * sigmas[i + 1];
+    }
+    return x;
+}
+
+static sd::Tensor<float> sample_euler_ancestral_cfg_pp(denoise_cb_t model,
+                                                       sd::Tensor<float> x,
+                                                       const std::vector<float>& sigmas,
+                                                       std::shared_ptr<RNG> rng,
+                                                       float eta) {
+    int steps = static_cast<int>(sigmas.size()) - 1;
+    for (int i = 0; i < steps; i++) {
+        float sigma = sigmas[i];
+        sd::Tensor<float> uncond_denoised;
+
+        auto denoised_opt = model(x, sigma, i + 1, &uncond_denoised);
+        if (denoised_opt.empty() || uncond_denoised.empty()) {
+            return {};
+        }
+
+        sd::Tensor<float> denoised = std::move(denoised_opt);
+        sd::Tensor<float> d        = (x - uncond_denoised) / sigma;
+
+        auto [sigma_down, sigma_up] = get_ancestral_step(sigmas[i], sigmas[i + 1], eta);
+
+        x = denoised + d * sigma_down;
+
+        if (sigmas[i + 1] > 0) {
+            x += sd::Tensor<float>::randn_like(x, rng) * sigma_up;
+        }
+    }
+    return x;
+}
+
 // k diffusion reverse ODE: dx = (x - D(x;\sigma)) / \sigma dt; \sigma(t) = t
 static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
                                             denoise_cb_t model,
@@ -1681,10 +1800,9 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
                                             std::vector<float> sigmas,
                                             std::shared_ptr<RNG> rng,
                                             float eta,
-                                            bool is_flow_denoiser) {
+                                            bool is_flow_denoiser,
+                                            const char* extra_sample_args) {
     switch (method) {
-        case EULER_FLOW_FLASH_SAMPLE_METHOD:
-            return sample_euler_flow_flash(model, std::move(x), sigmas, rng, eta);
         case EULER_A_SAMPLE_METHOD:
             if (is_flow_denoiser)
                 return sample_euler_flow(model, std::move(x), sigmas, rng, eta);
@@ -1706,7 +1824,7 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
         case DPMPP2Mv2_SAMPLE_METHOD:
             return sample_dpmpp_2m_v2(model, std::move(x), sigmas);
         case LCM_SAMPLE_METHOD:
-            return sample_lcm(model, std::move(x), sigmas, rng, is_flow_denoiser);
+            return sample_lcm(model, std::move(x), sigmas, rng, is_flow_denoiser, extra_sample_args);
         case IPNDM_SAMPLE_METHOD:
             return sample_ipndm(model, std::move(x), sigmas);
         case IPNDM_V_SAMPLE_METHOD:
@@ -1721,6 +1839,10 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
             return sample_ddim_trailing(model, std::move(x), sigmas, rng, eta);
         case TCD_SAMPLE_METHOD:
             return sample_tcd(model, std::move(x), sigmas, rng, eta);
+        case EULER_CFG_PP_SAMPLE_METHOD:
+            return sample_euler_cfg_pp(model, std::move(x), sigmas);
+        case EULER_A_CFG_PP_SAMPLE_METHOD:
+            return sample_euler_ancestral_cfg_pp(model, std::move(x), sigmas, rng, eta);
         default:
             return {};
     }

src/diffusion_model.hpp

@@ -26,9 +26,9 @@ struct DiffusionParams {
     const sd::Tensor<int32_t>* input_ids                               = nullptr;
     const sd::Tensor<int32_t>* input_pos                               = nullptr;
     const sd::Tensor<int32_t>* token_types                             = nullptr;
-    const sd::Tensor<int32_t>* image_embed_ranges     = nullptr;
     const sd::Tensor<int32_t>* vinput_mask                             = nullptr;
     const std::vector<sd::Tensor<float>>* vlm_images                   = nullptr;
+    const std::vector<std::pair<int, sd::Tensor<float>>>* image_embeds = nullptr;
     bool increase_ref_index                                            = false;
     int num_video_frames                                               = -1;
     const std::vector<sd::Tensor<float>>* controls                     = nullptr;
@@ -545,17 +545,16 @@ struct HiDreamO1Model : public DiffusionModel {
         GGML_ASSERT(diffusion_params.input_ids != nullptr);
         GGML_ASSERT(diffusion_params.input_pos != nullptr);
         GGML_ASSERT(diffusion_params.token_types != nullptr);
-        static const sd::Tensor<int32_t> empty_image_embed_ranges;
         static const std::vector<sd::Tensor<float>> empty_images;
+        static const std::vector<std::pair<int, sd::Tensor<float>>> empty_image_embeds;
         return hidream_o1.compute(n_threads,
                                   *diffusion_params.x,
                                   *diffusion_params.timesteps,
                                   *diffusion_params.input_ids,
                                   *diffusion_params.input_pos,
                                   *diffusion_params.token_types,
-                                  diffusion_params.image_embed_ranges ? *diffusion_params.image_embed_ranges : empty_image_embed_ranges,
+                                  tensor_or_empty(diffusion_params.vinput_mask),
-                                  diffusion_params.vinput_mask ? *diffusion_params.vinput_mask : empty_image_embed_ranges,
+                                  diffusion_params.image_embeds ? *diffusion_params.image_embeds : empty_image_embeds,
-                                  diffusion_params.vlm_images ? *diffusion_params.vlm_images : empty_images,
                                   diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_images);
     }
 };

src/ggml_extend.hpp

@@ -280,6 +280,9 @@ __STATIC_INLINE__ void print_sd_tensor(const sd::Tensor<T>& tensor, bool shape_o
     if (shape_only) {
         return;
     }
+    if (tensor.empty()) {
+        return;
+    }
     int range                  = 3;
     std::vector<int64_t> shape = tensor.shape();
     while (shape.size() < 4) {
@@ -2021,9 +2024,13 @@ protected:
             ggml_backend_buffer_t src_buf = sd::ggml_graph_cut::tensor_buffer(src);
             ggml_backend_buffer_t dst_buf = sd::ggml_graph_cut::tensor_buffer(dst);
             if (src_buf == nullptr || dst_buf == nullptr) {
-                LOG_ERROR("%s cache copy tensor buffer missing: name=%s src_buffer=%p src_view_src=%p src_view_src_buffer=%p dst_buffer=%p",
+                LOG_ERROR("%s cache copy tensor buffer missing: name=%s op=%s src0=%p src0_name=%s src0_buffer=%p src_buffer=%p src_view_src=%p src_view_src_buffer=%p dst_buffer=%p",
                           get_desc().c_str(),
                           src && src->name[0] != '\0' ? src->name : "<unnamed>",
+                          src ? ggml_op_name(src->op) : "<null>",
+                          src ? src->src[0] : nullptr,
+                          (src && src->src[0] && src->src[0]->name[0] != '\0') ? src->src[0]->name : "<unnamed>",
+                          (src && src->src[0]) ? sd::ggml_graph_cut::tensor_buffer(src->src[0]) : nullptr,
                           src ? src->buffer : nullptr,
                           src ? src->view_src : nullptr,
                           (src && src->view_src) ? src->view_src->buffer : nullptr,
@@ -2055,6 +2062,42 @@ protected:
         return true;
     }
 
+    template <typename T>
+    std::optional<sd::Tensor<T>> read_graph_tensor(ggml_tensor* tensor, const char* label) {
+        if (tensor == nullptr) {
+            LOG_ERROR("%s %s tensor is null", get_desc().c_str(), label);
+            return std::nullopt;
+        }
+        if (tensor->type != sd::GGMLTypeTraits<T>::type) {
+            LOG_ERROR("%s %s tensor type mismatch: got %s",
+                      get_desc().c_str(),
+                      label,
+                      ggml_type_name(tensor->type));
+            return std::nullopt;
+        }
+        ggml_backend_buffer_t buf = sd::ggml_graph_cut::tensor_buffer(tensor);
+        if (buf == nullptr) {
+            LOG_ERROR("%s %s tensor buffer missing: name=%s op=%s buffer=%p view_src=%p view_src_buffer=%p data=%p",
+                      get_desc().c_str(),
+                      label,
+                      tensor->name[0] != '\0' ? tensor->name : "<unnamed>",
+                      ggml_op_name(tensor->op),
+                      tensor->buffer,
+                      tensor->view_src,
+                      tensor->view_src ? tensor->view_src->buffer : nullptr,
+                      tensor->data);
+            return std::nullopt;
+        }
+
+        sd::Tensor<T> result(sd::shape_from_ggml(tensor));
+        if (tensor->view_src != nullptr || !ggml_is_contiguous(tensor) || tensor->buffer == nullptr) {
+            ggml_backend_tensor_get(tensor, result.data(), 0, ggml_nbytes(tensor));
+        } else {
+            ggml_backend_tensor_get(tensor, result.data(), 0, ggml_nbytes(tensor));
+        }
+        return result;
+    }
+
     void copy_data_to_backend_tensor(ggml_cgraph* gf, bool clear_after_copy = true) {
         GGML_ASSERT(gf != nullptr);
         std::unordered_set<const ggml_tensor*> graph_tensor_set;
@@ -2075,6 +2118,9 @@ protected:
                 continue;
             }
             const char* name = ggml_get_name(tensor);
+            if (graph_tensor_set.find(tensor) == graph_tensor_set.end()) {
+                continue;
+            }
             if (tensor->buffer == nullptr) {
                 LOG_WARN("%s skip backend tensor copy: tensor buffer not set, name='%s', ne=[%lld,%lld,%lld,%lld], type=%s",
                          get_desc().c_str(),
@@ -2087,10 +2133,6 @@ protected:
                 continue;
             }
 
-            if (graph_tensor_set.find(tensor) == graph_tensor_set.end()) {
-                continue;
-            }
-
             ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
             if (buf == nullptr) {
                 LOG_WARN("%s graph exec skip tensor copy: name=%s op=%s reason=buffer_not_set data=%p view_src=%p view_src_buffer=%p",
@@ -2476,11 +2518,32 @@ protected:
             return std::nullopt;
         }
 
+        std::unordered_set<const ggml_tensor*> debug_graph_tensor_set;
+        const int n_debug_leafs = sd::ggml_graph_cut::leaf_count(gf);
+        const int n_debug_nodes = ggml_graph_n_nodes(gf);
+        debug_graph_tensor_set.reserve(static_cast<size_t>(n_debug_leafs + n_debug_nodes));
+        for (int i = 0; i < n_debug_leafs; ++i) {
+            debug_graph_tensor_set.insert(sd::ggml_graph_cut::leaf_tensor(gf, i));
+        }
+        for (int i = 0; i < n_debug_nodes; ++i) {
+            debug_graph_tensor_set.insert(ggml_graph_node(gf, i));
+        }
+
         for (const auto& entry : debug_tensors) {
             auto tensor = entry.first;
             if (tensor == nullptr) {
                 continue;
             }
+            if (debug_graph_tensor_set.find(tensor) == debug_graph_tensor_set.end()) {
+                continue;
+            }
+            ggml_backend_buffer_t tensor_buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
+            if (tensor_buf == nullptr) {
+                LOG_WARN("%s skip debug tensor '%s': tensor buffer not set",
+                         get_desc().c_str(),
+                         entry.second.c_str());
+                continue;
+            }
             if (tensor->type != GGML_TYPE_F32) {
                 LOG_WARN("%s skip debug tensor '%s': only GGML_TYPE_F32 is supported, got %s",
                          get_desc().c_str(),
@@ -2505,7 +2568,15 @@ protected:
         auto result         = ggml_get_tensor(compute_ctx, final_result_name.c_str());
         std::optional<sd::Tensor<T>> output;
         if (!no_return) {
-            output = sd::make_sd_tensor_from_ggml<T>(result);
+            output = read_graph_tensor<T>(result, "output");
+            if (!output.has_value()) {
+                if (free_compute_buffer_immediately) {
+                    free_compute_buffer();
+                } else if (use_partial_param_offload) {
+                    restore_partial_params();
+                }
+                return std::nullopt;
+            }
         } else {
             output = sd::Tensor<T>();
         }
@@ -2645,6 +2716,23 @@ public:
 
     bool alloc_params_buffer() {
         size_t num_tensors = ggml_tensor_num(params_ctx);
+        if (num_tensors > 0) {
+            // ggml_backend_alloc_ctx_tensors fails when all tensors are already allocated
+            // (typical for memory-mapped weights). See ggml-alloc.c n_buffers==0 branch.
+            bool all_have_data = true;
+            for (ggml_tensor* t = ggml_get_first_tensor(params_ctx); t != nullptr; t = ggml_get_next_tensor(params_ctx, t)) {
+                if (t->data == nullptr) {
+                    all_have_data = false;
+                    break;
+                }
+            }
+            if (all_have_data) {
+                LOG_DEBUG("%s all params already mmap-allocated (no separate buffer needed)", get_desc().c_str());
+                params_buffer = nullptr;
+                rebuild_params_tensor_set();
+                return true;
+            }
+        }
         params_buffer = ggml_backend_alloc_ctx_tensors(params_ctx, params_backend);
         if (params_buffer == nullptr) {
             LOG_ERROR("%s alloc params backend buffer failed, num_tensors = %i",

src/ggml_graph_cut.cpp

@@ -45,6 +45,21 @@ namespace sd::ggml_graph_cut {
         return params_tensor_set.find(tensor) != params_tensor_set.end();
     }
 
+    static int graph_node_index_by_name(ggml_cgraph* gf, const char* name) {
+        GGML_ASSERT(gf != nullptr);
+        if (name == nullptr || name[0] == '\0') {
+            return -1;
+        }
+        const int n_nodes = ggml_graph_n_nodes(gf);
+        for (int i = 0; i < n_nodes; ++i) {
+            ggml_tensor* node = ggml_graph_node(gf, i);
+            if (node != nullptr && std::strcmp(node->name, name) == 0) {
+                return i;
+            }
+        }
+        return -1;
+    }
+
     static Plan::InputShape input_shape(const ggml_tensor* tensor) {
         Plan::InputShape shape;
         if (tensor == nullptr) {
@@ -244,6 +259,11 @@ namespace sd::ggml_graph_cut {
         if (tensor == nullptr) {
             return nullptr;
         }
+        if (tensor_buffer(tensor) == nullptr && tensor->src[0] != nullptr &&
+            ggml_nelements(tensor->src[0]) == ggml_nelements(tensor) &&
+            ggml_nbytes(tensor->src[0]) == ggml_nbytes(tensor)) {
+            return cache_source_tensor(tensor->src[0]);
+        }
         return tensor->view_src ? tensor->view_src : tensor;
     }
 
@@ -503,11 +523,15 @@ namespace sd::ggml_graph_cut {
                           log_desc);
         }
 
-        ggml_tensor* final_output = ggml_graph_node(gf, -1);
+        int final_output_index = graph_node_index_by_name(gf, "ggml_runner_final_result_tensor");
-        if (final_output != nullptr && available_cut_output_node_indices.find(n_nodes - 1) == available_cut_output_node_indices.end()) {
+        if (final_output_index < 0) {
+            final_output_index = n_nodes - 1;
+        }
+        ggml_tensor* final_output = final_output_index >= 0 ? ggml_graph_node(gf, final_output_index) : nullptr;
+        if (final_output != nullptr && available_cut_output_node_indices.find(final_output_index) == available_cut_output_node_indices.end()) {
             Segment final_segment;
             final_segment.group_name = "ggml_runner.final";
-            final_segment.output_node_indices.push_back(n_nodes - 1);
+            final_segment.output_node_indices.push_back(final_output_index);
             build_segment(gf,
                           plan,
                           final_segment,

src/hidream_o1.hpp

@@ -70,12 +70,16 @@ namespace HiDreamO1 {
         sd::Tensor<float> resized;
         if (s1 < s2) {
             int64_t resized_h = static_cast<int64_t>(std::llround(height / s1));
-            resized           = sd::ops::interpolate(image, {new_size.first, resized_h, image.shape()[2], image.shape()[3]});
+            resized           = sd::ops::interpolate(image,
+                                                     {new_size.first, resized_h, image.shape()[2], image.shape()[3]},
+                                                     sd::ops::InterpolateMode::Bicubic);
             int64_t top       = (resized_h - new_size.second) / 2;
             resized           = sd::ops::slice(resized, 1, top, top + new_size.second);
         } else {
             int64_t resized_w = static_cast<int64_t>(std::llround(width / s2));
-            resized           = sd::ops::interpolate(image, {resized_w, new_size.second, image.shape()[2], image.shape()[3]});
+            resized           = sd::ops::interpolate(image,
+                                                     {resized_w, new_size.second, image.shape()[2], image.shape()[3]},
+                                                     sd::ops::InterpolateMode::Bicubic);
             int64_t left      = (resized_w - new_size.first) / 2;
             resized           = sd::ops::slice(resized, 0, left, left + new_size.first);
         }
@@ -203,298 +207,34 @@ namespace HiDreamO1 {
     struct HiDreamO1Params {
         LLM::LLMParams llm;
         int patch_size = PATCH_SIZE;
-        int num_position_embeddings = 2304;
-        std::vector<int> deepstack_visual_indexes;
     };
 
-    struct VisionMLP : public GGMLBlock {
+    static inline HiDreamO1Params make_hidream_o1_params() {
-        VisionMLP(int64_t hidden_size, int64_t intermediate_size) {
+        HiDreamO1Params params;
-            blocks["linear_fc1"] = std::make_shared<Linear>(hidden_size, intermediate_size, true);
+        params.llm.arch                           = LLM::LLMArch::QWEN3_VL;
-            blocks["linear_fc2"] = std::make_shared<Linear>(intermediate_size, hidden_size, true);
+        params.llm.hidden_size                    = 4096;
+        params.llm.intermediate_size              = 12288;
+        params.llm.num_layers                     = 36;
+        params.llm.num_heads                      = 32;
+        params.llm.num_kv_heads                   = 8;
+        params.llm.head_dim                       = 128;
+        params.llm.qkv_bias                       = false;
+        params.llm.qk_norm                        = true;
+        params.llm.vocab_size                     = 151936;
+        params.llm.rms_norm_eps                   = 1e-6f;
+        params.llm.vision.arch                    = LLM::LLMVisionArch::QWEN3_VL;
+        params.llm.vision.num_layers              = 27;
+        params.llm.vision.hidden_size             = 1152;
+        params.llm.vision.intermediate_size       = 4304;
+        params.llm.vision.num_heads               = 16;
+        params.llm.vision.out_hidden_size         = 4096;
+        params.llm.vision.patch_size              = 16;
+        params.llm.vision.spatial_merge_size      = 2;
+        params.llm.vision.temporal_patch_size     = 2;
+        params.llm.vision.num_position_embeddings = 2304;
+        return params;
     }
 
-        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
-            auto linear_fc1 = std::dynamic_pointer_cast<Linear>(blocks["linear_fc1"]);
-            auto linear_fc2 = std::dynamic_pointer_cast<Linear>(blocks["linear_fc2"]);
-
-            x = linear_fc1->forward(ctx, x);
-            x = ggml_ext_gelu(ctx->ggml_ctx, x);
-            x = linear_fc2->forward(ctx, x);
-            return x;
-        }
-    };
-
-    struct VisionPatchEmbed : public GGMLBlock {
-        int patch_size;
-        int temporal_patch_size;
-        int64_t in_channels;
-        int64_t embed_dim;
-
-        VisionPatchEmbed(int patch_size,
-                         int temporal_patch_size,
-                         int64_t in_channels,
-                         int64_t embed_dim)
-            : patch_size(patch_size),
-              temporal_patch_size(temporal_patch_size),
-              in_channels(in_channels),
-              embed_dim(embed_dim) {
-            blocks["proj"] = std::make_shared<Conv3d>(in_channels,
-                                                      embed_dim,
-                                                      std::tuple<int, int, int>{temporal_patch_size, patch_size, patch_size},
-                                                      std::tuple<int, int, int>{temporal_patch_size, patch_size, patch_size},
-                                                      std::tuple<int, int, int>{0, 0, 0},
-                                                      std::tuple<int, int, int>{1, 1, 1},
-                                                      true);
-        }
-
-        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
-            auto proj = std::dynamic_pointer_cast<Conv3d>(blocks["proj"]);
-            x         = ggml_reshape_4d(ctx->ggml_ctx,
-                                        x,
-                                        patch_size,
-                                        patch_size,
-                                        temporal_patch_size,
-                                        ggml_nelements(x) / (temporal_patch_size * patch_size * patch_size));
-            x         = proj->forward(ctx, x);
-            x         = ggml_reshape_2d(ctx->ggml_ctx, x, embed_dim, ggml_nelements(x) / embed_dim);
-            return x;
-        }
-    };
-
-    struct VisionPatchMerger : public GGMLBlock {
-        int64_t hidden_size;
-        bool use_postshuffle_norm;
-
-        VisionPatchMerger(int64_t dim,
-                          int64_t context_dim,
-                          int spatial_merge_size,
-                          bool use_postshuffle_norm)
-            : hidden_size(context_dim * spatial_merge_size * spatial_merge_size),
-              use_postshuffle_norm(use_postshuffle_norm) {
-            blocks["norm"]       = std::make_shared<LayerNorm>(use_postshuffle_norm ? hidden_size : context_dim, 1e-6f);
-            blocks["linear_fc1"] = std::make_shared<Linear>(hidden_size, hidden_size, true);
-            blocks["linear_fc2"] = std::make_shared<Linear>(hidden_size, dim, true);
-        }
-
-        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
-            auto norm       = std::dynamic_pointer_cast<LayerNorm>(blocks["norm"]);
-            auto linear_fc1 = std::dynamic_pointer_cast<Linear>(blocks["linear_fc1"]);
-            auto linear_fc2 = std::dynamic_pointer_cast<Linear>(blocks["linear_fc2"]);
-
-            x = norm->forward(ctx, x);
-            x = ggml_reshape_2d(ctx->ggml_ctx, x, hidden_size, ggml_nelements(x) / hidden_size);
-            x = linear_fc1->forward(ctx, x);
-            x = ggml_ext_gelu(ctx->ggml_ctx, x);
-            x = linear_fc2->forward(ctx, x);
-            return x;
-        }
-    };
-
-    struct VisionAttention : public GGMLBlock {
-        int head_dim;
-        int num_heads;
-
-        VisionAttention(int64_t hidden_size, int num_heads)
-            : num_heads(num_heads) {
-            head_dim = static_cast<int>(hidden_size / num_heads);
-            GGML_ASSERT(num_heads * head_dim == hidden_size);
-            blocks["qkv"]  = std::make_shared<Linear>(hidden_size, hidden_size * 3, true);
-            blocks["proj"] = std::make_shared<Linear>(hidden_size, hidden_size, true);
-        }
-
-        ggml_tensor* forward(GGMLRunnerContext* ctx,
-                             ggml_tensor* x,
-                             ggml_tensor* pe) {
-            auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
-            auto proj     = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
-
-            auto qkv     = qkv_proj->forward(ctx, x);
-            auto qkv_vec = split_qkv(ctx->ggml_ctx, qkv);
-
-            auto q = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[0], head_dim, num_heads, qkv_vec[0]->ne[1], qkv_vec[0]->ne[2]);
-            auto k = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[1], head_dim, num_heads, qkv_vec[1]->ne[1], qkv_vec[1]->ne[2]);
-            auto v = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[2], head_dim, num_heads, qkv_vec[2]->ne[1], qkv_vec[2]->ne[2]);
-
-            x = Rope::attention(ctx, q, k, v, pe, nullptr, 1.f, false);
-            x = proj->forward(ctx, x);
-            return x;
-        }
-    };
-
-    struct VisionBlock : public GGMLBlock {
-        VisionBlock(int64_t hidden_size,
-                    int64_t intermediate_size,
-                    int num_heads) {
-            blocks["norm1"] = std::make_shared<LayerNorm>(hidden_size, 1e-6f);
-            blocks["norm2"] = std::make_shared<LayerNorm>(hidden_size, 1e-6f);
-            blocks["attn"]  = std::make_shared<VisionAttention>(hidden_size, num_heads);
-            blocks["mlp"]   = std::make_shared<VisionMLP>(hidden_size, intermediate_size);
-        }
-
-        ggml_tensor* forward(GGMLRunnerContext* ctx,
-                             ggml_tensor* x,
-                             ggml_tensor* pe) {
-            auto norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
-            auto norm2 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm2"]);
-            auto attn  = std::dynamic_pointer_cast<VisionAttention>(blocks["attn"]);
-            auto mlp   = std::dynamic_pointer_cast<VisionMLP>(blocks["mlp"]);
-
-            auto residual = x;
-            x             = norm1->forward(ctx, x);
-            x             = attn->forward(ctx, x, pe);
-            x             = ggml_add_inplace(ctx->ggml_ctx, x, residual);
-
-            residual = x;
-            x        = norm2->forward(ctx, x);
-            x        = mlp->forward(ctx, x);
-            x        = ggml_add_inplace(ctx->ggml_ctx, x, residual);
-            return x;
-        }
-    };
-
-    struct VisionOutput {
-        ggml_tensor* hidden_states = nullptr;
-        std::vector<ggml_tensor*> deepstack_hidden_states;
-    };
-
-    struct VisionModel : public GGMLBlock {
-        int num_layers;
-        int spatial_merge_size;
-        int num_grid_per_side;
-        std::vector<int> deepstack_visual_indexes;
-
-        VisionModel(int num_layers,
-                    int64_t in_channels,
-                    int64_t hidden_size,
-                    int64_t out_hidden_size,
-                    int64_t intermediate_size,
-                    int num_heads,
-                    int spatial_merge_size,
-                    int patch_size,
-                    int temporal_patch_size,
-                    int num_position_embeddings,
-                    std::vector<int> deepstack_visual_indexes)
-            : num_layers(num_layers),
-              spatial_merge_size(spatial_merge_size),
-              num_grid_per_side(static_cast<int>(std::sqrt(num_position_embeddings))),
-              deepstack_visual_indexes(std::move(deepstack_visual_indexes)) {
-            blocks["patch_embed"] = std::make_shared<VisionPatchEmbed>(patch_size,
-                                                                       temporal_patch_size,
-                                                                       in_channels,
-                                                                       hidden_size);
-            blocks["pos_embed"]   = std::make_shared<Embedding>(num_position_embeddings, hidden_size);
-            for (int i = 0; i < num_layers; ++i) {
-                blocks["blocks." + std::to_string(i)] = std::make_shared<VisionBlock>(hidden_size,
-                                                                                      intermediate_size,
-                                                                                      num_heads);
-            }
-            blocks["merger"] = std::make_shared<VisionPatchMerger>(out_hidden_size,
-                                                                   hidden_size,
-                                                                   spatial_merge_size,
-                                                                   false);
-            for (int i = 0; i < static_cast<int>(this->deepstack_visual_indexes.size()); ++i) {
-                blocks["deepstack_merger_list." + std::to_string(i)] = std::make_shared<VisionPatchMerger>(out_hidden_size,
-                                                                                                           hidden_size,
-                                                                                                           spatial_merge_size,
-                                                                                                           true);
-            }
-        }
-
-        ggml_tensor* fast_pos_embed_interpolate(GGMLRunnerContext* ctx,
-                                                int grid_h,
-                                                int grid_w) {
-            auto pos_embed = std::dynamic_pointer_cast<Embedding>(blocks["pos_embed"]);
-            std::vector<int32_t> idx_list[4];
-            std::vector<float> weight_list[4];
-            idx_list[0].reserve(static_cast<size_t>(grid_h * grid_w));
-            idx_list[1].reserve(static_cast<size_t>(grid_h * grid_w));
-            idx_list[2].reserve(static_cast<size_t>(grid_h * grid_w));
-            idx_list[3].reserve(static_cast<size_t>(grid_h * grid_w));
-            weight_list[0].reserve(static_cast<size_t>(grid_h * grid_w));
-            weight_list[1].reserve(static_cast<size_t>(grid_h * grid_w));
-            weight_list[2].reserve(static_cast<size_t>(grid_h * grid_w));
-            weight_list[3].reserve(static_cast<size_t>(grid_h * grid_w));
-
-            double max_index = static_cast<double>(num_grid_per_side - 1);
-            for (int h = 0; h < grid_h; ++h) {
-                double h_pos = grid_h == 1 ? 0.0 : max_index * h / static_cast<double>(grid_h - 1);
-                int h_floor  = static_cast<int>(std::floor(h_pos));
-                int h_ceil   = std::min(h_floor + 1, num_grid_per_side - 1);
-                double dh    = h_pos - h_floor;
-                for (int w = 0; w < grid_w; ++w) {
-                    double w_pos = grid_w == 1 ? 0.0 : max_index * w / static_cast<double>(grid_w - 1);
-                    int w_floor  = static_cast<int>(std::floor(w_pos));
-                    int w_ceil   = std::min(w_floor + 1, num_grid_per_side - 1);
-                    double dw    = w_pos - w_floor;
-
-                    idx_list[0].push_back(h_floor * num_grid_per_side + w_floor);
-                    idx_list[1].push_back(h_floor * num_grid_per_side + w_ceil);
-                    idx_list[2].push_back(h_ceil * num_grid_per_side + w_floor);
-                    idx_list[3].push_back(h_ceil * num_grid_per_side + w_ceil);
-
-                    weight_list[0].push_back(static_cast<float>((1.0 - dh) * (1.0 - dw)));
-                    weight_list[1].push_back(static_cast<float>((1.0 - dh) * dw));
-                    weight_list[2].push_back(static_cast<float>(dh * (1.0 - dw)));
-                    weight_list[3].push_back(static_cast<float>(dh * dw));
-                }
-            }
-
-            ggml_tensor* patch_pos_embeds = nullptr;
-            for (int i = 0; i < 4; ++i) {
-                auto idx_tensor = ggml_new_tensor_1d(ctx->ggml_ctx, GGML_TYPE_I32, static_cast<int64_t>(idx_list[i].size()));
-                std::memcpy(idx_tensor->data, idx_list[i].data(), idx_list[i].size() * sizeof(int32_t));
-                auto embed         = pos_embed->forward(ctx, idx_tensor);
-                auto weight_tensor = ggml_new_tensor_2d(ctx->ggml_ctx, GGML_TYPE_F32, 1, static_cast<int64_t>(weight_list[i].size()));
-                std::memcpy(weight_tensor->data, weight_list[i].data(), weight_list[i].size() * sizeof(float));
-                embed            = ggml_mul(ctx->ggml_ctx, embed, weight_tensor);
-                patch_pos_embeds = patch_pos_embeds == nullptr ? embed : ggml_add(ctx->ggml_ctx, patch_pos_embeds, embed);
-            }
-
-            patch_pos_embeds = ggml_reshape_4d(ctx->ggml_ctx,
-                                               patch_pos_embeds,
-                                               patch_pos_embeds->ne[0],
-                                               spatial_merge_size,
-                                               grid_w / spatial_merge_size,
-                                               grid_h * spatial_merge_size);
-            patch_pos_embeds = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, patch_pos_embeds, 0, 1, 3, 2));
-            patch_pos_embeds = ggml_reshape_2d(ctx->ggml_ctx,
-                                               patch_pos_embeds,
-                                               patch_pos_embeds->ne[0],
-                                               ggml_nelements(patch_pos_embeds) / patch_pos_embeds->ne[0]);
-            return patch_pos_embeds;
-        }
-
-        VisionOutput forward(GGMLRunnerContext* ctx,
-                             ggml_tensor* pixel_values,
-                             ggml_tensor* pe,
-                             int grid_h,
-                             int grid_w) {
-            auto patch_embed = std::dynamic_pointer_cast<VisionPatchEmbed>(blocks["patch_embed"]);
-            auto merger      = std::dynamic_pointer_cast<VisionPatchMerger>(blocks["merger"]);
-
-            auto x          = patch_embed->forward(ctx, pixel_values);
-            auto pos_embeds = fast_pos_embed_interpolate(ctx, grid_h, grid_w);
-            x               = ggml_add(ctx->ggml_ctx, x, pos_embeds);
-            x               = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0], x->ne[1], 1);
-
-            VisionOutput out;
-            for (int i = 0; i < num_layers; ++i) {
-                auto block = std::dynamic_pointer_cast<VisionBlock>(blocks["blocks." + std::to_string(i)]);
-                x          = block->forward(ctx, x, pe);
-                for (int j = 0; j < static_cast<int>(deepstack_visual_indexes.size()); ++j) {
-                    if (deepstack_visual_indexes[j] == i) {
-                        auto deepstack_merger = std::dynamic_pointer_cast<VisionPatchMerger>(blocks["deepstack_merger_list." + std::to_string(j)]);
-                        out.deepstack_hidden_states.push_back(deepstack_merger->forward(ctx, x));
-                        break;
-                    }
-                }
-            }
-
-            out.hidden_states = merger->forward(ctx, x);
-            return out;
-        }
-    };
-
     struct HiDreamO1Model : public GGMLBlock {
         HiDreamO1Params params;
 
@@ -502,17 +242,6 @@ namespace HiDreamO1 {
         explicit HiDreamO1Model(HiDreamO1Params params)
             : params(std::move(params)) {
             blocks["language_model"] = std::make_shared<LLM::TextModel>(this->params.llm);
-            blocks["visual"]         = std::make_shared<VisionModel>(this->params.llm.vision.num_layers,
-                                                             this->params.llm.vision.in_channels,
-                                                             this->params.llm.vision.hidden_size,
-                                                             this->params.llm.vision.out_hidden_size,
-                                                             this->params.llm.vision.intermediate_size,
-                                                             this->params.llm.vision.num_heads,
-                                                             this->params.llm.vision.spatial_merge_size,
-                                                             this->params.llm.vision.patch_size,
-                                                             this->params.llm.vision.temporal_patch_size,
-                                                             this->params.num_position_embeddings,
-                                                             this->params.deepstack_visual_indexes);
             blocks["t_embedder1"]    = std::make_shared<TimestepEmbedder>(this->params.llm.hidden_size);
             blocks["x_embedder"]     = std::make_shared<BottleneckPatchEmbed>(this->params.patch_size * this->params.patch_size * 3,
                                                                           this->params.llm.hidden_size / 4,
@@ -525,10 +254,6 @@ namespace HiDreamO1 {
             return std::dynamic_pointer_cast<LLM::TextModel>(blocks["language_model"]);
         }
 
-        std::shared_ptr<VisionModel> vision_model() {
-            return std::dynamic_pointer_cast<VisionModel>(blocks["visual"]);
-        }
-
         std::shared_ptr<TimestepEmbedder> timestep_embedder() {
             return std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder1"]);
         }
@@ -542,43 +267,80 @@ namespace HiDreamO1 {
         }
     };
 
-    struct HiDreamO1Runner : public GGMLRunner {
+    struct HiDreamO1VisionRunner : public GGMLRunner {
         HiDreamO1Params params;
-        HiDreamO1Model model;
+        std::shared_ptr<LLM::VisionModel> model;
 
         std::vector<int> window_index_vec;
         std::vector<int> window_inverse_index_vec;
         std::vector<float> window_mask_vec;
         std::vector<float> pe_vec;
+        std::array<std::vector<int32_t>, 4> pos_embed_idx_data_;
+        std::array<std::vector<float>, 4> pos_embed_weight_data_;
+
+        HiDreamO1VisionRunner(ggml_backend_t backend,
+                              bool offload_params_to_cpu,
+                              const String2TensorStorage& tensor_storage_map = {},
+                              const std::string& prefix                      = "model.visual")
+            : GGMLRunner(backend, offload_params_to_cpu),
+              params(make_hidream_o1_params()),
+              model(std::make_shared<LLM::VisionModel>(false, params.llm.vision)) {
+            model->init(params_ctx, tensor_storage_map, prefix);
+        }
+
+        std::string get_desc() override {
+            return "hidream_o1_vision";
+        }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix = "model.visual") {
+            model->get_param_tensors(tensors, prefix);
+        }
+
+        ggml_tensor* encode_image(GGMLRunnerContext* runner_ctx, ggml_tensor* image) {
+            return LLM::LLMRunner::encode_image_common(this,
+                                                       compute_ctx,
+                                                       runner_ctx,
+                                                       image,
+                                                       params.llm.vision,
+                                                       model,
+                                                       window_index_vec,
+                                                       window_inverse_index_vec,
+                                                       window_mask_vec,
+                                                       pe_vec,
+                                                       pos_embed_idx_data_,
+                                                       pos_embed_weight_data_);
+        }
+
+        ggml_cgraph* build_graph(const sd::Tensor<float>& image_tensor) {
+            ggml_cgraph* gf    = new_graph_custom(HIDREAM_O1_GRAPH_SIZE);
+            ggml_tensor* image = make_input(image_tensor);
+            auto runner_ctx    = get_context();
+            auto image_embeds  = encode_image(&runner_ctx, image);
+            ggml_build_forward_expand(gf, image_embeds);
+            return gf;
+        }
+
+        sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& image) {
+            auto get_graph = [&]() {
+                return build_graph(image);
+            };
+            auto output = GGMLRunner::compute<float>(get_graph, n_threads, false);
+            return output.has_value() ? std::move(output.value()) : sd::Tensor<float>();
+        }
+    };
+
+    struct HiDreamO1Runner : public GGMLRunner {
+        HiDreamO1Params params;
+        HiDreamO1Model model;
+
         std::vector<float> attention_mask_vec;
 
         HiDreamO1Runner(ggml_backend_t backend,
                         bool offload_params_to_cpu,
                         const String2TensorStorage& tensor_storage_map = {},
                         const std::string& prefix                      = "model")
-            : GGMLRunner(backend, offload_params_to_cpu) {
+            : GGMLRunner(backend, offload_params_to_cpu),
-            params.llm.arch                       = LLM::LLMArch::QWEN3_VL;
+              params(make_hidream_o1_params()) {
-            params.llm.hidden_size                = 4096;
-            params.llm.intermediate_size          = 12288;
-            params.llm.num_layers                 = 36;
-            params.llm.num_heads                  = 32;
-            params.llm.num_kv_heads               = 8;
-            params.llm.head_dim                   = 128;
-            params.llm.qkv_bias                   = false;
-            params.llm.qk_norm                    = true;
-            params.llm.vocab_size                 = 151936;
-            params.llm.rms_norm_eps               = 1e-6f;
-            params.llm.vision.num_layers          = 27;
-            params.llm.vision.hidden_size         = 1152;
-            params.llm.vision.intermediate_size   = 4304;
-            params.llm.vision.num_heads           = 16;
-            params.llm.vision.out_hidden_size     = 4096;
-            params.llm.vision.patch_size          = 16;
-            params.llm.vision.spatial_merge_size  = 2;
-            params.llm.vision.temporal_patch_size = 2;
-            params.num_position_embeddings        = 2304;
-            params.deepstack_visual_indexes       = {8, 16, 24};
-
             model = HiDreamO1Model(params);
             model.init(params_ctx, tensor_storage_map, prefix);
         }
@@ -591,105 +353,13 @@ namespace HiDreamO1 {
             model.get_param_tensors(tensors, prefix);
         }
 
-        ggml_tensor* process_image(ggml_context* ctx, ggml_tensor* image) {
-            int64_t C  = image->ne[2];
-            int64_t H  = image->ne[1];
-            int64_t W  = image->ne[0];
-            int64_t mh = params.llm.vision.spatial_merge_size;
-            int64_t mw = params.llm.vision.spatial_merge_size;
-            int64_t pt = params.llm.vision.temporal_patch_size;
-            int64_t ph = params.llm.vision.patch_size;
-            int64_t pw = params.llm.vision.patch_size;
-
-            image = ggml_reshape_4d(ctx, image, pw, mw, (W / mw / pw), H * C);
-            image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 3, 1));
-            image = ggml_reshape_4d(ctx, image, pw * (W / mw / pw), H, C, mw);
-            image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 3, 1));
-            image = ggml_reshape_4d(ctx, image, pw, (W / mw / pw) * C * mw, ph, mh * (H / mh / ph));
-            image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 1, 3));
-            image = ggml_reshape_4d(ctx, image, pw * ph, (W / mw / pw), C, mw * mh * (H / mh / ph));
-            image = ggml_concat(ctx, image, image, 0);
-            image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 1, 3));
-            image = ggml_reshape_4d(ctx, image, pw * ph * pt * C, (W / mw / pw), mw * mh, (H / mh / ph));
-            image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 1, 3));
-            image = ggml_reshape_2d(ctx, image, pw * ph * pt * C, mw * mh * (W / mw / pw) * (H / mh / ph));
-            return image;
-        }
-
-        ggml_tensor* concat_seq(GGMLRunnerContext* ctx, ggml_tensor* a, ggml_tensor* b) {
-            if (a == nullptr) {
-                return b;
-            }
-            if (b == nullptr) {
-                return a;
-            }
-            return ggml_concat(ctx->ggml_ctx, a, b, 1);
-        }
-
-        ggml_tensor* scatter_visual_embeds(GGMLRunnerContext* ctx,
-                                           ggml_tensor* inputs_embeds,
-                                           const sd::Tensor<int32_t>& image_embed_ranges_tensor,
-                                           ggml_tensor* visual_embeds) {
-            if (visual_embeds == nullptr || image_embed_ranges_tensor.empty()) {
-                return inputs_embeds;
-            }
-
-            ggml_tensor* output = nullptr;
-            int prev_end        = 0;
-            int n_ranges        = static_cast<int>(image_embed_ranges_tensor.shape()[1]);
-            int visual_offset   = 0;
-            for (int i = 0; i < n_ranges; ++i) {
-                int start = image_embed_ranges_tensor.values()[i * 2];
-                int len   = image_embed_ranges_tensor.values()[i * 2 + 1];
-
-                if (start > prev_end) {
-                    output = concat_seq(ctx, output, ggml_ext_slice(ctx->ggml_ctx, inputs_embeds, 1, prev_end, start));
-                }
-
-                output   = concat_seq(ctx,
-                                      output,
-                                      ggml_ext_slice(ctx->ggml_ctx, visual_embeds, 1, visual_offset, visual_offset + len));
-                prev_end = start + len;
-                visual_offset += len;
-            }
-
-            if (prev_end < inputs_embeds->ne[1]) {
-                output = concat_seq(ctx, output, ggml_ext_slice(ctx->ggml_ctx, inputs_embeds, 1, prev_end, inputs_embeds->ne[1]));
-            }
-            return output == nullptr ? inputs_embeds : output;
-        }
-
-        VisionOutput encode_image(GGMLRunnerContext* runner_ctx, ggml_tensor* image) {
-            auto vision = model.vision_model();
-            GGML_ASSERT(image->ne[1] % (params.llm.vision.patch_size * params.llm.vision.spatial_merge_size) == 0);
-            GGML_ASSERT(image->ne[0] % (params.llm.vision.patch_size * params.llm.vision.spatial_merge_size) == 0);
-
-            int grid_h = static_cast<int>(image->ne[1]) / params.llm.vision.patch_size;
-            int grid_w = static_cast<int>(image->ne[0]) / params.llm.vision.patch_size;
-
-            auto pixel_values = process_image(compute_ctx, image);
-
-            int head_dim = static_cast<int>(params.llm.vision.hidden_size / params.llm.vision.num_heads);
-            std::vector<int> window_index_vec(static_cast<size_t>((grid_h / params.llm.vision.spatial_merge_size) * (grid_w / params.llm.vision.spatial_merge_size)));
-            for (int i = 0; i < static_cast<int>(window_index_vec.size()); ++i) {
-                window_index_vec[static_cast<size_t>(i)] = i;
-            }
-            pe_vec      = Rope::gen_qwen2vl_pe(grid_h, grid_w, params.llm.vision.spatial_merge_size, window_index_vec, 10000, {head_dim / 2, head_dim / 2});
-            int pos_len = static_cast<int>(pe_vec.size() / head_dim / 2);
-            auto pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, head_dim / 2, pos_len);
-            set_backend_tensor_data(pe, pe_vec.data());
-
-            return vision->forward(runner_ctx, pixel_values, pe, grid_h, grid_w);
-        }
-
         ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
                                  const sd::Tensor<float>& timestep_tensor,
                                  const sd::Tensor<int32_t>& input_ids_tensor,
                                  const sd::Tensor<int32_t>& input_pos_tensor,
                                  const sd::Tensor<int32_t>& token_types_tensor,
-                                 const sd::Tensor<int32_t>& image_embed_ranges_tensor,
                                  const sd::Tensor<int32_t>& vinput_mask_tensor,
-                                 const std::vector<sd::Tensor<float>>& vlm_images,
+                                 const std::vector<std::pair<int, sd::Tensor<float>>>& image_embeds_tensor,
                                  const std::vector<sd::Tensor<float>>& ref_images) {
             ggml_cgraph* gf        = new_graph_custom(HIDREAM_O1_GRAPH_SIZE);
             ggml_tensor* x         = make_input(x_tensor);
@@ -702,11 +372,6 @@ namespace HiDreamO1 {
             auto x_embedder   = model.patch_embedder();
             auto final_layer2 = model.final_layer();
 
-            std::vector<ggml_tensor*> vlm_image_tensors;
-            for (const auto& image : vlm_images) {
-                vlm_image_tensors.push_back(make_input(image));
-            }
-
             std::vector<ggml_tensor*> ref_image_tensors;
             for (const auto& image : ref_images) {
                 ref_image_tensors.push_back(make_input(image));
@@ -726,14 +391,13 @@ namespace HiDreamO1 {
             set_backend_tensor_data(attention_mask, attention_mask_vec.data());
 
             auto runner_ctx = get_context();
-            ggml_tensor* visual_embeds = nullptr;
-            for (size_t i = 0; i < vlm_image_tensors.size(); ++i) {
-                auto image_output = encode_image(&runner_ctx, vlm_image_tensors[i]);
-                visual_embeds     = visual_embeds == nullptr ? image_output.hidden_states : ggml_concat(compute_ctx, visual_embeds, image_output.hidden_states, 1);
-            }
-
             auto txt        = text_model->embed(&runner_ctx, input_ids);
-            txt      = scatter_visual_embeds(&runner_ctx, txt, image_embed_ranges_tensor, visual_embeds);
+            std::vector<std::pair<int, ggml_tensor*>> image_embeds;
+            image_embeds.reserve(image_embeds_tensor.size());
+            for (const auto& image_embed : image_embeds_tensor) {
+                image_embeds.emplace_back(image_embed.first, make_input(image_embed.second));
+            }
+            txt = LLM::splice_image_embeds(&runner_ctx, txt, image_embeds);
 
             auto t_emb          = t_embedder1->forward(&runner_ctx, timestep);
             int64_t txt_seq_len = input_ids->ne[0];
@@ -765,15 +429,7 @@ namespace HiDreamO1 {
                 }
                 x_pred_start = first_vinput;
             }
-            auto x_pred = ggml_view_3d(compute_ctx,
+            auto x_pred = ggml_ext_slice(compute_ctx, x_pred_all, 1, x_pred_start, x_pred_start + target_tokens);
-                                       x_pred_all,
-                                       x_pred_all->ne[0],
-                                       target_tokens,
-                                       x_pred_all->ne[2],
-                                       x_pred_all->nb[1],
-                                       x_pred_all->nb[2],
-                                       x_pred_start * x_pred_all->nb[1]);
-            x_pred      = ggml_cont(compute_ctx, x_pred);
             x_pred      = DiT::unpatchify_and_crop(compute_ctx, x_pred, x->ne[1], x->ne[0], PATCH_SIZE, PATCH_SIZE);
 
             float sigma = 1.0f - timestep_tensor.values()[0];
@@ -790,12 +446,11 @@ namespace HiDreamO1 {
                                   const sd::Tensor<int32_t>& input_ids,
                                   const sd::Tensor<int32_t>& input_pos,
                                   const sd::Tensor<int32_t>& token_types,
-                                  const sd::Tensor<int32_t>& image_embed_ranges,
                                   const sd::Tensor<int32_t>& vinput_mask,
-                                  const std::vector<sd::Tensor<float>>& vlm_images,
+                                  const std::vector<std::pair<int, sd::Tensor<float>>>& image_embeds,
                                   const std::vector<sd::Tensor<float>>& ref_images) {
             auto get_graph = [&]() {
-                return build_graph(x, timestep, input_ids, input_pos, token_types, image_embed_ranges, vinput_mask, vlm_images, ref_images);
+                return build_graph(x, timestep, input_ids, input_pos, token_types, vinput_mask, image_embeds, ref_images);
             };
             return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
         }
@@ -803,19 +458,43 @@ namespace HiDreamO1 {
 
     struct HiDreamO1Conditioner : public Conditioner {
         Qwen2Tokenizer tokenizer;
+        std::shared_ptr<HiDreamO1VisionRunner> vision_runner;
+
+        HiDreamO1Conditioner(ggml_backend_t backend,
+                             bool offload_params_to_cpu,
+                             const String2TensorStorage& tensor_storage_map = {})
+            : vision_runner(std::make_shared<HiDreamO1VisionRunner>(backend, offload_params_to_cpu, tensor_storage_map)) {}
 
         void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
-            SD_UNUSED(tensors);
+            vision_runner->get_param_tensors(tensors);
         }
 
-        void alloc_params_buffer() override {}
+        void alloc_params_buffer() override {
-        void free_params_buffer() override {}
+            vision_runner->alloc_params_buffer();
-        size_t get_params_buffer_size() override { return 0; }
+        }
-        void set_flash_attention_enabled(bool enabled) override { SD_UNUSED(enabled); }
+
+        void free_params_buffer() override {
+            vision_runner->free_params_buffer();
+        }
+
+        size_t get_params_buffer_size() override {
+            return vision_runner->get_params_buffer_size();
+        }
+
+        void set_max_graph_vram_bytes(size_t max_graph_vram_bytes) override {
+            vision_runner->set_max_graph_vram_bytes(max_graph_vram_bytes);
+        }
+
+        void set_flash_attention_enabled(bool enabled) override {
+            vision_runner->set_flash_attention_enabled(enabled);
+        }
+
+        void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
+            vision_runner->set_weight_adapter(adapter);
+        }
 
         SDCondition get_learned_condition(int n_threads,
                                           const ConditionerParams& conditioner_params) override {
-            SD_UNUSED(n_threads);
             SDCondition result;
 
             int width                = conditioner_params.width;
@@ -827,12 +506,11 @@ namespace HiDreamO1 {
                 ref_images = *conditioner_params.ref_images;
             }
 
-            std::vector<sd::Tensor<float>> vlm_images;
+            std::vector<std::pair<int, sd::Tensor<float>>> vlm_images;
             std::vector<std::array<int32_t, 3>> image_grids;
             std::vector<int32_t> skip_vision_start;
 
             std::string prompt = "<|im_start|>user\n";
-            std::vector<int32_t> image_ranges;
 
             if (ref_images.empty()) {
                 prompt += conditioner_params.text;
@@ -849,12 +527,9 @@ namespace HiDreamO1 {
                 std::vector<int32_t> token_types(input_ids_pad.size(), 0);
                 int txt_seq_len = static_cast<int>(input_ids.size());
                 int bgn         = txt_seq_len - TIMESTEP_TOKEN_NUM;
-                for (int i = bgn; i < bgn + target_image_len + TIMESTEP_TOKEN_NUM; ++i) {
+                for (int i = bgn; i < static_cast<int>(token_types.size()); ++i) {
                     token_types[i] = 1;
                 }
-                for (int i = txt_seq_len - TIMESTEP_TOKEN_NUM; i < txt_seq_len; ++i) {
-                    token_types[i] = 3;
-                }
 
                 auto position_ids = build_position_ids(input_ids_pad, image_grids, skip_vision_start);
 
@@ -871,7 +546,6 @@ namespace HiDreamO1 {
                 result.c_position_ids = sd::Tensor<int32_t>(position_shape, position_ids);
                 result.c_token_types  = sd::Tensor<int32_t>(token_type_shape, std::move(token_types));
                 result.c_vinput_mask  = sd::Tensor<int32_t>(vinput_mask_shape, std::move(vinput_mask));
-                result.c_image_embed_ranges = sd::Tensor<int32_t>();
                 return result;
             }
 
@@ -899,21 +573,27 @@ namespace HiDreamO1 {
             }
 
             for (const auto& ref_image : ref_images) {
-                auto patch_img = resize_to_area(ref_image, max_size);
+                auto resized_ref = resize_to_area(ref_image, max_size);
-                patch_img      = sd::ops::clamp(patch_img, 0.0f, 1.0f);
+                resized_ref      = sd::ops::clamp(resized_ref, 0.0f, 1.0f);
-                patch_img      = patch_img * 2.0f - 1.0f;
-                result.c_ref_images.push_back(std::move(patch_img));
 
-                auto dims            = calculate_dimensions(cond_img_size, static_cast<double>(ref_image.shape()[0]) / static_cast<double>(ref_image.shape()[1]));
+                // VLM image: Qwen3-VL expects mean=[0.5]/std=[0.5] (i.e. range [-1,1]),
-                auto vlm_image       = clip_preprocess(ref_image, dims.first, dims.second);
+                // not CLIP normalization. Resize the already-resized ref directly to
+                // (cond_w, cond_h) to match the Python pipeline's pil_r.resize().
+                auto dims                   = calculate_dimensions(cond_img_size,
+                                                                   static_cast<double>(resized_ref.shape()[0]) / static_cast<double>(resized_ref.shape()[1]));
+                sd::Tensor<float> vlm_image = sd::ops::interpolate(
+                    resized_ref,
+                    {dims.first, dims.second, resized_ref.shape()[2], resized_ref.shape()[3]});
+                vlm_image            = vlm_image * 2.0f - 1.0f;
                 int64_t image_tokens = static_cast<int64_t>(dims.first / PATCH_SIZE) * static_cast<int64_t>(dims.second / PATCH_SIZE);
+
+                auto patch_img = resized_ref * 2.0f - 1.0f;
+                result.c_ref_images.push_back(std::move(patch_img));
                 int64_t prompt_start = static_cast<int64_t>(tokenizer.encode(prompt + "<|vision_start|>", nullptr).size());
                 prompt += "<|vision_start|>";
                 prompt += repeat_special_token("<|image_pad|>", image_tokens);
                 prompt += "<|vision_end|>";
-                image_ranges.push_back(static_cast<int32_t>(prompt_start));
+                vlm_images.emplace_back(static_cast<int>(prompt_start), std::move(vlm_image));
-                image_ranges.push_back(static_cast<int32_t>(image_tokens));
-                result.c_vlm_images.push_back(std::move(vlm_image));
                 image_grids.push_back({1, dims.second / PATCH_SIZE, dims.first / PATCH_SIZE});
                 skip_vision_start.push_back(0);
             }
@@ -928,10 +608,8 @@ namespace HiDreamO1 {
             image_grids.push_back({1, static_cast<int32_t>(height / PATCH_SIZE), static_cast<int32_t>(width / PATCH_SIZE)});
             skip_vision_start.push_back(1);
 
-            int64_t total_ref_len = 0;
             for (const auto& ref_image : result.c_ref_images) {
                 int64_t ref_len = static_cast<int64_t>(ref_image.shape()[0] / PATCH_SIZE) * static_cast<int64_t>(ref_image.shape()[1] / PATCH_SIZE);
-                total_ref_len += ref_len;
                 input_ids_pad.push_back(VISION_START_TOKEN_ID);
                 input_ids_pad.insert(input_ids_pad.end(), ref_len - 1, IMAGE_TOKEN_ID);
                 image_grids.push_back({1, static_cast<int32_t>(ref_image.shape()[1] / PATCH_SIZE), static_cast<int32_t>(ref_image.shape()[0] / PATCH_SIZE)});
@@ -941,21 +619,13 @@ namespace HiDreamO1 {
             std::vector<int32_t> token_types(input_ids_pad.size(), 0);
             int txt_seq_len = static_cast<int>(input_ids.size());
             int bgn         = txt_seq_len - TIMESTEP_TOKEN_NUM;
-            int end         = bgn + static_cast<int>(target_image_len) + TIMESTEP_TOKEN_NUM;
+            for (int i = bgn; i < static_cast<int>(token_types.size()); ++i) {
-            for (int i = bgn; i < end; ++i) {
                 token_types[i] = 1;
             }
-            for (int i = end; i < end + total_ref_len; ++i) {
-                token_types[i] = 2;
-            }
-            for (int i = txt_seq_len - TIMESTEP_TOKEN_NUM; i < txt_seq_len; ++i) {
-                token_types[i] = 3;
-            }
 
             std::vector<int64_t> input_shape{static_cast<int64_t>(input_ids.size())};
             std::vector<int64_t> position_shape{static_cast<int64_t>(input_ids_pad.size() * 4)};
             std::vector<int64_t> token_type_shape{static_cast<int64_t>(token_types.size())};
-            std::vector<int64_t> image_range_shape{2, static_cast<int64_t>(image_ranges.size() / 2)};
             std::vector<int32_t> vinput_mask(token_types.size(), 0);
             for (int i = txt_seq_len; i < static_cast<int>(vinput_mask.size()); ++i) {
                 vinput_mask[static_cast<size_t>(i)] = 1;
@@ -965,8 +635,16 @@ namespace HiDreamO1 {
             result.c_input_ids    = sd::Tensor<int32_t>(input_shape, std::move(input_ids));
             result.c_position_ids = sd::Tensor<int32_t>(position_shape, build_position_ids(input_ids_pad, image_grids, skip_vision_start));
             result.c_token_types  = sd::Tensor<int32_t>(token_type_shape, std::move(token_types));
-            result.c_image_embed_ranges = sd::Tensor<int32_t>(image_range_shape, std::move(image_ranges));
             result.c_vinput_mask  = sd::Tensor<int32_t>(vinput_mask_shape, std::move(vinput_mask));
+            result.c_image_embeds.reserve(vlm_images.size());
+            for (const auto& vlm_image : vlm_images) {
+                auto image_embed = vision_runner->compute(n_threads, vlm_image.second);
+                if (image_embed.empty()) {
+                    LOG_ERROR("hidream_o1 conditioner: encode VLM image failed");
+                    return SDCondition();
+                }
+                result.c_image_embeds.emplace_back(vlm_image.first, std::move(image_embed));
+            }
             return result;
         }
     };

src/llm.hpp

@@ -2,7 +2,10 @@
 #define __LLM_HPP__
 
 #include <algorithm>
+#include <array>
+#include <cmath>
 #include <fstream>
+#include <functional>
 #include <iostream>
 #include <map>
 #include <memory>
@@ -41,7 +44,13 @@ namespace LLM {
         "ministral3.3b",
     };
 
+    enum class LLMVisionArch {
+        QWEN2_5_VL,
+        QWEN3_VL,
+    };
+
     struct LLMVisionParams {
+        LLMVisionArch arch                  = LLMVisionArch::QWEN2_5_VL;
         int num_layers                      = 32;
         int64_t hidden_size                 = 1280;
         int64_t intermediate_size           = 3420;
@@ -52,6 +61,7 @@ namespace LLM {
         int patch_size                      = 14;
         int spatial_merge_size              = 2;
         int window_size                     = 112;
+        int num_position_embeddings         = 0;
         std::set<int> fullatt_block_indexes = {7, 15, 23, 31};
     };
 
@@ -92,6 +102,84 @@ namespace LLM {
         }
     };
 
+    static ggml_tensor* splice_image_embeds(GGMLRunnerContext* ctx,
+                                            ggml_tensor* x,
+                                            const std::vector<std::pair<int, ggml_tensor*>>& image_embeds) {
+        if (image_embeds.empty()) {
+            return x;
+        }
+
+        GGML_ASSERT(x->ne[2] == 1);  // N == 1
+
+        auto raw_x               = ggml_cast(ctx->ggml_ctx, x, image_embeds[0].second->type);
+        int64_t txt_token_start  = 0;
+        int64_t txt_token_end    = 0;
+        ggml_tensor* input_embed = nullptr;
+
+        for (int i = 0; i < image_embeds.size(); i++) {
+            if (i == 0) {
+                txt_token_start = 0;
+            } else {
+                txt_token_start = image_embeds[i - 1].first + image_embeds[i - 1].second->ne[1];
+            }
+            txt_token_end = image_embeds[i].first;
+
+            auto txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end);
+            if (input_embed == nullptr) {
+                input_embed = txt_embed;
+            } else {
+                input_embed = ggml_concat(ctx->ggml_ctx, input_embed, txt_embed, 1);
+            }
+
+            input_embed = ggml_concat(ctx->ggml_ctx, input_embed, image_embeds[i].second, 1);
+        }
+
+        txt_token_start = image_embeds[image_embeds.size() - 1].first + image_embeds[image_embeds.size() - 1].second->ne[1];
+        txt_token_end   = raw_x->ne[1];
+
+        auto final_txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end);
+        input_embed          = ggml_concat(ctx->ggml_ctx, input_embed, final_txt_embed, 1);
+        GGML_ASSERT(raw_x->ne[1] == input_embed->ne[1]);
+        return input_embed;
+    }
+
+    struct VisionMLP : public GGMLBlock {
+    protected:
+        LLMVisionArch arch_;
+
+    public:
+        VisionMLP(LLMVisionArch arch, int64_t hidden_size, int64_t intermediate_size)
+            : arch_(arch) {
+            if (arch_ == LLMVisionArch::QWEN3_VL) {
+                blocks["linear_fc1"] = std::make_shared<Linear>(hidden_size, intermediate_size, true);
+                blocks["linear_fc2"] = std::make_shared<Linear>(intermediate_size, hidden_size, true);
+            } else {
+                blocks["gate_proj"] = std::make_shared<Linear>(hidden_size, intermediate_size, true);
+                blocks["up_proj"]   = std::make_shared<Linear>(hidden_size, intermediate_size, true);
+                blocks["down_proj"] = std::make_shared<Linear>(intermediate_size, hidden_size, true);
+            }
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            if (arch_ == LLMVisionArch::QWEN3_VL) {
+                auto linear_fc1 = std::dynamic_pointer_cast<Linear>(blocks["linear_fc1"]);
+                auto linear_fc2 = std::dynamic_pointer_cast<Linear>(blocks["linear_fc2"]);
+                x               = linear_fc1->forward(ctx, x);
+                x               = ggml_ext_gelu(ctx->ggml_ctx, x);
+                x               = linear_fc2->forward(ctx, x);
+            } else {
+                auto gate_proj = std::dynamic_pointer_cast<Linear>(blocks["gate_proj"]);
+                auto up_proj   = std::dynamic_pointer_cast<Linear>(blocks["up_proj"]);
+                auto down_proj = std::dynamic_pointer_cast<Linear>(blocks["down_proj"]);
+                auto h         = gate_proj->forward(ctx, x);
+                h              = ggml_silu_inplace(ctx->ggml_ctx, h);
+                h              = ggml_mul_inplace(ctx->ggml_ctx, h, up_proj->forward(ctx, x));
+                x              = down_proj->forward(ctx, h);
+            }
+            return x;
+        }
+    };
+
     struct VisionPatchEmbed : public GGMLBlock {
     protected:
         bool llama_cpp_style;
@@ -102,6 +190,7 @@ namespace LLM {
 
     public:
         VisionPatchEmbed(bool llama_cpp_style,
+                         LLMVisionArch arch,
                          int patch_size          = 14,
                          int temporal_patch_size = 2,
                          int64_t in_channels     = 3,
@@ -111,36 +200,35 @@ namespace LLM {
               temporal_patch_size(temporal_patch_size),
               in_channels(in_channels),
               embed_dim(embed_dim) {
+            bool bias = arch == LLMVisionArch::QWEN3_VL;
             if (llama_cpp_style) {
                 blocks["proj.0"] = std::shared_ptr<GGMLBlock>(new Conv2d(in_channels,
                                                                          embed_dim,
                                                                          {patch_size, patch_size},
-                                                                         {patch_size, patch_size},  // stride
+                                                                         {patch_size, patch_size},
-                                                                         {0, 0},                    // padding
+                                                                         {0, 0},
-                                                                         {1, 1},                    // dilation
+                                                                         {1, 1},
-                                                                         false));
+                                                                         bias));
                 blocks["proj.1"] = std::shared_ptr<GGMLBlock>(new Conv2d(in_channels,
                                                                          embed_dim,
                                                                          {patch_size, patch_size},
-                                                                         {patch_size, patch_size},  // stride
+                                                                         {patch_size, patch_size},
-                                                                         {0, 0},                    // padding
+                                                                         {0, 0},
-                                                                         {1, 1},                    // dilation
+                                                                         {1, 1},
-                                                                         false));
+                                                                         bias));
             } else {
                 std::tuple<int, int, int> kernel_size = {(int)temporal_patch_size, (int)patch_size, (int)patch_size};
                 blocks["proj"]                        = std::shared_ptr<GGMLBlock>(new Conv3d(in_channels,
                                                                                               embed_dim,
                                                                                               kernel_size,
-                                                                                              kernel_size,  // stride
+                                                                                              kernel_size,
-                                                                                              {0, 0, 0},    // padding
+                                                                                              {0, 0, 0},
-                                                                                              {1, 1, 1},    // dilation
+                                                                                              {1, 1, 1},
-                                                                                              false));
+                                                                                              bias));
             }
         }
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
-            // x: [N*grid_t*grid_h*grid_w, in_channels, temporal_patch_size*patch_size*patch_size]
-            // return: [N*grid_t*grid_h*grid_w, embed_dim]
             x = ggml_reshape_4d(ctx->ggml_ctx,
                                 x,
                                 patch_size,
@@ -172,22 +260,43 @@ namespace LLM {
         }
     };
 
-    struct PatchMerger : public GGMLBlock {
+    struct VisionPatchMerger : public GGMLBlock {
     protected:
+        LLMVisionArch arch_;
         int64_t hidden_size;
 
     public:
-        PatchMerger(int64_t dim,
+        VisionPatchMerger(LLMVisionArch arch,
+                          int64_t dim,
                           int64_t context_dim,
-                    int64_t spatial_merge_size) {
+                          int64_t spatial_merge_size)
-            hidden_size     = context_dim * spatial_merge_size * spatial_merge_size;
+            : arch_(arch),
-            blocks["ln_q"]  = std::shared_ptr<GGMLBlock>(new RMSNorm(context_dim, 1e-6f));
+              hidden_size(context_dim * spatial_merge_size * spatial_merge_size) {
-            blocks["mlp.0"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, hidden_size));
+            if (arch_ == LLMVisionArch::QWEN3_VL) {
-            // mlp.1 is nn.GELU()
+                blocks["norm"]       = std::make_shared<LayerNorm>(context_dim, 1e-6f);
-            blocks["mlp.2"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, dim));
+                blocks["linear_fc1"] = std::make_shared<Linear>(hidden_size, hidden_size, true);
+                blocks["linear_fc2"] = std::make_shared<Linear>(hidden_size, dim, true);
+            } else {
+                blocks["ln_q"]  = std::make_shared<RMSNorm>(context_dim, 1e-6f);
+                blocks["mlp.0"] = std::make_shared<Linear>(hidden_size, hidden_size);
+                blocks["mlp.2"] = std::make_shared<Linear>(hidden_size, dim);
+            }
         }
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            if (arch_ == LLMVisionArch::QWEN3_VL) {
+                auto norm       = std::dynamic_pointer_cast<LayerNorm>(blocks["norm"]);
+                auto linear_fc1 = std::dynamic_pointer_cast<Linear>(blocks["linear_fc1"]);
+                auto linear_fc2 = std::dynamic_pointer_cast<Linear>(blocks["linear_fc2"]);
+
+                x = norm->forward(ctx, x);
+                x = ggml_reshape_2d(ctx->ggml_ctx, x, hidden_size, ggml_nelements(x) / hidden_size);
+                x = linear_fc1->forward(ctx, x);
+                x = ggml_gelu_erf(ctx->ggml_ctx, x);
+                x = linear_fc2->forward(ctx, x);
+                return x;
+            }
+
             auto ln_q  = std::dynamic_pointer_cast<RMSNorm>(blocks["ln_q"]);
             auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
             auto mlp_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
@@ -262,17 +371,36 @@ namespace LLM {
     };
 
     struct VisionBlock : public GGMLBlock {
+    protected:
+        LLMVisionArch arch_;
+
+        ggml_tensor* forward_norm(GGMLRunnerContext* ctx, const std::string& name, ggml_tensor* x) {
+            if (arch_ == LLMVisionArch::QWEN3_VL) {
+                auto norm = std::dynamic_pointer_cast<LayerNorm>(blocks[name]);
+                return norm->forward(ctx, x);
+            }
+            auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks[name]);
+            return norm->forward(ctx, x);
+        }
+
     public:
         VisionBlock(bool llama_cpp_style,
+                    LLMVisionArch arch,
                     int64_t hidden_size,
                     int64_t intermediate_size,
                     int num_heads,
-                    float eps = 1e-6f) {
+                    float eps = 1e-6f)
+            : arch_(arch) {
             blocks["attn"] = std::shared_ptr<GGMLBlock>(new VisionAttention(llama_cpp_style, hidden_size, num_heads));
-            blocks["mlp"]   = std::shared_ptr<GGMLBlock>(new MLP(hidden_size, intermediate_size, true));
+            blocks["mlp"]  = std::shared_ptr<GGMLBlock>(new VisionMLP(arch_, hidden_size, intermediate_size));
+            if (arch_ == LLMVisionArch::QWEN3_VL) {
+                blocks["norm1"] = std::shared_ptr<GGMLBlock>(new LayerNorm(hidden_size, eps));
+                blocks["norm2"] = std::shared_ptr<GGMLBlock>(new LayerNorm(hidden_size, eps));
+            } else {
                 blocks["norm1"] = std::shared_ptr<GGMLBlock>(new RMSNorm(hidden_size, eps));
                 blocks["norm2"] = std::shared_ptr<GGMLBlock>(new RMSNorm(hidden_size, eps));
             }
+        }
 
         ggml_tensor* forward(GGMLRunnerContext* ctx,
                              ggml_tensor* x,
@@ -280,17 +408,15 @@ namespace LLM {
                              ggml_tensor* mask = nullptr) {
             // x: [N, n_token, hidden_size]
             auto attn = std::dynamic_pointer_cast<VisionAttention>(blocks["attn"]);
-            auto mlp   = std::dynamic_pointer_cast<MLP>(blocks["mlp"]);
+            auto mlp  = std::dynamic_pointer_cast<VisionMLP>(blocks["mlp"]);
-            auto norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm1"]);
-            auto norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm2"]);
 
             auto residual = x;
-            x             = norm1->forward(ctx, x);
+            x             = forward_norm(ctx, "norm1", x);
             x             = attn->forward(ctx, x, pe, mask);
             x             = ggml_add_inplace(ctx->ggml_ctx, x, residual);
 
             residual = x;
-            x        = norm2->forward(ctx, x);
+            x        = forward_norm(ctx, "norm2", x);
             x        = mlp->forward(ctx, x);
             x        = ggml_add_inplace(ctx->ggml_ctx, x, residual);
 
@@ -300,38 +426,58 @@ namespace LLM {
 
     struct VisionModel : public GGMLBlock {
     protected:
+        LLMVisionArch arch_;
         int num_layers;
         int spatial_merge_size;
+        int num_grid_per_side;
         std::set<int> fullatt_block_indexes;
 
     public:
         VisionModel(bool llama_cpp_style,
-                    int num_layers,
+                    const LLMVisionParams& vision_params,
-                    int64_t in_channels,
-                    int64_t hidden_size,
-                    int64_t out_hidden_size,
-                    int64_t intermediate_size,
-                    int num_heads,
-                    int spatial_merge_size,
-                    int patch_size,
-                    int temporal_patch_size,
-                    int window_size,
-                    std::set<int> fullatt_block_indexes = {7, 15, 23, 31},
                     float eps = 1e-6f)
-            : num_layers(num_layers), fullatt_block_indexes(std::move(fullatt_block_indexes)), spatial_merge_size(spatial_merge_size) {
+            : arch_(vision_params.arch),
+              num_layers(vision_params.num_layers),
+              spatial_merge_size(vision_params.spatial_merge_size),
+              num_grid_per_side(vision_params.num_position_embeddings > 0 ? static_cast<int>(std::sqrt(vision_params.num_position_embeddings)) : 0),
+              fullatt_block_indexes(vision_params.fullatt_block_indexes) {
             blocks["patch_embed"] = std::shared_ptr<GGMLBlock>(new VisionPatchEmbed(llama_cpp_style,
-                                                                                    patch_size,
+                                                                                    arch_,
-                                                                                    temporal_patch_size,
+                                                                                    vision_params.patch_size,
-                                                                                    in_channels,
+                                                                                    vision_params.temporal_patch_size,
-                                                                                    hidden_size));
+                                                                                    vision_params.in_channels,
+                                                                                    vision_params.hidden_size));
+            if (vision_params.num_position_embeddings > 0) {
+                blocks["pos_embed"] = std::make_shared<Embedding>(vision_params.num_position_embeddings, vision_params.hidden_size);
+            }
             for (int i = 0; i < num_layers; i++) {
                 blocks["blocks." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new VisionBlock(llama_cpp_style,
-                                                                                                   hidden_size,
+                                                                                                   arch_,
-                                                                                                   intermediate_size,
+                                                                                                   vision_params.hidden_size,
-                                                                                                   num_heads,
+                                                                                                   vision_params.intermediate_size,
+                                                                                                   vision_params.num_heads,
                                                                                                    eps));
             }
-            blocks["merger"] = std::shared_ptr<GGMLBlock>(new PatchMerger(out_hidden_size, hidden_size, spatial_merge_size));
+            blocks["merger"] = std::shared_ptr<GGMLBlock>(new VisionPatchMerger(arch_,
+                                                                                vision_params.out_hidden_size,
+                                                                                vision_params.hidden_size,
+                                                                                spatial_merge_size));
+        }
+
+        std::shared_ptr<Embedding> pos_embedder() {
+            auto it = blocks.find("pos_embed");
+            if (it == blocks.end()) {
+                return nullptr;
+            }
+            return std::dynamic_pointer_cast<Embedding>(it->second);
+        }
+
+        int get_num_grid_per_side() const {
+            return num_grid_per_side;
+        }
+
+        int get_spatial_merge_size() const {
+            return spatial_merge_size;
         }
 
         ggml_tensor* forward(GGMLRunnerContext* ctx,
@@ -339,20 +485,26 @@ namespace LLM {
                              ggml_tensor* pe,
                              ggml_tensor* window_index,
                              ggml_tensor* window_inverse_index,
-                             ggml_tensor* window_mask) {
+                             ggml_tensor* window_mask,
+                             ggml_tensor* pos_embeds = nullptr) {
             // pixel_values: [grid_t*(H/mh/ph)*(W/mw/pw)*mh*mw, C*pt*ph*pw]
             // window_index: [grid_t*(H/mh/ph)*(W/mw/pw)]
             // window_inverse_index: [grid_t*(H/mh/ph)*(W/mw/pw)]
             // window_mask: [grid_h*grid_w, grid_h*grid_w]
             auto patch_embed = std::dynamic_pointer_cast<VisionPatchEmbed>(blocks["patch_embed"]);
-            auto merger      = std::dynamic_pointer_cast<PatchMerger>(blocks["merger"]);
+            auto merger      = std::dynamic_pointer_cast<VisionPatchMerger>(blocks["merger"]);
 
             auto x = patch_embed->forward(ctx, pixel_values);
             sd::ggml_graph_cut::mark_graph_cut(x, "llm.vision.prelude", "x");
+            if (pos_embeds != nullptr) {
+                x = ggml_add(ctx->ggml_ctx, x, pos_embeds);
+            }
 
+            if (window_index != nullptr) {
                 x = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0] * spatial_merge_size * spatial_merge_size, x->ne[1] / spatial_merge_size / spatial_merge_size, x->ne[2], x->ne[3]);
                 x = ggml_get_rows(ctx->ggml_ctx, x, window_index);
                 x = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0] / spatial_merge_size / spatial_merge_size, x->ne[1] * spatial_merge_size * spatial_merge_size, x->ne[2], x->ne[3]);
+            }
 
             for (int i = 0; i < num_layers; i++) {
                 auto block = std::dynamic_pointer_cast<VisionBlock>(blocks["blocks." + std::to_string(i)]);
@@ -362,13 +514,17 @@ namespace LLM {
                     mask = nullptr;
                 }
                 x = block->forward(ctx, x, pe, mask);
+                if (i == 0) {
+                }
                 sd::ggml_graph_cut::mark_graph_cut(x, "llm.vision.blocks." + std::to_string(i), "x");
             }
 
             x = merger->forward(ctx, x);
             sd::ggml_graph_cut::mark_graph_cut(x, "llm.vision.final", "x");
 
+            if (window_inverse_index != nullptr) {
                 x = ggml_get_rows(ctx->ggml_ctx, x, window_inverse_index);
+            }
 
             return x;
         }
@@ -510,47 +666,6 @@ namespace LLM {
             return x;
         }
 
-        ggml_tensor* splice_image_embeds(GGMLRunnerContext* ctx,
-                                         ggml_tensor* x,
-                                         std::vector<std::pair<int, ggml_tensor*>> image_embeds) {
-            if (image_embeds.empty()) {
-                return x;
-            }
-
-            GGML_ASSERT(x->ne[2] == 1);  // N == 1
-
-            auto raw_x               = ggml_cast(ctx->ggml_ctx, x, image_embeds[0].second->type);
-            int64_t txt_token_start  = 0;
-            int64_t txt_token_end    = 0;
-            ggml_tensor* input_embed = nullptr;
-
-            for (int i = 0; i < image_embeds.size(); i++) {
-                if (i == 0) {
-                    txt_token_start = 0;
-                } else {
-                    txt_token_start = image_embeds[i - 1].first + image_embeds[i - 1].second->ne[1];
-                }
-                txt_token_end = image_embeds[i].first;
-
-                auto txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end);
-                if (input_embed == nullptr) {
-                    input_embed = txt_embed;
-                } else {
-                    input_embed = ggml_concat(ctx->ggml_ctx, input_embed, txt_embed, 1);
-                }
-
-                input_embed = ggml_concat(ctx->ggml_ctx, input_embed, image_embeds[i].second, 1);
-            }
-
-            txt_token_start = image_embeds[image_embeds.size() - 1].first + image_embeds[image_embeds.size() - 1].second->ne[1];
-            txt_token_end   = raw_x->ne[1];
-
-            auto final_txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end);
-            input_embed          = ggml_concat(ctx->ggml_ctx, input_embed, final_txt_embed, 1);
-            GGML_ASSERT(raw_x->ne[1] == input_embed->ne[1]);
-            return input_embed;
-        }
-
         ggml_tensor* forward_embeds(GGMLRunnerContext* ctx,
                                     ggml_tensor* x,
                                     ggml_tensor* input_pos,
@@ -593,7 +708,7 @@ namespace LLM {
             // input_ids: [N, n_token]
             // return: [N, n_token, hidden_size]
             auto x = embed(ctx, input_ids);
-            x      = splice_image_embeds(ctx, x, std::move(image_embeds));
+            x      = splice_image_embeds(ctx, x, image_embeds);
             return forward_embeds(ctx, x, input_pos, attention_mask, std::move(out_layers));
         }
     };
@@ -608,18 +723,7 @@ namespace LLM {
             : enable_vision(enable_vision), params(params) {
             blocks["model"] = std::shared_ptr<GGMLBlock>(new TextModel(params));
             if (enable_vision) {
-                blocks["visual"] = std::shared_ptr<GGMLBlock>(new VisionModel(llama_cpp_style,
+                blocks["visual"] = std::shared_ptr<GGMLBlock>(new VisionModel(llama_cpp_style, params.vision));
-                                                                              params.vision.num_layers,
-                                                                              params.vision.in_channels,
-                                                                              params.vision.hidden_size,
-                                                                              params.vision.out_hidden_size,
-                                                                              params.vision.intermediate_size,
-                                                                              params.vision.num_heads,
-                                                                              params.vision.spatial_merge_size,
-                                                                              params.vision.patch_size,
-                                                                              params.vision.temporal_patch_size,
-                                                                              params.vision.window_size,
-                                                                              params.vision.fullatt_block_indexes));
             }
         }
 
@@ -636,15 +740,20 @@ namespace LLM {
             return x;
         }
 
+        std::shared_ptr<VisionModel> vision_model() {
+            GGML_ASSERT(enable_vision);
+            return std::dynamic_pointer_cast<VisionModel>(blocks["visual"]);
+        }
+
         ggml_tensor* vision_forward(GGMLRunnerContext* ctx,
                                     ggml_tensor* pixel_values,
                                     ggml_tensor* pe,
                                     ggml_tensor* window_index,
                                     ggml_tensor* window_inverse_index,
-                                    ggml_tensor* window_mask) {
+                                    ggml_tensor* window_mask,
+                                    ggml_tensor* pos_embeds = nullptr) {
             GGML_ASSERT(enable_vision);
-            auto vision_model = std::dynamic_pointer_cast<VisionModel>(blocks["visual"]);
+            return vision_model()->forward(ctx, pixel_values, pe, window_index, window_inverse_index, window_mask, pos_embeds);
-            return vision_model->forward(ctx, pixel_values, pe, window_index, window_inverse_index, window_mask);
         }
     };
 
@@ -659,7 +768,215 @@ namespace LLM {
         std::vector<int> window_index_vec;
         std::vector<int> window_inverse_index_vec;
         std::vector<float> pe_vec;
+        std::array<std::vector<int32_t>, 4> pos_embed_idx_data_;
+        std::array<std::vector<float>, 4> pos_embed_weight_data_;
 
+        static ggml_tensor* process_image_common(ggml_context* ctx,
+                                                 ggml_tensor* image,
+                                                 const LLMVisionParams& vision_params) {
+            // image: [C, H, W]
+            // return: [grid_t*(H/mh/ph)*(W/mw/pw)*mh*mw, C*pt*ph*pw], grid_t == 1
+            int64_t C  = image->ne[2];
+            int64_t H  = image->ne[1];
+            int64_t W  = image->ne[0];
+            int64_t mh = vision_params.spatial_merge_size;
+            int64_t mw = vision_params.spatial_merge_size;
+            int64_t pt = vision_params.temporal_patch_size;
+            int64_t ph = vision_params.patch_size;
+            int64_t pw = vision_params.patch_size;
+
+            image = ggml_reshape_4d(ctx, image, pw, mw, (W / mw / pw), H * C);                               // [C*H, (W/mw/pw), mw, pw]
+            image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 3, 1));                          // [mw, C*H, (W/mw/pw), pw]
+            image = ggml_reshape_4d(ctx, image, pw * (W / mw / pw), H, C, mw);                               // [mw, C, H, (W/mw/pw)*pw]
+            image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 3, 1));                          // [H, mw, C, (W/mw/pw)*pw]
+            image = ggml_reshape_4d(ctx, image, pw, (W / mw / pw) * C * mw, ph, mh * (H / mh / ph));         // [(H/mh/ph)*mh, ph, mw*C*(W/mw/pw), pw]
+            image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 1, 3));                          // [(H/mh/ph)*mh, mw*C*(W/mw/pw), ph, pw]
+            image = ggml_reshape_4d(ctx, image, pw * ph, (W / mw / pw), C, mw * mh * (H / mh / ph));         // [(H/mh/ph)*mh*mw, C, (W/mw/pw), ph*pw]
+            image = ggml_concat(ctx, image, image, 0);                                                       // [(H/mh/ph)*mh*mw, C, (W/mw/pw), pt*ph*pw]
+            image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 1, 3));                          // [(H/mh/ph)*mh*mw, (W/mw/pw), C, pt*ph*pw]
+            image = ggml_reshape_4d(ctx, image, pw * ph * pt * C, (W / mw / pw), mw * mh, (H / mh / ph));    // [(H/mh/ph), mh*mw, (W/mw/pw), C*pt*ph*pw]
+            image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 1, 3));                          // [(H/mh/ph), (W/mw/pw), mh*mw, C*pt*ph*pw]
+            image = ggml_reshape_2d(ctx, image, pw * ph * pt * C, mw * mh * (W / mw / pw) * (H / mh / ph));  // [(H/mh/ph)*(W/mw/pw)*mh*mw, C*pt*ph*pw]
+            return image;
+        }
+
+        static ggml_tensor* build_patch_pos_embeds_common(GGMLRunner* runner,
+                                                          ggml_context* compute_ctx,
+                                                          GGMLRunnerContext* runner_ctx,
+                                                          std::shared_ptr<VisionModel> vision,
+                                                          int grid_h,
+                                                          int grid_w,
+                                                          std::array<std::vector<int32_t>, 4>& pos_embed_idx_data,
+                                                          std::array<std::vector<float>, 4>& pos_embed_weight_data) {
+            auto pos_embed = vision->pos_embedder();
+            GGML_ASSERT(pos_embed != nullptr);
+            for (int i = 0; i < 4; ++i) {
+                pos_embed_idx_data[i].clear();
+                pos_embed_weight_data[i].clear();
+                pos_embed_idx_data[i].reserve(static_cast<size_t>(grid_h * grid_w));
+                pos_embed_weight_data[i].reserve(static_cast<size_t>(grid_h * grid_w));
+            }
+
+            int num_grid_per_side = vision->get_num_grid_per_side();
+            double max_index      = static_cast<double>(num_grid_per_side - 1);
+            int merge_size        = vision->get_spatial_merge_size();
+            GGML_ASSERT(grid_h % merge_size == 0);
+            GGML_ASSERT(grid_w % merge_size == 0);
+            for (int bh = 0; bh < grid_h / merge_size; ++bh) {
+                for (int bw = 0; bw < grid_w / merge_size; ++bw) {
+                    for (int ih = 0; ih < merge_size; ++ih) {
+                        int h        = bh * merge_size + ih;
+                        double h_pos = grid_h == 1 ? 0.0 : max_index * h / static_cast<double>(grid_h - 1);
+                        int h_floor  = static_cast<int>(std::floor(h_pos));
+                        int h_ceil   = std::min(h_floor + 1, num_grid_per_side - 1);
+                        double dh    = h_pos - h_floor;
+                        for (int iw = 0; iw < merge_size; ++iw) {
+                            int w        = bw * merge_size + iw;
+                            double w_pos = grid_w == 1 ? 0.0 : max_index * w / static_cast<double>(grid_w - 1);
+                            int w_floor  = static_cast<int>(std::floor(w_pos));
+                            int w_ceil   = std::min(w_floor + 1, num_grid_per_side - 1);
+                            double dw    = w_pos - w_floor;
+
+                            pos_embed_idx_data[0].push_back(h_floor * num_grid_per_side + w_floor);
+                            pos_embed_idx_data[1].push_back(h_floor * num_grid_per_side + w_ceil);
+                            pos_embed_idx_data[2].push_back(h_ceil * num_grid_per_side + w_floor);
+                            pos_embed_idx_data[3].push_back(h_ceil * num_grid_per_side + w_ceil);
+
+                            pos_embed_weight_data[0].push_back(static_cast<float>((1.0 - dh) * (1.0 - dw)));
+                            pos_embed_weight_data[1].push_back(static_cast<float>((1.0 - dh) * dw));
+                            pos_embed_weight_data[2].push_back(static_cast<float>(dh * (1.0 - dw)));
+                            pos_embed_weight_data[3].push_back(static_cast<float>(dh * dw));
+                        }
+                    }
+                }
+            }
+
+            ggml_tensor* patch_pos_embeds = nullptr;
+            for (int i = 0; i < 4; ++i) {
+                auto idx_tensor = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_I32, static_cast<int64_t>(pos_embed_idx_data[i].size()));
+                runner->set_backend_tensor_data(idx_tensor, pos_embed_idx_data[i].data());
+                auto embed         = pos_embed->forward(runner_ctx, idx_tensor);
+                auto weight_tensor = ggml_new_tensor_2d(compute_ctx, GGML_TYPE_F32, 1, static_cast<int64_t>(pos_embed_weight_data[i].size()));
+                runner->set_backend_tensor_data(weight_tensor, pos_embed_weight_data[i].data());
+                embed            = ggml_mul(compute_ctx, embed, weight_tensor);
+                patch_pos_embeds = patch_pos_embeds == nullptr ? embed : ggml_add(compute_ctx, patch_pos_embeds, embed);
+            }
+            return patch_pos_embeds;
+        }
+
+        static ggml_tensor* encode_image_common(GGMLRunner* runner,
+                                                ggml_context* compute_ctx,
+                                                GGMLRunnerContext* runner_ctx,
+                                                ggml_tensor* image,
+                                                const LLMVisionParams& vision_params,
+                                                std::shared_ptr<VisionModel> vision_model,
+                                                std::vector<int>& window_index_vec,
+                                                std::vector<int>& window_inverse_index_vec,
+                                                std::vector<float>& window_mask_vec,
+                                                std::vector<float>& pe_vec,
+                                                std::array<std::vector<int32_t>, 4>& pos_embed_idx_data,
+                                                std::array<std::vector<float>, 4>& pos_embed_weight_data) {
+            GGML_ASSERT(image->ne[1] % (vision_params.patch_size * vision_params.spatial_merge_size) == 0);
+            GGML_ASSERT(image->ne[0] % (vision_params.patch_size * vision_params.spatial_merge_size) == 0);
+
+            int grid_h = static_cast<int>(image->ne[1]) / vision_params.patch_size;
+            int grid_w = static_cast<int>(image->ne[0]) / vision_params.patch_size;
+
+            auto pixel_values = process_image_common(compute_ctx, image, vision_params);
+            int head_dim      = static_cast<int>(vision_params.hidden_size / vision_params.num_heads);
+
+            if (vision_params.arch == LLMVisionArch::QWEN3_VL) {
+                auto pos_embeds = build_patch_pos_embeds_common(runner,
+                                                                compute_ctx,
+                                                                runner_ctx,
+                                                                vision_model,
+                                                                grid_h,
+                                                                grid_w,
+                                                                pos_embed_idx_data,
+                                                                pos_embed_weight_data);
+                window_index_vec.resize(static_cast<size_t>((grid_h / vision_params.spatial_merge_size) * (grid_w / vision_params.spatial_merge_size)));
+                for (int i = 0; i < static_cast<int>(window_index_vec.size()); ++i) {
+                    window_index_vec[static_cast<size_t>(i)] = i;
+                }
+                pe_vec      = Rope::gen_qwen2vl_pe(grid_h,
+                                                   grid_w,
+                                                   vision_params.spatial_merge_size,
+                                                   window_index_vec,
+                                                   10000,
+                                                   {head_dim / 2, head_dim / 2});
+                int pos_len = static_cast<int>(pe_vec.size() / head_dim / 2);
+                auto pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, head_dim / 2, pos_len);
+                runner->set_backend_tensor_data(pe, pe_vec.data());
+                return vision_model->forward(runner_ctx, pixel_values, pe, nullptr, nullptr, nullptr, pos_embeds);
+            }
+
+            int llm_grid_h             = grid_h / vision_params.spatial_merge_size;
+            int llm_grid_w             = grid_w / vision_params.spatial_merge_size;
+            int vit_merger_window_size = vision_params.window_size / vision_params.patch_size / vision_params.spatial_merge_size;
+
+            int inverse_index = 0;
+            window_index_vec.resize(llm_grid_h * llm_grid_w);
+            window_inverse_index_vec.resize(llm_grid_h * llm_grid_w);
+            std::vector<int> seqlens;
+            for (int ih = 0; ih < llm_grid_h; ih += vit_merger_window_size) {
+                for (int iw = 0; iw < llm_grid_w; iw += vit_merger_window_size) {
+                    int win_h = std::min(vit_merger_window_size, llm_grid_h - ih);
+                    int win_w = std::min(vit_merger_window_size, llm_grid_w - iw);
+                    for (int iy = 0; iy < win_h; iy++) {
+                        for (int ix = 0; ix < win_w; ix++) {
+                            int index                       = (ih + iy) * llm_grid_w + iw + ix;
+                            window_index_vec[inverse_index] = index;
+                            window_inverse_index_vec[index] = inverse_index;
+                            inverse_index++;
+                        }
+                    }
+                    seqlens.push_back(win_h * win_w * vision_params.spatial_merge_size * vision_params.spatial_merge_size);
+                }
+            }
+            auto window_index         = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_I32, llm_grid_h * llm_grid_w);
+            auto window_inverse_index = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_I32, llm_grid_h * llm_grid_w);
+            runner->set_backend_tensor_data(window_index, window_index_vec.data());
+            runner->set_backend_tensor_data(window_inverse_index, window_inverse_index_vec.data());
+
+            window_mask_vec.resize((grid_h * grid_w) * (grid_h * grid_w));
+            int window_start_index = 0;
+            for (int seq_index = 0; seq_index < seqlens.size(); seq_index++) {
+                int window_end_index = window_start_index + seqlens[seq_index];
+                GGML_ASSERT(window_end_index <= grid_h * grid_w);
+                for (int i = window_start_index; i < window_end_index; i++) {
+                    for (int j = 0; j < grid_h * grid_w; j++) {
+                        float mask_value = -INFINITY;
+                        if (j >= window_start_index && j < window_end_index) {
+                            mask_value = 0;
+                        }
+                        GGML_ASSERT((i * (grid_h * grid_w) + j) < window_mask_vec.size());
+                        window_mask_vec[i * (grid_h * grid_w) + j] = mask_value;
+                    }
+                }
+                window_start_index = window_end_index;
+            }
+
+            auto window_mask = ggml_new_tensor_2d(compute_ctx,
+                                                  GGML_TYPE_F32,
+                                                  grid_h * grid_w,
+                                                  grid_h * grid_w);
+            runner->set_backend_tensor_data(window_mask, window_mask_vec.data());
+
+            pe_vec      = Rope::gen_qwen2vl_pe(grid_h,
+                                               grid_w,
+                                               vision_params.spatial_merge_size,
+                                               window_inverse_index_vec,
+                                               10000,
+                                               {head_dim / 2, head_dim / 2});
+            int pos_len = static_cast<int>(pe_vec.size() / head_dim / 2);
+
+            auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, head_dim / 2, pos_len);
+            runner->set_backend_tensor_data(pe, pe_vec.data());
+
+            return vision_model->forward(runner_ctx, pixel_values, pe, window_index, window_inverse_index, window_mask);
+        }
+
+    public:
         LLMRunner(LLMArch arch,
                   ggml_backend_t backend,
                   bool offload_params_to_cpu,
@@ -761,8 +1078,9 @@ namespace LLM {
                                     ggml_tensor* input_pos,
                                     ggml_tensor* window_index,
                                     ggml_tensor* window_inverse_index,
-                                    ggml_tensor* window_mask) {
+                                    ggml_tensor* window_mask,
-            auto hidden_states = model.vision_forward(ctx, pixel_values, input_pos, window_index, window_inverse_index, window_mask);
+                                    ggml_tensor* pos_embeds = nullptr) {
+            auto hidden_states = model.vision_forward(ctx, pixel_values, input_pos, window_index, window_inverse_index, window_mask, pos_embeds);
             return hidden_states;
         }
 
@@ -848,30 +1166,36 @@ namespace LLM {
         }
 
         ggml_tensor* process_image(ggml_context* ctx, ggml_tensor* image) {
-            // image: [C, H, W]
+            return process_image_common(ctx, image, params.vision);
-            // return: [grid_t*(H/mh/ph)*(W/mw/pw)*mh*mw, C*pt*ph*pw], grid_t == 1
+        }
-            int64_t C  = image->ne[2];
-            int64_t H  = image->ne[1];
-            int64_t W  = image->ne[0];
-            int64_t mh = params.vision.spatial_merge_size;
-            int64_t mw = params.vision.spatial_merge_size;
-            int64_t pt = params.vision.temporal_patch_size;
-            int64_t ph = params.vision.patch_size;
-            int64_t pw = params.vision.patch_size;
 
-            image = ggml_reshape_4d(ctx, image, pw, mw, (W / mw / pw), H * C);                               // [C*H, (W/mw/pw), mw, pw]
+        ggml_tensor* build_patch_pos_embeds(GGMLRunnerContext* runner_ctx,
-            image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 3, 1));                          // [mw, C*H, (W/mw/pw), pw]
+                                            std::shared_ptr<VisionModel> vision,
-            image = ggml_reshape_4d(ctx, image, pw * (W / mw / pw), H, C, mw);                               // [mw, C, H, (W/mw/pw)*pw]
+                                            int grid_h,
-            image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 3, 1));                          // [H, mw, C, (W/mw/pw)*pw]
+                                            int grid_w) {
-            image = ggml_reshape_4d(ctx, image, pw, (W / mw / pw) * C * mw, ph, mh * (H / mh / ph));         // [(H/mh/ph)*mh, ph, mw*C*(W/mw/pw), pw]
+            return build_patch_pos_embeds_common(this,
-            image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 1, 3));                          // [(H/mh/ph)*mh, mw*C*(W/mw/pw), ph, pw]
+                                                 compute_ctx,
-            image = ggml_reshape_4d(ctx, image, pw * ph, (W / mw / pw), C, mw * mh * (H / mh / ph));         // [(H/mh/ph)*mh*mw, C, (W/mw/pw), ph*pw]
+                                                 runner_ctx,
-            image = ggml_concat(ctx, image, image, 0);                                                       // [(H/mh/ph)*mh*mw, C, (W/mw/pw), pt*ph*pw]
+                                                 vision,
-            image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 1, 3));                          // [(H/mh/ph)*mh*mw, (W/mw/pw), C, pt*ph*pw]
+                                                 grid_h,
-            image = ggml_reshape_4d(ctx, image, pw * ph * pt * C, (W / mw / pw), mw * mh, (H / mh / ph));    // [(H/mh/ph), mh*mw, (W/mw/pw), C*pt*ph*pw]
+                                                 grid_w,
-            image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 1, 3));                          // [(H/mh/ph), (W/mw/pw), mh*mw, C*pt*ph*pw]
+                                                 pos_embed_idx_data_,
-            image = ggml_reshape_2d(ctx, image, pw * ph * pt * C, mw * mh * (W / mw / pw) * (H / mh / ph));  // [(H/mh/ph)*(W/mw/pw)*mh*mw, C*pt*ph*pw]
+                                                 pos_embed_weight_data_);
-            return image;
+        }
+
+        ggml_tensor* encode_image(GGMLRunnerContext* runner_ctx, ggml_tensor* image) {
+            return encode_image_common(this,
+                                       compute_ctx,
+                                       runner_ctx,
+                                       image,
+                                       params.vision,
+                                       model.vision_model(),
+                                       window_index_vec,
+                                       window_inverse_index_vec,
+                                       window_mask_vec,
+                                       pe_vec,
+                                       pos_embed_idx_data_,
+                                       pos_embed_weight_data_);
         }
 
         ggml_cgraph* build_encode_image_graph(const sd::Tensor<float>& image_tensor) {
@@ -881,116 +1205,8 @@ namespace LLM {
             GGML_ASSERT(image->ne[1] % (params.vision.patch_size * params.vision.spatial_merge_size) == 0);
             GGML_ASSERT(image->ne[0] % (params.vision.patch_size * params.vision.spatial_merge_size) == 0);
 
-            int grid_t                 = 1;
-            int grid_h                 = static_cast<int>(image->ne[1]) / params.vision.patch_size;
-            int grid_w                 = static_cast<int>(image->ne[0]) / params.vision.patch_size;
-            int llm_grid_h             = grid_h / params.vision.spatial_merge_size;
-            int llm_grid_w             = grid_w / params.vision.spatial_merge_size;
-            int vit_merger_window_size = params.vision.window_size / params.vision.patch_size / params.vision.spatial_merge_size;
-
-            auto pixel_values = process_image(compute_ctx, image);
-
-            // window index
-            int inverse_index = 0;
-            window_index_vec.resize(llm_grid_h * llm_grid_w);
-            window_inverse_index_vec.resize(llm_grid_h * llm_grid_w);
-            std::vector<int> seqlens;
-            for (int ih = 0; ih < llm_grid_h; ih += vit_merger_window_size) {
-                for (int iw = 0; iw < llm_grid_w; iw += vit_merger_window_size) {
-                    int win_h = std::min(vit_merger_window_size, llm_grid_h - ih);
-                    int win_w = std::min(vit_merger_window_size, llm_grid_w - iw);
-                    for (int iy = 0; iy < win_h; iy++) {
-                        for (int ix = 0; ix < win_w; ix++) {
-                            int index                       = (ih + iy) * llm_grid_w + iw + ix;
-                            window_index_vec[inverse_index] = index;
-                            window_inverse_index_vec[index] = inverse_index;
-                            inverse_index++;
-                        }
-                    }
-                    seqlens.push_back(win_h * win_w * params.vision.spatial_merge_size * params.vision.spatial_merge_size);
-                }
-            }
-            // printf("window_index: ");
-            // for (int i : window_index_vec) {
-            //     printf("%d ", i);
-            // }
-            // printf("\n");
-            // printf("window_inverse_index: ");
-            // for (int i : window_inverse_index_vec) {
-            //     printf("%d ", i);
-            // }
-            // printf("\n");
-            // printf("seqlens: ");
-            // for (int i : seqlens) {
-            //     printf("%d ", i);
-            // }
-            // printf("\n");
-            auto window_index         = ggml_new_tensor_1d(compute_ctx,
-                                                           GGML_TYPE_I32,
-                                                           llm_grid_h * llm_grid_w);
-            auto window_inverse_index = ggml_new_tensor_1d(compute_ctx,
-                                                           GGML_TYPE_I32,
-                                                           llm_grid_h * llm_grid_w);
-            set_backend_tensor_data(window_index, window_index_vec.data());
-            set_backend_tensor_data(window_inverse_index, window_inverse_index_vec.data());
-
-            // window mask
-            int seq_window_size = (vit_merger_window_size * params.vision.spatial_merge_size) * (vit_merger_window_size * params.vision.spatial_merge_size);
-            window_mask_vec.resize((grid_h * grid_w) * (grid_h * grid_w));
-            int window_start_index = 0;
-            for (int seq_index = 0; seq_index < seqlens.size(); seq_index++) {
-                int window_end_index = window_start_index + seqlens[seq_index];
-                // LOG_DEBUG("%d %d", window_start_index, window_end_index);
-                GGML_ASSERT(window_end_index <= grid_h * grid_w);
-                for (int i = window_start_index; i < window_end_index; i++) {
-                    for (int j = 0; j < grid_h * grid_w; j++) {
-                        float mask_value = -INFINITY;
-                        if (j >= window_start_index && j < window_end_index) {
-                            mask_value = 0;
-                        }
-                        GGML_ASSERT((i * (grid_h * grid_w) + j) < window_mask_vec.size());
-                        window_mask_vec[i * (grid_h * grid_w) + j] = mask_value;
-                    }
-                }
-                window_start_index = window_end_index;
-                // printf("\n");
-            }
-            // printf("window_mask: \n");
-            // for (int i = 0; i < grid_h*grid_w; i++) {
-            //     for (int j = 0; j < grid_h*grid_w; j++) {
-            //         printf("%f ", window_mask_vec[i * (grid_h * grid_w) + j]);
-            //     }
-            //     printf("\n");
-            // }
-            auto window_mask = ggml_new_tensor_2d(compute_ctx,
-                                                  GGML_TYPE_F32,
-                                                  grid_h * grid_w,
-                                                  grid_h * grid_w);
-            set_backend_tensor_data(window_mask, window_mask_vec.data());
-
-            // pe
-            int head_dim = static_cast<int>(params.vision.hidden_size / params.vision.num_heads);
-            pe_vec       = Rope::gen_qwen2vl_pe(grid_h,
-                                                grid_w,
-                                                params.vision.spatial_merge_size,
-                                                window_inverse_index_vec,
-                                                10000,
-                                                {head_dim / 2, head_dim / 2});
-            int pos_len  = static_cast<int>(pe_vec.size() / head_dim / 2);
-            // LOG_DEBUG("pos_len %d", pos_len);
-            auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, head_dim / 2, pos_len);
-            // pe->data = pe_vec.data();
-            // print_ggml_tensor(pe);
-            // pe->data = nullptr;
-            set_backend_tensor_data(pe, pe_vec.data());
-
             auto runnter_ctx           = get_context();
-            ggml_tensor* hidden_states = vision_forward(&runnter_ctx,
+            ggml_tensor* hidden_states = encode_image(&runnter_ctx, image);
-                                                        pixel_values,
-                                                        pe,
-                                                        window_index,
-                                                        window_inverse_index,
-                                                        window_mask);
             ggml_build_forward_expand(gf, hidden_states);
 
             return gf;

src/model.cpp

@@ -734,16 +734,10 @@ void ModelLoader::set_wtype_override(ggml_type wtype, std::string tensor_type_ru
     }
 }
 
-bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_threads_p, bool enable_mmap) {
+void ModelLoader::process_model_files(bool enable_mmap, bool writable_mmap) {
-    int64_t process_time_ms = 0;
+    if (model_files_processed) {
-    std::atomic<int64_t> read_time_ms(0);
+        return;
-    std::atomic<int64_t> memcpy_time_ms(0);
+    }
-    std::atomic<int64_t> copy_to_backend_time_ms(0);
-    std::atomic<int64_t> convert_time_ms(0);
-    std::atomic<uint64_t> bytes_processed(0);
-
-    int num_threads_to_use = n_threads_p > 0 ? n_threads_p : sd_get_num_physical_cores();
-    LOG_DEBUG("using %d threads for model loading", num_threads_to_use);
 
     int64_t start_time = ggml_time_ms();
 
@@ -755,22 +749,13 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
         processed_tensor_storages.push_back(tensor_storage);
     }
 
-    process_time_ms = ggml_time_ms() - start_time;
-
-    bool success                          = true;
-    size_t total_tensors_processed        = 0;
-    const size_t total_tensors_to_process = processed_tensor_storages.size();
-    const int64_t t_start                 = ggml_time_ms();
-    int last_n_threads                    = 1;
-
     for (size_t file_index = 0; file_index < file_paths_.size(); file_index++) {
         std::string file_path = file_paths_[file_index];
-        LOG_DEBUG("loading tensors from %s", file_path.c_str());
 
-        std::vector<const TensorStorage*> file_tensors;
+        std::vector<TensorStorage> file_tensors;
         for (const auto& ts : processed_tensor_storages) {
             if (ts.file_index == file_index) {
-                file_tensors.push_back(&ts);
+                file_tensors.push_back(ts);
             }
         }
         if (file_tensors.empty()) {
@@ -779,20 +764,168 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
 
         bool is_zip = false;
         for (auto const& ts : file_tensors) {
-            if (ts->index_in_zip >= 0) {
+            if (ts.index_in_zip >= 0) {
                 is_zip = true;
                 break;
             }
         }
 
-        std::unique_ptr<MmapWrapper> mmapped;
+        ModelFileData fdata = {};
+        fdata.path          = file_path;
+        fdata.is_zip        = is_zip;
+        fdata.tensors       = std::move(file_tensors);
+
         if (enable_mmap && !is_zip) {
             LOG_DEBUG("using mmap for I/O");
-            mmapped = MmapWrapper::create(file_path);
+            std::unique_ptr<MmapWrapper> mmapped = MmapWrapper::create(file_path, writable_mmap);
-            if (!mmapped) {
+            if (mmapped) {
-                LOG_WARN("failed to memory-map '%s'", file_path.c_str());
+                uint8_t* mmap_data             = static_cast<uint8_t*>(mmapped->writable_data());
+                ggml_backend_buffer_t buf_mmap = ggml_backend_cpu_buffer_from_ptr(mmap_data, mmapped->size());
+                if (buf_mmap) {
+                    LOG_INFO("using mmap for '%s'", file_path.c_str());
+                    fdata.mmbuffer = std::shared_ptr<struct ggml_backend_buffer>(buf_mmap, ggml_backend_buffer_free);
+                } else {
+                    LOG_WARN("mmap: failed to create backend buffer for file %s", fdata.path.c_str());
+                }
+                fdata.mmapped = std::shared_ptr<MmapWrapper>(std::move(mmapped));
+            } else {
+                LOG_WARN("failed to memory-map '%s' (falling back to read())", file_path.c_str());
+            }
+        } else if (!is_zip) {
+            LOG_INFO("NOT using mmap for '%s' (mmap disabled by caller)",
+                     file_path.c_str());
+        }
+
+        file_data.push_back(std::move(fdata));
+    }
+
+    model_files_processed = true;
+
+    int64_t end_time        = ggml_time_ms();
+    int64_t process_time_ms = end_time - start_time;
+
+    LOG_INFO("model files processing completed in %.2fs", process_time_ms / 1000.f);
+}
+
+std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggml_tensor*>& tensors,
+                                                       std::set<std::string> ignore_tensors,
+                                                       bool writable_mmap) {
+    process_model_files(true, writable_mmap);
+
+    std::vector<MmapTensorStore> result;
+    uint64_t mapped_bytes = 0;
+    size_t mapped_tensors = 0;
+
+    LOG_DEBUG("memory-mapping tensors...");
+
+    int64_t t_start = ggml_time_ms();
+
+    for (auto& fdata : file_data) {
+        if (!fdata.mmbuffer)
+            continue;
+
+        const std::vector<TensorStorage>& file_tensors = fdata.tensors;
+
+        size_t file_mapped_bytes   = 0;
+        size_t file_mapped_tensors = 0;
+
+        for (const auto& tensor_storage : file_tensors) {
+            const std::string& name = tensor_storage.name;
+
+            bool is_ignored = false;
+            for (const auto& ignore_prefix : ignore_tensors) {
+                if (starts_with(name, ignore_prefix)) {
+                    is_ignored = true;
+                    break;
                 }
             }
+            if (is_ignored)
+                continue;
+
+            auto it = tensors.find(name);
+            if (it == tensors.end())
+                continue;
+
+            ggml_tensor* dst_tensor = it->second;
+            if (dst_tensor == nullptr)
+                continue;
+
+            if (tensor_storage.type != dst_tensor->type)
+                continue;
+
+            size_t tensor_size   = tensor_storage.nbytes();
+            size_t tensor_offset = tensor_storage.offset;
+
+            if (tensor_storage.ne[0] != dst_tensor->ne[0] ||
+                tensor_storage.ne[1] != dst_tensor->ne[1] ||
+                tensor_storage.ne[2] != dst_tensor->ne[2] ||
+                tensor_storage.ne[3] != dst_tensor->ne[3] ||
+                tensor_size != ggml_nbytes(dst_tensor)) {
+                // let load_tensors worry about this
+                continue;
+            }
+
+            ggml_backend_buffer_t buf_mmap = fdata.mmbuffer.get();
+            uint8_t* mmap_data             = static_cast<uint8_t*>(ggml_backend_buffer_get_base(buf_mmap));
+            dst_tensor->buffer             = buf_mmap;
+            dst_tensor->data               = mmap_data + tensor_offset;
+
+            file_mapped_bytes += tensor_size;
+            file_mapped_tensors++;
+        }
+
+        if (file_mapped_bytes > 0) {
+            mapped_tensors += file_mapped_tensors;
+            mapped_bytes += file_mapped_bytes;
+            result.push_back({fdata.mmapped, fdata.mmbuffer});
+        }
+    }
+
+    int64_t t_end       = ggml_time_ms();
+    int64_t duration_ms = t_end - t_start;
+
+    LOG_INFO("memory-mapped %zu tensors in %zu files (%.2f MB), taking %.2fs",
+             mapped_tensors,
+             result.size(),
+             mapped_bytes / (1024.0 * 1024.0),
+             duration_ms / 1000.0);
+
+    return result;
+}
+
+bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_threads_p, bool enable_mmap) {
+    process_model_files(enable_mmap, false);
+
+    std::atomic<int64_t> read_time_ms(0);
+    std::atomic<int64_t> memcpy_time_ms(0);
+    std::atomic<int64_t> copy_to_backend_time_ms(0);
+    std::atomic<int64_t> convert_time_ms(0);
+    std::atomic<uint64_t> bytes_processed(0);
+
+    int num_threads_to_use = n_threads_p > 0 ? n_threads_p : sd_get_num_physical_cores();
+    LOG_DEBUG("using %d threads for model loading", num_threads_to_use);
+
+    int64_t start_time = ggml_time_ms();
+
+    size_t total_tensors_to_process = 0;
+    for (const auto& fdata : file_data) {
+        total_tensors_to_process += fdata.tensors.size();
+    }
+
+    bool success                   = true;
+    size_t total_tensors_processed = 0;
+    const int64_t t_start          = start_time;
+    int last_n_threads             = 1;
+
+    for (auto& fdata : file_data) {
+        const std::string& file_path = fdata.path;
+        LOG_DEBUG("loading tensors from %s", file_path.c_str());
+
+        const std::vector<TensorStorage>& file_tensors = fdata.tensors;
+
+        bool is_zip = fdata.is_zip;
+
+        std::shared_ptr<MmapWrapper> mmapped = fdata.mmapped;
 
         int n_threads = is_zip ? 1 : std::min(num_threads_to_use, (int)file_tensors.size());
         if (n_threads < 1) {
@@ -834,7 +967,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
                         break;
                     }
 
-                    const TensorStorage& tensor_storage = *file_tensors[idx];
+                    const TensorStorage& tensor_storage = file_tensors[idx];
                     ggml_tensor* dst_tensor             = nullptr;
 
                     t0 = ggml_time_ms();
@@ -851,6 +984,11 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
                         continue;
                     }
 
+                    // skip mmapped tensors
+                    if (dst_tensor->buffer != nullptr && dst_tensor->buffer == fdata.mmbuffer.get()) {
+                        continue;
+                    }
+
                     size_t nbytes_to_read = tensor_storage.nbytes_to_read();
 
                     auto read_data = [&](char* buf, size_t n) {
@@ -994,9 +1132,8 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
     }
 
     int64_t end_time = ggml_time_ms();
-    LOG_INFO("loading tensors completed, taking %.2fs (process: %.2fs, read: %.2fs, memcpy: %.2fs, convert: %.2fs, copy_to_backend: %.2fs)",
+    LOG_INFO("loading tensors completed, taking %.2fs (read: %.2fs, memcpy: %.2fs, convert: %.2fs, copy_to_backend: %.2fs)",
              (end_time - start_time) / 1000.f,
-             process_time_ms / 1000.f,
              (read_time_ms.load() / (float)last_n_threads) / 1000.f,
              (memcpy_time_ms.load() / (float)last_n_threads) / 1000.f,
              (convert_time_ms.load() / (float)last_n_threads) / 1000.f,

src/model.h

@@ -195,10 +195,27 @@ using TensorTypeRules = std::vector<std::pair<std::string, ggml_type>>;
 
 TensorTypeRules parse_tensor_type_rules(const std::string& tensor_type_rules);
 
+class MmapWrapper;
+
+struct ModelFileData {
+    std::string path;
+    std::vector<TensorStorage> tensors;
+    std::shared_ptr<MmapWrapper> mmapped;
+    std::shared_ptr<struct ggml_backend_buffer> mmbuffer;
+    bool is_zip;
+};
+
+struct MmapTensorStore {
+    std::shared_ptr<MmapWrapper> mmapped;
+    std::shared_ptr<struct ggml_backend_buffer> mmbuffer;
+};
+
 class ModelLoader {
 protected:
     SDVersion version_ = VERSION_COUNT;
     std::vector<std::string> file_paths_;
+    std::vector<ModelFileData> file_data;
+    bool model_files_processed = false;
     String2TensorStorage tensor_storage_map;
 
     void add_tensor_storage(const TensorStorage& tensor_storage);
@@ -222,6 +239,10 @@ public:
     std::map<ggml_type, uint32_t> get_vae_wtype_stat();
     String2TensorStorage& get_tensor_storage_map() { return tensor_storage_map; }
     void set_wtype_override(ggml_type wtype, std::string tensor_type_rules = "");
+    void process_model_files(bool enable_mmap = false, bool writable_mmap = true);
+    std::vector<MmapTensorStore> mmap_tensors(std::map<std::string, ggml_tensor*>& tensors,
+                                              std::set<std::string> ignore_tensors = {},
+                                              bool writable                        = true);
     bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_threads = 0, bool use_mmap = false);
     bool load_tensors(std::map<std::string, ggml_tensor*>& tensors,
                       std::set<std::string> ignore_tensors = {},

src/stable-diffusion.cpp

@@ -60,7 +60,6 @@ const char* model_version_to_str[] = {
 
 const char* sampling_methods_str[] = {
     "Euler",
-    "Euler Flow Flash",
     "Euler A",
     "Heun",
     "DPM2",
@@ -75,6 +74,8 @@ const char* sampling_methods_str[] = {
     "Res Multistep",
     "Res 2s",
     "ER-SDE",
+    "Euler CFG++",
+    "Euler A CFG++",
 };
 
 /*================================================== Helper Functions ================================================*/
@@ -110,6 +111,7 @@ static float get_cache_reuse_threshold(const sd_cache_params_t& params) {
 
 class StableDiffusionGGML {
 public:
+    std::vector<MmapTensorStore> mmap_tensor_store;
     ggml_backend_t backend             = nullptr;  // general backend
     ggml_backend_t clip_backend        = nullptr;
     ggml_backend_t control_net_backend = nullptr;
@@ -362,6 +364,51 @@ public:
             apply_lora_immediately = false;
         }
 
+        std::map<std::string, ggml_tensor*> mmap_able_tensors;
+        bool enable_mmap_tensors = false;
+        bool main_backend_mmap   = false;
+        bool needs_writable_mmap = false;
+        if (sd_ctx_params->enable_mmap) {
+            if (apply_lora_immediately) {
+                needs_writable_mmap = true;
+                LOG_WARN("in mode 'immediately', LoRAs will cause extra memory usage with mmap");
+            }
+            enable_mmap_tensors = true;
+            if (offload_params_to_cpu) {
+                main_backend_mmap = true;
+            } else {
+                ggml_backend_dev_t dev = ggml_backend_get_device(backend);
+                struct ggml_backend_dev_props props;
+                ggml_backend_dev_get_props(dev, &props);
+                main_backend_mmap = props.caps.buffer_from_host_ptr;
+            }
+        }
+
+        // split definition to avoid msvc choking on the extra parameter handling
+        auto get_param_tensors_p = [&](auto&& model, bool force_cpu, const char* prefix) {
+            std::map<std::string, ggml_tensor*> temp;
+            model->get_param_tensors(temp, prefix);
+            bool do_mmap = enable_mmap_tensors && (main_backend_mmap || force_cpu);
+            for (const auto& [key, tensor] : temp) {
+                tensors[key] = tensor;
+                if (do_mmap) {
+                    mmap_able_tensors[key] = tensor;
+                }
+            }
+        };
+
+        auto get_param_tensors = [&](auto&& model, bool force_cpu = false) {
+            std::map<std::string, ggml_tensor*> temp;
+            model->get_param_tensors(temp);
+            bool do_mmap = enable_mmap_tensors && (main_backend_mmap || force_cpu);
+            for (const auto& [key, tensor] : temp) {
+                tensors[key] = tensor;
+                if (do_mmap) {
+                    mmap_able_tensors[key] = tensor;
+                }
+            }
+        };
+
         if (sd_version_is_control(version)) {
             // Might need vae encode for control cond
             vae_decode_only = false;
@@ -473,8 +520,7 @@ public:
                                                                              offload_params_to_cpu,
                                                                              tensor_storage_map);
                     clip_vision->set_max_graph_vram_bytes(max_graph_vram_bytes);
-                    clip_vision->alloc_params_buffer();
+                    get_param_tensors(clip_vision);
-                    clip_vision->get_param_tensors(tensors);
                 }
             } else if (sd_version_is_qwen_image(version)) {
                 bool enable_vision = false;
@@ -494,7 +540,9 @@ public:
                                                                    version,
                                                                    sd_ctx_params->qwen_image_zero_cond_t);
             } else if (version == VERSION_HIDREAM_O1) {
-                cond_stage_model = std::make_shared<HiDreamO1::HiDreamO1Conditioner>();
+                cond_stage_model = std::make_shared<HiDreamO1::HiDreamO1Conditioner>(clip_backend,
+                                                                                     offload_params_to_cpu,
+                                                                                     tensor_storage_map);
                 diffusion_model  = std::make_shared<HiDreamO1Model>(backend,
                                                                    offload_params_to_cpu,
                                                                    tensor_storage_map,
@@ -556,12 +604,10 @@ public:
             }
 
             cond_stage_model->set_max_graph_vram_bytes(max_graph_vram_bytes);
-            cond_stage_model->alloc_params_buffer();
+            get_param_tensors(cond_stage_model, clip_on_cpu);
-            cond_stage_model->get_param_tensors(tensors);
 
             diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes);
-            diffusion_model->alloc_params_buffer();
+            get_param_tensors(diffusion_model);
-            diffusion_model->get_param_tensors(tensors);
 
             if (sd_version_is_unet_edit(version)) {
                 vae_decode_only = false;
@@ -569,8 +615,7 @@ public:
 
             if (high_noise_diffusion_model) {
                 high_noise_diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes);
-                high_noise_diffusion_model->alloc_params_buffer();
+                get_param_tensors(high_noise_diffusion_model);
-                high_noise_diffusion_model->get_param_tensors(tensors);
             }
 
             if (sd_ctx_params->keep_vae_on_cpu && !ggml_backend_is_cpu(backend)) {
@@ -633,6 +678,8 @@ public:
                 }
             };
 
+            bool force_vae_cpu = sd_ctx_params->keep_vae_on_cpu;
+
             if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1) {
                 LOG_INFO("using FakeVAE");
                 first_stage_model = std::make_shared<FakeVAE>(version,
@@ -642,20 +689,17 @@ public:
                 LOG_INFO("using TAE for encoding / decoding");
                 first_stage_model = create_tae();
                 first_stage_model->set_max_graph_vram_bytes(max_graph_vram_bytes);
-                first_stage_model->alloc_params_buffer();
+                get_param_tensors_p(first_stage_model, force_vae_cpu, "tae");
-                first_stage_model->get_param_tensors(tensors, "tae");
             } else {
                 LOG_INFO("using VAE for encoding / decoding");
                 first_stage_model = create_vae();
                 first_stage_model->set_max_graph_vram_bytes(max_graph_vram_bytes);
-                first_stage_model->alloc_params_buffer();
+                get_param_tensors_p(first_stage_model, force_vae_cpu, "first_stage_model");
-                first_stage_model->get_param_tensors(tensors, "first_stage_model");
                 if (use_tae && tae_preview_only) {
                     LOG_INFO("using TAE for preview");
                     preview_vae = create_tae();
                     preview_vae->set_max_graph_vram_bytes(max_graph_vram_bytes);
-                    preview_vae->alloc_params_buffer();
+                    get_param_tensors_p(first_stage_model, force_vae_cpu, "vae");
-                    preview_vae->get_param_tensors(tensors, "tae");
                 }
             }
 
@@ -720,11 +764,7 @@ public:
                 }
             }
             if (use_pmid) {
-                if (!pmid_model->alloc_params_buffer()) {
+                get_param_tensors_p(pmid_model, false, "pmid");
-                    LOG_ERROR(" pmid model params buffer allocation failed");
-                    return false;
-                }
-                pmid_model->get_param_tensors(tensors, "pmid");
             }
 
             if (sd_ctx_params->flash_attn) {
@@ -806,7 +846,43 @@ public:
         }
         if (version == VERSION_HIDREAM_O1) {
             ignore_tensors.insert("lm_head.");
+            ignore_tensors.insert("model.visual.deepstack_merger_list.");
         }
+
+        if (enable_mmap_tensors) {
+            if (mmap_able_tensors.empty()) {
+                LOG_DEBUG("no tensors could be memory-mapped");
+            } else {
+                mmap_tensor_store = model_loader.mmap_tensors(mmap_able_tensors, ignore_tensors, needs_writable_mmap);
+            }
+        }
+
+        if (clip_vision) {
+            clip_vision->alloc_params_buffer();
+        }
+        if (cond_stage_model) {
+            cond_stage_model->alloc_params_buffer();
+        }
+        if (diffusion_model) {
+            diffusion_model->alloc_params_buffer();
+        }
+        if (high_noise_diffusion_model) {
+            high_noise_diffusion_model->alloc_params_buffer();
+        }
+        if (first_stage_model) {
+            first_stage_model->alloc_params_buffer();
+        }
+        if (preview_vae) {
+            preview_vae->alloc_params_buffer();
+        }
+        if (use_pmid && pmid_model) {
+            if (!pmid_model->alloc_params_buffer()) {
+                LOG_ERROR(" pmid model params buffer allocation failed");
+                ggml_free(ctx);
+                return false;
+            }
+        }
+
         bool success = model_loader.load_tensors(tensors, ignore_tensors, n_threads, sd_ctx_params->enable_mmap);
         if (!success) {
             LOG_ERROR("load tensors from model loader failed");
@@ -1598,6 +1674,7 @@ public:
                              int shifted_timestep,
                              sample_method_t method,
                              bool is_flow_denoiser,
+                             const char* extra_sample_args,
                              const std::vector<float>& sigmas,
                              int start_merge_step,
                              const std::vector<sd::Tensor<float>>& ref_latents,
@@ -1615,6 +1692,15 @@ public:
                                                                                            cache_params,
                                                                                            denoiser.get(),
                                                                                            sigmas);
+
+        // Spectrum cache is not supported for CFG++ samplers
+        if (method == EULER_CFG_PP_SAMPLE_METHOD || method == EULER_A_CFG_PP_SAMPLE_METHOD) {
+            if (cache_runtime.spectrum_enabled) {
+                LOG_WARN("Spectrum cache requested but not supported for CFG++ samplers");
+                cache_runtime.spectrum_enabled = false;
+            }
+        }
+
         size_t steps       = sigmas.size() - 1;
         bool has_skiplayer = slg_scale != 0.0f && !skip_layers.empty();
         if (has_skiplayer && !sd_version_is_dit(version)) {
@@ -1633,7 +1719,7 @@ public:
         sd::Tensor<float> denoised   = x_t;
         SamplePreviewContext preview = prepare_sample_preview_context();
 
-        auto denoise = [&](const sd::Tensor<float>& x, float sigma, int step) -> sd::Tensor<float> {
+        auto denoise = [&](const sd::Tensor<float>& x, float sigma, int step, sd::Tensor<float>* out_uncond_denoised = nullptr) -> sd::Tensor<float> {
             if (step == 1 || step == -1) {
                 pretty_progress(0, (int)steps, 0);
             }
@@ -1656,6 +1742,7 @@ public:
             }
 
             if (cache_runtime.spectrum_enabled && cache_runtime.spectrum.should_predict()) {
+                if (out_uncond_denoised == nullptr) {
                     cache_runtime.spectrum.predict(&denoised);
                     if (!denoise_mask.empty()) {
                         denoised = denoised * denoise_mask + init_latent * (1.0f - denoise_mask);
@@ -1666,6 +1753,7 @@ public:
                     report_sample_progress(step, steps, t0);
                     return denoised;
                 }
+            }
 
             if (sd_should_preview_noisy() && preview.callback != nullptr) {
                 preview_image(step, noised_input, version, preview.mode, preview.callback, preview.data, true);
@@ -1706,9 +1794,8 @@ public:
                 diffusion_params.input_ids    = condition.c_input_ids.empty() ? nullptr : &condition.c_input_ids;
                 diffusion_params.input_pos    = condition.c_position_ids.empty() ? nullptr : &condition.c_position_ids;
                 diffusion_params.token_types  = condition.c_token_types.empty() ? nullptr : &condition.c_token_types;
-                diffusion_params.image_embed_ranges = condition.c_image_embed_ranges.empty() ? nullptr : &condition.c_image_embed_ranges;
                 diffusion_params.vinput_mask  = condition.c_vinput_mask.empty() ? nullptr : &condition.c_vinput_mask;
-                diffusion_params.vlm_images         = condition.c_vlm_images.empty() ? nullptr : &condition.c_vlm_images;
+                diffusion_params.image_embeds = condition.c_image_embeds.empty() ? nullptr : &condition.c_image_embeds;
                 diffusion_params.ref_latents  = condition.c_ref_images.empty() ? &ref_latents : &condition.c_ref_images;
                 diffusion_params.skip_layers  = local_skip_layers;
 
@@ -1794,6 +1881,10 @@ public:
                 latent_result += (cond_out - skip_cond_out) * slg_scale;
             }
             denoised = latent_result * c_out + x * c_skip;
+            if (out_uncond_denoised != nullptr) {
+                sd::Tensor<float> base_uncond = !uncond_out.empty() ? uncond_out : cond_out;
+                *out_uncond_denoised          = base_uncond * c_out + x * c_skip;
+            }
             if (cache_runtime.spectrum_enabled) {
                 cache_runtime.spectrum.update(denoised);
             }
@@ -1807,7 +1898,7 @@ public:
             return denoised;
         };
 
-        auto x0_opt = sample_k_diffusion(method, denoise, x_t, sigmas, sampler_rng, eta, is_flow_denoiser);
+        auto x0_opt = sample_k_diffusion(method, denoise, x_t, sigmas, sampler_rng, eta, is_flow_denoiser, extra_sample_args);
         if (x0_opt.empty()) {
             LOG_ERROR("Diffusion model sampling failed");
             if (control_net) {
@@ -1979,7 +2070,6 @@ enum rng_type_t str_to_rng_type(const char* str) {
 
 const char* sample_method_to_str[] = {
     "euler",
-    "euler_flow_flash",
     "euler_a",
     "heun",
     "dpm2",
@@ -1994,6 +2084,8 @@ const char* sample_method_to_str[] = {
     "res_multistep",
     "res_2s",
     "er_sde",
+    "euler_cfg_pp",
+    "euler_a_cfg_pp",
 };
 
 const char* sd_sample_method_name(enum sample_method_t sample_method) {
@@ -2299,6 +2391,7 @@ void sd_sample_params_init(sd_sample_params_t* sample_params) {
     sample_params->custom_sigmas               = nullptr;
     sample_params->custom_sigmas_count         = 0;
     sample_params->flow_shift                  = INFINITY;
+    sample_params->extra_sample_args           = nullptr;
 }
 
 char* sd_sample_params_to_str(const sd_sample_params_t* sample_params) {
@@ -2320,7 +2413,8 @@ char* sd_sample_params_to_str(const sd_sample_params_t* sample_params) {
              "sample_steps: %d, "
              "eta: %.2f, "
              "shifted_timestep: %d, "
-             "flow_shift: %.2f)",
+             "flow_shift: %.2f, "
+             "extra_sample_args: %s)",
              sample_params->guidance.txt_cfg,
              std::isfinite(sample_params->guidance.img_cfg)
                  ? sample_params->guidance.img_cfg
@@ -2335,7 +2429,8 @@ char* sd_sample_params_to_str(const sd_sample_params_t* sample_params) {
              sample_params->sample_steps,
              sample_params->eta,
              sample_params->shifted_timestep,
-             sample_params->flow_shift);
+             sample_params->flow_shift,
+             SAFE_STR(sample_params->extra_sample_args));
 
     return buf;
 }
@@ -2559,6 +2654,7 @@ static float resolve_eta(sd_ctx_t* sd_ctx,
             case EULER_A_SAMPLE_METHOD:
             case DPMPP2S_A_SAMPLE_METHOD:
             case ER_SDE_SAMPLE_METHOD:
+            case EULER_A_CFG_PP_SAMPLE_METHOD:
                 return 1.0f;
             default:;
         }
@@ -2768,6 +2864,8 @@ struct GenerationRequest {
 struct SamplePlan {
     enum sample_method_t sample_method            = SAMPLE_METHOD_COUNT;
     enum sample_method_t high_noise_sample_method = SAMPLE_METHOD_COUNT;
+    const char* extra_sample_args                 = nullptr;
+    const char* high_noise_extra_sample_args      = nullptr;
     float eta                                     = 0.f;
     float high_noise_eta                          = 0.f;
     int sample_steps                              = 0;
@@ -2781,6 +2879,7 @@ struct SamplePlan {
                const sd_img_gen_params_t* sd_img_gen_params,
                const GenerationRequest& request) {
         sample_method     = sd_img_gen_params->sample_params.sample_method;
+        extra_sample_args = sd_img_gen_params->sample_params.extra_sample_args;
         eta               = sd_img_gen_params->sample_params.eta;
         sample_steps      = sd_img_gen_params->sample_params.sample_steps;
         resolve(sd_ctx, &request, &sd_img_gen_params->sample_params);
@@ -2790,11 +2889,13 @@ struct SamplePlan {
                const sd_vid_gen_params_t* sd_vid_gen_params,
                const GenerationRequest& request) {
         sample_method     = sd_vid_gen_params->sample_params.sample_method;
+        extra_sample_args = sd_vid_gen_params->sample_params.extra_sample_args;
         eta               = sd_vid_gen_params->sample_params.eta;
         sample_steps      = sd_vid_gen_params->sample_params.sample_steps;
         if (sd_ctx->sd->high_noise_diffusion_model) {
             high_noise_sample_steps      = sd_vid_gen_params->high_noise_sample_params.sample_steps;
             high_noise_sample_method     = sd_vid_gen_params->high_noise_sample_params.sample_method;
+            high_noise_extra_sample_args = sd_vid_gen_params->high_noise_sample_params.extra_sample_args;
             high_noise_eta               = sd_vid_gen_params->high_noise_sample_params.eta;
         }
         moe_boundary = sd_vid_gen_params->moe_boundary;
@@ -3454,6 +3555,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
                                                    request.shifted_timestep,
                                                    plan.sample_method,
                                                    sd_ctx->sd->is_flow_denoiser(),
+                                                   plan.extra_sample_args,
                                                    plan.sigmas,
                                                    plan.start_merge_step,
                                                    latents.ref_latents,
@@ -3579,6 +3681,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
                                                             request.shifted_timestep,
                                                             plan.sample_method,
                                                             sd_ctx->sd->is_flow_denoiser(),
+                                                            plan.extra_sample_args,
                                                             hires_sigma_sched,
                                                             plan.start_merge_step,
                                                             latents.ref_latents,
@@ -3943,6 +4046,7 @@ SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* s
                                                            request.shifted_timestep,
                                                            plan.high_noise_sample_method,
                                                            sd_ctx->sd->is_flow_denoiser(),
+                                                           plan.high_noise_extra_sample_args,
                                                            high_noise_sigmas,
                                                            -1,
                                                            std::vector<sd::Tensor<float>>{},
@@ -3985,6 +4089,7 @@ SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* s
                                                         sd_vid_gen_params->sample_params.shifted_timestep,
                                                         plan.sample_method,
                                                         sd_ctx->sd->is_flow_denoiser(),
+                                                        plan.extra_sample_args,
                                                         plan.sigmas,
                                                         -1,
                                                         std::vector<sd::Tensor<float>>{},

src/util.cpp

@@ -112,7 +112,7 @@ private:
     HANDLE hmapping_;
 };
 
-std::unique_ptr<MmapWrapper> MmapWrapper::create(const std::string& filename) {
+std::unique_ptr<MmapWrapper> MmapWrapper::create(const std::string& filename, bool writable) {
     void* mapped_data = nullptr;
     size_t file_size  = 0;
 
@@ -137,14 +137,18 @@ std::unique_ptr<MmapWrapper> MmapWrapper::create(const std::string& filename) {
 
     file_size = static_cast<size_t>(size.QuadPart);
 
-    HANDLE mapping_handle = CreateFileMapping(file_handle, nullptr, PAGE_READONLY, 0, 0, nullptr);
+    DWORD page_prot = writable ? PAGE_WRITECOPY : PAGE_READONLY;
+
+    HANDLE mapping_handle = CreateFileMapping(file_handle, nullptr, page_prot, 0, 0, nullptr);
 
     if (mapping_handle == nullptr) {
         CloseHandle(file_handle);
         return nullptr;
     }
 
-    mapped_data = MapViewOfFile(mapping_handle, FILE_MAP_READ, 0, 0, file_size);
+    DWORD view_access = writable ? FILE_MAP_COPY : FILE_MAP_READ;
+
+    mapped_data = MapViewOfFile(mapping_handle, view_access, 0, 0, file_size);
 
     if (mapped_data == nullptr) {
         CloseHandle(mapping_handle);
@@ -172,28 +176,85 @@ bool is_directory(const std::string& path) {
     return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode));
 }
 
-class MmapWrapperImpl : public MmapWrapper {
+struct MmapFlags {
-public:
+    bool sequential;
-    MmapWrapperImpl(void* data, size_t size)
+    bool populate;
-        : MmapWrapper(data, size) {}
+    bool willneed;
-
+    bool dontneed;
-    ~MmapWrapperImpl() override {
-        munmap(data_, size_);
-    }
 };
 
-std::unique_ptr<MmapWrapper> MmapWrapper::create(const std::string& filename) {
+static MmapFlags get_mmap_flags() {
+    MmapFlags result          = {};
+    const char* SD_MMAP_FLAGS = std::getenv("SD_MMAP_FLAGS");
+    if (SD_MMAP_FLAGS && *SD_MMAP_FLAGS) {
+        std::stringstream ss(SD_MMAP_FLAGS);
+        std::string token;
+        while (std::getline(ss, token, ',')) {
+            std::string ntoken = trim(token);
+            std::transform(ntoken.begin(), ntoken.end(), ntoken.begin(), ::tolower);
+            if (ntoken == "sequential") {
+                result.sequential = true;
+            } else if (ntoken == "populate") {
+                result.populate = true;
+            } else if (ntoken == "willneed") {
+                result.willneed = true;
+            } else if (ntoken == "dontneed") {
+                result.dontneed = true;
+            }
+        }
+    }
+    return result;
+}
+
+class MmapWrapperImpl : public MmapWrapper {
+public:
+    MmapWrapperImpl(void* data, size_t size, int fd)
+        : MmapWrapper(data, size), fd_(fd) {}
+
+    ~MmapWrapperImpl() override {
+#ifdef __linux__
+        auto cfg_flags = get_mmap_flags();
+
+        // Drop the kernel pagecache pages for this file. madvise(DONTNEED)
+        // alone only unmaps from the process address space; pagecache
+        // entries persist (`free` reports them as buff/cache and the OOM
+        // killer doesn't touch them, but they ARE counted against
+        // overcommit and can starve other allocations on tight-RAM
+        // systems). posix_fadvise(POSIX_FADV_DONTNEED) is the documented
+        // way to evict pagecache for a specific fd's pages.
+        if (cfg_flags.dontneed) {
+            madvise(data_, size_, MADV_DONTNEED);
+            posix_fadvise(fd_, 0, 0, POSIX_FADV_DONTNEED);
+        }
+#endif
+        munmap(data_, size_);
+        close(fd_);
+    }
+
+private:
+    int fd_;
+};
+
+std::unique_ptr<MmapWrapper> MmapWrapper::create(const std::string& filename, bool writable) {
     int file_descriptor = open(filename.c_str(), O_RDONLY);
     if (file_descriptor == -1) {
         return nullptr;
     }
 
+    auto cfg_flags = get_mmap_flags();
+
     int mmap_flags = MAP_PRIVATE;
 
 #ifdef __linux__
-    // performance flags used by llama.cpp
+    // Sequential access hint helps the kernel read-ahead efficiently and
-    // posix_fadvise(file_descriptor, 0, 0, POSIX_FADV_SEQUENTIAL);
+    // also encourages eviction of already-read pages (the kernel keeps
-    // mmap_flags |= MAP_POPULATE;
+    // a smaller working set when this is set).
+    if (cfg_flags.sequential) {
+        posix_fadvise(file_descriptor, 0, 0, POSIX_FADV_SEQUENTIAL);
+    }
+    if (cfg_flags.populate) {
+        mmap_flags |= MAP_POPULATE;
+    }
 #endif
 
     struct stat sb;
@@ -204,20 +265,27 @@ std::unique_ptr<MmapWrapper> MmapWrapper::create(const std::string& filename) {
 
     size_t file_size = sb.st_size;
 
-    void* mapped_data = mmap(nullptr, file_size, PROT_READ, mmap_flags, file_descriptor, 0);
+    if (file_size == 0) {
-
         close(file_descriptor);
+        return nullptr;
+    }
+
+    int mmap_prot = PROT_READ | (writable ? PROT_WRITE : 0);
+
+    void* mapped_data = mmap(nullptr, file_size, mmap_prot, mmap_flags, file_descriptor, 0);
 
     if (mapped_data == MAP_FAILED) {
+        close(file_descriptor);
         return nullptr;
     }
 
 #ifdef __linux__
-    // performance flags used by llama.cpp
+    if (cfg_flags.willneed) {
-    // posix_madvise(mapped_data, file_size, POSIX_MADV_WILLNEED);
+        posix_madvise(mapped_data, file_size, POSIX_MADV_WILLNEED);
+    }
 #endif
 
-    return std::make_unique<MmapWrapperImpl>(mapped_data, file_size);
+    return std::make_unique<MmapWrapperImpl>(mapped_data, file_size, file_descriptor);
 }
 
 #endif

src/util.h

@@ -42,7 +42,7 @@ sd::Tensor<float> clip_preprocess(const sd::Tensor<float>& image, int target_wid
 
 class MmapWrapper {
 public:
-    static std::unique_ptr<MmapWrapper> create(const std::string& filename);
+    static std::unique_ptr<MmapWrapper> create(const std::string& filename, bool writable = false);
 
     virtual ~MmapWrapper() = default;
 
@@ -52,6 +52,7 @@ public:
     MmapWrapper& operator=(MmapWrapper&&)      = delete;
 
     const uint8_t* data() const { return static_cast<uint8_t*>(data_); }
+    uint8_t* writable_data() { return static_cast<uint8_t*>(data_); }
     size_t size() const { return size_; }
     bool copy_data(void* buf, size_t n, size_t offset) const;