fix(cmake): build HIP backend with PIC so the static-lib PIE link succeeds (#1593)

Co-authored-by: leejet <leejet714@gmail.com>

.github/workflows/build.yml

@@ -545,6 +545,30 @@ jobs:
       - name: Pack artifacts
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
         run: |
+          $ErrorActionPreference = "Stop"
+          $dst = "build\bin"
+          $rocmBin = Join-Path "${env:HIP_PATH}" "bin"
+          $requiredRocmPaths = @(
+            (Join-Path $rocmBin "rocblas.dll"),
+            (Join-Path $rocmBin "rocblas\library")
+          )
+          foreach ($path in $requiredRocmPaths) {
+            if (!(Test-Path $path)) {
+              throw "Missing ROCm runtime dependency: $path"
+            }
+          }
+
+          foreach ($pattern in @("rocblas*.dll", "hipblas*.dll", "libhipblas*.dll")) {
+            Copy-Item -Path (Join-Path $rocmBin $pattern) -Destination $dst -Force -ErrorAction SilentlyContinue
+          }
+
+          foreach ($dir in @("rocblas", "hipblaslt")) {
+            $src = Join-Path $rocmBin $dir
+            if (Test-Path $src) {
+              Copy-Item -Path $src -Destination $dst -Recurse -Force
+            }
+          }
+
           7z a sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-rocm-${{ env.ROCM_VERSION }}-x64.zip .\build\bin\*
 
       - name: Upload artifacts
@@ -687,16 +711,6 @@ jobs:
         with:
           submodules: recursive
 
-      - name: Setup Node
-        uses: actions/setup-node@v4
-        with:
-          node-version: 20
-
-      - name: Setup pnpm
-        uses: pnpm/action-setup@v4
-        with:
-          version: 10.15.1
-
       - name: ccache
         uses: ggml-org/ccache-action@v1.2.16
         with:
@@ -754,6 +768,18 @@ jobs:
           echo PATH=$PATH:$ROCM_PATH/bin >> $GITHUB_ENV
           echo LD_LIBRARY_PATH=$ROCM_PATH/lib:$ROCM_PATH/llvm/lib:$ROCM_PATH/lib/rocprofiler-systems >> $GITHUB_ENV
 
+      # setup-node installs into /opt/hostedtoolcache, which is removed above.
+      # Keep Node/pnpm setup after disk cleanup so the server frontend can be embedded.
+      - name: Setup Node
+        uses: actions/setup-node@v4
+        with:
+          node-version: 20
+
+      - name: Setup pnpm
+        uses: pnpm/action-setup@v4
+        with:
+          version: 10.15.1
+
       - name: Build
         id: cmake_build
         run: |

CMakeLists.txt

@@ -22,6 +22,31 @@ endif()
 set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin)
 set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin)
 
+if(APPLE)
+    function(sd_set_macos_rpaths target)
+        get_target_property(target_type ${target} TYPE)
+        if(target_type STREQUAL "EXECUTABLE")
+            set(runtime_paths "@executable_path" "@executable_path/../lib")
+        elseif(target_type STREQUAL "SHARED_LIBRARY" OR target_type STREQUAL "MODULE_LIBRARY")
+            set(runtime_paths "@loader_path" "@loader_path/../lib")
+            set_target_properties(${target} PROPERTIES
+                MACOSX_RPATH ON
+                INSTALL_NAME_DIR "@rpath"
+                BUILD_WITH_INSTALL_NAME_DIR ON
+            )
+        else()
+            return()
+        endif()
+
+        # Release artifacts zip the build output directly, so keep macOS rpaths relocatable.
+        set_target_properties(${target} PROPERTIES
+            BUILD_RPATH "${runtime_paths}"
+            INSTALL_RPATH "${runtime_paths}"
+            BUILD_WITH_INSTALL_RPATH ON
+        )
+    endfunction()
+endif()
+
 if(CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR)
     set(SD_STANDALONE ON)
 else()
@@ -100,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)
@@ -237,6 +268,10 @@ else()
     add_library(${SD_LIB} STATIC ${SD_LIB_SOURCES})
 endif()
 
+if(APPLE)
+    sd_set_macos_rpaths(${SD_LIB})
+endif()
+
 if(SD_SYCL)
     message("-- Use SYCL as backend stable-diffusion")
     set(GGML_SYCL ON)

README.md

@@ -15,6 +15,8 @@ API and command-line option may change frequently.***
 
 ## 🔥Important News
 
+* **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.  
 * **2026/01/18** 🚀 stable-diffusion.cpp now supports **FLUX.2-klein**  
@@ -37,9 +39,12 @@ API and command-line option may change frequently.***
     - [SD3/SD3.5](./docs/sd3.md)
     - [FLUX.1-dev/FLUX.1-schnell](./docs/flux.md)
     - [FLUX.2-dev/FLUX.2-klein](./docs/flux2.md)
+    - [Lens](./docs/lens.md)
     - [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)
@@ -48,6 +53,7 @@ API and command-line option may change frequently.***
   - Image Edit Models
     - [FLUX.1-Kontext-dev](./docs/kontext.md)
     - [Qwen Image Edit series](./docs/qwen_image_edit.md)
+    - [LongCat Image Edit](./docs/longcat_image.md)
   - Video Models
     - [Wan2.1/Wan2.2](./docs/wan.md)
     - [LTX-2.3](./docs/ltx2.md)
@@ -140,6 +146,8 @@ For runtime and parameter backend placement, see the [backend selection guide](.
 - [Anima](./docs/anima.md)
 - [ERNIE-Image](./docs/ernie_image.md)
 - [HiDream-O1-Image](./docs/hidream_o1_image.md)
+- [Lens](./docs/lens.md)
+- [LongCat Image / LongCat Image Edit](./docs/longcat_image.md)
 - [LoRA](./docs/lora.md)
 - [LCM/LCM-LoRA](./docs/lcm.md)
 - [Using PhotoMaker to personalize image generation](./docs/photo_maker.md)

docs/lens.md

@@ -0,0 +1,32 @@
+# How to Use
+
+Lens uses a Lens diffusion transformer, the FLUX.2 VAE, and GPT-OSS-20B as the LLM text encoder.
+
+## Download weights
+
+- Download Lens
+    - safetensors: https://huggingface.co/Comfy-Org/Lens/tree/main/diffusion_models
+- Download Lens Turbo
+    - safetensors: https://huggingface.co/Comfy-Org/Lens/tree/main/diffusion_models
+- Download vae
+    - safetensors: https://huggingface.co/black-forest-labs/FLUX.2-dev/tree/main
+- Download GPT-OSS-20B
+    - gguf: https://huggingface.co/unsloth/gpt-oss-20b-GGUF/tree/main
+
+## Examples
+
+### Lens
+
+```
+.\bin\Release\sd-cli.exe --diffusion-model ..\..\ComfyUI\models\diffusion_models\lens_bf16.safetensors --llm "..\..\llm\gpt-oss-20b-UD-Q8_K_XL.gguf" --vae ..\..\ComfyUI\models\vae\flux2_ae.safetensors --cfg-scale 5.0  -p "A crystal dragon soaring through an aurora borealis sky, its entire body made of transparent faceted crystal refracting the green and purple aurora light into rainbow spectra, ice particles trailing from its wings, high fantasy digital art" --diffusion-fa -v
+```
+
+<img width="256" alt="Lens example" src="../assets/lens/example.png" />
+
+### Lens Turbo
+
+```
+.\bin\Release\sd-cli.exe --diffusion-model ..\..\ComfyUI\models\diffusion_models\lens_turbo_bf16.safetensors --llm "..\..\llm\gpt-oss-20b-UD-Q8_K_XL.gguf" --vae ..\..\ComfyUI\models\vae\flux2_ae.safetensors --cfg-scale 1.0  -p "A crystal dragon soaring through an aurora borealis sky, its entire body made of transparent faceted crystal refracting the green and purple aurora light into rainbow spectra, ice particles trailing from its wings, high fantasy digital art" --diffusion-fa -v --steps 4
+```
+
+<img width="256" alt="Lens Turbo example" src="../assets/lens/turbo_example.png" />

docs/longcat_image.md

@@ -0,0 +1,30 @@
+# How to Use
+
+LongCat-Image uses a LongCat diffusion transformer, the FLUX VAE, and Qwen2.5-VL as the LLM text encoder.
+
+## Download weights
+
+- Download LongCat Image
+    - safetensors: https://huggingface.co/Comfy-Org/LongCat-Image/tree/main/split_files/diffusion_models
+    - gguf: https://huggingface.co/vantagewithai/LongCat-Image-GGUF/tree/main/comfy
+- Download LongCat Image Edit
+    - LongCat Image Edit Turbo: https://huggingface.co/meituan-longcat/LongCat-Image-Edit-Turbo
+    - gguf: https://huggingface.co/vantagewithai/LongCat-Image-Edit-GGUF/tree/main
+- Download vae
+    - safetensors: https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/ae.safetensors
+- Download qwen_2.5_vl 7b
+    - safetensors: https://huggingface.co/Comfy-Org/Qwen-Image_ComfyUI/tree/main/split_files/text_encoders
+    - gguf: https://huggingface.co/mradermacher/Qwen2.5-VL-7B-Instruct-GGUF/tree/main
+    - For image editing with GGUF text encoders, also download the matching mmproj file and pass it with `--llm_vision`.
+
+## Run
+
+LongCat uses quoted text for character-level text rendering. Put target text inside single quotes, double quotes, or Chinese quotes.
+
+### LongCat Image
+
+```
+.\bin\Release\sd-cli.exe --diffusion-model  ..\..\ComfyUI\models\diffusion_models\LongCat-Image-Q4_K_M.gguf --vae ..\..\ComfyUI\models\vae\ae.sft --llm ..\..\ComfyUI\models\text_encoders\Qwen2.5-VL-7B-Instruct-Q8_0.gguf -p "a lovely cat holding a sign says 'longcat.cpp'" --cfg-scale 5.0 --sampling-method euler --flow-shift 3 -v --offload-to-cpu --diffusion-fa
+```
+
+<img alt="longcat example" src="../assets/longcat/example.png" />

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/CMakeLists.txt

@@ -7,6 +7,9 @@ add_executable(${TARGET}
     image_metadata.cpp
     main.cpp
 )
+if(APPLE)
+    sd_set_macos_rpaths(${TARGET})
+endif()
 target_include_directories(${TARGET} PRIVATE
     "${CMAKE_CURRENT_SOURCE_DIR}/.."
     "${PROJECT_SOURCE_DIR}/src"

examples/cli/README.md

@@ -104,9 +104,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 +129,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 +140,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 +151,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/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())
@@ -348,6 +364,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",
@@ -639,6 +659,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;
 }
 
@@ -679,6 +704,7 @@ std::string SDContextParams::to_string() const {
         << "  high_noise_diffusion_model_path: \"" << high_noise_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"
@@ -772,6 +798,7 @@ 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,
         backend.c_str(),
         params_backend.c_str(),
@@ -833,7 +860,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 +940,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 +972,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",

examples/common/common.h

@@ -94,6 +94,7 @@ struct SDContextParams {
     std::string high_noise_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;

examples/server/CMakeLists.txt

@@ -74,6 +74,9 @@ add_executable(${TARGET}
     routes_sdapi.cpp
     routes_sdcpp.cpp
 )
+if(APPLE)
+    sd_set_macos_rpaths(${TARGET})
+endif()
 
 if(HAVE_FRONTEND_BUILD)
     add_dependencies(${TARGET} ${TARGET}_frontend)

examples/server/README.md

@@ -206,9 +206,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 +231,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 +242,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 +253,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

@@ -254,7 +254,7 @@ void refresh_lora_cache(ServerRuntime& rt) {
 
     fs::path lora_dir = rt.ctx_params->lora_model_dir;
     if (fs::exists(lora_dir) && fs::is_directory(lora_dir)) {
-        for (auto& entry : fs::recursive_directory_iterator(lora_dir)) {
+        for (auto& entry : fs::recursive_directory_iterator(lora_dir, fs::directory_options::skip_permission_denied)) {
             if (!entry.is_regular_file()) {
                 continue;
             }

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;
@@ -212,6 +220,7 @@ 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)
     const char* backend;
     const char* params_backend;

src/anima.hpp

@@ -7,6 +7,7 @@
 #include <vector>
 
 #include "common_block.hpp"
+#include "diffusion_model.hpp"
 #include "flux.hpp"
 #include "rope.hpp"
 
@@ -518,7 +519,7 @@ namespace Anima {
         }
     };
 
-    struct AnimaRunner : public GGMLRunner {
+    struct AnimaRunner : public DiffusionModelRunner {
     public:
         std::vector<float> image_pe_vec;
         std::vector<float> adapter_q_pe_vec;
@@ -529,7 +530,7 @@ namespace Anima {
                     ggml_backend_t params_backend,
                     const String2TensorStorage& tensor_storage_map = {},
                     const std::string prefix                       = "model.diffusion_model")
-            : GGMLRunner(backend, params_backend) {
+            : DiffusionModelRunner(backend, params_backend, prefix) {
             int64_t num_layers    = 0;
             std::string layer_tag = prefix + ".net.blocks.";
             for (const auto& kv : tensor_storage_map) {
@@ -559,7 +560,7 @@ namespace Anima {
             return "anima";
         }
 
-        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
             net.get_param_tensors(tensors, prefix + ".net");
         }
 
@@ -598,7 +599,8 @@ namespace Anima {
                                           {},
                                           empty_ref_latents,
                                           false,
-                                          1.0f);
+                                          1.0f,
+                                          false);
 
             std::vector<float> axis_thetas = {
                 static_cast<float>(theta) * calc_ntk_factor(t_extrapolation_ratio, axes_dim[0]),
@@ -683,6 +685,19 @@ namespace Anima {
             };
             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);
+            const auto* extra = diffusion_extra_as<AnimaDiffusionExtra>(diffusion_params);
+            return compute(n_threads,
+                           *diffusion_params.x,
+                           *diffusion_params.timesteps,
+                           tensor_or_empty(diffusion_params.context),
+                           tensor_or_empty(extra->t5_ids),
+                           tensor_or_empty(extra->t5_weights));
+        }
     };
 }  // namespace Anima
 

src/auto_encoder_kl.hpp

@@ -680,7 +680,7 @@ struct AutoEncoderKL : public VAE {
         } else if (sd_version_is_sd3(version)) {
             scale_factor = 1.5305f;
             shift_factor = 0.0609f;
-        } else if (sd_version_is_flux(version) || sd_version_is_z_image(version)) {
+        } else if (sd_version_is_flux(version) || sd_version_is_z_image(version) || sd_version_is_longcat(version)) {
             scale_factor = 0.3611f;
             shift_factor = 0.1159f;
         } else if (sd_version_uses_flux2_vae(version)) {

src/conditioner.hpp

@@ -102,7 +102,6 @@ struct ConditionerParams {
     int clip_skip                                    = -1;
     int width                                        = -1;
     int height                                       = -1;
-    int adm_in_channels                              = -1;
     bool zero_out_masked                             = false;
     int num_input_imgs                               = 0;        // for photomaker
     const std::vector<sd::Tensor<float>>* ref_images = nullptr;  // for qwen image edit
@@ -114,7 +113,7 @@ 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;
@@ -177,11 +176,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 {
@@ -502,7 +506,6 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
                                              int clip_skip,
                                              int width,
                                              int height,
-                                             int adm_in_channels  = -1,
                                              bool zero_out_masked = false) {
         int64_t t0 = ggml_time_ms();
         sd::Tensor<float> hidden_states;  // [n_token, hidden_size] or [n_token, hidden_size + hidden_size2]
@@ -588,7 +591,8 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
 
         sd::Tensor<float> vec;
         if (sd_version_is_sdxl(version)) {
-            int out_dim = 256;
+            int out_dim         = 256;
+            int adm_in_channels = 2816;
             GGML_ASSERT(!pooled.empty());
             vec = sd::Tensor<float>({adm_in_channels});
             vec.fill_(0.0f);
@@ -647,7 +651,6 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
                                                  conditioner_params.clip_skip,
                                                  conditioner_params.width,
                                                  conditioner_params.height,
-                                                 conditioner_params.adm_in_channels,
                                                  conditioner_params.zero_out_masked);
         return std::make_tuple(cond, clsm);
     }
@@ -674,7 +677,6 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
                                             conditioner_params.clip_skip,
                                             conditioner_params.width,
                                             conditioner_params.height,
-                                            conditioner_params.adm_in_channels,
                                             conditioner_params.zero_out_masked);
     }
 };
@@ -784,16 +786,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 {
@@ -1148,13 +1157,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 {
@@ -1391,10 +1405,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 {
@@ -1581,8 +1598,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 {
@@ -1696,11 +1716,19 @@ struct LLMEmbedder : public Conditioner {
             arch = LLM::LLMArch::MISTRAL_SMALL_3_2;
         } else if (sd_version_is_ernie_image(version)) {
             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_z_image(version) || version == VERSION_OVIS_IMAGE || version == VERSION_FLUX2_KLEIN) {
             arch = LLM::LLMArch::QWEN3;
         }
         if (arch == LLM::LLMArch::MISTRAL_SMALL_3_2 || arch == LLM::LLMArch::MINISTRAL_3_3B) {
             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>();
         }
@@ -1716,8 +1744,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 {
@@ -1747,7 +1778,8 @@ struct LLMEmbedder : public Conditioner {
     std::tuple<std::vector<int>, std::vector<float>, std::vector<float>> tokenize(std::string text,
                                                                                   const std::pair<int, int>& attn_range,
                                                                                   size_t min_length = 0,
-                                                                                  size_t max_length = 100000000) {
+                                                                                  size_t max_length = 100000000,
+                                                                                  bool spell_quotes = false) {
         std::vector<std::pair<std::string, float>> parsed_attention;
         if (attn_range.first >= 0 && attn_range.second > 0) {
             if (attn_range.first > 0) {
@@ -1755,6 +1787,9 @@ struct LLMEmbedder : public Conditioner {
             }
             if (attn_range.second - attn_range.first > 0) {
                 auto new_parsed_attention = parse_prompt_attention(text.substr(attn_range.first, attn_range.second - attn_range.first));
+                if (spell_quotes) {
+                    new_parsed_attention = split_quotation_attention(new_parsed_attention);
+                }
                 parsed_attention.insert(parsed_attention.end(),
                                         new_parsed_attention.begin(),
                                         new_parsed_attention.end());
@@ -1804,8 +1839,10 @@ struct LLMEmbedder : public Conditioner {
                                     int hidden_states_min_length,
                                     const std::vector<std::pair<int, sd::Tensor<float>>>& image_embeds,
                                     const std::set<int>& out_layers,
-                                    int prompt_template_encode_start_idx) {
-        auto tokens_weights_mask = tokenize(prompt, prompt_attn_range, min_length);
+                                    int prompt_template_encode_start_idx,
+                                    bool spell_quotes = false,
+                                    int max_length    = 100000000) {
+        auto tokens_weights_mask = tokenize(prompt, prompt_attn_range, min_length, max_length, spell_quotes);
         auto& tokens             = std::get<0>(tokens_weights_mask);
         auto& weights            = std::get<1>(tokens_weights_mask);
         auto& mask               = std::get<2>(tokens_weights_mask);
@@ -1813,12 +1850,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;
                 }
@@ -1865,7 +1906,9 @@ struct LLMEmbedder : public Conditioner {
         std::vector<std::pair<int, sd::Tensor<float>>> image_embeds;
         int prompt_template_encode_start_idx = 34;
         int min_length                       = 0;  // pad tokens
+        int max_length                       = 100000000;
         int hidden_states_min_length         = 0;  // zero pad hidden_states
+        bool spell_quotes                    = false;
         std::set<int> out_layers;
 
         int64_t t0 = ggml_time_ms();
@@ -1938,6 +1981,71 @@ struct LLMEmbedder : public Conditioner {
 
                 prompt += "<|im_end|>\n<|im_start|>assistant\n";
             }
+        } else if (sd_version_is_longcat(version)) {
+            spell_quotes = true;
+
+            if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
+                LOG_INFO("LongCatEditPipeline");
+                prompt_template_encode_start_idx = 67;
+                min_length                       = 512 + prompt_template_encode_start_idx;
+                int image_embed_idx              = 36 + 6;
+
+                int min_pixels          = 384 * 384;
+                int max_pixels          = 560 * 560;
+                std::string placeholder = "<|image_pad|>";
+                std::string img_prompt;
+
+                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;
+                    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);
+                    int w_bar         = static_cast<int>(std::round(width / factor) * factor);
+
+                    if (static_cast<double>(h_bar) * w_bar > max_pixels) {
+                        double beta = std::sqrt((height * width) / static_cast<double>(max_pixels));
+                        h_bar       = std::max(static_cast<int>(factor),
+                                               static_cast<int>(std::floor(height / beta / factor)) * static_cast<int>(factor));
+                        w_bar       = std::max(static_cast<int>(factor),
+                                               static_cast<int>(std::floor(width / beta / factor)) * static_cast<int>(factor));
+                    } else if (static_cast<double>(h_bar) * w_bar < min_pixels) {
+                        double beta = std::sqrt(static_cast<double>(min_pixels) / (height * width));
+                        h_bar       = static_cast<int>(std::ceil(height * beta / factor)) * static_cast<int>(factor);
+                        w_bar       = static_cast<int>(std::ceil(width * beta / factor)) * static_cast<int>(factor);
+                    }
+
+                    LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
+
+                    auto resized_image = clip_preprocess(image, w_bar, h_bar);
+                    auto image_embed   = llm->encode_image(n_threads, resized_image);
+                    GGML_ASSERT(!image_embed.empty());
+                    image_embeds.emplace_back(image_embed_idx, image_embed);
+                    image_embed_idx += 1 + static_cast<int>(image_embed.shape()[1]) + 6;
+
+                    img_prompt += "<|vision_start|>";
+                    int64_t num_image_tokens = image_embed.shape()[1];
+                    img_prompt.reserve(num_image_tokens * placeholder.size());
+                    for (int j = 0; j < num_image_tokens; j++) {
+                        img_prompt += placeholder;
+                    }
+                    img_prompt += "<|vision_end|>";
+                }
+
+                prompt = "<|im_start|>system\nAs an image editing expert, first analyze the content and attributes of the input image(s). Then, based on the user's editing instructions, clearly and precisely determine how to modify the given image(s), ensuring that only the specified parts are altered and all other aspects remain consistent with the original(s).<|im_end|>\n<|im_start|>user\n";
+                prompt += img_prompt;
+            } else {
+                prompt_template_encode_start_idx = 36;
+                min_length                       = 512 + prompt_template_encode_start_idx;
+
+                prompt = "<|im_start|>system\nAs an image captioning expert, generate a descriptive text prompt based on an image content, suitable for input to a text-to-image model.<|im_end|>\n<|im_start|>user\n";
+            }
+
+            prompt_attn_range.first = static_cast<int>(prompt.size());
+            prompt += conditioner_params.text;
+            prompt_attn_range.second = static_cast<int>(prompt.size());
+
+            prompt += "<|im_end|>\n<|im_start|>assistant\n";
         } else if (version == VERSION_FLUX2) {
             prompt_template_encode_start_idx = 0;
             hidden_states_min_length         = 512;
@@ -1957,6 +2065,30 @@ struct LLMEmbedder : public Conditioner {
             prompt_attn_range.first = 0;
             prompt += conditioner_params.text;
             prompt_attn_range.second = static_cast<int>(prompt.size());
+        } else if (sd_version_is_lens(version)) {
+            prompt_template_encode_start_idx = 97;
+            min_length                       = 0;
+            max_length                       = 512;
+            out_layers                       = {6, 12, 18, 24};
+
+            prompt =
+                "<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI.\n"
+                "Knowledge cutoff: 2024-06\n"
+                "Current date: 2026-05-26\n"  // fix for current date
+                "\n"
+                "Reasoning: medium\n"
+                "\n"
+                "# Valid channels: analysis, commentary, final. Channel must be included for every message.<|end|><|start|>developer<|message|># Instructions\n"
+                "\n"
+                "Describe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background.\n"
+                "\n"
+                "<|end|><|start|>user<|message|>";
+
+            prompt_attn_range.first = static_cast<int>(prompt.size());
+            prompt += conditioner_params.text;
+            prompt_attn_range.second = static_cast<int>(prompt.size());
+
+            prompt += "<|end|><|start|>assistant<|channel|>analysis<|message|>Need to generate one image according to the description.<|end|><|start|>assistant<|channel|>final<|message|>";
         } else if (sd_version_is_z_image(version)) {
             prompt_template_encode_start_idx = 0;
             out_layers                       = {35};  // -2
@@ -2001,6 +2133,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);
         }
@@ -2012,7 +2191,9 @@ struct LLMEmbedder : public Conditioner {
                                            hidden_states_min_length,
                                            image_embeds,
                                            out_layers,
-                                           prompt_template_encode_start_idx);
+                                           prompt_template_encode_start_idx,
+                                           spell_quotes,
+                                           max_length);
         std::vector<sd::Tensor<float>> extra_hidden_states_vec;
         for (int i = 0; i < extra_prompts.size(); i++) {
             auto extra_hidden_states = encode_prompt(n_threads,
@@ -2022,7 +2203,9 @@ struct LLMEmbedder : public Conditioner {
                                                      hidden_states_min_length,
                                                      image_embeds,
                                                      out_layers,
-                                                     prompt_template_encode_start_idx);
+                                                     prompt_template_encode_start_idx,
+                                                     spell_quotes,
+                                                     max_length);
             extra_hidden_states_vec.push_back(std::move(extra_hidden_states));
         }
 
@@ -2137,9 +2320,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 {
@@ -2156,6 +2344,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

@@ -496,84 +496,24 @@ struct LTX2Scheduler : SigmaScheduler {
         parse_extra_sample_args(extra_sample_args);
     }
 
-    static std::string trim(std::string value) {
-        const char* whitespace = " \t\r\n";
-        size_t begin           = value.find_first_not_of(whitespace);
-        if (begin == std::string::npos) {
-            return "";
-        }
-        size_t end = value.find_last_not_of(whitespace);
-        return value.substr(begin, end - begin + 1);
-    }
-
     void parse_extra_sample_args(const char* extra_sample_args) {
-        if (extra_sample_args == nullptr || extra_sample_args[0] == '\0') {
-            return;
-        }
-
-        std::string raw(extra_sample_args);
-        size_t start   = 0;
-        auto parse_arg = [&](const std::string& item) {
-            std::string token = trim(item);
-            if (token.empty()) {
-                return;
-            }
-            size_t eq = token.find('=');
-            if (eq == std::string::npos) {
-                LOG_WARN("ignoring invalid ltx2 scheduler arg '%s'", token.c_str());
-                return;
-            }
-
-            std::string key   = trim(token.substr(0, eq));
-            std::string value = trim(token.substr(eq + 1));
-            auto parse_float  = [&](float* out) -> bool {
-                try {
-                    size_t consumed = 0;
-                    float parsed    = std::stof(value, &consumed);
-                    if (!trim(value.substr(consumed)).empty()) {
-                        return false;
-                    }
-                    *out = parsed;
-                    return true;
-                } catch (const std::exception&) {
-                    return false;
+        for (const auto& [key, value] : parse_key_value_args(extra_sample_args, "ltx2 scheduler arg")) {
+            if (key == "max_shift") {
+                if (!parse_strict_float(value, max_shift)) {
+                    LOG_WARN("ignoring invalid ltx2 scheduler arg '%s=%s'", key.c_str(), value.c_str());
                 }
-            };
-            try {
-                if (key == "max_shift") {
-                    if (!parse_float(&max_shift)) {
-                        LOG_WARN("ignoring invalid ltx2 scheduler arg '%s'", token.c_str());
-                    }
-                } else if (key == "base_shift") {
-                    if (!parse_float(&base_shift)) {
-                        LOG_WARN("ignoring invalid ltx2 scheduler arg '%s'", token.c_str());
-                    }
-                } else if (key == "terminal") {
-                    if (!parse_float(&terminal)) {
-                        LOG_WARN("ignoring invalid ltx2 scheduler arg '%s'", token.c_str());
-                    }
-                } else if (key == "stretch") {
-                    std::string v = value;
-                    std::transform(v.begin(), v.end(), v.begin(), [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
-                    if (v == "1" || v == "true" || v == "yes" || v == "on") {
-                        stretch = true;
-                    } else if (v == "0" || v == "false" || v == "no" || v == "off") {
-                        stretch = false;
-                    } else {
-                        LOG_WARN("ignoring invalid ltx2 scheduler arg '%s'", token.c_str());
-                    }
-                } else {
-                    LOG_WARN("ignoring unknown ltx2 scheduler arg '%s'", key.c_str());
+            } else if (key == "base_shift") {
+                if (!parse_strict_float(value, base_shift)) {
+                    LOG_WARN("ignoring invalid ltx2 scheduler arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            } else if (key == "terminal") {
+                if (!parse_strict_float(value, terminal)) {
+                    LOG_WARN("ignoring invalid ltx2 scheduler arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            } else if (key == "stretch") {
+                if (!parse_strict_bool(value, stretch)) {
+                    LOG_WARN("ignoring invalid ltx2 scheduler arg '%s=%s'", key.c_str(), value.c_str());
                 }
-            } catch (const std::exception&) {
-                LOG_WARN("ignoring invalid ltx2 scheduler arg '%s'", token.c_str());
-            }
-        };
-
-        for (size_t pos = 0; pos <= raw.size(); ++pos) {
-            if (pos == raw.size() || raw[pos] == ',' || raw[pos] == ';') {
-                parse_arg(raw.substr(start, pos - start));
-                start = pos + 1;
             }
         }
     }
@@ -1276,7 +1216,7 @@ static sd::Tensor<float> sample_dpmpp_2m_v2(denoise_cb_t model,
     return x;
 }
 
-using SamplerExtraArgs = std::vector<std::pair<std::string, std::string>>;
+using SamplerExtraArgs = KeyValueArgs;
 
 static sd::Tensor<float> sample_lcm(denoise_cb_t model,
                                     sd::Tensor<float> x,
@@ -1296,27 +1236,26 @@ static sd::Tensor<float> sample_lcm(denoise_cb_t model,
 
     for (const auto& [key, value] : extra_sample_args) {
         float parsed = 0.0f;
-        try {
-            size_t consumed = 0;
-            parsed          = std::stof(value, &consumed);
-            if (trim(value.substr(consumed)).size() != 0) {
-                LOG_WARN("ignoring invalid lcm extra sample arg '%s'", key.c_str());
+        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;
             }
-        } catch (const std::exception&) {
-            LOG_WARN("ignoring invalid lcm extra sample arg '%s=%s'", key.c_str());
-            continue;
-        }
-        if (key == "noise_clip_std") {
             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());
         }
     }
 
@@ -1860,23 +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;
-        try {
-            size_t consumed = 0;
-            parsed          = std::stof(value, &consumed);
-            if (trim(value.substr(consumed)).size() != 0) {
-                LOG_WARN("ignoring invalid euler_ge extra sample arg '%s'", key.c_str());
+        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;
             }
-        } catch (const std::exception&) {
-            LOG_WARN("ignoring invalid euler_ge extra sample arg '%s'", key.c_str());
-            continue;
-        }
-        if (key == "gamma") {
             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());
         }
     }
 
@@ -1916,46 +1846,6 @@ static sd::Tensor<float> sample_gradient_estimation(denoise_cb_t model,
     return x;
 }
 
-static SamplerExtraArgs parse_sampler_args(const char* extra_sample_args) {
-    SamplerExtraArgs pairs;
-
-    if (extra_sample_args == nullptr || extra_sample_args[0] == '\0') {
-        return pairs;
-    }
-
-    auto trim = [](std::string value) -> std::string {
-        const char* whitespace = " \t\r\n";
-        size_t begin           = value.find_first_not_of(whitespace);
-        if (begin == std::string::npos) {
-            return "";
-        }
-        size_t end = value.find_last_not_of(whitespace);
-        return value.substr(begin, end - begin + 1);
-    };
-
-    std::string raw(extra_sample_args);
-    size_t start = 0;
-
-    for (size_t pos = 0; pos <= raw.size(); ++pos) {
-        if (pos == raw.size() || raw[pos] == ',' || raw[pos] == ';') {
-            std::string item  = raw.substr(start, pos - start);
-            std::string token = trim(item);
-
-            if (!token.empty()) {
-                size_t eq = token.find('=');
-                if (eq != std::string::npos) {
-                    std::string key   = trim(token.substr(0, eq));
-                    std::string value = trim(token.substr(eq + 1));
-                    pairs.emplace_back(std::move(key), std::move(value));
-                }
-            }
-            start = pos + 1;
-        }
-    }
-
-    return pairs;
-}
-
 // k diffusion reverse ODE: dx = (x - D(x;\sigma)) / \sigma dt; \sigma(t) = t
 static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
                                             denoise_cb_t model,
@@ -1965,7 +1855,7 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
                                             float eta,
                                             bool is_flow_denoiser,
                                             const char* extra_sample_args) {
-    SamplerExtraArgs extra_args = parse_sampler_args(extra_sample_args);
+    SamplerExtraArgs extra_args = parse_key_value_args(extra_sample_args, "extra sample arg");
     switch (method) {
         case EULER_A_SAMPLE_METHOD:
             return sample_euler_ancestral(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);

src/diffusion_model.hpp

@@ -1,775 +1,107 @@
 #ifndef __DIFFUSION_MODEL_H__
 #define __DIFFUSION_MODEL_H__
 
-#include <optional>
-#include "anima.hpp"
-#include "ernie_image.hpp"
-#include "flux.hpp"
-#include "hidream_o1.hpp"
-#include "ltxv.hpp"
-#include "mmdit.hpp"
-#include "qwen_image.hpp"
-#include "tensor_ggml.hpp"
-#include "unet.hpp"
-#include "wan.hpp"
-#include "z_image.hpp"
+#include <string>
+#include <utility>
+#include <variant>
 
-struct DiffusionParams {
-    const sd::Tensor<float>* x                                         = nullptr;
-    const sd::Tensor<float>* timesteps                                 = nullptr;
-    const sd::Tensor<float>* audio_x                                   = nullptr;
-    const sd::Tensor<float>* audio_timesteps                           = nullptr;
-    const sd::Tensor<float>* context                                   = nullptr;
-    const sd::Tensor<float>* c_concat                                  = nullptr;
-    const sd::Tensor<float>* y                                         = nullptr;
-    const sd::Tensor<int32_t>* t5_ids                                  = nullptr;
-    const sd::Tensor<float>* t5_weights                                = nullptr;
-    const sd::Tensor<float>* guidance                                  = nullptr;
-    const std::vector<sd::Tensor<float>>* ref_latents                  = nullptr;
+#include "ggml_extend.hpp"
+#include "tensor_ggml.hpp"
+
+struct UNetDiffusionExtra {
+    int num_video_frames                           = -1;
+    const std::vector<sd::Tensor<float>>* controls = nullptr;
+    float control_strength                         = 0.f;
+};
+
+struct SkipLayerDiffusionExtra {
+    const std::vector<int>* skip_layers = nullptr;
+};
+
+struct FluxDiffusionExtra {
+    const sd::Tensor<float>* guidance   = nullptr;
+    const std::vector<int>* skip_layers = nullptr;
+};
+
+struct AnimaDiffusionExtra {
+    const sd::Tensor<int32_t>* t5_ids   = nullptr;
+    const sd::Tensor<float>* t5_weights = nullptr;
+};
+
+struct WanDiffusionExtra {
+    const sd::Tensor<float>* vace_context = nullptr;
+    float vace_strength                   = 1.f;
+};
+
+struct HiDreamO1DiffusionExtra {
     const sd::Tensor<int32_t>* input_ids                               = nullptr;
     const sd::Tensor<int32_t>* input_pos                               = nullptr;
     const sd::Tensor<int32_t>* token_types                             = nullptr;
     const sd::Tensor<int32_t>* vinput_mask                             = nullptr;
-    const std::vector<sd::Tensor<float>>* vlm_images                   = nullptr;
     const std::vector<std::pair<int, sd::Tensor<float>>>* image_embeds = nullptr;
-    bool increase_ref_index                                            = false;
-    int num_video_frames                                               = -1;
-    const std::vector<sd::Tensor<float>>* controls                     = nullptr;
-    float control_strength                                             = 0.f;
-    const sd::Tensor<float>* vace_context                              = nullptr;
-    float vace_strength                                                = 1.f;
-    int audio_length                                                   = 0;
-    float frame_rate                                                   = 24.f;
-    const sd::Tensor<float>* video_positions                           = nullptr;
-    const std::vector<int>* skip_layers                                = nullptr;
 };
 
+struct LTXAVDiffusionExtra {
+    const sd::Tensor<float>* audio_x         = nullptr;
+    const sd::Tensor<float>* audio_timesteps = nullptr;
+    int audio_length                         = 0;
+    float frame_rate                         = 24.f;
+    const sd::Tensor<float>* video_positions = nullptr;
+};
+
+using DiffusionExtraParams = std::variant<std::monostate,
+                                          UNetDiffusionExtra,
+                                          SkipLayerDiffusionExtra,
+                                          FluxDiffusionExtra,
+                                          AnimaDiffusionExtra,
+                                          WanDiffusionExtra,
+                                          HiDreamO1DiffusionExtra,
+                                          LTXAVDiffusionExtra>;
+
+struct DiffusionParams {
+    const sd::Tensor<float>* x                        = nullptr;
+    const sd::Tensor<float>* timesteps                = nullptr;
+    const sd::Tensor<float>* context                  = nullptr;
+    const sd::Tensor<float>* c_concat                 = nullptr;
+    const sd::Tensor<float>* y                        = nullptr;
+    const std::vector<sd::Tensor<float>>* ref_latents = nullptr;
+    bool increase_ref_index                           = false;
+    DiffusionExtraParams extra                        = std::monostate{};
+};
+
+template <typename T>
+static inline const T* diffusion_extra_as(const DiffusionParams& params) {
+    const auto* extra = std::get_if<T>(&params.extra);
+    GGML_ASSERT(extra != nullptr);
+    return extra;
+}
+
 template <typename T>
 static inline const sd::Tensor<T>& tensor_or_empty(const sd::Tensor<T>* tensor) {
     static const sd::Tensor<T> kEmpty;
     return tensor != nullptr ? *tensor : kEmpty;
 }
 
-struct DiffusionModel {
-    virtual std::string get_desc()                                               = 0;
+struct DiffusionModelRunner : public GGMLRunner {
+protected:
+    std::string prefix;
+
+public:
+    DiffusionModelRunner(ggml_backend_t backend,
+                         ggml_backend_t params_backend,
+                         const std::string& prefix)
+        : GGMLRunner(backend, params_backend),
+          prefix(prefix) {}
+
     virtual sd::Tensor<float> compute(int n_threads,
-                                      const DiffusionParams& diffusion_params)   = 0;
-    virtual void alloc_params_buffer()                                           = 0;
-    virtual void free_params_buffer()                                            = 0;
-    virtual void free_compute_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_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter){};
-    virtual int64_t get_adm_in_channels()                            = 0;
-    virtual void set_flash_attention_enabled(bool enabled)           = 0;
-    virtual void set_max_graph_vram_bytes(size_t max_vram_bytes)     = 0;
-    virtual void set_circular_axes(bool circular_x, bool circular_y) = 0;
-};
+                                      const DiffusionParams& diffusion_params) = 0;
 
-struct UNetModel : public DiffusionModel {
-    UNetModelRunner unet;
-
-    UNetModel(ggml_backend_t backend,
-              ggml_backend_t params_backend,
-              const String2TensorStorage& tensor_storage_map = {},
-              SDVersion version                              = VERSION_SD1)
-        : unet(backend, params_backend, tensor_storage_map, "model.diffusion_model", version) {
-    }
-
-    std::string get_desc() override {
-        return unet.get_desc();
-    }
-
-    void alloc_params_buffer() override {
-        unet.alloc_params_buffer();
-    }
-
-    void free_params_buffer() override {
-        unet.free_params_buffer();
-    }
-
-    void free_compute_buffer() override {
-        unet.free_compute_buffer();
-    }
-
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
-        unet.get_param_tensors(tensors, "model.diffusion_model");
-    }
-
-    size_t get_params_buffer_size() override {
-        return unet.get_params_buffer_size();
-    }
-
-    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
-        unet.set_weight_adapter(adapter);
-    }
-
-    int64_t get_adm_in_channels() override {
-        return unet.unet.adm_in_channels;
-    }
-
-    void set_flash_attention_enabled(bool enabled) {
-        unet.set_flash_attention_enabled(enabled);
-    }
-
-    void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
-        unet.set_max_graph_vram_bytes(max_vram_bytes);
-    }
-
-    void set_circular_axes(bool circular_x, bool circular_y) override {
-        unet.set_circular_axes(circular_x, circular_y);
-    }
-
-    sd::Tensor<float> compute(int n_threads,
-                              const DiffusionParams& diffusion_params) override {
-        GGML_ASSERT(diffusion_params.x != nullptr);
-        GGML_ASSERT(diffusion_params.timesteps != nullptr);
-        static const std::vector<sd::Tensor<float>> empty_controls;
-        return unet.compute(n_threads,
-                            *diffusion_params.x,
-                            *diffusion_params.timesteps,
-                            tensor_or_empty(diffusion_params.context),
-                            tensor_or_empty(diffusion_params.c_concat),
-                            tensor_or_empty(diffusion_params.y),
-                            diffusion_params.num_video_frames,
-                            diffusion_params.controls ? *diffusion_params.controls : empty_controls,
-                            diffusion_params.control_strength);
-    }
-};
-
-struct MMDiTModel : public DiffusionModel {
-    MMDiTRunner mmdit;
-
-    MMDiTModel(ggml_backend_t backend,
-               ggml_backend_t params_backend,
-               const String2TensorStorage& tensor_storage_map = {})
-        : mmdit(backend, params_backend, tensor_storage_map, "model.diffusion_model") {
-    }
-
-    std::string get_desc() override {
-        return mmdit.get_desc();
-    }
-
-    void alloc_params_buffer() override {
-        mmdit.alloc_params_buffer();
-    }
-
-    void free_params_buffer() override {
-        mmdit.free_params_buffer();
-    }
-
-    void free_compute_buffer() override {
-        mmdit.free_compute_buffer();
-    }
-
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
-        mmdit.get_param_tensors(tensors, "model.diffusion_model");
-    }
-
-    size_t get_params_buffer_size() override {
-        return mmdit.get_params_buffer_size();
-    }
-
-    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
-        mmdit.set_weight_adapter(adapter);
-    }
-
-    int64_t get_adm_in_channels() override {
-        return 768 + 1280;
-    }
-
-    void set_flash_attention_enabled(bool enabled) {
-        mmdit.set_flash_attention_enabled(enabled);
-    }
-
-    void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
-        mmdit.set_max_graph_vram_bytes(max_vram_bytes);
-    }
-
-    void set_circular_axes(bool circular_x, bool circular_y) override {
-        mmdit.set_circular_axes(circular_x, circular_y);
-    }
-
-    sd::Tensor<float> compute(int n_threads,
-                              const DiffusionParams& diffusion_params) override {
-        GGML_ASSERT(diffusion_params.x != nullptr);
-        GGML_ASSERT(diffusion_params.timesteps != nullptr);
-        static const std::vector<int> empty_skip_layers;
-        return mmdit.compute(n_threads,
-                             *diffusion_params.x,
-                             *diffusion_params.timesteps,
-                             tensor_or_empty(diffusion_params.context),
-                             tensor_or_empty(diffusion_params.y),
-                             diffusion_params.skip_layers ? *diffusion_params.skip_layers : empty_skip_layers);
-    }
-};
-
-struct FluxModel : public DiffusionModel {
-    Flux::FluxRunner flux;
-
-    FluxModel(ggml_backend_t backend,
-              ggml_backend_t params_backend,
-              const String2TensorStorage& tensor_storage_map = {},
-              SDVersion version                              = VERSION_FLUX,
-              bool use_mask                                  = false)
-        : flux(backend, params_backend, tensor_storage_map, "model.diffusion_model", version, use_mask) {
-    }
-
-    std::string get_desc() override {
-        return flux.get_desc();
-    }
-
-    void alloc_params_buffer() override {
-        flux.alloc_params_buffer();
-    }
-
-    void free_params_buffer() override {
-        flux.free_params_buffer();
-    }
-
-    void free_compute_buffer() override {
-        flux.free_compute_buffer();
-    }
-
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
-        flux.get_param_tensors(tensors, "model.diffusion_model");
-    }
-
-    size_t get_params_buffer_size() override {
-        return flux.get_params_buffer_size();
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
+        get_param_tensors(tensors, prefix);
     }
 
-    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
-        flux.set_weight_adapter(adapter);
-    }
-
-    int64_t get_adm_in_channels() override {
-        return 768;
-    }
-
-    void set_flash_attention_enabled(bool enabled) {
-        flux.set_flash_attention_enabled(enabled);
-    }
-
-    void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
-        flux.set_max_graph_vram_bytes(max_vram_bytes);
-    }
-
-    void set_circular_axes(bool circular_x, bool circular_y) override {
-        flux.set_circular_axes(circular_x, circular_y);
-    }
-
-    sd::Tensor<float> compute(int n_threads,
-                              const DiffusionParams& diffusion_params) override {
-        GGML_ASSERT(diffusion_params.x != nullptr);
-        GGML_ASSERT(diffusion_params.timesteps != nullptr);
-        static const std::vector<sd::Tensor<float>> empty_ref_latents;
-        static const std::vector<int> empty_skip_layers;
-        return flux.compute(n_threads,
-                            *diffusion_params.x,
-                            *diffusion_params.timesteps,
-                            tensor_or_empty(diffusion_params.context),
-                            tensor_or_empty(diffusion_params.c_concat),
-                            tensor_or_empty(diffusion_params.y),
-                            tensor_or_empty(diffusion_params.guidance),
-                            diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
-                            diffusion_params.increase_ref_index,
-                            diffusion_params.skip_layers ? *diffusion_params.skip_layers : empty_skip_layers);
-    }
-};
-
-struct AnimaModel : public DiffusionModel {
-    std::string prefix;
-    Anima::AnimaRunner anima;
-
-    AnimaModel(ggml_backend_t backend,
-               ggml_backend_t params_backend,
-               const String2TensorStorage& tensor_storage_map = {},
-               const std::string prefix                       = "model.diffusion_model")
-        : prefix(prefix), anima(backend, params_backend, tensor_storage_map, prefix) {
-    }
-
-    std::string get_desc() override {
-        return anima.get_desc();
-    }
-
-    void alloc_params_buffer() override {
-        anima.alloc_params_buffer();
-    }
-
-    void free_params_buffer() override {
-        anima.free_params_buffer();
-    }
-
-    void free_compute_buffer() override {
-        anima.free_compute_buffer();
-    }
-
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
-        anima.get_param_tensors(tensors, prefix);
-    }
-
-    size_t get_params_buffer_size() override {
-        return anima.get_params_buffer_size();
-    }
-
-    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
-        anima.set_weight_adapter(adapter);
-    }
-
-    int64_t get_adm_in_channels() override {
-        return 768;
-    }
-
-    void set_flash_attention_enabled(bool enabled) {
-        anima.set_flash_attention_enabled(enabled);
-    }
-
-    void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
-        anima.set_max_graph_vram_bytes(max_vram_bytes);
-    }
-
-    void set_circular_axes(bool circular_x, bool circular_y) override {
-        anima.set_circular_axes(circular_x, circular_y);
-    }
-
-    sd::Tensor<float> compute(int n_threads,
-                              const DiffusionParams& diffusion_params) override {
-        GGML_ASSERT(diffusion_params.x != nullptr);
-        GGML_ASSERT(diffusion_params.timesteps != nullptr);
-        return anima.compute(n_threads,
-                             *diffusion_params.x,
-                             *diffusion_params.timesteps,
-                             tensor_or_empty(diffusion_params.context),
-                             tensor_or_empty(diffusion_params.t5_ids),
-                             tensor_or_empty(diffusion_params.t5_weights));
-    }
-};
-
-struct WanModel : public DiffusionModel {
-    std::string prefix;
-    WAN::WanRunner wan;
-
-    WanModel(ggml_backend_t backend,
-             ggml_backend_t params_backend,
-             const String2TensorStorage& tensor_storage_map = {},
-             const std::string prefix                       = "model.diffusion_model",
-             SDVersion version                              = VERSION_WAN2)
-        : prefix(prefix), wan(backend, params_backend, tensor_storage_map, prefix, version) {
-    }
-
-    std::string get_desc() override {
-        return wan.get_desc();
-    }
-
-    void alloc_params_buffer() override {
-        wan.alloc_params_buffer();
-    }
-
-    void free_params_buffer() override {
-        wan.free_params_buffer();
-    }
-
-    void free_compute_buffer() override {
-        wan.free_compute_buffer();
-    }
-
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
-        wan.get_param_tensors(tensors, prefix);
-    }
-
-    size_t get_params_buffer_size() override {
-        return wan.get_params_buffer_size();
-    }
-
-    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
-        wan.set_weight_adapter(adapter);
-    }
-
-    int64_t get_adm_in_channels() override {
-        return 768;
-    }
-
-    void set_flash_attention_enabled(bool enabled) {
-        wan.set_flash_attention_enabled(enabled);
-    }
-
-    void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
-        wan.set_max_graph_vram_bytes(max_vram_bytes);
-    }
-
-    void set_circular_axes(bool circular_x, bool circular_y) override {
-        wan.set_circular_axes(circular_x, circular_y);
-    }
-
-    sd::Tensor<float> compute(int n_threads,
-                              const DiffusionParams& diffusion_params) override {
-        GGML_ASSERT(diffusion_params.x != nullptr);
-        GGML_ASSERT(diffusion_params.timesteps != nullptr);
-        return wan.compute(n_threads,
-                           *diffusion_params.x,
-                           *diffusion_params.timesteps,
-                           tensor_or_empty(diffusion_params.context),
-                           tensor_or_empty(diffusion_params.y),
-                           tensor_or_empty(diffusion_params.c_concat),
-                           sd::Tensor<float>(),
-                           tensor_or_empty(diffusion_params.vace_context),
-                           diffusion_params.vace_strength);
-    }
-};
-
-struct QwenImageModel : public DiffusionModel {
-    std::string prefix;
-    Qwen::QwenImageRunner qwen_image;
-
-    QwenImageModel(ggml_backend_t backend,
-                   ggml_backend_t params_backend,
-                   const String2TensorStorage& tensor_storage_map = {},
-                   const std::string prefix                       = "model.diffusion_model",
-                   SDVersion version                              = VERSION_QWEN_IMAGE,
-                   bool zero_cond_t                               = false)
-        : prefix(prefix), qwen_image(backend, params_backend, tensor_storage_map, prefix, version, zero_cond_t) {
-    }
-
-    std::string get_desc() override {
-        return qwen_image.get_desc();
-    }
-
-    void alloc_params_buffer() override {
-        qwen_image.alloc_params_buffer();
-    }
-
-    void free_params_buffer() override {
-        qwen_image.free_params_buffer();
-    }
-
-    void free_compute_buffer() override {
-        qwen_image.free_compute_buffer();
-    }
-
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
-        qwen_image.get_param_tensors(tensors, prefix);
-    }
-
-    size_t get_params_buffer_size() override {
-        return qwen_image.get_params_buffer_size();
-    }
-
-    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
-        qwen_image.set_weight_adapter(adapter);
-    }
-
-    int64_t get_adm_in_channels() override {
-        return 768;
-    }
-
-    void set_flash_attention_enabled(bool enabled) {
-        qwen_image.set_flash_attention_enabled(enabled);
-    }
-
-    void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
-        qwen_image.set_max_graph_vram_bytes(max_vram_bytes);
-    }
-
-    void set_circular_axes(bool circular_x, bool circular_y) override {
-        qwen_image.set_circular_axes(circular_x, circular_y);
-    }
-
-    sd::Tensor<float> compute(int n_threads,
-                              const DiffusionParams& diffusion_params) override {
-        GGML_ASSERT(diffusion_params.x != nullptr);
-        GGML_ASSERT(diffusion_params.timesteps != nullptr);
-        static const std::vector<sd::Tensor<float>> empty_ref_latents;
-        return qwen_image.compute(n_threads,
-                                  *diffusion_params.x,
-                                  *diffusion_params.timesteps,
-                                  tensor_or_empty(diffusion_params.context),
-                                  diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
-                                  true);
-    }
-};
-
-struct HiDreamO1Model : public DiffusionModel {
-    std::string prefix;
-    HiDreamO1::HiDreamO1Runner hidream_o1;
-
-    HiDreamO1Model(ggml_backend_t backend,
-                   ggml_backend_t params_backend,
-                   const String2TensorStorage& tensor_storage_map = {},
-                   const std::string& prefix                      = "model")
-        : prefix(prefix), hidream_o1(backend, params_backend, tensor_storage_map, prefix) {
-    }
-
-    std::string get_desc() override {
-        return hidream_o1.get_desc();
-    }
-
-    void alloc_params_buffer() override {
-        hidream_o1.alloc_params_buffer();
-    }
-
-    void free_params_buffer() override {
-        hidream_o1.free_params_buffer();
-    }
-
-    void free_compute_buffer() override {
-        hidream_o1.free_compute_buffer();
-    }
-
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
-        hidream_o1.get_param_tensors(tensors, prefix);
-    }
-
-    size_t get_params_buffer_size() override {
-        return hidream_o1.get_params_buffer_size();
-    }
-
-    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
-        hidream_o1.set_weight_adapter(adapter);
-    }
-
-    int64_t get_adm_in_channels() override {
-        return 0;
-    }
-
-    void set_flash_attention_enabled(bool enabled) {
-        hidream_o1.set_flash_attention_enabled(enabled);
-    }
-
-    void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
-        hidream_o1.set_max_graph_vram_bytes(max_vram_bytes);
-    }
-
-    void set_circular_axes(bool circular_x, bool circular_y) override {
-        hidream_o1.set_circular_axes(circular_x, circular_y);
-    }
-
-    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.input_ids != nullptr);
-        GGML_ASSERT(diffusion_params.input_pos != nullptr);
-        GGML_ASSERT(diffusion_params.token_types != nullptr);
-        static const std::vector<sd::Tensor<float>> empty_images;
-        static const std::vector<std::pair<int, sd::Tensor<float>>> empty_image_embeds;
-        return hidream_o1.compute(n_threads,
-                                  *diffusion_params.x,
-                                  *diffusion_params.timesteps,
-                                  *diffusion_params.input_ids,
-                                  *diffusion_params.input_pos,
-                                  *diffusion_params.token_types,
-                                  tensor_or_empty(diffusion_params.vinput_mask),
-                                  diffusion_params.image_embeds ? *diffusion_params.image_embeds : empty_image_embeds,
-                                  diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_images);
-    }
-};
-
-struct ZImageModel : public DiffusionModel {
-    std::string prefix;
-    ZImage::ZImageRunner z_image;
-
-    ZImageModel(ggml_backend_t backend,
-                ggml_backend_t params_backend,
-                const String2TensorStorage& tensor_storage_map = {},
-                const std::string prefix                       = "model.diffusion_model",
-                SDVersion version                              = VERSION_Z_IMAGE)
-        : prefix(prefix), z_image(backend, params_backend, tensor_storage_map, prefix, version) {
-    }
-
-    std::string get_desc() override {
-        return z_image.get_desc();
-    }
-
-    void alloc_params_buffer() override {
-        z_image.alloc_params_buffer();
-    }
-
-    void free_params_buffer() override {
-        z_image.free_params_buffer();
-    }
-
-    void free_compute_buffer() override {
-        z_image.free_compute_buffer();
-    }
-
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
-        z_image.get_param_tensors(tensors, prefix);
-    }
-
-    size_t get_params_buffer_size() override {
-        return z_image.get_params_buffer_size();
-    }
-
-    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
-        z_image.set_weight_adapter(adapter);
-    }
-
-    int64_t get_adm_in_channels() override {
-        return 768;
-    }
-
-    void set_flash_attention_enabled(bool enabled) {
-        z_image.set_flash_attention_enabled(enabled);
-    }
-
-    void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
-        z_image.set_max_graph_vram_bytes(max_vram_bytes);
-    }
-
-    void set_circular_axes(bool circular_x, bool circular_y) override {
-        z_image.set_circular_axes(circular_x, circular_y);
-    }
-
-    sd::Tensor<float> compute(int n_threads,
-                              const DiffusionParams& diffusion_params) override {
-        GGML_ASSERT(diffusion_params.x != nullptr);
-        GGML_ASSERT(diffusion_params.timesteps != nullptr);
-        static const std::vector<sd::Tensor<float>> empty_ref_latents;
-        return z_image.compute(n_threads,
-                               *diffusion_params.x,
-                               *diffusion_params.timesteps,
-                               tensor_or_empty(diffusion_params.context),
-                               diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
-                               true);
-    }
-};
-
-struct ErnieImageModel : public DiffusionModel {
-    std::string prefix;
-    ErnieImage::ErnieImageRunner ernie_image;
-
-    ErnieImageModel(ggml_backend_t backend,
-                    ggml_backend_t params_backend,
-                    const String2TensorStorage& tensor_storage_map = {},
-                    const std::string prefix                       = "model.diffusion_model")
-        : prefix(prefix), ernie_image(backend, params_backend, tensor_storage_map, prefix) {
-    }
-
-    std::string get_desc() override {
-        return ernie_image.get_desc();
-    }
-
-    void alloc_params_buffer() override {
-        ernie_image.alloc_params_buffer();
-    }
-
-    void free_params_buffer() override {
-        ernie_image.free_params_buffer();
-    }
-
-    void free_compute_buffer() override {
-        ernie_image.free_compute_buffer();
-    }
-
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
-        ernie_image.get_param_tensors(tensors, prefix);
-    }
-
-    size_t get_params_buffer_size() override {
-        return ernie_image.get_params_buffer_size();
-    }
-
-    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
-        ernie_image.set_weight_adapter(adapter);
-    }
-
-    int64_t get_adm_in_channels() override {
-        return 768;
-    }
-
-    void set_flash_attention_enabled(bool enabled) {
-        ernie_image.set_flash_attention_enabled(enabled);
-    }
-
-    void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
-        ernie_image.set_max_graph_vram_bytes(max_vram_bytes);
-    }
-
-    void set_circular_axes(bool circular_x, bool circular_y) override {
-        ernie_image.set_circular_axes(circular_x, circular_y);
-    }
-
-    sd::Tensor<float> compute(int n_threads,
-                              const DiffusionParams& diffusion_params) override {
-        GGML_ASSERT(diffusion_params.x != nullptr);
-        GGML_ASSERT(diffusion_params.timesteps != nullptr);
-        return ernie_image.compute(n_threads,
-                                   *diffusion_params.x,
-                                   *diffusion_params.timesteps,
-                                   tensor_or_empty(diffusion_params.context));
-    }
-};
-
-struct LTXAVModel : public DiffusionModel {
-    std::string prefix;
-    LTXV::LTXAVRunner ltxav;
-
-    LTXAVModel(ggml_backend_t backend,
-               ggml_backend_t params_backend,
-               const String2TensorStorage& tensor_storage_map = {},
-               const std::string prefix                       = "model.diffusion_model")
-        : prefix(prefix), ltxav(backend, params_backend, tensor_storage_map, prefix) {
-    }
-
-    std::string get_desc() override {
-        return ltxav.get_desc();
-    }
-
-    void alloc_params_buffer() override {
-        ltxav.alloc_params_buffer();
-    }
-
-    void free_params_buffer() override {
-        ltxav.free_params_buffer();
-    }
-
-    void free_compute_buffer() override {
-        ltxav.free_compute_buffer();
-    }
-
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
-        ltxav.get_param_tensors(tensors, prefix);
-    }
-
-    size_t get_params_buffer_size() override {
-        return ltxav.get_params_buffer_size();
-    }
-
-    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
-        ltxav.set_weight_adapter(adapter);
-    }
-
-    int64_t get_adm_in_channels() override {
-        return 0;
-    }
-
-    void set_flash_attention_enabled(bool enabled) override {
-        ltxav.set_flash_attention_enabled(enabled);
-    }
-
-    void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
-        ltxav.set_max_graph_vram_bytes(max_vram_bytes);
-    }
-
-    void set_circular_axes(bool circular_x, bool circular_y) override {
-        ltxav.set_circular_axes(circular_x, circular_y);
-    }
-
-    sd::Tensor<float> compute(int n_threads,
-                              const DiffusionParams& diffusion_params) override {
-        GGML_ASSERT(diffusion_params.x != nullptr);
-        GGML_ASSERT(diffusion_params.timesteps != nullptr);
-        return ltxav.compute(n_threads,
-                             *diffusion_params.x,
-                             *diffusion_params.timesteps,
-                             tensor_or_empty(diffusion_params.context),
-                             tensor_or_empty(diffusion_params.audio_x),
-                             tensor_or_empty(diffusion_params.audio_timesteps),
-                             diffusion_params.audio_length,
-                             diffusion_params.frame_rate,
-                             tensor_or_empty(diffusion_params.video_positions));
-    }
+    virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors,
+                                   const std::string& prefix) = 0;
 };
 
 #endif

src/ernie_image.hpp

@@ -5,6 +5,7 @@
 #include <vector>
 
 #include "common_dit.hpp"
+#include "diffusion_model.hpp"
 #include "flux.hpp"
 #include "qwen_image.hpp"
 #include "rope.hpp"
@@ -325,7 +326,7 @@ namespace ErnieImage {
         }
     };
 
-    struct ErnieImageRunner : public GGMLRunner {
+    struct ErnieImageRunner : public DiffusionModelRunner {
         ErnieImageParams ernie_params;
         ErnieImageModel ernie_image;
         std::vector<float> pe_vec;
@@ -334,7 +335,7 @@ namespace ErnieImage {
                          ggml_backend_t params_backend,
                          const String2TensorStorage& tensor_storage_map = {},
                          const std::string prefix                       = "")
-            : GGMLRunner(backend, params_backend) {
+            : DiffusionModelRunner(backend, params_backend, prefix) {
             ernie_params.num_layers = 0;
             for (const auto& [name, tensor_storage] : tensor_storage_map) {
                 if (!starts_with(name, prefix)) {
@@ -393,7 +394,7 @@ namespace ErnieImage {
             return "ernie_image";
         }
 
-        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
             ernie_image.get_param_tensors(tensors, prefix);
         }
 
@@ -435,6 +436,16 @@ namespace ErnieImage {
             };
             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);
+            return compute(n_threads,
+                           *diffusion_params.x,
+                           *diffusion_params.timesteps,
+                           tensor_or_empty(diffusion_params.context));
+        }
     };
 }  // namespace ErnieImage
 

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

@@ -5,6 +5,7 @@
 #include <vector>
 
 #include "common_dit.hpp"
+#include "diffusion_model.hpp"
 #include "model.h"
 #include "rope.hpp"
 
@@ -446,7 +447,6 @@ namespace Flux {
             if (use_yak_mlp || use_mlp_silu_act) {
                 mlp_mult_factor = 2;
             }
-
             blocks["linear1"]  = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, hidden_size * 3 + mlp_hidden_dim * mlp_mult_factor, mlp_proj_bias));
             blocks["linear2"]  = std::shared_ptr<GGMLBlock>(new Linear(hidden_size + mlp_hidden_dim, hidden_size, mlp_proj_bias));
             blocks["norm"]     = std::shared_ptr<GGMLBlock>(new QKNorm(head_dim));
@@ -1177,7 +1177,7 @@ namespace Flux {
         }
     };
 
-    struct FluxRunner : public GGMLRunner {
+    struct FluxRunner : public DiffusionModelRunner {
     public:
         FluxParams flux_params;
         Flux flux;
@@ -1194,7 +1194,7 @@ namespace Flux {
                    const std::string prefix                       = "",
                    SDVersion version                              = VERSION_FLUX,
                    bool use_mask                                  = false)
-            : GGMLRunner(backend, params_backend), version(version), use_mask(use_mask) {
+            : DiffusionModelRunner(backend, params_backend, prefix), version(version), use_mask(use_mask) {
             flux_params.version             = version;
             flux_params.guidance_embed      = false;
             flux_params.depth               = 0;
@@ -1225,6 +1225,9 @@ namespace Flux {
                 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;
@@ -1306,7 +1309,7 @@ namespace Flux {
             return "flux";
         }
 
-        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
             flux.get_param_tensors(tensors, prefix);
         }
 
@@ -1412,7 +1415,6 @@ namespace Flux {
             } else if (version == VERSION_OVIS_IMAGE) {
                 txt_arange_dims = {1, 2};
             }
-
             pe_vec      = Rope::gen_flux_pe(static_cast<int>(x->ne[1]),
                                             static_cast<int>(x->ne[0]),
                                             flux_params.patch_size,
@@ -1425,7 +1427,8 @@ namespace Flux {
                                             flux_params.theta,
                                             circular_y_enabled,
                                             circular_x_enabled,
-                                            flux_params.axes_dim);
+                                            flux_params.axes_dim,
+                                            sd_version_is_longcat(version));
             int pos_len = static_cast<int>(pe_vec.size() / flux_params.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);
@@ -1488,6 +1491,25 @@ namespace Flux {
             return result;
         }
 
+        sd::Tensor<float> compute(int n_threads,
+                                  const DiffusionParams& diffusion_params) override {
+            GGML_ASSERT(diffusion_params.x != nullptr);
+            GGML_ASSERT(diffusion_params.timesteps != nullptr);
+            const auto* extra = diffusion_extra_as<FluxDiffusionExtra>(diffusion_params);
+            static const std::vector<sd::Tensor<float>> empty_ref_latents;
+            static const std::vector<int> empty_skip_layers;
+            return compute(n_threads,
+                           *diffusion_params.x,
+                           *diffusion_params.timesteps,
+                           tensor_or_empty(diffusion_params.context),
+                           tensor_or_empty(diffusion_params.c_concat),
+                           tensor_or_empty(diffusion_params.y),
+                           tensor_or_empty(extra->guidance),
+                           diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
+                           diffusion_params.increase_ref_index,
+                           extra->skip_layers ? *extra->skip_layers : empty_skip_layers);
+        }
+
         void test() {
             ggml_init_params params;
             params.mem_size   = static_cast<size_t>(1024 * 1024) * 1024;  // 1GB
@@ -1545,7 +1567,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;
@@ -1570,7 +1592,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

@@ -953,11 +953,17 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_group_norm_32(ggml_context* ctx,
     return ggml_group_norm(ctx, a, 32, eps);
 }
 
+__STATIC_INLINE__ bool ggml_ext_is_padded_1d(const ggml_tensor* x) {
+    return x->nb[0] == ggml_type_size(x->type) &&
+           x->nb[2] == x->nb[1] * x->ne[1] &&
+           x->nb[3] == x->nb[2] * x->ne[2];
+}
+
 __STATIC_INLINE__ ggml_tensor* ggml_ext_scale(ggml_context* ctx,
                                               ggml_tensor* x,
                                               float factor,
                                               bool inplace = false) {
-    if (!ggml_is_contiguous(x)) {
+    if (!ggml_ext_is_padded_1d(x)) {
         x = ggml_cont(ctx, x);
     }
     if (inplace) {
@@ -1323,13 +1329,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);
         }
@@ -1337,9 +1339,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);
         }
@@ -1347,26 +1346,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);
         }
 
@@ -1381,10 +1363,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;
@@ -1464,7 +1442,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;
@@ -1921,7 +1899,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;
     }
 
@@ -2008,7 +1986,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);
@@ -2315,13 +2293,13 @@ protected:
                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,
@@ -2458,8 +2436,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();
@@ -2701,7 +2679,7 @@ 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;
     }
@@ -2768,7 +2746,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);
 
@@ -3664,7 +3642,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_lokr_forward(
 
         ggml_tensor* hc  = ggml_transpose(ctx, hc_t);
         ggml_tensor* out = ggml_reshape_2d(ctx, ggml_cont(ctx, hc), up * vp, batch);
-        return ggml_scale(ctx, out, scale);
+        return ggml_ext_scale(ctx, out, scale);
     } else {
         int batch = (int)h->ne[3];
         // 1. Reshape input: [W, H, vq*uq, batch] -> [W, H, vq, uq * batch]
@@ -3747,7 +3725,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_lokr_forward(
         ggml_tensor* hc = ggml_transpose(ctx, hc_t);
         // ungroup
         ggml_tensor* out = ggml_reshape_4d(ctx, ggml_cont(ctx, hc), w_out, h_out, up * vp, batch);
-        return ggml_scale(ctx, out, scale);
+        return ggml_ext_scale(ctx, out, scale);
     }
 }
 

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/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),
@@ -43,6 +98,120 @@ namespace sd::guidance {
         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_cond,
+                                                      float guidance_scale,
+                                                      float image_guidance_scale) {
+        if (pred_img_cond != nullptr) {
+            if (pred_uncond != nullptr && guidance_scale == 1.0f) {
+                return *pred_img_cond - *pred_uncond;
+            }
+            if (pred_uncond != nullptr) {
+                return pred_cond +
+                       (*pred_uncond * (1.0f - image_guidance_scale) +
+                        *pred_img_cond * (image_guidance_scale - guidance_scale)) /
+                           (guidance_scale - 1.0f);
+            }
+            return pred_cond - *pred_img_cond;
+        }
+        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_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);
+            } 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);
+        }
+        if (!has_tensor(input.pred_uncond) && !has_tensor(input.pred_img_cond)) {
+            return output;
+        }
+
+        const sd::Tensor<float>* pred_uncond   = input.pred_uncond;
+        const sd::Tensor<float>* pred_img_cond = input.pred_img_cond;
+
+        sd::Tensor<float> deltas = calculate_guidance_delta(pred_cond,
+                                                            pred_uncond,
+                                                            pred_img_cond,
+                                                            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_cond != nullptr) {
+                output.pred = pred_cond + (image_guidance_scale_ - 1.0f) * deltas;
+            }
+        } else if (pred_img_cond != nullptr) {
+            output.pred = *pred_img_cond + guidance_scale_ * deltas;
+        }
+
+        return output;
+    }
+
     SkipLayerGuidance::SkipLayerGuidance(std::vector<int> layers,
                                          float scale,
                                          float start,
@@ -54,7 +223,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 +238,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,6 +17,17 @@ 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;
@@ -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

@@ -12,6 +12,7 @@
 
 #include "common_dit.hpp"
 #include "conditioner.hpp"
+#include "diffusion_model.hpp"
 #include "llm.hpp"
 #include "util.h"
 
@@ -329,7 +330,7 @@ namespace HiDreamO1 {
         }
     };
 
-    struct HiDreamO1Runner : public GGMLRunner {
+    struct HiDreamO1Runner : public DiffusionModelRunner {
         HiDreamO1Params params;
         HiDreamO1Model model;
 
@@ -339,7 +340,7 @@ namespace HiDreamO1 {
                         ggml_backend_t params_backend,
                         const String2TensorStorage& tensor_storage_map = {},
                         const std::string& prefix                      = "model")
-            : GGMLRunner(backend, params_backend),
+            : DiffusionModelRunner(backend, params_backend, prefix),
               params(make_hidream_o1_params()) {
             model = HiDreamO1Model(params);
             model.init(params_ctx, tensor_storage_map, prefix);
@@ -349,7 +350,7 @@ namespace HiDreamO1 {
             return "hidream_o1";
         }
 
-        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) {
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
             model.get_param_tensors(tensors, prefix);
         }
 
@@ -454,6 +455,28 @@ namespace HiDreamO1 {
             };
             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);
+            const auto* extra = diffusion_extra_as<HiDreamO1DiffusionExtra>(diffusion_params);
+            GGML_ASSERT(extra != nullptr);
+            GGML_ASSERT(extra->input_ids != nullptr);
+            GGML_ASSERT(extra->input_pos != nullptr);
+            GGML_ASSERT(extra->token_types != nullptr);
+            static const std::vector<sd::Tensor<float>> empty_images;
+            static const std::vector<std::pair<int, sd::Tensor<float>>> empty_image_embeds;
+            return compute(n_threads,
+                           *diffusion_params.x,
+                           *diffusion_params.timesteps,
+                           *extra->input_ids,
+                           *extra->input_pos,
+                           *extra->token_types,
+                           tensor_or_empty(extra->vinput_mask),
+                           extra->image_embeds ? *extra->image_embeds : empty_image_embeds,
+                           diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_images);
+        }
     };
 
     struct HiDreamO1Conditioner : public Conditioner {
@@ -469,8 +492,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/lens.hpp

@@ -0,0 +1,419 @@
+#ifndef __SD_LENS_HPP__
+#define __SD_LENS_HPP__
+
+#include <memory>
+#include <vector>
+
+#include "common_block.hpp"
+#include "diffusion_model.hpp"
+#include "flux.hpp"
+#include "qwen_image.hpp"
+#include "rope.hpp"
+
+namespace Lens {
+    constexpr int LENS_GRAPH_SIZE = 40960;
+
+    struct LensTimestepProjEmbeddings : public GGMLBlock {
+        LensTimestepProjEmbeddings(int64_t embedding_dim) {
+            blocks["timestep_embedder"] = std::make_shared<Qwen::TimestepEmbedding>(256, embedding_dim);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* timesteps) {
+            auto timestep_embedder = std::dynamic_pointer_cast<Qwen::TimestepEmbedding>(blocks["timestep_embedder"]);
+            auto timesteps_proj    = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, 256, 10000, 1000.f);
+            return timestep_embedder->forward(ctx, timesteps_proj);
+        }
+    };
+
+    struct LensGateMLP : public GGMLBlock {
+        LensGateMLP(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 gate = ggml_silu(ctx->ggml_ctx, w1->forward(ctx, x));
+            auto up   = w3->forward(ctx, x);
+            x         = ggml_mul(ctx->ggml_ctx, gate, up);
+            return w2->forward(ctx, x);
+        }
+    };
+
+    struct LensJointAttention : public GGMLBlock {
+        int64_t dim_head;
+        int64_t num_heads;
+
+        LensJointAttention(int64_t query_dim,
+                           int64_t dim_head,
+                           int64_t num_heads,
+                           float eps = 1e-5f)
+            : dim_head(dim_head), num_heads(num_heads) {
+            int64_t inner_dim = dim_head * num_heads;
+            blocks["img_qkv"] = std::make_shared<Linear>(query_dim, inner_dim * 3, true);
+            blocks["txt_qkv"] = std::make_shared<Linear>(query_dim, inner_dim * 3, true);
+
+            blocks["norm_q"]       = std::make_shared<RMSNorm>(dim_head, eps);
+            blocks["norm_k"]       = std::make_shared<RMSNorm>(dim_head, eps);
+            blocks["norm_added_q"] = std::make_shared<RMSNorm>(dim_head, eps);
+            blocks["norm_added_k"] = std::make_shared<RMSNorm>(dim_head, eps);
+
+            blocks["to_out.0"]   = std::make_shared<Linear>(inner_dim, query_dim, true);
+            blocks["to_add_out"] = std::make_shared<Linear>(inner_dim, query_dim, true);
+        }
+
+        std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
+                                                      ggml_tensor* img,
+                                                      ggml_tensor* txt,
+                                                      ggml_tensor* pe,
+                                                      ggml_tensor* mask = nullptr) {
+            auto img_qkv    = std::dynamic_pointer_cast<Linear>(blocks["img_qkv"]);
+            auto txt_qkv    = std::dynamic_pointer_cast<Linear>(blocks["txt_qkv"]);
+            auto norm_q     = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_q"]);
+            auto norm_k     = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"]);
+            auto norm_add_q = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_added_q"]);
+            auto norm_add_k = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_added_k"]);
+            auto to_out_0   = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
+            auto to_add_out = std::dynamic_pointer_cast<Linear>(blocks["to_add_out"]);
+            int64_t n_img   = img->ne[1];
+            int64_t n_txt   = txt->ne[1];
+            int64_t N       = img->ne[2];
+            int64_t inner   = dim_head * num_heads;
+
+            auto img_qkv_vec = split_qkv(ctx->ggml_ctx, img_qkv->forward(ctx, img));
+            auto txt_qkv_vec = split_qkv(ctx->ggml_ctx, txt_qkv->forward(ctx, txt));
+
+            auto img_q = ggml_reshape_4d(ctx->ggml_ctx, img_qkv_vec[0], dim_head, num_heads, n_img, N);
+            auto img_k = ggml_reshape_4d(ctx->ggml_ctx, img_qkv_vec[1], dim_head, num_heads, n_img, N);
+            auto img_v = ggml_reshape_4d(ctx->ggml_ctx, img_qkv_vec[2], dim_head, num_heads, n_img, N);
+
+            img_q = norm_q->forward(ctx, img_q);
+            img_k = norm_k->forward(ctx, img_k);
+
+            auto txt_q = ggml_reshape_4d(ctx->ggml_ctx, txt_qkv_vec[0], dim_head, num_heads, n_txt, N);
+            auto txt_k = ggml_reshape_4d(ctx->ggml_ctx, txt_qkv_vec[1], dim_head, num_heads, n_txt, N);
+            auto txt_v = ggml_reshape_4d(ctx->ggml_ctx, txt_qkv_vec[2], dim_head, num_heads, n_txt, N);
+
+            txt_q = norm_add_q->forward(ctx, txt_q);
+            txt_k = norm_add_k->forward(ctx, txt_k);
+
+            auto q = ggml_concat(ctx->ggml_ctx, img_q, txt_q, 2);
+            auto k = ggml_concat(ctx->ggml_ctx, img_k, txt_k, 2);
+            auto v = ggml_concat(ctx->ggml_ctx, img_v, txt_v, 2);
+
+            auto attn = Rope::attention(ctx, q, k, v, pe, mask, (1.0f / 128.f));
+
+            auto img_attn_out = ggml_view_3d(ctx->ggml_ctx,
+                                             attn,
+                                             inner,
+                                             n_img,
+                                             N,
+                                             attn->nb[1],
+                                             attn->nb[2],
+                                             0);
+            auto txt_attn_out = ggml_view_3d(ctx->ggml_ctx,
+                                             attn,
+                                             inner,
+                                             n_txt,
+                                             N,
+                                             attn->nb[1],
+                                             attn->nb[2],
+                                             n_img * attn->nb[1]);
+
+            img_attn_out = to_out_0->forward(ctx, ggml_cont(ctx->ggml_ctx, img_attn_out));
+            txt_attn_out = to_add_out->forward(ctx, ggml_cont(ctx->ggml_ctx, txt_attn_out));
+            return {img_attn_out, txt_attn_out};
+        }
+    };
+
+    struct LensTransformerBlock : public GGMLBlock {
+        LensTransformerBlock(int64_t dim,
+                             int64_t num_attention_heads,
+                             int64_t attention_head_dim,
+                             float eps = 1e-6f) {
+            int64_t mlp_hidden_dim = dim / 3 * 8;
+            blocks["img_mod.1"]    = std::make_shared<Linear>(dim, 6 * dim, true);
+            blocks["txt_mod.1"]    = std::make_shared<Linear>(dim, 6 * dim, true);
+            blocks["img_norm1"]    = std::make_shared<RMSNorm>(dim, eps);
+            blocks["img_norm2"]    = std::make_shared<RMSNorm>(dim, eps);
+            blocks["txt_norm1"]    = std::make_shared<RMSNorm>(dim, eps);
+            blocks["txt_norm2"]    = std::make_shared<RMSNorm>(dim, eps);
+            blocks["img_mlp"]      = std::make_shared<LensGateMLP>(dim, mlp_hidden_dim);
+            blocks["txt_mlp"]      = std::make_shared<LensGateMLP>(dim, mlp_hidden_dim);
+            blocks["attn"]         = std::make_shared<LensJointAttention>(dim, attention_head_dim, num_attention_heads);
+        }
+
+        std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
+                                                      ggml_tensor* img,
+                                                      ggml_tensor* txt,
+                                                      ggml_tensor* t_emb,
+                                                      ggml_tensor* pe) {
+            auto img_mod_1 = std::dynamic_pointer_cast<Linear>(blocks["img_mod.1"]);
+            auto txt_mod_1 = std::dynamic_pointer_cast<Linear>(blocks["txt_mod.1"]);
+            auto img_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["img_norm1"]);
+            auto img_norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["img_norm2"]);
+            auto txt_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["txt_norm1"]);
+            auto txt_norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["txt_norm2"]);
+            auto img_mlp   = std::dynamic_pointer_cast<LensGateMLP>(blocks["img_mlp"]);
+            auto txt_mlp   = std::dynamic_pointer_cast<LensGateMLP>(blocks["txt_mlp"]);
+            auto attn      = std::dynamic_pointer_cast<LensJointAttention>(blocks["attn"]);
+
+            auto temb = ggml_silu(ctx->ggml_ctx, t_emb);
+
+            auto img_mod_params = img_mod_1->forward(ctx, temb);
+            auto img_mod_vec    = ggml_ext_chunk(ctx->ggml_ctx, img_mod_params, 6, 0);
+            auto txt_mod_params = txt_mod_1->forward(ctx, temb);
+            auto txt_mod_vec    = ggml_ext_chunk(ctx->ggml_ctx, txt_mod_params, 6, 0);
+
+            auto img_normed    = img_norm1->forward(ctx, img);
+            auto img_modulated = Flux::modulate(ctx->ggml_ctx, img_normed, img_mod_vec[0], img_mod_vec[1]);
+            auto txt_normed    = txt_norm1->forward(ctx, txt);
+            auto txt_modulated = Flux::modulate(ctx->ggml_ctx, txt_normed, txt_mod_vec[0], txt_mod_vec[1]);
+
+            auto [img_attn_output, txt_attn_output] = attn->forward(ctx, img_modulated, txt_modulated, pe);
+
+            img = ggml_add(ctx->ggml_ctx, img, ggml_mul(ctx->ggml_ctx, img_attn_output, img_mod_vec[2]));
+            txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, txt_attn_output, txt_mod_vec[2]));
+
+            auto img_normed2    = img_norm2->forward(ctx, img);
+            auto img_modulated2 = Flux::modulate(ctx->ggml_ctx, img_normed2, img_mod_vec[3], img_mod_vec[4]);
+            auto txt_normed2    = txt_norm2->forward(ctx, txt);
+            auto txt_modulated2 = Flux::modulate(ctx->ggml_ctx, txt_normed2, txt_mod_vec[3], txt_mod_vec[4]);
+
+            img = ggml_add(ctx->ggml_ctx, img, ggml_mul(ctx->ggml_ctx, img_mlp->forward(ctx, img_modulated2), img_mod_vec[5]));
+            txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, txt_mlp->forward(ctx, txt_modulated2), txt_mod_vec[5]));
+            return {img, txt};
+        }
+    };
+
+    struct LensAdaLayerNormContinuous : public GGMLBlock {
+        int64_t hidden_size;
+        float eps;
+
+        LensAdaLayerNormContinuous(int64_t hidden_size, float eps = 1e-6f)
+            : hidden_size(hidden_size), eps(eps) {
+            blocks["linear"] = std::make_shared<Linear>(hidden_size, hidden_size * 2, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* conditioning) {
+            auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
+            auto mods   = ggml_ext_chunk(ctx->ggml_ctx, linear->forward(ctx, ggml_silu(ctx->ggml_ctx, conditioning)), 2, 0);
+            auto scale  = mods[0];
+            auto shift  = mods[1];
+            x           = ggml_norm(ctx->ggml_ctx, x, eps);
+            return Flux::modulate(ctx->ggml_ctx, x, shift, scale);
+        }
+    };
+
+    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;
+
+        LensModel() = default;
+        LensModel(LensParams params)
+            : params(params) {
+            int64_t inner_dim         = params.num_attention_heads * params.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);
+            }
+            for (int i = 0; i < params.num_layers; ++i) {
+                blocks["transformer_blocks." + std::to_string(i)] = std::make_shared<LensTransformerBlock>(inner_dim,
+                                                                                                           params.num_attention_heads,
+                                                                                                           params.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);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* timestep,
+                             ggml_tensor* context,
+                             ggml_tensor* pe) {
+            GGML_ASSERT(context != nullptr);
+            int64_t W = x->ne[0];
+            int64_t H = x->ne[1];
+            int64_t C = x->ne[2];
+            int64_t N = x->ne[3];
+
+            auto time_text_embed = std::dynamic_pointer_cast<LensTimestepProjEmbeddings>(blocks["time_text_embed"]);
+            auto img_in          = std::dynamic_pointer_cast<Linear>(blocks["img_in"]);
+            auto txt_in          = std::dynamic_pointer_cast<Linear>(blocks["txt_in"]);
+            auto norm_out        = std::dynamic_pointer_cast<LensAdaLayerNormContinuous>(blocks["norm_out"]);
+            auto proj_out        = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
+
+            auto t_emb = time_text_embed->forward(ctx, timestep);
+
+            auto img = ggml_reshape_3d(ctx->ggml_ctx, x, W * H, C, N);
+            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);
+            ggml_tensor* txt                     = nullptr;
+            for (int i = 0; i < params.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);
+            }
+            txt = txt_in->forward(ctx, txt);
+
+            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) {
+                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;
+                txt        = out.second;
+                sd::ggml_graph_cut::mark_graph_cut(img, "lens.transformer_blocks." + std::to_string(i), "img");
+                sd::ggml_graph_cut::mark_graph_cut(txt, "lens.transformer_blocks." + std::to_string(i), "txt");
+            }
+
+            img = norm_out->forward(ctx, img, t_emb);
+            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);
+            return out;
+        }
+    };
+
+    struct LensRunner : public DiffusionModelRunner {
+        LensParams lens_params;
+        LensModel lens;
+        std::vector<float> pe_vec;
+
+        LensRunner(ggml_backend_t backend,
+                   ggml_backend_t params_backend,
+                   const String2TensorStorage& tensor_storage_map = {},
+                   const std::string prefix                       = "")
+            : DiffusionModelRunner(backend, params_backend, prefix) {
+            lens_params.num_layers = 0;
+            for (const auto& [name, tensor_storage] : tensor_storage_map) {
+                if (!starts_with(name, prefix)) {
+                    continue;
+                }
+                if (ends_with(name, "img_in.weight") && tensor_storage.n_dims == 2) {
+                    lens_params.in_channels         = tensor_storage.ne[0];
+                    int64_t inner_dim               = tensor_storage.ne[1];
+                    lens_params.num_attention_heads = inner_dim / lens_params.attention_head_dim;
+                } else if (ends_with(name, "txt_in.weight") && tensor_storage.n_dims == 2) {
+                    lens_params.selected_layer_count = static_cast<int>(tensor_storage.ne[0] / lens_params.joint_attention_dim);
+                } else if (ends_with(name, "proj_out.weight") && tensor_storage.n_dims == 2) {
+                    lens_params.out_channels = tensor_storage.ne[1] / lens_params.patch_size / lens_params.patch_size;
+                } else if (ends_with(name, "transformer_blocks.0.attn.norm_q.weight") && tensor_storage.n_dims == 1) {
+                    lens_params.attention_head_dim = tensor_storage.ne[0];
+                }
+
+                size_t pos = name.find("transformer_blocks.");
+                if (pos != std::string::npos) {
+                    std::string layer_name = name.substr(pos);
+                    auto items             = split_string(layer_name, '.');
+                    if (items.size() > 1) {
+                        int block_index = atoi(items[1].c_str());
+                        if (block_index + 1 > lens_params.num_layers) {
+                            lens_params.num_layers = block_index + 1;
+                        }
+                    }
+                }
+            }
+            if (lens_params.num_layers == 0) {
+                lens_params.num_layers = 48;
+            }
+            lens_params.axes_dim_sum = 0;
+            for (int axis_dim : lens_params.axes_dim) {
+                lens_params.axes_dim_sum += axis_dim;
+            }
+
+            LOG_INFO("lens: layers = %d, in_channels = %" PRId64 ", out_channels = %" PRId64
+                     ", heads = %" PRId64 ", head_dim = %" PRId64,
+                     lens_params.num_layers,
+                     lens_params.in_channels,
+                     lens_params.out_channels,
+                     lens_params.num_attention_heads,
+                     lens_params.attention_head_dim);
+
+            lens = LensModel(lens_params);
+            lens.init(params_ctx, tensor_storage_map, prefix);
+        }
+
+        std::string get_desc() override {
+            return "lens";
+        }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
+            lens.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) {
+            ggml_cgraph* gf        = new_graph_custom(LENS_GRAPH_SIZE);
+            ggml_tensor* x         = make_input(x_tensor);
+            ggml_tensor* timesteps = make_input(timesteps_tensor);
+            GGML_ASSERT(x->ne[3] == 1);
+            GGML_ASSERT(!context_tensor.empty());
+            ggml_tensor* context = make_input(context_tensor);
+
+            pe_vec      = Rope::gen_lens_pe(static_cast<int>(x->ne[1]),
+                                            static_cast<int>(x->ne[0]),
+                                            static_cast<int>(x->ne[3]),
+                                            static_cast<int>(context->ne[1]),
+                                            lens_params.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);
+            set_backend_tensor_data(pe, pe_vec.data());
+
+            auto runner_ctx  = get_context();
+            ggml_tensor* out = lens.forward(&runner_ctx, x, timesteps, context, pe);
+            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) {
+            auto get_graph = [&]() -> ggml_cgraph* {
+                return build_graph(x, timesteps, context);
+            };
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+        }
+
+        sd::Tensor<float> compute(int n_threads,
+                                  const DiffusionParams& diffusion_params) override {
+            GGML_ASSERT(diffusion_params.x != nullptr);
+            GGML_ASSERT(diffusion_params.timesteps != nullptr);
+            return compute(n_threads,
+                           *diffusion_params.x,
+                           *diffusion_params.timesteps,
+                           tensor_or_empty(diffusion_params.context));
+        }
+    };
+}  // namespace Lens
+
+#endif  // __SD_LENS_HPP__

src/llm.hpp

@@ -23,11 +23,12 @@
 #include "rope.hpp"
 #include "tokenizers/bpe_tokenizer.h"
 #include "tokenizers/gemma_tokenizer.h"
+#include "tokenizers/gpt_oss_tokenizer.h"
 #include "tokenizers/mistral_tokenizer.h"
 #include "tokenizers/qwen2_tokenizer.h"
 
 namespace LLM {
-    constexpr int LLM_GRAPH_SIZE = 10240;
+    constexpr int LLM_GRAPH_SIZE = 65536;
 
     enum class LLMArch {
         QWEN2_5_VL,
@@ -36,6 +37,8 @@ namespace LLM {
         MISTRAL_SMALL_3_2,
         MINISTRAL_3_3B,
         GEMMA3_12B,
+        GEMMA2_2B,
+        GPT_OSS_20B,
         ARCH_COUNT,
     };
 
@@ -46,6 +49,8 @@ namespace LLM {
         "mistral_small3.2",
         "ministral3.3b",
         "gemma3_12b",
+        "gemma2_2b",
+        "gpt_oss_20b",
     };
 
     enum class MLPActivation {
@@ -83,6 +88,7 @@ namespace LLM {
         int num_kv_heads                = 4;
         int head_dim                    = 128;
         bool qkv_bias                   = true;
+        bool attention_out_bias         = false;
         bool qk_norm                    = false;
         bool rms_norm_add               = false;
         bool normalize_input            = false;
@@ -93,6 +99,8 @@ namespace LLM {
         std::vector<float> rope_thetas  = {1000000.f};
         std::vector<float> rope_scales  = {1.f};
         std::vector<int> sliding_attention;
+        int64_t num_experts         = 0;
+        int64_t num_experts_per_tok = 0;
         LLMVisionParams vision;
     };
 
@@ -163,6 +171,170 @@ namespace LLM {
         }
     };
 
+    struct GPTOSSMLP : public GGMLBlock {
+    protected:
+        int64_t hidden_size;
+        int64_t intermediate_size;
+        int64_t num_experts;
+        int64_t num_experts_per_tok;
+        bool has_combined_gate_up = false;
+
+        void init_params(ggml_context* ctx,
+                         const String2TensorStorage& tensor_storage_map = {},
+                         std::string prefix                             = "") override {
+            auto supported_type = [](ggml_type wtype, int64_t in_features) {
+                if (in_features % ggml_blck_size(wtype) != 0) {
+                    return GGML_TYPE_F32;
+                }
+                return wtype;
+            };
+
+            params["router.weight"] = ggml_new_tensor_2d(ctx,
+                                                         supported_type(get_type(prefix + "router.weight", tensor_storage_map, GGML_TYPE_F32), hidden_size),
+                                                         hidden_size,
+                                                         num_experts);
+            params["router.bias"]   = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_experts);
+
+            has_combined_gate_up = tensor_storage_map.find(prefix + "experts.gate_up_proj.weight") != tensor_storage_map.end();
+            if (has_combined_gate_up) {
+                ggml_type gate_up_type                = supported_type(get_type(prefix + "experts.gate_up_proj.weight", tensor_storage_map, GGML_TYPE_F32), hidden_size);
+                params["experts.gate_up_proj.weight"] = ggml_new_tensor_3d(ctx,
+                                                                           gate_up_type,
+                                                                           hidden_size,
+                                                                           intermediate_size * 2,
+                                                                           num_experts);
+                params["experts.gate_up_proj.bias"]   = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, intermediate_size * 2, num_experts);
+            } else {
+                ggml_type gate_type                = supported_type(get_type(prefix + "experts.gate_proj.weight", tensor_storage_map, GGML_TYPE_F32), hidden_size);
+                ggml_type up_type                  = supported_type(get_type(prefix + "experts.up_proj.weight", tensor_storage_map, GGML_TYPE_F32), hidden_size);
+                params["experts.gate_proj.weight"] = ggml_new_tensor_3d(ctx, gate_type, hidden_size, intermediate_size, num_experts);
+                params["experts.up_proj.weight"]   = ggml_new_tensor_3d(ctx, up_type, hidden_size, intermediate_size, num_experts);
+                params["experts.gate_proj.bias"]   = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, intermediate_size, num_experts);
+                params["experts.up_proj.bias"]     = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, intermediate_size, num_experts);
+            }
+
+            ggml_type down_type                = supported_type(get_type(prefix + "experts.down_proj.weight", tensor_storage_map, GGML_TYPE_F32), intermediate_size);
+            params["experts.down_proj.weight"] = ggml_new_tensor_3d(ctx, down_type, intermediate_size, hidden_size, num_experts);
+            params["experts.down_proj.bias"]   = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, num_experts);
+        }
+
+        ggml_tensor* expert_linear(GGMLRunnerContext* ctx,
+                                   const std::string& weight_name,
+                                   const std::string& bias_name,
+                                   ggml_tensor* x,
+                                   ggml_tensor* selected_experts) {
+            auto out = ggml_mul_mat_id(ctx->ggml_ctx, params[weight_name], x, selected_experts);
+            auto it  = params.find(bias_name);
+            if (it != params.end()) {
+                out = ggml_add_id(ctx->ggml_ctx, out, it->second, selected_experts);
+            }
+            return out;
+        }
+
+    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) {}
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            // x: [N, n_token, hidden_size]
+            GGML_ASSERT(num_experts > 0 && num_experts_per_tok > 0);
+
+            const int64_t n_token       = x->ne[1];
+            const int64_t N             = x->ne[2];
+            const int64_t n_token_total = n_token * N;
+            ggml_tensor* router_weight  = params["router.weight"];
+            ggml_tensor* router_bias    = params["router.bias"];
+            ggml_tensor* router_logits  = ggml_mul_mat(ctx->ggml_ctx, router_weight, x);
+            router_logits               = ggml_add(ctx->ggml_ctx, router_logits, router_bias);
+            router_logits               = ggml_reshape_2d(ctx->ggml_ctx, router_logits, num_experts, n_token_total);
+
+            ggml_tensor* selected_experts = ggml_argsort_top_k(ctx->ggml_ctx, router_logits, (int)num_experts_per_tok);  // [top_k, tokens]
+            ggml_tensor* probs            = ggml_reshape_3d(ctx->ggml_ctx, router_logits, 1, num_experts, n_token_total);
+            ggml_tensor* weights          = ggml_get_rows(ctx->ggml_ctx, probs, selected_experts);  // [1, top_k, tokens]
+            weights                       = ggml_reshape_2d(ctx->ggml_ctx, weights, num_experts_per_tok, n_token_total);
+            weights                       = ggml_soft_max(ctx->ggml_ctx, weights);
+            weights                       = ggml_reshape_3d(ctx->ggml_ctx, weights, 1, num_experts_per_tok, n_token_total);
+
+            x = ggml_reshape_3d(ctx->ggml_ctx, x, hidden_size, 1, n_token_total);
+
+            ggml_tensor* gate = nullptr;
+            ggml_tensor* up   = nullptr;
+            if (has_combined_gate_up) {
+                auto gate_up = expert_linear(ctx,
+                                             "experts.gate_up_proj.weight",
+                                             "experts.gate_up_proj.bias",
+                                             x,
+                                             selected_experts);  // [2 * intermediate, top_k, tokens]
+                gate_up      = ggml_reshape_4d(ctx->ggml_ctx,
+                                               gate_up,
+                                               2,
+                                               intermediate_size,
+                                               num_experts_per_tok,
+                                               n_token_total);
+                gate         = ggml_view_4d(ctx->ggml_ctx,
+                                            gate_up,
+                                            1,
+                                            intermediate_size,
+                                            num_experts_per_tok,
+                                            n_token_total,
+                                            gate_up->nb[1],
+                                            gate_up->nb[2],
+                                            gate_up->nb[3],
+                                            0);
+                up           = ggml_view_4d(ctx->ggml_ctx,
+                                            gate_up,
+                                            1,
+                                            intermediate_size,
+                                            num_experts_per_tok,
+                                            n_token_total,
+                                            gate_up->nb[1],
+                                            gate_up->nb[2],
+                                            gate_up->nb[3],
+                                            gate_up->nb[0]);
+                gate         = ggml_reshape_3d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, gate), intermediate_size, num_experts_per_tok, n_token_total);
+                up           = ggml_reshape_3d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, up), intermediate_size, num_experts_per_tok, n_token_total);
+            } else {
+                gate = expert_linear(ctx,
+                                     "experts.gate_proj.weight",
+                                     "experts.gate_proj.bias",
+                                     x,
+                                     selected_experts);
+                up   = expert_linear(ctx,
+                                     "experts.up_proj.weight",
+                                     "experts.up_proj.bias",
+                                     x,
+                                     selected_experts);
+            }
+
+            auto activated = ggml_swiglu_oai(ctx->ggml_ctx, gate, up, 1.702f, 7.0f);
+            auto experts   = expert_linear(ctx,
+                                           "experts.down_proj.weight",
+                                           "experts.down_proj.bias",
+                                           activated,
+                                           selected_experts);
+            experts        = ggml_mul(ctx->ggml_ctx, experts, weights);
+
+            ggml_tensor* out = nullptr;
+            for (int64_t i = 0; i < num_experts_per_tok; ++i) {
+                auto expert_out = ggml_view_2d(ctx->ggml_ctx,
+                                               experts,
+                                               hidden_size,
+                                               n_token_total,
+                                               experts->nb[2],
+                                               i * experts->nb[1]);
+                out             = out == nullptr ? expert_out : ggml_add(ctx->ggml_ctx, out, expert_out);
+            }
+            if (num_experts_per_tok == 1) {
+                out = ggml_cont(ctx->ggml_ctx, out);
+            }
+
+            return ggml_reshape_3d(ctx->ggml_ctx, out, hidden_size, n_token, N);
+        }
+    };
+
     static ggml_tensor* splice_image_embeds(GGMLRunnerContext* ctx,
                                             ggml_tensor* x,
                                             const std::vector<std::pair<int, ggml_tensor*>>& image_embeds) {
@@ -601,6 +773,15 @@ namespace LLM {
         int64_t max_position_embeddings;
         std::vector<float> rope_thetas;
         std::vector<float> rope_scales;
+        bool has_attention_sinks;
+
+        void init_params(ggml_context* ctx,
+                         const String2TensorStorage& tensor_storage_map = {},
+                         std::string prefix                             = "") override {
+            if (has_attention_sinks) {
+                params["sinks"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_heads);
+            }
+        }
 
     public:
         Attention(const LLMParams& params)
@@ -611,11 +792,12 @@ namespace LLM {
               qk_norm(params.qk_norm),
               max_position_embeddings(params.max_position_embeddings),
               rope_thetas(params.rope_thetas),
-              rope_scales(params.rope_scales) {
+              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, false);
+            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);
@@ -660,6 +842,36 @@ namespace LLM {
             } else if (arch == LLMArch::QWEN3) {
                 q = ggml_rope_ext(ctx->ggml_ctx, q, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 40960, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
                 k = ggml_rope_ext(ctx->ggml_ctx, k, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 40960, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
+            } else if (arch == LLMArch::GPT_OSS_20B) {
+                float rope_theta = rope_thetas.empty() ? 150000.f : rope_thetas[0];
+                float rope_scale = rope_scales.empty() ? 32.f : rope_scales[0];
+                float freq_scale = 1.f / rope_scale;
+                q                = ggml_rope_ext(ctx->ggml_ctx,
+                                                 q,
+                                                 input_pos,
+                                                 nullptr,
+                                                 head_dim,
+                                                 GGML_ROPE_TYPE_NEOX,
+                                                 4096,
+                                                 rope_theta,
+                                                 freq_scale,
+                                                 1.f,
+                                                 1.f,
+                                                 32.f,
+                                                 1.f);
+                k                = ggml_rope_ext(ctx->ggml_ctx,
+                                                 k,
+                                                 input_pos,
+                                                 nullptr,
+                                                 head_dim,
+                                                 GGML_ROPE_TYPE_NEOX,
+                                                 4096,
+                                                 rope_theta,
+                                                 freq_scale,
+                                                 1.f,
+                                                 1.f,
+                                                 32.f,
+                                                 1.f);
             } else if (arch == LLMArch::GEMMA3_12B) {
                 float rope_theta = (rope_index == 1 ? 10000.0f : 1000000.0f);
                 float rope_scale = (rope_index == 1 ? 1.f : 8.f);
@@ -669,8 +881,8 @@ namespace LLM {
                                                  input_pos,
                                                  nullptr,
                                                  head_dim,
-                                                 GGML_ROPE_TYPE_NORMAL,
-                                                 0,
+                                                 GGML_ROPE_TYPE_NEOX,
+                                                 131072,
                                                  rope_theta,
                                                  freq_scale,
                                                  0.f,
@@ -682,14 +894,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);
@@ -706,7 +945,28 @@ namespace LLM {
             k = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, k, 0, 2, 1, 3));  // [N, num_kv_heads, n_token, head_dim]
             k = ggml_reshape_3d(ctx->ggml_ctx, k, k->ne[0], k->ne[1], k->ne[2] * k->ne[3]);      // [N*num_kv_heads, n_token, head_dim]
 
-            x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, attention_mask, true, false);  // [N, n_token, hidden_size]
+            if (arch == LLMArch::GPT_OSS_20B) {
+                GGML_ASSERT(N == 1);
+                auto v_attn = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, v, 1, 2, 0, 3));  // [N, kv_heads, head_dim, tokens]
+                v_attn      = ggml_reshape_3d(ctx->ggml_ctx, v_attn, n_token, head_dim, num_kv_heads * N);
+
+                auto kq = ggml_mul_mat(ctx->ggml_ctx, k, q);
+                ggml_mul_mat_set_prec(kq, GGML_PREC_F32);
+                kq = ggml_scale_inplace(ctx->ggml_ctx, kq, 1.0f / std::sqrt(static_cast<float>(head_dim)));
+                if (attention_mask != nullptr) {
+                    kq = ggml_add_inplace(ctx->ggml_ctx, kq, attention_mask);
+                }
+                kq = ggml_soft_max_inplace(ctx->ggml_ctx, kq);
+                ggml_soft_max_add_sinks(kq, params["sinks"]);
+
+                auto kqv = ggml_mul_mat(ctx->ggml_ctx, v_attn, kq);
+                kqv      = ggml_reshape_4d(ctx->ggml_ctx, kqv, head_dim, n_token, num_heads, N);
+                kqv      = ggml_permute(ctx->ggml_ctx, kqv, 0, 2, 1, 3);
+                x        = ggml_ext_cont(ctx->ggml_ctx, kqv);
+                x        = ggml_reshape_3d(ctx->ggml_ctx, x, head_dim * num_heads, n_token, N);
+            } else {
+                x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, attention_mask, true, false);  // [N, n_token, hidden_size]
+            }
 
             x = out_proj->forward(ctx, x);  // [N, n_token, hidden_size]
             return x;
@@ -717,27 +977,44 @@ namespace LLM {
     protected:
         LLMArch arch;
         int sliding_attention;
-        bool has_post_attention_norm;
-        bool has_post_ffw_norm;
+        std::string post_attention_norm_name;
+        std::string pre_ffw_norm_name;
+        std::string post_ffw_norm_name;
 
     public:
         TransformerBlock(const LLMParams& params, int layer_index)
             : arch(params.arch),
-              sliding_attention(0),
-              has_post_attention_norm(params.arch == LLMArch::GEMMA3_12B),
-              has_post_ffw_norm(params.arch == LLMArch::GEMMA3_12B) {
-            blocks["self_attn"]                = std::make_shared<Attention>(params);
-            blocks["mlp"]                      = std::make_shared<MLP>(params.hidden_size,
-                                                  params.intermediate_size,
-                                                  false,
-                                                  params.mlp_activation);
-            blocks["input_layernorm"]          = std::make_shared<LLMRMSNorm>(params.hidden_size, params.rms_norm_eps, params.rms_norm_add);
-            blocks["post_attention_layernorm"] = std::make_shared<LLMRMSNorm>(params.hidden_size, params.rms_norm_eps, params.rms_norm_add);
-            if (has_post_attention_norm) {
-                blocks["post_attention_norm"] = std::make_shared<LLMRMSNorm>(params.hidden_size, params.rms_norm_eps, params.rms_norm_add);
+              sliding_attention(0) {
+            if (params.arch == LLMArch::GEMMA3_12B) {
+                post_attention_norm_name = "post_attention_norm";       // attn_post_norm
+                pre_ffw_norm_name        = "post_attention_layernorm";  // ffn_norm
+                post_ffw_norm_name       = "post_ffw_norm";             // ffn_post_norm
+            } else if (params.arch == LLMArch::GEMMA2_2B) {
+                post_attention_norm_name = "post_attention_layernorm";  // ffn_norm
+                pre_ffw_norm_name        = "pre_feedforward_layernorm";
+                post_ffw_norm_name       = "post_feedforward_layernorm";
+            } else if (params.arch == LLMArch::GPT_OSS_20B) {
+                pre_ffw_norm_name = "post_attention_norm";  // attn_post_norm
+            } else {
+                pre_ffw_norm_name = "post_attention_layernorm";  // ffn_norm
             }
-            if (has_post_ffw_norm) {
-                blocks["post_ffw_norm"] = std::make_shared<LLMRMSNorm>(params.hidden_size, params.rms_norm_eps, params.rms_norm_add);
+
+            blocks["self_attn"] = std::make_shared<Attention>(params);
+            if (params.arch == LLMArch::GPT_OSS_20B) {
+                blocks["mlp"] = std::make_shared<GPTOSSMLP>(params);
+            } else {
+                blocks["mlp"] = std::make_shared<MLP>(params.hidden_size,
+                                                      params.intermediate_size,
+                                                      false,
+                                                      params.mlp_activation);
+            }
+            blocks["input_layernorm"] = std::make_shared<LLMRMSNorm>(params.hidden_size, params.rms_norm_eps, params.rms_norm_add);
+            blocks[pre_ffw_norm_name] = std::make_shared<LLMRMSNorm>(params.hidden_size, params.rms_norm_eps, params.rms_norm_add);
+            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);
+            }
+            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);
             }
             if (!params.sliding_attention.empty()) {
                 sliding_attention = params.sliding_attention[layer_index % params.sliding_attention.size()];
@@ -751,20 +1028,19 @@ namespace LLM {
                              ggml_tensor* sliding_attention_mask = nullptr) {
             // x: [N, n_token, hidden_size]
             auto self_attn                                  = std::dynamic_pointer_cast<Attention>(blocks["self_attn"]);
-            auto mlp                                        = std::dynamic_pointer_cast<MLP>(blocks["mlp"]);
             auto input_layernorm                            = std::dynamic_pointer_cast<LLMRMSNorm>(blocks["input_layernorm"]);
-            auto post_attention_layernorm                   = std::dynamic_pointer_cast<LLMRMSNorm>(blocks["post_attention_layernorm"]);
+            auto pre_ffw_norm                               = std::dynamic_pointer_cast<LLMRMSNorm>(blocks[pre_ffw_norm_name]);
             std::shared_ptr<LLMRMSNorm> post_attention_norm = nullptr;
             std::shared_ptr<LLMRMSNorm> post_ffw_norm       = nullptr;
-            if (has_post_attention_norm) {
-                post_attention_norm = std::dynamic_pointer_cast<LLMRMSNorm>(blocks["post_attention_norm"]);
+            if (!post_attention_norm_name.empty()) {
+                post_attention_norm = std::dynamic_pointer_cast<LLMRMSNorm>(blocks[post_attention_norm_name]);
             }
-            if (has_post_ffw_norm) {
-                post_ffw_norm = std::dynamic_pointer_cast<LLMRMSNorm>(blocks["post_ffw_norm"]);
+            if (!post_ffw_norm_name.empty()) {
+                post_ffw_norm = std::dynamic_pointer_cast<LLMRMSNorm>(blocks[post_ffw_norm_name]);
             }
             ggml_tensor* block_attention_mask = attention_mask;
             int rope_index                    = 0;
-            if (arch == LLMArch::GEMMA3_12B && sliding_attention > 0) {
+            if ((arch == LLMArch::GEMMA3_12B || arch == LLMArch::GPT_OSS_20B) && sliding_attention > 0) {
                 block_attention_mask = sliding_attention_mask;
                 rope_index           = 1;
             }
@@ -778,8 +1054,14 @@ namespace LLM {
             x = ggml_add_inplace(ctx->ggml_ctx, x, residual);
 
             residual = x;
-            x        = post_attention_layernorm->forward(ctx, x);
-            x        = mlp->forward(ctx, x);
+            x        = pre_ffw_norm->forward(ctx, x);
+            if (arch == LLMArch::GPT_OSS_20B) {
+                auto mlp = std::dynamic_pointer_cast<GPTOSSMLP>(blocks["mlp"]);
+                x        = mlp->forward(ctx, x);
+            } else {
+                auto mlp = std::dynamic_pointer_cast<MLP>(blocks["mlp"]);
+                x        = mlp->forward(ctx, x);
+            }
             if (post_ffw_norm != nullptr) {
                 x = post_ffw_norm->forward(ctx, x);
             }
@@ -1202,6 +1484,39 @@ namespace LLM {
                 params.rope_thetas             = {1000000.f, 10000.f};
                 params.rope_scales             = {8.f, 1.f};
                 params.sliding_attention       = {1024, 1024, 1024, 1024, 1024, 0};
+            } else if (arch == LLMArch::GEMMA2_2B) {
+                params.head_dim                = 256;
+                params.num_heads               = 8;
+                params.num_kv_heads            = 4;
+                params.qkv_bias                = false;
+                params.qk_norm                 = false;
+                params.rms_norm_eps            = 1e-6f;
+                params.rms_norm_add            = true;
+                params.normalize_input         = true;
+                params.max_position_embeddings = 8192;
+                params.mlp_activation          = MLPActivation::GELU_TANH;
+                params.hidden_size             = 2304;
+                params.intermediate_size       = 9216;
+                params.num_layers              = 26;
+                params.vocab_size              = 256000;
+            } else if (arch == LLMArch::GPT_OSS_20B) {
+                params.head_dim                = 64;
+                params.num_heads               = 64;
+                params.num_kv_heads            = 8;
+                params.qkv_bias                = true;
+                params.attention_out_bias      = true;
+                params.qk_norm                 = false;
+                params.rms_norm_eps            = 1e-5f;
+                params.hidden_size             = 2880;
+                params.intermediate_size       = 2880;
+                params.num_layers              = 24;
+                params.vocab_size              = 201088;
+                params.max_position_embeddings = 131072;
+                params.rope_thetas             = {150000.f};
+                params.rope_scales             = {32.f};
+                params.sliding_attention       = {128, 0};
+                params.num_experts             = 32;
+                params.num_experts_per_tok     = 4;
             }
             bool have_vision_weight = false;
             bool llama_cpp_style    = false;
@@ -1236,6 +1551,12 @@ namespace LLM {
                 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;
@@ -1315,7 +1636,9 @@ namespace LLM {
             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::GEMMA3_12B ||
+                params.arch == LLMArch::GEMMA2_2B ||
+                params.arch == LLMArch::GPT_OSS_20B) {
                 input_pos_vec.resize(n_tokens);
                 for (int i = 0; i < n_tokens; ++i) {
                     input_pos_vec[i] = i;
@@ -1354,7 +1677,11 @@ namespace LLM {
                 set_backend_tensor_data(attention_mask, attention_mask_vec.data());
             }
 
-            if (params.arch == LLMArch::GEMMA3_12B) {
+            if (params.arch == LLMArch::GEMMA3_12B || params.arch == LLMArch::GPT_OSS_20B) {
+                int sliding_window = 0;
+                for (int window : params.sliding_attention) {
+                    sliding_window = std::max(sliding_window, window);
+                }
                 sliding_attention_mask_vec.resize(n_tokens * n_tokens);
                 if (!attention_mask_tensor.empty()) {
                     GGML_ASSERT(attention_mask_tensor.numel() == n_tokens * n_tokens);
@@ -1364,8 +1691,7 @@ namespace LLM {
                 }
                 for (int i0 = 0; i0 < n_tokens; i0++) {
                     for (int i1 = 0; i1 < n_tokens; i1++) {
-                        if (i0 + 1024 <= i1) {
-                            LOG_DEBUG("xxxxxxxxxxxxxx");
+                        if (sliding_window > 0 && i0 + sliding_window <= i1) {
                             sliding_attention_mask_vec[i1 * n_tokens + i0] = -INFINITY;
                         }
                     }
@@ -1485,6 +1811,8 @@ namespace LLM {
             : model(arch, backend, params_backend, tensor_storage_map, prefix, enable_vision) {
             if (arch == LLMArch::MISTRAL_SMALL_3_2 || arch == LLMArch::MINISTRAL_3_3B) {
                 tokenizer = std::make_shared<MistralTokenizer>();
+            } else if (arch == LLMArch::GPT_OSS_20B) {
+                tokenizer = std::make_shared<GPTOSSTokenizer>();
             } else {
                 tokenizer = std::make_shared<Qwen2Tokenizer>();
             }
@@ -1494,8 +1822,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,
@@ -1710,7 +2041,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;
@@ -1737,7 +2068,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

@@ -1052,7 +1052,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 +1068,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_latent_upscaler.hpp

@@ -6,8 +6,10 @@
 #include <cstdlib>
 #include <map>
 #include <memory>
+#include <set>
 #include <string>
 #include <utility>
+#include <vector>
 
 #include "common_dit.hpp"
 #include "ggml_extend.hpp"
@@ -26,6 +28,10 @@ namespace LTXVUpsampler {
         bool spatial_upsample    = true;
         bool temporal_upsample   = false;
         bool rational_resampler  = false;
+        float spatial_scale      = 2.f;
+        int spatial_up_num       = 2;
+        int spatial_down_den     = 1;
+        int temporal_up_factor   = 1;
     };
 
     static inline bool has_tensor(const String2TensorStorage& tensor_storage_map,
@@ -33,14 +39,21 @@ namespace LTXVUpsampler {
         return tensor_storage_map.find(name) != tensor_storage_map.end();
     }
 
+    static inline int64_t get_tensor_ne(const String2TensorStorage& tensor_storage_map,
+                                        const std::string& name,
+                                        int axis,
+                                        int64_t fallback) {
+        auto it = tensor_storage_map.find(name);
+        if (it == tensor_storage_map.end() || axis < 0 || axis >= GGML_MAX_DIMS) {
+            return fallback;
+        }
+        return it->second.ne[axis];
+    }
+
     static inline int64_t get_tensor_ne0(const String2TensorStorage& tensor_storage_map,
                                          const std::string& name,
                                          int64_t fallback) {
-        auto it = tensor_storage_map.find(name);
-        if (it == tensor_storage_map.end()) {
-            return fallback;
-        }
-        return it->second.ne[0];
+        return get_tensor_ne(tensor_storage_map, name, 0, fallback);
     }
 
     static inline int count_module_blocks(const String2TensorStorage& tensor_storage_map,
@@ -71,8 +84,36 @@ namespace LTXVUpsampler {
         if (detected_blocks > 0) {
             config.num_blocks_per_stage = detected_blocks;
         }
-        config.spatial_upsample  = has_tensor(tensor_storage_map, "upsampler.0.weight");
-        config.temporal_upsample = has_tensor(tensor_storage_map, "temporal_upsampler.0.weight");
+        config.rational_resampler      = has_tensor(tensor_storage_map, "upsampler.conv.weight");
+        int64_t upsampler_out_channels = get_tensor_ne0(tensor_storage_map, "upsampler.0.bias", 0);
+        config.spatial_upsample        = config.rational_resampler || upsampler_out_channels == 4 * config.mid_channels;
+        config.temporal_upsample       = upsampler_out_channels == 2 * config.mid_channels;
+        if (config.temporal_upsample) {
+            config.temporal_up_factor = 2;
+        }
+        if (config.rational_resampler) {
+            int64_t out_channels = get_tensor_ne(tensor_storage_map,
+                                                 "upsampler.conv.weight",
+                                                 3,
+                                                 config.mid_channels * 9);
+            if (config.mid_channels > 0 && out_channels % config.mid_channels == 0) {
+                int64_t ratio = out_channels / config.mid_channels;
+                int num       = static_cast<int>(std::round(std::sqrt(static_cast<double>(ratio))));
+                if (num > 0 && static_cast<int64_t>(num) * num == ratio) {
+                    config.spatial_up_num = num;
+                }
+            }
+            if (config.spatial_up_num == 3) {
+                config.spatial_down_den = 2;
+                config.spatial_scale    = 1.5f;
+            } else if (config.spatial_up_num == 4) {
+                config.spatial_down_den = 1;
+                config.spatial_scale    = 4.f;
+            } else {
+                config.spatial_down_den = 1;
+                config.spatial_scale    = static_cast<float>(config.spatial_up_num);
+            }
+        }
         return config;
     }
 
@@ -160,16 +201,135 @@ namespace LTXVUpsampler {
             : upscale_factor(upscale_factor) {}
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
-            GGML_ASSERT(upscale_factor == 2);
+            GGML_ASSERT(upscale_factor > 0);
             int64_t h = x->ne[1];
             int64_t w = x->ne[0];
-            // x: [b*f, c*4, h, w] -> [b*f, c, h*2, w*2]
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));  // [b*f, h, w, c*4]
-            x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0], x->ne[1] * x->ne[2], x->ne[3]);          // [b*f, h*w, c*4]
+            GGML_ASSERT(x->ne[2] % (upscale_factor * upscale_factor) == 0);
+            // x: [b*f, c*p1*p2, h, w] -> [b*f, c, h*p1, w*p2]
+            x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));  // [b*f, h, w, c*p1*p2]
+            x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0], x->ne[1] * x->ne[2], x->ne[3]);          // [b*f, h*w, c*p1*p2]
             return DiT::unpatchify(ctx->ggml_ctx, x, h, w, upscale_factor, upscale_factor, true);
         }
     };
 
+    class TemporalPixelShuffleND : public UnaryBlock {
+    protected:
+        int upscale_factor;
+
+    public:
+        explicit TemporalPixelShuffleND(int upscale_factor)
+            : upscale_factor(upscale_factor) {}
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+            GGML_ASSERT(upscale_factor > 0);
+            GGML_ASSERT(x->ne[3] % upscale_factor == 0);
+            const int64_t W = x->ne[0];
+            const int64_t H = x->ne[1];
+            const int64_t F = x->ne[2];
+            const int64_t C = x->ne[3] / upscale_factor;
+
+            // x: [b, c*p, f, h, w] -> [b, c, f*p, h, w]
+            x = ggml_ext_cont(ctx->ggml_ctx, x);
+            x = ggml_reshape_4d(ctx->ggml_ctx, x, W * H, F, upscale_factor, C);
+            x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));
+            return ggml_reshape_4d(ctx->ggml_ctx, x, W, H, F * upscale_factor, C);
+        }
+    };
+
+    class BlurDownsample : public GGMLBlock {
+    protected:
+        int64_t channels;
+        int stride;
+        ggml_tensor* kernel = nullptr;
+        std::vector<float> kernel_data;
+
+        void init_params(ggml_context* ctx,
+                         const String2TensorStorage& tensor_storage_map = {},
+                         const std::string prefix                       = "") override {
+            SD_UNUSED(tensor_storage_map);
+            if (stride == 1) {
+                return;
+            }
+            kernel           = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 5, 5, 1, channels);
+            std::string name = prefix + "kernel";
+            ggml_set_name(kernel, name.c_str());
+
+            static const float binomial[5] = {1.f, 4.f, 6.f, 4.f, 1.f};
+            kernel_data.resize(static_cast<size_t>(5 * 5 * channels));
+            for (int64_t c = 0; c < channels; ++c) {
+                for (int y = 0; y < 5; ++y) {
+                    for (int x = 0; x < 5; ++x) {
+                        kernel_data[static_cast<size_t>(x + 5 * (y + 5 * c))] =
+                            binomial[y] * binomial[x] / 256.f;
+                    }
+                }
+            }
+        }
+
+    public:
+        BlurDownsample(int64_t channels, int stride)
+            : channels(channels),
+              stride(stride) {
+            GGML_ASSERT(stride >= 1);
+        }
+
+        void load_fixed_tensors() {
+            if (kernel == nullptr || kernel_data.empty()) {
+                return;
+            }
+            ggml_backend_tensor_set(kernel, kernel_data.data(), 0, kernel_data.size() * sizeof(float));
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            if (stride == 1) {
+                return x;
+            }
+            GGML_ASSERT(kernel != nullptr);
+            GGML_ASSERT(x->ne[2] == channels);
+            if (ctx->conv2d_direct_enabled) {
+                return ggml_conv_2d_dw_direct(ctx->ggml_ctx, kernel, x, stride, stride, 2, 2, 1, 1);
+            }
+            return ggml_conv_2d_dw(ctx->ggml_ctx, kernel, x, stride, stride, 2, 2, 1, 1);
+        }
+    };
+
+    class SpatialRationalResampler : public GGMLBlock {
+    protected:
+        int64_t mid_channels;
+        int num;
+        int den;
+
+    public:
+        SpatialRationalResampler(int64_t mid_channels, int num, int den)
+            : mid_channels(mid_channels),
+              num(num),
+              den(den) {
+            GGML_ASSERT(num >= 1);
+            GGML_ASSERT(den >= 1);
+            blocks["conv"]          = std::shared_ptr<GGMLBlock>(new Conv2d(mid_channels, num * num * mid_channels, {3, 3}, {1, 1}, {1, 1}));
+            blocks["pixel_shuffle"] = std::shared_ptr<GGMLBlock>(new PixelShuffleND(num));
+            blocks["blur_down"]     = std::shared_ptr<GGMLBlock>(new BlurDownsample(mid_channels, den));
+        }
+
+        void load_fixed_tensors() {
+            auto blur_down = std::dynamic_pointer_cast<BlurDownsample>(blocks["blur_down"]);
+            blur_down->load_fixed_tensors();
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            auto conv          = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
+            auto pixel_shuffle = std::dynamic_pointer_cast<PixelShuffleND>(blocks["pixel_shuffle"]);
+            auto blur_down     = std::dynamic_pointer_cast<BlurDownsample>(blocks["blur_down"]);
+
+            // rearrange(x, "b c f h w -> (b f) c h w")
+            x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
+            x = conv->forward(ctx, x);
+            x = pixel_shuffle->forward(ctx, x);
+            x = blur_down->forward(ctx, x);
+            return ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
+        }
+    };
+
     class LatentUpsampler : public GGMLBlock {
     public:
         LatentUpsamplerConfig config;
@@ -177,9 +337,7 @@ namespace LTXVUpsampler {
         explicit LatentUpsampler(LatentUpsamplerConfig config)
             : config(std::move(config)) {
             GGML_ASSERT(this->config.dims == 3);
-            GGML_ASSERT(this->config.spatial_upsample);
-            GGML_ASSERT(!this->config.temporal_upsample);
-            GGML_ASSERT(!this->config.rational_resampler);
+            GGML_ASSERT(this->config.spatial_upsample || this->config.temporal_upsample);
 
             blocks["initial_conv"] = std::shared_ptr<GGMLBlock>(new Conv3d(this->config.in_channels,
                                                                            this->config.mid_channels,
@@ -190,12 +348,25 @@ namespace LTXVUpsampler {
             for (int i = 0; i < this->config.num_blocks_per_stage; ++i) {
                 blocks["res_blocks." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new ResBlock(this->config.mid_channels, this->config.dims));
             }
-            blocks["upsampler.0"] = std::shared_ptr<GGMLBlock>(new Conv2d(this->config.mid_channels,
-                                                                          4 * this->config.mid_channels,
-                                                                          {3, 3},
-                                                                          {1, 1},
-                                                                          {1, 1}));
-            blocks["upsampler.1"] = std::shared_ptr<GGMLBlock>(new PixelShuffleND(2));
+            if (this->config.rational_resampler) {
+                blocks["upsampler"] = std::shared_ptr<GGMLBlock>(new SpatialRationalResampler(this->config.mid_channels,
+                                                                                              this->config.spatial_up_num,
+                                                                                              this->config.spatial_down_den));
+            } else if (this->config.temporal_upsample) {
+                blocks["upsampler.0"] = std::shared_ptr<GGMLBlock>(new Conv3d(this->config.mid_channels,
+                                                                              this->config.temporal_up_factor * this->config.mid_channels,
+                                                                              {3, 3, 3},
+                                                                              {1, 1, 1},
+                                                                              {1, 1, 1}));
+                blocks["upsampler.1"] = std::shared_ptr<GGMLBlock>(new TemporalPixelShuffleND(this->config.temporal_up_factor));
+            } else {
+                blocks["upsampler.0"] = std::shared_ptr<GGMLBlock>(new Conv2d(this->config.mid_channels,
+                                                                              4 * this->config.mid_channels,
+                                                                              {3, 3},
+                                                                              {1, 1},
+                                                                              {1, 1}));
+                blocks["upsampler.1"] = std::shared_ptr<GGMLBlock>(new PixelShuffleND(2));
+            }
             for (int i = 0; i < this->config.num_blocks_per_stage; ++i) {
                 blocks["post_upsample_res_blocks." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new ResBlock(this->config.mid_channels, this->config.dims));
             }
@@ -207,13 +378,11 @@ namespace LTXVUpsampler {
         }
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
-            // x: [b*c, f, h, w]
-            // return: [b*c, f, h*2, w*2]
-            auto initial_conv  = std::dynamic_pointer_cast<Conv3d>(blocks["initial_conv"]);
-            auto initial_norm  = std::dynamic_pointer_cast<VideoGroupNorm>(blocks["initial_norm"]);
-            auto upsample_conv = std::dynamic_pointer_cast<Conv2d>(blocks["upsampler.0"]);
-            auto pixel_shuffle = std::dynamic_pointer_cast<PixelShuffleND>(blocks["upsampler.1"]);
-            auto final_conv    = std::dynamic_pointer_cast<Conv3d>(blocks["final_conv"]);
+            // x: [b, c, f, h, w]
+            // return: [b, c, scaled_f, scaled_h, scaled_w]
+            auto initial_conv = std::dynamic_pointer_cast<Conv3d>(blocks["initial_conv"]);
+            auto initial_norm = std::dynamic_pointer_cast<VideoGroupNorm>(blocks["initial_norm"]);
+            auto final_conv   = std::dynamic_pointer_cast<Conv3d>(blocks["final_conv"]);
 
             x = initial_conv->forward(ctx, x);
             x = initial_norm->forward(ctx, x);
@@ -226,11 +395,25 @@ namespace LTXVUpsampler {
                 sd::ggml_graph_cut::mark_graph_cut(x, "ltx_latent_upsampler.res_blocks." + std::to_string(i), "x");
             }
 
-            // rearrange(x, "b c f h w -> (b f) c h w"),
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));  // [b*f, c, h, w]
-            x = upsample_conv->forward(ctx, x);                                                      // [b*f, c*4, h, w]
-            x = pixel_shuffle->forward(ctx, x);                                                      // [b*f, c, h*2, w*2]
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));  // [b*c, f, h, w]
+            if (config.rational_resampler) {
+                auto upsampler = std::dynamic_pointer_cast<SpatialRationalResampler>(blocks["upsampler"]);
+                x              = upsampler->forward(ctx, x);
+            } else if (config.temporal_upsample) {
+                auto upsample_conv = std::dynamic_pointer_cast<Conv3d>(blocks["upsampler.0"]);
+                auto pixel_shuffle = std::dynamic_pointer_cast<TemporalPixelShuffleND>(blocks["upsampler.1"]);
+                x                  = upsample_conv->forward(ctx, x);                    // [b, c*2, f, h, w]
+                x                  = pixel_shuffle->forward(ctx, x);                    // [b, c, f*2, h, w]
+                x                  = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1, x->ne[2]);  // x[:, :, 1:, :, :]
+            } else {
+                auto upsample_conv = std::dynamic_pointer_cast<Conv2d>(blocks["upsampler.0"]);
+                auto pixel_shuffle = std::dynamic_pointer_cast<PixelShuffleND>(blocks["upsampler.1"]);
+
+                // rearrange(x, "b c f h w -> (b f) c h w"),
+                x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));  // [b*f, c, h, w]
+                x = upsample_conv->forward(ctx, x);                                                      // [b*f, c*4, h, w]
+                x = pixel_shuffle->forward(ctx, x);                                                      // [b*f, c, h*2, w*2]
+                x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));  // [b*c, f, h, w]
+            }
             sd::ggml_graph_cut::mark_graph_cut(x, "ltx_latent_upsampler.spatial_up", "x");
 
             for (int i = 0; i < config.num_blocks_per_stage; ++i) {
@@ -243,6 +426,14 @@ namespace LTXVUpsampler {
             sd::ggml_graph_cut::mark_graph_cut(x, "ltx_latent_upsampler.final", "x");
             return x;
         }
+
+        void load_fixed_tensors() {
+            if (!config.rational_resampler) {
+                return;
+            }
+            auto upsampler = std::dynamic_pointer_cast<SpatialRationalResampler>(blocks["upsampler"]);
+            upsampler->load_fixed_tensors();
+        }
     };
 
     struct LatentUpsamplerRunner : public GGMLRunner {
@@ -265,20 +456,24 @@ namespace LTXVUpsampler {
             }
 
             const auto& tensor_storage_map = model_loader.get_tensor_storage_map();
+            bool has_regular_upsampler     = has_tensor(tensor_storage_map, "upsampler.0.weight");
+            bool has_rational_spatial      = has_tensor(tensor_storage_map, "upsampler.conv.weight");
             if (!has_tensor(tensor_storage_map, "post_upsample_res_blocks.0.conv2.bias") ||
-                !has_tensor(tensor_storage_map, "upsampler.0.weight")) {
-                LOG_ERROR("unsupported LTX latent upsampler weights: expected spatial upsampler tensors");
+                (!has_regular_upsampler && !has_rational_spatial)) {
+                LOG_ERROR("unsupported LTX latent upsampler weights: expected upsampler tensors");
                 return false;
             }
 
             LatentUpsamplerConfig config = detect_config_from_weights(tensor_storage_map);
-            if (config.dims != 3 || !config.spatial_upsample || config.temporal_upsample ||
-                config.rational_resampler) {
-                LOG_ERROR("unsupported LTX latent upsampler config: dims=%d spatial=%d temporal=%d rational=%d",
+            if (config.dims != 3 || (!config.spatial_upsample && !config.temporal_upsample) ||
+                config.spatial_up_num < 1 || config.spatial_down_den < 1 || config.temporal_up_factor < 1) {
+                LOG_ERROR("unsupported LTX latent upsampler config: dims=%d spatial=%d temporal=%d rational=%d scale=%.3f temporal_factor=%d",
                           config.dims,
                           config.spatial_upsample,
                           config.temporal_upsample,
-                          config.rational_resampler);
+                          config.rational_resampler,
+                          config.spatial_scale,
+                          config.temporal_up_factor);
                 return false;
             }
 
@@ -291,15 +486,23 @@ namespace LTXVUpsampler {
 
             std::map<std::string, ggml_tensor*> tensors;
             model->get_param_tensors(tensors);
-            if (!model_loader.load_tensors(tensors, {}, n_threads)) {
+            std::set<std::string> ignore_tensors;
+            if (config.rational_resampler) {
+                ignore_tensors.insert("upsampler.blur_down.kernel");
+            }
+            if (!model_loader.load_tensors(tensors, ignore_tensors, n_threads)) {
                 LOG_ERROR("load LTX latent upsampler tensors failed");
                 return false;
             }
+            model->load_fixed_tensors();
 
-            LOG_INFO("LTX latent upsampler loaded: in_channels=%" PRId64 ", mid_channels=%" PRId64 ", blocks=%d",
+            LOG_INFO("LTX latent upsampler loaded: in_channels=%" PRId64 ", mid_channels=%" PRId64 ", blocks=%d, scale=%.3f, temporal_factor=%d, rational=%d",
                      config.in_channels,
                      config.mid_channels,
-                     config.num_blocks_per_stage);
+                     config.num_blocks_per_stage,
+                     config.spatial_scale,
+                     config.temporal_up_factor,
+                     config.rational_resampler);
             return true;
         }
 

src/ltx_vae.hpp

@@ -153,7 +153,7 @@ namespace LTXVAE {
 
             GGML_ASSERT(x->ne[2] >= temporal_pad);
 
-            int end_idx   = x->ne[2] - temporal_pad;
+            int end_idx   = (int)x->ne[2] - temporal_pad;
             int start_idx = std::max(end_idx - pad, 0);
 
             // Save a contiguous copy of the last `pad` frames so the large `x`
@@ -1251,65 +1251,22 @@ struct LTXVideoVAE : public VAE {
         temporal_tiling_enabled = enabled;
     }
 
-    static std::string trim_tiling_arg(std::string value) {
-        const char* whitespace = " \t\r\n";
-        size_t begin           = value.find_first_not_of(whitespace);
-        if (begin == std::string::npos) {
-            return "";
-        }
-        size_t end = value.find_last_not_of(whitespace);
-        return value.substr(begin, end - begin + 1);
-    }
-
-    static bool parse_tiling_int(const std::string& value, int& parsed) {
-        try {
-            size_t consumed = 0;
-            parsed          = std::stoi(value, &consumed);
-            return trim_tiling_arg(value.substr(consumed)).empty();
-        } catch (...) {
-            return false;
-        }
-    }
-
     void set_tiling_params(const sd_tiling_params_t& params) override {
         temporal_tiling_enabled = params.temporal_tiling;
         temporal_tile_frames    = DEFAULT_TEMPORAL_TILE_FRAMES;
         temporal_tile_overlap   = DEFAULT_TEMPORAL_TILE_OVERLAP;
 
-        const char* extra_tiling_args = params.extra_tiling_args;
-        if (extra_tiling_args == nullptr || extra_tiling_args[0] == '\0') {
-            return;
-        }
-
-        std::string raw(extra_tiling_args);
-        size_t start = 0;
-        for (size_t pos = 0; pos <= raw.size(); ++pos) {
-            if (pos != raw.size() && raw[pos] != ',' && raw[pos] != ';') {
-                continue;
+        for (const auto& [key, value] : parse_key_value_args(params.extra_tiling_args, "LTX VAE extra tiling arg")) {
+            int parsed = 0;
+            if (!parse_strict_int(value, parsed)) {
+                LOG_WARN("ignoring invalid LTX VAE extra tiling arg '%s=%s'", key.c_str(), value.c_str());
+            } else if (key == "temporal_tile_frames") {
+                temporal_tile_frames = std::max(1, parsed);
+            } else if (key == "temporal_tile_overlap") {
+                temporal_tile_overlap = std::max(0, parsed);
+            } else {
+                LOG_WARN("ignoring unknown LTX VAE extra tiling arg '%s'", key.c_str());
             }
-
-            std::string token = trim_tiling_arg(raw.substr(start, pos - start));
-            if (!token.empty()) {
-                size_t eq = token.find('=');
-                if (eq == std::string::npos) {
-                    LOG_WARN("ignoring malformed LTX VAE extra tiling arg '%s'", token.c_str());
-                } else {
-                    std::string key   = trim_tiling_arg(token.substr(0, eq));
-                    std::string value = trim_tiling_arg(token.substr(eq + 1));
-                    int parsed        = 0;
-                    if (!parse_tiling_int(value, parsed)) {
-                        LOG_WARN("ignoring invalid LTX VAE extra tiling arg '%s=%s'", key.c_str(), value.c_str());
-                    } else if (key == "temporal_tile_frames") {
-                        temporal_tile_frames = std::max(1, parsed);
-                    } else if (key == "temporal_tile_overlap") {
-                        temporal_tile_overlap = std::max(0, parsed);
-                    } else {
-                        LOG_WARN("ignoring unknown LTX VAE extra tiling arg '%s'", key.c_str());
-                    }
-                }
-            }
-
-            start = pos + 1;
         }
     }
 
@@ -1560,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;
@@ -1577,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

@@ -10,6 +10,7 @@
 #include <vector>
 
 #include "common_block.hpp"
+#include "diffusion_model.hpp"
 #include "flux.hpp"
 #include "rope.hpp"
 
@@ -1534,8 +1535,7 @@ namespace LTXV {
         }
     };
 
-    struct LTXAVRunner : public GGMLRunner {
-        std::string prefix;
+    struct LTXAVRunner : public DiffusionModelRunner {
         LTXAVParams params;
         LTXAVModelBlock model;
         std::vector<float> video_pe_vec;
@@ -1561,8 +1561,7 @@ namespace LTXV {
                     ggml_backend_t params_backend,
                     const String2TensorStorage& tensor_storage_map = {},
                     const std::string& prefix                      = "model.diffusion_model")
-            : GGMLRunner(backend, params_backend),
-              prefix(prefix),
+            : DiffusionModelRunner(backend, params_backend, prefix),
               params(),
               model(params) {
             auto patchify_proj_iter = tensor_storage_map.find(prefix + ".patchify_proj.weight");
@@ -1673,7 +1672,7 @@ namespace LTXV {
             return "ltxav";
         }
 
-        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) {
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
             model.get_param_tensors(tensors, prefix);
         }
 
@@ -1917,6 +1916,22 @@ namespace LTXV {
             return out;
         }
 
+        sd::Tensor<float> compute(int n_threads,
+                                  const DiffusionParams& diffusion_params) override {
+            GGML_ASSERT(diffusion_params.x != nullptr);
+            GGML_ASSERT(diffusion_params.timesteps != nullptr);
+            const auto* extra = diffusion_extra_as<LTXAVDiffusionExtra>(diffusion_params);
+            return compute(n_threads,
+                           *diffusion_params.x,
+                           *diffusion_params.timesteps,
+                           tensor_or_empty(diffusion_params.context),
+                           tensor_or_empty(extra->audio_x),
+                           tensor_or_empty(extra->audio_timesteps),
+                           extra->audio_length,
+                           extra->frame_rate,
+                           tensor_or_empty(extra->video_positions));
+        }
+
         void test(const std::string& x_path,
                   const std::string& timesteps_path       = "",
                   const std::string& context_path         = "",
@@ -1980,7 +1995,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;
@@ -2002,7 +2017,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

@@ -3,6 +3,7 @@
 
 #include <memory>
 
+#include "diffusion_model.hpp"
 #include "ggml_extend.hpp"
 #include "model.h"
 
@@ -824,14 +825,14 @@ public:
         return x;
     }
 };
-struct MMDiTRunner : public GGMLRunner {
+struct MMDiTRunner : public DiffusionModelRunner {
     MMDiT mmdit;
 
     MMDiTRunner(ggml_backend_t backend,
                 ggml_backend_t params_backend,
                 const String2TensorStorage& tensor_storage_map = {},
                 const std::string prefix                       = "")
-        : GGMLRunner(backend, params_backend), mmdit(tensor_storage_map) {
+        : DiffusionModelRunner(backend, params_backend, prefix), mmdit(tensor_storage_map) {
         mmdit.init(params_ctx, tensor_storage_map, prefix);
     }
 
@@ -839,7 +840,7 @@ struct MMDiTRunner : public GGMLRunner {
         return "mmdit";
     }
 
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
         mmdit.get_param_tensors(tensors, prefix);
     }
 
@@ -885,6 +886,20 @@ struct MMDiTRunner : public GGMLRunner {
         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);
+        const auto* extra = diffusion_extra_as<SkipLayerDiffusionExtra>(diffusion_params);
+        static const std::vector<int> empty_skip_layers;
+        return compute(n_threads,
+                       *diffusion_params.x,
+                       *diffusion_params.timesteps,
+                       tensor_or_empty(diffusion_params.context),
+                       tensor_or_empty(diffusion_params.y),
+                       extra->skip_layers ? *extra->skip_layers : empty_skip_layers);
+    }
+
     void test() {
         ggml_init_params params;
         params.mem_size   = static_cast<size_t>(10 * 1024 * 1024);  // 10 MB
@@ -932,13 +947,17 @@ struct MMDiTRunner : public GGMLRunner {
 
     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

@@ -410,7 +410,7 @@ bool ModelLoader::init_from_diffusers_file(const std::string& file_path, const s
 }
 
 SDVersion ModelLoader::get_sd_version() {
-    TensorStorage token_embedding_weight, input_block_weight;
+    TensorStorage token_embedding_weight, input_block_weight, context_ebedding_weight;
 
     bool has_multiple_encoders = false;
     bool is_unet               = false;
@@ -428,9 +428,13 @@ SDVersion ModelLoader::get_sd_version() {
     bool has_attn_1024               = false;
 
     for (auto& [name, tensor_storage] : tensor_storage_map) {
-        if (tensor_storage.name.find("model.diffusion_model.double_blocks.") != std::string::npos) {
+        if (tensor_storage.name.find("model.diffusion_model.double_blocks.") != std::string::npos ||
+            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("model.diffusion_model.nerf_final_layer_conv.") != std::string::npos) {
             return VERSION_CHROMA_RADIANCE;
         }
@@ -441,6 +445,10 @@ SDVersion ModelLoader::get_sd_version() {
             tensor_storage_map.find("model.language_model.layers.0.self_attn.q_proj.weight") != tensor_storage_map.end()) {
             return VERSION_HIDREAM_O1;
         }
+        if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.attn.norm_added_q.weight") != std::string::npos &&
+            tensor_storage_map.find("model.diffusion_model.transformer_blocks.0.img_mlp.w1.weight") != tensor_storage_map.end()) {
+            return VERSION_LENS;
+        }
         if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.img_mod.1.weight") != std::string::npos) {
             return VERSION_QWEN_IMAGE;
         }
@@ -522,6 +530,9 @@ SDVersion ModelLoader::get_sd_version() {
             tensor_storage.name == "unet.conv_in.weight") {
             input_block_weight = tensor_storage;
         }
+        if (tensor_storage.name == "model.diffusion_model.txt_in.weight" || tensor_storage.name == "model.diffusion_model.context_embedder.weight") {
+            context_ebedding_weight = tensor_storage;
+        }
     }
     if (is_wan) {
         LOG_DEBUG("patch_embedding_channels %d", patch_embedding_channels);
@@ -552,16 +563,20 @@ SDVersion ModelLoader::get_sd_version() {
     }
 
     if (is_flux && !is_flux2) {
-        if (input_block_weight.ne[0] == 384) {
-            return VERSION_FLUX_FILL;
+        if (context_ebedding_weight.ne[0] == 3584) {
+            return VERSION_LONGCAT;
+        } else {
+            if (input_block_weight.ne[0] == 384) {
+                return VERSION_FLUX_FILL;
+            }
+            if (input_block_weight.ne[0] == 128) {
+                return VERSION_FLUX_CONTROLS;
+            }
+            if (input_block_weight.ne[0] == 196) {
+                return VERSION_FLEX_2;
+            }
+            return VERSION_FLUX;
         }
-        if (input_block_weight.ne[0] == 128) {
-            return VERSION_FLUX_CONTROLS;
-        }
-        if (input_block_weight.ne[0] == 196) {
-            return VERSION_FLEX_2;
-        }
-        return VERSION_FLUX;
     }
 
     if (is_flux2) {
@@ -853,8 +868,13 @@ std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggm
             if (dst_tensor == nullptr)
                 continue;
 
-            if (tensor_storage.type != dst_tensor->type)
+            if (tensor_storage.is_f8_e4m3 ||
+                tensor_storage.is_f8_e5m2 ||
+                tensor_storage.is_f64 ||
+                tensor_storage.is_i64 ||
+                tensor_storage.type != dst_tensor->type) {
                 continue;
+            }
 
             size_t tensor_size   = tensor_storage.nbytes();
             size_t tensor_offset = tensor_storage.offset;
@@ -987,6 +1007,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;

src/model.h

@@ -47,6 +47,9 @@ enum SDVersion {
     VERSION_Z_IMAGE,
     VERSION_OVIS_IMAGE,
     VERSION_ERNIE_IMAGE,
+    VERSION_LENS,
+    VERSION_LONGCAT,
+    VERSION_PID,
     VERSION_COUNT,
 };
 
@@ -141,6 +144,13 @@ static inline bool sd_version_is_z_image(SDVersion version) {
     return false;
 }
 
+static inline bool sd_version_is_longcat(SDVersion version) {
+    if (version == VERSION_LONGCAT) {
+        return true;
+    }
+    return false;
+}
+
 static inline bool sd_version_is_ernie_image(SDVersion version) {
     if (version == VERSION_ERNIE_IMAGE) {
         return true;
@@ -148,8 +158,22 @@ static inline bool sd_version_is_ernie_image(SDVersion version) {
     return false;
 }
 
+static inline bool sd_version_is_lens(SDVersion version) {
+    if (version == VERSION_LENS) {
+        return true;
+    }
+    return false;
+}
+
+static inline bool sd_version_is_pid(SDVersion version) {
+    if (version == VERSION_PID) {
+        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)) {
+    if (sd_version_is_flux2(version) || sd_version_is_ernie_image(version) || sd_version_is_lens(version)) {
         return true;
     }
     return false;
@@ -176,7 +200,10 @@ static inline bool sd_version_is_dit(SDVersion version) {
         version == VERSION_HIDREAM_O1 ||
         sd_version_is_anima(version) ||
         sd_version_is_z_image(version) ||
-        sd_version_is_ernie_image(version)) {
+        sd_version_is_ernie_image(version) ||
+        sd_version_is_lens(version) ||
+        sd_version_is_longcat(version) ||
+        sd_version_is_pid(version)) {
         return true;
     }
     return false;

src/name_conversion.cpp

@@ -128,6 +128,7 @@ std::string convert_cond_stage_model_name(std::string name, std::string prefix)
     };
 
     static const std::vector<std::pair<std::string, std::string>> llm_name_map{
+        {"attn_sinks.weight", "self_attn.sinks"},
         {"token_embd.", "model.embed_tokens."},
         {"blk.", "model.layers."},
         {"attn_q.", "self_attn.q_proj."},
@@ -137,9 +138,15 @@ std::string convert_cond_stage_model_name(std::string name, std::string prefix)
         {"attn_k_norm.", "self_attn.k_norm."},
         {"attn_output.", "self_attn.o_proj."},
         {"attn_norm.", "input_layernorm."},
+        {"attn_post_norm.", "post_attention_norm."},
+        {"ffn_gate_inp.", "mlp.router."},
+        {"ffn_gate_exps.", "mlp.experts.gate_proj."},
+        {"ffn_up_exps.", "mlp.experts.up_proj."},
+        {"ffn_down_exps.", "mlp.experts.down_proj."},
         {"ffn_down.", "mlp.down_proj."},
         {"ffn_gate.", "mlp.gate_proj."},
         {"ffn_up.", "mlp.up_proj."},
+        {"ffn_post_norm.", "post_ffw_norm."},
         {"ffn_norm.", "post_attention_layernorm."},
         {"output_norm.", "model.norm."},
     };
@@ -508,6 +515,12 @@ std::string convert_diffusers_dit_to_original_flux(std::string name) {
     static std::unordered_map<std::string, std::string> flux_name_map;
 
     if (flux_name_map.empty()) {
+        // --- time_embed (longcat) ---
+        flux_name_map["time_embed.timestep_embedder.linear_1.weight"] = "time_in.in_layer.weight";
+        flux_name_map["time_embed.timestep_embedder.linear_1.bias"]   = "time_in.in_layer.bias";
+        flux_name_map["time_embed.timestep_embedder.linear_2.weight"] = "time_in.out_layer.weight";
+        flux_name_map["time_embed.timestep_embedder.linear_2.bias"]   = "time_in.out_layer.bias";
+
         // --- time_text_embed ---
         flux_name_map["time_text_embed.timestep_embedder.linear_1.weight"] = "time_in.in_layer.weight";
         flux_name_map["time_text_embed.timestep_embedder.linear_1.bias"]   = "time_in.in_layer.bias";
@@ -561,6 +574,11 @@ std::string convert_diffusers_dit_to_original_flux(std::string name) {
             flux_name_map[block_prefix + "attn.norm_k.weight"]       = dst_prefix + "img_attn.norm.key_norm.scale";
             flux_name_map[block_prefix + "attn.norm_added_q.weight"] = dst_prefix + "txt_attn.norm.query_norm.scale";
             flux_name_map[block_prefix + "attn.norm_added_k.weight"] = dst_prefix + "txt_attn.norm.key_norm.scale";
+            // Comfy-Org/LongCat-Image stores already-converted RMSNorm tensors as *.weight.
+            flux_name_map[dst_prefix + "img_attn.norm.query_norm.weight"] = dst_prefix + "img_attn.norm.query_norm.scale";
+            flux_name_map[dst_prefix + "img_attn.norm.key_norm.weight"]   = dst_prefix + "img_attn.norm.key_norm.scale";
+            flux_name_map[dst_prefix + "txt_attn.norm.query_norm.weight"] = dst_prefix + "txt_attn.norm.query_norm.scale";
+            flux_name_map[dst_prefix + "txt_attn.norm.key_norm.weight"]   = dst_prefix + "txt_attn.norm.key_norm.scale";
 
             // ff
             flux_name_map[block_prefix + "ff.net.0.proj.weight"] = dst_prefix + "img_mlp.0.weight";
@@ -599,8 +617,11 @@ std::string convert_diffusers_dit_to_original_flux(std::string name) {
 
             flux_name_map[block_prefix + "attn.norm_q.weight"] = dst_prefix + "norm.query_norm.scale";
             flux_name_map[block_prefix + "attn.norm_k.weight"] = dst_prefix + "norm.key_norm.scale";
-            flux_name_map[block_prefix + "proj_out.weight"]    = dst_prefix + "linear2.weight";
-            flux_name_map[block_prefix + "proj_out.bias"]      = dst_prefix + "linear2.bias";
+            // Comfy-Org/LongCat-Image stores already-converted RMSNorm tensors as *.weight.
+            flux_name_map[dst_prefix + "norm.query_norm.weight"] = dst_prefix + "norm.query_norm.scale";
+            flux_name_map[dst_prefix + "norm.key_norm.weight"]   = dst_prefix + "norm.key_norm.scale";
+            flux_name_map[block_prefix + "proj_out.weight"]      = dst_prefix + "linear2.weight";
+            flux_name_map[block_prefix + "proj_out.bias"]        = dst_prefix + "linear2.bias";
         }
 
         // --- final layers ---
@@ -668,7 +689,7 @@ std::string convert_diffusion_model_name(std::string name, std::string prefix, S
         name = convert_diffusers_unet_to_original_sdxl(name);
     } else if (sd_version_is_sd3(version)) {
         name = convert_diffusers_dit_to_original_sd3(name);
-    } else if (sd_version_is_flux(version) || sd_version_is_flux2(version)) {
+    } else if (sd_version_is_flux(version) || sd_version_is_flux2(version) || sd_version_is_longcat(version)) {
         name = convert_diffusers_dit_to_original_flux(name);
     } else if (sd_version_is_z_image(version)) {
         name = convert_diffusers_dit_to_original_lumina2(name);

src/pid.hpp

@@ -0,0 +1,842 @@
+#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 PixelDiTParams {
+        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;
+    };
+
+    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 {
+        PixelDiTParams params_cfg;
+
+        LQProjection2D(const PixelDiTParams& params_cfg)
+            : params_cfg(params_cfg) {
+            blocks["latent_proj.0"] = std::make_shared<Conv2d>(params_cfg.lq_latent_channels, params_cfg.lq_hidden_dim, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, std::pair<int, int>{1, 1});
+            blocks["latent_proj.2"] = std::make_shared<Conv2d>(params_cfg.lq_hidden_dim, params_cfg.lq_hidden_dim, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, std::pair<int, int>{1, 1});
+            for (int i = 0; i < params_cfg.lq_num_res_blocks; ++i) {
+                blocks["latent_proj." + std::to_string(3 + i)] = std::make_shared<PiDResBlock>(params_cfg.lq_hidden_dim);
+            }
+
+            int num_outputs = static_cast<int>((params_cfg.patch_depth + params_cfg.lq_interval - 1) / params_cfg.lq_interval);
+            for (int i = 0; i < num_outputs; ++i) {
+                blocks["output_heads." + std::to_string(i)] = std::make_shared<Linear>(params_cfg.lq_hidden_dim, params_cfg.hidden_size, true);
+                blocks["gate_modules." + std::to_string(i)] = std::make_shared<SigmaAwareGate>(params_cfg.hidden_size);
+            }
+        }
+
+        bool is_gate_active(int block_idx) const {
+            return block_idx % params_cfg.lq_interval == 0;
+        }
+
+        int get_output_index(int block_idx) const {
+            return block_idx / static_cast<int>(params_cfg.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>(params_cfg.lq_sr_scale * params_cfg.lq_latent_down_factor) /
+                                     static_cast<float>(params_cfg.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 < params_cfg.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>((params_cfg.patch_depth + params_cfg.lq_interval - 1) / params_cfg.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 {
+        PixelDiTParams params_cfg;
+
+        PixelDiT() = default;
+
+        PixelDiT(const PixelDiTParams& params_cfg)
+            : params_cfg(params_cfg) {
+            blocks["pixel_embedder"] = std::make_shared<PixelTokenEmbedder>(params_cfg.in_channels, params_cfg.pixel_hidden_size);
+            blocks["s_embedder"]     = std::make_shared<PatchTokenEmbedder>(params_cfg.in_channels * params_cfg.patch_size * params_cfg.patch_size, params_cfg.hidden_size, false, true);
+            blocks["t_embedder"]     = std::make_shared<PixelDiTTimestepEmbedder>(params_cfg.hidden_size);
+            blocks["y_embedder"]     = std::make_shared<PatchTokenEmbedder>(params_cfg.txt_embed_dim, params_cfg.hidden_size, true, true);
+            for (int i = 0; i < params_cfg.patch_depth; ++i) {
+                blocks["patch_blocks." + std::to_string(i)] = std::make_shared<MMDiTBlockT2I>(params_cfg.hidden_size, params_cfg.num_groups, params_cfg.patch_mlp_hidden_dim);
+            }
+            for (int i = 0; i < params_cfg.pixel_depth; ++i) {
+                blocks["pixel_blocks." + std::to_string(i)] = std::make_shared<PiTBlock>(params_cfg.pixel_hidden_size,
+                                                                                         params_cfg.hidden_size,
+                                                                                         params_cfg.patch_size,
+                                                                                         params_cfg.pixel_attn_hidden_size,
+                                                                                         params_cfg.pixel_num_groups);
+            }
+            blocks["final_layer"] = std::make_shared<FinalLayer>(params_cfg.pixel_hidden_size, params_cfg.in_channels);
+            blocks["lq_proj"]     = std::make_shared<LQProjection2D>(params_cfg);
+        }
+
+        void init_params(ggml_context* ctx,
+                         const String2TensorStorage& tensor_storage_map = {},
+                         std::string prefix                             = "") override {
+            params["y_pos_embedding"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, params_cfg.hidden_size, params_cfg.txt_max_length, 1);
+        }
+
+        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>(params_cfg.patch_size), static_cast<int>(params_cfg.patch_size));
+            int64_t W      = x->ne[0];
+            int64_t H      = x->ne[1];
+            int64_t B      = x->ne[3];
+            int64_t Hs     = H / params_cfg.patch_size;
+            int64_t Ws     = W / params_cfg.patch_size;
+            int64_t L      = Hs * Ws;
+            int64_t P2     = params_cfg.patch_size * params_cfg.patch_size;
+
+            auto x_patches = DiT::patchify(ctx->ggml_ctx, x, static_cast<int>(params_cfg.patch_size), static_cast<int>(params_cfg.patch_size), true);
+            auto t_emb     = t_embedder->forward(ctx, timesteps);
+            auto condition = ggml_silu(ctx->ggml_ctx, t_emb);
+
+            GGML_ASSERT(context != nullptr);
+            int64_t Ltxt = std::min<int64_t>(context->ne[1], params_cfg.txt_max_length);
+            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 < params_cfg.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, params_cfg.hidden_size, L * B);
+            auto pixels = pixel_embedder->forward(ctx, x, params_cfg.patch_size, pixel_pos_full);
+            for (int i = 0; i < params_cfg.pixel_depth; ++i) {
+                auto block = std::dynamic_pointer_cast<PiTBlock>(blocks["pixel_blocks." + std::to_string(i)]);
+                pixels     = block->forward(ctx, pixels, s_cond, H, W, pixel_pos_comp);
+                sd::ggml_graph_cut::mark_graph_cut(pixels, "pid.pixel_blocks." + std::to_string(i), "pixels");
+            }
+
+            pixels   = final_layer->forward(ctx, pixels);
+            pixels   = ggml_reshape_3d(ctx->ggml_ctx, pixels, params_cfg.in_channels * P2, L, B);
+            auto out = DiT::unpatchify(ctx->ggml_ctx,
+                                       pixels,
+                                       Hs,
+                                       Ws,
+                                       static_cast<int>(params_cfg.patch_size),
+                                       static_cast<int>(params_cfg.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 {
+        PixelDiTParams params_cfg;
+        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) {
+            for (const auto& pair : tensor_storage_map) {
+                const std::string& tensor_name = pair.first;
+                if (tensor_name.find(prefix) == std::string::npos) {
+                    continue;
+                }
+                size_t pos = tensor_name.find("patch_blocks.");
+                if (pos != std::string::npos) {
+                    auto items = split_string(tensor_name.substr(pos), '.');
+                    if (items.size() > 1) {
+                        int block_index        = atoi(items[1].c_str());
+                        params_cfg.patch_depth = std::max<int64_t>(params_cfg.patch_depth, block_index + 1);
+                    }
+                }
+                pos = tensor_name.find("pixel_blocks.");
+                if (pos != std::string::npos) {
+                    auto items = split_string(tensor_name.substr(pos), '.');
+                    if (items.size() > 1) {
+                        int block_index        = atoi(items[1].c_str());
+                        params_cfg.pixel_depth = std::max<int64_t>(params_cfg.pixel_depth, block_index + 1);
+                    }
+                }
+                if (tensor_name.find("lq_proj.latent_proj.0.weight") != std::string::npos) {
+                    params_cfg.lq_latent_channels    = pair.second.ne[2];
+                    params_cfg.lq_latent_down_factor = params_cfg.lq_latent_channels >= 64 ? 16 : 8;
+                }
+                if (tensor_name.find("patch_blocks.0.mlp_x.w1.weight") != std::string::npos) {
+                    params_cfg.patch_mlp_hidden_dim = pair.second.ne[1];
+                }
+            }
+            LOG_INFO("PiD params: patch_depth=%" PRId64 ", pixel_depth=%" PRId64 ", patch_mlp_hidden_dim=%" PRId64 ", lq_latent_channels=%" PRId64 ", lq_latent_down_factor=%" PRId64,
+                     params_cfg.patch_depth,
+                     params_cfg.pixel_depth,
+                     params_cfg.patch_mlp_hidden_dim,
+                     params_cfg.lq_latent_channels,
+                     params_cfg.lq_latent_down_factor);
+            model = PixelDiT(params_cfg);
+            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>(params_cfg.patch_size));
+            int64_t Hp = align_up(static_cast<int>(H), static_cast<int>(params_cfg.patch_size));
+            int64_t Hs = Hp / params_cfg.patch_size;
+            int64_t Ws = Wp / params_cfg.patch_size;
+
+            pos_img_vec  = make_rope_2d(static_cast<int>(Hs),
+                                        static_cast<int>(Ws),
+                                        static_cast<int>(params_cfg.hidden_size / params_cfg.num_groups),
+                                        10000.f,
+                                        16.f,
+                                        static_cast<int>(params_cfg.rope_ref_grid_h),
+                                        static_cast<int>(params_cfg.rope_ref_grid_w));
+            auto pos_img = ggml_new_tensor_4d(compute_ctx,
+                                              GGML_TYPE_F32,
+                                              2,
+                                              2,
+                                              params_cfg.hidden_size / params_cfg.num_groups / 2,
+                                              Hs * Ws);
+            set_backend_tensor_data(pos_img, pos_img_vec.data());
+
+            int64_t Ltxt = std::min<int64_t>(context->ne[1], params_cfg.txt_max_length);
+            pos_txt_vec  = make_rope_1d(static_cast<int>(Ltxt),
+                                        static_cast<int>(params_cfg.hidden_size / params_cfg.num_groups),
+                                        params_cfg.text_rope_theta);
+            auto pos_txt = ggml_new_tensor_4d(compute_ctx,
+                                              GGML_TYPE_F32,
+                                              2,
+                                              2,
+                                              params_cfg.hidden_size / params_cfg.num_groups / 2,
+                                              Ltxt);
+            set_backend_tensor_data(pos_txt, pos_txt_vec.data());
+
+            pixel_pos_vec  = make_pixel_abs_pos(static_cast<int>(Hp),
+                                                static_cast<int>(Wp),
+                                                static_cast<int>(params_cfg.pixel_hidden_size));
+            auto pixel_pos = ggml_new_tensor_3d(compute_ctx,
+                                                GGML_TYPE_F32,
+                                                params_cfg.pixel_hidden_size,
+                                                Wp * Hp,
+                                                1);
+            set_backend_tensor_data(pixel_pos, pixel_pos_vec.data());
+
+            pixel_pos_comp_vec  = make_rope_2d(static_cast<int>(Hs),
+                                               static_cast<int>(Ws),
+                                               static_cast<int>(params_cfg.pixel_attn_hidden_size / params_cfg.pixel_num_groups),
+                                               10000.f,
+                                               16.f,
+                                               static_cast<int>(params_cfg.rope_ref_grid_h),
+                                               static_cast<int>(params_cfg.rope_ref_grid_w));
+            auto pixel_pos_comp = ggml_new_tensor_4d(compute_ctx,
+                                                     GGML_TYPE_F32,
+                                                     2,
+                                                     2,
+                                                     params_cfg.pixel_attn_hidden_size / params_cfg.pixel_num_groups / 2,
+                                                     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

@@ -4,6 +4,7 @@
 #include <memory>
 
 #include "common_block.hpp"
+#include "diffusion_model.hpp"
 #include "flux.hpp"
 
 namespace Qwen {
@@ -479,7 +480,7 @@ namespace Qwen {
         }
     };
 
-    struct QwenImageRunner : public GGMLRunner {
+    struct QwenImageRunner : public DiffusionModelRunner {
     public:
         QwenImageParams qwen_image_params;
         QwenImageModel qwen_image;
@@ -493,7 +494,7 @@ namespace Qwen {
                         const std::string prefix                       = "",
                         SDVersion version                              = VERSION_QWEN_IMAGE,
                         bool zero_cond_t                               = false)
-            : GGMLRunner(backend, params_backend) {
+            : 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) {
@@ -528,7 +529,7 @@ namespace Qwen {
             return "qwen_image";
         }
 
-        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
             qwen_image.get_param_tensors(tensors, prefix);
         }
 
@@ -624,6 +625,19 @@ namespace Qwen {
             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);
+            static const std::vector<sd::Tensor<float>> empty_ref_latents;
+            return compute(n_threads,
+                           *diffusion_params.x,
+                           *diffusion_params.timesteps,
+                           tensor_or_empty(diffusion_params.context),
+                           diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
+                           diffusion_params.increase_ref_index);
+        }
+
         void test() {
             ggml_init_params params;
             params.mem_size   = static_cast<size_t>(1024 * 1024) * 1024;  // 1GB
@@ -669,7 +683,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;
@@ -691,7 +705,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

@@ -111,6 +111,16 @@ namespace Rope {
         return txt_ids;
     }
 
+    __STATIC_INLINE__ std::vector<std::vector<float>> gen_longcat_txt_ids(int bs, int context_len, int axes_dim_num) {
+        auto txt_ids = std::vector<std::vector<float>>(bs * context_len, std::vector<float>(axes_dim_num, 0.0f));
+        for (int i = 0; i < bs * context_len; i++) {
+            float token_index = static_cast<float>(i % context_len);
+            txt_ids[i][1]     = token_index;
+            txt_ids[i][2]     = token_index;
+        }
+        return txt_ids;
+    }
+
     __STATIC_INLINE__ std::vector<std::vector<float>> gen_flux_img_ids(int h,
                                                                        int w,
                                                                        int patch_size,
@@ -122,7 +132,6 @@ namespace Rope {
                                                                        bool scale_rope = false) {
         int h_len = (h + (patch_size / 2)) / patch_size;
         int w_len = (w + (patch_size / 2)) / patch_size;
-
         std::vector<std::vector<float>> img_ids(h_len * w_len, std::vector<float>(axes_dim_num, 0.0));
 
         int h_start = h_offset;
@@ -135,7 +144,6 @@ namespace Rope {
 
         std::vector<float> row_ids = linspace<float>(1.f * h_start, 1.f * h_start + h_len - 1, h_len);
         std::vector<float> col_ids = linspace<float>(1.f * w_start, 1.f * w_start + w_len - 1, w_len);
-
         for (int i = 0; i < h_len; ++i) {
             for (int j = 0; j < w_len; ++j) {
                 img_ids[i * w_len + j][0] = 1.f * index;
@@ -241,17 +249,72 @@ namespace Rope {
         return embed_nd(ids, bs, axis_thetas, axes_dim, wrap_dims, layout);
     }
 
+    __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,
+                                                                   int start_index,
                                                                    const std::vector<ggml_tensor*>& ref_latents,
                                                                    bool increase_ref_index,
                                                                    float ref_index_scale,
-                                                                   bool scale_rope) {
+                                                                   bool scale_rope,
+                                                                   int base_offset = 0) {
         std::vector<std::vector<float>> ids;
         int curr_h_offset = 0;
         int curr_w_offset = 0;
-        int index         = 1;
+        int index         = start_index;
         for (ggml_tensor* ref : ref_latents) {
             int h_offset = 0;
             int w_offset = 0;
@@ -270,8 +333,8 @@ namespace Rope {
                                             bs,
                                             axes_dim_num,
                                             static_cast<int>(index * ref_index_scale),
-                                            h_offset,
-                                            w_offset,
+                                            h_offset + base_offset,
+                                            w_offset + base_offset,
                                             scale_rope);
             ids          = concat_ids(ids, ref_ids, bs);
 
@@ -294,13 +357,17 @@ namespace Rope {
                                                                    std::set<int> txt_arange_dims,
                                                                    const std::vector<ggml_tensor*>& ref_latents,
                                                                    bool increase_ref_index,
-                                                                   float ref_index_scale) {
-        auto txt_ids = gen_flux_txt_ids(bs, context_len, axes_dim_num, txt_arange_dims);
-        auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num);
+                                                                   float ref_index_scale,
+                                                                   bool is_longcat) {
+        int x_index = is_longcat ? 1 : 0;
+
+        auto txt_ids = is_longcat ? gen_longcat_txt_ids(bs, context_len, axes_dim_num) : gen_flux_txt_ids(bs, context_len, axes_dim_num, txt_arange_dims);
+        int offset   = is_longcat ? context_len : 0;
+        auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, x_index, offset, offset);
 
         auto ids = concat_ids(txt_ids, img_ids, bs);
         if (ref_latents.size() > 0) {
-            auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, ref_latents, increase_ref_index, ref_index_scale, false);
+            auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, x_index + 1, ref_latents, increase_ref_index, ref_index_scale, false, offset);
             ids           = concat_ids(ids, refs_ids, bs);
         }
         return ids;
@@ -319,7 +386,8 @@ namespace Rope {
                                                      int theta,
                                                      bool circular_h,
                                                      bool circular_w,
-                                                     const std::vector<int>& axes_dim) {
+                                                     const std::vector<int>& axes_dim,
+                                                     bool is_longcat) {
         std::vector<std::vector<float>> ids = gen_flux_ids(h,
                                                            w,
                                                            patch_size,
@@ -329,7 +397,8 @@ namespace Rope {
                                                            txt_arange_dims,
                                                            ref_latents,
                                                            increase_ref_index,
-                                                           ref_index_scale);
+                                                           ref_index_scale,
+                                                           is_longcat);
         std::vector<std::vector<int>> wrap_dims;
         if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
             int h_len = (h + (patch_size / 2)) / patch_size;
@@ -394,7 +463,7 @@ namespace Rope {
         auto img_ids     = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, 0, 0, 0, true);
         auto ids         = concat_ids(txt_ids_repeated, img_ids, bs);
         if (ref_latents.size() > 0) {
-            auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, ref_latents, increase_ref_index, 1.f, true);
+            auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, 1, ref_latents, increase_ref_index, 1.f, true);
             ids           = concat_ids(ids, refs_ids, bs);
         }
         return ids;
@@ -462,6 +531,52 @@ namespace Rope {
         return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
     }
 
+    __STATIC_INLINE__ std::vector<std::vector<float>> gen_lens_ids(int h,
+                                                                   int w,
+                                                                   int bs,
+                                                                   int context_len,
+                                                                   bool scale_rope = true) {
+        auto img_ids_repeated = gen_flux_img_ids(h, w, 1, bs, 3, 0, 0, 0, scale_rope);
+
+        int txt_id_start = scale_rope ? std::max(h / 2, w / 2) : 0;
+        auto txt_ids     = linspace<float>(1.f * txt_id_start, 1.f * context_len + txt_id_start, context_len);
+        std::vector<std::vector<float>> txt_ids_repeated(bs * context_len, std::vector<float>(3));
+        for (int i = 0; i < bs; ++i) {
+            for (int j = 0; j < txt_ids.size(); ++j) {
+                txt_ids_repeated[i * txt_ids.size() + j] = {txt_ids[j], txt_ids[j], txt_ids[j]};
+            }
+        }
+
+        return concat_ids(img_ids_repeated, txt_ids_repeated, bs);
+    }
+
+    __STATIC_INLINE__ std::vector<float> gen_lens_pe(int h,
+                                                     int w,
+                                                     int bs,
+                                                     int context_len,
+                                                     int theta,
+                                                     bool circular_h,
+                                                     bool circular_w,
+                                                     const std::vector<int>& axes_dim) {
+        std::vector<std::vector<float>> ids = gen_lens_ids(h, w, bs, context_len, true);
+        std::vector<std::vector<int>> wrap_dims;
+        if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
+            size_t pos_len = ids.size() / bs;
+            wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
+            const size_t img_tokens = static_cast<size_t>(h) * static_cast<size_t>(w);
+            for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
+                if (circular_h) {
+                    wrap_dims[1][token_i] = h;
+                }
+                if (circular_w) {
+                    wrap_dims[2][token_i] = w;
+                }
+            }
+        }
+
+        return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
+    }
+
     __STATIC_INLINE__ std::vector<std::vector<float>> gen_ernie_image_ids(int h,
                                                                           int w,
                                                                           int patch_size,

src/t5.hpp

@@ -475,8 +475,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 +563,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 +581,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/tae.hpp

@@ -584,7 +584,7 @@ public:
     TAESD(bool decode_only = true, SDVersion version = VERSION_SD1)
         : decode_only(decode_only) {
         bool use_midblock_gn = false;
-        taef2                = sd_version_is_flux2(version);
+        taef2                = sd_version_uses_flux2_vae(version);
 
         if (sd_version_is_dit(version)) {
             z_channels = 16;

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/gpt_oss_tokenizer.cpp

@@ -0,0 +1,91 @@
+#include "gpt_oss_tokenizer.h"
+
+#include "json.hpp"
+#include "util.h"
+#include "vocab/vocab.h"
+
+void GPTOSSTokenizer::load_from_merges(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
+    auto byte_unicode_pairs = bytes_to_unicode();
+    byte_encoder            = std::map<int, std::u32string>(byte_unicode_pairs.begin(), byte_unicode_pairs.end());
+    for (auto& pair : byte_unicode_pairs) {
+        byte_decoder[pair.second] = pair.first;
+    }
+
+    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>(encoder.size());
+    for (auto& special_token : special_tokens) {
+        auto token           = utf8_to_utf32(special_token);
+        encoder[token]       = encoder_len;
+        decoder[encoder_len] = token;
+        encoder_len++;
+    }
+    encoder_len = static_cast<int>(encoder.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;
+}
+
+GPTOSSTokenizer::GPTOSSTokenizer(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
+    BOS_TOKEN = "<|startoftext|>";
+    UNK_TOKEN = "<|endoftext|>";
+    EOS_TOKEN = "<|endoftext|>";
+    PAD_TOKEN = "<|endoftext|>";
+
+    BOS_TOKEN_ID = 199998;
+    EOS_TOKEN_ID = 199999;
+    UNK_TOKEN_ID = 199999;
+    PAD_TOKEN_ID = 199999;
+
+    special_tokens = {
+        "<|startoftext|>",
+        "<|endoftext|>",
+        "<|reserved_200000|>",
+        "<|reserved_200001|>",
+        "<|return|>",
+        "<|constrain|>",
+        "<|reserved_200004|>",
+        "<|channel|>",
+        "<|start|>",
+        "<|end|>",
+        "<|message|>",
+        "<|reserved_200009|>",
+        "<|reserved_200010|>",
+        "<|reserved_200011|>",
+        "<|call|>",
+        "<|reserved_200013|>",
+        "<|reserved_200014|>",
+        "<|reserved_200015|>",
+        "<|reserved_200016|>",
+        "<|reserved_200017|>",
+        "<|endofprompt|>",
+    };
+
+    if (merges_utf8_str.size() > 0) {
+        load_from_merges(merges_utf8_str, vocab_utf8_str);
+    } else {
+        load_from_merges(load_gpt_oss_merges(), load_gpt_oss_vocab_json());
+    }
+}

src/tokenizers/gpt_oss_tokenizer.h

@@ -0,0 +1,16 @@
+#ifndef __SD_TOKENIZERS_GPT_OSS_TOKENIZER_H__
+#define __SD_TOKENIZERS_GPT_OSS_TOKENIZER_H__
+
+#include <string>
+
+#include "bpe_tokenizer.h"
+
+class GPTOSSTokenizer : public BPETokenizer {
+protected:
+    void load_from_merges(const std::string& merges_utf8_str, const std::string& vocab_utf8_str);
+
+public:
+    explicit GPTOSSTokenizer(const std::string& merges_utf8_str = "", const std::string& vocab_utf8_str = "");
+};
+
+#endif  // __SD_TOKENIZERS_GPT_OSS_TOKENIZER_H__

src/tokenizers/vocab/vocab.cpp

@@ -1,7 +1,11 @@
 #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"
+#include "gpt_oss_vocab.hpp"
 #include "mistral_merges.hpp"
 #include "mistral_vocab.hpp"
 #include "qwen_merges.hpp"
@@ -46,4 +50,24 @@ std::string load_gemma_merges() {
 std::string load_gemma_vocab_json() {
     std::string json_str(reinterpret_cast<const char*>(gemma_vocab_json_utf8_c_str), sizeof(gemma_vocab_json_utf8_c_str));
     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;
+}
+
+std::string load_gpt_oss_vocab_json() {
+    std::string json_str(reinterpret_cast<const char*>(gpt_oss_vocab_json_utf8_c_str), sizeof(gpt_oss_vocab_json_utf8_c_str));
+    return json_str;
 }

src/tokenizers/vocab/vocab.h

@@ -11,5 +11,9 @@ 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();
 
 #endif  // __SD_TOKENIZERS_VOCAB_VOCAB_H__

src/unet.hpp

@@ -2,6 +2,7 @@
 #define __UNET_HPP__
 
 #include "common_block.hpp"
+#include "diffusion_model.hpp"
 #include "model.h"
 
 /*==================================================== UnetModel =====================================================*/
@@ -599,7 +600,7 @@ public:
     }
 };
 
-struct UNetModelRunner : public GGMLRunner {
+struct UNetModelRunner : public DiffusionModelRunner {
     UnetModelBlock unet;
 
     UNetModelRunner(ggml_backend_t backend,
@@ -607,7 +608,7 @@ struct UNetModelRunner : public GGMLRunner {
                     const String2TensorStorage& tensor_storage_map,
                     const std::string prefix,
                     SDVersion version = VERSION_SD1)
-        : GGMLRunner(backend, params_backend), unet(version, tensor_storage_map) {
+        : DiffusionModelRunner(backend, params_backend, prefix), unet(version, tensor_storage_map) {
         unet.init(params_ctx, tensor_storage_map, prefix);
     }
 
@@ -615,7 +616,7 @@ struct UNetModelRunner : public GGMLRunner {
         return "unet";
     }
 
-    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
         unet.get_param_tensors(tensors, prefix);
     }
 
@@ -682,6 +683,23 @@ struct UNetModelRunner : public GGMLRunner {
         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);
+        const auto* extra = diffusion_extra_as<UNetDiffusionExtra>(diffusion_params);
+        static const std::vector<sd::Tensor<float>> empty_controls;
+        return compute(n_threads,
+                       *diffusion_params.x,
+                       *diffusion_params.timesteps,
+                       tensor_or_empty(diffusion_params.context),
+                       tensor_or_empty(diffusion_params.c_concat),
+                       tensor_or_empty(diffusion_params.y),
+                       extra->num_video_frames,
+                       extra->controls ? *extra->controls : empty_controls,
+                       extra->control_strength);
+    }
+
     void test() {
         ggml_init_params params;
         params.mem_size   = static_cast<size_t>(10 * 1024 * 1024);  // 10 MB

src/util.cpp

@@ -1,8 +1,10 @@
 #include "util.h"
 #include <algorithm>
+#include <cctype>
 #include <cmath>
 #include <codecvt>
 #include <cstdarg>
+#include <exception>
 #include <fstream>
 #include <locale>
 #include <regex>
@@ -23,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) {
@@ -406,6 +406,88 @@ std::vector<std::string> split_string(const std::string& str, char delimiter) {
     return result;
 }
 
+KeyValueArgs parse_key_value_args(const char* args, const char* context) {
+    KeyValueArgs pairs;
+
+    if (args == nullptr || args[0] == '\0') {
+        return pairs;
+    }
+
+    std::string raw(args);
+    size_t start = 0;
+    for (size_t pos = 0; pos <= raw.size(); ++pos) {
+        if (pos != raw.size() && raw[pos] != ',' && raw[pos] != ';') {
+            continue;
+        }
+
+        std::string token = trim(raw.substr(start, pos - start));
+        if (!token.empty()) {
+            size_t eq = token.find('=');
+            if (eq == std::string::npos) {
+                const char* log_context = context ? context : "key=value arg";
+                LOG_WARN("ignoring malformed %s '%s'", log_context, token.c_str());
+            } else {
+                std::string key   = trim(token.substr(0, eq));
+                std::string value = trim(token.substr(eq + 1));
+                pairs.emplace_back(std::move(key), std::move(value));
+            }
+        }
+
+        start = pos + 1;
+    }
+
+    return pairs;
+}
+
+KeyValueArgs parse_key_value_args(const std::string& args, const char* context) {
+    return parse_key_value_args(args.c_str(), context);
+}
+
+bool parse_strict_float(const std::string& text, float& value) {
+    try {
+        size_t consumed = 0;
+        float parsed    = std::stof(text, &consumed);
+        if (!trim(text.substr(consumed)).empty()) {
+            return false;
+        }
+        value = parsed;
+        return true;
+    } catch (const std::exception&) {
+        return false;
+    }
+}
+
+bool parse_strict_int(const std::string& text, int& value) {
+    try {
+        size_t consumed = 0;
+        int parsed      = std::stoi(text, &consumed);
+        if (!trim(text.substr(consumed)).empty()) {
+            return false;
+        }
+        value = parsed;
+        return true;
+    } catch (const std::exception&) {
+        return false;
+    }
+}
+
+bool parse_strict_bool(const std::string& text, bool& value) {
+    std::string lowered = trim(text);
+    std::transform(lowered.begin(), lowered.end(), lowered.begin(), [](unsigned char c) {
+        return static_cast<char>(std::tolower(c));
+    });
+
+    if (lowered == "1" || lowered == "true" || lowered == "yes" || lowered == "on") {
+        value = true;
+        return true;
+    }
+    if (lowered == "0" || lowered == "false" || lowered == "no" || lowered == "off") {
+        value = false;
+        return true;
+    }
+    return false;
+}
+
 static std::string build_progress_bar(int step, int steps) {
     std::string progress = "  |";
     int max_progress     = 50;
@@ -758,29 +840,135 @@ std::vector<std::pair<std::string, float>> parse_prompt_attention(const std::str
     return res;
 }
 
-// namespace is needed to avoid conflicts with ggml_backend_extend.hpp
-namespace ggml_cpu {
-#include "ggml-cpu.h"
+static size_t get_utf8_char_len(char c) {
+    unsigned char uc = static_cast<unsigned char>(c);
+    if ((uc & 0x80) == 0) {
+        return 1;
+    }
+    if ((uc & 0xE0) == 0xC0) {
+        return 2;
+    }
+    if ((uc & 0xF0) == 0xE0) {
+        return 3;
+    }
+    if ((uc & 0xF8) == 0xF0) {
+        return 4;
+    }
+    return 1;
 }
 
-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;
+static bool is_ascii_alpha(char c) {
+    return (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z');
+}
+
+static bool starts_with_at(const std::string& text, size_t pos, const std::string& needle) {
+    return pos + needle.size() <= text.size() && text.compare(pos, needle.size(), needle) == 0;
+}
+
+static bool is_word_internal_apostrophe(const std::string& text, size_t pos) {
+    return pos > 0 && pos + 1 < text.size() &&
+           is_ascii_alpha(text[pos - 1]) && is_ascii_alpha(text[pos + 1]);
+}
+
+static std::vector<std::pair<std::string, bool>> split_quotation(const std::string& text) {
+    static const std::vector<std::pair<std::string, std::string>> quote_pairs = {
+        {"'", "'"},
+        {"\"", "\""},
+        {"\xE2\x80\x98", "\xE2\x80\x99"},
+        {"\xE2\x80\x9C", "\xE2\x80\x9D"},
+    };
+
+    std::vector<std::pair<std::string, bool>> result;
+    size_t segment_start = 0;
+    size_t i             = 0;
+
+    auto push_segment = [&](size_t begin, size_t end, bool matched) {
+        if (end > begin) {
+            result.emplace_back(text.substr(begin, end - begin), matched);
+        }
+    };
+
+    while (i < text.size()) {
+        bool matched_quote = false;
+        for (const auto& quote_pair : quote_pairs) {
+            const std::string& open_quote  = quote_pair.first;
+            const std::string& close_quote = quote_pair.second;
+            if (!starts_with_at(text, i, open_quote)) {
+                continue;
+            }
+            if (open_quote == "'" && is_word_internal_apostrophe(text, i)) {
+                continue;
+            }
+
+            size_t search_pos = i + open_quote.size();
+            size_t close_pos  = std::string::npos;
+            bool invalid      = false;
+            while (search_pos < text.size()) {
+                if (open_quote != close_quote && starts_with_at(text, search_pos, open_quote)) {
+                    invalid = true;
+                    break;
+                }
+                if (starts_with_at(text, search_pos, close_quote)) {
+                    if (close_quote == "'" && is_word_internal_apostrophe(text, search_pos)) {
+                        search_pos += close_quote.size();
+                        continue;
+                    }
+                    close_pos = search_pos;
+                    break;
+                }
+
+                size_t char_len = get_utf8_char_len(text[search_pos]);
+                if (search_pos + char_len > text.size()) {
+                    char_len = 1;
+                }
+                search_pos += char_len;
+            }
+            if (invalid || close_pos == std::string::npos) {
+                continue;
+            }
+
+            size_t quote_start = i;
+            push_segment(segment_start, quote_start, false);
+            i = close_pos + close_quote.size();
+            push_segment(quote_start, i, true);
+            segment_start = i;
+            matched_quote = true;
+            break;
+        }
+        if (!matched_quote) {
+            size_t char_len = get_utf8_char_len(text[i]);
+            if (i + char_len > text.size()) {
+                char_len = 1;
+            }
+            i += char_len;
+        }
+    }
+
+    push_segment(segment_start, text.size(), false);
+    return result;
+}
+
+std::vector<std::pair<std::string, float>> split_quotation_attention(
+    const std::vector<std::pair<std::string, float>>& parsed_attention) {
+    std::vector<std::pair<std::string, float>> result;
+    for (const auto& item : parsed_attention) {
+        const std::string& text = item.first;
+        float weight            = item.second;
+        for (const auto& part : split_quotation(text)) {
+            if (part.second) {
+                size_t i = 0;
+                while (i < part.first.size()) {
+                    size_t char_len = get_utf8_char_len(part.first[i]);
+                    if (i + char_len > part.first.size()) {
+                        char_len = 1;
+                    }
+                    result.emplace_back(part.first.substr(i, char_len), weight);
+                    i += char_len;
+                }
+            } else {
+                result.emplace_back(part.first, weight);
+            }
+        }
+    }
+    return result;
 }

src/util.h

@@ -4,6 +4,7 @@
 #include <cstdint>
 #include <memory>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include "ggml-backend.h"
@@ -65,6 +66,15 @@ protected:
 
 std::string path_join(const std::string& p1, const std::string& p2);
 std::vector<std::string> split_string(const std::string& str, char delimiter);
+
+using KeyValueArgs = std::vector<std::pair<std::string, std::string>>;
+
+KeyValueArgs parse_key_value_args(const char* args, const char* context = "key=value arg");
+KeyValueArgs parse_key_value_args(const std::string& args, const char* context = "key=value arg");
+bool parse_strict_float(const std::string& text, float& value);
+bool parse_strict_int(const std::string& text, int& value);
+bool parse_strict_bool(const std::string& text, bool& value);
+
 void pretty_progress(int step, int steps, float time);
 void pretty_bytes_progress(int step, int steps, uint64_t bytes_processed, float elapsed_seconds);
 
@@ -73,6 +83,8 @@ void log_printf(sd_log_level_t level, const char* file, int line, const char* fo
 std::string trim(const std::string& s);
 
 std::vector<std::pair<std::string, float>> parse_prompt_attention(const std::string& text);
+std::vector<std::pair<std::string, float>> split_quotation_attention(
+    const std::vector<std::pair<std::string, float>>& parsed_attention);
 
 sd_progress_cb_t sd_get_progress_callback();
 void* sd_get_progress_callback_data();

src/wan.hpp

@@ -6,6 +6,7 @@
 #include <utility>
 
 #include "common_block.hpp"
+#include "diffusion_model.hpp"
 #include "flux.hpp"
 #include "rope.hpp"
 #include "vae.hpp"
@@ -1327,13 +1328,16 @@ namespace WAN {
 
         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<WanVAERunner> vae = std::make_shared<WanVAERunner>(backend, backend, String2TensorStorage{}, "", false, VERSION_WAN2_2_TI2V);
             {
                 LOG_INFO("loading from '%s'", file_path.c_str());
 
-                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");
 
@@ -2085,7 +2089,7 @@ namespace WAN {
         }
     };
 
-    struct WanRunner : public GGMLRunner {
+    struct WanRunner : public DiffusionModelRunner {
     public:
         std::string desc = "wan";
         WanParams wan_params;
@@ -2098,7 +2102,7 @@ namespace WAN {
                   const String2TensorStorage& tensor_storage_map = {},
                   const std::string prefix                       = "",
                   SDVersion version                              = VERSION_WAN2)
-            : GGMLRunner(backend, params_backend) {
+            : DiffusionModelRunner(backend, params_backend, prefix) {
             wan_params.num_layers = 0;
             for (auto pair : tensor_storage_map) {
                 std::string tensor_name = pair.first;
@@ -2208,7 +2212,7 @@ namespace WAN {
             return desc;
         }
 
-        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
             wan.get_param_tensors(tensors, prefix);
         }
 
@@ -2284,6 +2288,22 @@ namespace WAN {
             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);
+            const auto* extra = diffusion_extra_as<WanDiffusionExtra>(diffusion_params);
+            return compute(n_threads,
+                           *diffusion_params.x,
+                           *diffusion_params.timesteps,
+                           tensor_or_empty(diffusion_params.context),
+                           tensor_or_empty(diffusion_params.y),
+                           tensor_or_empty(diffusion_params.c_concat),
+                           sd::Tensor<float>(),
+                           tensor_or_empty(extra->vace_context),
+                           extra->vace_strength);
+        }
+
         void test() {
             ggml_init_params params;
             params.mem_size   = static_cast<size_t>(200 * 1024 * 1024);  // 200 MB
@@ -2328,7 +2348,7 @@ namespace WAN {
 
         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;
             LOG_INFO("loading from '%s'", file_path.c_str());
 
@@ -2351,7 +2371,11 @@ namespace WAN {
                                                                          "model.diffusion_model",
                                                                          VERSION_WAN2_2_TI2V);
 
-            wan->alloc_params_buffer();
+            if (!wan->alloc_params_buffer()) {
+                LOG_ERROR("wan buffer allocation failed");
+                return;
+            }
+
             std::map<std::string, ggml_tensor*> tensors;
             wan->get_param_tensors(tensors, "model.diffusion_model");
 

src/z_image.hpp

@@ -3,6 +3,7 @@
 
 #include <algorithm>
 
+#include "diffusion_model.hpp"
 #include "flux.hpp"
 #include "ggml_extend.hpp"
 #include "mmdit.hpp"
@@ -464,7 +465,7 @@ namespace ZImage {
         }
     };
 
-    struct ZImageRunner : public GGMLRunner {
+    struct ZImageRunner : public DiffusionModelRunner {
     public:
         ZImageParams z_image_params;
         ZImageModel z_image;
@@ -477,7 +478,7 @@ namespace ZImage {
                      const String2TensorStorage& tensor_storage_map = {},
                      const std::string prefix                       = "",
                      SDVersion version                              = VERSION_Z_IMAGE)
-            : GGMLRunner(backend, params_backend) {
+            : DiffusionModelRunner(backend, params_backend, prefix) {
             z_image = ZImageModel(z_image_params);
             z_image.init(params_ctx, tensor_storage_map, prefix);
         }
@@ -486,7 +487,7 @@ namespace ZImage {
             return "z_image";
         }
 
-        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
             z_image.get_param_tensors(tensors, prefix);
         }
 
@@ -556,6 +557,19 @@ namespace ZImage {
             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);
+            static const std::vector<sd::Tensor<float>> empty_ref_latents;
+            return compute(n_threads,
+                           *diffusion_params.x,
+                           *diffusion_params.timesteps,
+                           tensor_or_empty(diffusion_params.context),
+                           diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
+                           diffusion_params.increase_ref_index);
+        }
+
         void test() {
             ggml_init_params params;
             params.mem_size   = static_cast<size_t>(1024 * 1024) * 1024;  // 1GB
@@ -601,7 +615,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;
@@ -625,7 +639,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");