feat: add DPM++ 2M SDE sampler (#1742)

docs/backend.md

@@ -51,6 +51,97 @@ Module names are case-insensitive. Hyphens and underscores in module names are i
 sd-cli -m model.safetensors -p "a cat" --backend all=cuda0,te=cpu
 ```
 
+## Multiple devices per module (layer split)
+
+A `--backend` module assignment can list several devices separated by `&`:
+
+```shell
+sd-cli -m model.safetensors -p "a cat" --backend "diffusion=cuda0&cuda1"
+```
+
+The module's transformer blocks are then distributed across the listed devices
+in contiguous ranges sized proportionally to each device's free memory (minus a
+compute-buffer headroom of about 2 GiB per device), and the
+module's graphs are executed with a `ggml_backend_sched` that runs each block
+on the device holding its weights, copying the residual stream at the range
+boundaries. The first device in the list is the module's main device: it also
+holds the non-block tensors (embeddings, final norms, small sub-runners such as
+CLIP models or projectors) and the graph inputs/outputs.
+
+Layer split is supported for the `diffusion` and `te` modules. For `te` it
+applies to the dominant text encoder (`t5xxl` or the LLM); other modules accept
+only a single device. If the module has no recognizable transformer blocks, the
+assignment falls back to the first listed device.
+
+`--params-backend` accepts no device lists. If the module has no explicit
+params assignment, each block range's parameters are loaded directly to (and,
+with `--params-backend diffusion=disk`, released directly from) its own device;
+an explicit assignment such as `te=cpu` keeps the parameters on that backend
+and stages each range to its device on demand.
+
+Layer split cannot be combined with `--max-vram` graph-cut segmentation or
+`--stream-layers` for the split module; those are single-device mechanisms and
+are disabled for it.
+
+Use `--list-devices` to see the device names available on the system.
+
+### Row split (`--split-mode row`)
+
+`--split-mode` selects how a multi-device module distributes its weights:
+`layer` (the default, described above) or `row`. It accepts a single mode or
+per-module assignments:
+
+```shell
+sd-cli -m model.safetensors -p "a cat" --backend "diffusion=cuda0&cuda1" --split-mode row
+sd-cli -m model.safetensors -p "a cat" --backend "diffusion=cuda0&cuda1,te=cuda0&cuda1" --split-mode diffusion=row,te=layer
+```
+
+In row mode the module keeps executing on its main (first listed) device, but
+its transformer-block matmul weights are allocated in the backend's row-split
+buffer type, which slices each weight's rows across the listed devices in
+proportion to free memory and runs those matmuls on all devices in parallel.
+Compared to a layer split this uses all GPUs within every layer (instead of
+sequentially device by device) at the cost of a cross-device reduction per
+matmul - usually the faster option when the devices have fast interconnect.
+
+Row split requires backend support for split buffers and is currently
+available on CUDA only; on other backends (or when the listed devices belong
+to different backend registries) the module falls back to a layer split.
+Embeddings, normalization weights, biases and other non-block tensors stay in
+regular buffers on the main device.
+
+Direct ("immediately") LoRA application cannot patch row-split tensors; with
+`--split-mode row` the automatic LoRA mode selects runtime application, and an
+explicit `--lora-apply-mode immediately` skips the split tensors with a
+warning.
+
+## Automatic placement (`--auto-fit`)
+
+`--auto-fit` derives the `diffusion` / `te` / `vae` placements from the model
+metadata and the per-device memory budgets, then feeds them into the same
+backend assignment mechanism described above (the chosen specs are printed).
+`--backend` and `--params-backend` are ignored while auto-fit is enabled.
+
+```shell
+sd-cli -m model.safetensors -p "a cat" --auto-fit
+sd-cli -m model.safetensors -p "a cat" --auto-fit --max-vram cuda0=8,cuda1=14
+sd-cli -m model.safetensors -p "a cat" --auto-fit --split-mode row
+```
+
+Budgets reuse `--max-vram`: a positive per-device value caps what auto-fit
+plans with on that device, a negative value means "free memory minus that many
+GiB", and with no budget set each device's free memory minus a 512 MiB margin
+is used. (The same values still drive graph-cut segmented execution for
+modules that end up on a single device.)
+
+When everything fits resident, components are simply spread across the
+available GPUs. When it does not, auto-fit switches to time-share mode: the
+heavy components get `disk` params residency (loaded for their phase, freed
+after), and a component too large for any single device is split across all
+GPUs with the layer/row split mechanism (`--split-mode` selects which, layer
+by default). Components that fit nowhere fall back to the CPU. If a VAE decode
+still runs out of memory, tiling is enabled and the decode retried once.
+
 ## Modules
 
 | Module | Purpose | Accepted names |

examples/cli/main.cpp

@@ -481,7 +481,8 @@ bool save_results(const SDCliParams& cli_params,
         if (!img.data)
             return false;
 
-        const int64_t metadata_seed = cli_params.mode == VID_GEN ? gen_params.seed : gen_params.seed + idx;
+        int images_per_batch        = gen_params.batch_count > 0 ? std::max(1, num_results / gen_params.batch_count) : 1;
+        const int64_t metadata_seed = cli_params.mode == VID_GEN ? gen_params.seed : gen_params.seed + idx / images_per_batch;
         std::string params          = gen_params.embed_image_metadata
                                           ? get_image_params(ctx_params, gen_params, metadata_seed, cli_params.mode)
                                           : "";

examples/common/common.cpp

@@ -468,6 +468,13 @@ ArgOptions SDContextParams::get_options() {
          "parameter backend assignment, e.g. disk, cpu, or diffusion=disk,clip=cpu",
          (int)',',
          &params_backend},
+        {"",
+         "--split-mode",
+         "weight distribution for modules assigned multiple devices (--backend \"diffusion=cuda0&cuda1\"): "
+         "layer (whole transformer blocks per device, default) or row (matmul rows split across devices, CUDA only). "
+         "Accepts a single mode or per-module assignments, e.g. row or diffusion=row,te=layer",
+         (int)',',
+         &split_mode},
         {"",
          "--rpc-servers",
          "comma-separated list of RPC servers to connect to for offloading, in the format host:port, e.g. localhost:50052,192.168.1.3:50052",
@@ -501,6 +508,12 @@ ArgOptions SDContextParams::get_options() {
          "--eager-load",
          "load all params into the params backend at model-load time instead of lazily on first use (defaults to false)",
          true, &eager_load},
+        {"",
+         "--auto-fit",
+         "pick the diffusion/te/vae device placements automatically from the model size and the per-device "
+         "memory budgets (--max-vram; defaults to free memory minus a small margin). Overrides --backend and "
+         "--params-backend; may split modules across GPUs (--split-mode still selects layer or row)",
+         true, &auto_fit},
         {"",
          "--force-sdxl-vae-conv-scale",
          "force use of conv scale on sdxl vae",
@@ -663,6 +676,18 @@ ArgOptions SDContextParams::get_options() {
          "but it usually offers faster inference speed and, in some cases, lower memory usage. "
          "The at_runtime mode, on the other hand, is exactly the opposite.",
          on_lora_apply_mode_arg},
+        {"",
+         "--list-devices",
+         "list available ggml backend devices (one 'name<TAB>description' per line) and exit; "
+         "the names are the device names accepted by --backend and --params-backend",
+         [](int /*argc*/, const char** /*argv*/, int /*index*/) {
+             size_t device_list_size = sd_list_devices(nullptr, 0);
+             std::vector<char> devices(device_list_size + 1);
+             sd_list_devices(devices.data(), devices.size());
+             fputs(devices.data(), stdout);
+             std::exit(0);
+             return 0;
+         }},
     };
 
     return options;
@@ -818,6 +843,8 @@ std::string SDContextParams::to_string() const {
         << "  eager_load: " << (eager_load ? "true" : "false") << ",\n"
         << "  backend: \"" << backend << "\",\n"
         << "  params_backend: \"" << params_backend << "\",\n"
+        << "  split_mode: \"" << split_mode << "\",\n"
+        << "  auto_fit: " << (auto_fit ? "true" : "false") << ",\n"
         << "  enable_mmap: " << (enable_mmap ? "true" : "false") << ",\n"
         << "  control_net_cpu: " << (control_net_cpu ? "true" : "false") << ",\n"
         << "  clip_on_cpu: " << (clip_on_cpu ? "true" : "false") << ",\n"
@@ -898,6 +925,8 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) {
     sd_ctx_params.eager_load                      = eager_load;
     sd_ctx_params.backend                         = effective_backend.c_str();
     sd_ctx_params.params_backend                  = effective_params_backend.c_str();
+    sd_ctx_params.split_mode                      = split_mode.c_str();
+    sd_ctx_params.auto_fit                        = auto_fit;
     sd_ctx_params.rpc_servers                     = rpc_servers.c_str();
     return sd_ctx_params;
 }
@@ -1007,6 +1036,10 @@ ArgOptions SDGenerationParams::get_options() {
          "--batch-count",
          "batch count",
          &batch_count},
+        {"",
+         "--qwen-image-layers",
+         "number of Qwen Image Layered layers; latent/output count is layers + 1 (default: 3)",
+         &qwen_image_layers},
         {"",
          "--video-frames",
          "video frames (default: 1)",
@@ -1073,7 +1106,7 @@ ArgOptions SDGenerationParams::get_options() {
          &sample_params.guidance.slg.layer_end},
         {"",
          "--eta",
-         "noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a, er_sde and dpm++2s_a)",
+         "noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a, er_sde, dpm++2s_a and dpm++2m_sde)",
          &sample_params.eta},
         {"",
          "--flow-shift",
@@ -1105,7 +1138,7 @@ ArgOptions SDGenerationParams::get_options() {
          &high_noise_sample_params.guidance.slg.layer_end},
         {"",
          "--high-noise-eta",
-         "(high noise) noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a, er_sde and dpm++2s_a)",
+         "(high noise) noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and res_2s; 1 for euler_a, er_sde, dpm++2s_a and dpm++2m_sde)",
          &high_noise_sample_params.eta},
         {"",
          "--strength",
@@ -1476,12 +1509,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, euler_cfg_pp, euler_a_cfg_pp]"
+         "sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, dpm++2m_sde, 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, euler_cfg_pp, euler_a_cfg_pp]"
+         "(high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, dpm++2m_sde, 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},
         {"",
@@ -1827,6 +1860,7 @@ bool SDGenerationParams::from_json_str(
     load_if_exists("width", width);
     load_if_exists("height", height);
     load_if_exists("batch_count", batch_count);
+    load_if_exists("qwen_image_layers", qwen_image_layers);
     load_if_exists("video_frames", video_frames);
     load_if_exists("fps", fps);
     load_if_exists("upscale_repeats", upscale_repeats);
@@ -2251,6 +2285,11 @@ bool SDGenerationParams::validate(SDMode mode) {
         return false;
     }
 
+    if (qwen_image_layers < 0) {
+        LOG_ERROR("error: qwen_image_layers must be non-negative");
+        return false;
+    }
+
     if (sample_params.sample_steps <= 0) {
         LOG_ERROR("error: the sample_steps must be greater than 0\n");
         return false;
@@ -2417,6 +2456,7 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
     params.strength              = strength;
     params.seed                  = seed;
     params.batch_count           = batch_count;
+    params.qwen_image_layers     = qwen_image_layers;
     params.control_image         = control_image.get();
     params.control_strength      = control_strength;
     params.pm_params             = pm_params;
@@ -2542,6 +2582,7 @@ std::string SDGenerationParams::to_string() const {
         << "  width: " << width << ",\n"
         << "  height: " << height << ",\n"
         << "  batch_count: " << batch_count << ",\n"
+        << "  qwen_image_layers: " << qwen_image_layers << ",\n"
         << "  init_image_path: \"" << init_image_path << "\",\n"
         << "  end_image_path: \"" << end_image_path << "\",\n"
         << "  mask_image_path: \"" << mask_image_path << "\",\n"

examples/common/common.h

@@ -151,6 +151,8 @@ struct SDContextParams {
     bool eager_load             = false;
     std::string backend;
     std::string params_backend;
+    std::string split_mode;
+    bool auto_fit = false;
     std::string rpc_servers;
     std::string effective_backend;
     std::string effective_params_backend;
@@ -197,6 +199,7 @@ struct SDGenerationParams {
     int width                  = -1;
     int height                 = -1;
     int batch_count            = 1;
+    int qwen_image_layers      = 3;
     int64_t seed               = 42;
     float strength             = 0.75f;
     float control_strength     = 0.9f;

examples/server/async_jobs.cpp

@@ -2,6 +2,7 @@
 
 #include "async_jobs.h"
 
+#include <algorithm>
 #include <iomanip>
 #include <sstream>
 
@@ -195,6 +196,8 @@ bool execute_img_gen_job(ServerRuntime& runtime,
         encoded_format = EncodedImageFormat::WEBP;
     }
 
+    int batch_count      = job.img_gen.gen_params.batch_count;
+    int images_per_batch = batch_count > 0 ? std::max(1, num_results / batch_count) : 1;
     for (int i = 0; i < num_results; ++i) {
         if (results[i].data == nullptr) {
             continue;
@@ -203,7 +206,7 @@ bool execute_img_gen_job(ServerRuntime& runtime,
         const std::string metadata = job.img_gen.gen_params.embed_image_metadata
                                          ? get_image_params(*runtime.ctx_params,
                                                             job.img_gen.gen_params,
-                                                            job.img_gen.gen_params.seed + i)
+                                                            job.img_gen.gen_params.seed + i / images_per_batch)
                                          : "";
         auto image_bytes           = encode_image_to_vector(encoded_format,
                                                             results[i].data,

examples/server/routes_openai.cpp

@@ -284,14 +284,16 @@ void register_openai_api_endpoints(httplib::Server& svr, ServerRuntime& rt) {
             out["data"]          = json::array();
             out["output_format"] = request.output_format;
 
-            for (int i = 0; i < request.gen_params.batch_count; ++i) {
+            int result_count     = results.count();
+            int images_per_batch = request.gen_params.batch_count > 0 ? std::max(1, result_count / request.gen_params.batch_count) : 1;
+            for (int i = 0; i < result_count; ++i) {
                 if (results[i].data == nullptr) {
                     continue;
                 }
                 std::string params = request.gen_params.embed_image_metadata
                                          ? get_image_params(*runtime->ctx_params,
                                                             request.gen_params,
-                                                            request.gen_params.seed + i)
+                                                            request.gen_params.seed + i / images_per_batch)
                                          : "";
                 auto image_bytes   = encode_image_to_vector(request.output_format == "jpeg"
                                                                 ? EncodedImageFormat::JPEG
@@ -356,14 +358,16 @@ void register_openai_api_endpoints(httplib::Server& svr, ServerRuntime& rt) {
             out["data"]          = json::array();
             out["output_format"] = request.output_format;
 
-            for (int i = 0; i < request.gen_params.batch_count; ++i) {
+            int result_count     = results.count();
+            int images_per_batch = request.gen_params.batch_count > 0 ? std::max(1, result_count / request.gen_params.batch_count) : 1;
+            for (int i = 0; i < result_count; ++i) {
                 if (results[i].data == nullptr) {
                     continue;
                 }
                 std::string params = request.gen_params.embed_image_metadata
                                          ? get_image_params(*runtime->ctx_params,
                                                             request.gen_params,
-                                                            request.gen_params.seed + i)
+                                                            request.gen_params.seed + i / images_per_batch)
                                          : "";
                 auto image_bytes   = encode_image_to_vector(request.output_format == "jpeg" ? EncodedImageFormat::JPEG : EncodedImageFormat::PNG,
                                                           results[i].data,

examples/server/routes_sdapi.cpp

@@ -63,6 +63,8 @@ static enum sample_method_t get_sdapi_sample_method(std::string name) {
         {"ddim", DDIM_TRAILING_SAMPLE_METHOD},
         {"dpm++ 2m", DPMPP2M_SAMPLE_METHOD},
         {"k_dpmpp_2m", DPMPP2M_SAMPLE_METHOD},
+        {"dpm++ 2m sde", DPMPP2M_SDE_SAMPLE_METHOD},
+        {"k_dpmpp_2m_sde", DPMPP2M_SDE_SAMPLE_METHOD},
         {"res multistep", RES_MULTISTEP_SAMPLE_METHOD},
         {"k_res_multistep", RES_MULTISTEP_SAMPLE_METHOD},
         {"res 2s", RES_2S_SAMPLE_METHOD},
@@ -259,6 +261,48 @@ static bool build_sdapi_img_gen_request(const json& j,
     return true;
 }
 
+static nlohmann::json prepare_info_field(const SDContextParams& ctx_params,
+                                         const SDGenerationParams& gen_params,
+                                         bool img2img) {
+    nlohmann::json jsoninfo = nlohmann::json::object();
+    jsoninfo["prompt"]      = gen_params.prompt;
+    if (!gen_params.negative_prompt.empty()) {
+        jsoninfo["negative_prompt"] = gen_params.negative_prompt;
+    }
+    jsoninfo["seed"]         = gen_params.seed;
+    jsoninfo["cfg_scale"]    = gen_params.sample_params.guidance.txt_cfg;
+    jsoninfo["width"]        = gen_params.get_resolved_width();
+    jsoninfo["height"]       = gen_params.get_resolved_height();
+    jsoninfo["steps"]        = gen_params.sample_params.sample_steps;
+    jsoninfo["sampler_name"] = sd_sample_method_name(gen_params.sample_params.sample_method);
+    if (gen_params.clip_skip != -1) {
+        jsoninfo["clip_skip"] = gen_params.clip_skip;
+    }
+    if (gen_params.sample_params.scheduler != scheduler_t::SCHEDULER_COUNT) {
+        jsoninfo["extra_generation_params"] = nlohmann::json::object();
+        jsoninfo["extra_generation_params"]["Schedule type"] = sd_scheduler_name(gen_params.sample_params.scheduler);
+    }
+    if (img2img) {
+        jsoninfo["denoising_strength"] = gen_params.strength;
+    }
+    // not clear what should happen if we have both model and diffusion_model
+    if (!ctx_params.diffusion_model_path.empty()) {
+        jsoninfo["sd_model_name"] = sd_basename(ctx_params.diffusion_model_path);
+    } else if (!ctx_params.model_path.empty()) {
+        jsoninfo["sd_model_name"] = sd_basename(ctx_params.model_path);
+    }
+    if (!ctx_params.vae_path.empty()) {
+        jsoninfo["sd_vae_name"] = sd_basename(ctx_params.vae_path);
+    }
+    jsoninfo["version"] = "stable-diffusion.cpp";
+
+    jsoninfo["infotexts"]            = nlohmann::json::array();
+    jsoninfo["all_prompts"]          = nlohmann::json::array();
+    jsoninfo["all_negative_prompts"] = nlohmann::json::array();
+    jsoninfo["all_seeds"]            = nlohmann::json::array();
+    return jsoninfo;
+}
+
 void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) {
     ServerRuntime* runtime = &rt;
 
@@ -309,24 +353,26 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) {
             json out;
             out["images"]     = json::array();
             out["parameters"] = j;
-            out["info"]       = "";
+            json jsoninfo     = prepare_info_field(*runtime->ctx_params, request.gen_params, img2img);
 
+            int images_per_batch = request.gen_params.batch_count > 0 ? std::max(1, num_results / request.gen_params.batch_count) : 1;
             for (int i = 0; i < num_results; ++i) {
                 if (results[i].data == nullptr) {
                     continue;
                 }
 
-                std::string params = request.gen_params.embed_image_metadata
-                                         ? get_image_params(*runtime->ctx_params,
-                                                            request.gen_params,
-                                                            request.gen_params.seed + i)
-                                         : "";
+                bool embed_meta    = request.gen_params.embed_image_metadata;
+
+                std::string params = get_image_params(*runtime->ctx_params,
+                                                      request.gen_params,
+                                                      request.gen_params.seed + i / images_per_batch);
+
                 auto image_bytes   = encode_image_to_vector(EncodedImageFormat::PNG,
                                                             results[i].data,
                                                             results[i].width,
                                                             results[i].height,
                                                             results[i].channel,
-                                                            params);
+                                                            embed_meta ? params : "");
 
                 if (image_bytes.empty()) {
                     LOG_ERROR("write image to mem failed");
@@ -334,8 +380,16 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) {
                 }
 
                 out["images"].push_back(base64_encode(image_bytes));
+
+                jsoninfo["infotexts"][i]            = params;
+                jsoninfo["all_seeds"][i]            = request.gen_params.seed + i;
+                jsoninfo["all_prompts"][i]          = request.gen_params.prompt;
+                jsoninfo["all_negative_prompts"][i] = request.gen_params.negative_prompt;
             }
 
+            // not a mistake: it is supposed to be a string in json format
+            out["info"] = jsoninfo.dump();
+
             res.set_content(out.dump(), "application/json");
             res.status = 200;
 

examples/server/routes_sdcpp.cpp

@@ -126,6 +126,7 @@ static json make_img_gen_defaults_json(const SDGenerationParams& defaults, const
         {"strength", defaults.strength},
         {"seed", defaults.seed},
         {"batch_count", defaults.batch_count},
+        {"qwen_image_layers", defaults.qwen_image_layers},
         {"auto_resize_ref_image", defaults.auto_resize_ref_image},
         {"increase_ref_index", defaults.increase_ref_index},
         {"control_strength", defaults.control_strength},

include/stable-diffusion.h

@@ -54,6 +54,7 @@ enum sample_method_t {
     EULER_CFG_PP_SAMPLE_METHOD,
     EULER_A_CFG_PP_SAMPLE_METHOD,
     EULER_GE_SAMPLE_METHOD,
+    DPMPP2M_SDE_SAMPLE_METHOD,
     SAMPLE_METHOD_COUNT
 };
 
@@ -227,6 +228,8 @@ typedef struct {
     bool eager_load;  // Load all params into the params backend at model-load time instead of lazily on first use
     const char* backend;
     const char* params_backend;
+    const char* split_mode;  // weight distribution for multi-device modules: layer (default) or row, or per-module assignments e.g. "diffusion=row"
+    bool auto_fit;
     const char* rpc_servers;
 } sd_ctx_params_t;
 
@@ -380,6 +383,7 @@ typedef struct {
     sd_tiling_params_t vae_tiling_params;
     sd_cache_params_t cache;
     sd_hires_params_t hires;
+    int qwen_image_layers;
 } sd_img_gen_params_t;
 
 typedef struct {
@@ -533,6 +537,12 @@ SD_API void disable_imatrix_collection(void);
 SD_API const char* sd_commit(void);
 SD_API const char* sd_version(void);
 
+// List available ggml backend devices, one `name<TAB>description` per line.
+// The names are the device names accepted by the --backend / --params-backend
+// assignment specs. Returns the number of bytes required, excluding the null
+// terminator. Passing nullptr or buffer_size 0 only queries the required size.
+SD_API size_t sd_list_devices(char* buffer, size_t buffer_size);
+
 // for C API, caller needs to call free_sd_images to free the memory after use
 // This helps avoid CRT problems on Windows when memory is allocated in the library but freed in the caller, which may use a different CRT.
 SD_API void free_sd_images(sd_image_t* result_images, int num_images);

src/conditioning/conditioner.hpp

@@ -116,6 +116,8 @@ public:
     virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors)           = 0;
     virtual void set_max_graph_vram_bytes(size_t max_vram_bytes) {}
     virtual void set_stream_layers_enabled(bool enabled) {}
+    virtual void set_runtime_backends(const std::vector<ggml_backend_t>& backends) {}
+    virtual void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {}
     virtual void set_flash_attention_enabled(bool enabled) = 0;
     virtual void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) {}
     virtual void runner_done() {}
@@ -635,6 +637,18 @@ struct SD3CLIPEmbedder : public Conditioner {
         }
     }
 
+    void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
+        if (t5) {
+            t5->set_runtime_backends(backends);
+        }
+    }
+
+    void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        if (t5) {
+            t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
+        }
+    }
+
     void set_flash_attention_enabled(bool enabled) override {
         if (clip_l) {
             clip_l->set_flash_attention_enabled(enabled);
@@ -994,6 +1008,18 @@ struct FluxCLIPEmbedder : public Conditioner {
         }
     }
 
+    void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
+        if (t5) {
+            t5->set_runtime_backends(backends);
+        }
+    }
+
+    void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        if (t5) {
+            t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
+        }
+    }
+
     void set_flash_attention_enabled(bool enabled) override {
         if (clip_l) {
             clip_l->set_flash_attention_enabled(enabled);
@@ -1226,6 +1252,18 @@ struct T5CLIPEmbedder : public Conditioner {
         }
     }
 
+    void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
+        if (t5) {
+            t5->set_runtime_backends(backends);
+        }
+    }
+
+    void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        if (t5) {
+            t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
+        }
+    }
+
     void set_flash_attention_enabled(bool enabled) override {
         if (t5) {
             t5->set_flash_attention_enabled(enabled);
@@ -1418,6 +1456,18 @@ struct MiniT2IConditioner : public Conditioner {
         }
     }
 
+    void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
+        if (t5) {
+            t5->set_runtime_backends(backends);
+        }
+    }
+
+    void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        if (t5) {
+            t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
+        }
+    }
+
     void set_flash_attention_enabled(bool enabled) override {
         if (t5) {
             t5->set_flash_attention_enabled(enabled);
@@ -1502,6 +1552,14 @@ struct AnimaConditioner : public Conditioner {
         llm->set_stream_layers_enabled(enabled);
     }
 
+    void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
+        llm->set_runtime_backends(backends);
+    }
+
+    void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        llm->get_param_tensors(tensors, "text_encoders.llm");
+    }
+
     void set_flash_attention_enabled(bool enabled) override {
         llm->set_flash_attention_enabled(enabled);
     }
@@ -1647,6 +1705,14 @@ struct LLMEmbedder : public Conditioner {
         llm->set_stream_layers_enabled(enabled);
     }
 
+    void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
+        llm->set_runtime_backends(backends);
+    }
+
+    void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        llm->get_param_tensors(tensors, "text_encoders.llm");
+    }
+
     void set_flash_attention_enabled(bool enabled) override {
         llm->set_flash_attention_enabled(enabled);
     }
@@ -2316,6 +2382,14 @@ struct LTXAVEmbedder : public Conditioner {
         projector->set_max_graph_vram_bytes(max_vram_bytes);
     }
 
+    void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
+        llm->set_runtime_backends(backends);
+    }
+
+    void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        llm->get_param_tensors(tensors, "text_encoders.llm");
+    }
+
     void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
         llm->set_weight_adapter(adapter);
         projector->set_weight_adapter(adapter);

src/core/backend_fit.cpp

@@ -0,0 +1,390 @@
+#include "backend_fit.h"
+
+#include <algorithm>
+#include <cctype>
+#include <cstdint>
+#include <utility>
+#include <vector>
+
+#include "core/ggml_extend_backend.h"
+#include "core/util.h"
+#include "ggml-backend.h"
+
+namespace sd::backend_fit {
+    namespace {
+
+        constexpr int64_t MiB = 1024ll * 1024;
+
+        enum class ComponentKind {
+            DIT         = 0,
+            VAE         = 1,
+            CONDITIONER = 2,
+        };
+
+        struct Component {
+            ComponentKind kind;
+            const char* name;
+            int64_t params_bytes  = 0;
+            int64_t reserve_bytes = 0;
+            bool splittable       = false;
+        };
+
+        struct Device {
+            ggml_backend_dev_t dev = nullptr;
+            std::string name;
+            std::string description;
+            int64_t free_bytes   = 0;
+            int64_t total_bytes  = 0;
+            int64_t budget_bytes = 0;
+        };
+
+        struct Decision {
+            ComponentKind kind;
+            bool on_cpu = false;
+            std::vector<size_t> device_idxs;
+        };
+
+        struct Plan {
+            bool valid      = false;
+            bool time_share = false;
+            std::vector<Decision> decisions;
+        };
+
+        bool classify_tensor(const std::string& name, ComponentKind& out) {
+            auto contains = [&](const char* s) { return name.find(s) != std::string::npos; };
+
+            if (contains("model.diffusion_model.") || contains("unet.")) {
+                out = ComponentKind::DIT;
+                return true;
+            }
+            if (contains("first_stage_model.") ||
+                name.rfind("vae.", 0) == 0 ||
+                name.rfind("tae.", 0) == 0) {
+                out = ComponentKind::VAE;
+                return true;
+            }
+            if (contains("text_encoders") ||
+                contains("cond_stage_model") ||
+                contains("te.text_model.") ||
+                contains("conditioner") ||
+                name.rfind("text_encoder.", 0) == 0 ||
+                name.rfind("text_embedding_projection.", 0) == 0 ||
+                contains(".aggregate_embed.")) {
+                out = ComponentKind::CONDITIONER;
+                return true;
+            }
+            return false;
+        }
+
+        std::vector<Component> estimate_components(ModelLoader& loader, ggml_type override_wtype) {
+            const auto& storage = loader.get_tensor_storage_map();
+
+            int64_t bytes[3] = {0, 0, 0};
+            for (const auto& [name, ts_const] : storage) {
+                TensorStorage ts = ts_const;
+                if (is_unused_tensor(ts.name)) {
+                    continue;
+                }
+                ComponentKind kind;
+                if (!classify_tensor(ts.name, kind)) {
+                    continue;
+                }
+                if (override_wtype != GGML_TYPE_COUNT &&
+                    loader.tensor_should_be_converted(ts, override_wtype)) {
+                    ts.type = override_wtype;
+                } else if (ts.expected_type != GGML_TYPE_COUNT && ts.expected_type != ts.type) {
+                    ts.type = ts.expected_type;
+                }
+                bytes[int(kind)] += (int64_t)ts.nbytes() + 64;
+            }
+
+            std::vector<Component> out;
+            out.push_back({ComponentKind::DIT, "DiT", bytes[int(ComponentKind::DIT)], 2048 * MiB, true});
+            out.push_back({ComponentKind::VAE, "VAE", bytes[int(ComponentKind::VAE)], 1024 * MiB, false});
+            out.push_back({ComponentKind::CONDITIONER, "Conditioner", bytes[int(ComponentKind::CONDITIONER)], 2048 * MiB, true});
+            return out;
+        }
+
+        std::vector<Device> enumerate_gpu_devices(const sd::ggml_graph_cut::MaxVramAssignment& budgets) {
+            std::vector<Device> out;
+            for (size_t i = 0; i < ggml_backend_dev_count(); i++) {
+                ggml_backend_dev_t dev = ggml_backend_dev_get(i);
+                if (ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_GPU) {
+                    continue;
+                }
+                Device d;
+                d.dev             = dev;
+                d.name            = ggml_backend_dev_name(dev);
+                d.description     = ggml_backend_dev_description(dev);
+                size_t free_bytes = 0, total_bytes = 0;
+                ggml_backend_dev_memory(dev, &free_bytes, &total_bytes);
+                d.free_bytes  = (int64_t)free_bytes;
+                d.total_bytes = (int64_t)total_bytes;
+
+                std::string budget_key = d.name;
+                std::transform(budget_key.begin(), budget_key.end(), budget_key.begin(),
+                               [](unsigned char c) { return (char)std::tolower(c); });
+                float gib = budgets.default_gib;
+                auto it   = budgets.backend_gib.find(budget_key);
+                if (it != budgets.backend_gib.end()) {
+                    gib = it->second;
+                }
+                if (gib > 0.f) {
+                    d.budget_bytes = std::min<int64_t>((int64_t)(gib * 1024.0 * 1024.0 * 1024.0), d.free_bytes);
+                } else if (gib < 0.f) {
+                    d.budget_bytes = d.free_bytes + (int64_t)(gib * 1024.0 * 1024.0 * 1024.0);
+                } else {
+                    d.budget_bytes = d.free_bytes - 512 * MiB;
+                }
+                d.budget_bytes = std::max<int64_t>(d.budget_bytes, 0);
+                out.push_back(d);
+            }
+            return out;
+        }
+
+        Plan compute_plan(const std::vector<Component>& components, const std::vector<Device>& devices) {
+            Plan plan;
+            if (devices.empty()) {
+                return plan;
+            }
+
+            std::vector<size_t> order(components.size());
+            for (size_t i = 0; i < order.size(); i++) {
+                order[i] = i;
+            }
+            std::sort(order.begin(), order.end(), [&](size_t a, size_t b) {
+                return components[a].params_bytes > components[b].params_bytes;
+            });
+
+            {
+                std::vector<int64_t> params_sum(devices.size(), 0);
+                std::vector<int64_t> max_reserve(devices.size(), 0);
+                std::vector<Decision> decisions(components.size());
+                bool ok = true;
+                for (size_t ci : order) {
+                    const Component& comp = components[ci];
+                    decisions[ci].kind    = comp.kind;
+                    if (comp.params_bytes == 0) {
+                        continue;
+                    }
+                    int best = -1;
+                    for (size_t di = 0; di < devices.size(); di++) {
+                        int64_t need = params_sum[di] + comp.params_bytes + std::max(max_reserve[di], comp.reserve_bytes);
+                        if (need <= devices[di].budget_bytes &&
+                            (best < 0 || devices[di].budget_bytes - params_sum[di] > devices[best].budget_bytes - params_sum[best])) {
+                            best = (int)di;
+                        }
+                    }
+                    if (best < 0) {
+                        ok = false;
+                        break;
+                    }
+                    params_sum[best] += comp.params_bytes;
+                    max_reserve[best] = std::max(max_reserve[best], comp.reserve_bytes);
+                    decisions[ci].device_idxs.push_back((size_t)best);
+                }
+                if (ok) {
+                    plan.valid      = true;
+                    plan.time_share = false;
+                    plan.decisions  = std::move(decisions);
+                    return plan;
+                }
+            }
+
+            plan.decisions.assign(components.size(), {});
+            for (size_t ci : order) {
+                const Component& comp = components[ci];
+                Decision& decision    = plan.decisions[ci];
+                decision.kind         = comp.kind;
+                if (comp.params_bytes == 0) {
+                    continue;
+                }
+                int best = -1;
+                for (size_t di = 0; di < devices.size(); di++) {
+                    if (comp.params_bytes + comp.reserve_bytes <= devices[di].budget_bytes &&
+                        (best < 0 || devices[di].budget_bytes > devices[best].budget_bytes)) {
+                        best = (int)di;
+                    }
+                }
+                if (best >= 0) {
+                    decision.device_idxs.push_back((size_t)best);
+                    continue;
+                }
+                if (comp.splittable && devices.size() > 1) {
+                    int64_t capacity = 0;
+                    for (const Device& d : devices) {
+                        capacity += std::max<int64_t>(d.budget_bytes - comp.reserve_bytes, 0);
+                    }
+                    if (comp.params_bytes <= capacity) {
+                        std::vector<size_t> idxs(devices.size());
+                        for (size_t i = 0; i < idxs.size(); i++) {
+                            idxs[i] = i;
+                        }
+                        std::sort(idxs.begin(), idxs.end(), [&](size_t a, size_t b) {
+                            return devices[a].budget_bytes > devices[b].budget_bytes;
+                        });
+                        decision.device_idxs = std::move(idxs);
+                        continue;
+                    }
+                }
+                decision.on_cpu = true;
+            }
+            plan.valid      = true;
+            plan.time_share = true;
+            return plan;
+        }
+
+        void print_plan(const Plan& plan,
+                        const std::vector<Component>& components,
+                        const std::vector<Device>& devices) {
+            LOG_INFO("auto-fit plan%s:", plan.time_share ? " (time-share: params load per phase and free after)" : "");
+            LOG_INFO("  devices:");
+            for (const Device& d : devices) {
+                LOG_INFO("    %-12s %-32s free %6lld MiB, budget %6lld MiB",
+                         d.name.c_str(), d.description.c_str(),
+                         (long long)(d.free_bytes / MiB), (long long)(d.budget_bytes / MiB));
+            }
+            LOG_INFO("  components:");
+            for (size_t ci = 0; ci < components.size(); ci++) {
+                const Component& comp    = components[ci];
+                const Decision& decision = plan.decisions[ci];
+                std::string target;
+                if (comp.params_bytes == 0) {
+                    target = "(not present)";
+                } else if (decision.on_cpu) {
+                    target = "CPU";
+                } else {
+                    for (size_t k = 0; k < decision.device_idxs.size(); k++) {
+                        if (k > 0) {
+                            target += " & ";
+                        }
+                        target += devices[decision.device_idxs[k]].name;
+                    }
+                    if (decision.device_idxs.size() > 1) {
+                        target += " (split)";
+                    }
+                }
+                LOG_INFO("    %-12s params %6lld MiB, compute reserve %5lld MiB -> %s",
+                         comp.name,
+                         (long long)(comp.params_bytes / MiB),
+                         (long long)(comp.reserve_bytes / MiB),
+                         target.c_str());
+            }
+        }
+
+        void append_assignment(std::string& spec, const char* key, const std::string& value) {
+            if (!spec.empty()) {
+                spec += ",";
+            }
+            spec += key;
+            spec += "=";
+            spec += value;
+        }
+
+        void append_component_decision(const std::vector<Component>& components,
+                                       const std::vector<Device>& devices,
+                                       const Plan& plan,
+                                       ComponentKind kind,
+                                       const char* module_key,
+                                       std::string& runtime_spec,
+                                       std::string& params_spec) {
+            for (size_t ci = 0; ci < components.size(); ci++) {
+                if (components[ci].kind != kind || components[ci].params_bytes == 0) {
+                    continue;
+                }
+                const Decision& decision = plan.decisions[ci];
+                if (decision.on_cpu) {
+                    append_assignment(runtime_spec, module_key, "cpu");
+                    return;
+                }
+                if (decision.device_idxs.empty()) {
+                    return;
+                }
+                std::string device_list;
+                for (size_t k = 0; k < decision.device_idxs.size(); k++) {
+                    if (k > 0) {
+                        device_list += "&";
+                    }
+                    device_list += devices[decision.device_idxs[k]].name;
+                }
+                append_assignment(runtime_spec, module_key, device_list);
+                if (plan.time_share) {
+                    append_assignment(params_spec, module_key, "disk");
+                }
+                return;
+            }
+        }
+
+    }  // namespace
+
+    bool derive_backend_specs(ModelLoader& loader,
+                              ggml_type override_wtype,
+                              sd::ggml_graph_cut::MaxVramAssignment& budgets,
+                              std::string& runtime_spec,
+                              std::string& params_spec) {
+        if (!runtime_spec.empty() || !params_spec.empty()) {
+            LOG_WARN("--auto-fit is enabled; ignoring --backend / --params-backend");
+        }
+
+        {
+            std::string error;
+            if (!budgets.canonicalize_backend_keys(&error)) {
+                LOG_ERROR("%s", error.c_str());
+                return false;
+            }
+        }
+
+        auto components = estimate_components(loader, override_wtype);
+        auto devices    = enumerate_gpu_devices(budgets);
+        auto plan       = compute_plan(components, devices);
+        if (!plan.valid) {
+            LOG_WARN("auto-fit: no usable GPU devices; using the default backend");
+            runtime_spec.clear();
+            params_spec.clear();
+            return true;
+        }
+
+        print_plan(plan, components, devices);
+
+        std::string derived_runtime_spec;
+        std::string derived_params_spec;
+        append_component_decision(components, devices, plan, ComponentKind::DIT, "diffusion", derived_runtime_spec, derived_params_spec);
+        append_component_decision(components, devices, plan, ComponentKind::CONDITIONER, "te", derived_runtime_spec, derived_params_spec);
+        append_component_decision(components, devices, plan, ComponentKind::VAE, "vae", derived_runtime_spec, derived_params_spec);
+
+        runtime_spec = std::move(derived_runtime_spec);
+        params_spec  = std::move(derived_params_spec);
+
+        LOG_INFO("auto-fit: --backend \"%s\"%s%s%s",
+                 runtime_spec.empty() ? "(default)" : runtime_spec.c_str(),
+                 params_spec.empty() ? "" : " --params-backend \"",
+                 params_spec.c_str(),
+                 params_spec.empty() ? "" : "\"");
+        return true;
+    }
+
+    bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params, bool prefer_temporal_tiling) {
+        if (prefer_temporal_tiling) {
+            if (tiling_params.temporal_tiling) {
+                return false;
+            }
+            tiling_params.temporal_tiling = true;
+        } else {
+            if (tiling_params.enabled) {
+                return false;
+            }
+            tiling_params.enabled = true;
+            if (tiling_params.tile_size_x <= 0) {
+                tiling_params.tile_size_x = 256;
+            }
+            if (tiling_params.tile_size_y <= 0) {
+                tiling_params.tile_size_y = 256;
+            }
+        }
+
+        LOG_WARN("auto-fit: VAE decode failed (likely out of memory); retrying with %s tiling",
+                 tiling_params.temporal_tiling ? "temporal" : "spatial");
+        return true;
+    }
+
+}  // namespace sd::backend_fit

src/core/backend_fit.h

@@ -0,0 +1,23 @@
+#ifndef __SD_BACKEND_FIT_H__
+#define __SD_BACKEND_FIT_H__
+
+#include <string>
+
+#include "core/ggml_graph_cut.h"
+#include "model_loader.h"
+#include "stable-diffusion.h"
+
+namespace sd::backend_fit {
+
+    bool derive_backend_specs(ModelLoader& loader,
+                              ggml_type override_wtype,
+                              sd::ggml_graph_cut::MaxVramAssignment& budgets,
+                              std::string& runtime_spec,
+                              std::string& params_spec);
+
+    bool prepare_vae_decode_retry_tiling(sd_tiling_params_t& tiling_params,
+                                         bool prefer_temporal_tiling);
+
+}  // namespace sd::backend_fit
+
+#endif  // __SD_BACKEND_FIT_H__

src/core/ggml_extend.hpp

@@ -391,7 +391,7 @@ __STATIC_INLINE__ uint8_t* ggml_tensor_to_sd_image(ggml_tensor* input, uint8_t*
     int64_t width    = input->ne[0];
     int64_t height   = input->ne[1];
     int64_t channels = input->ne[2];
-    GGML_ASSERT(channels == 3 && input->type == GGML_TYPE_F32);
+    GGML_ASSERT(input->type == GGML_TYPE_F32);
     if (image_data == nullptr) {
         image_data = (uint8_t*)malloc(width * height * channels);
     }
@@ -1038,6 +1038,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_linear(ggml_context* ctx,
 }
 
 __STATIC_INLINE__ ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx,
+                                                ggml_backend_t backend,
                                                 ggml_tensor* x,
                                                 int lp0,
                                                 int rp0,
@@ -1063,7 +1064,17 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx,
     }
 
     if (lp0 != 0 || rp0 != 0 || lp1 != 0 || rp1 != 0 || lp2 != 0 || rp2 != 0 || lp3 != 0 || rp3 != 0) {
-        x = ggml_pad_ext(ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3);
+        ggml_tensor* padded = ggml_pad_ext(ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3);
+        if (backend == nullptr || ggml_backend_supports_op(backend, padded)) {
+            x = padded;
+        } else {
+            // Some backends (e.g. Metal) only implement right-padding for
+            // GGML_OP_PAD (see #850): pad right by lp+rp instead, then roll
+            // the padding around to the left. shift < ne always holds because
+            // ne grew by lp+rp.
+            x = ggml_pad_ext(ctx, x, 0, lp0 + rp0, 0, lp1 + rp1, 0, lp2 + rp2, 0, lp3 + rp3);
+            x = ggml_roll(ctx, x, lp0, lp1, lp2, lp3);
+        }
     }
     return x;
 }
@@ -1076,7 +1087,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_pad(ggml_context* ctx,
                                             int p3          = 0,
                                             bool circular_x = false,
                                             bool circular_y = false) {
-    return ggml_ext_pad_ext(ctx, x, 0, p0, 0, p1, 0, p2, 0, p3, circular_x, circular_y);
+    return ggml_ext_pad_ext(ctx, nullptr, x, 0, p0, 0, p1, 0, p2, 0, p3, circular_x, circular_y);
 }
 
 // w: [OC，IC, KH, KW]
@@ -1105,7 +1116,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx,
     }
 
     if ((p0 != 0 || p1 != 0) && (circular_x || circular_y)) {
-        x  = ggml_ext_pad_ext(ctx, x, p0, p0, p1, p1, 0, 0, 0, 0, circular_x, circular_y);
+        x  = ggml_ext_pad_ext(ctx, nullptr, x, p0, p0, p1, p1, 0, 0, 0, 0, circular_x, circular_y);
         p0 = 0;
         p1 = 0;
     }
@@ -1130,6 +1141,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx,
 // b: [OC,]
 // result: [N*OC, OD, OH, OW]
 __STATIC_INLINE__ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
+                                                ggml_backend_t backend,
                                                 ggml_tensor* x,
                                                 ggml_tensor* w,
                                                 ggml_tensor* b,
@@ -1159,7 +1171,21 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
         x          = ggml_cont(ctx, ggml_permute(ctx, x, 0, 1, 3, 2));
         x          = ggml_reshape_4d(ctx, x, im2col->ne[1], im2col->ne[2], OD, OC * N);
     } else {
-        x = ggml_conv_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2);
+        // ggml_conv_3d decomposes into GGML_OP_IM2COL_3D, which some backends
+        // (e.g. Metal, see #850) do not implement. Fall back to
+        // GGML_OP_CONV_3D on those backends.
+        bool im2col_3d_supported = true;
+        if (backend != nullptr) {
+            ggml_tensor* im2col = ggml_im2col_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, w->type);
+            im2col_3d_supported = ggml_backend_supports_op(backend, im2col);
+        }
+        if (im2col_3d_supported) {
+            x = ggml_conv_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2);
+        } else {
+            int64_t OC = w->ne[3] / IC;
+            int64_t N  = x->ne[3] / IC;
+            x          = ggml_conv_3d_direct(ctx, w, x, s0, s1, s2, p0, p1, p2, d0, d1, d2, (int)IC, (int)N, (int)OC);
+        }
     }
 
     if (b != nullptr) {
@@ -1362,6 +1388,9 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
         }
 
         auto out = ggml_flash_attn_ext(ctx, q_in, k_in, v_in, mask_in, scale / kv_scale, 0, 0);
+        if (!ggml_backend_supports_op(backend, out)) {
+            return nullptr;
+        }
         ggml_flash_attn_ext_set_prec(out, GGML_PREC_F32);
         if (kv_scale != 1.0f) {
             out = ggml_ext_scale(ctx, out, 1.0f / kv_scale);
@@ -1379,9 +1408,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
 
         if (can_use_flash_attn) {
             kqv = build_kqv(q, k, v, mask);
-            if (!ggml_backend_supports_op(backend, kqv)) {
-                kqv = nullptr;
-            } else {
+            if (kqv != nullptr) {
                 kqv = ggml_view_4d(ctx,
                                    kqv,
                                    d_head,
@@ -1719,6 +1746,11 @@ protected:
     bool stream_layers_enabled            = false;
     size_t observed_max_effective_budget_ = 0;
 
+    std::vector<ggml_backend_t> extra_runtime_backends;  // borrowed (SDBackendManager-owned)
+    ggml_backend_sched_t sched             = nullptr;    // owned, multi-device only
+    ggml_backend_t cpu_fallback_backend    = nullptr;    // owned, sched requires a trailing CPU backend
+    bool multi_device_eval_callback_warned = false;
+
     std::shared_ptr<WeightAdapter> weight_adapter = nullptr;
     std::weak_ptr<RunnerWeightManager> weight_manager;
     std::unordered_set<const ggml_tensor*> kept_compute_param_tensor_set;
@@ -1986,7 +2018,121 @@ protected:
         return true;
     }
 
+    // Pass explicit buffer types: synthesized defaults can make CUDA devices
+    // report supporting each other's buffers and skip a required copy.
+    bool ensure_sched(ggml_cgraph* gf) {
+        if (sched != nullptr) {
+            return true;
+        }
+        std::vector<ggml_backend_t> backends;
+        backends.reserve(extra_runtime_backends.size() + 2);
+        backends.push_back(runtime_backend);
+        for (ggml_backend_t backend : extra_runtime_backends) {
+            backends.push_back(backend);
+        }
+        if (cpu_fallback_backend == nullptr && !sd_backend_is_cpu(runtime_backend)) {
+            cpu_fallback_backend = sd_backend_cpu_init();
+        }
+        if (cpu_fallback_backend != nullptr) {
+            backends.push_back(cpu_fallback_backend);
+        }
+
+        std::vector<ggml_backend_buffer_type_t> bufts;
+        bufts.reserve(backends.size());
+        ggml_backend_dev_t main_dev = ggml_backend_get_device(runtime_backend);
+        for (ggml_backend_t backend : backends) {
+            ggml_backend_buffer_type_t buft = nullptr;
+            if (backend == cpu_fallback_backend && main_dev != nullptr) {
+                buft = ggml_backend_dev_host_buffer_type(main_dev);
+            }
+            if (buft == nullptr) {
+                buft = ggml_backend_get_default_buffer_type(backend);
+            }
+            bufts.push_back(buft);
+        }
+
+        size_t graph_size = MAX_GRAPH_SIZE;
+        if (gf != nullptr) {
+            graph_size = std::max<size_t>(graph_size, (size_t)ggml_graph_n_nodes(gf));
+        }
+        sched = ggml_backend_sched_new(backends.data(),
+                                       bufts.data(),
+                                       (int)backends.size(),
+                                       graph_size,
+                                       /*parallel=*/false,
+                                       /*op_offload=*/false);
+        if (sched == nullptr) {
+            LOG_ERROR("%s: failed to create backend sched", get_desc().c_str());
+            return false;
+        }
+        return true;
+    }
+
+    ggml_backend_t backend_for_weight(const ggml_tensor* tensor) const {
+        if (tensor == nullptr || tensor->buffer == nullptr) {
+            return nullptr;
+        }
+        if (ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_WEIGHTS ||
+            ggml_backend_buffer_is_host(tensor->buffer)) {
+            return nullptr;
+        }
+        ggml_backend_dev_t dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(tensor->buffer));
+        if (dev == nullptr) {
+            return nullptr;
+        }
+        if (ggml_backend_get_device(runtime_backend) == dev) {
+            return runtime_backend;
+        }
+        for (ggml_backend_t backend : extra_runtime_backends) {
+            if (ggml_backend_get_device(backend) == dev) {
+                return backend;
+            }
+        }
+        return nullptr;
+    }
+
+    // Weightless ops have no scheduler anchor, so pin them to the most recent
+    // weight device. Views must stay unpinned or cross-device copies can be
+    // skipped for their consumers.
+    void pin_multi_device_nodes(ggml_cgraph* gf) {
+        if (sched == nullptr || gf == nullptr) {
+            return;
+        }
+        ggml_backend_t current = runtime_backend;
+        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);
+            for (int s = 0; s < GGML_MAX_SRC; s++) {
+                ggml_backend_t weight_backend = backend_for_weight(node->src[s]);
+                if (weight_backend != nullptr) {
+                    current = weight_backend;
+                }
+            }
+            if (node->op == GGML_OP_NONE || node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE ||
+                node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE) {
+                continue;
+            }
+            if (ggml_backend_supports_op(current, node)) {
+                ggml_backend_sched_set_tensor_backend(sched, node, current);
+            }
+        }
+    }
+
+    bool is_multi_device() const {
+        return !extra_runtime_backends.empty();
+    }
+
     bool alloc_compute_buffer(ggml_cgraph* gf) {
+        if (is_multi_device()) {
+            // The sched replaces the gallocr. Do NOT ggml_backend_sched_reserve
+            // the graph here: reserve runs split_graph, which rewires the
+            // graph's src pointers to sched-internal copy tensors, and the
+            // later ggml_backend_sched_alloc_graph would split the already
+            // rewired graph, silently corrupting every cross-backend input. A
+            // graph must be split at most once; the alloc in execute_graph
+            // performs the real allocation.
+            return ensure_sched(gf);
+        }
         if (compute_allocr != nullptr) {
             return true;
         }
@@ -2202,12 +2348,14 @@ protected:
                plan.valid &&
                max_graph_vram_bytes > 0 &&
                plan.segments.size() > 1 &&
-               !sd_backend_is_cpu(runtime_backend);
+               !sd_backend_is_cpu(runtime_backend) &&
+               !is_multi_device();
     }
 
     bool can_attempt_graph_cut_segmented_compute() const {
         return max_graph_vram_bytes > 0 &&
-               !sd_backend_is_cpu(runtime_backend);
+               !sd_backend_is_cpu(runtime_backend) &&
+               !is_multi_device();
     }
 
     bool resolve_graph_cut_plan(ggml_cgraph* gf,
@@ -2463,7 +2611,14 @@ protected:
         };
         ComputeBufferGuard compute_buffer_guard(this, free_compute_buffer);
 
-        if (!ggml_gallocr_alloc_graph(compute_allocr, gf)) {
+        if (is_multi_device()) {
+            ggml_backend_sched_reset(sched);
+            pin_multi_device_nodes(gf);  // reset clears the pins; re-apply before alloc
+            if (!ggml_backend_sched_alloc_graph(sched, gf)) {
+                LOG_ERROR("%s sched alloc compute graph failed", get_desc().c_str());
+                return std::nullopt;
+            }
+        } else if (!ggml_gallocr_alloc_graph(compute_allocr, gf)) {
             LOG_ERROR("%s alloc compute graph failed", get_desc().c_str());
             return std::nullopt;
         }
@@ -2472,11 +2627,27 @@ protected:
         if (sd_backend_is_cpu(runtime_backend)) {
             sd_backend_cpu_set_n_threads(runtime_backend, n_threads);
         }
+        if (cpu_fallback_backend != nullptr) {
+            sd_backend_cpu_set_n_threads(cpu_fallback_backend, n_threads);
+        }
 
-        ggml_status status = sd_backend_graph_compute_with_eval_callback(runtime_backend,
-                                                                         gf,
-                                                                         sd_get_backend_eval_callback(),
-                                                                         sd_get_backend_eval_callback_data());
+        ggml_status status;
+        if (is_multi_device()) {
+            if (sd_get_backend_eval_callback() != nullptr && !multi_device_eval_callback_warned) {
+                LOG_WARN("%s: eval callback is not supported with multiple runtime backends; ignoring",
+                         get_desc().c_str());
+                multi_device_eval_callback_warned = true;
+            }
+            status = ggml_backend_sched_graph_compute(sched, gf);
+            if (status == GGML_STATUS_SUCCESS) {
+                ggml_backend_sched_synchronize(sched);
+            }
+        } else {
+            status = sd_backend_graph_compute_with_eval_callback(runtime_backend,
+                                                                 gf,
+                                                                 sd_get_backend_eval_callback(),
+                                                                 sd_get_backend_eval_callback_data());
+        }
         if (status != GGML_STATUS_SUCCESS) {
             LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status));
             return std::nullopt;
@@ -2653,6 +2824,10 @@ public:
         free_params_ctx();
         free_compute_ctx();
         free_cache_ctx_and_buffer();
+        if (cpu_fallback_backend != nullptr) {
+            ggml_backend_free(cpu_fallback_backend);
+            cpu_fallback_backend = nullptr;
+        }
     }
 
     virtual GGMLRunnerContext get_context() {
@@ -2693,10 +2868,20 @@ public:
             ggml_gallocr_free(compute_allocr);
             compute_allocr = nullptr;
         }
+        if (sched != nullptr) {
+            ggml_backend_sched_free(sched);
+            sched = nullptr;
+        }
     }
 
     // do copy after alloc graph
     void set_backend_tensor_data(ggml_tensor* tensor, const void* data) {
+        if (is_multi_device()) {
+            // The sched only assigns a backend (and thus a buffer) to tensors
+            // that participate in the graph; flag standalone data tensors as
+            // inputs so they get one.
+            ggml_set_input(tensor);
+        }
         backend_tensor_data_map[tensor] = data;
     }
 
@@ -2832,8 +3017,31 @@ public:
     }
 
     void set_stream_layers_enabled(bool enabled) {
+        if (enabled && is_multi_device()) {
+            LOG_WARN("%s: --stream-layers is not supported with multiple runtime backends; ignoring",
+                     get_desc().c_str());
+            return;
+        }
         stream_layers_enabled = enabled;
     }
+
+    void set_runtime_backends(const std::vector<ggml_backend_t>& backends) {
+        extra_runtime_backends.clear();
+        for (ggml_backend_t backend : backends) {
+            if (backend == nullptr || backend == runtime_backend) {
+                continue;
+            }
+            if (std::find(extra_runtime_backends.begin(), extra_runtime_backends.end(), backend) ==
+                extra_runtime_backends.end()) {
+                extra_runtime_backends.push_back(backend);
+            }
+        }
+        if (is_multi_device() && stream_layers_enabled) {
+            LOG_WARN("%s: --stream-layers is not supported with multiple runtime backends; ignoring",
+                     get_desc().c_str());
+            stream_layers_enabled = false;
+        }
+    }
 };
 
 class GGMLBlock {
@@ -3368,7 +3576,7 @@ public:
                 b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
             }
         }
-        return ggml_ext_conv_3d(ctx->ggml_ctx, x, w, b, in_channels,
+        return ggml_ext_conv_3d(ctx->ggml_ctx, ctx->backend, x, w, b, in_channels,
                                 std::get<2>(stride), std::get<1>(stride), std::get<0>(stride),
                                 std::get<2>(padding), std::get<1>(padding), std::get<0>(padding),
                                 std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation),

src/core/ggml_extend_backend.cpp

@@ -665,12 +665,52 @@ SDBackendManager::~SDBackendManager() {
 
 void SDBackendManager::reset() {
     backends_.clear();
-    runtime_assignment_ = {};
-    params_assignment_  = {};
+    runtime_assignment_    = {};
+    params_assignment_     = {};
+    split_mode_assignment_ = {};
+}
+
+static std::vector<std::string> split_device_list(const std::string& value) {
+    std::vector<std::string> names;
+    for (const std::string& raw : split_copy(value, '&')) {
+        const std::string name = trim_copy(raw);
+        if (!name.empty()) {
+            names.push_back(name);
+        }
+    }
+    return names;
+}
+
+static std::string primary_device_name(const std::string& value) {
+    std::vector<std::string> names = split_device_list(value);
+    return names.empty() ? std::string() : names.front();
 }
 
 ggml_backend_t SDBackendManager::runtime_backend(SDBackendModule module) {
-    return init_cached_backend(runtime_assignment_.get(module));
+    return init_cached_backend(primary_device_name(runtime_assignment_.get(module)));
+}
+
+std::vector<ggml_backend_t> SDBackendManager::runtime_backends(SDBackendModule module) {
+    std::vector<ggml_backend_t> backends;
+    for (const std::string& name : split_device_list(runtime_assignment_.get(module))) {
+        ggml_backend_t backend = init_cached_backend(name);
+        if (backend == nullptr) {
+            LOG_ERROR("failed to initialize backend '%s' for module %s",
+                      name.c_str(),
+                      sd_backend_module_name(module));
+            continue;
+        }
+        if (std::find(backends.begin(), backends.end(), backend) == backends.end()) {
+            backends.push_back(backend);
+        }
+    }
+    if (backends.empty()) {
+        ggml_backend_t backend = runtime_backend(module);
+        if (backend != nullptr) {
+            backends.push_back(backend);
+        }
+    }
+    return backends;
 }
 
 ggml_backend_t SDBackendManager::params_backend(SDBackendModule module) {
@@ -696,6 +736,10 @@ bool SDBackendManager::params_backend_is_disk(SDBackendModule module) const {
     return is_disk_backend_token(params_assignment_.get(module));
 }
 
+bool SDBackendManager::params_backend_follows_runtime(SDBackendModule module) const {
+    return params_assignment_.get(module).empty();
+}
+
 bool SDBackendManager::runtime_backend_supports_host_buffer(SDBackendModule module) {
     ggml_backend_t backend = runtime_backend(module);
     if (backend == nullptr) {
@@ -715,6 +759,7 @@ bool SDBackendManager::runtime_backend_supports_host_buffer(SDBackendModule modu
 
 bool SDBackendManager::init(const char* backend_spec,
                             const char* params_backend_spec,
+                            const char* split_mode_spec,
                             std::string* error) {
     reset();
 
@@ -724,12 +769,53 @@ bool SDBackendManager::init(const char* backend_spec,
     if (!sd_parse_backend_assignment(SAFE_STR(params_backend_spec), &params_assignment_, error)) {
         return false;
     }
+    if (!sd_parse_backend_assignment(SAFE_STR(split_mode_spec), &split_mode_assignment_, error)) {
+        return false;
+    }
 
     return validate(error);
 }
 
+SDSplitMode SDBackendManager::split_mode(SDBackendModule module) const {
+    return lower_copy(trim_copy(split_mode_assignment_.get(module))) == "row" ? SDSplitMode::ROW
+                                                                              : SDSplitMode::LAYER;
+}
+
+ggml_backend_buffer_type_t SDBackendManager::split_buffer_type(ggml_backend_t backend,
+                                                               const std::vector<float>& tensor_split) {
+    if (backend == nullptr) {
+        return nullptr;
+    }
+    ggml_backend_dev_t dev = ggml_backend_get_device(backend);
+    if (dev == nullptr) {
+        return nullptr;
+    }
+    ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev);
+    if (reg == nullptr) {
+        return nullptr;
+    }
+    auto fn = (ggml_backend_split_buffer_type_t)ggml_backend_reg_get_proc_address(reg, "ggml_backend_split_buffer_type");
+    if (fn == nullptr) {
+        return nullptr;
+    }
+    int main_device        = -1;
+    const size_t dev_count = ggml_backend_reg_dev_count(reg);
+    for (size_t i = 0; i < dev_count; ++i) {
+        if (ggml_backend_reg_dev_get(reg, i) == dev) {
+            main_device = (int)i;
+            break;
+        }
+    }
+    if (main_device < 0) {
+        return nullptr;
+    }
+    std::vector<float> padded_split(std::max<size_t>(tensor_split.size(), 64), 0.0f);
+    std::copy(tensor_split.begin(), tensor_split.end(), padded_split.begin());
+    return fn(main_device, padded_split.data());
+}
+
 bool SDBackendManager::validate(std::string* error) const {
-    auto validate_runtime_name = [&](const std::string& name) -> bool {
+    auto validate_single_runtime_name = [&](const std::string& name) -> bool {
         if (is_default_backend_token(name)) {
             return true;
         }
@@ -747,15 +833,56 @@ bool SDBackendManager::validate(std::string* error) const {
         }
         return false;
     };
+    auto validate_runtime_name = [&](const std::string& name) -> bool {
+        if (name.find('&') == std::string::npos) {
+            return validate_single_runtime_name(name);
+        }
+        std::vector<std::string> names = split_device_list(name);
+        if (names.empty()) {
+            if (error != nullptr) {
+                *error = "invalid backend device list '" + name + "'";
+            }
+            return false;
+        }
+        for (const std::string& entry : names) {
+            if (is_default_backend_token(entry)) {
+                if (error != nullptr) {
+                    *error = "default backend token is not allowed in a device list '" + name + "'";
+                }
+                return false;
+            }
+            if (!validate_single_runtime_name(entry)) {
+                return false;
+            }
+        }
+        return true;
+    };
     auto validate_params_name = [&](const std::string& name) -> bool {
         if (is_disk_backend_token(name)) {
             return true;
         }
-        return validate_runtime_name(name);
+        if (name.find('&') != std::string::npos) {
+            if (error != nullptr) {
+                *error = "params_backend does not accept device lists ('" + name + "')";
+            }
+            return false;
+        }
+        return validate_single_runtime_name(name);
+    };
+    auto validate_split_mode_name = [&](const std::string& name) -> bool {
+        const std::string lower = lower_copy(trim_copy(name));
+        if (lower.empty() || lower == "layer" || lower == "row") {
+            return true;
+        }
+        if (error != nullptr) {
+            *error = "invalid split mode '" + name + "' (expected layer or row)";
+        }
+        return false;
     };
 
     if (!validate_runtime_name(runtime_assignment_.default_name) ||
-        !validate_params_name(params_assignment_.default_name)) {
+        !validate_params_name(params_assignment_.default_name) ||
+        !validate_split_mode_name(split_mode_assignment_.default_name)) {
         return false;
     }
     for (const auto& kv : runtime_assignment_.module_names) {
@@ -768,6 +895,11 @@ bool SDBackendManager::validate(std::string* error) const {
             return false;
         }
     }
+    for (const auto& kv : split_mode_assignment_.module_names) {
+        if (!validate_split_mode_name(kv.second)) {
+            return false;
+        }
+    }
     return true;
 }
 

src/core/ggml_extend_backend.h

@@ -6,6 +6,7 @@
 #include <memory>
 #include <string>
 #include <unordered_map>
+#include <vector>
 
 #include "ggml-backend.h"
 #include "ggml.h"
@@ -37,10 +38,16 @@ struct SDBackendHandleDeleter {
 
 using SDBackendHandle = std::unique_ptr<struct ggml_backend, SDBackendHandleDeleter>;
 
+enum class SDSplitMode {
+    LAYER,
+    ROW,
+};
+
 class SDBackendManager {
 private:
     SDBackendAssignment runtime_assignment_;
     SDBackendAssignment params_assignment_;
+    SDBackendAssignment split_mode_assignment_;
     std::unordered_map<std::string, SDBackendHandle> backends_;
 
 public:
@@ -52,15 +59,23 @@ public:
 
     bool init(const char* backend_spec,
               const char* params_backend_spec,
+              const char* split_mode_spec,
               std::string* error);
     void reset();
 
     ggml_backend_t runtime_backend(SDBackendModule module);
     ggml_backend_t params_backend(SDBackendModule module);
 
+    std::vector<ggml_backend_t> runtime_backends(SDBackendModule module);
+
+    SDSplitMode split_mode(SDBackendModule module) const;
+    ggml_backend_buffer_type_t split_buffer_type(ggml_backend_t backend,
+                                                 const std::vector<float>& tensor_split);
+
     bool runtime_backend_is_cpu(SDBackendModule module);
     bool params_backend_is_cpu(SDBackendModule module);
     bool params_backend_is_disk(SDBackendModule module) const;
+    bool params_backend_follows_runtime(SDBackendModule module) const;
     bool runtime_backend_supports_host_buffer(SDBackendModule module);
 
 private:

src/core/layer_split_partition.cpp

@@ -0,0 +1,221 @@
+#include "core/layer_split_partition.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <cstdlib>
+#include <cstring>
+
+#include "core/util.h"
+
+namespace sd {
+
+    static bool layer_split_path_segment_starts_at(const std::string& name, size_t pos) {
+        return pos == 0 || name[pos - 1] == '.';
+    }
+
+    static bool layer_split_has_path_segment(const std::string& name, const char* segment) {
+        size_t pos = name.find(segment);
+        while (pos != std::string::npos) {
+            if (layer_split_path_segment_starts_at(name, pos)) {
+                return true;
+            }
+            pos = name.find(segment, pos + 1);
+        }
+        return false;
+    }
+
+    int layer_split_tensor_block_index(const std::string& name) {
+        static const char* unet_block_segments[] = {"input_blocks.", "output_blocks.", "middle_block.",
+                                                    "down_blocks.", "up_blocks.", "mid_block."};
+        for (const char* segment : unet_block_segments) {
+            if (layer_split_has_path_segment(name, segment)) {
+                return -1;
+            }
+        }
+
+        static const char* block_keywords[] = {"transformer_blocks.", "joint_blocks.", "double_blocks.",
+                                               "single_blocks.", "blocks.", "block.", "layers."};
+        for (const char* keyword : block_keywords) {
+            size_t pos = name.find(keyword);
+            while (pos != std::string::npos) {
+                if (!layer_split_path_segment_starts_at(name, pos)) {
+                    pos = name.find(keyword, pos + 1);
+                    continue;
+                }
+                pos += std::strlen(keyword);
+                size_t end = pos;
+                while (end < name.size() && name[end] >= '0' && name[end] <= '9') {
+                    end++;
+                }
+                if (end > pos && (end == name.size() || name[end] == '.')) {
+                    return std::atoi(name.substr(pos, end - pos).c_str());
+                }
+                break;
+            }
+        }
+        return -1;
+    }
+
+    std::string layer_split_backend_device_display_name(ggml_backend_t backend) {
+        ggml_backend_dev_t dev = ggml_backend_get_device(backend);
+        const char* name       = dev != nullptr ? ggml_backend_dev_name(dev) : ggml_backend_name(backend);
+        return name != nullptr ? name : "unknown";
+    }
+
+    static bool layer_split_backend_supports_tensor(ggml_backend_t backend, const ggml_tensor* tensor) {
+        return backend != nullptr && tensor != nullptr && ggml_backend_supports_op(backend, tensor);
+    }
+
+    static size_t layer_split_supported_target(const std::string& desc,
+                                               const std::string& tensor_name,
+                                               const ggml_tensor* tensor,
+                                               const std::vector<ggml_backend_t>& backends,
+                                               size_t preferred) {
+        if (tensor == nullptr || backends.empty()) {
+            return preferred;
+        }
+        size_t preferred_safe = std::min(preferred, backends.size() - 1);
+        if (layer_split_backend_supports_tensor(backends[preferred_safe], tensor)) {
+            return preferred_safe;
+        }
+        for (size_t i = 0; i < backends.size(); i++) {
+            if (layer_split_backend_supports_tensor(backends[i], tensor)) {
+                LOG_WARN("%s layer split: moving tensor '%s' from %s to %s because the preferred backend cannot run op=%s type=%s nbytes=%.2f MB",
+                         desc.c_str(),
+                         tensor_name.c_str(),
+                         layer_split_backend_device_display_name(backends[preferred_safe]).c_str(),
+                         layer_split_backend_device_display_name(backends[i]).c_str(),
+                         ggml_op_name(tensor->op),
+                         ggml_type_name(tensor->type),
+                         ggml_nbytes(tensor) / (1024.0 * 1024.0));
+                return i;
+            }
+        }
+        LOG_WARN("%s layer split: tensor '%s' is not supported by any split backend: op=%s type=%s nbytes=%.2f MB",
+                 desc.c_str(),
+                 tensor_name.c_str(),
+                 ggml_op_name(tensor->op),
+                 ggml_type_name(tensor->type),
+                 ggml_nbytes(tensor) / (1024.0 * 1024.0));
+        return preferred_safe;
+    }
+
+    std::vector<std::map<std::string, ggml_tensor*>> partition_layer_split_tensors(
+        const std::string& desc,
+        const std::map<std::string, ggml_tensor*>& tensors,
+        const std::map<std::string, ggml_tensor*>& split_tensors,
+        const std::vector<ggml_backend_t>& backends) {
+        std::vector<std::map<std::string, ggml_tensor*>> partitions(backends.size());
+        if (backends.empty()) {
+            LOG_WARN("%s: no backend available for a layer split", desc.c_str());
+            return partitions;
+        }
+
+        std::map<int, int64_t> block_bytes;
+        std::map<std::string, size_t> non_block_targets;
+        std::vector<int64_t> other_bytes_by_backend(backends.size(), 0);
+        int64_t total_block_bytes = 0;
+        int64_t total_other_bytes = 0;
+        int n_blocks              = 0;
+        for (const auto& kv : tensors) {
+            int64_t bytes = (int64_t)ggml_nbytes(kv.second);
+            int idx       = split_tensors.count(kv.first) != 0 ? layer_split_tensor_block_index(kv.first) : -1;
+            if (idx >= 0) {
+                block_bytes[idx] += bytes;
+                total_block_bytes += bytes;
+                n_blocks = std::max(n_blocks, idx + 1);
+            } else {
+                size_t target               = layer_split_supported_target(desc, kv.first, kv.second, backends, 0);
+                non_block_targets[kv.first] = target;
+                other_bytes_by_backend[target] += bytes;
+                total_other_bytes += bytes;
+            }
+        }
+        if (n_blocks == 0) {
+            LOG_WARN("%s: no transformer blocks found for a layer split; keeping tensors on compatible backends starting from %s",
+                     desc.c_str(),
+                     layer_split_backend_device_display_name(backends[0]).c_str());
+            for (const auto& kv : tensors) {
+                size_t target  = 0;
+                auto target_it = non_block_targets.find(kv.first);
+                if (target_it != non_block_targets.end()) {
+                    target = target_it->second;
+                }
+                partitions[target][kv.first] = kv.second;
+            }
+            return partitions;
+        }
+
+        // Reserve compute headroom and subtract each device's actual non-block
+        // bytes from its block budget.
+        constexpr int64_t compute_headroom_bytes = 2ll * 1024 * 1024 * 1024;
+        std::vector<double> device_weights(backends.size(), 1.0);
+        double weight_sum = 0.0;
+        for (size_t i = 0; i < backends.size(); i++) {
+            ggml_backend_dev_t dev = ggml_backend_get_device(backends[i]);
+            size_t free_bytes = 0, total_bytes = 0;
+            if (dev != nullptr) {
+                ggml_backend_dev_memory(dev, &free_bytes, &total_bytes);
+            }
+            // Keep a small share even for tight devices instead of dropping them.
+            int64_t usable_bytes = std::max<int64_t>((int64_t)free_bytes - compute_headroom_bytes,
+                                                     (int64_t)free_bytes / 8);
+            device_weights[i]    = usable_bytes > 0 ? (double)usable_bytes : 1.0;
+            weight_sum += device_weights[i];
+        }
+
+        std::vector<int64_t> block_budgets(backends.size(), 0);
+        const int64_t total_bytes = total_block_bytes + total_other_bytes;
+        for (size_t i = 0; i < backends.size(); i++) {
+            int64_t budget   = (int64_t)((double)total_bytes * device_weights[i] / weight_sum);
+            budget           = std::max<int64_t>(budget - other_bytes_by_backend[i], 0);
+            block_budgets[i] = budget;
+        }
+
+        std::vector<int> boundaries(backends.size(), n_blocks);
+        size_t current = 0;
+        int64_t used   = 0;
+        for (int b = 0; b < n_blocks; b++) {
+            int64_t bytes = block_bytes.count(b) != 0 ? block_bytes[b] : 0;
+            if (current + 1 < backends.size() && used > 0 && used + bytes > block_budgets[current]) {
+                boundaries[current] = b;
+                current++;
+                used = 0;
+            }
+            used += bytes;
+        }
+
+        for (const auto& kv : tensors) {
+            size_t target = 0;
+            int idx       = split_tensors.count(kv.first) != 0 ? layer_split_tensor_block_index(kv.first) : -1;
+            if (idx >= 0) {
+                while (target < boundaries.size() && idx >= boundaries[target]) {
+                    target++;
+                }
+                target = std::min(target, backends.size() - 1);
+                target = layer_split_supported_target(desc, kv.first, kv.second, backends, target);
+            } else {
+                auto target_it = non_block_targets.find(kv.first);
+                if (target_it != non_block_targets.end()) {
+                    target = target_it->second;
+                }
+            }
+            partitions[target][kv.first] = kv.second;
+        }
+
+        int range_start = 0;
+        for (size_t i = 0; i < backends.size(); i++) {
+            int range_end                = boundaries[i];
+            const char* non_block_suffix = other_bytes_by_backend[i] > 0 ? " + non-block tensors" : "";
+            LOG_INFO("%s layer split: %s <- blocks [%d, %d)%s",
+                     desc.c_str(),
+                     layer_split_backend_device_display_name(backends[i]).c_str(),
+                     range_start,
+                     range_end,
+                     non_block_suffix);
+            range_start = range_end;
+        }
+        return partitions;
+    }
+
+}  // namespace sd

src/core/layer_split_partition.h

@@ -0,0 +1,24 @@
+#ifndef __SD_CORE_LAYER_SPLIT_PARTITION_H__
+#define __SD_CORE_LAYER_SPLIT_PARTITION_H__
+
+#include <map>
+#include <string>
+#include <vector>
+
+#include "ggml-backend.h"
+#include "ggml.h"
+
+namespace sd {
+
+    std::string layer_split_backend_device_display_name(ggml_backend_t backend);
+    int layer_split_tensor_block_index(const std::string& name);
+
+    std::vector<std::map<std::string, ggml_tensor*>> partition_layer_split_tensors(
+        const std::string& desc,
+        const std::map<std::string, ggml_tensor*>& tensors,
+        const std::map<std::string, ggml_tensor*>& split_tensors,
+        const std::vector<ggml_backend_t>& backends);
+
+}  // namespace sd
+
+#endif  // __SD_CORE_LAYER_SPLIT_PARTITION_H__

src/core/util.cpp

@@ -4,6 +4,8 @@
 #include <cmath>
 #include <codecvt>
 #include <cstdarg>
+#include <cstdlib>
+#include <cstring>
 #include <exception>
 #include <fstream>
 #include <locale>
@@ -25,6 +27,7 @@
 #include <unistd.h>
 #endif
 
+#include "ggml-backend.h"
 #include "ggml.h"
 #include "stable-diffusion.h"
 
@@ -997,3 +1000,26 @@ std::vector<std::pair<std::string, float>> split_quotation_attention(
     }
     return result;
 }
+
+size_t sd_list_devices(char* buffer, size_t buffer_size) {
+    if (ggml_backend_dev_count() == 0) {
+        // dynamic-backend builds discover their backend modules at runtime
+        ggml_backend_load_all();
+    }
+
+    std::ostringstream oss;
+    for (size_t i = 0; i < ggml_backend_dev_count(); i++) {
+        ggml_backend_dev_t dev = ggml_backend_dev_get(i);
+        const char* name       = ggml_backend_dev_name(dev);
+        const char* desc       = ggml_backend_dev_description(dev);
+        oss << (name ? name : "") << '\t' << (desc ? desc : "") << '\n';
+    }
+
+    std::string devices = oss.str();
+    if (buffer != nullptr && buffer_size > 0) {
+        size_t copy_size = std::min(devices.size(), buffer_size - 1);
+        memcpy(buffer, devices.data(), copy_size);
+        buffer[copy_size] = '\0';
+    }
+    return devices.size();
+}

src/model.h

@@ -36,6 +36,7 @@ enum SDVersion {
     VERSION_WAN2_2_I2V,
     VERSION_WAN2_2_TI2V,
     VERSION_QWEN_IMAGE,
+    VERSION_QWEN_IMAGE_LAYERED,
     VERSION_ANIMA,
     VERSION_FLUX2,
     VERSION_FLUX2_KLEIN,
@@ -127,7 +128,7 @@ static inline bool sd_version_is_wan(SDVersion version) {
 }
 
 static inline bool sd_version_is_qwen_image(SDVersion version) {
-    if (version == VERSION_QWEN_IMAGE) {
+    if (version == VERSION_QWEN_IMAGE || version == VERSION_QWEN_IMAGE_LAYERED) {
         return true;
     }
     return false;

src/model/common/rope.hpp

@@ -12,6 +12,16 @@ namespace Rope {
         ErnieImage,
     };
 
+    enum class RefIndexMode {
+        FIXED,
+        INCREASE,
+        DECREASE,
+    };
+
+    __STATIC_INLINE__ RefIndexMode ref_index_mode_from_bool(bool increase_ref_index) {
+        return increase_ref_index ? RefIndexMode::INCREASE : RefIndexMode::FIXED;
+    }
+
     template <class T>
     __STATIC_INLINE__ std::vector<T> linspace(T start, T end, int num) {
         std::vector<T> result(num);
@@ -346,7 +356,7 @@ namespace Rope {
                                                                    int axes_dim_num,
                                                                    int start_index,
                                                                    const std::vector<ggml_tensor*>& ref_latents,
-                                                                   bool increase_ref_index,
+                                                                   RefIndexMode ref_index_mode,
                                                                    float ref_index_scale,
                                                                    bool scale_rope,
                                                                    int base_offset = 0) {
@@ -357,13 +367,15 @@ namespace Rope {
         for (ggml_tensor* ref : ref_latents) {
             int h_offset = 0;
             int w_offset = 0;
-            if (!increase_ref_index) {
+            if (ref_index_mode == RefIndexMode::FIXED) {
                 if (ref->ne[1] + curr_h_offset > ref->ne[0] + curr_w_offset) {
                     w_offset = curr_w_offset;
                 } else {
                     h_offset = curr_h_offset;
                 }
                 scale_rope = false;
+            } else if (ref_index_mode == RefIndexMode::DECREASE) {
+                index--;
             }
 
             auto ref_ids = gen_flux_img_ids(static_cast<int>(ref->ne[1]),
@@ -377,7 +389,7 @@ namespace Rope {
                                             scale_rope);
             ids          = concat_ids(ids, ref_ids, bs);
 
-            if (increase_ref_index) {
+            if (ref_index_mode == RefIndexMode::INCREASE) {
                 index++;
             }
 
@@ -395,7 +407,7 @@ namespace Rope {
                                                                    int context_len,
                                                                    std::set<int> txt_arange_dims,
                                                                    const std::vector<ggml_tensor*>& ref_latents,
-                                                                   bool increase_ref_index,
+                                                                   RefIndexMode ref_index_mode,
                                                                    float ref_index_scale,
                                                                    bool is_longcat) {
         int x_index = is_longcat ? 1 : 0;
@@ -406,7 +418,7 @@ namespace Rope {
 
         auto ids = concat_ids(txt_ids, img_ids, bs);
         if (ref_latents.size() > 0) {
-            auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, x_index + 1, ref_latents, increase_ref_index, ref_index_scale, false, offset);
+            auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, x_index + 1, ref_latents, ref_index_mode, ref_index_scale, false, offset);
             ids           = concat_ids(ids, refs_ids, bs);
         }
         return ids;
@@ -420,7 +432,7 @@ namespace Rope {
                                                      int context_len,
                                                      std::set<int> txt_arange_dims,
                                                      const std::vector<ggml_tensor*>& ref_latents,
-                                                     bool increase_ref_index,
+                                                     RefIndexMode ref_index_mode,
                                                      float ref_index_scale,
                                                      int theta,
                                                      bool circular_h,
@@ -435,7 +447,7 @@ namespace Rope {
                                                            context_len,
                                                            txt_arange_dims,
                                                            ref_latents,
-                                                           increase_ref_index,
+                                                           ref_index_mode,
                                                            ref_index_scale,
                                                            is_longcat);
         std::vector<std::vector<int>> wrap_dims;
@@ -481,17 +493,64 @@ namespace Rope {
         return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
     }
 
-    __STATIC_INLINE__ std::vector<std::vector<float>> gen_qwen_image_ids(int h,
+    __STATIC_INLINE__ std::vector<std::vector<float>> gen_vid_ids(int t,
+                                                                  int h,
+                                                                  int w,
+                                                                  int pt,
+                                                                  int ph,
+                                                                  int pw,
+                                                                  int bs,
+                                                                  int t_offset    = 0,
+                                                                  int h_offset    = 0,
+                                                                  int w_offset    = 0,
+                                                                  bool scale_rope = false) {
+        int t_len = (t + (pt / 2)) / pt;
+        int h_len = (h + (ph / 2)) / ph;
+        int w_len = (w + (pw / 2)) / pw;
+
+        std::vector<std::vector<float>> vid_ids(t_len * h_len * w_len, std::vector<float>(3, 0.0));
+
+        if (scale_rope) {
+            h_offset -= h_len / 2;
+            w_offset -= w_len / 2;
+        }
+
+        std::vector<float> t_ids = linspace<float>(1.f * t_offset, 1.f * t_len - 1 + t_offset, t_len);
+        std::vector<float> h_ids = linspace<float>(1.f * h_offset, 1.f * h_len - 1 + h_offset, h_len);
+        std::vector<float> w_ids = linspace<float>(1.f * w_offset, 1.f * w_len - 1 + w_offset, w_len);
+
+        for (int i = 0; i < t_len; ++i) {
+            for (int j = 0; j < h_len; ++j) {
+                for (int k = 0; k < w_len; ++k) {
+                    int idx         = i * h_len * w_len + j * w_len + k;
+                    vid_ids[idx][0] = t_ids[i];
+                    vid_ids[idx][1] = h_ids[j];
+                    vid_ids[idx][2] = w_ids[k];
+                }
+            }
+        }
+
+        std::vector<std::vector<float>> vid_ids_repeated(bs * vid_ids.size(), std::vector<float>(3));
+        for (int i = 0; i < bs; ++i) {
+            for (int j = 0; j < vid_ids.size(); ++j) {
+                vid_ids_repeated[i * vid_ids.size() + j] = vid_ids[j];
+            }
+        }
+        return vid_ids_repeated;
+    }
+
+    __STATIC_INLINE__ std::vector<std::vector<float>> gen_qwen_image_ids(int t,
+                                                                         int h,
                                                                          int w,
                                                                          int patch_size,
                                                                          int bs,
                                                                          int context_len,
                                                                          const std::vector<ggml_tensor*>& ref_latents,
-                                                                         bool increase_ref_index) {
+                                                                         RefIndexMode ref_index_mode) {
         int h_len        = (h + (patch_size / 2)) / patch_size;
         int w_len        = (w + (patch_size / 2)) / patch_size;
-        int txt_id_start = std::max(h_len, w_len);
-        auto txt_ids     = linspace<float>(1.f * txt_id_start, 1.f * context_len + txt_id_start, context_len);
+        int txt_id_start = std::max(h_len, w_len) / 2;
+        auto txt_ids     = linspace<float>(1.f * txt_id_start, 1.f * txt_id_start + context_len - 1, context_len);
         std::vector<std::vector<float>> txt_ids_repeated(bs * context_len, std::vector<float>(3));
         for (int i = 0; i < bs; ++i) {
             for (int j = 0; j < txt_ids.size(); ++j) {
@@ -499,28 +558,30 @@ namespace Rope {
             }
         }
         int axes_dim_num = 3;
-        auto img_ids     = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, 0, 0, 0, true);
+        auto img_ids     = gen_vid_ids(t, h, w, 1, patch_size, patch_size, bs, 0, 0, 0, true);
         auto ids         = concat_ids(txt_ids_repeated, img_ids, bs);
         if (ref_latents.size() > 0) {
-            auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, 1, ref_latents, increase_ref_index, 1.f, true);
-            ids           = concat_ids(ids, refs_ids, bs);
+            int ref_start_index = ref_index_mode == RefIndexMode::DECREASE ? 0 : 1;
+            auto refs_ids       = gen_refs_ids(patch_size, bs, axes_dim_num, ref_start_index, ref_latents, ref_index_mode, 1.f, true);
+            ids                 = concat_ids(ids, refs_ids, bs);
         }
         return ids;
     }
 
     // Generate qwen_image positional embeddings
-    __STATIC_INLINE__ std::vector<float> gen_qwen_image_pe(int h,
+    __STATIC_INLINE__ std::vector<float> gen_qwen_image_pe(int t,
+                                                           int h,
                                                            int w,
                                                            int patch_size,
                                                            int bs,
                                                            int context_len,
                                                            const std::vector<ggml_tensor*>& ref_latents,
-                                                           bool increase_ref_index,
+                                                           RefIndexMode ref_index_mode,
                                                            int theta,
                                                            bool circular_h,
                                                            bool circular_w,
                                                            const std::vector<int>& axes_dim) {
-        std::vector<std::vector<float>> ids = gen_qwen_image_ids(h, w, patch_size, bs, context_len, ref_latents, increase_ref_index);
+        std::vector<std::vector<float>> ids = gen_qwen_image_ids(t, h, w, patch_size, bs, context_len, ref_latents, ref_index_mode);
         std::vector<std::vector<int>> wrap_dims;
         // This logic simply stores the (pad and patch_adjusted) sizes of images so we can make sure rope correctly tiles
         if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
@@ -533,7 +594,7 @@ namespace Rope {
                 // Track per-token wrap lengths for the row/column axes so only spatial tokens become periodic.
                 wrap_dims.assign(axes_dim.size(), std::vector<int>(total_tokens / bs, 0));
                 size_t cursor           = context_len;  // ignore text tokens
-                const size_t img_tokens = static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
+                const size_t img_tokens = static_cast<size_t>(t) * static_cast<size_t>(h_len) * static_cast<size_t>(w_len);
                 for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
                     if (circular_h) {
                         wrap_dims[1][cursor + token_i] = h_len;
@@ -684,46 +745,6 @@ namespace Rope {
         return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims, EmbedNDLayout::ErnieImage);
     }
 
-    __STATIC_INLINE__ std::vector<std::vector<float>> gen_vid_ids(int t,
-                                                                  int h,
-                                                                  int w,
-                                                                  int pt,
-                                                                  int ph,
-                                                                  int pw,
-                                                                  int bs,
-                                                                  int t_offset = 0,
-                                                                  int h_offset = 0,
-                                                                  int w_offset = 0) {
-        int t_len = (t + (pt / 2)) / pt;
-        int h_len = (h + (ph / 2)) / ph;
-        int w_len = (w + (pw / 2)) / pw;
-
-        std::vector<std::vector<float>> vid_ids(t_len * h_len * w_len, std::vector<float>(3, 0.0));
-
-        std::vector<float> t_ids = linspace<float>(1.f * t_offset, 1.f * t_len - 1 + t_offset, t_len);
-        std::vector<float> h_ids = linspace<float>(1.f * h_offset, 1.f * h_len - 1 + h_offset, h_len);
-        std::vector<float> w_ids = linspace<float>(1.f * w_offset, 1.f * w_len - 1 + w_offset, w_len);
-
-        for (int i = 0; i < t_len; ++i) {
-            for (int j = 0; j < h_len; ++j) {
-                for (int k = 0; k < w_len; ++k) {
-                    int idx         = i * h_len * w_len + j * w_len + k;
-                    vid_ids[idx][0] = t_ids[i];
-                    vid_ids[idx][1] = h_ids[j];
-                    vid_ids[idx][2] = w_ids[k];
-                }
-            }
-        }
-
-        std::vector<std::vector<float>> vid_ids_repeated(bs * vid_ids.size(), std::vector<float>(3));
-        for (int i = 0; i < bs; ++i) {
-            for (int j = 0; j < vid_ids.size(); ++j) {
-                vid_ids_repeated[i * vid_ids.size() + j] = vid_ids[j];
-            }
-        }
-        return vid_ids_repeated;
-    }
-
     // Generate wan positional embeddings
     __STATIC_INLINE__ std::vector<float> gen_wan_pe(int t,
                                                     int h,
@@ -785,7 +806,8 @@ namespace Rope {
                                                                       int context_len,
                                                                       int seq_multi_of,
                                                                       const std::vector<ggml_tensor*>& ref_latents,
-                                                                      bool increase_ref_index) {
+                                                                      RefIndexMode ref_index_mode) {
+        SD_UNUSED(ref_index_mode);
         int padded_context_len = context_len + bound_mod(context_len, seq_multi_of);
         auto txt_ids           = std::vector<std::vector<float>>(bs * padded_context_len, std::vector<float>(3, 0.0f));
         for (int i = 0; i < bs * padded_context_len; i++) {
@@ -816,12 +838,12 @@ namespace Rope {
                                                         int context_len,
                                                         int seq_multi_of,
                                                         const std::vector<ggml_tensor*>& ref_latents,
-                                                        bool increase_ref_index,
+                                                        RefIndexMode ref_index_mode,
                                                         int theta,
                                                         bool circular_h,
                                                         bool circular_w,
                                                         const std::vector<int>& axes_dim) {
-        std::vector<std::vector<float>> ids = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, increase_ref_index);
+        std::vector<std::vector<float>> ids = gen_z_image_ids(h, w, patch_size, bs, context_len, seq_multi_of, ref_latents, ref_index_mode);
         std::vector<std::vector<int>> wrap_dims;
         if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
             int pad_h = (patch_size - (h % patch_size)) % patch_size;

src/model/diffusion/anima.hpp

@@ -612,7 +612,7 @@ namespace Anima {
                                           0,
                                           {},
                                           empty_ref_latents,
-                                          false,
+                                          Rope::RefIndexMode::FIXED,
                                           1.0f,
                                           false);
 

src/model/diffusion/dit.hpp

@@ -135,23 +135,23 @@ namespace DiT {
         return x;
     }
 
-    inline ggml_tensor* unpatchify(ggml_context* ctx,
-                                   ggml_tensor* x,
-                                   int64_t t_len,
-                                   int64_t h_len,
-                                   int64_t w_len,
-                                   int pt,
-                                   int ph,
-                                   int pw) {
-        // x: [N, t_len*h_len*w_len, pt*ph*pw*C]
+    inline ggml_tensor* unpatchify_3d(ggml_context* ctx,
+                                      ggml_tensor* x,
+                                      int64_t t_len,
+                                      int64_t h_len,
+                                      int64_t w_len,
+                                      int pt,
+                                      int ph,
+                                      int pw) {
+        // x: [N, t_len*h_len*w_len, C*pt*ph*pw]
         // return: [N*C, t_len*pt, h_len*ph, w_len*pw]
-        int64_t N = x->ne[3];
+        int64_t N = x->ne[2];
         int64_t C = x->ne[0] / pt / ph / pw;
 
         GGML_ASSERT(C * pt * ph * pw == x->ne[0]);
 
-        x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N);  // [N, t_len*h_len*w_len, pt*ph*pw, C]
-        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3));      // [N, C, t_len*h_len*w_len, pt*ph*pw]
+        x = ggml_reshape_4d(ctx, x, pw * ph * pt, C, w_len * h_len * t_len, N);  // [N, t_len*h_len*w_len, C, pt*ph*pw]
+        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N, C, t_len*h_len*w_len, pt*ph*pw]
         x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N);  // [N*C*t_len*h_len, w_len, pt*ph, pw]
         x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len*h_len, pt*ph, w_len, pw]
         x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N);  // [N*C*t_len*h_len, pt, ph, w_len*pw]

src/model/diffusion/flux.hpp

@@ -1485,7 +1485,7 @@ namespace Flux {
                                  const sd::Tensor<float>& y_tensor                        = {},
                                  const sd::Tensor<float>& guidance_tensor                 = {},
                                  const std::vector<sd::Tensor<float>>& ref_latents_tensor = {},
-                                 bool increase_ref_index                                  = false,
+                                 Rope::RefIndexMode ref_index_mode                        = Rope::RefIndexMode::FIXED,
                                  std::vector<int> skip_layers                             = {},
                                  const sd::Tensor<float>& pulid_id_tensor                 = {},
                                  float pulid_id_weight                                    = 1.0f) {
@@ -1527,8 +1527,8 @@ namespace Flux {
             }
             std::set<int> txt_arange_dims;
             if (sd_version_is_flux2(version) || sd_version_is_sefi_image(version)) {
-                txt_arange_dims    = {3};
-                increase_ref_index = true;
+                txt_arange_dims = {3};
+                ref_index_mode  = Rope::RefIndexMode::INCREASE;
             } else if (version == VERSION_OVIS_IMAGE) {
                 txt_arange_dims = {1, 2};
             }
@@ -1539,7 +1539,7 @@ namespace Flux {
                                             static_cast<int>(context->ne[1]),
                                             txt_arange_dims,
                                             ref_latents,
-                                            increase_ref_index,
+                                            ref_index_mode,
                                             config.ref_index_scale,
                                             config.theta,
                                             circular_y_enabled,
@@ -1599,7 +1599,7 @@ namespace Flux {
                                   const sd::Tensor<float>& y                        = {},
                                   const sd::Tensor<float>& guidance                 = {},
                                   const std::vector<sd::Tensor<float>>& ref_latents = {},
-                                  bool increase_ref_index                           = false,
+                                  Rope::RefIndexMode ref_index_mode                 = Rope::RefIndexMode::FIXED,
                                   std::vector<int> skip_layers                      = std::vector<int>(),
                                   const sd::Tensor<float>& pulid_id                 = {},
                                   float pulid_id_weight                             = 1.0f) {
@@ -1610,7 +1610,7 @@ namespace Flux {
             // guidance: [N, ]
             // pulid_id: empty (no injection) or [N, num_id_tokens=32, kv_dim=2048]
             auto get_graph = [&]() -> ggml_cgraph* {
-                return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, increase_ref_index, skip_layers, pulid_id, pulid_id_weight);
+                return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, ref_index_mode, skip_layers, pulid_id, pulid_id_weight);
             };
 
             auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
@@ -1632,7 +1632,7 @@ namespace Flux {
                            tensor_or_empty(diffusion_params.y),
                            tensor_or_empty(extra->guidance),
                            diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
-                           diffusion_params.increase_ref_index,
+                           diffusion_params.ref_index_mode,
                            extra->skip_layers ? *extra->skip_layers : empty_skip_layers,
                            tensor_or_empty(extra->pulid_id),
                            extra->pulid_id_weight);
@@ -1683,7 +1683,7 @@ namespace Flux {
                                        {},
                                        guidance,
                                        {},
-                                       false);
+                                       Rope::RefIndexMode::FIXED);
                 int64_t t1   = ggml_time_ms();
 
                 GGML_ASSERT(!out_opt.empty());

src/model/diffusion/model.hpp

@@ -7,6 +7,7 @@
 
 #include "core/ggml_extend.hpp"
 #include "core/tensor_ggml.hpp"
+#include "model/common/rope.hpp"
 #include "model_manager.h"
 
 struct UNetDiffusionExtra {
@@ -73,7 +74,7 @@ struct DiffusionParams {
     const sd::Tensor<float>* c_concat                 = nullptr;
     const sd::Tensor<float>* y                        = nullptr;
     const std::vector<sd::Tensor<float>>* ref_latents = nullptr;
-    bool increase_ref_index                           = false;
+    Rope::RefIndexMode ref_index_mode                 = Rope::RefIndexMode::FIXED;
     DiffusionExtraParams extra                        = std::monostate{};
 };
 

src/model/diffusion/qwen_image.hpp

@@ -4,6 +4,7 @@
 #include <memory>
 
 #include "model/common/block.hpp"
+#include "model/diffusion/dit.hpp"
 #include "model/diffusion/flux.hpp"
 #include "model/diffusion/model.hpp"
 #include "model_loader.h"
@@ -23,6 +24,7 @@ namespace Qwen {
         std::vector<int> axes_dim   = {16, 56, 56};
         int axes_dim_sum            = 128;
         bool zero_cond_t            = false;
+        bool use_additional_t_cond  = false;
 
         static QwenImageConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
             QwenImageConfig config;
@@ -88,19 +90,33 @@ namespace Qwen {
     };
 
     struct QwenTimestepProjEmbeddings : public GGMLBlock {
+    protected:
+        bool use_additional_t_cond = false;
+
     public:
-        QwenTimestepProjEmbeddings(int64_t embedding_dim) {
+        QwenTimestepProjEmbeddings(int64_t embedding_dim, bool use_additional_t_cond = false)
+            : use_additional_t_cond(use_additional_t_cond) {
             blocks["timestep_embedder"] = std::shared_ptr<GGMLBlock>(new TimestepEmbedding(256, embedding_dim));
+            if (use_additional_t_cond) {
+                blocks["addition_t_embedding"] = std::shared_ptr<GGMLBlock>(new Embedding(2, embedding_dim));
+            }
         }
 
         ggml_tensor* forward(GGMLRunnerContext* ctx,
-                             ggml_tensor* timesteps) {
+                             ggml_tensor* timesteps,
+                             ggml_tensor* addition_t_cond = nullptr) {
             // timesteps: [N,]
             // return: [N, embedding_dim]
             auto timestep_embedder = std::dynamic_pointer_cast<TimestepEmbedding>(blocks["timestep_embedder"]);
 
             auto timesteps_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, 256, 10000, 1.f);
             auto timesteps_emb  = timestep_embedder->forward(ctx, timesteps_proj);
+            if (use_additional_t_cond) {
+                GGML_ASSERT(addition_t_cond != nullptr);
+                auto addition_t_embedding = std::dynamic_pointer_cast<Embedding>(blocks["addition_t_embedding"]);
+                auto addition_t_emb       = addition_t_embedding->forward(ctx, addition_t_cond);
+                timesteps_emb             = ggml_add(ctx->ggml_ctx, timesteps_emb, addition_t_emb);
+            }
             return timesteps_emb;
         }
     };
@@ -402,7 +418,7 @@ namespace Qwen {
         QwenImageModel(QwenImageConfig config)
             : config(config) {
             int64_t inner_dim         = config.num_attention_heads * config.attention_head_dim;
-            blocks["time_text_embed"] = std::shared_ptr<GGMLBlock>(new QwenTimestepProjEmbeddings(inner_dim));
+            blocks["time_text_embed"] = std::shared_ptr<GGMLBlock>(new QwenTimestepProjEmbeddings(inner_dim, config.use_additional_t_cond));
             blocks["txt_norm"]        = std::shared_ptr<GGMLBlock>(new RMSNorm(config.joint_attention_dim, 1e-6f));
             blocks["img_in"]          = std::shared_ptr<GGMLBlock>(new Linear(config.in_channels, inner_dim));
             blocks["txt_in"]          = std::shared_ptr<GGMLBlock>(new Linear(config.joint_attention_dim, inner_dim));
@@ -424,6 +440,7 @@ namespace Qwen {
         ggml_tensor* forward_orig(GGMLRunnerContext* ctx,
                                   ggml_tensor* x,
                                   ggml_tensor* timestep,
+                                  ggml_tensor* addition_t_cond,
                                   ggml_tensor* context,
                                   ggml_tensor* pe,
                                   ggml_tensor* modulate_index = nullptr) {
@@ -434,9 +451,9 @@ namespace Qwen {
             auto norm_out        = std::dynamic_pointer_cast<AdaLayerNormContinuous>(blocks["norm_out"]);
             auto proj_out        = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
 
-            auto t_emb = time_text_embed->forward(ctx, timestep);
+            auto t_emb = time_text_embed->forward(ctx, timestep, addition_t_cond);
             if (config.zero_cond_t) {
-                auto t_emb_0 = time_text_embed->forward(ctx, ggml_ext_zeros_like(ctx->ggml_ctx, timestep));
+                auto t_emb_0 = time_text_embed->forward(ctx, ggml_ext_zeros_like(ctx->ggml_ctx, timestep), addition_t_cond);
                 t_emb        = ggml_concat(ctx->ggml_ctx, t_emb, t_emb_0, 1);
             }
             auto img = img_in->forward(ctx, x);
@@ -469,33 +486,50 @@ namespace Qwen {
         ggml_tensor* forward(GGMLRunnerContext* ctx,
                              ggml_tensor* x,
                              ggml_tensor* timestep,
+                             ggml_tensor* addition_t_cond,
                              ggml_tensor* context,
                              ggml_tensor* pe,
                              std::vector<ggml_tensor*> ref_latents = {},
                              ggml_tensor* modulate_index           = nullptr) {
             // Forward pass of DiT.
-            // x: [N, C, H, W]
+            // x: [N, C, H, W] or [N*C, T, H, W]
             // timestep: [N,]
             // context: [N, L, D]
             // pe: [L, d_head/2, 2, 2]
-            // return: [N, C, H, W]
+            // return: [N, C, H, W] or [N*C, T, H, W]
 
-            int64_t W = x->ne[0];
-            int64_t H = x->ne[1];
-            int64_t C = x->ne[2];
-            int64_t N = x->ne[3];
+            int64_t W          = x->ne[0];
+            int64_t H          = x->ne[1];
+            int64_t T          = 1;
+            int64_t N          = addition_t_cond != nullptr ? addition_t_cond->ne[0] : x->ne[3];
+            bool has_time_axis = false;
+            if (x->ne[3] != 1) {
+                T             = x->ne[2];
+                has_time_axis = true;
+            }
 
-            auto img           = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size);
+            auto patchify_input = [&](ggml_tensor* input) -> ggml_tensor* {
+                input = DiT::pad_to_patch_size(ctx, input, config.patch_size, config.patch_size);
+                if (!has_time_axis) {
+                    return DiT::patchify(ctx->ggml_ctx, input, config.patch_size, config.patch_size);
+                }
+                if (input->ne[3] == 1) {
+                    input = ggml_reshape_4d(ctx->ggml_ctx, input, input->ne[0], input->ne[1], 1, input->ne[2]);
+                }
+                return DiT::patchify(ctx->ggml_ctx, input, 1, config.patch_size, config.patch_size, N);
+            };
+
+            auto img           = patchify_input(x);
             int64_t img_tokens = img->ne[1];
 
             if (ref_latents.size() > 0) {
                 for (ggml_tensor* ref : ref_latents) {
-                    ref = DiT::pad_and_patchify(ctx, ref, config.patch_size, config.patch_size);
+                    ref = patchify_input(ref);
                     img = ggml_concat(ctx->ggml_ctx, img, ref, 1);
                 }
             }
 
-            auto out = forward_orig(ctx, img, timestep, context, pe, modulate_index);  // [N, h_len*w_len, ph*pw*C]
+            auto out = forward_orig(ctx, img, timestep, addition_t_cond, context, pe, modulate_index);  // [N, h_len*w_len, ph*pw*C]
 
             if (out->ne[1] > img_tokens) {
                 out = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, out, 0, 2, 1, 3));  // [num_tokens, N, C * patch_size * patch_size]
@@ -503,7 +537,17 @@ namespace Qwen {
                 out = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, out, 0, 2, 1, 3));  // [N, h*w, C * patch_size * patch_size]
             }
 
-            out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, config.patch_size, config.patch_size);  // [N, C, H, W]
+            if (has_time_axis) {
+                int pad_h = (config.patch_size - H % config.patch_size) % config.patch_size;
+                int pad_w = (config.patch_size - W % config.patch_size) % config.patch_size;
+                int h_len = static_cast<int>((H + pad_h) / config.patch_size);
+                int w_len = static_cast<int>((W + pad_w) / config.patch_size);
+                out       = DiT::unpatchify_3d(ctx->ggml_ctx, out, T, h_len, w_len, 1, config.patch_size, config.patch_size);
+                out       = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, H);  // [N*C, T, H, W + pad_w]
+                out       = ggml_ext_slice(ctx->ggml_ctx, out, 0, 0, W);  // [N*C, T, H, W]
+            } else {
+                out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, config.patch_size, config.patch_size);  // [N, C, H, W]
+            }
 
             return out;
         }
@@ -515,6 +559,7 @@ namespace Qwen {
         QwenImageModel qwen_image;
         std::vector<float> pe_vec;
         std::vector<float> modulate_index_vec;
+        std::vector<int32_t> additional_t_cond_vec;
         SDVersion version;
 
         QwenImageRunner(ggml_backend_t backend,
@@ -524,9 +569,13 @@ namespace Qwen {
                         bool zero_cond_t                                    = false,
                         std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
             : DiffusionModelRunner(backend, prefix, weight_manager),
-              config(QwenImageConfig::detect_from_weights(tensor_storage_map, prefix)) {
+              config(QwenImageConfig::detect_from_weights(tensor_storage_map, prefix)),
+              version(version) {
             config.zero_cond_t = config.zero_cond_t || zero_cond_t;
-            qwen_image         = QwenImageModel(config);
+            if (version == VERSION_QWEN_IMAGE_LAYERED) {
+                config.use_additional_t_cond = true;
+            }
+            qwen_image = QwenImageModel(config);
             qwen_image.init(params_ctx, tensor_storage_map, prefix);
         }
 
@@ -542,11 +591,11 @@ namespace Qwen {
                                  const sd::Tensor<float>& timesteps_tensor,
                                  const sd::Tensor<float>& context_tensor,
                                  const std::vector<sd::Tensor<float>>& ref_latents_tensor = {},
-                                 bool increase_ref_index                                  = false) {
+                                 Rope::RefIndexMode ref_index_mode                        = Rope::RefIndexMode::INCREASE) {
             ggml_cgraph* gf        = new_graph_custom(QWEN_IMAGE_GRAPH_SIZE);
             ggml_tensor* x         = make_input(x_tensor);
             ggml_tensor* timesteps = make_input(timesteps_tensor);
-            GGML_ASSERT(x->ne[3] == 1);
+            GGML_ASSERT(x->ne[3] == 1 || x_tensor.dim() == 5);
             GGML_ASSERT(!context_tensor.empty());
             ggml_tensor* context = make_input(context_tensor);
             std::vector<ggml_tensor*> ref_latents;
@@ -555,13 +604,29 @@ namespace Qwen {
                 ref_latents.push_back(make_input(ref_latent_tensor));
             }
 
-            pe_vec      = Rope::gen_qwen_image_pe(static_cast<int>(x->ne[1]),
+            int batch_size = static_cast<int>(x->ne[3]);
+            int time_len   = 1;
+            if (x_tensor.dim() == 5) {
+                time_len   = static_cast<int>(x_tensor.shape()[2]);
+                batch_size = static_cast<int>(x_tensor.shape()[4]);
+            }
+
+            ggml_tensor* addition_t_cond = nullptr;
+            if (version == VERSION_QWEN_IMAGE_LAYERED) {
+                additional_t_cond_vec.assign(static_cast<size_t>(batch_size), 0);
+                addition_t_cond = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_I32, batch_size);
+                set_backend_tensor_data(addition_t_cond, additional_t_cond_vec.data());
+                ref_index_mode = Rope::RefIndexMode::DECREASE;
+            }
+
+            pe_vec      = Rope::gen_qwen_image_pe(time_len,
+                                                  static_cast<int>(x->ne[1]),
                                                   static_cast<int>(x->ne[0]),
                                                   config.patch_size,
-                                                  static_cast<int>(x->ne[3]),
+                                                  batch_size,
                                                   static_cast<int>(context->ne[1]),
                                                   ref_latents,
-                                                  increase_ref_index,
+                                                  ref_index_mode,
                                                   config.theta,
                                                   circular_y_enabled,
                                                   circular_x_enabled,
@@ -604,6 +669,7 @@ namespace Qwen {
             ggml_tensor* out = qwen_image.forward(&runner_ctx,
                                                   x,
                                                   timesteps,
+                                                  addition_t_cond,
                                                   context,
                                                   pe,
                                                   ref_latents,
@@ -619,12 +685,12 @@ namespace Qwen {
                                   const sd::Tensor<float>& timesteps,
                                   const sd::Tensor<float>& context,
                                   const std::vector<sd::Tensor<float>>& ref_latents = {},
-                                  bool increase_ref_index                           = false) {
-            // x: [N, in_channels, h, w]
+                                  Rope::RefIndexMode ref_index_mode                 = Rope::RefIndexMode::INCREASE) {
+            // x: [N, C, H, W] or [N*C, T, H, W]
             // timesteps: [N, ]
             // context: [N, max_position, hidden_size]
             auto get_graph = [&]() -> ggml_cgraph* {
-                return build_graph(x, timesteps, context, ref_latents, increase_ref_index);
+                return build_graph(x, timesteps, context, ref_latents, ref_index_mode);
             };
 
             return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
@@ -640,7 +706,7 @@ namespace Qwen {
                            *diffusion_params.timesteps,
                            tensor_or_empty(diffusion_params.context),
                            diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
-                           diffusion_params.increase_ref_index);
+                           diffusion_params.ref_index_mode);
         }
 
         void test() {
@@ -674,7 +740,7 @@ namespace Qwen {
                                        timesteps,
                                        context,
                                        {},
-                                       false);
+                                       Rope::RefIndexMode::FIXED);
                 int64_t t1   = ggml_time_ms();
 
                 GGML_ASSERT(!out_opt.empty());

src/model/diffusion/z_image.hpp

@@ -575,7 +575,7 @@ namespace ZImage {
                                  const sd::Tensor<float>& timesteps_tensor,
                                  const sd::Tensor<float>& context_tensor,
                                  const std::vector<sd::Tensor<float>>& ref_latents_tensor = {},
-                                 bool increase_ref_index                                  = false) {
+                                 Rope::RefIndexMode ref_index_mode                        = Rope::RefIndexMode::FIXED) {
             ggml_cgraph* gf        = new_graph_custom(Z_IMAGE_GRAPH_SIZE);
             ggml_tensor* x         = make_input(x_tensor);
             ggml_tensor* timesteps = make_input(timesteps_tensor);
@@ -595,7 +595,7 @@ namespace ZImage {
                                                static_cast<int>(context->ne[1]),
                                                SEQ_MULTI_OF,
                                                ref_latents,
-                                               increase_ref_index,
+                                               ref_index_mode,
                                                config.theta,
                                                circular_y_enabled,
                                                circular_x_enabled,
@@ -626,12 +626,12 @@ namespace ZImage {
                                   const sd::Tensor<float>& timesteps,
                                   const sd::Tensor<float>& context,
                                   const std::vector<sd::Tensor<float>>& ref_latents = {},
-                                  bool increase_ref_index                           = false) {
+                                  Rope::RefIndexMode ref_index_mode                 = Rope::RefIndexMode::FIXED) {
             // x: [N, in_channels, h, w]
             // timesteps: [N, ]
             // context: [N, max_position, hidden_size]
             auto get_graph = [&]() -> ggml_cgraph* {
-                return build_graph(x, timesteps, context, ref_latents, increase_ref_index);
+                return build_graph(x, timesteps, context, ref_latents, ref_index_mode);
             };
 
             return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
@@ -647,7 +647,7 @@ namespace ZImage {
                            *diffusion_params.timesteps,
                            tensor_or_empty(diffusion_params.context),
                            diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
-                           diffusion_params.increase_ref_index);
+                           diffusion_params.ref_index_mode);
         }
 
         void test() {
@@ -681,7 +681,7 @@ namespace ZImage {
                                        timesteps,
                                        context,
                                        {},
-                                       false);
+                                       Rope::RefIndexMode::FIXED);
                 int64_t t1   = ggml_time_ms();
 
                 GGML_ASSERT(!out_opt.empty());

src/model/vae/ltx_audio_vae.hpp

@@ -214,7 +214,7 @@ namespace LTXV {
 
         auto x = ggml_reshape_3d(ctx, waveform, time, 1, channels * batch);
         if (left_pad > 0) {
-            x = ggml_pad_ext(ctx, x, static_cast<int>(left_pad), 0, 0, 0, 0, 0, 0, 0);
+            x = ggml_ext_pad_ext(ctx, runner_ctx->backend, x, static_cast<int>(left_pad), 0, 0, 0, 0, 0, 0, 0);
         }
 
         auto frames = ggml_conv_1d(ctx, forward_basis, x, hop_length, 0, 1);
@@ -451,6 +451,7 @@ namespace LTXV {
             int pad_h = kernel_size.first - 1;
             int pad_w = kernel_size.second - 1;
             x         = ggml_ext_pad_ext(ctx->ggml_ctx,
+                                         ctx->backend,
                                          x,
                                          pad_w / 2,
                                          pad_w - pad_w / 2,

src/model/vae/tae.hpp

@@ -408,7 +408,7 @@ public:
             h = conv->forward(ctx, h);
             for (int j = 0; j < num_blocks; j++) {
                 auto block = std::dynamic_pointer_cast<MemBlock>(blocks[std::to_string(index++)]);
-                auto mem   = ggml_pad_ext(ctx->ggml_ctx, h, 0, 0, 0, 0, 0, 0, 1, 0);
+                auto mem   = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, h, 0, 0, 0, 0, 0, 0, 1, 0);
                 mem        = ggml_view_4d(ctx->ggml_ctx, mem, h->ne[0], h->ne[1], h->ne[2], h->ne[3], h->nb[1], h->nb[2], h->nb[3], 0);
                 h          = block->forward(ctx, h, mem);
             }
@@ -479,7 +479,7 @@ public:
         int index = 3;
         for (int i = 0; i < num_layers; i++) {
             for (int j = 0; j < num_blocks; j++) {
-                auto mem = ggml_pad_ext(ctx->ggml_ctx, h, 0, 0, 0, 0, 0, 0, 1, 0);
+                auto mem = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, h, 0, 0, 0, 0, 0, 0, 1, 0);
                 mem      = ggml_view_4d(ctx->ggml_ctx, mem, h->ne[0], h->ne[1], h->ne[2], h->ne[3], h->nb[1], h->nb[2], h->nb[3], 0);
                 if (is_wide) {
                     auto block = std::dynamic_pointer_cast<WideMemBlock>(blocks[std::to_string(index++)]);

src/model/vae/wan_vae.hpp

@@ -72,8 +72,8 @@ namespace WAN {
                 lp2 -= (int)cache_x->ne[2];
             }
 
-            x = ggml_ext_pad_ext(ctx->ggml_ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
-            return ggml_ext_conv_3d(ctx->ggml_ctx, x, w, b, in_channels,
+            x = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, x, lp0, rp0, lp1, rp1, lp2, rp2, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
+            return ggml_ext_conv_3d(ctx->ggml_ctx, ctx->backend, x, w, b, in_channels,
                                     std::get<2>(stride), std::get<1>(stride), std::get<0>(stride),
                                     0, 0, 0,
                                     std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation));
@@ -195,7 +195,7 @@ namespace WAN {
                                                   2);
                         }
                         if (chunk_idx == 1 && cache_x->ne[2] < 2) {  // Rep
-                            cache_x = ggml_pad_ext(ctx->ggml_ctx, cache_x, 0, 0, 0, 0, (int)cache_x->ne[2], 0, 0, 0);
+                            cache_x = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, cache_x, 0, 0, 0, 0, (int)cache_x->ne[2], 0, 0, 0);
                             // aka cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device),cache_x],dim=2)
                         }
                         if (chunk_idx == 1) {
@@ -283,7 +283,7 @@ namespace WAN {
 
             int pad_t = (factor_t - T % factor_t) % factor_t;
 
-            x = ggml_pad_ext(ctx->ggml_ctx, x, 0, 0, 0, 0, pad_t, 0, 0, 0);
+            x = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, x, 0, 0, 0, 0, pad_t, 0, 0, 0);
             T = x->ne[2];
 
             x = ggml_reshape_4d(ctx->ggml_ctx, x, W * H, factor_t, T / factor_t, C);                                                  // [C, T/factor_t, factor_t, H*W]
@@ -631,6 +631,7 @@ namespace WAN {
     class Encoder3d : public GGMLBlock {
     protected:
         bool wan2_2;
+        int64_t in_channels;
         int64_t dim;
         int64_t z_dim;
         std::vector<int> dim_mult;
@@ -641,12 +642,14 @@ namespace WAN {
     public:
         Encoder3d(int64_t dim                           = 128,
                   int64_t z_dim                         = 4,
+                  int64_t in_channels                   = 3,
                   std::vector<int> dim_mult             = {1, 2, 4, 4},
                   int num_res_blocks                    = 2,
                   std::vector<bool> temperal_downsample = {false, true, true},
                   bool wan2_2                           = false,
                   bool is_2D                            = false)
-            : dim(dim),
+            : in_channels(in_channels),
+              dim(dim),
               z_dim(z_dim),
               dim_mult(dim_mult),
               num_res_blocks(num_res_blocks),
@@ -659,11 +662,10 @@ namespace WAN {
                 dims.push_back(dim * u);
             }
 
-            int64_t input_dim = wan2_2 ? 12 : 3;
             if (is_2D) {
-                blocks["conv1"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(input_dim, dims[0], {3, 3}, {1, 1}, {1, 1}));
+                blocks["conv1"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(in_channels, dims[0], {3, 3}, {1, 1}, {1, 1}));
             } else {
-                blocks["conv1"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(input_dim, dims[0], {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
+                blocks["conv1"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(in_channels, dims[0], {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
             }
 
             int index = 0;
@@ -812,6 +814,7 @@ namespace WAN {
     class Decoder3d : public GGMLBlock {
     protected:
         bool wan2_2;
+        int64_t out_channels;
         int64_t dim;
         int64_t z_dim;
         std::vector<int> dim_mult;
@@ -822,12 +825,14 @@ namespace WAN {
     public:
         Decoder3d(int64_t dim                         = 128,
                   int64_t z_dim                       = 4,
+                  int64_t out_channels                = 3,
                   std::vector<int> dim_mult           = {1, 2, 4, 4},
                   int num_res_blocks                  = 2,
                   std::vector<bool> temperal_upsample = {true, true, false},
                   bool wan2_2                         = false,
                   bool is_2D                          = false)
-            : dim(dim),
+            : out_channels(out_channels),
+              dim(dim),
               z_dim(z_dim),
               dim_mult(dim_mult),
               num_res_blocks(num_res_blocks),
@@ -889,7 +894,7 @@ namespace WAN {
 
             // output blocks
             blocks["head.0"]  = std::shared_ptr<GGMLBlock>(new RMS_norm(out_dim));
-            int64_t final_dim = wan2_2 ? 12 : 3;
+            int64_t final_dim = out_channels;
             // head.1 is nn.SiLU()
             if (is_2D) {
                 blocks["head.2"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(out_dim, final_dim, {3, 3}, {1, 1}, {1, 1}));
@@ -1002,6 +1007,8 @@ namespace WAN {
     public:
         bool wan2_2                           = false;
         bool decode_only                      = true;
+        int64_t input_channels                = 3;
+        int patch_size                        = 1;
         int64_t dim                           = 96;
         int64_t dec_dim                       = 96;
         int64_t z_dim                         = 16;
@@ -1026,16 +1033,22 @@ namespace WAN {
         }
 
     public:
-        WanVAE(bool decode_only = true, bool wan2_2 = false, bool is_2D = false)
-            : decode_only(decode_only), wan2_2(wan2_2), is_2D(is_2D) {
+        WanVAE(bool decode_only = true, SDVersion version = VERSION_WAN2, bool is_2D = false)
+            : decode_only(decode_only),
+              wan2_2(version == VERSION_WAN2_2_TI2V),
+              is_2D(is_2D) {
             // attn_scales is always []
             if (wan2_2) {
-                dim     = 160;
-                dec_dim = 256;
-                z_dim   = 48;
+                dim            = 160;
+                dec_dim        = 256;
+                z_dim          = 48;
+                input_channels = 12;
+                patch_size     = 2;
 
                 _conv_num     = 34;
                 _enc_conv_num = 26;
+            } else if (version == VERSION_QWEN_IMAGE_LAYERED) {
+                input_channels = 4;
             }
 
             if (is_2D) {
@@ -1044,14 +1057,14 @@ namespace WAN {
             }
 
             if (!decode_only) {
-                blocks["encoder"] = std::shared_ptr<GGMLBlock>(new Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks, temperal_downsample, wan2_2, is_2D));
+                blocks["encoder"] = std::shared_ptr<GGMLBlock>(new Encoder3d(dim, z_dim * 2, input_channels, dim_mult, num_res_blocks, temperal_downsample, wan2_2, is_2D));
                 if (is_2D) {
                     blocks["conv1"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(z_dim * 2, z_dim * 2, {1, 1}));
                 } else {
                     blocks["conv1"] = std::shared_ptr<GGMLBlock>(new CausalConv3d(z_dim * 2, z_dim * 2, {1, 1, 1}));
                 }
             }
-            blocks["decoder"] = std::shared_ptr<GGMLBlock>(new Decoder3d(dec_dim, z_dim, dim_mult, num_res_blocks, temperal_upsample, wan2_2, is_2D));
+            blocks["decoder"] = std::shared_ptr<GGMLBlock>(new Decoder3d(dec_dim, z_dim, input_channels, dim_mult, num_res_blocks, temperal_upsample, wan2_2, is_2D));
             if (is_2D) {
                 blocks["conv2"] = std::shared_ptr<GGMLBlock>(new Conv2dBut3d(z_dim, z_dim, {1, 1}));
             } else {
@@ -1125,9 +1138,7 @@ namespace WAN {
 
             clear_cache();
 
-            if (wan2_2) {
-                x = patchify(ctx->ggml_ctx, x, 2, b);
-            }
+            x = patchify(ctx->ggml_ctx, x, patch_size, b);
             // sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.encode.prelude", "x");
 
             auto encoder = std::dynamic_pointer_cast<Encoder3d>(blocks["encoder"]);
@@ -1202,9 +1213,7 @@ namespace WAN {
                     }
                 }
             }
-            if (wan2_2) {
-                out = unpatchify(ctx->ggml_ctx, out, 2, b);
-            }
+            out = unpatchify(ctx->ggml_ctx, out, patch_size, b);
             // sd::ggml_graph_cut::mark_graph_cut(out, "wan_vae.decode.final", "out");
             clear_cache();
             return out;
@@ -1225,9 +1234,7 @@ namespace WAN {
             auto in   = ggml_ext_slice(ctx->ggml_ctx, x, 2, i, i + 1);  // [b*c, 1, h, w]
             _conv_idx = 0;
             auto out  = decoder->forward(ctx, in, b, _feat_map, _conv_idx, i);
-            if (wan2_2) {
-                out = unpatchify(ctx->ggml_ctx, out, 2, b);
-            }
+            out       = unpatchify(ctx->ggml_ctx, out, patch_size, b);
             // sd::ggml_graph_cut::mark_graph_cut(out, "wan_vae.decode_partial.final", "out");
             return out;
         }
@@ -1257,7 +1264,7 @@ namespace WAN {
             if (is_2D) {
                 LOG_DEBUG("USING 2D VAE");
             }
-            ae = WanVAE(decode_only, version == VERSION_WAN2_2_TI2V, is_2D);
+            ae = WanVAE(decode_only, version, is_2D);
             ae.init(params_ctx, tensor_storage_map, prefix);
         }
 

src/model_loader.cpp

@@ -476,6 +476,9 @@ SDVersion ModelLoader::get_sd_version() {
             return VERSION_MINIT2I;
         }
         if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.img_mod.1.weight") != std::string::npos) {
+            if (tensor_storage_map.find("model.diffusion_model.time_text_embed.addition_t_embedding.weight") != tensor_storage_map.end()) {
+                return VERSION_QWEN_IMAGE_LAYERED;
+            }
             return VERSION_QWEN_IMAGE;
         }
         if (tensor_storage.name.find("llm_adapter.blocks.0.cross_attn.q_proj.weight") != std::string::npos) {

src/model_loader.h

@@ -27,6 +27,8 @@ struct MmapTensorStore {
     std::shared_ptr<struct ggml_backend_buffer> mmbuffer;
 };
 
+bool is_unused_tensor(const std::string& name);
+
 class ModelLoader {
 protected:
     SDVersion version_ = VERSION_COUNT;

src/model_manager.cpp

@@ -100,12 +100,41 @@ size_t estimate_tensors_size(const std::map<std::string, ggml_tensor*>& tensors)
     return size;
 }
 
+void ModelManager::set_split_buffer_type(ggml_backend_t compute_backend, ggml_backend_buffer_type_t split_buft) {
+    if (compute_backend == nullptr) {
+        return;
+    }
+    if (split_buft == nullptr) {
+        split_buffer_types_.erase(compute_backend);
+        return;
+    }
+    split_buffer_types_[compute_backend] = split_buft;
+}
+
+bool ModelManager::tensor_shape_supports_split_buffer(const ggml_tensor* tensor) {
+    return tensor != nullptr &&
+           tensor->view_src == nullptr &&
+           ggml_is_contiguous(tensor) &&
+           ggml_n_dims(tensor) == 2 &&
+           tensor->ne[0] >= 256 &&
+           tensor->ne[1] >= 256;
+}
+
+ggml_backend_buffer_type_t ModelManager::split_buffer_type_for(const TensorState& state) const {
+    if (!state.allow_split_buffer || !tensor_shape_supports_split_buffer(state.tensor)) {
+        return nullptr;
+    }
+    auto it = split_buffer_types_.find(state.compute_backend);
+    return it != split_buffer_types_.end() ? it->second : nullptr;
+}
+
 bool ModelManager::register_param_tensors(const std::string& desc,
                                           std::map<std::string, ggml_tensor*> tensors,
                                           ResidencyMode residency_mode,
                                           ggml_backend_t compute_backend,
                                           ggml_backend_t params_backend,
-                                          size_t* registered_tensor_size) {
+                                          size_t* registered_tensor_size,
+                                          bool allow_split_buffer) {
     if (desc.empty()) {
         LOG_ERROR("model manager tensor desc is empty");
         return false;
@@ -129,13 +158,14 @@ bool ModelManager::register_param_tensors(const std::string& desc,
         }
         ggml_set_name(tensor, name.c_str());
 
-        auto state             = std::make_unique<TensorState>();
-        state->name            = name;
-        state->tensor          = tensor;
-        state->desc            = desc;
-        state->residency_mode  = residency_mode;
-        state->compute_backend = compute_backend;
-        state->params_backend  = params_backend;
+        auto state                = std::make_unique<TensorState>();
+        state->name               = name;
+        state->tensor             = tensor;
+        state->desc               = desc;
+        state->residency_mode     = residency_mode;
+        state->compute_backend    = compute_backend;
+        state->params_backend     = params_backend;
+        state->allow_split_buffer = allow_split_buffer;
         new_states.push_back(std::move(state));
     }
 
@@ -237,7 +267,7 @@ bool ModelManager::load_tensors_to_params_backend(const std::vector<TensorState*
 }
 
 bool ModelManager::stage_tensors_to_compute_backend(const std::vector<TensorState*>& states) {
-    std::map<ggml_backend_t, std::vector<TensorState*>> states_by_compute_backend;
+    std::map<std::pair<ggml_backend_t, ggml_backend_buffer_type_t>, std::vector<TensorState*>> states_by_staging_target;
     for (TensorState* state : states) {
         if (state == nullptr || should_ignore(*state) || is_optional_missing_tensor(state->name)) {
             continue;
@@ -257,11 +287,16 @@ bool ModelManager::stage_tensors_to_compute_backend(const std::vector<TensorStat
             LOG_ERROR("model manager tensor '%s' is not loaded to params backend", state->name.c_str());
             return false;
         }
-        states_by_compute_backend[state->compute_backend].push_back(state);
+        ggml_backend_buffer_type_t staging_buft = split_buffer_type_for(*state);
+        if (staging_buft == nullptr) {
+            staging_buft = ggml_backend_get_default_buffer_type(state->compute_backend);
+        }
+        states_by_staging_target[{state->compute_backend, staging_buft}].push_back(state);
     }
 
-    for (const auto& pair : states_by_compute_backend) {
-        ggml_backend_t compute_backend          = pair.first;
+    for (const auto& pair : states_by_staging_target) {
+        ggml_backend_t compute_backend          = pair.first.first;
+        ggml_backend_buffer_type_t staging_buft = pair.first.second;
         const std::vector<TensorState*>& states = pair.second;
         if (states.empty()) {
             continue;
@@ -285,7 +320,7 @@ bool ModelManager::stage_tensors_to_compute_backend(const std::vector<TensorStat
             staged_tensors.push_back({state, staging_tensor});
         }
 
-        ggml_backend_buffer_t compute_buffer = ggml_backend_alloc_ctx_tensors(staging_ctx, compute_backend);
+        ggml_backend_buffer_t compute_buffer = ggml_backend_alloc_ctx_tensors_from_buft(staging_ctx, staging_buft);
         if (compute_buffer == nullptr) {
             LOG_ERROR("model manager alloc compute params backend buffer failed, num_tensors = %zu",
                       staged_tensors.size());
@@ -350,6 +385,17 @@ bool ModelManager::apply_loras_to_params(const std::vector<TensorState*>& states
             LOG_ERROR("model manager compute backend is null for lora target tensor '%s'", state->name.c_str());
             return false;
         }
+        if (state->tensor->buffer != nullptr &&
+            ggml_backend_buffer_get_type(state->tensor->buffer) == split_buffer_type_for(*state)) {
+            if (!warned_split_lora_skip_) {
+                LOG_WARN(
+                    "model manager skipping direct lora application to row-split tensors "
+                    "(use --lora-apply-mode at_runtime with row split)");
+                warned_split_lora_skip_ = true;
+            }
+            state->applied_lora_epoch = current_lora_epoch_;
+            continue;
+        }
         if (state->tensor->data == nullptr) {
             LOG_ERROR("model manager lora target tensor '%s' is not prepared", state->name.c_str());
             return false;
@@ -694,6 +740,8 @@ ggml_backend_buffer_type_t ModelManager::params_buffer_type_for(const TensorStat
         if (compute_dev != nullptr) {
             params_buft = ggml_backend_dev_host_buffer_type(compute_dev);
         }
+    } else if (state.params_backend == state.compute_backend) {
+        params_buft = split_buffer_type_for(state);
     }
     if (params_buft == nullptr) {
         params_buft = ggml_backend_get_default_buffer_type(state.params_backend);

src/model_manager.h

@@ -36,6 +36,7 @@ private:
         ResidencyMode residency_mode   = ResidencyMode::ParamBackend;
         ggml_backend_t compute_backend = nullptr;
         ggml_backend_t params_backend  = nullptr;
+        bool allow_split_buffer        = false;
         bool metadata_validated        = false;
 
         int active_prepare_count = 0;
@@ -63,6 +64,8 @@ private:
     std::map<std::string, TensorState*> tensor_states_by_name_;
     std::vector<std::unique_ptr<ParamsStorageBlock>> params_storage_blocks_;
     std::vector<std::unique_ptr<ComputeStagingBlock>> compute_staging_blocks_;
+    std::map<ggml_backend_t, ggml_backend_buffer_type_t> split_buffer_types_;
+    bool warned_split_lora_skip_ = false;
     std::set<std::string> common_ignore_tensors_;
     std::vector<LoraSpec> loras_;
     SDVersion lora_version_      = VERSION_COUNT;
@@ -91,6 +94,7 @@ private:
     bool stage_tensors_to_compute_backend(const std::vector<TensorState*>& states);
 
     ggml_backend_buffer_type_t params_buffer_type_for(const TensorState& state) const;
+    ggml_backend_buffer_type_t split_buffer_type_for(const TensorState& state) const;
     void release_compute_staging_blocks(bool force                                            = false,
                                         const std::unordered_set<TensorState*>* target_states = nullptr);
     void release_params_storage_blocks(bool force                                            = false,
@@ -114,6 +118,9 @@ public:
     void set_writable_mmap(bool writable_mmap) { writable_mmap_ = writable_mmap; }
     void set_common_ignore_tensors(std::set<std::string> ignore_tensors);
     void set_loras(std::vector<LoraSpec> loras, SDVersion version);
+    void set_split_buffer_type(ggml_backend_t compute_backend, ggml_backend_buffer_type_t split_buft);
+
+    static bool tensor_shape_supports_split_buffer(const ggml_tensor* tensor);
 
     std::set<std::string> tensor_names() const;
 
@@ -122,7 +129,8 @@ public:
                                 ResidencyMode residency_mode,
                                 ggml_backend_t compute_backend,
                                 ggml_backend_t params_backend,
-                                size_t* registered_tensor_size = nullptr);
+                                size_t* registered_tensor_size = nullptr,
+                                bool allow_split_buffer        = false);
 
     template <typename Runner>
     bool register_runner_params(const std::string& desc,

src/runtime/denoiser.hpp

@@ -1814,6 +1814,52 @@ static sd::Tensor<float> sample_dpmpp_2m_v2(denoise_cb_t model,
     return x;
 }
 
+// DPM-Solver++(2M) SDE, midpoint variant. Ref: Lu et al. arXiv:2211.01095;
+// k-diffusion sample_dpmpp_2m_sde.
+static sd::Tensor<float> sample_dpmpp_2m_sde(denoise_cb_t model,
+                                             sd::Tensor<float> x,
+                                             const std::vector<float>& sigmas,
+                                             std::shared_ptr<RNG> rng,
+                                             float eta) {
+    sd::Tensor<float> old_denoised;
+    bool have_old_denoised = false;
+    float h_last           = 0.f;
+
+    int steps = static_cast<int>(sigmas.size()) - 1;
+    for (int i = 0; i < steps; i++) {
+        auto denoised_opt = model(x, sigmas[i], i + 1);
+        if (denoised_opt.pred.empty()) {
+            return {};
+        }
+        sd::Tensor<float> denoised = std::move(denoised_opt.pred);
+
+        if (sigmas[i + 1] == 0.f) {
+            x = denoised;
+        } else {
+            float t     = -std::log(sigmas[i]);
+            float s     = -std::log(sigmas[i + 1]);
+            float h     = s - t;
+            float eta_h = eta * h;
+            float a     = sigmas[i + 1] / sigmas[i] * std::exp(-eta_h);
+            float b     = -std::expm1(-h - eta_h);
+
+            x = a * x + b * denoised;
+
+            if (have_old_denoised) {
+                float r = h_last / h;
+                x += (0.5f * b / r) * (denoised - old_denoised);
+            }
+            if (eta > 0.f) {
+                x += sd::Tensor<float>::randn_like(x, rng) * (sigmas[i + 1] * std::sqrt(-std::expm1(-2.f * eta_h)));
+            }
+            h_last = h;
+        }
+        old_denoised      = denoised;
+        have_old_denoised = true;
+    }
+    return x;
+}
+
 using SamplerExtraArgs = KeyValueArgs;
 
 static sd::Tensor<float> sample_lcm(denoise_cb_t model,
@@ -2490,6 +2536,8 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
             return sample_res_2s(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
         case ER_SDE_SAMPLE_METHOD:
             return sample_er_sde(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
+        case DPMPP2M_SDE_SAMPLE_METHOD:
+            return sample_dpmpp_2m_sde(model, std::move(x), sigmas, rng, eta);
         case DDIM_TRAILING_SAMPLE_METHOD:
             // DDIM is equivalent to Euler Ancestral with the Simple scheduler
             return sample_euler_ancestral(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);

src/stable-diffusion.cpp

@@ -2,11 +2,14 @@
 #include <cmath>
 #include <cstdlib>
 #include <set>
+#include <type_traits>
 #include <unordered_set>
+#include <utility>
 #include <vector>
 
 #include "core/ggml_extend.hpp"
 #include "core/ggml_graph_cut.h"
+#include "core/layer_split_partition.h"
 
 #include "core/rng.hpp"
 #include "core/rng_mt19937.hpp"
@@ -17,6 +20,7 @@
 #include "stable-diffusion.h"
 
 #include "conditioning/conditioner.hpp"
+#include "core/backend_fit.h"
 #include "extensions/generation_extension.h"
 #include "model/adapter/lora.hpp"
 #include "model/diffusion/anima.hpp"
@@ -84,6 +88,7 @@ const char* model_version_to_str[] = {
     "Wan 2.2 I2V",
     "Wan 2.2 TI2V",
     "Qwen Image",
+    "Qwen Image Layered",
     "Anima",
     "Flux.2",
     "Flux.2 klein",
@@ -169,6 +174,13 @@ static float get_cache_reuse_threshold(const sd_cache_params_t& params) {
 
 /*=============================================== StableDiffusionGGML ================================================*/
 
+template <typename T, typename = void>
+struct has_set_runtime_backends : std::false_type {};
+template <typename T>
+struct has_set_runtime_backends<T,
+                                std::void_t<decltype(std::declval<T&>().set_runtime_backends(
+                                    std::declval<const std::vector<ggml_backend_t>&>()))>> : std::true_type {};
+
 static_assert(std::atomic<sd_cancel_mode_t>::is_always_lock_free,
               "sd_cancel_mode_t must be lock-free");
 
@@ -207,6 +219,8 @@ public:
     bool eager_load    = false;
     std::string backend_spec;
     std::string params_backend_spec;
+    std::string split_mode_spec;
+    bool auto_fit_enabled = false;
 
     bool is_using_v_parameterization     = false;
     bool is_using_edm_v_parameterization = false;
@@ -274,18 +288,230 @@ public:
         if (model_manager == nullptr) {
             return true;
         }
+        ModelManager::ResidencyMode residency_mode =
+            backend_manager.params_backend_is_disk(module) ? ModelManager::ResidencyMode::Disk : ModelManager::ResidencyMode::ParamBackend;
+
+        std::vector<ggml_backend_t> module_backends = backend_manager.runtime_backends(module);
+        if (module_backends.size() > 1) {
+            if constexpr (has_set_runtime_backends<T>::value) {
+                if (module == SDBackendModule::DIFFUSION || module == SDBackendModule::TE) {
+                    if (backend_manager.split_mode(module) == SDSplitMode::ROW) {
+                        return register_row_split_runner_params(desc,
+                                                                model,
+                                                                module,
+                                                                module_backends,
+                                                                std::move(group_tensors),
+                                                                residency_mode,
+                                                                params_mem_size);
+                    }
+                    return register_layer_split_runner_params(desc,
+                                                              model,
+                                                              module,
+                                                              module_backends,
+                                                              std::move(group_tensors),
+                                                              residency_mode,
+                                                              params_mem_size);
+                }
+            }
+            LOG_WARN("%s module does not support multiple runtime backends; using %s",
+                     sd_backend_module_name(module),
+                     sd::layer_split_backend_device_display_name(module_backends[0]).c_str());
+        }
         return model_manager->register_param_tensors(desc,
                                                      std::move(group_tensors),
-                                                     backend_manager.params_backend_is_disk(module) ? ModelManager::ResidencyMode::Disk : ModelManager::ResidencyMode::ParamBackend,
+                                                     residency_mode,
                                                      backend_for(module),
                                                      params_backend_for(module),
                                                      params_mem_size);
     }
 
+    template <typename T>
+    bool register_row_split_runner_params(const std::string& desc,
+                                          const std::shared_ptr<T>& model,
+                                          SDBackendModule module,
+                                          const std::vector<ggml_backend_t>& module_backends,
+                                          std::map<std::string, ggml_tensor*> group_tensors,
+                                          ModelManager::ResidencyMode residency_mode,
+                                          size_t* params_mem_size) {
+        ggml_backend_t main_backend = module_backends[0];
+
+        auto fall_back_to_layer_split = [&](const char* reason) {
+            LOG_WARN("%s: row split unavailable (%s); falling back to layer split", desc.c_str(), reason);
+            return register_layer_split_runner_params(desc,
+                                                      model,
+                                                      module,
+                                                      module_backends,
+                                                      std::move(group_tensors),
+                                                      residency_mode,
+                                                      params_mem_size);
+        };
+
+        ggml_backend_dev_t main_dev = ggml_backend_get_device(main_backend);
+        ggml_backend_reg_t reg      = main_dev != nullptr ? ggml_backend_dev_backend_reg(main_dev) : nullptr;
+        if (reg == nullptr) {
+            return fall_back_to_layer_split("no backend registry");
+        }
+        const size_t reg_dev_count = ggml_backend_reg_dev_count(reg);
+        std::vector<float> tensor_split(reg_dev_count, 0.0f);
+        constexpr int64_t compute_headroom_bytes = 2ll * 1024 * 1024 * 1024;
+        for (ggml_backend_t backend : module_backends) {
+            ggml_backend_dev_t dev = ggml_backend_get_device(backend);
+            int reg_index          = -1;
+            for (size_t i = 0; i < reg_dev_count; i++) {
+                if (ggml_backend_reg_dev_get(reg, i) == dev) {
+                    reg_index = (int)i;
+                    break;
+                }
+            }
+            if (reg_index < 0) {
+                return fall_back_to_layer_split("devices span different backend registries");
+            }
+            size_t free_bytes = 0, total_bytes = 0;
+            ggml_backend_dev_memory(dev, &free_bytes, &total_bytes);
+            int64_t usable_bytes    = std::max<int64_t>((int64_t)free_bytes - compute_headroom_bytes,
+                                                     (int64_t)free_bytes / 8);
+            tensor_split[reg_index] = usable_bytes > 0 ? (float)((double)usable_bytes / (1024.0 * 1024.0)) : 1.0f;
+        }
+
+        ggml_backend_buffer_type_t split_buft = backend_manager.split_buffer_type(main_backend, tensor_split);
+        if (split_buft == nullptr) {
+            return fall_back_to_layer_split("backend has no split buffer type");
+        }
+        model_manager->set_split_buffer_type(main_backend, split_buft);
+
+        std::map<std::string, ggml_tensor*> split_tensors;
+        if constexpr (std::is_base_of_v<Conditioner, T>) {
+            model->get_layer_split_param_tensors(split_tensors);
+        } else {
+            split_tensors = group_tensors;
+        }
+
+        std::map<std::string, ggml_tensor*> row_split_map;
+        std::map<std::string, ggml_tensor*> regular_map;
+        size_t row_split_bytes = 0;
+        for (const auto& kv : group_tensors) {
+            if (split_tensors.count(kv.first) != 0 &&
+                sd::layer_split_tensor_block_index(kv.first) >= 0 &&
+                ModelManager::tensor_shape_supports_split_buffer(kv.second)) {
+                row_split_map[kv.first] = kv.second;
+                row_split_bytes += ggml_nbytes(kv.second);
+            } else {
+                regular_map[kv.first] = kv.second;
+            }
+        }
+        if (row_split_map.empty()) {
+            return fall_back_to_layer_split("no row-splittable transformer block weights found");
+        }
+
+        LOG_INFO("%s row split: %zu tensors (%.1f MB) split across %zu devices (main %s)",
+                 desc.c_str(),
+                 row_split_map.size(),
+                 row_split_bytes / (1024.f * 1024.f),
+                 module_backends.size(),
+                 sd::layer_split_backend_device_display_name(main_backend).c_str());
+
+        if (!model_manager->register_param_tensors(desc,
+                                                   std::move(row_split_map),
+                                                   residency_mode,
+                                                   main_backend,
+                                                   params_backend_for(module),
+                                                   params_mem_size,
+                                                   /*allow_split_buffer=*/true)) {
+            return false;
+        }
+        return model_manager->register_param_tensors(desc,
+                                                     std::move(regular_map),
+                                                     residency_mode,
+                                                     main_backend,
+                                                     params_backend_for(module),
+                                                     params_mem_size);
+    }
+
+    // Register each layer-split partition with its compute backend; the
+    // ModelManager handles allocation, staging, and LoRA by backend.
+    template <typename T>
+    bool register_layer_split_runner_params(const std::string& desc,
+                                            const std::shared_ptr<T>& model,
+                                            SDBackendModule module,
+                                            const std::vector<ggml_backend_t>& module_backends,
+                                            std::map<std::string, ggml_tensor*> group_tensors,
+                                            ModelManager::ResidencyMode residency_mode,
+                                            size_t* params_mem_size) {
+        bool has_cpu_device = false;
+        for (ggml_backend_t backend : module_backends) {
+            has_cpu_device = has_cpu_device || sd_backend_is_cpu(backend);
+        }
+        if (has_cpu_device) {
+            // The scheduler reserves the CPU slot for its fallback backend, and
+            // CPU weight participation is what --params-backend <module>=cpu is
+            // for; a CPU device in a split list is almost certainly a mistake.
+            LOG_WARN(
+                "%s: layer split across a CPU device is not supported; using %s "
+                "(use --params-backend %s=cpu to keep weights in RAM)",
+                desc.c_str(),
+                sd::layer_split_backend_device_display_name(module_backends[0]).c_str(),
+                sd_backend_module_name(module));
+            return model_manager->register_param_tensors(desc,
+                                                         std::move(group_tensors),
+                                                         residency_mode,
+                                                         module_backends[0],
+                                                         params_backend_for(module),
+                                                         params_mem_size);
+        }
+
+        std::map<std::string, ggml_tensor*> split_tensors;
+        if constexpr (std::is_base_of_v<Conditioner, T>) {
+            model->get_layer_split_param_tensors(split_tensors);
+        } else {
+            split_tensors = group_tensors;
+        }
+
+        auto partitions = sd::partition_layer_split_tensors(desc, group_tensors, split_tensors, module_backends);
+        bool is_split   = false;
+        for (size_t i = 1; i < partitions.size(); i++) {
+            if (!partitions[i].empty()) {
+                is_split = true;
+                break;
+            }
+        }
+        if (!is_split) {
+            return model_manager->register_param_tensors(desc,
+                                                         std::move(group_tensors),
+                                                         residency_mode,
+                                                         module_backends[0],
+                                                         params_backend_for(module),
+                                                         params_mem_size);
+        }
+
+        model->set_runtime_backends(module_backends);
+        const bool params_follow_runtime = backend_manager.params_backend_follows_runtime(module) ||
+                                           backend_manager.params_backend_is_disk(module);
+        for (size_t i = 0; i < module_backends.size(); i++) {
+            if (partitions[i].empty()) {
+                continue;
+            }
+            ggml_backend_t partition_params_backend =
+                params_follow_runtime ? module_backends[i] : params_backend_for(module);
+            if (partition_params_backend == nullptr) {
+                return false;
+            }
+            if (!model_manager->register_param_tensors(desc,
+                                                       std::move(partitions[i]),
+                                                       residency_mode,
+                                                       module_backends[i],
+                                                       partition_params_backend,
+                                                       params_mem_size)) {
+                return false;
+            }
+        }
+        return true;
+    }
+
     bool init_backend() {
         std::string error;
         if (!backend_manager.init(backend_spec.c_str(),
                                   params_backend_spec.c_str(),
+                                  split_mode_spec.c_str(),
                                   &error)) {
             LOG_ERROR("backend config failed: %s", error.c_str());
             return false;
@@ -293,6 +519,16 @@ public:
         return ensure_backend_pair(SDBackendModule::DIFFUSION);
     }
 
+    bool row_split_active() {
+        for (SDBackendModule module : {SDBackendModule::DIFFUSION, SDBackendModule::TE}) {
+            if (backend_manager.split_mode(module) == SDSplitMode::ROW &&
+                backend_manager.runtime_backends(module).size() > 1) {
+                return true;
+            }
+        }
+        return false;
+    }
+
     std::shared_ptr<RNG> get_rng(rng_type_t rng_type) {
         if (rng_type == STD_DEFAULT_RNG) {
             return std::make_shared<STDDefaultRNG>();
@@ -351,6 +587,8 @@ public:
         eager_load          = sd_ctx_params->eager_load;
         backend_spec        = SAFE_STR(sd_ctx_params->backend);
         params_backend_spec = SAFE_STR(sd_ctx_params->params_backend);
+        split_mode_spec     = SAFE_STR(sd_ctx_params->split_mode);
+        auto_fit_enabled    = sd_ctx_params->auto_fit;
         max_vram_assignment.reset(0.f);
         {
             std::string error;
@@ -376,21 +614,6 @@ public:
 
         ggml_log_set(ggml_log_callback_default, nullptr);
 
-        if (!init_backend()) {
-            return false;
-        }
-        {
-            std::string error;
-            if (!max_vram_assignment.canonicalize_backend_keys(&error)) {
-                LOG_ERROR("%s", error.c_str());
-                return false;
-            }
-        }
-        if (stream_layers && !backend_manager.params_backend_is_cpu(SDBackendModule::DIFFUSION)) {
-            LOG_WARN("--stream-layers has no effect unless diffusion params backend is cpu; ignoring");
-            stream_layers = false;
-        }
-
         model_manager = std::make_shared<ModelManager>();
         model_manager->set_n_threads(n_threads);
         model_manager->set_enable_mmap(enable_mmap);
@@ -538,6 +761,31 @@ public:
             model_loader.set_wtype_override(wtype, tensor_type_rules);
         }
 
+        if (auto_fit_enabled) {
+            if (!sd::backend_fit::derive_backend_specs(model_loader,
+                                                       wtype,
+                                                       max_vram_assignment,
+                                                       backend_spec,
+                                                       params_backend_spec)) {
+                return false;
+            }
+        }
+
+        if (!init_backend()) {
+            return false;
+        }
+        {
+            std::string error;
+            if (!max_vram_assignment.canonicalize_backend_keys(&error)) {
+                LOG_ERROR("%s", error.c_str());
+                return false;
+            }
+        }
+        if (stream_layers && !backend_manager.params_backend_is_cpu(SDBackendModule::DIFFUSION)) {
+            LOG_WARN("--stream-layers has no effect unless diffusion params backend is cpu; ignoring");
+            stream_layers = false;
+        }
+
         std::map<ggml_type, uint32_t> wtype_stat                 = model_loader.get_wtype_stat();
         std::map<ggml_type, uint32_t> conditioner_wtype_stat     = model_loader.get_conditioner_wtype_stat();
         std::map<ggml_type, uint32_t> diffusion_model_wtype_stat = model_loader.get_diffusion_model_wtype_stat();
@@ -579,12 +827,17 @@ public:
             // Avoid full-model LoRA merge buffers on constrained setups.
             const bool params_offloaded      = params_backend_for(SDBackendModule::DIFFUSION) != backend_for(SDBackendModule::DIFFUSION);
             const bool streaming_constrained = stream_layers || params_offloaded;
-            if (have_quantized_weight || streaming_constrained) {
+            if (have_quantized_weight || streaming_constrained || row_split_active()) {
                 apply_lora_immediately = false;
             } else {
                 apply_lora_immediately = true;
             }
         } else if (sd_ctx_params->lora_apply_mode == LORA_APPLY_IMMEDIATELY) {
+            if (row_split_active()) {
+                LOG_WARN(
+                    "row-split tensors do not support the immediately LoRA apply mode; "
+                    "LoRAs will not be applied to them (use --lora-apply-mode at_runtime)");
+            }
             apply_lora_immediately = true;
         } else {
             apply_lora_immediately = false;
@@ -754,13 +1007,14 @@ public:
                     }
                 }
             } else if (sd_version_is_qwen_image(version)) {
-                cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
+                bool enable_vision = version != VERSION_QWEN_IMAGE_LAYERED;
+                cond_stage_model   = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
                                                                  tensor_storage_map,
                                                                  version,
                                                                  "",
-                                                                 true,
+                                                                 enable_vision,
                                                                  model_manager);
-                diffusion_model  = std::make_shared<Qwen::QwenImageRunner>(backend_for(SDBackendModule::DIFFUSION),
+                diffusion_model    = std::make_shared<Qwen::QwenImageRunner>(backend_for(SDBackendModule::DIFFUSION),
                                                                           tensor_storage_map,
                                                                           "model.diffusion_model",
                                                                           version,
@@ -2118,9 +2372,9 @@ public:
             sd_sample::SampleStepCacheDispatcher step_cache(cache_runtime, step, sigma);
             std::vector<sd::Tensor<float>> controls;
             DiffusionParams diffusion_params;
-            diffusion_params.x                  = &noised_input;
-            diffusion_params.timesteps          = &timesteps_tensor;
-            diffusion_params.increase_ref_index = increase_ref_index;
+            diffusion_params.x              = &noised_input;
+            diffusion_params.timesteps      = &timesteps_tensor;
+            diffusion_params.ref_index_mode = Rope::ref_index_mode_from_bool(increase_ref_index);
             sd::guidance::GuidanceInput step_guidance_input;
             step_guidance_input.step          = step;
             step_guidance_input.schedule_size = sigmas.size();
@@ -2371,6 +2625,10 @@ public:
         return latent_channel;
     }
 
+    int get_image_channels() const {
+        return version == VERSION_QWEN_IMAGE_LAYERED ? 4 : 3;
+    }
+
     int get_image_seq_len(int h, int w) {
         int vae_scale_factor = get_vae_scale_factor();
         return (h / vae_scale_factor) * (w / vae_scale_factor);
@@ -2444,7 +2702,16 @@ public:
         }
         auto latents = first_stage_model->diffusion_to_vae_latents(x);
         first_stage_model->set_temporal_tiling_enabled(vae_tiling_params.temporal_tiling);
-        return first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y);
+        auto decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y);
+        if (decoded.empty() && auto_fit_enabled) {
+            bool prefer_temporal_tiling = decode_video && std::dynamic_pointer_cast<LTXVideoVAE>(first_stage_model) != nullptr;
+            if (sd::backend_fit::prepare_vae_decode_retry_tiling(vae_tiling_params, prefer_temporal_tiling)) {
+                first_stage_model->free_compute_buffer();
+                first_stage_model->set_temporal_tiling_enabled(vae_tiling_params.temporal_tiling);
+                decoded = first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y);
+            }
+        }
+        return decoded;
     }
 
     sd::Tensor<float> normalize_ltx_video_latents(const sd::Tensor<float>& x) {
@@ -2551,6 +2818,7 @@ const char* sample_method_to_str[] = {
     "euler_cfg_pp",
     "euler_a_cfg_pp",
     "euler_ge",
+    "dpm++2m_sde",
 };
 
 const char* sd_sample_method_name(enum sample_method_t sample_method) {
@@ -2800,6 +3068,8 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) {
     sd_ctx_params->vae_format           = SD_VAE_FORMAT_AUTO;
     sd_ctx_params->backend              = nullptr;
     sd_ctx_params->params_backend       = nullptr;
+    sd_ctx_params->split_mode           = nullptr;
+    sd_ctx_params->auto_fit             = false;
     sd_ctx_params->rpc_servers          = nullptr;
     sd_ctx_params->pulid_weights_path   = nullptr;
 }
@@ -2839,6 +3109,8 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
              "eager_load: %s\n"
              "backend: %s\n"
              "params_backend: %s\n"
+             "split_mode: %s\n"
+             "auto_fit: %s\n"
              "flash_attn: %s\n"
              "diffusion_flash_attn: %s\n"
              "circular_x: %s\n"
@@ -2875,6 +3147,8 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
              BOOL_STR(sd_ctx_params->eager_load),
              SAFE_STR(sd_ctx_params->backend),
              SAFE_STR(sd_ctx_params->params_backend),
+             SAFE_STR(sd_ctx_params->split_mode),
+             BOOL_STR(sd_ctx_params->auto_fit),
              BOOL_STR(sd_ctx_params->flash_attn),
              BOOL_STR(sd_ctx_params->diffusion_flash_attn),
              BOOL_STR(sd_ctx_params->circular_x),
@@ -2958,6 +3232,7 @@ void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params) {
     sd_img_gen_params->seed              = -1;
     sd_img_gen_params->batch_count       = 1;
     sd_img_gen_params->control_strength  = 0.9f;
+    sd_img_gen_params->qwen_image_layers = 3;
     sd_img_gen_params->pm_params         = {nullptr, 0, nullptr, 20.f};
     sd_img_gen_params->pulid_params      = {nullptr, 1.0f};
     sd_img_gen_params->vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr};
@@ -2984,6 +3259,7 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
              "seed: %" PRId64
              "\n"
              "batch_count: %d\n"
+             "qwen_image_layers: %d\n"
              "ref_images_count: %d\n"
              "auto_resize_ref_image: %s\n"
              "increase_ref_index: %s\n"
@@ -3000,6 +3276,7 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
              sd_img_gen_params->strength,
              sd_img_gen_params->seed,
              sd_img_gen_params->batch_count,
+             sd_img_gen_params->qwen_image_layers,
              sd_img_gen_params->ref_images_count,
              BOOL_STR(sd_img_gen_params->auto_resize_ref_image),
              BOOL_STR(sd_img_gen_params->increase_ref_index),
@@ -3243,6 +3520,7 @@ static float resolve_eta(sd_ctx_t* sd_ctx,
             case DPMPP2S_A_SAMPLE_METHOD:
             case ER_SDE_SAMPLE_METHOD:
             case EULER_A_CFG_PP_SAMPLE_METHOD:
+            case DPMPP2M_SDE_SAMPLE_METHOD:
                 return 1.0f;
             default:;
         }
@@ -3267,6 +3545,7 @@ struct GenerationRequest {
     bool has_ref_images                      = false;
     const sd_cache_params_t* cache_params    = nullptr;
     int batch_count                          = 1;
+    int qwen_image_layers                    = 3;
     int shifted_timestep                     = 0;
     float strength                           = 1.f;
     float control_strength                   = 0.f;
@@ -3292,6 +3571,7 @@ struct GenerationRequest {
         diffusion_model_down_factor = sd_ctx->sd->get_diffusion_model_down_factor();
         seed                        = sd_img_gen_params->seed;
         batch_count                 = sd_img_gen_params->batch_count;
+        qwen_image_layers           = std::max(0, sd_img_gen_params->qwen_image_layers);
         clip_skip                   = sd_img_gen_params->clip_skip;
         shifted_timestep            = sd_img_gen_params->sample_params.shifted_timestep;
         strength                    = sd_img_gen_params->strength;
@@ -3996,6 +4276,34 @@ public:
     ImageVaeAxesGuard& operator=(const ImageVaeAxesGuard&) = delete;
 };
 
+static sd::Tensor<float> ensure_image_tensor_channels(sd::Tensor<float> image, int channels) {
+    if (image.empty()) {
+        return image;
+    }
+    GGML_ASSERT(image.dim() == 4);
+    int64_t current_channels = image.shape()[2];
+    if (current_channels == channels) {
+        return image;
+    }
+    if (channels == 4) {
+        sd::Tensor<float> alpha = sd::full<float>({image.shape()[0], image.shape()[1], 1, image.shape()[3]}, 1.f);
+        if (current_channels == 3) {
+            return sd::ops::concat(image, alpha, 2);
+        }
+        if (current_channels == 1) {
+            sd::Tensor<float> rgb = sd::ops::concat(image, image, 2);
+            rgb                   = sd::ops::concat(rgb, image, 2);
+            return sd::ops::concat(rgb, alpha, 2);
+        }
+    }
+    if (channels == 3 && current_channels >= 3) {
+        return sd::ops::slice(image, 2, 0, 3);
+    }
+    GGML_ABORT("cannot convert image tensor from %lld to %d channels",
+               (long long)current_channels,
+               channels);
+}
+
 static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd_ctx_t* sd_ctx,
                                                                               const sd_img_gen_params_t* sd_img_gen_params,
                                                                               GenerationRequest* request,
@@ -4005,6 +4313,7 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
     sd::Tensor<float> init_image_tensor;
     sd::Tensor<float> control_image_tensor;
     sd::Tensor<float> mask_image_tensor;
+    int image_channels = sd_ctx->sd->get_image_channels();
 
     if (sd_img_gen_params->init_image.data != nullptr) {
         LOG_INFO("IMG2IMG");
@@ -4021,7 +4330,8 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
             plan->sample_steps = static_cast<int>(plan->sigmas.size() - 1);
         }
 
-        init_image_tensor = sd_image_to_tensor(sd_img_gen_params->init_image, request->width, request->height);
+        init_image_tensor = ensure_image_tensor_channels(sd_image_to_tensor(sd_img_gen_params->init_image, request->width, request->height),
+                                                         image_channels);
     }
 
     if (sd_img_gen_params->mask_image.data != nullptr) {
@@ -4053,7 +4363,11 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
     sd::Tensor<float> init_latent;
     sd::Tensor<float> control_latent;
     if (init_image_tensor.empty()) {
-        init_latent = sd_ctx->sd->generate_init_latent(request->width, request->height);
+        if (sd_ctx->sd->version == VERSION_QWEN_IMAGE_LAYERED) {
+            init_latent = sd_ctx->sd->generate_init_latent(request->width, request->height, request->qwen_image_layers + 1, true);
+        } else {
+            init_latent = sd_ctx->sd->generate_init_latent(request->width, request->height);
+        }
     } else {
         init_latent = sd_ctx->sd->encode_first_stage(init_image_tensor);
         if (init_latent.empty()) {
@@ -4072,12 +4386,13 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
 
     std::vector<sd::Tensor<float>> ref_images;
     for (int i = 0; i < sd_img_gen_params->ref_images_count; i++) {
-        ref_images.push_back(sd_image_to_tensor(sd_img_gen_params->ref_images[i]));
+        ref_images.push_back(ensure_image_tensor_channels(sd_image_to_tensor(sd_img_gen_params->ref_images[i]),
+                                                          image_channels));
     }
 
     if (ref_images.empty() && sd_version_is_unet_edit(sd_ctx->sd->version)) {
         LOG_WARN("This model needs at least one reference image; using an empty reference");
-        ref_images.push_back(sd::zeros<float>({request->width, request->height, 3, 1}));
+        ref_images.push_back(sd::zeros<float>({request->width, request->height, image_channels, 1}));
         request->guidance.img_cfg = request->guidance.txt_cfg;
         request->use_img_uncond   = false;
     }
@@ -4103,7 +4418,10 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
             vae_width  = round(vae_width / factor) * factor;
 
             auto resized_ref_img = sd::ops::interpolate(ref_images[i],
-                                                        {static_cast<int>(vae_width), static_cast<int>(vae_height), 3, 1});
+                                                        {static_cast<int>(vae_width),
+                                                         static_cast<int>(vae_height),
+                                                         ref_images[i].shape()[2],
+                                                         ref_images[i].shape()[3]});
 
             LOG_DEBUG("resize vae ref image %d from %" PRId64 "x%" PRId64 " to %" PRId64 "x%" PRId64,
                       static_cast<int>(i),
@@ -4333,13 +4651,41 @@ static sd_image_t* decode_image_outputs(sd_ctx_t* sd_ctx,
             cancelled = true;
             break;
         }
-        int64_t t1              = ggml_time_ms();
-        sd::Tensor<float> image = sd_ctx->sd->decode_first_stage(final_latents[i]);
-        if (image.empty()) {
-            LOG_ERROR("decode_first_stage failed for latent %" PRId64, i + 1);
-            return nullptr;
+        int64_t t1 = ggml_time_ms();
+        if (sd_ctx->sd->version == VERSION_QWEN_IMAGE_LAYERED) {
+            int qwen_image_latent_layers = request.qwen_image_layers + 1;
+            if (final_latents[i].dim() < 5 || final_latents[i].shape()[2] < qwen_image_latent_layers) {
+                LOG_ERROR("qwen image layered expected at least %d latent layers, got shape dim=%d",
+                          qwen_image_latent_layers,
+                          final_latents[i].dim());
+                return nullptr;
+            }
+            for (int layer_index = 0; layer_index < qwen_image_latent_layers; layer_index++) {
+                if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) {
+                    LOG_ERROR("cancelling latent decodings");
+                    cancelled = true;
+                    break;
+                }
+                sd::Tensor<float> layer_latent = sd::ops::slice(final_latents[i], 2, layer_index, layer_index + 1);
+                layer_latent.squeeze_(2);
+                sd::Tensor<float> image = sd_ctx->sd->decode_first_stage(layer_latent);
+                if (image.empty()) {
+                    LOG_ERROR("decode_first_stage failed for latent %zu layer %d", i + 1, layer_index + 1);
+                    return nullptr;
+                }
+                decoded_images.push_back(std::move(image));
+            }
+            if (cancelled) {
+                break;
+            }
+        } else {
+            sd::Tensor<float> image = sd_ctx->sd->decode_first_stage(final_latents[i]);
+            if (image.empty()) {
+                LOG_ERROR("decode_first_stage failed for latent %" PRId64, i + 1);
+                return nullptr;
+            }
+            decoded_images.push_back(std::move(image));
         }
-        decoded_images.push_back(std::move(image));
         int64_t t2 = ggml_time_ms();
         LOG_INFO("latent %zu decoded, taking %.2fs", i + 1, (t2 - t1) * 1.0f / 1000);
     }

src/upscaler.cpp

@@ -46,6 +46,7 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path,
     std::string error;
     if (!backend_manager.init(backend_spec.c_str(),
                               params_backend_spec.c_str(),
+                              /*split_mode_spec=*/nullptr,
                               &error)) {
         LOG_ERROR("upscaler backend config failed: %s", error.c_str());
         return false;