refactor: extract Wan VAE implementation (#1614)

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

CMakeLists.txt

@@ -125,6 +125,12 @@ endif ()
 if (SD_HIPBLAS)
     message("-- Use HIPBLAS as backend stable-diffusion")
     set(GGML_HIP ON)
+    # ggml-hip's device-stub objects must be position-independent, or the
+    # default-PIE sd-cli link fails with `relocation R_X86_64_32 ... cannot be
+    # used when making a PIE object` on distros that default to PIE
+    # (Ubuntu 24.04, Fedora 40+, Debian 12+). The shared-library branch below
+    # already sets this; the static build (the HIP default) did not.
+    set(CMAKE_POSITION_INDEPENDENT_CODE ON)
 endif ()
 
 if(SD_MUSA)

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.

Dockerfile

@@ -2,7 +2,18 @@ ARG UBUNTU_VERSION=24.04
 
 FROM ubuntu:$UBUNTU_VERSION AS build
 
-RUN apt-get update && apt-get install -y --no-install-recommends build-essential git cmake
+# sd-server embeds the web UI at build time, so the build image needs Node/pnpm.
+RUN apt-get update && apt-get install -y --no-install-recommends build-essential git cmake ca-certificates curl gnupg && \
+    mkdir -p /etc/apt/keyrings && \
+    curl -fsSL https://deb.nodesource.com/gpgkey/nodesource-repo.gpg.key -o /tmp/nodesource-repo.gpg.key && \
+    gpg --dearmor -o /etc/apt/keyrings/nodesource.gpg /tmp/nodesource-repo.gpg.key && \
+    rm /tmp/nodesource-repo.gpg.key && \
+    echo "deb [signed-by=/etc/apt/keyrings/nodesource.gpg] https://deb.nodesource.com/node_20.x nodistro main" > /etc/apt/sources.list.d/nodesource.list && \
+    apt-get update && \
+    apt-get install -y --no-install-recommends nodejs && \
+    npm install -g pnpm@10.15.1 && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
 
 WORKDIR /sd.cpp
 
@@ -20,4 +31,4 @@ RUN apt-get update && \
 COPY --from=build /sd.cpp/build/bin/sd-cli /sd-cli
 COPY --from=build /sd.cpp/build/bin/sd-server /sd-server
 
-ENTRYPOINT [ "/sd-cli" ]
+ENTRYPOINT [ "/sd-cli" ]

Dockerfile.cuda

@@ -3,7 +3,18 @@ ARG UBUNTU_VERSION=24.04
 
 FROM nvidia/cuda:${CUDA_VERSION}-cudnn-devel-ubuntu${UBUNTU_VERSION} AS build
 
-RUN apt-get update && apt-get install -y --no-install-recommends build-essential git ccache cmake
+# sd-server embeds the web UI at build time, so the build image needs Node/pnpm.
+RUN apt-get update && apt-get install -y --no-install-recommends build-essential git ccache cmake ca-certificates curl gnupg && \
+    mkdir -p /etc/apt/keyrings && \
+    curl -fsSL https://deb.nodesource.com/gpgkey/nodesource-repo.gpg.key -o /tmp/nodesource-repo.gpg.key && \
+    gpg --dearmor -o /etc/apt/keyrings/nodesource.gpg /tmp/nodesource-repo.gpg.key && \
+    rm /tmp/nodesource-repo.gpg.key && \
+    echo "deb [signed-by=/etc/apt/keyrings/nodesource.gpg] https://deb.nodesource.com/node_20.x nodistro main" > /etc/apt/sources.list.d/nodesource.list && \
+    apt-get update && \
+    apt-get install -y --no-install-recommends nodejs && \
+    npm install -g pnpm@10.15.1 && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
 
 WORKDIR /sd.cpp
 

Dockerfile.musa

@@ -3,7 +3,18 @@ ARG UBUNTU_VERSION=22.04
 
 FROM mthreads/musa:${MUSA_VERSION}-devel-ubuntu${UBUNTU_VERSION}-amd64 as build
 
-RUN apt-get update && apt-get install -y ccache cmake git
+# sd-server embeds the web UI at build time, so the build image needs Node/pnpm.
+RUN apt-get update && apt-get install -y --no-install-recommends ccache cmake git ca-certificates curl gnupg && \
+    mkdir -p /etc/apt/keyrings && \
+    curl -fsSL https://deb.nodesource.com/gpgkey/nodesource-repo.gpg.key -o /tmp/nodesource-repo.gpg.key && \
+    gpg --dearmor -o /etc/apt/keyrings/nodesource.gpg /tmp/nodesource-repo.gpg.key && \
+    rm /tmp/nodesource-repo.gpg.key && \
+    echo "deb [signed-by=/etc/apt/keyrings/nodesource.gpg] https://deb.nodesource.com/node_20.x nodistro main" > /etc/apt/sources.list.d/nodesource.list && \
+    apt-get update && \
+    apt-get install -y --no-install-recommends nodejs && \
+    npm install -g pnpm@10.15.1 && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
 
 WORKDIR /sd.cpp
 
@@ -21,4 +32,4 @@ FROM mthreads/musa:${MUSA_VERSION}-runtime-ubuntu${UBUNTU_VERSION}-amd64 as runt
 COPY --from=build /sd.cpp/build/bin/sd-cli /sd-cli
 COPY --from=build /sd.cpp/build/bin/sd-server /sd-server
 
-ENTRYPOINT [ "/sd-cli" ]
+ENTRYPOINT [ "/sd-cli" ]

Dockerfile.sycl

@@ -3,7 +3,18 @@ ARG SYCL_VERSION=2025.3.2-0
 
 FROM intel/oneapi-basekit:${SYCL_VERSION}-devel-ubuntu24.04 AS build
 
-RUN apt-get update && apt-get install -y cmake
+# sd-server embeds the web UI at build time, so the build image needs Node/pnpm.
+RUN apt-get update && apt-get install -y --no-install-recommends cmake ca-certificates curl gnupg && \
+    mkdir -p /etc/apt/keyrings && \
+    curl -fsSL https://deb.nodesource.com/gpgkey/nodesource-repo.gpg.key -o /tmp/nodesource-repo.gpg.key && \
+    gpg --dearmor -o /etc/apt/keyrings/nodesource.gpg /tmp/nodesource-repo.gpg.key && \
+    rm /tmp/nodesource-repo.gpg.key && \
+    echo "deb [signed-by=/etc/apt/keyrings/nodesource.gpg] https://deb.nodesource.com/node_20.x nodistro main" > /etc/apt/sources.list.d/nodesource.list && \
+    apt-get update && \
+    apt-get install -y --no-install-recommends nodejs && \
+    npm install -g pnpm@10.15.1 && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
 
 WORKDIR /sd.cpp
 

Dockerfile.vulkan

@@ -2,7 +2,18 @@ ARG UBUNTU_VERSION=24.04
 
 FROM ubuntu:$UBUNTU_VERSION AS build
 
-RUN apt-get update && apt-get install -y --no-install-recommends build-essential git cmake libvulkan-dev glslc spirv-headers
+# sd-server embeds the web UI at build time, so the build image needs Node/pnpm.
+RUN apt-get update && apt-get install -y --no-install-recommends build-essential git cmake libvulkan-dev glslc spirv-headers ca-certificates curl gnupg && \
+    mkdir -p /etc/apt/keyrings && \
+    curl -fsSL https://deb.nodesource.com/gpgkey/nodesource-repo.gpg.key -o /tmp/nodesource-repo.gpg.key && \
+    gpg --dearmor -o /etc/apt/keyrings/nodesource.gpg /tmp/nodesource-repo.gpg.key && \
+    rm /tmp/nodesource-repo.gpg.key && \
+    echo "deb [signed-by=/etc/apt/keyrings/nodesource.gpg] https://deb.nodesource.com/node_20.x nodistro main" > /etc/apt/sources.list.d/nodesource.list && \
+    apt-get update && \
+    apt-get install -y --no-install-recommends nodejs && \
+    npm install -g pnpm@10.15.1 && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
 
 WORKDIR /sd.cpp
 

README.md

@@ -15,6 +15,8 @@ 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**
 * **2026/04/11** 🚀 stable-diffusion.cpp now uses a brand-new embedded web UI.  
@@ -42,12 +44,14 @@ API and command-line option may change frequently.***
     - [Chroma](./docs/chroma.md)
     - [Chroma1-Radiance](./docs/chroma_radiance.md)
     - [Qwen Image](./docs/qwen_image.md)
+    - [PiD](./docs/pid.md)
     - [LongCat Image](./docs/longcat_image.md)
     - [Z-Image](./docs/z_image.md)
     - [Ovis-Image](./docs/ovis_image.md)
     - [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/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/pid.md

@@ -0,0 +1,39 @@
+# How to Use
+
+PiD is NVIDIA's Pixel Diffusion Decoder. It replaces the usual VAE decode or decode-then-upscale path with a pixel-space diffusion decoder conditioned on a
+source latent and text prompt.
+
+In stable-diffusion.cpp, PiD currently runs as an image edit pipeline: provide a reference image with `-r`/`--ref-image`, encode that image with a matching VAE, then let the PiD diffusion model decode/upscale directly to RGB.
+
+## Download weights
+
+- Download PiD
+    - safetensors: https://huggingface.co/Comfy-Org/PixelDiT/tree/main/diffusion_models
+- Download Gemma 2 2B
+    - safetensors: https://huggingface.co/Comfy-Org/PixelDiT/tree/main/text_encoders
+- Download the VAE that matches the PiD checkpoint backbone
+    - safetensors: https://huggingface.co/nvidia/PiD/tree/main/checkpoints
+    - Flux / Z-Image PiD: use the Flux VAE and pass `--vae-format flux`
+    - SD3 PiD: use the SD3 VAE and pass `--vae-format sd3`
+    - Flux.2 PiD: use the Flux.2 VAE and pass `--vae-format flux2`
+
+The official PiD model card should be checked before use. At the time of the initial PiD release, the official weights are under the NSCLv1 non-commercial license.
+
+## Examples
+
+```
+.\bin\Release\sd-cli.exe --diffusion-model ..\..\ComfyUI\models\diffusion_models\pid_flux1_512_to_2048_4step_bf16.safetensors --llm "..\..\ComfyUI\models\text_encoders\gemma_2_2b_it_elm_bf16.safetensors" --vae ..\..\ComfyUI\models\vae\ae.sft --vae-format flux --cfg-scale 1.0  -p "a lovely cat" -r ..\assets\ernie_image\turbo_example.png --diffusion-fa -v --steps 4 -H 2048 -W 2048 --rng cpu
+```
+
+Before:
+
+<img width="256" alt="ERNIE-Image Turbo example" src="../assets/ernie_image/turbo_example.png" />
+
+After:
+<img width="1024" alt="PiD example" src="../assets/pid/example.png" />
+
+## Notes
+
+- `-r`/`--ref-image` is required. PiD uses the first reference image as the source latent condition.
+- `--vae-format` should match the VAE latent layout used by the PiD checkpoint. This is important when using standalone VAE files because the PiD diffusion
+  checkpoint alone does not identify the VAE format.

examples/cli/README.md

@@ -41,6 +41,8 @@ Context Options:
   --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
+  --uncond-diffusion-model <string>        path to the standalone unconditional diffusion model, currently used by
+                                           Ideogram4 CFG
   --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
@@ -104,9 +106,10 @@ Generation Options:
   --hires-upscaler <string>                highres fix upscaler, Lanczos, Nearest, Latent, Latent (nearest), Latent
                                            (nearest-exact), Latent (antialiased), Latent (bicubic), Latent (bicubic
                                            antialiased), or a model name under --hires-upscalers-dir (default: Latent)
-  --extra-sample-args <string>             extra sampler/scheduler args, key=value list. lcm supports noise_clip_std,
-                                           noise_scale_start, noise_scale_end; ltx2 supports max_shift, base_shift,
-                                           stretch, terminal; euler_ge supports gamma
+  --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)
@@ -128,8 +131,8 @@ Generation Options:
   --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 instruct-pix2pix models: (default: same
-                                           as --cfg-scale)
+  --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
@@ -139,8 +142,8 @@ Generation Options:
                                            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 instruct-pix2pix models
-                                           (default: same as --cfg-scale)
+  --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:
@@ -150,7 +153,7 @@ Generation Options:
   --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>              
+  --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

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;
         };
 

examples/common/common.cpp

@@ -35,6 +35,22 @@ const char* const modes_str[] = {
     "metadata",
 };
 
+static sd_vae_format_t str_to_vae_format(const std::string& value) {
+    if (value == "auto") {
+        return SD_VAE_FORMAT_AUTO;
+    }
+    if (value == "flux") {
+        return SD_VAE_FORMAT_FLUX;
+    }
+    if (value == "sd3") {
+        return SD_VAE_FORMAT_SD3;
+    }
+    if (value == "flux2") {
+        return SD_VAE_FORMAT_FLUX2;
+    }
+    return SD_VAE_FORMAT_COUNT;
+}
+
 #if defined(_WIN32)
 static std::string utf16_to_utf8(const std::wstring& wstr) {
     if (wstr.empty())
@@ -229,6 +245,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;
@@ -271,7 +288,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;
@@ -283,7 +300,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) {
@@ -340,6 +359,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",
@@ -348,6 +371,10 @@ ArgOptions SDContextParams::get_options() {
          "--vae",
          "path to standalone vae model",
          &vae_path},
+        {"",
+         "--vae-format",
+         "VAE latent format override: auto, flux, sd3, or flux2 (default: auto)",
+         &vae_format},
         {"",
          "--audio-vae",
          "path to standalone LTX audio vae model",
@@ -418,6 +445,10 @@ ArgOptions SDContextParams::get_options() {
     };
 
     options.bool_options = {
+        {"",
+         "--stream-layers",
+         "enable residency+prefetch streaming on top of --max-vram (no effect without --max-vram; defaults to false)",
+         true, &stream_layers},
         {"",
          "--force-sdxl-vae-conv-scale",
          "force use of conv scale on sdxl vae",
@@ -639,6 +670,11 @@ bool SDContextParams::validate(SDMode mode) {
         }
     }
 
+    if (str_to_vae_format(vae_format) == SD_VAE_FORMAT_COUNT) {
+        LOG_ERROR("error: vae_format must be 'auto', 'flux', 'sd3', or 'flux2'");
+        return false;
+    }
+
     return true;
 }
 
@@ -677,8 +713,10 @@ 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"
         << "  audio_vae_path: \"" << audio_vae_path << "\",\n"
         << "  taesd_path: \"" << taesd_path << "\",\n"
         << "  esrgan_path: \"" << esrgan_path << "\",\n"
@@ -694,6 +732,7 @@ std::string SDContextParams::to_string() const {
         << "  sampler_rng_type: " << sd_rng_type_name(sampler_rng_type) << ",\n"
         << "  offload_params_to_cpu: " << (offload_params_to_cpu ? "true" : "false") << ",\n"
         << "  max_vram: " << max_vram << ",\n"
+        << "  stream_layers: " << (stream_layers ? "true" : "false") << ",\n"
         << "  backend: \"" << backend << "\",\n"
         << "  params_backend: \"" << params_backend << "\",\n"
         << "  enable_mmap: " << (enable_mmap ? "true" : "false") << ",\n"
@@ -738,6 +777,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool vae_decode_only, bool f
         llm_vision_path.c_str(),
         diffusion_model_path.c_str(),
         high_noise_diffusion_model_path.c_str(),
+        uncond_diffusion_model_path.c_str(),
         embeddings_connectors_path.c_str(),
         vae_path.c_str(),
         audio_vae_path.c_str(),
@@ -772,7 +812,9 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool vae_decode_only, bool f
         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(),
     };
@@ -833,7 +875,7 @@ ArgOptions SDGenerationParams::get_options() {
          &hires_upscaler},
         {"",
          "--extra-sample-args",
-         "extra sampler/scheduler args, key=value list. lcm supports noise_clip_std, noise_scale_start, noise_scale_end; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma",
+         "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_sample_args},
         {"",
          "--extra-tiling-args",
@@ -913,7 +955,7 @@ ArgOptions SDGenerationParams::get_options() {
          &sample_params.guidance.txt_cfg},
         {"",
          "--img-cfg-scale",
-         "image guidance scale for inpaint or instruct-pix2pix models: (default: same as --cfg-scale)",
+         "image guidance scale for inpaint or image edit models: (default: same as --cfg-scale)",
          &sample_params.guidance.img_cfg},
         {"",
          "--guidance",
@@ -945,7 +987,7 @@ ArgOptions SDGenerationParams::get_options() {
          &high_noise_sample_params.guidance.txt_cfg},
         {"",
          "--high-noise-img-cfg-scale",
-         "(high noise) image guidance scale for inpaint or instruct-pix2pix models (default: same as --cfg-scale)",
+         "(high noise) image guidance scale for inpaint or image edit models (default: same as --cfg-scale)",
          &high_noise_sample_params.guidance.img_cfg},
         {"",
          "--high-noise-guidance",
@@ -2486,6 +2528,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);
@@ -2653,6 +2696,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,8 +123,10 @@ 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";
     std::string audio_vae_path;
     std::string taesd_path;
     std::string esrgan_path;
@@ -112,6 +145,7 @@ struct SDContextParams {
     rng_type_t sampler_rng_type = RNG_TYPE_COUNT;
     bool offload_params_to_cpu  = false;
     float max_vram              = 0.f;
+    bool stream_layers          = false;
     std::string backend;
     std::string params_backend;
     bool enable_mmap           = false;

examples/server/README.md

@@ -143,6 +143,8 @@ Context Options:
   --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
+  --uncond-diffusion-model <string>        path to the standalone unconditional diffusion model, currently used by
+                                           Ideogram4 CFG
   --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
@@ -206,9 +208,10 @@ Default Generation Options:
   --hires-upscaler <string>                highres fix upscaler, Lanczos, Nearest, Latent, Latent (nearest), Latent
                                            (nearest-exact), Latent (antialiased), Latent (bicubic), Latent (bicubic
                                            antialiased), or a model name under --hires-upscalers-dir (default: Latent)
-  --extra-sample-args <string>             extra sampler/scheduler args, key=value list. lcm supports noise_clip_std,
-                                           noise_scale_start, noise_scale_end; ltx2 supports max_shift, base_shift,
-                                           stretch, terminal; euler_ge supports gamma
+  --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)
@@ -230,8 +233,8 @@ Default Generation Options:
   --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 instruct-pix2pix models: (default: same
-                                           as --cfg-scale)
+  --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
@@ -241,8 +244,8 @@ Default Generation Options:
                                            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 instruct-pix2pix models
-                                           (default: same as --cfg-scale)
+  --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:
@@ -252,7 +255,7 @@ Default Generation Options:
   --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>              
+  --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

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;
     };
 

include/stable-diffusion.h

@@ -168,6 +168,14 @@ typedef struct {
     const char* path;
 } sd_embedding_t;
 
+enum sd_vae_format_t {
+    SD_VAE_FORMAT_AUTO = -1,
+    SD_VAE_FORMAT_FLUX,
+    SD_VAE_FORMAT_SD3,
+    SD_VAE_FORMAT_FLUX2,
+    SD_VAE_FORMAT_COUNT,
+};
+
 typedef struct {
     const char* model_path;
     const char* clip_l_path;
@@ -178,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;
@@ -212,7 +221,9 @@ typedef struct {
     bool chroma_use_t5_mask;
     int chroma_t5_mask_pad;
     bool qwen_image_zero_cond_t;
+    enum sd_vae_format_t vae_format;
     float max_vram;  // GiB budget for graph-cut segmented param offload (0 = disabled, -1 = auto free VRAM minus 1 GiB)
+    bool stream_layers;  // Enable residency+prefetch streaming on top of --max-vram (no effect without --max-vram)
     const char* backend;
     const char* params_backend;
 } sd_ctx_params_t;

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/anima.hpp

@@ -1,6 +1,7 @@
 #ifndef __ANIMA_HPP__
 #define __ANIMA_HPP__
 
+#include <algorithm>
 #include <cmath>
 #include <memory>
 #include <utility>
@@ -14,6 +15,47 @@
 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);
+            : DiffusionModelRunner(backend, params_backend, prefix),
+              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;

src/conditioner.hpp

@@ -113,11 +113,12 @@ struct Conditioner {
 public:
     virtual SDCondition get_learned_condition(int n_threads,
                                               const ConditionerParams& conditioner_params) = 0;
-    virtual void alloc_params_buffer()                                                     = 0;
+    virtual bool alloc_params_buffer()                                                     = 0;
     virtual void free_params_buffer()                                                      = 0;
     virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors)           = 0;
     virtual size_t get_params_buffer_size()                                                = 0;
     virtual void set_max_graph_vram_bytes(size_t max_vram_bytes) {}
+    virtual void set_stream_layers_enabled(bool enabled) {}
     virtual void set_flash_attention_enabled(bool enabled) = 0;
     virtual void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) {}
     virtual std::tuple<SDCondition, std::vector<bool>> get_learned_condition_with_trigger(int n_threads,
@@ -176,11 +177,16 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
         }
     }
 
-    void alloc_params_buffer() override {
-        text_model->alloc_params_buffer();
-        if (sd_version_is_sdxl(version)) {
-            text_model2->alloc_params_buffer();
+    bool alloc_params_buffer() override {
+        if (!text_model->alloc_params_buffer()) {
+            return false;
         }
+        if (sd_version_is_sdxl(version)) {
+            if (!text_model2->alloc_params_buffer()) {
+                return false;
+            }
+        }
+        return true;
     }
 
     void free_params_buffer() override {
@@ -205,6 +211,13 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
         }
     }
 
+    void set_stream_layers_enabled(bool enabled) override {
+        text_model->set_stream_layers_enabled(enabled);
+        if (sd_version_is_sdxl(version)) {
+            text_model2->set_stream_layers_enabled(enabled);
+        }
+    }
+
     void set_flash_attention_enabled(bool enabled) override {
         text_model->set_flash_attention_enabled(enabled);
         if (sd_version_is_sdxl(version)) {
@@ -781,16 +794,23 @@ struct SD3CLIPEmbedder : public Conditioner {
         }
     }
 
-    void alloc_params_buffer() override {
+    bool alloc_params_buffer() override {
         if (clip_l) {
-            clip_l->alloc_params_buffer();
+            if (!clip_l->alloc_params_buffer()) {
+                return false;
+            }
         }
         if (clip_g) {
-            clip_g->alloc_params_buffer();
+            if (!clip_g->alloc_params_buffer()) {
+                return false;
+            }
         }
         if (t5) {
-            t5->alloc_params_buffer();
+            if (!t5->alloc_params_buffer()) {
+                return false;
+            }
         }
+        return true;
     }
 
     void free_params_buffer() override {
@@ -831,6 +851,18 @@ struct SD3CLIPEmbedder : public Conditioner {
         }
     }
 
+    void set_stream_layers_enabled(bool enabled) override {
+        if (clip_l) {
+            clip_l->set_stream_layers_enabled(enabled);
+        }
+        if (clip_g) {
+            clip_g->set_stream_layers_enabled(enabled);
+        }
+        if (t5) {
+            t5->set_stream_layers_enabled(enabled);
+        }
+    }
+
     void set_flash_attention_enabled(bool enabled) override {
         if (clip_l) {
             clip_l->set_flash_attention_enabled(enabled);
@@ -1145,13 +1177,18 @@ struct FluxCLIPEmbedder : public Conditioner {
         }
     }
 
-    void alloc_params_buffer() override {
+    bool alloc_params_buffer() override {
         if (clip_l) {
-            clip_l->alloc_params_buffer();
+            if (!clip_l->alloc_params_buffer()) {
+                return false;
+            }
         }
         if (t5) {
-            t5->alloc_params_buffer();
+            if (!t5->alloc_params_buffer()) {
+                return false;
+            }
         }
+        return true;
     }
 
     void free_params_buffer() override {
@@ -1183,6 +1220,15 @@ struct FluxCLIPEmbedder : public Conditioner {
         }
     }
 
+    void set_stream_layers_enabled(bool enabled) override {
+        if (clip_l) {
+            clip_l->set_stream_layers_enabled(enabled);
+        }
+        if (t5) {
+            t5->set_stream_layers_enabled(enabled);
+        }
+    }
+
     void set_flash_attention_enabled(bool enabled) override {
         if (clip_l) {
             clip_l->set_flash_attention_enabled(enabled);
@@ -1388,10 +1434,13 @@ struct T5CLIPEmbedder : public Conditioner {
         }
     }
 
-    void alloc_params_buffer() override {
+    bool alloc_params_buffer() override {
         if (t5) {
-            t5->alloc_params_buffer();
+            if (!t5->alloc_params_buffer()) {
+                return false;
+            }
         }
+        return true;
     }
 
     void free_params_buffer() override {
@@ -1414,6 +1463,12 @@ struct T5CLIPEmbedder : public Conditioner {
         }
     }
 
+    void set_stream_layers_enabled(bool enabled) override {
+        if (t5) {
+            t5->set_stream_layers_enabled(enabled);
+        }
+    }
+
     void set_flash_attention_enabled(bool enabled) override {
         if (t5) {
             t5->set_flash_attention_enabled(enabled);
@@ -1578,8 +1633,11 @@ struct AnimaConditioner : public Conditioner {
         llm->get_param_tensors(tensors, "text_encoders.llm");
     }
 
-    void alloc_params_buffer() override {
-        llm->alloc_params_buffer();
+    bool alloc_params_buffer() override {
+        if (!llm->alloc_params_buffer()) {
+            return false;
+        }
+        return true;
     }
 
     void free_params_buffer() override {
@@ -1594,6 +1652,10 @@ struct AnimaConditioner : public Conditioner {
         llm->set_max_graph_vram_bytes(max_vram_bytes);
     }
 
+    void set_stream_layers_enabled(bool enabled) override {
+        llm->set_stream_layers_enabled(enabled);
+    }
+
     void set_flash_attention_enabled(bool enabled) override {
         llm->set_flash_attention_enabled(enabled);
     }
@@ -1695,6 +1757,10 @@ struct LLMEmbedder : public Conditioner {
             arch = LLM::LLMArch::MINISTRAL_3_3B;
         } else if (sd_version_is_lens(version)) {
             arch = LLM::LLMArch::GPT_OSS_20B;
+        } else if (sd_version_is_pid(version)) {
+            arch = LLM::LLMArch::GEMMA2_2B;
+        } else if (sd_version_is_ideogram4(version)) {
+            arch = LLM::LLMArch::QWEN3_VL;
         } else if (sd_version_is_z_image(version) || version == VERSION_OVIS_IMAGE || version == VERSION_FLUX2_KLEIN) {
             arch = LLM::LLMArch::QWEN3;
         }
@@ -1702,6 +1768,8 @@ struct LLMEmbedder : public Conditioner {
             tokenizer = std::make_shared<MistralTokenizer>();
         } else if (arch == LLM::LLMArch::GPT_OSS_20B) {
             tokenizer = std::make_shared<GPTOSSTokenizer>();
+        } else if (arch == LLM::LLMArch::GEMMA2_2B) {
+            tokenizer = std::make_shared<Gemma2Tokenizer>();
         } else {
             tokenizer = std::make_shared<Qwen2Tokenizer>();
         }
@@ -1717,8 +1785,11 @@ struct LLMEmbedder : public Conditioner {
         llm->get_param_tensors(tensors, "text_encoders.llm");
     }
 
-    void alloc_params_buffer() override {
-        llm->alloc_params_buffer();
+    bool alloc_params_buffer() override {
+        if (!llm->alloc_params_buffer()) {
+            return false;
+        }
+        return true;
     }
 
     void free_params_buffer() override {
@@ -1735,6 +1806,10 @@ struct LLMEmbedder : public Conditioner {
         llm->set_max_graph_vram_bytes(max_vram_bytes);
     }
 
+    void set_stream_layers_enabled(bool enabled) override {
+        llm->set_stream_layers_enabled(enabled);
+    }
+
     void set_flash_attention_enabled(bool enabled) override {
         llm->set_flash_attention_enabled(enabled);
     }
@@ -1820,12 +1895,16 @@ struct LLMEmbedder : public Conditioner {
         sd::Tensor<int32_t> input_ids({static_cast<int64_t>(tokens.size())}, tokens);
         sd::Tensor<float> attention_mask;
         if (!mask.empty()) {
-            attention_mask = sd::Tensor<float>({static_cast<int64_t>(mask.size()), static_cast<int64_t>(mask.size())});
+            attention_mask                     = sd::Tensor<float>({static_cast<int64_t>(mask.size()), static_cast<int64_t>(mask.size())});
+            const float masked_attention_value = -std::numeric_limits<float>::max() / 4.0f;
             for (size_t i1 = 0; i1 < mask.size(); ++i1) {
                 for (size_t i0 = 0; i0 < mask.size(); ++i0) {
                     float value = 0.0f;
-                    if (mask[i0] == 0.0f || i0 > i1) {
-                        value = -INFINITY;
+                    if (mask[i0] == 0.0f) {
+                        value += masked_attention_value;
+                    }
+                    if (i0 > i1) {
+                        value += masked_attention_value;
                     }
                     attention_mask[static_cast<int64_t>(i0 + mask.size() * i1)] = value;
                 }
@@ -1892,7 +1971,7 @@ struct LLMEmbedder : public Conditioner {
 
                 for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
                     const auto& image = (*conditioner_params.ref_images)[i];
-                    double factor     = llm->params.vision.patch_size * llm->params.vision.spatial_merge_size;
+                    double factor     = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
                     int height        = static_cast<int>(image.shape()[1]);
                     int width         = static_cast<int>(image.shape()[0]);
                     int h_bar         = static_cast<int>(std::round(height / factor) * factor);
@@ -1963,7 +2042,7 @@ struct LLMEmbedder : public Conditioner {
 
                 for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
                     const auto& image = (*conditioner_params.ref_images)[i];
-                    double factor     = llm->params.vision.patch_size * llm->params.vision.spatial_merge_size;
+                    double factor     = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
                     int height        = static_cast<int>(image.shape()[1]);
                     int width         = static_cast<int>(image.shape()[0]);
                     int h_bar         = static_cast<int>(std::round(height / factor) * factor);
@@ -2024,6 +2103,14 @@ struct LLMEmbedder : public Conditioner {
             prompt_attn_range.second = static_cast<int>(prompt.size());
 
             prompt += "[/INST]";
+        } else if (sd_version_is_ideogram4(version)) {
+            prompt_template_encode_start_idx = 0;
+            out_layers                       = {1, 4, 7, 10, 13, 16, 19, 22, 25, 28, 31, 34, 36};
+
+            prompt = "<|im_start|>user\n";
+            prompt += conditioner_params.text;
+            prompt += "<|im_end|>\n<|im_start|>assistant\n";
+            prompt_attn_range = {0, 0};
         } else if (sd_version_is_ernie_image(version)) {
             prompt_template_encode_start_idx = 0;
             out_layers                       = {25};  // -2
@@ -2099,6 +2186,53 @@ struct LLMEmbedder : public Conditioner {
             prompt_attn_range.second = static_cast<int>(prompt.size());
 
             prompt += "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n";
+        } else if (sd_version_is_pid(version)) {
+            constexpr int pixeldit_max_length = 300;
+            const std::string chi_prompt =
+                "Given a user prompt, generate an \"Enhanced prompt\" that provides detailed visual descriptions suitable for image generation. Evaluate the level of detail in the user prompt:\n"
+                "- If the prompt is simple, focus on adding specifics about colors, shapes, sizes, textures, and spatial relationships to create vivid and concrete scenes.\n"
+                "- If the prompt is already detailed, refine and enhance the existing details slightly without overcomplicating.\n"
+                "Here are examples of how to transform or refine prompts:\n"
+                "- User Prompt: A cat sleeping -> Enhanced: A small, fluffy white cat curled up in a round shape, sleeping peacefully on a warm sunny windowsill, surrounded by pots of blooming red flowers.\n"
+                "- User Prompt: A busy city street -> Enhanced: A bustling city street scene at dusk, featuring glowing street lamps, a diverse crowd of people in colorful clothing, and a double-decker bus passing by towering glass skyscrapers.\n"
+                "Please generate only the enhanced description for the prompt below and avoid including any additional commentary or evaluations:\n"
+                "User Prompt: ";
+            auto chi_tokens       = std::get<0>(tokenize(chi_prompt, {0, 0}));
+            size_t num_chi_tokens = chi_tokens.size();
+            max_length            = (int)num_chi_tokens + pixeldit_max_length - 2;
+            min_length            = max_length;
+
+            prompt_attn_range.first = static_cast<int>(prompt.size());
+            prompt += " " + conditioner_params.text;
+            prompt_attn_range.second = static_cast<int>(prompt.size());
+
+            auto hidden_states = encode_prompt(n_threads,
+                                               prompt,
+                                               prompt_attn_range,
+                                               min_length,
+                                               0,
+                                               image_embeds,
+                                               out_layers,
+                                               0,
+                                               false,
+                                               max_length);
+            GGML_ASSERT(!hidden_states.empty());
+
+            if (hidden_states.shape()[1] > pixeldit_max_length) {
+                auto bos      = sd::ops::slice(hidden_states, 1, 0, 1);
+                auto tail     = sd::ops::slice(hidden_states,
+                                               1,
+                                               hidden_states.shape()[1] - (pixeldit_max_length - 1),
+                                               hidden_states.shape()[1]);
+                hidden_states = sd::ops::concat(bos, tail, 1);
+            }
+
+            int64_t t1 = ggml_time_ms();
+            LOG_DEBUG("computing condition graph completed, taking %" PRId64 " ms", t1 - t0);
+
+            SDCondition result;
+            result.c_crossattn = std::move(hidden_states);
+            return result;
         } else {
             GGML_ABORT("unknown version %d", version);
         }
@@ -2239,9 +2373,14 @@ struct LTXAVEmbedder : public Conditioner {
         projector->get_param_tensors(tensors, "text_embedding_projection");
     }
 
-    void alloc_params_buffer() override {
-        llm->alloc_params_buffer();
-        projector->alloc_params_buffer();
+    bool alloc_params_buffer() override {
+        if (!llm->alloc_params_buffer()) {
+            return false;
+        }
+        if (!projector->alloc_params_buffer()) {
+            return false;
+        }
+        return true;
     }
 
     void free_params_buffer() override {
@@ -2258,6 +2397,11 @@ struct LTXAVEmbedder : public Conditioner {
         projector->set_flash_attention_enabled(enabled);
     }
 
+    void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
+        llm->set_max_graph_vram_bytes(max_vram_bytes);
+        projector->set_max_graph_vram_bytes(max_vram_bytes);
+    }
+
     void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
         llm->set_weight_adapter(adapter);
         projector->set_weight_adapter(adapter);

src/control.hpp

@@ -457,7 +457,11 @@ struct ControlNet : public GGMLRunner {
 
     bool load_from_file(const std::string& file_path, int n_threads) {
         LOG_INFO("loading control net from '%s'", file_path.c_str());
-        alloc_params_buffer();
+        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;

src/convert.cpp

@@ -8,7 +8,7 @@
 #include "model_io/safetensors_io.h"
 #include "util.h"
 
-#include "ggml-cpu.h"
+#include "ggml_extend_backend.h"
 
 static ggml_type get_export_tensor_type(ModelLoader& model_loader,
                                         const TensorStorage& tensor_storage,
@@ -103,7 +103,7 @@ bool convert(const char* input_path,
     bool output_is_safetensors = ends_with(output_path, ".safetensors");
     TensorTypeRules type_rules = parse_tensor_type_rules(tensor_type_rules);
 
-    auto backend    = ggml_backend_cpu_init();
+    auto backend    = sd_backend_cpu_init();
     size_t mem_size = 1 * 1024 * 1024;  // for padding
     mem_size += model_loader.get_tensor_storage_map().size() * ggml_tensor_overhead();
     mem_size += model_loader.get_params_mem_size(backend, type);

src/denoiser.hpp

@@ -514,8 +514,6 @@ struct LTX2Scheduler : SigmaScheduler {
                 if (!parse_strict_bool(value, stretch)) {
                     LOG_WARN("ignoring invalid ltx2 scheduler arg '%s=%s'", key.c_str(), value.c_str());
                 }
-            } else {
-                LOG_WARN("ignoring unknown ltx2 scheduler arg '%s'", key.c_str());
             }
         }
     }
@@ -1238,20 +1236,26 @@ static sd::Tensor<float> sample_lcm(denoise_cb_t model,
 
     for (const auto& [key, value] : extra_sample_args) {
         float parsed = 0.0f;
-        if (!parse_strict_float(value, parsed)) {
-            LOG_WARN("ignoring invalid lcm extra sample arg '%s=%s'", key.c_str(), value.c_str());
-            continue;
-        }
         if (key == "noise_clip_std") {
+            if (!parse_strict_float(value, parsed)) {
+                LOG_WARN("ignoring invalid lcm extra sample arg '%s=%s'", key.c_str(), value.c_str());
+                continue;
+            }
             args.noise_clip_std = parsed;
         } else if (key == "noise_scale_start") {
+            if (!parse_strict_float(value, parsed)) {
+                LOG_WARN("ignoring invalid lcm extra sample arg '%s=%s'", key.c_str(), value.c_str());
+                continue;
+            }
             args.noise_scale_start    = parsed;
             noise_scale_start_was_set = true;
         } else if (key == "noise_scale_end") {
+            if (!parse_strict_float(value, parsed)) {
+                LOG_WARN("ignoring invalid lcm extra sample arg '%s=%s'", key.c_str(), value.c_str());
+                continue;
+            }
             args.noise_scale_end    = parsed;
             noise_scale_end_was_set = true;
-        } else {
-            LOG_WARN("ignoring unknown lcm extra sample arg '%s'", key.c_str());
         }
     }
 
@@ -1795,16 +1799,14 @@ static sd::Tensor<float> sample_gradient_estimation(denoise_cb_t model,
     float ge_gamma = 2.0f;
 
     for (const auto& [key, value] : extra_sample_args) {
-        float parsed = 0.0f;
-        if (!parse_strict_float(value, parsed)) {
-            LOG_WARN("ignoring invalid euler_ge extra sample arg '%s=%s'", key.c_str(), value.c_str());
-            continue;
-        }
         if (key == "gamma") {
+            float parsed = 0.0f;
+            if (!parse_strict_float(value, parsed)) {
+                LOG_WARN("ignoring invalid euler_ge extra sample arg '%s=%s'", key.c_str(), value.c_str());
+                continue;
+            }
             LOG_DEBUG("setting euler_ge gamma to %.2f", parsed);
             ge_gamma = parsed;
-        } else {
-            LOG_WARN("ignoring unknown euler_ge extra sample arg '%s'", key.c_str());
         }
     }
 

src/ernie_image.hpp

@@ -13,6 +13,76 @@
 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,15 +374,15 @@ 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;
 
@@ -335,58 +390,9 @@ namespace ErnieImage {
                          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);
+            : DiffusionModelRunner(backend, params_backend, prefix),
+              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();

src/esrgan.hpp

@@ -270,7 +270,11 @@ struct ESRGAN : public GGMLRunner {
         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, {}, "");
 
-        alloc_params_buffer();
+        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);
 

src/flux.hpp

@@ -13,6 +13,155 @@
 
 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;
@@ -1194,114 +1306,15 @@ namespace Flux {
                    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) {
+            : DiffusionModelRunner(backend, params_backend, prefix),
+              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();
@@ -1567,7 +1580,7 @@ namespace Flux {
 
         static void load_from_file_and_test(const std::string& file_path) {
             // ggml_backend_t backend = ggml_backend_cuda_init(0);
-            ggml_backend_t backend    = ggml_backend_cpu_init();
+            ggml_backend_t backend    = sd_backend_cpu_init();
             ggml_type model_data_type = GGML_TYPE_COUNT;
 
             ModelLoader model_loader;
@@ -1592,7 +1605,11 @@ namespace Flux {
                                                                             VERSION_FLUX2,
                                                                             false);
 
-            flux->alloc_params_buffer();
+            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");
 

src/ggml_extend.hpp

@@ -28,6 +28,7 @@
 #include "ggml.h"
 #include "ggml_extend_backend.h"
 #include "ggml_graph_cut.h"
+#include "layer_registry.h"
 
 #include "model.h"
 #include "tensor.hpp"
@@ -1329,13 +1330,9 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
 
     float scale = (1.0f / sqrt((float)d_head));
 
-    int kv_pad       = 0;
     ggml_tensor* kqv = nullptr;
 
     auto build_kqv = [&](ggml_tensor* q_in, ggml_tensor* k_in, ggml_tensor* v_in, ggml_tensor* mask_in) -> ggml_tensor* {
-        if (kv_pad != 0) {
-            k_in = ggml_pad(ctx, k_in, 0, kv_pad, 0, 0);
-        }
         if (kv_scale != 1.0f) {
             k_in = ggml_ext_scale(ctx, k_in, kv_scale);
         }
@@ -1343,9 +1340,6 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
 
         v_in = ggml_ext_cont(ctx, ggml_permute(ctx, v_in, 0, 2, 1, 3));
         v_in = ggml_reshape_3d(ctx, v_in, d_head, L_k, n_kv_head * N);
-        if (kv_pad != 0) {
-            v_in = ggml_pad(ctx, v_in, 0, kv_pad, 0, 0);
-        }
         if (kv_scale != 1.0f) {
             v_in = ggml_ext_scale(ctx, v_in, kv_scale);
         }
@@ -1353,26 +1347,9 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
 
         if (mask_in != nullptr) {
             mask_in = ggml_transpose(ctx, mask_in);
-        } else {
-            if (kv_pad > 0) {
-                mask_in         = ggml_ext_zeros(ctx, L_k, L_q, 1, 1);
-                auto pad_tensor = ggml_ext_full(ctx, -INFINITY, kv_pad, L_q, 1, 1);
-                mask_in         = ggml_concat(ctx, mask_in, pad_tensor, 0);
-            }
         }
 
         if (mask_in != nullptr) {
-            // the need for padding got removed in ggml 4767bda
-            // ensure we can still use the old version for now
-#ifdef GGML_KQ_MASK_PAD
-            int mask_pad = 0;
-            if (mask_in->ne[1] % GGML_KQ_MASK_PAD != 0) {
-                mask_pad = GGML_PAD(L_q, GGML_KQ_MASK_PAD) - mask_in->ne[1];
-            }
-            if (mask_pad > 0) {
-                mask_in = ggml_pad(ctx, mask_in, 0, mask_pad, 0, 0);
-            }
-#endif
             mask_in = ggml_cast(ctx, mask_in, GGML_TYPE_F16);
         }
 
@@ -1387,10 +1364,6 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx,
     if (flash_attn) {
         // LOG_DEBUG("attention_ext L_q:%d L_k:%d n_head:%d C:%d d_head:%d N:%d", L_q, L_k, n_head, C, d_head, N);
         bool can_use_flash_attn = true;
-        if (can_use_flash_attn && L_k % 256 != 0) {
-            kv_pad = GGML_PAD(L_k, 256) - static_cast<int>(L_k);
-        }
-
         if (mask != nullptr) {
             // TODO: figure out if we can bend t5 to work too
             can_use_flash_attn = can_use_flash_attn && mask->ne[3] == 1;
@@ -1470,7 +1443,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_group_norm(ggml_context* ctx,
 
 __STATIC_INLINE__ void ggml_ext_backend_tensor_get_and_sync(ggml_backend_t backend, const ggml_tensor* tensor, void* data, size_t offset, size_t size) {
     if ((sd_backend_is(backend, "ROCm") || sd_backend_is(backend, "CUDA") || sd_backend_is(backend, "SYCL")) &&
-        !ggml_backend_is_cpu(backend)) {
+        !sd_backend_is_cpu(backend)) {
         ggml_backend_tensor_get_async(backend, tensor, data, offset, size);
         ggml_backend_synchronize(backend);
         return;
@@ -1725,7 +1698,19 @@ protected:
     ggml_context* partial_offload_ctx                   = nullptr;
     ggml_backend_buffer_t partial_runtime_params_buffer = nullptr;
     std::vector<std::pair<ggml_tensor*, ggml_tensor*>> partial_offload_pairs;
-    size_t max_graph_vram_bytes = 0;
+
+    // Params kept on the runtime backend across streaming segments.
+    ggml_context* resident_offload_ctx = nullptr;
+    std::vector<std::pair<ggml_tensor*, ggml_tensor*>> resident_offload_pairs;
+    ggml_backend_buffer_t resident_runtime_params_buffer = nullptr;
+    std::unordered_set<ggml_tensor*> resident_param_set;
+    uint64_t resident_state_token = 0;
+
+    size_t max_graph_vram_bytes           = 0;
+    bool stream_layers_enabled            = false;
+    size_t observed_max_effective_budget_ = 0;
+
+    sd::layer_registry::LayerRegistry layer_registry_;
 
     std::shared_ptr<WeightAdapter> weight_adapter = nullptr;
 
@@ -1927,7 +1912,7 @@ protected:
         LOG_DEBUG("%s compute buffer size: %.2f MB(%s)",
                   get_desc().c_str(),
                   compute_buffer_size / 1024.0 / 1024.0,
-                  ggml_backend_is_cpu(runtime_backend) ? "RAM" : "VRAM");
+                  sd_backend_is_cpu(runtime_backend) ? "RAM" : "VRAM");
         return true;
     }
 
@@ -2014,7 +1999,7 @@ protected:
         LOG_DEBUG("%s cache backend buffer size = % 6.2f MB(%s) (%i tensors)",
                   get_desc().c_str(),
                   cache_buffer_size / (1024.f * 1024.f),
-                  ggml_backend_is_cpu(runtime_backend) ? "RAM" : "VRAM",
+                  sd_backend_is_cpu(runtime_backend) ? "RAM" : "VRAM",
                   num_tensors);
         if (old_cache_buffer != nullptr) {
             ggml_backend_buffer_free(old_cache_buffer);
@@ -2193,6 +2178,9 @@ protected:
             if (tensor == nullptr) {
                 continue;
             }
+            if (resident_param_set.find(tensor) != resident_param_set.end()) {
+                continue;
+            }
             if (seen_tensors.insert(tensor).second) {
                 unique_tensors.push_back(tensor);
             }
@@ -2315,36 +2303,225 @@ protected:
         }
     }
 
+    bool offload_resident_params(const std::vector<ggml_tensor*>& tensors) {
+        if (params_backend == runtime_backend) {
+            return true;
+        }
+        if (tensors.empty()) {
+            return true;
+        }
+        GGML_ASSERT(resident_runtime_params_buffer == nullptr);
+        GGML_ASSERT(resident_offload_ctx == nullptr);
+        GGML_ASSERT(resident_offload_pairs.empty());
+        GGML_ASSERT(resident_param_set.empty());
+
+        std::vector<ggml_tensor*> unique_tensors;
+        std::unordered_set<ggml_tensor*> seen;
+        unique_tensors.reserve(tensors.size());
+        seen.reserve(tensors.size());
+        for (ggml_tensor* t : tensors) {
+            if (t == nullptr)
+                continue;
+            if (seen.insert(t).second)
+                unique_tensors.push_back(t);
+        }
+        if (unique_tensors.empty())
+            return true;
+
+        ggml_init_params init = {};
+        init.mem_size         = std::max<size_t>(1, unique_tensors.size()) * ggml_tensor_overhead();
+        init.mem_buffer       = nullptr;
+        init.no_alloc         = true;
+        resident_offload_ctx  = ggml_init(init);
+        GGML_ASSERT(resident_offload_ctx != nullptr);
+
+        resident_offload_pairs.reserve(unique_tensors.size());
+        for (ggml_tensor* t : unique_tensors) {
+            GGML_ASSERT(t->view_src == nullptr);
+            ggml_tensor* twin = ggml_dup_tensor(resident_offload_ctx, t);
+            ggml_set_name(twin, t->name);
+            resident_offload_pairs.push_back({t, twin});
+        }
+
+        resident_runtime_params_buffer = ggml_backend_alloc_ctx_tensors(resident_offload_ctx, runtime_backend);
+        if (resident_runtime_params_buffer == nullptr) {
+            LOG_ERROR("%s alloc resident runtime params backend buffer failed, num_tensors = %zu",
+                      get_desc().c_str(), resident_offload_pairs.size());
+            ggml_free(resident_offload_ctx);
+            resident_offload_ctx = nullptr;
+            resident_offload_pairs.clear();
+            return false;
+        }
+        ggml_backend_buffer_set_usage(resident_runtime_params_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
+
+        for (auto& pair : resident_offload_pairs) {
+            ggml_tensor* t    = pair.first;
+            ggml_tensor* twin = pair.second;
+            ggml_backend_tensor_copy(t, twin);
+            std::swap(t->buffer, twin->buffer);
+            std::swap(t->data, twin->data);
+            std::swap(t->extra, twin->extra);
+            resident_param_set.insert(t);
+        }
+        ggml_backend_synchronize(runtime_backend);
+
+        size_t sz = ggml_backend_buffer_get_size(resident_runtime_params_buffer);
+        LOG_INFO("%s offload resident params (%6.2f MB, %zu tensors) to runtime backend (%s)",
+                 get_desc().c_str(),
+                 sz / (1024.f * 1024.f),
+                 resident_offload_pairs.size(),
+                 ggml_backend_name(runtime_backend));
+        return true;
+    }
+
+    void restore_resident_params() {
+        if (resident_offload_pairs.empty()) {
+            if (resident_runtime_params_buffer != nullptr) {
+                ggml_backend_buffer_free(resident_runtime_params_buffer);
+                resident_runtime_params_buffer = nullptr;
+            }
+            if (resident_offload_ctx != nullptr) {
+                ggml_free(resident_offload_ctx);
+                resident_offload_ctx = nullptr;
+            }
+            resident_param_set.clear();
+            resident_state_token = 0;
+            return;
+        }
+        for (auto& pair : resident_offload_pairs) {
+            ggml_tensor* t    = pair.first;
+            ggml_tensor* twin = pair.second;
+            t->buffer         = twin->buffer;
+            t->data           = twin->data;
+            t->extra          = twin->extra;
+            twin->buffer      = nullptr;
+            twin->data        = nullptr;
+            twin->extra       = nullptr;
+        }
+        if (resident_runtime_params_buffer != nullptr) {
+            ggml_backend_buffer_free(resident_runtime_params_buffer);
+            resident_runtime_params_buffer = nullptr;
+        }
+        resident_offload_pairs.clear();
+        if (resident_offload_ctx != nullptr) {
+            ggml_free(resident_offload_ctx);
+            resident_offload_ctx = nullptr;
+        }
+        resident_param_set.clear();
+        resident_state_token = 0;
+    }
+
     bool should_use_graph_cut_segmented_compute(const GraphCutPlan& plan) {
         return plan.has_cuts &&
                plan.valid &&
                max_graph_vram_bytes > 0 &&
                plan.segments.size() > 1 &&
                params_backend != runtime_backend &&
-               !ggml_backend_is_cpu(runtime_backend);
+               !sd_backend_is_cpu(runtime_backend);
     }
 
     bool can_attempt_graph_cut_segmented_compute() const {
         return max_graph_vram_bytes > 0 &&
                params_backend != runtime_backend &&
-               !ggml_backend_is_cpu(runtime_backend);
+               !sd_backend_is_cpu(runtime_backend);
     }
 
     bool resolve_graph_cut_plan(ggml_cgraph* gf,
-                                GraphCutPlan* plan_out) {
+                                GraphCutPlan* plan_out,
+                                size_t* effective_budget_out = nullptr) {
         GGML_ASSERT(plan_out != nullptr);
         GGML_ASSERT(gf != nullptr);
+
+        // Keep the plan and resident params under the same live-VRAM cap.
+        // Add back our own resident buffer so we don't see chunk-K's
+        // allocation as "taken" VRAM and shrink the budget on every step.
+        size_t effective_budget = max_graph_vram_bytes;
+        if (stream_layers_enabled && max_graph_vram_bytes > 0 && runtime_backend != nullptr) {
+            ggml_backend_dev_t dev = ggml_backend_get_device(runtime_backend);
+            if (dev != nullptr && ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_CPU) {
+                size_t free_vram = 0, total_vram = 0;
+                ggml_backend_dev_memory(dev, &free_vram, &total_vram);
+                if (resident_runtime_params_buffer != nullptr) {
+                    free_vram += ggml_backend_buffer_get_size(resident_runtime_params_buffer);
+                }
+                constexpr size_t safety_margin = 512ull * 1024 * 1024;
+                size_t free_clamp              = (free_vram > safety_margin) ? (free_vram - safety_margin) : 0;
+                if (free_clamp < effective_budget) {
+                    LOG_DEBUG("%s clamping streaming budget: actual free VRAM %.2f MB < user cap %.2f MB",
+                              get_desc().c_str(),
+                              free_clamp / (1024.0 * 1024.0),
+                              effective_budget / (1024.0 * 1024.0));
+                    effective_budget = free_clamp;
+                }
+            }
+        }
+
+        bool budget_increased = false;
+        if (stream_layers_enabled) {
+            if (effective_budget > observed_max_effective_budget_) {
+                observed_max_effective_budget_ = effective_budget;
+                budget_increased               = true;
+            } else {
+                effective_budget = observed_max_effective_budget_;
+            }
+        }
+
+        if (effective_budget_out != nullptr) {
+            *effective_budget_out = effective_budget;
+        }
+
         *plan_out = sd::ggml_graph_cut::resolve_plan(runtime_backend,
                                                      gf,
                                                      &graph_cut_plan_cache_,
-                                                     max_graph_vram_bytes,
+                                                     effective_budget,
                                                      params_tensor_set_,
                                                      get_desc().c_str());
+        if (stream_layers_enabled) {
+            if (budget_increased) {
+                LOG_INFO("%s streaming budget = %.2f MB",
+                         get_desc().c_str(),
+                         effective_budget / (1024.0 * 1024.0));
+            } else {
+                LOG_DEBUG("%s streaming budget = %.2f MB",
+                          get_desc().c_str(),
+                          effective_budget / (1024.0 * 1024.0));
+            }
+        }
         return true;
     }
 
+    struct PersistentExternalBinding {
+        ggml_backend_buffer_t buffer = nullptr;
+        void* data                   = nullptr;
+        void* extra                  = nullptr;
+    };
+
+    void snapshot_persistent_externals(const sd::ggml_graph_cut::Plan& plan,
+                                       ggml_cgraph* gf,
+                                       std::unordered_map<ggml_tensor*, PersistentExternalBinding>& out) {
+        GGML_ASSERT(gf != nullptr);
+        out.clear();
+        for (const auto& segment : plan.segments) {
+            for (const auto& input : segment.input_refs) {
+                if (input.type != GraphCutSegment::INPUT_EXTERNAL) {
+                    continue;
+                }
+                ggml_tensor* tensor = sd::ggml_graph_cut::input_tensor(gf, input);
+                if (tensor == nullptr || tensor->buffer == nullptr) {
+                    continue;
+                }
+                PersistentExternalBinding binding;
+                binding.buffer = tensor->buffer;
+                binding.data   = tensor->data;
+                binding.extra  = tensor->extra;
+                out[tensor]    = binding;
+            }
+        }
+    }
+
     void reset_segment_runtime_tensors(const GraphCutSegment& segment,
-                                       ggml_cgraph* gf) {
+                                       ggml_cgraph* gf,
+                                       const std::unordered_map<ggml_tensor*, PersistentExternalBinding>* persistent_externals = nullptr) {
         GGML_ASSERT(gf != nullptr);
 
         for (const auto& input : segment.input_refs) {
@@ -2354,11 +2531,25 @@ protected:
             }
             switch (input.type) {
                 case GraphCutSegment::INPUT_PREVIOUS_CUT:
-                case GraphCutSegment::INPUT_EXTERNAL:
                     input_tensor->buffer = nullptr;
                     input_tensor->data   = nullptr;
                     input_tensor->extra  = nullptr;
                     break;
+                case GraphCutSegment::INPUT_EXTERNAL: {
+                    if (persistent_externals != nullptr) {
+                        auto it = persistent_externals->find(input_tensor);
+                        if (it != persistent_externals->end()) {
+                            input_tensor->buffer = it->second.buffer;
+                            input_tensor->data   = it->second.data;
+                            input_tensor->extra  = it->second.extra;
+                            break;
+                        }
+                    }
+                    input_tensor->buffer = nullptr;
+                    input_tensor->data   = nullptr;
+                    input_tensor->extra  = nullptr;
+                    break;
+                }
                 case GraphCutSegment::INPUT_PARAM:
                     break;
             }
@@ -2464,8 +2655,8 @@ protected:
         int64_t t_copy_begin = ggml_time_ms();
         copy_data_to_backend_tensor(gf, !preserve_backend_tensor_data_map);
         int64_t t_copy_end = ggml_time_ms();
-        if (ggml_backend_is_cpu(runtime_backend)) {
-            ggml_backend_cpu_set_n_threads(runtime_backend, n_threads);
+        if (sd_backend_is_cpu(runtime_backend)) {
+            sd_backend_cpu_set_n_threads(runtime_backend, n_threads);
         }
 
         int64_t t_compute_begin = ggml_time_ms();
@@ -2573,6 +2764,9 @@ protected:
         free_compute_buffer();
         free_cache_ctx_and_buffer();
 
+        std::unordered_map<ggml_tensor*, PersistentExternalBinding> persistent_externals;
+        snapshot_persistent_externals(plan, gf, persistent_externals);
+
         std::optional<sd::Tensor<T>> output = sd::Tensor<T>();
         for (size_t seg_idx = 0; seg_idx < plan.segments.size(); ++seg_idx) {
             int64_t t_segment_begin = ggml_time_ms();
@@ -2584,7 +2778,7 @@ protected:
                       plan.segments.size(),
                       segment.group_name.c_str());
 
-            reset_segment_runtime_tensors(segment, gf);
+            reset_segment_runtime_tensors(segment, gf, &persistent_externals);
             if (!bind_segment_cached_inputs(gf, segment)) {
                 free_cache_ctx_and_buffer();
                 free_compute_buffer();
@@ -2629,6 +2823,150 @@ protected:
         return output;
     }
 
+public:
+    void release_streaming_residency() {
+        restore_resident_params();
+    }
+
+    template <typename T>
+    std::optional<sd::Tensor<T>> compute_streaming_segments(ggml_cgraph* gf,
+                                                            const GraphCutPlan& plan,
+                                                            size_t residency_budget_bytes,
+                                                            int n_threads,
+                                                            bool free_compute_buffer_immediately,
+                                                            bool no_return = false) {
+        GGML_ASSERT(gf != nullptr);
+
+        // Runtime LoRA composes `weight + diff` in the compute graph via
+        // ggml_add; the resident weight tensor's data is never mutated, so
+        // chunk-K residency stays valid across sampling steps.
+        // Reserve room for the worst merged segment so chunk-K can't grow
+        // large enough to starve later partial-param allocations.
+        size_t worst_merged_segment_footprint = 0;
+        for (const auto& seg : plan.segments) {
+            const size_t fp = seg.input_param_bytes +
+                              seg.compute_buffer_size +
+                              seg.output_bytes +
+                              seg.input_previous_cut_bytes +
+                              seg.input_external_bytes;
+            if (fp > worst_merged_segment_footprint) {
+                worst_merged_segment_footprint = fp;
+            }
+        }
+        const size_t residency_budget_for_annotate =
+            residency_budget_bytes > worst_merged_segment_footprint
+                ? residency_budget_bytes - worst_merged_segment_footprint
+                : 0;
+
+        sd::ggml_graph_cut::Plan& base_plan = graph_cut_plan_cache_.graph_cut_plan;
+        if (base_plan.available) {
+            sd::ggml_graph_cut::annotate_residency(base_plan, residency_budget_for_annotate);
+
+            std::vector<ggml_tensor*> resident_params;
+            uint64_t token = 0;
+            for (const auto& segment : base_plan.segments) {
+                if (segment.residency != sd::ggml_graph_cut::SegmentResidency::RESIDENT) {
+                    continue;
+                }
+                auto seg_params = sd::ggml_graph_cut::param_tensors(gf, segment);
+                for (ggml_tensor* t : seg_params) {
+                    if (t == nullptr)
+                        continue;
+                    resident_params.push_back(t);
+                    token ^= reinterpret_cast<uintptr_t>(t) * 0x9E3779B97F4A7C15ull;
+                }
+            }
+            if (token != resident_state_token) {
+                restore_resident_params();
+                if (!resident_params.empty()) {
+                    if (offload_resident_params(resident_params)) {
+                        resident_state_token = token;
+                    } else {
+                        LOG_ERROR("%s chunk-K: resident offload failed; continuing with per-segment streaming",
+                                  get_desc().c_str());
+                        restore_resident_params();
+                    }
+                }
+            }
+        }
+
+        free_compute_buffer();
+        free_cache_ctx_and_buffer();
+
+        layer_registry_.move_layer_to_gpu("_global");
+
+        std::unordered_map<ggml_tensor*, PersistentExternalBinding> persistent_externals;
+        snapshot_persistent_externals(plan, gf, persistent_externals);
+
+        std::optional<sd::Tensor<T>> output = sd::Tensor<T>();
+        for (size_t seg_idx = 0; seg_idx < plan.segments.size(); ++seg_idx) {
+            int64_t t_segment_begin = ggml_time_ms();
+            const auto& segment     = plan.segments[seg_idx];
+            const bool is_last      = seg_idx + 1 == plan.segments.size();
+            auto future_cut_names   = sd::ggml_graph_cut::collect_future_input_names(gf, plan, seg_idx);
+
+            LOG_DEBUG("%s streaming-cut executing segment %zu/%zu: %s (residency=%s)",
+                      get_desc().c_str(),
+                      seg_idx + 1,
+                      plan.segments.size(),
+                      segment.group_name.c_str(),
+                      segment.residency == sd::ggml_graph_cut::SegmentResidency::RESIDENT ? "RESIDENT" : "STREAMED");
+
+            if (!layer_registry_.move_layer_to_gpu(segment.group_name)) {
+                LOG_DEBUG("%s streaming: no registry entry for group '%s' (using upstream offload path)",
+                          get_desc().c_str(),
+                          segment.group_name.c_str());
+            }
+
+            reset_segment_runtime_tensors(segment, gf, &persistent_externals);
+            if (!bind_segment_cached_inputs(gf, segment)) {
+                free_cache_ctx_and_buffer();
+                free_compute_buffer();
+                free_compute_ctx();
+                return std::nullopt;
+            }
+
+            if (!is_last) {
+                for (size_t output_idx = 0; output_idx < segment.output_node_indices.size(); ++output_idx) {
+                    ggml_tensor* out_tensor = sd::ggml_graph_cut::output_tensor(gf, segment, output_idx);
+                    if (out_tensor != nullptr &&
+                        sd::ggml_graph_cut::is_graph_cut_tensor(out_tensor) &&
+                        future_cut_names.find(out_tensor->name) != future_cut_names.end()) {
+                        cache(out_tensor->name, out_tensor);
+                    }
+                }
+            }
+
+            ggml_context* segment_graph_ctx = nullptr;
+            ggml_cgraph* segment_graph      = sd::ggml_graph_cut::build_segment_graph(gf, segment, &segment_graph_ctx);
+            auto segment_output             = execute_graph<T>(segment_graph,
+                                                   n_threads,
+                                                   /*free_compute_buffer_immediately=*/true,
+                                                   sd::ggml_graph_cut::runtime_param_tensors(gf, segment, get_desc().c_str()),
+                                                   /*preserve_backend_tensor_data_map=*/true,
+                                                   /*no_return=*/!is_last || no_return,
+                                                   &future_cut_names);
+            ggml_free(segment_graph_ctx);
+            if (!segment_output.has_value()) {
+                free_cache_ctx_and_buffer();
+                free_compute_buffer();
+                free_compute_ctx();
+                return std::nullopt;
+            }
+            output = std::move(segment_output);
+
+            if (segment.residency == sd::ggml_graph_cut::SegmentResidency::STREAMED) {
+                layer_registry_.move_layer_to_cpu(segment.group_name);
+            }
+            (void)t_segment_begin;
+        }
+
+        backend_tensor_data_map.clear();
+        free_cache_ctx_and_buffer();
+        free_compute_ctx();
+        return output;
+    }
+
 public:
     virtual std::string get_desc() = 0;
 
@@ -2638,9 +2976,11 @@ public:
         GGML_ASSERT(runtime_backend != nullptr);
         GGML_ASSERT(params_backend != nullptr);
         alloc_params_ctx();
+        layer_registry_.set_backends(runtime_backend, params_backend);
     }
 
     virtual ~GGMLRunner() {
+        restore_resident_params();
         free_params_buffer();
         free_compute_buffer();
         free_params_ctx();
@@ -2694,7 +3034,18 @@ public:
             LOG_DEBUG("%s skipping params allocation (no tensors)", get_desc().c_str());
             return true;
         }
-        params_buffer = ggml_backend_alloc_ctx_tensors(params_ctx, params_backend);
+        // Pinned host buffer when CPU-offloaded for DMA-direct H2D.
+        ggml_backend_buffer_type_t params_buft = nullptr;
+        if (params_backend != runtime_backend) {
+            ggml_backend_dev_t runtime_dev = ggml_backend_get_device(runtime_backend);
+            if (runtime_dev != nullptr) {
+                params_buft = ggml_backend_dev_host_buffer_type(runtime_dev);
+            }
+        }
+        if (params_buft == nullptr) {
+            params_buft = ggml_backend_get_default_buffer_type(params_backend);
+        }
+        params_buffer = ggml_backend_alloc_ctx_tensors_from_buft(params_ctx, params_buft);
         if (params_buffer == nullptr) {
             LOG_ERROR("%s alloc params backend buffer failed, num_tensors = %i",
                       get_desc().c_str(),
@@ -2707,16 +3058,19 @@ public:
         LOG_DEBUG("%s params backend buffer size = % 6.2f MB(%s) (%i tensors)",
                   get_desc().c_str(),
                   params_buffer_size / (1024.f * 1024.f),
-                  ggml_backend_is_cpu(params_backend) ? "RAM" : "VRAM",
+                  sd_backend_is_cpu(params_backend) ? "RAM" : "VRAM",
                   num_tensors);
         return true;
     }
 
     void free_params_buffer() {
+        // Restore swapped resident params before freeing their backing buffer.
+        restore_resident_params();
         if (params_buffer != nullptr) {
             ggml_backend_buffer_free(params_buffer);
             params_buffer = nullptr;
         }
+        observed_max_effective_budget_ = 0;
     }
 
     size_t get_params_buffer_size() {
@@ -2774,7 +3128,7 @@ public:
             return nullptr;
         }
         // it's performing a compute, check if backend isn't cpu
-        if (!ggml_backend_is_cpu(runtime_backend) && (tensor->buffer == nullptr || ggml_backend_buffer_is_host(tensor->buffer))) {
+        if (!sd_backend_is_cpu(runtime_backend) && (tensor->buffer == nullptr || ggml_backend_buffer_is_host(tensor->buffer))) {
             // pass input tensors to gpu memory
             auto backend_tensor = ggml_dup_tensor(compute_ctx, tensor);
 
@@ -2812,11 +3166,20 @@ public:
 
         if (can_attempt_graph_cut_segmented_compute()) {
             GraphCutPlan plan;
-            if (!resolve_graph_cut_plan(gf, &plan)) {
+            size_t effective_graph_vram_bytes = 0;
+            if (!resolve_graph_cut_plan(gf, &plan, &effective_graph_vram_bytes)) {
                 free_compute_ctx();
                 return std::nullopt;
             }
             if (should_use_graph_cut_segmented_compute(plan)) {
+                if (stream_layers_enabled) {
+                    return compute_streaming_segments<T>(gf,
+                                                         plan,
+                                                         effective_graph_vram_bytes,
+                                                         n_threads,
+                                                         free_compute_buffer_immediately,
+                                                         no_return);
+                }
                 return compute_with_graph_cuts<T>(gf,
                                                   plan,
                                                   n_threads,
@@ -2857,6 +3220,12 @@ public:
         max_graph_vram_bytes = max_vram_bytes;
     }
 
+    void set_stream_layers_enabled(bool enabled) {
+        stream_layers_enabled = enabled;
+    }
+
+    sd::layer_registry::LayerRegistry& get_layer_registry() { return layer_registry_; }
+
     ggml_backend_t get_runtime_backend() {
         return runtime_backend;
     }
@@ -2978,11 +3347,14 @@ protected:
     bool bias;
     bool force_f32;
     bool force_prec_f32;
+    bool allow_weight_scale;
+    bool has_weight_scale = false;
     float scale;
     std::string prefix;
 
     void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
         this->prefix         = prefix;
+        has_weight_scale     = false;
         enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32);
         if (in_features % ggml_blck_size(wtype) != 0 || force_f32) {
             wtype = GGML_TYPE_F32;
@@ -2992,20 +3364,26 @@ protected:
             enum ggml_type wtype = GGML_TYPE_F32;
             params["bias"]       = ggml_new_tensor_1d(ctx, wtype, out_features);
         }
+        if (allow_weight_scale && tensor_storage_map.find(prefix + "weight_scale") != tensor_storage_map.end()) {
+            params["weight_scale"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_features);
+            has_weight_scale       = true;
+        }
     }
 
 public:
     Linear(int64_t in_features,
            int64_t out_features,
-           bool bias           = true,
-           bool force_f32      = false,
-           bool force_prec_f32 = false,
-           float scale         = 1.f)
+           bool bias               = true,
+           bool force_f32          = false,
+           bool force_prec_f32     = false,
+           float scale             = 1.f,
+           bool allow_weight_scale = false)
         : in_features(in_features),
           out_features(out_features),
           bias(bias),
           force_f32(force_f32),
           force_prec_f32(force_prec_f32),
+          allow_weight_scale(allow_weight_scale),
           scale(scale) {}
 
     void set_scale(float scale_) {
@@ -3022,14 +3400,24 @@ public:
         if (bias) {
             b = params["bias"];
         }
+        ggml_tensor* linear_bias = has_weight_scale ? nullptr : b;
+        ggml_tensor* out         = nullptr;
         if (ctx->weight_adapter) {
             WeightAdapter::ForwardParams forward_params;
             forward_params.op_type               = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR;
             forward_params.linear.force_prec_f32 = force_prec_f32;
             forward_params.linear.scale          = scale;
-            return ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x, w, b, prefix, forward_params);
+            out                                  = ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x, w, linear_bias, prefix, forward_params);
+        } else {
+            out = ggml_ext_linear(ctx->ggml_ctx, x, w, linear_bias, force_prec_f32, scale);
         }
-        return ggml_ext_linear(ctx->ggml_ctx, x, w, b, force_prec_f32, scale);
+        if (has_weight_scale) {
+            out = ggml_mul(ctx->ggml_ctx, out, params["weight_scale"]);
+            if (b != nullptr) {
+                out = ggml_add_inplace(ctx->ggml_ctx, out, b);
+            }
+        }
+        return out;
     }
 };
 

src/ggml_extend_backend.cpp

@@ -8,6 +8,7 @@
 #include <stdexcept>
 #include <vector>
 
+#include "stable-diffusion.h"
 #include "util.h"
 
 static std::string trim_copy(const std::string& value) {
@@ -300,6 +301,61 @@ static ggml_backend_t init_named_backend(const std::string& name) {
     return ggml_backend_init_by_name(resolved.c_str(), nullptr);
 }
 
+bool sd_backend_is_cpu(ggml_backend_t backend) {
+    if (backend == nullptr) {
+        return false;
+    }
+    auto dev = ggml_backend_get_device(backend);
+    return dev != nullptr && ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU;
+}
+
+ggml_backend_t sd_backend_cpu_init() {
+    ggml_backend_load_all_once();
+    return ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr);
+}
+
+bool sd_backend_cpu_set_n_threads(ggml_backend_t backend, int n_threads) {
+    if (backend == nullptr) {
+        return false;
+    }
+    auto dev = ggml_backend_get_device(backend);
+    if (dev != nullptr && ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU) {
+        auto reg                           = ggml_backend_dev_backend_reg(dev);
+        auto ggml_backend_set_n_threads_fn = (ggml_backend_set_n_threads_t)ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads");
+        if (ggml_backend_set_n_threads_fn != nullptr) {
+            ggml_backend_set_n_threads_fn(backend, n_threads);
+            return true;
+        }
+    }
+    return false;
+}
+
+const char* sd_get_system_info() {
+    static std::string cache_info = []() -> std::string {
+        ggml_backend_load_all_once();
+        std::stringstream ss;
+        ss << "System Info: \n";
+        auto dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
+        if (dev != nullptr) {
+            auto reg                          = ggml_backend_dev_backend_reg(dev);
+            auto ggml_backend_get_features_fn = (ggml_backend_get_features_t)ggml_backend_reg_get_proc_address(reg, "ggml_backend_get_features");
+            if (ggml_backend_get_features_fn != nullptr) {
+                ggml_backend_feature* feat = ggml_backend_get_features_fn(reg);
+                while (feat->name && feat->value) {
+                    ss << "   " << feat->name << " = " << feat->value << " | ";
+                    feat++;
+                }
+            } else {
+                LOG_WARN("unable to get CPU features");
+            }
+        } else {
+            LOG_WARN("unable to get CPU features");
+        }
+        return ss.str();
+    }();
+    return cache_info.c_str();
+}
+
 static ggml_backend_t sd_get_default_backend() {
     ggml_backend_load_all_once();
     static std::once_flag once;
@@ -349,10 +405,10 @@ static ggml_backend_t sd_get_default_backend() {
 
     if (!backend) {
         LOG_WARN("loading CPU backend");
-        backend = ggml_backend_cpu_init();
+        backend = sd_backend_cpu_init();
     }
 
-    if (ggml_backend_is_cpu(backend)) {
+    if (sd_backend_is_cpu(backend)) {
         LOG_DEBUG("Using CPU backend");
     }
 
@@ -452,11 +508,11 @@ ggml_backend_t SDBackendManager::params_backend(SDBackendModule module) {
 }
 
 bool SDBackendManager::runtime_backend_is_cpu(SDBackendModule module) {
-    return ggml_backend_is_cpu(runtime_backend(module));
+    return sd_backend_is_cpu(runtime_backend(module));
 }
 
 bool SDBackendManager::params_backend_is_cpu(SDBackendModule module) {
-    return ggml_backend_is_cpu(params_backend(module));
+    return sd_backend_is_cpu(params_backend(module));
 }
 
 bool SDBackendManager::runtime_backend_supports_host_buffer(SDBackendModule module) {
@@ -464,7 +520,7 @@ bool SDBackendManager::runtime_backend_supports_host_buffer(SDBackendModule modu
     if (backend == nullptr) {
         return false;
     }
-    if (ggml_backend_is_cpu(backend)) {
+    if (sd_backend_is_cpu(backend)) {
         return true;
     }
     ggml_backend_dev_t dev = ggml_backend_get_device(backend);

src/ggml_extend_backend.h

@@ -8,7 +8,6 @@
 #include <unordered_map>
 
 #include "ggml-backend.h"
-#include "ggml-cpu.h"
 #include "ggml.h"
 
 enum class SDBackendModule {
@@ -72,6 +71,9 @@ private:
 };
 
 bool sd_backend_is(ggml_backend_t backend, const std::string& name);
+bool sd_backend_is_cpu(ggml_backend_t backend);
+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

src/ggml_graph_cut.cpp

@@ -699,9 +699,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;
@@ -753,4 +753,54 @@ namespace sd::ggml_graph_cut {
         return resolved_plan;
     }
 
+    void annotate_residency(Plan& plan, size_t max_graph_vram_bytes) {
+        // Cached plans may be reused with a smaller live budget.
+        for (auto& seg : plan.segments) {
+            seg.residency = SegmentResidency::STREAMED;
+        }
+        if (max_graph_vram_bytes == 0 || plan.segments.size() < 2) {
+            return;
+        }
+
+        bool any_param_bearing = false;
+        for (const auto& seg : plan.segments) {
+            if (seg.input_param_bytes > 0) {
+                any_param_bearing = true;
+                break;
+            }
+        }
+        if (!any_param_bearing) {
+            return;
+        }
+
+        // Leave room for the largest active streamed segment.
+        size_t worst_streamed_footprint = 0;
+        for (const auto& seg : plan.segments) {
+            const size_t seg_footprint = seg.input_param_bytes +
+                                         seg.compute_buffer_size +
+                                         seg.output_bytes +
+                                         seg.input_previous_cut_bytes +
+                                         seg.input_external_bytes;
+            if (seg_footprint > worst_streamed_footprint) {
+                worst_streamed_footprint = seg_footprint;
+            }
+        }
+        constexpr size_t safety = 512ull * 1024 * 1024;
+        const size_t reserved   = safety + worst_streamed_footprint;
+
+        if (max_graph_vram_bytes <= reserved) {
+            return;
+        }
+        const size_t available = max_graph_vram_bytes - reserved;
+
+        size_t cumulative = 0;
+        for (auto& seg : plan.segments) {
+            if (cumulative + seg.input_param_bytes > available) {
+                break;
+            }
+            seg.residency = SegmentResidency::RESIDENT;
+            cumulative += seg.input_param_bytes;
+        }
+    }
+
 }  // namespace sd::ggml_graph_cut

src/ggml_graph_cut.h

@@ -2,6 +2,7 @@
 #define __SD_GGML_GRAPH_CUT_H__
 
 #include <array>
+#include <cstdint>
 #include <string>
 #include <unordered_set>
 #include <vector>
@@ -11,6 +12,12 @@
 
 namespace sd::ggml_graph_cut {
 
+    // Streaming residency for a segment's params.
+    enum class SegmentResidency : uint8_t {
+        STREAMED = 0,
+        RESIDENT = 1,
+    };
+
     struct Segment {
         enum InputType {
             INPUT_EXTERNAL = 0,
@@ -34,6 +41,7 @@ namespace sd::ggml_graph_cut {
         std::vector<int> internal_node_indices;
         std::vector<int> output_node_indices;
         std::vector<InputRef> input_refs;
+        SegmentResidency residency = SegmentResidency::STREAMED;
     };
 
     struct Plan {
@@ -101,6 +109,9 @@ namespace sd::ggml_graph_cut {
                       size_t max_graph_vram_bytes,
                       const std::unordered_set<const ggml_tensor*>& params_tensor_set,
                       const char* log_desc);
+
+    // Mark leading segments resident when they fit after streamed-segment headroom.
+    void annotate_residency(Plan& plan, size_t max_graph_vram_bytes);
 }  // namespace sd::ggml_graph_cut
 
 #endif

src/guidance.cpp

@@ -1,13 +1,68 @@
 #include "guidance.h"
 
+#include <algorithm>
+#include <cmath>
+#include <cstdlib>
+#include <string>
 #include <utility>
 
+#include "util.h"
+
 namespace sd::guidance {
 
     static bool has_tensor(const sd::Tensor<float>* tensor) {
         return tensor != nullptr && !tensor->empty();
     }
 
+    bool is_adaptive_projected_guidance_enabled(const AdaptiveProjectedGuidanceParams& params) {
+        return params.eta != 1.0f || params.momentum != 0.0f || params.norm_threshold > 0.0f;
+    }
+
+    AdaptiveProjectedGuidanceParams parse_adaptive_projected_guidance_args(const char* extra_sample_args) {
+        AdaptiveProjectedGuidanceParams params;
+        for (const auto& [key, value] : parse_key_value_args(extra_sample_args, "extra sample arg")) {
+            float parsed = 0.0f;
+            if (key == "apg_eta") {
+                if (!parse_strict_float(value, parsed)) {
+                    LOG_WARN("ignoring invalid APG extra sample arg '%s=%s'", key.c_str(), value.c_str());
+                    continue;
+                }
+                params.eta = parsed;
+            } else if (key == "apg_momentum") {
+                if (!parse_strict_float(value, parsed)) {
+                    LOG_WARN("ignoring invalid APG extra sample arg '%s=%s'", key.c_str(), value.c_str());
+                    continue;
+                }
+                params.momentum = parsed;
+            } else if (key == "apg_norm_threshold") {
+                if (!parse_strict_float(value, parsed)) {
+                    LOG_WARN("ignoring invalid APG extra sample arg '%s=%s'", key.c_str(), value.c_str());
+                    continue;
+                }
+                params.norm_threshold = parsed;
+            } else if (key == "apg_norm_threshold_smoothing") {
+                if (!parse_strict_float(value, parsed)) {
+                    LOG_WARN("ignoring invalid APG extra sample arg '%s=%s'", key.c_str(), value.c_str());
+                    continue;
+                }
+                params.norm_threshold_smoothing = parsed;
+            }
+        }
+        return params;
+    }
+
+    bool parse_skip_layer_guidance_uncond_arg(const char* extra_sample_args) {
+        bool uncond = false;
+        for (const auto& [key, value] : parse_key_value_args(extra_sample_args, "extra sample arg")) {
+            if (key == "slg_uncond") {
+                if (!parse_strict_bool(value, uncond)) {
+                    LOG_WARN("ignoring invalid SLG extra sample arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            }
+        }
+        return uncond;
+    }
+
     ClassifierFreeGuidance::ClassifierFreeGuidance(float guidance_scale,
                                                    float image_guidance_scale)
         : guidance_scale_(guidance_scale),
@@ -27,17 +82,132 @@ 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;
+    }
+
+    AdaptiveProjectedGuidance::AdaptiveProjectedGuidance(float guidance_scale,
+                                                         float image_guidance_scale,
+                                                         AdaptiveProjectedGuidanceParams params)
+        : guidance_scale_(guidance_scale),
+          image_guidance_scale_(image_guidance_scale),
+          params_(params) {
+    }
+
+    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_uncond,
+                                                      float guidance_scale,
+                                                      float image_guidance_scale) {
+        if (pred_img_uncond != nullptr) {
+            if (pred_uncond != nullptr && guidance_scale == 1.0f) {
+                return *pred_uncond - *pred_img_uncond;
+            }
+            if (pred_uncond != nullptr) {
+                return pred_cond +
+                       (*pred_uncond * (image_guidance_scale - guidance_scale) +
+                        *pred_img_uncond * (1.0f - image_guidance_scale)) /
+                           (guidance_scale - 1.0f);
+            }
+            return pred_cond - *pred_img_uncond;
+        }
+        return pred_cond - *pred_uncond;
+    }
+
+    GuiderOutput AdaptiveProjectedGuidance::forward(const GuidanceInput& input,
+                                                    GuiderOutput previous) const {
+        (void)previous;
+
+        GuiderOutput output;
+        if (!has_tensor(input.pred_cond)) {
+            return output;
+        }
+
+        const sd::Tensor<float>& pred_cond = *input.pred_cond;
+        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_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_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_uncond)) {
+            return output;
+        }
+
+        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_uncond,
+                                                            guidance_scale_,
+                                                            image_guidance_scale_);
+        if (params_.momentum != 0.0f) {
+            if (momentum_buffer_.shape() != deltas.shape()) {
+                momentum_buffer_ = sd::Tensor<float>::zeros_like(deltas);
+            }
+            deltas += params_.momentum * momentum_buffer_;
+            momentum_buffer_ = deltas;
+        }
+
+        float diff_norm = 0.0f;
+        if (params_.norm_threshold > 0.0f) {
+            diff_norm = std::sqrt((deltas * deltas).sum());
+        }
+
+        float apg_scale_factor = 1.0f;
+        if (params_.norm_threshold > 0.0f) {
+            if (diff_norm > 0.0f) {
+                if (params_.norm_threshold_smoothing <= 0.0f) {
+                    apg_scale_factor = std::min(1.0f, params_.norm_threshold / diff_norm);
+                } else {
+                    float x          = params_.norm_threshold / diff_norm;
+                    apg_scale_factor = x / std::pow(1.0f + std::pow(x, 1.0f / params_.norm_threshold_smoothing),
+                                                    params_.norm_threshold_smoothing);
+                }
+            }
+        }
+
+        deltas *= apg_scale_factor;
+        if (params_.eta != 1.0f) {
+            float cond_norm_sq = (pred_cond * pred_cond).sum();
+            if (cond_norm_sq != 0.0f) {
+                float projection_scale = (pred_cond * deltas).sum() / cond_norm_sq;
+                deltas += (params_.eta - 1.0f) * (projection_scale * pred_cond);
+            }
+        }
+
+        output.pred = pred_cond;
+        if (pred_uncond != nullptr) {
+            if (guidance_scale_ != 1.0f) {
+                output.pred = pred_cond + (guidance_scale_ - 1.0f) * deltas;
+            } else if (pred_img_uncond != nullptr) {
+                output.pred = pred_cond + (image_guidance_scale_ - 1.0f) * deltas;
+            }
+        } else if (pred_img_uncond != nullptr) {
+            output.pred = *pred_img_uncond + guidance_scale_ * deltas;
         }
 
         return output;
@@ -54,7 +224,7 @@ namespace sd::guidance {
     }
 
     bool SkipLayerGuidance::is_enabled_for_step(const GuidanceInput& input) const {
-        if (scale_ == 0.0f || layers_.empty() || input.schedule_size == 0) {
+        if (layers_.empty() || input.schedule_size == 0) {
             return false;
         }
 
@@ -69,7 +239,7 @@ namespace sd::guidance {
 
     GuiderOutput SkipLayerGuidance::forward(const GuidanceInput& input,
                                             GuiderOutput output) const {
-        if (!is_enabled_for_step(input) || !input.predict_skip_layer) {
+        if (scale_ == 0.0f || !is_enabled_for_step(input) || !input.predict_skip_layer) {
             return output;
         }
 

src/guidance.h

@@ -17,12 +17,23 @@ namespace sd::guidance {
         sd::Tensor<float> pred_skip_layer;
     };
 
+    struct AdaptiveProjectedGuidanceParams {
+        float eta                      = 1.0f;
+        float momentum                 = 0.0f;
+        float norm_threshold           = 0.0f;
+        float norm_threshold_smoothing = 0.0f;
+    };
+
+    AdaptiveProjectedGuidanceParams parse_adaptive_projected_guidance_args(const char* extra_sample_args);
+    bool is_adaptive_projected_guidance_enabled(const AdaptiveProjectedGuidanceParams& params);
+    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;
     };
@@ -46,6 +57,21 @@ namespace sd::guidance {
                              GuiderOutput previous) const override;
     };
 
+    class AdaptiveProjectedGuidance : public BaseGuidance {
+        float guidance_scale_       = 1.0f;
+        float image_guidance_scale_ = 1.0f;
+        AdaptiveProjectedGuidanceParams params_;
+        mutable sd::Tensor<float> momentum_buffer_;
+
+    public:
+        AdaptiveProjectedGuidance(float guidance_scale,
+                                  float image_guidance_scale,
+                                  AdaptiveProjectedGuidanceParams params);
+
+        GuiderOutput forward(const GuidanceInput& input,
+                             GuiderOutput previous) const override;
+    };
+
     class SkipLayerGuidance : public BaseGuidance {
         std::vector<int> layers_;
         float scale_ = 0.0f;

src/hidream_o1.hpp

@@ -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;
@@ -284,8 +286,8 @@ namespace HiDreamO1 {
                               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)) {
+              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,
@@ -331,7 +333,7 @@ namespace HiDreamO1 {
     };
 
     struct HiDreamO1Runner : public DiffusionModelRunner {
-        HiDreamO1Params params;
+        HiDreamO1Config config;
         HiDreamO1Model model;
 
         std::vector<float> attention_mask_vec;
@@ -341,8 +343,8 @@ namespace HiDreamO1 {
                         const String2TensorStorage& tensor_storage_map = {},
                         const std::string& prefix                      = "model")
             : DiffusionModelRunner(backend, params_backend, prefix),
-              params(make_hidream_o1_params()) {
-            model = HiDreamO1Model(params);
+              config(HiDreamO1Config::detect_from_weights(tensor_storage_map, prefix)) {
+            model = HiDreamO1Model(config);
             model.init(params_ctx, tensor_storage_map, prefix);
         }
 
@@ -492,8 +494,11 @@ namespace HiDreamO1 {
             vision_runner->get_param_tensors(tensors);
         }
 
-        void alloc_params_buffer() override {
-            vision_runner->alloc_params_buffer();
+        bool alloc_params_buffer() override {
+            if (!vision_runner->alloc_params_buffer()) {
+                return false;
+            }
+            return true;
         }
 
         void free_params_buffer() override {

src/ideogram4.hpp

@@ -0,0 +1,531 @@
+#ifndef __IDEOGRAM4_HPP__
+#define __IDEOGRAM4_HPP__
+
+#include <algorithm>
+#include <cmath>
+#include <cstdlib>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "diffusion_model.hpp"
+#include "ggml_extend.hpp"
+#include "ggml_graph_cut.h"
+#include "rope.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) {
+        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)};
+            }
+        }
+
+        return Rope::embed_interleaved_mrope(ids, bs, static_cast<float>(rope_theta), head_dim, mrope_section);
+    }
+
+    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,
+                        ggml_backend_t params_backend,
+                        const String2TensorStorage& tensor_storage_map = {},
+                        const std::string prefix                       = "")
+            : DiffusionModelRunner(backend, params_backend, prefix),
+              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;
+
+            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);
+            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());
+
+            auto runner_ctx  = get_context();
+            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), 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  // __IDEOGRAM4_HPP__

src/layer_registry.cpp

@@ -0,0 +1,132 @@
+#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

@@ -0,0 +1,50 @@
+#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/lens.hpp

@@ -13,6 +13,71 @@
 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,13 +345,13 @@ 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;
 
@@ -308,53 +359,9 @@ namespace Lens {
                    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);
+            : DiffusionModelRunner(backend, params_backend, prefix),
+              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();

src/llm.hpp

@@ -37,6 +37,7 @@ namespace LLM {
         MISTRAL_SMALL_3_2,
         MINISTRAL_3_3B,
         GEMMA3_12B,
+        GEMMA2_2B,
         GPT_OSS_20B,
         ARCH_COUNT,
     };
@@ -48,6 +49,7 @@ namespace LLM {
         "mistral_small3.2",
         "ministral3.3b",
         "gemma3_12b",
+        "gemma2_2b",
         "gpt_oss_20b",
     };
 
@@ -61,7 +63,7 @@ namespace LLM {
         QWEN3_VL,
     };
 
-    struct LLMVisionParams {
+    struct LLMVisionConfig {
         LLMVisionArch arch                  = LLMVisionArch::QWEN2_5_VL;
         int num_layers                      = 32;
         int64_t hidden_size                 = 1280;
@@ -77,7 +79,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;
@@ -99,7 +101,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 {
@@ -230,11 +354,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]
@@ -665,7 +789,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),
@@ -782,23 +906,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);
             }
         }
 
@@ -879,8 +1003,8 @@ namespace LLM {
                                                  input_pos,
                                                  nullptr,
                                                  head_dim,
-                                                 GGML_ROPE_TYPE_NORMAL,
-                                                 0,
+                                                 GGML_ROPE_TYPE_NEOX,
+                                                 131072,
                                                  rope_theta,
                                                  freq_scale,
                                                  0.f,
@@ -892,14 +1016,41 @@ namespace LLM {
                                                  input_pos,
                                                  nullptr,
                                                  head_dim,
-                                                 GGML_ROPE_TYPE_NORMAL,
-                                                 0,
+                                                 GGML_ROPE_TYPE_NEOX,
+                                                 131072,
                                                  rope_theta,
                                                  freq_scale,
                                                  0.f,
                                                  1.f,
                                                  32.f,
                                                  1.f);
+            } else if (arch == LLMArch::GEMMA2_2B) {
+                q = ggml_rope_ext(ctx->ggml_ctx,
+                                  q,
+                                  input_pos,
+                                  nullptr,
+                                  head_dim,
+                                  GGML_ROPE_TYPE_NEOX,
+                                  8192,
+                                  10000.f,
+                                  1.f,
+                                  0.f,
+                                  1.f,
+                                  32.f,
+                                  1.f);
+                k = ggml_rope_ext(ctx->ggml_ctx,
+                                  k,
+                                  input_pos,
+                                  nullptr,
+                                  head_dim,
+                                  GGML_ROPE_TYPE_NEOX,
+                                  8192,
+                                  10000.f,
+                                  1.f,
+                                  0.f,
+                                  1.f,
+                                  32.f,
+                                  1.f);
             } else if (arch == LLMArch::QWEN3_VL) {
                 int sections[4] = {24, 20, 20, 0};
                 q               = ggml_rope_multi(ctx->ggml_ctx, q, input_pos, nullptr, head_dim, sections, GGML_ROPE_TYPE_IMROPE, 262144, 5000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
@@ -953,34 +1104,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) {
-                post_attention_norm_name = "post_attention_norm";
-                post_ffw_norm_name       = "post_ffw_norm";
-            }
-            pre_ffw_norm_name = params.arch == LLMArch::GPT_OSS_20B ? "post_attention_norm" : "post_attention_layernorm";
-
-            blocks["self_attn"] = std::make_shared<Attention>(params);
-            if (params.arch == LLMArch::GPT_OSS_20B) {
-                blocks["mlp"] = std::make_shared<GPTOSSMLP>(params);
+            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 (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 (config.arch == LLMArch::GPT_OSS_20B) {
+                pre_ffw_norm_name = "post_attention_norm";  // attn_post_norm
             } else {
-                blocks["mlp"] = std::make_shared<MLP>(params.hidden_size,
-                                                      params.intermediate_size,
-                                                      false,
-                                                      params.mlp_activation);
+                pre_ffw_norm_name = "post_attention_layernorm";  // ffn_norm
             }
-            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["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>(config.hidden_size,
+                                                      config.intermediate_size,
+                                                      false,
+                                                      config.mlp_activation);
+            }
+            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()];
             }
         }
 
@@ -1037,16 +1196,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,
@@ -1066,8 +1225,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);
@@ -1137,15 +1296,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));
             }
         }
 
@@ -1189,7 +1348,7 @@ namespace LLM {
     };
 
     struct LLMRunner : public GGMLRunner {
-        LLMParams params;
+        LLMConfig config;
         bool enable_vision;
         LLM model;
 
@@ -1205,7 +1364,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];
@@ -1300,7 +1459,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,
@@ -1415,121 +1574,25 @@ namespace LLM {
                   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::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) {
+            : GGMLRunner(backend, params_backend),
+              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) {
@@ -1581,11 +1644,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::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;
@@ -1624,9 +1688,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);
@@ -1682,15 +1746,15 @@ namespace LLM {
 
         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,
@@ -1712,7 +1776,7 @@ namespace LLM {
                                        compute_ctx,
                                        runner_ctx,
                                        image,
-                                       params.vision,
+                                       config.vision,
                                        model.vision_model(),
                                        window_index_vec,
                                        window_inverse_index_vec,
@@ -1726,8 +1790,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);
@@ -1769,8 +1833,11 @@ namespace LLM {
             model.get_param_tensors(tensors, prefix);
         }
 
-        void alloc_params_buffer() {
-            model.alloc_params_buffer();
+        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,
@@ -1985,7 +2052,7 @@ namespace LLM {
         static void load_from_file_and_test(const std::string& file_path) {
             // cpu f16: pass
             // ggml_backend_t backend = ggml_backend_cuda_init(0);
-            ggml_backend_t backend    = ggml_backend_cpu_init();
+            ggml_backend_t backend    = sd_backend_cpu_init();
             ggml_type model_data_type = GGML_TYPE_COUNT;
 
             ModelLoader model_loader;
@@ -2012,7 +2079,11 @@ namespace LLM {
                                                                              "text_encoders.llm",
                                                                              true);
 
-            llm->alloc_params_buffer();
+            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");
 

src/lora.hpp

@@ -86,7 +86,10 @@ struct LoraModel : public GGMLRunner {
             lora_tensors[name] = real;
         }
 
-        alloc_params_buffer();
+        if (!alloc_params_buffer()) {
+            LOG_ERROR("lora model buffer allocation failed");
+            return false;
+        }
 
         dry_run = false;
         model_loader.load_tensors(on_new_tensor_cb, n_threads);
@@ -769,7 +772,7 @@ struct LoraModel : public GGMLRunner {
             }
 
             ggml_tensor* original_tensor = model_tensor;
-            if (!ggml_backend_is_cpu(runtime_backend) && ggml_backend_buffer_is_host(original_tensor->buffer)) {
+            if (!sd_backend_is_cpu(runtime_backend) && ggml_backend_buffer_is_host(original_tensor->buffer)) {
                 model_tensor = ggml_dup_tensor(compute_ctx, model_tensor);
                 set_backend_tensor_data(model_tensor, original_tensor->data);
             }
@@ -783,7 +786,7 @@ struct LoraModel : public GGMLRunner {
                 final_tensor = ggml_add_inplace(compute_ctx, model_tensor, diff);
             }
             ggml_build_forward_expand(gf, final_tensor);
-            if (!ggml_backend_is_cpu(runtime_backend) && ggml_backend_buffer_is_host(original_tensor->buffer)) {
+            if (!sd_backend_is_cpu(runtime_backend) && ggml_backend_buffer_is_host(original_tensor->buffer)) {
                 original_tensor_to_final_tensor[original_tensor] = final_tensor;
             }
         }

src/ltx_audio_vae.h

@@ -58,11 +58,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 +169,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;
         }
     };
@@ -1052,7 +1059,7 @@ namespace LTXV {
         static void load_from_file_and_test(const std::string& model_path,
                                             const std::string& input_path,
                                             const std::string& prefix = "") {
-            ggml_backend_t backend = ggml_backend_cpu_init();
+            ggml_backend_t backend = sd_backend_cpu_init();
             // ggml_backend_t backend = ggml_backend_cuda_init(0);
             LOG_INFO("loading ltx audio vae from '%s'", model_path.c_str());
 
@@ -1068,7 +1075,11 @@ namespace LTXV {
                                                                      tensor_storage_map,
                                                                      prefix);
 
-            ltx_audio_vae->alloc_params_buffer();
+            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, "");
 

src/ltx_vae.hpp

@@ -11,7 +11,7 @@
 
 #include "ltxv.hpp"
 #include "vae.hpp"
-#include "wan.hpp"
+#include "wan_vae.hpp"
 
 namespace LTXVAE {
 
@@ -1517,7 +1517,7 @@ struct LTXVideoVAE : public VAE {
     static void load_from_file_and_test(const std::string& model_path,
                                         const std::string& input_path) {
         // ggml_backend_t backend = ggml_backend_cuda_init(0);
-        ggml_backend_t backend = ggml_backend_cpu_init();
+        ggml_backend_t backend = sd_backend_cpu_init();
         LOG_INFO("loading ltx vae from '%s'", model_path.c_str());
 
         ModelLoader model_loader;
@@ -1534,7 +1534,11 @@ struct LTXVideoVAE : public VAE {
                                                                          true,
                                                                          VERSION_LTXAV);
 
-        vae->alloc_params_buffer();
+        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");
 

src/ltxv.hpp

@@ -72,6 +72,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 +943,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 +1306,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 +1430,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 +1445,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 +1463,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 +1496,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 +1565,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 +1589,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 +1611,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 +1628,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 +1654,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 +1673,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 +1684,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);
+              config(LTXAVConfig::detect_from_weights(tensor_storage_map, prefix)),
+              model(config) {
             model.init(params_ctx, tensor_storage_map, prefix);
         }
 
@@ -1692,21 +1718,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 +1793,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 +1820,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 +1890,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());
             }
@@ -1995,7 +2021,7 @@ namespace LTXV {
                                             const std::string& audio_x_path         = "",
                                             const std::string& audio_timesteps_path = "") {
             // ggml_backend_t backend = ggml_backend_cuda_init(0);
-            ggml_backend_t backend = ggml_backend_cpu_init();
+            ggml_backend_t backend = sd_backend_cpu_init();
             LOG_INFO("loading ltxav from '%s'", model_path.c_str());
 
             ModelLoader model_loader;
@@ -2017,7 +2043,10 @@ namespace LTXV {
                                                                                tensor_storage_map,
                                                                                "model.diffusion_model");
 
-            ltxav->alloc_params_buffer();
+            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");
 

src/mmdit.hpp

@@ -1,7 +1,10 @@
 #ifndef __MMDIT_HPP__
 #define __MMDIT_HPP__
 
+#include <algorithm>
 #include <memory>
+#include <string>
+#include <vector>
 
 #include "diffusion_model.hpp"
 #include "ggml_extend.hpp"
@@ -9,6 +12,128 @@
 
 #define MMDIT_GRAPH_SIZE 10240
 
+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;
+
+    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;
+
+        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;
+    }
+};
+
 struct Mlp : public GGMLBlock {
 public:
     Mlp(int64_t in_features,
@@ -612,28 +737,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 +759,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 +791,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 +814,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 +836,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 +879,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 +899,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) {
+        : DiffusionModelRunner(backend, params_backend, prefix),
+          config(MMDiTConfig::detect_from_weights(tensor_storage_map, prefix)),
+          mmdit(config) {
         mmdit.init(params_ctx, tensor_storage_map, prefix);
     }
 
@@ -947,13 +1029,17 @@ struct MMDiTRunner : public DiffusionModelRunner {
 
     static void load_from_file_and_test(const std::string& file_path) {
         // ggml_backend_t backend    = ggml_backend_cuda_init(0);
-        ggml_backend_t backend             = ggml_backend_cpu_init();
+        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);
         {
             LOG_INFO("loading from '%s'", file_path.c_str());
 
-            mmdit->alloc_params_buffer();
+            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");
 

src/model.cpp

@@ -432,6 +432,12 @@ SDVersion ModelLoader::get_sd_version() {
             tensor_storage.name.find("model.diffusion_model.single_transformer_blocks.") != std::string::npos) {
             is_flux = true;
         }
+        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;
         }
@@ -1004,6 +1010,12 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
                         continue;
                     }
 
+                    if (dst_tensor->data == nullptr) {
+                        LOG_ERROR("process tensor data failed: '%s'", tensor_storage.name.c_str());
+                        failed = true;
+                        break;
+                    }
+
                     // skip mmapped tensors
                     if (dst_tensor->buffer != nullptr && dst_tensor->buffer == fdata.mmbuffer.get()) {
                         continue;
@@ -1245,6 +1257,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.h

@@ -49,6 +49,8 @@ enum SDVersion {
     VERSION_ERNIE_IMAGE,
     VERSION_LENS,
     VERSION_LONGCAT,
+    VERSION_PID,
+    VERSION_IDEOGRAM4,
     VERSION_COUNT,
 };
 
@@ -164,8 +166,22 @@ static inline bool sd_version_is_lens(SDVersion version) {
     return false;
 }
 
+static inline bool sd_version_is_pid(SDVersion version) {
+    if (version == VERSION_PID) {
+        return true;
+    }
+    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;
@@ -194,7 +210,9 @@ static inline bool sd_version_is_dit(SDVersion version) {
         sd_version_is_z_image(version) ||
         sd_version_is_ernie_image(version) ||
         sd_version_is_lens(version) ||
-        sd_version_is_longcat(version)) {
+        sd_version_is_longcat(version) ||
+        sd_version_is_pid(version) ||
+        sd_version_is_ideogram4(version)) {
         return true;
     }
     return false;

src/pid.hpp

@@ -0,0 +1,847 @@
+#ifndef __SD_PID_HPP__
+#define __SD_PID_HPP__
+
+#include <cmath>
+#include <cstdlib>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "common_dit.hpp"
+#include "ggml_extend.hpp"
+#include "mmdit.hpp"
+#include "rope.hpp"
+
+namespace Pid {
+    constexpr int PID_GRAPH_SIZE = 196608;
+    constexpr float PID_PI       = 3.14159265358979323846f;
+
+    struct PixelDiTConfig {
+        int64_t in_channels            = 3;
+        int64_t hidden_size            = 1536;
+        int64_t num_groups             = 24;
+        int64_t patch_mlp_hidden_dim   = 4096;
+        int64_t pixel_hidden_size      = 16;
+        int64_t pixel_attn_hidden_size = 1152;
+        int64_t pixel_num_groups       = 16;
+        int64_t patch_depth            = 14;
+        int64_t pixel_depth            = 2;
+        int64_t patch_size             = 16;
+        int64_t txt_embed_dim          = 2304;
+        int64_t txt_max_length         = 300;
+        float text_rope_theta          = 10000.f;
+        int64_t lq_latent_channels     = 16;
+        int64_t lq_hidden_dim          = 512;
+        int64_t lq_num_res_blocks      = 4;
+        int64_t lq_interval            = 2;
+        int64_t lq_sr_scale            = 4;
+        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,
+                                           int dim,
+                                           float theta) {
+        GGML_ASSERT(dim % 2 == 0);
+        return Rope::flatten(Rope::rope(Rope::linspace(0.f, static_cast<float>(length - 1), length), dim, theta));
+    }
+
+    inline std::vector<float> make_rope_2d(int height,
+                                           int width,
+                                           int dim,
+                                           float theta    = 10000.f,
+                                           float scale    = 16.f,
+                                           int ref_grid_h = 0,
+                                           int ref_grid_w = 0) {
+        GGML_ASSERT(dim % 4 == 0);
+        return Rope::embed_2d_interleaved(height, width, dim, theta, scale, ref_grid_h, ref_grid_w);
+    }
+
+    inline std::vector<float> make_pixel_abs_pos(int height,
+                                                 int width,
+                                                 int dim) {
+        GGML_ASSERT(dim % 4 == 0);
+        int half_dim = dim / 2;
+        std::vector<float> x_pos;
+        std::vector<float> y_pos;
+        x_pos.reserve(static_cast<size_t>(height) * width);
+        y_pos.reserve(static_cast<size_t>(height) * width);
+        for (int iy = 0; iy < height; ++iy) {
+            for (int ix = 0; ix < width; ++ix) {
+                x_pos.push_back(static_cast<float>(ix));
+                y_pos.push_back(static_cast<float>(iy));
+            }
+        }
+
+        auto x_emb = timestep_embedding(x_pos, half_dim, 10000, false);
+        auto y_emb = timestep_embedding(y_pos, half_dim, 10000, false);
+
+        std::vector<float> out(static_cast<size_t>(dim) * height * width);
+        for (int pos = 0; pos < height * width; ++pos) {
+            size_t out_base = static_cast<size_t>(pos) * dim;
+            size_t emb_base = static_cast<size_t>(pos) * half_dim;
+            for (int i = 0; i < half_dim; ++i) {
+                out[out_base + i]            = x_emb[emb_base + i];
+                out[out_base + half_dim + i] = y_emb[emb_base + i];
+            }
+        }
+        return out;
+    }
+
+    inline ggml_tensor* apply_adaln(ggml_context* ctx,
+                                    ggml_tensor* x,
+                                    ggml_tensor* shift,
+                                    ggml_tensor* scale) {
+        return ggml_add(ctx, ggml_add(ctx, x, ggml_mul(ctx, x, scale)), shift);
+    }
+
+    struct PatchTokenEmbedder : public GGMLBlock {
+        bool use_rms_norm;
+
+        PatchTokenEmbedder(int64_t in_chans,
+                           int64_t embed_dim,
+                           bool use_rms_norm = false,
+                           bool bias         = true)
+            : use_rms_norm(use_rms_norm) {
+            blocks["proj"] = std::make_shared<Linear>(in_chans, embed_dim, bias);
+            if (use_rms_norm) {
+                blocks["norm"] = std::make_shared<RMSNorm>(embed_dim, 1e-6f);
+            }
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
+            x         = proj->forward(ctx, x);
+            if (use_rms_norm) {
+                auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
+                x         = norm->forward(ctx, x);
+            }
+            return x;
+        }
+    };
+
+    struct PixelDiTTimestepEmbedder : public GGMLBlock {
+        int frequency_embedding_size;
+
+        PixelDiTTimestepEmbedder(int64_t hidden_size,
+                                 int frequency_embedding_size = 256)
+            : frequency_embedding_size(frequency_embedding_size) {
+            blocks["mlp.0"] = std::make_shared<Linear>(frequency_embedding_size, hidden_size, true, true);
+            blocks["mlp.2"] = std::make_shared<Linear>(hidden_size, hidden_size, true, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* t) {
+            auto mlp_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
+            auto mlp_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
+            auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, t, frequency_embedding_size, 10);
+            t_emb      = mlp_0->forward(ctx, t_emb);
+            t_emb      = ggml_silu_inplace(ctx->ggml_ctx, t_emb);
+            return mlp_2->forward(ctx, t_emb);
+        }
+    };
+
+    struct FeedForward : public GGMLBlock {
+        FeedForward(int64_t dim, int64_t hidden_dim) {
+            blocks["w1"] = std::make_shared<Linear>(dim, hidden_dim, false);
+            blocks["w2"] = std::make_shared<Linear>(hidden_dim, dim, false);
+            blocks["w3"] = std::make_shared<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 h  = ggml_silu_inplace(ctx->ggml_ctx, w1->forward(ctx, x));
+            h       = ggml_mul_inplace(ctx->ggml_ctx, h, w3->forward(ctx, x));
+            return w2->forward(ctx, h);
+        }
+    };
+
+    struct FinalLayer : public GGMLBlock {
+        FinalLayer(int64_t hidden_size, int64_t out_channels) {
+            blocks["norm"]   = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
+            blocks["linear"] = std::make_shared<Linear>(hidden_size, out_channels, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            auto norm   = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
+            auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
+            return linear->forward(ctx, norm->forward(ctx, x));
+        }
+    };
+
+    struct RotaryAttention : public GGMLBlock {
+        int64_t dim;
+        int64_t num_heads;
+
+        RotaryAttention(int64_t dim, int64_t num_heads)
+            : dim(dim), num_heads(num_heads) {
+            int64_t head_dim = dim / num_heads;
+            blocks["qkv"]    = std::make_shared<Linear>(dim, dim * 3, false);
+            blocks["q_norm"] = std::make_shared<RMSNorm>(head_dim, 1e-6f);
+            blocks["k_norm"] = std::make_shared<RMSNorm>(head_dim, 1e-6f);
+            blocks["proj"]   = std::make_shared<Linear>(dim, dim, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pos) {
+            auto qkv_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
+            auto q_norm   = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm"]);
+            auto k_norm   = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm"]);
+            auto proj     = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
+
+            auto qkv         = qkv_proj->forward(ctx, x);
+            auto qkv_vec     = split_qkv(ctx->ggml_ctx, qkv);
+            int64_t L        = x->ne[1];
+            int64_t N        = x->ne[2];
+            int64_t head_dim = dim / num_heads;
+            auto q           = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[0], head_dim, num_heads, L, N);
+            auto k           = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[1], head_dim, num_heads, L, N);
+            auto v           = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[2], head_dim, num_heads, L, N);
+            q                = q_norm->forward(ctx, q);
+            k                = k_norm->forward(ctx, k);
+            x                = Rope::attention(ctx, q, k, v, pos, nullptr, 1.0f / 128.f, true);
+            return proj->forward(ctx, x);
+        }
+    };
+
+    struct MMDiTJointAttention : public GGMLBlock {
+        int64_t dim;
+        int64_t num_heads;
+
+        MMDiTJointAttention(int64_t dim, int64_t num_heads)
+            : dim(dim), num_heads(num_heads) {
+            int64_t head_dim   = dim / num_heads;
+            blocks["qkv_x"]    = std::make_shared<Linear>(dim, dim * 3, false);
+            blocks["qkv_y"]    = std::make_shared<Linear>(dim, dim * 3, false);
+            blocks["q_norm_x"] = std::make_shared<RMSNorm>(head_dim, 1e-6f);
+            blocks["k_norm_x"] = std::make_shared<RMSNorm>(head_dim, 1e-6f);
+            blocks["q_norm_y"] = std::make_shared<RMSNorm>(head_dim, 1e-6f);
+            blocks["k_norm_y"] = std::make_shared<RMSNorm>(head_dim, 1e-6f);
+            blocks["proj_x"]   = std::make_shared<Linear>(dim, dim, true);
+            blocks["proj_y"]   = std::make_shared<Linear>(dim, dim, true);
+        }
+
+        std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
+                                                      ggml_tensor* x,
+                                                      ggml_tensor* y,
+                                                      ggml_tensor* pos_img,
+                                                      ggml_tensor* pos_txt) {
+            auto qkv_x_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv_x"]);
+            auto qkv_y_proj = std::dynamic_pointer_cast<Linear>(blocks["qkv_y"]);
+            auto q_norm_x   = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm_x"]);
+            auto k_norm_x   = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm_x"]);
+            auto q_norm_y   = std::dynamic_pointer_cast<RMSNorm>(blocks["q_norm_y"]);
+            auto k_norm_y   = std::dynamic_pointer_cast<RMSNorm>(blocks["k_norm_y"]);
+            auto proj_x     = std::dynamic_pointer_cast<Linear>(blocks["proj_x"]);
+            auto proj_y     = std::dynamic_pointer_cast<Linear>(blocks["proj_y"]);
+
+            int64_t Nx       = x->ne[1];
+            int64_t Ny       = y->ne[1];
+            int64_t N        = x->ne[2];
+            int64_t head_dim = dim / num_heads;
+
+            auto qkv_x = split_qkv(ctx->ggml_ctx, qkv_x_proj->forward(ctx, x));
+            auto qx    = ggml_reshape_4d(ctx->ggml_ctx, qkv_x[0], head_dim, num_heads, Nx, N);
+            auto kx    = ggml_reshape_4d(ctx->ggml_ctx, qkv_x[1], head_dim, num_heads, Nx, N);
+            auto vx    = ggml_reshape_4d(ctx->ggml_ctx, qkv_x[2], head_dim, num_heads, Nx, N);
+            qx         = q_norm_x->forward(ctx, qx);
+            kx         = k_norm_x->forward(ctx, kx);
+
+            auto qkv_y = split_qkv(ctx->ggml_ctx, qkv_y_proj->forward(ctx, y));
+            auto qy    = ggml_reshape_4d(ctx->ggml_ctx, qkv_y[0], head_dim, num_heads, Ny, N);
+            auto ky    = ggml_reshape_4d(ctx->ggml_ctx, qkv_y[1], head_dim, num_heads, Ny, N);
+            auto vy    = ggml_reshape_4d(ctx->ggml_ctx, qkv_y[2], head_dim, num_heads, Ny, N);
+            qy         = q_norm_y->forward(ctx, qy);
+            ky         = k_norm_y->forward(ctx, ky);
+
+            auto q_joint   = ggml_concat(ctx->ggml_ctx, qy, qx, 2);
+            auto k_joint   = ggml_concat(ctx->ggml_ctx, ky, kx, 2);
+            auto v_joint   = ggml_concat(ctx->ggml_ctx, vy, vx, 2);
+            auto pos_joint = ggml_concat(ctx->ggml_ctx, pos_txt, pos_img, 3);
+            auto out       = Rope::attention(ctx, q_joint, k_joint, v_joint, pos_joint, nullptr, 1.0f, true);
+
+            auto out_y = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, Ny);
+            auto out_x = ggml_ext_slice(ctx->ggml_ctx, out, 1, Ny, Ny + Nx);
+            return {proj_x->forward(ctx, out_x), proj_y->forward(ctx, out_y)};
+        }
+    };
+
+    struct MMDiTBlockT2I : public GGMLBlock {
+        int64_t hidden_size;
+
+        MMDiTBlockT2I(int64_t hidden_size, int64_t groups, int64_t mlp_hidden_dim)
+            : hidden_size(hidden_size) {
+            blocks["norm_x1"]                = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
+            blocks["norm_y1"]                = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
+            blocks["attn"]                   = std::make_shared<MMDiTJointAttention>(hidden_size, groups);
+            blocks["norm_x2"]                = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
+            blocks["norm_y2"]                = std::make_shared<RMSNorm>(hidden_size, 1e-6f);
+            blocks["mlp_x"]                  = std::make_shared<FeedForward>(hidden_size, mlp_hidden_dim);
+            blocks["mlp_y"]                  = std::make_shared<FeedForward>(hidden_size, mlp_hidden_dim);
+            blocks["adaLN_modulation_img.0"] = std::make_shared<Linear>(hidden_size, 6 * hidden_size, true);
+            blocks["adaLN_modulation_txt.0"] = std::make_shared<Linear>(hidden_size, 6 * hidden_size, true);
+        }
+
+        std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
+                                                      ggml_tensor* x,
+                                                      ggml_tensor* y,
+                                                      ggml_tensor* c,
+                                                      ggml_tensor* pos_img,
+                                                      ggml_tensor* pos_txt) {
+            auto norm_x1 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_x1"]);
+            auto norm_y1 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_y1"]);
+            auto attn    = std::dynamic_pointer_cast<MMDiTJointAttention>(blocks["attn"]);
+            auto norm_x2 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_x2"]);
+            auto norm_y2 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_y2"]);
+            auto mlp_x   = std::dynamic_pointer_cast<FeedForward>(blocks["mlp_x"]);
+            auto mlp_y   = std::dynamic_pointer_cast<FeedForward>(blocks["mlp_y"]);
+            auto ada_img = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation_img.0"]);
+            auto ada_txt = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation_txt.0"]);
+
+            auto mx = ggml_ext_chunk(ctx->ggml_ctx, ada_img->forward(ctx, c), 6, 0);
+            auto my = ggml_ext_chunk(ctx->ggml_ctx, ada_txt->forward(ctx, c), 6, 0);
+
+            auto x_norm   = apply_adaln(ctx->ggml_ctx, norm_x1->forward(ctx, x), mx[0], mx[1]);
+            auto y_norm   = apply_adaln(ctx->ggml_ctx, norm_y1->forward(ctx, y), my[0], my[1]);
+            auto attn_out = attn->forward(ctx, x_norm, y_norm, pos_img, pos_txt);
+
+            x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out.first, mx[2]));
+            y = ggml_add(ctx->ggml_ctx, y, ggml_mul(ctx->ggml_ctx, attn_out.second, my[2]));
+
+            auto x_mlp = mlp_x->forward(ctx, apply_adaln(ctx->ggml_ctx, norm_x2->forward(ctx, x), mx[3], mx[4]));
+            auto y_mlp = mlp_y->forward(ctx, apply_adaln(ctx->ggml_ctx, norm_y2->forward(ctx, y), my[3], my[4]));
+            x          = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, x_mlp, mx[5]));
+            y          = ggml_add(ctx->ggml_ctx, y, ggml_mul(ctx->ggml_ctx, y_mlp, my[5]));
+            return {x, y};
+        }
+    };
+
+    struct PixelTokenEmbedder : public GGMLBlock {
+        int64_t in_channels;
+        int64_t hidden_size_output;
+
+        PixelTokenEmbedder(int64_t in_channels, int64_t hidden_size_output)
+            : in_channels(in_channels), hidden_size_output(hidden_size_output) {
+            blocks["proj"] = std::make_shared<Linear>(in_channels, hidden_size_output, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* inputs,
+                             int64_t patch_size,
+                             ggml_tensor* pos_full) {
+            auto proj  = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
+            int64_t W  = inputs->ne[0];
+            int64_t H  = inputs->ne[1];
+            int64_t B  = inputs->ne[3];
+            int64_t L  = (W / patch_size) * (H / patch_size);
+            int64_t P2 = patch_size * patch_size;
+
+            auto x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, inputs, 2, 0, 1, 3));
+            x      = ggml_reshape_3d(ctx->ggml_ctx, x, in_channels, W * H, B);
+            x      = proj->forward(ctx, x);
+            x      = ggml_add(ctx->ggml_ctx, x, pos_full);
+            x      = ggml_reshape_4d(ctx->ggml_ctx, x, hidden_size_output, W, H, B);
+            x      = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 2, 0, 3));
+            x      = DiT::patchify(ctx->ggml_ctx, x, static_cast<int>(patch_size), static_cast<int>(patch_size), false);
+            x      = ggml_reshape_3d(ctx->ggml_ctx, x, hidden_size_output, P2, L * B);
+            return x;
+        }
+    };
+
+    struct PiTBlock : public GGMLBlock {
+        int64_t pixel_dim;
+        int64_t context_dim;
+        int64_t attn_dim;
+        int64_t num_heads;
+        int64_t patch_size;
+
+        PiTBlock(int64_t pixel_dim,
+                 int64_t context_dim,
+                 int64_t patch_size,
+                 int64_t attn_dim,
+                 int64_t num_heads)
+            : pixel_dim(pixel_dim),
+              context_dim(context_dim),
+              attn_dim(attn_dim),
+              num_heads(num_heads),
+              patch_size(patch_size) {
+            int64_t p2                   = patch_size * patch_size;
+            blocks["compress_to_attn"]   = std::make_shared<Linear>(p2 * pixel_dim, attn_dim, true);
+            blocks["expand_from_attn"]   = std::make_shared<Linear>(attn_dim, p2 * pixel_dim, true);
+            blocks["norm1"]              = std::make_shared<RMSNorm>(pixel_dim, 1e-6f);
+            blocks["attn"]               = std::make_shared<RotaryAttention>(attn_dim, num_heads);
+            blocks["norm2"]              = std::make_shared<RMSNorm>(pixel_dim, 1e-6f);
+            blocks["mlp"]                = std::make_shared<Mlp>(pixel_dim, pixel_dim * 4);
+            blocks["adaLN_modulation.0"] = std::make_shared<Linear>(context_dim, 6 * pixel_dim * p2, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* s_cond,
+                             int64_t image_height,
+                             int64_t image_width,
+                             ggml_tensor* pos_comp) {
+            auto compress = std::dynamic_pointer_cast<Linear>(blocks["compress_to_attn"]);
+            auto expand   = std::dynamic_pointer_cast<Linear>(blocks["expand_from_attn"]);
+            auto norm1    = std::dynamic_pointer_cast<RMSNorm>(blocks["norm1"]);
+            auto attn     = std::dynamic_pointer_cast<RotaryAttention>(blocks["attn"]);
+            auto norm2    = std::dynamic_pointer_cast<RMSNorm>(blocks["norm2"]);
+            auto mlp      = std::dynamic_pointer_cast<Mlp>(blocks["mlp"]);
+            auto ada      = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.0"]);
+
+            int64_t Hs = image_height / patch_size;
+            int64_t Ws = image_width / patch_size;
+            int64_t L  = Hs * Ws;
+            int64_t BL = x->ne[2];
+            int64_t B  = BL / L;
+            int64_t P2 = patch_size * patch_size;
+
+            auto ada_params = ada->forward(ctx, s_cond);
+            ada_params      = ggml_reshape_3d(ctx->ggml_ctx, ada_params, 6 * pixel_dim, P2, BL);
+            auto mod        = ggml_ext_chunk(ctx->ggml_ctx, ada_params, 6, 0);
+
+            auto x_norm    = apply_adaln(ctx->ggml_ctx, norm1->forward(ctx, x), mod[0], mod[1]);
+            auto x_flat    = ggml_reshape_2d(ctx->ggml_ctx, x_norm, P2 * pixel_dim, BL);
+            auto x_comp    = compress->forward(ctx, x_flat);
+            x_comp         = ggml_reshape_3d(ctx->ggml_ctx, x_comp, attn_dim, L, B);
+            auto attn_out  = attn->forward(ctx, x_comp, pos_comp);
+            auto attn_flat = expand->forward(ctx, ggml_reshape_2d(ctx->ggml_ctx, attn_out, attn_dim, BL));
+            auto attn_exp  = ggml_reshape_3d(ctx->ggml_ctx, attn_flat, pixel_dim, P2, BL);
+            x              = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_exp, mod[2]));
+
+            auto mlp_out = mlp->forward(ctx, apply_adaln(ctx->ggml_ctx, norm2->forward(ctx, x), mod[3], mod[4]));
+            return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, mlp_out, mod[5]));
+        }
+    };
+
+    struct SigmaAwareGate : public GGMLBlock {
+        int64_t dim;
+
+        SigmaAwareGate(int64_t dim)
+            : dim(dim) {
+            blocks["content_proj"] = std::make_shared<Linear>(dim * 2, dim, true);
+        }
+
+        void init_params(ggml_context* ctx,
+                         const String2TensorStorage& tensor_storage_map = {},
+                         std::string prefix                             = "") override {
+            params["log_alpha"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* lq,
+                             ggml_tensor* sigma) {
+            auto content_proj = std::dynamic_pointer_cast<Linear>(blocks["content_proj"]);
+
+            auto content_logit = content_proj->forward(ctx, ggml_concat(ctx->ggml_ctx, x, lq, 0));
+            sigma              = ggml_reshape_3d(ctx->ggml_ctx, sigma, 1, 1, sigma->ne[0]);
+            auto alpha         = ggml_exp(ctx->ggml_ctx, params["log_alpha"]);
+            auto offset        = ggml_neg(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, alpha, sigma));
+            auto gate          = ggml_sigmoid(ctx->ggml_ctx, ggml_add(ctx->ggml_ctx, content_logit, offset));
+            return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, gate, lq));
+        }
+    };
+
+    struct PiDResBlock : public GGMLBlock {
+        PiDResBlock(int64_t channels) {
+            blocks["block.0"] = std::make_shared<GroupNorm>(4, channels, 1e-5f);
+            blocks["block.2"] = std::make_shared<Conv2d>(channels, channels, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, std::pair<int, int>{1, 1});
+            blocks["block.3"] = std::make_shared<GroupNorm>(4, channels, 1e-5f);
+            blocks["block.5"] = std::make_shared<Conv2d>(channels, channels, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, std::pair<int, int>{1, 1});
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            auto norm1 = std::dynamic_pointer_cast<GroupNorm>(blocks["block.0"]);
+            auto conv1 = std::dynamic_pointer_cast<Conv2d>(blocks["block.2"]);
+            auto norm2 = std::dynamic_pointer_cast<GroupNorm>(blocks["block.3"]);
+            auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["block.5"]);
+            auto h     = ggml_silu_inplace(ctx->ggml_ctx, norm1->forward(ctx, x));
+            h          = conv1->forward(ctx, h);
+            h          = ggml_silu_inplace(ctx->ggml_ctx, norm2->forward(ctx, h));
+            h          = conv2->forward(ctx, h);
+            return ggml_add(ctx->ggml_ctx, x, h);
+        }
+    };
+
+    struct LQProjection2D : public GGMLBlock {
+        PixelDiTConfig config;
+
+        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>((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>(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 % config.lq_interval == 0;
+        }
+
+        int get_output_index(int block_idx) const {
+            return block_idx / static_cast<int>(config.lq_interval);
+        }
+
+        ggml_tensor* gate(GGMLRunnerContext* ctx,
+                          ggml_tensor* x,
+                          ggml_tensor* lq,
+                          ggml_tensor* sigma,
+                          int out_idx) {
+            auto gate_module = std::dynamic_pointer_cast<SigmaAwareGate>(blocks["gate_modules." + std::to_string(out_idx)]);
+            return gate_module->forward(ctx, x, lq, sigma);
+        }
+
+        std::vector<ggml_tensor*> forward(GGMLRunnerContext* ctx,
+                                          ggml_tensor* lq_latent,
+                                          int64_t target_pH,
+                                          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>(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,
+                                             lq_latent,
+                                             target_pW,
+                                             target_pH,
+                                             lq_latent->ne[2],
+                                             lq_latent->ne[3],
+                                             GGML_SCALE_MODE_NEAREST);
+            }
+
+            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 < 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);
+            }
+
+            int64_t B   = feat->ne[3];
+            int64_t C   = feat->ne[2];
+            int64_t L   = target_pH * target_pW;
+            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>((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) {
+                auto head = std::dynamic_pointer_cast<Linear>(blocks["output_heads." + std::to_string(i)]);
+                outputs.push_back(head->forward(ctx, tokens));
+            }
+            return outputs;
+        }
+    };
+
+    struct PixelDiT : public GGMLBlock {
+        PixelDiTConfig config;
+
+        PixelDiT() = default;
+
+        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 < 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>(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, config.hidden_size, config.txt_max_length, 1);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* timesteps,
+                             ggml_tensor* context,
+                             ggml_tensor* lq_latent,
+                             ggml_tensor* degrade_sigma,
+                             ggml_tensor* pos_img,
+                             ggml_tensor* pos_txt,
+                             ggml_tensor* pixel_pos_full,
+                             ggml_tensor* pixel_pos_comp) {
+            auto pixel_embedder = std::dynamic_pointer_cast<PixelTokenEmbedder>(blocks["pixel_embedder"]);
+            auto s_embedder     = std::dynamic_pointer_cast<PatchTokenEmbedder>(blocks["s_embedder"]);
+            auto t_embedder     = std::dynamic_pointer_cast<PixelDiTTimestepEmbedder>(blocks["t_embedder"]);
+            auto y_embedder     = std::dynamic_pointer_cast<PatchTokenEmbedder>(blocks["y_embedder"]);
+            auto final_layer    = std::dynamic_pointer_cast<FinalLayer>(blocks["final_layer"]);
+            auto lq_proj        = std::dynamic_pointer_cast<LQProjection2D>(blocks["lq_proj"]);
+
+            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>(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 / config.patch_size;
+            int64_t Ws     = W / config.patch_size;
+            int64_t L      = Hs * Ws;
+            int64_t P2     = config.patch_size * config.patch_size;
+
+            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], 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);
+            y_emb        = ggml_add(ctx->ggml_ctx, y_emb, y_pos);
+
+            std::vector<ggml_tensor*> lq_features = lq_proj->forward(ctx, lq_latent, Hs, Ws);
+
+            auto s = s_embedder->forward(ctx, x_patches);
+
+            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())) {
+                        s = lq_proj->gate(ctx, s, lq_features[out_idx], degrade_sigma, out_idx);
+                    }
+                }
+                auto block = std::dynamic_pointer_cast<MMDiTBlockT2I>(blocks["patch_blocks." + std::to_string(i)]);
+                auto out   = block->forward(ctx,
+                                            s,
+                                            y_emb,
+                                            condition,
+                                            pos_img,
+                                            pos_txt);
+                s          = out.first;
+                y_emb      = out.second;
+                sd::ggml_graph_cut::mark_graph_cut(s, "pid.patch_blocks." + std::to_string(i), "s");
+                sd::ggml_graph_cut::mark_graph_cut(y_emb, "pid.patch_blocks." + std::to_string(i), "y");
+            }
+            s = ggml_silu(ctx->ggml_ctx, ggml_add(ctx->ggml_ctx, s, t_emb));
+
+            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, config.in_channels * P2, L, B);
+            auto out = DiT::unpatchify(ctx->ggml_ctx,
+                                       pixels,
+                                       Hs,
+                                       Ws,
+                                       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);
+            return out;
+        }
+    };
+
+    struct PiDRunner : public DiffusionModelRunner {
+        PixelDiTConfig config;
+        PixelDiT model;
+        std::vector<float> pos_img_vec;
+        std::vector<float> pos_txt_vec;
+        std::vector<float> pixel_pos_vec;
+        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),
+              config(PixelDiTConfig::detect_from_weights(tensor_storage_map, prefix)) {
+            model = PixelDiT(config);
+            model.init(params_ctx, tensor_storage_map, prefix);
+        }
+
+        std::string get_desc() override {
+            return "PiD";
+        }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
+            model.get_param_tensors(tensors, prefix);
+        }
+
+        ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
+                                 const sd::Tensor<float>& timesteps_tensor,
+                                 const sd::Tensor<float>& context_tensor,
+                                 const sd::Tensor<float>& lq_latent_tensor,
+                                 const sd::Tensor<float>& degrade_sigma_tensor) {
+            ggml_cgraph* gf            = new_graph_custom(PID_GRAPH_SIZE);
+            ggml_tensor* x             = make_input(x_tensor);
+            ggml_tensor* timesteps     = make_input(timesteps_tensor);
+            ggml_tensor* context       = make_input(context_tensor);
+            ggml_tensor* lq_latent     = make_input(lq_latent_tensor);
+            ggml_tensor* degrade_sigma = make_input(degrade_sigma_tensor);
+
+            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>(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>(config.hidden_size / config.num_groups),
+                                        10000.f,
+                                        16.f,
+                                        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,
+                                              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], config.txt_max_length);
+            pos_txt_vec  = make_rope_1d(static_cast<int>(Ltxt),
+                                        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,
+                                              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>(config.pixel_hidden_size));
+            auto pixel_pos = ggml_new_tensor_3d(compute_ctx,
+                                                GGML_TYPE_F32,
+                                                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>(config.pixel_attn_hidden_size / config.pixel_num_groups),
+                                               10000.f,
+                                               16.f,
+                                               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,
+                                                     config.pixel_attn_hidden_size / config.pixel_num_groups / 2,
+                                                     Hs * Ws);
+            set_backend_tensor_data(pixel_pos_comp, pixel_pos_comp_vec.data());
+
+            auto runner_ctx = get_context();
+            auto out        = model.forward(&runner_ctx,
+                                            x,
+                                            timesteps,
+                                            context,
+                                            lq_latent,
+                                            degrade_sigma,
+                                            pos_img,
+                                            pos_txt,
+                                            pixel_pos,
+                                            pixel_pos_comp);
+            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,
+                                  const sd::Tensor<float>& lq_latent,
+                                  const sd::Tensor<float>& degrade_sigma) {
+            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());
+        }
+
+        sd::Tensor<float> compute(int n_threads,
+                                  const DiffusionParams& diffusion_params) override {
+            GGML_ASSERT(diffusion_params.x != nullptr);
+            GGML_ASSERT(diffusion_params.timesteps != nullptr);
+            GGML_ASSERT(diffusion_params.context != nullptr);
+            GGML_ASSERT(diffusion_params.ref_latents != nullptr);
+            GGML_ASSERT(!diffusion_params.ref_latents->empty());
+            auto degrade_sigma = sd::Tensor<float>::from_vector({0.0f});
+            return compute(n_threads,
+                           *diffusion_params.x,
+                           *diffusion_params.timesteps,
+                           *diffusion_params.context,
+                           diffusion_params.ref_latents->front(),
+                           degrade_sigma);
+        }
+    };
+}  // namespace Pid
+
+#endif  // __SD_PID_HPP__

src/pmid.hpp

@@ -615,7 +615,10 @@ struct PhotoMakerIDEmbed : public GGMLRunner {
         };
 
         model_loader->load_tensors(on_new_tensor_cb, n_threads);
-        alloc_params_buffer();
+        if (!alloc_params_buffer()) {
+            LOG_ERROR("PhotoMaker ID embeds buffer allocation failed");
+            return false;
+        }
 
         dry_run = false;
         model_loader->load_tensors(on_new_tensor_cb, n_threads);

src/qwen_image.hpp

@@ -10,6 +10,48 @@
 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 +392,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 +434,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 +445,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 +455,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 +484,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 +502,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,7 +510,7 @@ 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;
@@ -494,34 +522,10 @@ namespace Qwen {
                         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);
+            : DiffusionModelRunner(backend, params_backend, prefix),
+              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 +556,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;
                 }
@@ -683,7 +687,7 @@ namespace Qwen {
             // cuda q8: pass
             // cuda q8 fa: pass
             // ggml_backend_t backend    = ggml_backend_cuda_init(0);
-            ggml_backend_t backend    = ggml_backend_cpu_init();
+            ggml_backend_t backend    = sd_backend_cpu_init();
             ggml_type model_data_type = GGML_TYPE_Q8_0;
 
             ModelLoader model_loader;
@@ -705,7 +709,11 @@ namespace Qwen {
                                                                                             "model.diffusion_model",
                                                                                             VERSION_QWEN_IMAGE);
 
-            qwen_image->alloc_params_buffer();
+            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");
 

src/rope.hpp

@@ -249,6 +249,93 @@ 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) {
+        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) {
+            axis_embs.push_back(rope(trans_ids[axis], head_dim, theta));
+        }
+
+        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,
+                                                              float theta    = 10000.f,
+                                                              float scale    = 16.f,
+                                                              int ref_grid_h = 0,
+                                                              int ref_grid_w = 0) {
+        assert(dim % 4 == 0);
+        int half_dim      = dim / 2;
+        int dim_axis      = dim / 2;
+        int axis_half_dim = dim_axis / 2;
+
+        float h_ntk = 1.f;
+        float w_ntk = 1.f;
+        if (ref_grid_h > 0 && ref_grid_w > 0 && dim_axis > 2) {
+            float power = static_cast<float>(dim_axis) / static_cast<float>(dim_axis - 2);
+            h_ntk       = std::pow(static_cast<float>(height) / static_cast<float>(ref_grid_h), power);
+            w_ntk       = std::pow(static_cast<float>(width) / static_cast<float>(ref_grid_w), power);
+        }
+
+        std::vector<float> x_pos;
+        std::vector<float> y_pos;
+        x_pos.reserve(static_cast<size_t>(height) * width);
+        y_pos.reserve(static_cast<size_t>(height) * width);
+        for (int iy = 0; iy < height; ++iy) {
+            float y = height == 1 ? 0.f : scale * static_cast<float>(iy) / static_cast<float>(height - 1);
+            for (int ix = 0; ix < width; ++ix) {
+                float x = width == 1 ? 0.f : scale * static_cast<float>(ix) / static_cast<float>(width - 1);
+                x_pos.push_back(x);
+                y_pos.push_back(y);
+            }
+        }
+
+        auto x_emb = rope(x_pos, dim_axis, theta * w_ntk);
+        auto y_emb = rope(y_pos, dim_axis, theta * h_ntk);
+
+        std::vector<float> out(static_cast<size_t>(height) * width * half_dim * 4);
+        for (int pos = 0; pos < height * width; ++pos) {
+            for (int i = 0; i < axis_half_dim; ++i) {
+                int jx        = 2 * i;
+                int jy        = 2 * i + 1;
+                size_t base_x = static_cast<size_t>(pos) * half_dim * 4 + static_cast<size_t>(jx) * 4;
+                size_t base_y = static_cast<size_t>(pos) * half_dim * 4 + static_cast<size_t>(jy) * 4;
+                size_t axis   = static_cast<size_t>(i) * 4;
+                for (int k = 0; k < 4; ++k) {
+                    out[base_x + k] = x_emb[pos][axis + k];
+                    out[base_y + k] = y_emb[pos][axis + k];
+                }
+            }
+        }
+        return out;
+    }
+
     __STATIC_INLINE__ std::vector<std::vector<float>> gen_refs_ids(int patch_size,
                                                                    int bs,
                                                                    int axes_dim_num,

src/stable-diffusion.cpp

@@ -1,3 +1,7 @@
+#include <algorithm>
+#include <cmath>
+#include <cstdlib>
+
 #include "ggml_extend.hpp"
 #include "ggml_graph_cut.h"
 
@@ -19,6 +23,7 @@
 #include "flux.hpp"
 #include "guidance.h"
 #include "hidream_o1.hpp"
+#include "ideogram4.hpp"
 #include "lens.hpp"
 #include "lora.hpp"
 #include "ltx_audio_vae.h"
@@ -26,6 +31,7 @@
 #include "ltx_vae.hpp"
 #include "ltxv.hpp"
 #include "mmdit.hpp"
+#include "pid.hpp"
 #include "pmid.hpp"
 #include "qwen_image.hpp"
 #include "sample-cache.h"
@@ -34,11 +40,15 @@
 #include "upscaler.h"
 #include "vae.hpp"
 #include "wan.hpp"
+#include "wan_vae.hpp"
 #include "z_image.hpp"
 
 #include "latent-preview.h"
 #include "name_conversion.h"
 
+const char* sd_vae_format_name(enum sd_vae_format_t format);
+static SDVersion sd_vae_format_to_version(enum sd_vae_format_t format, SDVersion fallback);
+
 const char* model_version_to_str[] = {
     "SD 1.x",
     "SD 1.x Inpaint",
@@ -75,6 +85,8 @@ const char* model_version_to_str[] = {
     "Ernie Image",
     "Lens",
     "Longcat-Image",
+    "PiD",
+    "Ideogram 4",
 };
 
 const char* sampling_methods_str[] = {
@@ -100,6 +112,19 @@ const char* sampling_methods_str[] = {
 
 /*================================================== Helper Functions ================================================*/
 
+static bool sd_version_supports_ref_latent_img_cfg(SDVersion version) {
+    return version == VERSION_FLUX ||
+           sd_version_is_flux2(version) ||
+           sd_version_is_qwen_image(version) ||
+           sd_version_is_longcat(version) ||
+           sd_version_is_z_image(version);
+}
+
+static bool sd_version_supports_img_cfg(SDVersion version, bool has_ref_images) {
+    return sd_version_is_inpaint_or_unet_edit(version) ||
+           (has_ref_images && sd_version_supports_ref_latent_img_cfg(version));
+}
+
 void calculate_alphas_cumprod(float* alphas_cumprod,
                               float linear_start = 0.00085f,
                               float linear_end   = 0.0120f,
@@ -167,6 +192,7 @@ public:
     sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0, 0, nullptr};
     bool offload_params_to_cpu           = false;
     float max_vram                       = 0.f;
+    bool stream_layers                   = false;
     bool use_pmid                        = false;
     std::string backend_spec;
     std::string params_backend_spec;
@@ -212,7 +238,7 @@ public:
         std::string error;
         if (!backend_manager.init(sd_ctx_params->backend,
                                   sd_ctx_params->params_backend,
-                                  sd_ctx_params->offload_params_to_cpu,
+                                  offload_params_to_cpu,
                                   sd_ctx_params->keep_clip_on_cpu,
                                   sd_ctx_params->keep_vae_on_cpu,
                                   sd_ctx_params->keep_control_net_on_cpu,
@@ -239,8 +265,18 @@ public:
         free_params_immediately = sd_ctx_params->free_params_immediately;
         offload_params_to_cpu   = sd_ctx_params->offload_params_to_cpu;
         max_vram                = sd_ctx_params->max_vram;
+        stream_layers           = sd_ctx_params->stream_layers;
         backend_spec            = SAFE_STR(sd_ctx_params->backend);
         params_backend_spec     = SAFE_STR(sd_ctx_params->params_backend);
+        if (stream_layers && max_vram == 0.f) {
+            LOG_WARN("--stream-layers has no effect without --max-vram set; ignoring");
+            stream_layers = false;
+        }
+        if (stream_layers && !offload_params_to_cpu && params_backend_spec.empty()) {
+            // Streaming needs CPU-resident params.
+            LOG_WARN("--stream-layers has no effect without --offload-to-cpu (or --params-backend); ignoring");
+            stream_layers = false;
+        }
 
         bool use_tae       = false;
         bool use_audio_vae = false;
@@ -282,6 +318,13 @@ public:
             }
         }
 
+        if (strlen(SAFE_STR(sd_ctx_params->uncond_diffusion_model_path)) > 0) {
+            LOG_INFO("loading unconditional diffusion model from '%s'", sd_ctx_params->uncond_diffusion_model_path);
+            if (!model_loader.init_from_file(sd_ctx_params->uncond_diffusion_model_path, "model.diffusion_model.uncond.")) {
+                LOG_WARN("loading unconditional diffusion model from '%s' failed", sd_ctx_params->uncond_diffusion_model_path);
+            }
+        }
+
         bool is_unet = sd_version_is_unet(model_loader.get_sd_version());
 
         if (strlen(SAFE_STR(sd_ctx_params->clip_l_path)) > 0) {
@@ -419,7 +462,10 @@ public:
                     }
                 }
             }
-            if (have_quantized_weight) {
+            // Avoid full-model LoRA merge buffers on constrained setups.
+            const bool streaming_constrained = stream_layers ||
+                                               sd_ctx_params->offload_params_to_cpu;
+            if (have_quantized_weight || streaming_constrained) {
                 apply_lora_immediately = false;
             } else {
                 apply_lora_immediately = true;
@@ -501,6 +547,27 @@ public:
                                                                 params_backend_for(SDBackendModule::DIFFUSION),
                                                                 tensor_storage_map,
                                                                 "model.diffusion_model");
+            } else if (sd_version_is_pid(version)) {
+                vae_decode_only  = false;
+                cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
+                                                                 params_backend_for(SDBackendModule::TE),
+                                                                 tensor_storage_map,
+                                                                 version);
+                diffusion_model  = std::make_shared<Pid::PiDRunner>(backend_for(SDBackendModule::DIFFUSION),
+                                                                   params_backend_for(SDBackendModule::DIFFUSION),
+                                                                   tensor_storage_map,
+                                                                   "model.diffusion_model.net");
+            } else if (sd_version_is_ideogram4(version)) {
+                cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
+                                                                 params_backend_for(SDBackendModule::TE),
+                                                                 tensor_storage_map,
+                                                                 version,
+                                                                 "",
+                                                                 false);
+                diffusion_model  = std::make_shared<Ideogram4::Ideogram4Runner>(backend_for(SDBackendModule::DIFFUSION),
+                                                                               params_backend_for(SDBackendModule::DIFFUSION),
+                                                                               tensor_storage_map,
+                                                                               "model.diffusion_model");
             } else if (sd_version_is_flux(version)) {
                 bool is_chroma = false;
                 for (auto pair : tensor_storage_map) {
@@ -705,6 +772,7 @@ public:
             get_param_tensors(cond_stage_model, module_can_mmap(SDBackendModule::TE));
 
             diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes);
+            diffusion_model->set_stream_layers_enabled(stream_layers);
             get_param_tensors(diffusion_model, module_can_mmap(SDBackendModule::DIFFUSION));
 
             if (sd_version_is_unet_edit(version)) {
@@ -713,6 +781,7 @@ public:
 
             if (high_noise_diffusion_model) {
                 high_noise_diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes);
+                high_noise_diffusion_model->set_stream_layers_enabled(stream_layers);
                 get_param_tensors(high_noise_diffusion_model, module_can_mmap(SDBackendModule::DIFFUSION));
             }
 
@@ -743,6 +812,16 @@ public:
                 }
             };
 
+            sd_vae_format_t vae_format = sd_ctx_params->vae_format;
+            if (vae_format < SD_VAE_FORMAT_AUTO || vae_format >= SD_VAE_FORMAT_COUNT) {
+                LOG_WARN("invalid VAE format override, using auto");
+                vae_format = SD_VAE_FORMAT_AUTO;
+            }
+            SDVersion vae_version = version;
+            if (sd_version_is_pid(version) && vae_format != SD_VAE_FORMAT_AUTO) {
+                vae_version = sd_vae_format_to_version(vae_format, vae_version);
+            }
+
             auto create_vae = [&]() -> std::shared_ptr<VAE> {
                 if (sd_version_is_ltxav(version)) {
                     return std::make_shared<LTXVideoVAE>(backend_for(SDBackendModule::VAE),
@@ -767,7 +846,7 @@ public:
                                                                  "first_stage_model",
                                                                  vae_decode_only,
                                                                  false,
-                                                                 version);
+                                                                 vae_version);
                     if (sd_version_is_sdxl(version) &&
                         (strlen(SAFE_STR(sd_ctx_params->vae_path)) == 0 || sd_ctx_params->force_sdxl_vae_conv_scale || external_vae_is_invalid)) {
                         float vae_conv_2d_scale = 1.f / 32.f;
@@ -966,6 +1045,12 @@ public:
             ignore_tensors.insert("text_encoders.llm.model.layers.0.mlp.experts.gate_up_proj.weight_scale_2");
             ignore_tensors.insert("text_encoders.llm.model.layers.0.mlp.experts.down_proj.weight_scale_2");
         }
+        if (sd_version_is_ideogram4(version)) {
+            ignore_tensors.insert("text_encoders.llm.lm_head.");
+            ignore_tensors.insert("text_encoders.llm.visual.");
+            ignore_tensors.insert("text_encoders.llm.vision_model.");
+            ignore_tensors.insert("text_encoders.llm.tokenizer_json");
+        }
         if (version == VERSION_HIDREAM_O1) {
             ignore_tensors.insert("lm_head.");
             ignore_tensors.insert("model.visual.deepstack_merger_list.");
@@ -984,14 +1069,20 @@ public:
             ggml_free(ctx);
             return false;
         }
-        if (cond_stage_model) {
-            cond_stage_model->alloc_params_buffer();
+        if (cond_stage_model && !cond_stage_model->alloc_params_buffer()) {
+            LOG_ERROR("Conditioner model params buffer allocation failed");
+            ggml_free(ctx);
+            return false;
         }
-        if (diffusion_model) {
-            diffusion_model->alloc_params_buffer();
+        if (diffusion_model && !diffusion_model->alloc_params_buffer()) {
+            LOG_ERROR("Diffusion model params buffer allocation failed");
+            ggml_free(ctx);
+            return false;
         }
-        if (high_noise_diffusion_model) {
-            high_noise_diffusion_model->alloc_params_buffer();
+        if (high_noise_diffusion_model && !high_noise_diffusion_model->alloc_params_buffer()) {
+            LOG_ERROR("High noise diffusion model params buffer allocation failed");
+            ggml_free(ctx);
+            return false;
         }
         if (first_stage_model && !first_stage_model->alloc_params_buffer()) {
             LOG_ERROR("VAE params buffer allocation failed");
@@ -1060,7 +1151,7 @@ public:
                 if (module_backend == nullptr) {
                     return false;
                 }
-                if (ggml_backend_is_cpu(module_backend)) {
+                if (sd_backend_is_cpu(module_backend)) {
                     total_params_ram_size += size;
                 } else {
                     total_params_vram_size += size;
@@ -1075,7 +1166,7 @@ public:
                 if (module_backend == nullptr) {
                     return "N/A";
                 }
-                return ggml_backend_is_cpu(module_backend) ? "RAM" : "VRAM";
+                return sd_backend_is_cpu(module_backend) ? "RAM" : "VRAM";
             };
 
             if (!add_params_memory(clip_params_mem_size, SDBackendModule::TE) ||
@@ -1134,12 +1225,18 @@ public:
                            version == VERSION_HIDREAM_O1 ||
                            sd_version_is_anima(version) ||
                            sd_version_is_ernie_image(version) ||
-                           sd_version_is_z_image(version)) {
+                           sd_version_is_z_image(version) ||
+                           sd_version_is_pid(version) ||
+                           sd_version_is_ideogram4(version)) {
                     pred_type = FLOW_PRED;
                     if (sd_version_is_wan(version)) {
                         default_flow_shift = 5.f;
                     } else if (sd_version_is_ernie_image(version)) {
                         default_flow_shift = 4.f;
+                    } else if (sd_version_is_pid(version)) {
+                        default_flow_shift = 1.5f;
+                    } else if (sd_version_is_ideogram4(version)) {
+                        default_flow_shift = 1.0f;
                     } else {
                         default_flow_shift = 3.f;
                     }
@@ -1613,12 +1710,15 @@ public:
                                          const sd::Tensor<float>& init_latent,
                                          const sd::Tensor<float>& denoise_mask) {
         if (diffusion_model->get_desc() == "Wan2.2-TI2V-5B") {
-            auto new_timesteps = std::vector<float>(static_cast<size_t>(init_latent.shape()[2]), timesteps[0]);
+            int64_t frame_count = init_latent.shape()[2];
+            auto new_timesteps  = std::vector<float>(static_cast<size_t>(frame_count), timesteps[0]);
 
-            if (!denoise_mask.empty()) {
-                float value = denoise_mask.dim() == 5 ? denoise_mask.index(0, 0, 0, 0, 0) : denoise_mask.index(0, 0, 0, 0);
-                if (value == 0.f) {
-                    new_timesteps[0] = 0.f;
+            if (!denoise_mask.empty() && denoise_mask.dim() >= 4 && denoise_mask.shape()[2] == frame_count) {
+                for (int64_t frame = 0; frame < frame_count; ++frame) {
+                    float value = denoise_mask.dim() == 5 ? denoise_mask.index(0, 0, frame, 0, 0) : denoise_mask.index(0, 0, frame, 0);
+                    if (value == 0.f) {
+                        new_timesteps[static_cast<size_t>(frame)] = 0.f;
+                    }
                 }
             }
             return new_timesteps;
@@ -1799,7 +1899,7 @@ public:
         if (version == VERSION_HIDREAM_O1) {
             return std::vector<float>{1.0f - (t / static_cast<float>(TIMESTEPS))};
         }
-        if (sd_version_is_z_image(version)) {
+        if (sd_version_is_z_image(version) || sd_version_is_ideogram4(version)) {
             return std::vector<float>{1000.f - t};
         }
         return std::vector<float>{t};
@@ -1878,7 +1978,7 @@ public:
                              sd::Tensor<float> noise,
                              const SDCondition& cond,
                              const SDCondition& uncond,
-                             const SDCondition& img_cond,
+                             const SDCondition& img_uncond,
                              const SDCondition& id_cond,
                              const sd::Tensor<float>& control_image,
                              float control_strength,
@@ -1903,6 +2003,7 @@ public:
         float cfg_scale     = guidance.txt_cfg;
         float img_cfg_scale = guidance.img_cfg;
         float slg_scale     = guidance.slg.scale;
+        bool slg_uncond     = sd::guidance::parse_skip_layer_guidance_uncond_arg(extra_sample_args);
 
         sd_sample::SampleCacheRuntime cache_runtime = sd_sample::init_sample_cache_runtime(version,
                                                                                            cache_params,
@@ -1919,12 +2020,21 @@ public:
         }
 
         size_t steps       = sigmas.size() - 1;
-        bool has_skiplayer = slg_scale != 0.0f && !skip_layers.empty();
+        bool has_skiplayer = (slg_scale != 0.0f || slg_uncond) && !skip_layers.empty();
         if (has_skiplayer && !sd_version_is_dit(version)) {
             has_skiplayer = false;
             LOG_WARN("SLG is incompatible with this model type");
         }
+        sd::guidance::AdaptiveProjectedGuidanceParams apg_params = sd::guidance::parse_adaptive_projected_guidance_args(extra_sample_args);
+        bool use_apg_guidance                                    = sd::guidance::is_adaptive_projected_guidance_enabled(apg_params);
+        if (use_apg_guidance) {
+            LOG_INFO("using Adaptive Projected Guidance (APG)");
+        }
         sd::guidance::ClassifierFreeGuidance classifier_free_guidance(cfg_scale, img_cfg_scale);
+        sd::guidance::AdaptiveProjectedGuidance adaptive_projected_guidance(cfg_scale, img_cfg_scale, apg_params);
+        const sd::guidance::BaseGuidance& primary_guidance = use_apg_guidance
+                                                                 ? static_cast<const sd::guidance::BaseGuidance&>(adaptive_projected_guidance)
+                                                                 : static_cast<const sd::guidance::BaseGuidance&>(classifier_free_guidance);
         sd::guidance::SkipLayerGuidance skip_layer_guidance(has_skiplayer ? skip_layers : std::vector<int>(),
                                                             has_skiplayer ? slg_scale : 0.0f,
                                                             guidance.slg.layer_start,
@@ -1993,13 +2103,17 @@ public:
 
             sd::Tensor<float> cond_out;
             sd::Tensor<float> uncond_out;
-            sd::Tensor<float> img_cond_out;
+            sd::Tensor<float> img_uncond_out;
             sd_sample::SampleStepCacheDispatcher step_cache(cache_runtime, step, sigma);
             std::vector<sd::Tensor<float>> controls;
             DiffusionParams diffusion_params;
             diffusion_params.x                  = &noised_input;
             diffusion_params.timesteps          = &timesteps_tensor;
             diffusion_params.increase_ref_index = increase_ref_index;
+            sd::guidance::GuidanceInput step_guidance_input;
+            step_guidance_input.step          = step;
+            step_guidance_input.schedule_size = sigmas.size();
+            bool is_skiplayer_step            = skip_layer_guidance.is_enabled_for_step(step_guidance_input);
 
             compute_sample_controls(control_image,
                                     noised_input,
@@ -2007,13 +2121,19 @@ public:
                                     cond,
                                     &controls);
 
+            static const std::vector<sd::Tensor<float>> empty_ref_latents;
+            bool uncond_without_ref_latents = !img_uncond.empty() &&
+                                              !ref_latents.empty() &&
+                                              sd_version_supports_ref_latent_img_cfg(version);
+
             auto run_condition = [&](const SDCondition& condition,
-                                     const sd::Tensor<float>* c_concat_override = nullptr,
-                                     const std::vector<int>* local_skip_layers  = nullptr) -> sd::Tensor<float> {
+                                     const sd::Tensor<float>* c_concat_override                 = nullptr,
+                                     const std::vector<int>* local_skip_layers                  = nullptr,
+                                     const std::vector<sd::Tensor<float>>* ref_latents_override = nullptr) -> sd::Tensor<float> {
                 diffusion_params.context     = condition.c_crossattn.empty() ? nullptr : &condition.c_crossattn;
                 diffusion_params.c_concat    = c_concat_override != nullptr ? c_concat_override : (condition.c_concat.empty() ? nullptr : &condition.c_concat);
                 diffusion_params.y           = condition.c_vector.empty() ? nullptr : &condition.c_vector;
-                diffusion_params.ref_latents = condition.c_ref_images.empty() ? &ref_latents : &condition.c_ref_images;
+                diffusion_params.ref_latents = ref_latents_override != nullptr ? ref_latents_override : (condition.c_ref_images.empty() ? &ref_latents : &condition.c_ref_images);
 
                 if (sd_version_is_unet(version)) {
                     diffusion_params.extra = UNetDiffusionExtra{-1, &controls, control_strength};
@@ -2083,31 +2203,40 @@ public:
                                             uncond,
                                             &controls);
                 }
-                uncond_out = run_condition(uncond);
+                const std::vector<int>* uncond_skip_layers = nullptr;
+                if (is_skiplayer_step && slg_uncond) {
+                    LOG_DEBUG("Skipping layers at uncond step %d\n", step);
+                    uncond_skip_layers = &skip_layer_guidance.layers();
+                }
+                uncond_out = run_condition(uncond,
+                                           uncond.c_concat.empty() ? nullptr : &uncond.c_concat,
+                                           uncond_skip_layers);
                 if (uncond_out.empty()) {
                     return {};
                 }
             }
-            if (!img_cond.empty()) {
-                img_cond_out = run_condition(img_cond,
-                                             cond.c_concat.empty() ? nullptr : &cond.c_concat);
-                if (img_cond_out.empty()) {
+            if (!img_uncond.empty()) {
+                img_uncond_out = run_condition(img_uncond,
+                                               img_uncond.c_concat.empty() ? nullptr : &img_uncond.c_concat,
+                                               nullptr,
+                                               uncond_without_ref_latents ? &empty_ref_latents : nullptr);
+                if (img_uncond_out.empty()) {
                     return {};
                 }
             }
             sd::guidance::GuidanceInput guidance_input;
-            guidance_input.step          = step;
-            guidance_input.schedule_size = sigmas.size();
-            guidance_input.pred_cond     = &cond_out;
-            guidance_input.pred_uncond   = uncond_out.empty() ? nullptr : &uncond_out;
-            guidance_input.pred_img_cond = img_cond_out.empty() ? nullptr : &img_cond_out;
+            guidance_input.step            = step;
+            guidance_input.schedule_size   = sigmas.size();
+            guidance_input.pred_cond       = &cond_out;
+            guidance_input.pred_uncond     = uncond_out.empty() ? nullptr : &uncond_out;
+            guidance_input.pred_img_uncond = img_uncond_out.empty() ? nullptr : &img_uncond_out;
 
-            sd::guidance::GuiderOutput guided = classifier_free_guidance.forward(guidance_input, {});
+            sd::guidance::GuiderOutput guided = primary_guidance.forward(guidance_input, {});
             if (guided.pred.empty()) {
                 return {};
             }
 
-            if (skip_layer_guidance.is_enabled_for_step(guidance_input)) {
+            if (is_skiplayer_step && slg_scale != 0.0f) {
                 LOG_DEBUG("Skipping layers at step %d\n", step);
                 if (!step_cache.is_step_skipped()) {
                     guidance_input.predict_skip_layer = [&]() -> sd::Tensor<float> {
@@ -2127,7 +2256,9 @@ public:
             sd::guidance::GuiderOutput output;
             output.pred = denoised;
             if (needs_uncond_denoised) {
-                const sd::Tensor<float>& base_uncond = !uncond_out.empty() ? uncond_out : cond_out;
+                const sd::Tensor<float>& base_uncond = !img_uncond_out.empty()
+                                                           ? img_uncond_out
+                                                           : (!uncond_out.empty() ? uncond_out : cond_out);
                 output.pred_uncond                   = base_uncond * c_out + x * c_skip;
             }
             if (cache_runtime.spectrum_enabled) {
@@ -2174,6 +2305,9 @@ public:
     }
 
     int get_vae_scale_factor() {
+        if (sd_version_is_pid(version)) {
+            return 1;
+        }
         return first_stage_model->get_scale_factor();
     }
 
@@ -2200,6 +2334,8 @@ public:
                 latent_channel = 3;
             } else if (version == VERSION_CHROMA_RADIANCE) {
                 latent_channel = 3;
+            } else if (sd_version_is_pid(version)) {
+                latent_channel = 3;
             } else if (sd_version_uses_flux2_vae(version)) {
                 latent_channel = 128;
             } else {
@@ -2277,6 +2413,18 @@ public:
     }
 
     sd::Tensor<float> decode_first_stage(const sd::Tensor<float>& x, bool decode_video = false) {
+        if (sd_version_is_pid(version)) {
+            return sd::ops::clamp((x + 1.f) * 0.5f, 0.0f, 1.0f);
+        }
+        // Free resident diffusion params before VAE allocates its compute buffer.
+        if (stream_layers) {
+            if (diffusion_model) {
+                diffusion_model->release_streaming_residency();
+            }
+            if (high_noise_diffusion_model) {
+                high_noise_diffusion_model->release_streaming_residency();
+            }
+        }
         auto latents = first_stage_model->diffusion_to_vae_latents(x);
         first_stage_model->set_temporal_tiling_enabled(vae_tiling_params.temporal_tiling);
         return first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y);
@@ -2537,6 +2685,35 @@ enum sd_hires_upscaler_t str_to_sd_hires_upscaler(const char* str) {
     return SD_HIRES_UPSCALER_COUNT;
 }
 
+const char* sd_vae_format_name(enum sd_vae_format_t format) {
+    switch (format) {
+        case SD_VAE_FORMAT_AUTO:
+            return "auto";
+        case SD_VAE_FORMAT_FLUX:
+            return "flux";
+        case SD_VAE_FORMAT_SD3:
+            return "sd3";
+        case SD_VAE_FORMAT_FLUX2:
+            return "flux2";
+        default:
+            return NONE_STR;
+    }
+}
+
+static SDVersion sd_vae_format_to_version(enum sd_vae_format_t format, SDVersion fallback) {
+    switch (format) {
+        case SD_VAE_FORMAT_FLUX:
+            return VERSION_FLUX;
+        case SD_VAE_FORMAT_SD3:
+            return VERSION_SD3;
+        case SD_VAE_FORMAT_FLUX2:
+            return VERSION_FLUX2;
+        case SD_VAE_FORMAT_AUTO:
+        default:
+            return fallback;
+    }
+}
+
 void sd_cache_params_init(sd_cache_params_t* cache_params) {
     *cache_params                             = {};
     cache_params->mode                        = SD_CACHE_DISABLED;
@@ -2592,6 +2769,7 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) {
     sd_ctx_params->lora_apply_mode         = LORA_APPLY_AUTO;
     sd_ctx_params->offload_params_to_cpu   = false;
     sd_ctx_params->max_vram                = 0.f;
+    sd_ctx_params->stream_layers           = false;
     sd_ctx_params->enable_mmap             = false;
     sd_ctx_params->keep_clip_on_cpu        = false;
     sd_ctx_params->keep_control_net_on_cpu = false;
@@ -2602,6 +2780,7 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) {
     sd_ctx_params->chroma_use_dit_mask     = true;
     sd_ctx_params->chroma_use_t5_mask      = false;
     sd_ctx_params->chroma_t5_mask_pad      = 1;
+    sd_ctx_params->vae_format              = SD_VAE_FORMAT_AUTO;
     sd_ctx_params->backend                 = nullptr;
     sd_ctx_params->params_backend          = nullptr;
 }
@@ -2622,6 +2801,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
              "llm_vision_path: %s\n"
              "diffusion_model_path: %s\n"
              "high_noise_diffusion_model_path: %s\n"
+             "uncond_diffusion_model_path: %s\n"
              "embeddings_connectors_path: %s\n"
              "vae_path: %s\n"
              "audio_vae_path: %s\n"
@@ -2638,6 +2818,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
              "prediction: %s\n"
              "offload_params_to_cpu: %s\n"
              "max_vram: %.3f\n"
+             "stream_layers: %s\n"
              "backend: %s\n"
              "params_backend: %s\n"
              "keep_clip_on_cpu: %s\n"
@@ -2649,7 +2830,8 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
              "circular_y: %s\n"
              "chroma_use_dit_mask: %s\n"
              "chroma_use_t5_mask: %s\n"
-             "chroma_t5_mask_pad: %d\n",
+             "chroma_t5_mask_pad: %d\n"
+             "vae_format: %s\n",
              SAFE_STR(sd_ctx_params->model_path),
              SAFE_STR(sd_ctx_params->clip_l_path),
              SAFE_STR(sd_ctx_params->clip_g_path),
@@ -2659,6 +2841,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
              SAFE_STR(sd_ctx_params->llm_vision_path),
              SAFE_STR(sd_ctx_params->diffusion_model_path),
              SAFE_STR(sd_ctx_params->high_noise_diffusion_model_path),
+             SAFE_STR(sd_ctx_params->uncond_diffusion_model_path),
              SAFE_STR(sd_ctx_params->embeddings_connectors_path),
              SAFE_STR(sd_ctx_params->vae_path),
              SAFE_STR(sd_ctx_params->audio_vae_path),
@@ -2675,6 +2858,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
              sd_prediction_name(sd_ctx_params->prediction),
              BOOL_STR(sd_ctx_params->offload_params_to_cpu),
              sd_ctx_params->max_vram,
+             BOOL_STR(sd_ctx_params->stream_layers),
              SAFE_STR(sd_ctx_params->backend),
              SAFE_STR(sd_ctx_params->params_backend),
              BOOL_STR(sd_ctx_params->keep_clip_on_cpu),
@@ -2686,7 +2870,8 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
              BOOL_STR(sd_ctx_params->circular_y),
              BOOL_STR(sd_ctx_params->chroma_use_dit_mask),
              BOOL_STR(sd_ctx_params->chroma_use_t5_mask),
-             sd_ctx_params->chroma_t5_mask_pad);
+             sd_ctx_params->chroma_t5_mask_pad,
+             sd_vae_format_name(sd_ctx_params->vae_format));
 
     return buf;
 }
@@ -2963,6 +3148,9 @@ SD_API bool sd_ctx_supports_video_generation(const sd_ctx_t* sd_ctx) {
 
 enum sample_method_t sd_get_default_sample_method(const sd_ctx_t* sd_ctx) {
     if (sd_ctx != nullptr && sd_ctx->sd != nullptr) {
+        if (sd_version_is_pid(sd_ctx->sd->version)) {
+            return LCM_SAMPLE_METHOD;
+        }
         if (sd_version_is_dit(sd_ctx->sd->version)) {
             return EULER_SAMPLE_METHOD;
         }
@@ -3046,9 +3234,10 @@ struct GenerationRequest {
     int diffusion_model_down_factor          = -1;
     int64_t seed                             = -1;
     bool use_uncond                          = false;
-    bool use_img_cond                        = false;
+    bool use_img_uncond                      = false;
     bool use_high_noise_uncond               = false;
-    bool use_high_noise_img_cond             = false;
+    bool use_high_noise_img_uncond           = false;
+    bool has_ref_images                      = false;
     const sd_cache_params_t* cache_params    = nullptr;
     int batch_count                          = 1;
     int shifted_timestep                     = 0;
@@ -3082,6 +3271,7 @@ struct GenerationRequest {
         eta                         = sd_img_gen_params->sample_params.eta;
         increase_ref_index          = sd_img_gen_params->increase_ref_index;
         auto_resize_ref_image       = sd_img_gen_params->auto_resize_ref_image;
+        has_ref_images              = sd_img_gen_params->ref_images_count > 0;
         guidance                    = sd_img_gen_params->sample_params.guidance;
         pm_params                   = sd_img_gen_params->pm_params;
         hires                       = sd_img_gen_params->hires;
@@ -3204,28 +3394,36 @@ struct GenerationRequest {
     static void resolve_guidance(sd_ctx_t* sd_ctx,
                                  sd_guidance_params_t* guidance,
                                  bool* use_uncond,
-                                 bool* use_img_cond,
+                                 bool* use_img_uncond,
+                                 bool has_ref_images,
                                  const char* stage_name = nullptr) {
         GGML_ASSERT(guidance != nullptr);
         GGML_ASSERT(use_uncond != nullptr);
-        GGML_ASSERT(use_img_cond != nullptr);
-        // out_uncond + text_cfg_scale * (out_cond - out_img_cond) + image_cfg_scale * (out_img_cond - out_uncond)
-        // img_cfg == txt_cfg means that img_cfg is not used
-        if (!std::isfinite(guidance->img_cfg)) {
-            guidance->img_cfg = guidance->txt_cfg;
+        GGML_ASSERT(use_img_uncond != nullptr);
+        // out_img_uncond + text_cfg_scale * (out_cond - out_uncond) + image_cfg_scale * (out_uncond - out_img_uncond)
+        // -> text_cfg_scale * out_cond + (image_cfg_scale - text_cfg_scale) * out_uncond + (1 - image_cfg_scale) * out_img_uncond
+        // out_cond       : prompt, image latent
+        // out_uncond     : negative prompt, image latent
+        // out_img_uncond : negative prompt, zero image latent
+        // image_cfg_scale == 1 reduces 3-cond CFG to 2-cond CFG.
+        bool img_cfg_was_set = std::isfinite(guidance->img_cfg);
+        if (!img_cfg_was_set) {
+            guidance->img_cfg = 1.f;
         }
 
-        if (!sd_version_is_inpaint_or_unet_edit(sd_ctx->sd->version)) {
-            guidance->img_cfg = guidance->txt_cfg;
-        }
-
-        if (guidance->txt_cfg != 1.f) {
-            *use_uncond = true;
+        if (!sd_version_supports_img_cfg(sd_ctx->sd->version, has_ref_images)) {
+            if (img_cfg_was_set && guidance->img_cfg != 1.f) {
+                LOG_WARN("3-conditioning CFG is not supported with this model, disabling it for better performance");
+            }
+            guidance->img_cfg = 1.f;
         }
 
         if (guidance->img_cfg != guidance->txt_cfg) {
-            *use_img_cond = true;
-            *use_uncond   = true;
+            *use_uncond = true;
+        }
+
+        if (guidance->img_cfg != 1.f) {
+            *use_img_uncond = true;
         }
 
         if (guidance->txt_cfg < 1.f) {
@@ -3244,12 +3442,13 @@ struct GenerationRequest {
         resolve_hires();
         seed = resolve_seed(seed);
 
-        resolve_guidance(sd_ctx, &guidance, &use_uncond, &use_img_cond);
+        resolve_guidance(sd_ctx, &guidance, &use_uncond, &use_img_uncond, has_ref_images);
         if (sd_ctx->sd->high_noise_diffusion_model) {
             resolve_guidance(sd_ctx,
                              &high_noise_guidance,
                              &use_high_noise_uncond,
-                             &use_high_noise_img_cond,
+                             &use_high_noise_img_uncond,
+                             has_ref_images,
                              "high noise: ");
         }
 
@@ -3367,7 +3566,7 @@ struct SamplePlan {
 struct ImageGenerationLatents {
     sd::Tensor<float> init_latent;
     sd::Tensor<float> concat_latent;
-    sd::Tensor<float> uncond_concat_latent;
+    sd::Tensor<float> img_uncond_concat_latent;
     sd::Tensor<float> audio_latent;
     sd::Tensor<float> video_positions;
     sd::Tensor<float> control_image;
@@ -3690,7 +3889,7 @@ static int get_ltxav_num_audio_latents(int frames, int fps) {
 struct ImageGenerationEmbeds {
     SDCondition cond;
     SDCondition uncond;
-    SDCondition img_cond;
+    SDCondition img_uncond;
     SDCondition id_cond;
 };
 
@@ -3849,6 +4048,7 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
         LOG_WARN("This model needs at least one reference image; using an empty reference");
         ref_images.push_back(sd::zeros<float>({request->width, request->height, 3, 1}));
         request->guidance.img_cfg = request->guidance.txt_cfg;
+        request->use_img_uncond   = false;
     }
 
     if (!ref_images.empty()) {
@@ -3861,7 +4061,7 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
             continue;
         }
         sd::Tensor<float> ref_latent;
-        if (request->auto_resize_ref_image) {
+        if (request->auto_resize_ref_image && !sd_version_is_pid(sd_ctx->sd->version)) {
             LOG_DEBUG("auto resize ref images");
             int vae_image_size = std::min(1024 * 1024, request->width * request->height);
             double vae_width   = sqrt(vae_image_size * ref_images[i].shape()[0] / ref_images[i].shape()[1]);
@@ -3893,8 +4093,15 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
         ref_latents.push_back(std::move(ref_latent));
     }
 
+    if (sd_version_is_pid(sd_ctx->sd->version)) {
+        if (ref_latents.empty()) {
+            LOG_ERROR("PiD requires a reference image");
+            return std::nullopt;
+        }
+    }
+
     sd::Tensor<float> concat_latent;
-    sd::Tensor<float> uncond_concat_latent;
+    sd::Tensor<float> img_uncond_concat_latent;
     if (sd_version_is_inpaint(sd_ctx->sd->version)) {
         sd::Tensor<float> masked_init_latent;
 
@@ -3922,8 +4129,8 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
                                                    request->height / request->vae_scale_factor});
             mask      = mask.permute({1, 3, 0, 2}).reshape({request->width / request->vae_scale_factor, request->height / request->vae_scale_factor, request->vae_scale_factor * request->vae_scale_factor, 1});
 
-            concat_latent        = sd::ops::concat(masked_init_latent, mask, 2);
-            uncond_concat_latent = sd::ops::concat(uncond_masked_init_latent, mask, 2);
+            concat_latent            = sd::ops::concat(masked_init_latent, mask, 2);
+            img_uncond_concat_latent = sd::ops::concat(uncond_masked_init_latent, mask, 2);
         } else if (sd_ctx->sd->version == VERSION_FLEX_2) {
             concat_latent = sd::ops::concat(masked_init_latent, latent_mask, 2);
             if (!control_latent.empty()) {
@@ -3932,16 +4139,16 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
                 concat_latent = sd::ops::concat(concat_latent, sd::Tensor<float>::zeros_like(masked_init_latent), 2);
             }
 
-            uncond_concat_latent = sd::ops::concat(uncond_masked_init_latent, latent_mask, 2);
-            uncond_concat_latent = sd::ops::concat(uncond_concat_latent, sd::Tensor<float>::zeros_like(masked_init_latent), 2);
+            img_uncond_concat_latent = sd::ops::concat(uncond_masked_init_latent, latent_mask, 2);
+            img_uncond_concat_latent = sd::ops::concat(img_uncond_concat_latent, sd::Tensor<float>::zeros_like(masked_init_latent), 2);
         } else {  // SD1.x SD2.x SDXL inpaint
-            concat_latent        = sd::ops::concat(latent_mask, masked_init_latent, 2);
-            uncond_concat_latent = sd::ops::concat(latent_mask, uncond_masked_init_latent, 2);
+            concat_latent            = sd::ops::concat(latent_mask, masked_init_latent, 2);
+            img_uncond_concat_latent = sd::ops::concat(latent_mask, uncond_masked_init_latent, 2);
         }
     }
     if (sd_version_is_unet_edit(sd_ctx->sd->version)) {
-        concat_latent        = sd::ops::interpolate<float>(ref_latents[0], init_latent.shape());
-        uncond_concat_latent = sd::Tensor<float>::zeros_like(concat_latent);
+        concat_latent            = sd::ops::interpolate<float>(ref_latents[0], init_latent.shape());
+        img_uncond_concat_latent = sd::Tensor<float>::zeros_like(concat_latent);
     }
     if (sd_ctx->sd->version == VERSION_FLUX_CONTROLS) {
         if (!control_latent.empty()) {
@@ -3949,7 +4156,7 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
         } else {
             concat_latent = sd::Tensor<float>::zeros_like(init_latent);
         }
-        uncond_concat_latent = sd::Tensor<float>::zeros_like(concat_latent);
+        img_uncond_concat_latent = sd::Tensor<float>::zeros_like(concat_latent);
     }
 
     if (sd_img_gen_params->init_image.data != nullptr || sd_img_gen_params->ref_images_count > 0) {
@@ -3958,12 +4165,12 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
     }
 
     ImageGenerationLatents latents;
-    latents.init_latent          = std::move(init_latent);
-    latents.concat_latent        = std::move(concat_latent);
-    latents.uncond_concat_latent = std::move(uncond_concat_latent);
-    latents.control_image        = std::move(control_image_tensor);
-    latents.ref_images           = std::move(ref_images);
-    latents.ref_latents          = std::move(ref_latents);
+    latents.init_latent              = std::move(init_latent);
+    latents.concat_latent            = std::move(concat_latent);
+    latents.img_uncond_concat_latent = std::move(img_uncond_concat_latent);
+    latents.control_image            = std::move(control_image_tensor);
+    latents.ref_images               = std::move(ref_images);
+    latents.ref_latents              = std::move(ref_latents);
 
     if (sd_version_is_inpaint(sd_ctx->sd->version)) {
         latent_mask = sd::ops::max_pool_2d(latent_mask,
@@ -3997,20 +4204,53 @@ static std::optional<ImageGenerationEmbeds> prepare_image_generation_embeds(sd_c
         cond.c_concat = latents->concat_latent;  // TODO: optimize
     }
 
+    bool use_ref_latent_img_cfg = request->use_img_uncond &&
+                                  !latents->ref_images.empty() &&
+                                  sd_version_supports_ref_latent_img_cfg(sd_ctx->sd->version);
+
     SDCondition uncond;
     if (request->use_uncond || request->use_high_noise_uncond) {
-        bool zero_out_masked = false;
-        if (sd_version_is_sdxl(sd_ctx->sd->version) &&
-            request->negative_prompt.empty() &&
-            !sd_ctx->sd->is_using_edm_v_parameterization) {
-            zero_out_masked = true;
+        if (sd_version_is_ideogram4(sd_ctx->sd->version)) {
+            uncond.c_vector = sd::Tensor<float>::from_vector({1.0f});
+        } else {
+            bool zero_out_masked = false;
+            if (sd_version_is_sdxl(sd_ctx->sd->version) &&
+                request->negative_prompt.empty() &&
+                !sd_ctx->sd->is_using_edm_v_parameterization) {
+                zero_out_masked = true;
+            }
+            condition_params.text            = request->negative_prompt;
+            condition_params.zero_out_masked = zero_out_masked;
+            uncond                           = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads,
+                                                                                                   condition_params);
         }
-        condition_params.text            = request->negative_prompt;
-        condition_params.zero_out_masked = zero_out_masked;
-        uncond                           = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads,
-                                                                                               condition_params);
         if (uncond.c_concat.empty()) {
-            uncond.c_concat = latents->uncond_concat_latent;  // TODO: optimize
+            uncond.c_concat = latents->concat_latent;  // TODO: optimize
+        }
+    }
+
+    SDCondition img_uncond;
+    if (request->use_img_uncond) {
+        if ((request->use_uncond || request->use_high_noise_uncond) && (latents->ref_images.empty() || !use_ref_latent_img_cfg)) {
+            img_uncond = SDCondition(uncond.c_crossattn, uncond.c_vector, latents->img_uncond_concat_latent);
+        } else {
+            bool zero_out_masked = false;
+            if (sd_version_is_sdxl(sd_ctx->sd->version) &&
+                request->negative_prompt.empty() &&
+                !sd_ctx->sd->is_using_edm_v_parameterization) {
+                zero_out_masked = true;
+            }
+            condition_params.text            = request->negative_prompt;
+            condition_params.zero_out_masked = zero_out_masked;
+            if (use_ref_latent_img_cfg) {
+                std::vector<sd::Tensor<float>> empty_ref_images;
+                condition_params.ref_images = &empty_ref_images;
+            }
+            img_uncond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads,
+                                                                             condition_params);
+            if (img_uncond.c_concat.empty()) {
+                img_uncond.c_concat = latents->img_uncond_concat_latent;  // TODO: optimize
+            }
         }
     }
 
@@ -4022,12 +4262,10 @@ static std::optional<ImageGenerationEmbeds> prepare_image_generation_embeds(sd_c
     }
 
     ImageGenerationEmbeds embeds;
-    if (request->use_img_cond) {
-        embeds.img_cond = SDCondition(uncond.c_crossattn, uncond.c_vector, cond.c_concat);
-    }
-    embeds.cond    = std::move(cond);
-    embeds.uncond  = std::move(uncond);
-    embeds.id_cond = std::move(id_cond);
+    embeds.img_uncond = std::move(img_uncond);
+    embeds.cond       = std::move(cond);
+    embeds.uncond     = std::move(uncond);
+    embeds.id_cond    = std::move(id_cond);
 
     return embeds;
 }
@@ -4307,7 +4545,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
                                                    std::move(noise),
                                                    embeds.cond,
                                                    embeds.uncond,
-                                                   embeds.img_cond,
+                                                   embeds.img_uncond,
                                                    embeds.id_cond,
                                                    latents.control_image,
                                                    request.control_strength,
@@ -4427,7 +4665,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s
                                                             std::move(noise),
                                                             embeds.cond,
                                                             embeds.uncond,
-                                                            embeds.img_cond,
+                                                            embeds.img_uncond,
                                                             embeds.id_cond,
                                                             latents.control_image,
                                                             request.control_strength,
@@ -4695,6 +4933,17 @@ static std::optional<ImageGenerationLatents> prepare_video_generation_latents(sd
         latents.denoise_mask = sd::full<float>({latents.init_latent.shape()[0], latents.init_latent.shape()[1], latents.init_latent.shape()[2], 1, 1}, 1.f);
         sd::ops::fill_slice(&latents.denoise_mask, 2, 0, init_image_latent.shape()[2], 0.0f);
 
+        if (!end_image.empty()) {
+            auto end_img          = end_image.reshape({end_image.shape()[0], end_image.shape()[1], 1, end_image.shape()[2], 1});
+            auto end_image_latent = sd_ctx->sd->encode_first_stage(end_img);  // [b, c, 1, h/vae_scale_factor, w/vae_scale_factor]
+            if (end_image_latent.empty()) {
+                LOG_ERROR("failed to encode end video frame");
+                return std::nullopt;
+            }
+            sd::ops::slice_assign(&latents.init_latent, 2, latents.init_latent.shape()[2] - 1, latents.init_latent.shape()[2], end_image_latent);
+            sd::ops::fill_slice(&latents.denoise_mask, 2, latents.init_latent.shape()[2] - 1, latents.init_latent.shape()[2], 0.0f);
+        }
+
         int64_t t2 = ggml_time_ms();
         LOG_INFO("encode_first_stage completed, taking %" PRId64 " ms", t2 - t1);
     } else if (sd_ctx->sd->diffusion_model->get_desc() == "Wan2.1-VACE-1.3B" ||
@@ -5131,7 +5380,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
                                                            std::move(noise),
                                                            embeds.cond,
                                                            request.use_high_noise_uncond ? embeds.uncond : SDCondition(),
-                                                           embeds.img_cond,
+                                                           embeds.img_uncond,
                                                            embeds.id_cond,
                                                            sd::Tensor<float>(),
                                                            0.f,
@@ -5177,7 +5426,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
                                                         std::move(noise),
                                                         embeds.cond,
                                                         request.use_uncond ? embeds.uncond : SDCondition(),
-                                                        embeds.img_cond,
+                                                        embeds.img_uncond,
                                                         embeds.id_cond,
                                                         sd::Tensor<float>(),
                                                         0.f,
@@ -5321,7 +5570,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
                                             std::move(noise),
                                             embeds.cond,
                                           hires_request.use_uncond ? embeds.uncond : SDCondition(),
-                                            embeds.img_cond,
+                                            embeds.img_uncond,
                                             embeds.id_cond,
                                             sd::Tensor<float>(),
                                             0.f,

src/t5.hpp

@@ -14,6 +14,28 @@
 #include "model.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 +294,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,7 +329,7 @@ public:
 };
 
 struct T5Runner : public GGMLRunner {
-    T5Params params;
+    T5Config config;
     T5 model;
     std::vector<int> relative_position_bucket_vec;
 
@@ -325,12 +338,9 @@ struct T5Runner : public GGMLRunner {
              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);
+        : GGMLRunner(backend, params_backend),
+          config(T5Config::detect_from_weights(tensor_storage_map, prefix, is_umt5)) {
+        model = T5(config);
         model.init(params_ctx, tensor_storage_map, prefix);
     }
 
@@ -475,8 +485,11 @@ struct T5Embedder {
         model.get_param_tensors(tensors, prefix);
     }
 
-    void alloc_params_buffer() {
-        model.alloc_params_buffer();
+    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,
@@ -560,7 +573,7 @@ struct T5Embedder {
         // cuda f32: pass
         // cuda q8_0: pass
         // ggml_backend_t backend = ggml_backend_cuda_init(0);
-        ggml_backend_t backend    = ggml_backend_cpu_init();
+        ggml_backend_t backend    = sd_backend_cpu_init();
         ggml_type model_data_type = GGML_TYPE_F16;
 
         ModelLoader model_loader;
@@ -578,7 +591,10 @@ struct T5Embedder {
 
         std::shared_ptr<T5Embedder> t5 = std::make_shared<T5Embedder>(backend, backend, tensor_storage_map, "", true);
 
-        t5->alloc_params_buffer();
+        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, "");
 

src/tensor.hpp

@@ -235,6 +235,7 @@ namespace sd {
 
         Tensor& masked_fill_(const Tensor<uint8_t>& mask, const T& value);
 
+        T sum() const;
         T mean() const;
 
         static Tensor zeros(std::vector<int64_t> shape) {
@@ -327,6 +328,24 @@ namespace sd {
         std::vector<int64_t> shape_;
     };
 
+    template <typename T>
+    inline T Tensor<T>::sum() const {
+        T total = T{};
+        for (const T& value : data_) {
+            total += value;
+        }
+        return total;
+    }
+
+    template <>
+    inline float Tensor<float>::sum() const {
+        double total = 0.0;
+        for (float value : data_) {
+            total += static_cast<double>(value);
+        }
+        return static_cast<float>(total);
+    }
+
     template <typename T>
     inline T Tensor<T>::mean() const {
         if (empty()) {

src/tokenizers/bpe_tokenizer.cpp

@@ -131,11 +131,10 @@ std::vector<std::u32string> BPETokenizer::bpe(const std::u32string& token) const
 }
 
 std::vector<int> BPETokenizer::encode(const std::string& text, on_new_token_cb_t on_new_token_cb) {
-    std::string normalized_text = normalize(text);
     std::vector<int32_t> bpe_tokens;
     std::vector<std::string> token_strs;
 
-    auto splited_texts = split_with_special_tokens(normalized_text, special_tokens);
+    auto splited_texts = split_with_special_tokens(text, special_tokens);
 
     for (auto& splited_text : splited_texts) {
         if (is_special_token(splited_text)) {
@@ -160,7 +159,7 @@ std::vector<int> BPETokenizer::encode(const std::string& text, on_new_token_cb_t
                 }
             }
 
-            std::string token_str = token;
+            std::string token_str = normalize(token);
             std::u32string utf32_token;
             if (byte_level_bpe) {
                 for (int i = 0; i < token_str.length(); i++) {
@@ -183,7 +182,8 @@ std::vector<int> BPETokenizer::encode(const std::string& text, on_new_token_cb_t
                             unsigned char b = utf8_token_str[i];
                             char hex_buf[16];
                             snprintf(hex_buf, sizeof(hex_buf), "<0x%02X>", b);
-                            iter = encoder.find(utf8_to_utf32(hex_buf));
+                            iter     = encoder.find(utf8_to_utf32(hex_buf));
+                            token_id = iter != encoder.end() ? iter->second : UNK_TOKEN_ID;
                             bpe_tokens.push_back(token_id);
                             token_strs.push_back(hex_buf);
                         }

src/tokenizers/gemma_tokenizer.cpp

@@ -189,3 +189,164 @@ GemmaTokenizer::GemmaTokenizer(const std::string& merges_utf8_str, const std::st
         load_from_merges(load_gemma_merges(), load_gemma_vocab_json());
     }
 }
+
+std::string Gemma2Tokenizer::normalize(const std::string& text) const {
+    std::string normalized = text;
+    size_t pos             = 0;
+    while ((pos = normalized.find(' ', pos)) != std::string::npos) {
+        normalized.replace(pos, 1, "\xE2\x96\x81");
+        pos += 3;
+    }
+    return normalized;
+}
+
+void Gemma2Tokenizer::load_from_merges(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
+    nlohmann::json vocab;
+    try {
+        vocab = nlohmann::json::parse(vocab_utf8_str);
+    } catch (const nlohmann::json::parse_error&) {
+        GGML_ABORT("invalid vocab json str");
+    }
+    for (const auto& [key, value] : vocab.items()) {
+        std::u32string token = utf8_to_utf32(key);
+        int i                = value;
+        encoder[token]       = i;
+        decoder[i]           = token;
+    }
+    encoder_len = static_cast<int>(vocab.size());
+    LOG_DEBUG("vocab size: %d", encoder_len);
+
+    std::vector<std::u32string> merges = split_utf32(merges_utf8_str);
+    std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;
+    for (const auto& merge : merges) {
+        size_t space_pos = merge.find(' ');
+        merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1));
+    }
+    LOG_DEBUG("merges size %zu", merge_pairs.size());
+
+    int rank = 0;
+    for (const auto& merge : merge_pairs) {
+        bpe_ranks[merge] = rank++;
+    }
+    bpe_len = rank;
+}
+
+Gemma2Tokenizer::Gemma2Tokenizer(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
+    byte_level_bpe = false;
+    byte_fallback  = true;
+    add_bos_token  = true;
+    PAD_TOKEN      = "<pad>";
+    EOS_TOKEN      = "<eos>";
+    BOS_TOKEN      = "<bos>";
+    UNK_TOKEN      = "<unk>";
+
+    PAD_TOKEN_ID = 0;
+    EOS_TOKEN_ID = 1;
+    BOS_TOKEN_ID = 2;
+    UNK_TOKEN_ID = 3;
+
+    std::vector<std::string> special_tokens_before_merge = {
+        PAD_TOKEN,
+        EOS_TOKEN,
+        BOS_TOKEN,
+        UNK_TOKEN,
+        "<mask>",
+        "<2mass>",
+        "[@BOS@]",
+    };
+    for (int i = 0; i <= 98; i++) {
+        special_tokens_before_merge.push_back("<unused" + std::to_string(i) + ">");
+    }
+    special_tokens_before_merge.push_back("<start_of_turn>");
+    special_tokens_before_merge.push_back("<end_of_turn>");
+    for (int i = 1; i <= 31; i++) {
+        special_tokens_before_merge.push_back(std::string(i, '\n'));
+    }
+    for (int i = 2; i <= 31; i++) {
+        std::string whitespace_token;
+        for (int j = 0; j < i; j++) {
+            whitespace_token += "\xE2\x96\x81";
+        }
+        special_tokens_before_merge.push_back(whitespace_token);
+    }
+    std::vector<std::string> html_tokens = {
+        "<table>",
+        "<caption>",
+        "<thead>",
+        "<tbody>",
+        "<tfoot>",
+        "<tr>",
+        "<th>",
+        "<td>",
+        "</table>",
+        "</caption>",
+        "</thead>",
+        "</tbody>",
+        "</tfoot>",
+        "</tr>",
+        "</th>",
+        "</td>",
+        "<h1>",
+        "<h2>",
+        "<h3>",
+        "<h4>",
+        "<h5>",
+        "<h6>",
+        "<blockquote>",
+        "</h1>",
+        "</h2>",
+        "</h3>",
+        "</h4>",
+        "</h5>",
+        "</h6>",
+        "</blockquote>",
+        "<strong>",
+        "<em>",
+        "<b>",
+        "<i>",
+        "<u>",
+        "<s>",
+        "<sub>",
+        "<sup>",
+        "<code>",
+        "</strong>",
+        "</em>",
+        "</b>",
+        "</i>",
+        "</u>",
+        "</s>",
+        "</sub>",
+        "</sup>",
+        "</code>",
+    };
+    special_tokens_before_merge.insert(special_tokens_before_merge.end(),
+                                       html_tokens.begin(),
+                                       html_tokens.end());
+    for (int i = 0; i <= 0xFF; i++) {
+        char hex_buf[16];
+        snprintf(hex_buf, sizeof(hex_buf), "<0x%02X>", i);
+        special_tokens_before_merge.push_back(hex_buf);
+    }
+
+    std::vector<std::string> special_tokens_after_merge = {
+        "[toxicity=0]",
+    };
+    for (int i = 1; i <= 31; i++) {
+        special_tokens_after_merge.insert(special_tokens_after_merge.begin() + i - 1,
+                                          std::string(i, '\t'));
+    }
+    for (int i = 99; i <= 99; i++) {
+        special_tokens_after_merge.push_back("<unused" + std::to_string(i) + ">");
+    }
+
+    special_tokens = special_tokens_before_merge;
+    special_tokens.insert(special_tokens.end(),
+                          special_tokens_after_merge.begin(),
+                          special_tokens_after_merge.end());
+
+    if (merges_utf8_str.size() > 0 && vocab_utf8_str.size() > 0) {
+        load_from_merges(merges_utf8_str, vocab_utf8_str);
+    } else {
+        load_from_merges(load_gemma2_merges(), load_gemma2_vocab_json());
+    }
+}

src/tokenizers/gemma_tokenizer.h

@@ -14,4 +14,13 @@ public:
     explicit GemmaTokenizer(const std::string& merges_utf8_str = "", const std::string& vocab_utf8_str = "");
 };
 
+class Gemma2Tokenizer : public BPETokenizer {
+protected:
+    void load_from_merges(const std::string& merges_utf8_str, const std::string& vocab_utf8_str);
+    std::string normalize(const std::string& text) const override;
+
+public:
+    explicit Gemma2Tokenizer(const std::string& merges_utf8_str = "", const std::string& vocab_utf8_str = "");
+};
+
 #endif  // __SD_TOKENIZERS_GEMMA_TOKENIZER_H__

src/tokenizers/vocab/vocab.cpp

@@ -1,5 +1,7 @@
 #include "vocab.h"
 #include "clip_merges.hpp"
+#include "gemma2_merges.hpp"
+#include "gemma2_vocab.hpp"
 #include "gemma_merges.hpp"
 #include "gemma_vocab.hpp"
 #include "gpt_oss_merges.hpp"
@@ -50,6 +52,16 @@ std::string load_gemma_vocab_json() {
     return json_str;
 }
 
+std::string load_gemma2_merges() {
+    std::string merges_utf8_str(reinterpret_cast<const char*>(gemma2_merges_utf8_c_str), sizeof(gemma2_merges_utf8_c_str));
+    return merges_utf8_str;
+}
+
+std::string load_gemma2_vocab_json() {
+    std::string json_str(reinterpret_cast<const char*>(gemma2_vocab_json_utf8_c_str), sizeof(gemma2_vocab_json_utf8_c_str));
+    return json_str;
+}
+
 std::string load_gpt_oss_merges() {
     std::string merges_utf8_str(reinterpret_cast<const char*>(gpt_oss_merges_utf8_c_str), sizeof(gpt_oss_merges_utf8_c_str));
     return merges_utf8_str;

src/tokenizers/vocab/vocab.h

@@ -11,6 +11,8 @@ std::string load_t5_tokenizer_json();
 std::string load_umt5_tokenizer_json();
 std::string load_gemma_merges();
 std::string load_gemma_vocab_json();
+std::string load_gemma2_merges();
+std::string load_gemma2_vocab_json();
 std::string load_gpt_oss_merges();
 std::string load_gpt_oss_vocab_json();
 

src/unet.hpp

@@ -1,6 +1,9 @@
 #ifndef __UNET_HPP__
 #define __UNET_HPP__
 
+#include <algorithm>
+#include <vector>
+
 #include "common_block.hpp"
 #include "diffusion_model.hpp"
 #include "model.h"
@@ -9,6 +12,125 @@
 
 #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,6 +690,7 @@ public:
 };
 
 struct UNetModelRunner : public DiffusionModelRunner {
+    UNetConfig config;
     UnetModelBlock unet;
 
     UNetModelRunner(ggml_backend_t backend,
@@ -608,7 +698,9 @@ struct UNetModelRunner : public DiffusionModelRunner {
                     const String2TensorStorage& tensor_storage_map,
                     const std::string prefix,
                     SDVersion version = VERSION_SD1)
-        : DiffusionModelRunner(backend, params_backend, prefix), unet(version, tensor_storage_map) {
+        : DiffusionModelRunner(backend, params_backend, prefix),
+          config(UNetConfig::detect_from_weights(tensor_storage_map, prefix, version)),
+          unet(config) {
         unet.init(params_ctx, tensor_storage_map, prefix);
     }
 

src/upscaler.cpp

@@ -25,6 +25,13 @@ void UpscalerGGML::set_max_graph_vram_bytes(size_t max_vram_bytes) {
     }
 }
 
+void UpscalerGGML::set_stream_layers_enabled(bool enabled) {
+    stream_layers_enabled = enabled;
+    if (esrgan_upscaler) {
+        esrgan_upscaler->set_stream_layers_enabled(enabled);
+    }
+}
+
 bool UpscalerGGML::load_from_file(const std::string& esrgan_path,
                                   bool offload_params_to_cpu,
                                   int n_threads) {
@@ -76,6 +83,7 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path,
                                                tile_size,
                                                model_loader.get_tensor_storage_map());
     esrgan_upscaler->set_max_graph_vram_bytes(max_graph_vram_bytes);
+    esrgan_upscaler->set_stream_layers_enabled(stream_layers_enabled);
     if (direct) {
         esrgan_upscaler->set_conv2d_direct_enabled(true);
     }

src/upscaler.h

@@ -18,6 +18,7 @@ struct UpscalerGGML {
     bool direct                 = false;
     int tile_size               = 128;
     size_t max_graph_vram_bytes = 0;
+    bool stream_layers_enabled  = false;
     std::string backend_spec;
     std::string params_backend_spec;
 
@@ -31,6 +32,7 @@ struct UpscalerGGML {
                         bool offload_params_to_cpu,
                         int n_threads);
     void set_max_graph_vram_bytes(size_t max_vram_bytes);
+    void set_stream_layers_enabled(bool enabled);
     sd::Tensor<float> upscale_tensor(const sd::Tensor<float>& input_tensor);
     sd_image_t upscale(sd_image_t input_image, uint32_t upscale_factor);
 };

src/util.cpp

@@ -25,9 +25,7 @@
 #include <unistd.h>
 #endif
 
-#include "ggml-backend.h"
 #include "ggml.h"
-#include "ggml_extend_backend.h"
 #include "stable-diffusion.h"
 
 bool ends_with(const std::string& str, const std::string& ending) {
@@ -974,30 +972,3 @@ std::vector<std::pair<std::string, float>> split_quotation_attention(
     }
     return result;
 }
-
-// namespace is needed to avoid conflicts with ggml_backend_extend.hpp
-namespace ggml_cpu {
-#include "ggml-cpu.h"
-}
-
-const char* sd_get_system_info() {
-    using namespace ggml_cpu;
-    static char buffer[1024];
-    std::stringstream ss;
-    ss << "System Info: \n";
-    ss << "    SSE3 = " << ggml_cpu_has_sse3() << " | ";
-    ss << "    AVX = " << ggml_cpu_has_avx() << " | ";
-    ss << "    AVX2 = " << ggml_cpu_has_avx2() << " | ";
-    ss << "    AVX512 = " << ggml_cpu_has_avx512() << " | ";
-    ss << "    AVX512_VBMI = " << ggml_cpu_has_avx512_vbmi() << " | ";
-    ss << "    AVX512_VNNI = " << ggml_cpu_has_avx512_vnni() << " | ";
-    ss << "    FMA = " << ggml_cpu_has_fma() << " | ";
-    ss << "    NEON = " << ggml_cpu_has_neon() << " | ";
-    ss << "    ARM_FMA = " << ggml_cpu_has_arm_fma() << " | ";
-    ss << "    F16C = " << ggml_cpu_has_f16c() << " | ";
-    ss << "    FP16_VA = " << ggml_cpu_has_fp16_va() << " | ";
-    ss << "    WASM_SIMD = " << ggml_cpu_has_wasm_simd() << " | ";
-    ss << "    VSX = " << ggml_cpu_has_vsx() << " | ";
-    snprintf(buffer, sizeof(buffer), "%s", ss.str().c_str());
-    return buffer;
-}

src/z_image.hpp

@@ -20,6 +20,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 +361,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 +471,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 +488,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 +520,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,7 +545,7 @@ 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;
@@ -478,8 +556,9 @@ namespace ZImage {
                      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);
+            : DiffusionModelRunner(backend, params_backend, prefix),
+              config(ZImageConfig::detect_from_weights(tensor_storage_map, prefix)) {
+            z_image = ZImageModel(config);
             z_image.init(params_ctx, tensor_storage_map, prefix);
         }
 
@@ -510,19 +589,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;
@@ -615,7 +694,7 @@ namespace ZImage {
             // cuda q8: pass
             // cuda q8 fa: pass
             // ggml_backend_t backend = ggml_backend_cuda_init(0);
-            ggml_backend_t backend    = ggml_backend_cpu_init();
+            ggml_backend_t backend    = sd_backend_cpu_init();
             ggml_type model_data_type = GGML_TYPE_Q8_0;
 
             ModelLoader model_loader;
@@ -639,7 +718,10 @@ namespace ZImage {
                                                                                    "model.diffusion_model",
                                                                                    VERSION_QWEN_IMAGE);
 
-            z_image->alloc_params_buffer();
+            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");