feat: hot-reload ControlNet - swap without rebuilding the context (#1768)

.github/workflows/close-inactive-issues.yml

@@ -0,0 +1,48 @@
+name: Close inactive issues
+
+on:
+  schedule:
+    # Run daily. GitHub cron schedules use UTC.
+    - cron: "30 1 * * *"
+  workflow_dispatch:
+    inputs:
+      debug_only:
+        description: "Dry run: log intended actions without changing issues"
+        required: false
+        default: false
+        type: boolean
+
+permissions:
+  issues: write
+
+concurrency:
+  group: ${{ github.workflow }}
+  cancel-in-progress: false
+
+jobs:
+  close-inactive-issues:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Comment and close inactive issues
+        uses: actions/stale@v10
+        with:
+          days-before-issue-stale: 365
+          days-before-issue-close: 0
+
+          days-before-pr-stale: -1
+          days-before-pr-close: -1
+
+          stale-issue-label: issue:inactive
+          close-issue-label: issue:auto-closed
+          close-issue-reason: completed
+          stale-issue-message: ""
+          close-issue-message: >
+            This issue has had no activity for one year. The latest version of
+            the code may already have fixed the problem.
+
+            If the issue still exists in the latest version, you can reopen
+            this issue at any time with updated reproduction details.
+
+          remove-issue-stale-when-updated: true
+          operations-per-run: 1000
+          debug-only: ${{ github.event_name == 'workflow_dispatch' && inputs.debug_only || false }}

.gitignore

@@ -13,3 +13,7 @@ output*.png
 models*
 *.log
 preview.png
+.claude/
+CLAUDE.local.md
+.agents/
+.codex/

AGENTS.md

@@ -0,0 +1,183 @@
+# Instructions for stable-diffusion.cpp
+
+This document is for AI coding agents working in this repository. It should
+describe agent-specific workflow, repository routing, editing boundaries, and
+project-specific pitfalls.
+
+For general contribution rules, including PR scope, commit conventions, code
+style, dependency updates, security hygiene, and AI-assisted contribution policy,
+see `CONTRIBUTING.md`.
+
+---
+
+## Agent Operating Rules
+
+Before analyzing or modifying the repository:
+
+1. Read this file.
+2. Use `rg` / `rg --files` or directory listing commands to confirm the current
+   tree before relying on a path.
+3. Start from `src/` and relevant `docs/` for runtime behavior.
+4. Read the relevant code before editing.
+5. Prefer the smallest change that fits the existing architecture.
+6. Report focused verification and mention any tests not run.
+
+Agents must not:
+
+* Run `git push`, create PRs, or submit issue/PR comments on the user's behalf.
+* Create commits unless the user explicitly requests that specific commit.
+* Modify `ggml/`, `thirdparty/`, or `examples/server/frontend/` unless
+  explicitly requested and necessary.
+* Read large local model files or tokenizer vocabulary files.
+* Rewrite unrelated code for style-only reasons.
+* Add secrets, model weights, generated binaries, local absolute paths, or
+  machine-specific output.
+
+When a change is large, architectural, or likely to affect public behavior,
+pause and present a short plan before editing.
+
+---
+
+## Repository Map and Editing Boundaries
+
+This is a routing map for agents, not a full architecture document. The layout
+can change, so verify paths before using them. Do not inspect excluded
+large-data directories while checking the tree.
+
+### Primary Project Code
+
+Core implementation lives under `src/`.
+
+Current source layout includes:
+
+* `src/core/` - shared tensor, ggml integration, backend, graph, RNG, and utility
+  code.
+* `src/model/` - model families and model components.
+* `src/model_io/` - model file loading, GGUF, safetensors, pickle, and related
+  serialization helpers.
+* `src/runtime/` - sampling, denoising, guidance, caching, preprocessing, and
+  runtime execution helpers.
+* `src/tokenizers/` - tokenizer implementations.
+* `src/conditioning/` - conditioning and prompt-related implementation.
+* `src/extensions/` - optional feature extensions.
+* top-level `src/*.cpp` and `src/*.h` files - public implementation entry
+  points, model loading, conversion, versioning, and shared managers.
+
+`src/tokenizers/vocab/` contains large tokenizer vocabulary data. Do not read or
+parse files in this directory; reference the path only when necessary.
+
+### Public API
+
+`include/` contains the C API exposed by the project. Currently the primary
+public header is `include/stable-diffusion.h`.
+
+Treat public headers as stable API. Avoid breaking compatibility unless the user
+explicitly requests it. If public behavior changes, update relevant examples or
+documentation.
+
+### Examples
+
+`examples/` contains programs demonstrating library usage.
+
+* `examples/cli/` - command line program for running models, testing features,
+  and debugging.
+* `examples/common/` - shared example support code.
+* `examples/server/` - server application built on top of the library.
+* `examples/server/frontend/` - git submodule containing independent frontend
+  code. Avoid modifying it unless explicitly requested.
+
+### Documentation and Tooling
+
+* `docs/` - documentation for supported models, build options, behavior, and
+  workflows.
+* `scripts/` - development, model processing, build automation, formatting, and
+  tooling scripts.
+* `cmake/` - CMake support modules.
+* `docker/` - Docker-related project files.
+* `assets/` - documentation assets; not runtime code.
+
+### External, Local, and Generated State
+
+* `ggml/` - git submodule for the ggml dependency.
+* `thirdparty/` - vendored third-party dependencies.
+* `models/` - local model storage. Ignore this directory and do not read model
+  files.
+* `test/` - local testing scripts. Use only when relevant to the task.
+* `build/`, `build_*`, and similar directories - generated build outputs.
+  Inspect them only when debugging a build result.
+
+---
+
+## Agent Workflow for Code Changes
+
+1. Identify the relevant modules under `src/`.
+2. Check whether the change touches the public API in `include/`.
+3. Consult relevant `docs/` and examples before changing user-facing behavior.
+4. Follow existing local patterns before adding new abstractions.
+5. Keep edits scoped to the requested behavior.
+6. Run the narrowest useful build, test, or inspection command available.
+
+Follow `CONTRIBUTING.md` for formatting, naming, PR expectations, dependency
+update policy, and security rules.
+
+---
+
+## Code Comments
+
+Keep comments rare and useful.
+
+Do not add comments that only describe what the code does. Add comments only
+when the code cannot fully express the logic, the logic is unusually complex, or
+there are historical reasons, invariants, constraints, compatibility concerns,
+or known pitfalls that future maintainers need to understand.
+
+Do not add task-specific comments that will be meaningless after review.
+
+Examples from the current codebase:
+
+```cpp
+// GOOD: explains a safety constraint that is not obvious from the assignment.
+// From src/model_io/pickle_io.cpp.
+// Non-tensor checkpoint metadata can use REDUCE for arbitrary
+// Python objects. Do not execute it; keep stack shape only.
+stack.push_back(make_none_value());
+
+// BAD: describes only what the next line does.
+// Set the token count to zero.
+token_count = 0;
+```
+
+---
+
+## Text File Encoding
+
+When reading or editing repository text files:
+
+* Prefer UTF-8 with LF for Markdown, frontend source, JSON, and other text-first
+  project files unless the file already clearly uses a different encoding.
+* Do not assume terminal output encoding matches file encoding on Windows.
+* A file that looks garbled in PowerShell output may still be valid UTF-8.
+* When inspecting UTF-8 files in PowerShell, prefer explicit UTF-8 reads such as:
+  * `[Console]::OutputEncoding = [System.Text.Encoding]::UTF8`
+  * `Get-Content -Encoding utf8 <path>`
+* Avoid rewriting a file purely because console output looked garbled; verify
+  the actual file encoding first.
+
+---
+
+## Tensor and Layout Notes
+
+Additional tensor/layout rules for this codebase:
+
+* `sd::Tensor` shape order is not PyTorch/NumPy-style. `shape()[0]` is the
+  lowest and most contiguous dimension, and higher indices are higher
+  dimensions.
+* Broadcasting for `sd::Tensor` must align dimensions from low to high dimension
+  indices. If one tensor has fewer dimensions, append implicit `1`s at the
+  higher-dimension end.
+* `ggml_n_dims` / `ggml_n_dims(tensor)` can drop trailing singleton high
+  dimensions. Do not assume a logical trailing dimension of `1` will still be
+  counted in ggml metadata.
+* Internal tensor-returning interfaces use an empty `sd::Tensor` to represent
+  null, absent, or failure states. Do not add `std::optional<sd::Tensor<...>>`
+  for internal APIs unless a distinct semantic state is truly required.

CLAUDE.md

@@ -0,0 +1,11 @@
+@AGENTS.md
+
+## Claude Code
+
+Follow `AGENTS.md` as the shared repository instructions.
+
+Do not duplicate contribution, style, PR, dependency, or security policy here;
+use `CONTRIBUTING.md` as the canonical source for those rules.
+
+Keep Claude-specific project notes in this file only when they do not apply to
+other coding agents.

docs/chroma.md

@@ -24,10 +24,9 @@ You can download the preconverted gguf weights from [silveroxides/Chroma-GGUF](h
 For example:
 
 ```
- .\bin\Release\sd-cli.exe --diffusion-model  ..\models\chroma-unlocked-v40-q8_0.gguf --vae ..\models\ae.sft --t5xxl ..\models\t5xxl_fp16.safetensors  -p "a lovely cat holding a sign says 'chroma.cpp'" --cfg-scale 4.0 --sampling-method euler -v --chroma-disable-dit-mask --clip-on-cpu
+ .\bin\Release\sd-cli.exe --diffusion-model  ..\models\chroma-unlocked-v40-q8_0.gguf --vae ..\models\ae.sft --t5xxl ..\models\t5xxl_fp16.safetensors  -p "a lovely cat holding a sign says 'chroma.cpp'" --cfg-scale 4.0 --sampling-method euler -v --model-args chroma_use_dit_mask=false --clip-on-cpu
 ```
 
 ![](../assets/flux/chroma_v40.png)
 
 
-

docs/performance.md

@@ -53,6 +53,26 @@ Per-module assignments can target only the largest modules:
 
 See [backend selection](./backend.md) for full syntax.
 
+## Run models that don't fit in VRAM (CPU streaming).
+
+`--offload-to-cpu` alone keeps every parameter in system RAM and stages it to the runtime backend on first use, then leaves it resident there. If the diffusion model is larger than the runtime backend's free memory (e.g. Flux dev at bf16 on an 8 GiB GPU), that residency stops fitting during the sampling loop and generation fails. Two additional flags make it fit by trading a small amount of speed for room:
+
+- `--max-vram <GiB>` sets a VRAM budget the graph-cut segmenter respects. It cuts each forward pass into segments sized to fit the budget, running them in sequence and freeing intermediate activations between them. Negative values auto-detect free VRAM and spare the given amount (`--max-vram -1` uses most of the free VRAM and keeps ~1 GiB headroom), a positive value caps the budget, `0` disables segmentation.
+- `--stream-layers` streams the diffusion model's transformer blocks one at a time. Each block's parameters are copied from the CPU to the runtime backend just before it runs and evicted when the residency budget is reached. Prefetching hides most of the copy latency behind compute. This flag only takes effect when the diffusion params backend is CPU, so it must be combined with `--offload-to-cpu` (or an explicit `--params-backend diffusion=cpu`); a warning is logged and the flag is ignored otherwise.
+
+The three flags stack. The recommended shape for "biggest model my card can host":
+
+```shell
+sd-cli --diffusion-model flux1-dev.safetensors ... \
+       --offload-to-cpu --max-vram -1 --stream-layers
+```
+
+- `--offload-to-cpu`: params in RAM, staged as needed.
+- `--max-vram -1`: use most of the free VRAM as the compute budget, spare 1 GiB headroom, let the graph-cut segmenter split each forward pass to fit.
+- `--stream-layers`: on top of the segmenter, stream individual transformer blocks so their weights don't all need to be resident at once.
+
+Ordered from fastest to smallest-VRAM: no flags → `--offload-to-cpu` → `--offload-to-cpu --max-vram <N>` → `--offload-to-cpu --max-vram <N> --stream-layers`. Each step down costs a few percent of throughput to buy more room; combined they can run models roughly 3-4x larger than the raw VRAM would allow.
+
 ## Use quantization to reduce memory usage.
 
 [quantization](./quantization_and_gguf.md)

docs/qwen_image_edit.md

@@ -39,10 +39,10 @@
 
 ### Qwen Image Edit 2511
 
-To use the new Qwen Image Edit 2511 mode, the  `--qwen-image-zero-cond-t` flag must be enabled; otherwise, image editing quality will degrade significantly.
+To use the new Qwen Image Edit 2511 mode, `--model-args qwen_image_zero_cond_t=true` must be enabled; otherwise, image editing quality will degrade significantly.
 
 ```
-.\bin\Release\sd-cli.exe --diffusion-model  ..\models\diffusion_models\qwen-image-edit-2511-Q4_K_M.gguf --vae ..\models\vae\qwen_image_vae.safetensors  --llm ..\models\text_encoders\qwen_2.5_vl_7b.safetensors --cfg-scale 2.5 --sampling-method euler -v --offload-to-cpu --diffusion-fa --flow-shift 3 -r ..\assets\flux\flux1-dev-q8_0.png -p "change 'flux.cpp' to 'edit.cpp'"  --qwen-image-zero-cond-t
+.\bin\Release\sd-cli.exe --diffusion-model  ..\models\diffusion_models\qwen-image-edit-2511-Q4_K_M.gguf --vae ..\models\vae\qwen_image_vae.safetensors  --llm ..\models\text_encoders\qwen_2.5_vl_7b.safetensors --cfg-scale 2.5 --sampling-method euler -v --offload-to-cpu --diffusion-fa --flow-shift 3 -r ..\assets\flux\flux1-dev-q8_0.png -p "change 'flux.cpp' to 'edit.cpp'" --model-args qwen_image_zero_cond_t=true
 ```
 
-<img alt="qwen_image_edit_2509" src="../assets/qwen/qwen_image_edit_2511.png" />
+<img alt="qwen_image_edit_2509" src="../assets/qwen/qwen_image_edit_2511.png" />

examples/common/common.cpp

@@ -443,6 +443,12 @@ ArgOptions SDContextParams::get_options() {
          "weight type per tensor pattern (example: \"^vae\\.=f16,model\\.=q8_0\")",
          (int)',',
          &tensor_type_rules},
+        {"",
+         "--model-args",
+         "extra model args, key=value list. Supports chroma_use_dit_mask, chroma_use_t5_mask, "
+         "chroma_t5_mask_pad, qwen_image_zero_cond_t",
+         (int)',',
+         &model_args},
         {"",
          "--photo-maker",
          "path to PHOTOMAKER model",
@@ -493,10 +499,6 @@ ArgOptions SDContextParams::get_options() {
          "number of threads to use during computation (default: -1). "
          "If threads <= 0, then threads will be set to the number of CPU physical cores",
          &n_threads},
-        {"",
-         "--chroma-t5-mask-pad",
-         "t5 mask pad size of chroma",
-         &chroma_t5_mask_pad},
     };
 
     options.bool_options = {
@@ -554,18 +556,6 @@ ArgOptions SDContextParams::get_options() {
          "--vae-conv-direct",
          "use ggml_conv2d_direct in the vae model",
          true, &vae_conv_direct},
-        {"",
-         "--chroma-disable-dit-mask",
-         "disable dit mask for chroma",
-         false, &chroma_use_dit_mask},
-        {"",
-         "--qwen-image-zero-cond-t",
-         "enable zero_cond_t for qwen image",
-         true, &qwen_image_zero_cond_t},
-        {"",
-         "--chroma-enable-t5-mask",
-         "enable t5 mask for chroma",
-         true, &chroma_use_t5_mask},
     };
 
     auto on_type_arg = [&](int argc, const char** argv, int index) {
@@ -832,6 +822,7 @@ std::string SDContextParams::to_string() const {
         << "  backend: \"" << backend << "\",\n"
         << "  params_backend: \"" << params_backend << "\",\n"
         << "  split_mode: \"" << split_mode << "\",\n"
+        << "  model_args: \"" << model_args << "\",\n"
         << "  auto_fit: " << (auto_fit ? "true" : "false") << ",\n"
         << "  enable_mmap: " << (enable_mmap ? "true" : "false") << ",\n"
         << "  control_net_cpu: " << (control_net_cpu ? "true" : "false") << ",\n"
@@ -841,10 +832,6 @@ std::string SDContextParams::to_string() const {
         << "  diffusion_flash_attn: " << (diffusion_flash_attn ? "true" : "false") << ",\n"
         << "  diffusion_conv_direct: " << (diffusion_conv_direct ? "true" : "false") << ",\n"
         << "  vae_conv_direct: " << (vae_conv_direct ? "true" : "false") << ",\n"
-        << "  chroma_use_dit_mask: " << (chroma_use_dit_mask ? "true" : "false") << ",\n"
-        << "  qwen_image_zero_cond_t: " << (qwen_image_zero_cond_t ? "true" : "false") << ",\n"
-        << "  chroma_use_t5_mask: " << (chroma_use_t5_mask ? "true" : "false") << ",\n"
-        << "  chroma_t5_mask_pad: " << chroma_t5_mask_pad << ",\n"
         << "  prediction: " << sd_prediction_name(prediction) << ",\n"
         << "  lora_apply_mode: " << sd_lora_apply_mode_name(lora_apply_mode) << ",\n"
         << "  force_sdxl_vae_conv_scale: " << (force_sdxl_vae_conv_scale ? "true" : "false") << "\n"
@@ -898,10 +885,6 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) {
     sd_ctx_params.diffusion_conv_direct           = diffusion_conv_direct;
     sd_ctx_params.vae_conv_direct                 = vae_conv_direct;
     sd_ctx_params.force_sdxl_vae_conv_scale       = force_sdxl_vae_conv_scale;
-    sd_ctx_params.chroma_use_dit_mask             = chroma_use_dit_mask;
-    sd_ctx_params.chroma_use_t5_mask              = chroma_use_t5_mask;
-    sd_ctx_params.chroma_t5_mask_pad              = chroma_t5_mask_pad;
-    sd_ctx_params.qwen_image_zero_cond_t          = qwen_image_zero_cond_t;
     sd_ctx_params.vae_format                      = str_to_vae_format(vae_format);
     sd_ctx_params.max_vram                        = max_vram.c_str();
     sd_ctx_params.stream_layers                   = stream_layers;
@@ -911,6 +894,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) {
     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();
+    sd_ctx_params.model_args                      = model_args.empty() ? nullptr : model_args.c_str();
     return sd_ctx_params;
 }
 

examples/common/common.h

@@ -152,6 +152,7 @@ struct SDContextParams {
     std::string backend;
     std::string params_backend;
     std::string split_mode;
+    std::string model_args;
     bool auto_fit = false;
     std::string rpc_servers;
     std::string effective_backend;
@@ -165,12 +166,6 @@ struct SDContextParams {
     bool diffusion_conv_direct = false;
     bool vae_conv_direct       = false;
 
-    bool chroma_use_dit_mask = true;
-    bool chroma_use_t5_mask  = false;
-    int chroma_t5_mask_pad   = 1;
-
-    bool qwen_image_zero_cond_t = false;
-
     prediction_t prediction           = PREDICTION_COUNT;
     lora_apply_mode_t lora_apply_mode = LORA_APPLY_AUTO;
 

include/stable-diffusion.h

@@ -217,10 +217,6 @@ typedef struct {
     bool diffusion_conv_direct;
     bool vae_conv_direct;
     bool force_sdxl_vae_conv_scale;
-    bool chroma_use_dit_mask;
-    bool chroma_use_t5_mask;
-    int chroma_t5_mask_pad;
-    bool qwen_image_zero_cond_t;
     enum sd_vae_format_t vae_format;
     const char* max_vram;  // GiB budget or backend assignment spec for graph-cut segmented param offload (0 = disabled, -1 = auto)
     bool stream_layers;  // Enable residency+prefetch streaming on top of --max-vram (no effect without --max-vram)
@@ -230,6 +226,7 @@ typedef struct {
     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;
+    const char* model_args;
 } sd_ctx_params_t;
 
 typedef struct {
@@ -431,6 +428,11 @@ SD_API const char* sd_get_system_info();
 SD_API bool sd_ctx_supports_image_generation(const sd_ctx_t* sd_ctx);
 SD_API bool sd_ctx_supports_video_generation(const sd_ctx_t* sd_ctx);
 
+// ControlNet hot-swap APIs are not safe to call while generation is in flight.
+SD_API bool sd_ctx_load_control_net(sd_ctx_t* sd_ctx, const char* path);
+SD_API bool sd_ctx_unload_control_net(sd_ctx_t* sd_ctx);
+SD_API bool sd_ctx_has_control_net(const sd_ctx_t* sd_ctx);
+
 SD_API const char* sd_type_name(enum sd_type_t type);
 SD_API enum sd_type_t str_to_sd_type(const char* str);
 SD_API const char* sd_rng_type_name(enum rng_type_t rng_type);

src/conditioning/conditioner.hpp

@@ -6,6 +6,7 @@
 #include <optional>
 
 #include "core/tensor_ggml.hpp"
+#include "core/util.h"
 #include "model/te/clip.hpp"
 #include "model/te/llm.hpp"
 #include "model/te/t5.hpp"
@@ -117,6 +118,8 @@ public:
     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 set_graph_cut_layer_split_enabled(bool enabled) {}
+    virtual void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) {}
     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) {}
@@ -180,6 +183,27 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
         }
     }
 
+    void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
+        text_model->set_runtime_backends(backends);
+        if (sd_version_is_sdxl(version)) {
+            text_model2->set_runtime_backends(backends);
+        }
+    }
+
+    void set_graph_cut_layer_split_enabled(bool enabled) override {
+        text_model->set_graph_cut_layer_split_enabled(enabled);
+        if (sd_version_is_sdxl(version)) {
+            text_model2->set_graph_cut_layer_split_enabled(enabled);
+        }
+    }
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
+        text_model->set_graph_cut_layer_split_backend_vram_limits(limits);
+        if (sd_version_is_sdxl(version)) {
+            text_model2->set_graph_cut_layer_split_backend_vram_limits(limits);
+        }
+    }
+
     void set_flash_attention_enabled(bool enabled) override {
         text_model->set_flash_attention_enabled(enabled);
         if (sd_version_is_sdxl(version)) {
@@ -638,11 +662,41 @@ struct SD3CLIPEmbedder : public Conditioner {
     }
 
     void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
+        if (clip_l) {
+            clip_l->set_runtime_backends(backends);
+        }
+        if (clip_g) {
+            clip_g->set_runtime_backends(backends);
+        }
         if (t5) {
             t5->set_runtime_backends(backends);
         }
     }
 
+    void set_graph_cut_layer_split_enabled(bool enabled) override {
+        if (clip_l) {
+            clip_l->set_graph_cut_layer_split_enabled(enabled);
+        }
+        if (clip_g) {
+            clip_g->set_graph_cut_layer_split_enabled(enabled);
+        }
+        if (t5) {
+            t5->set_graph_cut_layer_split_enabled(enabled);
+        }
+    }
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
+        if (clip_l) {
+            clip_l->set_graph_cut_layer_split_backend_vram_limits(limits);
+        }
+        if (clip_g) {
+            clip_g->set_graph_cut_layer_split_backend_vram_limits(limits);
+        }
+        if (t5) {
+            t5->set_graph_cut_layer_split_backend_vram_limits(limits);
+        }
+    }
+
     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");
@@ -1009,11 +1063,32 @@ struct FluxCLIPEmbedder : public Conditioner {
     }
 
     void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
+        if (clip_l) {
+            clip_l->set_runtime_backends(backends);
+        }
         if (t5) {
             t5->set_runtime_backends(backends);
         }
     }
 
+    void set_graph_cut_layer_split_enabled(bool enabled) override {
+        if (clip_l) {
+            clip_l->set_graph_cut_layer_split_enabled(enabled);
+        }
+        if (t5) {
+            t5->set_graph_cut_layer_split_enabled(enabled);
+        }
+    }
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
+        if (clip_l) {
+            clip_l->set_graph_cut_layer_split_backend_vram_limits(limits);
+        }
+        if (t5) {
+            t5->set_graph_cut_layer_split_backend_vram_limits(limits);
+        }
+    }
+
     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");
@@ -1217,8 +1292,27 @@ struct T5CLIPEmbedder : public Conditioner {
                    bool use_mask                                       = false,
                    int mask_pad                                        = 0,
                    bool is_umt5                                        = false,
-                   std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+                   std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
+                   const char* model_args                              = nullptr)
         : use_mask(use_mask), mask_pad(mask_pad), t5_tokenizer(is_umt5) {
+        for (const auto& [key, value] : parse_key_value_args(model_args, "model arg")) {
+            if (key == "chroma_use_t5_mask") {
+                bool parsed = false;
+                if (parse_strict_bool(value, parsed)) {
+                    this->use_mask = parsed;
+                } else {
+                    LOG_WARN("ignoring invalid Chroma T5 model arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            } else if (key == "chroma_t5_mask_pad") {
+                int parsed = 0;
+                if (parse_strict_int(value, parsed)) {
+                    this->mask_pad = parsed;
+                } else {
+                    LOG_WARN("ignoring invalid Chroma T5 model arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            }
+        }
+
         bool use_t5 = false;
         for (auto pair : tensor_storage_map) {
             if (pair.first.find("text_encoders.t5xxl") != std::string::npos) {
@@ -1258,6 +1352,18 @@ struct T5CLIPEmbedder : public Conditioner {
         }
     }
 
+    void set_graph_cut_layer_split_enabled(bool enabled) override {
+        if (t5) {
+            t5->set_graph_cut_layer_split_enabled(enabled);
+        }
+    }
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
+        if (t5) {
+            t5->set_graph_cut_layer_split_backend_vram_limits(limits);
+        }
+    }
+
     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");
@@ -1462,6 +1568,18 @@ struct MiniT2IConditioner : public Conditioner {
         }
     }
 
+    void set_graph_cut_layer_split_enabled(bool enabled) override {
+        if (t5) {
+            t5->set_graph_cut_layer_split_enabled(enabled);
+        }
+    }
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
+        if (t5) {
+            t5->set_graph_cut_layer_split_backend_vram_limits(limits);
+        }
+    }
+
     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");
@@ -1556,6 +1674,14 @@ struct AnimaConditioner : public Conditioner {
         llm->set_runtime_backends(backends);
     }
 
+    void set_graph_cut_layer_split_enabled(bool enabled) override {
+        llm->set_graph_cut_layer_split_enabled(enabled);
+    }
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
+        llm->set_graph_cut_layer_split_backend_vram_limits(limits);
+    }
+
     void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
         llm->get_param_tensors(tensors, "text_encoders.llm");
     }
@@ -1709,6 +1835,18 @@ struct LLMEmbedder : public Conditioner {
         llm->set_runtime_backends(backends);
     }
 
+    void set_graph_cut_layer_split_enabled(bool enabled) override {
+        if (llm) {
+            llm->set_graph_cut_layer_split_enabled(enabled);
+        }
+    }
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
+        if (llm) {
+            llm->set_graph_cut_layer_split_backend_vram_limits(limits);
+        }
+    }
+
     void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
         llm->get_param_tensors(tensors, "text_encoders.llm");
     }
@@ -2386,6 +2524,14 @@ struct LTXAVEmbedder : public Conditioner {
         llm->set_runtime_backends(backends);
     }
 
+    void set_graph_cut_layer_split_enabled(bool enabled) override {
+        llm->set_graph_cut_layer_split_enabled(enabled);
+    }
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
+        llm->set_graph_cut_layer_split_backend_vram_limits(limits);
+    }
+
     void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
         llm->get_param_tensors(tensors, "text_encoders.llm");
     }

src/core/ggml_extend.hpp

@@ -21,10 +21,12 @@
 #include <sstream>
 #include <string>
 #include <unordered_map>
+#include <unordered_set>
 #include <vector>
 
 #include "core/ggml_extend_backend.h"
 #include "core/ggml_graph_cut.h"
+#include "core/layer_split_partition.h"
 #include "ggml-alloc.h"
 #include "ggml-backend.h"
 #include "ggml.h"
@@ -1745,6 +1747,8 @@ protected:
     size_t max_graph_vram_bytes           = 0;
     bool stream_layers_enabled            = false;
     size_t observed_max_effective_budget_ = 0;
+    bool graph_cut_layer_split_enabled    = false;
+    std::vector<size_t> graph_cut_layer_split_backend_vram_limits_;
 
     std::vector<ggml_backend_t> extra_runtime_backends;  // borrowed (SDBackendManager-owned)
     ggml_backend_sched_t sched             = nullptr;    // owned, multi-device only
@@ -1776,6 +1780,9 @@ protected:
 
     sd::ggml_graph_cut::PlanCache graph_cut_plan_cache_;
     std::unordered_set<const ggml_tensor*> params_tensor_set_;
+    std::unordered_map<const ggml_tensor*, ggml_backend_t> graph_cut_layer_split_assignments_;
+    std::unordered_map<const ggml_tensor*, ggml_backend_t> graph_cut_layer_split_node_assignments_;
+    bool graph_cut_layer_split_primary_notice_logged_ = false;
 
     template <typename T>
     static sd::Tensor<T> take_or_empty(std::optional<sd::Tensor<T>> tensor) {
@@ -1874,6 +1881,20 @@ protected:
         params_tensor_set_dirty_ = false;
     }
 
+    ggml_tensor* canonical_param_tensor(ggml_tensor* tensor) {
+        if (tensor == nullptr) {
+            return nullptr;
+        }
+        if (params_tensor_set_.find(tensor) != params_tensor_set_.end()) {
+            return tensor;
+        }
+        if (tensor->view_src != nullptr &&
+            params_tensor_set_.find(tensor->view_src) != params_tensor_set_.end()) {
+            return tensor->view_src;
+        }
+        return nullptr;
+    }
+
     std::vector<ggml_tensor*> collect_used_param_tensors(ggml_cgraph* gf) {
         std::vector<ggml_tensor*> used_params;
         rebuild_params_tensor_set();
@@ -1886,12 +1907,8 @@ protected:
         seen_params.reserve(static_cast<size_t>(n_leafs));
         for (int i = 0; i < n_leafs; ++i) {
             ggml_tensor* leaf       = sd::ggml_graph_cut::leaf_tensor(gf, i);
-            ggml_tensor* param_leaf = leaf;
-            if (param_leaf != nullptr && params_tensor_set_.find(param_leaf) == params_tensor_set_.end()) {
-                param_leaf = param_leaf->view_src;
-            }
+            ggml_tensor* param_leaf = canonical_param_tensor(leaf);
             if (param_leaf != nullptr &&
-                params_tensor_set_.find(param_leaf) != params_tensor_set_.end() &&
                 seen_params.insert(param_leaf).second) {
                 used_params.push_back(param_leaf);
             }
@@ -2101,11 +2118,17 @@ protected:
         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);
+            ggml_tensor* node    = ggml_graph_node(gf, i);
+            auto node_assignment = graph_cut_layer_split_node_assignments_.find(node);
+            if (node_assignment != graph_cut_layer_split_node_assignments_.end()) {
+                current = node_assignment->second;
+            }
             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_assignment == graph_cut_layer_split_node_assignments_.end()) {
+                        current = weight_backend;
+                    }
                 }
             }
             if (node->op == GGML_OP_NONE || node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE ||
@@ -2435,6 +2458,123 @@ protected:
         return true;
     }
 
+    bool resolve_graph_cut_layer_split_plan(ggml_cgraph* gf,
+                                            GraphCutPlan* plan_out) {
+        GGML_ASSERT(plan_out != nullptr);
+        GGML_ASSERT(gf != nullptr);
+        *plan_out = sd::ggml_graph_cut::resolve_plan(runtime_backend,
+                                                     gf,
+                                                     &graph_cut_plan_cache_,
+                                                     0,
+                                                     params_tensor_set_,
+                                                     get_desc().c_str());
+        return true;
+    }
+
+    bool assign_graph_cut_layer_split_backends(ggml_cgraph* gf) {
+        graph_cut_layer_split_node_assignments_.clear();
+        if (!graph_cut_layer_split_enabled) {
+            return true;
+        }
+        if (!is_multi_device()) {
+            LOG_ERROR("%s graph-cut layer split requires multiple runtime backends", get_desc().c_str());
+            return false;
+        }
+
+        GraphCutPlan plan;
+        if (!resolve_graph_cut_layer_split_plan(gf, &plan)) {
+            return false;
+        }
+        if (!plan.valid || !plan.has_cuts || plan.segments.size() <= 1) {
+            auto manager = weight_manager.lock();
+            if (manager == nullptr) {
+                LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str());
+                return false;
+            }
+            std::vector<ggml_tensor*> graph_params = collect_used_param_tensors(gf);
+            if (!graph_params.empty() &&
+                !manager->assign_compute_backend(graph_params, runtime_backend)) {
+                LOG_ERROR("%s graph-cut layer split failed to assign unmarked graph params to %s",
+                          get_desc().c_str(),
+                          sd::layer_split_backend_device_display_name(runtime_backend).c_str());
+                return false;
+            }
+            for (ggml_tensor* param : graph_params) {
+                if (param != nullptr) {
+                    graph_cut_layer_split_assignments_[param] = 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);
+                if (node != nullptr) {
+                    graph_cut_layer_split_node_assignments_[node] = runtime_backend;
+                }
+            }
+            if (!graph_cut_layer_split_primary_notice_logged_) {
+                LOG_WARN("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params",
+                         get_desc().c_str(),
+                         sd::layer_split_backend_device_display_name(runtime_backend).c_str(),
+                         graph_params.size());
+                graph_cut_layer_split_primary_notice_logged_ = true;
+            } else {
+                LOG_DEBUG("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params",
+                          get_desc().c_str(),
+                          sd::layer_split_backend_device_display_name(runtime_backend).c_str(),
+                          graph_params.size());
+            }
+            return true;
+        }
+
+        std::vector<ggml_backend_t> split_backends;
+        split_backends.reserve(extra_runtime_backends.size() + 1);
+        split_backends.push_back(runtime_backend);
+        for (ggml_backend_t backend : extra_runtime_backends) {
+            if (backend != nullptr) {
+                split_backends.push_back(backend);
+            }
+        }
+
+        auto manager = weight_manager.lock();
+        if (manager == nullptr) {
+            LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str());
+            return false;
+        }
+
+        sd::GraphCutLayerSplitAssignment assignment;
+        auto canonicalize_param = [this](ggml_tensor* tensor) {
+            return canonical_param_tensor(tensor);
+        };
+        if (!sd::partition_graph_cut_layer_split(get_desc().c_str(),
+                                                 gf,
+                                                 plan,
+                                                 split_backends,
+                                                 graph_cut_layer_split_backend_vram_limits_,
+                                                 max_graph_vram_bytes,
+                                                 graph_cut_layer_split_assignments_,
+                                                 canonicalize_param,
+                                                 &assignment)) {
+            return false;
+        }
+
+        for (size_t i = 0; i < split_backends.size(); i++) {
+            if (assignment.tensors_by_backend[i].empty()) {
+                continue;
+            }
+            if (!manager->assign_compute_backend(assignment.tensors_by_backend[i], split_backends[i])) {
+                LOG_ERROR("%s graph-cut layer split failed to assign params to %s",
+                          get_desc().c_str(),
+                          sd::layer_split_backend_device_display_name(split_backends[i]).c_str());
+                return false;
+            }
+        }
+
+        graph_cut_layer_split_node_assignments_ = std::move(assignment.node_assignments);
+        sd::log_graph_cut_layer_split_assignment(get_desc().c_str(), split_backends, assignment);
+
+        return true;
+    }
+
     struct PersistentExternalBinding {
         ggml_backend_buffer_t buffer = nullptr;
         void* data                   = nullptr;
@@ -2972,6 +3112,11 @@ public:
         GGML_ASSERT(gf != nullptr);
         rebuild_params_tensor_set();
 
+        if (!assign_graph_cut_layer_split_backends(gf)) {
+            free_compute_ctx();
+            return std::nullopt;
+        }
+
         if (can_attempt_graph_cut_segmented_compute()) {
             GraphCutPlan plan;
             if (!resolve_graph_cut_plan(gf, &plan)) {
@@ -3025,6 +3170,22 @@ public:
         stream_layers_enabled = enabled;
     }
 
+    void set_graph_cut_layer_split_enabled(bool enabled) {
+        graph_cut_layer_split_enabled = enabled;
+        if (!enabled) {
+            graph_cut_layer_split_assignments_.clear();
+            graph_cut_layer_split_node_assignments_.clear();
+            graph_cut_layer_split_primary_notice_logged_ = false;
+        }
+    }
+
+    void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) {
+        graph_cut_layer_split_backend_vram_limits_ = limits;
+        graph_cut_layer_split_assignments_.clear();
+        graph_cut_layer_split_node_assignments_.clear();
+        graph_cut_layer_split_primary_notice_logged_ = false;
+    }
+
     void set_runtime_backends(const std::vector<ggml_backend_t>& backends) {
         extra_runtime_backends.clear();
         for (ggml_backend_t backend : backends) {
@@ -3036,6 +3197,9 @@ public:
                 extra_runtime_backends.push_back(backend);
             }
         }
+        graph_cut_layer_split_assignments_.clear();
+        graph_cut_layer_split_node_assignments_.clear();
+        graph_cut_layer_split_primary_notice_logged_ = false;
         if (is_multi_device() && stream_layers_enabled) {
             LOG_WARN("%s: --stream-layers is not supported with multiple runtime backends; ignoring",
                      get_desc().c_str());

src/core/layer_split_partition.cpp

@@ -1,9 +1,11 @@
 #include "core/layer_split_partition.h"
 
 #include <algorithm>
-#include <cstdint>
 #include <cstdlib>
 #include <cstring>
+#include <limits>
+#include <unordered_set>
+#include <utility>
 
 #include "core/util.h"
 
@@ -62,160 +64,194 @@ namespace sd {
         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 graph_cut_layer_split_backend_vram_limit(const std::vector<size_t>& backend_vram_limits,
+                                                           size_t backend_index,
+                                                           size_t primary_backend_vram_limit) {
+        if (backend_index < backend_vram_limits.size()) {
+            return backend_vram_limits[backend_index];
+        }
+        return backend_index == 0 ? primary_backend_vram_limit : 0;
     }
 
-    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.
+    static std::vector<int64_t> graph_cut_layer_split_backend_capacities(const std::vector<ggml_backend_t>& backends,
+                                                                         const std::vector<size_t>& backend_vram_limits,
+                                                                         size_t primary_backend_vram_limit) {
+        std::vector<int64_t> capacities(backends.size(), std::numeric_limits<int64_t>::max() / 4);
         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];
+            if (free_bytes > 0) {
+                capacities[i] = std::max<int64_t>((int64_t)free_bytes - compute_headroom_bytes, 0);
+            }
+            size_t limit_bytes = graph_cut_layer_split_backend_vram_limit(backend_vram_limits,
+                                                                          i,
+                                                                          primary_backend_vram_limit);
+            if (limit_bytes > 0) {
+                capacities[i] = std::min<int64_t>(capacities[i], (int64_t)limit_bytes);
+            }
         }
+        return capacities;
+    }
 
-        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;
-        }
+    bool partition_graph_cut_layer_split(const char* desc,
+                                         ggml_cgraph* gf,
+                                         const sd::ggml_graph_cut::Plan& plan,
+                                         const std::vector<ggml_backend_t>& split_backends,
+                                         const std::vector<size_t>& backend_vram_limits,
+                                         size_t primary_backend_vram_limit,
+                                         std::unordered_map<const ggml_tensor*, ggml_backend_t>& param_assignments,
+                                         const std::function<ggml_tensor*(ggml_tensor*)>& canonical_param_tensor,
+                                         GraphCutLayerSplitAssignment* assignment_out) {
+        GGML_ASSERT(gf != nullptr);
+        GGML_ASSERT(assignment_out != nullptr);
+        GGML_ASSERT(canonical_param_tensor != nullptr);
+        GGML_ASSERT(!split_backends.empty());
+
+        GraphCutLayerSplitAssignment assignment;
+        assignment.segment_count = plan.segments.size();
+        assignment.tensors_by_backend.resize(split_backends.size());
+        assignment.bytes_by_backend.resize(split_backends.size(), 0);
+        assignment.first_segment_by_backend.resize(split_backends.size(), plan.segments.size());
+        assignment.last_segment_by_backend.resize(split_backends.size(), 0);
 
-        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;
+        std::vector<std::vector<ggml_tensor*>> segment_params(plan.segments.size());
+        std::vector<int64_t> segment_param_bytes(plan.segments.size(), 0);
+        std::unordered_set<ggml_tensor*> seen_params;
+        for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) {
+            std::vector<ggml_tensor*> params = sd::ggml_graph_cut::param_tensors(gf, plan.segments[seg_idx]);
+            for (ggml_tensor* raw_param : params) {
+                ggml_tensor* param = canonical_param_tensor(raw_param);
+                if (param == nullptr || !seen_params.insert(param).second) {
+                    continue;
+                }
+                segment_params[seg_idx].push_back(param);
+                segment_param_bytes[seg_idx] += (int64_t)ggml_nbytes(param);
             }
-            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++;
+        int64_t total_param_bytes = 0;
+        for (int64_t bytes : segment_param_bytes) {
+            total_param_bytes += bytes;
+        }
+        if (total_param_bytes <= 0) {
+            LOG_ERROR("%s graph-cut layer split found no graph params to assign", desc);
+            return false;
+        }
+
+        std::vector<int64_t> backend_capacities = graph_cut_layer_split_backend_capacities(split_backends,
+                                                                                           backend_vram_limits,
+                                                                                           primary_backend_vram_limit);
+
+        std::vector<ggml_backend_t> backend_by_segment(plan.segments.size(), split_backends[0]);
+        size_t current_backend = 0;
+        int64_t current_used   = 0;
+        for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) {
+            int64_t bytes = segment_param_bytes[seg_idx];
+            while (current_backend + 1 < split_backends.size() &&
+                   bytes > 0 &&
+                   current_used + bytes > backend_capacities[current_backend]) {
+                current_backend++;
+                current_used = 0;
+            }
+            if (bytes > 0 && current_used + bytes > backend_capacities[current_backend]) {
+                LOG_ERROR("%s graph-cut layer split: segment %zu needs %.1f MB on %s, but only %.1f MB is available under current VRAM limits",
+                          desc,
+                          seg_idx,
+                          (current_used + bytes) / (1024.0 * 1024.0),
+                          layer_split_backend_device_display_name(split_backends[current_backend]).c_str(),
+                          backend_capacities[current_backend] / (1024.0 * 1024.0));
+                return false;
+            }
+            current_used += bytes;
+            backend_by_segment[seg_idx] = split_backends[current_backend];
+
+            for (ggml_tensor* param : segment_params[seg_idx]) {
+                ggml_backend_t target_backend = split_backends[current_backend];
+                auto assigned_it              = param_assignments.find(param);
+                if (assigned_it == param_assignments.end()) {
+                    param_assignments[param]            = target_backend;
+                    assignment.has_new_param_assignment = true;
+                } else {
+                    target_backend = assigned_it->second;
                 }
-                target = std::min(target, backends.size() - 1);
-                target = layer_split_supported_target(desc, kv.first, kv.second, backends, target);
+
+                auto backend_it = std::find(split_backends.begin(), split_backends.end(), target_backend);
+                if (backend_it == split_backends.end()) {
+                    LOG_ERROR("%s graph-cut layer split tensor '%s' is assigned to an unavailable backend",
+                              desc,
+                              ggml_get_name(param));
+                    return false;
+                }
+                size_t backend_idx                               = (size_t)std::distance(split_backends.begin(), backend_it);
+                assignment.first_segment_by_backend[backend_idx] = std::min(assignment.first_segment_by_backend[backend_idx], seg_idx);
+                assignment.last_segment_by_backend[backend_idx]  = std::max(assignment.last_segment_by_backend[backend_idx], seg_idx + 1);
+                assignment.tensors_by_backend[backend_idx].push_back(param);
+                assignment.bytes_by_backend[backend_idx] += (int64_t)ggml_nbytes(param);
+            }
+        }
+
+        const int n_nodes = ggml_graph_n_nodes(gf);
+        for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) {
+            ggml_backend_t backend = backend_by_segment[seg_idx];
+            const auto& segment    = plan.segments[seg_idx];
+            for (int node_index : segment.internal_node_indices) {
+                if (node_index < 0 || node_index >= n_nodes) {
+                    continue;
+                }
+                ggml_tensor* node = ggml_graph_node(gf, node_index);
+                if (node != nullptr) {
+                    assignment.node_assignments[node] = backend;
+                }
+            }
+            for (int node_index : segment.output_node_indices) {
+                if (node_index < 0 || node_index >= n_nodes) {
+                    continue;
+                }
+                ggml_tensor* node = ggml_graph_node(gf, node_index);
+                if (node != nullptr) {
+                    assignment.node_assignments[node] = backend;
+                }
+            }
+        }
+
+        *assignment_out = std::move(assignment);
+        return true;
+    }
+
+    void log_graph_cut_layer_split_assignment(const char* desc,
+                                              const std::vector<ggml_backend_t>& split_backends,
+                                              const GraphCutLayerSplitAssignment& assignment) {
+        for (size_t i = 0; i < split_backends.size(); i++) {
+            if (i >= assignment.tensors_by_backend.size() ||
+                assignment.tensors_by_backend[i].empty()) {
+                continue;
+            }
+            size_t first_segment = assignment.first_segment_by_backend[i] == assignment.segment_count
+                                       ? 0
+                                       : assignment.first_segment_by_backend[i];
+            size_t last_segment  = assignment.last_segment_by_backend[i];
+            if (assignment.has_new_param_assignment) {
+                LOG_INFO("%s graph-cut layer split: %s <- segments [%zu, %zu), %zu tensors, %.1f MB",
+                         desc,
+                         layer_split_backend_device_display_name(split_backends[i]).c_str(),
+                         first_segment,
+                         last_segment,
+                         assignment.tensors_by_backend[i].size(),
+                         assignment.bytes_by_backend[i] / (1024.0 * 1024.0));
             } else {
-                auto target_it = non_block_targets.find(kv.first);
-                if (target_it != non_block_targets.end()) {
-                    target = target_it->second;
-                }
+                LOG_DEBUG("%s graph-cut layer split: %s <- segments [%zu, %zu), %zu tensors, %.1f MB",
+                          desc,
+                          layer_split_backend_device_display_name(split_backends[i]).c_str(),
+                          first_segment,
+                          last_segment,
+                          assignment.tensors_by_backend[i].size(),
+                          assignment.bytes_by_backend[i] / (1024.0 * 1024.0));
             }
-            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

@@ -1,23 +1,43 @@
 #ifndef __SD_CORE_LAYER_SPLIT_PARTITION_H__
 #define __SD_CORE_LAYER_SPLIT_PARTITION_H__
 
-#include <map>
+#include <cstdint>
+#include <functional>
 #include <string>
+#include <unordered_map>
 #include <vector>
 
 #include "ggml-backend.h"
 #include "ggml.h"
 
+#include "core/ggml_graph_cut.h"
+
 namespace sd {
 
+    struct GraphCutLayerSplitAssignment {
+        std::vector<std::vector<ggml_tensor*>> tensors_by_backend;
+        std::vector<int64_t> bytes_by_backend;
+        std::vector<size_t> first_segment_by_backend;
+        std::vector<size_t> last_segment_by_backend;
+        std::unordered_map<const ggml_tensor*, ggml_backend_t> node_assignments;
+        size_t segment_count          = 0;
+        bool has_new_param_assignment = false;
+    };
+
     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);
+    bool partition_graph_cut_layer_split(const char* desc,
+                                         ggml_cgraph* gf,
+                                         const sd::ggml_graph_cut::Plan& plan,
+                                         const std::vector<ggml_backend_t>& split_backends,
+                                         const std::vector<size_t>& backend_vram_limits,
+                                         size_t primary_backend_vram_limit,
+                                         std::unordered_map<const ggml_tensor*, ggml_backend_t>& param_assignments,
+                                         const std::function<ggml_tensor*(ggml_tensor*)>& canonical_param_tensor,
+                                         GraphCutLayerSplitAssignment* assignment_out);
+    void log_graph_cut_layer_split_assignment(const char* desc,
+                                              const std::vector<ggml_backend_t>& split_backends,
+                                              const GraphCutLayerSplitAssignment& assignment);
 
 }  // namespace sd
 

src/model/diffusion/control.hpp

@@ -5,7 +5,8 @@
 #include "model_loader.h"
 #include "model_manager.h"
 
-#define CONTROL_NET_GRAPH_SIZE 1536
+// Match main UNet's MAX_GRAPH_SIZE so SDXL ControlNet (transformer_depth={1,2,10}) fits.
+#define CONTROL_NET_GRAPH_SIZE MAX_GRAPH_SIZE
 
 /*
     =================================== ControlNet ===================================

src/model/diffusion/flux.hpp

@@ -4,6 +4,7 @@
 #include <memory>
 #include <vector>
 
+#include "core/util.h"
 #include "model/adapter/pulid.hpp"
 #include "model/common/rope.hpp"
 #include "model/diffusion/dit.hpp"
@@ -1400,18 +1401,28 @@ namespace Flux {
         std::vector<float> dct_vec;
         sd::Tensor<float> guidance_tensor;
         SDVersion version;
-        bool use_mask = false;
+        bool use_mask = true;
 
         FluxRunner(ggml_backend_t backend,
                    const String2TensorStorage& tensor_storage_map      = {},
                    const std::string prefix                            = "",
                    SDVersion version                                   = VERSION_FLUX,
-                   bool use_mask                                       = false,
-                   std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+                   std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
+                   const char* model_args                              = nullptr)
             : DiffusionModelRunner(backend, prefix, weight_manager),
               config(FluxConfig::detect_from_weights(tensor_storage_map, prefix, version)),
-              version(version),
-              use_mask(use_mask) {
+              version(version) {
+            for (const auto& [key, value] : parse_key_value_args(model_args, "model arg")) {
+                if (key == "chroma_use_dit_mask") {
+                    bool parsed = true;
+                    if (parse_strict_bool(value, parsed)) {
+                        use_mask = parsed;
+                    } else {
+                        LOG_WARN("ignoring invalid Chroma DiT model arg '%s=%s'", key.c_str(), value.c_str());
+                    }
+                }
+            }
+
             if (config.is_chroma) {
                 LOG_INFO("Using pruned modulation (Chroma)");
             }
@@ -1718,7 +1729,6 @@ namespace Flux {
                                                                             tensor_storage_map,
                                                                             "model.diffusion_model",
                                                                             VERSION_FLUX2,
-                                                                            false,
                                                                             model_manager);
 
             if (!model_manager->register_runner_params("Flux test",

src/model/diffusion/qwen_image.hpp

@@ -3,6 +3,7 @@
 
 #include <memory>
 
+#include "core/util.h"
 #include "model/common/block.hpp"
 #include "model/diffusion/dit.hpp"
 #include "model/diffusion/flux.hpp"
@@ -566,12 +567,21 @@ namespace Qwen {
                         const String2TensorStorage& tensor_storage_map      = {},
                         const std::string prefix                            = "",
                         SDVersion version                                   = VERSION_QWEN_IMAGE,
-                        bool zero_cond_t                                    = false,
-                        std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+                        std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
+                        const char* model_args                              = nullptr)
             : DiffusionModelRunner(backend, prefix, weight_manager),
               config(QwenImageConfig::detect_from_weights(tensor_storage_map, prefix)),
               version(version) {
-            config.zero_cond_t = config.zero_cond_t || zero_cond_t;
+            for (const auto& [key, value] : parse_key_value_args(model_args, "model arg")) {
+                if (key == "qwen_image_zero_cond_t") {
+                    bool parsed = false;
+                    if (parse_strict_bool(value, parsed)) {
+                        config.zero_cond_t = config.zero_cond_t || parsed;
+                    } else {
+                        LOG_WARN("ignoring invalid Qwen Image model arg '%s=%s'", key.c_str(), value.c_str());
+                    }
+                }
+            }
             if (version == VERSION_QWEN_IMAGE_LAYERED) {
                 config.use_additional_t_cond = true;
             }
@@ -775,7 +785,6 @@ namespace Qwen {
                                                                                             tensor_storage_map,
                                                                                             "model.diffusion_model",
                                                                                             VERSION_QWEN_IMAGE,
-                                                                                            false,
                                                                                             model_manager);
 
             if (!model_manager->register_runner_params("Qwen image test",

src/model_io/safetensors_io.cpp

@@ -43,7 +43,7 @@ bool is_safetensors_file(const std::string& file_path) {
     }
 
     size_t header_size_ = model_io::read_u64(header_size_buf);
-    if (header_size_ >= file_size_ || header_size_ <= 2) {
+    if (header_size_ > file_size_ - ST_HEADER_SIZE_LEN || header_size_ <= 2) {
         return false;
     }
 
@@ -114,10 +114,11 @@ bool read_safetensors_file(const std::string& file_path,
     }
 
     size_t header_size_ = model_io::read_u64(header_size_buf);
-    if (header_size_ >= file_size_) {
+    if (header_size_ > file_size_ - ST_HEADER_SIZE_LEN) {
         set_error(error, "invalid safetensor file '" + file_path + "'");
         return false;
     }
+    const size_t data_start = ST_HEADER_SIZE_LEN + header_size_;
 
     // read header
     std::vector<char> header_buf;
@@ -156,6 +157,10 @@ bool read_safetensors_file(const std::string& file_path,
 
         size_t begin = tensor_info["data_offsets"][0].get<size_t>();
         size_t end   = tensor_info["data_offsets"][1].get<size_t>();
+        if (begin > end || end > file_size_ - data_start) {
+            set_error(error, "data offsets out of bounds for tensor '" + name + "'");
+            return false;
+        }
 
         ggml_type type = safetensors_dtype_to_ggml_type(dtype);
         if (type == GGML_TYPE_COUNT) {
@@ -187,7 +192,7 @@ bool read_safetensors_file(const std::string& file_path,
             n_dims = 1;
         }
 
-        TensorStorage tensor_storage(name, type, ne, n_dims, 0, ST_HEADER_SIZE_LEN + header_size_ + begin);
+        TensorStorage tensor_storage(name, type, ne, n_dims, 0, data_start + begin);
         tensor_storage.reverse_ne();
 
         size_t tensor_data_size = end - begin;

src/model_manager.cpp

@@ -134,7 +134,8 @@ bool ModelManager::register_param_tensors(const std::string& desc,
                                           ggml_backend_t compute_backend,
                                           ggml_backend_t params_backend,
                                           size_t* registered_tensor_size,
-                                          bool allow_split_buffer) {
+                                          bool allow_split_buffer,
+                                          bool params_follow_compute_backend) {
     if (desc.empty()) {
         LOG_ERROR("model manager tensor desc is empty");
         return false;
@@ -158,14 +159,15 @@ 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;
-        state->allow_split_buffer = allow_split_buffer;
+        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;
+        state->params_follow_compute_backend = params_follow_compute_backend;
         new_states.push_back(std::move(state));
     }
 
@@ -177,6 +179,102 @@ bool ModelManager::register_param_tensors(const std::string& desc,
     return true;
 }
 
+bool ModelManager::unregister_param_tensors(const std::string& desc, size_t* registered_tensor_size) {
+    if (desc.empty()) {
+        return true;
+    }
+
+    std::unordered_set<TensorState*> target_states;
+    size_t released_size = 0;
+    for (auto& state : tensor_states_) {
+        if (state == nullptr || state->desc != desc) {
+            continue;
+        }
+        if (state->active_prepare_count > 0) {
+            LOG_ERROR("model manager cannot unregister active %s tensor '%s'",
+                      desc.c_str(),
+                      state->name.c_str());
+            return false;
+        }
+        target_states.insert(state.get());
+        if (state->tensor != nullptr) {
+            released_size += ggml_nbytes(state->tensor);
+        }
+    }
+
+    if (target_states.empty()) {
+        return true;
+    }
+
+    release_compute_staging_blocks(false);
+
+    std::vector<ParamsStorageBlock*> storage_blocks_to_release;
+    std::unordered_set<TensorState*> affected_storage_states;
+    for (const auto& block : params_storage_blocks_) {
+        if (block == nullptr) {
+            continue;
+        }
+        bool has_target_state = false;
+        for (TensorState* state : block->states) {
+            if (state != nullptr && target_states.count(state) > 0) {
+                has_target_state = true;
+                break;
+            }
+        }
+        if (!has_target_state) {
+            continue;
+        }
+        storage_blocks_to_release.push_back(block.get());
+        for (TensorState* state : block->states) {
+            if (state != nullptr) {
+                affected_storage_states.insert(state);
+            }
+        }
+    }
+
+    for (TensorState* state : affected_storage_states) {
+        if (state == nullptr) {
+            continue;
+        }
+        if (state->active_prepare_count > 0 || state->staged_to_compute_backend) {
+            LOG_ERROR("model manager cannot unregister %s while tensor '%s' is active",
+                      desc.c_str(),
+                      state->name.c_str());
+            return false;
+        }
+    }
+
+    for (ParamsStorageBlock* block : storage_blocks_to_release) {
+        if (block != nullptr) {
+            free_params_storage_block(*block);
+            erase_params_storage_block(block);
+        }
+    }
+
+    for (auto it = tensor_states_by_name_.begin(); it != tensor_states_by_name_.end();) {
+        if (target_states.count(it->second) > 0) {
+            it = tensor_states_by_name_.erase(it);
+        } else {
+            ++it;
+        }
+    }
+    tensor_states_.erase(std::remove_if(tensor_states_.begin(),
+                                        tensor_states_.end(),
+                                        [&](const std::unique_ptr<TensorState>& s) {
+                                            return s == nullptr || target_states.count(s.get()) > 0;
+                                        }),
+                         tensor_states_.end());
+
+    if (registered_tensor_size != nullptr) {
+        if (released_size > *registered_tensor_size) {
+            *registered_tensor_size = 0;
+        } else {
+            *registered_tensor_size -= released_size;
+        }
+    }
+    return true;
+}
+
 bool ModelManager::load_all_params_eagerly() {
     std::vector<TensorState*> all_states;
     all_states.reserve(tensor_states_.size());
@@ -919,6 +1017,54 @@ bool ModelManager::resolve_required_tensor_states(const std::vector<ggml_tensor*
     return true;
 }
 
+bool ModelManager::assign_compute_backend(const std::vector<ggml_tensor*>& tensors,
+                                          ggml_backend_t compute_backend) {
+    if (tensors.empty()) {
+        return true;
+    }
+    if (compute_backend == nullptr) {
+        LOG_ERROR("model manager cannot assign tensors to a null compute backend");
+        return false;
+    }
+
+    std::vector<TensorState*> required_states;
+    if (!resolve_required_tensor_states(tensors, required_states)) {
+        return false;
+    }
+
+    for (TensorState* state : required_states) {
+        if (state == nullptr || state->tensor == nullptr) {
+            continue;
+        }
+
+        const bool params_follow_compute = state->params_follow_compute_backend ||
+                                           state->residency_mode == ResidencyMode::Disk;
+        const bool compute_changes = state->compute_backend != compute_backend;
+        const bool params_changes  = params_follow_compute && state->params_backend != compute_backend;
+        if (!compute_changes && !params_changes) {
+            continue;
+        }
+
+        if (state->active_prepare_count > 0 || state->staged_to_compute_backend) {
+            LOG_ERROR("model manager cannot move active tensor '%s' to another compute backend",
+                      state->name.c_str());
+            return false;
+        }
+        if (params_changes && state->loaded_to_params_backend) {
+            LOG_ERROR("model manager cannot move loaded tensor '%s' to another params backend",
+                      state->name.c_str());
+            return false;
+        }
+
+        state->compute_backend = compute_backend;
+        if (params_follow_compute) {
+            state->params_backend = compute_backend;
+        }
+    }
+
+    return true;
+}
+
 bool ModelManager::prepare_params(const std::vector<ggml_tensor*>& tensors) {
     if (tensors.empty()) {
         return true;

src/model_manager.h

@@ -33,11 +33,12 @@ private:
         ggml_tensor* tensor = nullptr;
         std::string desc;
 
-        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;
+        ResidencyMode residency_mode       = ResidencyMode::ParamBackend;
+        ggml_backend_t compute_backend     = nullptr;
+        ggml_backend_t params_backend      = nullptr;
+        bool allow_split_buffer            = false;
+        bool params_follow_compute_backend = false;
+        bool metadata_validated            = false;
 
         int active_prepare_count = 0;
 
@@ -129,8 +130,12 @@ public:
                                 ResidencyMode residency_mode,
                                 ggml_backend_t compute_backend,
                                 ggml_backend_t params_backend,
-                                size_t* registered_tensor_size = nullptr,
-                                bool allow_split_buffer        = false);
+                                size_t* registered_tensor_size     = nullptr,
+                                bool allow_split_buffer            = false,
+                                bool params_follow_compute_backend = false);
+
+    bool unregister_param_tensors(const std::string& desc,
+                                  size_t* registered_tensor_size = nullptr);
 
     template <typename Runner>
     bool register_runner_params(const std::string& desc,
@@ -170,6 +175,8 @@ public:
     bool validate_registered_tensors();
     bool load_all_params_eagerly();
 
+    bool assign_compute_backend(const std::vector<ggml_tensor*>& tensors,
+                                ggml_backend_t compute_backend) override;
     bool prepare_params(const std::vector<ggml_tensor*>& tensors) override;
     void release_compute_backend_params(const std::vector<ggml_tensor*>& tensors) override;
     void release_params_backend_params(const std::vector<ggml_tensor*>& tensors) override;

src/name_conversion.cpp

@@ -736,6 +736,62 @@ std::string convert_diffusers_dit_to_original_krea2(std::string name) {
     return name;
 }
 
+// Convert a diffusers-format ControlNet tensor name to the original (LDM/lllyasviel) layout
+// declared by ControlNetBlock. Reuses the UNet down/mid conversion for the shared encoder
+// (down_blocks, mid_block, time_embedding, add_embedding, conv_in) and adds the ControlNet-only
+// mappings: input_hint_block, zero_convs, middle_block_out.
+std::string convert_diffusers_controlnet_to_original_sdxl(std::string name) {
+    name = convert_diffusers_unet_to_original_sdxl(std::move(name));
+
+    static const std::vector<std::pair<std::string, std::string>> prefix_map = {
+        {"controlnet_cond_embedding.conv_in.", "input_hint_block.0."},
+        {"controlnet_cond_embedding.blocks.0.", "input_hint_block.2."},
+        {"controlnet_cond_embedding.blocks.1.", "input_hint_block.4."},
+        {"controlnet_cond_embedding.blocks.2.", "input_hint_block.6."},
+        {"controlnet_cond_embedding.blocks.3.", "input_hint_block.8."},
+        {"controlnet_cond_embedding.blocks.4.", "input_hint_block.10."},
+        {"controlnet_cond_embedding.blocks.5.", "input_hint_block.12."},
+        {"controlnet_cond_embedding.conv_out.", "input_hint_block.14."},
+        {"controlnet_mid_block.", "middle_block_out.0."},
+    };
+    for (const auto& p : prefix_map) {
+        if (starts_with(name, p.first)) {
+            return p.second + name.substr(p.first.size());
+        }
+    }
+
+    static const std::string controlnet_down_prefix = "controlnet_down_blocks.";
+    if (starts_with(name, controlnet_down_prefix)) {
+        size_t rest_start = controlnet_down_prefix.size();
+        size_t dot        = name.find('.', rest_start);
+        if (dot != std::string::npos) {
+            std::string idx = name.substr(rest_start, dot - rest_start);
+            return "zero_convs." + idx + ".0" + name.substr(dot);
+        }
+    }
+
+    return name;
+}
+
+static bool is_diffusers_controlnet_name(const std::string& name) {
+    static const std::vector<std::string> heads = {
+        "controlnet_cond_embedding.",
+        "controlnet_down_blocks.",
+        "controlnet_mid_block.",
+        "down_blocks.",
+        "mid_block.",
+        "time_embedding.",
+        "add_embedding.",
+        "conv_in.",
+    };
+    for (const auto& h : heads) {
+        if (starts_with(name, h)) {
+            return true;
+        }
+    }
+    return false;
+}
+
 std::string convert_diffusion_model_name(std::string name, std::string prefix, SDVersion version) {
     if (sd_version_is_sd1(version) || sd_version_is_sd2(version)) {
         name = convert_diffusers_unet_to_original_sd1(name);
@@ -1338,6 +1394,9 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
                 name = name.substr(pos + 1);
             }
         }
+        if (sd_version_is_sdxl(version) && is_diffusers_controlnet_name(name)) {
+            name = convert_diffusers_controlnet_to_original_sdxl(name);
+        }
     }
 
     if (is_lora) {

src/runtime/denoiser.hpp

@@ -1897,10 +1897,10 @@ public:
     }
 
     sd::Tensor<float> operator()(double sigma_a, double sigma_b) {
-        double a          = clamp(std::min(sigma_a, sigma_b));
-        double b          = clamp(std::max(sigma_a, sigma_b));
-        auto dW           = w(b) - w(a);
-        float span        = static_cast<float>(std::max(std::abs(sigma_b - sigma_a), 1e-12));
+        double a   = clamp(std::min(sigma_a, sigma_b));
+        double b   = clamp(std::max(sigma_a, sigma_b));
+        auto dW    = w(b) - w(a);
+        float span = static_cast<float>(std::max(std::abs(sigma_b - sigma_a), 1e-12));
         return dW * (1.0f / std::sqrt(span));
     }
 

src/runtime/imatrix.cpp

@@ -287,6 +287,10 @@ bool IMatrixCollector::load_imatrix(const char* fname) {
         if (e.values.empty()) {
             e.values.resize(nval, 0);
             e.counts.resize(nval, 0);
+        } else if (e.values.size() != (size_t)nval) {
+            LOG_ERROR("inconsistent size for a repeated entry (%d vs %d)\n", (int)e.values.size(), nval);
+            stats_ = {};
+            return false;
         }
 
         std::vector<float> tmp(nval);

src/stable-diffusion.cpp

@@ -57,11 +57,11 @@
 #include "name_conversion.h"
 #include "runtime/latent-preview.h"
 
+#include <atomic>
+
 const char* sd_vae_format_name(enum sd_vae_format_t format);
 static SDVersion sd_vae_format_to_version(enum sd_vae_format_t format, SDVersion fallback);
 
-#include <atomic>
-
 const char* model_version_to_str[] = {
     "SD 1.x",
     "SD 1.x Inpaint",
@@ -222,9 +222,13 @@ public:
     std::string split_mode_spec;
     bool auto_fit_enabled = false;
 
+    bool diffusion_conv_direct = false;
+
     bool is_using_v_parameterization     = false;
     bool is_using_edm_v_parameterization = false;
 
+    size_t control_net_params_mem_size = 0;
+
     std::shared_ptr<ModelManager> model_manager;
 
     std::shared_ptr<Denoiser> denoiser = std::make_shared<CompVisDenoiser>();
@@ -268,6 +272,15 @@ public:
         return max_vram_assignment.bytes_for_backend(backend_for(module));
     }
 
+    std::vector<size_t> layer_split_vram_limits_for_backends(const std::vector<ggml_backend_t>& backends) {
+        std::vector<size_t> limits;
+        limits.reserve(backends.size());
+        for (ggml_backend_t backend : backends) {
+            limits.push_back(max_vram_assignment.bytes_for_backend(backend));
+        }
+        return limits;
+    }
+
     bool ensure_backend_pair(SDBackendModule module) {
         if (backend_for(module) == nullptr) {
             return false;
@@ -427,8 +440,9 @@ public:
                                                      params_mem_size);
     }
 
-    // Register each layer-split partition with its compute backend; the
-    // ModelManager handles allocation, staging, and LoRA by backend.
+    // Register graph-cut layer-split tensors on the primary backend first.
+    // The first real graph assigns each param tensor to a runtime backend
+    // before weights are loaded or staged.
     template <typename T>
     bool register_layer_split_runner_params(const std::string& desc,
                                             const std::shared_ptr<T>& model,
@@ -459,51 +473,98 @@ public:
                                                          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);
+        model->set_graph_cut_layer_split_backend_vram_limits(layer_split_vram_limits_for_backends(module_backends));
+        model->set_graph_cut_layer_split_enabled(true);
         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)) {
+        ggml_backend_t initial_params_backend = params_follow_runtime ? module_backends[0] : params_backend_for(module);
+        if (initial_params_backend == nullptr) {
+            return false;
+        }
+
+        LOG_INFO("%s graph-cut layer split: deferring %zu tensors across %zu runtime backends until first graph",
+                 desc.c_str(),
+                 group_tensors.size(),
+                 module_backends.size());
+
+        return model_manager->register_param_tensors(desc,
+                                                     std::move(group_tensors),
+                                                     residency_mode,
+                                                     module_backends[0],
+                                                     initial_params_backend,
+                                                     params_mem_size,
+                                                     false,
+                                                     params_follow_runtime);
+    }
+
+    bool unload_control_net() {
+        if (control_net == nullptr) {
+            return true;
+        }
+        if (model_manager != nullptr) {
+            if (!model_manager->unregister_param_tensors("ControlNet", &control_net_params_mem_size)) {
                 return false;
             }
         }
+        control_net.reset();
+        control_net_params_mem_size = 0;
+        return true;
+    }
+
+    bool load_control_net_from_file(const std::string& path) {
+        if (path.empty()) {
+            LOG_ERROR("sd_ctx_load_control_net: empty path");
+            return false;
+        }
+        if (model_manager == nullptr) {
+            LOG_ERROR("sd_ctx_load_control_net: model_manager not initialized");
+            return false;
+        }
+
+        if (!unload_control_net()) {
+            return false;
+        }
+
+        ModelLoader& shared_loader = model_manager->loader();
+        if (!shared_loader.init_from_file(path)) {
+            LOG_ERROR("sd_ctx_load_control_net: failed to load '%s'", path.c_str());
+            return false;
+        }
+        shared_loader.convert_tensors_name();
+
+        if (!ensure_backend_pair(SDBackendModule::CONTROL_NET)) {
+            LOG_ERROR("sd_ctx_load_control_net: control_net backend unavailable");
+            return false;
+        }
+
+        control_net = std::make_shared<ControlNet>(backend_for(SDBackendModule::CONTROL_NET),
+                                                   params_backend_for(SDBackendModule::CONTROL_NET),
+                                                   shared_loader.get_tensor_storage_map(),
+                                                   version,
+                                                   "",
+                                                   model_manager);
+        if (diffusion_conv_direct) {
+            LOG_INFO("Using Conv2d direct in the control net");
+            control_net->set_conv2d_direct_enabled(true);
+        }
+        if (!register_runner_params("ControlNet",
+                                    control_net,
+                                    SDBackendModule::CONTROL_NET,
+                                    &control_net_params_mem_size)) {
+            LOG_ERROR("sd_ctx_load_control_net: register_runner_params failed");
+            control_net.reset();
+            control_net_params_mem_size = 0;
+            return false;
+        }
+        if (!model_manager->validate_registered_tensors()) {
+            LOG_ERROR("sd_ctx_load_control_net: registered tensors validation failed");
+            unload_control_net();
+            return false;
+        }
+        LOG_INFO("sd_ctx_load_control_net: loaded '%s' (%.2f MB)",
+                 path.c_str(),
+                 control_net_params_mem_size / 1024.0 / 1024.0);
         return true;
     }
 
@@ -529,6 +590,16 @@ public:
         return false;
     }
 
+    bool graph_cut_layer_split_active() {
+        for (SDBackendModule module : {SDBackendModule::DIFFUSION, SDBackendModule::TE}) {
+            if (backend_manager.split_mode(module) == SDSplitMode::LAYER &&
+                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>();
@@ -785,6 +856,10 @@ public:
             LOG_WARN("--stream-layers has no effect unless diffusion params backend is cpu; ignoring");
             stream_layers = false;
         }
+        if (eager_load && graph_cut_layer_split_active()) {
+            LOG_WARN("--eager-load is not supported with graph-cut layer split; weights will be prepared lazily");
+            eager_load = 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();
@@ -851,10 +926,12 @@ public:
         model_loader.process_model_files(enable_mmap, needs_writable_mmap);
         load_alphas_cumprod(model_loader);
 
+        diffusion_conv_direct = sd_ctx_params->diffusion_conv_direct;
+
         size_t text_encoder_params_mem_size = 0;
         size_t unet_params_mem_size         = 0;
         size_t vae_params_mem_size          = 0;
-        size_t control_net_params_mem_size  = 0;
+        control_net_params_mem_size         = 0;
         size_t extension_params_mem_size    = 0;
 
         bool tae_preview_only = sd_ctx_params->tae_preview_only;
@@ -921,10 +998,11 @@ public:
                 if (is_chroma) {
                     cond_stage_model = std::make_shared<T5CLIPEmbedder>(backend_for(SDBackendModule::TE),
                                                                         tensor_storage_map,
-                                                                        sd_ctx_params->chroma_use_t5_mask,
-                                                                        sd_ctx_params->chroma_t5_mask_pad,
                                                                         false,
-                                                                        model_manager);
+                                                                        1,
+                                                                        false,
+                                                                        model_manager,
+                                                                        sd_ctx_params->model_args);
                 } else if (version == VERSION_OVIS_IMAGE) {
                     cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
                                                                      tensor_storage_map,
@@ -941,8 +1019,8 @@ public:
                                                                      tensor_storage_map,
                                                                      "model.diffusion_model",
                                                                      version,
-                                                                     sd_ctx_params->chroma_use_dit_mask,
-                                                                     model_manager);
+                                                                     model_manager,
+                                                                     sd_ctx_params->model_args);
             } else if (sd_version_is_flux2(version) || sd_version_is_sefi_image(version)) {
                 bool is_chroma   = false;
                 cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
@@ -955,8 +1033,8 @@ public:
                                                                      tensor_storage_map,
                                                                      "model.diffusion_model",
                                                                      version,
-                                                                     sd_ctx_params->chroma_use_dit_mask,
-                                                                     model_manager);
+                                                                     model_manager,
+                                                                     sd_ctx_params->model_args);
             } else if (sd_version_is_ltxav(version)) {
                 cond_stage_model = std::make_shared<LTXAVEmbedder>(backend_for(SDBackendModule::TE),
                                                                    tensor_storage_map,
@@ -1014,8 +1092,8 @@ public:
                                                                           tensor_storage_map,
                                                                           "model.diffusion_model",
                                                                           version,
-                                                                          sd_ctx_params->qwen_image_zero_cond_t,
-                                                                          model_manager);
+                                                                          model_manager,
+                                                                          sd_ctx_params->model_args);
             } else if (sd_version_is_longcat(version)) {
                 cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
                                                                  tensor_storage_map,
@@ -1027,8 +1105,8 @@ public:
                                                                      tensor_storage_map,
                                                                      "model.diffusion_model",
                                                                      version,
-                                                                     sd_ctx_params->chroma_use_dit_mask,
-                                                                     model_manager);
+                                                                     model_manager,
+                                                                     sd_ctx_params->model_args);
             } else if (version == VERSION_HIDREAM_O1) {
                 cond_stage_model = std::make_shared<HiDreamO1::HiDreamO1Conditioner>(backend_for(SDBackendModule::TE),
                                                                                      tensor_storage_map,
@@ -1364,7 +1442,6 @@ public:
                     high_noise_diffusion_model->set_flash_attention_enabled(true);
                 }
             }
-
         }
 
         LOG_DEBUG("validating model metadata");
@@ -3048,15 +3125,13 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) {
     sd_ctx_params->eager_load           = false;
     sd_ctx_params->enable_mmap          = false;
     sd_ctx_params->diffusion_flash_attn = false;
-    sd_ctx_params->chroma_use_dit_mask  = true;
-    sd_ctx_params->chroma_use_t5_mask   = false;
-    sd_ctx_params->chroma_t5_mask_pad   = 1;
     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->model_args           = nullptr;
     sd_ctx_params->pulid_weights_path   = nullptr;
 }
 
@@ -3096,12 +3171,10 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
              "backend: %s\n"
              "params_backend: %s\n"
              "split_mode: %s\n"
+             "model_args: %s\n"
              "auto_fit: %s\n"
              "flash_attn: %s\n"
              "diffusion_flash_attn: %s\n"
-             "chroma_use_dit_mask: %s\n"
-             "chroma_use_t5_mask: %s\n"
-             "chroma_t5_mask_pad: %d\n"
              "vae_format: %s\n",
              SAFE_STR(sd_ctx_params->model_path),
              SAFE_STR(sd_ctx_params->clip_l_path),
@@ -3132,12 +3205,10 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
              SAFE_STR(sd_ctx_params->backend),
              SAFE_STR(sd_ctx_params->params_backend),
              SAFE_STR(sd_ctx_params->split_mode),
+             SAFE_STR(sd_ctx_params->model_args),
              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->chroma_use_dit_mask),
-             BOOL_STR(sd_ctx_params->chroma_use_t5_mask),
-             sd_ctx_params->chroma_t5_mask_pad,
              sd_vae_format_name(sd_ctx_params->vae_format));
 
     return buf;
@@ -3434,6 +3505,27 @@ SD_API bool sd_ctx_supports_video_generation(const sd_ctx_t* sd_ctx) {
     return sd_version_supports_video_generation(sd_ctx->sd->version);
 }
 
+SD_API bool sd_ctx_load_control_net(sd_ctx_t* sd_ctx, const char* path) {
+    if (sd_ctx == nullptr || sd_ctx->sd == nullptr || path == nullptr) {
+        return false;
+    }
+    return sd_ctx->sd->load_control_net_from_file(path);
+}
+
+SD_API bool sd_ctx_unload_control_net(sd_ctx_t* sd_ctx) {
+    if (sd_ctx == nullptr || sd_ctx->sd == nullptr) {
+        return false;
+    }
+    return sd_ctx->sd->unload_control_net();
+}
+
+SD_API bool sd_ctx_has_control_net(const sd_ctx_t* sd_ctx) {
+    if (sd_ctx == nullptr || sd_ctx->sd == nullptr) {
+        return false;
+    }
+    return sd_ctx->sd->control_net != nullptr;
+}
+
 enum sample_method_t sd_get_default_sample_method(const sd_ctx_t* sd_ctx) {
     if (sd_ctx != nullptr && sd_ctx->sd != nullptr) {
         if (sd_version_is_pid(sd_ctx->sd->version)) {

src/weight_manager.h

@@ -3,10 +3,14 @@
 
 #include <vector>
 
+#include "ggml-backend.h"
+
 struct ggml_tensor;
 
 struct RunnerWeightManager {
     virtual ~RunnerWeightManager()                                                        = default;
+    virtual bool assign_compute_backend(const std::vector<ggml_tensor*>& tensors,
+                                        ggml_backend_t compute_backend)                   = 0;
     virtual bool prepare_params(const std::vector<ggml_tensor*>& tensors)                 = 0;
     virtual void release_compute_backend_params(const std::vector<ggml_tensor*>& tensors) = 0;
     virtual void release_params_backend_params(const std::vector<ggml_tensor*>& tensors)  = 0;