feat: add RPC support (#1629)

Signed-off-by: kei-g <km.8k6ce+github@gmail.com>

CMakeLists.txt

@@ -204,12 +204,33 @@ if(SD_WEBM)
     endif()
 endif()
 
+if (SD_RPC)
+    message("-- Use RPC as backend stable-diffusion")
+    set(GGML_RPC ON)
+    add_definitions(-DSD_USE_RPC)
+endif ()
+
 set(SD_LIB stable-diffusion)
 
 file(GLOB SD_LIB_SOURCES CONFIGURE_DEPENDS
     "src/*.h"
     "src/*.cpp"
     "src/*.hpp"
+    "src/conditioning/*.h"
+    "src/conditioning/*.cpp"
+    "src/conditioning/*.hpp"
+    "src/core/*.h"
+    "src/core/*.cpp"
+    "src/core/*.hpp"
+    "src/extensions/*.h"
+    "src/extensions/*.cpp"
+    "src/extensions/*.hpp"
+    "src/model/*/*.h"
+    "src/model/*/*.cpp"
+    "src/model/*/*.hpp"
+    "src/runtime/*.h"
+    "src/runtime/*.cpp"
+    "src/runtime/*.hpp"
     "src/model_io/*.h"
     "src/model_io/*.cpp"
     "src/tokenizers/*.h"
@@ -312,6 +333,7 @@ add_subdirectory(thirdparty)
 
 target_link_libraries(${SD_LIB} PUBLIC ggml zip)
 target_include_directories(${SD_LIB} PUBLIC . src include)
+target_include_directories(${SD_LIB} PRIVATE src/core)
 target_include_directories(${SD_LIB} PUBLIC . thirdparty)
 target_compile_features(${SD_LIB} PUBLIC c_std_11 cxx_std_17)
 

CONTRIBUTING.md

@@ -44,6 +44,8 @@ Naming conventions:
 
 Some older code in the project may not fully follow the current conventions. Please do not submit PRs that only rewrite existing code to match style rules.
 
+When adding or modifying model implementations, follow the model config and weight detection conventions in [docs/model_config.md](docs/model_config.md).
+
 ## AI-Assisted Contributions
 
 AI tools may be used to assist development, but contributors are responsible for the quality and correctness of the submitted code.

README.md

@@ -15,6 +15,7 @@ API and command-line option may change frequently.***
 
 ## 🔥Important News
 
+* **2026/06/04** 🚀 stable-diffusion.cpp now supports **Ideogram4**
 * **2026/05/31** 🚀 stable-diffusion.cpp now supports **PiD**
 * **2026/05/27** 🚀 stable-diffusion.cpp now supports **Lens**
 * **2026/05/17** 🚀 stable-diffusion.cpp now supports **LTX-2.3**
@@ -50,6 +51,7 @@ API and command-line option may change frequently.***
     - [Anima](./docs/anima.md)
     - [ERNIE-Image](./docs/ernie_image.md)
     - [HiDream-O1-Image](./docs/hidream_o1_image.md)
+    - [Ideogram4](./docs/ideogram4.md)
   - Image Edit Models
     - [FLUX.1-Kontext-dev](./docs/kontext.md)
     - [Qwen Image Edit series](./docs/qwen_image_edit.md)

docs/backend.md

@@ -3,7 +3,7 @@
 `stable-diffusion.cpp` has two backend assignments:
 
 - `--backend` selects the runtime backend used to execute model graphs.
-- `--params-backend` selects the backend used to allocate model parameters.
+- `--params-backend` selects where model parameters are kept.
 
 If `--params-backend` is not set, parameters use the same backend as their module runtime backend.
 
@@ -29,6 +29,12 @@ The same syntax is used for parameter placement:
 sd-cli -m model.safetensors -p "a cat" --backend cuda0 --params-backend te=cpu,vae=cpu
 ```
 
+`--params-backend` also accepts the special value `disk`:
+
+```shell
+sd-cli -m model.safetensors -p "a cat" --backend cuda0 --params-backend disk
+```
+
 Module names are case-insensitive. Hyphens and underscores in module names are ignored, so `clip_vision`, `clip-vision`, and `clipvision` are equivalent.
 
 `all=`, `default=`, and `*=` can be used to set the default backend inside a mixed assignment:
@@ -64,9 +70,11 @@ The special values `auto`, `default`, and an empty backend name select the defau
 
 The special value `gpu` selects the first GPU backend, falling back to the first integrated GPU backend.
 
+The special value `disk` is accepted only by `--params-backend`. `--backend disk` is invalid because `disk` is a parameter residency mode, not a runtime compute backend.
+
 ## Runtime backend vs. parameter backend
 
-The runtime backend controls where graph execution runs. The parameter backend controls where model weights are allocated.
+The runtime backend controls where graph execution runs. The parameter backend controls where model weights are allocated or whether they are reloaded from disk on demand.
 
 For example:
 
@@ -76,6 +84,16 @@ sd-cli -m model.safetensors -p "a cat" --backend cuda0 --params-backend cpu
 
 This runs all modules on `cuda0`, but stores parameters in CPU RAM. During execution, parameters are moved to the runtime backend as needed.
 
+For example:
+
+```shell
+sd-cli -m model.safetensors -p "a cat" --backend cuda0 --params-backend disk
+```
+
+This runs all modules on `cuda0`, reloads parameters from the model file as needed, and releases those parameter buffers after use.
+
+`disk` is never selected implicitly. If `--params-backend` is not set, parameters use the runtime backend.
+
 Per-module assignments can be mixed:
 
 ```shell
@@ -100,23 +118,27 @@ uses one shared CPU backend for both `te` and `vae` runtime execution.
 
 Runtime and parameter assignments also share the same backend cache. If `--backend diffusion=cuda0` and `--params-backend diffusion=cuda0` resolve to the same device, both use the same backend instance.
 
+`--params-backend disk` does not create a separate backend instance. Parameters are loaded lazily using the module runtime backend.
+
 `SDBackendManager` owns the backend instances and frees them when the context or upscaler is destroyed. Model runners receive non-owning runtime and parameter backend pointers and do not free them.
 
 ## Compatibility flags
 
-The older CPU placement flags are still supported:
+The example CLI/server still accepts these older CPU placement flags as compatibility aliases:
 
 - `--clip-on-cpu`
 - `--vae-on-cpu`
 - `--control-net-cpu`
 - `--offload-to-cpu`
 
-`--clip-on-cpu`, `--vae-on-cpu`, and `--control-net-cpu` affect runtime backend assignment only when `--backend` is not set. They map to `te=cpu`, `vae=cpu`, and `controlnet=cpu`.
+`--clip-on-cpu`, `--vae-on-cpu`, and `--control-net-cpu` are deprecated. The example argument layer prepends `te=cpu`, `vae=cpu`, and `controlnet=cpu` to `--backend` before creating the context.
 
-`--offload-to-cpu` affects parameter backend assignment only when `--params-backend` is not set. It is equivalent to:
+`--offload-to-cpu` prepends a CPU default to the parameter assignment in the caller before creating the context:
 
 ```shell
---params-backend cpu
+--params-backend '*=cpu'
 ```
 
-Explicit `--backend` and `--params-backend` assignments are preferred for new commands.
+Because this default is inserted first, later explicit `--params-backend` entries can still override it, for example `--offload-to-cpu --params-backend te=disk` keeps non-TE parameters on CPU and reloads TE parameters from disk.
+
+Library callers should set `backend` and `params_backend` directly. The old CPU/offload fields are no longer part of the C API. Explicit `--backend` and `--params-backend` assignments are preferred for new commands.

docs/ideogram4.md

@@ -0,0 +1,40 @@
+# How to Use
+
+## Download weights
+
+- Download Ideogram4
+    - safetensors: https://huggingface.co/ideogram-ai/ideogram-4-fp8/tree/main/transformer
+- Download Ideogram4 uncond
+    - safetensors: https://huggingface.co/ideogram-ai/ideogram-4-fp8/tree/main/unconditional_transformer
+- Download vae
+    - safetensors: https://huggingface.co/black-forest-labs/FLUX.2-dev/tree/main
+- Download Qwen3-VL-8B-Instruct
+    - gguf: https://huggingface.co/unsloth/Qwen3-VL-8B-Instruct-GGUF/tree/main
+
+## Convert weights
+
+fp8 scale -> bf16
+
+```
+python .\convert_fp8_scale_to_bf16.py --input .\ideogram4_fp8.safetensors --output ideogram4_bf16.safetensors
+python .\convert_fp8_scale_to_bf16.py --input .\ideogram4_uncond_fp8.safetensors --output ideogram4_uncond_bf16.safetensors
+```
+
+bf16 -> q8
+
+```
+.\bin\Release\sd-cli.exe -M convert -m ideogram4_bf16.safetensors -o ideogram4-Q8_0.gguf --tensor-type-rules "^layers.*adaln_modulation.*weight=q8_0,layers.*attention.o.*weight=q8_0,layers.*attention.qkv.*weight=q8_0,layers.*feed_forward.*weight=q8_0" -v
+
+.\bin\Release\sd-cli.exe -M convert -m ideogram4_uncond_bf16.safetensors -o ideogram4_uncond-Q8_0.gguf --tensor-type-rules "^layers.*adaln_modulation.*weight=q8_0,layers.*attention.o.*weight=q8_0,layers.*attention.qkv.*weight=q8_0,layers.*feed_forward.*weight=q8_0" -v
+```
+
+If you want lower VRAM usage, you can change the quantization from q8_0 to a lower-level quantization, such as q4_0.
+
+
+## Examples
+
+```sh
+.\bin\Release\sd-cli.exe --diffusion-model ideogram4-Q8_0.gguf --uncond-diffusion-model ideogram4_uncond-Q8_0.gguf --llm ..\..\llm\Qwen3VL-8B-Instruct-Q4_K_M.gguf --vae ..\..\ComfyUI\models\vae\flux2_ae.safetensors -p '{"high_level_description":"A square 1024 x 1024 luxury fashion magazine cover featuring exactly one short chubby fluffy cat as the main model. The cat sits on a soft ivory studio floor, facing the viewer with a stylish calm expression, wearing tiny black sunglasses, a red silk scarf, and a small gold collar charm. In front of the cat on the floor is a wide horizontal luxury nameplate that clearly reads ideogram4.cpp. The whole design feels premium, fashionable, clean, and editorial.","style_description":{"aesthetics":"luxury fashion magazine cover, high-end pet couture campaign, minimalist editorial design, elegant studio photography, soft paper texture, refined typography, fashionable and polished","lighting":"Soft diffused studio lighting, gentle spotlight on the cat, subtle floor shadow, warm ivory highlights, clean separation between subject and background","photo":"high-resolution fashion editorial photography look, front-facing cat portrait, crisp fur details, glossy sunglasses, clear readable nameplate text, shallow depth of field","medium":"mixed media fashion photography and premium editorial graphic design","color_palette":["#F4EFE7","#111111","#D8B56D","#B73A3A","#FFFFFF","#8A7A6A"]},"compositional_deconstruction":{"canvas":"Square 1024 x 1024 canvas with a normal upright orientation. Do not rotate the poster or any text. Use a clean fashion magazine cover layout.","background":"Warm ivory studio backdrop with subtle paper grain, a soft spotlight gradient, faint floor shadow, and a few minimal gold editorial lines. The background is spacious, premium, and uncluttered.","layout":"Top center has a small elegant headline. Center area features one cat as the main fashion model. Lower foreground has a wide horizontal luxury nameplate placed on the floor in front of the cat. Bottom center has a small footer. All text is horizontal, upright, and readable left to right.","elements":[{"type":"text","desc":"Top center headline reading LOOK WHAT I FOUND in a refined high-fashion serif font. The headline is horizontal, centered, elegant, and secondary to the nameplate text."},{"type":"obj","desc":"Exactly one short chubby fluffy cat sitting in the center like a luxury fashion model. The cat has a large round head, compact body, short legs, soft detailed fur, expressive eyes, and a calm confident pose. The cat is cute and rounded, not tall, not stretched, not duplicated."},{"type":"obj","desc":"Tiny glossy black sunglasses worn naturally by the cat, slightly oversized but still showing the cat face clearly. The sunglasses add a chic fashion-editorial attitude."},{"type":"obj","desc":"A red silk scarf tied neatly around the cat neck, with soft folds and a couture feeling. The scarf must not cover the cat face or the nameplate."},{"type":"obj","desc":"A small gold collar charm or fashion accessory under the scarf, subtle and premium, adding a luxury campaign detail."},{"type":"obj","desc":"In the lower foreground, place a wide horizontal luxury nameplate on the floor in front of the cat. The nameplate is low, flat, landscape-oriented, much wider than tall, like a fashion show seat card or premium display plaque. It is centered, front-facing, level, and fully visible. It must not become vertical, tall, standing, rotated, or side-facing."},{"type":"text","desc":"Print the exact text ideogram4.cpp only on the wide horizontal nameplate. Use clean bold black lettering, perfectly spelled, lowercase, with the number 4 and .cpp extension. The text must fit completely inside the nameplate, stay horizontal, and be readable from left to right."},{"type":"obj","desc":"Add sparse premium editorial accents around the edges: thin gold lines, small code brackets, tiny cursor marks, subtle dots, and minimal geometric details. No extra cats, no stickers, no animal faces, no busy decorations."},{"type":"text","desc":"Bottom center footer reading tiny paws, big compile energy in a small refined monospace or editorial font. The footer is horizontal, centered, understated, and much smaller than the nameplate text."}]}}'  --diffusion-fa -v --offload-to-cpu -H 1024 -W 1024
+```
+
+<img alt="ideogram4 image example" src="../assets/ideogram4/example.png" />

docs/model_config.md

@@ -0,0 +1,118 @@
+# Model Configuration Conventions
+
+This document describes the conventions for model configuration structs and
+weight-based configuration detection.
+
+## Config Types
+
+Model configuration should live in a model-specific `*Config` struct.
+
+Examples:
+
+- `ZImageConfig`
+- `UNetConfig`
+- `MMDiTConfig`
+- `LLMConfig`
+
+Preserve established acronym casing in type names, such as `UNet`, `MMDiT`,
+`LLM`, `VAE`, and `T5`.
+
+Place the config struct near the top of the model header, before the main model
+blocks and runner types that consume it.
+
+## Config Variables
+
+Variables and members that hold a config should be named `config`.
+
+Examples:
+
+```cpp
+UNetConfig config;
+UnetModelBlock unet;
+
+MMDiTRunner(...)
+    : DiffusionModelRunner(backend, params_backend, prefix),
+      config(MMDiTConfig::detect_from_weights(tensor_storage_map, prefix)),
+      mmdit(config) {
+}
+```
+
+Avoid alternate names such as `params`, `params_cfg`, `model_params`, or
+model-specific aliases unless an existing public API requires them.
+
+## Weight Detection
+
+If a model can derive configuration from loaded weight metadata, expose that
+logic as a static method on the config type:
+
+```cpp
+static XxxConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
+                                     const std::string& prefix);
+```
+
+Additional selector arguments are allowed when required by an existing model
+family, for example `SDVersion version` or an architecture enum:
+
+```cpp
+static UNetConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
+                                      const std::string& prefix,
+                                      SDVersion version = VERSION_SD1);
+```
+
+Use `TensorStorage` metadata, especially `n_dims` and `ne`, to infer shapes.
+Do not load or parse tensor data for config detection.
+
+Detection should respect `prefix`. For nested weights, construct full names from
+`prefix + "." + suffix` or filter entries with `starts_with(name, prefix)`.
+
+Do not add persistent config fields such as `inferred_from_weights` only to
+record whether detection happened. If the function needs to decide whether to
+print a debug line, keep that as local control flow inside `detect_from_weights`.
+
+## Logging
+
+When config values are inferred from weights, print one `LOG_DEBUG` line at the
+end of `detect_from_weights`.
+
+Example:
+
+```cpp
+LOG_DEBUG("llm: num_layers = %" PRId64 ", vocab_size = %" PRId64 ", hidden_size = %" PRId64 ", intermediate_size = %" PRId64,
+          config.num_layers,
+          config.vocab_size,
+          config.hidden_size,
+          config.intermediate_size);
+```
+
+Only print the config detection log when the function actually inferred values
+from weights. Do not duplicate the same config summary in runner constructors or
+model loading code.
+
+Use the correct format specifiers for field types, such as `%" PRId64 "` for
+`int64_t` and `%d` for `int`.
+
+## Runner And Model Responsibilities
+
+Runners should detect the config once and pass it into the model block:
+
+```cpp
+struct XxxRunner : public DiffusionModelRunner {
+    XxxConfig config;
+    XxxModel model;
+
+    XxxRunner(..., const String2TensorStorage& tensor_storage_map, const std::string prefix)
+        : DiffusionModelRunner(backend, params_backend, prefix),
+          config(XxxConfig::detect_from_weights(tensor_storage_map, prefix)),
+          model(config) {
+        model.init(params_ctx, tensor_storage_map, prefix);
+    }
+};
+```
+
+Model blocks should consume `config` directly instead of re-scanning weights in
+their constructors. Keep config-derived behavior centralized in the config
+struct.
+
+If a model has no weight-derived config today, it may still provide
+`detect_from_weights` for API consistency, but it should not print a config
+detection log unless it actually derives values from weights.

docs/performance.md

@@ -21,6 +21,38 @@ and the compute buffer shrink in the debug log:
 
 Using `--offload-to-cpu` allows you to offload weights to the CPU, saving VRAM without reducing generation speed.
 
+## Use params backend to reduce VRAM or RAM usage.
+
+`--params-backend` controls where model parameters are kept. If it is not set, parameters use the same backend as `--backend`, so a GPU runtime backend also keeps parameters in VRAM.
+
+Use CPU params to reduce VRAM usage:
+
+```shell
+--backend cuda0 --params-backend cpu
+```
+
+This keeps model weights in system RAM and moves them to the runtime backend when needed. In the example CLI/server, `--offload-to-cpu` is a compatibility shortcut that prepends `*=cpu` to `--params-backend` before creating the context, so explicit module assignments can still override it:
+
+```shell
+--offload-to-cpu --params-backend te=disk
+```
+
+Use disk params to reduce both VRAM and RAM usage:
+
+```shell
+--backend cuda0 --params-backend disk
+```
+
+This reloads parameters from the model file on demand and releases them after use. It has the lowest memory residency, but can be slower because weights must be read again. `disk` is never selected implicitly; set it explicitly when RAM usage matters more than reload cost.
+
+Per-module assignments can target only the largest modules:
+
+```shell
+--backend cuda0 --params-backend diffusion=disk,te=cpu,vae=cpu
+```
+
+See [backend selection](./backend.md) for full syntax.
+
 ## Use quantization to reduce memory usage.
 
-[quantization](./quantization_and_gguf.md)
+[quantization](./quantization_and_gguf.md)

docs/rpc.md

@@ -0,0 +1,220 @@
+# Building and Using the RPC Server with `stable-diffusion.cpp`
+
+This guide covers how to build a version of [the RPC server from `llama.cpp`](https://github.com/ggml-org/llama.cpp/blob/master/tools/rpc/README.md) that is compatible with your version of `stable-diffusion.cpp` to manage multi-backends setups. RPC allows you to offload specific model components to a remote server.
+
+> **Note on Model Location:** The model files (e.g., `.safetensors` or `.gguf`) remain on the **Client** machine. The client parses the file and transmits the necessary tensor data and computational graphs to the server. The server does not need to store the model files locally.
+
+## 1. Building `stable-diffusion.cpp` with RPC client
+
+First, you should build the client application from source. It requires `SD_RPC=ON` to include the RPC backend to your client.
+
+```bash
+mkdir build
+cd build
+cmake .. \
+    -DSD_RPC=ON \
+    # Add other build flags here (e.g., -DSD_VULKAN=ON)
+cmake --build . --config Release -j $(nproc)
+```
+
+> **Note:** Ensure you add the other flags you would normally use (e.g., `-DSD_VULKAN=ON`, `-DSD_CUDA=ON`, `-DSD_HIPBLAS=ON`, or `-DGGML_METAL=ON`), for more information about building `stable-diffusion.cpp` from source, please refer to the [build.md](build.md) documentation.
+
+## 2. Ensure `llama.cpp` is at the correct commit
+
+`stable-diffusion.cpp`'s RPC client is designed to work with a specific version of `llama.cpp` (compatible with the `ggml` submodule) to ensure API compatibility. The commit hash for `llama.cpp` is stored in `ggml/scripts/sync-llama.last`.
+
+> **Start from Root:** Perform these steps from the root of your `stable-diffusion.cpp` directory.
+
+1.  Read the target commit hash from the submodule tracker:
+
+    ```bash
+    # Linux / WSL / MacOS
+    HASH=$(cat ggml/scripts/sync-llama.last)
+
+    # Windows (PowerShell)
+    $HASH = Get-Content -Path "ggml\scripts\sync-llama.last"
+    ```
+
+2.  Clone `llama.cpp` at the target commit .
+    ```bash
+    git clone https://github.com/ggml-org/llama.cpp.git
+    cd llama.cpp
+    git checkout $HASH
+    ```
+    To save on download time and storage, you can use a shallow clone to download only the target commit:
+    ```bash
+    mkdir -p llama.cpp
+    cd llama.cpp
+    git init
+    git remote add origin https://github.com/ggml-org/llama.cpp.git
+    git fetch --depth 1 origin $HASH
+    git checkout FETCH_HEAD
+    ```
+
+## 3. Build `llama.cpp` (RPC Server)
+
+The RPC server acts as the worker. You must explicitly enable the **backend** (the hardware interface, such as CUDA for Nvidia, Metal for Apple Silicon, or Vulkan) when building, otherwise the server will default to using only the CPU.
+
+To find the correct flags for your system, refer to the official documentation for the [`llama.cpp`](https://github.com/ggml-org/llama.cpp/blob/master/docs/build.md) repository.
+
+> **Crucial:** You must include the compiler flags required to satisfy the API compatibility with `stable-diffusion.cpp` (`-DGGML_MAX_NAME=128`). Without this flag, `GGML_MAX_NAME` will default to `64` for the server, and data transfers between the client and server will fail. Of course, `-DGGML_RPC` must also be enabled.
+>
+> I recommend disabling the `LLAMA_CURL` flag to avoid unnecessary dependencies, and disabling shared library builds to avoid potential conflicts.
+
+> **Build Target:** We are specifically building the `rpc-server` target. This prevents the build system from compiling the entire `llama.cpp` suite (like `llama-server`), making the build significantly faster.
+
+### Linux / WSL (Vulkan)
+
+```bash
+mkdir build
+cd build
+cmake .. -DGGML_RPC=ON \
+    -DGGML_VULKAN=ON \        # Ensure backend is enabled
+    -DGGML_BUILD_SHARED_LIBS=OFF \
+    -DLLAMA_CURL=OFF \
+    -DCMAKE_C_FLAGS=-DGGML_MAX_NAME=128 \
+    -DCMAKE_CXX_FLAGS=-DGGML_MAX_NAME=128
+cmake --build . --config Release --target rpc-server -j $(nproc)
+```
+
+### macOS (Metal)
+
+```bash
+mkdir build
+cd build
+cmake .. -DGGML_RPC=ON \
+    -DGGML_METAL=ON \
+    -DGGML_BUILD_SHARED_LIBS=OFF \
+    -DLLAMA_CURL=OFF \
+    -DCMAKE_C_FLAGS=-DGGML_MAX_NAME=128 \
+    -DCMAKE_CXX_FLAGS=-DGGML_MAX_NAME=128
+cmake --build . --config Release --target rpc-server
+```
+
+### Windows (Visual Studio 2022, Vulkan)
+
+```powershell
+mkdir build
+cd build
+cmake .. -G "Visual Studio 17 2022" -A x64 `
+    -DGGML_RPC=ON `
+    -DGGML_VULKAN=ON `
+    -DGGML_BUILD_SHARED_LIBS=OFF `
+    -DLLAMA_CURL=OFF `
+    -DCMAKE_C_FLAGS=-DGGML_MAX_NAME=128 `
+    -DCMAKE_CXX_FLAGS=-DGGML_MAX_NAME=128
+cmake --build . --config Release --target rpc-server
+```
+
+## 4. Usage
+
+Once both applications are built, you can run the server and the client to manage your GPU allocation.
+
+### Step A: Run the RPC Server
+
+Start the server. It listens for connections on the default address (usually `localhost:50052`). If your server is on a different machine, ensure the server binds to the correct interface and your firewall allows the connection.
+
+**On the Server :**
+If running on the same machine, you can use the default address:
+
+```bash
+./rpc-server
+```
+
+If you want to allow connections from other machines on the network:
+
+```bash
+./rpc-server --host 0.0.0.0
+```
+
+> **Security Warning:** The RPC server does not currently support authentication or encryption. **Only run the server on trusted local networks**. Never expose the RPC server directly to the open internet.
+
+> **Drivers & Hardware:** Ensure the Server machine has the necessary drivers installed and functional (e.g., Nvidia Drivers for CUDA, Vulkan SDK, or Metal). If no devices are found, the server will simply fallback to CPU usage.
+
+<!-- ### Step B: Check if the client is able to connect to the server and see the available devices
+
+We're assuming the server is running on your local machine, and listening on the default port `50052`. If it's running on a different machine, you can replace `localhost` with the IP address of the server.
+
+**On the Client:**
+
+```bash
+./sd-cli --rpc-servers localhost:50052 --list-devices
+```
+
+If the server is running and the client is able to connect, you should see `RPC0    localhost:50052` in the list of devices.
+
+Example output:
+(Client built without GPU acceleration, two GPUs available on the server)
+
+```
+List of available GGML devices:
+Name    Description
+-------------------
+CPU     AMD Ryzen 9 5900X 12-Core Processor
+RPC0    localhost:50052
+RPC1    localhost:50052
+``` -->
+
+### Step B: Run with RPC device
+
+If everything is working correctly, you can now run the client while offloading some or all of the work to the RPC server.
+
+Example: Setting the main backend to the RPC0 device for doing all the work on the server.
+
+```bash
+./sd-cli -m models/sd1.5.safetensors -p "A cat" --rpc-servers localhost:50052  --backend RPC0
+```
+
+---
+
+## 5. Scaling: Multiple RPC Servers
+
+You can connect the client to multiple RPC servers simultaneously to scale out your hardware usage.
+
+Example: A main machine (192.168.1.10) with 3 GPUs, with one GPU running CUDA and the other two running Vulkan, and a second machine (192.168.1.11) only one GPU.
+
+**On the first machine (Running two server instances):**
+
+**Terminal 1 (CUDA):**
+
+```bash
+# Linux / WSL
+export CUDA_VISIBLE_DEVICES=0
+cd ./build_cuda/bin/Release
+./rpc-server --host 0.0.0.0
+
+# Windows PowerShell
+$env:CUDA_VISIBLE_DEVICES="0"
+cd .\build_cuda\bin\Release
+./rpc-server --host 0.0.0.0
+```
+
+**Terminal 2 (Vulkan):**
+
+```bash
+cd ./build_vulkan/bin/Release
+# ignore the first GPU (used by CUDA server)
+./rpc-server --host 0.0.0.0 --port 50053 -d Vulkan1,Vulkan2
+```
+
+**On the second machine:**
+
+```bash
+cd ./build/bin/Release
+./rpc-server --host 0.0.0.0
+```
+
+**On the Client:**
+Pass multiple server addresses separated by commas.
+
+```bash
+./sd-cli --rpc-servers 192.168.1.10:50052,192.168.1.10:50053,192.168.1.11:50052 [...]
+```
+
+The client will map these servers to sequential device IDs (e.g., RPC0 from the first server, RPC2, RPC3 from the second, and RPC4 from the third). With this setup, you could for example use RPC0 for the main backend, RPC1 and RPC2 for the text encoders, and RPC3 for the VAE.
+
+---
+
+## 6. Performance Considerations
+
+RPC performance is heavily dependent on network bandwidth, as large weights and activations must be transferred back and forth over the network, especially for large models, or when using high resolutions. For best results, ensure your network connection is stable and has sufficient bandwidth (>1Gbps recommended). This shoumd not be a concern if you are running the server and client on the same machine, as the data transfer will happen over the loopback interface.

examples/cli/README.md

@@ -1,202 +1,9 @@
-# Run
+# Usage
 
-```
-usage: ./bin/sd-cli  [options]
+For detailed command-line arguments, run:
 
-CLI Options:
-  -o, --output <string>         path to write result image to. you can use printf-style %d format specifiers for image
-                                sequences (default: ./output.png) (eg. output_%03d.png). Single-file video outputs
-                                support .avi, .webm, and animated .webp
-  --image <string>              path to the image to inspect (for metadata mode)
-  --metadata-format <string>    metadata output format, one of [text, json] (default: text)
-  --preview-path <string>       path to write preview image to (default: ./preview.png). Multi-frame previews support
-                                .avi, .webm, and animated .webp
-  --preview-interval <int>      interval in denoising steps between consecutive updates of the image preview file
-                                (default is 1, meaning updating at every step)
-  --output-begin-idx <int>      starting index for output image sequence, must be non-negative (default 0 if specified
-                                %d in output path, 1 otherwise)
-  --canny                       apply canny preprocessor (edge detection)
-  --convert-name                convert tensor name (for convert mode)
-  -v, --verbose                 print extra info
-  --color                       colors the logging tags according to level
-  --taesd-preview-only          prevents usage of taesd for decoding the final image. (for use with --preview tae)
-  --preview-noisy               enables previewing noisy inputs of the models rather than the denoised outputs
-  --metadata-raw                include raw hex previews for unparsed metadata payloads
-  --metadata-brief              truncate long metadata text values in text output
-  --metadata-all                include structural/container entries such as IHDR, IDAT, and non-metadata JPEG segments
-  -M, --mode                    run mode, one of [img_gen, vid_gen, upscale, convert, metadata], default: img_gen
-  --preview                     preview method. must be one of the following [none, proj, tae, vae] (default is none)
-  -h, --help                    show this help message and exit
-
-Context Options:
-  -m, --model <string>                     path to full model
-  --clip_l <string>                        path to the clip-l text encoder
-  --clip_g <string>                        path to the clip-g text encoder
-  --clip_vision <string>                   path to the clip-vision encoder
-  --t5xxl <string>                         path to the t5xxl text encoder
-  --llm <string>                           path to the llm text encoder. For example: (qwenvl2.5 for qwen-image,
-                                           mistral-small3.2 for flux2, ...)
-  --llm_vision <string>                    path to the llm vit
-  --qwen2vl <string>                       alias of --llm. Deprecated.
-  --qwen2vl_vision <string>                alias of --llm_vision. Deprecated.
-  --diffusion-model <string>               path to the standalone diffusion model
-  --high-noise-diffusion-model <string>    path to the standalone high noise diffusion model
-  --vae <string>                           path to standalone vae model
-  --taesd <string>                         path to taesd. Using Tiny AutoEncoder for fast decoding (low quality)
-  --tae <string>                           alias of --taesd
-  --control-net <string>                   path to control net model
-  --embd-dir <string>                      embeddings directory
-  --lora-model-dir <string>                lora model directory
-  --hires-upscalers-dir <string>           highres fix upscaler model directory
-  --tensor-type-rules <string>             weight type per tensor pattern (example: "^vae\.=f16,model\.=q8_0")
-  --photo-maker <string>                   path to PHOTOMAKER model
-  --upscale-model <string>                 path to esrgan model.
-  -t, --threads <int>                      number of threads to use during computation (default: -1). If threads <= 0,
-                                           then threads will be set to the number of CPU physical cores
-  --chroma-t5-mask-pad <int>               t5 mask pad size of chroma
-  --max-vram <float>                       maximum VRAM budget in GiB for graph-cut segmented execution. 0 disables
-                                           graph splitting; a negative value auto-detects free VRAM, sparing the
-                                           specified value (e.g. -0.5 will keep at least 0.5 GiB free)
-  --force-sdxl-vae-conv-scale              force use of conv scale on sdxl vae
-  --offload-to-cpu                         place the weights in RAM to save VRAM, and automatically load them into VRAM
-                                           when needed
-  --mmap                                   whether to memory-map model
-  --control-net-cpu                        keep controlnet in cpu (for low vram)
-  --clip-on-cpu                            keep clip in cpu (for low vram)
-  --vae-on-cpu                             keep vae in cpu (for low vram)
-  --fa                                     use flash attention
-  --diffusion-fa                           use flash attention in the diffusion model only
-  --diffusion-conv-direct                  use ggml_conv2d_direct in the diffusion model
-  --vae-conv-direct                        use ggml_conv2d_direct in the vae model
-  --circular                               enable circular padding for convolutions
-  --circularx                              enable circular RoPE wrapping on x-axis (width) only
-  --circulary                              enable circular RoPE wrapping on y-axis (height) only
-  --chroma-disable-dit-mask                disable dit mask for chroma
-  --qwen-image-zero-cond-t                 enable zero_cond_t for qwen image
-  --chroma-enable-t5-mask                  enable t5 mask for chroma
-  --type                                   weight type (examples: f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0, q2_K, q3_K,
-                                           q4_K). If not specified, the default is the type of the weight file
-  --rng                                    RNG, one of [std_default, cuda, cpu], default: cuda(sd-webui), cpu(comfyui)
-  --sampler-rng                            sampler RNG, one of [std_default, cuda, cpu]. If not specified, use --rng
-  --prediction                             prediction type override, one of [eps, v, edm_v, sd3_flow, flux_flow,
-                                           flux2_flow]
-  --lora-apply-mode                        the way to apply LoRA, one of [auto, immediately, at_runtime], default is
-                                           auto. In auto mode, if the model weights contain any quantized parameters,
-                                           the at_runtime mode will be used; otherwise, immediately will be used.The
-                                           immediately mode may have precision and compatibility issues with quantized
-                                           parameters, but it usually offers faster inference speed and, in some cases,
-                                           lower memory usage. The at_runtime mode, on the other hand, is exactly the
-                                           opposite.
-
-Generation Options:
-  -p, --prompt <string>                    the prompt to render
-  -n, --negative-prompt <string>           the negative prompt (default: "")
-  -i, --init-img <string>                  path to the init image
-  --end-img <string>                       path to the end image, required by flf2v
-  --mask <string>                          path to the mask image
-  --control-image <string>                 path to control image, control net
-  --control-video <string>                 path to control video frames, It must be a directory path. The video frames
-                                           inside should be stored as images in lexicographical (character) order. For
-                                           example, if the control video path is `frames`, the directory contain images
-                                           such as 00.png, 01.png, ... etc.
-  --pm-id-images-dir <string>              path to PHOTOMAKER input id images dir
-  --pm-id-embed-path <string>              path to PHOTOMAKER v2 id embed
-  --hires-upscaler <string>                highres fix upscaler, Lanczos, Nearest, Latent, Latent (nearest), Latent
-                                           (nearest-exact), Latent (antialiased), Latent (bicubic), Latent (bicubic
-                                           antialiased), or a model name under --hires-upscalers-dir (default: Latent)
-  --extra-sample-args <string>             extra sampler/scheduler/guidance args, key=value list. APG supports apg_eta,
-                                           apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports
-                                           slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end;
-                                           ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma
-  --extra-tiling-args <string>             extra VAE tiling args, key=value list. LTX video VAE supports
-                                           temporal_tile_frames (default: 4), temporal_tile_overlap (default: 1)
-  -H, --height <int>                       image height, in pixel space (default: 512)
-  -W, --width <int>                        image width, in pixel space (default: 512)
-  --steps <int>                            number of sample steps (default: 20)
-  --high-noise-steps <int>                 (high noise) number of sample steps (default: -1 = auto)
-  --clip-skip <int>                        ignore last layers of CLIP network; 1 ignores none, 2 ignores one layer
-                                           (default: -1). <= 0 represents unspecified, will be 1 for SD1.x, 2 for SD2.x
-  -b, --batch-count <int>                  batch count
-  --video-frames <int>                     video frames (default: 1)
-  --fps <int>                              fps (default: 24)
-  --timestep-shift <int>                   shift timestep for NitroFusion models (default: 0). recommended N for
-                                           NitroSD-Realism around 250 and 500 for NitroSD-Vibrant
-  --upscale-repeats <int>                  Run the ESRGAN upscaler this many times (default: 1)
-  --upscale-tile-size <int>                tile size for ESRGAN upscaling (default: 128)
-  --hires-width <int>                      highres fix target width, 0 to use --hires-scale (default: 0)
-  --hires-height <int>                     highres fix target height, 0 to use --hires-scale (default: 0)
-  --hires-steps <int>                      highres fix second pass sample steps, 0 to reuse --steps (default: 0)
-  --hires-upscale-tile-size <int>          highres fix upscaler tile size, reserved for model-backed upscalers (default:
-                                           128)
-  --cfg-scale <float>                      unconditional guidance scale: (default: 7.0)
-  --img-cfg-scale <float>                  image guidance scale for inpaint or image edit models: (default: same as
-                                           --cfg-scale)
-  --guidance <float>                       distilled guidance scale for models with guidance input (default: 3.5)
-  --slg-scale <float>                      skip layer guidance (SLG) scale, only for DiT models: (default: 0). 0 means
-                                           disabled, a value of 2.5 is nice for sd3.5 medium
-  --skip-layer-start <float>               SLG enabling point (default: 0.01)
-  --skip-layer-end <float>                 SLG disabling point (default: 0.2)
-  --eta <float>                            noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and
-                                           res_2s; 1 for euler_a, er_sde and dpm++2s_a)
-  --flow-shift <float>                     shift value for Flow models like SD3.x or WAN (default: auto)
-  --high-noise-cfg-scale <float>           (high noise) unconditional guidance scale: (default: 7.0)
-  --high-noise-img-cfg-scale <float>       (high noise) image guidance scale for inpaint or image edit models (default:
-                                           same as --cfg-scale)
-  --high-noise-guidance <float>            (high noise) distilled guidance scale for models with guidance input
-                                           (default: 3.5)
-  --high-noise-slg-scale <float>           (high noise) skip layer guidance (SLG) scale, only for DiT models: (default:
-                                           0)
-  --high-noise-skip-layer-start <float>    (high noise) SLG enabling point (default: 0.01)
-  --high-noise-skip-layer-end <float>      (high noise) SLG disabling point (default: 0.2)
-  --high-noise-eta <float>                 (high noise) noise multiplier (default: 0 for ddim_trailing, tcd,
-                                           res_multistep and res_2s; 1 for euler_a, er_sde and dpm++2s_a)
-  --strength <float>                       strength for noising/unnoising (default: 0.75)
-  --pm-style-strength <float>
-  --control-strength <float>               strength to apply Control Net (default: 0.9). 1.0 corresponds to full
-                                           destruction of information in init image
-  --moe-boundary <float>                   timestep boundary for Wan2.2 MoE model. (default: 0.875). Only enabled if
-                                           `--high-noise-steps` is set to -1
-  --vace-strength <float>                  wan vace strength
-  --vae-tile-overlap <float>               tile overlap for vae tiling, in fraction of tile size (default: 0.5)
-  --hires-scale <float>                    highres fix scale when target size is not set (default: 2.0)
-  --hires-denoising-strength <float>       highres fix second pass denoising strength (default: 0.7)
-  --increase-ref-index                     automatically increase the indices of references images based on the order
-                                           they are listed (starting with 1).
-  --disable-auto-resize-ref-image          disable auto resize of ref images
-  --disable-image-metadata                 do not embed generation metadata on image files
-  --vae-tiling                             process vae in tiles to reduce memory usage
-  --temporal-tiling                        enable temporal tiling for LTX video VAE decode
-  --hires                                  enable highres fix
-  -s, --seed                               RNG seed (default: 42, use random seed for < 0)
-  --sampling-method                        sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m,
-                                           dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s,
-                                           er_sde, euler_cfg_pp, euler_a_cfg_pp] (default: euler for Flux/SD3/Wan, euler_a otherwise)
-  --high-noise-sampling-method             (high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a,
-                                           dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep,
-                                           res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp] default: euler for Flux/SD3/Wan, euler_a otherwise
-  --scheduler                              denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits,
-                                           smoothstep, sgm_uniform, simple, kl_optimal, lcm, bong_tangent, ltx2], default:
-                                           model-specific
-  --sigmas                                 custom sigma values for the sampler, comma-separated (e.g.,
-                                           "14.61,7.8,3.5,0.0").
-  --hires-sigmas                           custom sigma values for the highres fix second pass, comma-separated (e.g.,
-                                           "0.85,0.725,0.421875,0.0").
-  --skip-layers                            layers to skip for SLG steps (default: [7,8,9])
-  --high-noise-skip-layers                 (high noise) layers to skip for SLG steps (default: [7,8,9])
-  -r, --ref-image                          reference image for Flux Kontext models (can be used multiple times)
-  --cache-mode                             caching method: 'easycache' (DiT), 'ucache' (UNET),
-                                           'dbcache'/'taylorseer'/'cache-dit' (DiT block-level), 'spectrum' (UNET/DiT
-                                           Chebyshev+Taylor forecasting)
-  --cache-option                           named cache params (key=value format, comma-separated). easycache/ucache:
-                                           threshold=,start=,end=,decay=,relative=,reset=; dbcache/taylorseer/cache-dit:
-                                           Fn=,Bn=,threshold=,warmup=; spectrum: w=,m=,lam=,window=,flex=,warmup=,stop=.
-                                           Examples: "threshold=0.25" or "threshold=1.5,reset=0"
-  --scm-mask                               SCM steps mask for cache-dit: comma-separated 0/1 (e.g.,
-                                           "1,1,1,0,0,1,0,0,1,0") - 1=compute, 0=can cache
-  --scm-policy                             SCM policy: 'dynamic' (default) or 'static'
-  --vae-tile-size                          tile size for vae tiling, format [X]x[Y] (default: 32x32)
-  --vae-relative-tile-size                 relative tile size for vae tiling, format [X]x[Y], in fraction of image size
-                                           if < 1, in number of tiles per dim if >=1 (overrides --vae-tile-size)
+```bash
+./bin/sd-cli -h
 ```
 
 Metadata mode inspects PNG/JPEG container metadata without loading any model:

examples/cli/main.cpp

@@ -169,8 +169,9 @@ struct SDCliParams {
             return 1;
         };
 
-        auto on_help_arg = [&](int argc, const char** argv, int index) {
+        auto on_help_arg = [&](int argc, const char** argv, int index, bool& valid) {
             normal_exit = true;
+            valid       = true;
             return -1;
         };
 
@@ -622,8 +623,6 @@ int main(int argc, const char* argv[]) {
         }
     }
 
-    bool vae_decode_only = true;
-
     auto load_image_and_update_size = [&](const std::string& path,
                                           SDImageOwner& image,
                                           bool resize_image    = true,
@@ -645,21 +644,18 @@ int main(int argc, const char* argv[]) {
     };
 
     if (gen_params.init_image_path.size() > 0) {
-        vae_decode_only = false;
         if (!load_image_and_update_size(gen_params.init_image_path, gen_params.init_image)) {
             return 1;
         }
     }
 
     if (gen_params.end_image_path.size() > 0) {
-        vae_decode_only = false;
         if (!load_image_and_update_size(gen_params.end_image_path, gen_params.end_image)) {
             return 1;
         }
     }
 
     if (gen_params.ref_image_paths.size() > 0) {
-        vae_decode_only = false;
         gen_params.ref_images.clear();
         for (auto& path : gen_params.ref_image_paths) {
             SDImageOwner ref_image({0, 0, 3, nullptr});
@@ -734,18 +730,7 @@ int main(int argc, const char* argv[]) {
         }
     }
 
-    if (cli_params.mode == VID_GEN) {
-        vae_decode_only = false;
-    }
-
-    if (gen_params.hires_enabled &&
-        (gen_params.resolved_hires_upscaler == SD_HIRES_UPSCALER_MODEL ||
-         gen_params.resolved_hires_upscaler == SD_HIRES_UPSCALER_LANCZOS ||
-         gen_params.resolved_hires_upscaler == SD_HIRES_UPSCALER_NEAREST)) {
-        vae_decode_only = false;
-    }
-
-    sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(vae_decode_only, true, cli_params.taesd_preview);
+    sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(cli_params.taesd_preview);
 
     SDImageVec results;
     int num_results             = 0;
@@ -797,12 +782,11 @@ int main(int argc, const char* argv[]) {
     int upscale_factor = 4;  // unused for RealESRGAN_x4plus_anime_6B.pth
     if (ctx_params.esrgan_path.size() > 0 && gen_params.upscale_repeats > 0) {
         UpscalerCtxPtr upscaler_ctx(new_upscaler_ctx(ctx_params.esrgan_path.c_str(),
-                                                     ctx_params.offload_params_to_cpu,
                                                      ctx_params.diffusion_conv_direct,
                                                      ctx_params.n_threads,
                                                      gen_params.upscale_tile_size,
-                                                     ctx_params.backend.c_str(),
-                                                     ctx_params.params_backend.c_str()));
+                                                     sd_ctx_params.backend,
+                                                     sd_ctx_params.params_backend));
 
         if (upscaler_ctx == nullptr) {
             LOG_ERROR("new_upscaler_ctx failed");

examples/common/common.cpp

@@ -51,6 +51,10 @@ static sd_vae_format_t str_to_vae_format(const std::string& value) {
     return SD_VAE_FORMAT_COUNT;
 }
 
+static void prepend_backend_assignment(std::string& spec, const char* assignment) {
+    spec = spec.empty() ? assignment : std::string(assignment) + "," + spec;
+}
+
 #if defined(_WIN32)
 static std::string utf16_to_utf8(const std::wstring& wstr) {
     if (wstr.empty())
@@ -245,6 +249,7 @@ bool parse_options(int argc, const char** argv, const std::vector<ArgOptions>& o
         return false;
     };
 
+    bool valid = false;
     for (int i = 1; i < argc; i++) {
         arg            = argv[i];
         bool found_arg = false;
@@ -287,7 +292,7 @@ bool parse_options(int argc, const char** argv, const std::vector<ArgOptions>& o
                 break;
 
             if (match_and_apply(options.manual_options, [&](auto& option) {
-                    int ret = option.cb(argc, argv, i);
+                    int ret = option.cb(argc, argv, i, valid);
                     if (ret < 0) {
                         invalid_arg = true;
                         return;
@@ -299,7 +304,9 @@ bool parse_options(int argc, const char** argv, const std::vector<ArgOptions>& o
         }
 
         if (invalid_arg) {
-            LOG_ERROR("error: invalid parameter for argument: %s", arg.c_str());
+            if (!valid) {
+                LOG_ERROR("error: invalid parameter for argument: %s", arg.c_str());
+            }
             return false;
         }
         if (!found_arg) {
@@ -356,6 +363,10 @@ ArgOptions SDContextParams::get_options() {
          "--high-noise-diffusion-model",
          "path to the standalone high noise diffusion model",
          &high_noise_diffusion_model_path},
+        {"",
+         "--uncond-diffusion-model",
+         "path to the standalone unconditional diffusion model, currently used by Ideogram4 CFG",
+         &uncond_diffusion_model_path},
         {"",
          "--embeddings-connectors",
          "path to LTXAV embeddings connectors",
@@ -414,8 +425,12 @@ ArgOptions SDContextParams::get_options() {
          &backend},
         {"",
          "--params-backend",
-         "parameter backend assignment, e.g. cpu or diffusion=cpu,clip=cpu",
+         "parameter backend assignment, e.g. disk, cpu, or diffusion=disk,clip=cpu",
          &params_backend},
+        {"",
+         "--rpc-servers",
+         "comma-separated list of RPC servers to connect to for offloading, in the format host:port, e.g. localhost:50052,192.168.1.3:50052",
+         &rpc_servers},
     };
 
     options.int_options = {
@@ -456,15 +471,15 @@ ArgOptions SDContextParams::get_options() {
          true, &enable_mmap},
         {"",
          "--control-net-cpu",
-         "keep controlnet in cpu (for low vram)",
+         "deprecated; use --backend controlnet=cpu",
          true, &control_net_cpu},
         {"",
          "--clip-on-cpu",
-         "keep clip in cpu (for low vram)",
+         "deprecated; use --backend te=cpu",
          true, &clip_on_cpu},
         {"",
          "--vae-on-cpu",
-         "keep vae in cpu (for low vram)",
+         "deprecated; use --backend vae=cpu",
          true, &vae_on_cpu},
         {"",
          "--fa",
@@ -681,6 +696,25 @@ bool SDContextParams::resolve_and_validate(SDMode mode) {
     return true;
 }
 
+void SDContextParams::prepare_backend_assignments() {
+    effective_backend        = backend;
+    effective_params_backend = params_backend;
+
+    if (offload_params_to_cpu) {
+        prepend_backend_assignment(effective_params_backend, "*=cpu");
+    }
+
+    if (clip_on_cpu) {
+        prepend_backend_assignment(effective_backend, "te=cpu");
+    }
+    if (vae_on_cpu) {
+        prepend_backend_assignment(effective_backend, "vae=cpu");
+    }
+    if (control_net_cpu) {
+        prepend_backend_assignment(effective_backend, "controlnet=cpu");
+    }
+}
+
 std::string SDContextParams::to_string() const {
     std::ostringstream emb_ss;
     emb_ss << "{\n";
@@ -706,6 +740,7 @@ std::string SDContextParams::to_string() const {
         << "  llm_vision_path: \"" << llm_vision_path << "\",\n"
         << "  diffusion_model_path: \"" << diffusion_model_path << "\",\n"
         << "  high_noise_diffusion_model_path: \"" << high_noise_diffusion_model_path << "\",\n"
+        << "  uncond_diffusion_model_path: \"" << uncond_diffusion_model_path << "\",\n"
         << "  embeddings_connectors_path: \"" << embeddings_connectors_path << "\",\n"
         << "  vae_path: \"" << vae_path << "\",\n"
         << "  vae_format: \"" << vae_format << "\",\n"
@@ -749,7 +784,8 @@ std::string SDContextParams::to_string() const {
     return oss.str();
 }
 
-sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool vae_decode_only, bool free_params_immediately, bool taesd_preview) {
+sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) {
+    prepare_backend_assignments();
     embedding_vec.clear();
     embedding_vec.reserve(embedding_map.size());
     for (const auto& kv : embedding_map) {
@@ -759,56 +795,52 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool vae_decode_only, bool f
         embedding_vec.emplace_back(item);
     }
 
-    sd_ctx_params_t sd_ctx_params = {
-        model_path.c_str(),
-        clip_l_path.c_str(),
-        clip_g_path.c_str(),
-        clip_vision_path.c_str(),
-        t5xxl_path.c_str(),
-        llm_path.c_str(),
-        llm_vision_path.c_str(),
-        diffusion_model_path.c_str(),
-        high_noise_diffusion_model_path.c_str(),
-        embeddings_connectors_path.c_str(),
-        vae_path.c_str(),
-        audio_vae_path.c_str(),
-        taesd_path.c_str(),
-        control_net_path.c_str(),
-        embedding_vec.data(),
-        static_cast<uint32_t>(embedding_vec.size()),
-        photo_maker_path.c_str(),
-        tensor_type_rules.c_str(),
-        vae_decode_only,
-        free_params_immediately,
-        n_threads,
-        wtype,
-        rng_type,
-        sampler_rng_type,
-        prediction,
-        lora_apply_mode,
-        offload_params_to_cpu,
-        enable_mmap,
-        clip_on_cpu,
-        control_net_cpu,
-        vae_on_cpu,
-        flash_attn,
-        diffusion_flash_attn,
-        taesd_preview,
-        diffusion_conv_direct,
-        vae_conv_direct,
-        circular || circular_x,
-        circular || circular_y,
-        force_sdxl_vae_conv_scale,
-        chroma_use_dit_mask,
-        chroma_use_t5_mask,
-        chroma_t5_mask_pad,
-        qwen_image_zero_cond_t,
-        str_to_vae_format(vae_format),
-        max_vram,
-        stream_layers,
-        backend.c_str(),
-        params_backend.c_str(),
-    };
+    sd_ctx_params_t sd_ctx_params;
+    sd_ctx_params_init(&sd_ctx_params);
+    sd_ctx_params.model_path                      = model_path.c_str();
+    sd_ctx_params.clip_l_path                     = clip_l_path.c_str();
+    sd_ctx_params.clip_g_path                     = clip_g_path.c_str();
+    sd_ctx_params.clip_vision_path                = clip_vision_path.c_str();
+    sd_ctx_params.t5xxl_path                      = t5xxl_path.c_str();
+    sd_ctx_params.llm_path                        = llm_path.c_str();
+    sd_ctx_params.llm_vision_path                 = llm_vision_path.c_str();
+    sd_ctx_params.diffusion_model_path            = diffusion_model_path.c_str();
+    sd_ctx_params.high_noise_diffusion_model_path = high_noise_diffusion_model_path.c_str();
+    sd_ctx_params.uncond_diffusion_model_path     = uncond_diffusion_model_path.c_str();
+    sd_ctx_params.embeddings_connectors_path      = embeddings_connectors_path.c_str();
+    sd_ctx_params.vae_path                        = vae_path.c_str();
+    sd_ctx_params.audio_vae_path                  = audio_vae_path.c_str();
+    sd_ctx_params.taesd_path                      = taesd_path.c_str();
+    sd_ctx_params.control_net_path                = control_net_path.c_str();
+    sd_ctx_params.embeddings                      = embedding_vec.data();
+    sd_ctx_params.embedding_count                 = static_cast<uint32_t>(embedding_vec.size());
+    sd_ctx_params.photo_maker_path                = photo_maker_path.c_str();
+    sd_ctx_params.tensor_type_rules               = tensor_type_rules.c_str();
+    sd_ctx_params.n_threads                       = n_threads;
+    sd_ctx_params.wtype                           = wtype;
+    sd_ctx_params.rng_type                        = rng_type;
+    sd_ctx_params.sampler_rng_type                = sampler_rng_type;
+    sd_ctx_params.prediction                      = prediction;
+    sd_ctx_params.lora_apply_mode                 = lora_apply_mode;
+    sd_ctx_params.enable_mmap                     = enable_mmap;
+    sd_ctx_params.flash_attn                      = flash_attn;
+    sd_ctx_params.diffusion_flash_attn            = diffusion_flash_attn;
+    sd_ctx_params.tae_preview_only                = taesd_preview;
+    sd_ctx_params.diffusion_conv_direct           = diffusion_conv_direct;
+    sd_ctx_params.vae_conv_direct                 = vae_conv_direct;
+    sd_ctx_params.circular_x                      = circular || circular_x;
+    sd_ctx_params.circular_y                      = circular || circular_y;
+    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;
+    sd_ctx_params.stream_layers                   = stream_layers;
+    sd_ctx_params.backend                         = effective_backend.c_str();
+    sd_ctx_params.params_backend                  = effective_params_backend.c_str();
+    sd_ctx_params.rpc_servers                     = rpc_servers.c_str();
     return sd_ctx_params;
 }
 
@@ -2519,6 +2551,7 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params,
     set_json_basename_if_not_empty(models, "llm_vision", ctx_params.llm_vision_path);
     set_json_basename_if_not_empty(models, "diffusion_model", ctx_params.diffusion_model_path);
     set_json_basename_if_not_empty(models, "high_noise_diffusion_model", ctx_params.high_noise_diffusion_model_path);
+    set_json_basename_if_not_empty(models, "uncond_diffusion_model", ctx_params.uncond_diffusion_model_path);
     set_json_basename_if_not_empty(models, "vae", ctx_params.vae_path);
     set_json_basename_if_not_empty(models, "taesd", ctx_params.taesd_path);
     set_json_basename_if_not_empty(models, "control_net", ctx_params.control_net_path);
@@ -2686,6 +2719,9 @@ std::string get_image_params(const SDContextParams& ctx_params,
     if (!ctx_params.diffusion_model_path.empty()) {
         parameter_string += "Unet: " + sd_basename(ctx_params.diffusion_model_path) + ", ";
     }
+    if (!ctx_params.uncond_diffusion_model_path.empty()) {
+        parameter_string += "Uncond Unet: " + sd_basename(ctx_params.uncond_diffusion_model_path) + ", ";
+    }
     if (!ctx_params.vae_path.empty()) {
         parameter_string += "VAE: " + sd_basename(ctx_params.vae_path) + ", ";
     }

examples/common/common.h

@@ -56,11 +56,42 @@ struct BoolOption {
     bool* target;
 };
 
+struct ManualFunction {
+    std::function<int(int, const char**, int, bool&)> _func;
+
+    ManualFunction() = default;
+
+    ManualFunction(std::function<int(int argc, const char** argv, int index, bool& valid)> func)
+        : _func(std::move(func)) {
+    }
+
+    template <typename F>
+    ManualFunction(F func)
+        : _func(make_function(func)) {
+    }
+
+    int operator()(int argc, const char** argv, int index, bool& valid) const {
+        return _func(argc, argv, index, valid);
+    }
+
+private:
+    template <typename F>
+    static std::function<int(int, const char**, int, bool&)> make_function(F func) {
+        if constexpr (std::is_invocable_v<F, int, const char**, int, bool&>) {
+            return func;
+        } else {
+            return [func](int argc, const char** argv, int index, bool&) {
+                return func(argc, argv, index);
+            };
+        }
+    }
+};
+
 struct ManualOption {
     std::string short_name;
     std::string long_name;
     std::string desc;
-    std::function<int(int argc, const char** argv, int index)> cb;
+    ManualFunction cb;
 };
 
 struct ArgOptions {
@@ -92,6 +123,7 @@ struct SDContextParams {
     std::string llm_vision_path;
     std::string diffusion_model_path;
     std::string high_noise_diffusion_model_path;
+    std::string uncond_diffusion_model_path;
     std::string embeddings_connectors_path;
     std::string vae_path;
     std::string vae_format = "auto";
@@ -116,6 +148,9 @@ struct SDContextParams {
     bool stream_layers          = false;
     std::string backend;
     std::string params_backend;
+    std::string rpc_servers;
+    std::string effective_backend;
+    std::string effective_params_backend;
     bool enable_mmap           = false;
     bool control_net_cpu       = false;
     bool clip_on_cpu           = false;
@@ -143,11 +178,12 @@ struct SDContextParams {
     float flow_shift = INFINITY;
     ArgOptions get_options();
     void build_embedding_map();
+    void prepare_backend_assignments();
     bool resolve(SDMode mode);
     bool validate(SDMode mode);
     bool resolve_and_validate(SDMode mode);
     std::string to_string() const;
-    sd_ctx_params_t to_sd_ctx_params_t(bool vae_decode_only, bool free_params_immediately, bool taesd_preview);
+    sd_ctx_params_t to_sd_ctx_params_t(bool taesd_preview);
 };
 
 struct SDGenerationParams {

examples/server/README.md

@@ -117,186 +117,10 @@ In this case, the server will load and serve the specified `index.html` file ins
 * using a custom UI
 * avoiding rebuilding the binary after frontend modifications
 
-# Run
+# Usage
 
-```
-usage: ./bin/sd-server  [options]
-
-Svr Options:
-  -l, --listen-ip <string>      server listen ip (default: 127.0.0.1)
-  --serve-html-path <string>    path to HTML file to serve at root (optional)
-  --listen-port <int>           server listen port (default: 1234)
-  -v, --verbose                 print extra info
-  --color                       colors the logging tags according to level
-  -h, --help                    show this help message and exit
-
-Context Options:
-  -m, --model <string>                     path to full model
-  --clip_l <string>                        path to the clip-l text encoder
-  --clip_g <string>                        path to the clip-g text encoder
-  --clip_vision <string>                   path to the clip-vision encoder
-  --t5xxl <string>                         path to the t5xxl text encoder
-  --llm <string>                           path to the llm text encoder. For example: (qwenvl2.5 for qwen-image,
-                                           mistral-small3.2 for flux2, ...)
-  --llm_vision <string>                    path to the llm vit
-  --qwen2vl <string>                       alias of --llm. Deprecated.
-  --qwen2vl_vision <string>                alias of --llm_vision. Deprecated.
-  --diffusion-model <string>               path to the standalone diffusion model
-  --high-noise-diffusion-model <string>    path to the standalone high noise diffusion model
-  --vae <string>                           path to standalone vae model
-  --taesd <string>                         path to taesd. Using Tiny AutoEncoder for fast decoding (low quality)
-  --tae <string>                           alias of --taesd
-  --control-net <string>                   path to control net model
-  --embd-dir <string>                      embeddings directory
-  --lora-model-dir <string>                lora model directory
-  --hires-upscalers-dir <string>           highres fix upscaler model directory
-  --tensor-type-rules <string>             weight type per tensor pattern (example: "^vae\.=f16,model\.=q8_0")
-  --photo-maker <string>                   path to PHOTOMAKER model
-  --upscale-model <string>                 path to esrgan model.
-  -t, --threads <int>                      number of threads to use during computation (default: -1). If threads <= 0,
-                                           then threads will be set to the number of CPU physical cores
-  --chroma-t5-mask-pad <int>               t5 mask pad size of chroma
-  --max-vram <float>                       maximum VRAM budget in GiB for graph-cut segmented execution. 0 disables
-                                           graph splitting; a negative value auto-detects free VRAM, sparing the
-                                           specified value (e.g. -0.5 will keep at least 0.5 GiB free)
-  --force-sdxl-vae-conv-scale              force use of conv scale on sdxl vae
-  --offload-to-cpu                         place the weights in RAM to save VRAM, and automatically load them into VRAM
-                                           when needed
-  --mmap                                   whether to memory-map model
-  --control-net-cpu                        keep controlnet in cpu (for low vram)
-  --clip-on-cpu                            keep clip in cpu (for low vram)
-  --vae-on-cpu                             keep vae in cpu (for low vram)
-  --fa                                     use flash attention
-  --diffusion-fa                           use flash attention in the diffusion model only
-  --diffusion-conv-direct                  use ggml_conv2d_direct in the diffusion model
-  --vae-conv-direct                        use ggml_conv2d_direct in the vae model
-  --circular                               enable circular padding for convolutions
-  --circularx                              enable circular RoPE wrapping on x-axis (width) only
-  --circulary                              enable circular RoPE wrapping on y-axis (height) only
-  --chroma-disable-dit-mask                disable dit mask for chroma
-  --qwen-image-zero-cond-t                 enable zero_cond_t for qwen image
-  --chroma-enable-t5-mask                  enable t5 mask for chroma
-  --type                                   weight type (examples: f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0, q2_K, q3_K,
-                                           q4_K). If not specified, the default is the type of the weight file
-  --rng                                    RNG, one of [std_default, cuda, cpu], default: cuda(sd-webui), cpu(comfyui)
-  --sampler-rng                            sampler RNG, one of [std_default, cuda, cpu]. If not specified, use --rng
-  --prediction                             prediction type override, one of [eps, v, edm_v, sd3_flow, flux_flow,
-                                           flux2_flow]
-  --lora-apply-mode                        the way to apply LoRA, one of [auto, immediately, at_runtime], default is
-                                           auto. In auto mode, if the model weights contain any quantized parameters,
-                                           the at_runtime mode will be used; otherwise, immediately will be used.The
-                                           immediately mode may have precision and compatibility issues with quantized
-                                           parameters, but it usually offers faster inference speed and, in some cases,
-                                           lower memory usage. The at_runtime mode, on the other hand, is exactly the
-                                           opposite.
-
-Default Generation Options:
-  -p, --prompt <string>                    the prompt to render
-  -n, --negative-prompt <string>           the negative prompt (default: "")
-  -i, --init-img <string>                  path to the init image
-  --end-img <string>                       path to the end image, required by flf2v
-  --mask <string>                          path to the mask image
-  --control-image <string>                 path to control image, control net
-  --control-video <string>                 path to control video frames, It must be a directory path. The video frames
-                                           inside should be stored as images in lexicographical (character) order. For
-                                           example, if the control video path is `frames`, the directory contain images
-                                           such as 00.png, 01.png, ... etc.
-  --pm-id-images-dir <string>              path to PHOTOMAKER input id images dir
-  --pm-id-embed-path <string>              path to PHOTOMAKER v2 id embed
-  --hires-upscaler <string>                highres fix upscaler, Lanczos, Nearest, Latent, Latent (nearest), Latent
-                                           (nearest-exact), Latent (antialiased), Latent (bicubic), Latent (bicubic
-                                           antialiased), or a model name under --hires-upscalers-dir (default: Latent)
-  --extra-sample-args <string>             extra sampler/scheduler/guidance args, key=value list. APG supports apg_eta,
-                                           apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports
-                                           slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end;
-                                           ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma
-  --extra-tiling-args <string>             extra VAE tiling args, key=value list. LTX video VAE supports
-                                           temporal_tile_frames (default: 4), temporal_tile_overlap (default: 1)
-  -H, --height <int>                       image height, in pixel space (default: 512)
-  -W, --width <int>                        image width, in pixel space (default: 512)
-  --steps <int>                            number of sample steps (default: 20)
-  --high-noise-steps <int>                 (high noise) number of sample steps (default: -1 = auto)
-  --clip-skip <int>                        ignore last layers of CLIP network; 1 ignores none, 2 ignores one layer
-                                           (default: -1). <= 0 represents unspecified, will be 1 for SD1.x, 2 for SD2.x
-  -b, --batch-count <int>                  batch count
-  --video-frames <int>                     video frames (default: 1)
-  --fps <int>                              fps (default: 24)
-  --timestep-shift <int>                   shift timestep for NitroFusion models (default: 0). recommended N for
-                                           NitroSD-Realism around 250 and 500 for NitroSD-Vibrant
-  --upscale-repeats <int>                  Run the ESRGAN upscaler this many times (default: 1)
-  --upscale-tile-size <int>                tile size for ESRGAN upscaling (default: 128)
-  --hires-width <int>                      highres fix target width, 0 to use --hires-scale (default: 0)
-  --hires-height <int>                     highres fix target height, 0 to use --hires-scale (default: 0)
-  --hires-steps <int>                      highres fix second pass sample steps, 0 to reuse --steps (default: 0)
-  --hires-upscale-tile-size <int>          highres fix upscaler tile size, reserved for model-backed upscalers (default:
-                                           128)
-  --cfg-scale <float>                      unconditional guidance scale: (default: 7.0)
-  --img-cfg-scale <float>                  image guidance scale for inpaint or image edit models: (default: same as
-                                           --cfg-scale)
-  --guidance <float>                       distilled guidance scale for models with guidance input (default: 3.5)
-  --slg-scale <float>                      skip layer guidance (SLG) scale, only for DiT models: (default: 0). 0 means
-                                           disabled, a value of 2.5 is nice for sd3.5 medium
-  --skip-layer-start <float>               SLG enabling point (default: 0.01)
-  --skip-layer-end <float>                 SLG disabling point (default: 0.2)
-  --eta <float>                            noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and
-                                           res_2s; 1 for euler_a, er_sde and dpm++2s_a)
-  --flow-shift <float>                     shift value for Flow models like SD3.x or WAN (default: auto)
-  --high-noise-cfg-scale <float>           (high noise) unconditional guidance scale: (default: 7.0)
-  --high-noise-img-cfg-scale <float>       (high noise) image guidance scale for inpaint or image edit models (default:
-                                           same as --cfg-scale)
-  --high-noise-guidance <float>            (high noise) distilled guidance scale for models with guidance input
-                                           (default: 3.5)
-  --high-noise-slg-scale <float>           (high noise) skip layer guidance (SLG) scale, only for DiT models: (default:
-                                           0)
-  --high-noise-skip-layer-start <float>    (high noise) SLG enabling point (default: 0.01)
-  --high-noise-skip-layer-end <float>      (high noise) SLG disabling point (default: 0.2)
-  --high-noise-eta <float>                 (high noise) noise multiplier (default: 0 for ddim_trailing, tcd,
-                                           res_multistep and res_2s; 1 for euler_a, er_sde and dpm++2s_a)
-  --strength <float>                       strength for noising/unnoising (default: 0.75)
-  --pm-style-strength <float>
-  --control-strength <float>               strength to apply Control Net (default: 0.9). 1.0 corresponds to full
-                                           destruction of information in init image
-  --moe-boundary <float>                   timestep boundary for Wan2.2 MoE model. (default: 0.875). Only enabled if
-                                           `--high-noise-steps` is set to -1
-  --vace-strength <float>                  wan vace strength
-  --vae-tile-overlap <float>               tile overlap for vae tiling, in fraction of tile size (default: 0.5)
-  --hires-scale <float>                    highres fix scale when target size is not set (default: 2.0)
-  --hires-denoising-strength <float>       highres fix second pass denoising strength (default: 0.7)
-  --increase-ref-index                     automatically increase the indices of references images based on the order
-                                           they are listed (starting with 1).
-  --disable-auto-resize-ref-image          disable auto resize of ref images
-  --disable-image-metadata                 do not embed generation metadata on image files
-  --vae-tiling                             process vae in tiles to reduce memory usage
-  --temporal-tiling                        enable temporal tiling for LTX video VAE decode
-  --hires                                  enable highres fix
-  -s, --seed                               RNG seed (default: 42, use random seed for < 0)
-  --sampling-method                        sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m,
-                                           dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s,
-                                           er_sde, euler_cfg_pp, euler_a_cfg_pp] (default: euler for Flux/SD3/Wan, euler_a otherwise)
-  --high-noise-sampling-method             (high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a,
-                                           dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep,
-                                           res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp] default: euler for Flux/SD3/Wan, euler_a otherwise
-  --scheduler                              denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits,
-                                           smoothstep, sgm_uniform, simple, kl_optimal, lcm, bong_tangent, ltx2], default:
-                                           model-specific
-  --sigmas                                 custom sigma values for the sampler, comma-separated (e.g.,
-                                           "14.61,7.8,3.5,0.0").
-  --hires-sigmas                           custom sigma values for the highres fix second pass, comma-separated (e.g.,
-                                           "0.85,0.725,0.421875,0.0").
-  --skip-layers                            layers to skip for SLG steps (default: [7,8,9])
-  --high-noise-skip-layers                 (high noise) layers to skip for SLG steps (default: [7,8,9])
-  -r, --ref-image                          reference image for Flux Kontext models (can be used multiple times)
-  --cache-mode                             caching method: 'easycache' (DiT), 'ucache' (UNET),
-                                           'dbcache'/'taylorseer'/'cache-dit' (DiT block-level), 'spectrum' (UNET/DiT
-                                           Chebyshev+Taylor forecasting)
-  --cache-option                           named cache params (key=value format, comma-separated). easycache/ucache:
-                                           threshold=,start=,end=,decay=,relative=,reset=; dbcache/taylorseer/cache-dit:
-                                           Fn=,Bn=,threshold=,warmup=; spectrum: w=,m=,lam=,window=,flex=,warmup=,stop=.
-                                           Examples: "threshold=0.25" or "threshold=1.5,reset=0"
-  --scm-mask                               SCM steps mask for cache-dit: comma-separated 0/1 (e.g.,
-                                           "1,1,1,0,0,1,0,0,1,0") - 1=compute, 0=can cache
-  --scm-policy                             SCM policy: 'dynamic' (default) or 'static'
-  --vae-tile-size                          tile size for vae tiling, format [X]x[Y] (default: 32x32)
-  --vae-relative-tile-size                 relative tile size for vae tiling, format [X]x[Y], in fraction of image size
-                                           if < 1, in number of tiles per dim if >=1 (overrides --vae-tile-size)
+For detailed command-line arguments, run:
+
+```bash
+./bin/sd-server -h
 ```

examples/server/main.cpp

@@ -85,7 +85,7 @@ int main(int argc, const char** argv) {
     LOG_DEBUG("%s", ctx_params.to_string().c_str());
     LOG_DEBUG("%s", default_gen_params.to_string().c_str());
 
-    sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(false, false, false);
+    sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(false);
     SDCtxPtr sd_ctx(new_sd_ctx(&sd_ctx_params));
 
     if (sd_ctx == nullptr) {

examples/server/runtime.cpp

@@ -203,8 +203,9 @@ ArgOptions SDSvrParams::get_options() {
         {"", "--color", "colors the logging tags according to level", true, &color},
     };
 
-    auto on_help_arg = [&](int, const char**, int) {
+    auto on_help_arg = [&](int, const char**, int, bool& valid) {
         normal_exit = true;
+        valid       = true;
         return -1;
     };
 

format-code.ps1

@@ -0,0 +1,54 @@
+$patterns = @(
+    "src/*.cpp"
+    "src/*.h"
+    "src/*.hpp"
+    "src/conditioning/*.cpp"
+    "src/conditioning/*.h"
+    "src/conditioning/*.hpp"
+    "src/core/*.cpp"
+    "src/core/*.h"
+    "src/core/*.hpp"
+    "src/extensions/*.cpp"
+    "src/extensions/*.h"
+    "src/extensions/*.hpp"
+    "src/runtime/*.cpp"
+    "src/runtime/*.h"
+    "src/runtime/*.hpp"
+    "src/model/*/*.cpp"
+    "src/model/*/*.h"
+    "src/model/*/*.hpp"
+    "src/tokenizers/*.h"
+    "src/tokenizers/*.cpp"
+    "src/tokenizers/vocab/*.h"
+    "src/tokenizers/vocab/*.cpp"
+    "src/model_io/*.h"
+    "src/model_io/*.cpp"
+    "examples/cli/*.cpp"
+    "examples/cli/*.h"
+    "examples/server/*.cpp"
+    "examples/common/*.hpp"
+    "examples/common/*.h"
+    "examples/common/*.cpp"
+)
+
+$root = (Get-Location).Path
+
+foreach ($pattern in $patterns) {
+    $files = Get-ChildItem -Path $pattern -File -ErrorAction SilentlyContinue | Sort-Object FullName
+
+    foreach ($file in $files) {
+        $relativePath = $file.FullName.Substring($root.Length).TrimStart('\', '/') -replace '\\', '/'
+
+        if ($relativePath -like "vocab*") {
+            continue
+        }
+
+        Write-Host "formatting '$relativePath'"
+
+        # if ($relativePath -ne "stable-diffusion.h") {
+        #     clang-tidy -fix -p build_linux/ "$relativePath"
+        # }
+
+        & clang-format -style=file -i $relativePath
+    }
+}

format-code.sh

@@ -1,10 +1,17 @@
-for f in src/*.cpp src/*.h src/*.hpp src/tokenizers/*.h src/tokenizers/*.cpp src/tokenizers/vocab/*.h src/tokenizers/vocab/*.cpp \
+for f in src/*.cpp src/*.h src/*.hpp \
+         src/conditioning/*.cpp src/conditioning/*.h src/conditioning/*.hpp \
+         src/core/*.cpp src/core/*.h src/core/*.hpp \
+         src/extensions/*.cpp src/extensions/*.h src/extensions/*.hpp \
+         src/runtime/*.cpp src/runtime/*.h src/runtime/*.hpp \
+         src/model/*/*.cpp src/model/*/*.h src/model/*/*.hpp \
+         src/tokenizers/*.h src/tokenizers/*.cpp src/tokenizers/vocab/*.h src/tokenizers/vocab/*.cpp \
          src/model_io/*.h src/model_io/*.cpp examples/cli/*.cpp examples/cli/*.h examples/server/*.cpp \
          examples/common/*.hpp examples/common/*.h examples/common/*.cpp; do
+  [[ -e "$f" ]] || continue
   [[ "$f" == vocab* ]] && continue
   echo "formatting '$f'"
   # if [ "$f" != "stable-diffusion.h" ]; then
   #   clang-tidy -fix -p build_linux/ "$f"
   # fi
   clang-format -style=file -i "$f"
-done
+done

include/stable-diffusion.h

@@ -186,6 +186,7 @@ typedef struct {
     const char* llm_vision_path;
     const char* diffusion_model_path;
     const char* high_noise_diffusion_model_path;
+    const char* uncond_diffusion_model_path;
     const char* embeddings_connectors_path;
     const char* vae_path;
     const char* audio_vae_path;
@@ -195,19 +196,13 @@ typedef struct {
     uint32_t embedding_count;
     const char* photo_maker_path;
     const char* tensor_type_rules;
-    bool vae_decode_only;
-    bool free_params_immediately;
     int n_threads;
     enum sd_type_t wtype;
     enum rng_type_t rng_type;
     enum rng_type_t sampler_rng_type;
     enum prediction_t prediction;
     enum lora_apply_mode_t lora_apply_mode;
-    bool offload_params_to_cpu;
     bool enable_mmap;
-    bool keep_clip_on_cpu;
-    bool keep_control_net_on_cpu;
-    bool keep_vae_on_cpu;
     bool flash_attn;
     bool diffusion_flash_attn;
     bool tae_preview_only;
@@ -225,6 +220,7 @@ typedef struct {
     bool stream_layers;  // Enable residency+prefetch streaming on top of --max-vram (no effect without --max-vram)
     const char* backend;
     const char* params_backend;
+    const char* rpc_servers;
 } sd_ctx_params_t;
 
 typedef struct {
@@ -459,7 +455,6 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
 typedef struct upscaler_ctx_t upscaler_ctx_t;
 
 SD_API upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path,
-                                        bool offload_params_to_cpu,
                                         bool direct,
                                         int n_threads,
                                         int tile_size,
@@ -490,6 +485,10 @@ SD_API bool preprocess_canny(sd_image_t image,
 SD_API const char* sd_commit(void);
 SD_API const char* sd_version(void);
 
+// for C API, caller needs to call free_sd_images to free the memory after use
+// This helps avoid CRT problems on Windows when memory is allocated in the library but freed in the caller, which may use a different CRT.
+SD_API void free_sd_images(sd_image_t* result_images, int num_images);
+
 #ifdef __cplusplus
 }
 #endif

script/convert_fp8_scale_to_bf16.py

@@ -0,0 +1,283 @@
+#!/usr/bin/env python
+import argparse
+import json
+import math
+import os
+import struct
+from collections import Counter
+from pathlib import Path
+
+import torch
+from safetensors import safe_open
+
+
+FLOAT_DTYPES = {
+    "BF16",
+    "F16",
+    "F32",
+    "F64",
+    "F8_E4M3",
+    "F8_E4M3FN",
+    "F8_E5M2",
+}
+
+FP8_DTYPES = {
+    "F8_E4M3",
+    "F8_E4M3FN",
+    "F8_E5M2",
+}
+
+DTYPE_SIZES = {
+    "BOOL": 1,
+    "U8": 1,
+    "I8": 1,
+    "F8_E4M3": 1,
+    "F8_E4M3FN": 1,
+    "F8_E5M2": 1,
+    "U16": 2,
+    "I16": 2,
+    "F16": 2,
+    "BF16": 2,
+    "U32": 4,
+    "I32": 4,
+    "F32": 4,
+    "U64": 8,
+    "I64": 8,
+    "F64": 8,
+}
+
+
+def read_safetensors_header(path: Path):
+    with path.open("rb") as f:
+        header_len = struct.unpack("<Q", f.read(8))[0]
+        header = f.read(header_len).decode("utf-8").rstrip()
+    return json.loads(header)
+
+
+def numel(shape):
+    return math.prod(shape) if shape else 1
+
+
+def scale_key_for_weight(name: str):
+    if name.endswith(".weight"):
+        return name[:-len(".weight")] + ".weight_scale"
+    if name.endswith("weight"):
+        return name + "_scale"
+    return None
+
+
+def tensor_nbytes(dtype: str, shape):
+    return numel(shape) * DTYPE_SIZES[dtype]
+
+
+def build_output_plan(header):
+    entries = {k: v for k, v in header.items() if k != "__metadata__"}
+    paired_scale_keys = set()
+    plan = []
+
+    for name, info in entries.items():
+        scale_key = scale_key_for_weight(name)
+        if info["dtype"] in FP8_DTYPES and scale_key in entries:
+            paired_scale_keys.add(scale_key)
+
+    for name, info in entries.items():
+        if name in paired_scale_keys:
+            continue
+
+        dtype = info["dtype"]
+        shape = info["shape"]
+        scale_key = scale_key_for_weight(name)
+
+        if dtype in FP8_DTYPES and scale_key in entries:
+            scale_info = entries[scale_key]
+            plan.append(
+                {
+                    "name": name,
+                    "source_dtype": dtype,
+                    "output_dtype": "BF16",
+                    "shape": shape,
+                    "mode": "fp8_scaled_weight",
+                    "scale_key": scale_key,
+                }
+            )
+            continue
+
+        if dtype in FLOAT_DTYPES:
+            plan.append(
+                {
+                    "name": name,
+                    "source_dtype": dtype,
+                    "output_dtype": "BF16",
+                    "shape": shape,
+                    "mode": "float_to_bf16",
+                }
+            )
+        else:
+            plan.append(
+                {
+                    "name": name,
+                    "source_dtype": dtype,
+                    "output_dtype": dtype,
+                    "shape": shape,
+                    "mode": "copy",
+                }
+            )
+
+    metadata = dict(header.get("__metadata__", {}) or {})
+    metadata["format"] = "pt"
+    metadata["conversion"] = "fp8_weight_scale_to_bf16"
+
+    output_header = {"__metadata__": metadata}
+    offset = 0
+    for item in plan:
+        size = tensor_nbytes(item["output_dtype"], item["shape"])
+        output_header[item["name"]] = {
+            "dtype": item["output_dtype"],
+            "shape": item["shape"],
+            "data_offsets": [offset, offset + size],
+        }
+        offset += size
+
+    return plan, output_header, offset
+
+
+def write_tensor_bytes(out, tensor):
+    tensor = tensor.detach().cpu().contiguous()
+    if tensor.numel() == 0:
+        return
+    if tensor.dtype == torch.bfloat16:
+        tensor.view(torch.uint16).numpy().tofile(out)
+    elif tensor.dtype in (getattr(torch, "float8_e4m3fn", None), getattr(torch, "float8_e5m2", None)):
+        tensor.view(torch.uint8).numpy().tofile(out)
+    else:
+        tensor.numpy().tofile(out)
+
+
+def scale_view_for_chunk(scale, chunk, first_dim_start=0, first_dim_end=None):
+    scale = scale.to(torch.float32)
+
+    if scale.numel() == 1:
+        return scale.reshape((1,) * chunk.ndim)
+
+    if chunk.ndim > 0 and scale.ndim == 1:
+        if first_dim_end is not None and scale.shape[0] >= first_dim_end:
+            scale = scale[first_dim_start:first_dim_end]
+        if scale.shape[0] == chunk.shape[0]:
+            return scale.reshape((scale.shape[0],) + (1,) * (chunk.ndim - 1))
+
+    return scale
+
+
+def write_scaled_fp8_weight(out, weight, scale, chunk_rows):
+    if weight.ndim == 0:
+        result = weight.to(torch.float32) * scale_view_for_chunk(scale, weight)
+        write_tensor_bytes(out, result.to(torch.bfloat16))
+        return
+
+    rows = weight.shape[0]
+    for start in range(0, rows, chunk_rows):
+        end = min(start + chunk_rows, rows)
+        chunk = weight[start:end].to(torch.float32)
+        scale_view = scale_view_for_chunk(scale, chunk, start, end)
+        result = chunk * scale_view
+        write_tensor_bytes(out, result.to(torch.bfloat16))
+
+
+def write_float_as_bf16(out, tensor, chunk_rows):
+    if tensor.dtype == torch.bfloat16:
+        write_tensor_bytes(out, tensor)
+        return
+
+    if tensor.ndim == 0:
+        write_tensor_bytes(out, tensor.to(torch.bfloat16))
+        return
+
+    rows = tensor.shape[0]
+    for start in range(0, rows, chunk_rows):
+        end = min(start + chunk_rows, rows)
+        write_tensor_bytes(out, tensor[start:end].to(torch.bfloat16))
+
+
+def convert(input_path: Path, output_path: Path, chunk_rows: int, dry_run: bool):
+    header = read_safetensors_header(input_path)
+    plan, output_header, data_size = build_output_plan(header)
+
+    source_counts = Counter(item["source_dtype"] for item in plan)
+    output_counts = Counter(item["output_dtype"] for item in plan)
+    scaled_count = sum(item["mode"] == "fp8_scaled_weight" for item in plan)
+    dropped_scales = sum(item["mode"] == "fp8_scaled_weight" for item in plan)
+    header_bytes = json.dumps(output_header, separators=(",", ":")).encode("utf-8")
+    expected_size = 8 + len(header_bytes) + data_size
+
+    print(f"input:  {input_path}")
+    print(f"output: {output_path}")
+    print(f"tensors written: {len(plan)}")
+    print(f"scaled fp8 weights dequantized: {scaled_count}")
+    print(f"weight_scale tensors dropped: {dropped_scales}")
+    print(f"source dtypes: {dict(sorted(source_counts.items()))}")
+    print(f"output dtypes: {dict(sorted(output_counts.items()))}")
+    print(f"expected output size: {expected_size / (1024 ** 3):.2f} GiB")
+
+    if dry_run:
+        return
+
+    if output_path.exists():
+        raise FileExistsError(f"{output_path} already exists; pass --overwrite to replace it")
+
+    tmp_path = output_path.with_suffix(output_path.suffix + ".tmp")
+    if tmp_path.exists():
+        raise FileExistsError(f"{tmp_path} already exists; remove it or choose another output")
+
+    with safe_open(str(input_path), framework="pt", device="cpu") as sf, tmp_path.open("wb") as out:
+        out.write(struct.pack("<Q", len(header_bytes)))
+        out.write(header_bytes)
+
+        for index, item in enumerate(plan, 1):
+            name = item["name"]
+            print(f"[{index:04d}/{len(plan):04d}] {name} -> {item['output_dtype']}")
+
+            tensor = sf.get_tensor(name)
+            if item["mode"] == "fp8_scaled_weight":
+                scale = sf.get_tensor(item["scale_key"])
+                write_scaled_fp8_weight(out, tensor, scale, chunk_rows)
+            elif item["mode"] == "float_to_bf16":
+                write_float_as_bf16(out, tensor, chunk_rows)
+            else:
+                write_tensor_bytes(out, tensor)
+
+        actual_size = out.tell()
+
+    if actual_size != expected_size:
+        tmp_path.unlink(missing_ok=True)
+        raise RuntimeError(f"wrote {actual_size} bytes, expected {expected_size} bytes")
+
+    tmp_path.replace(output_path)
+    print("done")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Convert an fp8 safetensors checkpoint with weight_scale tensors to bf16."
+    )
+    parser.add_argument("--input", default="ideogram4_fp8.safetensors", type=Path)
+    parser.add_argument("--output", default="ideogram4_bf16.safetensors", type=Path)
+    parser.add_argument("--chunk-rows", default=1024, type=int)
+    parser.add_argument("--dry-run", action="store_true")
+    parser.add_argument("--overwrite", action="store_true")
+    args = parser.parse_args()
+
+    input_path = args.input.resolve()
+    output_path = args.output.resolve()
+
+    if args.chunk_rows < 1:
+        raise ValueError("--chunk-rows must be >= 1")
+    if not input_path.exists():
+        raise FileNotFoundError(input_path)
+    if args.overwrite and output_path.exists():
+        output_path.unlink()
+
+    convert(input_path, output_path, args.chunk_rows, args.dry_run)
+
+
+if __name__ == "__main__":
+    main()

src/convert.cpp

@@ -3,9 +3,9 @@
 #include <regex>
 #include <vector>
 
-#include "model.h"
 #include "model_io/gguf_io.h"
 #include "model_io/safetensors_io.h"
+#include "model_loader.h"
 #include "util.h"
 
 #include "ggml_extend_backend.h"

src/core/ggml_extend_backend.cpp

@@ -1,4 +1,4 @@
-#include "ggml_extend_backend.h"
+#include "core/ggml_extend_backend.h"
 
 #include <algorithm>
 #include <cctype>
@@ -8,8 +8,8 @@
 #include <stdexcept>
 #include <vector>
 
+#include "core/util.h"
 #include "stable-diffusion.h"
-#include "util.h"
 
 static std::string trim_copy(const std::string& value) {
     size_t begin = 0;
@@ -45,6 +45,10 @@ static bool is_default_backend_token(const std::string& name) {
     return lower.empty() || lower == "default" || lower == "auto";
 }
 
+static bool is_disk_backend_token(const std::string& name) {
+    return lower_copy(trim_copy(name)) == "disk";
+}
+
 static bool parse_backend_module(const std::string& raw_name, SDBackendModule* module) {
     std::string name = lower_copy(trim_copy(raw_name));
     name.erase(std::remove(name.begin(), name.end(), '-'), name.end());
@@ -200,6 +204,36 @@ void ggml_ext_im_set_f32_1d(const struct ggml_tensor* tensor, int i, float value
     }
 }
 
+bool add_rpc_devices(const std::string& servers) {
+    const std::string in = trim_copy(servers);
+    if (in.empty()) {
+        return true;
+    }
+    auto rpc_servers = split_copy(in, ',');
+    if (rpc_servers.empty()) {
+        LOG_ERROR("invalid RPC servers specification: '%s'", servers.c_str());
+        return false;
+    }
+    ggml_backend_reg_t rpc_reg = ggml_backend_reg_by_name("RPC");
+    if (!rpc_reg) {
+        LOG_ERROR("RPC backend not found, cannot add RPC servers");
+        return false;
+    }
+    typedef ggml_backend_reg_t (*ggml_backend_rpc_add_server_t)(const char* endpoint);
+    ggml_backend_rpc_add_server_t ggml_backend_rpc_add_server_fn = (ggml_backend_rpc_add_server_t)ggml_backend_reg_get_proc_address(rpc_reg, "ggml_backend_rpc_add_server");
+    if (!ggml_backend_rpc_add_server_fn) {
+        LOG_ERROR("RPC backend does not have ggml_backend_rpc_add_server function, cannot add RPC servers");
+        return false;
+    }
+    for (const auto& server : rpc_servers) {
+        LOG_INFO("Adding RPC server: %s", server.c_str());
+        auto reg = ggml_backend_rpc_add_server_fn(server.c_str());
+        // no return value to check for success but should print errors from the RPC backend if it fails to add the server
+        ggml_backend_register(reg);
+    }
+    return true;
+}
+
 static void ggml_backend_load_all_once() {
     // If the registry already has devices and the CPU backend is present,
     // assume either static registration or explicit host-side preloading has
@@ -504,6 +538,9 @@ ggml_backend_t SDBackendManager::params_backend(SDBackendModule module) {
     if (name.empty()) {
         return runtime_backend(module);
     }
+    if (is_disk_backend_token(name)) {
+        return runtime_backend(module);
+    }
     return init_cached_backend(name);
 }
 
@@ -515,6 +552,10 @@ bool SDBackendManager::params_backend_is_cpu(SDBackendModule module) {
     return sd_backend_is_cpu(params_backend(module));
 }
 
+bool SDBackendManager::params_backend_is_disk(SDBackendModule module) const {
+    return is_disk_backend_token(params_assignment_.get(module));
+}
+
 bool SDBackendManager::runtime_backend_supports_host_buffer(SDBackendModule module) {
     ggml_backend_t backend = runtime_backend(module);
     if (backend == nullptr) {
@@ -534,10 +575,6 @@ bool SDBackendManager::runtime_backend_supports_host_buffer(SDBackendModule modu
 
 bool SDBackendManager::init(const char* backend_spec,
                             const char* params_backend_spec,
-                            bool offload_params_to_cpu,
-                            bool keep_clip_on_cpu,
-                            bool keep_vae_on_cpu,
-                            bool keep_control_net_on_cpu,
                             std::string* error) {
     reset();
 
@@ -548,30 +585,20 @@ bool SDBackendManager::init(const char* backend_spec,
         return false;
     }
 
-    if (runtime_assignment_.empty()) {
-        if (keep_clip_on_cpu) {
-            runtime_assignment_.set_module(SDBackendModule::TE, "cpu");
-        }
-        if (keep_vae_on_cpu) {
-            runtime_assignment_.set_module(SDBackendModule::VAE, "cpu");
-        }
-        if (keep_control_net_on_cpu) {
-            runtime_assignment_.set_module(SDBackendModule::CONTROL_NET, "cpu");
-        }
-    }
-
-    if (params_assignment_.empty() && offload_params_to_cpu) {
-        params_assignment_.set_default("cpu");
-    }
-
     return validate(error);
 }
 
 bool SDBackendManager::validate(std::string* error) const {
-    auto validate_name = [&](const std::string& name) -> bool {
+    auto validate_runtime_name = [&](const std::string& name) -> bool {
         if (is_default_backend_token(name)) {
             return true;
         }
+        if (is_disk_backend_token(name)) {
+            if (error != nullptr) {
+                *error = "backend 'disk' is only supported by params_backend";
+            }
+            return false;
+        }
         if (!sd_resolve_backend_name(name).empty()) {
             return true;
         }
@@ -580,18 +607,24 @@ bool SDBackendManager::validate(std::string* error) const {
         }
         return false;
     };
+    auto validate_params_name = [&](const std::string& name) -> bool {
+        if (is_disk_backend_token(name)) {
+            return true;
+        }
+        return validate_runtime_name(name);
+    };
 
-    if (!validate_name(runtime_assignment_.default_name) ||
-        !validate_name(params_assignment_.default_name)) {
+    if (!validate_runtime_name(runtime_assignment_.default_name) ||
+        !validate_params_name(params_assignment_.default_name)) {
         return false;
     }
     for (const auto& kv : runtime_assignment_.module_names) {
-        if (!validate_name(kv.second)) {
+        if (!validate_runtime_name(kv.second)) {
             return false;
         }
     }
     for (const auto& kv : params_assignment_.module_names) {
-        if (!validate_name(kv.second)) {
+        if (!validate_params_name(kv.second)) {
             return false;
         }
     }

src/core/ggml_extend_backend.h

@@ -1,5 +1,5 @@
-#ifndef __SD_GGML_EXTEND_BACKEND_H__
-#define __SD_GGML_EXTEND_BACKEND_H__
+#ifndef __SD_CORE_GGML_EXTEND_BACKEND_H__
+#define __SD_CORE_GGML_EXTEND_BACKEND_H__
 
 #include <cstdint>
 #include <cstring>
@@ -51,10 +51,6 @@ public:
 
     bool init(const char* backend_spec,
               const char* params_backend_spec,
-              bool offload_params_to_cpu,
-              bool keep_clip_on_cpu,
-              bool keep_vae_on_cpu,
-              bool keep_control_net_on_cpu,
               std::string* error);
     void reset();
 
@@ -63,6 +59,7 @@ public:
 
     bool runtime_backend_is_cpu(SDBackendModule module);
     bool params_backend_is_cpu(SDBackendModule module);
+    bool params_backend_is_disk(SDBackendModule module) const;
     bool runtime_backend_supports_host_buffer(SDBackendModule module);
 
 private:
@@ -76,4 +73,5 @@ ggml_backend_t sd_backend_cpu_init();
 bool sd_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_threads);
 const char* sd_backend_module_name(SDBackendModule module);
 void ggml_ext_im_set_f32_1d(const struct ggml_tensor* tensor, int i, float value);
-#endif
+bool add_rpc_devices(const std::string& servers);
+#endif  // __SD_CORE_GGML_EXTEND_BACKEND_H__

src/core/ggml_graph_cut.cpp

@@ -1,4 +1,4 @@
-#include "ggml_graph_cut.h"
+#include "core/ggml_graph_cut.h"
 
 #include <algorithm>
 #include <cstring>
@@ -8,11 +8,11 @@
 #include <stack>
 #include <unordered_map>
 
+#include "core/util.h"
 #include "ggml-alloc.h"
 #include "ggml-backend.h"
-#include "util.h"
 
-#include "../ggml/src/ggml-impl.h"
+#include "ggml/src/ggml-impl.h"
 
 namespace sd::ggml_graph_cut {
 
@@ -44,7 +44,9 @@ namespace sd::ggml_graph_cut {
         if (tensor == nullptr) {
             return false;
         }
-        return params_tensor_set.find(tensor) != params_tensor_set.end();
+        return params_tensor_set.find(tensor) != params_tensor_set.end() ||
+               (tensor->view_src != nullptr &&
+                params_tensor_set.find(tensor->view_src) != params_tensor_set.end());
     }
 
     static int graph_node_index_by_name(ggml_cgraph* gf, const char* name) {
@@ -135,6 +137,24 @@ namespace sd::ggml_graph_cut {
         return max_vram_bytes_to_gib(resolve_auto_max_vram_bytes(-max_vram, backend));
     }
 
+    static bool is_segment_output_needed_after(const Plan& plan,
+                                               size_t end_segment_index,
+                                               int output_node_index) {
+        if (end_segment_index + 1 >= plan.segments.size()) {
+            return false;
+        }
+        for (size_t seg_idx = end_segment_index + 1; seg_idx < plan.segments.size(); ++seg_idx) {
+            const auto& segment = plan.segments[seg_idx];
+            for (const auto& input_ref : segment.input_refs) {
+                if (input_ref.type == Segment::INPUT_PREVIOUS_CUT &&
+                    input_ref.node_index == output_node_index) {
+                    return true;
+                }
+            }
+        }
+        return false;
+    }
+
     static Segment make_segment_seed(const Plan& plan,
                                      size_t start_segment_index,
                                      size_t end_segment_index) {
@@ -147,8 +167,11 @@ namespace sd::ggml_graph_cut {
         const auto& target_segment = plan.segments[end_segment_index];
         std::unordered_set<int> seen_output_node_indices;
         for (size_t seg_idx = start_segment_index; seg_idx <= end_segment_index; ++seg_idx) {
+            const bool is_boundary_segment = seg_idx == end_segment_index;
             for (int output_node_index : plan.segments[seg_idx].output_node_indices) {
-                if (seen_output_node_indices.insert(output_node_index).second) {
+                if ((is_boundary_segment ||
+                     is_segment_output_needed_after(plan, end_segment_index, output_node_index)) &&
+                    seen_output_node_indices.insert(output_node_index).second) {
                     seed.output_node_indices.push_back(output_node_index);
                 }
             }
@@ -400,23 +423,6 @@ namespace sd::ggml_graph_cut {
         return tensors;
     }
 
-    std::vector<ggml_tensor*> runtime_param_tensors(ggml_cgraph* gf, const Segment& segment, const char* log_desc) {
-        std::vector<ggml_tensor*> tensors = param_tensors(gf, segment);
-        std::vector<ggml_tensor*> filtered_tensors;
-        filtered_tensors.reserve(tensors.size());
-        for (ggml_tensor* tensor : tensors) {
-            if (tensor_buffer(tensor) == nullptr) {
-                LOG_WARN("%s graph cut skipping param input without buffer: segment=%s tensor=%s",
-                         log_desc == nullptr ? "unknown" : log_desc,
-                         segment.group_name.c_str(),
-                         tensor->name);
-                continue;
-            }
-            filtered_tensors.push_back(tensor);
-        }
-        return filtered_tensors;
-    }
-
     std::unordered_set<std::string> collect_future_input_names(ggml_cgraph* gf,
                                                                const Plan& plan,
                                                                size_t current_segment_index) {
@@ -487,6 +493,44 @@ namespace sd::ggml_graph_cut {
             return 0;
         }
 
+        struct TensorRuntimeBinding {
+            ggml_backend_buffer_t buffer = nullptr;
+            void* data                   = nullptr;
+            void* extra                  = nullptr;
+        };
+        std::unordered_map<ggml_tensor*, TensorRuntimeBinding> saved_bindings;
+        auto mark_measurement_external = [&](ggml_tensor* tensor) {
+            if (tensor == nullptr) {
+                return;
+            }
+            auto save_tensor = [&](ggml_tensor* t) {
+                if (t == nullptr || saved_bindings.find(t) != saved_bindings.end()) {
+                    return;
+                }
+                saved_bindings[t] = {t->buffer, t->data, t->extra};
+                // During real execution params and previous-cut inputs already
+                // have backend/cache buffers, so gallocr must not reserve them.
+                t->data = reinterpret_cast<void*>(static_cast<uintptr_t>(1));
+            };
+            save_tensor(tensor);
+            save_tensor(tensor->view_src);
+        };
+        for (const auto& input : segment.input_refs) {
+            if (input.type != Segment::INPUT_PARAM &&
+                input.type != Segment::INPUT_PREVIOUS_CUT) {
+                continue;
+            }
+            mark_measurement_external(input_tensor(gf, input));
+        }
+
+        std::unordered_map<ggml_tensor*, int32_t> saved_output_flags;
+        for (int output_node_index : segment.output_node_indices) {
+            ggml_tensor* output = ggml_graph_node(gf, output_node_index);
+            if (output != nullptr && saved_output_flags.find(output) == saved_output_flags.end()) {
+                saved_output_flags[output] = output->flags;
+            }
+        }
+
         ggml_context* graph_ctx    = nullptr;
         ggml_cgraph* segment_graph = build_segment_graph(gf, segment, &graph_ctx);
         ggml_gallocr_t allocr      = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend));
@@ -502,6 +546,14 @@ namespace sd::ggml_graph_cut {
 
         ggml_gallocr_free(allocr);
         ggml_free(graph_ctx);
+        for (const auto& kv : saved_output_flags) {
+            kv.first->flags = kv.second;
+        }
+        for (const auto& kv : saved_bindings) {
+            kv.first->buffer = kv.second.buffer;
+            kv.first->data   = kv.second.data;
+            kv.first->extra  = kv.second.extra;
+        }
         return buffer_size;
     }
 
@@ -669,7 +721,8 @@ namespace sd::ggml_graph_cut {
                 GGML_ASSERT(!candidate_plan.segments.empty());
 
                 const auto& candidate_segment = candidate_plan.segments.back();
-                if (graph_cut_segment_vram_bytes(candidate_segment) > max_graph_vram_bytes) {
+                const size_t candidate_bytes  = graph_cut_segment_vram_bytes(candidate_segment);
+                if (candidate_bytes > max_graph_vram_bytes) {
                     break;
                 }
 
@@ -699,9 +752,9 @@ namespace sd::ggml_graph_cut {
         }
 
         if (log_desc != nullptr) {
-            LOG_INFO("%s graph cut max_vram budget merge took %lld ms",
-                     log_desc,
-                     ggml_time_ms() - t_budget_begin);
+            LOG_DEBUG("%s graph cut max_vram budget merge took %lld ms",
+                      log_desc,
+                      ggml_time_ms() - t_budget_begin);
         }
 
         return merged_plan;

src/core/ggml_graph_cut.h

@@ -1,5 +1,5 @@
-#ifndef __SD_GGML_GRAPH_CUT_H__
-#define __SD_GGML_GRAPH_CUT_H__
+#ifndef __SD_CORE_GGML_GRAPH_CUT_H__
+#define __SD_CORE_GGML_GRAPH_CUT_H__
 
 #include <array>
 #include <cstdint>
@@ -80,7 +80,6 @@ namespace sd::ggml_graph_cut {
     ggml_tensor* output_tensor(ggml_cgraph* gf, const Segment& segment, size_t output_index);
     ggml_tensor* input_tensor(ggml_cgraph* gf, const Segment::InputRef& input_ref);
     std::vector<ggml_tensor*> param_tensors(ggml_cgraph* gf, const Segment& segment);
-    std::vector<ggml_tensor*> runtime_param_tensors(ggml_cgraph* gf, const Segment& segment, const char* log_desc);
     std::unordered_set<std::string> collect_future_input_names(ggml_cgraph* gf,
                                                                const Plan& plan,
                                                                size_t current_segment_index);
@@ -114,4 +113,4 @@ namespace sd::ggml_graph_cut {
     void annotate_residency(Plan& plan, size_t max_graph_vram_bytes);
 }  // namespace sd::ggml_graph_cut
 
-#endif
+#endif  // __SD_CORE_GGML_GRAPH_CUT_H__

src/core/ordered_map.hpp

@@ -1,5 +1,5 @@
-#ifndef __ORDERED_MAP_HPP__
-#define __ORDERED_MAP_HPP__
+#ifndef __SD_CORE_ORDERED_MAP_HPP__
+#define __SD_CORE_ORDERED_MAP_HPP__
 
 #include <iostream>
 #include <list>
@@ -174,4 +174,4 @@ public:
     }
 };
 
-#endif  // __ORDERED_MAP_HPP__
+#endif  // __SD_CORE_ORDERED_MAP_HPP__

src/core/rng.hpp

@@ -1,5 +1,5 @@
-#ifndef __RNG_H__
-#define __RNG_H__
+#ifndef __SD_CORE_RNG_HPP__
+#define __SD_CORE_RNG_HPP__
 
 #include <random>
 #include <vector>
@@ -32,4 +32,4 @@ public:
     }
 };
 
-#endif  // __RNG_H__
+#endif  // __SD_CORE_RNG_HPP__

src/core/rng_mt19937.hpp

@@ -1,10 +1,10 @@
-#ifndef __RNG_MT19937_HPP__
-#define __RNG_MT19937_HPP__
+#ifndef __SD_CORE_RNG_MT19937_HPP__
+#define __SD_CORE_RNG_MT19937_HPP__
 
 #include <cmath>
 #include <vector>
 
-#include "rng.hpp"
+#include "core/rng.hpp"
 
 // RNG imitiating torch cpu randn on CPU.
 // Port from pytorch, original license: https://github.com/pytorch/pytorch/blob/d01a7b0241ed1c4cded7e7ca097249feb343f072/LICENSE
@@ -144,4 +144,4 @@ public:
     }
 };
 
-#endif  // __RNG_MT19937_HPP__
+#endif  // __SD_CORE_RNG_MT19937_HPP__

src/core/rng_philox.hpp

@@ -1,10 +1,10 @@
-#ifndef __RNG_PHILOX_H__
-#define __RNG_PHILOX_H__
+#ifndef __SD_CORE_RNG_PHILOX_HPP__
+#define __SD_CORE_RNG_PHILOX_HPP__
 
 #include <cmath>
 #include <vector>
 
-#include "rng.hpp"
+#include "core/rng.hpp"
 
 // RNG imitiating torch cuda randn on CPU.
 // Port from: https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/5ef669de080814067961f28357256e8fe27544f4/modules/rng_philox.py
@@ -122,4 +122,4 @@ public:
     }
 };
 
-#endif  // __RNG_PHILOX_H__
+#endif  // __SD_CORE_RNG_PHILOX_HPP__

src/core/tensor.hpp

@@ -1,5 +1,5 @@
-#ifndef __SD_TENSOR_HPP__
-#define __SD_TENSOR_HPP__
+#ifndef __SD_CORE_TENSOR_HPP__
+#define __SD_CORE_TENSOR_HPP__
 
 #include <algorithm>
 #include <cmath>
@@ -16,7 +16,7 @@
 #include <utility>
 #include <vector>
 
-#include "rng.hpp"
+#include "core/rng.hpp"
 
 namespace sd {
 
@@ -1661,4 +1661,4 @@ namespace sd {
 
 }  // namespace sd
 
-#endif
+#endif  // __SD_CORE_TENSOR_HPP__

src/core/tensor_ggml.hpp

@@ -1,5 +1,5 @@
-#ifndef __SD_TENSOR_GGML_HPP__
-#define __SD_TENSOR_GGML_HPP__
+#ifndef __SD_CORE_TENSOR_GGML_HPP__
+#define __SD_CORE_TENSOR_GGML_HPP__
 
 #include <array>
 #include <cstring>
@@ -8,8 +8,8 @@
 #include <string>
 #include <type_traits>
 
+#include "core/tensor.hpp"
 #include "ggml.h"
-#include "tensor.hpp"
 
 namespace sd {
 
@@ -124,4 +124,4 @@ namespace sd {
 
 }  // namespace sd
 
-#endif
+#endif  // __SD_CORE_TENSOR_GGML_HPP__

src/core/util.cpp

@@ -1,4 +1,4 @@
-#include "util.h"
+#include "core/util.h"
 #include <algorithm>
 #include <cctype>
 #include <cmath>
@@ -13,7 +13,7 @@
 #include <thread>
 #include <unordered_set>
 #include <vector>
-#include "preprocessing.hpp"
+#include "runtime/preprocessing.hpp"
 
 #if defined(__APPLE__) && defined(__MACH__)
 #include <sys/sysctl.h>
@@ -488,7 +488,7 @@ bool parse_strict_bool(const std::string& text, bool& value) {
     return false;
 }
 
-static std::string build_progress_bar(int step, int steps) {
+static std::string build_progress_bar(int step, int steps, char progress_char = '=', bool show_head = true) {
     std::string progress = "  |";
     int max_progress     = 50;
     int32_t current      = 0;
@@ -498,21 +498,21 @@ static std::string build_progress_bar(int step, int steps) {
     for (int i = 0; i < 50; i++) {
         if (i > current) {
             progress += " ";
-        } else if (i == current && i != max_progress - 1) {
+        } else if (show_head && i == current && i != max_progress - 1) {
             progress += ">";
         } else {
-            progress += "=";
+            progress += progress_char;
         }
     }
     progress += "|";
     return progress;
 }
 
-static void print_progress_line(int step, int steps, const std::string& speed_text) {
+static void print_progress_line(int step, int steps, const std::string& speed_text, char progress_char = '=', bool show_head = true) {
     if (step == 0) {
         return;
     }
-    std::string progress = build_progress_bar(step, steps);
+    std::string progress = build_progress_bar(step, steps, progress_char, show_head);
     const char* lf       = (step == steps ? "\n" : "");
     printf("\r%s %i/%i - %s\033[K%s", progress.c_str(), step, steps, speed_text.c_str(), lf);
     fflush(stdout);  // for linux
@@ -552,9 +552,9 @@ void pretty_bytes_progress(int step, int steps, uint64_t bytes_processed, float
 
     double speed_mb = bytes_per_second / (1024.0 * 1024.0);
     if (speed_mb >= 1024.0) {
-        print_progress_line(step, steps, sd_format("%.2fGB/s", speed_mb / 1024.0));
+        print_progress_line(step, steps, sd_format("%.2fGB/s", speed_mb / 1024.0), '#', false);
     } else {
-        print_progress_line(step, steps, sd_format("%.2fMB/s", speed_mb));
+        print_progress_line(step, steps, sd_format("%.2fMB/s", speed_mb), '#', false);
     }
 }
 

src/core/util.h

@@ -1,5 +1,5 @@
-#ifndef __UTIL_H__
-#define __UTIL_H__
+#ifndef __SD_CORE_UTIL_H__
+#define __SD_CORE_UTIL_H__
 
 #include <cstdint>
 #include <memory>
@@ -7,9 +7,9 @@
 #include <utility>
 #include <vector>
 
+#include "core/tensor.hpp"
 #include "ggml-backend.h"
 #include "stable-diffusion.h"
-#include "tensor.hpp"
 
 #define SAFE_STR(s) ((s) ? (s) : "")
 #define BOOL_STR(b) ((b) ? "true" : "false")
@@ -103,4 +103,4 @@ bool sd_backend_is(ggml_backend_t backend, const std::string& name);
 #define LOG_INFO(format, ...) log_printf(SD_LOG_INFO, __FILE__, __LINE__, format, ##__VA_ARGS__)
 #define LOG_WARN(format, ...) log_printf(SD_LOG_WARN, __FILE__, __LINE__, format, ##__VA_ARGS__)
 #define LOG_ERROR(format, ...) log_printf(SD_LOG_ERROR, __FILE__, __LINE__, format, ##__VA_ARGS__)
-#endif  // __UTIL_H__
+#endif  // __SD_CORE_UTIL_H__

src/esrgan.hpp

@@ -1,375 +0,0 @@
-#ifndef __ESRGAN_HPP__
-#define __ESRGAN_HPP__
-
-#include "ggml_extend.hpp"
-#include "model.h"
-
-/*
-    ===================================    ESRGAN  ===================================
-    References:
-    https://github.com/xinntao/Real-ESRGAN/blob/master/inference_realesrgan.py
-    https://github.com/XPixelGroup/BasicSR/blob/v1.4.2/basicsr/archs/rrdbnet_arch.py
-
-*/
-
-class ResidualDenseBlock : public GGMLBlock {
-protected:
-    int num_feat;
-    int num_grow_ch;
-
-public:
-    ResidualDenseBlock(int num_feat = 64, int num_grow_ch = 32)
-        : num_feat(num_feat), num_grow_ch(num_grow_ch) {
-        blocks["conv1"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_grow_ch, {3, 3}, {1, 1}, {1, 1}));
-        blocks["conv2"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat + num_grow_ch, num_grow_ch, {3, 3}, {1, 1}, {1, 1}));
-        blocks["conv3"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat + 2 * num_grow_ch, num_grow_ch, {3, 3}, {1, 1}, {1, 1}));
-        blocks["conv4"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat + 3 * num_grow_ch, num_grow_ch, {3, 3}, {1, 1}, {1, 1}));
-        blocks["conv5"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat + 4 * num_grow_ch, num_feat, {3, 3}, {1, 1}, {1, 1}));
-    }
-
-    ggml_tensor* lrelu(GGMLRunnerContext* ctx, ggml_tensor* x) {
-        return ggml_leaky_relu(ctx->ggml_ctx, x, 0.2f, true);
-    }
-
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
-        // x: [n, num_feat, h, w]
-        // return: [n, num_feat, h, w]
-
-        auto conv1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv1"]);
-        auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["conv2"]);
-        auto conv3 = std::dynamic_pointer_cast<Conv2d>(blocks["conv3"]);
-        auto conv4 = std::dynamic_pointer_cast<Conv2d>(blocks["conv4"]);
-        auto conv5 = std::dynamic_pointer_cast<Conv2d>(blocks["conv5"]);
-
-        auto x1    = lrelu(ctx, conv1->forward(ctx, x));
-        auto x_cat = ggml_concat(ctx->ggml_ctx, x, x1, 2);
-        auto x2    = lrelu(ctx, conv2->forward(ctx, x_cat));
-        x_cat      = ggml_concat(ctx->ggml_ctx, x_cat, x2, 2);
-        auto x3    = lrelu(ctx, conv3->forward(ctx, x_cat));
-        x_cat      = ggml_concat(ctx->ggml_ctx, x_cat, x3, 2);
-        auto x4    = lrelu(ctx, conv4->forward(ctx, x_cat));
-        x_cat      = ggml_concat(ctx->ggml_ctx, x_cat, x4, 2);
-        auto x5    = conv5->forward(ctx, x_cat);
-
-        x5 = ggml_add(ctx->ggml_ctx, ggml_ext_scale(ctx->ggml_ctx, x5, 0.2f), x);
-        return x5;
-    }
-};
-
-class RRDB : public GGMLBlock {
-public:
-    RRDB(int num_feat, int num_grow_ch = 32) {
-        blocks["rdb1"] = std::shared_ptr<GGMLBlock>(new ResidualDenseBlock(num_feat, num_grow_ch));
-        blocks["rdb2"] = std::shared_ptr<GGMLBlock>(new ResidualDenseBlock(num_feat, num_grow_ch));
-        blocks["rdb3"] = std::shared_ptr<GGMLBlock>(new ResidualDenseBlock(num_feat, num_grow_ch));
-    }
-
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
-        // x: [n, num_feat, h, w]
-        // return: [n, num_feat, h, w]
-
-        auto rdb1 = std::dynamic_pointer_cast<ResidualDenseBlock>(blocks["rdb1"]);
-        auto rdb2 = std::dynamic_pointer_cast<ResidualDenseBlock>(blocks["rdb2"]);
-        auto rdb3 = std::dynamic_pointer_cast<ResidualDenseBlock>(blocks["rdb3"]);
-
-        auto out = rdb1->forward(ctx, x);
-        out      = rdb2->forward(ctx, out);
-        out      = rdb3->forward(ctx, out);
-
-        out = ggml_add(ctx->ggml_ctx, ggml_ext_scale(ctx->ggml_ctx, out, 0.2f), x);
-        return out;
-    }
-};
-
-class RRDBNet : public GGMLBlock {
-protected:
-    int scale       = 4;
-    int num_block   = 23;
-    int num_in_ch   = 3;
-    int num_out_ch  = 3;
-    int num_feat    = 64;
-    int num_grow_ch = 32;
-
-public:
-    RRDBNet(int scale, int num_block, int num_in_ch, int num_out_ch, int num_feat, int num_grow_ch)
-        : scale(scale), num_block(num_block), num_in_ch(num_in_ch), num_out_ch(num_out_ch), num_feat(num_feat), num_grow_ch(num_grow_ch) {
-        blocks["conv_first"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_in_ch, num_feat, {3, 3}, {1, 1}, {1, 1}));
-        for (int i = 0; i < num_block; i++) {
-            std::string name = "body." + std::to_string(i);
-            blocks[name]     = std::shared_ptr<GGMLBlock>(new RRDB(num_feat, num_grow_ch));
-        }
-        blocks["conv_body"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_feat, {3, 3}, {1, 1}, {1, 1}));
-        if (scale >= 2) {
-            blocks["conv_up1"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_feat, {3, 3}, {1, 1}, {1, 1}));
-        }
-        if (scale == 4) {
-            blocks["conv_up2"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_feat, {3, 3}, {1, 1}, {1, 1}));
-        }
-        blocks["conv_hr"]   = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_feat, {3, 3}, {1, 1}, {1, 1}));
-        blocks["conv_last"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_out_ch, {3, 3}, {1, 1}, {1, 1}));
-    }
-
-    int get_scale() { return scale; }
-    int get_num_block() { return num_block; }
-
-    ggml_tensor* lrelu(GGMLRunnerContext* ctx, ggml_tensor* x) {
-        return ggml_leaky_relu(ctx->ggml_ctx, x, 0.2f, true);
-    }
-
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
-        // x: [n, num_in_ch, h, w]
-        // return: [n, num_out_ch, h*scale, w*scale]
-        auto conv_first = std::dynamic_pointer_cast<Conv2d>(blocks["conv_first"]);
-        auto conv_body  = std::dynamic_pointer_cast<Conv2d>(blocks["conv_body"]);
-        auto conv_hr    = std::dynamic_pointer_cast<Conv2d>(blocks["conv_hr"]);
-        auto conv_last  = std::dynamic_pointer_cast<Conv2d>(blocks["conv_last"]);
-
-        auto feat = conv_first->forward(ctx, x);
-        sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.prelude", "feat");
-        auto body_feat = feat;
-        for (int i = 0; i < num_block; i++) {
-            std::string name = "body." + std::to_string(i);
-            auto block       = std::dynamic_pointer_cast<RRDB>(blocks[name]);
-
-            body_feat = block->forward(ctx, body_feat);
-            sd::ggml_graph_cut::mark_graph_cut(body_feat, "esrgan.body." + std::to_string(i), "feat");
-        }
-        body_feat = conv_body->forward(ctx, body_feat);
-        feat      = ggml_add(ctx->ggml_ctx, feat, body_feat);
-        sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.body.out", "feat");
-        // upsample
-        if (scale >= 2) {
-            auto conv_up1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv_up1"]);
-            feat          = lrelu(ctx, conv_up1->forward(ctx, ggml_upscale(ctx->ggml_ctx, feat, 2, GGML_SCALE_MODE_NEAREST)));
-            sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.up1", "feat");
-            if (scale == 4) {
-                auto conv_up2 = std::dynamic_pointer_cast<Conv2d>(blocks["conv_up2"]);
-                feat          = lrelu(ctx, conv_up2->forward(ctx, ggml_upscale(ctx->ggml_ctx, feat, 2, GGML_SCALE_MODE_NEAREST)));
-                sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.up2", "feat");
-            }
-        }
-        // for all scales
-        auto out = conv_last->forward(ctx, lrelu(ctx, conv_hr->forward(ctx, feat)));
-        sd::ggml_graph_cut::mark_graph_cut(out, "esrgan.final", "out");
-        return out;
-    }
-};
-
-struct ESRGAN : public GGMLRunner {
-    std::unique_ptr<RRDBNet> rrdb_net;
-    int scale     = 4;
-    int tile_size = 128;  // avoid cuda OOM for 4gb VRAM
-
-    ESRGAN(ggml_backend_t backend,
-           ggml_backend_t params_backend,
-           int tile_size                                  = 128,
-           const String2TensorStorage& tensor_storage_map = {})
-        : GGMLRunner(backend, params_backend) {
-        this->tile_size = tile_size;
-    }
-
-    std::string get_desc() override {
-        return "esrgan";
-    }
-
-    bool load_from_file(const std::string& file_path, int n_threads) {
-        LOG_INFO("loading esrgan from '%s'", file_path.c_str());
-
-        ModelLoader model_loader;
-        if (!model_loader.init_from_file_and_convert_name(file_path)) {
-            LOG_ERROR("init esrgan model loader from file failed: '%s'", file_path.c_str());
-            return false;
-        }
-
-        // Get tensor names
-        auto tensor_names = model_loader.get_tensor_names();
-
-        // Detect if it's ESRGAN format
-        bool is_ESRGAN = std::find(tensor_names.begin(), tensor_names.end(), "model.0.weight") != tensor_names.end();
-
-        // Detect parameters from tensor names
-        int detected_num_block = 0;
-        if (is_ESRGAN) {
-            for (const auto& name : tensor_names) {
-                if (name.find("model.1.sub.") == 0) {
-                    size_t first_dot = name.find('.', 12);
-                    if (first_dot != std::string::npos) {
-                        size_t second_dot = name.find('.', first_dot + 1);
-                        if (second_dot != std::string::npos && name.substr(first_dot + 1, 3) == "RDB") {
-                            try {
-                                int idx            = std::stoi(name.substr(12, first_dot - 12));
-                                detected_num_block = std::max(detected_num_block, idx + 1);
-                            } catch (...) {
-                            }
-                        }
-                    }
-                }
-            }
-        } else {
-            // Original format
-            for (const auto& name : tensor_names) {
-                if (name.find("body.") == 0) {
-                    size_t pos = name.find('.', 5);
-                    if (pos != std::string::npos) {
-                        try {
-                            int idx            = std::stoi(name.substr(5, pos - 5));
-                            detected_num_block = std::max(detected_num_block, idx + 1);
-                        } catch (...) {
-                        }
-                    }
-                }
-            }
-        }
-
-        int detected_scale = 4;  // default
-        if (is_ESRGAN) {
-            // For ESRGAN format, detect scale by highest model number
-            int max_model_num = 0;
-            for (const auto& name : tensor_names) {
-                if (name.find("model.") == 0) {
-                    size_t dot_pos = name.find('.', 6);
-                    if (dot_pos != std::string::npos) {
-                        try {
-                            int num       = std::stoi(name.substr(6, dot_pos - 6));
-                            max_model_num = std::max(max_model_num, num);
-                        } catch (...) {
-                        }
-                    }
-                }
-            }
-            if (max_model_num <= 4) {
-                detected_scale = 1;
-            } else if (max_model_num <= 7) {
-                detected_scale = 2;
-            } else {
-                detected_scale = 4;
-            }
-        } else {
-            // Original format
-            bool has_conv_up2 = std::any_of(tensor_names.begin(), tensor_names.end(), [](const std::string& name) {
-                return name == "conv_up2.weight";
-            });
-            bool has_conv_up1 = std::any_of(tensor_names.begin(), tensor_names.end(), [](const std::string& name) {
-                return name == "conv_up1.weight";
-            });
-            if (has_conv_up2) {
-                detected_scale = 4;
-            } else if (has_conv_up1) {
-                detected_scale = 2;
-            } else {
-                detected_scale = 1;
-            }
-        }
-
-        int detected_num_in_ch   = 3;
-        int detected_num_out_ch  = 3;
-        int detected_num_feat    = 64;
-        int detected_num_grow_ch = 32;
-
-        // Create RRDBNet with detected parameters
-        rrdb_net = std::make_unique<RRDBNet>(detected_scale, detected_num_block, detected_num_in_ch, detected_num_out_ch, detected_num_feat, detected_num_grow_ch);
-        rrdb_net->init(params_ctx, {}, "");
-
-        if (!alloc_params_buffer()) {
-            LOG_ERROR("esrgan model buffer allocation failed");
-            return false;
-        }
-
-        std::map<std::string, ggml_tensor*> esrgan_tensors;
-        rrdb_net->get_param_tensors(esrgan_tensors);
-
-        bool success;
-        if (is_ESRGAN) {
-            // Build name mapping for ESRGAN format
-            std::map<std::string, std::string> expected_to_model;
-            expected_to_model["conv_first.weight"] = "model.0.weight";
-            expected_to_model["conv_first.bias"]   = "model.0.bias";
-
-            for (int i = 0; i < detected_num_block; i++) {
-                for (int j = 1; j <= 3; j++) {
-                    for (int k = 1; k <= 5; k++) {
-                        std::string expected_weight        = "body." + std::to_string(i) + ".rdb" + std::to_string(j) + ".conv" + std::to_string(k) + ".weight";
-                        std::string model_weight           = "model.1.sub." + std::to_string(i) + ".RDB" + std::to_string(j) + ".conv" + std::to_string(k) + ".0.weight";
-                        expected_to_model[expected_weight] = model_weight;
-
-                        std::string expected_bias        = "body." + std::to_string(i) + ".rdb" + std::to_string(j) + ".conv" + std::to_string(k) + ".bias";
-                        std::string model_bias           = "model.1.sub." + std::to_string(i) + ".RDB" + std::to_string(j) + ".conv" + std::to_string(k) + ".0.bias";
-                        expected_to_model[expected_bias] = model_bias;
-                    }
-                }
-            }
-
-            if (detected_scale == 1) {
-                expected_to_model["conv_body.weight"] = "model.1.sub." + std::to_string(detected_num_block) + ".weight";
-                expected_to_model["conv_body.bias"]   = "model.1.sub." + std::to_string(detected_num_block) + ".bias";
-                expected_to_model["conv_hr.weight"]   = "model.2.weight";
-                expected_to_model["conv_hr.bias"]     = "model.2.bias";
-                expected_to_model["conv_last.weight"] = "model.4.weight";
-                expected_to_model["conv_last.bias"]   = "model.4.bias";
-            } else {
-                expected_to_model["conv_body.weight"] = "model.1.sub." + std::to_string(detected_num_block) + ".weight";
-                expected_to_model["conv_body.bias"]   = "model.1.sub." + std::to_string(detected_num_block) + ".bias";
-                if (detected_scale >= 2) {
-                    expected_to_model["conv_up1.weight"] = "model.3.weight";
-                    expected_to_model["conv_up1.bias"]   = "model.3.bias";
-                }
-                if (detected_scale == 4) {
-                    expected_to_model["conv_up2.weight"]  = "model.6.weight";
-                    expected_to_model["conv_up2.bias"]    = "model.6.bias";
-                    expected_to_model["conv_hr.weight"]   = "model.8.weight";
-                    expected_to_model["conv_hr.bias"]     = "model.8.bias";
-                    expected_to_model["conv_last.weight"] = "model.10.weight";
-                    expected_to_model["conv_last.bias"]   = "model.10.bias";
-                } else if (detected_scale == 2) {
-                    expected_to_model["conv_hr.weight"]   = "model.5.weight";
-                    expected_to_model["conv_hr.bias"]     = "model.5.bias";
-                    expected_to_model["conv_last.weight"] = "model.7.weight";
-                    expected_to_model["conv_last.bias"]   = "model.7.bias";
-                }
-            }
-
-            std::map<std::string, ggml_tensor*> model_tensors;
-            for (auto& p : esrgan_tensors) {
-                auto it = expected_to_model.find(p.first);
-                if (it != expected_to_model.end()) {
-                    model_tensors[it->second] = p.second;
-                }
-            }
-
-            success = model_loader.load_tensors(model_tensors, {}, n_threads);
-        } else {
-            success = model_loader.load_tensors(esrgan_tensors, {}, n_threads);
-        }
-
-        if (!success) {
-            LOG_ERROR("load esrgan tensors from model loader failed");
-            return false;
-        }
-
-        scale = rrdb_net->get_scale();
-        LOG_INFO("esrgan model loaded with scale=%d, num_block=%d", scale, detected_num_block);
-        return success;
-    }
-
-    ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor) {
-        if (!rrdb_net)
-            return nullptr;
-        constexpr int kGraphNodes = 1 << 16;  // 65k
-        ggml_cgraph* gf           = new_graph_custom(kGraphNodes);
-        ggml_tensor* x            = make_input(x_tensor);
-
-        auto runner_ctx  = get_context();
-        ggml_tensor* out = rrdb_net->forward(&runner_ctx, x);
-        ggml_build_forward_expand(gf, out);
-        return gf;
-    }
-
-    sd::Tensor<float> compute(const int n_threads,
-                              const sd::Tensor<float>& x) {
-        auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x); };
-        auto result    = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
-        return result;
-    }
-};
-
-#endif  // __ESRGAN_HPP__

src/extensions/generation_extension.h

@@ -0,0 +1,63 @@
+#ifndef __SD_EXTENSIONS_GENERATION_EXTENSION_H__
+#define __SD_EXTENSIONS_GENERATION_EXTENSION_H__
+
+#include <functional>
+#include <map>
+#include <memory>
+#include <set>
+#include <string>
+#include <vector>
+
+#include "conditioning/conditioner.hpp"
+#include "core/ggml_extend_backend.h"
+#include "model_loader.h"
+#include "model_manager.h"
+#include "stable-diffusion.h"
+
+struct GenerationExtensionInitContext {
+    const sd_ctx_params_t* params;
+    SDVersion version;
+    const String2TensorStorage& tensor_storage_map;
+    ModelLoader& model_loader;
+    std::shared_ptr<ModelManager> model_manager;
+    int n_threads;
+    std::function<bool(SDBackendModule)> ensure_backend_pair;
+    std::function<ggml_backend_t(SDBackendModule)> backend_for;
+    std::function<ggml_backend_t(SDBackendModule)> params_backend_for;
+};
+
+struct GenerationExtensionConditionContext {
+    Conditioner* conditioner;
+    ConditionerParams& condition_params;
+    const sd_pm_params_t& pm_params;
+    int n_threads;
+    int total_steps;
+};
+
+struct GenerationExtension {
+    virtual ~GenerationExtension() = default;
+
+    virtual const char* name() const = 0;
+    virtual bool is_enabled() const {
+        return false;
+    }
+    virtual bool init(const GenerationExtensionInitContext&) {
+        return true;
+    }
+    virtual void get_param_tensors(std::map<std::string, ggml_tensor*>&) {}
+    virtual void collect_loras(std::vector<ModelManager::LoraSpec>&) {}
+    virtual void add_ignore_tensors(std::set<std::string>&) const {}
+    virtual void runner_done() {}
+    virtual void reset_runtime_condition() {}
+    virtual bool prepare_condition(GenerationExtensionConditionContext&) {
+        return false;
+    }
+    virtual const SDCondition& before_condition(int step,
+                                                const SDCondition& condition) const {
+        return condition;
+    }
+};
+
+std::shared_ptr<GenerationExtension> create_photomaker_extension();
+
+#endif

src/extensions/photomaker_extension.cpp

@@ -0,0 +1,292 @@
+#include "extensions/generation_extension.h"
+
+#include <algorithm>
+#include <cstring>
+#include <tuple>
+#include <utility>
+
+#include "core/tensor_ggml.hpp"
+#include "core/util.h"
+#include "model/adapter/pmid.hpp"
+
+static std::tuple<std::vector<int>, std::vector<float>, std::vector<bool>>
+tokenize_photomaker_trigger(FrozenCLIPEmbedderWithCustomWords& clip_conditioner,
+                            const std::string& text,
+                            int trigger_token_count,
+                            int32_t image_token) {
+    auto tokens_and_weights           = clip_conditioner.tokenize(text);
+    std::vector<int> source_tokens    = std::move(tokens_and_weights.first);
+    std::vector<float> source_weights = std::move(tokens_and_weights.second);
+
+    if (!source_tokens.empty() && source_tokens.front() == clip_conditioner.tokenizer.BOS_TOKEN_ID) {
+        source_tokens.erase(source_tokens.begin());
+        source_weights.erase(source_weights.begin());
+    }
+    if (!source_tokens.empty() && source_tokens.back() == clip_conditioner.tokenizer.EOS_TOKEN_ID) {
+        source_tokens.pop_back();
+        source_weights.pop_back();
+    }
+
+    std::vector<int> tokens;
+    std::vector<float> weights;
+    int32_t class_idx = -1;
+    for (size_t i = 0; i < source_tokens.size(); i++) {
+        int token = source_tokens[i];
+        if (token == image_token) {
+            if (!tokens.empty()) {
+                class_idx          = static_cast<int32_t>(tokens.size()) - 1;
+                int class_token    = tokens.back();
+                float class_weight = weights.back();
+                for (int j = 1; j < trigger_token_count; j++) {
+                    tokens.push_back(class_token);
+                    weights.push_back(class_weight);
+                }
+            }
+            continue;
+        }
+        tokens.push_back(token);
+        weights.push_back(source_weights[i]);
+    }
+
+    clip_conditioner.tokenizer.pad_tokens(tokens,
+                                          &weights,
+                                          nullptr,
+                                          clip_conditioner.text_model->model.n_token,
+                                          clip_conditioner.text_model->model.n_token,
+                                          true);
+    std::vector<bool> class_token_mask;
+    for (int i = 0; i < tokens.size(); i++) {
+        class_token_mask.push_back(class_idx + 1 <= i && i < class_idx + 1 + trigger_token_count);
+    }
+
+    return std::make_tuple(tokens, weights, class_token_mask);
+}
+
+static std::tuple<SDCondition, std::vector<bool>>
+get_photomaker_condition_with_trigger(FrozenCLIPEmbedderWithCustomWords& clip_conditioner,
+                                      int n_threads,
+                                      const ConditionerParams& conditioner_params,
+                                      const std::string& trigger_word,
+                                      int trigger_token_count) {
+    auto image_tokens = clip_conditioner.convert_token_to_id(trigger_word);
+    GGML_ASSERT(image_tokens.size() == 1);
+    auto tokens_and_weights         = tokenize_photomaker_trigger(clip_conditioner,
+                                                                  conditioner_params.text,
+                                                                  trigger_token_count,
+                                                                  image_tokens[0]);
+    std::vector<int>& tokens        = std::get<0>(tokens_and_weights);
+    std::vector<float>& weights     = std::get<1>(tokens_and_weights);
+    std::vector<bool>& trigger_mask = std::get<2>(tokens_and_weights);
+    auto cond                       = clip_conditioner.get_learned_condition_common(n_threads,
+                                                                                    tokens,
+                                                                                    weights,
+                                                                                    conditioner_params.clip_skip,
+                                                                                    conditioner_params.width,
+                                                                                    conditioner_params.height,
+                                                                                    conditioner_params.zero_out_masked);
+    return std::make_tuple(std::move(cond), trigger_mask);
+}
+
+static std::string remove_photomaker_trigger_from_prompt(FrozenCLIPEmbedderWithCustomWords& clip_conditioner,
+                                                         const std::string& prompt,
+                                                         const std::string& trigger_word) {
+    auto image_tokens = clip_conditioner.convert_token_to_id(trigger_word);
+    GGML_ASSERT(image_tokens.size() == 1);
+    auto tokens_and_weights  = clip_conditioner.tokenize(prompt);
+    std::vector<int>& tokens = tokens_and_weights.first;
+    auto it                  = std::find(tokens.begin(), tokens.end(), image_tokens[0]);
+    GGML_ASSERT(it != tokens.end());
+    tokens.erase(it);
+    return clip_conditioner.decode(tokens);
+}
+
+struct PhotoMakerExtension : public GenerationExtension {
+    std::shared_ptr<PhotoMakerIDEncoder> pmid_model;
+    bool enabled = false;
+    std::string model_path;
+    std::string trigger_word = "img";
+    SDCondition id_condition;
+    int start_merge_step = -1;
+
+    const char* name() const override {
+        return "photomaker";
+    }
+
+    bool is_enabled() const override {
+        return enabled;
+    }
+
+    bool init(const GenerationExtensionInitContext& ctx) override {
+        model_path = SAFE_STR(ctx.params->photo_maker_path);
+        if (model_path.empty()) {
+            return true;
+        }
+
+        if (!ctx.ensure_backend_pair(SDBackendModule::PHOTOMAKER)) {
+            return false;
+        }
+
+        PMVersion pm_version = std::strstr(model_path.c_str(), "v2") != nullptr ? PM_VERSION_2 : PM_VERSION_1;
+        LOG_INFO("loading stacked ID embedding (PHOTOMAKER) model file from '%s'", model_path.c_str());
+        if (!ctx.model_loader.init_from_file_and_convert_name(model_path, "pmid.")) {
+            LOG_WARN("loading stacked ID embedding from '%s' failed", model_path.c_str());
+            return true;
+        }
+
+        pmid_model = std::make_shared<PhotoMakerIDEncoder>(ctx.backend_for(SDBackendModule::PHOTOMAKER),
+                                                           ctx.tensor_storage_map,
+                                                           "pmid",
+                                                           ctx.version,
+                                                           pm_version,
+                                                           20.f,
+                                                           ctx.model_manager);
+        if (pm_version == PM_VERSION_2) {
+            LOG_INFO("using PhotoMaker Version 2");
+        }
+
+        enabled = true;
+        return true;
+    }
+
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        if (!enabled || pmid_model == nullptr) {
+            return;
+        }
+
+        pmid_model->get_param_tensors(tensors, "pmid");
+    }
+
+    void collect_loras(std::vector<ModelManager::LoraSpec>& loras) override {
+        if (!enabled || model_path.empty()) {
+            return;
+        }
+        ModelManager::LoraSpec lora;
+        lora.path                      = model_path;
+        lora.multiplier                = 1.0f;
+        lora.tensor_name_prefix_filter = "lora.model";
+        lora.required                  = true;
+        loras.push_back(std::move(lora));
+    }
+
+    void add_ignore_tensors(std::set<std::string>& ignore_tensors) const override {
+        if (!enabled) {
+            return;
+        }
+        ignore_tensors.insert("pmid.unet.");
+    }
+
+    void runner_done() override {
+        if (pmid_model != nullptr) {
+            pmid_model->runner_done();
+        }
+    }
+
+    void reset_runtime_condition() override {
+        id_condition     = {};
+        start_merge_step = -1;
+    }
+
+    bool prepare_condition(GenerationExtensionConditionContext& ctx) override {
+        reset_runtime_condition();
+        if (!enabled || pmid_model == nullptr) {
+            return false;
+        }
+
+        bool pmv2 = pmid_model->get_version() == PM_VERSION_2;
+        if (ctx.pm_params.id_images_count <= 0 || ctx.pm_params.id_images == nullptr) {
+            LOG_WARN("Provided PhotoMaker model file, but NO input ID images");
+            LOG_WARN("Turn off PhotoMaker for this request");
+            return false;
+        }
+        auto* clip_conditioner = dynamic_cast<FrozenCLIPEmbedderWithCustomWords*>(ctx.conditioner);
+        if (clip_conditioner == nullptr) {
+            LOG_WARN("PhotoMaker requires FrozenCLIPEmbedderWithCustomWords conditioner");
+            LOG_WARN("Turn off PhotoMaker for this request");
+            return false;
+        }
+
+        int clip_image_size        = 224;
+        pmid_model->style_strength = ctx.pm_params.style_strength;
+        sd::Tensor<float> id_image_tensor;
+        for (int i = 0; i < ctx.pm_params.id_images_count; i++) {
+            auto id_image           = sd_image_to_tensor(ctx.pm_params.id_images[i]);
+            auto processed_id_image = clip_preprocess(id_image, clip_image_size, clip_image_size);
+            if (id_image_tensor.empty()) {
+                id_image_tensor = processed_id_image;
+            } else {
+                id_image_tensor = sd::ops::concat(id_image_tensor, processed_id_image, 3);
+            }
+        }
+
+        int64_t t0                        = ggml_time_ms();
+        int trigger_token_count           = pmv2 ? 2 * ctx.pm_params.id_images_count : ctx.pm_params.id_images_count;
+        auto cond_tup                     = get_photomaker_condition_with_trigger(*clip_conditioner,
+                                                                                  ctx.n_threads,
+                                                                                  ctx.condition_params,
+                                                                                  trigger_word,
+                                                                                  trigger_token_count);
+        SDCondition prepared_id_condition = std::get<0>(cond_tup);
+        auto class_tokens_mask            = std::get<1>(cond_tup);
+        if (std::find(class_tokens_mask.begin(), class_tokens_mask.end(), true) == class_tokens_mask.end()) {
+            LOG_WARN("PhotoMaker trigger word '%s' was not found in prompt", trigger_word.c_str());
+            LOG_WARN("Turn off PhotoMaker for this request");
+            return false;
+        }
+
+        sd::Tensor<float> id_embeds;
+        if (pmv2 && ctx.pm_params.id_embed_path != nullptr) {
+            try {
+                id_embeds = sd::load_tensor_from_file_as_tensor<float>(ctx.pm_params.id_embed_path);
+            } catch (const std::exception&) {
+                id_embeds = {};
+            }
+        }
+        if (pmv2 && id_embeds.empty()) {
+            LOG_WARN("Provided PhotoMaker images, but NO valid ID embeds file for PM v2");
+            LOG_WARN("Turn off PhotoMaker for this request");
+            return false;
+        }
+        if (pmv2 && ctx.pm_params.id_images_count != id_embeds.shape()[1]) {
+            LOG_WARN("PhotoMaker image count (%d) does NOT match ID embeds (%d). You should run face_detect.py again.",
+                     ctx.pm_params.id_images_count,
+                     static_cast<int>(id_embeds.shape()[1]));
+            LOG_WARN("Turn off PhotoMaker for this request");
+            return false;
+        }
+
+        auto res = pmid_model->compute(ctx.n_threads,
+                                       id_image_tensor,
+                                       prepared_id_condition.c_crossattn,
+                                       id_embeds,
+                                       class_tokens_mask);
+        if (res.empty()) {
+            LOG_ERROR("Photomaker ID Stacking failed");
+            LOG_WARN("Turn off PhotoMaker for this request");
+            return false;
+        }
+
+        prepared_id_condition.c_crossattn = std::move(res);
+        int64_t t1                        = ggml_time_ms();
+        id_condition                      = std::move(prepared_id_condition);
+        start_merge_step                  = int(ctx.pm_params.style_strength / 100.f * ctx.total_steps);
+        ctx.condition_params.text         = remove_photomaker_trigger_from_prompt(*clip_conditioner,
+                                                                                  ctx.condition_params.text,
+                                                                                  trigger_word);
+        LOG_INFO("Photomaker ID Stacking, taking %" PRId64 " ms", t1 - t0);
+        LOG_INFO("PHOTOMAKER: start_merge_step: %d", start_merge_step);
+
+        return true;
+    }
+
+    const SDCondition& before_condition(int step,
+                                        const SDCondition& condition) const override {
+        if (!id_condition.empty() && start_merge_step != -1 && step > start_merge_step) {
+            return id_condition;
+        }
+        return condition;
+    }
+};
+
+std::shared_ptr<GenerationExtension> create_photomaker_extension() {
+    return std::make_shared<PhotoMakerExtension>();
+}

src/gits_noise.inl

@@ -1,349 +0,0 @@
-#ifndef GITS_NOISE_INL
-#define GITS_NOISE_INL
-
-const std::vector<std::vector<float>> GITS_NOISE_0_80 = {
-    { 14.61464119f, 7.49001646f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 6.77309084f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 7.49001646f, 3.07277966f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 2.05039096f, 0.02916753f },
-    { 14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.85520077f, 2.05039096f, 0.02916753f },
-    { 14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.85520077f, 3.07277966f, 1.56271636f, 0.02916753f },
-    { 14.61464119f, 12.96784878f, 11.54541874f, 8.75849152f, 7.49001646f, 5.85520077f, 3.07277966f, 1.56271636f, 0.02916753f },
-    { 14.61464119f, 13.76078796f, 12.23089790f, 10.90732002f, 8.75849152f, 7.49001646f, 5.85520077f, 3.07277966f, 1.56271636f, 0.02916753f },
-    { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 10.90732002f, 8.75849152f, 7.49001646f, 5.85520077f, 3.07277966f, 1.56271636f, 0.02916753f },
-    { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 5.85520077f, 3.07277966f, 1.56271636f, 0.02916753f },
-    { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.07277966f, 1.56271636f, 0.02916753f },
-    { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.07277966f, 1.56271636f, 0.02916753f },
-    { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.07277966f, 1.56271636f, 0.02916753f },
-    { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.07277966f, 1.56271636f, 0.02916753f },
-    { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.75859547f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.19567990f, 1.98035145f, 0.86115354f, 0.02916753f },
-    { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.75859547f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.19567990f, 1.98035145f, 0.86115354f, 0.02916753f },
-    { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.75859547f, 9.24142551f, 8.75849152f, 8.30717278f, 7.88507891f, 7.49001646f, 6.77309084f, 5.85520077f, 4.65472794f, 3.07277966f, 1.84880662f, 0.83188516f, 0.02916753f }
-};
-
-const std::vector<std::vector<float>> GITS_NOISE_0_85 = {
-    { 14.61464119f, 7.49001646f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 1.84880662f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 6.77309084f, 1.56271636f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 7.11996698f, 3.07277966f, 1.24153244f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 7.49001646f, 5.09240818f, 2.84484982f, 0.95350921f, 0.02916753f },
-    { 14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.09240818f, 2.84484982f, 0.95350921f, 0.02916753f },
-    { 14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.58536053f, 3.19567990f, 1.84880662f, 0.803307f, 0.02916753f },
-    { 14.61464119f, 12.96784878f, 11.54541874f, 8.75849152f, 7.49001646f, 5.58536053f, 3.19567990f, 1.84880662f, 0.803307f, 0.02916753f },
-    { 14.61464119f, 12.96784878f, 11.54541874f, 8.75849152f, 7.49001646f, 6.14220476f, 4.65472794f, 3.07277966f, 1.84880662f, 0.803307f, 0.02916753f },
-    { 14.61464119f, 13.76078796f, 12.23089790f, 10.90732002f, 8.75849152f, 7.49001646f, 6.14220476f, 4.65472794f, 3.07277966f, 1.84880662f, 0.803307f, 0.02916753f },
-    { 14.61464119f, 13.76078796f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.65472794f, 3.07277966f, 1.84880662f, 0.803307f, 0.02916753f },
-    { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.65472794f, 3.07277966f, 1.84880662f, 0.803307f, 0.02916753f },
-    { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.65472794f, 3.07277966f, 1.84880662f, 0.803307f, 0.02916753f },
-    { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.60512662f, 2.63833880f, 1.56271636f, 0.72133851f, 0.02916753f },
-    { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 4.65472794f, 3.46139455f, 2.45070267f, 1.56271636f, 0.72133851f, 0.02916753f },
-    { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 4.65472794f, 3.46139455f, 2.45070267f, 1.56271636f, 0.72133851f, 0.02916753f },
-    { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 4.65472794f, 3.46139455f, 2.45070267f, 1.56271636f, 0.72133851f, 0.02916753f },
-    { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.75859547f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 4.65472794f, 3.46139455f, 2.45070267f, 1.56271636f, 0.72133851f, 0.02916753f },
-    { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.75859547f, 9.24142551f, 8.75849152f, 8.30717278f, 7.88507891f, 7.49001646f, 6.77309084f, 5.85520077f, 4.65472794f, 3.46139455f, 2.45070267f, 1.56271636f, 0.72133851f, 0.02916753f }
-};
-
-const std::vector<std::vector<float>> GITS_NOISE_0_90 = {
-    { 14.61464119f, 6.77309084f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 1.56271636f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 3.07277966f, 0.95350921f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 4.86714602f, 2.54230714f, 0.89115214f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 7.49001646f, 4.86714602f, 2.54230714f, 0.89115214f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 7.49001646f, 5.09240818f, 3.07277966f, 1.61558151f, 0.69515091f, 0.02916753f },
-    { 14.61464119f, 12.23089790f, 8.75849152f, 7.11996698f, 4.86714602f, 3.07277966f, 1.61558151f, 0.69515091f, 0.02916753f },
-    { 14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.85520077f, 4.45427561f, 2.95596409f, 1.61558151f, 0.69515091f, 0.02916753f },
-    { 14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.19988537f, 1.24153244f, 0.57119018f, 0.02916753f },
-    { 14.61464119f, 12.96784878f, 10.90732002f, 8.75849152f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.19988537f, 1.24153244f, 0.57119018f, 0.02916753f },
-    { 14.61464119f, 12.96784878f, 11.54541874f, 9.24142551f, 8.30717278f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.19988537f, 1.24153244f, 0.57119018f, 0.02916753f },
-    { 14.61464119f, 12.96784878f, 11.54541874f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.75677586f, 2.84484982f, 1.84880662f, 1.08895338f, 0.52423614f, 0.02916753f },
-    { 14.61464119f, 13.76078796f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.75677586f, 2.84484982f, 1.84880662f, 1.08895338f, 0.52423614f, 0.02916753f },
-    { 14.61464119f, 13.76078796f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.44769001f, 5.58536053f, 4.45427561f, 3.32507086f, 2.45070267f, 1.61558151f, 0.95350921f, 0.45573691f, 0.02916753f },
-    { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.44769001f, 5.58536053f, 4.45427561f, 3.32507086f, 2.45070267f, 1.61558151f, 0.95350921f, 0.45573691f, 0.02916753f },
-    { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 4.86714602f, 3.91689563f, 3.07277966f, 2.27973175f, 1.56271636f, 0.95350921f, 0.45573691f, 0.02916753f },
-    { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 4.86714602f, 3.91689563f, 3.07277966f, 2.27973175f, 1.56271636f, 0.95350921f, 0.45573691f, 0.02916753f },
-    { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 4.86714602f, 3.91689563f, 3.07277966f, 2.27973175f, 1.56271636f, 0.95350921f, 0.45573691f, 0.02916753f },
-    { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 5.09240818f, 4.45427561f, 3.60512662f, 2.95596409f, 2.19988537f, 1.51179266f, 0.89115214f, 0.43325692f, 0.02916753f }
-};
-
-const std::vector<std::vector<float>> GITS_NOISE_0_95 = {
-    { 14.61464119f, 6.77309084f, 0.02916753f },
-    { 14.61464119f, 6.77309084f, 1.56271636f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 2.84484982f, 0.89115214f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 4.86714602f, 2.36326075f, 0.803307f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 4.86714602f, 2.95596409f, 1.56271636f, 0.64427125f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 7.49001646f, 4.86714602f, 2.95596409f, 1.56271636f, 0.64427125f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 7.49001646f, 4.86714602f, 3.07277966f, 1.91321158f, 1.08895338f, 0.50118381f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.07277966f, 1.91321158f, 1.08895338f, 0.50118381f, 0.02916753f },
-    { 14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.85520077f, 4.45427561f, 3.07277966f, 1.91321158f, 1.08895338f, 0.50118381f, 0.02916753f },
-    { 14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.19988537f, 1.41535246f, 0.803307f, 0.38853383f, 0.02916753f },
-    { 14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.85520077f, 4.65472794f, 3.46139455f, 2.63833880f, 1.84880662f, 1.24153244f, 0.72133851f, 0.34370604f, 0.02916753f },
-    { 14.61464119f, 12.96784878f, 10.90732002f, 8.75849152f, 7.49001646f, 5.85520077f, 4.65472794f, 3.46139455f, 2.63833880f, 1.84880662f, 1.24153244f, 0.72133851f, 0.34370604f, 0.02916753f },
-    { 14.61464119f, 12.96784878f, 10.90732002f, 8.75849152f, 7.49001646f, 6.14220476f, 4.86714602f, 3.75677586f, 2.95596409f, 2.19988537f, 1.56271636f, 1.05362725f, 0.64427125f, 0.32104823f, 0.02916753f },
-    { 14.61464119f, 12.96784878f, 10.90732002f, 8.75849152f, 7.49001646f, 6.44769001f, 5.58536053f, 4.65472794f, 3.60512662f, 2.95596409f, 2.19988537f, 1.56271636f, 1.05362725f, 0.64427125f, 0.32104823f, 0.02916753f },
-    { 14.61464119f, 12.96784878f, 11.54541874f, 9.24142551f, 8.30717278f, 7.49001646f, 6.44769001f, 5.58536053f, 4.65472794f, 3.60512662f, 2.95596409f, 2.19988537f, 1.56271636f, 1.05362725f, 0.64427125f, 0.32104823f, 0.02916753f },
-    { 14.61464119f, 12.96784878f, 11.54541874f, 9.24142551f, 8.30717278f, 7.49001646f, 6.44769001f, 5.58536053f, 4.65472794f, 3.75677586f, 3.07277966f, 2.45070267f, 1.78698075f, 1.24153244f, 0.83188516f, 0.50118381f, 0.22545385f, 0.02916753f },
-    { 14.61464119f, 12.96784878f, 11.54541874f, 9.24142551f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 5.09240818f, 4.45427561f, 3.60512662f, 2.95596409f, 2.36326075f, 1.72759056f, 1.24153244f, 0.83188516f, 0.50118381f, 0.22545385f, 0.02916753f },
-    { 14.61464119f, 13.76078796f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 5.09240818f, 4.45427561f, 3.60512662f, 2.95596409f, 2.36326075f, 1.72759056f, 1.24153244f, 0.83188516f, 0.50118381f, 0.22545385f, 0.02916753f },
-    { 14.61464119f, 13.76078796f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 5.09240818f, 4.45427561f, 3.75677586f, 3.07277966f, 2.45070267f, 1.91321158f, 1.46270394f, 1.05362725f, 0.72133851f, 0.43325692f, 0.19894916f, 0.02916753f }
-};
-
-const std::vector<std::vector<float>> GITS_NOISE_1_00 = {
-    { 14.61464119f, 1.56271636f, 0.02916753f },
-    { 14.61464119f, 6.77309084f, 0.95350921f, 0.02916753f },
-    { 14.61464119f, 6.77309084f, 2.36326075f, 0.803307f, 0.02916753f },
-    { 14.61464119f, 7.11996698f, 3.07277966f, 1.56271636f, 0.59516323f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 4.86714602f, 2.84484982f, 1.41535246f, 0.57119018f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 4.86714602f, 2.84484982f, 1.61558151f, 0.86115354f, 0.38853383f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 7.49001646f, 4.86714602f, 2.84484982f, 1.61558151f, 0.86115354f, 0.38853383f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 7.49001646f, 4.86714602f, 3.07277966f, 1.98035145f, 1.24153244f, 0.72133851f, 0.34370604f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.07277966f, 1.98035145f, 1.24153244f, 0.72133851f, 0.34370604f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.27973175f, 1.51179266f, 0.95350921f, 0.54755926f, 0.25053367f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.36326075f, 1.61558151f, 1.08895338f, 0.72133851f, 0.41087446f, 0.17026083f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 8.75849152f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.36326075f, 1.61558151f, 1.08895338f, 0.72133851f, 0.41087446f, 0.17026083f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 8.75849152f, 7.49001646f, 5.85520077f, 4.65472794f, 3.60512662f, 2.84484982f, 2.12350607f, 1.56271636f, 1.08895338f, 0.72133851f, 0.41087446f, 0.17026083f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 8.75849152f, 7.49001646f, 5.85520077f, 4.65472794f, 3.60512662f, 2.84484982f, 2.19988537f, 1.61558151f, 1.162866f, 0.803307f, 0.50118381f, 0.27464288f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 8.75849152f, 7.49001646f, 5.85520077f, 4.65472794f, 3.75677586f, 3.07277966f, 2.45070267f, 1.84880662f, 1.36964464f, 1.01931262f, 0.72133851f, 0.45573691f, 0.25053367f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 8.75849152f, 7.49001646f, 6.14220476f, 5.09240818f, 4.26497746f, 3.46139455f, 2.84484982f, 2.19988537f, 1.67050016f, 1.24153244f, 0.92192322f, 0.64427125f, 0.43325692f, 0.25053367f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 8.75849152f, 7.49001646f, 6.14220476f, 5.09240818f, 4.26497746f, 3.60512662f, 2.95596409f, 2.45070267f, 1.91321158f, 1.51179266f, 1.12534678f, 0.83188516f, 0.59516323f, 0.38853383f, 0.22545385f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 12.23089790f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 5.09240818f, 4.26497746f, 3.60512662f, 2.95596409f, 2.45070267f, 1.91321158f, 1.51179266f, 1.12534678f, 0.83188516f, 0.59516323f, 0.38853383f, 0.22545385f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 12.23089790f, 9.24142551f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 5.09240818f, 4.26497746f, 3.60512662f, 2.95596409f, 2.45070267f, 1.91321158f, 1.51179266f, 1.12534678f, 0.83188516f, 0.59516323f, 0.38853383f, 0.22545385f, 0.09824532f, 0.02916753f }
-};
-
-const std::vector<std::vector<float>> GITS_NOISE_1_05 = {
-    { 14.61464119f, 0.95350921f, 0.02916753f },
-    { 14.61464119f, 6.77309084f, 0.89115214f, 0.02916753f },
-    { 14.61464119f, 6.77309084f, 2.05039096f, 0.72133851f, 0.02916753f },
-    { 14.61464119f, 6.77309084f, 2.84484982f, 1.28281462f, 0.52423614f, 0.02916753f },
-    { 14.61464119f, 6.77309084f, 3.07277966f, 1.61558151f, 0.803307f, 0.34370604f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 4.86714602f, 2.84484982f, 1.56271636f, 0.803307f, 0.34370604f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 4.86714602f, 2.84484982f, 1.61558151f, 0.95350921f, 0.52423614f, 0.22545385f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 4.86714602f, 3.07277966f, 1.98035145f, 1.24153244f, 0.74807048f, 0.41087446f, 0.17026083f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 4.86714602f, 3.19567990f, 2.27973175f, 1.51179266f, 0.95350921f, 0.59516323f, 0.34370604f, 0.13792117f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 5.09240818f, 3.46139455f, 2.45070267f, 1.61558151f, 1.08895338f, 0.72133851f, 0.45573691f, 0.25053367f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 7.49001646f, 5.09240818f, 3.46139455f, 2.45070267f, 1.61558151f, 1.08895338f, 0.72133851f, 0.45573691f, 0.25053367f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.36326075f, 1.61558151f, 1.08895338f, 0.72133851f, 0.45573691f, 0.25053367f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.45070267f, 1.72759056f, 1.24153244f, 0.86115354f, 0.59516323f, 0.38853383f, 0.22545385f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.65472794f, 3.60512662f, 2.84484982f, 2.19988537f, 1.61558151f, 1.162866f, 0.83188516f, 0.59516323f, 0.38853383f, 0.22545385f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.65472794f, 3.60512662f, 2.84484982f, 2.19988537f, 1.67050016f, 1.28281462f, 0.95350921f, 0.72133851f, 0.52423614f, 0.34370604f, 0.19894916f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.65472794f, 3.60512662f, 2.95596409f, 2.36326075f, 1.84880662f, 1.41535246f, 1.08895338f, 0.83188516f, 0.61951244f, 0.45573691f, 0.32104823f, 0.19894916f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.65472794f, 3.60512662f, 2.95596409f, 2.45070267f, 1.91321158f, 1.51179266f, 1.20157266f, 0.95350921f, 0.74807048f, 0.57119018f, 0.43325692f, 0.29807833f, 0.19894916f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 8.30717278f, 7.11996698f, 5.85520077f, 4.65472794f, 3.60512662f, 2.95596409f, 2.45070267f, 1.91321158f, 1.51179266f, 1.20157266f, 0.95350921f, 0.74807048f, 0.57119018f, 0.43325692f, 0.29807833f, 0.19894916f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 8.30717278f, 7.11996698f, 5.85520077f, 4.65472794f, 3.60512662f, 2.95596409f, 2.45070267f, 1.98035145f, 1.61558151f, 1.32549286f, 1.08895338f, 0.86115354f, 0.69515091f, 0.54755926f, 0.41087446f, 0.29807833f, 0.19894916f, 0.09824532f, 0.02916753f }
-};
-
-const std::vector<std::vector<float>> GITS_NOISE_1_10 = {
-    { 14.61464119f, 0.89115214f, 0.02916753f },
-    { 14.61464119f, 2.36326075f, 0.72133851f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 1.61558151f, 0.57119018f, 0.02916753f },
-    { 14.61464119f, 6.77309084f, 2.45070267f, 1.08895338f, 0.45573691f, 0.02916753f },
-    { 14.61464119f, 6.77309084f, 2.95596409f, 1.56271636f, 0.803307f, 0.34370604f, 0.02916753f },
-    { 14.61464119f, 6.77309084f, 3.07277966f, 1.61558151f, 0.89115214f, 0.4783645f, 0.19894916f, 0.02916753f },
-    { 14.61464119f, 6.77309084f, 3.07277966f, 1.84880662f, 1.08895338f, 0.64427125f, 0.34370604f, 0.13792117f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 4.86714602f, 2.84484982f, 1.61558151f, 0.95350921f, 0.54755926f, 0.27464288f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 4.86714602f, 2.95596409f, 1.91321158f, 1.24153244f, 0.803307f, 0.4783645f, 0.25053367f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 4.86714602f, 3.07277966f, 2.05039096f, 1.41535246f, 0.95350921f, 0.64427125f, 0.41087446f, 0.22545385f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 4.86714602f, 3.19567990f, 2.27973175f, 1.61558151f, 1.12534678f, 0.803307f, 0.54755926f, 0.36617002f, 0.22545385f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 4.86714602f, 3.32507086f, 2.45070267f, 1.72759056f, 1.24153244f, 0.89115214f, 0.64427125f, 0.45573691f, 0.32104823f, 0.19894916f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 5.09240818f, 3.60512662f, 2.84484982f, 2.05039096f, 1.51179266f, 1.08895338f, 0.803307f, 0.59516323f, 0.43325692f, 0.29807833f, 0.19894916f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 5.09240818f, 3.60512662f, 2.84484982f, 2.12350607f, 1.61558151f, 1.24153244f, 0.95350921f, 0.72133851f, 0.54755926f, 0.41087446f, 0.29807833f, 0.19894916f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.45070267f, 1.84880662f, 1.41535246f, 1.08895338f, 0.83188516f, 0.64427125f, 0.50118381f, 0.36617002f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.45070267f, 1.91321158f, 1.51179266f, 1.20157266f, 0.95350921f, 0.74807048f, 0.59516323f, 0.45573691f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 5.85520077f, 4.45427561f, 3.46139455f, 2.84484982f, 2.19988537f, 1.72759056f, 1.36964464f, 1.08895338f, 0.86115354f, 0.69515091f, 0.54755926f, 0.43325692f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.46139455f, 2.84484982f, 2.19988537f, 1.72759056f, 1.36964464f, 1.08895338f, 0.86115354f, 0.69515091f, 0.54755926f, 0.43325692f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.46139455f, 2.84484982f, 2.19988537f, 1.72759056f, 1.36964464f, 1.08895338f, 0.89115214f, 0.72133851f, 0.59516323f, 0.4783645f, 0.38853383f, 0.29807833f, 0.22545385f, 0.17026083f, 0.09824532f, 0.02916753f }
-};
-
-const std::vector<std::vector<float>> GITS_NOISE_1_15 = {
-    { 14.61464119f, 0.83188516f, 0.02916753f },
-    { 14.61464119f, 1.84880662f, 0.59516323f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 1.56271636f, 0.52423614f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 1.91321158f, 0.83188516f, 0.34370604f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.45070267f, 1.24153244f, 0.59516323f, 0.25053367f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.84484982f, 1.51179266f, 0.803307f, 0.41087446f, 0.17026083f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.84484982f, 1.56271636f, 0.89115214f, 0.50118381f, 0.25053367f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 6.77309084f, 3.07277966f, 1.84880662f, 1.12534678f, 0.72133851f, 0.43325692f, 0.22545385f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 6.77309084f, 3.07277966f, 1.91321158f, 1.24153244f, 0.803307f, 0.52423614f, 0.34370604f, 0.19894916f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 4.86714602f, 2.95596409f, 1.91321158f, 1.24153244f, 0.803307f, 0.52423614f, 0.34370604f, 0.19894916f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 4.86714602f, 3.07277966f, 2.05039096f, 1.36964464f, 0.95350921f, 0.69515091f, 0.4783645f, 0.32104823f, 0.19894916f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 4.86714602f, 3.07277966f, 2.12350607f, 1.51179266f, 1.08895338f, 0.803307f, 0.59516323f, 0.43325692f, 0.29807833f, 0.19894916f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 4.86714602f, 3.07277966f, 2.12350607f, 1.51179266f, 1.08895338f, 0.803307f, 0.59516323f, 0.45573691f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 4.86714602f, 3.07277966f, 2.19988537f, 1.61558151f, 1.24153244f, 0.95350921f, 0.74807048f, 0.59516323f, 0.45573691f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 4.86714602f, 3.19567990f, 2.45070267f, 1.78698075f, 1.32549286f, 1.01931262f, 0.803307f, 0.64427125f, 0.50118381f, 0.38853383f, 0.29807833f, 0.22545385f, 0.17026083f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 4.86714602f, 3.19567990f, 2.45070267f, 1.78698075f, 1.32549286f, 1.01931262f, 0.803307f, 0.64427125f, 0.52423614f, 0.41087446f, 0.32104823f, 0.25053367f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 4.86714602f, 3.19567990f, 2.45070267f, 1.84880662f, 1.41535246f, 1.12534678f, 0.89115214f, 0.72133851f, 0.59516323f, 0.4783645f, 0.38853383f, 0.32104823f, 0.25053367f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 4.86714602f, 3.19567990f, 2.45070267f, 1.84880662f, 1.41535246f, 1.12534678f, 0.89115214f, 0.72133851f, 0.59516323f, 0.50118381f, 0.41087446f, 0.34370604f, 0.29807833f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }
-};
-
-const std::vector<std::vector<float>> GITS_NOISE_1_20 = {
-    { 14.61464119f, 0.803307f, 0.02916753f },
-    { 14.61464119f, 1.56271636f, 0.52423614f, 0.02916753f },
-    { 14.61464119f, 2.36326075f, 0.92192322f, 0.36617002f, 0.02916753f },
-    { 14.61464119f, 2.84484982f, 1.24153244f, 0.59516323f, 0.25053367f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.05039096f, 0.95350921f, 0.45573691f, 0.17026083f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.45070267f, 1.24153244f, 0.64427125f, 0.29807833f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.45070267f, 1.36964464f, 0.803307f, 0.45573691f, 0.25053367f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.84484982f, 1.61558151f, 0.95350921f, 0.59516323f, 0.36617002f, 0.19894916f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.84484982f, 1.67050016f, 1.08895338f, 0.74807048f, 0.50118381f, 0.32104823f, 0.19894916f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.95596409f, 1.84880662f, 1.24153244f, 0.83188516f, 0.59516323f, 0.41087446f, 0.27464288f, 0.17026083f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 3.07277966f, 1.98035145f, 1.36964464f, 0.95350921f, 0.69515091f, 0.50118381f, 0.36617002f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 6.77309084f, 3.46139455f, 2.36326075f, 1.56271636f, 1.08895338f, 0.803307f, 0.59516323f, 0.45573691f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 6.77309084f, 3.46139455f, 2.45070267f, 1.61558151f, 1.162866f, 0.86115354f, 0.64427125f, 0.50118381f, 0.38853383f, 0.29807833f, 0.22545385f, 0.17026083f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 4.65472794f, 3.07277966f, 2.12350607f, 1.51179266f, 1.08895338f, 0.83188516f, 0.64427125f, 0.50118381f, 0.38853383f, 0.29807833f, 0.22545385f, 0.17026083f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 4.65472794f, 3.07277966f, 2.12350607f, 1.51179266f, 1.08895338f, 0.83188516f, 0.64427125f, 0.50118381f, 0.41087446f, 0.32104823f, 0.25053367f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 4.65472794f, 3.07277966f, 2.12350607f, 1.51179266f, 1.08895338f, 0.83188516f, 0.64427125f, 0.50118381f, 0.41087446f, 0.34370604f, 0.27464288f, 0.22545385f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 4.65472794f, 3.07277966f, 2.19988537f, 1.61558151f, 1.20157266f, 0.92192322f, 0.72133851f, 0.57119018f, 0.45573691f, 0.36617002f, 0.29807833f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 4.65472794f, 3.07277966f, 2.19988537f, 1.61558151f, 1.24153244f, 0.95350921f, 0.74807048f, 0.59516323f, 0.4783645f, 0.38853383f, 0.32104823f, 0.27464288f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 7.49001646f, 4.65472794f, 3.07277966f, 2.19988537f, 1.61558151f, 1.24153244f, 0.95350921f, 0.74807048f, 0.59516323f, 0.50118381f, 0.41087446f, 0.34370604f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }
-};
-
-const std::vector<std::vector<float>> GITS_NOISE_1_25 = {
-    { 14.61464119f, 0.72133851f, 0.02916753f },
-    { 14.61464119f, 1.56271636f, 0.50118381f, 0.02916753f },
-    { 14.61464119f, 2.05039096f, 0.803307f, 0.32104823f, 0.02916753f },
-    { 14.61464119f, 2.36326075f, 0.95350921f, 0.43325692f, 0.17026083f, 0.02916753f },
-    { 14.61464119f, 2.84484982f, 1.24153244f, 0.59516323f, 0.27464288f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 3.07277966f, 1.51179266f, 0.803307f, 0.43325692f, 0.22545385f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.36326075f, 1.24153244f, 0.72133851f, 0.41087446f, 0.22545385f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.45070267f, 1.36964464f, 0.83188516f, 0.52423614f, 0.34370604f, 0.19894916f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.84484982f, 1.61558151f, 0.98595673f, 0.64427125f, 0.43325692f, 0.27464288f, 0.17026083f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.84484982f, 1.67050016f, 1.08895338f, 0.74807048f, 0.52423614f, 0.36617002f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.84484982f, 1.72759056f, 1.162866f, 0.803307f, 0.59516323f, 0.45573691f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.95596409f, 1.84880662f, 1.24153244f, 0.86115354f, 0.64427125f, 0.4783645f, 0.36617002f, 0.27464288f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.95596409f, 1.84880662f, 1.28281462f, 0.92192322f, 0.69515091f, 0.52423614f, 0.41087446f, 0.32104823f, 0.25053367f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.95596409f, 1.91321158f, 1.32549286f, 0.95350921f, 0.72133851f, 0.54755926f, 0.43325692f, 0.34370604f, 0.27464288f, 0.22545385f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.95596409f, 1.91321158f, 1.32549286f, 0.95350921f, 0.72133851f, 0.57119018f, 0.45573691f, 0.36617002f, 0.29807833f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.95596409f, 1.91321158f, 1.32549286f, 0.95350921f, 0.74807048f, 0.59516323f, 0.4783645f, 0.38853383f, 0.32104823f, 0.27464288f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 3.07277966f, 2.05039096f, 1.41535246f, 1.05362725f, 0.803307f, 0.61951244f, 0.50118381f, 0.41087446f, 0.34370604f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 3.07277966f, 2.05039096f, 1.41535246f, 1.05362725f, 0.803307f, 0.64427125f, 0.52423614f, 0.43325692f, 0.36617002f, 0.32104823f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 3.07277966f, 2.05039096f, 1.46270394f, 1.08895338f, 0.83188516f, 0.66947293f, 0.54755926f, 0.45573691f, 0.38853383f, 0.34370604f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }
-};
-
-const std::vector<std::vector<float>> GITS_NOISE_1_30 = {
-    { 14.61464119f, 0.72133851f, 0.02916753f },
-    { 14.61464119f, 1.24153244f, 0.43325692f, 0.02916753f },
-    { 14.61464119f, 1.56271636f, 0.59516323f, 0.22545385f, 0.02916753f },
-    { 14.61464119f, 1.84880662f, 0.803307f, 0.36617002f, 0.13792117f, 0.02916753f },
-    { 14.61464119f, 2.36326075f, 1.01931262f, 0.52423614f, 0.25053367f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.84484982f, 1.36964464f, 0.74807048f, 0.41087446f, 0.22545385f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 3.07277966f, 1.56271636f, 0.89115214f, 0.54755926f, 0.34370604f, 0.19894916f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 3.07277966f, 1.61558151f, 0.95350921f, 0.61951244f, 0.41087446f, 0.27464288f, 0.17026083f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.45070267f, 1.36964464f, 0.83188516f, 0.54755926f, 0.36617002f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.45070267f, 1.41535246f, 0.92192322f, 0.64427125f, 0.45573691f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.6383388f, 1.56271636f, 1.01931262f, 0.72133851f, 0.50118381f, 0.36617002f, 0.27464288f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.84484982f, 1.61558151f, 1.05362725f, 0.74807048f, 0.54755926f, 0.41087446f, 0.32104823f, 0.25053367f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.84484982f, 1.61558151f, 1.08895338f, 0.77538133f, 0.57119018f, 0.43325692f, 0.34370604f, 0.27464288f, 0.22545385f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.84484982f, 1.61558151f, 1.08895338f, 0.803307f, 0.59516323f, 0.45573691f, 0.36617002f, 0.29807833f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.84484982f, 1.61558151f, 1.08895338f, 0.803307f, 0.59516323f, 0.4783645f, 0.38853383f, 0.32104823f, 0.27464288f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.84484982f, 1.72759056f, 1.162866f, 0.83188516f, 0.64427125f, 0.50118381f, 0.41087446f, 0.34370604f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.84484982f, 1.72759056f, 1.162866f, 0.83188516f, 0.64427125f, 0.52423614f, 0.43325692f, 0.36617002f, 0.32104823f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.84484982f, 1.78698075f, 1.24153244f, 0.92192322f, 0.72133851f, 0.57119018f, 0.45573691f, 0.38853383f, 0.34370604f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.84484982f, 1.78698075f, 1.24153244f, 0.92192322f, 0.72133851f, 0.57119018f, 0.4783645f, 0.41087446f, 0.36617002f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }
-};
-
-const std::vector<std::vector<float>> GITS_NOISE_1_35 = {
-    { 14.61464119f, 0.69515091f, 0.02916753f },
-    { 14.61464119f, 0.95350921f, 0.34370604f, 0.02916753f },
-    { 14.61464119f, 1.56271636f, 0.57119018f, 0.19894916f, 0.02916753f },
-    { 14.61464119f, 1.61558151f, 0.69515091f, 0.29807833f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 1.84880662f, 0.83188516f, 0.43325692f, 0.22545385f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.45070267f, 1.162866f, 0.64427125f, 0.36617002f, 0.19894916f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.84484982f, 1.36964464f, 0.803307f, 0.50118381f, 0.32104823f, 0.19894916f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.84484982f, 1.41535246f, 0.83188516f, 0.54755926f, 0.36617002f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.84484982f, 1.56271636f, 0.95350921f, 0.64427125f, 0.45573691f, 0.32104823f, 0.22545385f, 0.17026083f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.84484982f, 1.56271636f, 0.95350921f, 0.64427125f, 0.45573691f, 0.34370604f, 0.25053367f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 3.07277966f, 1.61558151f, 1.01931262f, 0.72133851f, 0.52423614f, 0.38853383f, 0.29807833f, 0.22545385f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 3.07277966f, 1.61558151f, 1.01931262f, 0.72133851f, 0.52423614f, 0.41087446f, 0.32104823f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 3.07277966f, 1.61558151f, 1.05362725f, 0.74807048f, 0.54755926f, 0.43325692f, 0.34370604f, 0.27464288f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 3.07277966f, 1.72759056f, 1.12534678f, 0.803307f, 0.59516323f, 0.45573691f, 0.36617002f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 3.07277966f, 1.72759056f, 1.12534678f, 0.803307f, 0.59516323f, 0.4783645f, 0.38853383f, 0.32104823f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.45070267f, 1.51179266f, 1.01931262f, 0.74807048f, 0.57119018f, 0.45573691f, 0.36617002f, 0.32104823f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.6383388f, 1.61558151f, 1.08895338f, 0.803307f, 0.61951244f, 0.50118381f, 0.41087446f, 0.34370604f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.6383388f, 1.61558151f, 1.08895338f, 0.803307f, 0.64427125f, 0.52423614f, 0.43325692f, 0.36617002f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 5.85520077f, 2.6383388f, 1.61558151f, 1.08895338f, 0.803307f, 0.64427125f, 0.52423614f, 0.45573691f, 0.38853383f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }
-};
-
-const std::vector<std::vector<float>> GITS_NOISE_1_40 = {
-    { 14.61464119f, 0.59516323f, 0.02916753f },
-    { 14.61464119f, 0.95350921f, 0.34370604f, 0.02916753f },
-    { 14.61464119f, 1.08895338f, 0.43325692f, 0.13792117f, 0.02916753f },
-    { 14.61464119f, 1.56271636f, 0.64427125f, 0.27464288f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 1.61558151f, 0.803307f, 0.43325692f, 0.22545385f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.05039096f, 0.95350921f, 0.54755926f, 0.34370604f, 0.19894916f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.45070267f, 1.24153244f, 0.72133851f, 0.43325692f, 0.27464288f, 0.17026083f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.45070267f, 1.24153244f, 0.74807048f, 0.50118381f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.45070267f, 1.28281462f, 0.803307f, 0.52423614f, 0.36617002f, 0.27464288f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.45070267f, 1.28281462f, 0.803307f, 0.54755926f, 0.38853383f, 0.29807833f, 0.22545385f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.84484982f, 1.41535246f, 0.86115354f, 0.59516323f, 0.43325692f, 0.32104823f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.84484982f, 1.51179266f, 0.95350921f, 0.64427125f, 0.45573691f, 0.34370604f, 0.27464288f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.84484982f, 1.51179266f, 0.95350921f, 0.64427125f, 0.4783645f, 0.36617002f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.84484982f, 1.56271636f, 0.98595673f, 0.69515091f, 0.52423614f, 0.41087446f, 0.34370604f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.84484982f, 1.56271636f, 1.01931262f, 0.72133851f, 0.54755926f, 0.43325692f, 0.36617002f, 0.32104823f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.84484982f, 1.61558151f, 1.05362725f, 0.74807048f, 0.57119018f, 0.45573691f, 0.38853383f, 0.34370604f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.84484982f, 1.61558151f, 1.08895338f, 0.803307f, 0.61951244f, 0.50118381f, 0.41087446f, 0.36617002f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.84484982f, 1.61558151f, 1.08895338f, 0.803307f, 0.61951244f, 0.50118381f, 0.43325692f, 0.38853383f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.84484982f, 1.61558151f, 1.08895338f, 0.803307f, 0.64427125f, 0.52423614f, 0.45573691f, 0.41087446f, 0.36617002f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }
-};
-
-const std::vector<std::vector<float>> GITS_NOISE_1_45 = {
-    { 14.61464119f, 0.59516323f, 0.02916753f },
-    { 14.61464119f, 0.803307f, 0.25053367f, 0.02916753f },
-    { 14.61464119f, 0.95350921f, 0.34370604f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 1.24153244f, 0.54755926f, 0.25053367f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 1.56271636f, 0.72133851f, 0.36617002f, 0.19894916f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 1.61558151f, 0.803307f, 0.45573691f, 0.27464288f, 0.17026083f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 1.91321158f, 0.95350921f, 0.57119018f, 0.36617002f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.19988537f, 1.08895338f, 0.64427125f, 0.41087446f, 0.27464288f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.45070267f, 1.24153244f, 0.74807048f, 0.50118381f, 0.34370604f, 0.25053367f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.45070267f, 1.24153244f, 0.74807048f, 0.50118381f, 0.36617002f, 0.27464288f, 0.22545385f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.45070267f, 1.28281462f, 0.803307f, 0.54755926f, 0.41087446f, 0.32104823f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.45070267f, 1.28281462f, 0.803307f, 0.57119018f, 0.43325692f, 0.34370604f, 0.27464288f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.45070267f, 1.28281462f, 0.83188516f, 0.59516323f, 0.45573691f, 0.36617002f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.45070267f, 1.28281462f, 0.83188516f, 0.59516323f, 0.45573691f, 0.36617002f, 0.32104823f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.84484982f, 1.51179266f, 0.95350921f, 0.69515091f, 0.52423614f, 0.41087446f, 0.34370604f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.84484982f, 1.51179266f, 0.95350921f, 0.69515091f, 0.52423614f, 0.43325692f, 0.36617002f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.84484982f, 1.56271636f, 0.98595673f, 0.72133851f, 0.54755926f, 0.45573691f, 0.38853383f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.84484982f, 1.56271636f, 1.01931262f, 0.74807048f, 0.57119018f, 0.4783645f, 0.41087446f, 0.36617002f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.84484982f, 1.56271636f, 1.01931262f, 0.74807048f, 0.59516323f, 0.50118381f, 0.43325692f, 0.38853383f, 0.36617002f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }
-};
-
-const std::vector<std::vector<float>> GITS_NOISE_1_50 = {
-    { 14.61464119f, 0.54755926f, 0.02916753f },
-    { 14.61464119f, 0.803307f, 0.25053367f, 0.02916753f },
-    { 14.61464119f, 0.86115354f, 0.32104823f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 1.24153244f, 0.54755926f, 0.25053367f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 1.56271636f, 0.72133851f, 0.36617002f, 0.19894916f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 1.61558151f, 0.803307f, 0.45573691f, 0.27464288f, 0.17026083f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 1.61558151f, 0.83188516f, 0.52423614f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 1.84880662f, 0.95350921f, 0.59516323f, 0.38853383f, 0.27464288f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 1.84880662f, 0.95350921f, 0.59516323f, 0.41087446f, 0.29807833f, 0.22545385f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 1.84880662f, 0.95350921f, 0.61951244f, 0.43325692f, 0.32104823f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.19988537f, 1.12534678f, 0.72133851f, 0.50118381f, 0.36617002f, 0.27464288f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.19988537f, 1.12534678f, 0.72133851f, 0.50118381f, 0.36617002f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.36326075f, 1.24153244f, 0.803307f, 0.57119018f, 0.43325692f, 0.34370604f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.36326075f, 1.24153244f, 0.803307f, 0.57119018f, 0.43325692f, 0.34370604f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.36326075f, 1.24153244f, 0.803307f, 0.59516323f, 0.45573691f, 0.36617002f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.36326075f, 1.24153244f, 0.803307f, 0.59516323f, 0.45573691f, 0.38853383f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.45070267f, 1.32549286f, 0.86115354f, 0.64427125f, 0.50118381f, 0.41087446f, 0.36617002f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.45070267f, 1.36964464f, 0.92192322f, 0.69515091f, 0.54755926f, 0.45573691f, 0.41087446f, 0.36617002f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f },
-    { 14.61464119f, 2.45070267f, 1.41535246f, 0.95350921f, 0.72133851f, 0.57119018f, 0.4783645f, 0.43325692f, 0.38853383f, 0.36617002f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }
-};
-
-const std::vector<const std::vector<std::vector<float>>*> GITS_NOISE = {
-    &GITS_NOISE_0_80,
-    &GITS_NOISE_0_85,
-    &GITS_NOISE_0_90,
-    &GITS_NOISE_0_95,
-    &GITS_NOISE_1_00,
-    &GITS_NOISE_1_05,
-    &GITS_NOISE_1_10,
-    &GITS_NOISE_1_15,
-    &GITS_NOISE_1_20,
-    &GITS_NOISE_1_25,
-    &GITS_NOISE_1_30,
-    &GITS_NOISE_1_35,
-    &GITS_NOISE_1_40,
-    &GITS_NOISE_1_45,
-    &GITS_NOISE_1_50
-};
-
-#endif // GITS_NOISE_INL

src/layer_registry.cpp

@@ -1,132 +0,0 @@
-#include "layer_registry.h"
-
-#include <utility>
-
-#include "util.h"
-
-namespace sd::layer_registry {
-
-    void LayerRegistry::register_layer(const std::string& name, ggml_tensor* tensor) {
-        auto& info = layers_[name];
-        info.tensors.push_back(tensor);
-        info.bytes += ggml_nbytes(tensor);
-    }
-
-    bool LayerRegistry::move_layer_to_gpu(const std::string& name) {
-        auto it = layers_.find(name);
-        if (it == layers_.end())
-            return false;
-
-        LayerInfo& info = it->second;
-        if (info.on_gpu)
-            return true;
-        if (gpu_backend_ == nullptr || cpu_backend_ == nullptr) {
-            LOG_ERROR("layer_registry: backends not set; cannot move '%s' to GPU",
-                      name.c_str());
-            return false;
-        }
-        if (info.tensors.empty()) {
-            info.on_gpu = true;
-            return true;
-        }
-
-        // 1. Build a no_alloc context big enough to hold one twin tensor per CPU
-        //    tensor, plus a little overhead.
-        const size_t ctx_size = info.tensors.size() * ggml_tensor_overhead() + 1024;
-        ggml_init_params ctx_params{ctx_size, /*mem_buffer=*/nullptr, /*no_alloc=*/true};
-        ggml_context* twin_ctx = ggml_init(ctx_params);
-        if (twin_ctx == nullptr) {
-            LOG_ERROR("layer_registry: failed to allocate twin context for '%s'",
-                      name.c_str());
-            return false;
-        }
-
-        // 2. Create one GPU twin per CPU tensor. The twin shares the original
-        //    name so any name-based lookup keeps working.
-        std::vector<ggml_tensor*> gpu_twins;
-        gpu_twins.reserve(info.tensors.size());
-        for (ggml_tensor* cpu_t : info.tensors) {
-            ggml_tensor* twin = ggml_dup_tensor(twin_ctx, cpu_t);
-            if (cpu_t->name[0] != '\0') {
-                ggml_set_name(twin, cpu_t->name);
-            }
-            gpu_twins.push_back(twin);
-        }
-
-        // 3. Back the twins with a GPU buffer in one alloc call.
-        ggml_backend_buffer_t gpu_buffer = ggml_backend_alloc_ctx_tensors(twin_ctx, gpu_backend_);
-        if (gpu_buffer == nullptr) {
-            LOG_ERROR("layer_registry: failed to allocate GPU buffer for '%s'",
-                      name.c_str());
-            ggml_free(twin_ctx);
-            return false;
-        }
-
-        // 4. H2D copy + sync.
-        for (size_t i = 0; i < info.tensors.size(); ++i) {
-            ggml_backend_tensor_copy(info.tensors[i], gpu_twins[i]);
-        }
-        ggml_backend_synchronize(gpu_backend_);
-
-        // 5. Swap buffer/data/extra so the originals now point at GPU memory.
-        for (size_t i = 0; i < info.tensors.size(); ++i) {
-            std::swap(info.tensors[i]->buffer, gpu_twins[i]->buffer);
-            std::swap(info.tensors[i]->data, gpu_twins[i]->data);
-            std::swap(info.tensors[i]->extra, gpu_twins[i]->extra);
-        }
-
-        info.gpu_twins  = std::move(gpu_twins);
-        info.twin_ctx   = twin_ctx;
-        info.gpu_buffer = gpu_buffer;
-        info.on_gpu     = true;
-        return true;
-    }
-
-    bool LayerRegistry::move_layer_to_cpu(const std::string& name) {
-        auto it = layers_.find(name);
-        if (it == layers_.end())
-            return false;
-
-        LayerInfo& info = it->second;
-        if (!info.on_gpu)
-            return true;
-        if (info.tensors.size() != info.gpu_twins.size()) {
-            LOG_ERROR("layer_registry: twin/tensor count mismatch for '%s'",
-                      name.c_str());
-            return false;
-        }
-
-        // 1. Swap back: originals point at CPU memory again.
-        for (size_t i = 0; i < info.tensors.size(); ++i) {
-            if (info.gpu_twins[i] == nullptr)
-                continue;
-            std::swap(info.tensors[i]->buffer, info.gpu_twins[i]->buffer);
-            std::swap(info.tensors[i]->data, info.gpu_twins[i]->data);
-            std::swap(info.tensors[i]->extra, info.gpu_twins[i]->extra);
-        }
-
-        // 2. Free the GPU buffer + twin context.
-        if (info.gpu_buffer != nullptr) {
-            ggml_backend_buffer_free(info.gpu_buffer);
-            info.gpu_buffer = nullptr;
-        }
-        if (info.twin_ctx != nullptr) {
-            ggml_free(info.twin_ctx);
-            info.twin_ctx = nullptr;
-        }
-        info.gpu_twins.clear();
-        info.on_gpu = false;
-        return true;
-    }
-
-    bool LayerRegistry::is_layer_on_gpu(const std::string& name) const {
-        auto it = layers_.find(name);
-        return it != layers_.end() && it->second.on_gpu;
-    }
-
-    size_t LayerRegistry::get_layer_size(const std::string& name) const {
-        auto it = layers_.find(name);
-        return it != layers_.end() ? it->second.bytes : 0;
-    }
-
-}  // namespace sd::layer_registry

src/layer_registry.h

@@ -1,50 +0,0 @@
-#ifndef __LAYER_REGISTRY_H__
-#define __LAYER_REGISTRY_H__
-
-#include <map>
-#include <set>
-#include <string>
-#include <vector>
-
-#include "ggml-backend.h"
-#include "ggml.h"
-
-namespace sd::layer_registry {
-
-    struct LayerInfo {
-        std::vector<ggml_tensor*> tensors;
-        std::vector<ggml_tensor*> gpu_twins;
-        ggml_context* twin_ctx           = nullptr;
-        ggml_backend_buffer_t gpu_buffer = nullptr;
-        bool on_gpu                      = false;
-        size_t bytes                     = 0;
-    };
-
-    class LayerRegistry {
-    public:
-        LayerRegistry() = default;
-        LayerRegistry(ggml_backend_t gpu_backend, ggml_backend_t cpu_backend)
-            : gpu_backend_(gpu_backend), cpu_backend_(cpu_backend) {}
-
-        void set_backends(ggml_backend_t gpu_backend, ggml_backend_t cpu_backend) {
-            gpu_backend_ = gpu_backend;
-            cpu_backend_ = cpu_backend;
-        }
-        void register_layer(const std::string& name, ggml_tensor* tensor);
-        bool move_layer_to_gpu(const std::string& name);
-        bool move_layer_to_cpu(const std::string& name);
-        bool is_layer_on_gpu(const std::string& name) const;
-        size_t get_layer_size(const std::string& name) const;
-        size_t get_layer_count() const { return layers_.size(); }
-
-        const std::map<std::string, LayerInfo>& layers() const { return layers_; }
-
-    private:
-        ggml_backend_t gpu_backend_ = nullptr;
-        ggml_backend_t cpu_backend_ = nullptr;
-        std::map<std::string, LayerInfo> layers_;
-    };
-
-}  // namespace sd::layer_registry
-
-#endif

src/model.h

@@ -1,17 +1,14 @@
 #ifndef __MODEL_H__
 #define __MODEL_H__
 
-#include <functional>
-#include <map>
-#include <memory>
-#include <set>
 #include <string>
+#include <utility>
 #include <vector>
 
+#include "core/ordered_map.hpp"
 #include "ggml-backend.h"
 #include "ggml.h"
 #include "model_io/tensor_storage.h"
-#include "ordered_map.hpp"
 
 enum SDVersion {
     VERSION_SD1,
@@ -50,6 +47,8 @@ enum SDVersion {
     VERSION_LENS,
     VERSION_LONGCAT,
     VERSION_PID,
+    VERSION_IDEOGRAM4,
+    VERSION_ESRGAN,
     VERSION_COUNT,
 };
 
@@ -172,8 +171,15 @@ static inline bool sd_version_is_pid(SDVersion version) {
     return false;
 }
 
+static inline bool sd_version_is_ideogram4(SDVersion version) {
+    if (version == VERSION_IDEOGRAM4) {
+        return true;
+    }
+    return false;
+}
+
 static inline bool sd_version_uses_flux2_vae(SDVersion version) {
-    if (sd_version_is_flux2(version) || sd_version_is_ernie_image(version) || sd_version_is_lens(version)) {
+    if (sd_version_is_flux2(version) || sd_version_is_ernie_image(version) || sd_version_is_lens(version) || sd_version_is_ideogram4(version)) {
         return true;
     }
     return false;
@@ -203,7 +209,8 @@ static inline bool sd_version_is_dit(SDVersion version) {
         sd_version_is_ernie_image(version) ||
         sd_version_is_lens(version) ||
         sd_version_is_longcat(version) ||
-        sd_version_is_pid(version)) {
+        sd_version_is_pid(version) ||
+        sd_version_is_ideogram4(version)) {
         return true;
     }
     return false;
@@ -229,73 +236,4 @@ enum PMVersion {
 typedef OrderedMap<std::string, TensorStorage> String2TensorStorage;
 using TensorTypeRules = std::vector<std::pair<std::string, ggml_type>>;
 
-TensorTypeRules parse_tensor_type_rules(const std::string& tensor_type_rules);
-
-class MmapWrapper;
-
-struct ModelFileData {
-    std::string path;
-    std::vector<TensorStorage> tensors;
-    std::shared_ptr<MmapWrapper> mmapped;
-    std::shared_ptr<struct ggml_backend_buffer> mmbuffer;
-    bool is_zip;
-};
-
-struct MmapTensorStore {
-    std::shared_ptr<MmapWrapper> mmapped;
-    std::shared_ptr<struct ggml_backend_buffer> mmbuffer;
-};
-
-class ModelLoader {
-protected:
-    SDVersion version_ = VERSION_COUNT;
-    std::vector<std::string> file_paths_;
-    std::vector<ModelFileData> file_data;
-    bool model_files_processed = false;
-    String2TensorStorage tensor_storage_map;
-
-    void add_tensor_storage(const TensorStorage& tensor_storage);
-
-    bool init_from_gguf_file(const std::string& file_path, const std::string& prefix = "");
-    bool init_from_safetensors_file(const std::string& file_path, const std::string& prefix = "");
-    bool init_from_torch_zip_file(const std::string& file_path, const std::string& prefix = "");
-    bool init_from_torch_legacy_file(const std::string& file_path, const std::string& prefix = "");
-    bool init_from_diffusers_file(const std::string& file_path, const std::string& prefix = "");
-
-public:
-    bool init_from_file(const std::string& file_path, const std::string& prefix = "");
-    void convert_tensors_name();
-    bool init_from_file_and_convert_name(const std::string& file_path,
-                                         const std::string& prefix = "",
-                                         SDVersion version         = VERSION_COUNT);
-    SDVersion get_sd_version();
-    std::map<ggml_type, uint32_t> get_wtype_stat();
-    std::map<ggml_type, uint32_t> get_conditioner_wtype_stat();
-    std::map<ggml_type, uint32_t> get_diffusion_model_wtype_stat();
-    std::map<ggml_type, uint32_t> get_vae_wtype_stat();
-    String2TensorStorage& get_tensor_storage_map() { return tensor_storage_map; }
-    void set_wtype_override(ggml_type wtype, std::string tensor_type_rules = "");
-    void process_model_files(bool enable_mmap = false, bool writable_mmap = true);
-    std::vector<MmapTensorStore> mmap_tensors(std::map<std::string, ggml_tensor*>& tensors,
-                                              std::set<std::string> ignore_tensors = {},
-                                              bool writable                        = true);
-    bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_threads = 0, bool use_mmap = false);
-    bool load_tensors(std::map<std::string, ggml_tensor*>& tensors,
-                      std::set<std::string> ignore_tensors = {},
-                      int n_threads                        = 0,
-                      bool use_mmap                        = false);
-
-    std::vector<std::string> get_tensor_names() const {
-        std::vector<std::string> names;
-        for (const auto& [name, tensor_storage] : tensor_storage_map) {
-            names.push_back(name);
-        }
-        return names;
-    }
-
-    bool tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type);
-    int64_t get_params_mem_size(ggml_backend_t backend, ggml_type type = GGML_TYPE_COUNT);
-    ~ModelLoader() = default;
-};
-
 #endif  // __MODEL_H__

src/model/adapter/lora.hpp

@@ -1,8 +1,10 @@
-#ifndef __LORA_HPP__
-#define __LORA_HPP__
+#ifndef __SD_MODEL_ADAPTER_LORA_HPP__
+#define __SD_MODEL_ADAPTER_LORA_HPP__
 
 #include <mutex>
-#include "ggml_extend.hpp"
+#include "core/ggml_extend.hpp"
+#include "model_loader.h"
+#include "model_manager.h"
 
 #define LORA_GRAPH_BASE_SIZE 10240
 
@@ -13,22 +15,24 @@ struct LoraModel : public GGMLRunner {
     std::map<ggml_tensor*, ggml_tensor*> original_tensor_to_final_tensor;
     std::set<std::string> applied_lora_tensors;
     std::string file_path;
-    ModelLoader model_loader;
-    bool load_failed         = false;
-    bool applied             = false;
-    bool tensor_preprocessed = false;
+    std::shared_ptr<ModelManager> model_manager;
+    ggml_backend_t params_backend = nullptr;
+    bool load_failed              = false;
+    bool applied                  = false;
+    bool tensor_preprocessed      = false;
 
     typedef std::function<bool(const std::string&)> filter_t;
 
     LoraModel(const std::string& lora_id,
               ggml_backend_t backend,
-              ggml_backend_t params_backend,
-              const std::string& file_path = "",
-              std::string prefix           = "",
-              SDVersion version            = VERSION_COUNT)
-        : lora_id(lora_id), file_path(file_path), GGMLRunner(backend, params_backend) {
+              ggml_backend_t params_backend_,
+              const std::string& file_path          = "",
+              std::string prefix                    = "",
+              SDVersion version                     = VERSION_COUNT,
+              std::shared_ptr<ModelManager> manager = std::make_shared<ModelManager>())
+        : GGMLRunner(backend, manager), lora_id(lora_id), file_path(file_path), model_manager(std::move(manager)), params_backend(params_backend_) {
         prefix = "lora." + prefix;
-        if (!model_loader.init_from_file_and_convert_name(file_path, prefix, version)) {
+        if (model_manager == nullptr || !model_manager->loader().init_from_file_and_convert_name(file_path, prefix, version)) {
             load_failed = true;
         }
     }
@@ -70,7 +74,11 @@ struct LoraModel : public GGMLRunner {
             return true;
         };
 
-        model_loader.load_tensors(on_new_tensor_cb, n_threads);
+        if (model_manager != nullptr) {
+            model_manager->set_n_threads(n_threads);
+        }
+        ModelLoader& model_loader = model_manager->loader();
+        model_loader.load_tensors(on_new_tensor_cb);
 
         if (tensors_to_create.empty()) {
             return true;
@@ -86,25 +94,64 @@ struct LoraModel : public GGMLRunner {
             lora_tensors[name] = real;
         }
 
-        if (!alloc_params_buffer()) {
-            LOG_ERROR("lora model buffer allocation failed");
+        std::map<std::string, ggml_tensor*> tensors;
+        for (const auto& pair : lora_tensors) {
+            tensors[pair.first] = pair.second;
+        }
+        if (model_manager == nullptr ||
+            !model_manager->register_param_tensors("LoRA",
+                                                   std::move(tensors),
+                                                   ModelManager::ResidencyMode::ParamBackend,
+                                                   runtime_backend,
+                                                   params_backend) ||
+            !model_manager->validate_registered_tensors()) {
+            LOG_ERROR("lora model manager registration failed");
+            return false;
+        }
+        std::vector<ggml_tensor*> lora_params;
+        lora_params.reserve(lora_tensors.size());
+        for (const auto& pair : lora_tensors) {
+            lora_params.push_back(pair.second);
+        }
+        if (!model_manager->prepare_params(lora_params)) {
+            LOG_ERROR("lora model manager prepare params failed");
             return false;
         }
-
-        dry_run = false;
-        model_loader.load_tensors(on_new_tensor_cb, n_threads);
 
         LOG_DEBUG("finished loaded lora");
         return true;
     }
 
-    void preprocess_lora_tensors(const std::map<std::string, ggml_tensor*>& model_tensors) {
+    void release_loaded_tensors() {
+        runner_done();
+        free_compute_buffer();
+        model_manager.reset();
+        free_params_ctx();
+        alloc_params_ctx();
+        model_manager  = std::make_shared<ModelManager>();
+        weight_manager = model_manager;
+        lora_tensors.clear();
+        original_tensor_to_final_tensor.clear();
+        applied_lora_tensors.clear();
+        applied             = false;
+        tensor_preprocessed = false;
+    }
+
+    static std::set<std::string> tensor_names(const std::map<std::string, ggml_tensor*>& model_tensors) {
+        std::set<std::string> names;
+        for (const auto& item : model_tensors) {
+            names.insert(item.first);
+        }
+        return names;
+    }
+
+    void preprocess_lora_tensors(const std::set<std::string>& model_tensor_names) {
         if (tensor_preprocessed) {
             return;
         }
         tensor_preprocessed = true;
         // I really hate these hardcoded processes.
-        if (model_tensors.find("cond_stage_model.1.transformer.text_model.encoder.layers.0.self_attn.in_proj.weight") != model_tensors.end()) {
+        if (model_tensor_names.find("cond_stage_model.1.transformer.text_model.encoder.layers.0.self_attn.in_proj.weight") != model_tensor_names.end()) {
             std::unordered_map<std::string, ggml_tensor*> new_lora_tensors;
             for (auto& [old_name, tensor] : lora_tensors) {
                 std::string new_name = old_name;
@@ -611,7 +658,7 @@ struct LoraModel : public GGMLRunner {
                 if (lokr_w2)
                     applied_lora_tensors.insert(lokr_w2_name);
                 if (lokr_w2_a)
-                    applied_lora_tensors.insert(lokr_w2_name);
+                    applied_lora_tensors.insert(lokr_w2_a_name);
                 if (lokr_w2_b)
                     applied_lora_tensors.insert(lokr_w2_b_name);
                 applied_lora_tensors.insert(alpha_name);
@@ -752,11 +799,13 @@ struct LoraModel : public GGMLRunner {
         return out_diff;
     }
 
-    ggml_cgraph* build_lora_graph(const std::map<std::string, ggml_tensor*>& model_tensors, SDVersion version) {
+    ggml_cgraph* build_lora_graph(const std::map<std::string, ggml_tensor*>& model_tensors,
+                                  const std::set<std::string>& model_tensor_names,
+                                  SDVersion version) {
         size_t lora_graph_size = LORA_GRAPH_BASE_SIZE + lora_tensors.size() * 10;
         ggml_cgraph* gf        = ggml_new_graph_custom(compute_ctx, lora_graph_size, false);
 
-        preprocess_lora_tensors(model_tensors);
+        preprocess_lora_tensors(model_tensor_names);
 
         original_tensor_to_final_tensor.clear();
         applied_lora_tensors.clear();
@@ -793,12 +842,16 @@ struct LoraModel : public GGMLRunner {
         return gf;
     }
 
-    void apply(std::map<std::string, ggml_tensor*> model_tensors, SDVersion version, int n_threads) {
+    void apply(std::map<std::string, ggml_tensor*> model_tensors,
+               const std::set<std::string>& model_tensor_names,
+               SDVersion version,
+               int n_threads,
+               bool warn_unused = true) {
         auto get_graph = [&]() -> ggml_cgraph* {
-            return build_lora_graph(model_tensors, version);
+            return build_lora_graph(model_tensors, model_tensor_names, version);
         };
-        GGMLRunner::compute<float>(get_graph, n_threads, false, true);
-        stat();
+        GGMLRunner::compute<float>(get_graph, n_threads, false, false, false, true);
+        stat(!warn_unused);
         for (auto item : original_tensor_to_final_tensor) {
             ggml_tensor* original_tensor = item.first;
             ggml_tensor* final_tensor    = item.second;
@@ -809,6 +862,10 @@ struct LoraModel : public GGMLRunner {
         GGMLRunner::free_compute_buffer();
     }
 
+    void apply(std::map<std::string, ggml_tensor*> model_tensors, SDVersion version, int n_threads, bool warn_unused = true) {
+        apply(model_tensors, tensor_names(model_tensors), version, n_threads, warn_unused);
+    }
+
     void stat(bool at_runntime = false) {
         size_t total_lora_tensors_count   = 0;
         size_t applied_lora_tensors_count = 0;
@@ -914,4 +971,4 @@ public:
     }
 };
 
-#endif  // __LORA_HPP__
+#endif  // __SD_MODEL_ADAPTER_LORA_HPP__

src/model/adapter/pmid.hpp

@@ -1,10 +1,12 @@
-#ifndef __PMI_HPP__
-#define __PMI_HPP__
+#ifndef __SD_MODEL_ADAPTER_PMID_HPP__
+#define __SD_MODEL_ADAPTER_PMID_HPP__
 
-#include "ggml_extend.hpp"
+#include "core/ggml_extend.hpp"
 
-#include "clip.hpp"
-#include "lora.hpp"
+#include "model/adapter/lora.hpp"
+#include "model/common/block.hpp"
+#include "model/te/clip.hpp"
+#include "model_loader.h"
 
 struct FuseBlock : public GGMLBlock {
     // network hparams
@@ -411,13 +413,13 @@ public:
 
 public:
     PhotoMakerIDEncoder(ggml_backend_t backend,
-                        ggml_backend_t params_backend,
                         const String2TensorStorage& tensor_storage_map,
                         const std::string prefix,
-                        SDVersion version = VERSION_SDXL,
-                        PMVersion pm_v    = PM_VERSION_1,
-                        float sty         = 20.f)
-        : GGMLRunner(backend, params_backend),
+                        SDVersion version                                   = VERSION_SDXL,
+                        PMVersion pm_v                                      = PM_VERSION_1,
+                        float sty                                           = 20.f,
+                        std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+        : GGMLRunner(backend, weight_manager),
           version(version),
           pm_version(pm_v),
           style_strength(sty) {
@@ -556,24 +558,25 @@ public:
             return build_graph(id_pixel_values, prompt_embeds, class_tokens_mask, id_embeds);
         };
 
-        return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true));
+        return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true, true, true));
     }
 };
 
 struct PhotoMakerIDEmbed : public GGMLRunner {
     std::map<std::string, ggml_tensor*> tensors;
     std::string file_path;
-    ModelLoader* model_loader;
-    bool load_failed = false;
-    bool applied     = false;
+    std::shared_ptr<ModelManager> model_manager;
+    ggml_backend_t params_backend = nullptr;
+    bool load_failed              = false;
+    bool applied                  = false;
 
     PhotoMakerIDEmbed(ggml_backend_t backend,
-                      ggml_backend_t params_backend,
-                      ModelLoader* ml,
-                      const std::string& file_path = "",
-                      const std::string& prefix    = "")
-        : file_path(file_path), GGMLRunner(backend, params_backend), model_loader(ml) {
-        if (!model_loader->init_from_file_and_convert_name(file_path, prefix)) {
+                      ggml_backend_t params_backend_,
+                      std::shared_ptr<ModelManager> manager = std::make_shared<ModelManager>(),
+                      const std::string& file_path          = "",
+                      const std::string& prefix             = "")
+        : GGMLRunner(backend, manager), file_path(file_path), model_manager(std::move(manager)), params_backend(params_backend_) {
+        if (model_manager == nullptr || !model_manager->loader().init_from_file_and_convert_name(file_path, prefix)) {
             load_failed = true;
         }
     }
@@ -614,14 +617,27 @@ struct PhotoMakerIDEmbed : public GGMLRunner {
             return true;
         };
 
-        model_loader->load_tensors(on_new_tensor_cb, n_threads);
-        if (!alloc_params_buffer()) {
-            LOG_ERROR("PhotoMaker ID embeds buffer allocation failed");
+        model_manager->set_n_threads(n_threads);
+        ModelLoader& model_loader = model_manager->loader();
+        model_loader.load_tensors(on_new_tensor_cb);
+        if (!model_manager->register_param_tensors("PhotoMaker ID embeds",
+                                                   tensors,
+                                                   ModelManager::ResidencyMode::ParamBackend,
+                                                   runtime_backend,
+                                                   params_backend) ||
+            !model_manager->validate_registered_tensors()) {
+            LOG_ERROR("PhotoMaker ID embeds model manager registration failed");
+            return false;
+        }
+        std::vector<ggml_tensor*> id_embed_params;
+        id_embed_params.reserve(tensors.size());
+        for (const auto& pair : tensors) {
+            id_embed_params.push_back(pair.second);
+        }
+        if (!model_manager->prepare_params(id_embed_params)) {
+            LOG_ERROR("PhotoMaker ID embeds model manager prepare params failed");
             return false;
         }
-
-        dry_run = false;
-        model_loader->load_tensors(on_new_tensor_cb, n_threads);
 
         LOG_DEBUG("finished loading PhotoMaker ID Embeds ");
         return true;
@@ -636,4 +652,4 @@ struct PhotoMakerIDEmbed : public GGMLRunner {
     }
 };
 
-#endif  // __PMI_HPP__
+#endif  // __SD_MODEL_ADAPTER_PMID_HPP__

src/model/common/block.hpp

@@ -1,9 +1,9 @@
-#ifndef __COMMON_BLOCK_HPP__
-#define __COMMON_BLOCK_HPP__
+#ifndef __SD_MODEL_COMMON_BLOCK_HPP__
+#define __SD_MODEL_COMMON_BLOCK_HPP__
 
+#include "core/ggml_extend.hpp"
+#include "core/util.h"
 #include "ggml-backend.h"
-#include "ggml_extend.hpp"
-#include "util.h"
 
 class DownSampleBlock : public GGMLBlock {
 protected:
@@ -227,6 +227,37 @@ public:
     }
 };
 
+struct Mlp : public GGMLBlock {
+public:
+    Mlp(int64_t in_features,
+        int64_t hidden_features = -1,
+        int64_t out_features    = -1,
+        bool bias               = true) {
+        // act_layer is always lambda: nn.GELU(approximate="tanh")
+        // norm_layer is always None
+        // use_conv is always False
+        if (hidden_features == -1) {
+            hidden_features = in_features;
+        }
+        if (out_features == -1) {
+            out_features = in_features;
+        }
+        blocks["fc1"] = std::shared_ptr<GGMLBlock>(new Linear(in_features, hidden_features, bias));
+        blocks["fc2"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_features, out_features, bias));
+    }
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+        // x: [N, n_token, in_features]
+        auto fc1 = std::dynamic_pointer_cast<Linear>(blocks["fc1"]);
+        auto fc2 = std::dynamic_pointer_cast<Linear>(blocks["fc2"]);
+
+        x = fc1->forward(ctx, x);
+        x = ggml_ext_gelu(ctx->ggml_ctx, x, true);
+        x = fc2->forward(ctx, x);
+        return x;
+    }
+};
+
 class FeedForward : public GGMLBlock {
 public:
     enum class Activation {
@@ -529,11 +560,11 @@ protected:
         params["mix_factor"] = ggml_new_tensor_1d(ctx, wtype, 1);
     }
 
-    float get_alpha() {
+    ggml_tensor* get_alpha(GGMLRunnerContext* ctx) {
         // image_only_indicator is always tensor([0.]) and since mix_factor.shape is [1,]
         // so learned_with_images is same as learned
-        float alpha = ggml_ext_backend_tensor_get_f32(params["mix_factor"]);
-        return sigmoid(alpha);
+        auto mix_factor = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, params["mix_factor"]);
+        return ggml_sigmoid(ctx->ggml_ctx, mix_factor);
     }
 
 public:
@@ -547,11 +578,12 @@ public:
                          ggml_tensor* x_spatial,
                          ggml_tensor* x_temporal) {
         // image_only_indicator is always tensor([0.])
-        float alpha = get_alpha();
-        auto x      = ggml_add(ctx->ggml_ctx,
-                               ggml_ext_scale(ctx->ggml_ctx, x_spatial, alpha),
-                               ggml_ext_scale(ctx->ggml_ctx, x_temporal, 1.0f - alpha));
-        return x;
+        auto alpha = get_alpha(ctx);
+        return ggml_add(ctx->ggml_ctx,
+                        x_temporal,
+                        ggml_mul(ctx->ggml_ctx,
+                                 ggml_sub(ctx->ggml_ctx, x_spatial, x_temporal),
+                                 alpha));
     }
 };
 
@@ -603,4 +635,4 @@ public:
     }
 };
 
-#endif  // __COMMON_BLOCK_HPP__
+#endif  // __SD_MODEL_COMMON_BLOCK_HPP__

src/model/common/rope.hpp

@@ -1,10 +1,10 @@
-#ifndef __ROPE_HPP__
-#define __ROPE_HPP__
+#ifndef __SD_MODEL_COMMON_ROPE_HPP__
+#define __SD_MODEL_COMMON_ROPE_HPP__
 
 #include <algorithm>
 #include <cmath>
 #include <vector>
-#include "ggml_extend.hpp"
+#include "core/ggml_extend.hpp"
 
 namespace Rope {
     enum class EmbedNDLayout {
@@ -249,6 +249,45 @@ namespace Rope {
         return embed_nd(ids, bs, axis_thetas, axes_dim, wrap_dims, layout);
     }
 
+    __STATIC_INLINE__ std::vector<float> embed_interleaved_mrope(const std::vector<std::vector<float>>& ids,
+                                                                 int bs,
+                                                                 float theta,
+                                                                 int head_dim,
+                                                                 const std::vector<int>& mrope_section,
+                                                                 const std::vector<std::vector<int>>& axis_wrap_dims = {}) {
+        GGML_ASSERT(bs > 0);
+        GGML_ASSERT(head_dim % 2 == 0);
+        GGML_ASSERT(mrope_section.size() >= 3);
+
+        std::vector<std::vector<float>> trans_ids = transpose(ids);
+        size_t pos_len                            = ids.size() / bs;
+        int half_dim                              = head_dim / 2;
+
+        std::vector<std::vector<std::vector<float>>> axis_embs;
+        axis_embs.reserve(3);
+        for (int axis = 0; axis < 3; ++axis) {
+            std::vector<int> axis_wrap;
+            if (axis < static_cast<int>(axis_wrap_dims.size())) {
+                axis_wrap = axis_wrap_dims[axis];
+            }
+            axis_embs.push_back(rope(trans_ids[axis], head_dim, theta, axis_wrap));
+        }
+
+        std::vector<std::vector<float>> emb = axis_embs[0];
+        for (int axis = 1; axis < 3; ++axis) {
+            int length = std::min<int>(mrope_section[axis] * 3, half_dim);
+            for (int freq_idx = axis; freq_idx < length; freq_idx += 3) {
+                for (size_t pos_idx = 0; pos_idx < bs * pos_len; ++pos_idx) {
+                    for (int k = 0; k < 4; ++k) {
+                        emb[pos_idx][4 * freq_idx + k] = axis_embs[axis][pos_idx][4 * freq_idx + k];
+                    }
+                }
+            }
+        }
+
+        return flatten(emb);
+    }
+
     __STATIC_INLINE__ std::vector<float> embed_2d_interleaved(int height,
                                                               int width,
                                                               int dim,
@@ -868,4 +907,4 @@ namespace Rope {
     }
 };  // namespace Rope
 
-#endif  // __ROPE_HPP__
+#endif  // __SD_MODEL_COMMON_ROPE_HPP__

src/model/diffusion/anima.hpp

@@ -1,19 +1,61 @@
-#ifndef __ANIMA_HPP__
-#define __ANIMA_HPP__
+#ifndef __SD_MODEL_DIFFUSION_ANIMA_HPP__
+#define __SD_MODEL_DIFFUSION_ANIMA_HPP__
 
+#include <algorithm>
 #include <cmath>
 #include <memory>
 #include <utility>
 #include <vector>
 
-#include "common_block.hpp"
-#include "diffusion_model.hpp"
-#include "flux.hpp"
-#include "rope.hpp"
+#include "model/common/block.hpp"
+#include "model/common/rope.hpp"
+#include "model/diffusion/flux.hpp"
+#include "model/diffusion/model.hpp"
 
 namespace Anima {
     constexpr int ANIMA_GRAPH_SIZE = 65536;
 
+    struct AnimaConfig {
+        int64_t in_channels       = 16;
+        int64_t out_channels      = 16;
+        int64_t hidden_size       = 2048;
+        int64_t text_embed_dim    = 1024;
+        int64_t num_heads         = 16;
+        int64_t head_dim          = 128;
+        int patch_size            = 2;
+        int64_t num_layers        = 28;
+        std::vector<int> axes_dim = {44, 42, 42};
+        int theta                 = 10000;
+
+        static AnimaConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
+            AnimaConfig config;
+            int64_t detected_layers = 0;
+            std::string layer_tag   = prefix.empty() ? "blocks." : prefix + ".blocks.";
+            for (const auto& [name, _] : tensor_storage_map) {
+                size_t pos = name.find(layer_tag);
+                if (pos == std::string::npos) {
+                    continue;
+                }
+                size_t start = pos + layer_tag.size();
+                size_t end   = name.find('.', start);
+                if (end == std::string::npos) {
+                    continue;
+                }
+                int64_t layer_id = atoll(name.substr(start, end - start).c_str());
+                detected_layers  = std::max(detected_layers, layer_id + 1);
+            }
+            if (detected_layers > 0) {
+                config.num_layers = detected_layers;
+                LOG_DEBUG("anima: num_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %" PRId64 ", head_dim = %" PRId64,
+                          config.num_layers,
+                          config.hidden_size,
+                          config.num_heads,
+                          config.head_dim);
+            }
+            return config;
+        }
+    };
+
     __STATIC_INLINE__ ggml_tensor* apply_gate(ggml_context* ctx,
                                               ggml_tensor* x,
                                               ggml_tensor* gate) {
@@ -418,31 +460,22 @@ namespace Anima {
 
     struct AnimaNet : public GGMLBlock {
     public:
-        int64_t in_channels       = 16;
-        int64_t out_channels      = 16;
-        int64_t hidden_size       = 2048;
-        int64_t text_embed_dim    = 1024;
-        int64_t num_heads         = 16;
-        int64_t head_dim          = 128;
-        int patch_size            = 2;
-        int64_t num_layers        = 28;
-        std::vector<int> axes_dim = {44, 42, 42};
-        int theta                 = 10000;
+        AnimaConfig config;
 
     public:
         AnimaNet() = default;
-        explicit AnimaNet(int64_t num_layers)
-            : num_layers(num_layers) {
-            blocks["x_embedder"]       = std::make_shared<XEmbedder>((in_channels + 1) * patch_size * patch_size, hidden_size);
-            blocks["t_embedder"]       = std::make_shared<TimestepEmbedder>(hidden_size, hidden_size * 3);
-            blocks["t_embedding_norm"] = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
-            for (int i = 0; i < num_layers; i++) {
-                blocks["blocks." + std::to_string(i)] = std::make_shared<TransformerBlock>(hidden_size,
-                                                                                           text_embed_dim,
-                                                                                           num_heads,
-                                                                                           head_dim);
+        explicit AnimaNet(AnimaConfig config)
+            : config(config) {
+            blocks["x_embedder"]       = std::make_shared<XEmbedder>((config.in_channels + 1) * config.patch_size * config.patch_size, config.hidden_size);
+            blocks["t_embedder"]       = std::make_shared<TimestepEmbedder>(config.hidden_size, config.hidden_size * 3);
+            blocks["t_embedding_norm"] = std::make_shared<RMSNorm>(config.hidden_size, 1e-6f);
+            for (int i = 0; i < config.num_layers; i++) {
+                blocks["blocks." + std::to_string(i)] = std::make_shared<TransformerBlock>(config.hidden_size,
+                                                                                           config.text_embed_dim,
+                                                                                           config.num_heads,
+                                                                                           config.head_dim);
             }
-            blocks["final_layer"] = std::make_shared<FinalLayer>(hidden_size, patch_size, out_channels);
+            blocks["final_layer"] = std::make_shared<FinalLayer>(config.hidden_size, config.patch_size, config.out_channels);
             blocks["llm_adapter"] = std::make_shared<LLMAdapter>(1024, 1024, 1024, 6, 16);
         }
 
@@ -469,11 +502,11 @@ namespace Anima {
             auto padding_mask = ggml_ext_zeros(ctx->ggml_ctx, x->ne[0], x->ne[1], 1, x->ne[3]);
             x                 = ggml_concat(ctx->ggml_ctx, x, padding_mask, 2);  // [N, C + 1, H, W]
 
-            x = DiT::pad_and_patchify(ctx, x, patch_size, patch_size);  // [N, h*w, (C+1)*ph*pw]
+            x = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size);  // [N, h*w, (C+1)*ph*pw]
 
             x = x_embedder->forward(ctx, x);
 
-            auto timestep_proj     = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast<int>(hidden_size));
+            auto timestep_proj     = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast<int>(config.hidden_size));
             auto temb              = t_embedder->forward(ctx, timestep_proj);
             auto embedded_timestep = t_embedding_norm->forward(ctx, timestep_proj);
 
@@ -505,7 +538,7 @@ namespace Anima {
             sd::ggml_graph_cut::mark_graph_cut(temb, "anima.prelude", "temb");
             sd::ggml_graph_cut::mark_graph_cut(encoder_hidden_states, "anima.prelude", "context");
 
-            for (int i = 0; i < num_layers; i++) {
+            for (int i = 0; i < config.num_layers; i++) {
                 auto block = std::dynamic_pointer_cast<TransformerBlock>(blocks["blocks." + std::to_string(i)]);
                 x          = block->forward(ctx, x, encoder_hidden_states, embedded_timestep, temb, image_pe);
                 sd::ggml_graph_cut::mark_graph_cut(x, "anima.blocks." + std::to_string(i), "x");
@@ -513,7 +546,7 @@ namespace Anima {
 
             x = final_layer->forward(ctx, x, embedded_timestep, temb);  // [N, h*w, ph*pw*C]
 
-            x = DiT::unpatchify_and_crop(ctx->ggml_ctx, x, H, W, patch_size, patch_size, false);  // [N, C, H, W]
+            x = DiT::unpatchify_and_crop(ctx->ggml_ctx, x, H, W, config.patch_size, config.patch_size, false);  // [N, C, H, W]
 
             return x;
         }
@@ -524,35 +557,16 @@ namespace Anima {
         std::vector<float> image_pe_vec;
         std::vector<float> adapter_q_pe_vec;
         std::vector<float> adapter_k_pe_vec;
+        AnimaConfig config;
         AnimaNet net;
 
         AnimaRunner(ggml_backend_t backend,
-                    ggml_backend_t params_backend,
-                    const String2TensorStorage& tensor_storage_map = {},
-                    const std::string prefix                       = "model.diffusion_model")
-            : DiffusionModelRunner(backend, params_backend, prefix) {
-            int64_t num_layers    = 0;
-            std::string layer_tag = prefix + ".net.blocks.";
-            for (const auto& kv : tensor_storage_map) {
-                const std::string& tensor_name = kv.first;
-                size_t pos                     = tensor_name.find(layer_tag);
-                if (pos == std::string::npos) {
-                    continue;
-                }
-                size_t start = pos + layer_tag.size();
-                size_t end   = tensor_name.find('.', start);
-                if (end == std::string::npos) {
-                    continue;
-                }
-                int64_t layer_id = atoll(tensor_name.substr(start, end - start).c_str());
-                num_layers       = std::max(num_layers, layer_id + 1);
-            }
-            if (num_layers <= 0) {
-                num_layers = 28;
-            }
-            LOG_INFO("anima net layers: %" PRId64, num_layers);
-
-            net = AnimaNet(num_layers);
+                    const String2TensorStorage& tensor_storage_map      = {},
+                    const std::string prefix                            = "model.diffusion_model",
+                    std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : DiffusionModelRunner(backend, prefix, weight_manager),
+              config(AnimaConfig::detect_from_weights(tensor_storage_map, prefix + ".net")) {
+            net = AnimaNet(config);
             net.init(params_ctx, tensor_storage_map, prefix + ".net");
         }
 
@@ -623,22 +637,22 @@ namespace Anima {
             GGML_ASSERT(x->ne[3] == 1);
             ggml_cgraph* gf = new_graph_custom(ANIMA_GRAPH_SIZE);
 
-            int64_t pad_h = (net.patch_size - x->ne[1] % net.patch_size) % net.patch_size;
-            int64_t pad_w = (net.patch_size - x->ne[0] % net.patch_size) % net.patch_size;
+            int64_t pad_h = (config.patch_size - x->ne[1] % config.patch_size) % config.patch_size;
+            int64_t pad_w = (config.patch_size - x->ne[0] % config.patch_size) % config.patch_size;
             int64_t h_pad = x->ne[1] + pad_h;
             int64_t w_pad = x->ne[0] + pad_w;
 
             image_pe_vec          = gen_anima_image_pe_vec(1,
                                                            static_cast<int>(h_pad),
                                                            static_cast<int>(w_pad),
-                                                           static_cast<int>(net.patch_size),
-                                                           net.theta,
-                                                           net.axes_dim,
+                                                           static_cast<int>(config.patch_size),
+                                                           config.theta,
+                                                           config.axes_dim,
                                                            4.0f,
                                                            4.0f,
                                                            1.0f);
-            int64_t image_pos_len = static_cast<int64_t>(image_pe_vec.size()) / (2 * 2 * (net.head_dim / 2));
-            auto image_pe         = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, net.head_dim / 2, image_pos_len);
+            int64_t image_pos_len = static_cast<int64_t>(image_pe_vec.size()) / (2 * 2 * (config.head_dim / 2));
+            auto image_pe         = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, image_pos_len);
             set_backend_tensor_data(image_pe, image_pe_vec.data());
 
             ggml_tensor* adapter_q_pe = nullptr;
@@ -683,7 +697,7 @@ namespace Anima {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, t5_ids, t5_weights);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,
@@ -701,4 +715,4 @@ namespace Anima {
     };
 }  // namespace Anima
 
-#endif  // __ANIMA_HPP__
+#endif  // __SD_MODEL_DIFFUSION_ANIMA_HPP__

src/model/diffusion/control.hpp

@@ -1,8 +1,9 @@
-#ifndef __CONTROL_HPP__
-#define __CONTROL_HPP__
+#ifndef __SD_MODEL_DIFFUSION_CONTROL_HPP__
+#define __SD_MODEL_DIFFUSION_CONTROL_HPP__
 
-#include "common_block.hpp"
-#include "model.h"
+#include "model/common/block.hpp"
+#include "model_loader.h"
+#include "model_manager.h"
 
 #define CONTROL_NET_GRAPH_SIZE 1536
 
@@ -309,73 +310,47 @@ public:
 struct ControlNet : public GGMLRunner {
     SDVersion version = VERSION_SD1;
     ControlNetBlock control_net;
+    std::string weight_prefix;
 
-    ggml_backend_buffer_t control_buffer = nullptr;
-    ggml_context* control_ctx            = nullptr;
     std::vector<ggml_tensor*> control_outputs_ggml;
     ggml_tensor* guided_hint_output_ggml = nullptr;
     std::vector<sd::Tensor<float>> controls;
-    sd::Tensor<float> guided_hint;
     bool guided_hint_cached = false;
+    std::shared_ptr<ModelManager> owned_model_manager;
+    ggml_backend_t params_backend = nullptr;
+
+    static const char* guided_hint_cache_name() {
+        return "controlnet.guided_hint";
+    }
 
     ControlNet(ggml_backend_t backend,
-               ggml_backend_t params_backend,
-               const String2TensorStorage& tensor_storage_map = {},
-               SDVersion version                              = VERSION_SD1)
-        : GGMLRunner(backend, params_backend), control_net(version) {
-        control_net.init(params_ctx, tensor_storage_map, "");
+               ggml_backend_t params_backend_,
+               const String2TensorStorage& tensor_storage_map      = {},
+               SDVersion version                                   = VERSION_SD1,
+               const std::string& prefix                           = "",
+               std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+        : GGMLRunner(backend, weight_manager), version(version), control_net(version), weight_prefix(prefix), params_backend(params_backend_) {
+        control_net.init(params_ctx, tensor_storage_map, prefix);
     }
 
     ~ControlNet() override {
         free_control_ctx();
     }
 
-    void alloc_control_ctx(std::vector<ggml_tensor*> outs) {
-        ggml_init_params params;
-        params.mem_size   = static_cast<size_t>(outs.size() * ggml_tensor_overhead()) + 1024 * 1024;
-        params.mem_buffer = nullptr;
-        params.no_alloc   = true;
-        control_ctx       = ggml_init(params);
-
-        control_outputs_ggml.resize(outs.size() - 1);
-
-        size_t control_buffer_size = 0;
-
-        guided_hint_output_ggml = ggml_dup_tensor(control_ctx, outs[0]);
-        control_buffer_size += ggml_nbytes(guided_hint_output_ggml);
-
-        for (int i = 0; i < outs.size() - 1; i++) {
-            control_outputs_ggml[i] = ggml_dup_tensor(control_ctx, outs[i + 1]);
-            control_buffer_size += ggml_nbytes(control_outputs_ggml[i]);
-        }
-
-        control_buffer = ggml_backend_alloc_ctx_tensors(control_ctx, runtime_backend);
-
-        LOG_DEBUG("control buffer size %.2fMB", control_buffer_size * 1.f / 1024.f / 1024.f);
-    }
-
     void free_control_ctx() {
-        if (control_buffer != nullptr) {
-            ggml_backend_buffer_free(control_buffer);
-            control_buffer = nullptr;
-        }
-        if (control_ctx != nullptr) {
-            ggml_free(control_ctx);
-            control_ctx = nullptr;
-        }
         guided_hint_output_ggml = nullptr;
         guided_hint_cached      = false;
-        guided_hint             = {};
         control_outputs_ggml.clear();
         controls.clear();
+        free_cache_ctx_and_buffer();
     }
 
     std::string get_desc() override {
         return "control_net";
     }
 
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
-        control_net.get_param_tensors(tensors, prefix);
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
+        control_net.get_param_tensors(tensors, weight_prefix);
     }
 
     ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
@@ -391,11 +366,17 @@ struct ControlNet : public GGMLRunner {
         ggml_tensor* context   = make_optional_input(context_tensor);
         ggml_tensor* y         = make_optional_input(y_tensor);
 
+        guided_hint_output_ggml = nullptr;
+        control_outputs_ggml.clear();
+
         ggml_tensor* guided_hint_input = nullptr;
-        if (guided_hint_cached && !guided_hint.empty()) {
-            guided_hint_input = make_input(guided_hint);
-            hint              = nullptr;
-        } else {
+        if (guided_hint_cached) {
+            guided_hint_input = get_cache_tensor_by_name(guided_hint_cache_name());
+            if (guided_hint_input == nullptr) {
+                guided_hint_cached = false;
+            }
+        }
+        if (guided_hint_input == nullptr) {
             hint = make_input(hint_tensor);
         }
 
@@ -409,13 +390,19 @@ struct ControlNet : public GGMLRunner {
                                         context,
                                         y);
 
-        if (control_ctx == nullptr) {
-            alloc_control_ctx(outs);
+        if (guided_hint_input == nullptr && !outs.empty()) {
+            guided_hint_output_ggml = outs[0];
+            ggml_set_output(guided_hint_output_ggml);
+            cache(guided_hint_cache_name(), guided_hint_output_ggml);
+            ggml_build_forward_expand(gf, guided_hint_output_ggml);
         }
 
-        ggml_build_forward_expand(gf, ggml_cpy(compute_ctx, outs[0], guided_hint_output_ggml));
-        for (int i = 0; i < outs.size() - 1; i++) {
-            ggml_build_forward_expand(gf, ggml_cpy(compute_ctx, outs[i + 1], control_outputs_ggml[i]));
+        control_outputs_ggml.reserve(outs.size() > 0 ? outs.size() - 1 : 0);
+        for (size_t i = 1; i < outs.size(); i++) {
+            ggml_tensor* control_output = outs[i];
+            ggml_set_output(control_output);
+            ggml_build_forward_expand(gf, control_output);
+            control_outputs_ggml.push_back(control_output);
         }
 
         return gf;
@@ -435,15 +422,12 @@ struct ControlNet : public GGMLRunner {
             return build_graph(x, hint, timesteps, context, y);
         };
 
-        auto compute_result = GGMLRunner::compute<float>(get_graph, n_threads, false);
+        auto compute_result = GGMLRunner::compute<float>(get_graph, n_threads, false, false, false, true);
         if (!compute_result.has_value()) {
             return std::nullopt;
         }
 
-        if (guided_hint_output_ggml != nullptr) {
-            guided_hint = restore_trailing_singleton_dims(sd::make_sd_tensor_from_ggml<float>(guided_hint_output_ggml),
-                                                          4);
-        }
+        guided_hint_cached = get_cache_tensor_by_name(guided_hint_cache_name()) != nullptr;
         controls.clear();
         controls.reserve(control_outputs_ggml.size());
         for (ggml_tensor* control : control_outputs_ggml) {
@@ -451,37 +435,41 @@ struct ControlNet : public GGMLRunner {
             GGML_ASSERT(!control_host.empty());
             controls.push_back(std::move(control_host));
         }
-        guided_hint_cached = true;
         return controls;
     }
 
     bool load_from_file(const std::string& file_path, int n_threads) {
         LOG_INFO("loading control net from '%s'", file_path.c_str());
-        if (!alloc_params_buffer()) {
-            LOG_ERROR("control net model buffer allocation failed");
-            return false;
-        }
-
         std::map<std::string, ggml_tensor*> tensors;
         control_net.get_param_tensors(tensors);
-        std::set<std::string> ignore_tensors;
 
-        ModelLoader model_loader;
+        auto manager = std::dynamic_pointer_cast<ModelManager>(weight_manager.lock());
+        if (manager == nullptr) {
+            owned_model_manager = std::make_shared<ModelManager>();
+            weight_manager      = owned_model_manager;
+            manager             = owned_model_manager;
+        }
+
+        ModelLoader& model_loader = manager->loader();
         if (!model_loader.init_from_file_and_convert_name(file_path)) {
             LOG_ERROR("init control net model loader from file failed: '%s'", file_path.c_str());
             return false;
         }
 
-        bool success = model_loader.load_tensors(tensors, ignore_tensors, n_threads);
-
-        if (!success) {
-            LOG_ERROR("load control net tensors from model loader failed");
+        manager->set_n_threads(n_threads);
+        if (!manager->register_param_tensors("ControlNet",
+                                             std::move(tensors),
+                                             ModelManager::ResidencyMode::ParamBackend,
+                                             runtime_backend,
+                                             params_backend) ||
+            !manager->validate_registered_tensors()) {
+            LOG_ERROR("register control net tensors with model manager failed");
             return false;
         }
 
         LOG_INFO("control net model loaded");
-        return success;
+        return true;
     }
 };
 
-#endif  // __CONTROL_HPP__
+#endif  // __SD_MODEL_DIFFUSION_CONTROL_HPP__

src/model/diffusion/dit.hpp

@@ -1,7 +1,7 @@
-#ifndef __COMMON_DIT_HPP__
-#define __COMMON_DIT_HPP__
+#ifndef __SD_MODEL_DIFFUSION_DIT_HPP__
+#define __SD_MODEL_DIFFUSION_DIT_HPP__
 
-#include "ggml_extend.hpp"
+#include "core/ggml_extend.hpp"
 
 namespace DiT {
     inline ggml_tensor* patchify(ggml_context* ctx,
@@ -163,4 +163,4 @@ namespace DiT {
     }
 }  // namespace DiT
 
-#endif  // __COMMON_DIT_HPP__
+#endif  // __SD_MODEL_DIFFUSION_DIT_HPP__

src/model/diffusion/ernie_image.hpp

@@ -1,18 +1,88 @@
-#ifndef __SD_ERNIE_IMAGE_HPP__
-#define __SD_ERNIE_IMAGE_HPP__
+#ifndef __SD_MODEL_DIFFUSION_ERNIE_IMAGE_HPP__
+#define __SD_MODEL_DIFFUSION_ERNIE_IMAGE_HPP__
 
 #include <memory>
 #include <vector>
 
-#include "common_dit.hpp"
-#include "diffusion_model.hpp"
-#include "flux.hpp"
-#include "qwen_image.hpp"
-#include "rope.hpp"
+#include "model/common/rope.hpp"
+#include "model/diffusion/dit.hpp"
+#include "model/diffusion/flux.hpp"
+#include "model/diffusion/model.hpp"
+#include "model/diffusion/qwen_image.hpp"
 
 namespace ErnieImage {
     constexpr int ERNIE_IMAGE_GRAPH_SIZE = 40960;
 
+    struct ErnieImageConfig {
+        int64_t hidden_size       = 4096;
+        int64_t num_heads         = 32;
+        int64_t num_layers        = 36;
+        int64_t ffn_hidden_size   = 12288;
+        int64_t in_channels       = 128;
+        int64_t out_channels      = 128;
+        int patch_size            = 1;
+        int64_t text_in_dim       = 3072;
+        int theta                 = 256;
+        std::vector<int> axes_dim = {32, 48, 48};
+        int axes_dim_sum          = 128;
+        float eps                 = 1e-6f;
+
+        static ErnieImageConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
+            ErnieImageConfig config;
+            config.num_layers         = 0;
+            int64_t detected_head_dim = 0;
+            for (const auto& [name, tensor_storage] : tensor_storage_map) {
+                if (!starts_with(name, prefix)) {
+                    continue;
+                }
+                if (ends_with(name, "x_embedder.proj.weight") && tensor_storage.n_dims == 4) {
+                    config.patch_size  = static_cast<int>(tensor_storage.ne[0]);
+                    config.in_channels = tensor_storage.ne[2];
+                    config.hidden_size = tensor_storage.ne[3];
+                } else if (ends_with(name, "text_proj.weight") && tensor_storage.n_dims == 2) {
+                    config.text_in_dim = tensor_storage.ne[0];
+                } else if (ends_with(name, "layers.0.self_attention.norm_q.weight")) {
+                    detected_head_dim = tensor_storage.ne[0];
+                } else if (ends_with(name, "layers.0.mlp.gate_proj.weight") && tensor_storage.n_dims == 2) {
+                    config.ffn_hidden_size = tensor_storage.ne[1];
+                } else if (ends_with(name, "final_linear.weight") && tensor_storage.n_dims == 2) {
+                    int64_t out_dim     = tensor_storage.ne[1];
+                    int64_t patch_area  = config.patch_size * config.patch_size;
+                    config.out_channels = out_dim / patch_area;
+                }
+
+                size_t pos = name.find("layers.");
+                if (pos != std::string::npos) {
+                    auto items = split_string(name.substr(pos), '.');
+                    if (items.size() > 1) {
+                        int block_index = atoi(items[1].c_str());
+                        if (block_index + 1 > config.num_layers) {
+                            config.num_layers = block_index + 1;
+                        }
+                    }
+                }
+            }
+            if (config.num_layers == 0) {
+                config.num_layers = 36;
+            }
+            if (detected_head_dim > 0) {
+                config.num_heads = config.hidden_size / detected_head_dim;
+            }
+            config.axes_dim_sum = 0;
+            for (int axis_dim : config.axes_dim) {
+                config.axes_dim_sum += axis_dim;
+            }
+            LOG_DEBUG("ernie_image: num_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %" PRId64 ", ffn_hidden_size = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64,
+                      config.num_layers,
+                      config.hidden_size,
+                      config.num_heads,
+                      config.ffn_hidden_size,
+                      config.in_channels,
+                      config.out_channels);
+            return config;
+        }
+    };
+
     __STATIC_INLINE__ ggml_tensor* timestep_embedding_sin_cos(ggml_context* ctx,
                                                               ggml_tensor* timesteps,
                                                               int dim,
@@ -208,51 +278,36 @@ namespace ErnieImage {
         }
     };
 
-    struct ErnieImageParams {
-        int64_t hidden_size       = 4096;
-        int64_t num_heads         = 32;
-        int64_t num_layers        = 36;
-        int64_t ffn_hidden_size   = 12288;
-        int64_t in_channels       = 128;
-        int64_t out_channels      = 128;
-        int patch_size            = 1;
-        int64_t text_in_dim       = 3072;
-        int theta                 = 256;
-        std::vector<int> axes_dim = {32, 48, 48};
-        int axes_dim_sum          = 128;
-        float eps                 = 1e-6f;
-    };
-
     class ErnieImageModel : public GGMLBlock {
     public:
-        ErnieImageParams params;
+        ErnieImageConfig config;
 
         ErnieImageModel() = default;
-        ErnieImageModel(ErnieImageParams params)
-            : params(params) {
-            blocks["x_embedder.proj"] = std::make_shared<Conv2d>(params.in_channels,
-                                                                 params.hidden_size,
-                                                                 std::pair<int, int>{params.patch_size, params.patch_size},
-                                                                 std::pair<int, int>{params.patch_size, params.patch_size},
+        ErnieImageModel(ErnieImageConfig config)
+            : config(config) {
+            blocks["x_embedder.proj"] = std::make_shared<Conv2d>(config.in_channels,
+                                                                 config.hidden_size,
+                                                                 std::pair<int, int>{config.patch_size, config.patch_size},
+                                                                 std::pair<int, int>{config.patch_size, config.patch_size},
                                                                  std::pair<int, int>{0, 0},
                                                                  std::pair<int, int>{1, 1},
                                                                  true);
-            if (params.text_in_dim != params.hidden_size) {
-                blocks["text_proj"] = std::make_shared<Linear>(params.text_in_dim, params.hidden_size, false);
+            if (config.text_in_dim != config.hidden_size) {
+                blocks["text_proj"] = std::make_shared<Linear>(config.text_in_dim, config.hidden_size, false);
             }
-            blocks["time_embedding"]     = std::make_shared<Qwen::TimestepEmbedding>(params.hidden_size, params.hidden_size);
-            blocks["adaLN_modulation.1"] = std::make_shared<Linear>(params.hidden_size, 6 * params.hidden_size, true);
+            blocks["time_embedding"]     = std::make_shared<Qwen::TimestepEmbedding>(config.hidden_size, config.hidden_size);
+            blocks["adaLN_modulation.1"] = std::make_shared<Linear>(config.hidden_size, 6 * config.hidden_size, true);
 
-            for (int i = 0; i < params.num_layers; i++) {
-                blocks["layers." + std::to_string(i)] = std::make_shared<ErnieImageSharedAdaLNBlock>(params.hidden_size,
-                                                                                                     params.num_heads,
-                                                                                                     params.ffn_hidden_size,
-                                                                                                     params.eps);
+            for (int i = 0; i < config.num_layers; i++) {
+                blocks["layers." + std::to_string(i)] = std::make_shared<ErnieImageSharedAdaLNBlock>(config.hidden_size,
+                                                                                                     config.num_heads,
+                                                                                                     config.ffn_hidden_size,
+                                                                                                     config.eps);
             }
 
-            blocks["final_norm"]   = std::make_shared<ErnieImageAdaLNContinuous>(params.hidden_size, params.eps);
-            blocks["final_linear"] = std::make_shared<Linear>(params.hidden_size,
-                                                              params.patch_size * params.patch_size * params.out_channels,
+            blocks["final_norm"]   = std::make_shared<ErnieImageAdaLNContinuous>(config.hidden_size, config.eps);
+            blocks["final_linear"] = std::make_shared<Linear>(config.hidden_size,
+                                                              config.patch_size * config.patch_size * config.out_channels,
                                                               true);
         }
 
@@ -265,12 +320,12 @@ namespace ErnieImage {
             // context: [N, text_tokens, 3072]
             // pe: [image_tokens + text_tokens, head_dim/2, 2, 2]
             GGML_ASSERT(context != nullptr);
-            GGML_ASSERT(x->ne[1] % params.patch_size == 0 && x->ne[0] % params.patch_size == 0);
+            GGML_ASSERT(x->ne[1] % config.patch_size == 0 && x->ne[0] % config.patch_size == 0);
 
             int64_t W     = x->ne[0];
             int64_t H     = x->ne[1];
-            int64_t Hp    = H / params.patch_size;
-            int64_t Wp    = W / params.patch_size;
+            int64_t Hp    = H / config.patch_size;
+            int64_t Wp    = W / config.patch_size;
             int64_t n_img = Hp * Wp;
             int64_t N     = x->ne[3];
 
@@ -292,7 +347,7 @@ namespace ErnieImage {
 
             auto hidden_states = ggml_concat(ctx->ggml_ctx, img, txt, 1);  // [N, image_tokens + text_tokens, hidden_size]
 
-            auto sample = timestep_embedding_sin_cos(ctx->ggml_ctx, timestep, static_cast<int>(params.hidden_size));
+            auto sample = timestep_embedding_sin_cos(ctx->ggml_ctx, timestep, static_cast<int>(config.hidden_size));
             auto c      = time_embedding->forward(ctx, sample);  // [N, hidden_size]
 
             auto mod_params = adaLN_mod->forward(ctx, ggml_silu(ctx->ggml_ctx, c));  // [N, 6 * hidden_size]
@@ -305,7 +360,7 @@ namespace ErnieImage {
                 temb.push_back(ggml_reshape_3d(ctx->ggml_ctx, chunk, chunk->ne[0], 1, chunk->ne[1]));  // [N, 1, hidden_size]
             }
 
-            for (int i = 0; i < params.num_layers; i++) {
+            for (int i = 0; i < config.num_layers; i++) {
                 auto layer    = std::dynamic_pointer_cast<ErnieImageSharedAdaLNBlock>(blocks["layers." + std::to_string(i)]);
                 hidden_states = layer->forward(ctx, hidden_states, pe, temb);
                 sd::ggml_graph_cut::mark_graph_cut(hidden_states, "ernie_image.layers." + std::to_string(i), "hidden_states");
@@ -319,74 +374,25 @@ namespace ErnieImage {
                                        patches,
                                        Hp,
                                        Wp,
-                                       params.patch_size,
-                                       params.patch_size,
+                                       config.patch_size,
+                                       config.patch_size,
                                        false);  // [N, out_channels, H, W]
             return out;
         }
     };
 
     struct ErnieImageRunner : public DiffusionModelRunner {
-        ErnieImageParams ernie_params;
+        ErnieImageConfig config;
         ErnieImageModel ernie_image;
         std::vector<float> pe_vec;
 
         ErnieImageRunner(ggml_backend_t backend,
-                         ggml_backend_t params_backend,
-                         const String2TensorStorage& tensor_storage_map = {},
-                         const std::string prefix                       = "")
-            : DiffusionModelRunner(backend, params_backend, prefix) {
-            ernie_params.num_layers = 0;
-            for (const auto& [name, tensor_storage] : tensor_storage_map) {
-                if (!starts_with(name, prefix)) {
-                    continue;
-                }
-                if (ends_with(name, "x_embedder.proj.weight") && tensor_storage.n_dims == 4) {
-                    ernie_params.patch_size  = static_cast<int>(tensor_storage.ne[0]);
-                    ernie_params.in_channels = tensor_storage.ne[2];
-                    ernie_params.hidden_size = tensor_storage.ne[3];
-                } else if (ends_with(name, "text_proj.weight") && tensor_storage.n_dims == 2) {
-                    ernie_params.text_in_dim = tensor_storage.ne[0];
-                } else if (ends_with(name, "layers.0.self_attention.norm_q.weight")) {
-                    int64_t head_dim       = tensor_storage.ne[0];
-                    ernie_params.num_heads = ernie_params.hidden_size / head_dim;
-                } else if (ends_with(name, "layers.0.mlp.gate_proj.weight") && tensor_storage.n_dims == 2) {
-                    ernie_params.ffn_hidden_size = tensor_storage.ne[1];
-                } else if (ends_with(name, "final_linear.weight") && tensor_storage.n_dims == 2) {
-                    int64_t out_dim           = tensor_storage.ne[1];
-                    ernie_params.out_channels = out_dim / ernie_params.patch_size / ernie_params.patch_size;
-                }
-
-                size_t pos = name.find("layers.");
-                if (pos != std::string::npos) {
-                    std::string layer_name = name.substr(pos);
-                    auto items             = split_string(layer_name, '.');
-                    if (items.size() > 1) {
-                        int block_index = atoi(items[1].c_str());
-                        if (block_index + 1 > ernie_params.num_layers) {
-                            ernie_params.num_layers = block_index + 1;
-                        }
-                    }
-                }
-            }
-            if (ernie_params.num_layers == 0) {
-                ernie_params.num_layers = 36;
-            }
-            ernie_params.axes_dim_sum = 0;
-            for (int axis_dim : ernie_params.axes_dim) {
-                ernie_params.axes_dim_sum += axis_dim;
-            }
-
-            LOG_INFO("ernie_image: layers = %" PRId64 ", hidden_size = %" PRId64 ", heads = %" PRId64
-                     ", ffn_hidden_size = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64,
-                     ernie_params.num_layers,
-                     ernie_params.hidden_size,
-                     ernie_params.num_heads,
-                     ernie_params.ffn_hidden_size,
-                     ernie_params.in_channels,
-                     ernie_params.out_channels);
-
-            ernie_image = ErnieImageModel(ernie_params);
+                         const String2TensorStorage& tensor_storage_map      = {},
+                         const std::string prefix                            = "",
+                         std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : DiffusionModelRunner(backend, prefix, weight_manager),
+              config(ErnieImageConfig::detect_from_weights(tensor_storage_map, prefix)) {
+            ernie_image = ErnieImageModel(config);
             ernie_image.init(params_ctx, tensor_storage_map, prefix);
         }
 
@@ -410,15 +416,15 @@ namespace ErnieImage {
 
             pe_vec      = Rope::gen_ernie_image_pe(static_cast<int>(x->ne[1]),
                                                    static_cast<int>(x->ne[0]),
-                                                   ernie_params.patch_size,
+                                                   config.patch_size,
                                                    static_cast<int>(x->ne[3]),
                                                    static_cast<int>(context->ne[1]),
-                                                   ernie_params.theta,
+                                                   config.theta,
                                                    circular_y_enabled,
                                                    circular_x_enabled,
-                                                   ernie_params.axes_dim);
-            int pos_len = static_cast<int>(pe_vec.size() / ernie_params.axes_dim_sum / 2);
-            auto pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, ernie_params.axes_dim_sum, 1, pos_len, 2);
+                                                   config.axes_dim);
+            int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
+            auto pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, config.axes_dim_sum, 1, pos_len, 2);
             set_backend_tensor_data(pe, pe_vec.data());
 
             auto runner_ctx  = get_context();
@@ -434,7 +440,7 @@ namespace ErnieImage {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,
@@ -449,4 +455,4 @@ namespace ErnieImage {
     };
 }  // namespace ErnieImage
 
-#endif  // __SD_ERNIE_IMAGE_HPP__
+#endif  // __SD_MODEL_DIFFUSION_ERNIE_IMAGE_HPP__

src/model/diffusion/flux.hpp

@@ -1,18 +1,167 @@
-#ifndef __FLUX_HPP__
-#define __FLUX_HPP__
+#ifndef __SD_MODEL_DIFFUSION_FLUX_HPP__
+#define __SD_MODEL_DIFFUSION_FLUX_HPP__
 
 #include <memory>
 #include <vector>
 
-#include "common_dit.hpp"
-#include "diffusion_model.hpp"
-#include "model.h"
-#include "rope.hpp"
+#include "model/common/rope.hpp"
+#include "model/diffusion/dit.hpp"
+#include "model/diffusion/model.hpp"
+#include "model_loader.h"
 
 #define FLUX_GRAPH_SIZE 10240
 
 namespace Flux {
 
+    struct ChromaRadianceConfig {
+        int64_t nerf_hidden_size = 64;
+        int nerf_mlp_ratio       = 4;
+        int nerf_depth           = 4;
+        int nerf_max_freqs       = 8;
+        bool use_x0              = false;
+        bool fake_patch_size_x2  = false;
+    };
+
+    struct FluxConfig {
+        SDVersion version         = VERSION_FLUX;
+        bool is_chroma            = false;
+        int patch_size            = 2;
+        int64_t in_channels       = 64;
+        int64_t out_channels      = 64;
+        int64_t vec_in_dim        = 768;
+        int64_t context_in_dim    = 4096;
+        int64_t hidden_size       = 3072;
+        float mlp_ratio           = 4.0f;
+        int num_heads             = 24;
+        int depth                 = 19;
+        int depth_single_blocks   = 38;
+        std::vector<int> axes_dim = {16, 56, 56};
+        int axes_dim_sum          = 128;
+        int theta                 = 10000;
+        bool qkv_bias             = true;
+        bool guidance_embed       = true;
+        int64_t in_dim            = 64;
+        bool disable_bias         = false;
+        bool share_modulation     = false;
+        bool semantic_txt_norm    = false;
+        bool use_yak_mlp          = false;
+        bool use_mlp_silu_act     = false;
+        float ref_index_scale     = 1.f;
+        ChromaRadianceConfig chroma_radiance_params;
+
+        static FluxConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
+                                              const std::string& prefix,
+                                              SDVersion version = VERSION_FLUX) {
+            FluxConfig config;
+            config.version             = version;
+            config.guidance_embed      = false;
+            config.depth               = 0;
+            config.depth_single_blocks = 0;
+            if (version == VERSION_FLUX_FILL) {
+                config.in_channels = 384;
+            } else if (version == VERSION_FLUX_CONTROLS) {
+                config.in_channels = 128;
+            } else if (version == VERSION_FLEX_2) {
+                config.in_channels = 196;
+            } else if (version == VERSION_CHROMA_RADIANCE) {
+                config.in_channels = 3;
+                config.patch_size  = 16;
+            } else if (version == VERSION_OVIS_IMAGE) {
+                config.semantic_txt_norm = true;
+                config.use_yak_mlp       = true;
+                config.vec_in_dim        = 0;
+            } else if (sd_version_is_flux2(version)) {
+                config.in_channels      = 128;
+                config.patch_size       = 1;
+                config.out_channels     = 128;
+                config.mlp_ratio        = 3.f;
+                config.theta            = 2000;
+                config.axes_dim         = {32, 32, 32, 32};
+                config.vec_in_dim       = 0;
+                config.qkv_bias         = false;
+                config.disable_bias     = true;
+                config.share_modulation = true;
+                config.ref_index_scale  = 10.f;
+                config.use_mlp_silu_act = true;
+            } else if (sd_version_is_longcat(version)) {
+                config.context_in_dim = 3584;
+                config.vec_in_dim     = 0;
+            }
+
+            int64_t head_dim                   = 0;
+            int64_t actual_radiance_patch_size = -1;
+            for (const auto& [name, tensor_storage] : tensor_storage_map) {
+                if (!starts_with(name, prefix)) {
+                    continue;
+                }
+                if (name.find("guidance_in.in_layer.weight") != std::string::npos) {
+                    config.guidance_embed = true;
+                }
+                if (name.find("__x0__") != std::string::npos) {
+                    LOG_DEBUG("using x0 prediction");
+                    config.chroma_radiance_params.use_x0 = true;
+                }
+                if (name.find("__32x32__") != std::string::npos) {
+                    LOG_DEBUG("using patch size 32");
+                    config.patch_size = 32;
+                }
+                if (name.find("img_in_patch.weight") != std::string::npos) {
+                    actual_radiance_patch_size = tensor_storage.ne[0];
+                    LOG_DEBUG("actual radiance patch size: %" PRId64, actual_radiance_patch_size);
+                }
+                if (name.find("distilled_guidance_layer.in_proj.weight") != std::string::npos) {
+                    config.is_chroma = true;
+                }
+                size_t db = name.find("double_blocks.");
+                if (db != std::string::npos) {
+                    std::string block_name = name.substr(db);
+                    int block_depth        = atoi(block_name.substr(14, block_name.find(".", 14)).c_str());
+                    if (block_depth + 1 > config.depth) {
+                        config.depth = block_depth + 1;
+                    }
+                }
+                size_t sb = name.find("single_blocks.");
+                if (sb != std::string::npos) {
+                    std::string block_name = name.substr(sb);
+                    int block_depth        = atoi(block_name.substr(14, block_name.find(".", 14)).c_str());
+                    if (block_depth + 1 > config.depth_single_blocks) {
+                        config.depth_single_blocks = block_depth + 1;
+                    }
+                }
+                if (ends_with(name, "txt_in.weight")) {
+                    config.context_in_dim = tensor_storage.ne[0];
+                    config.hidden_size    = tensor_storage.ne[1];
+                }
+                if (ends_with(name, "single_blocks.0.norm.key_norm.scale")) {
+                    head_dim = tensor_storage.ne[0];
+                }
+                if (ends_with(name, "double_blocks.0.txt_attn.norm.key_norm.scale")) {
+                    head_dim = tensor_storage.ne[0];
+                }
+            }
+            if (actual_radiance_patch_size > 0 && actual_radiance_patch_size != config.patch_size) {
+                GGML_ASSERT(config.patch_size == 2 * actual_radiance_patch_size);
+                LOG_DEBUG("using fake x2 patch size");
+                config.chroma_radiance_params.fake_patch_size_x2 = true;
+            }
+            if (head_dim > 0) {
+                config.num_heads = static_cast<int>(config.hidden_size / head_dim);
+            }
+            config.axes_dim_sum = 0;
+            for (int axis_dim : config.axes_dim) {
+                config.axes_dim_sum += axis_dim;
+            }
+            LOG_DEBUG("flux: depth = %d, depth_single_blocks = %d, guidance_embed = %s, context_in_dim = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %d",
+                      config.depth,
+                      config.depth_single_blocks,
+                      config.guidance_embed ? "true" : "false",
+                      config.context_in_dim,
+                      config.hidden_size,
+                      config.num_heads);
+            return config;
+        }
+    };
+
     struct MLPEmbedder : public UnaryBlock {
     public:
         MLPEmbedder(int64_t in_dim, int64_t hidden_dim, bool bias = true) {
@@ -723,127 +872,90 @@ namespace Flux {
         }
     };
 
-    struct ChromaRadianceParams {
-        int64_t nerf_hidden_size = 64;
-        int nerf_mlp_ratio       = 4;
-        int nerf_depth           = 4;
-        int nerf_max_freqs       = 8;
-        bool use_x0              = false;
-        bool fake_patch_size_x2  = false;
-    };
-
-    struct FluxParams {
-        SDVersion version         = VERSION_FLUX;
-        bool is_chroma            = false;
-        int patch_size            = 2;
-        int64_t in_channels       = 64;
-        int64_t out_channels      = 64;
-        int64_t vec_in_dim        = 768;
-        int64_t context_in_dim    = 4096;
-        int64_t hidden_size       = 3072;
-        float mlp_ratio           = 4.0f;
-        int num_heads             = 24;
-        int depth                 = 19;
-        int depth_single_blocks   = 38;
-        std::vector<int> axes_dim = {16, 56, 56};
-        int axes_dim_sum          = 128;
-        int theta                 = 10000;
-        bool qkv_bias             = true;
-        bool guidance_embed       = true;
-        int64_t in_dim            = 64;
-        bool disable_bias         = false;
-        bool share_modulation     = false;
-        bool semantic_txt_norm    = false;
-        bool use_yak_mlp          = false;
-        bool use_mlp_silu_act     = false;
-        float ref_index_scale     = 1.f;
-        ChromaRadianceParams chroma_radiance_params;
-    };
-
     struct Flux : public GGMLBlock {
     public:
-        FluxParams params;
+        FluxConfig config;
         Flux() {}
-        Flux(FluxParams params)
-            : params(params) {
-            if (params.version == VERSION_CHROMA_RADIANCE) {
-                std::pair<int, int> kernel_size = {params.patch_size, params.patch_size};
-                if (params.chroma_radiance_params.fake_patch_size_x2) {
-                    kernel_size = {params.patch_size / 2, params.patch_size / 2};
+        Flux(FluxConfig config)
+            : config(config) {
+            if (config.version == VERSION_CHROMA_RADIANCE) {
+                std::pair<int, int> kernel_size = {config.patch_size, config.patch_size};
+                if (config.chroma_radiance_params.fake_patch_size_x2) {
+                    kernel_size = {config.patch_size / 2, config.patch_size / 2};
                 }
                 std::pair<int, int> stride = kernel_size;
 
-                blocks["img_in_patch"] = std::make_shared<Conv2d>(params.in_channels,
-                                                                  params.hidden_size,
+                blocks["img_in_patch"] = std::make_shared<Conv2d>(config.in_channels,
+                                                                  config.hidden_size,
                                                                   kernel_size,
                                                                   stride);
             } else {
-                blocks["img_in"] = std::make_shared<Linear>(params.in_channels, params.hidden_size, !params.disable_bias);
+                blocks["img_in"] = std::make_shared<Linear>(config.in_channels, config.hidden_size, !config.disable_bias);
             }
-            if (params.is_chroma) {
-                blocks["distilled_guidance_layer"] = std::make_shared<ChromaApproximator>(params.in_dim, params.hidden_size);
+            if (config.is_chroma) {
+                blocks["distilled_guidance_layer"] = std::make_shared<ChromaApproximator>(config.in_dim, config.hidden_size);
             } else {
-                blocks["time_in"] = std::make_shared<MLPEmbedder>(256, params.hidden_size, !params.disable_bias);
-                if (params.vec_in_dim > 0) {
-                    blocks["vector_in"] = std::make_shared<MLPEmbedder>(params.vec_in_dim, params.hidden_size, !params.disable_bias);
+                blocks["time_in"] = std::make_shared<MLPEmbedder>(256, config.hidden_size, !config.disable_bias);
+                if (config.vec_in_dim > 0) {
+                    blocks["vector_in"] = std::make_shared<MLPEmbedder>(config.vec_in_dim, config.hidden_size, !config.disable_bias);
                 }
-                if (params.guidance_embed) {
-                    blocks["guidance_in"] = std::make_shared<MLPEmbedder>(256, params.hidden_size, !params.disable_bias);
+                if (config.guidance_embed) {
+                    blocks["guidance_in"] = std::make_shared<MLPEmbedder>(256, config.hidden_size, !config.disable_bias);
                 }
             }
-            if (params.semantic_txt_norm) {
-                blocks["txt_norm"] = std::make_shared<RMSNorm>(params.context_in_dim);
+            if (config.semantic_txt_norm) {
+                blocks["txt_norm"] = std::make_shared<RMSNorm>(config.context_in_dim);
             }
-            blocks["txt_in"] = std::make_shared<Linear>(params.context_in_dim, params.hidden_size, !params.disable_bias);
+            blocks["txt_in"] = std::make_shared<Linear>(config.context_in_dim, config.hidden_size, !config.disable_bias);
 
-            for (int i = 0; i < params.depth; i++) {
-                blocks["double_blocks." + std::to_string(i)] = std::make_shared<DoubleStreamBlock>(params.hidden_size,
-                                                                                                   params.num_heads,
-                                                                                                   params.mlp_ratio,
+            for (int i = 0; i < config.depth; i++) {
+                blocks["double_blocks." + std::to_string(i)] = std::make_shared<DoubleStreamBlock>(config.hidden_size,
+                                                                                                   config.num_heads,
+                                                                                                   config.mlp_ratio,
                                                                                                    i,
-                                                                                                   params.qkv_bias,
-                                                                                                   params.is_chroma,
-                                                                                                   params.share_modulation,
-                                                                                                   !params.disable_bias,
-                                                                                                   params.use_yak_mlp,
-                                                                                                   params.use_mlp_silu_act);
+                                                                                                   config.qkv_bias,
+                                                                                                   config.is_chroma,
+                                                                                                   config.share_modulation,
+                                                                                                   !config.disable_bias,
+                                                                                                   config.use_yak_mlp,
+                                                                                                   config.use_mlp_silu_act);
             }
 
-            for (int i = 0; i < params.depth_single_blocks; i++) {
-                blocks["single_blocks." + std::to_string(i)] = std::make_shared<SingleStreamBlock>(params.hidden_size,
-                                                                                                   params.num_heads,
-                                                                                                   params.mlp_ratio,
+            for (int i = 0; i < config.depth_single_blocks; i++) {
+                blocks["single_blocks." + std::to_string(i)] = std::make_shared<SingleStreamBlock>(config.hidden_size,
+                                                                                                   config.num_heads,
+                                                                                                   config.mlp_ratio,
                                                                                                    i,
                                                                                                    0.f,
-                                                                                                   params.is_chroma,
-                                                                                                   params.share_modulation,
-                                                                                                   !params.disable_bias,
-                                                                                                   params.use_yak_mlp,
-                                                                                                   params.use_mlp_silu_act);
+                                                                                                   config.is_chroma,
+                                                                                                   config.share_modulation,
+                                                                                                   !config.disable_bias,
+                                                                                                   config.use_yak_mlp,
+                                                                                                   config.use_mlp_silu_act);
             }
 
-            if (params.version == VERSION_CHROMA_RADIANCE) {
-                blocks["nerf_image_embedder"] = std::make_shared<NerfEmbedder>(params.in_channels,
-                                                                               params.chroma_radiance_params.nerf_hidden_size,
-                                                                               params.chroma_radiance_params.nerf_max_freqs);
+            if (config.version == VERSION_CHROMA_RADIANCE) {
+                blocks["nerf_image_embedder"] = std::make_shared<NerfEmbedder>(config.in_channels,
+                                                                               config.chroma_radiance_params.nerf_hidden_size,
+                                                                               config.chroma_radiance_params.nerf_max_freqs);
 
-                for (int i = 0; i < params.chroma_radiance_params.nerf_depth; i++) {
-                    blocks["nerf_blocks." + std::to_string(i)] = std::make_shared<NerfGLUBlock>(params.hidden_size,
-                                                                                                params.chroma_radiance_params.nerf_hidden_size,
-                                                                                                params.chroma_radiance_params.nerf_mlp_ratio);
+                for (int i = 0; i < config.chroma_radiance_params.nerf_depth; i++) {
+                    blocks["nerf_blocks." + std::to_string(i)] = std::make_shared<NerfGLUBlock>(config.hidden_size,
+                                                                                                config.chroma_radiance_params.nerf_hidden_size,
+                                                                                                config.chroma_radiance_params.nerf_mlp_ratio);
                 }
 
-                blocks["nerf_final_layer_conv"] = std::make_shared<NerfFinalLayerConv>(params.chroma_radiance_params.nerf_hidden_size,
-                                                                                       params.in_channels);
+                blocks["nerf_final_layer_conv"] = std::make_shared<NerfFinalLayerConv>(config.chroma_radiance_params.nerf_hidden_size,
+                                                                                       config.in_channels);
 
             } else {
-                blocks["final_layer"] = std::make_shared<LastLayer>(params.hidden_size, 1, params.out_channels, params.is_chroma, !params.disable_bias);
+                blocks["final_layer"] = std::make_shared<LastLayer>(config.hidden_size, 1, config.out_channels, config.is_chroma, !config.disable_bias);
             }
 
-            if (params.share_modulation) {
-                blocks["double_stream_modulation_img"] = std::make_shared<Modulation>(params.hidden_size, true, !params.disable_bias);
-                blocks["double_stream_modulation_txt"] = std::make_shared<Modulation>(params.hidden_size, true, !params.disable_bias);
-                blocks["single_stream_modulation"]     = std::make_shared<Modulation>(params.hidden_size, false, !params.disable_bias);
+            if (config.share_modulation) {
+                blocks["double_stream_modulation_img"] = std::make_shared<Modulation>(config.hidden_size, true, !config.disable_bias);
+                blocks["double_stream_modulation_txt"] = std::make_shared<Modulation>(config.hidden_size, true, !config.disable_bias);
+                blocks["single_stream_modulation"]     = std::make_shared<Modulation>(config.hidden_size, false, !config.disable_bias);
             }
         }
 
@@ -866,7 +978,7 @@ namespace Flux {
 
             ggml_tensor* vec;
             ggml_tensor* txt_img_mask = nullptr;
-            if (params.is_chroma) {
+            if (config.is_chroma) {
                 int64_t mod_index_length = 344;
                 auto approx              = std::dynamic_pointer_cast<ChromaApproximator>(blocks["distilled_guidance_layer"]);
                 auto distill_timestep    = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, 16, 10000, 1000.f);
@@ -894,7 +1006,7 @@ namespace Flux {
             } else {
                 auto time_in = std::dynamic_pointer_cast<MLPEmbedder>(blocks["time_in"]);
                 vec          = time_in->forward(ctx, ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, 256, 10000, 1000.f));
-                if (params.guidance_embed) {
+                if (config.guidance_embed) {
                     GGML_ASSERT(guidance != nullptr);
                     auto guidance_in = std::dynamic_pointer_cast<MLPEmbedder>(blocks["guidance_in"]);
                     // bf16 and fp16 result is different
@@ -902,7 +1014,7 @@ namespace Flux {
                     vec       = ggml_add(ctx->ggml_ctx, vec, guidance_in->forward(ctx, g_in));
                 }
 
-                if (params.vec_in_dim > 0) {
+                if (config.vec_in_dim > 0) {
                     auto vector_in = std::dynamic_pointer_cast<MLPEmbedder>(blocks["vector_in"]);
                     vec            = ggml_add(ctx->ggml_ctx, vec, vector_in->forward(ctx, y));
                 }
@@ -911,7 +1023,7 @@ namespace Flux {
             std::vector<ModulationOut> ds_img_mods;
             std::vector<ModulationOut> ds_txt_mods;
             std::vector<ModulationOut> ss_mods;
-            if (params.share_modulation) {
+            if (config.share_modulation) {
                 auto double_stream_modulation_img = std::dynamic_pointer_cast<Modulation>(blocks["double_stream_modulation_img"]);
                 auto double_stream_modulation_txt = std::dynamic_pointer_cast<Modulation>(blocks["double_stream_modulation_txt"]);
                 auto single_stream_modulation     = std::dynamic_pointer_cast<Modulation>(blocks["single_stream_modulation"]);
@@ -921,7 +1033,7 @@ namespace Flux {
                 ss_mods     = single_stream_modulation->forward(ctx, vec);
             }
 
-            if (params.semantic_txt_norm) {
+            if (config.semantic_txt_norm) {
                 auto semantic_txt_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["txt_norm"]);
 
                 txt = semantic_txt_norm->forward(ctx, txt);
@@ -932,7 +1044,7 @@ namespace Flux {
             sd::ggml_graph_cut::mark_graph_cut(txt, "flux.prelude", "txt");
             sd::ggml_graph_cut::mark_graph_cut(vec, "flux.prelude", "vec");
 
-            for (int i = 0; i < params.depth; i++) {
+            for (int i = 0; i < config.depth; i++) {
                 if (skip_layers.size() > 0 && std::find(skip_layers.begin(), skip_layers.end(), i) != skip_layers.end()) {
                     continue;
                 }
@@ -947,8 +1059,8 @@ namespace Flux {
             }
 
             auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1);  // [N, n_txt_token + n_img_token, hidden_size]
-            for (int i = 0; i < params.depth_single_blocks; i++) {
-                if (skip_layers.size() > 0 && std::find(skip_layers.begin(), skip_layers.end(), i + params.depth) != skip_layers.end()) {
+            for (int i = 0; i < config.depth_single_blocks; i++) {
+                if (skip_layers.size() > 0 && std::find(skip_layers.begin(), skip_layers.end(), i + config.depth) != skip_layers.end()) {
                     continue;
                 }
                 auto block = std::dynamic_pointer_cast<SingleStreamBlock>(blocks["single_blocks." + std::to_string(i)]);
@@ -999,14 +1111,14 @@ namespace Flux {
             int64_t W      = x->ne[0];
             int64_t H      = x->ne[1];
             int64_t C      = x->ne[2];
-            int patch_size = params.patch_size;
+            int patch_size = config.patch_size;
             int pad_h      = (patch_size - H % patch_size) % patch_size;
             int pad_w      = (patch_size - W % patch_size) % patch_size;
 
-            auto img      = DiT::pad_to_patch_size(ctx, x, params.patch_size, params.patch_size);
+            auto img      = DiT::pad_to_patch_size(ctx, x, config.patch_size, config.patch_size);
             auto orig_img = img;
 
-            if (params.chroma_radiance_params.fake_patch_size_x2) {
+            if (config.chroma_radiance_params.fake_patch_size_x2) {
                 // It's supposed to be using GGML_SCALE_MODE_NEAREST, but this seems more stable
                 // Maybe the implementation of nearest-neighbor interpolation in ggml behaves differently than the one in PyTorch?
                 // img = F.interpolate(img, size=(H//2, W//2), mode="nearest")
@@ -1037,7 +1149,7 @@ namespace Flux {
             auto nerf_hidden = ggml_reshape_2d(ctx->ggml_ctx, out, out->ne[0], out->ne[1] * out->ne[2]);  // [N*num_patches, hidden_size]
             auto img_dct     = nerf_image_embedder->forward(ctx, nerf_pixels, dct);                       // [N*num_patches, patch_size*patch_size, nerf_hidden_size]
 
-            for (int i = 0; i < params.chroma_radiance_params.nerf_depth; i++) {
+            for (int i = 0; i < config.chroma_radiance_params.nerf_depth; i++) {
                 auto block = std::dynamic_pointer_cast<NerfGLUBlock>(blocks["nerf_blocks." + std::to_string(i)]);
 
                 img_dct = block->forward(ctx, img_dct, nerf_hidden);
@@ -1049,7 +1161,7 @@ namespace Flux {
 
             out = nerf_final_layer_conv->forward(ctx, img_dct);  // [N, C, H, W]
 
-            if (params.chroma_radiance_params.use_x0) {
+            if (config.chroma_radiance_params.use_x0) {
                 out = _apply_x0_residual(ctx, out, orig_img, timestep);
             }
 
@@ -1073,14 +1185,14 @@ namespace Flux {
             int64_t W      = x->ne[0];
             int64_t H      = x->ne[1];
             int64_t C      = x->ne[2];
-            int patch_size = params.patch_size;
+            int patch_size = config.patch_size;
             int pad_h      = (patch_size - H % patch_size) % patch_size;
             int pad_w      = (patch_size - W % patch_size) % patch_size;
 
             auto img           = DiT::pad_and_patchify(ctx, x, patch_size, patch_size);
             int64_t img_tokens = img->ne[1];
 
-            if (params.version == VERSION_FLUX_FILL) {
+            if (config.version == VERSION_FLUX_FILL) {
                 GGML_ASSERT(c_concat != nullptr);
                 ggml_tensor* masked = ggml_view_4d(ctx->ggml_ctx, c_concat, c_concat->ne[0], c_concat->ne[1], C, 1, c_concat->nb[1], c_concat->nb[2], c_concat->nb[3], 0);
                 ggml_tensor* mask   = ggml_view_4d(ctx->ggml_ctx, c_concat, c_concat->ne[0], c_concat->ne[1], 8 * 8, 1, c_concat->nb[1], c_concat->nb[2], c_concat->nb[3], c_concat->nb[2] * C);
@@ -1089,7 +1201,7 @@ namespace Flux {
                 mask   = DiT::pad_and_patchify(ctx, mask, patch_size, patch_size);
 
                 img = ggml_concat(ctx->ggml_ctx, img, ggml_concat(ctx->ggml_ctx, masked, mask, 0), 0);
-            } else if (params.version == VERSION_FLEX_2) {
+            } else if (config.version == VERSION_FLEX_2) {
                 GGML_ASSERT(c_concat != nullptr);
                 ggml_tensor* masked  = ggml_view_4d(ctx->ggml_ctx, c_concat, c_concat->ne[0], c_concat->ne[1], C, 1, c_concat->nb[1], c_concat->nb[2], c_concat->nb[3], 0);
                 ggml_tensor* mask    = ggml_view_4d(ctx->ggml_ctx, c_concat, c_concat->ne[0], c_concat->ne[1], 1, 1, c_concat->nb[1], c_concat->nb[2], c_concat->nb[3], c_concat->nb[2] * C);
@@ -1100,7 +1212,7 @@ namespace Flux {
                 control = DiT::pad_and_patchify(ctx, control, patch_size, patch_size);
 
                 img = ggml_concat(ctx->ggml_ctx, img, ggml_concat(ctx->ggml_ctx, ggml_concat(ctx->ggml_ctx, masked, mask, 0), control, 0), 0);
-            } else if (params.version == VERSION_FLUX_CONTROLS) {
+            } else if (config.version == VERSION_FLUX_CONTROLS) {
                 GGML_ASSERT(c_concat != nullptr);
 
                 auto control = DiT::pad_and_patchify(ctx, c_concat, patch_size, patch_size);
@@ -1147,7 +1259,7 @@ namespace Flux {
             // pe: (L, d_head/2, 2, 2)
             // return: (N, C, H, W)
 
-            if (params.version == VERSION_CHROMA_RADIANCE) {
+            if (config.version == VERSION_CHROMA_RADIANCE) {
                 return forward_chroma_radiance(ctx,
                                                x,
                                                timestep,
@@ -1179,7 +1291,7 @@ namespace Flux {
 
     struct FluxRunner : public DiffusionModelRunner {
     public:
-        FluxParams flux_params;
+        FluxConfig config;
         Flux flux;
         std::vector<float> pe_vec;
         std::vector<float> mod_index_arange_vec;
@@ -1189,119 +1301,20 @@ namespace Flux {
         bool use_mask = false;
 
         FluxRunner(ggml_backend_t backend,
-                   ggml_backend_t params_backend,
-                   const String2TensorStorage& tensor_storage_map = {},
-                   const std::string prefix                       = "",
-                   SDVersion version                              = VERSION_FLUX,
-                   bool use_mask                                  = false)
-            : DiffusionModelRunner(backend, params_backend, prefix), version(version), use_mask(use_mask) {
-            flux_params.version             = version;
-            flux_params.guidance_embed      = false;
-            flux_params.depth               = 0;
-            flux_params.depth_single_blocks = 0;
-            if (version == VERSION_FLUX_FILL) {
-                flux_params.in_channels = 384;
-            } else if (version == VERSION_FLUX_CONTROLS) {
-                flux_params.in_channels = 128;
-            } else if (version == VERSION_FLEX_2) {
-                flux_params.in_channels = 196;
-            } else if (version == VERSION_CHROMA_RADIANCE) {
-                flux_params.in_channels = 3;
-                flux_params.patch_size  = 16;
-            } else if (version == VERSION_OVIS_IMAGE) {
-                flux_params.semantic_txt_norm = true;
-                flux_params.use_yak_mlp       = true;
-                flux_params.vec_in_dim        = 0;
-            } else if (sd_version_is_flux2(version)) {
-                flux_params.in_channels      = 128;
-                flux_params.patch_size       = 1;
-                flux_params.out_channels     = 128;
-                flux_params.mlp_ratio        = 3.f;
-                flux_params.theta            = 2000;
-                flux_params.axes_dim         = {32, 32, 32, 32};
-                flux_params.vec_in_dim       = 0;
-                flux_params.qkv_bias         = false;
-                flux_params.disable_bias     = true;
-                flux_params.share_modulation = true;
-                flux_params.ref_index_scale  = 10.f;
-                flux_params.use_mlp_silu_act = true;
-            } else if (sd_version_is_longcat(version)) {
-                flux_params.context_in_dim = 3584;
-                flux_params.vec_in_dim     = 0;
-            }
-            int64_t head_dim                   = 0;
-            int64_t actual_radiance_patch_size = -1;
-            for (auto pair : tensor_storage_map) {
-                std::string tensor_name = pair.first;
-                if (!starts_with(tensor_name, prefix))
-                    continue;
-                if (tensor_name.find("guidance_in.in_layer.weight") != std::string::npos) {
-                    flux_params.guidance_embed = true;
-                }
-                if (tensor_name.find("__x0__") != std::string::npos) {
-                    LOG_DEBUG("using x0 prediction");
-                    flux_params.chroma_radiance_params.use_x0 = true;
-                }
-                if (tensor_name.find("__32x32__") != std::string::npos) {
-                    LOG_DEBUG("using patch size 32");
-                    flux_params.patch_size = 32;
-                }
-                if (tensor_name.find("img_in_patch.weight") != std::string::npos) {
-                    actual_radiance_patch_size = pair.second.ne[0];
-                    LOG_DEBUG("actual radiance patch size: %d", actual_radiance_patch_size);
-                }
-                if (tensor_name.find("distilled_guidance_layer.in_proj.weight") != std::string::npos) {
-                    // Chroma
-                    flux_params.is_chroma = true;
-                }
-                size_t db = tensor_name.find("double_blocks.");
-                if (db != std::string::npos) {
-                    tensor_name     = tensor_name.substr(db);  // remove prefix
-                    int block_depth = atoi(tensor_name.substr(14, tensor_name.find(".", 14)).c_str());
-                    if (block_depth + 1 > flux_params.depth) {
-                        flux_params.depth = block_depth + 1;
-                    }
-                }
-                size_t sb = tensor_name.find("single_blocks.");
-                if (sb != std::string::npos) {
-                    tensor_name     = tensor_name.substr(sb);  // remove prefix
-                    int block_depth = atoi(tensor_name.substr(14, tensor_name.find(".", 14)).c_str());
-                    if (block_depth + 1 > flux_params.depth_single_blocks) {
-                        flux_params.depth_single_blocks = block_depth + 1;
-                    }
-                }
-                if (ends_with(tensor_name, "txt_in.weight")) {
-                    flux_params.context_in_dim = pair.second.ne[0];
-                    flux_params.hidden_size    = pair.second.ne[1];
-                }
-                if (ends_with(tensor_name, "single_blocks.0.norm.key_norm.scale")) {
-                    head_dim = pair.second.ne[0];
-                }
-                if (ends_with(tensor_name, "double_blocks.0.txt_attn.norm.key_norm.scale")) {
-                    head_dim = pair.second.ne[0];
-                }
-            }
-            if (actual_radiance_patch_size > 0 && actual_radiance_patch_size != flux_params.patch_size) {
-                GGML_ASSERT(flux_params.patch_size == 2 * actual_radiance_patch_size);
-                LOG_DEBUG("using fake x2 patch size");
-                flux_params.chroma_radiance_params.fake_patch_size_x2 = true;
-            }
-
-            flux_params.num_heads = static_cast<int>(flux_params.hidden_size / head_dim);
-
-            LOG_INFO("flux: depth = %d, depth_single_blocks = %d, guidance_embed = %s, context_in_dim = %" PRId64
-                     ", hidden_size = %" PRId64 ", num_heads = %d",
-                     flux_params.depth,
-                     flux_params.depth_single_blocks,
-                     flux_params.guidance_embed ? "true" : "false",
-                     flux_params.context_in_dim,
-                     flux_params.hidden_size,
-                     flux_params.num_heads);
-            if (flux_params.is_chroma) {
+                   const String2TensorStorage& tensor_storage_map      = {},
+                   const std::string prefix                            = "",
+                   SDVersion version                                   = VERSION_FLUX,
+                   bool use_mask                                       = false,
+                   std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : DiffusionModelRunner(backend, prefix, weight_manager),
+              config(FluxConfig::detect_from_weights(tensor_storage_map, prefix, version)),
+              version(version),
+              use_mask(use_mask) {
+            if (config.is_chroma) {
                 LOG_INFO("Using pruned modulation (Chroma)");
             }
 
-            flux = Flux(flux_params);
+            flux = Flux(config);
             flux.init(params_ctx, tensor_storage_map, prefix);
         }
 
@@ -1377,10 +1390,10 @@ namespace Flux {
             ggml_tensor* context   = make_optional_input(context_tensor);
             ggml_tensor* c_concat  = make_optional_input(c_concat_tensor);
             ggml_tensor* y         = make_optional_input(y_tensor);
-            if (flux_params.guidance_embed || flux_params.is_chroma) {
+            if (config.guidance_embed || config.is_chroma) {
                 if (!guidance_tensor.empty()) {
                     this->guidance_tensor = guidance_tensor;
-                    if (flux_params.is_chroma) {
+                    if (config.is_chroma) {
                         this->guidance_tensor.fill_(0.f);
                     }
                 }
@@ -1398,7 +1411,7 @@ namespace Flux {
             ggml_tensor* mod_index_arange = nullptr;
             ggml_tensor* dct              = nullptr;  // for chroma radiance
 
-            if (flux_params.is_chroma) {
+            if (config.is_chroma) {
                 if (!use_mask) {
                     y = nullptr;
                 }
@@ -1417,29 +1430,29 @@ namespace Flux {
             }
             pe_vec      = Rope::gen_flux_pe(static_cast<int>(x->ne[1]),
                                             static_cast<int>(x->ne[0]),
-                                            flux_params.patch_size,
+                                            config.patch_size,
                                             static_cast<int>(x->ne[3]),
                                             static_cast<int>(context->ne[1]),
                                             txt_arange_dims,
                                             ref_latents,
                                             increase_ref_index,
-                                            flux_params.ref_index_scale,
-                                            flux_params.theta,
+                                            config.ref_index_scale,
+                                            config.theta,
                                             circular_y_enabled,
                                             circular_x_enabled,
-                                            flux_params.axes_dim,
+                                            config.axes_dim,
                                             sd_version_is_longcat(version));
-            int pos_len = static_cast<int>(pe_vec.size() / flux_params.axes_dim_sum / 2);
+            int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
             // LOG_DEBUG("pos_len %d", pos_len);
-            auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, flux_params.axes_dim_sum / 2, pos_len);
+            auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
             // pe->data = pe_vec.data();
             // print_ggml_tensor(pe);
             // pe->data = nullptr;
             set_backend_tensor_data(pe, pe_vec.data());
 
             if (version == VERSION_CHROMA_RADIANCE) {
-                int patch_size     = flux_params.patch_size;
-                int nerf_max_freqs = flux_params.chroma_radiance_params.nerf_max_freqs;
+                int patch_size     = config.patch_size;
+                int nerf_max_freqs = config.chroma_radiance_params.nerf_max_freqs;
                 dct_vec            = fetch_dct_pos(patch_size, nerf_max_freqs);
                 dct                = ggml_new_tensor_2d(compute_ctx, GGML_TYPE_F32, nerf_max_freqs * nerf_max_freqs, patch_size * patch_size);
                 // dct->data = dct_vec.data();
@@ -1487,7 +1500,7 @@ namespace Flux {
                 return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, increase_ref_index, skip_layers);
             };
 
-            auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
             return result;
         }
 
@@ -1570,7 +1583,8 @@ namespace Flux {
             ggml_backend_t backend    = sd_backend_cpu_init();
             ggml_type model_data_type = GGML_TYPE_COUNT;
 
-            ModelLoader model_loader;
+            auto model_manager        = std::make_shared<ModelManager>();
+            ModelLoader& model_loader = model_manager->loader();
             if (!model_loader.init_from_file_and_convert_name(file_path, "model.diffusion_model.")) {
                 LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
                 return;
@@ -1586,24 +1600,20 @@ namespace Flux {
             }
 
             std::shared_ptr<FluxRunner> flux = std::make_shared<FluxRunner>(backend,
-                                                                            backend,
                                                                             tensor_storage_map,
                                                                             "model.diffusion_model",
                                                                             VERSION_FLUX2,
-                                                                            false);
+                                                                            false,
+                                                                            model_manager);
 
-            if (!flux->alloc_params_buffer()) {
-                LOG_ERROR("flux model allocation failed");
-                return;
-            }
-
-            std::map<std::string, ggml_tensor*> tensors;
-            flux->get_param_tensors(tensors, "model.diffusion_model");
-
-            bool success = model_loader.load_tensors(tensors);
-
-            if (!success) {
-                LOG_ERROR("load tensors from model loader failed");
+            if (!model_manager->register_runner_params("Flux test",
+                                                       *flux,
+                                                       "model.diffusion_model",
+                                                       ModelManager::ResidencyMode::ParamBackend,
+                                                       backend,
+                                                       backend) ||
+                !model_manager->validate_registered_tensors()) {
+                LOG_ERROR("register flux tensors with model manager failed");
                 return;
             }
 
@@ -1614,4 +1624,4 @@ namespace Flux {
 
 }  // namespace Flux
 
-#endif  // __FLUX_HPP__
+#endif  // __SD_MODEL_DIFFUSION_FLUX_HPP__

src/model/diffusion/hidream_o1.hpp

@@ -1,5 +1,5 @@
-#ifndef __SD_HIDREAM_O1_H__
-#define __SD_HIDREAM_O1_H__
+#ifndef __SD_MODEL_DIFFUSION_HIDREAM_O1_HPP__
+#define __SD_MODEL_DIFFUSION_HIDREAM_O1_HPP__
 
 #include <algorithm>
 #include <array>
@@ -10,11 +10,11 @@
 #include <utility>
 #include <vector>
 
-#include "common_dit.hpp"
-#include "conditioner.hpp"
-#include "diffusion_model.hpp"
-#include "llm.hpp"
-#include "util.h"
+#include "conditioning/conditioner.hpp"
+#include "core/util.h"
+#include "model/diffusion/dit.hpp"
+#include "model/diffusion/model.hpp"
+#include "model/te/llm.hpp"
 
 namespace HiDreamO1 {
     constexpr int HIDREAM_O1_GRAPH_SIZE = 32768;
@@ -23,6 +23,39 @@ namespace HiDreamO1 {
     constexpr int IMAGE_TOKEN_ID        = 151655;
     constexpr int VISION_START_TOKEN_ID = 151652;
 
+    struct HiDreamO1Config {
+        LLM::LLMConfig llm;
+        int patch_size = PATCH_SIZE;
+
+        static HiDreamO1Config detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
+            (void)tensor_storage_map;
+            (void)prefix;
+            HiDreamO1Config config;
+            config.llm.arch                           = LLM::LLMArch::QWEN3_VL;
+            config.llm.hidden_size                    = 4096;
+            config.llm.intermediate_size              = 12288;
+            config.llm.num_layers                     = 36;
+            config.llm.num_heads                      = 32;
+            config.llm.num_kv_heads                   = 8;
+            config.llm.head_dim                       = 128;
+            config.llm.qkv_bias                       = false;
+            config.llm.qk_norm                        = true;
+            config.llm.vocab_size                     = 151936;
+            config.llm.rms_norm_eps                   = 1e-6f;
+            config.llm.vision.arch                    = LLM::LLMVisionArch::QWEN3_VL;
+            config.llm.vision.num_layers              = 27;
+            config.llm.vision.hidden_size             = 1152;
+            config.llm.vision.intermediate_size       = 4304;
+            config.llm.vision.num_heads               = 16;
+            config.llm.vision.out_hidden_size         = 4096;
+            config.llm.vision.patch_size              = 16;
+            config.llm.vision.spatial_merge_size      = 2;
+            config.llm.vision.temporal_patch_size     = 2;
+            config.llm.vision.num_position_embeddings = 2304;
+            return config;
+        }
+    };
+
     static inline std::string repeat_special_token(const std::string& token, int64_t count) {
         std::string out;
         out.reserve(static_cast<size_t>(count) * token.size());
@@ -205,50 +238,19 @@ namespace HiDreamO1 {
         }
     };
 
-    struct HiDreamO1Params {
-        LLM::LLMParams llm;
-        int patch_size = PATCH_SIZE;
-    };
-
-    static inline HiDreamO1Params make_hidream_o1_params() {
-        HiDreamO1Params params;
-        params.llm.arch                           = LLM::LLMArch::QWEN3_VL;
-        params.llm.hidden_size                    = 4096;
-        params.llm.intermediate_size              = 12288;
-        params.llm.num_layers                     = 36;
-        params.llm.num_heads                      = 32;
-        params.llm.num_kv_heads                   = 8;
-        params.llm.head_dim                       = 128;
-        params.llm.qkv_bias                       = false;
-        params.llm.qk_norm                        = true;
-        params.llm.vocab_size                     = 151936;
-        params.llm.rms_norm_eps                   = 1e-6f;
-        params.llm.vision.arch                    = LLM::LLMVisionArch::QWEN3_VL;
-        params.llm.vision.num_layers              = 27;
-        params.llm.vision.hidden_size             = 1152;
-        params.llm.vision.intermediate_size       = 4304;
-        params.llm.vision.num_heads               = 16;
-        params.llm.vision.out_hidden_size         = 4096;
-        params.llm.vision.patch_size              = 16;
-        params.llm.vision.spatial_merge_size      = 2;
-        params.llm.vision.temporal_patch_size     = 2;
-        params.llm.vision.num_position_embeddings = 2304;
-        return params;
-    }
-
     struct HiDreamO1Model : public GGMLBlock {
-        HiDreamO1Params params;
+        HiDreamO1Config config;
 
         HiDreamO1Model() = default;
-        explicit HiDreamO1Model(HiDreamO1Params params)
-            : params(std::move(params)) {
-            blocks["language_model"] = std::make_shared<LLM::TextModel>(this->params.llm);
-            blocks["t_embedder1"]    = std::make_shared<TimestepEmbedder>(this->params.llm.hidden_size);
-            blocks["x_embedder"]     = std::make_shared<BottleneckPatchEmbed>(this->params.patch_size * this->params.patch_size * 3,
-                                                                          this->params.llm.hidden_size / 4,
-                                                                          this->params.llm.hidden_size);
-            blocks["final_layer2"]   = std::make_shared<FinalLayer>(this->params.llm.hidden_size,
-                                                                  this->params.patch_size * this->params.patch_size * 3);
+        explicit HiDreamO1Model(HiDreamO1Config config)
+            : config(std::move(config)) {
+            blocks["language_model"] = std::make_shared<LLM::TextModel>(this->config.llm);
+            blocks["t_embedder1"]    = std::make_shared<TimestepEmbedder>(this->config.llm.hidden_size);
+            blocks["x_embedder"]     = std::make_shared<BottleneckPatchEmbed>(this->config.patch_size * this->config.patch_size * 3,
+                                                                          this->config.llm.hidden_size / 4,
+                                                                          this->config.llm.hidden_size);
+            blocks["final_layer2"]   = std::make_shared<FinalLayer>(this->config.llm.hidden_size,
+                                                                  this->config.patch_size * this->config.patch_size * 3);
         }
 
         std::shared_ptr<LLM::TextModel> text_model() {
@@ -269,7 +271,7 @@ namespace HiDreamO1 {
     };
 
     struct HiDreamO1VisionRunner : public GGMLRunner {
-        HiDreamO1Params params;
+        HiDreamO1Config config;
         std::shared_ptr<LLM::VisionModel> model;
 
         std::vector<int> window_index_vec;
@@ -280,12 +282,12 @@ namespace HiDreamO1 {
         std::array<std::vector<float>, 4> pos_embed_weight_data_;
 
         HiDreamO1VisionRunner(ggml_backend_t backend,
-                              ggml_backend_t params_backend,
-                              const String2TensorStorage& tensor_storage_map = {},
-                              const std::string& prefix                      = "model.visual")
-            : GGMLRunner(backend, params_backend),
-              params(make_hidream_o1_params()),
-              model(std::make_shared<LLM::VisionModel>(false, params.llm.vision)) {
+                              const String2TensorStorage& tensor_storage_map      = {},
+                              const std::string& prefix                           = "model.visual",
+                              std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : GGMLRunner(backend, weight_manager),
+              config(HiDreamO1Config::detect_from_weights(tensor_storage_map, prefix)),
+              model(std::make_shared<LLM::VisionModel>(false, config.llm.vision)) {
             model->init(params_ctx, tensor_storage_map, prefix);
         }
 
@@ -302,7 +304,7 @@ namespace HiDreamO1 {
                                                        compute_ctx,
                                                        runner_ctx,
                                                        image,
-                                                       params.llm.vision,
+                                                       config.llm.vision,
                                                        model,
                                                        window_index_vec,
                                                        window_inverse_index_vec,
@@ -321,28 +323,32 @@ namespace HiDreamO1 {
             return gf;
         }
 
-        sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& image) {
+        sd::Tensor<float> compute(int n_threads,
+                                  const sd::Tensor<float>& image,
+                                  bool auto_free           = true,
+                                  bool free_compute_buffer = true,
+                                  bool free_compute_params = true) {
             auto get_graph = [&]() {
                 return build_graph(image);
             };
-            auto output = GGMLRunner::compute<float>(get_graph, n_threads, false);
+            auto output = GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params);
             return output.has_value() ? std::move(output.value()) : sd::Tensor<float>();
         }
     };
 
     struct HiDreamO1Runner : public DiffusionModelRunner {
-        HiDreamO1Params params;
+        HiDreamO1Config config;
         HiDreamO1Model model;
 
         std::vector<float> attention_mask_vec;
 
         HiDreamO1Runner(ggml_backend_t backend,
-                        ggml_backend_t params_backend,
-                        const String2TensorStorage& tensor_storage_map = {},
-                        const std::string& prefix                      = "model")
-            : DiffusionModelRunner(backend, params_backend, prefix),
-              params(make_hidream_o1_params()) {
-            model = HiDreamO1Model(params);
+                        const String2TensorStorage& tensor_storage_map      = {},
+                        const std::string& prefix                           = "model",
+                        std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : DiffusionModelRunner(backend, prefix, weight_manager),
+              config(HiDreamO1Config::detect_from_weights(tensor_storage_map, prefix)) {
+            model = HiDreamO1Model(config);
             model.init(params_ctx, tensor_storage_map, prefix);
         }
 
@@ -453,7 +459,7 @@ namespace HiDreamO1 {
             auto get_graph = [&]() {
                 return build_graph(x, timestep, input_ids, input_pos, token_types, vinput_mask, image_embeds, ref_images);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,
@@ -484,29 +490,14 @@ namespace HiDreamO1 {
         std::shared_ptr<HiDreamO1VisionRunner> vision_runner;
 
         HiDreamO1Conditioner(ggml_backend_t backend,
-                             ggml_backend_t params_backend,
-                             const String2TensorStorage& tensor_storage_map = {})
-            : vision_runner(std::make_shared<HiDreamO1VisionRunner>(backend, params_backend, tensor_storage_map)) {}
+                             const String2TensorStorage& tensor_storage_map      = {},
+                             std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : vision_runner(std::make_shared<HiDreamO1VisionRunner>(backend, tensor_storage_map, "model.visual", weight_manager)) {}
 
         void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
             vision_runner->get_param_tensors(tensors);
         }
 
-        bool alloc_params_buffer() override {
-            if (!vision_runner->alloc_params_buffer()) {
-                return false;
-            }
-            return true;
-        }
-
-        void free_params_buffer() override {
-            vision_runner->free_params_buffer();
-        }
-
-        size_t get_params_buffer_size() override {
-            return vision_runner->get_params_buffer_size();
-        }
-
         void set_max_graph_vram_bytes(size_t max_graph_vram_bytes) override {
             vision_runner->set_max_graph_vram_bytes(max_graph_vram_bytes);
         }
@@ -519,6 +510,10 @@ namespace HiDreamO1 {
             vision_runner->set_weight_adapter(adapter);
         }
 
+        void runner_done() override {
+            vision_runner->runner_done();
+        }
+
         SDCondition get_learned_condition(int n_threads,
                                           const ConditionerParams& conditioner_params) override {
             SDCondition result;
@@ -664,7 +659,7 @@ namespace HiDreamO1 {
             result.c_vinput_mask  = sd::Tensor<int32_t>(vinput_mask_shape, std::move(vinput_mask));
             result.c_image_embeds.reserve(vlm_images.size());
             for (const auto& vlm_image : vlm_images) {
-                auto image_embed = vision_runner->compute(n_threads, vlm_image.second);
+                auto image_embed = vision_runner->compute(n_threads, vlm_image.second, false, true, true);
                 if (image_embed.empty()) {
                     LOG_ERROR("hidream_o1 conditioner: encode VLM image failed");
                     return SDCondition();
@@ -676,4 +671,4 @@ namespace HiDreamO1 {
     };
 }  // namespace HiDreamO1
 
-#endif  // __SD_HIDREAM_O1_H__
+#endif  // __SD_MODEL_DIFFUSION_HIDREAM_O1_HPP__

src/model/diffusion/ideogram4.hpp

@@ -0,0 +1,557 @@
+#ifndef __SD_MODEL_DIFFUSION_IDEOGRAM4_HPP__
+#define __SD_MODEL_DIFFUSION_IDEOGRAM4_HPP__
+
+#include <algorithm>
+#include <cmath>
+#include <cstdlib>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "core/ggml_extend.hpp"
+#include "core/ggml_graph_cut.h"
+#include "model/common/rope.hpp"
+#include "model/diffusion/model.hpp"
+
+namespace Ideogram4 {
+    constexpr int IDEOGRAM4_GRAPH_SIZE    = 65536;
+    constexpr int OUTPUT_IMAGE_INDICATOR  = 2;
+    constexpr int IMAGE_POSITION_OFFSET   = 65536;
+    constexpr int DEFAULT_MROPE_SECTION_T = 24;
+    constexpr int DEFAULT_MROPE_SECTION_H = 20;
+    constexpr int DEFAULT_MROPE_SECTION_W = 20;
+    constexpr int TIMESTEP_MAX_PERIOD     = 10000;
+    constexpr int LLM_HIDDEN_STATE_LAYERS = 13;
+
+    struct Ideogram4Config {
+        int64_t emb_dim                = 4608;
+        int64_t num_layers             = 34;
+        int64_t num_heads              = 18;
+        int64_t intermediate_size      = 12288;
+        int64_t adanln_dim             = 512;
+        int64_t in_channels            = 128;
+        int64_t llm_features_dim       = 53248;
+        int64_t rope_theta             = 5000000;
+        float norm_eps                 = 1e-5f;
+        int patch_size                 = 2;
+        int ae_channels                = 32;
+        std::vector<int> mrope_section = {DEFAULT_MROPE_SECTION_T,
+                                          DEFAULT_MROPE_SECTION_H,
+                                          DEFAULT_MROPE_SECTION_W};
+
+        static Ideogram4Config detect_from_weights(const String2TensorStorage& tensor_storage_map,
+                                                   const std::string& prefix) {
+            Ideogram4Config config;
+            int64_t detected_layers  = 0;
+            std::string layer_prefix = prefix.empty() ? "layers." : prefix + ".layers.";
+            for (const auto& [name, _] : tensor_storage_map) {
+                if (name.find(layer_prefix) != 0) {
+                    continue;
+                }
+                std::string tail = name.substr(layer_prefix.size());
+                size_t dot       = tail.find('.');
+                if (dot == std::string::npos) {
+                    continue;
+                }
+                int layer_idx   = std::atoi(tail.substr(0, dot).c_str());
+                detected_layers = std::max<int64_t>(detected_layers, layer_idx + 1);
+            }
+            if (detected_layers > 0) {
+                config.num_layers = detected_layers;
+                LOG_DEBUG("ideogram4: num_layers = %" PRId64 ", emb_dim = %" PRId64 ", num_heads = %" PRId64 ", intermediate_size = %" PRId64,
+                          config.num_layers,
+                          config.emb_dim,
+                          config.num_heads,
+                          config.intermediate_size);
+            }
+            return config;
+        }
+    };
+
+    __STATIC_INLINE__ ggml_tensor* timestep_embedding_sin_cos(ggml_context* ctx,
+                                                              ggml_tensor* timesteps,
+                                                              int dim) {
+        GGML_ASSERT(dim % 2 == 0);
+        auto embedding = ggml_ext_timestep_embedding(ctx, timesteps, dim, TIMESTEP_MAX_PERIOD, 10.f);
+        auto chunks    = ggml_ext_chunk(ctx, embedding, 2, 0);
+        return ggml_concat(ctx, chunks[1], chunks[0], 0);
+    }
+
+    __STATIC_INLINE__ ggml_tensor* to_token_modulation(ggml_context* ctx, ggml_tensor* x) {
+        // [N, C] -> [N, 1, C] in PyTorch layout.
+        if (ggml_n_dims(x) < 3 || x->ne[1] != 1) {
+            x = ggml_reshape_3d(ctx, x, x->ne[0], 1, x->ne[1]);
+        }
+        return x;
+    }
+
+    __STATIC_INLINE__ ggml_tensor* interleave_hidden_state_layers(ggml_context* ctx, ggml_tensor* x) {
+        // Match upstream stack(...).permute(1, 2, 3, 0).reshape(...):
+        // [layers * hidden, tokens, batch] -> [hidden * layers, tokens, batch].
+        GGML_ASSERT(x->ne[0] % LLM_HIDDEN_STATE_LAYERS == 0);
+        const int64_t hidden_size = x->ne[0] / LLM_HIDDEN_STATE_LAYERS;
+        const int64_t token_count = x->ne[1];
+        const int64_t batch_count = x->ne[2];
+
+        x = ggml_reshape_4d(ctx, x, hidden_size, LLM_HIDDEN_STATE_LAYERS, token_count, batch_count);
+        x = ggml_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3));
+        return ggml_reshape_3d(ctx, x, hidden_size * LLM_HIDDEN_STATE_LAYERS, token_count, batch_count);
+    }
+
+    __STATIC_INLINE__ ggml_tensor* modulate(ggml_context* ctx, ggml_tensor* x, ggml_tensor* scale) {
+        scale = to_token_modulation(ctx, scale);
+        return ggml_add(ctx, x, ggml_mul(ctx, x, scale));
+    }
+
+    __STATIC_INLINE__ ggml_tensor* patchify(ggml_context* ctx, ggml_tensor* x, const Ideogram4Config& config) {
+        // x: [N, 128, H, W] with channel order [ae, ph, pw].
+        // return: [N, H*W, 128] with token channel order [ph, pw, ae].
+        const int64_t W = x->ne[0];
+        const int64_t H = x->ne[1];
+        const int64_t C = x->ne[2];
+        const int64_t N = x->ne[3];
+
+        GGML_ASSERT(N == 1);
+        GGML_ASSERT(C == config.ae_channels * config.patch_size * config.patch_size);
+
+        x = ggml_cont(ctx, x);
+        x = ggml_reshape_4d(ctx, x, W * H, config.patch_size, config.patch_size, config.ae_channels);
+        x = ggml_cont(ctx, ggml_permute(ctx, x, 3, 1, 2, 0));
+        x = ggml_reshape_3d(ctx, x, C, W * H, N);
+        return x;
+    }
+
+    __STATIC_INLINE__ ggml_tensor* unpatchify(ggml_context* ctx,
+                                              ggml_tensor* x,
+                                              int64_t H,
+                                              int64_t W,
+                                              const Ideogram4Config& config) {
+        const int64_t C = x->ne[0];
+        const int64_t N = x->ne[2];
+
+        GGML_ASSERT(N == 1);
+        GGML_ASSERT(C == config.ae_channels * config.patch_size * config.patch_size);
+        GGML_ASSERT(x->ne[1] == H * W);
+
+        x = ggml_reshape_4d(ctx, x, config.ae_channels, config.patch_size, config.patch_size, H * W);
+        x = ggml_cont(ctx, ggml_permute(ctx, x, 3, 1, 2, 0));
+        x = ggml_reshape_4d(ctx, x, W, H, C, N);
+        return x;
+    }
+
+    __STATIC_INLINE__ std::shared_ptr<Linear> make_linear(int64_t in_features,
+                                                          int64_t out_features,
+                                                          bool bias = true) {
+        return std::make_shared<Linear>(in_features, out_features, bias, false, false, 1.f, true);
+    }
+
+    __STATIC_INLINE__ std::vector<float> gen_ideogram4_pe(int grid_h,
+                                                          int grid_w,
+                                                          int bs,
+                                                          int context_len,
+                                                          int head_dim,
+                                                          int rope_theta,
+                                                          const std::vector<int>& mrope_section,
+                                                          bool circular_x = false,
+                                                          bool circular_y = false) {
+        GGML_ASSERT(bs == 1);
+        std::vector<std::vector<float>> ids(static_cast<size_t>(bs) * (context_len + grid_h * grid_w),
+                                            std::vector<float>(3, 0.f));
+
+        for (int i = 0; i < context_len; ++i) {
+            ids[i] = {static_cast<float>(i), static_cast<float>(i), static_cast<float>(i)};
+        }
+
+        int cursor = context_len;
+        for (int y = 0; y < grid_h; ++y) {
+            for (int x = 0; x < grid_w; ++x) {
+                ids[cursor++] = {static_cast<float>(IMAGE_POSITION_OFFSET),
+                                 static_cast<float>(IMAGE_POSITION_OFFSET + y),
+                                 static_cast<float>(IMAGE_POSITION_OFFSET + x)};
+            }
+        }
+
+        std::vector<std::vector<int>> axis_wrap_dims(3);
+        if (circular_y || circular_x) {
+            size_t total_len = static_cast<size_t>(bs) * (context_len + grid_h * grid_w);
+            axis_wrap_dims[1].assign(total_len, 0);
+            axis_wrap_dims[2].assign(total_len, 0);
+            if (circular_y) {
+                for (size_t idx = static_cast<size_t>(context_len); idx < total_len; ++idx) {
+                    axis_wrap_dims[1][idx] = grid_h;
+                }
+            }
+            if (circular_x) {
+                for (size_t idx = static_cast<size_t>(context_len); idx < total_len; ++idx) {
+                    axis_wrap_dims[2][idx] = grid_w;
+                }
+            }
+        }
+
+        return Rope::embed_interleaved_mrope(ids,
+                                             bs,
+                                             static_cast<float>(rope_theta),
+                                             head_dim,
+                                             mrope_section,
+                                             axis_wrap_dims);
+    }
+
+    class Ideogram4Attention : public GGMLBlock {
+    protected:
+        int64_t hidden_size;
+        int64_t num_heads;
+        int64_t head_dim;
+
+    public:
+        Ideogram4Attention(int64_t hidden_size, int64_t num_heads, float eps)
+            : hidden_size(hidden_size), num_heads(num_heads), head_dim(hidden_size / num_heads) {
+            GGML_ASSERT(hidden_size % num_heads == 0);
+            blocks["qkv"]    = make_linear(hidden_size, hidden_size * 3, false);
+            blocks["norm_q"] = std::make_shared<RMSNorm>(head_dim, eps);
+            blocks["norm_k"] = std::make_shared<RMSNorm>(head_dim, eps);
+            blocks["o"]      = make_linear(hidden_size, hidden_size, false);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* pe,
+                             ggml_tensor* mask = nullptr) {
+            int64_t n_token = x->ne[1];
+            int64_t N       = x->ne[2];
+
+            auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
+            auto norm_q   = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_q"]);
+            auto norm_k   = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"]);
+            auto out_proj = std::dynamic_pointer_cast<Linear>(blocks["o"]);
+
+            auto qkv     = qkv_proj->forward(ctx, x);
+            auto qkv_vec = split_qkv(ctx->ggml_ctx, qkv);
+            auto q       = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[0], head_dim, num_heads, n_token, N);
+            auto k       = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[1], head_dim, num_heads, n_token, N);
+            auto v       = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[2], head_dim, num_heads, n_token, N);
+
+            q = norm_q->forward(ctx, q);
+            k = norm_k->forward(ctx, k);
+
+            x = Rope::attention(ctx, q, k, v, pe, mask, 1.f / 128.f, false);
+            x = out_proj->forward(ctx, x);
+            return x;
+        }
+    };
+
+    class Ideogram4MLP : public GGMLBlock {
+    public:
+        Ideogram4MLP(int64_t dim, int64_t hidden_dim) {
+            blocks["w1"] = make_linear(dim, hidden_dim, false);
+            blocks["w2"] = make_linear(hidden_dim, dim, false);
+            blocks["w3"] = make_linear(dim, hidden_dim, false);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            auto w1 = std::dynamic_pointer_cast<Linear>(blocks["w1"]);
+            auto w2 = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
+            auto w3 = std::dynamic_pointer_cast<Linear>(blocks["w3"]);
+
+            auto x1 = ggml_silu(ctx->ggml_ctx, w1->forward(ctx, x));
+            auto x3 = w3->forward(ctx, x);
+            x       = ggml_mul(ctx->ggml_ctx, x1, x3);
+            x       = w2->forward(ctx, x);
+            return x;
+        }
+    };
+
+    class Ideogram4TransformerBlock : public GGMLBlock {
+    public:
+        Ideogram4TransformerBlock(const Ideogram4Config& config) {
+            blocks["attention"]        = std::make_shared<Ideogram4Attention>(config.emb_dim, config.num_heads, config.norm_eps);
+            blocks["feed_forward"]     = std::make_shared<Ideogram4MLP>(config.emb_dim, config.intermediate_size);
+            blocks["attention_norm1"]  = std::make_shared<RMSNorm>(config.emb_dim, config.norm_eps);
+            blocks["ffn_norm1"]        = std::make_shared<RMSNorm>(config.emb_dim, config.norm_eps);
+            blocks["attention_norm2"]  = std::make_shared<RMSNorm>(config.emb_dim, config.norm_eps);
+            blocks["ffn_norm2"]        = std::make_shared<RMSNorm>(config.emb_dim, config.norm_eps);
+            blocks["adaln_modulation"] = make_linear(config.adanln_dim, 4 * config.emb_dim, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* pe,
+                             ggml_tensor* adaln_input,
+                             ggml_tensor* mask = nullptr) {
+            auto attention        = std::dynamic_pointer_cast<Ideogram4Attention>(blocks["attention"]);
+            auto feed_forward     = std::dynamic_pointer_cast<Ideogram4MLP>(blocks["feed_forward"]);
+            auto attention_norm1  = std::dynamic_pointer_cast<RMSNorm>(blocks["attention_norm1"]);
+            auto ffn_norm1        = std::dynamic_pointer_cast<RMSNorm>(blocks["ffn_norm1"]);
+            auto attention_norm2  = std::dynamic_pointer_cast<RMSNorm>(blocks["attention_norm2"]);
+            auto ffn_norm2        = std::dynamic_pointer_cast<RMSNorm>(blocks["ffn_norm2"]);
+            auto adaln_modulation = std::dynamic_pointer_cast<Linear>(blocks["adaln_modulation"]);
+
+            auto mod       = adaln_modulation->forward(ctx, adaln_input);
+            auto mods      = ggml_ext_chunk(ctx->ggml_ctx, mod, 4, 0);
+            auto scale_msa = mods[0];
+            auto gate_msa  = to_token_modulation(ctx->ggml_ctx, ggml_tanh(ctx->ggml_ctx, mods[1]));
+            auto scale_mlp = mods[2];
+            auto gate_mlp  = to_token_modulation(ctx->ggml_ctx, ggml_tanh(ctx->ggml_ctx, mods[3]));
+
+            auto attn_out = attention_norm1->forward(ctx, x);
+            attn_out      = modulate(ctx->ggml_ctx, attn_out, scale_msa);
+            attn_out      = attention->forward(ctx, attn_out, pe, mask);
+            attn_out      = attention_norm2->forward(ctx, attn_out);
+            x             = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out, gate_msa));
+
+            auto ffn_out = ffn_norm1->forward(ctx, x);
+            ffn_out      = modulate(ctx->ggml_ctx, ffn_out, scale_mlp);
+            ffn_out      = feed_forward->forward(ctx, ffn_out);
+            ffn_out      = ffn_norm2->forward(ctx, ffn_out);
+            x            = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, ffn_out, gate_mlp));
+
+            return x;
+        }
+    };
+
+    class Ideogram4EmbedScalar : public GGMLBlock {
+    protected:
+        int64_t dim;
+
+    public:
+        Ideogram4EmbedScalar(int64_t dim)
+            : dim(dim) {
+            blocks["mlp_in"]  = make_linear(dim, dim, true);
+            blocks["mlp_out"] = make_linear(dim, dim, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            auto mlp_in  = std::dynamic_pointer_cast<Linear>(blocks["mlp_in"]);
+            auto mlp_out = std::dynamic_pointer_cast<Linear>(blocks["mlp_out"]);
+
+            x = timestep_embedding_sin_cos(ctx->ggml_ctx, x, static_cast<int>(dim));
+            x = ggml_silu(ctx->ggml_ctx, mlp_in->forward(ctx, x));
+            x = mlp_out->forward(ctx, x);
+            return x;
+        }
+    };
+
+    class Ideogram4FinalLayer : public GGMLBlock {
+    public:
+        Ideogram4FinalLayer(const Ideogram4Config& config) {
+            blocks["norm_final"]       = std::make_shared<LayerNorm>(config.emb_dim, 1e-6f, false);
+            blocks["linear"]           = make_linear(config.emb_dim, config.in_channels, true);
+            blocks["adaln_modulation"] = make_linear(config.adanln_dim, config.emb_dim, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* c) {
+            auto norm_final       = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_final"]);
+            auto linear           = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
+            auto adaln_modulation = std::dynamic_pointer_cast<Linear>(blocks["adaln_modulation"]);
+
+            auto scale = adaln_modulation->forward(ctx, ggml_silu(ctx->ggml_ctx, c));
+            x          = norm_final->forward(ctx, x);
+            x          = modulate(ctx->ggml_ctx, x, scale);
+            x          = linear->forward(ctx, x);
+            return x;
+        }
+    };
+
+    class Ideogram4Transformer : public GGMLBlock {
+    protected:
+        Ideogram4Config config;
+
+    public:
+        Ideogram4Transformer() = default;
+        explicit Ideogram4Transformer(Ideogram4Config config)
+            : config(std::move(config)) {
+            blocks["input_proj"]            = make_linear(this->config.in_channels, this->config.emb_dim, true);
+            blocks["llm_cond_norm"]         = std::make_shared<RMSNorm>(this->config.llm_features_dim, 1e-6f);
+            blocks["llm_cond_proj"]         = make_linear(this->config.llm_features_dim, this->config.emb_dim, true);
+            blocks["t_embedding"]           = std::make_shared<Ideogram4EmbedScalar>(this->config.emb_dim);
+            blocks["adaln_proj"]            = make_linear(this->config.emb_dim, this->config.adanln_dim, true);
+            blocks["embed_image_indicator"] = std::make_shared<Embedding>(2, this->config.emb_dim);
+
+            for (int i = 0; i < this->config.num_layers; ++i) {
+                blocks["layers." + std::to_string(i)] = std::make_shared<Ideogram4TransformerBlock>(this->config);
+            }
+            blocks["final_layer"] = std::make_shared<Ideogram4FinalLayer>(this->config);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* timestep,
+                             ggml_tensor* context,
+                             ggml_tensor* pe,
+                             ggml_tensor* image_indicator_ids) {
+            int64_t W = x->ne[0];
+            int64_t H = x->ne[1];
+            int64_t N = x->ne[3];
+            GGML_ASSERT(N == 1);
+
+            auto input_proj            = std::dynamic_pointer_cast<Linear>(blocks["input_proj"]);
+            auto llm_cond_norm         = std::dynamic_pointer_cast<RMSNorm>(blocks["llm_cond_norm"]);
+            auto llm_cond_proj         = std::dynamic_pointer_cast<Linear>(blocks["llm_cond_proj"]);
+            auto t_embedding           = std::dynamic_pointer_cast<Ideogram4EmbedScalar>(blocks["t_embedding"]);
+            auto adaln_proj            = std::dynamic_pointer_cast<Linear>(blocks["adaln_proj"]);
+            auto embed_image_indicator = std::dynamic_pointer_cast<Embedding>(blocks["embed_image_indicator"]);
+            auto final_layer           = std::dynamic_pointer_cast<Ideogram4FinalLayer>(blocks["final_layer"]);
+
+            auto img = patchify(ctx->ggml_ctx, x, config);
+            img      = input_proj->forward(ctx, img);
+
+            ggml_tensor* h      = img;
+            int64_t context_len = 0;
+            if (context != nullptr) {
+                if (ggml_n_dims(context) < 3) {
+                    context = ggml_reshape_3d(ctx->ggml_ctx, context, context->ne[0], context->ne[1], 1);
+                }
+                context     = interleave_hidden_state_layers(ctx->ggml_ctx, context);
+                context_len = context->ne[1];
+                auto txt    = llm_cond_norm->forward(ctx, context);
+                txt         = llm_cond_proj->forward(ctx, txt);
+                h           = ggml_concat(ctx->ggml_ctx, txt, img, 1);
+            }
+
+            auto indicator_embedding = embed_image_indicator->forward(ctx, image_indicator_ids);
+            h                        = ggml_add(ctx->ggml_ctx, h, indicator_embedding);
+
+            auto t_cond      = t_embedding->forward(ctx, timestep);
+            auto adaln_input = ggml_silu(ctx->ggml_ctx, adaln_proj->forward(ctx, t_cond));
+
+            for (int i = 0; i < config.num_layers; ++i) {
+                auto block = std::dynamic_pointer_cast<Ideogram4TransformerBlock>(blocks["layers." + std::to_string(i)]);
+                h          = block->forward(ctx, h, pe, adaln_input, nullptr);
+                sd::ggml_graph_cut::mark_graph_cut(h, "ideogram4.layers." + std::to_string(i), "hidden");
+            }
+
+            h = final_layer->forward(ctx, h, adaln_input);
+            if (context_len > 0) {
+                h = ggml_ext_slice(ctx->ggml_ctx, h, 1, context_len, h->ne[1]);
+            }
+
+            h = unpatchify(ctx->ggml_ctx, h, H, W, config);
+            h = ggml_ext_scale(ctx->ggml_ctx, h, -1.f);
+            return h;
+        }
+    };
+
+    class Ideogram4Runner : public DiffusionModelRunner {
+    protected:
+        bool should_use_uncond_model(const DiffusionParams& diffusion_params) const {
+            return has_uncond_model &&
+                   diffusion_params.context == nullptr &&
+                   diffusion_params.y != nullptr &&
+                   !diffusion_params.y->empty();
+        }
+
+    public:
+        Ideogram4Config config;
+        Ideogram4Transformer model;
+        Ideogram4Transformer uncond_model;
+        bool has_uncond_model = false;
+        std::string uncond_prefix;
+        std::vector<float> pe_vec;
+        std::vector<int32_t> image_indicator_vec;
+
+        Ideogram4Runner(ggml_backend_t backend,
+                        const String2TensorStorage& tensor_storage_map      = {},
+                        const std::string prefix                            = "",
+                        std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : DiffusionModelRunner(backend, prefix, weight_manager),
+              config(Ideogram4Config::detect_from_weights(tensor_storage_map, prefix)),
+              uncond_prefix(prefix + ".uncond") {
+            model = Ideogram4Transformer(config);
+            model.init(params_ctx, tensor_storage_map, prefix);
+            for (const auto& pair : tensor_storage_map) {
+                const std::string& name = pair.first;
+                if (starts_with(name, uncond_prefix)) {
+                    has_uncond_model = true;
+                    break;
+                }
+            }
+            if (has_uncond_model) {
+                LOG_DEBUG("using uncond model");
+                uncond_model = Ideogram4Transformer(config);
+                uncond_model.init(params_ctx, tensor_storage_map, uncond_prefix);
+            }
+        }
+
+        std::string get_desc() override {
+            return "ideogram4";
+        }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
+            model.get_param_tensors(tensors, prefix);
+            if (has_uncond_model) {
+                uncond_model.get_param_tensors(tensors, this->uncond_prefix);
+            }
+        }
+
+        ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
+                                 const sd::Tensor<float>& timesteps_tensor,
+                                 const sd::Tensor<float>& context_tensor,
+                                 bool use_uncond_model = false) {
+            ggml_cgraph* gf        = new_graph_custom(IDEOGRAM4_GRAPH_SIZE);
+            ggml_tensor* x         = make_input(x_tensor);
+            ggml_tensor* timesteps = make_input(timesteps_tensor);
+            GGML_ASSERT(x->ne[3] == 1);
+            Ideogram4Transformer& active_model = use_uncond_model ? uncond_model : model;
+
+            ggml_tensor* context = nullptr;
+            int64_t context_len  = 0;
+            if (!context_tensor.empty()) {
+                context     = make_input(context_tensor);
+                context_len = context->ne[1];
+            }
+
+            int64_t grid_w   = x->ne[0];
+            int64_t grid_h   = x->ne[1];
+            int64_t pos_len  = context_len + grid_h * grid_w;
+            int64_t head_dim = config.emb_dim / config.num_heads;
+
+            auto runner_ctx = get_context();
+            pe_vec          = gen_ideogram4_pe(static_cast<int>(grid_h),
+                                               static_cast<int>(grid_w),
+                                               static_cast<int>(x->ne[3]),
+                                               static_cast<int>(context_len),
+                                               static_cast<int>(head_dim),
+                                               static_cast<int>(config.rope_theta),
+                                               config.mrope_section,
+                                               runner_ctx.circular_x_enabled,
+                                               runner_ctx.circular_y_enabled);
+            auto pe         = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, head_dim / 2, pos_len);
+            set_backend_tensor_data(pe, pe_vec.data());
+
+            image_indicator_vec.assign(static_cast<size_t>(pos_len), 1);
+            for (int64_t i = 0; i < context_len; ++i) {
+                image_indicator_vec[static_cast<size_t>(i)] = 0;
+            }
+            auto indicator = ggml_new_tensor_2d(compute_ctx, GGML_TYPE_I32, pos_len, x->ne[3]);
+            set_backend_tensor_data(indicator, image_indicator_vec.data());
+
+            ggml_tensor* out = active_model.forward(&runner_ctx, x, timesteps, context, pe, indicator);
+            ggml_build_forward_expand(gf, out);
+            return gf;
+        }
+
+        sd::Tensor<float> compute(int n_threads,
+                                  const sd::Tensor<float>& x,
+                                  const sd::Tensor<float>& timesteps,
+                                  const sd::Tensor<float>& context,
+                                  bool use_uncond_model = false) {
+            auto get_graph = [&]() -> ggml_cgraph* {
+                return build_graph(x, timesteps, context, use_uncond_model);
+            };
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
+        }
+
+        sd::Tensor<float> compute(int n_threads,
+                                  const DiffusionParams& diffusion_params) override {
+            GGML_ASSERT(diffusion_params.x != nullptr);
+            GGML_ASSERT(diffusion_params.timesteps != nullptr);
+            bool use_uncond_model = should_use_uncond_model(diffusion_params);
+            return compute(n_threads,
+                           *diffusion_params.x,
+                           *diffusion_params.timesteps,
+                           tensor_or_empty(diffusion_params.context),
+                           use_uncond_model);
+        }
+    };
+}  // namespace Ideogram4
+
+#endif  // __SD_MODEL_DIFFUSION_IDEOGRAM4_HPP__

src/model/diffusion/lens.hpp

@@ -1,18 +1,83 @@
-#ifndef __SD_LENS_HPP__
-#define __SD_LENS_HPP__
+#ifndef __SD_MODEL_DIFFUSION_LENS_HPP__
+#define __SD_MODEL_DIFFUSION_LENS_HPP__
 
 #include <memory>
 #include <vector>
 
-#include "common_block.hpp"
-#include "diffusion_model.hpp"
-#include "flux.hpp"
-#include "qwen_image.hpp"
-#include "rope.hpp"
+#include "model/common/block.hpp"
+#include "model/common/rope.hpp"
+#include "model/diffusion/flux.hpp"
+#include "model/diffusion/model.hpp"
+#include "model/diffusion/qwen_image.hpp"
 
 namespace Lens {
     constexpr int LENS_GRAPH_SIZE = 40960;
 
+    struct LensConfig {
+        int patch_size              = 2;
+        int64_t in_channels         = 128;
+        int64_t out_channels        = 32;
+        int num_layers              = 48;
+        int64_t attention_head_dim  = 64;
+        int64_t num_attention_heads = 24;
+        int64_t joint_attention_dim = 2880;
+        int selected_layer_count    = 4;
+        int theta                   = 10000;
+        std::vector<int> axes_dim   = {8, 28, 28};
+        int axes_dim_sum            = 64;
+
+        static LensConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
+            LensConfig config;
+            config.num_layers = 0;
+            for (const auto& [name, tensor_storage] : tensor_storage_map) {
+                if (!starts_with(name, prefix)) {
+                    continue;
+                }
+                if (ends_with(name, "img_in.weight") && tensor_storage.n_dims == 2) {
+                    config.in_channels = tensor_storage.ne[0];
+                    int64_t inner_dim  = tensor_storage.ne[1];
+                    if (config.attention_head_dim > 0) {
+                        config.num_attention_heads = inner_dim / config.attention_head_dim;
+                    }
+                } else if (ends_with(name, "txt_in.weight") && tensor_storage.n_dims == 2) {
+                    config.selected_layer_count = static_cast<int>(tensor_storage.ne[0] / config.joint_attention_dim);
+                } else if (ends_with(name, "proj_out.weight") && tensor_storage.n_dims == 2) {
+                    int64_t patch_area  = config.patch_size * config.patch_size;
+                    config.out_channels = tensor_storage.ne[1] / patch_area;
+                } else if (ends_with(name, "transformer_blocks.0.attn.norm_q.weight") && tensor_storage.n_dims == 1) {
+                    config.attention_head_dim = tensor_storage.ne[0];
+                }
+
+                size_t pos = name.find("transformer_blocks.");
+                if (pos != std::string::npos) {
+                    auto items = split_string(name.substr(pos), '.');
+                    if (items.size() > 1) {
+                        int block_index = atoi(items[1].c_str());
+                        if (block_index + 1 > config.num_layers) {
+                            config.num_layers = block_index + 1;
+                        }
+                    }
+                }
+            }
+            if (config.num_layers == 0) {
+                config.num_layers = 48;
+            }
+            config.axes_dim_sum = 0;
+            for (int axis_dim : config.axes_dim) {
+                config.axes_dim_sum += axis_dim;
+            }
+            LOG_DEBUG("lens: num_layers = %d, selected_layer_count = %d, hidden_size = %" PRId64 ", num_attention_heads = %" PRId64 ", attention_head_dim = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64,
+                      config.num_layers,
+                      config.selected_layer_count,
+                      config.num_attention_heads * config.attention_head_dim,
+                      config.num_attention_heads,
+                      config.attention_head_dim,
+                      config.in_channels,
+                      config.out_channels);
+            return config;
+        }
+    };
+
     struct LensTimestepProjEmbeddings : public GGMLBlock {
         LensTimestepProjEmbeddings(int64_t embedding_dim) {
             blocks["timestep_embedder"] = std::make_shared<Qwen::TimestepEmbedding>(256, embedding_dim);
@@ -209,41 +274,27 @@ namespace Lens {
         }
     };
 
-    struct LensParams {
-        int patch_size              = 2;
-        int64_t in_channels         = 128;
-        int64_t out_channels        = 32;
-        int num_layers              = 48;
-        int64_t attention_head_dim  = 64;
-        int64_t num_attention_heads = 24;
-        int64_t joint_attention_dim = 2880;
-        int selected_layer_count    = 4;
-        int theta                   = 10000;
-        std::vector<int> axes_dim   = {8, 28, 28};
-        int axes_dim_sum            = 64;
-    };
-
     class LensModel : public GGMLBlock {
     public:
-        LensParams params;
+        LensConfig config;
 
         LensModel() = default;
-        LensModel(LensParams params)
-            : params(params) {
-            int64_t inner_dim         = params.num_attention_heads * params.attention_head_dim;
+        LensModel(LensConfig config)
+            : config(config) {
+            int64_t inner_dim         = config.num_attention_heads * config.attention_head_dim;
             blocks["time_text_embed"] = std::make_shared<LensTimestepProjEmbeddings>(inner_dim);
-            blocks["img_in"]          = std::make_shared<Linear>(params.in_channels, inner_dim, true);
-            blocks["txt_in"]          = std::make_shared<Linear>(params.joint_attention_dim * params.selected_layer_count, inner_dim, true);
-            for (int i = 0; i < params.selected_layer_count; ++i) {
-                blocks["txt_norm." + std::to_string(i)] = std::make_shared<RMSNorm>(params.joint_attention_dim, 1e-5f);
+            blocks["img_in"]          = std::make_shared<Linear>(config.in_channels, inner_dim, true);
+            blocks["txt_in"]          = std::make_shared<Linear>(config.joint_attention_dim * config.selected_layer_count, inner_dim, true);
+            for (int i = 0; i < config.selected_layer_count; ++i) {
+                blocks["txt_norm." + std::to_string(i)] = std::make_shared<RMSNorm>(config.joint_attention_dim, 1e-5f);
             }
-            for (int i = 0; i < params.num_layers; ++i) {
+            for (int i = 0; i < config.num_layers; ++i) {
                 blocks["transformer_blocks." + std::to_string(i)] = std::make_shared<LensTransformerBlock>(inner_dim,
-                                                                                                           params.num_attention_heads,
-                                                                                                           params.attention_head_dim);
+                                                                                                           config.num_attention_heads,
+                                                                                                           config.attention_head_dim);
             }
             blocks["norm_out"] = std::make_shared<LensAdaLayerNormContinuous>(inner_dim, 1e-6f);
-            blocks["proj_out"] = std::make_shared<Linear>(inner_dim, params.patch_size * params.patch_size * params.out_channels, true);
+            blocks["proj_out"] = std::make_shared<Linear>(inner_dim, config.patch_size * config.patch_size * config.out_channels, true);
         }
 
         ggml_tensor* forward(GGMLRunnerContext* ctx,
@@ -269,9 +320,9 @@ namespace Lens {
             img      = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, img, 1, 0, 2, 3));
             img      = img_in->forward(ctx, img);
 
-            std::vector<ggml_tensor*> txt_chunks = ggml_ext_chunk(ctx->ggml_ctx, context, params.selected_layer_count, 0);
+            std::vector<ggml_tensor*> txt_chunks = ggml_ext_chunk(ctx->ggml_ctx, context, config.selected_layer_count, 0);
             ggml_tensor* txt                     = nullptr;
-            for (int i = 0; i < params.selected_layer_count; ++i) {
+            for (int i = 0; i < config.selected_layer_count; ++i) {
                 auto txt_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["txt_norm." + std::to_string(i)]);
                 auto chunk    = txt_norm->forward(ctx, txt_chunks[i]);
                 txt           = txt == nullptr ? chunk : ggml_concat(ctx->ggml_ctx, txt, chunk, 0);
@@ -281,7 +332,7 @@ namespace Lens {
             sd::ggml_graph_cut::mark_graph_cut(img, "lens.prelude", "img");
             sd::ggml_graph_cut::mark_graph_cut(txt, "lens.prelude", "txt");
 
-            for (int i = 0; i < params.num_layers; ++i) {
+            for (int i = 0; i < config.num_layers; ++i) {
                 auto block = std::dynamic_pointer_cast<LensTransformerBlock>(blocks["transformer_blocks." + std::to_string(i)]);
                 auto out   = block->forward(ctx, img, txt, t_emb, pe);
                 img        = out.first;
@@ -294,67 +345,23 @@ namespace Lens {
             img = proj_out->forward(ctx, img);
 
             auto out = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, img, 1, 0, 2, 3));
-            out      = ggml_reshape_4d(ctx->ggml_ctx, out, W, H, params.patch_size * params.patch_size * params.out_channels, N);
+            out      = ggml_reshape_4d(ctx->ggml_ctx, out, W, H, config.patch_size * config.patch_size * config.out_channels, N);
             return out;
         }
     };
 
     struct LensRunner : public DiffusionModelRunner {
-        LensParams lens_params;
+        LensConfig config;
         LensModel lens;
         std::vector<float> pe_vec;
 
         LensRunner(ggml_backend_t backend,
-                   ggml_backend_t params_backend,
-                   const String2TensorStorage& tensor_storage_map = {},
-                   const std::string prefix                       = "")
-            : DiffusionModelRunner(backend, params_backend, prefix) {
-            lens_params.num_layers = 0;
-            for (const auto& [name, tensor_storage] : tensor_storage_map) {
-                if (!starts_with(name, prefix)) {
-                    continue;
-                }
-                if (ends_with(name, "img_in.weight") && tensor_storage.n_dims == 2) {
-                    lens_params.in_channels         = tensor_storage.ne[0];
-                    int64_t inner_dim               = tensor_storage.ne[1];
-                    lens_params.num_attention_heads = inner_dim / lens_params.attention_head_dim;
-                } else if (ends_with(name, "txt_in.weight") && tensor_storage.n_dims == 2) {
-                    lens_params.selected_layer_count = static_cast<int>(tensor_storage.ne[0] / lens_params.joint_attention_dim);
-                } else if (ends_with(name, "proj_out.weight") && tensor_storage.n_dims == 2) {
-                    lens_params.out_channels = tensor_storage.ne[1] / lens_params.patch_size / lens_params.patch_size;
-                } else if (ends_with(name, "transformer_blocks.0.attn.norm_q.weight") && tensor_storage.n_dims == 1) {
-                    lens_params.attention_head_dim = tensor_storage.ne[0];
-                }
-
-                size_t pos = name.find("transformer_blocks.");
-                if (pos != std::string::npos) {
-                    std::string layer_name = name.substr(pos);
-                    auto items             = split_string(layer_name, '.');
-                    if (items.size() > 1) {
-                        int block_index = atoi(items[1].c_str());
-                        if (block_index + 1 > lens_params.num_layers) {
-                            lens_params.num_layers = block_index + 1;
-                        }
-                    }
-                }
-            }
-            if (lens_params.num_layers == 0) {
-                lens_params.num_layers = 48;
-            }
-            lens_params.axes_dim_sum = 0;
-            for (int axis_dim : lens_params.axes_dim) {
-                lens_params.axes_dim_sum += axis_dim;
-            }
-
-            LOG_INFO("lens: layers = %d, in_channels = %" PRId64 ", out_channels = %" PRId64
-                     ", heads = %" PRId64 ", head_dim = %" PRId64,
-                     lens_params.num_layers,
-                     lens_params.in_channels,
-                     lens_params.out_channels,
-                     lens_params.num_attention_heads,
-                     lens_params.attention_head_dim);
-
-            lens = LensModel(lens_params);
+                   const String2TensorStorage& tensor_storage_map      = {},
+                   const std::string prefix                            = "",
+                   std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : DiffusionModelRunner(backend, prefix, weight_manager),
+              config(LensConfig::detect_from_weights(tensor_storage_map, prefix)) {
+            lens = LensModel(config);
             lens.init(params_ctx, tensor_storage_map, prefix);
         }
 
@@ -380,12 +387,12 @@ namespace Lens {
                                             static_cast<int>(x->ne[0]),
                                             static_cast<int>(x->ne[3]),
                                             static_cast<int>(context->ne[1]),
-                                            lens_params.theta,
+                                            config.theta,
                                             circular_y_enabled,
                                             circular_x_enabled,
-                                            lens_params.axes_dim);
-            int pos_len = static_cast<int>(pe_vec.size() / lens_params.axes_dim_sum / 2);
-            auto pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, lens_params.axes_dim_sum / 2, pos_len);
+                                            config.axes_dim);
+            int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
+            auto pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
             set_backend_tensor_data(pe, pe_vec.data());
 
             auto runner_ctx  = get_context();
@@ -401,7 +408,7 @@ namespace Lens {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,
@@ -416,4 +423,4 @@ namespace Lens {
     };
 }  // namespace Lens
 
-#endif  // __SD_LENS_HPP__
+#endif  // __SD_MODEL_DIFFUSION_LENS_HPP__

src/model/diffusion/ltxv.hpp

@@ -1,5 +1,5 @@
-#ifndef __SD_LTXV_HPP__
-#define __SD_LTXV_HPP__
+#ifndef __SD_MODEL_DIFFUSION_LTXV_HPP__
+#define __SD_MODEL_DIFFUSION_LTXV_HPP__
 
 #include <algorithm>
 #include <cmath>
@@ -9,10 +9,11 @@
 #include <utility>
 #include <vector>
 
-#include "common_block.hpp"
-#include "diffusion_model.hpp"
-#include "flux.hpp"
-#include "rope.hpp"
+#include "model/common/block.hpp"
+#include "model/common/rope.hpp"
+#include "model/diffusion/flux.hpp"
+#include "model/diffusion/model.hpp"
+#include "model_loader.h"
 
 namespace LTXV {
 
@@ -72,6 +73,200 @@ namespace LTXV {
         return max_block + 1;
     }
 
+    struct LTXAVConfig {
+        int64_t in_channels                           = 128;
+        int64_t out_channels                          = 128;
+        int64_t hidden_size                           = 3840;
+        int64_t cross_attention_dim                   = 4096;
+        int64_t caption_channels                      = 3840;
+        int64_t num_attention_heads                   = 30;
+        int64_t attention_head_dim                    = 128;
+        int64_t num_layers                            = 28;
+        float positional_embedding_theta              = 10000.f;
+        std::vector<int> positional_embedding_max_pos = {20, 2048, 2048};
+        std::tuple<int, int, int> vae_scale_factors   = {8, 32, 32};
+        bool causal_temporal_positioning              = true;
+        float timestep_scale_multiplier               = 1000.f;
+
+        int64_t audio_in_channels                           = 128;
+        int64_t audio_out_channels                          = 128;
+        int64_t audio_hidden_size                           = 2048;
+        int64_t audio_cross_attention_dim                   = 2048;
+        int64_t audio_num_attention_heads                   = 32;
+        int64_t audio_attention_head_dim                    = 64;
+        std::vector<int> audio_positional_embedding_max_pos = {20};
+        float av_ca_timestep_scale_multiplier               = 1000.f;
+        int64_t num_audio_channels                          = 8;
+        int64_t audio_frequency_bins                        = 16;
+
+        bool use_connector                   = false;
+        int64_t connector_hidden_size        = 3840;
+        int64_t connector_num_heads          = 30;
+        int64_t connector_head_dim           = 128;
+        int64_t connector_num_layers         = 2;
+        int64_t connector_num_registers      = 128;
+        bool connector_rope_interleaved      = false;
+        bool connector_apply_gated_attention = false;
+
+        bool use_audio_connector                   = false;
+        int64_t audio_connector_hidden_size        = 2048;
+        int64_t audio_connector_num_heads          = 32;
+        int64_t audio_connector_head_dim           = 64;
+        int64_t audio_connector_num_layers         = 2;
+        int64_t audio_connector_num_registers      = 128;
+        bool audio_connector_rope_interleaved      = false;
+        bool audio_connector_apply_gated_attention = false;
+
+        bool video_rope_interleaved  = false;
+        bool use_middle_indices_grid = true;
+        bool cross_attention_adaln   = false;
+
+        bool use_caption_projection          = true;
+        bool use_audio_caption_projection    = true;
+        bool caption_proj_before_connector   = true;
+        bool caption_projection_first_linear = false;
+
+        bool self_attention_gated  = false;
+        bool cross_attention_gated = false;
+
+        static std::pair<int64_t, int64_t> infer_attention_layout(int64_t hidden_size,
+                                                                  int64_t preferred_heads = -1) {
+            if (preferred_heads > 0 && hidden_size % preferred_heads == 0) {
+                return {preferred_heads, hidden_size / preferred_heads};
+            }
+            const int candidates[] = {128, 96, 80, 64, 48, 40, 32};
+            for (int head_dim : candidates) {
+                if (hidden_size % head_dim == 0) {
+                    int64_t heads = hidden_size / head_dim;
+                    if (heads >= 8 && heads <= 64) {
+                        return {heads, head_dim};
+                    }
+                }
+            }
+            return {32, hidden_size / 32};
+        }
+
+        static int64_t infer_gate_heads(const String2TensorStorage& tensor_storage_map,
+                                        const std::string& bias_name,
+                                        int64_t fallback_heads) {
+            auto it = tensor_storage_map.find(bias_name);
+            if (it != tensor_storage_map.end()) {
+                return it->second.ne[0];
+            }
+            return fallback_heads;
+        }
+
+        static LTXAVConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
+            LTXAVConfig config;
+            auto patchify_proj_iter = tensor_storage_map.find(prefix + ".patchify_proj.weight");
+            if (patchify_proj_iter != tensor_storage_map.end()) {
+                config.in_channels         = patchify_proj_iter->second.ne[0];
+                config.hidden_size         = patchify_proj_iter->second.ne[1];
+                int64_t video_heads        = infer_gate_heads(tensor_storage_map, prefix + ".transformer_blocks.0.attn1.to_gate_logits.bias", 32);
+                auto attn_layout           = infer_attention_layout(config.hidden_size, video_heads);
+                config.num_attention_heads = attn_layout.first;
+                config.attention_head_dim  = attn_layout.second;
+            }
+
+            auto audio_patchify_proj_iter = tensor_storage_map.find(prefix + ".audio_patchify_proj.weight");
+            if (audio_patchify_proj_iter != tensor_storage_map.end()) {
+                config.audio_in_channels         = audio_patchify_proj_iter->second.ne[0];
+                config.audio_hidden_size         = audio_patchify_proj_iter->second.ne[1];
+                config.audio_out_channels        = config.audio_in_channels;
+                int64_t audio_heads              = infer_gate_heads(tensor_storage_map, prefix + ".transformer_blocks.0.audio_attn1.to_gate_logits.bias", 32);
+                auto audio_attn_layout           = infer_attention_layout(config.audio_hidden_size, audio_heads);
+                config.audio_num_attention_heads = audio_attn_layout.first;
+                config.audio_attention_head_dim  = audio_attn_layout.second;
+            }
+
+            auto proj_out_iter = tensor_storage_map.find(prefix + ".proj_out.weight");
+            if (proj_out_iter != tensor_storage_map.end()) {
+                config.out_channels = proj_out_iter->second.ne[1];
+            }
+            auto audio_proj_out_iter = tensor_storage_map.find(prefix + ".audio_proj_out.weight");
+            if (audio_proj_out_iter != tensor_storage_map.end()) {
+                config.audio_out_channels = audio_proj_out_iter->second.ne[1];
+            }
+
+            auto attn2_iter = tensor_storage_map.find(prefix + ".transformer_blocks.0.attn2.to_k.weight");
+            if (attn2_iter != tensor_storage_map.end()) {
+                config.cross_attention_dim = attn2_iter->second.ne[0];
+            }
+            auto audio_attn2_iter = tensor_storage_map.find(prefix + ".transformer_blocks.0.audio_attn2.to_k.weight");
+            if (audio_attn2_iter != tensor_storage_map.end()) {
+                config.audio_cross_attention_dim = audio_attn2_iter->second.ne[0];
+            }
+            if (tensor_storage_map.find(prefix + ".transformer_blocks.0.prompt_scale_shift_table") != tensor_storage_map.end()) {
+                config.cross_attention_adaln = true;
+            }
+            if (tensor_storage_map.find(prefix + ".transformer_blocks.0.attn1.to_gate_logits.weight") != tensor_storage_map.end() ||
+                tensor_storage_map.find(prefix + ".transformer_blocks.0.audio_attn1.to_gate_logits.weight") != tensor_storage_map.end()) {
+                config.self_attention_gated = true;
+            }
+            if (tensor_storage_map.find(prefix + ".transformer_blocks.0.attn2.to_gate_logits.weight") != tensor_storage_map.end() ||
+                tensor_storage_map.find(prefix + ".transformer_blocks.0.audio_attn2.to_gate_logits.weight") != tensor_storage_map.end()) {
+                config.cross_attention_gated = true;
+            }
+            if (tensor_storage_map.find(prefix + ".caption_projection.linear_1.weight") == tensor_storage_map.end() &&
+                tensor_storage_map.find(prefix + ".caption_projection.linear_2.weight") == tensor_storage_map.end()) {
+                config.use_caption_projection = false;
+            }
+            if (tensor_storage_map.find(prefix + ".audio_caption_projection.linear_1.weight") == tensor_storage_map.end() &&
+                tensor_storage_map.find(prefix + ".audio_caption_projection.linear_2.weight") == tensor_storage_map.end()) {
+                config.use_audio_caption_projection = false;
+            }
+
+            config.num_layers = count_prefix_blocks(tensor_storage_map, prefix + ".", "transformer_blocks.");
+
+            auto connector_iter = tensor_storage_map.find(prefix + ".video_embeddings_connector.transformer_1d_blocks.0.attn1.to_q.weight");
+            if (connector_iter != tensor_storage_map.end()) {
+                config.use_connector         = true;
+                config.connector_hidden_size = connector_iter->second.ne[1];
+                int64_t connector_heads      = infer_gate_heads(tensor_storage_map,
+                                                                prefix + ".video_embeddings_connector.transformer_1d_blocks.0.attn1.to_gate_logits.bias",
+                                                                32);
+                auto connector_layout        = infer_attention_layout(config.connector_hidden_size, connector_heads);
+                config.connector_num_heads   = connector_layout.first;
+                config.connector_head_dim    = connector_layout.second;
+                config.connector_num_layers  = count_prefix_blocks(tensor_storage_map, prefix + ".video_embeddings_connector.", "transformer_1d_blocks.");
+                auto register_iter           = tensor_storage_map.find(prefix + ".video_embeddings_connector.learnable_registers");
+                if (register_iter != tensor_storage_map.end()) {
+                    config.connector_num_registers = register_iter->second.ne[1];
+                }
+                if (tensor_storage_map.find(prefix + ".video_embeddings_connector.transformer_1d_blocks.0.attn1.to_gate_logits.weight") != tensor_storage_map.end()) {
+                    config.connector_apply_gated_attention = true;
+                }
+            }
+
+            auto audio_connector_iter = tensor_storage_map.find(prefix + ".audio_embeddings_connector.transformer_1d_blocks.0.attn1.to_q.weight");
+            if (audio_connector_iter != tensor_storage_map.end()) {
+                config.use_audio_connector         = true;
+                config.audio_connector_hidden_size = audio_connector_iter->second.ne[1];
+                int64_t connector_heads            = infer_gate_heads(tensor_storage_map,
+                                                                      prefix + ".audio_embeddings_connector.transformer_1d_blocks.0.attn1.to_gate_logits.bias",
+                                                                      32);
+                auto connector_layout              = infer_attention_layout(config.audio_connector_hidden_size, connector_heads);
+                config.audio_connector_num_heads   = connector_layout.first;
+                config.audio_connector_head_dim    = connector_layout.second;
+                config.audio_connector_num_layers  = count_prefix_blocks(tensor_storage_map, prefix + ".audio_embeddings_connector.", "transformer_1d_blocks.");
+                auto register_iter                 = tensor_storage_map.find(prefix + ".audio_embeddings_connector.learnable_registers");
+                if (register_iter != tensor_storage_map.end()) {
+                    config.audio_connector_num_registers = register_iter->second.ne[1];
+                }
+                if (tensor_storage_map.find(prefix + ".audio_embeddings_connector.transformer_1d_blocks.0.attn1.to_gate_logits.weight") != tensor_storage_map.end()) {
+                    config.audio_connector_apply_gated_attention = true;
+                }
+            }
+            LOG_DEBUG("ltxav: num_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_attention_heads = %" PRId64 ", audio_hidden_size = %" PRId64 ", audio_num_attention_heads = %" PRId64,
+                      config.num_layers,
+                      config.hidden_size,
+                      config.num_attention_heads,
+                      config.audio_hidden_size,
+                      config.audio_num_attention_heads);
+            return config;
+        }
+    };
+
     __STATIC_INLINE__ std::vector<float> generate_freq_grid(float theta,
                                                             int positional_dims,
                                                             int dim) {
@@ -749,63 +944,6 @@ namespace LTXV {
         }
     };
 
-    struct LTXAVParams {
-        int64_t in_channels                           = 128;
-        int64_t out_channels                          = 128;
-        int64_t hidden_size                           = 3840;
-        int64_t cross_attention_dim                   = 4096;
-        int64_t caption_channels                      = 3840;
-        int64_t num_attention_heads                   = 30;
-        int64_t attention_head_dim                    = 128;
-        int64_t num_layers                            = 28;
-        float positional_embedding_theta              = 10000.f;
-        std::vector<int> positional_embedding_max_pos = {20, 2048, 2048};
-        std::tuple<int, int, int> vae_scale_factors   = {8, 32, 32};
-        bool causal_temporal_positioning              = true;
-        float timestep_scale_multiplier               = 1000.f;
-
-        int64_t audio_in_channels                           = 128;
-        int64_t audio_out_channels                          = 128;
-        int64_t audio_hidden_size                           = 2048;
-        int64_t audio_cross_attention_dim                   = 2048;
-        int64_t audio_num_attention_heads                   = 32;
-        int64_t audio_attention_head_dim                    = 64;
-        std::vector<int> audio_positional_embedding_max_pos = {20};
-        float av_ca_timestep_scale_multiplier               = 1000.f;
-        int64_t num_audio_channels                          = 8;
-        int64_t audio_frequency_bins                        = 16;
-
-        bool use_connector                   = false;
-        int64_t connector_hidden_size        = 3840;
-        int64_t connector_num_heads          = 30;
-        int64_t connector_head_dim           = 128;
-        int64_t connector_num_layers         = 2;
-        int64_t connector_num_registers      = 128;
-        bool connector_rope_interleaved      = false;
-        bool connector_apply_gated_attention = false;
-
-        bool use_audio_connector                   = false;
-        int64_t audio_connector_hidden_size        = 2048;
-        int64_t audio_connector_num_heads          = 32;
-        int64_t audio_connector_head_dim           = 64;
-        int64_t audio_connector_num_layers         = 2;
-        int64_t audio_connector_num_registers      = 128;
-        bool audio_connector_rope_interleaved      = false;
-        bool audio_connector_apply_gated_attention = false;
-
-        bool video_rope_interleaved  = false;
-        bool use_middle_indices_grid = true;
-        bool cross_attention_adaln   = false;
-
-        bool use_caption_projection          = true;
-        bool use_audio_caption_projection    = true;
-        bool caption_proj_before_connector   = true;
-        bool caption_projection_first_linear = false;
-
-        bool self_attention_gated  = false;
-        bool cross_attention_gated = false;
-    };
-
     __STATIC_INLINE__ std::pair<int64_t, int64_t> infer_attention_layout(int64_t hidden_size,
                                                                          int64_t preferred_heads = -1) {
         if (preferred_heads > 0 && hidden_size % preferred_heads == 0) {
@@ -1169,92 +1307,92 @@ namespace LTXV {
     };
 
     struct LTXAVModelBlock : public GGMLBlock {
-        LTXAVParams cfg;
+        LTXAVConfig config;
 
         void init_params(ggml_context* ctx,
                          const String2TensorStorage& tensor_storage_map = {},
                          const std::string prefix                       = "") override {
             params["scale_shift_table"]       = ggml_new_tensor_2d(ctx,
                                                                    get_type(prefix + "scale_shift_table", tensor_storage_map, GGML_TYPE_F32),
-                                                                   cfg.hidden_size,
+                                                                   config.hidden_size,
                                                                    2);
             params["audio_scale_shift_table"] = ggml_new_tensor_2d(ctx,
                                                                    get_type(prefix + "audio_scale_shift_table", tensor_storage_map, GGML_TYPE_F32),
-                                                                   cfg.audio_hidden_size,
+                                                                   config.audio_hidden_size,
                                                                    2);
         }
 
-        LTXAVModelBlock(const LTXAVParams& params)
-            : cfg(params) {
-            blocks["patchify_proj"]       = std::make_shared<Linear>(cfg.in_channels, cfg.hidden_size, true, true);
-            blocks["audio_patchify_proj"] = std::make_shared<Linear>(cfg.audio_in_channels, cfg.audio_hidden_size, true, true);
-            blocks["adaln_single"]        = std::make_shared<AdaLayerNormSingle>(cfg.hidden_size, cfg.cross_attention_adaln ? 9 : 6);
-            blocks["audio_adaln_single"]  = std::make_shared<AdaLayerNormSingle>(cfg.audio_hidden_size, cfg.cross_attention_adaln ? 9 : 6);
-            if (cfg.cross_attention_adaln) {
-                blocks["prompt_adaln_single"]       = std::make_shared<AdaLayerNormSingle>(cfg.hidden_size, 2);
-                blocks["audio_prompt_adaln_single"] = std::make_shared<AdaLayerNormSingle>(cfg.audio_hidden_size, 2);
+        LTXAVModelBlock(const LTXAVConfig& config)
+            : config(config) {
+            blocks["patchify_proj"]       = std::make_shared<Linear>(config.in_channels, config.hidden_size, true, true);
+            blocks["audio_patchify_proj"] = std::make_shared<Linear>(config.audio_in_channels, config.audio_hidden_size, true, true);
+            blocks["adaln_single"]        = std::make_shared<AdaLayerNormSingle>(config.hidden_size, config.cross_attention_adaln ? 9 : 6);
+            blocks["audio_adaln_single"]  = std::make_shared<AdaLayerNormSingle>(config.audio_hidden_size, config.cross_attention_adaln ? 9 : 6);
+            if (config.cross_attention_adaln) {
+                blocks["prompt_adaln_single"]       = std::make_shared<AdaLayerNormSingle>(config.hidden_size, 2);
+                blocks["audio_prompt_adaln_single"] = std::make_shared<AdaLayerNormSingle>(config.audio_hidden_size, 2);
             }
-            blocks["av_ca_video_scale_shift_adaln_single"] = std::make_shared<AdaLayerNormSingle>(cfg.hidden_size, 4);
-            blocks["av_ca_a2v_gate_adaln_single"]          = std::make_shared<AdaLayerNormSingle>(cfg.hidden_size, 1);
-            blocks["av_ca_audio_scale_shift_adaln_single"] = std::make_shared<AdaLayerNormSingle>(cfg.audio_hidden_size, 4);
-            blocks["av_ca_v2a_gate_adaln_single"]          = std::make_shared<AdaLayerNormSingle>(cfg.audio_hidden_size, 1);
+            blocks["av_ca_video_scale_shift_adaln_single"] = std::make_shared<AdaLayerNormSingle>(config.hidden_size, 4);
+            blocks["av_ca_a2v_gate_adaln_single"]          = std::make_shared<AdaLayerNormSingle>(config.hidden_size, 1);
+            blocks["av_ca_audio_scale_shift_adaln_single"] = std::make_shared<AdaLayerNormSingle>(config.audio_hidden_size, 4);
+            blocks["av_ca_v2a_gate_adaln_single"]          = std::make_shared<AdaLayerNormSingle>(config.audio_hidden_size, 1);
 
-            if (cfg.use_caption_projection) {
-                if (cfg.caption_proj_before_connector) {
-                    if (cfg.caption_projection_first_linear) {
-                        blocks["caption_projection"] = std::make_shared<NormSingleLinearTextProjection>(cfg.caption_channels, cfg.hidden_size);
+            if (config.use_caption_projection) {
+                if (config.caption_proj_before_connector) {
+                    if (config.caption_projection_first_linear) {
+                        blocks["caption_projection"] = std::make_shared<NormSingleLinearTextProjection>(config.caption_channels, config.hidden_size);
                     }
                 } else {
-                    blocks["caption_projection"] = std::make_shared<PixArtAlphaTextProjection>(cfg.caption_channels, cfg.hidden_size, cfg.hidden_size);
+                    blocks["caption_projection"] = std::make_shared<PixArtAlphaTextProjection>(config.caption_channels, config.hidden_size, config.hidden_size);
                 }
             }
-            if (cfg.use_audio_caption_projection) {
-                if (cfg.caption_proj_before_connector) {
-                    if (cfg.caption_projection_first_linear) {
-                        blocks["audio_caption_projection"] = std::make_shared<NormSingleLinearTextProjection>(cfg.caption_channels, cfg.audio_hidden_size);
+            if (config.use_audio_caption_projection) {
+                if (config.caption_proj_before_connector) {
+                    if (config.caption_projection_first_linear) {
+                        blocks["audio_caption_projection"] = std::make_shared<NormSingleLinearTextProjection>(config.caption_channels, config.audio_hidden_size);
                     }
                 } else {
-                    blocks["audio_caption_projection"] = std::make_shared<PixArtAlphaTextProjection>(cfg.caption_channels, cfg.audio_hidden_size, cfg.audio_hidden_size);
+                    blocks["audio_caption_projection"] = std::make_shared<PixArtAlphaTextProjection>(config.caption_channels, config.audio_hidden_size, config.audio_hidden_size);
                 }
             }
 
-            if (cfg.use_connector) {
-                blocks["video_embeddings_connector"] = std::make_shared<Embeddings1DConnector>(cfg.connector_hidden_size,
-                                                                                               cfg.connector_num_heads,
-                                                                                               cfg.connector_head_dim,
-                                                                                               cfg.connector_num_layers,
-                                                                                               cfg.connector_num_registers,
-                                                                                               cfg.connector_rope_interleaved,
-                                                                                               cfg.connector_apply_gated_attention);
+            if (config.use_connector) {
+                blocks["video_embeddings_connector"] = std::make_shared<Embeddings1DConnector>(config.connector_hidden_size,
+                                                                                               config.connector_num_heads,
+                                                                                               config.connector_head_dim,
+                                                                                               config.connector_num_layers,
+                                                                                               config.connector_num_registers,
+                                                                                               config.connector_rope_interleaved,
+                                                                                               config.connector_apply_gated_attention);
             }
-            if (cfg.use_audio_connector) {
-                blocks["audio_embeddings_connector"] = std::make_shared<Embeddings1DConnector>(cfg.audio_connector_hidden_size,
-                                                                                               cfg.audio_connector_num_heads,
-                                                                                               cfg.audio_connector_head_dim,
-                                                                                               cfg.audio_connector_num_layers,
-                                                                                               cfg.audio_connector_num_registers,
-                                                                                               cfg.audio_connector_rope_interleaved,
-                                                                                               cfg.audio_connector_apply_gated_attention);
+            if (config.use_audio_connector) {
+                blocks["audio_embeddings_connector"] = std::make_shared<Embeddings1DConnector>(config.audio_connector_hidden_size,
+                                                                                               config.audio_connector_num_heads,
+                                                                                               config.audio_connector_head_dim,
+                                                                                               config.audio_connector_num_layers,
+                                                                                               config.audio_connector_num_registers,
+                                                                                               config.audio_connector_rope_interleaved,
+                                                                                               config.audio_connector_apply_gated_attention);
             }
 
-            for (int i = 0; i < cfg.num_layers; i++) {
-                blocks["transformer_blocks." + std::to_string(i)] = std::make_shared<BasicAVTransformerBlock>(cfg.hidden_size,
-                                                                                                              cfg.audio_hidden_size,
-                                                                                                              cfg.num_attention_heads,
-                                                                                                              cfg.audio_num_attention_heads,
-                                                                                                              cfg.attention_head_dim,
-                                                                                                              cfg.audio_attention_head_dim,
-                                                                                                              cfg.cross_attention_dim,
-                                                                                                              cfg.audio_cross_attention_dim,
-                                                                                                              cfg.self_attention_gated || cfg.cross_attention_gated,
-                                                                                                              cfg.cross_attention_adaln,
-                                                                                                              cfg.video_rope_interleaved);
+            for (int i = 0; i < config.num_layers; i++) {
+                blocks["transformer_blocks." + std::to_string(i)] = std::make_shared<BasicAVTransformerBlock>(config.hidden_size,
+                                                                                                              config.audio_hidden_size,
+                                                                                                              config.num_attention_heads,
+                                                                                                              config.audio_num_attention_heads,
+                                                                                                              config.attention_head_dim,
+                                                                                                              config.audio_attention_head_dim,
+                                                                                                              config.cross_attention_dim,
+                                                                                                              config.audio_cross_attention_dim,
+                                                                                                              config.self_attention_gated || config.cross_attention_gated,
+                                                                                                              config.cross_attention_adaln,
+                                                                                                              config.video_rope_interleaved);
             }
 
-            blocks["norm_out"]       = std::make_shared<LayerNorm>(cfg.hidden_size, 1e-6f, false);
-            blocks["proj_out"]       = std::make_shared<Linear>(cfg.hidden_size, cfg.out_channels, true, true);
-            blocks["audio_norm_out"] = std::make_shared<LayerNorm>(cfg.audio_hidden_size, 1e-6f, false);
-            blocks["audio_proj_out"] = std::make_shared<Linear>(cfg.audio_hidden_size, cfg.audio_out_channels, true, true);
+            blocks["norm_out"]       = std::make_shared<LayerNorm>(config.hidden_size, 1e-6f, false);
+            blocks["proj_out"]       = std::make_shared<Linear>(config.hidden_size, config.out_channels, true, true);
+            blocks["audio_norm_out"] = std::make_shared<LayerNorm>(config.audio_hidden_size, 1e-6f, false);
+            blocks["audio_proj_out"] = std::make_shared<Linear>(config.audio_hidden_size, config.audio_out_channels, true, true);
         }
 
         ggml_tensor* patchify_video(GGMLRunnerContext* ctx, ggml_tensor* x, int64_t n) {
@@ -1293,8 +1431,8 @@ namespace LTXV {
             if (ax == nullptr) {
                 return nullptr;
             }
-            ax = ggml_reshape_4d(ctx->ggml_ctx, ax, cfg.audio_frequency_bins, cfg.num_audio_channels, audio_length, ax->ne[2]);  // [b, t, c, f]
-            ax = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, ax, 0, 2, 1, 3));                                // [b, c, t, f]
+            ax = ggml_reshape_4d(ctx->ggml_ctx, ax, config.audio_frequency_bins, config.num_audio_channels, audio_length, ax->ne[2]);  // [b, t, c, f]
+            ax = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, ax, 0, 2, 1, 3));                                      // [b, c, t, f]
             return ax;
         }
 
@@ -1308,17 +1446,17 @@ namespace LTXV {
             }
 
             bool is_fully_processed_context =
-                context->ne[0] == cfg.cross_attention_dim + cfg.audio_cross_attention_dim &&
+                context->ne[0] == config.cross_attention_dim + config.audio_cross_attention_dim &&
                 context->ne[1] >= 1024;
             bool is_unprocessed_dual_context =
-                context->ne[0] == cfg.cross_attention_dim + cfg.audio_cross_attention_dim &&
+                context->ne[0] == config.cross_attention_dim + config.audio_cross_attention_dim &&
                 context->ne[1] < 1024;
 
             if (is_fully_processed_context) {
-                auto v_context         = ggml_ext_slice(ctx->ggml_ctx, context, 0, 0, cfg.cross_attention_dim);
+                auto v_context         = ggml_ext_slice(ctx->ggml_ctx, context, 0, 0, config.cross_attention_dim);
                 ggml_tensor* a_context = nullptr;
                 if (process_audio_context) {
-                    a_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, cfg.cross_attention_dim, cfg.cross_attention_dim + cfg.audio_cross_attention_dim);
+                    a_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, config.cross_attention_dim, config.cross_attention_dim + config.audio_cross_attention_dim);
                 }
                 return {v_context, a_context};
             }
@@ -1326,32 +1464,32 @@ namespace LTXV {
             ggml_tensor* v_context = context;
             ggml_tensor* a_context = process_audio_context ? context : nullptr;
             if (is_unprocessed_dual_context) {
-                v_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, 0, cfg.cross_attention_dim);
+                v_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, 0, config.cross_attention_dim);
                 if (process_audio_context) {
-                    a_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, cfg.cross_attention_dim, cfg.cross_attention_dim + cfg.audio_cross_attention_dim);
+                    a_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, config.cross_attention_dim, config.cross_attention_dim + config.audio_cross_attention_dim);
                 }
-            } else if (context->ne[0] == cfg.caption_channels * 2) {
-                v_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, 0, cfg.caption_channels);
+            } else if (context->ne[0] == config.caption_channels * 2) {
+                v_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, 0, config.caption_channels);
                 if (process_audio_context) {
-                    a_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, cfg.caption_channels, cfg.caption_channels * 2);
+                    a_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, config.caption_channels, config.caption_channels * 2);
                 }
             }
 
-            if (cfg.caption_proj_before_connector) {
-                if (cfg.use_caption_projection &&
+            if (config.caption_proj_before_connector) {
+                if (config.use_caption_projection &&
                     blocks.count("caption_projection") > 0 &&
                     v_context != nullptr &&
-                    v_context->ne[0] == cfg.caption_channels) {
+                    v_context->ne[0] == config.caption_channels) {
                     auto caption_projection = std::dynamic_pointer_cast<NormSingleLinearTextProjection>(blocks["caption_projection"]);
                     if (caption_projection != nullptr) {
                         v_context = caption_projection->forward(ctx, v_context);
                     }
                 }
                 if (process_audio_context &&
-                    cfg.use_audio_caption_projection &&
+                    config.use_audio_caption_projection &&
                     blocks.count("audio_caption_projection") > 0 &&
                     a_context != nullptr &&
-                    a_context->ne[0] == cfg.caption_channels) {
+                    a_context->ne[0] == config.caption_channels) {
                     auto caption_projection = std::dynamic_pointer_cast<NormSingleLinearTextProjection>(blocks["audio_caption_projection"]);
                     if (caption_projection != nullptr) {
                         a_context = caption_projection->forward(ctx, a_context);
@@ -1359,34 +1497,34 @@ namespace LTXV {
                 }
             }
 
-            if (cfg.use_connector && v_context != nullptr && v_context->ne[0] == cfg.connector_hidden_size) {
+            if (config.use_connector && v_context != nullptr && v_context->ne[0] == config.connector_hidden_size) {
                 auto connector = std::dynamic_pointer_cast<Embeddings1DConnector>(blocks["video_embeddings_connector"]);
                 v_context      = connector->forward(ctx, v_context, video_connector_pe);
             }
             if (process_audio_context &&
-                cfg.use_audio_connector &&
+                config.use_audio_connector &&
                 a_context != nullptr &&
-                a_context->ne[0] == cfg.audio_connector_hidden_size) {
+                a_context->ne[0] == config.audio_connector_hidden_size) {
                 auto connector = std::dynamic_pointer_cast<Embeddings1DConnector>(blocks["audio_embeddings_connector"]);
                 a_context      = connector->forward(ctx, a_context, audio_connector_pe);
             }
 
-            if (!cfg.caption_proj_before_connector &&
-                cfg.use_caption_projection &&
+            if (!config.caption_proj_before_connector &&
+                config.use_caption_projection &&
                 blocks.count("caption_projection") > 0 &&
                 v_context != nullptr &&
-                v_context->ne[0] == cfg.caption_channels) {
+                v_context->ne[0] == config.caption_channels) {
                 auto caption_projection = std::dynamic_pointer_cast<PixArtAlphaTextProjection>(blocks["caption_projection"]);
                 if (caption_projection != nullptr) {
                     v_context = caption_projection->forward(ctx, v_context);
                 }
             }
             if (process_audio_context &&
-                !cfg.caption_proj_before_connector &&
-                cfg.use_audio_caption_projection &&
+                !config.caption_proj_before_connector &&
+                config.use_audio_caption_projection &&
                 blocks.count("audio_caption_projection") > 0 &&
                 a_context != nullptr &&
-                a_context->ne[0] == cfg.caption_channels) {
+                a_context->ne[0] == config.caption_channels) {
                 auto caption_projection = std::dynamic_pointer_cast<PixArtAlphaTextProjection>(blocks["audio_caption_projection"]);
                 if (caption_projection != nullptr) {
                     a_context = caption_projection->forward(ctx, a_context);
@@ -1428,8 +1566,8 @@ namespace LTXV {
             auto audio_norm_out      = std::dynamic_pointer_cast<LayerNorm>(blocks["audio_norm_out"]);
             auto audio_proj_out      = std::dynamic_pointer_cast<Linear>(blocks["audio_proj_out"]);
 
-            GGML_ASSERT(vx->ne[3] % cfg.in_channels == 0);
-            int64_t n          = vx->ne[3] / cfg.in_channels;
+            GGML_ASSERT(vx->ne[3] % config.in_channels == 0);
+            int64_t n          = vx->ne[3] / config.in_channels;
             int64_t width      = vx->ne[0];
             int64_t height     = vx->ne[1];
             int64_t frames     = vx->ne[2];
@@ -1452,20 +1590,20 @@ namespace LTXV {
                 a_context = ggml_cont(ctx->ggml_ctx, a_context);
             }
 
-            auto v_timestep_scaled = ggml_ext_scale(ctx->ggml_ctx, timestep, cfg.timestep_scale_multiplier);
+            auto v_timestep_scaled = ggml_ext_scale(ctx->ggml_ctx, timestep, config.timestep_scale_multiplier);
             auto v_pair            = adaln_single->forward(ctx, v_timestep_scaled);
             auto v_timestep_mod    = v_pair.first;
             auto v_embedded_time   = v_pair.second;
 
             ggml_tensor* effective_audio_timestep = audio_timestep != nullptr ? audio_timestep : timestep;
-            auto a_timestep_scaled                = ggml_ext_scale(ctx->ggml_ctx, effective_audio_timestep, cfg.timestep_scale_multiplier);
+            auto a_timestep_scaled                = ggml_ext_scale(ctx->ggml_ctx, effective_audio_timestep, config.timestep_scale_multiplier);
             auto a_pair                           = audio_adaln_single->forward(ctx, a_timestep_scaled);
             auto a_timestep_mod                   = a_pair.first;
             auto a_embedded_time                  = a_pair.second;
 
             ggml_tensor* v_prompt_timestep_mod = nullptr;
             ggml_tensor* a_prompt_timestep_mod = nullptr;
-            if (cfg.cross_attention_adaln) {
+            if (config.cross_attention_adaln) {
                 auto prompt_adaln_single       = std::dynamic_pointer_cast<AdaLayerNormSingle>(blocks["prompt_adaln_single"]);
                 auto audio_prompt_adaln_single = std::dynamic_pointer_cast<AdaLayerNormSingle>(blocks["audio_prompt_adaln_single"]);
                 v_prompt_timestep_mod          = prompt_adaln_single->forward(ctx, a_timestep_scaled).first;
@@ -1474,7 +1612,7 @@ namespace LTXV {
 
             auto av_ca_video_timestep = repeat_scalar_timestep_like(ctx, effective_audio_timestep, timestep);
             auto av_ca_audio_timestep = effective_audio_timestep;
-            auto av_ca_factor         = cfg.av_ca_timestep_scale_multiplier / cfg.timestep_scale_multiplier;
+            auto av_ca_factor         = config.av_ca_timestep_scale_multiplier / config.timestep_scale_multiplier;
             auto av_ca_video_scale_shift_timestep =
                 std::dynamic_pointer_cast<AdaLayerNormSingle>(blocks["av_ca_video_scale_shift_adaln_single"])->forward(ctx, av_ca_video_timestep).first;
             auto av_ca_a2v_gate_noise_timestep =
@@ -1491,7 +1629,7 @@ namespace LTXV {
             sd::ggml_graph_cut::mark_graph_cut(vx, "ltxav.prelude", "vx");
             sd::ggml_graph_cut::mark_graph_cut(ax, "ltxav.prelude", "ax");
 
-            for (int i = 0; i < cfg.num_layers; i++) {
+            for (int i = 0; i < config.num_layers; i++) {
                 auto block = std::dynamic_pointer_cast<BasicAVTransformerBlock>(blocks["transformer_blocks." + std::to_string(i)]);
                 auto out   = block->forward(ctx,
                                             vx,
@@ -1517,14 +1655,14 @@ namespace LTXV {
                 sd::ggml_graph_cut::mark_graph_cut(ax, "ltxav.transformer_blocks." + std::to_string(i), "ax");
             }
 
-            auto v_shift_scale = get_output_scale_shift(ctx, params["scale_shift_table"], v_embedded_time, cfg.hidden_size);
+            auto v_shift_scale = get_output_scale_shift(ctx, params["scale_shift_table"], v_embedded_time, config.hidden_size);
             vx                 = norm_out->forward(ctx, vx);
             vx                 = modulate(ctx->ggml_ctx, vx, v_shift_scale[0], v_shift_scale[1]);
             vx                 = proj_out->forward(ctx, vx);
             vx                 = unpatchify_video(ctx, vx, width, height, frames);
 
             if (ax != nullptr && audio_time > 0) {
-                auto a_shift_scale = get_output_scale_shift(ctx, params["audio_scale_shift_table"], a_embedded_time, cfg.audio_hidden_size);
+                auto a_shift_scale = get_output_scale_shift(ctx, params["audio_scale_shift_table"], a_embedded_time, config.audio_hidden_size);
                 ax                 = audio_norm_out->forward(ctx, ax);
                 ax                 = modulate(ctx->ggml_ctx, ax, a_shift_scale[0], a_shift_scale[1]);
                 ax                 = audio_proj_out->forward(ctx, ax);
@@ -1536,7 +1674,7 @@ namespace LTXV {
     };
 
     struct LTXAVRunner : public DiffusionModelRunner {
-        LTXAVParams params;
+        LTXAVConfig config;
         LTXAVModelBlock model;
         std::vector<float> video_pe_vec;
         std::vector<float> audio_pe_vec;
@@ -1547,124 +1685,13 @@ namespace LTXV {
         sd::Tensor<float> vx_input_cache;
         sd::Tensor<float> ax_input_cache;
 
-        static int64_t infer_gate_heads(const String2TensorStorage& tensor_storage_map,
-                                        const std::string& bias_name,
-                                        int64_t fallback_heads) {
-            auto it = tensor_storage_map.find(bias_name);
-            if (it != tensor_storage_map.end()) {
-                return it->second.ne[0];
-            }
-            return fallback_heads;
-        }
-
         LTXAVRunner(ggml_backend_t backend,
-                    ggml_backend_t params_backend,
-                    const String2TensorStorage& tensor_storage_map = {},
-                    const std::string& prefix                      = "model.diffusion_model")
-            : DiffusionModelRunner(backend, params_backend, prefix),
-              params(),
-              model(params) {
-            auto patchify_proj_iter = tensor_storage_map.find(prefix + ".patchify_proj.weight");
-            if (patchify_proj_iter != tensor_storage_map.end()) {
-                params.in_channels         = patchify_proj_iter->second.ne[0];
-                params.hidden_size         = patchify_proj_iter->second.ne[1];
-                int64_t video_heads        = infer_gate_heads(tensor_storage_map, prefix + ".transformer_blocks.0.attn1.to_gate_logits.bias", 32);
-                auto attn_layout           = infer_attention_layout(params.hidden_size, video_heads);
-                params.num_attention_heads = attn_layout.first;
-                params.attention_head_dim  = attn_layout.second;
-            }
-
-            auto audio_patchify_proj_iter = tensor_storage_map.find(prefix + ".audio_patchify_proj.weight");
-            if (audio_patchify_proj_iter != tensor_storage_map.end()) {
-                params.audio_in_channels         = audio_patchify_proj_iter->second.ne[0];
-                params.audio_hidden_size         = audio_patchify_proj_iter->second.ne[1];
-                params.audio_out_channels        = params.audio_in_channels;
-                int64_t audio_heads              = infer_gate_heads(tensor_storage_map, prefix + ".transformer_blocks.0.audio_attn1.to_gate_logits.bias", 32);
-                auto audio_attn_layout           = infer_attention_layout(params.audio_hidden_size, audio_heads);
-                params.audio_num_attention_heads = audio_attn_layout.first;
-                params.audio_attention_head_dim  = audio_attn_layout.second;
-            }
-
-            auto proj_out_iter = tensor_storage_map.find(prefix + ".proj_out.weight");
-            if (proj_out_iter != tensor_storage_map.end()) {
-                params.out_channels = proj_out_iter->second.ne[1];
-            }
-            auto audio_proj_out_iter = tensor_storage_map.find(prefix + ".audio_proj_out.weight");
-            if (audio_proj_out_iter != tensor_storage_map.end()) {
-                params.audio_out_channels = audio_proj_out_iter->second.ne[1];
-            }
-
-            auto attn2_iter = tensor_storage_map.find(prefix + ".transformer_blocks.0.attn2.to_k.weight");
-            if (attn2_iter != tensor_storage_map.end()) {
-                params.cross_attention_dim = attn2_iter->second.ne[0];
-            }
-            auto audio_attn2_iter = tensor_storage_map.find(prefix + ".transformer_blocks.0.audio_attn2.to_k.weight");
-            if (audio_attn2_iter != tensor_storage_map.end()) {
-                params.audio_cross_attention_dim = audio_attn2_iter->second.ne[0];
-            }
-            if (tensor_storage_map.find(prefix + ".transformer_blocks.0.prompt_scale_shift_table") != tensor_storage_map.end()) {
-                params.cross_attention_adaln = true;
-            }
-            if (tensor_storage_map.find(prefix + ".transformer_blocks.0.attn1.to_gate_logits.weight") != tensor_storage_map.end() ||
-                tensor_storage_map.find(prefix + ".transformer_blocks.0.audio_attn1.to_gate_logits.weight") != tensor_storage_map.end()) {
-                params.self_attention_gated = true;
-            }
-            if (tensor_storage_map.find(prefix + ".transformer_blocks.0.attn2.to_gate_logits.weight") != tensor_storage_map.end() ||
-                tensor_storage_map.find(prefix + ".transformer_blocks.0.audio_attn2.to_gate_logits.weight") != tensor_storage_map.end()) {
-                params.cross_attention_gated = true;
-            }
-            if (tensor_storage_map.find(prefix + ".caption_projection.linear_1.weight") == tensor_storage_map.end() &&
-                tensor_storage_map.find(prefix + ".caption_projection.linear_2.weight") == tensor_storage_map.end()) {
-                params.use_caption_projection = false;
-            }
-            if (tensor_storage_map.find(prefix + ".audio_caption_projection.linear_1.weight") == tensor_storage_map.end() &&
-                tensor_storage_map.find(prefix + ".audio_caption_projection.linear_2.weight") == tensor_storage_map.end()) {
-                params.use_audio_caption_projection = false;
-            }
-
-            params.num_layers = count_prefix_blocks(tensor_storage_map, prefix + ".", "transformer_blocks.");
-
-            auto connector_iter = tensor_storage_map.find(prefix + ".video_embeddings_connector.transformer_1d_blocks.0.attn1.to_q.weight");
-            if (connector_iter != tensor_storage_map.end()) {
-                params.use_connector         = true;
-                params.connector_hidden_size = connector_iter->second.ne[1];
-                int64_t connector_heads      = infer_gate_heads(tensor_storage_map,
-                                                                prefix + ".video_embeddings_connector.transformer_1d_blocks.0.attn1.to_gate_logits.bias",
-                                                                32);
-                auto connector_layout        = infer_attention_layout(params.connector_hidden_size, connector_heads);
-                params.connector_num_heads   = connector_layout.first;
-                params.connector_head_dim    = connector_layout.second;
-                params.connector_num_layers  = count_prefix_blocks(tensor_storage_map, prefix + ".video_embeddings_connector.", "transformer_1d_blocks.");
-                auto register_iter           = tensor_storage_map.find(prefix + ".video_embeddings_connector.learnable_registers");
-                if (register_iter != tensor_storage_map.end()) {
-                    params.connector_num_registers = register_iter->second.ne[1];
-                }
-                if (tensor_storage_map.find(prefix + ".video_embeddings_connector.transformer_1d_blocks.0.attn1.to_gate_logits.weight") != tensor_storage_map.end()) {
-                    params.connector_apply_gated_attention = true;
-                }
-            }
-
-            auto audio_connector_iter = tensor_storage_map.find(prefix + ".audio_embeddings_connector.transformer_1d_blocks.0.attn1.to_q.weight");
-            if (audio_connector_iter != tensor_storage_map.end()) {
-                params.use_audio_connector         = true;
-                params.audio_connector_hidden_size = audio_connector_iter->second.ne[1];
-                int64_t connector_heads            = infer_gate_heads(tensor_storage_map,
-                                                                      prefix + ".audio_embeddings_connector.transformer_1d_blocks.0.attn1.to_gate_logits.bias",
-                                                                      32);
-                auto connector_layout              = infer_attention_layout(params.audio_connector_hidden_size, connector_heads);
-                params.audio_connector_num_heads   = connector_layout.first;
-                params.audio_connector_head_dim    = connector_layout.second;
-                params.audio_connector_num_layers  = count_prefix_blocks(tensor_storage_map, prefix + ".audio_embeddings_connector.", "transformer_1d_blocks.");
-                auto register_iter                 = tensor_storage_map.find(prefix + ".audio_embeddings_connector.learnable_registers");
-                if (register_iter != tensor_storage_map.end()) {
-                    params.audio_connector_num_registers = register_iter->second.ne[1];
-                }
-                if (tensor_storage_map.find(prefix + ".audio_embeddings_connector.transformer_1d_blocks.0.attn1.to_gate_logits.weight") != tensor_storage_map.end()) {
-                    params.audio_connector_apply_gated_attention = true;
-                }
-            }
-
-            model = LTXAVModelBlock(params);
+                    const String2TensorStorage& tensor_storage_map      = {},
+                    const std::string& prefix                           = "model.diffusion_model",
+                    std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : DiffusionModelRunner(backend, prefix, weight_manager),
+              config(LTXAVConfig::detect_from_weights(tensor_storage_map, prefix)),
+              model(config) {
             model.init(params_ctx, tensor_storage_map, prefix);
         }
 
@@ -1692,21 +1719,21 @@ namespace LTXV {
             int64_t total_channels = x_tensor.shape()[3];
             int64_t spatial_size   = width * height * frames;
 
-            GGML_ASSERT(total_channels >= params.in_channels);
+            GGML_ASSERT(total_channels >= config.in_channels);
 
-            sd::Tensor<float> vx({width, height, frames, params.in_channels});
-            size_t video_values = static_cast<size_t>(params.in_channels * spatial_size);
+            sd::Tensor<float> vx({width, height, frames, config.in_channels});
+            size_t video_values = static_cast<size_t>(config.in_channels * spatial_size);
             std::copy_n(x_tensor.data(), video_values, vx.data());
 
-            if (audio_length <= 0 || total_channels == params.in_channels) {
+            if (audio_length <= 0 || total_channels == config.in_channels) {
                 return {vx, {}};
             }
 
-            int64_t needed_audio_values = static_cast<int64_t>(audio_length) * params.num_audio_channels * params.audio_frequency_bins;
-            int64_t packed_audio_values = (total_channels - params.in_channels) * spatial_size;
+            int64_t needed_audio_values = static_cast<int64_t>(audio_length) * config.num_audio_channels * config.audio_frequency_bins;
+            int64_t packed_audio_values = (total_channels - config.in_channels) * spatial_size;
             GGML_ASSERT(packed_audio_values >= needed_audio_values);
 
-            sd::Tensor<float> ax({params.audio_frequency_bins, audio_length, params.num_audio_channels, 1});
+            sd::Tensor<float> ax({config.audio_frequency_bins, audio_length, config.num_audio_channels, 1});
             const float* audio_src = x_tensor.data() + video_values;
             std::copy_n(audio_src, static_cast<size_t>(needed_audio_values), ax.data());
             return {vx, ax};
@@ -1767,25 +1794,25 @@ namespace LTXV {
             if (has_video_positions) {
                 GGML_ASSERT(video_positions_tensor.shape()[2] == video_token_count);
                 video_pe_vec = build_video_rope_matrix_from_positions(video_positions_tensor,
-                                                                      static_cast<int>(params.hidden_size),
-                                                                      static_cast<int>(params.num_attention_heads),
-                                                                      params.positional_embedding_theta,
-                                                                      params.positional_embedding_max_pos,
-                                                                      params.use_middle_indices_grid);
+                                                                      static_cast<int>(config.hidden_size),
+                                                                      static_cast<int>(config.num_attention_heads),
+                                                                      config.positional_embedding_theta,
+                                                                      config.positional_embedding_max_pos,
+                                                                      config.use_middle_indices_grid);
             } else {
                 video_pe_vec = build_video_rope_matrix(vx->ne[0],
                                                        vx->ne[1],
                                                        vx->ne[2],
-                                                       static_cast<int>(params.hidden_size),
-                                                       static_cast<int>(params.num_attention_heads),
+                                                       static_cast<int>(config.hidden_size),
+                                                       static_cast<int>(config.num_attention_heads),
                                                        video_frame_rate,
-                                                       params.positional_embedding_theta,
-                                                       params.positional_embedding_max_pos,
-                                                       params.vae_scale_factors,
-                                                       params.causal_temporal_positioning,
-                                                       params.use_middle_indices_grid);
+                                                       config.positional_embedding_theta,
+                                                       config.positional_embedding_max_pos,
+                                                       config.vae_scale_factors,
+                                                       config.causal_temporal_positioning,
+                                                       config.use_middle_indices_grid);
             }
-            auto video_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, params.attention_head_dim / 2, video_token_count * params.num_attention_heads);
+            auto video_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.attention_head_dim / 2, video_token_count * config.num_attention_heads);
             ggml_set_name(video_pe, "ltxav_video_pe");
             set_backend_tensor_data(video_pe, video_pe_vec.data());
 
@@ -1794,66 +1821,66 @@ namespace LTXV {
             ggml_tensor* audio_cross_pe = nullptr;
             if (ax != nullptr && ggml_nelements(ax) > 0 && ax->ne[1] > 0) {
                 audio_pe_vec = build_audio_rope_matrix(ax->ne[1],
-                                                       static_cast<int>(params.audio_hidden_size),
-                                                       static_cast<int>(params.audio_num_attention_heads),
-                                                       params.positional_embedding_theta,
-                                                       params.audio_positional_embedding_max_pos[0],
-                                                       params.use_middle_indices_grid);
-                audio_pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, params.audio_attention_head_dim / 2, ax->ne[1] * params.audio_num_attention_heads);
+                                                       static_cast<int>(config.audio_hidden_size),
+                                                       static_cast<int>(config.audio_num_attention_heads),
+                                                       config.positional_embedding_theta,
+                                                       config.audio_positional_embedding_max_pos[0],
+                                                       config.use_middle_indices_grid);
+                audio_pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.audio_attention_head_dim / 2, ax->ne[1] * config.audio_num_attention_heads);
                 ggml_set_name(audio_pe, "ltxav_audio_pe");
                 set_backend_tensor_data(audio_pe, audio_pe_vec.data());
 
-                int temporal_max_pos = std::max(params.positional_embedding_max_pos[0], params.audio_positional_embedding_max_pos[0]);
+                int temporal_max_pos = std::max(config.positional_embedding_max_pos[0], config.audio_positional_embedding_max_pos[0]);
                 if (has_video_positions) {
                     video_cross_pe_vec = build_video_temporal_rope_matrix_from_positions(video_positions_tensor,
-                                                                                         static_cast<int>(params.audio_cross_attention_dim),
-                                                                                         static_cast<int>(params.audio_num_attention_heads),
-                                                                                         params.positional_embedding_theta,
+                                                                                         static_cast<int>(config.audio_cross_attention_dim),
+                                                                                         static_cast<int>(config.audio_num_attention_heads),
+                                                                                         config.positional_embedding_theta,
                                                                                          temporal_max_pos,
                                                                                          true);
                 } else {
                     video_cross_pe_vec = build_video_temporal_rope_matrix(vx->ne[0],
                                                                           vx->ne[1],
                                                                           vx->ne[2],
-                                                                          static_cast<int>(params.audio_cross_attention_dim),
-                                                                          static_cast<int>(params.audio_num_attention_heads),
+                                                                          static_cast<int>(config.audio_cross_attention_dim),
+                                                                          static_cast<int>(config.audio_num_attention_heads),
                                                                           video_frame_rate,
-                                                                          params.positional_embedding_theta,
+                                                                          config.positional_embedding_theta,
                                                                           temporal_max_pos,
-                                                                          std::get<0>(params.vae_scale_factors),
-                                                                          params.causal_temporal_positioning,
+                                                                          std::get<0>(config.vae_scale_factors),
+                                                                          config.causal_temporal_positioning,
                                                                           true);
                 }
-                video_cross_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, params.audio_attention_head_dim / 2, video_token_count * params.audio_num_attention_heads);
+                video_cross_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.audio_attention_head_dim / 2, video_token_count * config.audio_num_attention_heads);
                 ggml_set_name(video_cross_pe, "ltxav_video_cross_pe");
                 set_backend_tensor_data(video_cross_pe, video_cross_pe_vec.data());
 
                 audio_cross_pe_vec = build_audio_rope_matrix(ax->ne[1],
-                                                             static_cast<int>(params.audio_cross_attention_dim),
-                                                             static_cast<int>(params.audio_num_attention_heads),
-                                                             params.positional_embedding_theta,
+                                                             static_cast<int>(config.audio_cross_attention_dim),
+                                                             static_cast<int>(config.audio_num_attention_heads),
+                                                             config.positional_embedding_theta,
                                                              temporal_max_pos,
                                                              true);
-                audio_cross_pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, params.audio_attention_head_dim / 2, ax->ne[1] * params.audio_num_attention_heads);
+                audio_cross_pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.audio_attention_head_dim / 2, ax->ne[1] * config.audio_num_attention_heads);
                 ggml_set_name(audio_cross_pe, "ltxav_audio_cross_pe");
                 set_backend_tensor_data(audio_cross_pe, audio_cross_pe_vec.data());
             }
 
             bool needs_video_connector_pe =
-                params.use_connector &&
+                config.use_connector &&
                 context != nullptr &&
-                (context->ne[0] == params.connector_hidden_size ||
-                 ((context->ne[0] == params.cross_attention_dim + params.audio_cross_attention_dim ||
-                   context->ne[0] == params.caption_channels * 2) &&
+                (context->ne[0] == config.connector_hidden_size ||
+                 ((context->ne[0] == config.cross_attention_dim + config.audio_cross_attention_dim ||
+                   context->ne[0] == config.caption_channels * 2) &&
                   context->ne[1] < 1024));
             ggml_tensor* video_connector_pe = nullptr;
             if (needs_video_connector_pe) {
                 int64_t seq_len      = context->ne[1];
                 int64_t target_len   = std::max<int64_t>(1024, seq_len);
-                int64_t duplications = (target_len + params.connector_num_registers - 1) / params.connector_num_registers;
-                int64_t full_len     = seq_len + duplications * params.connector_num_registers - seq_len;
-                connector_pe_vec     = build_1d_rope_matrix(full_len, static_cast<int>(params.connector_hidden_size), static_cast<int>(params.connector_num_heads), 10000.f, 4096.f, true);
-                video_connector_pe   = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, params.connector_head_dim / 2, full_len * params.connector_num_heads);
+                int64_t duplications = (target_len + config.connector_num_registers - 1) / config.connector_num_registers;
+                int64_t full_len     = seq_len + duplications * config.connector_num_registers - seq_len;
+                connector_pe_vec     = build_1d_rope_matrix(full_len, static_cast<int>(config.connector_hidden_size), static_cast<int>(config.connector_num_heads), 10000.f, 4096.f, true);
+                video_connector_pe   = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.connector_head_dim / 2, full_len * config.connector_num_heads);
                 ggml_set_name(video_connector_pe, "ltxav_video_connector_pe");
                 set_backend_tensor_data(video_connector_pe, connector_pe_vec.data());
             }
@@ -1864,20 +1891,20 @@ namespace LTXV {
                 ax->ne[1] > 0;
             bool needs_audio_connector_pe =
                 run_audio_context &&
-                params.use_audio_connector &&
+                config.use_audio_connector &&
                 context != nullptr &&
-                (context->ne[0] == params.audio_connector_hidden_size ||
-                 ((context->ne[0] == params.cross_attention_dim + params.audio_cross_attention_dim ||
-                   context->ne[0] == params.caption_channels * 2) &&
+                (context->ne[0] == config.audio_connector_hidden_size ||
+                 ((context->ne[0] == config.cross_attention_dim + config.audio_cross_attention_dim ||
+                   context->ne[0] == config.caption_channels * 2) &&
                   context->ne[1] < 1024));
             ggml_tensor* audio_connector_pe = nullptr;
             if (needs_audio_connector_pe) {
                 int64_t seq_len        = context->ne[1];
                 int64_t target_len     = std::max<int64_t>(1024, seq_len);
-                int64_t duplications   = (target_len + params.audio_connector_num_registers - 1) / params.audio_connector_num_registers;
-                int64_t full_len       = seq_len + duplications * params.audio_connector_num_registers - seq_len;
-                audio_connector_pe_vec = build_1d_rope_matrix(full_len, static_cast<int>(params.audio_connector_hidden_size), static_cast<int>(params.audio_connector_num_heads), 10000.f, 4096.f, true);
-                audio_connector_pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, params.audio_connector_head_dim / 2, full_len * params.audio_connector_num_heads);
+                int64_t duplications   = (target_len + config.audio_connector_num_registers - 1) / config.audio_connector_num_registers;
+                int64_t full_len       = seq_len + duplications * config.audio_connector_num_registers - seq_len;
+                audio_connector_pe_vec = build_1d_rope_matrix(full_len, static_cast<int>(config.audio_connector_hidden_size), static_cast<int>(config.audio_connector_num_heads), 10000.f, 4096.f, true);
+                audio_connector_pe     = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.audio_connector_head_dim / 2, full_len * config.audio_connector_num_heads);
                 ggml_set_name(audio_connector_pe, "ltxav_audio_connector_pe");
                 set_backend_tensor_data(audio_connector_pe, audio_connector_pe_vec.data());
             }
@@ -1912,7 +1939,7 @@ namespace LTXV {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, audio_x, audio_timesteps, audio_length, frame_rate, video_positions);
             };
-            auto out = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            auto out = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
             return out;
         }
 
@@ -1998,7 +2025,8 @@ namespace LTXV {
             ggml_backend_t backend = sd_backend_cpu_init();
             LOG_INFO("loading ltxav from '%s'", model_path.c_str());
 
-            ModelLoader model_loader;
+            auto model_manager        = std::make_shared<ModelManager>();
+            ModelLoader& model_loader = model_manager->loader();
             if (!model_loader.init_from_file_and_convert_name(model_path, "model.diffusion_model.")) {
                 LOG_ERROR("init model loader from file failed: '%s'", model_path.c_str());
                 return;
@@ -2013,19 +2041,18 @@ namespace LTXV {
 
             auto& tensor_storage_map           = model_loader.get_tensor_storage_map();
             std::shared_ptr<LTXAVRunner> ltxav = std::make_shared<LTXAVRunner>(backend,
-                                                                               backend,
                                                                                tensor_storage_map,
-                                                                               "model.diffusion_model");
+                                                                               "model.diffusion_model",
+                                                                               model_manager);
 
-            if (!ltxav->alloc_params_buffer()) {
-                LOG_ERROR("ltxav buffer allocation failed");
-                return;
-            }
-            std::map<std::string, ggml_tensor*> tensors;
-            ltxav->get_param_tensors(tensors, "model.diffusion_model");
-
-            if (!model_loader.load_tensors(tensors)) {
-                LOG_ERROR("load tensors from model loader failed");
+            if (!model_manager->register_runner_params("LTXAV test",
+                                                       *ltxav,
+                                                       "model.diffusion_model",
+                                                       ModelManager::ResidencyMode::ParamBackend,
+                                                       backend,
+                                                       backend) ||
+                !model_manager->validate_registered_tensors()) {
+                LOG_ERROR("register ltxav tensors with model manager failed");
                 return;
             }
 
@@ -2036,4 +2063,4 @@ namespace LTXV {
 
 };  // namespace LTXV
 
-#endif
+#endif  // __SD_MODEL_DIFFUSION_LTXV_HPP__

src/model/diffusion/mmdit.hpp

@@ -1,42 +1,137 @@
-#ifndef __MMDIT_HPP__
-#define __MMDIT_HPP__
+#ifndef __SD_MODEL_DIFFUSION_MMDIT_HPP__
+#define __SD_MODEL_DIFFUSION_MMDIT_HPP__
 
+#include <algorithm>
 #include <memory>
+#include <string>
+#include <vector>
 
-#include "diffusion_model.hpp"
-#include "ggml_extend.hpp"
-#include "model.h"
+#include "core/ggml_extend.hpp"
+#include "model/common/block.hpp"
+#include "model/diffusion/model.hpp"
+#include "model_loader.h"
 
 #define MMDIT_GRAPH_SIZE 10240
 
-struct Mlp : public GGMLBlock {
-public:
-    Mlp(int64_t in_features,
-        int64_t hidden_features = -1,
-        int64_t out_features    = -1,
-        bool bias               = true) {
-        // act_layer is always lambda: nn.GELU(approximate="tanh")
-        // norm_layer is always None
-        // use_conv is always False
-        if (hidden_features == -1) {
-            hidden_features = in_features;
-        }
-        if (out_features == -1) {
-            out_features = in_features;
-        }
-        blocks["fc1"] = std::shared_ptr<GGMLBlock>(new Linear(in_features, hidden_features, bias));
-        blocks["fc2"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_features, out_features, bias));
-    }
+struct MMDiTConfig {
+    int64_t input_size               = -1;
+    int patch_size                   = 2;
+    int64_t in_channels              = 16;
+    int64_t d_self                   = -1;  // >=0 for MMdiT-X
+    int64_t depth                    = 24;
+    float mlp_ratio                  = 4.0f;
+    int64_t adm_in_channels          = 2048;
+    int64_t out_channels             = 16;
+    int64_t pos_embed_max_size       = 192;
+    int64_t num_patches              = 36864;  // 192 * 192
+    int64_t context_size             = 4096;
+    int64_t context_embedder_out_dim = 1536;
+    int64_t hidden_size              = 1536;
+    std::string qk_norm;
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
-        // x: [N, n_token, in_features]
-        auto fc1 = std::dynamic_pointer_cast<Linear>(blocks["fc1"]);
-        auto fc2 = std::dynamic_pointer_cast<Linear>(blocks["fc2"]);
+    static MMDiTConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
+        MMDiTConfig config;
+        bool has_weight_config = false;
+        bool has_pos_embed     = false;
+        bool has_hidden_size   = false;
+        bool has_context_embed = false;
 
-        x = fc1->forward(ctx, x);
-        x = ggml_ext_gelu(ctx->ggml_ctx, x, true);
-        x = fc2->forward(ctx, x);
-        return x;
+        for (const auto& [name, tensor_storage] : tensor_storage_map) {
+            if (!starts_with(name, prefix)) {
+                continue;
+            }
+
+            if (name.find("x_embedder.proj.weight") != std::string::npos && tensor_storage.n_dims == 4) {
+                has_weight_config  = true;
+                has_hidden_size    = true;
+                config.patch_size  = static_cast<int>(tensor_storage.ne[0]);
+                config.in_channels = tensor_storage.ne[2];
+                config.hidden_size = tensor_storage.ne[3];
+            } else if (name.find("t_embedder.mlp.0.weight") != std::string::npos && tensor_storage.n_dims == 2) {
+                has_weight_config  = true;
+                has_hidden_size    = true;
+                config.hidden_size = tensor_storage.ne[1];
+            } else if (name.find("y_embedder.mlp.0.weight") != std::string::npos && tensor_storage.n_dims == 2) {
+                has_weight_config      = true;
+                has_hidden_size        = true;
+                config.adm_in_channels = tensor_storage.ne[0];
+                config.hidden_size     = tensor_storage.ne[1];
+            } else if (name.find("context_embedder.weight") != std::string::npos && tensor_storage.n_dims == 2) {
+                has_weight_config               = true;
+                has_context_embed               = true;
+                config.context_size             = tensor_storage.ne[0];
+                config.context_embedder_out_dim = tensor_storage.ne[1];
+            } else if (name.find("final_layer.linear.weight") != std::string::npos && tensor_storage.n_dims == 2) {
+                has_weight_config  = true;
+                has_hidden_size    = true;
+                config.hidden_size = tensor_storage.ne[0];
+                int64_t patch_area = static_cast<int64_t>(config.patch_size) * config.patch_size;
+                if (patch_area > 0) {
+                    config.out_channels = tensor_storage.ne[1] / patch_area;
+                }
+            } else if (name.find("pos_embed") != std::string::npos && tensor_storage.n_dims == 3) {
+                has_weight_config  = true;
+                has_pos_embed      = true;
+                has_hidden_size    = true;
+                config.hidden_size = tensor_storage.ne[0];
+                config.num_patches = tensor_storage.ne[1];
+                for (int64_t size = 1; size * size <= config.num_patches; size++) {
+                    if (size * size == config.num_patches) {
+                        config.pos_embed_max_size = size;
+                        break;
+                    }
+                }
+            }
+
+            size_t jb = name.find("joint_blocks.");
+            if (jb == std::string::npos) {
+                continue;
+            }
+
+            has_weight_config      = true;
+            std::string block_name = name.substr(jb);
+            int64_t block_depth    = atoi(block_name.substr(13, block_name.find(".", 13)).c_str());
+            if (block_depth + 1 > config.depth) {
+                config.depth = block_depth + 1;
+            }
+            if (block_name.find("attn.ln") != std::string::npos) {
+                if (block_name.find(".bias") != std::string::npos) {
+                    config.qk_norm = "ln";
+                } else {
+                    config.qk_norm = "rms";
+                }
+            }
+            if (block_name.find("attn2") != std::string::npos) {
+                if (block_depth > config.d_self) {
+                    config.d_self = block_depth;
+                }
+            }
+        }
+
+        if (!has_pos_embed && config.d_self >= 0) {
+            config.pos_embed_max_size *= 2;
+            config.num_patches *= 4;
+        }
+        if (!has_hidden_size || config.hidden_size <= 0) {
+            config.hidden_size = 64 * config.depth;
+        }
+        if (!has_context_embed || config.context_embedder_out_dim <= 0) {
+            config.context_embedder_out_dim = config.hidden_size;
+        }
+
+        if (has_weight_config) {
+            LOG_DEBUG("mmdit: num_layers = %" PRId64 ", num_mmdit_x_layers = %" PRId64 ", hidden_size = %" PRId64 ", patch_size = %d, in_channels = %" PRId64 ", out_channels = %" PRId64 ", context_size = %" PRId64 ", adm_in_channels = %" PRId64 ", qk_norm = %s",
+                      config.depth,
+                      config.d_self + 1,
+                      config.hidden_size,
+                      config.patch_size,
+                      config.in_channels,
+                      config.out_channels,
+                      config.context_size,
+                      config.adm_in_channels,
+                      config.qk_norm.empty() ? "none" : config.qk_norm.c_str());
+        }
+        return config;
     }
 };
 
@@ -612,28 +707,16 @@ public:
 struct MMDiT : public GGMLBlock {
     // Diffusion model with a Transformer backbone.
 protected:
-    int64_t input_size               = -1;
-    int patch_size                   = 2;
-    int64_t in_channels              = 16;
-    int64_t d_self                   = -1;  // >=0 for MMdiT-X
-    int64_t depth                    = 24;
-    float mlp_ratio                  = 4.0f;
-    int64_t adm_in_channels          = 2048;
-    int64_t out_channels             = 16;
-    int64_t pos_embed_max_size       = 192;
-    int64_t num_patchs               = 36864;  // 192 * 192
-    int64_t context_size             = 4096;
-    int64_t context_embedder_out_dim = 1536;
-    int64_t hidden_size;
-    std::string qk_norm;
-
     void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") override {
         enum ggml_type wtype = GGML_TYPE_F32;
-        params["pos_embed"]  = ggml_new_tensor_3d(ctx, wtype, hidden_size, num_patchs, 1);
+        params["pos_embed"]  = ggml_new_tensor_3d(ctx, wtype, config.hidden_size, config.num_patches, 1);
     }
 
 public:
-    MMDiT(const String2TensorStorage& tensor_storage_map = {}) {
+    MMDiTConfig config;
+
+    explicit MMDiT(MMDiTConfig config = {})
+        : config(config) {
         // input_size is always None
         // learn_sigma is always False
         // register_length is alwalys 0
@@ -646,64 +729,30 @@ public:
         // pos_embed_offset is not used
         // context_embedder_config is always {'target': 'torch.nn.Linear', 'params': {'in_features': 4096, 'out_features': 1536}}
 
-        for (auto pair : tensor_storage_map) {
-            std::string tensor_name = pair.first;
-            if (tensor_name.find("model.diffusion_model.") == std::string::npos)
-                continue;
-            size_t jb = tensor_name.find("joint_blocks.");
-            if (jb != std::string::npos) {
-                tensor_name     = tensor_name.substr(jb);  // remove prefix
-                int block_depth = atoi(tensor_name.substr(13, tensor_name.find(".", 13)).c_str());
-                if (block_depth + 1 > depth) {
-                    depth = block_depth + 1;
-                }
-                if (tensor_name.find("attn.ln") != std::string::npos) {
-                    if (tensor_name.find(".bias") != std::string::npos) {
-                        qk_norm = "ln";
-                    } else {
-                        qk_norm = "rms";
-                    }
-                }
-                if (tensor_name.find("attn2") != std::string::npos) {
-                    if (block_depth > d_self) {
-                        d_self = block_depth;
-                    }
-                }
-            }
+        blocks["x_embedder"] = std::shared_ptr<GGMLBlock>(new PatchEmbed(config.input_size,
+                                                                         config.patch_size,
+                                                                         config.in_channels,
+                                                                         config.hidden_size,
+                                                                         true));
+        blocks["t_embedder"] = std::shared_ptr<GGMLBlock>(new TimestepEmbedder(config.hidden_size));
+
+        if (config.adm_in_channels != -1) {
+            blocks["y_embedder"] = std::shared_ptr<GGMLBlock>(new VectorEmbedder(config.adm_in_channels, config.hidden_size));
         }
 
-        if (d_self >= 0) {
-            pos_embed_max_size *= 2;
-            num_patchs *= 4;
-        }
+        blocks["context_embedder"] = std::shared_ptr<GGMLBlock>(new Linear(config.context_size, config.context_embedder_out_dim, true, true));
 
-        LOG_INFO("MMDiT layers: %d (including %d MMDiT-x layers)", depth, d_self + 1);
-
-        int64_t default_out_channels = in_channels;
-        hidden_size                  = 64 * depth;
-        context_embedder_out_dim     = 64 * depth;
-        int64_t num_heads            = depth;
-
-        blocks["x_embedder"] = std::shared_ptr<GGMLBlock>(new PatchEmbed(input_size, patch_size, in_channels, hidden_size, true));
-        blocks["t_embedder"] = std::shared_ptr<GGMLBlock>(new TimestepEmbedder(hidden_size));
-
-        if (adm_in_channels != -1) {
-            blocks["y_embedder"] = std::shared_ptr<GGMLBlock>(new VectorEmbedder(adm_in_channels, hidden_size));
-        }
-
-        blocks["context_embedder"] = std::shared_ptr<GGMLBlock>(new Linear(4096, context_embedder_out_dim, true, true));
-
-        for (int i = 0; i < depth; i++) {
-            blocks["joint_blocks." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new JointBlock(hidden_size,
-                                                                                                    num_heads,
-                                                                                                    mlp_ratio,
-                                                                                                    qk_norm,
+        for (int i = 0; i < config.depth; i++) {
+            blocks["joint_blocks." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new JointBlock(config.hidden_size,
+                                                                                                    config.depth,
+                                                                                                    config.mlp_ratio,
+                                                                                                    config.qk_norm,
                                                                                                     true,
-                                                                                                    i == depth - 1,
-                                                                                                    i <= d_self));
+                                                                                                    i == config.depth - 1,
+                                                                                                    i <= config.d_self));
         }
 
-        blocks["final_layer"] = std::shared_ptr<GGMLBlock>(new FinalLayer(hidden_size, patch_size, out_channels));
+        blocks["final_layer"] = std::shared_ptr<GGMLBlock>(new FinalLayer(config.hidden_size, config.patch_size, config.out_channels));
     }
 
     ggml_tensor*
@@ -712,22 +761,22 @@ public:
                       int64_t w) {
         auto pos_embed = params["pos_embed"];
 
-        h = (h + 1) / patch_size;
-        w = (w + 1) / patch_size;
+        h = (h + 1) / config.patch_size;
+        w = (w + 1) / config.patch_size;
 
-        GGML_ASSERT(h <= pos_embed_max_size && h > 0);
-        GGML_ASSERT(w <= pos_embed_max_size && w > 0);
+        GGML_ASSERT(h <= config.pos_embed_max_size && h > 0);
+        GGML_ASSERT(w <= config.pos_embed_max_size && w > 0);
 
-        int64_t top  = (pos_embed_max_size - h) / 2;
-        int64_t left = (pos_embed_max_size - w) / 2;
+        int64_t top  = (config.pos_embed_max_size - h) / 2;
+        int64_t left = (config.pos_embed_max_size - w) / 2;
 
-        auto spatial_pos_embed = ggml_reshape_3d(ctx, pos_embed, hidden_size, pos_embed_max_size, pos_embed_max_size);
+        auto spatial_pos_embed = ggml_reshape_3d(ctx, pos_embed, config.hidden_size, config.pos_embed_max_size, config.pos_embed_max_size);
 
         // spatial_pos_embed = spatial_pos_embed[:, top : top + h, left : left + w, :]
         spatial_pos_embed = ggml_view_3d(ctx,
                                          spatial_pos_embed,
-                                         hidden_size,
-                                         pos_embed_max_size,
+                                         config.hidden_size,
+                                         config.pos_embed_max_size,
                                          h,
                                          spatial_pos_embed->nb[1],
                                          spatial_pos_embed->nb[2],
@@ -735,14 +784,14 @@ public:
         spatial_pos_embed = ggml_cont(ctx, ggml_permute(ctx, spatial_pos_embed, 0, 2, 1, 3));  // [pos_embed_max_size, h, hidden_size]
         spatial_pos_embed = ggml_view_3d(ctx,
                                          spatial_pos_embed,
-                                         hidden_size,
+                                         config.hidden_size,
                                          h,
                                          w,
                                          spatial_pos_embed->nb[1],
                                          spatial_pos_embed->nb[2],
-                                         spatial_pos_embed->nb[2] * left);                     // [w, h, hidden_size]
-        spatial_pos_embed = ggml_cont(ctx, ggml_permute(ctx, spatial_pos_embed, 0, 2, 1, 3));  // [h, w, hidden_size]
-        spatial_pos_embed = ggml_reshape_3d(ctx, spatial_pos_embed, hidden_size, h * w, 1);    // [1, h*w, hidden_size]
+                                         spatial_pos_embed->nb[2] * left);                          // [w, h, hidden_size]
+        spatial_pos_embed = ggml_cont(ctx, ggml_permute(ctx, spatial_pos_embed, 0, 2, 1, 3));       // [h, w, hidden_size]
+        spatial_pos_embed = ggml_reshape_3d(ctx, spatial_pos_embed, config.hidden_size, h * w, 1);  // [1, h*w, hidden_size]
         return spatial_pos_embed;
     }
 
@@ -757,7 +806,7 @@ public:
         // return: [N, N*W, patch_size * patch_size * out_channels]
         auto final_layer = std::dynamic_pointer_cast<FinalLayer>(blocks["final_layer"]);
 
-        for (int i = 0; i < depth; i++) {
+        for (int i = 0; i < config.depth; i++) {
             // skip iteration if i is in skip_layers
             if (skip_layers.size() > 0 && std::find(skip_layers.begin(), skip_layers.end(), i) != skip_layers.end()) {
                 continue;
@@ -800,7 +849,7 @@ public:
         x                = ggml_add(ctx->ggml_ctx, patch_embed, pos_embed);  // [N, H*W, hidden_size]
 
         auto c = t_embedder->forward(ctx, t);  // [N, hidden_size]
-        if (y != nullptr && adm_in_channels != -1) {
+        if (y != nullptr && config.adm_in_channels != -1) {
             auto y_embedder = std::dynamic_pointer_cast<VectorEmbedder>(blocks["y_embedder"]);
 
             y = y_embedder->forward(ctx, y);  // [N, hidden_size]
@@ -820,19 +869,22 @@ public:
 
         x = forward_core_with_concat(ctx, x, c, context, skip_layers);  // (N, H*W, patch_size ** 2 * out_channels)
 
-        x = DiT::unpatchify_and_crop(ctx->ggml_ctx, x, H, W, patch_size, patch_size, /*patch_last*/ false);  // [N, C, H, W]
+        x = DiT::unpatchify_and_crop(ctx->ggml_ctx, x, H, W, config.patch_size, config.patch_size, /*patch_last*/ false);  // [N, C, H, W]
 
         return x;
     }
 };
 struct MMDiTRunner : public DiffusionModelRunner {
+    MMDiTConfig config;
     MMDiT mmdit;
 
     MMDiTRunner(ggml_backend_t backend,
-                ggml_backend_t params_backend,
-                const String2TensorStorage& tensor_storage_map = {},
-                const std::string prefix                       = "")
-        : DiffusionModelRunner(backend, params_backend, prefix), mmdit(tensor_storage_map) {
+                const String2TensorStorage& tensor_storage_map      = {},
+                const std::string prefix                            = "",
+                std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+        : DiffusionModelRunner(backend, prefix, weight_manager),
+          config(MMDiTConfig::detect_from_weights(tensor_storage_map, prefix)),
+          mmdit(config) {
         mmdit.init(params_ctx, tensor_storage_map, prefix);
     }
 
@@ -883,7 +935,7 @@ struct MMDiTRunner : public DiffusionModelRunner {
             return build_graph(x, timesteps, context, y, skip_layers);
         };
 
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
     }
 
     sd::Tensor<float> compute(int n_threads,
@@ -949,28 +1001,25 @@ struct MMDiTRunner : public DiffusionModelRunner {
         // ggml_backend_t backend    = ggml_backend_cuda_init(0);
         ggml_backend_t backend             = sd_backend_cpu_init();
         ggml_type model_data_type          = GGML_TYPE_F16;
-        std::shared_ptr<MMDiTRunner> mmdit = std::make_shared<MMDiTRunner>(backend, backend);
+        auto model_manager                 = std::make_shared<ModelManager>();
+        std::shared_ptr<MMDiTRunner> mmdit = std::make_shared<MMDiTRunner>(backend, String2TensorStorage{}, "", model_manager);
         {
             LOG_INFO("loading from '%s'", file_path.c_str());
 
-            if (!mmdit->alloc_params_buffer()) {
-                LOG_ERROR("mmdit embeds buffer allocation failed");
-                return;
-            }
-
-            std::map<std::string, ggml_tensor*> tensors;
-            mmdit->get_param_tensors(tensors, "model.diffusion_model");
-
-            ModelLoader model_loader;
+            ModelLoader& model_loader = model_manager->loader();
             if (!model_loader.init_from_file_and_convert_name(file_path)) {
                 LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
                 return;
             }
 
-            bool success = model_loader.load_tensors(tensors);
-
-            if (!success) {
-                LOG_ERROR("load tensors from model loader failed");
+            if (!model_manager->register_runner_params("MMDiT test",
+                                                       *mmdit,
+                                                       "model.diffusion_model",
+                                                       ModelManager::ResidencyMode::ParamBackend,
+                                                       backend,
+                                                       backend) ||
+                !model_manager->validate_registered_tensors()) {
+                LOG_ERROR("register mmdit tensors with model manager failed");
                 return;
             }
 
@@ -980,4 +1029,4 @@ struct MMDiTRunner : public DiffusionModelRunner {
     }
 };
 
-#endif
+#endif  // __SD_MODEL_DIFFUSION_MMDIT_HPP__

src/model/diffusion/model.hpp

@@ -1,12 +1,13 @@
-#ifndef __DIFFUSION_MODEL_H__
-#define __DIFFUSION_MODEL_H__
+#ifndef __SD_MODEL_DIFFUSION_MODEL_HPP__
+#define __SD_MODEL_DIFFUSION_MODEL_HPP__
 
 #include <string>
 #include <utility>
 #include <variant>
 
-#include "ggml_extend.hpp"
-#include "tensor_ggml.hpp"
+#include "core/ggml_extend.hpp"
+#include "core/tensor_ggml.hpp"
+#include "model_manager.h"
 
 struct UNetDiffusionExtra {
     int num_video_frames                           = -1;
@@ -88,9 +89,9 @@ protected:
 
 public:
     DiffusionModelRunner(ggml_backend_t backend,
-                         ggml_backend_t params_backend,
-                         const std::string& prefix)
-        : GGMLRunner(backend, params_backend),
+                         const std::string& prefix,
+                         std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+        : GGMLRunner(backend, weight_manager),
           prefix(prefix) {}
 
     virtual sd::Tensor<float> compute(int n_threads,
@@ -104,4 +105,4 @@ public:
                                    const std::string& prefix) = 0;
 };
 
-#endif
+#endif  // __SD_MODEL_DIFFUSION_MODEL_HPP__

src/model/diffusion/pid.hpp

@@ -1,5 +1,5 @@
-#ifndef __SD_PID_HPP__
-#define __SD_PID_HPP__
+#ifndef __SD_MODEL_DIFFUSION_PID_HPP__
+#define __SD_MODEL_DIFFUSION_PID_HPP__
 
 #include <cmath>
 #include <cstdlib>
@@ -7,16 +7,16 @@
 #include <string>
 #include <vector>
 
-#include "common_dit.hpp"
-#include "ggml_extend.hpp"
-#include "mmdit.hpp"
-#include "rope.hpp"
+#include "core/ggml_extend.hpp"
+#include "model/common/rope.hpp"
+#include "model/diffusion/dit.hpp"
+#include "model/diffusion/mmdit.hpp"
 
 namespace Pid {
     constexpr int PID_GRAPH_SIZE = 196608;
     constexpr float PID_PI       = 3.14159265358979323846f;
 
-    struct PixelDiTParams {
+    struct PixelDiTConfig {
         int64_t in_channels            = 3;
         int64_t hidden_size            = 1536;
         int64_t num_groups             = 24;
@@ -38,6 +38,45 @@ namespace Pid {
         int64_t lq_latent_down_factor  = 8;
         int64_t rope_ref_grid_h        = 64;
         int64_t rope_ref_grid_w        = 64;
+
+        static PixelDiTConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
+            PixelDiTConfig config;
+            for (const auto& [name, tensor_storage] : tensor_storage_map) {
+                if (!starts_with(name, prefix)) {
+                    continue;
+                }
+                size_t pos = name.find("patch_blocks.");
+                if (pos != std::string::npos) {
+                    auto items = split_string(name.substr(pos), '.');
+                    if (items.size() > 1) {
+                        int block_index    = atoi(items[1].c_str());
+                        config.patch_depth = std::max<int64_t>(config.patch_depth, block_index + 1);
+                    }
+                }
+                pos = name.find("pixel_blocks.");
+                if (pos != std::string::npos) {
+                    auto items = split_string(name.substr(pos), '.');
+                    if (items.size() > 1) {
+                        int block_index    = atoi(items[1].c_str());
+                        config.pixel_depth = std::max<int64_t>(config.pixel_depth, block_index + 1);
+                    }
+                }
+                if (name.find("lq_proj.latent_proj.0.weight") != std::string::npos) {
+                    config.lq_latent_channels    = tensor_storage.ne[2];
+                    config.lq_latent_down_factor = config.lq_latent_channels >= 64 ? 16 : 8;
+                }
+                if (name.find("patch_blocks.0.mlp_x.w1.weight") != std::string::npos) {
+                    config.patch_mlp_hidden_dim = tensor_storage.ne[1];
+                }
+            }
+            LOG_DEBUG("pid: patch_depth = %" PRId64 ", pixel_depth = %" PRId64 ", patch_mlp_hidden_dim = %" PRId64 ", lq_latent_channels = %" PRId64 ", lq_latent_down_factor = %" PRId64,
+                      config.patch_depth,
+                      config.pixel_depth,
+                      config.patch_mlp_hidden_dim,
+                      config.lq_latent_channels,
+                      config.lq_latent_down_factor);
+            return config;
+        }
     };
 
     inline std::vector<float> make_rope_1d(int length,
@@ -466,29 +505,29 @@ namespace Pid {
     };
 
     struct LQProjection2D : public GGMLBlock {
-        PixelDiTParams params_cfg;
+        PixelDiTConfig config;
 
-        LQProjection2D(const PixelDiTParams& params_cfg)
-            : params_cfg(params_cfg) {
-            blocks["latent_proj.0"] = std::make_shared<Conv2d>(params_cfg.lq_latent_channels, params_cfg.lq_hidden_dim, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, std::pair<int, int>{1, 1});
-            blocks["latent_proj.2"] = std::make_shared<Conv2d>(params_cfg.lq_hidden_dim, params_cfg.lq_hidden_dim, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, std::pair<int, int>{1, 1});
-            for (int i = 0; i < params_cfg.lq_num_res_blocks; ++i) {
-                blocks["latent_proj." + std::to_string(3 + i)] = std::make_shared<PiDResBlock>(params_cfg.lq_hidden_dim);
+        LQProjection2D(const PixelDiTConfig& config)
+            : config(config) {
+            blocks["latent_proj.0"] = std::make_shared<Conv2d>(config.lq_latent_channels, config.lq_hidden_dim, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, std::pair<int, int>{1, 1});
+            blocks["latent_proj.2"] = std::make_shared<Conv2d>(config.lq_hidden_dim, config.lq_hidden_dim, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, std::pair<int, int>{1, 1});
+            for (int i = 0; i < config.lq_num_res_blocks; ++i) {
+                blocks["latent_proj." + std::to_string(3 + i)] = std::make_shared<PiDResBlock>(config.lq_hidden_dim);
             }
 
-            int num_outputs = static_cast<int>((params_cfg.patch_depth + params_cfg.lq_interval - 1) / params_cfg.lq_interval);
+            int num_outputs = static_cast<int>((config.patch_depth + config.lq_interval - 1) / config.lq_interval);
             for (int i = 0; i < num_outputs; ++i) {
-                blocks["output_heads." + std::to_string(i)] = std::make_shared<Linear>(params_cfg.lq_hidden_dim, params_cfg.hidden_size, true);
-                blocks["gate_modules." + std::to_string(i)] = std::make_shared<SigmaAwareGate>(params_cfg.hidden_size);
+                blocks["output_heads." + std::to_string(i)] = std::make_shared<Linear>(config.lq_hidden_dim, config.hidden_size, true);
+                blocks["gate_modules." + std::to_string(i)] = std::make_shared<SigmaAwareGate>(config.hidden_size);
             }
         }
 
         bool is_gate_active(int block_idx) const {
-            return block_idx % params_cfg.lq_interval == 0;
+            return block_idx % config.lq_interval == 0;
         }
 
         int get_output_index(int block_idx) const {
-            return block_idx / static_cast<int>(params_cfg.lq_interval);
+            return block_idx / static_cast<int>(config.lq_interval);
         }
 
         ggml_tensor* gate(GGMLRunnerContext* ctx,
@@ -506,8 +545,8 @@ namespace Pid {
                                           int64_t target_pW) {
             auto conv0             = std::dynamic_pointer_cast<Conv2d>(blocks["latent_proj.0"]);
             auto conv2             = std::dynamic_pointer_cast<Conv2d>(blocks["latent_proj.2"]);
-            float z_to_patch_ratio = static_cast<float>(params_cfg.lq_sr_scale * params_cfg.lq_latent_down_factor) /
-                                     static_cast<float>(params_cfg.patch_size);
+            float z_to_patch_ratio = static_cast<float>(config.lq_sr_scale * config.lq_latent_down_factor) /
+                                     static_cast<float>(config.patch_size);
             GGML_ASSERT(z_to_patch_ratio >= 1.0f);
             if (lq_latent->ne[0] != target_pW || lq_latent->ne[1] != target_pH) {
                 lq_latent = ggml_interpolate(ctx->ggml_ctx,
@@ -522,7 +561,7 @@ namespace Pid {
             auto feat = conv0->forward(ctx, lq_latent);
             feat      = ggml_silu_inplace(ctx->ggml_ctx, feat);
             feat      = conv2->forward(ctx, feat);
-            for (int i = 0; i < params_cfg.lq_num_res_blocks; ++i) {
+            for (int i = 0; i < config.lq_num_res_blocks; ++i) {
                 auto block = std::dynamic_pointer_cast<PiDResBlock>(blocks["latent_proj." + std::to_string(3 + i)]);
                 feat       = block->forward(ctx, feat);
             }
@@ -533,7 +572,7 @@ namespace Pid {
             auto tokens = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, feat, 2, 0, 1, 3));
             tokens      = ggml_reshape_3d(ctx->ggml_ctx, tokens, C, L, B);
 
-            int num_outputs = static_cast<int>((params_cfg.patch_depth + params_cfg.lq_interval - 1) / params_cfg.lq_interval);
+            int num_outputs = static_cast<int>((config.patch_depth + config.lq_interval - 1) / config.lq_interval);
             std::vector<ggml_tensor*> outputs;
             outputs.reserve(num_outputs);
             for (int i = 0; i < num_outputs; ++i) {
@@ -545,34 +584,34 @@ namespace Pid {
     };
 
     struct PixelDiT : public GGMLBlock {
-        PixelDiTParams params_cfg;
+        PixelDiTConfig config;
 
         PixelDiT() = default;
 
-        PixelDiT(const PixelDiTParams& params_cfg)
-            : params_cfg(params_cfg) {
-            blocks["pixel_embedder"] = std::make_shared<PixelTokenEmbedder>(params_cfg.in_channels, params_cfg.pixel_hidden_size);
-            blocks["s_embedder"]     = std::make_shared<PatchTokenEmbedder>(params_cfg.in_channels * params_cfg.patch_size * params_cfg.patch_size, params_cfg.hidden_size, false, true);
-            blocks["t_embedder"]     = std::make_shared<PixelDiTTimestepEmbedder>(params_cfg.hidden_size);
-            blocks["y_embedder"]     = std::make_shared<PatchTokenEmbedder>(params_cfg.txt_embed_dim, params_cfg.hidden_size, true, true);
-            for (int i = 0; i < params_cfg.patch_depth; ++i) {
-                blocks["patch_blocks." + std::to_string(i)] = std::make_shared<MMDiTBlockT2I>(params_cfg.hidden_size, params_cfg.num_groups, params_cfg.patch_mlp_hidden_dim);
+        PixelDiT(const PixelDiTConfig& config)
+            : config(config) {
+            blocks["pixel_embedder"] = std::make_shared<PixelTokenEmbedder>(config.in_channels, config.pixel_hidden_size);
+            blocks["s_embedder"]     = std::make_shared<PatchTokenEmbedder>(config.in_channels * config.patch_size * config.patch_size, config.hidden_size, false, true);
+            blocks["t_embedder"]     = std::make_shared<PixelDiTTimestepEmbedder>(config.hidden_size);
+            blocks["y_embedder"]     = std::make_shared<PatchTokenEmbedder>(config.txt_embed_dim, config.hidden_size, true, true);
+            for (int i = 0; i < config.patch_depth; ++i) {
+                blocks["patch_blocks." + std::to_string(i)] = std::make_shared<MMDiTBlockT2I>(config.hidden_size, config.num_groups, config.patch_mlp_hidden_dim);
             }
-            for (int i = 0; i < params_cfg.pixel_depth; ++i) {
-                blocks["pixel_blocks." + std::to_string(i)] = std::make_shared<PiTBlock>(params_cfg.pixel_hidden_size,
-                                                                                         params_cfg.hidden_size,
-                                                                                         params_cfg.patch_size,
-                                                                                         params_cfg.pixel_attn_hidden_size,
-                                                                                         params_cfg.pixel_num_groups);
+            for (int i = 0; i < config.pixel_depth; ++i) {
+                blocks["pixel_blocks." + std::to_string(i)] = std::make_shared<PiTBlock>(config.pixel_hidden_size,
+                                                                                         config.hidden_size,
+                                                                                         config.patch_size,
+                                                                                         config.pixel_attn_hidden_size,
+                                                                                         config.pixel_num_groups);
             }
-            blocks["final_layer"] = std::make_shared<FinalLayer>(params_cfg.pixel_hidden_size, params_cfg.in_channels);
-            blocks["lq_proj"]     = std::make_shared<LQProjection2D>(params_cfg);
+            blocks["final_layer"] = std::make_shared<FinalLayer>(config.pixel_hidden_size, config.in_channels);
+            blocks["lq_proj"]     = std::make_shared<LQProjection2D>(config);
         }
 
         void init_params(ggml_context* ctx,
                          const String2TensorStorage& tensor_storage_map = {},
                          std::string prefix                             = "") override {
-            params["y_pos_embedding"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, params_cfg.hidden_size, params_cfg.txt_max_length, 1);
+            params["y_pos_embedding"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, config.hidden_size, config.txt_max_length, 1);
         }
 
         ggml_tensor* forward(GGMLRunnerContext* ctx,
@@ -594,21 +633,21 @@ namespace Pid {
 
             int64_t W_orig = x->ne[0];
             int64_t H_orig = x->ne[1];
-            x              = DiT::pad_to_patch_size(ctx, x, static_cast<int>(params_cfg.patch_size), static_cast<int>(params_cfg.patch_size));
+            x              = DiT::pad_to_patch_size(ctx, x, static_cast<int>(config.patch_size), static_cast<int>(config.patch_size));
             int64_t W      = x->ne[0];
             int64_t H      = x->ne[1];
             int64_t B      = x->ne[3];
-            int64_t Hs     = H / params_cfg.patch_size;
-            int64_t Ws     = W / params_cfg.patch_size;
+            int64_t Hs     = H / config.patch_size;
+            int64_t Ws     = W / config.patch_size;
             int64_t L      = Hs * Ws;
-            int64_t P2     = params_cfg.patch_size * params_cfg.patch_size;
+            int64_t P2     = config.patch_size * config.patch_size;
 
-            auto x_patches = DiT::patchify(ctx->ggml_ctx, x, static_cast<int>(params_cfg.patch_size), static_cast<int>(params_cfg.patch_size), true);
+            auto x_patches = DiT::patchify(ctx->ggml_ctx, x, static_cast<int>(config.patch_size), static_cast<int>(config.patch_size), true);
             auto t_emb     = t_embedder->forward(ctx, timesteps);
             auto condition = ggml_silu(ctx->ggml_ctx, t_emb);
 
             GGML_ASSERT(context != nullptr);
-            int64_t Ltxt = std::min<int64_t>(context->ne[1], params_cfg.txt_max_length);
+            int64_t Ltxt = std::min<int64_t>(context->ne[1], config.txt_max_length);
             auto y       = ggml_ext_slice(ctx->ggml_ctx, context, 1, 0, Ltxt);
             auto y_emb   = y_embedder->forward(ctx, y);
             auto y_pos   = ggml_ext_slice(ctx->ggml_ctx, params["y_pos_embedding"], 1, 0, Ltxt);
@@ -618,7 +657,7 @@ namespace Pid {
 
             auto s = s_embedder->forward(ctx, x_patches);
 
-            for (int i = 0; i < params_cfg.patch_depth; ++i) {
+            for (int i = 0; i < config.patch_depth; ++i) {
                 if (lq_proj->is_gate_active(i)) {
                     int out_idx = lq_proj->get_output_index(i);
                     if (out_idx < static_cast<int>(lq_features.size())) {
@@ -639,22 +678,22 @@ namespace Pid {
             }
             s = ggml_silu(ctx->ggml_ctx, ggml_add(ctx->ggml_ctx, s, t_emb));
 
-            auto s_cond = ggml_reshape_2d(ctx->ggml_ctx, s, params_cfg.hidden_size, L * B);
-            auto pixels = pixel_embedder->forward(ctx, x, params_cfg.patch_size, pixel_pos_full);
-            for (int i = 0; i < params_cfg.pixel_depth; ++i) {
+            auto s_cond = ggml_reshape_2d(ctx->ggml_ctx, s, config.hidden_size, L * B);
+            auto pixels = pixel_embedder->forward(ctx, x, config.patch_size, pixel_pos_full);
+            for (int i = 0; i < config.pixel_depth; ++i) {
                 auto block = std::dynamic_pointer_cast<PiTBlock>(blocks["pixel_blocks." + std::to_string(i)]);
                 pixels     = block->forward(ctx, pixels, s_cond, H, W, pixel_pos_comp);
                 sd::ggml_graph_cut::mark_graph_cut(pixels, "pid.pixel_blocks." + std::to_string(i), "pixels");
             }
 
             pixels   = final_layer->forward(ctx, pixels);
-            pixels   = ggml_reshape_3d(ctx->ggml_ctx, pixels, params_cfg.in_channels * P2, L, B);
+            pixels   = ggml_reshape_3d(ctx->ggml_ctx, pixels, config.in_channels * P2, L, B);
             auto out = DiT::unpatchify(ctx->ggml_ctx,
                                        pixels,
                                        Hs,
                                        Ws,
-                                       static_cast<int>(params_cfg.patch_size),
-                                       static_cast<int>(params_cfg.patch_size),
+                                       static_cast<int>(config.patch_size),
+                                       static_cast<int>(config.patch_size),
                                        false);
             out      = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, H_orig);
             out      = ggml_ext_slice(ctx->ggml_ctx, out, 0, 0, W_orig);
@@ -663,7 +702,7 @@ namespace Pid {
     };
 
     struct PiDRunner : public DiffusionModelRunner {
-        PixelDiTParams params_cfg;
+        PixelDiTConfig config;
         PixelDiT model;
         std::vector<float> pos_img_vec;
         std::vector<float> pos_txt_vec;
@@ -671,46 +710,12 @@ namespace Pid {
         std::vector<float> pixel_pos_comp_vec;
 
         PiDRunner(ggml_backend_t backend,
-                  ggml_backend_t params_backend,
                   const String2TensorStorage& tensor_storage_map,
-                  const std::string prefix = "model.diffusion_model")
-            : DiffusionModelRunner(backend, params_backend, prefix) {
-            for (const auto& pair : tensor_storage_map) {
-                const std::string& tensor_name = pair.first;
-                if (tensor_name.find(prefix) == std::string::npos) {
-                    continue;
-                }
-                size_t pos = tensor_name.find("patch_blocks.");
-                if (pos != std::string::npos) {
-                    auto items = split_string(tensor_name.substr(pos), '.');
-                    if (items.size() > 1) {
-                        int block_index        = atoi(items[1].c_str());
-                        params_cfg.patch_depth = std::max<int64_t>(params_cfg.patch_depth, block_index + 1);
-                    }
-                }
-                pos = tensor_name.find("pixel_blocks.");
-                if (pos != std::string::npos) {
-                    auto items = split_string(tensor_name.substr(pos), '.');
-                    if (items.size() > 1) {
-                        int block_index        = atoi(items[1].c_str());
-                        params_cfg.pixel_depth = std::max<int64_t>(params_cfg.pixel_depth, block_index + 1);
-                    }
-                }
-                if (tensor_name.find("lq_proj.latent_proj.0.weight") != std::string::npos) {
-                    params_cfg.lq_latent_channels    = pair.second.ne[2];
-                    params_cfg.lq_latent_down_factor = params_cfg.lq_latent_channels >= 64 ? 16 : 8;
-                }
-                if (tensor_name.find("patch_blocks.0.mlp_x.w1.weight") != std::string::npos) {
-                    params_cfg.patch_mlp_hidden_dim = pair.second.ne[1];
-                }
-            }
-            LOG_INFO("PiD params: patch_depth=%" PRId64 ", pixel_depth=%" PRId64 ", patch_mlp_hidden_dim=%" PRId64 ", lq_latent_channels=%" PRId64 ", lq_latent_down_factor=%" PRId64,
-                     params_cfg.patch_depth,
-                     params_cfg.pixel_depth,
-                     params_cfg.patch_mlp_hidden_dim,
-                     params_cfg.lq_latent_channels,
-                     params_cfg.lq_latent_down_factor);
-            model = PixelDiT(params_cfg);
+                  const std::string prefix                            = "model.diffusion_model",
+                  std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : DiffusionModelRunner(backend, prefix, weight_manager),
+              config(PixelDiTConfig::detect_from_weights(tensor_storage_map, prefix)) {
+            model = PixelDiT(config);
             model.init(params_ctx, tensor_storage_map, prefix);
         }
 
@@ -737,60 +742,60 @@ namespace Pid {
             int64_t W  = x->ne[0];
             int64_t H  = x->ne[1];
             int64_t B  = x->ne[3];
-            int64_t Wp = align_up(static_cast<int>(W), static_cast<int>(params_cfg.patch_size));
-            int64_t Hp = align_up(static_cast<int>(H), static_cast<int>(params_cfg.patch_size));
-            int64_t Hs = Hp / params_cfg.patch_size;
-            int64_t Ws = Wp / params_cfg.patch_size;
+            int64_t Wp = align_up(static_cast<int>(W), static_cast<int>(config.patch_size));
+            int64_t Hp = align_up(static_cast<int>(H), static_cast<int>(config.patch_size));
+            int64_t Hs = Hp / config.patch_size;
+            int64_t Ws = Wp / config.patch_size;
 
             pos_img_vec  = make_rope_2d(static_cast<int>(Hs),
                                         static_cast<int>(Ws),
-                                        static_cast<int>(params_cfg.hidden_size / params_cfg.num_groups),
+                                        static_cast<int>(config.hidden_size / config.num_groups),
                                         10000.f,
                                         16.f,
-                                        static_cast<int>(params_cfg.rope_ref_grid_h),
-                                        static_cast<int>(params_cfg.rope_ref_grid_w));
+                                        static_cast<int>(config.rope_ref_grid_h),
+                                        static_cast<int>(config.rope_ref_grid_w));
             auto pos_img = ggml_new_tensor_4d(compute_ctx,
                                               GGML_TYPE_F32,
                                               2,
                                               2,
-                                              params_cfg.hidden_size / params_cfg.num_groups / 2,
+                                              config.hidden_size / config.num_groups / 2,
                                               Hs * Ws);
             set_backend_tensor_data(pos_img, pos_img_vec.data());
 
-            int64_t Ltxt = std::min<int64_t>(context->ne[1], params_cfg.txt_max_length);
+            int64_t Ltxt = std::min<int64_t>(context->ne[1], config.txt_max_length);
             pos_txt_vec  = make_rope_1d(static_cast<int>(Ltxt),
-                                        static_cast<int>(params_cfg.hidden_size / params_cfg.num_groups),
-                                        params_cfg.text_rope_theta);
+                                        static_cast<int>(config.hidden_size / config.num_groups),
+                                        config.text_rope_theta);
             auto pos_txt = ggml_new_tensor_4d(compute_ctx,
                                               GGML_TYPE_F32,
                                               2,
                                               2,
-                                              params_cfg.hidden_size / params_cfg.num_groups / 2,
+                                              config.hidden_size / config.num_groups / 2,
                                               Ltxt);
             set_backend_tensor_data(pos_txt, pos_txt_vec.data());
 
             pixel_pos_vec  = make_pixel_abs_pos(static_cast<int>(Hp),
                                                 static_cast<int>(Wp),
-                                                static_cast<int>(params_cfg.pixel_hidden_size));
+                                                static_cast<int>(config.pixel_hidden_size));
             auto pixel_pos = ggml_new_tensor_3d(compute_ctx,
                                                 GGML_TYPE_F32,
-                                                params_cfg.pixel_hidden_size,
+                                                config.pixel_hidden_size,
                                                 Wp * Hp,
                                                 1);
             set_backend_tensor_data(pixel_pos, pixel_pos_vec.data());
 
             pixel_pos_comp_vec  = make_rope_2d(static_cast<int>(Hs),
                                                static_cast<int>(Ws),
-                                               static_cast<int>(params_cfg.pixel_attn_hidden_size / params_cfg.pixel_num_groups),
+                                               static_cast<int>(config.pixel_attn_hidden_size / config.pixel_num_groups),
                                                10000.f,
                                                16.f,
-                                               static_cast<int>(params_cfg.rope_ref_grid_h),
-                                               static_cast<int>(params_cfg.rope_ref_grid_w));
+                                               static_cast<int>(config.rope_ref_grid_h),
+                                               static_cast<int>(config.rope_ref_grid_w));
             auto pixel_pos_comp = ggml_new_tensor_4d(compute_ctx,
                                                      GGML_TYPE_F32,
                                                      2,
                                                      2,
-                                                     params_cfg.pixel_attn_hidden_size / params_cfg.pixel_num_groups / 2,
+                                                     config.pixel_attn_hidden_size / config.pixel_num_groups / 2,
                                                      Hs * Ws);
             set_backend_tensor_data(pixel_pos_comp, pixel_pos_comp_vec.data());
 
@@ -818,7 +823,7 @@ namespace Pid {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(x, timesteps, context, lq_latent, degrade_sigma);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,
@@ -839,4 +844,4 @@ namespace Pid {
     };
 }  // namespace Pid
 
-#endif  // __SD_PID_HPP__
+#endif  // __SD_MODEL_DIFFUSION_PID_HPP__

src/model/diffusion/qwen_image.hpp

@@ -1,15 +1,58 @@
-#ifndef __QWEN_IMAGE_HPP__
-#define __QWEN_IMAGE_HPP__
+#ifndef __SD_MODEL_DIFFUSION_QWEN_IMAGE_HPP__
+#define __SD_MODEL_DIFFUSION_QWEN_IMAGE_HPP__
 
 #include <memory>
 
-#include "common_block.hpp"
-#include "diffusion_model.hpp"
-#include "flux.hpp"
+#include "model/common/block.hpp"
+#include "model/diffusion/flux.hpp"
+#include "model/diffusion/model.hpp"
+#include "model_loader.h"
 
 namespace Qwen {
     constexpr int QWEN_IMAGE_GRAPH_SIZE = 20480;
 
+    struct QwenImageConfig {
+        int patch_size              = 2;
+        int64_t in_channels         = 64;
+        int64_t out_channels        = 16;
+        int num_layers              = 60;
+        int64_t attention_head_dim  = 128;
+        int64_t num_attention_heads = 24;
+        int64_t joint_attention_dim = 3584;
+        int theta                   = 10000;
+        std::vector<int> axes_dim   = {16, 56, 56};
+        int axes_dim_sum            = 128;
+        bool zero_cond_t            = false;
+
+        static QwenImageConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
+            QwenImageConfig config;
+            config.num_layers = 0;
+            for (const auto& [name, _] : tensor_storage_map) {
+                if (!starts_with(name, prefix)) {
+                    continue;
+                }
+                if (name.find("__index_timestep_zero__") != std::string::npos) {
+                    config.zero_cond_t = true;
+                }
+                size_t pos = name.find("transformer_blocks.");
+                if (pos == std::string::npos) {
+                    continue;
+                }
+                auto items = split_string(name.substr(pos), '.');
+                if (items.size() > 1) {
+                    int block_index = atoi(items[1].c_str());
+                    if (block_index + 1 > config.num_layers) {
+                        config.num_layers = block_index + 1;
+                    }
+                }
+            }
+            LOG_DEBUG("qwen_image: num_layers = %d, zero_cond_t = %s",
+                      config.num_layers,
+                      config.zero_cond_t ? "true" : "false");
+            return config;
+        }
+    };
+
     struct TimestepEmbedding : public GGMLBlock {
     public:
         TimestepEmbedding(int64_t in_channels,
@@ -350,46 +393,32 @@ namespace Qwen {
         }
     };
 
-    struct QwenImageParams {
-        int patch_size              = 2;
-        int64_t in_channels         = 64;
-        int64_t out_channels        = 16;
-        int num_layers              = 60;
-        int64_t attention_head_dim  = 128;
-        int64_t num_attention_heads = 24;
-        int64_t joint_attention_dim = 3584;
-        int theta                   = 10000;
-        std::vector<int> axes_dim   = {16, 56, 56};
-        int axes_dim_sum            = 128;
-        bool zero_cond_t            = false;
-    };
-
     class QwenImageModel : public GGMLBlock {
     protected:
-        QwenImageParams params;
+        QwenImageConfig config;
 
     public:
         QwenImageModel() {}
-        QwenImageModel(QwenImageParams params)
-            : params(params) {
-            int64_t inner_dim         = params.num_attention_heads * params.attention_head_dim;
+        QwenImageModel(QwenImageConfig config)
+            : config(config) {
+            int64_t inner_dim         = config.num_attention_heads * config.attention_head_dim;
             blocks["time_text_embed"] = std::shared_ptr<GGMLBlock>(new QwenTimestepProjEmbeddings(inner_dim));
-            blocks["txt_norm"]        = std::shared_ptr<GGMLBlock>(new RMSNorm(params.joint_attention_dim, 1e-6f));
-            blocks["img_in"]          = std::shared_ptr<GGMLBlock>(new Linear(params.in_channels, inner_dim));
-            blocks["txt_in"]          = std::shared_ptr<GGMLBlock>(new Linear(params.joint_attention_dim, inner_dim));
+            blocks["txt_norm"]        = std::shared_ptr<GGMLBlock>(new RMSNorm(config.joint_attention_dim, 1e-6f));
+            blocks["img_in"]          = std::shared_ptr<GGMLBlock>(new Linear(config.in_channels, inner_dim));
+            blocks["txt_in"]          = std::shared_ptr<GGMLBlock>(new Linear(config.joint_attention_dim, inner_dim));
 
             // blocks
-            for (int i = 0; i < params.num_layers; i++) {
+            for (int i = 0; i < config.num_layers; i++) {
                 auto block                                        = std::shared_ptr<GGMLBlock>(new QwenImageTransformerBlock(inner_dim,
-                                                                                                                             params.num_attention_heads,
-                                                                                                                             params.attention_head_dim,
+                                                                                                                             config.num_attention_heads,
+                                                                                                                             config.attention_head_dim,
                                                                                                                              1e-6f,
-                                                                                                                             params.zero_cond_t));
+                                                                                                                             config.zero_cond_t));
                 blocks["transformer_blocks." + std::to_string(i)] = block;
             }
 
             blocks["norm_out"] = std::shared_ptr<GGMLBlock>(new AdaLayerNormContinuous(inner_dim, inner_dim, false, 1e-6f));
-            blocks["proj_out"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, params.patch_size * params.patch_size * params.out_channels));
+            blocks["proj_out"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, config.patch_size * config.patch_size * config.out_channels));
         }
 
         ggml_tensor* forward_orig(GGMLRunnerContext* ctx,
@@ -406,7 +435,7 @@ namespace Qwen {
             auto proj_out        = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
 
             auto t_emb = time_text_embed->forward(ctx, timestep);
-            if (params.zero_cond_t) {
+            if (config.zero_cond_t) {
                 auto t_emb_0 = time_text_embed->forward(ctx, ggml_ext_zeros_like(ctx->ggml_ctx, timestep));
                 t_emb        = ggml_concat(ctx->ggml_ctx, t_emb, t_emb_0, 1);
             }
@@ -417,7 +446,7 @@ namespace Qwen {
             sd::ggml_graph_cut::mark_graph_cut(txt, "qwen_image.prelude", "txt");
             // sd::ggml_graph_cut::mark_graph_cut(t_emb, "qwen_image.prelude", "t_emb");
 
-            for (int i = 0; i < params.num_layers; i++) {
+            for (int i = 0; i < config.num_layers; i++) {
                 auto block = std::dynamic_pointer_cast<QwenImageTransformerBlock>(blocks["transformer_blocks." + std::to_string(i)]);
 
                 auto result = block->forward(ctx, img, txt, t_emb, pe, modulate_index);
@@ -427,7 +456,7 @@ namespace Qwen {
                 sd::ggml_graph_cut::mark_graph_cut(txt, "qwen_image.transformer_blocks." + std::to_string(i), "txt");
             }
 
-            if (params.zero_cond_t) {
+            if (config.zero_cond_t) {
                 t_emb = ggml_ext_chunk(ctx->ggml_ctx, t_emb, 2, 1)[0];
             }
 
@@ -456,12 +485,12 @@ namespace Qwen {
             int64_t C = x->ne[2];
             int64_t N = x->ne[3];
 
-            auto img           = DiT::pad_and_patchify(ctx, x, params.patch_size, params.patch_size);
+            auto img           = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size);
             int64_t img_tokens = img->ne[1];
 
             if (ref_latents.size() > 0) {
                 for (ggml_tensor* ref : ref_latents) {
-                    ref = DiT::pad_and_patchify(ctx, ref, params.patch_size, params.patch_size);
+                    ref = DiT::pad_and_patchify(ctx, ref, config.patch_size, config.patch_size);
                     img = ggml_concat(ctx->ggml_ctx, img, ref, 1);
                 }
             }
@@ -474,7 +503,7 @@ namespace Qwen {
                 out = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, out, 0, 2, 1, 3));  // [N, h*w, C * patch_size * patch_size]
             }
 
-            out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, params.patch_size, params.patch_size);  // [N, C, H, W]
+            out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, config.patch_size, config.patch_size);  // [N, C, H, W]
 
             return out;
         }
@@ -482,46 +511,22 @@ namespace Qwen {
 
     struct QwenImageRunner : public DiffusionModelRunner {
     public:
-        QwenImageParams qwen_image_params;
+        QwenImageConfig config;
         QwenImageModel qwen_image;
         std::vector<float> pe_vec;
         std::vector<float> modulate_index_vec;
         SDVersion version;
 
         QwenImageRunner(ggml_backend_t backend,
-                        ggml_backend_t params_backend,
-                        const String2TensorStorage& tensor_storage_map = {},
-                        const std::string prefix                       = "",
-                        SDVersion version                              = VERSION_QWEN_IMAGE,
-                        bool zero_cond_t                               = false)
-            : DiffusionModelRunner(backend, params_backend, prefix) {
-            qwen_image_params.num_layers  = 0;
-            qwen_image_params.zero_cond_t = zero_cond_t;
-            for (auto pair : tensor_storage_map) {
-                std::string tensor_name = pair.first;
-                if (tensor_name.find(prefix) == std::string::npos)
-                    continue;
-                if (tensor_name.find("__index_timestep_zero__") != std::string::npos) {
-                    qwen_image_params.zero_cond_t = true;
-                }
-                size_t pos = tensor_name.find("transformer_blocks.");
-                if (pos != std::string::npos) {
-                    tensor_name = tensor_name.substr(pos);  // remove prefix
-                    auto items  = split_string(tensor_name, '.');
-                    if (items.size() > 1) {
-                        int block_index = atoi(items[1].c_str());
-                        if (block_index + 1 > qwen_image_params.num_layers) {
-                            qwen_image_params.num_layers = block_index + 1;
-                        }
-                    }
-                    continue;
-                }
-            }
-            LOG_INFO("qwen_image_params.num_layers: %ld", qwen_image_params.num_layers);
-            if (qwen_image_params.zero_cond_t) {
-                LOG_INFO("use zero_cond_t");
-            }
-            qwen_image = QwenImageModel(qwen_image_params);
+                        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)
+            : DiffusionModelRunner(backend, prefix, weight_manager),
+              config(QwenImageConfig::detect_from_weights(tensor_storage_map, prefix)) {
+            config.zero_cond_t = config.zero_cond_t || zero_cond_t;
+            qwen_image         = QwenImageModel(config);
             qwen_image.init(params_ctx, tensor_storage_map, prefix);
         }
 
@@ -552,36 +557,36 @@ namespace Qwen {
 
             pe_vec      = Rope::gen_qwen_image_pe(static_cast<int>(x->ne[1]),
                                                   static_cast<int>(x->ne[0]),
-                                                  qwen_image_params.patch_size,
+                                                  config.patch_size,
                                                   static_cast<int>(x->ne[3]),
                                                   static_cast<int>(context->ne[1]),
                                                   ref_latents,
                                                   increase_ref_index,
-                                                  qwen_image_params.theta,
+                                                  config.theta,
                                                   circular_y_enabled,
                                                   circular_x_enabled,
-                                                  qwen_image_params.axes_dim);
-            int pos_len = static_cast<int>(pe_vec.size() / qwen_image_params.axes_dim_sum / 2);
+                                                  config.axes_dim);
+            int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
             // LOG_DEBUG("pos_len %d", pos_len);
-            auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, qwen_image_params.axes_dim_sum / 2, pos_len);
+            auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
             // pe->data = pe_vec.data();
             // print_ggml_tensor(pe, true, "pe");
             // pe->data = nullptr;
             set_backend_tensor_data(pe, pe_vec.data());
 
             ggml_tensor* modulate_index = nullptr;
-            if (qwen_image_params.zero_cond_t) {
+            if (config.zero_cond_t) {
                 modulate_index_vec.clear();
 
-                int64_t h_len          = ((x->ne[1] + (qwen_image_params.patch_size / 2)) / qwen_image_params.patch_size);
-                int64_t w_len          = ((x->ne[0] + (qwen_image_params.patch_size / 2)) / qwen_image_params.patch_size);
+                int64_t h_len          = ((x->ne[1] + (config.patch_size / 2)) / config.patch_size);
+                int64_t w_len          = ((x->ne[0] + (config.patch_size / 2)) / config.patch_size);
                 int64_t num_img_tokens = h_len * w_len;
 
                 modulate_index_vec.insert(modulate_index_vec.end(), num_img_tokens, 0.f);
                 int64_t num_ref_img_tokens = 0;
                 for (ggml_tensor* ref : ref_latents) {
-                    int64_t h_len = ((ref->ne[1] + (qwen_image_params.patch_size / 2)) / qwen_image_params.patch_size);
-                    int64_t w_len = ((ref->ne[0] + (qwen_image_params.patch_size / 2)) / qwen_image_params.patch_size);
+                    int64_t h_len = ((ref->ne[1] + (config.patch_size / 2)) / config.patch_size);
+                    int64_t w_len = ((ref->ne[0] + (config.patch_size / 2)) / config.patch_size);
 
                     num_ref_img_tokens += h_len * w_len;
                 }
@@ -622,7 +627,7 @@ namespace Qwen {
                 return build_graph(x, timesteps, context, ref_latents, increase_ref_index);
             };
 
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,
@@ -686,7 +691,8 @@ namespace Qwen {
             ggml_backend_t backend    = sd_backend_cpu_init();
             ggml_type model_data_type = GGML_TYPE_Q8_0;
 
-            ModelLoader model_loader;
+            auto model_manager        = std::make_shared<ModelManager>();
+            ModelLoader& model_loader = model_manager->loader();
             if (!model_loader.init_from_file_and_convert_name(file_path, "model.diffusion_model.")) {
                 LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
                 return;
@@ -700,23 +706,20 @@ namespace Qwen {
             }
 
             std::shared_ptr<QwenImageRunner> qwen_image = std::make_shared<QwenImageRunner>(backend,
-                                                                                            backend,
                                                                                             tensor_storage_map,
                                                                                             "model.diffusion_model",
-                                                                                            VERSION_QWEN_IMAGE);
+                                                                                            VERSION_QWEN_IMAGE,
+                                                                                            false,
+                                                                                            model_manager);
 
-            if (!qwen_image->alloc_params_buffer()) {
-                LOG_ERROR("qwen_image buffer allocation failed");
-                return;
-            }
-
-            std::map<std::string, ggml_tensor*> tensors;
-            qwen_image->get_param_tensors(tensors, "model.diffusion_model");
-
-            bool success = model_loader.load_tensors(tensors);
-
-            if (!success) {
-                LOG_ERROR("load tensors from model loader failed");
+            if (!model_manager->register_runner_params("Qwen image test",
+                                                       *qwen_image,
+                                                       "model.diffusion_model",
+                                                       ModelManager::ResidencyMode::ParamBackend,
+                                                       backend,
+                                                       backend) ||
+                !model_manager->validate_registered_tensors()) {
+                LOG_ERROR("register qwen_image tensors with model manager failed");
                 return;
             }
 
@@ -727,4 +730,4 @@ namespace Qwen {
 
 }  // namespace name
 
-#endif  // __QWEN_IMAGE_HPP__
+#endif  // __SD_MODEL_DIFFUSION_QWEN_IMAGE_HPP__

src/model/diffusion/unet.hpp

@@ -1,14 +1,136 @@
-#ifndef __UNET_HPP__
-#define __UNET_HPP__
+#ifndef __SD_MODEL_DIFFUSION_UNET_HPP__
+#define __SD_MODEL_DIFFUSION_UNET_HPP__
+
+#include <algorithm>
+#include <vector>
 
-#include "common_block.hpp"
-#include "diffusion_model.hpp"
 #include "model.h"
+#include "model/common/block.hpp"
+#include "model/diffusion/model.hpp"
 
 /*==================================================== UnetModel =====================================================*/
 
 #define UNET_GRAPH_SIZE 102400
 
+struct UNetConfig {
+    SDVersion version = VERSION_SD1;
+    // network hparams
+    int in_channels                        = 4;
+    int out_channels                       = 4;
+    int num_res_blocks                     = 2;
+    std::vector<int> attention_resolutions = {4, 2, 1};
+    std::vector<int> channel_mult          = {1, 2, 4, 4};
+    std::vector<int> transformer_depth     = {1, 1, 1, 1};
+    int time_embed_dim                     = 1280;  // model_channels*4
+    int num_heads                          = 8;
+    int num_head_channels                  = -1;   // channels // num_heads
+    int context_dim                        = 768;  // 1024 for VERSION_SD2, 2048 for VERSION_SDXL
+    bool use_linear_projection             = false;
+    bool tiny_unet                         = false;
+    int model_channels                     = 320;
+    int adm_in_channels                    = 2816;  // only for VERSION_SDXL/SVD
+
+    static UNetConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
+                                          const std::string& prefix,
+                                          SDVersion version = VERSION_SD1) {
+        UNetConfig config;
+        config.version = version;
+
+        if (sd_version_is_sd2(version)) {
+            config.context_dim           = 1024;
+            config.num_head_channels     = 64;
+            config.num_heads             = -1;
+            config.use_linear_projection = true;
+        } else if (sd_version_is_sdxl(version)) {
+            config.context_dim           = 2048;
+            config.attention_resolutions = {4, 2};
+            config.channel_mult          = {1, 2, 4};
+            config.transformer_depth     = {1, 2, 10};
+            config.num_head_channels     = 64;
+            config.num_heads             = -1;
+            config.use_linear_projection = true;
+            if (version == VERSION_SDXL_VEGA) {
+                config.transformer_depth = {1, 1, 2};
+            }
+        } else if (version == VERSION_SVD) {
+            config.in_channels           = 8;
+            config.out_channels          = 4;
+            config.context_dim           = 1024;
+            config.adm_in_channels       = 768;
+            config.num_head_channels     = 64;
+            config.num_heads             = -1;
+            config.use_linear_projection = true;
+        }
+        if (sd_version_is_inpaint(version)) {
+            config.in_channels = 9;
+        } else if (sd_version_is_unet_edit(version)) {
+            config.in_channels = 8;
+        }
+        if (version == VERSION_SD1_TINY_UNET || version == VERSION_SD2_TINY_UNET || version == VERSION_SDXS_512_DS || version == VERSION_SDXS_09) {
+            config.num_res_blocks = 1;
+            config.channel_mult   = {1, 2, 4};
+            config.tiny_unet      = true;
+            if (version == VERSION_SDXS_512_DS) {
+                config.attention_resolutions = {4, 2};  // here just like SDXL
+            }
+        }
+
+        auto find_weight = [&](const std::string& suffix) -> const TensorStorage* {
+            std::string name = prefix.empty() ? suffix : prefix + "." + suffix;
+            auto it          = tensor_storage_map.find(name);
+            if (it == tensor_storage_map.end()) {
+                return nullptr;
+            }
+            return &it->second;
+        };
+
+        if (const TensorStorage* input = find_weight("input_blocks.0.0.weight")) {
+            if (input->n_dims == 4) {
+                config.in_channels    = static_cast<int>(input->ne[2]);
+                config.model_channels = static_cast<int>(input->ne[3]);
+                config.time_embed_dim = config.model_channels * 4;
+            }
+        }
+        if (const TensorStorage* time_embed = find_weight("time_embed.0.weight")) {
+            if (time_embed->n_dims == 2) {
+                config.model_channels = static_cast<int>(time_embed->ne[0]);
+                config.time_embed_dim = static_cast<int>(time_embed->ne[1]);
+            }
+        }
+        if (const TensorStorage* label_emb = find_weight("label_emb.0.0.weight")) {
+            if (label_emb->n_dims == 2) {
+                config.adm_in_channels = static_cast<int>(label_emb->ne[0]);
+                config.time_embed_dim  = static_cast<int>(label_emb->ne[1]);
+            }
+        }
+        if (const TensorStorage* out = find_weight("out.2.weight")) {
+            if (out->n_dims == 4) {
+                config.out_channels = static_cast<int>(out->ne[3]);
+            }
+        }
+        for (const auto& [name, tensor_storage] : tensor_storage_map) {
+            if (!starts_with(name, prefix)) {
+                continue;
+            }
+            if (name.find("attn2.to_k.weight") != std::string::npos && tensor_storage.n_dims == 2) {
+                config.context_dim = static_cast<int>(tensor_storage.ne[0]);
+                break;
+            }
+        }
+
+        LOG_DEBUG("unet: in_channels = %d, out_channels = %d, model_channels = %d, time_embed_dim = %d, context_dim = %d, adm_in_channels = %d, num_res_blocks = %d, tiny_unet = %s",
+                  config.in_channels,
+                  config.out_channels,
+                  config.model_channels,
+                  config.time_embed_dim,
+                  config.context_dim,
+                  config.adm_in_channels,
+                  config.num_res_blocks,
+                  config.tiny_unet ? "true" : "false");
+        return config;
+    }
+};
+
 class SpatialVideoTransformer : public SpatialTransformer {
 protected:
     int64_t time_depth;
@@ -166,66 +288,26 @@ public:
 
 // ldm.modules.diffusionmodules.openaimodel.UNetModel
 class UnetModelBlock : public GGMLBlock {
-protected:
-    SDVersion version = VERSION_SD1;
-    // network hparams
-    int in_channels                        = 4;
-    int out_channels                       = 4;
-    int num_res_blocks                     = 2;
-    std::vector<int> attention_resolutions = {4, 2, 1};
-    std::vector<int> channel_mult          = {1, 2, 4, 4};
-    std::vector<int> transformer_depth     = {1, 1, 1, 1};
-    int time_embed_dim                     = 1280;  // model_channels*4
-    int num_heads                          = 8;
-    int num_head_channels                  = -1;   // channels // num_heads
-    int context_dim                        = 768;  // 1024 for VERSION_SD2, 2048 for VERSION_SDXL
-    bool use_linear_projection             = false;
-    bool tiny_unet                         = false;
-
 public:
-    int model_channels  = 320;
-    int adm_in_channels = 2816;  // only for VERSION_SDXL/SVD
+    UNetConfig config;
 
-    UnetModelBlock(SDVersion version = VERSION_SD1, const String2TensorStorage& tensor_storage_map = {})
-        : version(version) {
-        if (sd_version_is_sd2(version)) {
-            context_dim           = 1024;
-            num_head_channels     = 64;
-            num_heads             = -1;
-            use_linear_projection = true;
-        } else if (sd_version_is_sdxl(version)) {
-            context_dim           = 2048;
-            attention_resolutions = {4, 2};
-            channel_mult          = {1, 2, 4};
-            transformer_depth     = {1, 2, 10};
-            num_head_channels     = 64;
-            num_heads             = -1;
-            use_linear_projection = true;
-            if (version == VERSION_SDXL_VEGA) {
-                transformer_depth = {1, 1, 2};
-            }
-        } else if (version == VERSION_SVD) {
-            in_channels           = 8;
-            out_channels          = 4;
-            context_dim           = 1024;
-            adm_in_channels       = 768;
-            num_head_channels     = 64;
-            num_heads             = -1;
-            use_linear_projection = true;
-        }
-        if (sd_version_is_inpaint(version)) {
-            in_channels = 9;
-        } else if (sd_version_is_unet_edit(version)) {
-            in_channels = 8;
-        }
-        if (version == VERSION_SD1_TINY_UNET || version == VERSION_SD2_TINY_UNET || version == VERSION_SDXS_512_DS || version == VERSION_SDXS_09) {
-            num_res_blocks = 1;
-            channel_mult   = {1, 2, 4};
-            tiny_unet      = true;
-            if (version == VERSION_SDXS_512_DS) {
-                attention_resolutions = {4, 2};  // here just like SDXL
-            }
-        }
+    explicit UnetModelBlock(UNetConfig config = {})
+        : config(config) {
+        const SDVersion version           = this->config.version;
+        const int in_channels             = this->config.in_channels;
+        const int out_channels            = this->config.out_channels;
+        const int num_res_blocks          = this->config.num_res_blocks;
+        const auto& attention_resolutions = this->config.attention_resolutions;
+        const auto& channel_mult          = this->config.channel_mult;
+        const auto& transformer_depth     = this->config.transformer_depth;
+        const int time_embed_dim          = this->config.time_embed_dim;
+        const int num_heads               = this->config.num_heads;
+        const int num_head_channels       = this->config.num_head_channels;
+        const int context_dim             = this->config.context_dim;
+        const bool use_linear_projection  = this->config.use_linear_projection;
+        const bool tiny_unet              = this->config.tiny_unet;
+        const int model_channels          = this->config.model_channels;
+        const int adm_in_channels         = this->config.adm_in_channels;
 
         // dims is always 2
         // use_temporal_attention is always True for SVD
@@ -398,7 +480,7 @@ public:
                                   ggml_tensor* x,
                                   ggml_tensor* emb,
                                   int num_video_frames) {
-        if (version == VERSION_SVD) {
+        if (config.version == VERSION_SVD) {
             auto block = std::dynamic_pointer_cast<VideoResBlock>(blocks[name]);
 
             return block->forward(ctx, x, emb, num_video_frames);
@@ -414,7 +496,7 @@ public:
                                          ggml_tensor* x,
                                          ggml_tensor* context,
                                          int timesteps) {
-        if (version == VERSION_SVD) {
+        if (config.version == VERSION_SVD) {
             auto block = std::dynamic_pointer_cast<SpatialVideoTransformer>(blocks[name]);
 
             return block->forward(ctx, x, context, timesteps);
@@ -440,6 +522,13 @@ public:
         // c_concat: [N, in_channels, h, w] or [1, in_channels, h, w]
         // y: [N, adm_in_channels] or [1, adm_in_channels]
         // return: [N, out_channels, h, w]
+        const SDVersion version           = config.version;
+        const int model_channels          = config.model_channels;
+        const int num_res_blocks          = config.num_res_blocks;
+        const auto& attention_resolutions = config.attention_resolutions;
+        const auto& channel_mult          = config.channel_mult;
+        const bool tiny_unet              = config.tiny_unet;
+
         if (context != nullptr) {
             if (context->ne[2] != x->ne[3]) {
                 context = ggml_repeat(ctx->ggml_ctx, context, ggml_new_tensor_3d(ctx->ggml_ctx, GGML_TYPE_F32, context->ne[0], context->ne[1], x->ne[3]));
@@ -601,14 +690,17 @@ public:
 };
 
 struct UNetModelRunner : public DiffusionModelRunner {
+    UNetConfig config;
     UnetModelBlock unet;
 
     UNetModelRunner(ggml_backend_t backend,
-                    ggml_backend_t params_backend,
                     const String2TensorStorage& tensor_storage_map,
                     const std::string prefix,
-                    SDVersion version = VERSION_SD1)
-        : DiffusionModelRunner(backend, params_backend, prefix), unet(version, tensor_storage_map) {
+                    SDVersion version                                   = VERSION_SD1,
+                    std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+        : DiffusionModelRunner(backend, prefix, weight_manager),
+          config(UNetConfig::detect_from_weights(tensor_storage_map, prefix, version)),
+          unet(config) {
         unet.init(params_ctx, tensor_storage_map, prefix);
     }
 
@@ -680,7 +772,7 @@ struct UNetModelRunner : public DiffusionModelRunner {
             return build_graph(x, timesteps, context, c_concat, y, num_video_frames, controls, control_strength);
         };
 
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
     }
 
     sd::Tensor<float> compute(int n_threads,
@@ -752,4 +844,4 @@ struct UNetModelRunner : public DiffusionModelRunner {
     }
 };
 
-#endif  // __UNET_HPP__
+#endif  // __SD_MODEL_DIFFUSION_UNET_HPP__

src/model/diffusion/z_image.hpp

@@ -1,14 +1,15 @@
-#ifndef __Z_IMAGE_HPP__
-#define __Z_IMAGE_HPP__
+#ifndef __SD_MODEL_DIFFUSION_Z_IMAGE_HPP__
+#define __SD_MODEL_DIFFUSION_Z_IMAGE_HPP__
 
 #include <algorithm>
 
-#include "diffusion_model.hpp"
-#include "flux.hpp"
-#include "ggml_extend.hpp"
-#include "mmdit.hpp"
+#include "core/ggml_extend.hpp"
+#include "model/diffusion/flux.hpp"
+#include "model/diffusion/mmdit.hpp"
+#include "model/diffusion/model.hpp"
+#include "model_loader.h"
 
-// Ref: https://github.com/Alpha-VLLM/Lumina-Image-2.0/blob/main/models/model.py
+// Ref: https://github.com/Alpha-VLLM/Lumina-Image-2.0/blob/main/model/model.py
 // Ref: https://github.com/huggingface/diffusers/pull/12703
 
 #ifndef MIN
@@ -20,6 +21,104 @@ namespace ZImage {
     constexpr int ADALN_EMBED_DIM    = 256;
     constexpr int SEQ_MULTI_OF       = 32;
 
+    struct ZImageConfig {
+        int patch_size             = 2;
+        int64_t hidden_size        = 3840;
+        int64_t in_channels        = 16;
+        int64_t out_channels       = 16;
+        int64_t num_layers         = 30;
+        int64_t num_refiner_layers = 2;
+        int64_t head_dim           = 128;
+        int64_t num_heads          = 30;
+        int64_t num_kv_heads       = 30;
+        int64_t multiple_of        = 256;
+        float ffn_dim_multiplier   = 8.0f / 3.0f;
+        float norm_eps             = 1e-5f;
+        bool qk_norm               = true;
+        int64_t cap_feat_dim       = 2560;
+        int theta                  = 256;
+        std::vector<int> axes_dim  = {32, 48, 48};
+        int64_t axes_dim_sum       = 128;
+
+        static ZImageConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
+            ZImageConfig config;
+            int64_t detected_layers          = 0;
+            int64_t detected_refiner_layers  = 0;
+            int64_t detected_context_refiner = 0;
+            int64_t detected_head_dim        = 0;
+            int64_t detected_qkv_dim         = 0;
+
+            for (const auto& [name, tensor_storage] : tensor_storage_map) {
+                if (!starts_with(name, prefix)) {
+                    continue;
+                }
+                if (ends_with(name, "x_embedder.weight") && tensor_storage.n_dims == 2) {
+                    int64_t patch_area = config.patch_size * config.patch_size;
+                    config.in_channels = tensor_storage.ne[0] / patch_area;
+                    config.hidden_size = tensor_storage.ne[1];
+                } else if (ends_with(name, "cap_embedder.1.weight") && tensor_storage.n_dims == 2) {
+                    config.cap_feat_dim = tensor_storage.ne[0];
+                    config.hidden_size  = tensor_storage.ne[1];
+                } else if (ends_with(name, "layers.0.attention.q_norm.weight") && tensor_storage.n_dims == 1) {
+                    detected_head_dim = tensor_storage.ne[0];
+                } else if (ends_with(name, "layers.0.attention.qkv.weight") && tensor_storage.n_dims == 2) {
+                    detected_qkv_dim = tensor_storage.ne[1];
+                } else if (ends_with(name, "final_layer.linear.weight") && tensor_storage.n_dims == 2) {
+                    int64_t patch_area  = config.patch_size * config.patch_size;
+                    config.out_channels = tensor_storage.ne[1] / patch_area;
+                }
+
+                size_t pos = name.find("layers.");
+                if (pos != std::string::npos) {
+                    auto items = split_string(name.substr(pos), '.');
+                    if (items.size() > 1) {
+                        int block_index = atoi(items[1].c_str());
+                        detected_layers = std::max<int64_t>(detected_layers, block_index + 1);
+                    }
+                }
+                pos = name.find("noise_refiner.");
+                if (pos != std::string::npos) {
+                    auto items = split_string(name.substr(pos), '.');
+                    if (items.size() > 1) {
+                        int block_index         = atoi(items[1].c_str());
+                        detected_refiner_layers = std::max<int64_t>(detected_refiner_layers, block_index + 1);
+                    }
+                }
+                pos = name.find("context_refiner.");
+                if (pos != std::string::npos) {
+                    auto items = split_string(name.substr(pos), '.');
+                    if (items.size() > 1) {
+                        int block_index          = atoi(items[1].c_str());
+                        detected_context_refiner = std::max<int64_t>(detected_context_refiner, block_index + 1);
+                    }
+                }
+            }
+            if (detected_layers > 0) {
+                config.num_layers = detected_layers;
+            }
+            if (detected_refiner_layers > 0 || detected_context_refiner > 0) {
+                config.num_refiner_layers = std::max(detected_refiner_layers, detected_context_refiner);
+            }
+            if (detected_head_dim > 0) {
+                config.head_dim  = detected_head_dim;
+                config.num_heads = config.hidden_size / config.head_dim;
+                if (detected_qkv_dim > 0) {
+                    int64_t qkv_heads   = detected_qkv_dim / config.head_dim;
+                    config.num_kv_heads = std::max<int64_t>(1, (qkv_heads - config.num_heads) / 2);
+                }
+            }
+            LOG_DEBUG("z_image: num_layers = %" PRId64 ", num_refiner_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %" PRId64 ", num_kv_heads = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64,
+                      config.num_layers,
+                      config.num_refiner_layers,
+                      config.hidden_size,
+                      config.num_heads,
+                      config.num_kv_heads,
+                      config.in_channels,
+                      config.out_channels);
+            return config;
+        }
+    };
+
     struct JointAttention : public GGMLBlock {
     protected:
         int64_t head_dim;
@@ -263,90 +362,70 @@ namespace ZImage {
         }
     };
 
-    struct ZImageParams {
-        int patch_size             = 2;
-        int64_t hidden_size        = 3840;
-        int64_t in_channels        = 16;
-        int64_t out_channels       = 16;
-        int64_t num_layers         = 30;
-        int64_t num_refiner_layers = 2;
-        int64_t head_dim           = 128;
-        int64_t num_heads          = 30;
-        int64_t num_kv_heads       = 30;
-        int64_t multiple_of        = 256;
-        float ffn_dim_multiplier   = 8.0f / 3.0f;
-        float norm_eps             = 1e-5f;
-        bool qk_norm               = true;
-        int64_t cap_feat_dim       = 2560;
-        int theta                  = 256;
-        std::vector<int> axes_dim  = {32, 48, 48};
-        int64_t axes_dim_sum       = 128;
-    };
-
     class ZImageModel : public GGMLBlock {
     protected:
-        ZImageParams z_image_params;
+        ZImageConfig config;
 
         void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
-            params["cap_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, z_image_params.hidden_size);
-            params["x_pad_token"]   = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, z_image_params.hidden_size);
+            params["cap_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size);
+            params["x_pad_token"]   = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size);
         }
 
     public:
         ZImageModel() = default;
-        ZImageModel(ZImageParams z_image_params)
-            : z_image_params(z_image_params) {
-            blocks["x_embedder"]     = std::make_shared<Linear>(z_image_params.patch_size * z_image_params.patch_size * z_image_params.in_channels, z_image_params.hidden_size);
-            blocks["t_embedder"]     = std::make_shared<TimestepEmbedder>(MIN(z_image_params.hidden_size, 1024), 256, 256);
-            blocks["cap_embedder.0"] = std::make_shared<RMSNorm>(z_image_params.cap_feat_dim, z_image_params.norm_eps);
-            blocks["cap_embedder.1"] = std::make_shared<Linear>(z_image_params.cap_feat_dim, z_image_params.hidden_size);
+        ZImageModel(ZImageConfig config)
+            : config(config) {
+            blocks["x_embedder"]     = std::make_shared<Linear>(config.patch_size * config.patch_size * config.in_channels, config.hidden_size);
+            blocks["t_embedder"]     = std::make_shared<TimestepEmbedder>(MIN(config.hidden_size, 1024), 256, 256);
+            blocks["cap_embedder.0"] = std::make_shared<RMSNorm>(config.cap_feat_dim, config.norm_eps);
+            blocks["cap_embedder.1"] = std::make_shared<Linear>(config.cap_feat_dim, config.hidden_size);
 
-            for (int i = 0; i < z_image_params.num_refiner_layers; i++) {
+            for (int i = 0; i < config.num_refiner_layers; i++) {
                 auto block = std::make_shared<JointTransformerBlock>(i,
-                                                                     z_image_params.hidden_size,
-                                                                     z_image_params.head_dim,
-                                                                     z_image_params.num_heads,
-                                                                     z_image_params.num_kv_heads,
-                                                                     z_image_params.multiple_of,
-                                                                     z_image_params.ffn_dim_multiplier,
-                                                                     z_image_params.norm_eps,
-                                                                     z_image_params.qk_norm,
+                                                                     config.hidden_size,
+                                                                     config.head_dim,
+                                                                     config.num_heads,
+                                                                     config.num_kv_heads,
+                                                                     config.multiple_of,
+                                                                     config.ffn_dim_multiplier,
+                                                                     config.norm_eps,
+                                                                     config.qk_norm,
                                                                      true);
 
                 blocks["noise_refiner." + std::to_string(i)] = block;
             }
 
-            for (int i = 0; i < z_image_params.num_refiner_layers; i++) {
+            for (int i = 0; i < config.num_refiner_layers; i++) {
                 auto block = std::make_shared<JointTransformerBlock>(i,
-                                                                     z_image_params.hidden_size,
-                                                                     z_image_params.head_dim,
-                                                                     z_image_params.num_heads,
-                                                                     z_image_params.num_kv_heads,
-                                                                     z_image_params.multiple_of,
-                                                                     z_image_params.ffn_dim_multiplier,
-                                                                     z_image_params.norm_eps,
-                                                                     z_image_params.qk_norm,
+                                                                     config.hidden_size,
+                                                                     config.head_dim,
+                                                                     config.num_heads,
+                                                                     config.num_kv_heads,
+                                                                     config.multiple_of,
+                                                                     config.ffn_dim_multiplier,
+                                                                     config.norm_eps,
+                                                                     config.qk_norm,
                                                                      false);
 
                 blocks["context_refiner." + std::to_string(i)] = block;
             }
 
-            for (int i = 0; i < z_image_params.num_layers; i++) {
+            for (int i = 0; i < config.num_layers; i++) {
                 auto block = std::make_shared<JointTransformerBlock>(i,
-                                                                     z_image_params.hidden_size,
-                                                                     z_image_params.head_dim,
-                                                                     z_image_params.num_heads,
-                                                                     z_image_params.num_kv_heads,
-                                                                     z_image_params.multiple_of,
-                                                                     z_image_params.ffn_dim_multiplier,
-                                                                     z_image_params.norm_eps,
-                                                                     z_image_params.qk_norm,
+                                                                     config.hidden_size,
+                                                                     config.head_dim,
+                                                                     config.num_heads,
+                                                                     config.num_kv_heads,
+                                                                     config.multiple_of,
+                                                                     config.ffn_dim_multiplier,
+                                                                     config.norm_eps,
+                                                                     config.qk_norm,
                                                                      true);
 
                 blocks["layers." + std::to_string(i)] = block;
             }
 
-            blocks["final_layer"] = std::make_shared<FinalLayer>(z_image_params.hidden_size, z_image_params.patch_size, z_image_params.out_channels);
+            blocks["final_layer"] = std::make_shared<FinalLayer>(config.hidden_size, config.patch_size, config.out_channels);
         }
 
         ggml_tensor* forward_core(GGMLRunnerContext* ctx,
@@ -393,14 +472,14 @@ namespace ZImage {
             auto txt_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, 0, txt->ne[1]);
             auto img_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt->ne[1], pe->ne[3]);
 
-            for (int i = 0; i < z_image_params.num_refiner_layers; i++) {
+            for (int i = 0; i < config.num_refiner_layers; i++) {
                 auto block = std::dynamic_pointer_cast<JointTransformerBlock>(blocks["context_refiner." + std::to_string(i)]);
 
                 txt = block->forward(ctx, txt, txt_pe, nullptr, nullptr);
                 sd::ggml_graph_cut::mark_graph_cut(txt, "z_image.context_refiner." + std::to_string(i), "txt");
             }
 
-            for (int i = 0; i < z_image_params.num_refiner_layers; i++) {
+            for (int i = 0; i < config.num_refiner_layers; i++) {
                 auto block = std::dynamic_pointer_cast<JointTransformerBlock>(blocks["noise_refiner." + std::to_string(i)]);
 
                 img = block->forward(ctx, img, img_pe, nullptr, t_emb);
@@ -410,7 +489,7 @@ namespace ZImage {
             auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1);  // [N, n_txt_token + n_txt_pad_token + n_img_token + n_img_pad_token, hidden_size]
             sd::ggml_graph_cut::mark_graph_cut(txt_img, "z_image.prelude", "txt_img");
 
-            for (int i = 0; i < z_image_params.num_layers; i++) {
+            for (int i = 0; i < config.num_layers; i++) {
                 auto block = std::dynamic_pointer_cast<JointTransformerBlock>(blocks["layers." + std::to_string(i)]);
 
                 txt_img = block->forward(ctx, txt_img, pe, nullptr, t_emb);
@@ -442,7 +521,7 @@ namespace ZImage {
             int64_t C = x->ne[2];
             int64_t N = x->ne[3];
 
-            int patch_size = z_image_params.patch_size;
+            int patch_size = config.patch_size;
 
             auto img             = DiT::pad_and_patchify(ctx, x, patch_size, patch_size, false);
             uint64_t n_img_token = img->ne[1];
@@ -467,19 +546,20 @@ namespace ZImage {
 
     struct ZImageRunner : public DiffusionModelRunner {
     public:
-        ZImageParams z_image_params;
+        ZImageConfig config;
         ZImageModel z_image;
         std::vector<float> pe_vec;
         std::vector<float> timestep_vec;
         SDVersion version;
 
         ZImageRunner(ggml_backend_t backend,
-                     ggml_backend_t params_backend,
-                     const String2TensorStorage& tensor_storage_map = {},
-                     const std::string prefix                       = "",
-                     SDVersion version                              = VERSION_Z_IMAGE)
-            : DiffusionModelRunner(backend, params_backend, prefix) {
-            z_image = ZImageModel(z_image_params);
+                     const String2TensorStorage& tensor_storage_map      = {},
+                     const std::string prefix                            = "",
+                     SDVersion version                                   = VERSION_Z_IMAGE,
+                     std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : DiffusionModelRunner(backend, prefix, weight_manager),
+              config(ZImageConfig::detect_from_weights(tensor_storage_map, prefix)) {
+            z_image = ZImageModel(config);
             z_image.init(params_ctx, tensor_storage_map, prefix);
         }
 
@@ -510,19 +590,19 @@ namespace ZImage {
 
             pe_vec      = Rope::gen_z_image_pe(static_cast<int>(x->ne[1]),
                                                static_cast<int>(x->ne[0]),
-                                               z_image_params.patch_size,
+                                               config.patch_size,
                                                static_cast<int>(x->ne[3]),
                                                static_cast<int>(context->ne[1]),
                                                SEQ_MULTI_OF,
                                                ref_latents,
                                                increase_ref_index,
-                                               z_image_params.theta,
+                                               config.theta,
                                                circular_y_enabled,
                                                circular_x_enabled,
-                                               z_image_params.axes_dim);
-            int pos_len = static_cast<int>(pe_vec.size() / z_image_params.axes_dim_sum / 2);
+                                               config.axes_dim);
+            int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
             // LOG_DEBUG("pos_len %d", pos_len);
-            auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, z_image_params.axes_dim_sum / 2, pos_len);
+            auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
             // pe->data = pe_vec.data();
             // print_ggml_tensor(pe, true, "pe");
             // pe->data = nullptr;
@@ -554,7 +634,7 @@ namespace ZImage {
                 return build_graph(x, timesteps, context, ref_latents, increase_ref_index);
             };
 
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
         }
 
         sd::Tensor<float> compute(int n_threads,
@@ -618,7 +698,8 @@ namespace ZImage {
             ggml_backend_t backend    = sd_backend_cpu_init();
             ggml_type model_data_type = GGML_TYPE_Q8_0;
 
-            ModelLoader model_loader;
+            auto model_manager        = std::make_shared<ModelManager>();
+            ModelLoader& model_loader = model_manager->loader();
             if (!model_loader.init_from_file_and_convert_name(file_path, "model.diffusion_model.")) {
                 LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
                 return;
@@ -634,22 +715,19 @@ namespace ZImage {
             }
 
             std::shared_ptr<ZImageRunner> z_image = std::make_shared<ZImageRunner>(backend,
-                                                                                   backend,
                                                                                    tensor_storage_map,
                                                                                    "model.diffusion_model",
-                                                                                   VERSION_QWEN_IMAGE);
+                                                                                   VERSION_QWEN_IMAGE,
+                                                                                   model_manager);
 
-            if (!z_image->alloc_params_buffer()) {
-                LOG_ERROR("z_image buffer allocation failed");
-                return;
-            }
-            std::map<std::string, ggml_tensor*> tensors;
-            z_image->get_param_tensors(tensors, "model.diffusion_model");
-
-            bool success = model_loader.load_tensors(tensors);
-
-            if (!success) {
-                LOG_ERROR("load tensors from model loader failed");
+            if (!model_manager->register_runner_params("ZImage test",
+                                                       *z_image,
+                                                       "model.diffusion_model",
+                                                       ModelManager::ResidencyMode::ParamBackend,
+                                                       backend,
+                                                       backend) ||
+                !model_manager->validate_registered_tensors()) {
+                LOG_ERROR("register z_image tensors with model manager failed");
                 return;
             }
 
@@ -660,4 +738,4 @@ namespace ZImage {
 
 }  // namespace ZImage
 
-#endif  // __Z_IMAGE_HPP__
+#endif  // __SD_MODEL_DIFFUSION_Z_IMAGE_HPP__

src/model/te/clip.hpp

@@ -1,13 +1,13 @@
-#ifndef __CLIP_HPP__
-#define __CLIP_HPP__
+#ifndef __SD_MODEL_TE_CLIP_HPP__
+#define __SD_MODEL_TE_CLIP_HPP__
 
-#include "ggml_extend.hpp"
+#include "core/ggml_extend.hpp"
 #include "model.h"
 #include "tokenizers/clip_tokenizer.h"
 
 /*================================================ FrozenCLIPEmbedder ================================================*/
 
-// Ref: https://github.com/huggingface/transformers/blob/main/src/transformers/models/clip/modeling_clip.py
+// Ref: https://github.com/huggingface/transformers/blob/main/src/transformers/model/clip/modeling_clip.py
 
 struct CLIPMLP : public GGMLBlock {
 protected:
@@ -469,13 +469,13 @@ struct CLIPTextModelRunner : public GGMLRunner {
     std::vector<float> attention_mask_vec;
 
     CLIPTextModelRunner(ggml_backend_t backend,
-                        ggml_backend_t params_backend,
                         const String2TensorStorage& tensor_storage_map,
                         const std::string prefix,
-                        CLIPVersion version = OPENAI_CLIP_VIT_L_14,
-                        bool with_final_ln  = true,
-                        bool force_clip_f32 = false)
-        : GGMLRunner(backend, params_backend) {
+                        CLIPVersion version                                 = OPENAI_CLIP_VIT_L_14,
+                        bool with_final_ln                                  = true,
+                        bool force_clip_f32                                 = false,
+                        std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+        : GGMLRunner(backend, weight_manager) {
         bool proj_in = false;
         for (const auto& [name, tensor_storage] : tensor_storage_map) {
             if (!starts_with(name, prefix)) {
@@ -567,11 +567,14 @@ struct CLIPTextModelRunner : public GGMLRunner {
                               void* custom_embeddings_data,
                               size_t max_token_idx,
                               bool return_pooled,
-                              int clip_skip) {
+                              int clip_skip,
+                              bool auto_free           = true,
+                              bool free_compute_buffer = true,
+                              bool free_compute_params = true) {
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(input_ids, num_custom_embeddings, custom_embeddings_data, max_token_idx, return_pooled, clip_skip);
         };
-        auto result = GGMLRunner::compute<float>(get_graph, n_threads, true);
+        auto result = GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params);
         if (return_pooled) {
             return take_or_empty(std::move(result));
         }
@@ -579,4 +582,4 @@ struct CLIPTextModelRunner : public GGMLRunner {
     }
 };
 
-#endif  // __CLIP_HPP__
+#endif  // __SD_MODEL_TE_CLIP_HPP__

src/model/te/llm.hpp

@@ -1,5 +1,5 @@
-#ifndef __LLM_HPP__
-#define __LLM_HPP__
+#ifndef __SD_MODEL_TE_LLM_HPP__
+#define __SD_MODEL_TE_LLM_HPP__
 
 #include <algorithm>
 #include <array>
@@ -18,9 +18,11 @@
 #include <utility>
 #include <vector>
 
-#include "ggml_extend.hpp"
+#include "core/ggml_extend.hpp"
 #include "json.hpp"
-#include "rope.hpp"
+#include "model/common/rope.hpp"
+#include "model_loader.h"
+#include "model_manager.h"
 #include "tokenizers/bpe_tokenizer.h"
 #include "tokenizers/gemma_tokenizer.h"
 #include "tokenizers/gpt_oss_tokenizer.h"
@@ -63,7 +65,7 @@ namespace LLM {
         QWEN3_VL,
     };
 
-    struct LLMVisionParams {
+    struct LLMVisionConfig {
         LLMVisionArch arch                  = LLMVisionArch::QWEN2_5_VL;
         int num_layers                      = 32;
         int64_t hidden_size                 = 1280;
@@ -79,7 +81,7 @@ namespace LLM {
         std::set<int> fullatt_block_indexes = {7, 15, 23, 31};
     };
 
-    struct LLMParams {
+    struct LLMConfig {
         LLMArch arch                    = LLMArch::QWEN2_5_VL;
         int64_t num_layers              = 28;
         int64_t hidden_size             = 3584;
@@ -101,7 +103,129 @@ namespace LLM {
         std::vector<int> sliding_attention;
         int64_t num_experts         = 0;
         int64_t num_experts_per_tok = 0;
-        LLMVisionParams vision;
+        LLMVisionConfig vision;
+        bool have_vision_weight = false;
+        bool llama_cpp_style    = false;
+
+        static LLMConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
+                                             const std::string& prefix,
+                                             LLMArch arch) {
+            LLMConfig config;
+            config.arch = arch;
+            if (arch == LLMArch::MISTRAL_SMALL_3_2 || arch == LLMArch::MINISTRAL_3_3B) {
+                config.head_dim     = 128;
+                config.num_heads    = 32;
+                config.num_kv_heads = 8;
+                config.qkv_bias     = false;
+                config.rms_norm_eps = 1e-5f;
+            } else if (arch == LLMArch::QWEN3 || arch == LLMArch::QWEN3_VL) {
+                config.head_dim     = 128;
+                config.num_heads    = 32;
+                config.num_kv_heads = 8;
+                config.qkv_bias     = false;
+                config.qk_norm      = true;
+                config.rms_norm_eps = 1e-6f;
+                if (arch == LLMArch::QWEN3_VL) {
+                    config.max_position_embeddings = 262144;
+                    config.rope_thetas             = {5000000.f};
+                    config.vision.arch             = LLMVisionArch::QWEN3_VL;
+                }
+            } else if (arch == LLMArch::GEMMA3_12B) {
+                config.head_dim                = 256;
+                config.num_heads               = 16;
+                config.num_kv_heads            = 8;
+                config.qkv_bias                = false;
+                config.qk_norm                 = true;
+                config.rms_norm_eps            = 1e-6f;
+                config.rms_norm_add            = false;
+                config.normalize_input         = true;
+                config.max_position_embeddings = 131072;
+                config.mlp_activation          = MLPActivation::GELU_TANH;
+                config.rope_thetas             = {1000000.f, 10000.f};
+                config.rope_scales             = {8.f, 1.f};
+                config.sliding_attention       = {1024, 1024, 1024, 1024, 1024, 0};
+            } else if (arch == LLMArch::GEMMA2_2B) {
+                config.head_dim                = 256;
+                config.num_heads               = 8;
+                config.num_kv_heads            = 4;
+                config.qkv_bias                = false;
+                config.qk_norm                 = false;
+                config.rms_norm_eps            = 1e-6f;
+                config.rms_norm_add            = true;
+                config.normalize_input         = true;
+                config.max_position_embeddings = 8192;
+                config.mlp_activation          = MLPActivation::GELU_TANH;
+                config.hidden_size             = 2304;
+                config.intermediate_size       = 9216;
+                config.num_layers              = 26;
+                config.vocab_size              = 256000;
+            } else if (arch == LLMArch::GPT_OSS_20B) {
+                config.head_dim                = 64;
+                config.num_heads               = 64;
+                config.num_kv_heads            = 8;
+                config.qkv_bias                = true;
+                config.attention_out_bias      = true;
+                config.qk_norm                 = false;
+                config.rms_norm_eps            = 1e-5f;
+                config.hidden_size             = 2880;
+                config.intermediate_size       = 2880;
+                config.num_layers              = 24;
+                config.vocab_size              = 201088;
+                config.max_position_embeddings = 131072;
+                config.rope_thetas             = {150000.f};
+                config.rope_scales             = {32.f};
+                config.sliding_attention       = {128, 0};
+                config.num_experts             = 32;
+                config.num_experts_per_tok     = 4;
+            }
+
+            config.num_layers = 0;
+            for (const auto& [name, tensor_storage] : tensor_storage_map) {
+                if (!starts_with(name, prefix)) {
+                    continue;
+                }
+                size_t pos = name.find("visual.");
+                if (pos != std::string::npos) {
+                    config.have_vision_weight = true;
+                    if (contains(name, "attn.q_proj")) {
+                        config.llama_cpp_style = true;
+                    }
+                    continue;
+                }
+                pos = name.find("layers.");
+                if (pos != std::string::npos) {
+                    auto items = split_string(name.substr(pos), '.');
+                    if (items.size() > 1) {
+                        int block_index = atoi(items[1].c_str());
+                        if (block_index + 1 > config.num_layers) {
+                            config.num_layers = block_index + 1;
+                        }
+                    }
+                }
+                if (contains(name, "embed_tokens.weight")) {
+                    config.hidden_size = tensor_storage.ne[0];
+                    config.vocab_size  = tensor_storage.ne[1];
+                }
+                if (contains(name, "layers.0.mlp.gate_proj.weight")) {
+                    config.intermediate_size = tensor_storage.ne[1];
+                }
+                if (contains(name, "layers.0.mlp.experts.gate_up_proj.weight")) {
+                    config.intermediate_size = tensor_storage.ne[1] / 2;
+                }
+                if (contains(name, "layers.0.mlp.experts.gate_proj.weight")) {
+                    config.intermediate_size = tensor_storage.ne[1];
+                }
+            }
+            if (arch == LLMArch::QWEN3 && config.num_layers == 28) {
+                config.num_heads = 16;
+            }
+            LOG_DEBUG("llm: num_layers = %" PRId64 ", vocab_size = %" PRId64 ", hidden_size = %" PRId64 ", intermediate_size = %" PRId64,
+                      config.num_layers,
+                      config.vocab_size,
+                      config.hidden_size,
+                      config.intermediate_size);
+            return config;
+        }
     };
 
     struct LLMRMSNorm : public UnaryBlock {
@@ -232,11 +356,11 @@ namespace LLM {
         }
 
     public:
-        GPTOSSMLP(const LLMParams& params)
-            : hidden_size(params.hidden_size),
-              intermediate_size(params.intermediate_size),
-              num_experts(params.num_experts),
-              num_experts_per_tok(params.num_experts_per_tok) {}
+        GPTOSSMLP(const LLMConfig& config)
+            : hidden_size(config.hidden_size),
+              intermediate_size(config.intermediate_size),
+              num_experts(config.num_experts),
+              num_experts_per_tok(config.num_experts_per_tok) {}
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
             // x: [N, n_token, hidden_size]
@@ -667,7 +791,7 @@ namespace LLM {
 
     public:
         VisionModel(bool llama_cpp_style,
-                    const LLMVisionParams& vision_params,
+                    const LLMVisionConfig& vision_params,
                     float eps = 1e-6f)
             : arch_(vision_params.arch),
               num_layers(vision_params.num_layers),
@@ -784,23 +908,23 @@ namespace LLM {
         }
 
     public:
-        Attention(const LLMParams& params)
-            : arch(params.arch),
-              num_heads(params.num_heads),
-              num_kv_heads(params.num_kv_heads),
-              head_dim(params.head_dim),
-              qk_norm(params.qk_norm),
-              max_position_embeddings(params.max_position_embeddings),
-              rope_thetas(params.rope_thetas),
-              rope_scales(params.rope_scales),
-              has_attention_sinks(params.arch == LLMArch::GPT_OSS_20B) {
-            blocks["q_proj"] = std::make_shared<Linear>(params.hidden_size, num_heads * head_dim, params.qkv_bias);
-            blocks["k_proj"] = std::make_shared<Linear>(params.hidden_size, num_kv_heads * head_dim, params.qkv_bias);
-            blocks["v_proj"] = std::make_shared<Linear>(params.hidden_size, num_kv_heads * head_dim, params.qkv_bias);
-            blocks["o_proj"] = std::make_shared<Linear>(num_heads * head_dim, params.hidden_size, params.attention_out_bias);
-            if (params.qk_norm) {
-                blocks["q_norm"] = std::make_shared<LLMRMSNorm>(head_dim, params.rms_norm_eps, params.rms_norm_add);
-                blocks["k_norm"] = std::make_shared<LLMRMSNorm>(head_dim, params.rms_norm_eps, params.rms_norm_add);
+        Attention(const LLMConfig& config)
+            : arch(config.arch),
+              num_heads(config.num_heads),
+              num_kv_heads(config.num_kv_heads),
+              head_dim(config.head_dim),
+              qk_norm(config.qk_norm),
+              max_position_embeddings(config.max_position_embeddings),
+              rope_thetas(config.rope_thetas),
+              rope_scales(config.rope_scales),
+              has_attention_sinks(config.arch == LLMArch::GPT_OSS_20B) {
+            blocks["q_proj"] = std::make_shared<Linear>(config.hidden_size, num_heads * head_dim, config.qkv_bias);
+            blocks["k_proj"] = std::make_shared<Linear>(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias);
+            blocks["v_proj"] = std::make_shared<Linear>(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias);
+            blocks["o_proj"] = std::make_shared<Linear>(num_heads * head_dim, config.hidden_size, config.attention_out_bias);
+            if (config.qk_norm) {
+                blocks["q_norm"] = std::make_shared<LLMRMSNorm>(head_dim, config.rms_norm_eps, config.rms_norm_add);
+                blocks["k_norm"] = std::make_shared<LLMRMSNorm>(head_dim, config.rms_norm_eps, config.rms_norm_add);
             }
         }
 
@@ -982,42 +1106,42 @@ namespace LLM {
         std::string post_ffw_norm_name;
 
     public:
-        TransformerBlock(const LLMParams& params, int layer_index)
-            : arch(params.arch),
+        TransformerBlock(const LLMConfig& config, int layer_index)
+            : arch(config.arch),
               sliding_attention(0) {
-            if (params.arch == LLMArch::GEMMA3_12B) {
+            if (config.arch == LLMArch::GEMMA3_12B) {
                 post_attention_norm_name = "post_attention_norm";       // attn_post_norm
                 pre_ffw_norm_name        = "post_attention_layernorm";  // ffn_norm
                 post_ffw_norm_name       = "post_ffw_norm";             // ffn_post_norm
-            } else if (params.arch == LLMArch::GEMMA2_2B) {
+            } else if (config.arch == LLMArch::GEMMA2_2B) {
                 post_attention_norm_name = "post_attention_layernorm";  // ffn_norm
                 pre_ffw_norm_name        = "pre_feedforward_layernorm";
                 post_ffw_norm_name       = "post_feedforward_layernorm";
-            } else if (params.arch == LLMArch::GPT_OSS_20B) {
+            } else if (config.arch == LLMArch::GPT_OSS_20B) {
                 pre_ffw_norm_name = "post_attention_norm";  // attn_post_norm
             } else {
                 pre_ffw_norm_name = "post_attention_layernorm";  // ffn_norm
             }
 
-            blocks["self_attn"] = std::make_shared<Attention>(params);
-            if (params.arch == LLMArch::GPT_OSS_20B) {
-                blocks["mlp"] = std::make_shared<GPTOSSMLP>(params);
+            blocks["self_attn"] = std::make_shared<Attention>(config);
+            if (config.arch == LLMArch::GPT_OSS_20B) {
+                blocks["mlp"] = std::make_shared<GPTOSSMLP>(config);
             } else {
-                blocks["mlp"] = std::make_shared<MLP>(params.hidden_size,
-                                                      params.intermediate_size,
+                blocks["mlp"] = std::make_shared<MLP>(config.hidden_size,
+                                                      config.intermediate_size,
                                                       false,
-                                                      params.mlp_activation);
+                                                      config.mlp_activation);
             }
-            blocks["input_layernorm"] = std::make_shared<LLMRMSNorm>(params.hidden_size, params.rms_norm_eps, params.rms_norm_add);
-            blocks[pre_ffw_norm_name] = std::make_shared<LLMRMSNorm>(params.hidden_size, params.rms_norm_eps, params.rms_norm_add);
+            blocks["input_layernorm"] = std::make_shared<LLMRMSNorm>(config.hidden_size, config.rms_norm_eps, config.rms_norm_add);
+            blocks[pre_ffw_norm_name] = std::make_shared<LLMRMSNorm>(config.hidden_size, config.rms_norm_eps, config.rms_norm_add);
             if (!post_attention_norm_name.empty()) {
-                blocks[post_attention_norm_name] = std::make_shared<LLMRMSNorm>(params.hidden_size, params.rms_norm_eps, params.rms_norm_add);
+                blocks[post_attention_norm_name] = std::make_shared<LLMRMSNorm>(config.hidden_size, config.rms_norm_eps, config.rms_norm_add);
             }
             if (!post_ffw_norm_name.empty()) {
-                blocks[post_ffw_norm_name] = std::make_shared<LLMRMSNorm>(params.hidden_size, params.rms_norm_eps, params.rms_norm_add);
+                blocks[post_ffw_norm_name] = std::make_shared<LLMRMSNorm>(config.hidden_size, config.rms_norm_eps, config.rms_norm_add);
             }
-            if (!params.sliding_attention.empty()) {
-                sliding_attention = params.sliding_attention[layer_index % params.sliding_attention.size()];
+            if (!config.sliding_attention.empty()) {
+                sliding_attention = config.sliding_attention[layer_index % config.sliding_attention.size()];
             }
         }
 
@@ -1074,16 +1198,16 @@ namespace LLM {
     struct TextModel : public GGMLBlock {
     protected:
         int64_t num_layers;
-        LLMParams params;
+        LLMConfig config;
 
     public:
-        TextModel(const LLMParams& params)
-            : num_layers(params.num_layers), params(params) {
-            blocks["embed_tokens"] = std::shared_ptr<GGMLBlock>(new Embedding(params.vocab_size, params.hidden_size));
+        TextModel(const LLMConfig& config)
+            : num_layers(config.num_layers), config(config) {
+            blocks["embed_tokens"] = std::shared_ptr<GGMLBlock>(new Embedding(config.vocab_size, config.hidden_size));
             for (int i = 0; i < num_layers; i++) {
-                blocks["layers." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new TransformerBlock(params, i));
+                blocks["layers." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new TransformerBlock(config, i));
             }
-            blocks["norm"] = std::shared_ptr<GGMLBlock>(new LLMRMSNorm(params.hidden_size, params.rms_norm_eps, params.rms_norm_add));
+            blocks["norm"] = std::shared_ptr<GGMLBlock>(new LLMRMSNorm(config.hidden_size, config.rms_norm_eps, config.rms_norm_add));
         }
 
         ggml_tensor* embed(GGMLRunnerContext* ctx,
@@ -1103,8 +1227,8 @@ namespace LLM {
             auto norm = std::dynamic_pointer_cast<LLMRMSNorm>(blocks["norm"]);
             std::vector<ggml_tensor*> intermediate_outputs;
 
-            if (params.normalize_input) {
-                x = ggml_ext_scale(ctx->ggml_ctx, x, std::sqrt(static_cast<float>(params.hidden_size)), true);
+            if (config.normalize_input) {
+                x = ggml_ext_scale(ctx->ggml_ctx, x, std::sqrt(static_cast<float>(config.hidden_size)), true);
             }
             if (return_all_hidden_states) {
                 intermediate_outputs.push_back(x);
@@ -1174,15 +1298,15 @@ namespace LLM {
 
     struct LLM : public GGMLBlock {
         bool enable_vision;
-        LLMParams params;
+        LLMConfig config;
 
     public:
         LLM() = default;
-        LLM(LLMParams params, bool enable_vision = false, bool llama_cpp_style = false)
-            : enable_vision(enable_vision), params(params) {
-            blocks["model"] = std::shared_ptr<GGMLBlock>(new TextModel(params));
+        LLM(LLMConfig config, bool enable_vision = false, bool llama_cpp_style = false)
+            : enable_vision(enable_vision), config(config) {
+            blocks["model"] = std::shared_ptr<GGMLBlock>(new TextModel(config));
             if (enable_vision) {
-                blocks["visual"] = std::shared_ptr<GGMLBlock>(new VisionModel(llama_cpp_style, params.vision));
+                blocks["visual"] = std::shared_ptr<GGMLBlock>(new VisionModel(llama_cpp_style, config.vision));
             }
         }
 
@@ -1226,7 +1350,7 @@ namespace LLM {
     };
 
     struct LLMRunner : public GGMLRunner {
-        LLMParams params;
+        LLMConfig config;
         bool enable_vision;
         LLM model;
 
@@ -1242,7 +1366,7 @@ namespace LLM {
 
         static ggml_tensor* process_image_common(ggml_context* ctx,
                                                  ggml_tensor* image,
-                                                 const LLMVisionParams& vision_params) {
+                                                 const LLMVisionConfig& vision_params) {
             // image: [C, H, W]
             // return: [grid_t*(H/mh/ph)*(W/mw/pw)*mh*mw, C*pt*ph*pw], grid_t == 1
             int64_t C  = image->ne[2];
@@ -1337,7 +1461,7 @@ namespace LLM {
                                                 ggml_context* compute_ctx,
                                                 GGMLRunnerContext* runner_ctx,
                                                 ggml_tensor* image,
-                                                const LLMVisionParams& vision_params,
+                                                const LLMVisionConfig& vision_params,
                                                 std::shared_ptr<VisionModel> vision_model,
                                                 std::vector<int>& window_index_vec,
                                                 std::vector<int>& window_inverse_index_vec,
@@ -1448,140 +1572,29 @@ namespace LLM {
     public:
         LLMRunner(LLMArch arch,
                   ggml_backend_t backend,
-                  ggml_backend_t params_backend,
                   const String2TensorStorage& tensor_storage_map,
                   const std::string prefix,
-                  bool enable_vision_ = false)
-            : GGMLRunner(backend, params_backend), enable_vision(enable_vision_) {
-            params.arch = arch;
-            if (arch == LLMArch::MISTRAL_SMALL_3_2 || arch == LLMArch::MINISTRAL_3_3B) {
-                params.head_dim     = 128;
-                params.num_heads    = 32;
-                params.num_kv_heads = 8;
-                params.qkv_bias     = false;
-                params.rms_norm_eps = 1e-5f;
-            } else if (arch == LLMArch::QWEN3) {
-                params.head_dim     = 128;
-                params.num_heads    = 32;
-                params.num_kv_heads = 8;
-                params.qkv_bias     = false;
-                params.qk_norm      = true;
-                params.rms_norm_eps = 1e-6f;
-            } else if (arch == LLMArch::GEMMA3_12B) {
-                params.head_dim     = 256;
-                params.num_heads    = 16;
-                params.num_kv_heads = 8;
-                params.qkv_bias     = false;
-                params.qk_norm      = true;
-                params.rms_norm_eps = 1e-6f;
-                // llama.cpp adds +1 to Gemma3 norm.weight when exporting GGUF, so GGUF loading
-                // must keep rms_norm_add disabled here or the offset gets applied twice.
-                // Convenient for the converter, less convenient for whoever gets to debug it later.
-                params.rms_norm_add            = false;
-                params.normalize_input         = true;
-                params.max_position_embeddings = 131072;
-                params.mlp_activation          = MLPActivation::GELU_TANH;
-                params.rope_thetas             = {1000000.f, 10000.f};
-                params.rope_scales             = {8.f, 1.f};
-                params.sliding_attention       = {1024, 1024, 1024, 1024, 1024, 0};
-            } else if (arch == LLMArch::GEMMA2_2B) {
-                params.head_dim                = 256;
-                params.num_heads               = 8;
-                params.num_kv_heads            = 4;
-                params.qkv_bias                = false;
-                params.qk_norm                 = false;
-                params.rms_norm_eps            = 1e-6f;
-                params.rms_norm_add            = true;
-                params.normalize_input         = true;
-                params.max_position_embeddings = 8192;
-                params.mlp_activation          = MLPActivation::GELU_TANH;
-                params.hidden_size             = 2304;
-                params.intermediate_size       = 9216;
-                params.num_layers              = 26;
-                params.vocab_size              = 256000;
-            } else if (arch == LLMArch::GPT_OSS_20B) {
-                params.head_dim                = 64;
-                params.num_heads               = 64;
-                params.num_kv_heads            = 8;
-                params.qkv_bias                = true;
-                params.attention_out_bias      = true;
-                params.qk_norm                 = false;
-                params.rms_norm_eps            = 1e-5f;
-                params.hidden_size             = 2880;
-                params.intermediate_size       = 2880;
-                params.num_layers              = 24;
-                params.vocab_size              = 201088;
-                params.max_position_embeddings = 131072;
-                params.rope_thetas             = {150000.f};
-                params.rope_scales             = {32.f};
-                params.sliding_attention       = {128, 0};
-                params.num_experts             = 32;
-                params.num_experts_per_tok     = 4;
-            }
-            bool have_vision_weight = false;
-            bool llama_cpp_style    = false;
-            params.num_layers       = 0;
-            for (auto pair : tensor_storage_map) {
-                std::string tensor_name = pair.first;
-                if (tensor_name.find(prefix) == std::string::npos)
-                    continue;
-                size_t pos = tensor_name.find("visual.");
-                if (pos != std::string::npos) {
-                    have_vision_weight = true;
-                    if (contains(tensor_name, "attn.q_proj")) {
-                        llama_cpp_style = true;
-                    }
-                    continue;
-                }
-                pos = tensor_name.find("layers.");
-                if (pos != std::string::npos) {
-                    tensor_name = tensor_name.substr(pos);  // remove prefix
-                    auto items  = split_string(tensor_name, '.');
-                    if (items.size() > 1) {
-                        int block_index = atoi(items[1].c_str());
-                        if (block_index + 1 > params.num_layers) {
-                            params.num_layers = block_index + 1;
-                        }
-                    }
-                }
-                if (contains(tensor_name, "embed_tokens.weight")) {
-                    params.hidden_size = pair.second.ne[0];
-                    params.vocab_size  = pair.second.ne[1];
-                }
-                if (contains(tensor_name, "layers.0.mlp.gate_proj.weight")) {
-                    params.intermediate_size = pair.second.ne[1];
-                }
-                if (contains(tensor_name, "layers.0.mlp.experts.gate_up_proj.weight")) {
-                    params.intermediate_size = pair.second.ne[1] / 2;
-                }
-                if (contains(tensor_name, "layers.0.mlp.experts.gate_proj.weight")) {
-                    params.intermediate_size = pair.second.ne[1];
-                }
-            }
-            if (arch == LLMArch::QWEN3 && params.num_layers == 28) {  // Qwen3 2B
-                params.num_heads = 16;
-            }
-            LOG_DEBUG("llm: num_layers = %" PRId64 ", vocab_size = %" PRId64 ", hidden_size = %" PRId64 ", intermediate_size = %" PRId64,
-                      params.num_layers,
-                      params.vocab_size,
-                      params.hidden_size,
-                      params.intermediate_size);
-            if (enable_vision && !have_vision_weight) {
+                  bool enable_vision_                                 = false,
+                  std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : GGMLRunner(backend, weight_manager),
+              config(LLMConfig::detect_from_weights(tensor_storage_map, prefix, arch)),
+              enable_vision(enable_vision_) {
+            if (enable_vision && !config.have_vision_weight) {
                 LOG_WARN("no vision weights detected, vision disabled");
                 enable_vision = false;
             }
             if (enable_vision) {
                 LOG_DEBUG("enable llm vision");
-                if (llama_cpp_style) {
+                if (config.llama_cpp_style) {
                     LOG_DEBUG("llama.cpp style vision weight");
                 }
             }
-            model = LLM(params, enable_vision, llama_cpp_style);
+            model = LLM(config, enable_vision, config.llama_cpp_style);
             model.init(params_ctx, tensor_storage_map, prefix);
         }
 
         std::string get_desc() override {
-            return llm_arch_to_str[static_cast<int>(params.arch)];
+            return llm_arch_to_str[static_cast<int>(config.arch)];
         }
 
         void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
@@ -1633,12 +1646,12 @@ namespace LLM {
             }
 
             int64_t n_tokens = input_ids->ne[0];
-            if (params.arch == LLMArch::MISTRAL_SMALL_3_2 ||
-                params.arch == LLMArch::MINISTRAL_3_3B ||
-                params.arch == LLMArch::QWEN3 ||
-                params.arch == LLMArch::GEMMA3_12B ||
-                params.arch == LLMArch::GEMMA2_2B ||
-                params.arch == LLMArch::GPT_OSS_20B) {
+            if (config.arch == LLMArch::MISTRAL_SMALL_3_2 ||
+                config.arch == LLMArch::MINISTRAL_3_3B ||
+                config.arch == LLMArch::QWEN3 ||
+                config.arch == LLMArch::GEMMA3_12B ||
+                config.arch == LLMArch::GEMMA2_2B ||
+                config.arch == LLMArch::GPT_OSS_20B) {
                 input_pos_vec.resize(n_tokens);
                 for (int i = 0; i < n_tokens; ++i) {
                     input_pos_vec[i] = i;
@@ -1677,9 +1690,9 @@ namespace LLM {
                 set_backend_tensor_data(attention_mask, attention_mask_vec.data());
             }
 
-            if (params.arch == LLMArch::GEMMA3_12B || params.arch == LLMArch::GPT_OSS_20B) {
+            if (config.arch == LLMArch::GEMMA3_12B || config.arch == LLMArch::GPT_OSS_20B) {
                 int sliding_window = 0;
-                for (int window : params.sliding_attention) {
+                for (int window : config.sliding_attention) {
                     sliding_window = std::max(sliding_window, window);
                 }
                 sliding_attention_mask_vec.resize(n_tokens * n_tokens);
@@ -1721,7 +1734,10 @@ namespace LLM {
                                   const sd::Tensor<float>& attention_mask,
                                   const std::vector<std::pair<int, sd::Tensor<float>>>& image_embeds,
                                   std::set<int> out_layers,
-                                  bool return_all_hidden_states = false) {
+                                  bool return_all_hidden_states = false,
+                                  bool auto_free                = true,
+                                  bool free_compute_buffer      = true,
+                                  bool free_compute_params      = true) {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_graph(input_ids,
                                    attention_mask,
@@ -1729,21 +1745,21 @@ namespace LLM {
                                    out_layers,
                                    return_all_hidden_states);
             };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, true),
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params),
                                                    input_ids.dim() + 1);
         }
 
         int64_t get_num_image_tokens(int64_t t, int64_t h, int64_t w) {
             int64_t grid_t     = 1;
-            int64_t grid_h     = h / params.vision.patch_size;
-            int64_t grid_w     = w / params.vision.patch_size;
-            int64_t llm_grid_h = grid_h / params.vision.spatial_merge_size;
-            int64_t llm_grid_w = grid_w / params.vision.spatial_merge_size;
+            int64_t grid_h     = h / config.vision.patch_size;
+            int64_t grid_w     = w / config.vision.patch_size;
+            int64_t llm_grid_h = grid_h / config.vision.spatial_merge_size;
+            int64_t llm_grid_w = grid_w / config.vision.spatial_merge_size;
             return grid_t * grid_h * grid_w;
         }
 
         ggml_tensor* process_image(ggml_context* ctx, ggml_tensor* image) {
-            return process_image_common(ctx, image, params.vision);
+            return process_image_common(ctx, image, config.vision);
         }
 
         ggml_tensor* build_patch_pos_embeds(GGMLRunnerContext* runner_ctx,
@@ -1765,7 +1781,7 @@ namespace LLM {
                                        compute_ctx,
                                        runner_ctx,
                                        image,
-                                       params.vision,
+                                       config.vision,
                                        model.vision_model(),
                                        window_index_vec,
                                        window_inverse_index_vec,
@@ -1779,8 +1795,8 @@ namespace LLM {
             ggml_cgraph* gf    = new_graph_custom(LLM_GRAPH_SIZE);
             ggml_tensor* image = make_input(image_tensor);
 
-            GGML_ASSERT(image->ne[1] % (params.vision.patch_size * params.vision.spatial_merge_size) == 0);
-            GGML_ASSERT(image->ne[0] % (params.vision.patch_size * params.vision.spatial_merge_size) == 0);
+            GGML_ASSERT(image->ne[1] % (config.vision.patch_size * config.vision.spatial_merge_size) == 0);
+            GGML_ASSERT(image->ne[0] % (config.vision.patch_size * config.vision.spatial_merge_size) == 0);
 
             auto runnter_ctx           = get_context();
             ggml_tensor* hidden_states = encode_image(&runnter_ctx, image);
@@ -1790,11 +1806,14 @@ namespace LLM {
         }
 
         sd::Tensor<float> encode_image(const int n_threads,
-                                       const sd::Tensor<float>& image) {
+                                       const sd::Tensor<float>& image,
+                                       bool auto_free           = false,
+                                       bool free_compute_buffer = false,
+                                       bool free_compute_params = false) {
             auto get_graph = [&]() -> ggml_cgraph* {
                 return build_encode_image_graph(image);
             };
-            return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, false));
+            return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params));
         }
     };
 
@@ -1804,11 +1823,11 @@ namespace LLM {
 
         LLMEmbedder(LLMArch arch,
                     ggml_backend_t backend,
-                    ggml_backend_t params_backend,
-                    const String2TensorStorage& tensor_storage_map = {},
-                    const std::string prefix                       = "",
-                    bool enable_vision                             = false)
-            : model(arch, backend, params_backend, tensor_storage_map, prefix, enable_vision) {
+                    const String2TensorStorage& tensor_storage_map      = {},
+                    const std::string prefix                            = "",
+                    bool enable_vision                                  = false,
+                    std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : model(arch, backend, tensor_storage_map, prefix, enable_vision, weight_manager) {
             if (arch == LLMArch::MISTRAL_SMALL_3_2 || arch == LLMArch::MINISTRAL_3_3B) {
                 tokenizer = std::make_shared<MistralTokenizer>();
             } else if (arch == LLMArch::GPT_OSS_20B) {
@@ -1822,13 +1841,6 @@ namespace LLM {
             model.get_param_tensors(tensors, prefix);
         }
 
-        bool alloc_params_buffer() {
-            if (!model.alloc_params_buffer()) {
-                return false;
-            }
-            return true;
-        }
-
         std::tuple<std::vector<int>, std::vector<float>> tokenize(std::string text,
                                                                   std::pair<int, int> attn_range,
                                                                   size_t max_length = 0,
@@ -2044,7 +2056,8 @@ namespace LLM {
             ggml_backend_t backend    = sd_backend_cpu_init();
             ggml_type model_data_type = GGML_TYPE_COUNT;
 
-            ModelLoader model_loader;
+            auto model_manager        = std::make_shared<ModelManager>();
+            ModelLoader& model_loader = model_manager->loader();
             if (!model_loader.init_from_file_and_convert_name(file_path, "text_encoders.llm.")) {
                 LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
                 return;
@@ -2062,24 +2075,20 @@ namespace LLM {
             LLMArch arch = LLMArch::QWEN3;
 
             std::shared_ptr<LLMEmbedder> llm = std::make_shared<LLMEmbedder>(arch,
-                                                                             backend,
                                                                              backend,
                                                                              tensor_storage_map,
                                                                              "text_encoders.llm",
-                                                                             true);
+                                                                             true,
+                                                                             model_manager);
 
-            if (!llm->alloc_params_buffer()) {
-                LOG_ERROR("llm model allocation failed");
-                return;
-            }
-
-            std::map<std::string, ggml_tensor*> tensors;
-            llm->get_param_tensors(tensors, "text_encoders.llm");
-
-            bool success = model_loader.load_tensors(tensors);
-
-            if (!success) {
-                LOG_ERROR("load tensors from model loader failed");
+            if (!model_manager->register_runner_params("LLM test",
+                                                       *llm,
+                                                       "text_encoders.llm",
+                                                       ModelManager::ResidencyMode::ParamBackend,
+                                                       backend,
+                                                       backend) ||
+                !model_manager->validate_registered_tensors()) {
+                LOG_ERROR("register llm tensors with model manager failed");
                 return;
             }
 
@@ -2089,4 +2098,4 @@ namespace LLM {
     };
 };  // LLM
 
-#endif  // __LLM_HPP__
+#endif  // __SD_MODEL_TE_LLM_HPP__

src/model/te/t5.hpp

@@ -1,5 +1,5 @@
-#ifndef __T5_HPP__
-#define __T5_HPP__
+#ifndef __SD_MODEL_TE_T5_HPP__
+#define __SD_MODEL_TE_T5_HPP__
 
 #include <cfloat>
 #include <limits>
@@ -10,10 +10,33 @@
 #include <string>
 #include <unordered_map>
 
-#include "ggml_extend.hpp"
-#include "model.h"
+#include "core/ggml_extend.hpp"
+#include "model_loader.h"
+#include "model_manager.h"
 #include "tokenizers/t5_unigram_tokenizer.h"
 
+struct T5Config {
+    int64_t num_layers      = 24;
+    int64_t model_dim       = 4096;
+    int64_t ff_dim          = 10240;
+    int64_t num_heads       = 64;
+    int64_t vocab_size      = 32128;
+    bool relative_attention = true;
+
+    static T5Config detect_from_weights(const String2TensorStorage& tensor_storage_map,
+                                        const std::string& prefix,
+                                        bool is_umt5 = false) {
+        (void)tensor_storage_map;
+        (void)prefix;
+        T5Config config;
+        if (is_umt5) {
+            config.vocab_size         = 256384;
+            config.relative_attention = false;
+        }
+        return config;
+    }
+};
+
 class T5LayerNorm : public UnaryBlock {
 protected:
     int64_t hidden_size;
@@ -272,30 +295,21 @@ public:
     }
 };
 
-struct T5Params {
-    int64_t num_layers      = 24;
-    int64_t model_dim       = 4096;
-    int64_t ff_dim          = 10240;
-    int64_t num_heads       = 64;
-    int64_t vocab_size      = 32128;
-    bool relative_attention = true;
-};
-
 struct T5 : public GGMLBlock {
-    T5Params params;
+    T5Config config;
 
 public:
     T5() {}
-    T5(T5Params params)
-        : params(params) {
-        blocks["encoder"] = std::shared_ptr<GGMLBlock>(new T5Stack(params.num_layers,
-                                                                   params.model_dim,
-                                                                   params.model_dim,
-                                                                   params.ff_dim,
-                                                                   params.num_heads,
-                                                                   params.relative_attention));
-        blocks["shared"]  = std::shared_ptr<GGMLBlock>(new Embedding(params.vocab_size,
-                                                                     params.model_dim));
+    T5(T5Config config)
+        : config(config) {
+        blocks["encoder"] = std::shared_ptr<GGMLBlock>(new T5Stack(config.num_layers,
+                                                                   config.model_dim,
+                                                                   config.model_dim,
+                                                                   config.ff_dim,
+                                                                   config.num_heads,
+                                                                   config.relative_attention));
+        blocks["shared"]  = std::shared_ptr<GGMLBlock>(new Embedding(config.vocab_size,
+                                                                     config.model_dim));
     }
 
     ggml_tensor* forward(GGMLRunnerContext* ctx,
@@ -316,21 +330,18 @@ public:
 };
 
 struct T5Runner : public GGMLRunner {
-    T5Params params;
+    T5Config config;
     T5 model;
     std::vector<int> relative_position_bucket_vec;
 
     T5Runner(ggml_backend_t backend,
-             ggml_backend_t params_backend,
              const String2TensorStorage& tensor_storage_map,
              const std::string prefix,
-             bool is_umt5 = false)
-        : GGMLRunner(backend, params_backend) {
-        if (is_umt5) {
-            params.vocab_size         = 256384;
-            params.relative_attention = false;
-        }
-        model = T5(params);
+             bool is_umt5                                        = false,
+             std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+        : GGMLRunner(backend, weight_manager),
+          config(T5Config::detect_from_weights(tensor_storage_map, prefix, is_umt5)) {
+        model = T5(config);
         model.init(params_ctx, tensor_storage_map, prefix);
     }
 
@@ -384,11 +395,14 @@ struct T5Runner : public GGMLRunner {
 
     sd::Tensor<float> compute(const int n_threads,
                               const sd::Tensor<int32_t>& input_ids,
-                              const sd::Tensor<float>& attention_mask) {
+                              const sd::Tensor<float>& attention_mask,
+                              bool auto_free           = true,
+                              bool free_compute_buffer = true,
+                              bool free_compute_params = true) {
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(input_ids, attention_mask);
         };
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, true), 3);
+        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params), 3);
     }
 
     static std::vector<int> _relative_position_bucket(const std::vector<int>& relative_position,
@@ -464,24 +478,17 @@ struct T5Embedder {
     T5Runner model;
 
     T5Embedder(ggml_backend_t backend,
-               ggml_backend_t params_backend,
-               const String2TensorStorage& tensor_storage_map = {},
-               const std::string prefix                       = "",
-               bool is_umt5                                   = false)
-        : model(backend, params_backend, tensor_storage_map, prefix, is_umt5), tokenizer(is_umt5) {
+               const String2TensorStorage& tensor_storage_map      = {},
+               const std::string prefix                            = "",
+               bool is_umt5                                        = false,
+               std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+        : model(backend, tensor_storage_map, prefix, is_umt5, weight_manager), tokenizer(is_umt5) {
     }
 
     void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
         model.get_param_tensors(tensors, prefix);
     }
 
-    bool alloc_params_buffer() {
-        if (!model.alloc_params_buffer()) {
-            return false;
-        }
-        return true;
-    }
-
     std::tuple<std::vector<int>, std::vector<float>, std::vector<float>> tokenize(std::string text,
                                                                                   size_t max_length = 0,
                                                                                   bool padding      = false) {
@@ -566,7 +573,8 @@ struct T5Embedder {
         ggml_backend_t backend    = sd_backend_cpu_init();
         ggml_type model_data_type = GGML_TYPE_F16;
 
-        ModelLoader model_loader;
+        auto model_manager        = std::make_shared<ModelManager>();
+        ModelLoader& model_loader = model_manager->loader();
         if (!model_loader.init_from_file_and_convert_name(file_path)) {
             LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str());
             return;
@@ -579,19 +587,16 @@ struct T5Embedder {
             }
         }
 
-        std::shared_ptr<T5Embedder> t5 = std::make_shared<T5Embedder>(backend, backend, tensor_storage_map, "", true);
+        std::shared_ptr<T5Embedder> t5 = std::make_shared<T5Embedder>(backend, tensor_storage_map, "", true, model_manager);
 
-        if (!t5->alloc_params_buffer()) {
-            LOG_ERROR("t5 params buffer allocation failed");
-            return;
-        }
-        std::map<std::string, ggml_tensor*> tensors;
-        t5->get_param_tensors(tensors, "");
-
-        bool success = model_loader.load_tensors(tensors);
-
-        if (!success) {
-            LOG_ERROR("load tensors from model loader failed");
+        if (!model_manager->register_runner_params("T5 test",
+                                                   *t5,
+                                                   "",
+                                                   ModelManager::ResidencyMode::ParamBackend,
+                                                   backend,
+                                                   backend) ||
+            !model_manager->validate_registered_tensors()) {
+            LOG_ERROR("register t5 tensors with model manager failed");
             return;
         }
 
@@ -600,4 +605,4 @@ struct T5Embedder {
     }
 };
 
-#endif  // __T5_HPP__
+#endif  // __SD_MODEL_TE_T5_HPP__

src/model/upscaler/esrgan.hpp

@@ -0,0 +1,273 @@
+#ifndef __SD_MODEL_UPSCALER_ESRGAN_HPP__
+#define __SD_MODEL_UPSCALER_ESRGAN_HPP__
+
+#include <algorithm>
+#include <map>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "core/ggml_extend.hpp"
+#include "core/util.h"
+
+/*
+    ===================================    ESRGAN  ===================================
+    References:
+    https://github.com/xinntao/Real-ESRGAN/blob/master/inference_realesrgan.py
+    https://github.com/XPixelGroup/BasicSR/blob/v1.4.2/basicsr/archs/rrdbnet_arch.py
+
+*/
+
+struct ESRGANConfig {
+    int scale       = 4;
+    int num_block   = 23;
+    int num_in_ch   = 3;
+    int num_out_ch  = 3;
+    int num_feat    = 64;
+    int num_grow_ch = 32;
+
+    static ESRGANConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
+                                            const std::string& prefix = "") {
+        ESRGANConfig config;
+        auto find_weight = [&](const std::string& suffix) -> const TensorStorage* {
+            std::string name = prefix.empty() ? suffix : prefix + "." + suffix;
+            auto iter        = tensor_storage_map.find(name);
+            if (iter == tensor_storage_map.end()) {
+                return nullptr;
+            }
+            return &iter->second;
+        };
+
+        int detected_num_block        = 0;
+        const std::string body_prefix = prefix.empty() ? "body." : prefix + ".body.";
+        for (const auto& [name, _] : tensor_storage_map) {
+            if (!starts_with(name, body_prefix)) {
+                continue;
+            }
+            size_t pos = name.find('.', body_prefix.size());
+            if (pos == std::string::npos) {
+                continue;
+            }
+            try {
+                int idx            = std::stoi(name.substr(body_prefix.size(), pos - body_prefix.size()));
+                detected_num_block = std::max(detected_num_block, idx + 1);
+            } catch (...) {
+            }
+        }
+        if (detected_num_block > 0) {
+            config.num_block = detected_num_block;
+        }
+
+        bool has_conv_up2 = find_weight("conv_up2.weight") != nullptr;
+        bool has_conv_up1 = find_weight("conv_up1.weight") != nullptr;
+        bool has_model_tensor =
+            detected_num_block > 0 ||
+            find_weight("conv_first.weight") != nullptr ||
+            find_weight("conv_hr.weight") != nullptr ||
+            find_weight("conv_last.weight") != nullptr;
+        if (has_conv_up2) {
+            config.scale = 4;
+        } else if (has_conv_up1) {
+            config.scale = 2;
+        } else if (has_model_tensor) {
+            config.scale = 1;
+        }
+
+        if (has_model_tensor || has_conv_up1 || has_conv_up2) {
+            LOG_DEBUG("esrgan: scale = %d, num_block = %d, num_in_ch = %d, num_out_ch = %d, num_feat = %d, num_grow_ch = %d",
+                      config.scale,
+                      config.num_block,
+                      config.num_in_ch,
+                      config.num_out_ch,
+                      config.num_feat,
+                      config.num_grow_ch);
+        }
+        return config;
+    }
+};
+
+class ResidualDenseBlock : public GGMLBlock {
+protected:
+    int num_feat;
+    int num_grow_ch;
+
+public:
+    ResidualDenseBlock(int num_feat = 64, int num_grow_ch = 32)
+        : num_feat(num_feat), num_grow_ch(num_grow_ch) {
+        blocks["conv1"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat, num_grow_ch, {3, 3}, {1, 1}, {1, 1}));
+        blocks["conv2"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat + num_grow_ch, num_grow_ch, {3, 3}, {1, 1}, {1, 1}));
+        blocks["conv3"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat + 2 * num_grow_ch, num_grow_ch, {3, 3}, {1, 1}, {1, 1}));
+        blocks["conv4"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat + 3 * num_grow_ch, num_grow_ch, {3, 3}, {1, 1}, {1, 1}));
+        blocks["conv5"] = std::shared_ptr<GGMLBlock>(new Conv2d(num_feat + 4 * num_grow_ch, num_feat, {3, 3}, {1, 1}, {1, 1}));
+    }
+
+    ggml_tensor* lrelu(GGMLRunnerContext* ctx, ggml_tensor* x) {
+        return ggml_leaky_relu(ctx->ggml_ctx, x, 0.2f, true);
+    }
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+        // x: [n, num_feat, h, w]
+        // return: [n, num_feat, h, w]
+
+        auto conv1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv1"]);
+        auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["conv2"]);
+        auto conv3 = std::dynamic_pointer_cast<Conv2d>(blocks["conv3"]);
+        auto conv4 = std::dynamic_pointer_cast<Conv2d>(blocks["conv4"]);
+        auto conv5 = std::dynamic_pointer_cast<Conv2d>(blocks["conv5"]);
+
+        auto x1    = lrelu(ctx, conv1->forward(ctx, x));
+        auto x_cat = ggml_concat(ctx->ggml_ctx, x, x1, 2);
+        auto x2    = lrelu(ctx, conv2->forward(ctx, x_cat));
+        x_cat      = ggml_concat(ctx->ggml_ctx, x_cat, x2, 2);
+        auto x3    = lrelu(ctx, conv3->forward(ctx, x_cat));
+        x_cat      = ggml_concat(ctx->ggml_ctx, x_cat, x3, 2);
+        auto x4    = lrelu(ctx, conv4->forward(ctx, x_cat));
+        x_cat      = ggml_concat(ctx->ggml_ctx, x_cat, x4, 2);
+        auto x5    = conv5->forward(ctx, x_cat);
+
+        x5 = ggml_add(ctx->ggml_ctx, ggml_ext_scale(ctx->ggml_ctx, x5, 0.2f), x);
+        return x5;
+    }
+};
+
+class RRDB : public GGMLBlock {
+public:
+    RRDB(int num_feat, int num_grow_ch = 32) {
+        blocks["rdb1"] = std::shared_ptr<GGMLBlock>(new ResidualDenseBlock(num_feat, num_grow_ch));
+        blocks["rdb2"] = std::shared_ptr<GGMLBlock>(new ResidualDenseBlock(num_feat, num_grow_ch));
+        blocks["rdb3"] = std::shared_ptr<GGMLBlock>(new ResidualDenseBlock(num_feat, num_grow_ch));
+    }
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+        // x: [n, num_feat, h, w]
+        // return: [n, num_feat, h, w]
+
+        auto rdb1 = std::dynamic_pointer_cast<ResidualDenseBlock>(blocks["rdb1"]);
+        auto rdb2 = std::dynamic_pointer_cast<ResidualDenseBlock>(blocks["rdb2"]);
+        auto rdb3 = std::dynamic_pointer_cast<ResidualDenseBlock>(blocks["rdb3"]);
+
+        auto out = rdb1->forward(ctx, x);
+        out      = rdb2->forward(ctx, out);
+        out      = rdb3->forward(ctx, out);
+
+        out = ggml_add(ctx->ggml_ctx, ggml_ext_scale(ctx->ggml_ctx, out, 0.2f), x);
+        return out;
+    }
+};
+
+class RRDBNet : public GGMLBlock {
+protected:
+    ESRGANConfig config;
+
+public:
+    explicit RRDBNet(ESRGANConfig config)
+        : config(std::move(config)) {
+        blocks["conv_first"] = std::shared_ptr<GGMLBlock>(new Conv2d(this->config.num_in_ch, this->config.num_feat, {3, 3}, {1, 1}, {1, 1}));
+        for (int i = 0; i < this->config.num_block; i++) {
+            std::string name = "body." + std::to_string(i);
+            blocks[name]     = std::shared_ptr<GGMLBlock>(new RRDB(this->config.num_feat, this->config.num_grow_ch));
+        }
+        blocks["conv_body"] = std::shared_ptr<GGMLBlock>(new Conv2d(this->config.num_feat, this->config.num_feat, {3, 3}, {1, 1}, {1, 1}));
+        if (this->config.scale >= 2) {
+            blocks["conv_up1"] = std::shared_ptr<GGMLBlock>(new Conv2d(this->config.num_feat, this->config.num_feat, {3, 3}, {1, 1}, {1, 1}));
+        }
+        if (this->config.scale == 4) {
+            blocks["conv_up2"] = std::shared_ptr<GGMLBlock>(new Conv2d(this->config.num_feat, this->config.num_feat, {3, 3}, {1, 1}, {1, 1}));
+        }
+        blocks["conv_hr"]   = std::shared_ptr<GGMLBlock>(new Conv2d(this->config.num_feat, this->config.num_feat, {3, 3}, {1, 1}, {1, 1}));
+        blocks["conv_last"] = std::shared_ptr<GGMLBlock>(new Conv2d(this->config.num_feat, this->config.num_out_ch, {3, 3}, {1, 1}, {1, 1}));
+    }
+
+    int get_scale() { return config.scale; }
+    int get_num_block() { return config.num_block; }
+
+    ggml_tensor* lrelu(GGMLRunnerContext* ctx, ggml_tensor* x) {
+        return ggml_leaky_relu(ctx->ggml_ctx, x, 0.2f, true);
+    }
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+        // x: [n, num_in_ch, h, w]
+        // return: [n, num_out_ch, h*scale, w*scale]
+        auto conv_first = std::dynamic_pointer_cast<Conv2d>(blocks["conv_first"]);
+        auto conv_body  = std::dynamic_pointer_cast<Conv2d>(blocks["conv_body"]);
+        auto conv_hr    = std::dynamic_pointer_cast<Conv2d>(blocks["conv_hr"]);
+        auto conv_last  = std::dynamic_pointer_cast<Conv2d>(blocks["conv_last"]);
+
+        auto feat = conv_first->forward(ctx, x);
+        sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.prelude", "feat");
+        auto body_feat = feat;
+        for (int i = 0; i < config.num_block; i++) {
+            std::string name = "body." + std::to_string(i);
+            auto block       = std::dynamic_pointer_cast<RRDB>(blocks[name]);
+
+            body_feat = block->forward(ctx, body_feat);
+            sd::ggml_graph_cut::mark_graph_cut(body_feat, "esrgan.body." + std::to_string(i), "feat");
+        }
+        body_feat = conv_body->forward(ctx, body_feat);
+        feat      = ggml_add(ctx->ggml_ctx, feat, body_feat);
+        sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.body.out", "feat");
+        // upsample
+        if (config.scale >= 2) {
+            auto conv_up1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv_up1"]);
+            feat          = lrelu(ctx, conv_up1->forward(ctx, ggml_upscale(ctx->ggml_ctx, feat, 2, GGML_SCALE_MODE_NEAREST)));
+            sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.up1", "feat");
+            if (config.scale == 4) {
+                auto conv_up2 = std::dynamic_pointer_cast<Conv2d>(blocks["conv_up2"]);
+                feat          = lrelu(ctx, conv_up2->forward(ctx, ggml_upscale(ctx->ggml_ctx, feat, 2, GGML_SCALE_MODE_NEAREST)));
+                sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.up2", "feat");
+            }
+        }
+        // for all scales
+        auto out = conv_last->forward(ctx, lrelu(ctx, conv_hr->forward(ctx, feat)));
+        sd::ggml_graph_cut::mark_graph_cut(out, "esrgan.final", "out");
+        return out;
+    }
+};
+
+struct ESRGAN : public GGMLRunner {
+    ESRGANConfig config;
+    std::unique_ptr<RRDBNet> rrdb_net;
+
+    ESRGAN(ggml_backend_t backend,
+           const String2TensorStorage& tensor_storage_map      = {},
+           std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+        : GGMLRunner(backend, weight_manager),
+          config(ESRGANConfig::detect_from_weights(tensor_storage_map)),
+          rrdb_net(std::make_unique<RRDBNet>(config)) {
+        rrdb_net->init(params_ctx, tensor_storage_map, "");
+    }
+
+    std::string get_desc() override {
+        return "esrgan";
+    }
+
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
+        if (!rrdb_net) {
+            return;
+        }
+
+        rrdb_net->get_param_tensors(tensors);
+    }
+
+    ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor) {
+        if (!rrdb_net)
+            return nullptr;
+        constexpr int kGraphNodes = 1 << 16;  // 65k
+        ggml_cgraph* gf           = new_graph_custom(kGraphNodes);
+        ggml_tensor* x            = make_input(x_tensor);
+
+        auto runner_ctx  = get_context();
+        ggml_tensor* out = rrdb_net->forward(&runner_ctx, x);
+        ggml_build_forward_expand(gf, out);
+        return gf;
+    }
+
+    sd::Tensor<float> compute(const int n_threads,
+                              const sd::Tensor<float>& x) {
+        auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x); };
+        auto result    = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
+        return result;
+    }
+};
+
+#endif  // __SD_MODEL_UPSCALER_ESRGAN_HPP__

src/model/upscaler/ltx_latent_upscaler.hpp

@@ -1,9 +1,9 @@
-#ifndef __SD_LTX_LATENT_UPSCALER_HPP__
-#define __SD_LTX_LATENT_UPSCALER_HPP__
+#ifndef __SD_MODEL_UPSCALER_LTX_LATENT_UPSCALER_HPP__
+#define __SD_MODEL_UPSCALER_LTX_LATENT_UPSCALER_HPP__
 
+#include <algorithm>
 #include <cinttypes>
 #include <cmath>
-#include <cstdlib>
 #include <map>
 #include <memory>
 #include <set>
@@ -11,11 +11,11 @@
 #include <utility>
 #include <vector>
 
-#include "common_dit.hpp"
-#include "ggml_extend.hpp"
-#include "ggml_graph_cut.h"
-#include "model.h"
-#include "util.h"
+#include "core/ggml_extend.hpp"
+#include "core/ggml_graph_cut.h"
+#include "core/util.h"
+#include "model/diffusion/dit.hpp"
+#include "model_loader.h"
 
 namespace LTXVUpsampler {
     constexpr int LTX_UPSAMPLER_GRAPH_SIZE = 10240;
@@ -32,90 +32,100 @@ namespace LTXVUpsampler {
         int spatial_up_num       = 2;
         int spatial_down_den     = 1;
         int temporal_up_factor   = 1;
-    };
 
-    static inline bool has_tensor(const String2TensorStorage& tensor_storage_map,
-                                  const std::string& name) {
-        return tensor_storage_map.find(name) != tensor_storage_map.end();
-    }
+        static LatentUpsamplerConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
+                                                         const std::string& prefix = "") {
+            LatentUpsamplerConfig config;
+            auto find_weight = [&](const std::string& suffix) -> const TensorStorage* {
+                std::string name = prefix.empty() ? suffix : prefix + "." + suffix;
+                auto iter        = tensor_storage_map.find(name);
+                if (iter == tensor_storage_map.end()) {
+                    return nullptr;
+                }
+                return &iter->second;
+            };
 
-    static inline int64_t get_tensor_ne(const String2TensorStorage& tensor_storage_map,
-                                        const std::string& name,
-                                        int axis,
-                                        int64_t fallback) {
-        auto it = tensor_storage_map.find(name);
-        if (it == tensor_storage_map.end() || axis < 0 || axis >= GGML_MAX_DIMS) {
-            return fallback;
-        }
-        return it->second.ne[axis];
-    }
+            bool inferred = false;
 
-    static inline int64_t get_tensor_ne0(const String2TensorStorage& tensor_storage_map,
-                                         const std::string& name,
-                                         int64_t fallback) {
-        return get_tensor_ne(tensor_storage_map, name, 0, fallback);
-    }
-
-    static inline int count_module_blocks(const String2TensorStorage& tensor_storage_map,
-                                          const std::string& module_name) {
-        int max_block            = -1;
-        const std::string prefix = module_name + ".";
-        for (const auto& pair : tensor_storage_map) {
-            const std::string& name = pair.first;
-            if (name.find(prefix) != 0) {
-                continue;
+            const TensorStorage* initial_norm = find_weight("initial_norm.weight");
+            if (initial_norm != nullptr) {
+                config.mid_channels = initial_norm->ne[0];
+                inferred            = true;
             }
-            size_t begin = prefix.size();
-            size_t end   = name.find('.', begin);
-            if (end == std::string::npos) {
-                continue;
-            }
-            int index = atoi(name.substr(begin, end - begin).c_str());
-            max_block = std::max(max_block, index);
-        }
-        return max_block + 1;
-    }
 
-    static inline LatentUpsamplerConfig detect_config_from_weights(const String2TensorStorage& tensor_storage_map) {
-        LatentUpsamplerConfig config;
-        config.mid_channels = get_tensor_ne0(tensor_storage_map, "initial_norm.weight", config.mid_channels);
-        config.in_channels  = get_tensor_ne0(tensor_storage_map, "final_conv.bias", config.in_channels);
-        int detected_blocks = count_module_blocks(tensor_storage_map, "res_blocks");
-        if (detected_blocks > 0) {
-            config.num_blocks_per_stage = detected_blocks;
-        }
-        config.rational_resampler      = has_tensor(tensor_storage_map, "upsampler.conv.weight");
-        int64_t upsampler_out_channels = get_tensor_ne0(tensor_storage_map, "upsampler.0.bias", 0);
-        config.spatial_upsample        = config.rational_resampler || upsampler_out_channels == 4 * config.mid_channels;
-        config.temporal_upsample       = upsampler_out_channels == 2 * config.mid_channels;
-        if (config.temporal_upsample) {
-            config.temporal_up_factor = 2;
-        }
-        if (config.rational_resampler) {
-            int64_t out_channels = get_tensor_ne(tensor_storage_map,
-                                                 "upsampler.conv.weight",
-                                                 3,
-                                                 config.mid_channels * 9);
-            if (config.mid_channels > 0 && out_channels % config.mid_channels == 0) {
-                int64_t ratio = out_channels / config.mid_channels;
-                int num       = static_cast<int>(std::round(std::sqrt(static_cast<double>(ratio))));
-                if (num > 0 && static_cast<int64_t>(num) * num == ratio) {
-                    config.spatial_up_num = num;
+            const TensorStorage* final_conv = find_weight("final_conv.bias");
+            if (final_conv != nullptr) {
+                config.in_channels = final_conv->ne[0];
+                inferred           = true;
+            }
+
+            int detected_blocks                 = 0;
+            const std::string res_blocks_prefix = prefix.empty() ? "res_blocks." : prefix + ".res_blocks.";
+            for (const auto& [name, _] : tensor_storage_map) {
+                if (!starts_with(name, res_blocks_prefix)) {
+                    continue;
+                }
+                size_t begin = res_blocks_prefix.size();
+                size_t end   = name.find('.', begin);
+                if (end == std::string::npos) {
+                    continue;
+                }
+                try {
+                    int idx         = std::stoi(name.substr(begin, end - begin));
+                    detected_blocks = std::max(detected_blocks, idx + 1);
+                } catch (...) {
                 }
             }
-            if (config.spatial_up_num == 3) {
-                config.spatial_down_den = 2;
-                config.spatial_scale    = 1.5f;
-            } else if (config.spatial_up_num == 4) {
-                config.spatial_down_den = 1;
-                config.spatial_scale    = 4.f;
-            } else {
-                config.spatial_down_den = 1;
-                config.spatial_scale    = static_cast<float>(config.spatial_up_num);
+            if (detected_blocks > 0) {
+                config.num_blocks_per_stage = detected_blocks;
+                inferred                    = true;
             }
+
+            const TensorStorage* rational_upsampler_weight = find_weight("upsampler.conv.weight");
+            const TensorStorage* upsampler_bias            = find_weight("upsampler.0.bias");
+            config.rational_resampler                      = rational_upsampler_weight != nullptr;
+            int64_t upsampler_out_channels                 = upsampler_bias == nullptr ? 0 : upsampler_bias->ne[0];
+            config.spatial_upsample                        = config.rational_resampler || upsampler_out_channels == 4 * config.mid_channels;
+            config.temporal_upsample                       = upsampler_out_channels == 2 * config.mid_channels;
+            if (config.rational_resampler || upsampler_out_channels > 0) {
+                inferred = true;
+            }
+            if (config.temporal_upsample) {
+                config.temporal_up_factor = 2;
+            }
+            if (rational_upsampler_weight != nullptr) {
+                int64_t out_channels = rational_upsampler_weight->ne[3];
+                if (config.mid_channels > 0 && out_channels % config.mid_channels == 0) {
+                    int64_t ratio = out_channels / config.mid_channels;
+                    int num       = static_cast<int>(std::round(std::sqrt(static_cast<double>(ratio))));
+                    if (num > 0 && static_cast<int64_t>(num) * num == ratio) {
+                        config.spatial_up_num = num;
+                    }
+                }
+                if (config.spatial_up_num == 3) {
+                    config.spatial_down_den = 2;
+                    config.spatial_scale    = 1.5f;
+                } else if (config.spatial_up_num == 4) {
+                    config.spatial_down_den = 1;
+                    config.spatial_scale    = 4.f;
+                } else {
+                    config.spatial_down_den = 1;
+                    config.spatial_scale    = static_cast<float>(config.spatial_up_num);
+                }
+            }
+
+            if (inferred) {
+                LOG_DEBUG("ltx latent upsampler: in_channels = %" PRId64 ", mid_channels = %" PRId64 ", num_blocks_per_stage = %d, spatial_scale = %.3f, temporal_up_factor = %d, rational_resampler = %d",
+                          config.in_channels,
+                          config.mid_channels,
+                          config.num_blocks_per_stage,
+                          config.spatial_scale,
+                          config.temporal_up_factor,
+                          config.rational_resampler);
+            }
+            return config;
         }
-        return config;
-    }
+    };
 
     class VideoGroupNorm : public GGMLBlock {
     protected:
@@ -240,20 +250,25 @@ namespace LTXVUpsampler {
     protected:
         int64_t channels;
         int stride;
-        ggml_tensor* kernel = nullptr;
         std::vector<float> kernel_data;
+        std::string kernel_name;
 
         void init_params(ggml_context* ctx,
                          const String2TensorStorage& tensor_storage_map = {},
                          const std::string prefix                       = "") override {
+            SD_UNUSED(ctx);
             SD_UNUSED(tensor_storage_map);
             if (stride == 1) {
                 return;
             }
-            kernel           = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 5, 5, 1, channels);
-            std::string name = prefix + "kernel";
-            ggml_set_name(kernel, name.c_str());
+            kernel_name = prefix + "kernel";
+        }
 
+    public:
+        BlurDownsample(int64_t channels, int stride)
+            : channels(channels),
+              stride(stride) {
+            GGML_ASSERT(stride >= 1);
             static const float binomial[5] = {1.f, 4.f, 6.f, 4.f, 1.f};
             kernel_data.resize(static_cast<size_t>(5 * 5 * channels));
             for (int64_t c = 0; c < channels; ++c) {
@@ -266,26 +281,16 @@ namespace LTXVUpsampler {
             }
         }
 
-    public:
-        BlurDownsample(int64_t channels, int stride)
-            : channels(channels),
-              stride(stride) {
-            GGML_ASSERT(stride >= 1);
-        }
-
-        void load_fixed_tensors() {
-            if (kernel == nullptr || kernel_data.empty()) {
-                return;
-            }
-            ggml_backend_tensor_set(kernel, kernel_data.data(), 0, kernel_data.size() * sizeof(float));
-        }
-
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
             if (stride == 1) {
                 return x;
             }
-            GGML_ASSERT(kernel != nullptr);
+            GGML_ASSERT(ctx != nullptr);
+            GGML_ASSERT(!kernel_data.empty());
             GGML_ASSERT(x->ne[2] == channels);
+            ggml_tensor* kernel = ggml_new_tensor_4d(ctx->ggml_ctx, GGML_TYPE_F32, 5, 5, 1, channels);
+            ggml_set_name(kernel, kernel_name.empty() ? "blur_down.kernel" : kernel_name.c_str());
+            ctx->bind_backend_tensor_data(kernel, kernel_data.data());
             if (ctx->conv2d_direct_enabled) {
                 return ggml_conv_2d_dw_direct(ctx->ggml_ctx, kernel, x, stride, stride, 2, 2, 1, 1);
             }
@@ -311,11 +316,6 @@ namespace LTXVUpsampler {
             blocks["blur_down"]     = std::shared_ptr<GGMLBlock>(new BlurDownsample(mid_channels, den));
         }
 
-        void load_fixed_tensors() {
-            auto blur_down = std::dynamic_pointer_cast<BlurDownsample>(blocks["blur_down"]);
-            blur_down->load_fixed_tensors();
-        }
-
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
             auto conv          = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
             auto pixel_shuffle = std::dynamic_pointer_cast<PixelShuffleND>(blocks["pixel_shuffle"]);
@@ -426,45 +426,17 @@ namespace LTXVUpsampler {
             sd::ggml_graph_cut::mark_graph_cut(x, "ltx_latent_upsampler.final", "x");
             return x;
         }
-
-        void load_fixed_tensors() {
-            if (!config.rational_resampler) {
-                return;
-            }
-            auto upsampler = std::dynamic_pointer_cast<SpatialRationalResampler>(blocks["upsampler"]);
-            upsampler->load_fixed_tensors();
-        }
     };
 
     struct LatentUpsamplerRunner : public GGMLRunner {
+        LatentUpsamplerConfig config;
         std::unique_ptr<LatentUpsampler> model;
 
         LatentUpsamplerRunner(ggml_backend_t backend,
-                              ggml_backend_t params_backend)
-            : GGMLRunner(backend, params_backend) {}
-
-        std::string get_desc() override {
-            return "ltx_latent_upsampler";
-        }
-
-        bool load_from_file(const std::string& file_path, int n_threads) {
-            LOG_INFO("loading LTX latent upsampler from '%s'", file_path.c_str());
-            ModelLoader model_loader;
-            if (!model_loader.init_from_file(file_path)) {
-                LOG_ERROR("init LTX latent upsampler model loader from file failed: '%s'", file_path.c_str());
-                return false;
-            }
-
-            const auto& tensor_storage_map = model_loader.get_tensor_storage_map();
-            bool has_regular_upsampler     = has_tensor(tensor_storage_map, "upsampler.0.weight");
-            bool has_rational_spatial      = has_tensor(tensor_storage_map, "upsampler.conv.weight");
-            if (!has_tensor(tensor_storage_map, "post_upsample_res_blocks.0.conv2.bias") ||
-                (!has_regular_upsampler && !has_rational_spatial)) {
-                LOG_ERROR("unsupported LTX latent upsampler weights: expected upsampler tensors");
-                return false;
-            }
-
-            LatentUpsamplerConfig config = detect_config_from_weights(tensor_storage_map);
+                              const String2TensorStorage& tensor_storage_map,
+                              std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : GGMLRunner(backend, weight_manager),
+              config(LatentUpsamplerConfig::detect_from_weights(tensor_storage_map)) {
             if (config.dims != 3 || (!config.spatial_upsample && !config.temporal_upsample) ||
                 config.spatial_up_num < 1 || config.spatial_down_den < 1 || config.temporal_up_factor < 1) {
                 LOG_ERROR("unsupported LTX latent upsampler config: dims=%d spatial=%d temporal=%d rational=%d scale=%.3f temporal_factor=%d",
@@ -474,36 +446,21 @@ namespace LTXVUpsampler {
                           config.rational_resampler,
                           config.spatial_scale,
                           config.temporal_up_factor);
-                return false;
+                return;
             }
 
             model = std::make_unique<LatentUpsampler>(config);
             model->init(params_ctx, tensor_storage_map, "");
-            if (!alloc_params_buffer()) {
-                LOG_ERROR("LTX latent upsampler params buffer allocation failed");
-                return false;
-            }
+        }
 
-            std::map<std::string, ggml_tensor*> tensors;
-            model->get_param_tensors(tensors);
-            std::set<std::string> ignore_tensors;
-            if (config.rational_resampler) {
-                ignore_tensors.insert("upsampler.blur_down.kernel");
-            }
-            if (!model_loader.load_tensors(tensors, ignore_tensors, n_threads)) {
-                LOG_ERROR("load LTX latent upsampler tensors failed");
-                return false;
-            }
-            model->load_fixed_tensors();
+        std::string get_desc() override {
+            return "ltx_latent_upsampler";
+        }
 
-            LOG_INFO("LTX latent upsampler loaded: in_channels=%" PRId64 ", mid_channels=%" PRId64 ", blocks=%d, scale=%.3f, temporal_factor=%d, rational=%d",
-                     config.in_channels,
-                     config.mid_channels,
-                     config.num_blocks_per_stage,
-                     config.spatial_scale,
-                     config.temporal_up_factor,
-                     config.rational_resampler);
-            return true;
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
+            if (model) {
+                model->get_param_tensors(tensors);
+            }
         }
 
         ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor) {
@@ -534,18 +491,18 @@ namespace LTXVUpsampler {
                           (long long)x.shape()[4]);
                 return {};
             }
-            if (x.shape()[3] != model->config.in_channels) {
+            if (x.shape()[3] != config.in_channels) {
                 LOG_ERROR("LTX latent upsampler expected %" PRId64 " channels, got %lld",
-                          model->config.in_channels,
+                          config.in_channels,
                           (long long)x.shape()[3]);
                 return {};
             }
             size_t expected_dim = static_cast<size_t>(x.dim());
             auto get_graph      = [&]() -> ggml_cgraph* { return build_graph(x); };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), expected_dim);
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), expected_dim);
         }
     };
 
 }  // namespace LTXVUpsampler
 
-#endif  // __SD_LTX_LATENT_UPSCALER_HPP__
+#endif  // __SD_MODEL_UPSCALER_LTX_LATENT_UPSCALER_HPP__

src/model/vae/auto_encoder_kl.hpp

@@ -1,7 +1,7 @@
-#ifndef __AUTO_ENCODER_KL_HPP__
-#define __AUTO_ENCODER_KL_HPP__
+#ifndef __SD_MODEL_VAE_AUTO_ENCODER_KL_HPP__
+#define __SD_MODEL_VAE_AUTO_ENCODER_KL_HPP__
 
-#include "vae.hpp"
+#include "model/vae/vae.hpp"
 
 /*================================================== AutoEncoderKL ===================================================*/
 
@@ -213,9 +213,9 @@ protected:
         params["mix_factor"] = ggml_new_tensor_1d(ctx, wtype, 1);
     }
 
-    float get_alpha() {
-        float alpha = ggml_ext_backend_tensor_get_f32(params["mix_factor"]);
-        return sigmoid(alpha);
+    ggml_tensor* get_alpha(GGMLRunnerContext* ctx) {
+        auto mix_factor = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, params["mix_factor"]);
+        return ggml_sigmoid(ctx->ggml_ctx, mix_factor);
     }
 
 public:
@@ -250,10 +250,12 @@ public:
 
         x = time_stack->forward(ctx, x);  // b t c (h w)
 
-        float alpha = get_alpha();
-        x           = ggml_add(ctx->ggml_ctx,
-                               ggml_ext_scale(ctx->ggml_ctx, x, alpha),
-                               ggml_ext_scale(ctx->ggml_ctx, x_mix, 1.0f - alpha));
+        auto alpha = get_alpha(ctx);
+        x          = ggml_add(ctx->ggml_ctx,
+                              x_mix,
+                              ggml_mul(ctx->ggml_ctx,
+                                       ggml_sub(ctx->ggml_ctx, x, x_mix),
+                                       alpha));
 
         x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));  // b c t (h w) -> b t c (h w)
         x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H, C, T * B);                     // b t c (h w) -> (b t) c h w
@@ -664,13 +666,13 @@ struct AutoEncoderKL : public VAE {
     AutoEncoderKLModel ae;
 
     AutoEncoderKL(ggml_backend_t backend,
-                  ggml_backend_t params_backend,
                   const String2TensorStorage& tensor_storage_map,
                   const std::string prefix,
-                  bool decode_only       = false,
-                  bool use_video_decoder = false,
-                  SDVersion version      = VERSION_SD1)
-        : decode_only(decode_only), VAE(version, backend, params_backend) {
+                  bool decode_only                                    = false,
+                  bool use_video_decoder                              = false,
+                  SDVersion version                                   = VERSION_SD1,
+                  std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+        : VAE(version, backend, prefix, weight_manager), decode_only(decode_only) {
         if (sd_version_is_sd1(version) || sd_version_is_sd2(version)) {
             scale_factor = 0.18215f;
             shift_factor = 0.f;
@@ -718,8 +720,8 @@ struct AutoEncoderKL : public VAE {
         return "vae";
     }
 
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) override {
-        ae.get_param_tensors(tensors, prefix);
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        ae.get_param_tensors(tensors, weight_prefix);
     }
 
     ggml_cgraph* build_graph(const sd::Tensor<float>& z_tensor, bool decode_graph) {
@@ -742,7 +744,7 @@ struct AutoEncoderKL : public VAE {
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(z, decode_graph);
         };
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), z.dim());
+        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), z.dim());
     }
 
     sd::Tensor<float> gaussian_latent_sample(const sd::Tensor<float>& moments, std::shared_ptr<RNG> rng) {
@@ -886,4 +888,4 @@ struct AutoEncoderKL : public VAE {
     };
 };
 
-#endif  // __AUTO_ENCODER_KL_HPP__
+#endif  // __SD_MODEL_VAE_AUTO_ENCODER_KL_HPP__

src/model/vae/ltx_audio_vae.hpp

@@ -1,5 +1,5 @@
-#ifndef __SD_LTX_AUDIO_VAE_H__
-#define __SD_LTX_AUDIO_VAE_H__
+#ifndef __SD_MODEL_VAE_LTX_AUDIO_VAE_HPP__
+#define __SD_MODEL_VAE_LTX_AUDIO_VAE_HPP__
 
 #include <cmath>
 #include <limits>
@@ -7,7 +7,9 @@
 #include <string>
 #include <vector>
 
-#include "ggml_extend.hpp"
+#include "core/ggml_extend.hpp"
+#include "model_loader.h"
+#include "model_manager.h"
 
 namespace LTXV {
 
@@ -58,11 +60,12 @@ namespace LTXV {
             return base_output_sample_rate();
         }
 
-        static LTXAudioVAEConfig detect_from_weights(const String2TensorStorage& tensor_storage_map) {
+        static LTXAudioVAEConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix = "") {
             LTXAudioVAEConfig config;
 
             auto require = [&](const std::string& name) -> const TensorStorage* {
-                auto iter = tensor_storage_map.find(name);
+                std::string tensor_name = prefix.empty() ? name : prefix + "." + name;
+                auto iter               = tensor_storage_map.find(tensor_name);
                 if (iter == tensor_storage_map.end()) {
                     return nullptr;
                 }
@@ -168,6 +171,12 @@ namespace LTXV {
             if (config.audio_channels != 2 || config.latent_channels != 8 || config.mel_bins != 64) {
                 return config;
             }
+            LOG_DEBUG("ltx_audio_vae: sample_rate = %d, mel_bins = %d, latent_channels = %d, latent_frequency_bins = %d, has_bwe = %s",
+                      config.sample_rate,
+                      config.mel_bins,
+                      config.latent_channels,
+                      config.latent_frequency_bins,
+                      config.has_bwe ? "true" : "false");
             return config;
         }
     };
@@ -989,13 +998,15 @@ namespace LTXV {
     struct LTXAudioVAERunner : public GGMLRunner {
         LTXAudioVAEConfig config;
         LTXAudioVAE model;
+        std::string weight_prefix;
         sd::Tensor<float> bwe_skip_filter_tensor;
 
         LTXAudioVAERunner(ggml_backend_t backend,
-                          ggml_backend_t params_backend,
                           const String2TensorStorage& tensor_storage_map,
-                          const std::string& prefix = "")
-            : GGMLRunner(backend, params_backend),
+                          const std::string& prefix                           = "",
+                          std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : GGMLRunner(backend, weight_manager),
+              weight_prefix(prefix),
               config(LTXAudioVAEConfig::detect_from_weights(tensor_storage_map)),
               model(config) {
             model.init(params_ctx, tensor_storage_map, prefix);
@@ -1005,11 +1016,11 @@ namespace LTXV {
             }
         }
 
-        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
-            model.get_param_tensors(tensors, prefix);
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
+            model.get_param_tensors(tensors, weight_prefix);
         }
 
-        size_t get_params_buffer_size() {
+        size_t get_params_mem_size() {
             return model.get_params_mem_size();
         }
 
@@ -1029,7 +1040,7 @@ namespace LTXV {
                 ggml_build_forward_expand(gf, waveform);
                 return gf;
             };
-            auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), 4);
+            auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), 4);
             int64_t t1  = ggml_time_ms();
             LOG_INFO("ltx audio vae decode completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
             return result;
@@ -1056,7 +1067,8 @@ namespace LTXV {
             // ggml_backend_t backend = ggml_backend_cuda_init(0);
             LOG_INFO("loading ltx audio vae from '%s'", model_path.c_str());
 
-            ModelLoader model_loader;
+            auto model_manager        = std::make_shared<ModelManager>();
+            ModelLoader& model_loader = model_manager->loader();
             if (!model_loader.init_from_file(model_path)) {
                 LOG_ERROR("init model loader from file failed: '%s'", model_path.c_str());
                 return;
@@ -1064,20 +1076,17 @@ namespace LTXV {
 
             auto& tensor_storage_map = model_loader.get_tensor_storage_map();
             auto ltx_audio_vae       = std::make_shared<LTXAudioVAERunner>(backend,
-                                                                     backend,
                                                                      tensor_storage_map,
-                                                                     prefix);
+                                                                     prefix,
+                                                                     model_manager);
 
-            if (!ltx_audio_vae->alloc_params_buffer()) {
-                LOG_ERROR("ltx audio vae buffer allocation failed");
-                return;
-            }
-
-            std::map<std::string, ggml_tensor*> tensors;
-            ltx_audio_vae->get_param_tensors(tensors, "");
-
-            if (!model_loader.load_tensors(tensors)) {
-                LOG_ERROR("load tensors from model loader failed");
+            if (!model_manager->register_runner_params("LTX audio VAE test",
+                                                       *ltx_audio_vae,
+                                                       ModelManager::ResidencyMode::ParamBackend,
+                                                       backend,
+                                                       backend) ||
+                !model_manager->validate_registered_tensors()) {
+                LOG_ERROR("register ltx audio vae tensors with model manager failed");
                 return;
             }
 
@@ -1088,4 +1097,4 @@ namespace LTXV {
 
 }  // namespace LTXV
 
-#endif  // __SD_LTX_AUDIO_VAE_H__
+#endif  // __SD_MODEL_VAE_LTX_AUDIO_VAE_HPP__

src/model/vae/ltx_vae.hpp

@@ -1,5 +1,5 @@
-#ifndef __SD_LTX_VAE_HPP__
-#define __SD_LTX_VAE_HPP__
+#ifndef __SD_MODEL_VAE_LTX_VAE_HPP__
+#define __SD_MODEL_VAE_LTX_VAE_HPP__
 
 #include <algorithm>
 #include <fstream>
@@ -9,9 +9,10 @@
 #include <utility>
 #include <vector>
 
-#include "ltxv.hpp"
-#include "vae.hpp"
-#include "wan.hpp"
+#include "model/diffusion/ltxv.hpp"
+#include "model/vae/vae.hpp"
+#include "model/vae/wan_vae.hpp"
+#include "model_loader.h"
 
 namespace LTXVAE {
 
@@ -956,8 +957,8 @@ namespace LTXVAE {
 
             ggml_tensor* scaled_timestep = timestep;
             if (timestep_conditioning) {
-                auto multiplier = ggml_ext_backend_tensor_get_f32(params["timestep_scale_multiplier"]);
-                scaled_timestep = ggml_ext_scale(ctx->ggml_ctx, timestep, multiplier);
+                auto multiplier = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, params["timestep_scale_multiplier"]);
+                scaled_timestep = ggml_mul(ctx->ggml_ctx, timestep, multiplier);
             }
 
             x = conv_in->forward(ctx, x, causal_decoder);
@@ -1007,8 +1008,8 @@ namespace LTXVAE {
 
             ggml_tensor* scaled_timestep = timestep;
             if (timestep_conditioning && timestep != nullptr) {
-                auto multiplier = ggml_ext_backend_tensor_get_f32(params["timestep_scale_multiplier"]);
-                scaled_timestep = ggml_ext_scale(ctx->ggml_ctx, timestep, multiplier);
+                auto multiplier = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, params["timestep_scale_multiplier"]);
+                scaled_timestep = ggml_mul(ctx->ggml_ctx, timestep, multiplier);
             }
 
             // conv_in with feat_map for left temporal context
@@ -1222,11 +1223,11 @@ struct LTXVideoVAE : public VAE {
     LTXVAE::VideoVAE vae;
 
     LTXVideoVAE(ggml_backend_t backend,
-                ggml_backend_t params_backend,
                 const String2TensorStorage& tensor_storage_map,
                 const std::string& prefix,
-                bool decode_only  = true,
-                SDVersion version = VERSION_LTXAV)
+                bool decode_only                                    = true,
+                SDVersion version                                   = VERSION_LTXAV,
+                std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
         : decode_only(decode_only),
           ltx_vae_version(LTXVAE::detect_ltx_vae_version(tensor_storage_map, prefix)),
           timestep_conditioning(LTXVAE::detect_ltx_vae_timestep_conditioning(tensor_storage_map, prefix)),
@@ -1238,7 +1239,7 @@ struct LTXVideoVAE : public VAE {
               patch_size,
               tensor_storage_map,
               prefix),
-          VAE(version, backend, params_backend) {
+          VAE(version, backend, prefix, weight_manager) {
         vae.init(params_ctx, tensor_storage_map, prefix);
         decode_timestep_tensor.values()[0] = vae.decode_timestep;
     }
@@ -1270,8 +1271,8 @@ struct LTXVideoVAE : public VAE {
         }
     }
 
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) override {
-        vae.get_param_tensors(tensors, prefix);
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        vae.get_param_tensors(tensors, weight_prefix);
     }
 
     struct TemporalTilePlan {
@@ -1395,7 +1396,7 @@ struct LTXVideoVAE : public VAE {
                                                  static_cast<int>(start),
                                                  chunk_overlap);
             };
-            auto chunk = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, true),
+            auto chunk = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, true, true, true),
                                                          expected_dim);
             if (chunk.empty()) {
                 free_cache_ctx_and_buffer();
@@ -1425,7 +1426,7 @@ struct LTXVideoVAE : public VAE {
                                const sd::Tensor<float>& z,
                                bool decode_graph) override {
         if (!decode_graph && decode_only) {
-            LOG_ERROR("LTX video VAE encode requires encoder weights; create the context with vae_decode_only=false");
+            LOG_ERROR("LTX video VAE encode requires encoder weights");
             return {};
         }
         sd::Tensor<float> input = z;
@@ -1451,7 +1452,7 @@ struct LTXVideoVAE : public VAE {
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(input, decode_graph);
         };
-        auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), expected_dim);
+        auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), expected_dim);
         if (result.empty()) {
             return {};
         }
@@ -1464,7 +1465,7 @@ struct LTXVideoVAE : public VAE {
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_latent_statistics_graph(z, normalize);
         };
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false),
+        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false),
                                                static_cast<size_t>(z.dim()));
     }
 
@@ -1520,7 +1521,8 @@ struct LTXVideoVAE : public VAE {
         ggml_backend_t backend = sd_backend_cpu_init();
         LOG_INFO("loading ltx vae from '%s'", model_path.c_str());
 
-        ModelLoader model_loader;
+        auto model_manager        = std::make_shared<ModelManager>();
+        ModelLoader& model_loader = model_manager->loader();
         if (!model_loader.init_from_file_and_convert_name(model_path, "vae.")) {
             LOG_ERROR("init model loader from file failed: '%s'", model_path.c_str());
             return;
@@ -1528,22 +1530,19 @@ struct LTXVideoVAE : public VAE {
 
         auto& tensor_storage_map         = model_loader.get_tensor_storage_map();
         std::shared_ptr<LTXVideoVAE> vae = std::make_shared<LTXVideoVAE>(backend,
-                                                                         backend,
                                                                          tensor_storage_map,
                                                                          "first_stage_model",
                                                                          true,
-                                                                         VERSION_LTXAV);
+                                                                         VERSION_LTXAV,
+                                                                         model_manager);
 
-        if (!vae->alloc_params_buffer()) {
-            LOG_ERROR("vae buffer allocation failed");
-            return;
-        }
-
-        std::map<std::string, ggml_tensor*> tensors;
-        vae->get_param_tensors(tensors, "first_stage_model");
-
-        if (!model_loader.load_tensors(tensors)) {
-            LOG_ERROR("load tensors from model loader failed");
+        if (!model_manager->register_runner_params("LTX VAE test",
+                                                   *vae,
+                                                   ModelManager::ResidencyMode::ParamBackend,
+                                                   backend,
+                                                   backend) ||
+            !model_manager->validate_registered_tensors()) {
+            LOG_ERROR("register ltx vae tensors with model manager failed");
             return;
         }
 
@@ -1552,4 +1551,4 @@ struct LTXVideoVAE : public VAE {
     }
 };
 
-#endif  // __SD_LTX_VAE_HPP__
+#endif  // __SD_MODEL_VAE_LTX_VAE_HPP__

src/model/vae/tae.hpp

@@ -1,13 +1,13 @@
-#ifndef __TAE_HPP__
-#define __TAE_HPP__
+#ifndef __SD_MODEL_VAE_TAE_HPP__
+#define __SD_MODEL_VAE_TAE_HPP__
 
-#include "ggml_extend.hpp"
+#include "core/ggml_extend.hpp"
 #include "model.h"
 
 /*
     ===================================    TinyAutoEncoder  ===================================
     References:
-    https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/autoencoders/vae.py
+    https://github.com/huggingface/diffusers/blob/main/src/diffusers/model/autoencoders/vae.py
     https://github.com/madebyollin/taesd/blob/main/taesd.py
 
 */
@@ -623,14 +623,14 @@ struct TinyImageAutoEncoder : public VAE {
     bool decode_only = false;
 
     TinyImageAutoEncoder(ggml_backend_t backend,
-                         ggml_backend_t params_backend,
                          const String2TensorStorage& tensor_storage_map,
                          const std::string prefix,
-                         bool decoder_only = true,
-                         SDVersion version = VERSION_SD1)
-        : decode_only(decoder_only),
-          taesd(decoder_only, version),
-          VAE(version, backend, params_backend) {
+                         bool decoder_only                                   = true,
+                         SDVersion version                                   = VERSION_SD1,
+                         std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+        : VAE(version, backend, "tae", weight_manager),
+          decode_only(decoder_only),
+          taesd(decoder_only, version) {
         scale_input = false;
         taesd.init(params_ctx, tensor_storage_map, prefix);
     }
@@ -639,8 +639,8 @@ struct TinyImageAutoEncoder : public VAE {
         return "taesd";
     }
 
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
-        taesd.get_param_tensors(tensors, prefix);
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        taesd.get_param_tensors(tensors, weight_prefix);
     }
 
     sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) override {
@@ -676,7 +676,7 @@ struct TinyImageAutoEncoder : public VAE {
             return build_graph(z_tensor, decode_graph);
         };
 
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), z_tensor.dim());
+        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), z_tensor.dim());
     }
 };
 
@@ -686,13 +686,13 @@ struct TinyVideoAutoEncoder : public VAE {
     bool is_wide     = false;
 
     TinyVideoAutoEncoder(ggml_backend_t backend,
-                         ggml_backend_t params_backend,
                          const String2TensorStorage& tensor_storage_map,
                          const std::string prefix,
-                         bool decoder_only = true,
-                         SDVersion version = VERSION_WAN2)
-        : decode_only(decoder_only),
-          VAE(version, backend, params_backend) {
+                         bool decoder_only                                   = true,
+                         SDVersion version                                   = VERSION_WAN2,
+                         std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+        : VAE(version, backend, "tae", weight_manager),
+          decode_only(decoder_only) {
         for (auto tensor_storage : tensor_storage_map) {
             if (tensor_storage.first.find(prefix + ".3.conv.6.weight") != std::string::npos) {
                 is_wide = true;
@@ -708,8 +708,8 @@ struct TinyVideoAutoEncoder : public VAE {
         return "taehv";
     }
 
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
-        taehv.get_param_tensors(tensors, prefix);
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        taehv.get_param_tensors(tensors, weight_prefix);
     }
 
     sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) override {
@@ -746,8 +746,8 @@ struct TinyVideoAutoEncoder : public VAE {
             return build_graph(z_tensor, decode_graph);
         };
 
-        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), z_tensor.dim());
+        return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), z_tensor.dim());
     }
 };
 
-#endif  // __TAE_HPP__
+#endif  // __SD_MODEL_VAE_TAE_HPP__

src/model/vae/vae.hpp

@@ -1,12 +1,14 @@
-#ifndef __VAE_HPP__
-#define __VAE_HPP__
+#ifndef __SD_MODEL_VAE_VAE_HPP__
+#define __SD_MODEL_VAE_VAE_HPP__
 
-#include "common_block.hpp"
-#include "tensor_ggml.hpp"
+#include "core/tensor_ggml.hpp"
+#include "model/common/block.hpp"
+#include "model_manager.h"
 
 struct VAE : public GGMLRunner {
 protected:
     SDVersion version;
+    std::string weight_prefix;
     bool scale_input                                      = true;
     virtual sd::Tensor<float> _compute(const int n_threads,
                                        const sd::Tensor<float>& z,
@@ -62,8 +64,11 @@ protected:
     }
 
 public:
-    VAE(SDVersion version, ggml_backend_t backend, ggml_backend_t params_backend)
-        : version(version), GGMLRunner(backend, params_backend) {}
+    VAE(SDVersion version,
+        ggml_backend_t backend,
+        const std::string& weight_prefix                    = "",
+        std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+        : version(version), weight_prefix(weight_prefix), GGMLRunner(backend, weight_manager) {}
 
     int get_scale_factor() {
         int scale_factor = 8;
@@ -214,7 +219,7 @@ public:
     virtual sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) = 0;
     virtual sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents)                           = 0;
     virtual sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents)                           = 0;
-    virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix)         = 0;
+    virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors)                                   = 0;
     virtual void set_conv2d_scale(float scale) { SD_UNUSED(scale); };
     virtual void set_temporal_tiling_enabled(bool enabled) { SD_UNUSED(enabled); };
     virtual void set_tiling_params(const sd_tiling_params_t& params) {
@@ -223,8 +228,10 @@ public:
 };
 
 struct FakeVAE : public VAE {
-    FakeVAE(SDVersion version, ggml_backend_t backend, ggml_backend_t params_backend)
-        : VAE(version, backend, params_backend) {}
+    FakeVAE(SDVersion version,
+            ggml_backend_t backend,
+            std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+        : VAE(version, backend, "", weight_manager) {}
 
     int get_encoder_output_channels(int input_channels) {
         return input_channels;
@@ -251,11 +258,11 @@ struct FakeVAE : public VAE {
         return latents;
     }
 
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) override {}
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {}
 
     std::string get_desc() override {
         return "fake_vae";
     }
 };
 
-#endif  // __VAE_HPP__
+#endif  // __SD_MODEL_VAE_VAE_HPP__

src/model_io/gguf_io.cpp

@@ -5,9 +5,9 @@
 #include <string>
 #include <vector>
 
+#include "core/util.h"
 #include "gguf.h"
 #include "gguf_reader_ext.h"
-#include "util.h"
 
 static void set_error(std::string* error, const std::string& message) {
     if (error != nullptr) {

src/model_io/gguf_reader_ext.h

@@ -6,8 +6,8 @@
 #include <string>
 #include <vector>
 
+#include "core/util.h"
 #include "ggml.h"
-#include "util.h"
 
 struct GGUFTensorInfo {
     std::string name;

src/model_io/pickle_io.cpp

@@ -8,7 +8,7 @@
 #include <vector>
 
 #include "binary_io.h"
-#include "util.h"
+#include "core/util.h"
 
 // $ python -m pickletools sd-v1-4/archive/data.pkl | head -n 100
 //     0: \x80 PROTO      2

src/model_io/safetensors_io.cpp

@@ -7,8 +7,8 @@
 #include <vector>
 
 #include "binary_io.h"
+#include "core/util.h"
 #include "json.hpp"
-#include "util.h"
 
 static constexpr size_t ST_HEADER_SIZE_LEN = 8;
 

src/model_io/torch_legacy_io.cpp

@@ -7,8 +7,8 @@
 #include <unordered_map>
 #include <vector>
 
+#include "core/util.h"
 #include "pickle_io.h"
-#include "util.h"
 
 // torch.save format background:
 //

src/model_loader.cpp

@@ -1,6 +1,7 @@
 #include <algorithm>
 #include <atomic>
 #include <chrono>
+#include <cinttypes>
 #include <cstdarg>
 #include <cstdlib>
 #include <fstream>
@@ -13,18 +14,18 @@
 #include <unordered_map>
 #include <vector>
 
-#include "model.h"
+#include "core/util.h"
 #include "model_io/gguf_io.h"
 #include "model_io/safetensors_io.h"
 #include "model_io/torch_legacy_io.h"
 #include "model_io/torch_zip_io.h"
+#include "model_loader.h"
 #include "stable-diffusion.h"
-#include "util.h"
 
+#include "core/ggml_extend_backend.h"
 #include "ggml-alloc.h"
 #include "ggml-backend.h"
 #include "ggml.h"
-#include "ggml_extend_backend.h"
 #include "zip.h"
 
 #include "name_conversion.h"
@@ -204,10 +205,28 @@ void convert_tensor(void* src,
 
 /*================================================= ModelLoader ==================================================*/
 
+ModelLoader::ModelLoader()
+    : n_threads_(sd_get_num_physical_cores()) {
+}
+
+size_t ModelLoader::add_file_path(const std::string& file_path) {
+    if (model_files_processed) {
+        file_data.clear();
+        model_files_processed = false;
+    }
+    file_paths_.push_back(file_path);
+    return file_paths_.size() - 1;
+}
+
 void ModelLoader::add_tensor_storage(const TensorStorage& tensor_storage) {
     tensor_storage_map[tensor_storage.name] = tensor_storage;
 }
 
+void ModelLoader::set_n_threads(int n_threads) {
+    n_threads_ = n_threads > 0 ? n_threads : sd_get_num_physical_cores();
+    LOG_DEBUG("using %d threads for model loading", n_threads_);
+}
+
 bool ModelLoader::init_from_file(const std::string& file_path, const std::string& prefix) {
     if (is_directory(file_path)) {
         LOG_INFO("load %s using diffusers format", file_path.c_str());
@@ -271,8 +290,7 @@ bool ModelLoader::init_from_gguf_file(const std::string& file_path, const std::s
         return false;
     }
 
-    file_paths_.push_back(file_path);
-    size_t file_index = file_paths_.size() - 1;
+    size_t file_index = add_file_path(file_path);
 
     for (auto& tensor_storage : tensor_storages) {
         // LOG_DEBUG("%s", tensor_storage.name.c_str());
@@ -300,8 +318,7 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const
         return false;
     }
 
-    file_paths_.push_back(file_path);
-    size_t file_index = file_paths_.size() - 1;
+    size_t file_index = add_file_path(file_path);
 
     for (auto& tensor_storage : tensor_storages) {
         if (is_unused_tensor(tensor_storage.name)) {
@@ -335,8 +352,7 @@ bool ModelLoader::init_from_torch_legacy_file(const std::string& file_path, cons
         return false;
     }
 
-    file_paths_.push_back(file_path);
-    size_t file_index = file_paths_.size() - 1;
+    size_t file_index = add_file_path(file_path);
 
     for (auto& tensor_storage : tensor_storages) {
         if (is_unused_tensor(tensor_storage.name)) {
@@ -366,8 +382,7 @@ bool ModelLoader::init_from_torch_zip_file(const std::string& file_path, const s
         return false;
     }
 
-    file_paths_.push_back(file_path);
-    size_t file_index = file_paths_.size() - 1;
+    size_t file_index = add_file_path(file_path);
 
     for (auto& tensor_storage : tensor_storages) {
         if (!starts_with(tensor_storage.name, prefix)) {
@@ -435,6 +450,9 @@ SDVersion ModelLoader::get_sd_version() {
         if (tensor_storage.name.find("model.diffusion_model.net.lq_proj.latent_proj.0.weight") != std::string::npos) {
             return VERSION_PID;
         }
+        if (tensor_storage.name.find("embed_image_indicator.weight") != std::string::npos) {
+            return VERSION_IDEOGRAM4;
+        }
         if (tensor_storage.name.find("model.diffusion_model.nerf_final_layer_conv.") != std::string::npos) {
             return VERSION_CHROMA_RADIANCE;
         }
@@ -757,8 +775,6 @@ void ModelLoader::process_model_files(bool enable_mmap, bool writable_mmap) {
         return;
     }
 
-    int64_t start_time = ggml_time_ms();
-
     std::vector<TensorStorage> processed_tensor_storages;
     for (const auto& [name, tensor_storage] : tensor_storage_map) {
         if (is_unused_tensor(tensor_storage.name)) {
@@ -809,20 +825,12 @@ void ModelLoader::process_model_files(bool enable_mmap, bool writable_mmap) {
             } else {
                 LOG_WARN("failed to memory-map '%s' (falling back to read())", file_path.c_str());
             }
-        } else if (!is_zip) {
-            LOG_INFO("NOT using mmap for '%s' (mmap disabled by caller)",
-                     file_path.c_str());
         }
 
         file_data.push_back(std::move(fdata));
     }
 
     model_files_processed = true;
-
-    int64_t end_time        = ggml_time_ms();
-    int64_t process_time_ms = end_time - start_time;
-
-    LOG_INFO("model files processing completed in %.2fs", process_time_ms / 1000.f);
 }
 
 std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggml_tensor*>& tensors,
@@ -916,7 +924,9 @@ std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggm
     return result;
 }
 
-bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_threads_p, bool enable_mmap) {
+bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
+                               bool enable_mmap,
+                               const std::set<std::string>* target_tensor_names) {
     process_model_files(enable_mmap, false);
 
     std::atomic<int64_t> read_time_ms(0);
@@ -925,14 +935,26 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
     std::atomic<int64_t> convert_time_ms(0);
     std::atomic<uint64_t> bytes_processed(0);
 
-    int num_threads_to_use = n_threads_p > 0 ? n_threads_p : sd_get_num_physical_cores();
-    LOG_DEBUG("using %d threads for model loading", num_threads_to_use);
+    int num_threads_to_use = n_threads_;
 
     int64_t start_time = ggml_time_ms();
 
     size_t total_tensors_to_process = 0;
+    std::vector<size_t> file_tensors_to_process_counts;
+    file_tensors_to_process_counts.reserve(file_data.size());
     for (const auto& fdata : file_data) {
-        total_tensors_to_process += fdata.tensors.size();
+        size_t file_tensors_to_process = 0;
+        if (target_tensor_names == nullptr) {
+            file_tensors_to_process = fdata.tensors.size();
+        } else {
+            for (const TensorStorage& tensor_storage : fdata.tensors) {
+                if (target_tensor_names->find(tensor_storage.name) != target_tensor_names->end()) {
+                    file_tensors_to_process++;
+                }
+            }
+        }
+        file_tensors_to_process_counts.push_back(file_tensors_to_process);
+        total_tensors_to_process += file_tensors_to_process;
     }
 
     bool success                   = true;
@@ -940,17 +962,38 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
     const int64_t t_start          = start_time;
     int last_n_threads             = 1;
 
-    for (auto& fdata : file_data) {
+    for (size_t file_index = 0; file_index < file_data.size(); ++file_index) {
+        auto& fdata                  = file_data[file_index];
         const std::string& file_path = fdata.path;
-        LOG_DEBUG("loading tensors from %s", file_path.c_str());
 
         const std::vector<TensorStorage>& file_tensors = fdata.tensors;
+        std::vector<const TensorStorage*> tensors_to_process;
+        size_t file_tensors_to_process = file_tensors_to_process_counts[file_index];
+        tensors_to_process.reserve(file_tensors_to_process);
+        if (target_tensor_names == nullptr) {
+            for (const TensorStorage& tensor_storage : file_tensors) {
+                tensors_to_process.push_back(&tensor_storage);
+            }
+        } else {
+            for (const TensorStorage& tensor_storage : file_tensors) {
+                if (target_tensor_names->find(tensor_storage.name) != target_tensor_names->end()) {
+                    tensors_to_process.push_back(&tensor_storage);
+                }
+            }
+        }
+        if (tensors_to_process.empty()) {
+            continue;
+        }
+        LOG_DEBUG("loading %zu/%zu tensors from %s",
+                  tensors_to_process.size(),
+                  file_tensors.size(),
+                  file_path.c_str());
 
         bool is_zip = fdata.is_zip;
 
         std::shared_ptr<MmapWrapper> mmapped = fdata.mmapped;
 
-        int n_threads = is_zip ? 1 : std::min(num_threads_to_use, (int)file_tensors.size());
+        int n_threads = is_zip ? 1 : std::min(num_threads_to_use, (int)tensors_to_process.size());
         if (n_threads < 1) {
             n_threads = 1;
         }
@@ -959,6 +1002,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
         std::atomic<size_t> tensor_idx(0);
         std::atomic<bool> failed(false);
         std::vector<std::thread> workers;
+        std::mutex rpc_backend_mutex;
 
         for (int i = 0; i < n_threads; ++i) {
             workers.emplace_back([&, file_path, is_zip]() {
@@ -986,11 +1030,11 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
                 while (true) {
                     int64_t t0, t1;
                     size_t idx = tensor_idx.fetch_add(1);
-                    if (idx >= file_tensors.size() || failed) {
+                    if (idx >= tensors_to_process.size() || failed) {
                         break;
                     }
 
-                    const TensorStorage& tensor_storage = file_tensors[idx];
+                    const TensorStorage& tensor_storage = *tensors_to_process[idx];
                     ggml_tensor* dst_tensor             = nullptr;
 
                     t0 = ggml_time_ms();
@@ -1115,7 +1159,19 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
 
                     if (dst_tensor->buffer != nullptr && !ggml_backend_buffer_is_host(dst_tensor->buffer)) {
                         t0 = ggml_time_ms();
-                        ggml_backend_tensor_set(dst_tensor, convert_buf, 0, ggml_nbytes(dst_tensor));
+
+                        // RPC backends require serialized access to prevent concurrency issues
+                        const char* buffer_type_name = ggml_backend_buft_name(ggml_backend_buffer_get_type(dst_tensor->buffer));
+                        bool is_rpc_buffer           = buffer_type_name != nullptr &&
+                                             std::string(buffer_type_name).find("RPC") != std::string::npos;
+
+                        if (is_rpc_buffer) {
+                            std::lock_guard<std::mutex> lock(rpc_backend_mutex);
+                            ggml_backend_tensor_set(dst_tensor, convert_buf, 0, ggml_nbytes(dst_tensor));
+                        } else {
+                            ggml_backend_tensor_set(dst_tensor, convert_buf, 0, ggml_nbytes(dst_tensor));
+                        }
+
                         t1 = ggml_time_ms();
                         copy_to_backend_time_ms.fetch_add(t1 - t0);
                     }
@@ -1130,16 +1186,18 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
 
         while (true) {
             size_t current_idx = tensor_idx.load();
-            if (current_idx >= file_tensors.size() || failed) {
+            if (current_idx >= tensors_to_process.size() || failed) {
                 break;
             }
             size_t curr_num       = total_tensors_processed + current_idx;
             float elapsed_seconds = (ggml_time_ms() - t_start) / 1000.0f;
-            pretty_bytes_progress(static_cast<int>(curr_num),
-                                  static_cast<int>(total_tensors_to_process),
-                                  bytes_processed.load(),
-                                  elapsed_seconds);
-            std::this_thread::sleep_for(std::chrono::milliseconds(200));
+            if (total_tensors_to_process > 0) {
+                pretty_bytes_progress(static_cast<int>(curr_num),
+                                      static_cast<int>(total_tensors_to_process),
+                                      bytes_processed.load(),
+                                      elapsed_seconds);
+            }
+            std::this_thread::sleep_for(std::chrono::milliseconds(total_tensors_to_process <= 4 ? 10 : 200));
         }
 
         for (auto& w : workers) {
@@ -1150,12 +1208,14 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
             success = false;
             break;
         }
-        total_tensors_processed += file_tensors.size();
-        pretty_bytes_progress(static_cast<int>(total_tensors_processed),
-                              static_cast<int>(total_tensors_to_process),
-                              bytes_processed.load(),
-                              (ggml_time_ms() - t_start) / 1000.0f);
-        if (total_tensors_processed < total_tensors_to_process) {
+        total_tensors_processed += tensors_to_process.size();
+        if (total_tensors_to_process > 0) {
+            pretty_bytes_progress(static_cast<int>(total_tensors_processed),
+                                  static_cast<int>(total_tensors_to_process),
+                                  bytes_processed.load(),
+                                  (ggml_time_ms() - t_start) / 1000.0f);
+        }
+        if (total_tensors_processed < total_tensors_to_process && total_tensors_to_process > 0) {
             printf("\n");
         }
     }
@@ -1170,9 +1230,77 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
     return success;
 }
 
+bool ModelLoader::load_float_tensor(const std::string& name,
+                                    std::vector<float>& data,
+                                    int n_threads,
+                                    bool use_mmap) {
+    data.clear();
+
+    auto tensor_storage_it = tensor_storage_map.find(name);
+    if (tensor_storage_it == tensor_storage_map.end()) {
+        return false;
+    }
+
+    const TensorStorage& tensor_storage = tensor_storage_it->second;
+    int64_t n_elements                  = tensor_storage.nelements();
+    if (n_elements <= 0) {
+        LOG_ERROR("tensor '%s' has invalid element count: %" PRId64, name.c_str(), n_elements);
+        return false;
+    }
+    if (tensor_storage.n_dims <= 0 || tensor_storage.n_dims > GGML_MAX_DIMS) {
+        LOG_ERROR("tensor '%s' has unsupported dims: %d", name.c_str(), tensor_storage.n_dims);
+        return false;
+    }
+
+    std::vector<float> loaded_data(static_cast<size_t>(n_elements));
+    ggml_init_params params;
+    params.mem_size   = ggml_tensor_overhead();
+    params.mem_buffer = nullptr;
+    params.no_alloc   = true;
+
+    ggml_context* ctx = ggml_init(params);
+    if (ctx == nullptr) {
+        LOG_ERROR("failed to create context for tensor '%s'", name.c_str());
+        return false;
+    }
+
+    ggml_tensor* tensor = ggml_new_tensor(ctx, GGML_TYPE_F32, tensor_storage.n_dims, tensor_storage.ne);
+    ggml_set_name(tensor, name.c_str());
+    tensor->data = loaded_data.data();
+
+    bool loaded           = false;
+    auto on_new_tensor_cb = [&](const TensorStorage& current_tensor_storage, ggml_tensor** dst_tensor) -> bool {
+        *dst_tensor = nullptr;
+        if (current_tensor_storage.name != name) {
+            return true;
+        }
+        if (current_tensor_storage.nelements() != n_elements) {
+            LOG_ERROR("tensor '%s' element count changed during load", name.c_str());
+            return false;
+        }
+        *dst_tensor = tensor;
+        loaded      = true;
+        return true;
+    };
+
+    std::set<std::string> target_tensor_names{name};
+    if (n_threads > 0) {
+        set_n_threads(n_threads);
+    }
+    bool success = load_tensors(on_new_tensor_cb, use_mmap, &target_tensor_names);
+    ggml_free(ctx);
+
+    if (!success || !loaded) {
+        data.clear();
+        return false;
+    }
+
+    data = std::move(loaded_data);
+    return true;
+}
+
 bool ModelLoader::load_tensors(std::map<std::string, ggml_tensor*>& tensors,
                                std::set<std::string> ignore_tensors,
-                               int n_threads,
                                bool enable_mmap) {
     std::set<std::string> tensor_names_in_file;
     std::mutex tensor_names_mutex;
@@ -1216,7 +1344,7 @@ bool ModelLoader::load_tensors(std::map<std::string, ggml_tensor*>& tensors,
         return true;
     };
 
-    bool success = load_tensors(on_new_tensor_cb, n_threads, enable_mmap);
+    bool success = load_tensors(on_new_tensor_cb, enable_mmap);
     if (!success) {
         LOG_ERROR("load tensors from file failed");
         return false;
@@ -1254,6 +1382,8 @@ bool ModelLoader::tensor_should_be_converted(const TensorStorage& tensor_storage
             // Pass, do not convert
         } else if (ends_with(name, ".scale")) {
             // Pass, do not convert
+        } else if (ends_with(name, ".weight_scale")) {
+            // Pass, do not convert
         } else if (contains(name, "img_in.") ||
                    contains(name, "txt_in.") ||
                    contains(name, "time_in.") ||

src/model_loader.h

@@ -0,0 +1,93 @@
+#ifndef __MODEL_LOADER_H__
+#define __MODEL_LOADER_H__
+
+#include <cstdint>
+#include <map>
+#include <memory>
+#include <set>
+#include <string>
+#include <vector>
+
+#include "model.h"
+
+TensorTypeRules parse_tensor_type_rules(const std::string& tensor_type_rules);
+
+class MmapWrapper;
+
+struct ModelFileData {
+    std::string path;
+    std::vector<TensorStorage> tensors;
+    std::shared_ptr<MmapWrapper> mmapped;
+    std::shared_ptr<struct ggml_backend_buffer> mmbuffer;
+    bool is_zip;
+};
+
+struct MmapTensorStore {
+    std::shared_ptr<MmapWrapper> mmapped;
+    std::shared_ptr<struct ggml_backend_buffer> mmbuffer;
+};
+
+class ModelLoader {
+protected:
+    SDVersion version_ = VERSION_COUNT;
+    std::vector<std::string> file_paths_;
+    std::vector<ModelFileData> file_data;
+    bool model_files_processed = false;
+    String2TensorStorage tensor_storage_map;
+    int n_threads_;
+
+    size_t add_file_path(const std::string& file_path);
+    void add_tensor_storage(const TensorStorage& tensor_storage);
+
+    bool init_from_gguf_file(const std::string& file_path, const std::string& prefix = "");
+    bool init_from_safetensors_file(const std::string& file_path, const std::string& prefix = "");
+    bool init_from_torch_zip_file(const std::string& file_path, const std::string& prefix = "");
+    bool init_from_torch_legacy_file(const std::string& file_path, const std::string& prefix = "");
+    bool init_from_diffusers_file(const std::string& file_path, const std::string& prefix = "");
+
+public:
+    ModelLoader();
+
+    bool init_from_file(const std::string& file_path, const std::string& prefix = "");
+    void convert_tensors_name();
+    bool init_from_file_and_convert_name(const std::string& file_path,
+                                         const std::string& prefix = "",
+                                         SDVersion version         = VERSION_COUNT);
+    SDVersion get_sd_version();
+    std::map<ggml_type, uint32_t> get_wtype_stat();
+    std::map<ggml_type, uint32_t> get_conditioner_wtype_stat();
+    std::map<ggml_type, uint32_t> get_diffusion_model_wtype_stat();
+    std::map<ggml_type, uint32_t> get_vae_wtype_stat();
+    String2TensorStorage& get_tensor_storage_map() { return tensor_storage_map; }
+    const String2TensorStorage& get_tensor_storage_map() const { return tensor_storage_map; }
+    void set_n_threads(int n_threads);
+    void set_wtype_override(ggml_type wtype, std::string tensor_type_rules = "");
+    void process_model_files(bool enable_mmap = false, bool writable_mmap = true);
+    std::vector<MmapTensorStore> mmap_tensors(std::map<std::string, ggml_tensor*>& tensors,
+                                              std::set<std::string> ignore_tensors = {},
+                                              bool writable                        = true);
+    bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
+                      bool use_mmap                                    = false,
+                      const std::set<std::string>* target_tensor_names = nullptr);
+    bool load_tensors(std::map<std::string, ggml_tensor*>& tensors,
+                      std::set<std::string> ignore_tensors = {},
+                      bool use_mmap                        = false);
+    bool load_float_tensor(const std::string& name,
+                           std::vector<float>& data,
+                           int n_threads = 0,
+                           bool use_mmap = false);
+
+    std::vector<std::string> get_tensor_names() const {
+        std::vector<std::string> names;
+        for (const auto& [name, tensor_storage] : tensor_storage_map) {
+            names.push_back(name);
+        }
+        return names;
+    }
+
+    bool tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type);
+    int64_t get_params_mem_size(ggml_backend_t backend, ggml_type type = GGML_TYPE_COUNT);
+    ~ModelLoader() = default;
+};
+
+#endif  // __MODEL_LOADER_H__

src/model_manager.cpp

@@ -0,0 +1,944 @@
+#include "model_manager.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <iterator>
+#include <mutex>
+#include <unordered_set>
+
+#include "core/ggml_extend_backend.h"
+#include "core/util.h"
+#include "model/adapter/lora.hpp"
+
+static size_t aligned_offset(const void* buffer, size_t offset, size_t alignment) {
+    GGML_ASSERT(alignment != 0 && (alignment & (alignment - 1)) == 0);
+    size_t align = (alignment - ((reinterpret_cast<uintptr_t>(buffer) + offset) % alignment)) % alignment;
+    return offset + align;
+}
+
+static bool lora_specs_equal(const std::vector<ModelManager::LoraSpec>& lhs,
+                             const std::vector<ModelManager::LoraSpec>& rhs) {
+    if (lhs.size() != rhs.size()) {
+        return false;
+    }
+    for (size_t i = 0; i < lhs.size(); ++i) {
+        if (lhs[i].path != rhs[i].path ||
+            lhs[i].multiplier != rhs[i].multiplier ||
+            lhs[i].is_high_noise != rhs[i].is_high_noise ||
+            lhs[i].tensor_name_prefix_filter != rhs[i].tensor_name_prefix_filter ||
+            lhs[i].required != rhs[i].required) {
+            return false;
+        }
+    }
+    return true;
+}
+
+static std::string lora_id(const ModelManager::LoraSpec& lora) {
+    return lora.is_high_noise ? "|high_noise|" + lora.path : lora.path;
+}
+
+static bool backend_supports_host_buffer(ggml_backend_t backend) {
+    if (backend == nullptr) {
+        return false;
+    }
+    if (sd_backend_is_cpu(backend)) {
+        return true;
+    }
+    ggml_backend_dev_t dev = ggml_backend_get_device(backend);
+    if (dev == nullptr) {
+        return false;
+    }
+    ggml_backend_dev_props props;
+    ggml_backend_dev_get_props(dev, &props);
+    return props.caps.buffer_from_host_ptr;
+}
+
+ModelManager::~ModelManager() {
+    release_all();
+}
+
+void ModelManager::set_common_ignore_tensors(std::set<std::string> ignore_tensors) {
+    common_ignore_tensors_ = std::move(ignore_tensors);
+}
+
+void ModelManager::set_loras(std::vector<LoraSpec> loras, SDVersion version) {
+    if (loras.empty() && loras_.empty()) {
+        lora_version_ = version;
+        return;
+    }
+    if (lora_version_ == version && lora_specs_equal(loras_, loras)) {
+        return;
+    }
+
+    loras_        = std::move(loras);
+    lora_version_ = version;
+    current_lora_epoch_++;
+    reset_lora_applied_params();
+}
+
+std::set<std::string> ModelManager::tensor_names() const {
+    std::set<std::string> names;
+    for (const auto& state : tensor_states_) {
+        if (state != nullptr) {
+            names.insert(state->name);
+        }
+    }
+    return names;
+}
+
+size_t estimate_tensors_size(const std::map<std::string, ggml_tensor*>& tensors) {
+    size_t size = 0;
+    std::unordered_set<ggml_tensor*> seen;
+    for (const auto& pair : tensors) {
+        ggml_tensor* tensor = pair.second;
+        if (tensor == nullptr || seen.find(tensor) != seen.end()) {
+            continue;
+        }
+        seen.insert(tensor);
+        size += ggml_nbytes(tensor);
+    }
+    return size;
+}
+
+bool ModelManager::register_param_tensors(const std::string& desc,
+                                          std::map<std::string, ggml_tensor*> tensors,
+                                          ResidencyMode residency_mode,
+                                          ggml_backend_t compute_backend,
+                                          ggml_backend_t params_backend,
+                                          size_t* registered_tensor_size) {
+    if (desc.empty()) {
+        LOG_ERROR("model manager tensor desc is empty");
+        return false;
+    }
+    if (registered_tensor_size != nullptr) {
+        *registered_tensor_size += estimate_tensors_size(tensors);
+    }
+
+    std::vector<std::unique_ptr<TensorState>> new_states;
+    new_states.reserve(tensors.size());
+
+    for (const auto& pair : tensors) {
+        const std::string& name = pair.first;
+        ggml_tensor* tensor     = pair.second;
+        if (tensor == nullptr) {
+            continue;
+        }
+        if (tensor_states_by_name_.find(name) != tensor_states_by_name_.end()) {
+            LOG_ERROR("model manager tensor name '%s' is already registered", name.c_str());
+            return false;
+        }
+        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;
+        new_states.push_back(std::move(state));
+    }
+
+    for (auto& state : new_states) {
+        TensorState* registered_state                  = state.get();
+        tensor_states_by_name_[registered_state->name] = registered_state;
+        tensor_states_.push_back(std::move(state));
+    }
+    return true;
+}
+
+bool ModelManager::validate_registered_tensors() {
+    bool ok = true;
+    for (const auto& state : tensor_states_) {
+        if (state == nullptr) {
+            ok = false;
+            continue;
+        }
+        bool state_ok = validate_tensor(*state);
+        if (state_ok) {
+            state->metadata_validated = true;
+        }
+        ok = state_ok && ok;
+    }
+    return ok;
+}
+
+bool ModelManager::load_tensors_to_params_backend(const std::vector<TensorState*>& states) {
+    std::vector<TensorState*> need_load;
+    need_load.reserve(states.size());
+    for (TensorState* state : states) {
+        if (state == nullptr || should_ignore(*state) || is_optional_missing_tensor(state->name)) {
+            continue;
+        }
+        if (!state->metadata_validated) {
+            if (!validate_tensor(*state)) {
+                return false;
+            }
+            state->metadata_validated = true;
+        }
+        if (!state->loaded_to_params_backend) {
+            need_load.push_back(state);
+        }
+    }
+    if (need_load.empty()) {
+        return true;
+    }
+
+    std::vector<ParamsStorageBlock*> created_storage_blocks;
+    if (!mmap_params(need_load, created_storage_blocks)) {
+        for (ParamsStorageBlock* block : created_storage_blocks) {
+            if (block != nullptr) {
+                free_params_storage_block(*block);
+                erase_params_storage_block(block);
+            }
+        }
+        return false;
+    }
+
+    std::vector<TensorState*> need_alloc;
+    need_alloc.reserve(need_load.size());
+    for (TensorState* state : need_load) {
+        if (state->tensor != nullptr && state->tensor->data == nullptr && state->tensor->view_src == nullptr) {
+            need_alloc.push_back(state);
+        }
+    }
+
+    if (!alloc_params_buffers(need_alloc, created_storage_blocks) ||
+        !load_tensors(need_load)) {
+        for (ParamsStorageBlock* block : created_storage_blocks) {
+            if (block != nullptr) {
+                free_params_storage_block(*block);
+                erase_params_storage_block(block);
+            }
+        }
+        return false;
+    }
+    for (ParamsStorageBlock* block : created_storage_blocks) {
+        if (block != nullptr && block->buffer != nullptr) {
+            LOG_DEBUG("model manager prepared params backend buffer (%6.2f MB, %zu tensors, %s)",
+                      ggml_backend_buffer_get_size(block->buffer) / (1024.f * 1024.f),
+                      block->states.size(),
+                      ggml_backend_buffer_is_host(block->buffer) ? "RAM" : "VRAM");
+        }
+    }
+
+    return true;
+}
+
+bool ModelManager::stage_tensors_to_compute_backend(const std::vector<TensorState*>& states) {
+    std::map<ggml_backend_t, std::vector<TensorState*>> states_by_compute_backend;
+    for (TensorState* state : states) {
+        if (state == nullptr || should_ignore(*state) || is_optional_missing_tensor(state->name)) {
+            continue;
+        }
+        if (state->compute_backend == nullptr) {
+            LOG_ERROR("model manager compute backend is null for tensor '%s'", state->name.c_str());
+            return false;
+        }
+        if (state->params_backend == nullptr) {
+            LOG_ERROR("model manager params backend is null for tensor '%s'", state->name.c_str());
+            return false;
+        }
+        if (state->compute_backend == state->params_backend || state->staged_to_compute_backend) {
+            continue;
+        }
+        if (!state->loaded_to_params_backend || state->tensor == nullptr || state->tensor->data == nullptr) {
+            LOG_ERROR("model manager tensor '%s' is not loaded to params backend", state->name.c_str());
+            return false;
+        }
+        states_by_compute_backend[state->compute_backend].push_back(state);
+    }
+
+    for (const auto& pair : states_by_compute_backend) {
+        ggml_backend_t compute_backend          = pair.first;
+        const std::vector<TensorState*>& states = pair.second;
+        if (states.empty()) {
+            continue;
+        }
+
+        int64_t t0 = ggml_time_ms();
+
+        ggml_init_params init_params;
+        init_params.mem_size   = std::max<size_t>(1, states.size()) * ggml_tensor_overhead();
+        init_params.mem_buffer = nullptr;
+        init_params.no_alloc   = true;
+
+        ggml_context* staging_ctx = ggml_init(init_params);
+        GGML_ASSERT(staging_ctx != nullptr);
+
+        std::vector<std::pair<TensorState*, ggml_tensor*>> staged_tensors;
+        staged_tensors.reserve(states.size());
+        for (TensorState* state : states) {
+            ggml_tensor* staging_tensor = ggml_dup_tensor(staging_ctx, state->tensor);
+            ggml_set_name(staging_tensor, state->tensor->name);
+            staged_tensors.push_back({state, staging_tensor});
+        }
+
+        ggml_backend_buffer_t compute_buffer = ggml_backend_alloc_ctx_tensors(staging_ctx, compute_backend);
+        if (compute_buffer == nullptr) {
+            LOG_ERROR("model manager alloc compute params backend buffer failed, num_tensors = %zu",
+                      staged_tensors.size());
+            ggml_free(staging_ctx);
+            return false;
+        }
+        ggml_backend_buffer_set_usage(compute_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
+
+        for (auto& staged_tensor : staged_tensors) {
+            TensorState* state          = staged_tensor.first;
+            ggml_tensor* managed_tensor = state->tensor;
+            ggml_tensor* staging_tensor = staged_tensor.second;
+            ggml_backend_tensor_copy(managed_tensor, staging_tensor);
+            std::swap(managed_tensor->buffer, staging_tensor->buffer);
+            std::swap(managed_tensor->data, staging_tensor->data);
+            std::swap(managed_tensor->extra, staging_tensor->extra);
+        }
+        ggml_backend_synchronize(compute_backend);
+
+        auto block             = std::make_unique<ComputeStagingBlock>();
+        block->compute_backend = compute_backend;
+        block->buffer          = compute_buffer;
+        block->staging_ctx     = staging_ctx;
+        block->staged_tensors  = std::move(staged_tensors);
+        for (auto& staged_tensor : block->staged_tensors) {
+            TensorState* state               = staged_tensor.first;
+            state->staged_to_compute_backend = true;
+        }
+        compute_staging_blocks_.push_back(std::move(block));
+
+        int64_t t1 = ggml_time_ms();
+        LOG_DEBUG("model manager staged compute params (%6.2f MB, %zu tensors) to %s, taking %.2fs",
+                  ggml_backend_buffer_get_size(compute_buffer) / (1024.f * 1024.f),
+                  states.size(),
+                  ggml_backend_name(compute_backend),
+                  (t1 - t0) * 1.0f / 1000);
+    }
+
+    return true;
+}
+
+bool ModelManager::apply_loras_to_params(const std::vector<TensorState*>& states) {
+    if (loras_.empty()) {
+        return true;
+    }
+
+    struct LoraApplyGroup {
+        std::map<std::string, ggml_tensor*> model_tensors;
+        std::vector<TensorState*> states;
+    };
+
+    std::map<ggml_backend_t, LoraApplyGroup> groups;
+    for (TensorState* state : states) {
+        if (state == nullptr || state->tensor == nullptr ||
+            should_ignore(*state) || is_optional_missing_tensor(state->name)) {
+            continue;
+        }
+        if (state->applied_lora_epoch == current_lora_epoch_) {
+            continue;
+        }
+        if (state->compute_backend == nullptr) {
+            LOG_ERROR("model manager compute backend is null for lora target tensor '%s'", state->name.c_str());
+            return false;
+        }
+        if (state->tensor->data == nullptr) {
+            LOG_ERROR("model manager lora target tensor '%s' is not prepared", state->name.c_str());
+            return false;
+        }
+        LoraApplyGroup& group            = groups[state->compute_backend];
+        group.model_tensors[state->name] = state->tensor;
+        group.states.push_back(state);
+    }
+
+    if (groups.empty()) {
+        return true;
+    }
+
+    std::set<std::string> all_tensor_names = tensor_names();
+    for (auto& group_pair : groups) {
+        ggml_backend_t compute_backend = group_pair.first;
+        LoraApplyGroup& group          = group_pair.second;
+        for (const LoraSpec& lora_spec : loras_) {
+            if (group.model_tensors.empty()) {
+                continue;
+            }
+
+            std::string id = lora_id(lora_spec);
+            auto lora      = std::make_shared<LoraModel>(id,
+                                                    compute_backend,
+                                                    compute_backend,
+                                                    lora_spec.path,
+                                                    lora_spec.is_high_noise ? "model.high_noise_" : "",
+                                                    lora_version_);
+
+            LoraModel::filter_t lora_tensor_filter = nullptr;
+            if (!lora_spec.tensor_name_prefix_filter.empty()) {
+                lora_tensor_filter = [&](const std::string& tensor_name) {
+                    return starts_with(tensor_name, lora_spec.tensor_name_prefix_filter);
+                };
+            }
+            if (!lora->load_from_file(n_threads_, lora_tensor_filter)) {
+                LOG_WARN("load lora tensors from %s failed", lora_spec.path.c_str());
+                if (lora_spec.required) {
+                    return false;
+                }
+                continue;
+            }
+            if (lora->lora_tensors.empty()) {
+                if (lora_spec.required) {
+                    LOG_ERROR("required lora has no tensors: %s", lora_spec.path.c_str());
+                    return false;
+                }
+                continue;
+            }
+            lora->multiplier = lora_spec.multiplier;
+            lora->apply(group.model_tensors, all_tensor_names, lora_version_, n_threads_, false);
+            lora->release_loaded_tensors();
+        }
+
+        for (TensorState* state : group.states) {
+            if (state != nullptr) {
+                state->applied_lora_epoch = current_lora_epoch_;
+            }
+        }
+    }
+    return true;
+}
+
+void ModelManager::reset_lora_applied_params() {
+    release_compute_staging_blocks(true);
+    release_params_storage_blocks(true);
+    for (auto& state : tensor_states_) {
+        state->applied_lora_epoch = UINT64_MAX;
+    }
+}
+
+bool ModelManager::should_ignore(const TensorState& state) const {
+    for (const auto& ignore_prefix : common_ignore_tensors_) {
+        if (starts_with(state.name, ignore_prefix)) {
+            return true;
+        }
+    }
+    return false;
+}
+
+bool ModelManager::is_optional_missing_tensor(const std::string& name) const {
+    return name.find("cond_stage_model.transformer.text_model.encoder.layers.23") != std::string::npos ||
+           name.find("alphas_cumprod") != std::string::npos;
+}
+
+bool ModelManager::validate_tensor(const TensorState& state) const {
+    if (state.tensor == nullptr || should_ignore(state) || is_optional_missing_tensor(state.name)) {
+        return true;
+    }
+
+    const auto& tensor_storage_map = model_loader_.get_tensor_storage_map();
+    auto ts_it                     = tensor_storage_map.find(state.name);
+    if (ts_it == tensor_storage_map.end()) {
+        LOG_ERROR("%s tensor '%s' not in model metadata", state.desc.c_str(), state.name.c_str());
+        return false;
+    }
+
+    const TensorStorage& tensor_storage = ts_it->second;
+    if (state.tensor->ne[0] != tensor_storage.ne[0] ||
+        state.tensor->ne[1] != tensor_storage.ne[1] ||
+        state.tensor->ne[2] != tensor_storage.ne[2] ||
+        state.tensor->ne[3] != tensor_storage.ne[3]) {
+        LOG_ERROR(
+            "%s tensor '%s' has wrong shape in model metadata: got [%d, %d, %d, %d], expected [%d, %d, %d, %d]",
+            state.desc.c_str(),
+            state.name.c_str(),
+            (int)tensor_storage.ne[0], (int)tensor_storage.ne[1], (int)tensor_storage.ne[2], (int)tensor_storage.ne[3],
+            (int)state.tensor->ne[0], (int)state.tensor->ne[1], (int)state.tensor->ne[2], (int)state.tensor->ne[3]);
+        return false;
+    }
+    return true;
+}
+
+bool ModelManager::mmap_params(const std::vector<TensorState*>& states,
+                               std::vector<ParamsStorageBlock*>& created_storage_blocks) {
+    std::map<std::string, ggml_tensor*> mmap_candidates;
+    std::map<std::string, TensorState*> mmap_states;
+    for (TensorState* state : states) {
+        if (state == nullptr || !can_mmap_storage(*state) || state->tensor == nullptr ||
+            state->tensor->data != nullptr || state->tensor->view_src != nullptr) {
+            continue;
+        }
+        mmap_candidates[state->name] = state->tensor;
+        mmap_states[state->name]     = state;
+    }
+    if (mmap_candidates.empty()) {
+        return true;
+    }
+
+    auto mmap_store = model_loader_.mmap_tensors(mmap_candidates, {}, true);
+    if (mmap_store.empty()) {
+        return true;
+    }
+
+    auto block                = std::make_unique<ParamsStorageBlock>();
+    block->mmap_tensor_stores = std::move(mmap_store);
+    ParamsStorageBlock* raw   = block.get();
+    for (const auto& pair : mmap_states) {
+        TensorState* state = pair.second;
+        if (state != nullptr && state->tensor != nullptr && state->tensor->data != nullptr) {
+            block->states.push_back(state);
+        }
+    }
+
+    if (!block->states.empty()) {
+        params_storage_blocks_.push_back(std::move(block));
+        created_storage_blocks.push_back(raw);
+    }
+    return true;
+}
+
+bool ModelManager::can_mmap_storage(const TensorState& state) const {
+    if (!enable_mmap_ || state.residency_mode != ResidencyMode::ParamBackend) {
+        return false;
+    }
+    if (state.compute_backend == nullptr || state.params_backend == nullptr) {
+        return false;
+    }
+    return sd_backend_is_cpu(state.compute_backend) ||
+           sd_backend_is_cpu(state.params_backend) ||
+           backend_supports_host_buffer(state.compute_backend);
+}
+
+bool ModelManager::alloc_params_buffers(const std::vector<TensorState*>& states,
+                                        std::vector<ParamsStorageBlock*>& created_storage_blocks) {
+    std::map<std::pair<ggml_backend_buffer_type_t, int>, std::vector<TensorState*>> states_by_buffer_type;
+    for (TensorState* state : states) {
+        if (state == nullptr || state->tensor == nullptr) {
+            continue;
+        }
+        ggml_backend_buffer_type_t params_buft = params_buffer_type_for(*state);
+        if (params_buft == nullptr) {
+            return false;
+        }
+        states_by_buffer_type[{params_buft, static_cast<int>(state->residency_mode)}].push_back(state);
+    }
+
+    for (const auto& pair : states_by_buffer_type) {
+        ggml_backend_buffer_type_t params_buft  = pair.first.first;
+        const std::vector<TensorState*>& states = pair.second;
+        size_t alignment                        = ggml_backend_buft_get_alignment(params_buft);
+        size_t max_size                         = ggml_backend_buft_get_max_size(params_buft);
+
+        auto alloc_chunk = [&](const std::vector<TensorState*>& chunk, size_t chunk_size) -> bool {
+            if (chunk.empty() || chunk_size == 0) {
+                return true;
+            }
+
+            ggml_backend_buffer_t buffer = ggml_backend_buft_alloc_buffer(params_buft, chunk_size);
+            if (buffer == nullptr) {
+                LOG_ERROR("model manager alloc params backend buffer failed, size = %.2fMB",
+                          chunk_size / (1024.0 * 1024.0));
+                return false;
+            }
+            ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
+
+            std::vector<ggml_tensor*> initialized_tensors;
+            void* base    = ggml_backend_buffer_get_base(buffer);
+            size_t offset = aligned_offset(base, 0, ggml_backend_buffer_get_alignment(buffer));
+            for (TensorState* state : chunk) {
+                ggml_tensor* tensor     = state->tensor;
+                size_t tensor_size      = GGML_PAD(ggml_backend_buffer_get_alloc_size(buffer, tensor),
+                                                   ggml_backend_buffer_get_alignment(buffer));
+                enum ggml_status status = ggml_backend_tensor_alloc(buffer, tensor, static_cast<char*>(base) + offset);
+                if (status != GGML_STATUS_SUCCESS) {
+                    LOG_ERROR("model manager failed to initialize params tensor '%s'", ggml_get_name(tensor));
+                    for (ggml_tensor* initialized : initialized_tensors) {
+                        initialized->buffer = nullptr;
+                        initialized->data   = nullptr;
+                        initialized->extra  = nullptr;
+                    }
+                    LOG_DEBUG("model manager releasing params backend buffer (%6.2f MB, %zu tensors, %s)",
+                              ggml_backend_buffer_get_size(buffer) / (1024.f * 1024.f),
+                              initialized_tensors.size(),
+                              ggml_backend_buffer_is_host(buffer) ? "RAM" : "VRAM");
+                    ggml_backend_buffer_free(buffer);
+                    return false;
+                }
+                initialized_tensors.push_back(tensor);
+                offset += tensor_size;
+            }
+
+            auto block              = std::make_unique<ParamsStorageBlock>();
+            block->buffer           = buffer;
+            block->states           = chunk;
+            ParamsStorageBlock* raw = block.get();
+            params_storage_blocks_.push_back(std::move(block));
+            created_storage_blocks.push_back(raw);
+
+            return true;
+        };
+
+        std::vector<TensorState*> chunk;
+        size_t chunk_size = 0;
+        for (TensorState* state : states) {
+            ggml_tensor* tensor = state->tensor;
+            size_t tensor_size  = GGML_PAD(ggml_backend_buft_get_alloc_size(params_buft, tensor), alignment);
+            if (max_size > 0 && tensor_size > max_size) {
+                LOG_ERROR("model manager tensor '%s' is too large for params buffer: %zu > %zu",
+                          ggml_get_name(tensor),
+                          tensor_size,
+                          max_size);
+                return false;
+            }
+            if (!chunk.empty() && max_size > 0 && chunk_size + tensor_size > max_size) {
+                if (!alloc_chunk(chunk, chunk_size)) {
+                    return false;
+                }
+                chunk.clear();
+                chunk_size = 0;
+            }
+            chunk.push_back(state);
+            chunk_size += tensor_size;
+        }
+
+        if (!alloc_chunk(chunk, chunk_size)) {
+            return false;
+        }
+    }
+
+    return true;
+}
+
+bool ModelManager::load_tensors(const std::vector<TensorState*>& states) {
+    std::map<std::string, TensorState*> states_by_name;
+    std::set<std::string> target_tensor_names;
+    for (TensorState* state : states) {
+        if (state == nullptr) {
+            continue;
+        }
+        states_by_name[state->name] = state;
+        target_tensor_names.insert(state->name);
+    }
+    if (states_by_name.empty()) {
+        return true;
+    }
+
+    std::set<std::string> loaded_names;
+    std::mutex loaded_names_mutex;
+    auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool {
+        const std::string& name = tensor_storage.name;
+        *dst_tensor             = nullptr;
+
+        auto state_it = states_by_name.find(name);
+        if (state_it == states_by_name.end()) {
+            return true;
+        }
+
+        TensorState* state = state_it->second;
+        if (state == nullptr || state->tensor == nullptr) {
+            LOG_ERROR("model manager tensor '%s' is null", name.c_str());
+            return false;
+        }
+
+        if (state->tensor->ne[0] != tensor_storage.ne[0] ||
+            state->tensor->ne[1] != tensor_storage.ne[1] ||
+            state->tensor->ne[2] != tensor_storage.ne[2] ||
+            state->tensor->ne[3] != tensor_storage.ne[3]) {
+            LOG_ERROR(
+                "model manager tensor '%s' has wrong shape in model file: got [%d, %d, %d, %d], expected [%d, %d, %d, %d]",
+                name.c_str(),
+                (int)tensor_storage.ne[0], (int)tensor_storage.ne[1], (int)tensor_storage.ne[2], (int)tensor_storage.ne[3],
+                (int)state->tensor->ne[0], (int)state->tensor->ne[1], (int)state->tensor->ne[2], (int)state->tensor->ne[3]);
+            return false;
+        }
+
+        {
+            std::lock_guard<std::mutex> lock(loaded_names_mutex);
+            loaded_names.insert(name);
+        }
+        *dst_tensor = state->tensor;
+        return true;
+    };
+
+    if (!model_loader_.load_tensors(on_new_tensor_cb, enable_mmap_, &target_tensor_names)) {
+        LOG_ERROR("model manager load tensors failed");
+        return false;
+    }
+
+    bool missing = false;
+    for (const auto& pair : states_by_name) {
+        const std::string& name = pair.first;
+        if (loaded_names.find(name) == loaded_names.end()) {
+            LOG_ERROR("model manager tensor '%s' was not loaded", name.c_str());
+            missing = true;
+        }
+    }
+    if (missing) {
+        return false;
+    }
+
+    for (const auto& pair : states_by_name) {
+        pair.second->loaded_to_params_backend = true;
+    }
+    return true;
+}
+
+ggml_backend_buffer_type_t ModelManager::params_buffer_type_for(const TensorState& state) const {
+    if (state.params_backend == nullptr) {
+        LOG_ERROR("model manager params backend is null for tensor '%s'", state.name.c_str());
+        return nullptr;
+    }
+    ggml_backend_buffer_type_t params_buft = nullptr;
+    if (state.compute_backend != nullptr && state.params_backend != state.compute_backend) {
+        ggml_backend_dev_t compute_dev = ggml_backend_get_device(state.compute_backend);
+        if (compute_dev != nullptr) {
+            params_buft = ggml_backend_dev_host_buffer_type(compute_dev);
+        }
+    }
+    if (params_buft == nullptr) {
+        params_buft = ggml_backend_get_default_buffer_type(state.params_backend);
+    }
+    return params_buft;
+}
+
+void ModelManager::free_compute_staging_block(ComputeStagingBlock& block) {
+    for (auto& staged_tensor : block.staged_tensors) {
+        TensorState* state          = staged_tensor.first;
+        ggml_tensor* staging_tensor = staged_tensor.second;
+        if (state == nullptr || state->tensor == nullptr || staging_tensor == nullptr) {
+            continue;
+        }
+        ggml_tensor* managed_tensor = state->tensor;
+        managed_tensor->buffer      = staging_tensor->buffer;
+        managed_tensor->data        = staging_tensor->data;
+        managed_tensor->extra       = staging_tensor->extra;
+        staging_tensor->buffer      = nullptr;
+        staging_tensor->data        = nullptr;
+        staging_tensor->extra       = nullptr;
+
+        state->staged_to_compute_backend = false;
+        state->applied_lora_epoch        = UINT64_MAX;
+    }
+
+    if (block.buffer != nullptr) {
+        LOG_DEBUG("model manager releasing compute params (%6.2f MB, %zu tensors) from %s",
+                  ggml_backend_buffer_get_size(block.buffer) / (1024.f * 1024.f),
+                  block.staged_tensors.size(),
+                  block.compute_backend != nullptr ? ggml_backend_name(block.compute_backend) : "unknown");
+        ggml_backend_buffer_free(block.buffer);
+        block.buffer = nullptr;
+    }
+    if (block.staging_ctx != nullptr) {
+        ggml_free(block.staging_ctx);
+        block.staging_ctx = nullptr;
+    }
+    block.staged_tensors.clear();
+}
+
+void ModelManager::release_compute_staging_blocks(bool force,
+                                                  const std::unordered_set<TensorState*>* target_states) {
+    for (auto it = compute_staging_blocks_.begin(); it != compute_staging_blocks_.end();) {
+        ComputeStagingBlock* block = it->get();
+        bool can_release           = force;
+        if (!can_release) {
+            can_release = std::all_of(block->staged_tensors.begin(),
+                                      block->staged_tensors.end(),
+                                      [target_states](const std::pair<TensorState*, ggml_tensor*>& pair) {
+                                          TensorState* state = pair.first;
+                                          if (state == nullptr) {
+                                              return true;
+                                          }
+                                          if (target_states != nullptr &&
+                                              target_states->find(state) == target_states->end()) {
+                                              return false;
+                                          }
+                                          return state->active_prepare_count == 0;
+                                      });
+        }
+
+        if (can_release) {
+            free_compute_staging_block(*block);
+            it = compute_staging_blocks_.erase(it);
+        } else {
+            ++it;
+        }
+    }
+}
+
+void ModelManager::free_params_storage_block(ParamsStorageBlock& block) {
+    if (block.buffer != nullptr) {
+        LOG_DEBUG("model manager releasing params backend buffer (%6.2f MB, %zu tensors, %s)",
+                  ggml_backend_buffer_get_size(block.buffer) / (1024.f * 1024.f),
+                  block.states.size(),
+                  ggml_backend_buffer_is_host(block.buffer) ? "RAM" : "VRAM");
+        ggml_backend_buffer_free(block.buffer);
+        block.buffer = nullptr;
+    }
+    block.mmap_tensor_stores.clear();
+
+    for (TensorState* state : block.states) {
+        if (state == nullptr || state->tensor == nullptr) {
+            continue;
+        }
+        state->tensor->buffer = nullptr;
+        state->tensor->data   = nullptr;
+        state->tensor->extra  = nullptr;
+
+        state->loaded_to_params_backend = false;
+        state->applied_lora_epoch       = UINT64_MAX;
+    }
+    block.states.clear();
+}
+
+void ModelManager::release_params_storage_blocks(bool force,
+                                                 const std::unordered_set<TensorState*>* target_states) {
+    for (auto it = params_storage_blocks_.begin(); it != params_storage_blocks_.end();) {
+        ParamsStorageBlock* block = it->get();
+        bool can_release          = force;
+        if (!can_release) {
+            can_release = std::all_of(block->states.begin(),
+                                      block->states.end(),
+                                      [target_states](TensorState* state) {
+                                          if (state == nullptr) {
+                                              return true;
+                                          }
+                                          if (target_states != nullptr &&
+                                              target_states->find(state) == target_states->end()) {
+                                              return false;
+                                          }
+                                          return state->active_prepare_count == 0 &&
+                                                 !state->staged_to_compute_backend &&
+                                                 state->residency_mode == ResidencyMode::Disk;
+                                      });
+        }
+
+        if (can_release) {
+            free_params_storage_block(*block);
+            it = params_storage_blocks_.erase(it);
+        } else {
+            ++it;
+        }
+    }
+}
+
+void ModelManager::erase_params_storage_block(ParamsStorageBlock* block) {
+    auto it = std::find_if(params_storage_blocks_.begin(),
+                           params_storage_blocks_.end(),
+                           [block](const std::unique_ptr<ParamsStorageBlock>& item) {
+                               return item.get() == block;
+                           });
+    if (it != params_storage_blocks_.end()) {
+        params_storage_blocks_.erase(it);
+    }
+}
+
+void ModelManager::release_all() {
+    for (auto& state : tensor_states_) {
+        state->active_prepare_count = 0;
+        state->applied_lora_epoch   = UINT64_MAX;
+    }
+    release_compute_staging_blocks(true);
+    release_params_storage_blocks(true);
+}
+
+bool ModelManager::resolve_required_tensor_states(const std::vector<ggml_tensor*>& tensors,
+                                                  std::vector<TensorState*>& required_states) const {
+    required_states.clear();
+    std::unordered_set<TensorState*> seen;
+    for (ggml_tensor* tensor : tensors) {
+        if (tensor == nullptr) {
+            continue;
+        }
+        const char* raw_name = ggml_get_name(tensor);
+        if (raw_name == nullptr || raw_name[0] == '\0') {
+            LOG_ERROR("model manager unnamed tensor is not registered");
+            return false;
+        }
+        auto state_it = tensor_states_by_name_.find(raw_name);
+        if (state_it == tensor_states_by_name_.end()) {
+            LOG_ERROR("model manager tensor '%s' is not registered", raw_name);
+            return false;
+        }
+        TensorState* state = state_it->second;
+        if (state == nullptr) {
+            LOG_ERROR("model manager tensor '%s' has no tensor state", raw_name);
+            return false;
+        }
+        if (seen.insert(state).second) {
+            required_states.push_back(state);
+        }
+    }
+    return true;
+}
+
+bool ModelManager::prepare_params(const std::vector<ggml_tensor*>& tensors) {
+    if (tensors.empty()) {
+        return true;
+    }
+
+    std::vector<TensorState*> required_states;
+    if (!resolve_required_tensor_states(tensors, required_states)) {
+        return false;
+    }
+
+    if (!load_tensors_to_params_backend(required_states)) {
+        return false;
+    }
+
+    if (!stage_tensors_to_compute_backend(required_states)) {
+        release_compute_staging_blocks(false);
+        release_params_storage_blocks(false);
+        return false;
+    }
+
+    if (!apply_loras_to_params(required_states)) {
+        release_compute_staging_blocks(false);
+        release_params_storage_blocks(false);
+        return false;
+    }
+
+    for (TensorState* state : required_states) {
+        if (state == nullptr) {
+            continue;
+        }
+        state->active_prepare_count++;
+    }
+    return true;
+}
+
+void ModelManager::finish_compute_backend_usage(const std::vector<TensorState*>& states) {
+    if (states.empty()) {
+        return;
+    }
+
+    std::unordered_set<TensorState*> target_states;
+    for (TensorState* state : states) {
+        if (state == nullptr || !target_states.insert(state).second) {
+            continue;
+        }
+        if (state->active_prepare_count > 0) {
+            state->active_prepare_count--;
+        }
+    }
+    release_compute_staging_blocks(false, &target_states);
+}
+
+void ModelManager::release_compute_backend_params(const std::vector<ggml_tensor*>& tensors) {
+    if (tensors.empty()) {
+        return;
+    }
+    std::vector<TensorState*> required_states;
+    if (!resolve_required_tensor_states(tensors, required_states)) {
+        return;
+    }
+    finish_compute_backend_usage(required_states);
+}
+
+void ModelManager::release_params_backend_params(const std::vector<ggml_tensor*>& tensors) {
+    if (tensors.empty()) {
+        return;
+    }
+    std::vector<TensorState*> required_states;
+    if (!resolve_required_tensor_states(tensors, required_states)) {
+        return;
+    }
+    if (required_states.empty()) {
+        return;
+    }
+    std::unordered_set<TensorState*> target_states(required_states.begin(), required_states.end());
+    release_params_storage_blocks(false, &target_states);
+}

src/model_manager.h

@@ -0,0 +1,167 @@
+#ifndef __MODEL_MANAGER_H__
+#define __MODEL_MANAGER_H__
+
+#include <cstdint>
+#include <map>
+#include <memory>
+#include <set>
+#include <string>
+#include <unordered_set>
+#include <vector>
+
+#include "model_loader.h"
+#include "weight_manager.h"
+
+class ModelManager : public RunnerWeightManager {
+public:
+    enum class ResidencyMode {
+        Disk,
+        ParamBackend,
+    };
+
+    struct LoraSpec {
+        std::string path;
+        float multiplier   = 1.0f;
+        bool is_high_noise = false;
+        std::string tensor_name_prefix_filter;
+        bool required = false;
+    };
+
+private:
+    struct TensorState {
+        std::string name;
+        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 metadata_validated        = false;
+
+        int active_prepare_count = 0;
+
+        bool loaded_to_params_backend  = false;
+        bool staged_to_compute_backend = false;
+        uint64_t applied_lora_epoch    = UINT64_MAX;
+    };
+
+    struct ParamsStorageBlock {
+        ggml_backend_buffer_t buffer = nullptr;
+        std::vector<MmapTensorStore> mmap_tensor_stores;
+        std::vector<TensorState*> states;
+    };
+
+    struct ComputeStagingBlock {
+        ggml_backend_t compute_backend = nullptr;
+        ggml_backend_buffer_t buffer   = nullptr;
+        ggml_context* staging_ctx      = nullptr;
+        std::vector<std::pair<TensorState*, ggml_tensor*>> staged_tensors;
+    };
+
+    ModelLoader model_loader_;
+    std::vector<std::unique_ptr<TensorState>> tensor_states_;
+    std::map<std::string, TensorState*> tensor_states_by_name_;
+    std::vector<std::unique_ptr<ParamsStorageBlock>> params_storage_blocks_;
+    std::vector<std::unique_ptr<ComputeStagingBlock>> compute_staging_blocks_;
+    std::set<std::string> common_ignore_tensors_;
+    std::vector<LoraSpec> loras_;
+    SDVersion lora_version_      = VERSION_COUNT;
+    uint64_t current_lora_epoch_ = 0;
+    int n_threads_               = 0;
+    bool enable_mmap_            = false;
+
+    void finish_compute_backend_usage(const std::vector<TensorState*>& states);
+    void release_all();
+
+    bool resolve_required_tensor_states(const std::vector<ggml_tensor*>& tensors,
+                                        std::vector<TensorState*>& required_states) const;
+    bool should_ignore(const TensorState& state) const;
+    bool is_optional_missing_tensor(const std::string& name) const;
+    bool validate_tensor(const TensorState& state) const;
+
+    bool load_tensors_to_params_backend(const std::vector<TensorState*>& states);
+    bool apply_loras_to_params(const std::vector<TensorState*>& states);
+    bool mmap_params(const std::vector<TensorState*>& states,
+                     std::vector<ParamsStorageBlock*>& created_storage_blocks);
+    bool can_mmap_storage(const TensorState& state) const;
+    bool alloc_params_buffers(const std::vector<TensorState*>& states,
+                              std::vector<ParamsStorageBlock*>& created_storage_blocks);
+    bool load_tensors(const std::vector<TensorState*>& states);
+    bool stage_tensors_to_compute_backend(const std::vector<TensorState*>& states);
+
+    ggml_backend_buffer_type_t params_buffer_type_for(const TensorState& state) const;
+    void release_compute_staging_blocks(bool force                                            = false,
+                                        const std::unordered_set<TensorState*>* target_states = nullptr);
+    void release_params_storage_blocks(bool force                                            = false,
+                                       const std::unordered_set<TensorState*>* target_states = nullptr);
+    void free_compute_staging_block(ComputeStagingBlock& block);
+    void free_params_storage_block(ParamsStorageBlock& block);
+    void erase_params_storage_block(ParamsStorageBlock* block);
+    void reset_lora_applied_params();
+
+public:
+    ~ModelManager() override;
+
+    ModelLoader& loader() { return model_loader_; }
+    const ModelLoader& loader() const { return model_loader_; }
+
+    void set_n_threads(int n_threads) {
+        n_threads_ = n_threads;
+        model_loader_.set_n_threads(n_threads);
+    }
+    void set_enable_mmap(bool enable_mmap) { enable_mmap_ = enable_mmap; }
+    void set_common_ignore_tensors(std::set<std::string> ignore_tensors);
+    void set_loras(std::vector<LoraSpec> loras, SDVersion version);
+
+    std::set<std::string> tensor_names() const;
+
+    bool register_param_tensors(const std::string& desc,
+                                std::map<std::string, ggml_tensor*> tensors,
+                                ResidencyMode residency_mode,
+                                ggml_backend_t compute_backend,
+                                ggml_backend_t params_backend,
+                                size_t* registered_tensor_size = nullptr);
+
+    template <typename Runner>
+    bool register_runner_params(const std::string& desc,
+                                Runner& runner,
+                                ResidencyMode residency_mode,
+                                ggml_backend_t compute_backend,
+                                ggml_backend_t params_backend,
+                                size_t* registered_tensor_size = nullptr) {
+        std::map<std::string, ggml_tensor*> tensors;
+        runner.get_param_tensors(tensors);
+        return register_param_tensors(desc,
+                                      std::move(tensors),
+                                      residency_mode,
+                                      compute_backend,
+                                      params_backend,
+                                      registered_tensor_size);
+    }
+
+    template <typename Runner>
+    bool register_runner_params(const std::string& desc,
+                                Runner& runner,
+                                const std::string& prefix,
+                                ResidencyMode residency_mode,
+                                ggml_backend_t compute_backend,
+                                ggml_backend_t params_backend,
+                                size_t* registered_tensor_size = nullptr) {
+        std::map<std::string, ggml_tensor*> tensors;
+        runner.get_param_tensors(tensors, prefix);
+        return register_param_tensors(desc,
+                                      std::move(tensors),
+                                      residency_mode,
+                                      compute_backend,
+                                      params_backend,
+                                      registered_tensor_size);
+    }
+
+    bool validate_registered_tensors();
+
+    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;
+};
+
+#endif  // __MODEL_MANAGER_H__

src/name_conversion.cpp

@@ -1,8 +1,9 @@
+#include <map>
 #include <unordered_map>
 #include <unordered_set>
 
+#include "core/util.h"
 #include "name_conversion.h"
-#include "util.h"
 
 void replace_with_name_map(std::string& name, const std::vector<std::pair<std::string, std::string>>& name_map) {
     for (auto kv : name_map) {
@@ -989,7 +990,46 @@ bool is_first_stage_model_name(const std::string& name) {
     return false;
 }
 
+static std::string convert_esrgan_tensor_name(std::string name) {
+    static std::unordered_map<std::string, std::string> esrgan_name_map;
+
+    if (esrgan_name_map.empty()) {
+        esrgan_name_map["model.0."] = "conv_first.";
+
+        constexpr int max_num_blocks = 64;
+        for (int i = 0; i < max_num_blocks; i++) {
+            std::string block_prefix = "model.1.sub." + std::to_string(i) + ".";
+            for (int rdb = 1; rdb <= 3; rdb++) {
+                for (int conv = 1; conv <= 5; conv++) {
+                    esrgan_name_map[block_prefix + "RDB" + std::to_string(rdb) + ".conv" + std::to_string(conv) + ".0."] =
+                        "body." + std::to_string(i) + ".rdb" + std::to_string(rdb) + ".conv" + std::to_string(conv) + ".";
+                }
+            }
+            esrgan_name_map[block_prefix + "weight"] = "conv_body.weight";
+            esrgan_name_map[block_prefix + "bias"]   = "conv_body.bias";
+        }
+
+        // RealESRGAN stores only the learned layers in a Sequential. These indices
+        // cover the common x1, x2 and x4 layouts.
+        esrgan_name_map["model.2."]  = "conv_hr.";
+        esrgan_name_map["model.3."]  = "conv_up1.";
+        esrgan_name_map["model.4."]  = "conv_last.";
+        esrgan_name_map["model.5."]  = "conv_hr.";
+        esrgan_name_map["model.6."]  = "conv_up2.";
+        esrgan_name_map["model.7."]  = "conv_last.";
+        esrgan_name_map["model.8."]  = "conv_hr.";
+        esrgan_name_map["model.10."] = "conv_last.";
+    }
+
+    replace_with_prefix_map(name, esrgan_name_map);
+    return name;
+}
+
 std::string convert_tensor_name(std::string name, SDVersion version) {
+    if (version == VERSION_ESRGAN) {
+        return convert_esrgan_tensor_name(std::move(name));
+    }
+
     bool is_lora                             = false;
     bool is_lycoris_underline                = false;
     bool is_underline                        = false;

src/runtime/cache_dit.hpp

@@ -1,5 +1,5 @@
-#ifndef __CACHE_DIT_HPP__
-#define __CACHE_DIT_HPP__
+#ifndef __SD_RUNTIME_CACHE_DIT_HPP__
+#define __SD_RUNTIME_CACHE_DIT_HPP__
 
 #include <algorithm>
 #include <cmath>
@@ -8,9 +8,9 @@
 #include <unordered_map>
 #include <vector>
 
-#include "condition_cache_utils.hpp"
-#include "ggml_extend.hpp"
-#include "tensor.hpp"
+#include "core/ggml_extend.hpp"
+#include "core/tensor.hpp"
+#include "runtime/condition_cache_utils.hpp"
 
 struct DBCacheConfig {
     bool enabled                        = false;
@@ -893,4 +893,4 @@ struct CacheDitConditionState {
     }
 };
 
-#endif
+#endif  // __SD_RUNTIME_CACHE_DIT_HPP__

src/runtime/condition_cache_utils.hpp

@@ -1,9 +1,9 @@
-#ifndef __CONDITION_CACHE_UTILS_HPP__
-#define __CONDITION_CACHE_UTILS_HPP__
+#ifndef __SD_RUNTIME_CONDITION_CACHE_UTILS_HPP__
+#define __SD_RUNTIME_CONDITION_CACHE_UTILS_HPP__
 
 #include <vector>
 
-#include "tensor.hpp"
+#include "core/tensor.hpp"
 
 namespace sd {
 
@@ -61,4 +61,4 @@ namespace sd {
 
 }  // namespace sd
 
-#endif  // __CONDITION_CACHE_UTILS_HPP__
+#endif  // __SD_RUNTIME_CONDITION_CACHE_UTILS_HPP__

src/runtime/denoiser.hpp

@@ -1,5 +1,5 @@
-#ifndef __DENOISER_HPP__
-#define __DENOISER_HPP__
+#ifndef __SD_RUNTIME_DENOISER_HPP__
+#define __SD_RUNTIME_DENOISER_HPP__
 
 #include <algorithm>
 #include <cctype>
@@ -8,10 +8,10 @@
 #include <string>
 #include <utility>
 
-#include "ggml_extend.hpp"
-#include "gits_noise.inl"
-#include "guidance.h"
-#include "tensor.hpp"
+#include "core/ggml_extend.hpp"
+#include "core/tensor.hpp"
+#include "runtime/gits_noise.h"
+#include "runtime/guidance.h"
 
 /*================================================= CompVisDenoiser ==================================================*/
 
@@ -1902,4 +1902,4 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
     }
 }
 
-#endif  // __DENOISER_HPP__
+#endif  // __SD_RUNTIME_DENOISER_HPP__

src/runtime/easycache.hpp

@@ -1,15 +1,15 @@
-#ifndef __EASYCACHE_HPP__
-#define __EASYCACHE_HPP__
+#ifndef __SD_RUNTIME_EASYCACHE_HPP__
+#define __SD_RUNTIME_EASYCACHE_HPP__
 
 #include <cmath>
 #include <limits>
 #include <unordered_map>
 #include <vector>
 
-#include "condition_cache_utils.hpp"
-#include "denoiser.hpp"
-#include "ggml_extend.hpp"
-#include "tensor.hpp"
+#include "core/ggml_extend.hpp"
+#include "core/tensor.hpp"
+#include "runtime/condition_cache_utils.hpp"
+#include "runtime/denoiser.hpp"
 
 struct EasyCacheConfig {
     bool enabled          = false;
@@ -258,4 +258,4 @@ struct EasyCacheState {
     }
 };
 
-#endif
+#endif  // __SD_RUNTIME_EASYCACHE_HPP__

src/runtime/gits_noise.h

@@ -0,0 +1,333 @@
+#ifndef __SD_RUNTIME_GITS_NOISE_H__
+#define __SD_RUNTIME_GITS_NOISE_H__
+
+const std::vector<std::vector<float>> GITS_NOISE_0_80 = {
+    {14.61464119f, 7.49001646f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 6.77309084f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 7.49001646f, 3.07277966f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 2.05039096f, 0.02916753f},
+    {14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.85520077f, 2.05039096f, 0.02916753f},
+    {14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.85520077f, 3.07277966f, 1.56271636f, 0.02916753f},
+    {14.61464119f, 12.96784878f, 11.54541874f, 8.75849152f, 7.49001646f, 5.85520077f, 3.07277966f, 1.56271636f, 0.02916753f},
+    {14.61464119f, 13.76078796f, 12.23089790f, 10.90732002f, 8.75849152f, 7.49001646f, 5.85520077f, 3.07277966f, 1.56271636f, 0.02916753f},
+    {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 10.90732002f, 8.75849152f, 7.49001646f, 5.85520077f, 3.07277966f, 1.56271636f, 0.02916753f},
+    {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 5.85520077f, 3.07277966f, 1.56271636f, 0.02916753f},
+    {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.07277966f, 1.56271636f, 0.02916753f},
+    {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.07277966f, 1.56271636f, 0.02916753f},
+    {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.07277966f, 1.56271636f, 0.02916753f},
+    {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.07277966f, 1.56271636f, 0.02916753f},
+    {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.75859547f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.19567990f, 1.98035145f, 0.86115354f, 0.02916753f},
+    {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.75859547f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.19567990f, 1.98035145f, 0.86115354f, 0.02916753f},
+    {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.75859547f, 9.24142551f, 8.75849152f, 8.30717278f, 7.88507891f, 7.49001646f, 6.77309084f, 5.85520077f, 4.65472794f, 3.07277966f, 1.84880662f, 0.83188516f, 0.02916753f}};
+
+const std::vector<std::vector<float>> GITS_NOISE_0_85 = {
+    {14.61464119f, 7.49001646f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 1.84880662f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 6.77309084f, 1.56271636f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 7.11996698f, 3.07277966f, 1.24153244f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 7.49001646f, 5.09240818f, 2.84484982f, 0.95350921f, 0.02916753f},
+    {14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.09240818f, 2.84484982f, 0.95350921f, 0.02916753f},
+    {14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.58536053f, 3.19567990f, 1.84880662f, 0.803307f, 0.02916753f},
+    {14.61464119f, 12.96784878f, 11.54541874f, 8.75849152f, 7.49001646f, 5.58536053f, 3.19567990f, 1.84880662f, 0.803307f, 0.02916753f},
+    {14.61464119f, 12.96784878f, 11.54541874f, 8.75849152f, 7.49001646f, 6.14220476f, 4.65472794f, 3.07277966f, 1.84880662f, 0.803307f, 0.02916753f},
+    {14.61464119f, 13.76078796f, 12.23089790f, 10.90732002f, 8.75849152f, 7.49001646f, 6.14220476f, 4.65472794f, 3.07277966f, 1.84880662f, 0.803307f, 0.02916753f},
+    {14.61464119f, 13.76078796f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.65472794f, 3.07277966f, 1.84880662f, 0.803307f, 0.02916753f},
+    {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.65472794f, 3.07277966f, 1.84880662f, 0.803307f, 0.02916753f},
+    {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.65472794f, 3.07277966f, 1.84880662f, 0.803307f, 0.02916753f},
+    {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.60512662f, 2.63833880f, 1.56271636f, 0.72133851f, 0.02916753f},
+    {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 4.65472794f, 3.46139455f, 2.45070267f, 1.56271636f, 0.72133851f, 0.02916753f},
+    {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 4.65472794f, 3.46139455f, 2.45070267f, 1.56271636f, 0.72133851f, 0.02916753f},
+    {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 4.65472794f, 3.46139455f, 2.45070267f, 1.56271636f, 0.72133851f, 0.02916753f},
+    {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.75859547f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 4.65472794f, 3.46139455f, 2.45070267f, 1.56271636f, 0.72133851f, 0.02916753f},
+    {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.75859547f, 9.24142551f, 8.75849152f, 8.30717278f, 7.88507891f, 7.49001646f, 6.77309084f, 5.85520077f, 4.65472794f, 3.46139455f, 2.45070267f, 1.56271636f, 0.72133851f, 0.02916753f}};
+
+const std::vector<std::vector<float>> GITS_NOISE_0_90 = {
+    {14.61464119f, 6.77309084f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 1.56271636f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 3.07277966f, 0.95350921f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 4.86714602f, 2.54230714f, 0.89115214f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 7.49001646f, 4.86714602f, 2.54230714f, 0.89115214f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 7.49001646f, 5.09240818f, 3.07277966f, 1.61558151f, 0.69515091f, 0.02916753f},
+    {14.61464119f, 12.23089790f, 8.75849152f, 7.11996698f, 4.86714602f, 3.07277966f, 1.61558151f, 0.69515091f, 0.02916753f},
+    {14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.85520077f, 4.45427561f, 2.95596409f, 1.61558151f, 0.69515091f, 0.02916753f},
+    {14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.19988537f, 1.24153244f, 0.57119018f, 0.02916753f},
+    {14.61464119f, 12.96784878f, 10.90732002f, 8.75849152f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.19988537f, 1.24153244f, 0.57119018f, 0.02916753f},
+    {14.61464119f, 12.96784878f, 11.54541874f, 9.24142551f, 8.30717278f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.19988537f, 1.24153244f, 0.57119018f, 0.02916753f},
+    {14.61464119f, 12.96784878f, 11.54541874f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.75677586f, 2.84484982f, 1.84880662f, 1.08895338f, 0.52423614f, 0.02916753f},
+    {14.61464119f, 13.76078796f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.75677586f, 2.84484982f, 1.84880662f, 1.08895338f, 0.52423614f, 0.02916753f},
+    {14.61464119f, 13.76078796f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.44769001f, 5.58536053f, 4.45427561f, 3.32507086f, 2.45070267f, 1.61558151f, 0.95350921f, 0.45573691f, 0.02916753f},
+    {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.44769001f, 5.58536053f, 4.45427561f, 3.32507086f, 2.45070267f, 1.61558151f, 0.95350921f, 0.45573691f, 0.02916753f},
+    {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 4.86714602f, 3.91689563f, 3.07277966f, 2.27973175f, 1.56271636f, 0.95350921f, 0.45573691f, 0.02916753f},
+    {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 4.86714602f, 3.91689563f, 3.07277966f, 2.27973175f, 1.56271636f, 0.95350921f, 0.45573691f, 0.02916753f},
+    {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 4.86714602f, 3.91689563f, 3.07277966f, 2.27973175f, 1.56271636f, 0.95350921f, 0.45573691f, 0.02916753f},
+    {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 5.09240818f, 4.45427561f, 3.60512662f, 2.95596409f, 2.19988537f, 1.51179266f, 0.89115214f, 0.43325692f, 0.02916753f}};
+
+const std::vector<std::vector<float>> GITS_NOISE_0_95 = {
+    {14.61464119f, 6.77309084f, 0.02916753f},
+    {14.61464119f, 6.77309084f, 1.56271636f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 2.84484982f, 0.89115214f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 4.86714602f, 2.36326075f, 0.803307f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 4.86714602f, 2.95596409f, 1.56271636f, 0.64427125f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 7.49001646f, 4.86714602f, 2.95596409f, 1.56271636f, 0.64427125f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 7.49001646f, 4.86714602f, 3.07277966f, 1.91321158f, 1.08895338f, 0.50118381f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.07277966f, 1.91321158f, 1.08895338f, 0.50118381f, 0.02916753f},
+    {14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.85520077f, 4.45427561f, 3.07277966f, 1.91321158f, 1.08895338f, 0.50118381f, 0.02916753f},
+    {14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.19988537f, 1.41535246f, 0.803307f, 0.38853383f, 0.02916753f},
+    {14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.85520077f, 4.65472794f, 3.46139455f, 2.63833880f, 1.84880662f, 1.24153244f, 0.72133851f, 0.34370604f, 0.02916753f},
+    {14.61464119f, 12.96784878f, 10.90732002f, 8.75849152f, 7.49001646f, 5.85520077f, 4.65472794f, 3.46139455f, 2.63833880f, 1.84880662f, 1.24153244f, 0.72133851f, 0.34370604f, 0.02916753f},
+    {14.61464119f, 12.96784878f, 10.90732002f, 8.75849152f, 7.49001646f, 6.14220476f, 4.86714602f, 3.75677586f, 2.95596409f, 2.19988537f, 1.56271636f, 1.05362725f, 0.64427125f, 0.32104823f, 0.02916753f},
+    {14.61464119f, 12.96784878f, 10.90732002f, 8.75849152f, 7.49001646f, 6.44769001f, 5.58536053f, 4.65472794f, 3.60512662f, 2.95596409f, 2.19988537f, 1.56271636f, 1.05362725f, 0.64427125f, 0.32104823f, 0.02916753f},
+    {14.61464119f, 12.96784878f, 11.54541874f, 9.24142551f, 8.30717278f, 7.49001646f, 6.44769001f, 5.58536053f, 4.65472794f, 3.60512662f, 2.95596409f, 2.19988537f, 1.56271636f, 1.05362725f, 0.64427125f, 0.32104823f, 0.02916753f},
+    {14.61464119f, 12.96784878f, 11.54541874f, 9.24142551f, 8.30717278f, 7.49001646f, 6.44769001f, 5.58536053f, 4.65472794f, 3.75677586f, 3.07277966f, 2.45070267f, 1.78698075f, 1.24153244f, 0.83188516f, 0.50118381f, 0.22545385f, 0.02916753f},
+    {14.61464119f, 12.96784878f, 11.54541874f, 9.24142551f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 5.09240818f, 4.45427561f, 3.60512662f, 2.95596409f, 2.36326075f, 1.72759056f, 1.24153244f, 0.83188516f, 0.50118381f, 0.22545385f, 0.02916753f},
+    {14.61464119f, 13.76078796f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 5.09240818f, 4.45427561f, 3.60512662f, 2.95596409f, 2.36326075f, 1.72759056f, 1.24153244f, 0.83188516f, 0.50118381f, 0.22545385f, 0.02916753f},
+    {14.61464119f, 13.76078796f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 5.09240818f, 4.45427561f, 3.75677586f, 3.07277966f, 2.45070267f, 1.91321158f, 1.46270394f, 1.05362725f, 0.72133851f, 0.43325692f, 0.19894916f, 0.02916753f}};
+
+const std::vector<std::vector<float>> GITS_NOISE_1_00 = {
+    {14.61464119f, 1.56271636f, 0.02916753f},
+    {14.61464119f, 6.77309084f, 0.95350921f, 0.02916753f},
+    {14.61464119f, 6.77309084f, 2.36326075f, 0.803307f, 0.02916753f},
+    {14.61464119f, 7.11996698f, 3.07277966f, 1.56271636f, 0.59516323f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 4.86714602f, 2.84484982f, 1.41535246f, 0.57119018f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 4.86714602f, 2.84484982f, 1.61558151f, 0.86115354f, 0.38853383f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 7.49001646f, 4.86714602f, 2.84484982f, 1.61558151f, 0.86115354f, 0.38853383f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 7.49001646f, 4.86714602f, 3.07277966f, 1.98035145f, 1.24153244f, 0.72133851f, 0.34370604f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.07277966f, 1.98035145f, 1.24153244f, 0.72133851f, 0.34370604f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.27973175f, 1.51179266f, 0.95350921f, 0.54755926f, 0.25053367f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.36326075f, 1.61558151f, 1.08895338f, 0.72133851f, 0.41087446f, 0.17026083f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 8.75849152f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.36326075f, 1.61558151f, 1.08895338f, 0.72133851f, 0.41087446f, 0.17026083f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 8.75849152f, 7.49001646f, 5.85520077f, 4.65472794f, 3.60512662f, 2.84484982f, 2.12350607f, 1.56271636f, 1.08895338f, 0.72133851f, 0.41087446f, 0.17026083f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 8.75849152f, 7.49001646f, 5.85520077f, 4.65472794f, 3.60512662f, 2.84484982f, 2.19988537f, 1.61558151f, 1.162866f, 0.803307f, 0.50118381f, 0.27464288f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 8.75849152f, 7.49001646f, 5.85520077f, 4.65472794f, 3.75677586f, 3.07277966f, 2.45070267f, 1.84880662f, 1.36964464f, 1.01931262f, 0.72133851f, 0.45573691f, 0.25053367f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 8.75849152f, 7.49001646f, 6.14220476f, 5.09240818f, 4.26497746f, 3.46139455f, 2.84484982f, 2.19988537f, 1.67050016f, 1.24153244f, 0.92192322f, 0.64427125f, 0.43325692f, 0.25053367f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 8.75849152f, 7.49001646f, 6.14220476f, 5.09240818f, 4.26497746f, 3.60512662f, 2.95596409f, 2.45070267f, 1.91321158f, 1.51179266f, 1.12534678f, 0.83188516f, 0.59516323f, 0.38853383f, 0.22545385f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 12.23089790f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 5.09240818f, 4.26497746f, 3.60512662f, 2.95596409f, 2.45070267f, 1.91321158f, 1.51179266f, 1.12534678f, 0.83188516f, 0.59516323f, 0.38853383f, 0.22545385f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 12.23089790f, 9.24142551f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 5.09240818f, 4.26497746f, 3.60512662f, 2.95596409f, 2.45070267f, 1.91321158f, 1.51179266f, 1.12534678f, 0.83188516f, 0.59516323f, 0.38853383f, 0.22545385f, 0.09824532f, 0.02916753f}};
+
+const std::vector<std::vector<float>> GITS_NOISE_1_05 = {
+    {14.61464119f, 0.95350921f, 0.02916753f},
+    {14.61464119f, 6.77309084f, 0.89115214f, 0.02916753f},
+    {14.61464119f, 6.77309084f, 2.05039096f, 0.72133851f, 0.02916753f},
+    {14.61464119f, 6.77309084f, 2.84484982f, 1.28281462f, 0.52423614f, 0.02916753f},
+    {14.61464119f, 6.77309084f, 3.07277966f, 1.61558151f, 0.803307f, 0.34370604f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 4.86714602f, 2.84484982f, 1.56271636f, 0.803307f, 0.34370604f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 4.86714602f, 2.84484982f, 1.61558151f, 0.95350921f, 0.52423614f, 0.22545385f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 4.86714602f, 3.07277966f, 1.98035145f, 1.24153244f, 0.74807048f, 0.41087446f, 0.17026083f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 4.86714602f, 3.19567990f, 2.27973175f, 1.51179266f, 0.95350921f, 0.59516323f, 0.34370604f, 0.13792117f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 5.09240818f, 3.46139455f, 2.45070267f, 1.61558151f, 1.08895338f, 0.72133851f, 0.45573691f, 0.25053367f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 7.49001646f, 5.09240818f, 3.46139455f, 2.45070267f, 1.61558151f, 1.08895338f, 0.72133851f, 0.45573691f, 0.25053367f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.36326075f, 1.61558151f, 1.08895338f, 0.72133851f, 0.45573691f, 0.25053367f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.45070267f, 1.72759056f, 1.24153244f, 0.86115354f, 0.59516323f, 0.38853383f, 0.22545385f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.65472794f, 3.60512662f, 2.84484982f, 2.19988537f, 1.61558151f, 1.162866f, 0.83188516f, 0.59516323f, 0.38853383f, 0.22545385f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.65472794f, 3.60512662f, 2.84484982f, 2.19988537f, 1.67050016f, 1.28281462f, 0.95350921f, 0.72133851f, 0.52423614f, 0.34370604f, 0.19894916f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.65472794f, 3.60512662f, 2.95596409f, 2.36326075f, 1.84880662f, 1.41535246f, 1.08895338f, 0.83188516f, 0.61951244f, 0.45573691f, 0.32104823f, 0.19894916f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.65472794f, 3.60512662f, 2.95596409f, 2.45070267f, 1.91321158f, 1.51179266f, 1.20157266f, 0.95350921f, 0.74807048f, 0.57119018f, 0.43325692f, 0.29807833f, 0.19894916f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 8.30717278f, 7.11996698f, 5.85520077f, 4.65472794f, 3.60512662f, 2.95596409f, 2.45070267f, 1.91321158f, 1.51179266f, 1.20157266f, 0.95350921f, 0.74807048f, 0.57119018f, 0.43325692f, 0.29807833f, 0.19894916f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 8.30717278f, 7.11996698f, 5.85520077f, 4.65472794f, 3.60512662f, 2.95596409f, 2.45070267f, 1.98035145f, 1.61558151f, 1.32549286f, 1.08895338f, 0.86115354f, 0.69515091f, 0.54755926f, 0.41087446f, 0.29807833f, 0.19894916f, 0.09824532f, 0.02916753f}};
+
+const std::vector<std::vector<float>> GITS_NOISE_1_10 = {
+    {14.61464119f, 0.89115214f, 0.02916753f},
+    {14.61464119f, 2.36326075f, 0.72133851f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 1.61558151f, 0.57119018f, 0.02916753f},
+    {14.61464119f, 6.77309084f, 2.45070267f, 1.08895338f, 0.45573691f, 0.02916753f},
+    {14.61464119f, 6.77309084f, 2.95596409f, 1.56271636f, 0.803307f, 0.34370604f, 0.02916753f},
+    {14.61464119f, 6.77309084f, 3.07277966f, 1.61558151f, 0.89115214f, 0.4783645f, 0.19894916f, 0.02916753f},
+    {14.61464119f, 6.77309084f, 3.07277966f, 1.84880662f, 1.08895338f, 0.64427125f, 0.34370604f, 0.13792117f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 4.86714602f, 2.84484982f, 1.61558151f, 0.95350921f, 0.54755926f, 0.27464288f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 4.86714602f, 2.95596409f, 1.91321158f, 1.24153244f, 0.803307f, 0.4783645f, 0.25053367f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 4.86714602f, 3.07277966f, 2.05039096f, 1.41535246f, 0.95350921f, 0.64427125f, 0.41087446f, 0.22545385f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 4.86714602f, 3.19567990f, 2.27973175f, 1.61558151f, 1.12534678f, 0.803307f, 0.54755926f, 0.36617002f, 0.22545385f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 4.86714602f, 3.32507086f, 2.45070267f, 1.72759056f, 1.24153244f, 0.89115214f, 0.64427125f, 0.45573691f, 0.32104823f, 0.19894916f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 5.09240818f, 3.60512662f, 2.84484982f, 2.05039096f, 1.51179266f, 1.08895338f, 0.803307f, 0.59516323f, 0.43325692f, 0.29807833f, 0.19894916f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 5.09240818f, 3.60512662f, 2.84484982f, 2.12350607f, 1.61558151f, 1.24153244f, 0.95350921f, 0.72133851f, 0.54755926f, 0.41087446f, 0.29807833f, 0.19894916f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.45070267f, 1.84880662f, 1.41535246f, 1.08895338f, 0.83188516f, 0.64427125f, 0.50118381f, 0.36617002f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.45070267f, 1.91321158f, 1.51179266f, 1.20157266f, 0.95350921f, 0.74807048f, 0.59516323f, 0.45573691f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 5.85520077f, 4.45427561f, 3.46139455f, 2.84484982f, 2.19988537f, 1.72759056f, 1.36964464f, 1.08895338f, 0.86115354f, 0.69515091f, 0.54755926f, 0.43325692f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.46139455f, 2.84484982f, 2.19988537f, 1.72759056f, 1.36964464f, 1.08895338f, 0.86115354f, 0.69515091f, 0.54755926f, 0.43325692f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.46139455f, 2.84484982f, 2.19988537f, 1.72759056f, 1.36964464f, 1.08895338f, 0.89115214f, 0.72133851f, 0.59516323f, 0.4783645f, 0.38853383f, 0.29807833f, 0.22545385f, 0.17026083f, 0.09824532f, 0.02916753f}};
+
+const std::vector<std::vector<float>> GITS_NOISE_1_15 = {
+    {14.61464119f, 0.83188516f, 0.02916753f},
+    {14.61464119f, 1.84880662f, 0.59516323f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 1.56271636f, 0.52423614f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 1.91321158f, 0.83188516f, 0.34370604f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.45070267f, 1.24153244f, 0.59516323f, 0.25053367f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.84484982f, 1.51179266f, 0.803307f, 0.41087446f, 0.17026083f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.84484982f, 1.56271636f, 0.89115214f, 0.50118381f, 0.25053367f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 6.77309084f, 3.07277966f, 1.84880662f, 1.12534678f, 0.72133851f, 0.43325692f, 0.22545385f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 6.77309084f, 3.07277966f, 1.91321158f, 1.24153244f, 0.803307f, 0.52423614f, 0.34370604f, 0.19894916f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 4.86714602f, 2.95596409f, 1.91321158f, 1.24153244f, 0.803307f, 0.52423614f, 0.34370604f, 0.19894916f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 4.86714602f, 3.07277966f, 2.05039096f, 1.36964464f, 0.95350921f, 0.69515091f, 0.4783645f, 0.32104823f, 0.19894916f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 4.86714602f, 3.07277966f, 2.12350607f, 1.51179266f, 1.08895338f, 0.803307f, 0.59516323f, 0.43325692f, 0.29807833f, 0.19894916f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 4.86714602f, 3.07277966f, 2.12350607f, 1.51179266f, 1.08895338f, 0.803307f, 0.59516323f, 0.45573691f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 4.86714602f, 3.07277966f, 2.19988537f, 1.61558151f, 1.24153244f, 0.95350921f, 0.74807048f, 0.59516323f, 0.45573691f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 4.86714602f, 3.19567990f, 2.45070267f, 1.78698075f, 1.32549286f, 1.01931262f, 0.803307f, 0.64427125f, 0.50118381f, 0.38853383f, 0.29807833f, 0.22545385f, 0.17026083f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 4.86714602f, 3.19567990f, 2.45070267f, 1.78698075f, 1.32549286f, 1.01931262f, 0.803307f, 0.64427125f, 0.52423614f, 0.41087446f, 0.32104823f, 0.25053367f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 4.86714602f, 3.19567990f, 2.45070267f, 1.84880662f, 1.41535246f, 1.12534678f, 0.89115214f, 0.72133851f, 0.59516323f, 0.4783645f, 0.38853383f, 0.32104823f, 0.25053367f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 4.86714602f, 3.19567990f, 2.45070267f, 1.84880662f, 1.41535246f, 1.12534678f, 0.89115214f, 0.72133851f, 0.59516323f, 0.50118381f, 0.41087446f, 0.34370604f, 0.29807833f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}};
+
+const std::vector<std::vector<float>> GITS_NOISE_1_20 = {
+    {14.61464119f, 0.803307f, 0.02916753f},
+    {14.61464119f, 1.56271636f, 0.52423614f, 0.02916753f},
+    {14.61464119f, 2.36326075f, 0.92192322f, 0.36617002f, 0.02916753f},
+    {14.61464119f, 2.84484982f, 1.24153244f, 0.59516323f, 0.25053367f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.05039096f, 0.95350921f, 0.45573691f, 0.17026083f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.45070267f, 1.24153244f, 0.64427125f, 0.29807833f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.45070267f, 1.36964464f, 0.803307f, 0.45573691f, 0.25053367f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.84484982f, 1.61558151f, 0.95350921f, 0.59516323f, 0.36617002f, 0.19894916f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.84484982f, 1.67050016f, 1.08895338f, 0.74807048f, 0.50118381f, 0.32104823f, 0.19894916f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.95596409f, 1.84880662f, 1.24153244f, 0.83188516f, 0.59516323f, 0.41087446f, 0.27464288f, 0.17026083f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 3.07277966f, 1.98035145f, 1.36964464f, 0.95350921f, 0.69515091f, 0.50118381f, 0.36617002f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 6.77309084f, 3.46139455f, 2.36326075f, 1.56271636f, 1.08895338f, 0.803307f, 0.59516323f, 0.45573691f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 6.77309084f, 3.46139455f, 2.45070267f, 1.61558151f, 1.162866f, 0.86115354f, 0.64427125f, 0.50118381f, 0.38853383f, 0.29807833f, 0.22545385f, 0.17026083f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 4.65472794f, 3.07277966f, 2.12350607f, 1.51179266f, 1.08895338f, 0.83188516f, 0.64427125f, 0.50118381f, 0.38853383f, 0.29807833f, 0.22545385f, 0.17026083f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 4.65472794f, 3.07277966f, 2.12350607f, 1.51179266f, 1.08895338f, 0.83188516f, 0.64427125f, 0.50118381f, 0.41087446f, 0.32104823f, 0.25053367f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 4.65472794f, 3.07277966f, 2.12350607f, 1.51179266f, 1.08895338f, 0.83188516f, 0.64427125f, 0.50118381f, 0.41087446f, 0.34370604f, 0.27464288f, 0.22545385f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 4.65472794f, 3.07277966f, 2.19988537f, 1.61558151f, 1.20157266f, 0.92192322f, 0.72133851f, 0.57119018f, 0.45573691f, 0.36617002f, 0.29807833f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 4.65472794f, 3.07277966f, 2.19988537f, 1.61558151f, 1.24153244f, 0.95350921f, 0.74807048f, 0.59516323f, 0.4783645f, 0.38853383f, 0.32104823f, 0.27464288f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 7.49001646f, 4.65472794f, 3.07277966f, 2.19988537f, 1.61558151f, 1.24153244f, 0.95350921f, 0.74807048f, 0.59516323f, 0.50118381f, 0.41087446f, 0.34370604f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}};
+
+const std::vector<std::vector<float>> GITS_NOISE_1_25 = {
+    {14.61464119f, 0.72133851f, 0.02916753f},
+    {14.61464119f, 1.56271636f, 0.50118381f, 0.02916753f},
+    {14.61464119f, 2.05039096f, 0.803307f, 0.32104823f, 0.02916753f},
+    {14.61464119f, 2.36326075f, 0.95350921f, 0.43325692f, 0.17026083f, 0.02916753f},
+    {14.61464119f, 2.84484982f, 1.24153244f, 0.59516323f, 0.27464288f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 3.07277966f, 1.51179266f, 0.803307f, 0.43325692f, 0.22545385f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.36326075f, 1.24153244f, 0.72133851f, 0.41087446f, 0.22545385f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.45070267f, 1.36964464f, 0.83188516f, 0.52423614f, 0.34370604f, 0.19894916f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.84484982f, 1.61558151f, 0.98595673f, 0.64427125f, 0.43325692f, 0.27464288f, 0.17026083f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.84484982f, 1.67050016f, 1.08895338f, 0.74807048f, 0.52423614f, 0.36617002f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.84484982f, 1.72759056f, 1.162866f, 0.803307f, 0.59516323f, 0.45573691f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.95596409f, 1.84880662f, 1.24153244f, 0.86115354f, 0.64427125f, 0.4783645f, 0.36617002f, 0.27464288f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.95596409f, 1.84880662f, 1.28281462f, 0.92192322f, 0.69515091f, 0.52423614f, 0.41087446f, 0.32104823f, 0.25053367f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.95596409f, 1.91321158f, 1.32549286f, 0.95350921f, 0.72133851f, 0.54755926f, 0.43325692f, 0.34370604f, 0.27464288f, 0.22545385f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.95596409f, 1.91321158f, 1.32549286f, 0.95350921f, 0.72133851f, 0.57119018f, 0.45573691f, 0.36617002f, 0.29807833f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.95596409f, 1.91321158f, 1.32549286f, 0.95350921f, 0.74807048f, 0.59516323f, 0.4783645f, 0.38853383f, 0.32104823f, 0.27464288f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 3.07277966f, 2.05039096f, 1.41535246f, 1.05362725f, 0.803307f, 0.61951244f, 0.50118381f, 0.41087446f, 0.34370604f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 3.07277966f, 2.05039096f, 1.41535246f, 1.05362725f, 0.803307f, 0.64427125f, 0.52423614f, 0.43325692f, 0.36617002f, 0.32104823f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 3.07277966f, 2.05039096f, 1.46270394f, 1.08895338f, 0.83188516f, 0.66947293f, 0.54755926f, 0.45573691f, 0.38853383f, 0.34370604f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}};
+
+const std::vector<std::vector<float>> GITS_NOISE_1_30 = {
+    {14.61464119f, 0.72133851f, 0.02916753f},
+    {14.61464119f, 1.24153244f, 0.43325692f, 0.02916753f},
+    {14.61464119f, 1.56271636f, 0.59516323f, 0.22545385f, 0.02916753f},
+    {14.61464119f, 1.84880662f, 0.803307f, 0.36617002f, 0.13792117f, 0.02916753f},
+    {14.61464119f, 2.36326075f, 1.01931262f, 0.52423614f, 0.25053367f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.84484982f, 1.36964464f, 0.74807048f, 0.41087446f, 0.22545385f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 3.07277966f, 1.56271636f, 0.89115214f, 0.54755926f, 0.34370604f, 0.19894916f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 3.07277966f, 1.61558151f, 0.95350921f, 0.61951244f, 0.41087446f, 0.27464288f, 0.17026083f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.45070267f, 1.36964464f, 0.83188516f, 0.54755926f, 0.36617002f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.45070267f, 1.41535246f, 0.92192322f, 0.64427125f, 0.45573691f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.6383388f, 1.56271636f, 1.01931262f, 0.72133851f, 0.50118381f, 0.36617002f, 0.27464288f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.84484982f, 1.61558151f, 1.05362725f, 0.74807048f, 0.54755926f, 0.41087446f, 0.32104823f, 0.25053367f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.84484982f, 1.61558151f, 1.08895338f, 0.77538133f, 0.57119018f, 0.43325692f, 0.34370604f, 0.27464288f, 0.22545385f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.84484982f, 1.61558151f, 1.08895338f, 0.803307f, 0.59516323f, 0.45573691f, 0.36617002f, 0.29807833f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.84484982f, 1.61558151f, 1.08895338f, 0.803307f, 0.59516323f, 0.4783645f, 0.38853383f, 0.32104823f, 0.27464288f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.84484982f, 1.72759056f, 1.162866f, 0.83188516f, 0.64427125f, 0.50118381f, 0.41087446f, 0.34370604f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.84484982f, 1.72759056f, 1.162866f, 0.83188516f, 0.64427125f, 0.52423614f, 0.43325692f, 0.36617002f, 0.32104823f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.84484982f, 1.78698075f, 1.24153244f, 0.92192322f, 0.72133851f, 0.57119018f, 0.45573691f, 0.38853383f, 0.34370604f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.84484982f, 1.78698075f, 1.24153244f, 0.92192322f, 0.72133851f, 0.57119018f, 0.4783645f, 0.41087446f, 0.36617002f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}};
+
+const std::vector<std::vector<float>> GITS_NOISE_1_35 = {
+    {14.61464119f, 0.69515091f, 0.02916753f},
+    {14.61464119f, 0.95350921f, 0.34370604f, 0.02916753f},
+    {14.61464119f, 1.56271636f, 0.57119018f, 0.19894916f, 0.02916753f},
+    {14.61464119f, 1.61558151f, 0.69515091f, 0.29807833f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 1.84880662f, 0.83188516f, 0.43325692f, 0.22545385f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.45070267f, 1.162866f, 0.64427125f, 0.36617002f, 0.19894916f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.84484982f, 1.36964464f, 0.803307f, 0.50118381f, 0.32104823f, 0.19894916f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.84484982f, 1.41535246f, 0.83188516f, 0.54755926f, 0.36617002f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.84484982f, 1.56271636f, 0.95350921f, 0.64427125f, 0.45573691f, 0.32104823f, 0.22545385f, 0.17026083f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.84484982f, 1.56271636f, 0.95350921f, 0.64427125f, 0.45573691f, 0.34370604f, 0.25053367f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 3.07277966f, 1.61558151f, 1.01931262f, 0.72133851f, 0.52423614f, 0.38853383f, 0.29807833f, 0.22545385f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 3.07277966f, 1.61558151f, 1.01931262f, 0.72133851f, 0.52423614f, 0.41087446f, 0.32104823f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 3.07277966f, 1.61558151f, 1.05362725f, 0.74807048f, 0.54755926f, 0.43325692f, 0.34370604f, 0.27464288f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 3.07277966f, 1.72759056f, 1.12534678f, 0.803307f, 0.59516323f, 0.45573691f, 0.36617002f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 3.07277966f, 1.72759056f, 1.12534678f, 0.803307f, 0.59516323f, 0.4783645f, 0.38853383f, 0.32104823f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.45070267f, 1.51179266f, 1.01931262f, 0.74807048f, 0.57119018f, 0.45573691f, 0.36617002f, 0.32104823f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.6383388f, 1.61558151f, 1.08895338f, 0.803307f, 0.61951244f, 0.50118381f, 0.41087446f, 0.34370604f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.6383388f, 1.61558151f, 1.08895338f, 0.803307f, 0.64427125f, 0.52423614f, 0.43325692f, 0.36617002f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 5.85520077f, 2.6383388f, 1.61558151f, 1.08895338f, 0.803307f, 0.64427125f, 0.52423614f, 0.45573691f, 0.38853383f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}};
+
+const std::vector<std::vector<float>> GITS_NOISE_1_40 = {
+    {14.61464119f, 0.59516323f, 0.02916753f},
+    {14.61464119f, 0.95350921f, 0.34370604f, 0.02916753f},
+    {14.61464119f, 1.08895338f, 0.43325692f, 0.13792117f, 0.02916753f},
+    {14.61464119f, 1.56271636f, 0.64427125f, 0.27464288f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 1.61558151f, 0.803307f, 0.43325692f, 0.22545385f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.05039096f, 0.95350921f, 0.54755926f, 0.34370604f, 0.19894916f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.45070267f, 1.24153244f, 0.72133851f, 0.43325692f, 0.27464288f, 0.17026083f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.45070267f, 1.24153244f, 0.74807048f, 0.50118381f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.45070267f, 1.28281462f, 0.803307f, 0.52423614f, 0.36617002f, 0.27464288f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.45070267f, 1.28281462f, 0.803307f, 0.54755926f, 0.38853383f, 0.29807833f, 0.22545385f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.84484982f, 1.41535246f, 0.86115354f, 0.59516323f, 0.43325692f, 0.32104823f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.84484982f, 1.51179266f, 0.95350921f, 0.64427125f, 0.45573691f, 0.34370604f, 0.27464288f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.84484982f, 1.51179266f, 0.95350921f, 0.64427125f, 0.4783645f, 0.36617002f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.84484982f, 1.56271636f, 0.98595673f, 0.69515091f, 0.52423614f, 0.41087446f, 0.34370604f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.84484982f, 1.56271636f, 1.01931262f, 0.72133851f, 0.54755926f, 0.43325692f, 0.36617002f, 0.32104823f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.84484982f, 1.61558151f, 1.05362725f, 0.74807048f, 0.57119018f, 0.45573691f, 0.38853383f, 0.34370604f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.84484982f, 1.61558151f, 1.08895338f, 0.803307f, 0.61951244f, 0.50118381f, 0.41087446f, 0.36617002f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.84484982f, 1.61558151f, 1.08895338f, 0.803307f, 0.61951244f, 0.50118381f, 0.43325692f, 0.38853383f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.84484982f, 1.61558151f, 1.08895338f, 0.803307f, 0.64427125f, 0.52423614f, 0.45573691f, 0.41087446f, 0.36617002f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}};
+
+const std::vector<std::vector<float>> GITS_NOISE_1_45 = {
+    {14.61464119f, 0.59516323f, 0.02916753f},
+    {14.61464119f, 0.803307f, 0.25053367f, 0.02916753f},
+    {14.61464119f, 0.95350921f, 0.34370604f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 1.24153244f, 0.54755926f, 0.25053367f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 1.56271636f, 0.72133851f, 0.36617002f, 0.19894916f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 1.61558151f, 0.803307f, 0.45573691f, 0.27464288f, 0.17026083f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 1.91321158f, 0.95350921f, 0.57119018f, 0.36617002f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.19988537f, 1.08895338f, 0.64427125f, 0.41087446f, 0.27464288f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.45070267f, 1.24153244f, 0.74807048f, 0.50118381f, 0.34370604f, 0.25053367f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.45070267f, 1.24153244f, 0.74807048f, 0.50118381f, 0.36617002f, 0.27464288f, 0.22545385f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.45070267f, 1.28281462f, 0.803307f, 0.54755926f, 0.41087446f, 0.32104823f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.45070267f, 1.28281462f, 0.803307f, 0.57119018f, 0.43325692f, 0.34370604f, 0.27464288f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.45070267f, 1.28281462f, 0.83188516f, 0.59516323f, 0.45573691f, 0.36617002f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.45070267f, 1.28281462f, 0.83188516f, 0.59516323f, 0.45573691f, 0.36617002f, 0.32104823f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.84484982f, 1.51179266f, 0.95350921f, 0.69515091f, 0.52423614f, 0.41087446f, 0.34370604f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.84484982f, 1.51179266f, 0.95350921f, 0.69515091f, 0.52423614f, 0.43325692f, 0.36617002f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.84484982f, 1.56271636f, 0.98595673f, 0.72133851f, 0.54755926f, 0.45573691f, 0.38853383f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.84484982f, 1.56271636f, 1.01931262f, 0.74807048f, 0.57119018f, 0.4783645f, 0.41087446f, 0.36617002f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.84484982f, 1.56271636f, 1.01931262f, 0.74807048f, 0.59516323f, 0.50118381f, 0.43325692f, 0.38853383f, 0.36617002f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}};
+
+const std::vector<std::vector<float>> GITS_NOISE_1_50 = {
+    {14.61464119f, 0.54755926f, 0.02916753f},
+    {14.61464119f, 0.803307f, 0.25053367f, 0.02916753f},
+    {14.61464119f, 0.86115354f, 0.32104823f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 1.24153244f, 0.54755926f, 0.25053367f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 1.56271636f, 0.72133851f, 0.36617002f, 0.19894916f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 1.61558151f, 0.803307f, 0.45573691f, 0.27464288f, 0.17026083f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 1.61558151f, 0.83188516f, 0.52423614f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 1.84880662f, 0.95350921f, 0.59516323f, 0.38853383f, 0.27464288f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 1.84880662f, 0.95350921f, 0.59516323f, 0.41087446f, 0.29807833f, 0.22545385f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 1.84880662f, 0.95350921f, 0.61951244f, 0.43325692f, 0.32104823f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.19988537f, 1.12534678f, 0.72133851f, 0.50118381f, 0.36617002f, 0.27464288f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.19988537f, 1.12534678f, 0.72133851f, 0.50118381f, 0.36617002f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.36326075f, 1.24153244f, 0.803307f, 0.57119018f, 0.43325692f, 0.34370604f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.36326075f, 1.24153244f, 0.803307f, 0.57119018f, 0.43325692f, 0.34370604f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.36326075f, 1.24153244f, 0.803307f, 0.59516323f, 0.45573691f, 0.36617002f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.36326075f, 1.24153244f, 0.803307f, 0.59516323f, 0.45573691f, 0.38853383f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.45070267f, 1.32549286f, 0.86115354f, 0.64427125f, 0.50118381f, 0.41087446f, 0.36617002f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.45070267f, 1.36964464f, 0.92192322f, 0.69515091f, 0.54755926f, 0.45573691f, 0.41087446f, 0.36617002f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f},
+    {14.61464119f, 2.45070267f, 1.41535246f, 0.95350921f, 0.72133851f, 0.57119018f, 0.4783645f, 0.43325692f, 0.38853383f, 0.36617002f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}};
+
+const std::vector<const std::vector<std::vector<float>>*> GITS_NOISE = {
+    &GITS_NOISE_0_80,
+    &GITS_NOISE_0_85,
+    &GITS_NOISE_0_90,
+    &GITS_NOISE_0_95,
+    &GITS_NOISE_1_00,
+    &GITS_NOISE_1_05,
+    &GITS_NOISE_1_10,
+    &GITS_NOISE_1_15,
+    &GITS_NOISE_1_20,
+    &GITS_NOISE_1_25,
+    &GITS_NOISE_1_30,
+    &GITS_NOISE_1_35,
+    &GITS_NOISE_1_40,
+    &GITS_NOISE_1_45,
+    &GITS_NOISE_1_50};
+
+#endif  // __SD_RUNTIME_GITS_NOISE_H__

src/runtime/guidance.cpp

@@ -1,4 +1,4 @@
-#include "guidance.h"
+#include "runtime/guidance.h"
 
 #include <algorithm>
 #include <cmath>
@@ -6,7 +6,7 @@
 #include <string>
 #include <utility>
 
-#include "util.h"
+#include "core/util.h"
 
 namespace sd::guidance {
 
@@ -82,17 +82,18 @@ namespace sd::guidance {
         output.pred                        = pred_cond;
         if (has_tensor(input.pred_uncond)) {
             const sd::Tensor<float>& pred_uncond = *input.pred_uncond;
-            if (has_tensor(input.pred_img_cond)) {
-                const sd::Tensor<float>& pred_img_cond = *input.pred_img_cond;
-                output.pred                            = pred_uncond +
-                              image_guidance_scale_ * (pred_img_cond - pred_uncond) +
-                              guidance_scale_ * (pred_cond - pred_img_cond);
+            if (has_tensor(input.pred_img_uncond)) {
+                const sd::Tensor<float>& pred_img_uncond = *input.pred_img_uncond;
+                output.pred                              = pred_img_uncond +
+                              image_guidance_scale_ * (pred_uncond - pred_img_uncond) +
+                              guidance_scale_ * (pred_cond - pred_uncond);
+
             } else {
                 output.pred = pred_uncond + guidance_scale_ * (pred_cond - pred_uncond);
             }
-        } else if (has_tensor(input.pred_img_cond)) {
-            const sd::Tensor<float>& pred_img_cond = *input.pred_img_cond;
-            output.pred                            = pred_img_cond + guidance_scale_ * (pred_cond - pred_img_cond);
+        } else if (has_tensor(input.pred_img_uncond)) {
+            const sd::Tensor<float>& pred_img_uncond = *input.pred_img_uncond;
+            output.pred                              = pred_img_uncond + guidance_scale_ * (pred_cond - pred_img_uncond);
         }
 
         return output;
@@ -108,20 +109,20 @@ namespace sd::guidance {
 
     static sd::Tensor<float> calculate_guidance_delta(const sd::Tensor<float>& pred_cond,
                                                       const sd::Tensor<float>* pred_uncond,
-                                                      const sd::Tensor<float>* pred_img_cond,
+                                                      const sd::Tensor<float>* pred_img_uncond,
                                                       float guidance_scale,
                                                       float image_guidance_scale) {
-        if (pred_img_cond != nullptr) {
+        if (pred_img_uncond != nullptr) {
             if (pred_uncond != nullptr && guidance_scale == 1.0f) {
-                return *pred_img_cond - *pred_uncond;
+                return *pred_uncond - *pred_img_uncond;
             }
             if (pred_uncond != nullptr) {
                 return pred_cond +
-                       (*pred_uncond * (1.0f - image_guidance_scale) +
-                        *pred_img_cond * (image_guidance_scale - guidance_scale)) /
+                       (*pred_uncond * (image_guidance_scale - guidance_scale) +
+                        *pred_img_uncond * (1.0f - image_guidance_scale)) /
                            (guidance_scale - 1.0f);
             }
-            return pred_cond - *pred_img_cond;
+            return pred_cond - *pred_img_uncond;
         }
         return pred_cond - *pred_uncond;
     }
@@ -139,28 +140,28 @@ namespace sd::guidance {
         output.pred                        = pred_cond;
         if (has_tensor(input.pred_uncond)) {
             const sd::Tensor<float>& pred_uncond = *input.pred_uncond;
-            if (has_tensor(input.pred_img_cond)) {
-                const sd::Tensor<float>& pred_img_cond = *input.pred_img_cond;
-                output.pred                            = pred_uncond +
-                              image_guidance_scale_ * (pred_img_cond - pred_uncond) +
-                              guidance_scale_ * (pred_cond - pred_img_cond);
+            if (has_tensor(input.pred_img_uncond)) {
+                const sd::Tensor<float>& pred_img_uncond = *input.pred_img_uncond;
+                output.pred                              = pred_img_uncond +
+                              image_guidance_scale_ * (pred_uncond - pred_img_uncond) +
+                              guidance_scale_ * (pred_cond - pred_uncond);
             } else {
                 output.pred = pred_uncond + guidance_scale_ * (pred_cond - pred_uncond);
             }
-        } else if (has_tensor(input.pred_img_cond)) {
-            const sd::Tensor<float>& pred_img_cond = *input.pred_img_cond;
-            output.pred                            = pred_img_cond + guidance_scale_ * (pred_cond - pred_img_cond);
+        } else if (has_tensor(input.pred_img_uncond)) {
+            const sd::Tensor<float>& pred_img_uncond = *input.pred_img_uncond;
+            output.pred                              = pred_img_uncond + guidance_scale_ * (pred_cond - pred_img_uncond);
         }
-        if (!has_tensor(input.pred_uncond) && !has_tensor(input.pred_img_cond)) {
+        if (!has_tensor(input.pred_uncond) && !has_tensor(input.pred_img_uncond)) {
             return output;
         }
 
-        const sd::Tensor<float>* pred_uncond   = input.pred_uncond;
-        const sd::Tensor<float>* pred_img_cond = input.pred_img_cond;
+        const sd::Tensor<float>* pred_uncond     = input.pred_uncond;
+        const sd::Tensor<float>* pred_img_uncond = input.pred_img_uncond;
 
         sd::Tensor<float> deltas = calculate_guidance_delta(pred_cond,
                                                             pred_uncond,
-                                                            pred_img_cond,
+                                                            pred_img_uncond,
                                                             guidance_scale_,
                                                             image_guidance_scale_);
         if (params_.momentum != 0.0f) {
@@ -171,9 +172,10 @@ namespace sd::guidance {
             momentum_buffer_ = deltas;
         }
 
-        float diff_norm = 0.0f;
+        float diff_norm        = 0.0f;
+        const int standard_res = 2 * 1024 / 8;  // Use SDXL as the standard resolution (1024x1024, 8x8 patches, 4=2x2 channels)
         if (params_.norm_threshold > 0.0f) {
-            diff_norm = std::sqrt((deltas * deltas).sum());
+            diff_norm = std::sqrt((deltas * deltas).sum()) * standard_res / std::sqrt(static_cast<float>(deltas.numel()));
         }
 
         float apg_scale_factor = 1.0f;
@@ -202,11 +204,11 @@ namespace sd::guidance {
         if (pred_uncond != nullptr) {
             if (guidance_scale_ != 1.0f) {
                 output.pred = pred_cond + (guidance_scale_ - 1.0f) * deltas;
-            } else if (pred_img_cond != nullptr) {
+            } else if (pred_img_uncond != nullptr) {
                 output.pred = pred_cond + (image_guidance_scale_ - 1.0f) * deltas;
             }
-        } else if (pred_img_cond != nullptr) {
-            output.pred = *pred_img_cond + guidance_scale_ * deltas;
+        } else if (pred_img_uncond != nullptr) {
+            output.pred = *pred_img_uncond + guidance_scale_ * deltas;
         }
 
         return output;

src/runtime/guidance.h

@@ -1,11 +1,11 @@
-#ifndef __SD_GUIDANCE_H__
-#define __SD_GUIDANCE_H__
+#ifndef __SD_RUNTIME_GUIDANCE_H__
+#define __SD_RUNTIME_GUIDANCE_H__
 
 #include <cstddef>
 #include <functional>
 #include <vector>
 
-#include "tensor.hpp"
+#include "core/tensor.hpp"
 
 namespace sd::guidance {
 
@@ -29,11 +29,11 @@ namespace sd::guidance {
     bool parse_skip_layer_guidance_uncond_arg(const char* extra_sample_args);
 
     struct GuidanceInput {
-        int step                               = 0;
-        size_t schedule_size                   = 0;
-        const sd::Tensor<float>* pred_cond     = nullptr;
-        const sd::Tensor<float>* pred_uncond   = nullptr;
-        const sd::Tensor<float>* pred_img_cond = nullptr;
+        int step                                 = 0;
+        size_t schedule_size                     = 0;
+        const sd::Tensor<float>* pred_cond       = nullptr;
+        const sd::Tensor<float>* pred_uncond     = nullptr;
+        const sd::Tensor<float>* pred_img_uncond = nullptr;
 
         std::function<sd::Tensor<float>()> predict_skip_layer;
     };
@@ -93,4 +93,4 @@ namespace sd::guidance {
 
 }  // namespace sd::guidance
 
-#endif  // __SD_GUIDANCE_H__
+#endif  // __SD_RUNTIME_GUIDANCE_H__

src/runtime/latent-preview.h

@@ -1,8 +1,8 @@
 #include <algorithm>
 #include <cstddef>
 #include <cstdint>
+#include "core/tensor.hpp"
 #include "ggml.h"
-#include "tensor.hpp"
 
 const float ltxav_latent_rgb_proj[128][3] = {
     {-0.0293802f, -0.0362516f, -0.0291386f},

src/runtime/preprocessing.hpp

@@ -1,10 +1,10 @@
-#ifndef __PREPROCESSING_HPP__
-#define __PREPROCESSING_HPP__
+#ifndef __SD_RUNTIME_PREPROCESSING_HPP__
+#define __SD_RUNTIME_PREPROCESSING_HPP__
 
 #include <cmath>
 #include <limits>
 
-#include "ggml_extend.hpp"
+#include "core/ggml_extend.hpp"
 
 #define M_PI_ 3.14159265358979323846f
 
@@ -331,4 +331,4 @@ bool preprocess_canny(sd_image_t img, float high_threshold, float low_threshold,
     return true;
 }
 
-#endif  // __PREPROCESSING_HPP__
+#endif  // __SD_RUNTIME_PREPROCESSING_HPP__

src/runtime/sample-cache.cpp

@@ -1,4 +1,4 @@
-#include "sample-cache.h"
+#include "runtime/sample-cache.h"
 
 namespace sd_sample {
 

src/runtime/sample-cache.h

@@ -1,16 +1,16 @@
-#ifndef __SAMPLE_CACHE_H__
-#define __SAMPLE_CACHE_H__
+#ifndef __SD_RUNTIME_SAMPLE_CACHE_H__
+#define __SD_RUNTIME_SAMPLE_CACHE_H__
 
 #include <vector>
 
-#include "cache_dit.hpp"
-#include "denoiser.hpp"
-#include "easycache.hpp"
+#include "core/tensor.hpp"
+#include "core/util.h"
 #include "model.h"
-#include "spectrum.hpp"
-#include "tensor.hpp"
-#include "ucache.hpp"
-#include "util.h"
+#include "runtime/cache_dit.hpp"
+#include "runtime/denoiser.hpp"
+#include "runtime/easycache.hpp"
+#include "runtime/spectrum.hpp"
+#include "runtime/ucache.hpp"
 
 namespace sd_sample {
 
@@ -58,4 +58,4 @@ namespace sd_sample {
 
 }  // namespace sd_sample
 
-#endif  // __SAMPLE_CACHE_H__
+#endif  // __SD_RUNTIME_SAMPLE_CACHE_H__

src/runtime/spectrum.hpp

@@ -1,12 +1,12 @@
-#ifndef __SPECTRUM_HPP__
-#define __SPECTRUM_HPP__
+#ifndef __SD_RUNTIME_SPECTRUM_HPP__
+#define __SD_RUNTIME_SPECTRUM_HPP__
 
 #include <cmath>
 #include <cstring>
 #include <vector>
 
-#include "ggml_extend.hpp"
-#include "tensor.hpp"
+#include "core/ggml_extend.hpp"
+#include "core/tensor.hpp"
 
 struct SpectrumConfig {
     float w            = 0.40f;
@@ -184,4 +184,4 @@ private:
     }
 };
 
-#endif  // __SPECTRUM_HPP__
+#endif  // __SD_RUNTIME_SPECTRUM_HPP__

src/runtime/ucache.hpp

@@ -1,15 +1,15 @@
-#ifndef __UCACHE_HPP__
-#define __UCACHE_HPP__
+#ifndef __SD_RUNTIME_UCACHE_HPP__
+#define __SD_RUNTIME_UCACHE_HPP__
 
 #include <cmath>
 #include <limits>
 #include <unordered_map>
 #include <vector>
 
-#include "condition_cache_utils.hpp"
-#include "denoiser.hpp"
-#include "ggml_extend.hpp"
-#include "tensor.hpp"
+#include "core/ggml_extend.hpp"
+#include "core/tensor.hpp"
+#include "runtime/condition_cache_utils.hpp"
+#include "runtime/denoiser.hpp"
 
 struct UCacheConfig {
     bool enabled                = false;
@@ -420,4 +420,4 @@ struct UCacheState {
     }
 };
 
-#endif  // __UCACHE_HPP__
+#endif  // __SD_RUNTIME_UCACHE_HPP__

src/tokenizers/bpe_tokenizer.cpp

@@ -3,8 +3,8 @@
 #include <algorithm>
 #include <sstream>
 
+#include "core/util.h"
 #include "tokenize_util.h"
-#include "util.h"
 
 std::vector<std::pair<int, std::u32string>> BPETokenizer::bytes_to_unicode() {
     std::vector<std::pair<int, std::u32string>> byte_unicode_pairs;

src/tokenizers/clip_tokenizer.cpp

@@ -6,9 +6,9 @@
 #include <regex>
 #include <set>
 
+#include "core/util.h"
 #include "ggml.h"
 #include "tokenize_util.h"
-#include "util.h"
 #include "vocab/vocab.h"
 
 CLIPTokenizer::CLIPTokenizer(int pad_token_id, const std::string& merges_utf8_str) {

src/tokenizers/gemma_tokenizer.cpp

@@ -1,8 +1,8 @@
 #include "gemma_tokenizer.h"
 
+#include "core/util.h"
 #include "ggml.h"
 #include "json.hpp"
-#include "util.h"
 #include "vocab/vocab.h"
 
 std::string GemmaTokenizer::normalize(const std::string& text) const {