feat: add ltx2.3 flf2v support (#1505)

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

.github/pull_request_template.md

@@ -0,0 +1,15 @@
+## Summary
+
+<!-- Describe what changed and why. Keep the PR focused on one clear change. -->
+
+## Related Issue / Discussion
+
+<!-- Link related issues, discussions, or previous PRs if applicable. -->
+
+## Additional Information
+
+<!-- Add verification notes, screenshots, sample output, or other context when applicable. -->
+
+## Checklist
+
+- [ ] I have read and confirmed this PR follows the [contribution guidelines](https://github.com/leejet/stable-diffusion.cpp/blob/master/CONTRIBUTING.md).

.github/workflows/build.yml

@@ -135,7 +135,7 @@ jobs:
         id: depends
         run: |
           sudo apt-get update
-          sudo apt-get install build-essential libvulkan-dev glslc
+          sudo apt-get install build-essential libvulkan-dev glslc spirv-headers
 
       - name: Build
         id: cmake_build
@@ -444,12 +444,95 @@ jobs:
           path: |
             sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-${{ matrix.build }}-x64.zip
 
+  windows-latest-rocm:
+    runs-on: windows-2022
+
+    env:
+      ROCM_VERSION: "7.13.0"
+      GPU_TARGETS: "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1102;gfx1150;gfx1151;gfx1200;gfx1201"
+
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          submodules: recursive
+
+      - name: Cache ROCm Installation
+        id: cache-rocm
+        uses: actions/cache@v4
+        with:
+          path: C:\TheRock\build
+          key: rocm-${{ env.ROCM_VERSION }}-gfx1151-${{ runner.os }}
+
+      - name: ccache
+        uses: ggml-org/ccache-action@v1.2.16
+        with:
+          key: windows-latest-rocm-${{ env.ROCM_VERSION }}-x64
+          evict-old-files: 1d
+
+      - name: Install ROCm
+        if: steps.cache-rocm.outputs.cache-hit != 'true'
+        run: |
+          $ErrorActionPreference = "Stop"
+          write-host "Downloading AMD ROCm ${{ env.ROCM_VERSION }} tarball"
+          Invoke-WebRequest -Uri "https://repo.amd.com/rocm/tarball/therock-dist-windows-gfx1151-${{ env.ROCM_VERSION }}.tar.gz" -OutFile "${env:RUNNER_TEMP}\rocm.tar.gz"
+          write-host "Extracting ROCm tarball"
+          mkdir C:\TheRock\build -Force
+          tar -xzf "${env:RUNNER_TEMP}\rocm.tar.gz" -C C:\TheRock\build --strip-components=1
+          write-host "Completed ROCm extraction"
+
+      - name: Setup ROCm Environment
+        run: |
+          $rocmPath = "C:\TheRock\build"
+          echo "HIP_PATH=$rocmPath" >> $env:GITHUB_ENV
+          echo "HIP_DEVICE_LIB_PATH=$rocmPath\lib\llvm\amdgcn\bitcode" >> $env:GITHUB_ENV
+          echo "HIP_PLATFORM=amd" >> $env:GITHUB_ENV
+          echo "LLVM_PATH=$rocmPath\lib\llvm" >> $env:GITHUB_ENV
+          echo "$rocmPath\bin" >> $env:GITHUB_PATH
+          echo "$rocmPath\lib\llvm\bin" >> $env:GITHUB_PATH
+
+      - name: Build
+        run: |
+          mkdir build
+          cd build
+          cmake .. `
+            -G "Unix Makefiles" `
+            -DCMAKE_PREFIX_PATH="${env:HIP_PATH}" `
+            -DSD_HIPBLAS=ON `
+            -DSD_BUILD_SHARED_LIBS=ON `
+            -DGGML_NATIVE=OFF `
+            -DCMAKE_C_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang.exe" `
+            -DCMAKE_CXX_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang++.exe" `
+            -DCMAKE_HIP_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang.exe" `
+            -DHIP_PATH="${env:HIP_PATH}" `
+            -DCMAKE_BUILD_TYPE=Release `
+            -DGPU_TARGETS="${{ env.GPU_TARGETS }}"
+          cmake --build . --config Release --parallel ${env:NUMBER_OF_PROCESSORS}
+
+      - name: Get commit hash
+        id: commit
+        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
+        uses: pr-mpt/actions-commit-hash@v2
+
+      - name: Pack artifacts
+        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
+        run: |
+          7z a sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-rocm-${{ env.ROCM_VERSION }}-x64.zip .\build\bin\*
+
+      - name: Upload artifacts
+        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
+        uses: actions/upload-artifact@v4
+        with:
+          name: sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-rocm-${{ env.ROCM_VERSION }}-x64.zip
+          path: |
+            sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-rocm-${{ env.ROCM_VERSION }}-x64.zip
+
   windows-latest-cmake-hip:
     runs-on: windows-2022
 
     env:
-      HIPSDK_INSTALLER_VERSION: "25.Q3"
-      GPU_TARGETS: "gfx1151;gfx1200;gfx1201;gfx1100;gfx1101;gfx1102;gfx1030;gfx1031;gfx1032"
+      HIPSDK_INSTALLER_VERSION: "26.Q1"
+      ROCM_VERSION: "7.1.1"
+      GPU_TARGETS: "gfx1150;gfx1151;gfx1200;gfx1201;gfx1100;gfx1101;gfx1102;gfx1030;gfx1031;gfx1032"
 
     steps:
       - uses: actions/checkout@v3
@@ -484,7 +567,7 @@ jobs:
         run: |
           $ErrorActionPreference = "Stop"
           write-host "Downloading AMD HIP SDK Installer"
-          Invoke-WebRequest -Uri "https://download.amd.com/developer/eula/rocm-hub/AMD-Software-PRO-Edition-${{ env.HIPSDK_INSTALLER_VERSION }}-WinSvr2022-For-HIP.exe" -OutFile "${env:RUNNER_TEMP}\rocm-install.exe"
+          Invoke-WebRequest -Uri "https://download.amd.com/developer/eula/rocm-hub/AMD-Software-PRO-Edition-${{ env.HIPSDK_INSTALLER_VERSION }}-Win11-For-HIP.exe" -OutFile "${env:RUNNER_TEMP}\rocm-install.exe"
           write-host "Installing AMD HIP SDK"
           $proc = Start-Process "${env:RUNNER_TEMP}\rocm-install.exe" -ArgumentList '-install' -NoNewWindow -PassThru
           $completed = $proc.WaitForExit(600000)
@@ -537,32 +620,38 @@ jobs:
         run: |
           md "build\bin\rocblas\library\"
           md "build\bin\hipblaslt\library"
-          cp "${env:HIP_PATH}\bin\hipblas.dll" "build\bin\"
-          cp "${env:HIP_PATH}\bin\hipblaslt.dll" "build\bin\"
+          cp "${env:HIP_PATH}\bin\libhipblas.dll" "build\bin\"
+          cp "${env:HIP_PATH}\bin\libhipblaslt.dll" "build\bin\"
           cp "${env:HIP_PATH}\bin\rocblas.dll" "build\bin\"
           cp "${env:HIP_PATH}\bin\rocblas\library\*" "build\bin\rocblas\library\"
           cp "${env:HIP_PATH}\bin\hipblaslt\library\*" "build\bin\hipblaslt\library\"
-          7z a sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-rocm-x64.zip .\build\bin\*
+          7z a sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-rocm-${{ env.ROCM_VERSION }}-x64.zip .\build\bin\*
 
       - name: Upload artifacts
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
         uses: actions/upload-artifact@v4
         with:
-          name: sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-rocm-x64.zip
+          name: sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-rocm-${{ env.ROCM_VERSION }}-x64.zip
           path: |
-            sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-rocm-x64.zip
+            sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-rocm-${{ env.ROCM_VERSION }}-x64.zip
 
   ubuntu-latest-rocm:
-    runs-on: ubuntu-latest
-    container: rocm/dev-ubuntu-24.04:7.2
+    runs-on: ubuntu-24.04
 
     env:
-      ROCM_VERSION: "7.2"
       UBUNTU_VERSION: "24.04"
-      GPU_TARGETS: "gfx1151;gfx1150;gfx1100;gfx1101;gfx1102;gfx1200;gfx1201"
+
+    strategy:
+      matrix:
+        include:
+          - ROCM_VERSION: "7.2.1"
+            gpu_targets: "gfx908;gfx90a;gfx942;gfx1030;gfx1031;gfx1032;gfx1100;gfx1101;gfx1102;gfx1151;gfx1150;gfx1200;gfx1201"
+            build: 'x64'
+          - ROCM_VERSION: "7.13.0"
+            gpu_targets: "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1102;gfx1150;gfx1151;gfx1200;gfx1201"
+            build: x64
 
     steps:
-      - run: apt-get update && apt-get install -y git
       - name: Clone
         id: checkout
         uses: actions/checkout@v6
@@ -579,6 +668,38 @@ jobs:
         with:
           version: 10.15.1
 
+      - name: ccache
+        uses: ggml-org/ccache-action@v1.2.16
+        with:
+          key: ubuntu-rocm-cmake-${{ matrix.ROCM_VERSION }}-${{ matrix.build }}
+          evict-old-files: 1d
+
+      - name: Dependencies
+        id: depends
+        run: |
+          sudo apt install -y build-essential cmake wget zip ninja-build
+
+      - name: Setup Legacy ROCm
+        if: matrix.ROCM_VERSION == '7.2.1'
+        id: legacy_env
+        run: |
+          sudo mkdir --parents --mode=0755 /etc/apt/keyrings
+          wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | \
+            gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null
+
+          sudo tee /etc/apt/sources.list.d/rocm.list << EOF
+          deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/${{ matrix.ROCM_VERSION }} noble main
+          EOF
+
+          sudo tee /etc/apt/preferences.d/rocm-pin-600 << EOF
+          Package: *
+          Pin: release o=repo.radeon.com
+          Pin-Priority: 600
+          EOF
+
+          sudo apt update
+          sudo apt-get install -y libssl-dev rocm-hip-sdk
+
       - name: Free disk space
         run: |
           # Remove preinstalled SDKs and caches not needed for this job
@@ -592,51 +713,17 @@ jobs:
           sudo rm -rf /var/lib/apt/lists/* || true
           sudo apt clean
 
-      - name: Dependencies
-        id: depends
+      - name: Setup TheRock
+        if: matrix.ROCM_VERSION != '7.2.1'
+        id: therock_env
         run: |
-          sudo apt-get update
-          sudo apt install -y \
-            cmake \
-            hip-dev \
-            hipblas-dev \
-            ninja-build \
-            rocm-dev \
-            zip
-          # Clean apt caches to recover disk space
-          sudo apt clean
-          sudo rm -rf /var/lib/apt/lists/* || true
-
-      - name: Setup ROCm Environment
-        run: |
-          # Add ROCm to PATH for current session
-          echo "/opt/rocm/bin" >> $GITHUB_PATH
-
-          # Build regex pattern from ${{ env.GPU_TARGETS }} (match target as substring)
-          TARGET_REGEX="($(printf '%s' "${{ env.GPU_TARGETS }}" | sed 's/;/|/g'))"
-
-          # Remove library files for architectures we're not building for to save disk space
-          echo "Cleaning up unneeded architecture files..."
-          cd /opt/rocm/lib/rocblas/library
-          # Keep only our target architectures
-          for file in *; do
-            if printf '%s' "$file" | grep -q 'gfx'; then
-              if ! printf '%s' "$file" | grep -Eq "$TARGET_REGEX"; then
-                echo "Removing $file" &&
-                sudo rm -f "$file";
-              fi
-            fi
-          done
-
-          cd /opt/rocm/lib/hipblaslt/library
-          for file in *; do
-            if printf '%s' "$file" | grep -q 'gfx'; then
-              if ! printf '%s' "$file" | grep -Eq "$TARGET_REGEX"; then
-                echo "Removing $file" &&
-                sudo rm -f "$file";
-              fi
-            fi
-          done
+          wget https://repo.amd.com/rocm/tarball/therock-dist-linux-gfx1151-${{ matrix.ROCM_VERSION }}.tar.gz
+          mkdir install
+          tar -xf *.tar.gz -C install
+          export ROCM_PATH=$(pwd)/install
+          echo ROCM_PATH=$ROCM_PATH >> $GITHUB_ENV
+          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
 
       - name: Build
         id: cmake_build
@@ -644,12 +731,12 @@ jobs:
           mkdir build
           cd build
           cmake .. -G Ninja \
-            -DCMAKE_CXX_COMPILER=amdclang++ \
-            -DCMAKE_C_COMPILER=amdclang \
+            -DCMAKE_HIP_COMPILER="$(hipconfig -l)/clang" \
+            -DCMAKE_HIP_FLAGS="-mllvm --amdgpu-unroll-threshold-local=600" \
             -DCMAKE_BUILD_TYPE=Release \
             -DSD_HIPBLAS=ON \
-            -DGPU_TARGETS="${{ env.GPU_TARGETS }}" \
-            -DAMDGPU_TARGETS="${{ env.GPU_TARGETS }}" \
+            -DHIP_PLATFORM=amd \
+            -DGPU_TARGETS="${{ matrix.gpu_targets }}" \
             -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \
             -DCMAKE_POSITION_INDEPENDENT_CODE=ON \
             -DSD_BUILD_SHARED_LIBS=ON
@@ -668,16 +755,6 @@ jobs:
           cp ggml/LICENSE ./build/bin/ggml.txt
           cp LICENSE ./build/bin/stable-diffusion.cpp.txt
 
-          # Move ROCm runtime libraries (to avoid double space consumption)
-          sudo mv /opt/rocm/lib/librocsparse.so* ./build/bin/
-          sudo mv /opt/rocm/lib/libhsa-runtime64.so* ./build/bin/
-          sudo mv /opt/rocm/lib/libamdhip64.so* ./build/bin/
-          sudo mv /opt/rocm/lib/libhipblas.so* ./build/bin/
-          sudo mv /opt/rocm/lib/libhipblaslt.so* ./build/bin/
-          sudo mv /opt/rocm/lib/librocblas.so* ./build/bin/
-          sudo mv /opt/rocm/lib/rocblas/ ./build/bin/
-          sudo mv /opt/rocm/lib/hipblaslt/ ./build/bin/
-
       - name: Fetch system info
         id: system-info
         run: |
@@ -692,15 +769,15 @@ jobs:
         run: |
           cp ggml/LICENSE ./build/bin/ggml.txt
           cp LICENSE ./build/bin/stable-diffusion.cpp.txt
-          zip -y -r sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-${{ steps.system-info.outputs.OS_TYPE }}-Ubuntu-${{ env.UBUNTU_VERSION }}-${{ steps.system-info.outputs.CPU_ARCH }}-rocm.zip ./build/bin
+          zip -y -r sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-${{ steps.system-info.outputs.OS_TYPE }}-Ubuntu-${{ env.UBUNTU_VERSION }}-${{ steps.system-info.outputs.CPU_ARCH }}-rocm-${{ matrix.ROCM_VERSION }}.zip ./build/bin
 
       - name: Upload artifacts
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
         uses: actions/upload-artifact@v4
         with:
-          name: sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-${{ steps.system-info.outputs.OS_TYPE }}-Ubuntu-${{ env.UBUNTU_VERSION }}-${{ steps.system-info.outputs.CPU_ARCH }}-rocm.zip
+          name: sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-${{ steps.system-info.outputs.OS_TYPE }}-Ubuntu-${{ env.UBUNTU_VERSION }}-${{ steps.system-info.outputs.CPU_ARCH }}-rocm-${{ matrix.ROCM_VERSION }}.zip
           path: |
-            sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-${{ steps.system-info.outputs.OS_TYPE }}-Ubuntu-${{ env.UBUNTU_VERSION }}-${{ steps.system-info.outputs.CPU_ARCH }}-rocm.zip
+            sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-${{ steps.system-info.outputs.OS_TYPE }}-Ubuntu-${{ env.UBUNTU_VERSION }}-${{ steps.system-info.outputs.CPU_ARCH }}-rocm-${{ matrix.ROCM_VERSION }}.zip
 
   release:
     if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
@@ -715,6 +792,7 @@ jobs:
       - macOS-latest-cmake
       - windows-latest-cmake
       - windows-latest-cmake-hip
+      - windows-latest-rocm
 
     steps:
       - name: Clone

.gitmodules

@@ -1,6 +1,6 @@
 [submodule "ggml"]
     path = ggml
-	url = https://github.com/ggml-org/ggml.git
+	url = https://github.com/leejet/ggml.git
 [submodule "examples/server/frontend"]
 	path = examples/server/frontend
 	url = https://github.com/leejet/sdcpp-webui.git

CMakeLists.txt

@@ -13,7 +13,9 @@ if (MSVC)
     add_compile_definitions(_SILENCE_CXX17_CODECVT_HEADER_DEPRECATION_WARNING)
     add_compile_options(
         $<$<COMPILE_LANGUAGE:C>:/MP>
+        $<$<COMPILE_LANGUAGE:C>:/utf-8>
         $<$<COMPILE_LANGUAGE:CXX>:/MP>
+        $<$<COMPILE_LANGUAGE:CXX>:/utf-8>
     )
 endif()
 
@@ -69,40 +71,40 @@ option(SD_BUILD_SHARED_GGML_LIB      "sd: build ggml as a separate shared lib" O
 option(SD_USE_SYSTEM_GGML            "sd: use system-installed GGML library" OFF)
 #option(SD_BUILD_SERVER               "sd: build server example"                           ON)
 
+set(CMAKE_C_STANDARD 11)
+set(CMAKE_C_STANDARD_REQUIRED true)
+
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED true)
+
 if(SD_CUDA)
     message("-- Use CUDA as backend stable-diffusion")
     set(GGML_CUDA ON)
-    add_definitions(-DSD_USE_CUDA)
 endif()
 
 if(SD_METAL)
     message("-- Use Metal as backend stable-diffusion")
     set(GGML_METAL ON)
-    add_definitions(-DSD_USE_METAL)
 endif()
 
 if (SD_VULKAN)
     message("-- Use Vulkan as backend stable-diffusion")
     set(GGML_VULKAN ON)
-    add_definitions(-DSD_USE_VULKAN)
 endif ()
 
 if (SD_OPENCL)
     message("-- Use OpenCL as backend stable-diffusion")
     set(GGML_OPENCL ON)
-    add_definitions(-DSD_USE_OPENCL)
 endif ()
 
 if (SD_HIPBLAS)
     message("-- Use HIPBLAS as backend stable-diffusion")
     set(GGML_HIP ON)
-    add_definitions(-DSD_USE_CUDA)
 endif ()
 
 if(SD_MUSA)
     message("-- Use MUSA as backend stable-diffusion")
     set(GGML_MUSA ON)
-    add_definitions(-DSD_USE_CUDA)
 endif()
 
 if(SD_WEBP)
@@ -112,19 +114,28 @@ if(SD_WEBP)
           "Or link against system library:\n  cmake (...) -DSD_USE_SYSTEM_WEBP=ON")
     endif()
     if(SD_USE_SYSTEM_WEBP)
-        find_package(WebP REQUIRED)
-        add_library(webp ALIAS WebP::webp)
-        # libwebp CMake target naming is not consistent across versions/distros.
-        # Some export WebP::libwebpmux, others export WebP::webpmux.
-        if(TARGET WebP::libwebpmux)
-            add_library(libwebpmux ALIAS WebP::libwebpmux)
-        elseif(TARGET WebP::webpmux)
-            add_library(libwebpmux ALIAS WebP::webpmux)
+        find_package(WebP)
+        if(WebP_FOUND)
+            add_library(webp ALIAS WebP::webp)
+            # libwebp CMake target naming is not consistent across versions/distros.
+            # Some export WebP::libwebpmux, others export WebP::webpmux.
+            if(TARGET WebP::libwebpmux)
+                add_library(libwebpmux ALIAS WebP::libwebpmux)
+            elseif(TARGET WebP::webpmux)
+                add_library(libwebpmux ALIAS WebP::webpmux)
+            else()
+                message(FATAL_ERROR
+                    "Could not find a compatible webpmux target in system WebP package. "
+                    "Expected WebP::libwebpmux or WebP::webpmux."
+                )
+            endif()
         else()
-            message(FATAL_ERROR
-                "Could not find a compatible webpmux target in system WebP package. "
-                "Expected WebP::libwebpmux or WebP::webpmux."
-            )
+            find_package(PkgConfig REQUIRED)
+            pkg_check_modules(WebP REQUIRED IMPORTED_TARGET GLOBAL libwebp)
+            pkg_check_modules(WebPMux REQUIRED IMPORTED_TARGET GLOBAL libwebpmux)
+            link_libraries(PkgConfig::WebP)
+            link_libraries(PkgConfig::WebPMux)
+            add_library(libwebpmux ALIAS PkgConfig::WebPMux)
         endif()
     endif()
 endif()
@@ -139,18 +150,26 @@ if(SD_WEBM)
           "Or link against system library:\n  cmake (...) -DSD_USE_SYSTEM_WEBM=ON")
     endif()
     if(SD_USE_SYSTEM_WEBM)
-        find_path(WEBM_INCLUDE_DIR
-            NAMES mkvmuxer/mkvmuxer.h mkvparser/mkvparser.h common/webmids.h
-            PATH_SUFFIXES webm
-            REQUIRED)
-        find_library(WEBM_LIBRARY
-            NAMES webm libwebm
-            REQUIRED)
+        find_package(PkgConfig)
+        if(PkgConfig_FOUND)
+            pkg_check_modules(WebM REQUIRED IMPORTED_TARGET GLOBAL libwebm)
+        endif()
+        if(PkgConfig_FOUND AND WebM_FOUND)
+            link_libraries(PkgConfig::WebM)
+        else()
+            find_path(WEBM_INCLUDE_DIR
+                NAMES mkvmuxer/mkvmuxer.h mkvparser/mkvparser.h common/webmids.h
+                PATH_SUFFIXES webm
+                REQUIRED)
+            find_library(WEBM_LIBRARY
+                NAMES webm libwebm
+                REQUIRED)
 
-        add_library(webm UNKNOWN IMPORTED)
-        set_target_properties(webm PROPERTIES
-            IMPORTED_LOCATION "${WEBM_LIBRARY}"
-            INTERFACE_INCLUDE_DIRECTORIES "${WEBM_INCLUDE_DIR}")
+            add_library(webm UNKNOWN IMPORTED)
+            set_target_properties(webm PROPERTIES
+                IMPORTED_LOCATION "${WEBM_LIBRARY}"
+                INTERFACE_INCLUDE_DIRECTORIES "${WEBM_INCLUDE_DIR}")
+        endif()
     endif()
 endif()
 
@@ -222,7 +241,6 @@ if(SD_SYCL)
     message("-- Use SYCL as backend stable-diffusion")
     set(GGML_SYCL ON)
     set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-narrowing -fsycl")
-    add_definitions(-DSD_USE_SYCL)
     # disable fast-math on host, see:
     # https://www.intel.com/content/www/us/en/docs/cpp-compiler/developer-guide-reference/2021-10/fp-model-fp.html
     if (WIN32)

CONTRIBUTING.md

@@ -0,0 +1,65 @@
+# Contributing
+
+This document collects general contribution conventions for this repository.
+
+## Before You Start
+
+Before opening a PR, please search existing PRs to avoid duplicating ongoing work.
+
+For large-scale refactors or changes with broad impact, please open an issue first to discuss the approach before submitting a PR.
+
+If you want to update a third-party dependency, please open an issue first instead of submitting a direct PR. See [Dependency Updates](#dependency-updates) for details.
+
+## Pull Requests
+
+Keep each PR focused on one clear change. Large or overly complex PRs are harder to review and may not be merged.
+
+Follow Conventional Commit-style subjects seen in history: `feat:`, `fix:`, `refactor:`, `ci:`, `docs:`, `chore:`. Keep subjects imperative and scoped.
+
+PRs should include:
+
+- What changed and why (short problem/solution summary).
+- Verification evidence when applicable (commands and key outputs).
+- Linked issue/PR context when applicable.
+- Screenshots or sample outputs for UI/visual behavior changes.
+
+## Code Style
+
+Format code according to the repository style before submitting changes.
+
+Formatting follows `.clang-format` (Chromium base, 4-space indent, no tabs). Run `format-code.sh` before opening a PR. Keep C++ standard at C++17-compatible patterns used in this repo.
+
+Naming conventions:
+
+- Use `PascalCase` for class/struct/type names.
+- In `PascalCase` names, preserve common abbreviations in uppercase, for example `SD`, `API`, `HTTP`, `JSON`, `RGB`, `VAE`, `TAE`, `LoRA`, and `WebP`.
+- Use `snake_case` for functions, methods, variables, and file names unless an existing API requires a different style.
+- Use a trailing underscore for private data member names, for example `hidden_size_` or `tokenizer_`.
+- Use `.h` for C and C++ header files. Do not introduce new `.hpp` headers.
+- Use macro-based header include guards instead of `#pragma once`.
+- Format header include guards as `__SD_{PATH}__`, where `{PATH}` is the header path in uppercase snake case without the file extension. For example, `src/sample.h` should use `__SD_SAMPLE_H__`.
+- Do not introduce anonymous namespaces in new or modified code; prefer `static` file-local functions/variables or an explicit named namespace when scoping is needed.
+- In `class`/`struct` definitions, place data members before member functions unless an existing type already clearly follows a different pattern.
+- Keep `test_*.cpp` / `test_*.py` naming for tests.
+
+Some older code in the project may not fully follow the current conventions. Please do not submit PRs that only rewrite existing code to match style rules.
+
+## AI-Assisted Contributions
+
+AI tools may be used to assist development, but contributors are responsible for the quality and correctness of the submitted code.
+
+If any part of a contribution was generated with AI assistance, the contributor must perform a thorough human review before submitting the PR and understand every changed line.
+
+Do not list AI tools as co-authors. The human contributor is the sole responsible author of the submitted code.
+
+Please do not submit AI-generated code that you do not understand, and do not include meaningless experiments, temporary test code, or unrelated generated output in a PR.
+
+## Dependency Updates
+
+Do not submit PRs that update `ggml`. `ggml` updates are performed only after local validation by the maintainer.
+
+Other third-party dependencies are not updated unless necessary. If you want to update a dependency, please open an issue first instead of submitting a direct PR.
+
+## Security & Configuration
+
+Do not commit model weights, secrets, or local absolute paths. Keep large binaries out of git unless intentionally tracked release assets.

Dockerfile.vulkan

@@ -2,7 +2,7 @@ ARG UBUNTU_VERSION=24.04
 
 FROM ubuntu:$UBUNTU_VERSION AS build
 
-RUN apt-get update && apt-get install -y --no-install-recommends build-essential git cmake libvulkan-dev glslc
+RUN apt-get update && apt-get install -y --no-install-recommends build-essential git cmake libvulkan-dev glslc spirv-headers
 
 WORKDIR /sd.cpp
 

README.md

@@ -58,11 +58,13 @@ API and command-line option may change frequently.***
     - [Ovis-Image](./docs/ovis_image.md)
     - [Anima](./docs/anima.md)
     - [ERNIE-Image](./docs/ernie_image.md)
+    - [HiDream-O1-Image](./docs/hidream_o1_image.md)
   - Image Edit Models
     - [FLUX.1-Kontext-dev](./docs/kontext.md)
     - [Qwen Image Edit series](./docs/qwen_image_edit.md)
   - Video Models
     - [Wan2.1/Wan2.2](./docs/wan.md)
+    - [LTX-2.3](./docs/ltx2.md)
   - [PhotoMaker](https://github.com/TencentARC/PhotoMaker) support.
   - Control Net support with SD 1.5
   - LoRA support, same as [stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#lora)
@@ -132,9 +134,11 @@ API and command-line option may change frequently.***
 ## Performance
 
 If you want to improve performance or reduce VRAM/RAM usage, please refer to [performance guide](./docs/performance.md).
+For runtime and parameter backend placement, see the [backend selection guide](./docs/backend.md).
 
 ## More Guides
 
+- [Backend selection](./docs/backend.md)
 - [SD1.x/SD2.x/SDXL](./docs/sd.md)
 - [SD3/SD3.5](./docs/sd3.md)
 - [FLUX.1-dev/FLUX.1-schnell](./docs/flux.md)
@@ -144,10 +148,12 @@ If you want to improve performance or reduce VRAM/RAM usage, please refer to [pe
 - [🔥Qwen Image](./docs/qwen_image.md)
 - [🔥Qwen Image Edit series](./docs/qwen_image_edit.md)
 - [🔥Wan2.1/Wan2.2](./docs/wan.md)
+- [🔥LTX-2.3](./docs/ltx2.md)
 - [🔥Z-Image](./docs/z_image.md)
 - [Ovis-Image](./docs/ovis_image.md)
 - [Anima](./docs/anima.md)
 - [ERNIE-Image](./docs/ernie_image.md)
+- [HiDream-O1-Image](./docs/hidream_o1_image.md)
 - [LoRA](./docs/lora.md)
 - [LCM/LCM-LoRA](./docs/lcm.md)
 - [Using PhotoMaker to personalize image generation](./docs/photo_maker.md)
@@ -163,6 +169,7 @@ These projects wrap `stable-diffusion.cpp` for easier use in other languages/fra
 
 * Golang (non-cgo): [seasonjs/stable-diffusion](https://github.com/seasonjs/stable-diffusion)
 * Golang (cgo): [Binozo/GoStableDiffusion](https://github.com/Binozo/GoStableDiffusion)
+* Golang (non-cgo): [l8bloom/gosd](https://github.com/l8bloom/gosd)
 * C#: [DarthAffe/StableDiffusion.NET](https://github.com/DarthAffe/StableDiffusion.NET)
 * Python: [william-murray1204/stable-diffusion-cpp-python](https://github.com/william-murray1204/stable-diffusion-cpp-python)
 * Rust: [newfla/diffusion-rs](https://github.com/newfla/diffusion-rs)

docs/backend.md

@@ -0,0 +1,122 @@
+# Backend selection
+
+`stable-diffusion.cpp` has two backend assignments:
+
+- `--backend` selects the runtime backend used to execute model graphs.
+- `--params-backend` selects the backend used to allocate model parameters.
+
+If `--params-backend` is not set, parameters use the same backend as their module runtime backend.
+
+## Syntax
+
+A backend assignment can be a single backend name:
+
+```shell
+sd-cli -m model.safetensors -p "a cat" --backend cpu
+```
+
+This applies to every module that does not have a more specific assignment.
+
+Assignments can also target individual modules:
+
+```shell
+sd-cli -m model.safetensors -p "a cat" --backend te=cpu,vae=cuda0,diffusion=vulkan0
+```
+
+The same syntax is used for parameter placement:
+
+```shell
+sd-cli -m model.safetensors -p "a cat" --backend cuda0 --params-backend te=cpu,vae=cpu
+```
+
+Module names are case-insensitive. Hyphens and underscores in module names are ignored, so `clip_vision`, `clip-vision`, and `clipvision` are equivalent.
+
+`all=`, `default=`, and `*=` can be used to set the default backend inside a mixed assignment:
+
+```shell
+sd-cli -m model.safetensors -p "a cat" --backend all=cuda0,te=cpu
+```
+
+## Modules
+
+| Module | Purpose | Accepted names |
+| --- | --- | --- |
+| `diffusion` | UNet, DiT, MMDiT, Flux, Wan, Qwen Image, and other diffusion models | `diffusion`, `model`, `unet`, `dit` |
+| `te` | Text encoders and conditioners | `te`, `clip`, `text`, `textencoder`, `textencoders`, `conditioner`, `cond`, `llm`, `t5`, `t5xxl` |
+| `clip_vision` | CLIP vision encoder | `clip_vision`, `clipvision`, `clip-vision`, `vision` |
+| `vae` | VAE and TAE | `vae`, `firststage`, `autoencoder`, `tae` |
+| `controlnet` | ControlNet | `controlnet`, `control` |
+| `photomaker` | PhotoMaker ID encoder and PhotoMaker LoRA | `photomaker`, `photomakerid`, `pmid`, `photo` |
+| `upscaler` | ESRGAN upscaler | `upscaler`, `esrgan`, `hires` |
+
+`te` is the preferred module name for text encoders. `clip` is kept as an accepted alias because many existing commands and model names use CLIP terminology.
+
+## Backend names
+
+Backend names are resolved against the GGML backend device list. Matching is case-insensitive and accepts exact names or unique prefixes, so common values include names such as:
+
+- `cpu`
+- `cuda0`
+- `vulkan0`
+- `metal`
+
+The special values `auto`, `default`, and an empty backend name select the default backend. The default preference is GPU, then integrated GPU, then CPU.
+
+The special value `gpu` selects the first GPU backend, falling back to the first integrated GPU backend.
+
+## Runtime backend vs. parameter backend
+
+The runtime backend controls where graph execution runs. The parameter backend controls where model weights are allocated.
+
+For example:
+
+```shell
+sd-cli -m model.safetensors -p "a cat" --backend cuda0 --params-backend cpu
+```
+
+This runs all modules on `cuda0`, but stores parameters in CPU RAM. During execution, parameters are moved to the runtime backend as needed.
+
+Per-module assignments can be mixed:
+
+```shell
+sd-cli -m model.safetensors -p "a cat" --backend diffusion=cuda0,te=cpu,vae=cpu --params-backend diffusion=cuda0,te=cpu,vae=cpu
+```
+
+This keeps text encoding and VAE execution on CPU while the diffusion model runs on GPU.
+
+## Backend sharing and lifetime
+
+Backends are managed by `SDBackendManager`.
+
+Within one manager, backend instances are cached by resolved backend device name. If multiple modules request the same backend, they share the same `ggml_backend_t`.
+
+For example:
+
+```shell
+--backend te=cpu,vae=cpu
+```
+
+uses one shared CPU backend for both `te` and `vae` runtime execution.
+
+Runtime and parameter assignments also share the same backend cache. If `--backend diffusion=cuda0` and `--params-backend diffusion=cuda0` resolve to the same device, both use the same backend instance.
+
+`SDBackendManager` owns the backend instances and frees them when the context or upscaler is destroyed. Model runners receive non-owning runtime and parameter backend pointers and do not free them.
+
+## Compatibility flags
+
+The older CPU placement flags are still supported:
+
+- `--clip-on-cpu`
+- `--vae-on-cpu`
+- `--control-net-cpu`
+- `--offload-to-cpu`
+
+`--clip-on-cpu`, `--vae-on-cpu`, and `--control-net-cpu` affect runtime backend assignment only when `--backend` is not set. They map to `te=cpu`, `vae=cpu`, and `controlnet=cpu`.
+
+`--offload-to-cpu` affects parameter backend assignment only when `--params-backend` is not set. It is equivalent to:
+
+```shell
+--params-backend cpu
+```
+
+Explicit `--backend` and `--params-backend` assignments are preferred for new commands.

docs/build.md

@@ -102,6 +102,11 @@ cmake --build . --config Release
 ## Build with Vulkan
 
 Install Vulkan SDK from https://www.lunarg.com/vulkan-sdk/.
+On Ubuntu, install the Vulkan development packages and SPIR-V headers:
+
+```shell
+sudo apt-get install build-essential libvulkan-dev glslc spirv-headers
+```
 
 ```shell
 mkdir build && cd build

docs/hidream_o1_image.md

@@ -0,0 +1,20 @@
+# How to Use
+
+## Download weights
+
+- Download HiDream-O1-Image-Dev
+    - safetensors: https://huggingface.co/Comfy-Org/HiDream-O1-Image/tree/main/checkpoints
+- Download HiDream-O1-Image
+    - safetensors: https://huggingface.co/Comfy-Org/HiDream-O1-Image/tree/main/checkpoints
+
+## Examples
+
+### HiDream-O1-Image-Dev
+
+```
+.\bin\Release\sd-cli.exe -m  ..\..\ComfyUI\models\diffusion_models\hidream_o1_image_dev_bf16.safetensors -p "a lovely cat holding a sign says 
+'hidream o1 cpp'" --cfg-scale 1.0  -v -H 1024 -W 1024
+```
+
+<img width="256" alt="HiDream-O1-Image-Dev example" src="../assets/hidream-o1/dev_example.png" />
+

docs/hipBLAS_on_Windows.md

@@ -26,12 +26,12 @@ Fortunately, `AMD` provides complete help documentation, you can use the help do
 
 Then we must set `ROCM` as environment variables before running cmake.
 
-Usually if you install according to the official tutorial and do not modify the ROCM path, then there is a high probability that it is here `C:\Program Files\AMD\ROCm\5.5\bin`
+Usually if you install according to the official tutorial and do not modify the ROCM path, then there is a high probability that it is here `C:\Program Files\AMD\ROCm\7.1.1\bin`
 
 This is what I use to set the clang:
 ```Commandline
-set CC=C:\Program Files\AMD\ROCm\5.5\bin\clang.exe
-set CXX=C:\Program Files\AMD\ROCm\5.5\bin\clang++.exe
+set CC=C:\Program Files\AMD\ROCm\7.1.1\bin\clang.exe
+set CXX=C:\Program Files\AMD\ROCm\7.1.1\bin\clang++.exe
 ```
 
 ## Ninja
@@ -46,7 +46,7 @@ set ninja=C:\Program Files\ninja\ninja.exe
 ## Building stable-diffusion.cpp
 
 The thing different from the regular CPU build is `-DSD_HIPBLAS=ON` ,
-`-G "Ninja"`, `-DCMAKE_C_COMPILER=clang`, `-DCMAKE_CXX_COMPILER=clang++`, `-DAMDGPU_TARGETS=gfx1100`
+`-G "Ninja"`, `-DCMAKE_C_COMPILER=clang`, `-DCMAKE_CXX_COMPILER=clang++`, `-DAMDGPU_TARGETS=gfx1150;gfx1151;gfx1200;gfx1201;gfx1100;gfx1101;gfx1102;gfx1030;gfx1031;gfx1032`
 
 >**Notice**: check the `clang` and `clang++` information:
 ```Commandline
@@ -59,26 +59,29 @@ If you see like this, we can continue:
 clang version 17.0.0 (git@github.amd.com:Compute-Mirrors/llvm-project e3201662d21c48894f2156d302276eb1cf47c7be)
 Target: x86_64-pc-windows-msvc
 Thread model: posix
-InstalledDir: C:\Program Files\AMD\ROCm\5.5\bin
+InstalledDir: C:\Program Files\AMD\ROCm\7.1.1\bin
 ```
 
 ```
 clang version 17.0.0 (git@github.amd.com:Compute-Mirrors/llvm-project e3201662d21c48894f2156d302276eb1cf47c7be)
 Target: x86_64-pc-windows-msvc
 Thread model: posix
-InstalledDir: C:\Program Files\AMD\ROCm\5.5\bin
+InstalledDir: C:\Program Files\AMD\ROCm\7.1.1\bin
 ```
 
->**Notice** that the `gfx1100` is the GPU architecture of my GPU, you can change it to your GPU architecture. Click here to see your architecture [LLVM Target](https://rocm.docs.amd.com/en/latest/release/windows_support.html#windows-supported-gpus)
+>**Notice** that the GPU targets are now compatible with multiple GPU architectures (ROCm 7.1.1 targets). You can change them to match your GPU architecture. Click here to see your architecture [LLVM Target](https://rocm.docs.amd.com/en/latest/release/windows_support.html#windows-supported-gpus)
 
-My GPU is AMD Radeon™ RX 7900 XTX Graphics, so I set it to `gfx1100`.
+Examples:
+- AMD Radeon™ RX 7900 XTX Graphics: `gfx1100`
+- AMD Radeon™ RX 7900 XT Graphics: `gfx1101`
+- AMD Radeon™ RX 7900 GRE Graphics: `gfx1102`
 
 option:
 
 ```commandline
 mkdir build
 cd build
-cmake .. -G "Ninja" -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DSD_HIPBLAS=ON -DCMAKE_BUILD_TYPE=Release -DAMDGPU_TARGETS=gfx1100
+cmake .. -G "Ninja" -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DSD_HIPBLAS=ON -DCMAKE_BUILD_TYPE=Release -DAMDGPU_TARGETS="gfx1150;gfx1151;gfx1200;gfx1201;gfx1100;gfx1101;gfx1102;gfx1030;gfx1031;gfx1032"
 cmake --build . --config Release
 ```
 

docs/ltx2.md

@@ -0,0 +1,53 @@
+# How to Use
+
+## Download weights
+
+- Download LTX-2.3
+    - safetensors: https://huggingface.co/Kijai/LTX2.3_comfy/tree/main/diffusion_models
+    - gguf: https://huggingface.co/unsloth/LTX-2.3-GGUF/tree/main
+- Download gemma-3-12b-it
+    - gguf: https://huggingface.co/unsloth/gemma-3-12b-it-GGUF/tree/main
+- Download embeddings connectors
+    - safetensors: https://huggingface.co/unsloth/LTX-2.3-GGUF/tree/main/text_encoders
+- Download vae
+    - safetensors: https://huggingface.co/unsloth/LTX-2.3-GGUF/tree/main/vae
+- Download audio vae
+    - safetensors: https://huggingface.co/unsloth/LTX-2.3-GGUF/tree/main/vae
+
+## Examples
+
+### LTX-2.3 dev T2V
+
+```
+.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model  ..\..\ComfyUI\models\diffusion_models\ltx-2.3-22b-dev-UD-Q4_K_M.gguf --vae ..\..\ComfyUI\models\vae\ltx-2.3-22b-dev_video_vae.safetensors --audio-vae ..\..\ComfyUI\models\vae\ltx-2.3-22b-dev_audio_vae.safetensors --llm ..\..\ComfyUI\models\text_encoders\gemma-3-12b-it-qat-UD-Q4_K_XL.gguf --embeddings-connectors ..\..\ComfyUI\models\text_encoders\ltx-2.3-22b-dev_embeddings_connectors.safetensors  -p "a lovely cat" --cfg-scale 6.0 --sampling-method euler -v -n "worst quality, low quality, blurry, distorted, artifacts" -W 1280 -H 720 --diffusion-fa --offload-to-cpu --video-frames 33 --fps 24 -o t2v.webm
+```
+
+<video
+  src="../assets/ltx2/t2v.webm"
+  controls
+  muted
+  style="max-width: 100%; height: auto;"></video>
+
+### LTX-2.3 dev I2V
+
+```
+.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model  ..\..\ComfyUI\models\diffusion_models\ltx-2.3-22b-dev-UD-Q4_K_M.gguf --vae ..\..\ComfyUI\models\vae\ltx-2.3-22b-dev_video_vae.safetensors --audio-vae ..\..\ComfyUI\models\vae\ltx-2.3-22b-dev_audio_vae.safetensors --llm ..\..\ComfyUI\models\text_encoders\gemma-3-12b-it-qat-UD-Q4_K_XL.gguf --embeddings-connectors ..\..\ComfyUI\models\text_encoders\ltx-2.3-22b-dev_embeddings_connectors.safetensors  -p "a lovely cat" --cfg-scale 6.0 --sampling-method euler -v  -W 1280 -H 720 --diffusion-fa --offload-to-cpu --video-frames 33 -i ..\assets\ernie_image\turbo_example.png -o i2v.webm
+```
+
+<video
+  src="../assets/ltx2/i2v.webm"
+  controls
+  muted
+  style="max-width: 100%; height: auto;"></video>
+
+### LTX-2.3 dev FLF2V
+
+```
+.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model  ..\..\ComfyUI\models\diffusion_models\ltx-2.3-22b-dev-UD-Q4_K_M.gguf --vae ..\..\ComfyUI\models\vae\ltx-2.3-22b-dev_video_vae.safetensors --audio-vae ..\..\ComfyUI\models\vae\ltx-2.3-22b-dev_audio_vae.safetensors --llm ..\..\ComfyUI\models\text_encoders\gemma-3-12b-it-qat-UD-Q4_K_XL.gguf --embeddings-connectors ..\..\ComfyUI\models\text_encoders\ltx-2.3-22b-dev_embeddings_connectors.safetensors  -p "glass flower blossom" --cfg-scale 6.0 --sampling-method euler -v  -W 1280 -H 720 --diffusion-fa --offload-to-cpu --video-frames 33 --init-img ..\..\ComfyUI\input\start_image.png --end-img ..\..\ComfyUI\input\end_image.png -o flf2v.webm
+```
+
+<video
+  src="../assets/ltx2/flf2v.webm"
+  controls
+  muted
+  style="max-width: 100%; height: auto;"></video>

docs/z_image.md

@@ -21,7 +21,7 @@ You can run Z-Image with stable-diffusion.cpp on GPUs with 4GB of VRAM — or ev
 ### Z-Image-Turbo
 
 ```
-.\bin\Release\sd-cli.exe --diffusion-model  z_image_turbo-Q3_K.gguf --vae ..\..\ComfyUI\models\vae\ae.sft  --llm ..\..\ComfyUI\models\text_encoders\Qwen3-4B-Instruct-2507-Q4_K_M.gguf -p "A cinematic, melancholic photograph of a solitary hooded figure walking through a sprawling, rain-slicked metropolis at night. The city lights are a chaotic blur of neon orange and cool blue, reflecting on the wet asphalt. The scene evokes a sense of being a single component in a vast machine. Superimposed over the image in a sleek, modern, slightly glitched font is the philosophical quote: 'THE CITY IS A CIRCUIT BOARD, AND I AM A BROKEN TRANSISTOR.' -- moody, atmospheric, profound, dark academic" --cfg-scale 1.0 -v --offload-to-cpu --diffusion-fa -H 1024 -W 512
+.\bin\Release\sd-cli.exe --diffusion-model  z_image_turbo-Q3_K.gguf --vae ..\..\ComfyUI\models\vae\ae.sft  --llm ..\..\ComfyUI\models\text_encoders\Qwen3-4B-Instruct-2507-Q4_K_M.gguf -p "A cinematic, melancholic photograph of a solitary hooded figure walking through a sprawling, rain-slicked metropolis at night. The city lights are a chaotic blur of neon orange and cool blue, reflecting on the wet asphalt. The scene evokes a sense of being a single component in a vast machine. Superimposed over the image in a sleek, modern, slightly glitched font is the philosophical quote: 'THE CITY IS A CIRCUIT BOARD, AND I AM A BROKEN TRANSISTOR.' -- moody, atmospheric, profound, dark academic" --cfg-scale 1.0 -v --offload-to-cpu --diffusion-fa -H 1024 -W 512 --steps 8
 ```
 
 <img width="256" alt="z-image example" src="../assets/z_image/q3_K.png" />

examples/cli/CMakeLists.txt

@@ -7,6 +7,10 @@ add_executable(${TARGET}
     image_metadata.cpp
     main.cpp
 )
+target_include_directories(${TARGET} PRIVATE
+    "${CMAKE_CURRENT_SOURCE_DIR}/.."
+    "${PROJECT_SOURCE_DIR}/src"
+)
 install(TARGETS ${TARGET} RUNTIME)
 target_link_libraries(${TARGET} PRIVATE stable-diffusion zip ${CMAKE_THREAD_LIBS_INIT})
 if(SD_WEBP)

examples/cli/README.md

@@ -54,6 +54,8 @@ Context Options:
   -t, --threads <int>                      number of threads to use during computation (default: -1). If threads <= 0,
                                            then threads will be set to the number of CPU physical cores
   --chroma-t5-mask-pad <int>               t5 mask pad size of chroma
+  --max-vram <float>                       maximum VRAM budget in GiB for graph-cut segmented execution. 0 disables
+                                           graph splitting; -1 auto-detects free VRAM minus 1 GiB
   --force-sdxl-vae-conv-scale              force use of conv scale on sdxl vae
   --offload-to-cpu                         place the weights in RAM to save VRAM, and automatically load them into VRAM
                                            when needed
@@ -101,6 +103,9 @@ 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
   -H, --height <int>                       image height, in pixel space (default: 512)
   -W, --width <int>                        image width, in pixel space (default: 512)
   --steps <int>                            number of sample steps (default: 20)
@@ -156,17 +161,18 @@ Generation Options:
   --disable-auto-resize-ref-image          disable auto resize of ref images
   --disable-image-metadata                 do not embed generation metadata on image files
   --vae-tiling                             process vae in tiles to reduce memory usage
+  --temporal-tiling                        enable temporal tiling for LTX video VAE decode
   --hires                                  enable highres fix
   -s, --seed                               RNG seed (default: 42, use random seed for < 0)
   --sampling-method                        sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m,
                                            dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s,
-                                           er_sde] (default: euler for Flux/SD3/Wan, euler_a otherwise)
+                                           er_sde, euler_cfg_pp, euler_a_cfg_pp] (default: euler for Flux/SD3/Wan, euler_a otherwise)
   --high-noise-sampling-method             (high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a,
                                            dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep,
-                                           res_2s, er_sde] default: euler for Flux/SD3/Wan, euler_a otherwise
+                                           res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp] default: euler for Flux/SD3/Wan, euler_a otherwise
   --scheduler                              denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits,
-                                           smoothstep, sgm_uniform, simple, kl_optimal, lcm, bong_tangent], default:
-                                           discrete
+                                           smoothstep, sgm_uniform, simple, kl_optimal, lcm, bong_tangent, ltx2], default:
+                                           model-specific
   --sigmas                                 custom sigma values for the sampler, comma-separated (e.g.,
                                            "14.61,7.8,3.5,0.0").
   --skip-layers                            layers to skip for SLG steps (default: [7,8,9])

examples/cli/main.cpp

@@ -385,11 +385,32 @@ std::string format_frame_idx(std::string pattern, int frame_idx) {
     return result;
 }
 
+static fs::path get_video_audio_sidecar_path(const SDCliParams& cli_params) {
+    fs::path out_path     = cli_params.output_path;
+    fs::path base_path    = out_path;
+    fs::path ext          = out_path.has_extension() ? out_path.extension() : fs::path{};
+    std::string ext_lower = ext.string();
+    std::transform(ext_lower.begin(), ext_lower.end(), ext_lower.begin(), ::tolower);
+    const EncodedImageFormat output_format = encoded_image_format_from_path(out_path.string());
+    if (!ext.empty()) {
+        if (output_format == EncodedImageFormat::JPEG ||
+            output_format == EncodedImageFormat::PNG ||
+            output_format == EncodedImageFormat::WEBP ||
+            ext_lower == ".avi" ||
+            ext_lower == ".webm") {
+            base_path.replace_extension();
+        }
+    }
+    base_path += ".wav";
+    return base_path;
+}
+
 bool save_results(const SDCliParams& cli_params,
                   const SDContextParams& ctx_params,
                   const SDGenerationParams& gen_params,
                   sd_image_t* results,
-                  int num_results) {
+                  int num_results,
+                  const sd_audio_t* generated_audio = nullptr) {
     if (results == nullptr || num_results <= 0) {
         return false;
     }
@@ -442,6 +463,21 @@ bool save_results(const SDCliParams& cli_params,
         return ok;
     };
 
+    auto write_audio_sidecar = [&](const fs::path& wav_path) {
+        if (generated_audio == nullptr) {
+            return;
+        }
+        if (write_wav_to_file(wav_path.string(),
+                              generated_audio->data,
+                              generated_audio->sample_count,
+                              generated_audio->channels,
+                              generated_audio->sample_rate)) {
+            LOG_INFO("save result audio to '%s'", wav_path.string().c_str());
+        } else {
+            LOG_WARN("failed to save result audio to '%s'", wav_path.string().c_str());
+        }
+    };
+
     int sucessful_reults = 0;
 
     if (std::regex_search(cli_params.output_path, format_specifier_regex)) {
@@ -465,8 +501,16 @@ bool save_results(const SDCliParams& cli_params,
             ext = ".avi";
         fs::path video_path = base_path;
         video_path += ext;
-        if (create_video_from_sd_images(video_path.string().c_str(), results, num_results, gen_params.fps) == 0) {
+        std::string final_ext_lower = ext.string();
+        std::transform(final_ext_lower.begin(), final_ext_lower.end(), final_ext_lower.begin(), ::tolower);
+        const bool mux_audio = generated_audio != nullptr && (final_ext_lower == ".avi" || final_ext_lower == ".webm");
+        if (create_video_from_sd_images(video_path.string().c_str(), results, num_results, gen_params.fps, 90, mux_audio ? generated_audio : nullptr) == 0) {
             LOG_INFO("save result video to '%s'", video_path.string().c_str());
+            if (generated_audio != nullptr && !mux_audio) {
+                fs::path wav_path = video_path;
+                wav_path.replace_extension(".wav");
+                write_audio_sidecar(wav_path);
+            }
             return true;
         } else {
             LOG_ERROR("Failed to save result video to '%s'", video_path.string().c_str());
@@ -488,6 +532,9 @@ bool save_results(const SDCliParams& cli_params,
         }
     }
     LOG_INFO("%d/%d images saved", sucessful_reults, num_results);
+    if (generated_audio != nullptr) {
+        write_audio_sidecar(get_video_audio_sidecar_path(cli_params));
+    }
     return sucessful_reults != 0;
 }
 
@@ -701,7 +748,8 @@ int main(int argc, const char* argv[]) {
     sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(vae_decode_only, true, cli_params.taesd_preview);
 
     SDImageVec results;
-    int num_results = 0;
+    int num_results             = 0;
+    sd_audio_t* generated_audio = nullptr;
 
     if (cli_params.mode == UPSCALE) {
         num_results = 1;
@@ -733,7 +781,10 @@ int main(int argc, const char* argv[]) {
             results.adopt(generate_image(sd_ctx.get(), &img_gen_params), num_results);
         } else if (cli_params.mode == VID_GEN) {
             sd_vid_gen_params_t vid_gen_params = gen_params.to_sd_vid_gen_params_t();
-            sd_image_t* generated_video        = generate_video(sd_ctx.get(), &vid_gen_params, &num_results);
+            sd_image_t* generated_video        = nullptr;
+            if (!generate_video(sd_ctx.get(), &vid_gen_params, &generated_video, &num_results, &generated_audio)) {
+                generated_video = nullptr;
+            }
             results.adopt(generated_video, num_results);
         }
 
@@ -749,7 +800,9 @@ int main(int argc, const char* argv[]) {
                                                      ctx_params.offload_params_to_cpu,
                                                      ctx_params.diffusion_conv_direct,
                                                      ctx_params.n_threads,
-                                                     gen_params.upscale_tile_size));
+                                                     gen_params.upscale_tile_size,
+                                                     ctx_params.backend.c_str(),
+                                                     ctx_params.params_backend.c_str()));
 
         if (upscaler_ctx == nullptr) {
             LOG_ERROR("new_upscaler_ctx failed");
@@ -773,9 +826,12 @@ int main(int argc, const char* argv[]) {
         }
     }
 
-    if (!save_results(cli_params, ctx_params, gen_params, results.data(), num_results)) {
+    if (!save_results(cli_params, ctx_params, gen_params, results.data(), num_results, generated_audio)) {
+        free_sd_audio(generated_audio);
         return 1;
     }
 
+    free_sd_audio(generated_audio);
+
     return 0;
 }

examples/common/common.cpp

@@ -340,10 +340,18 @@ ArgOptions SDContextParams::get_options() {
          "--high-noise-diffusion-model",
          "path to the standalone high noise diffusion model",
          &high_noise_diffusion_model_path},
+        {"",
+         "--embeddings-connectors",
+         "path to LTXAV embeddings connectors",
+         &embeddings_connectors_path},
         {"",
          "--vae",
          "path to standalone vae model",
          &vae_path},
+        {"",
+         "--audio-vae",
+         "path to standalone LTX audio vae model",
+         &audio_vae_path},
         {"",
          "--taesd",
          "path to taesd. Using Tiny AutoEncoder for fast decoding (low quality)",
@@ -380,6 +388,14 @@ ArgOptions SDContextParams::get_options() {
          "--upscale-model",
          "path to esrgan model.",
          &esrgan_path},
+        {"",
+         "--backend",
+         "runtime backend assignment, e.g. cpu or clip=cpu,vae=cuda0,diffusion=vulkan0",
+         &backend},
+        {"",
+         "--params-backend",
+         "parameter backend assignment, e.g. cpu or diffusion=cpu,clip=cpu",
+         &params_backend},
     };
 
     options.int_options = {
@@ -394,7 +410,12 @@ ArgOptions SDContextParams::get_options() {
          &chroma_t5_mask_pad},
     };
 
-    options.float_options = {};
+    options.float_options = {
+        {"",
+         "--max-vram",
+         "maximum VRAM budget in GiB for graph-cut segmented execution. 0 disables graph splitting; -1 auto-detects free VRAM minus 1 GiB",
+         &max_vram},
+    };
 
     options.bool_options = {
         {"",
@@ -656,7 +677,9 @@ std::string SDContextParams::to_string() const {
         << "  llm_vision_path: \"" << llm_vision_path << "\",\n"
         << "  diffusion_model_path: \"" << diffusion_model_path << "\",\n"
         << "  high_noise_diffusion_model_path: \"" << high_noise_diffusion_model_path << "\",\n"
+        << "  embeddings_connectors_path: \"" << embeddings_connectors_path << "\",\n"
         << "  vae_path: \"" << vae_path << "\",\n"
+        << "  audio_vae_path: \"" << audio_vae_path << "\",\n"
         << "  taesd_path: \"" << taesd_path << "\",\n"
         << "  esrgan_path: \"" << esrgan_path << "\",\n"
         << "  control_net_path: \"" << control_net_path << "\",\n"
@@ -670,6 +693,9 @@ std::string SDContextParams::to_string() const {
         << "  rng_type: " << sd_rng_type_name(rng_type) << ",\n"
         << "  sampler_rng_type: " << sd_rng_type_name(sampler_rng_type) << ",\n"
         << "  offload_params_to_cpu: " << (offload_params_to_cpu ? "true" : "false") << ",\n"
+        << "  max_vram: " << max_vram << ",\n"
+        << "  backend: \"" << backend << "\",\n"
+        << "  params_backend: \"" << params_backend << "\",\n"
         << "  enable_mmap: " << (enable_mmap ? "true" : "false") << ",\n"
         << "  control_net_cpu: " << (control_net_cpu ? "true" : "false") << ",\n"
         << "  clip_on_cpu: " << (clip_on_cpu ? "true" : "false") << ",\n"
@@ -712,7 +738,9 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool vae_decode_only, bool f
         llm_vision_path.c_str(),
         diffusion_model_path.c_str(),
         high_noise_diffusion_model_path.c_str(),
+        embeddings_connectors_path.c_str(),
         vae_path.c_str(),
+        audio_vae_path.c_str(),
         taesd_path.c_str(),
         control_net_path.c_str(),
         embedding_vec.data(),
@@ -744,6 +772,9 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool vae_decode_only, bool f
         chroma_use_t5_mask,
         chroma_t5_mask_pad,
         qwen_image_zero_cond_t,
+        max_vram,
+        backend.c_str(),
+        params_backend.c_str(),
     };
     return sd_ctx_params;
 }
@@ -800,6 +831,10 @@ ArgOptions SDGenerationParams::get_options() {
          "Latent (antialiased), Latent (bicubic), Latent (bicubic antialiased), or a model name "
          "under --hires-upscalers-dir (default: Latent)",
          &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",
+         &extra_sample_args},
     };
 
     options.int_options = {
@@ -983,6 +1018,11 @@ ArgOptions SDGenerationParams::get_options() {
          "process vae in tiles to reduce memory usage",
          true,
          &vae_tiling_params.enabled},
+        {"",
+         "--temporal-tiling",
+         "enable temporal tiling for LTX video VAE decode",
+         true,
+         &vae_tiling_params.temporal_tiling},
         {"",
          "--hires",
          "enable highres fix",
@@ -1237,17 +1277,17 @@ ArgOptions SDGenerationParams::get_options() {
          on_seed_arg},
         {"",
          "--sampling-method",
-         "sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde] "
+         "sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp]"
          "(default: euler for Flux/SD3/Wan, euler_a otherwise)",
          on_sample_method_arg},
         {"",
          "--high-noise-sampling-method",
-         "(high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde]"
+         "(high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp]"
          " default: euler for Flux/SD3/Wan, euler_a otherwise",
          on_high_noise_sample_method_arg},
         {"",
          "--scheduler",
-         "denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple, kl_optimal, lcm, bong_tangent], default: discrete",
+         "denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple, kl_optimal, lcm, bong_tangent, ltx2], default: model-specific",
          on_scheduler_arg},
         {"",
          "--sigmas",
@@ -1600,6 +1640,7 @@ bool SDGenerationParams::from_json_str(
 
     auto parse_sample_params_json = [&](const json& sample_json,
                                         sd_sample_params_t& target_params,
+                                        std::string& target_extra_sample_args,
                                         std::vector<int>& target_skip_layers,
                                         std::vector<float>* target_custom_sigmas) {
         if (sample_json.contains("sample_steps") && sample_json["sample_steps"].is_number_integer()) {
@@ -1614,6 +1655,9 @@ bool SDGenerationParams::from_json_str(
         if (sample_json.contains("flow_shift") && sample_json["flow_shift"].is_number()) {
             target_params.flow_shift = sample_json["flow_shift"];
         }
+        if (sample_json.contains("extra_sample_args") && sample_json["extra_sample_args"].is_string()) {
+            target_extra_sample_args = sample_json["extra_sample_args"].get<std::string>();
+        }
         if (target_custom_sigmas != nullptr &&
             sample_json.contains("custom_sigmas") &&
             sample_json["custom_sigmas"].is_array()) {
@@ -1661,11 +1705,12 @@ bool SDGenerationParams::from_json_str(
     };
 
     if (j.contains("sample_params") && j["sample_params"].is_object()) {
-        parse_sample_params_json(j["sample_params"], sample_params, skip_layers, &custom_sigmas);
+        parse_sample_params_json(j["sample_params"], sample_params, extra_sample_args, skip_layers, &custom_sigmas);
     }
     if (j.contains("high_noise_sample_params") && j["high_noise_sample_params"].is_object()) {
         parse_sample_params_json(j["high_noise_sample_params"],
                                  high_noise_sample_params,
+                                 high_noise_extra_sample_args,
                                  high_noise_skip_layers,
                                  nullptr);
     }
@@ -1675,6 +1720,9 @@ bool SDGenerationParams::from_json_str(
         if (tiling_json.contains("enabled") && tiling_json["enabled"].is_boolean()) {
             vae_tiling_params.enabled = tiling_json["enabled"];
         }
+        if (tiling_json.contains("temporal_tiling") && tiling_json["temporal_tiling"].is_boolean()) {
+            vae_tiling_params.temporal_tiling = tiling_json["temporal_tiling"];
+        }
         if (tiling_json.contains("tile_size_x") && tiling_json["tile_size_x"].is_number_integer()) {
             vae_tiling_params.tile_size_x = tiling_json["tile_size_x"];
         }
@@ -2092,6 +2140,8 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
     high_noise_sample_params.guidance.slg.layer_count = high_noise_skip_layers.size();
     sample_params.custom_sigmas                       = custom_sigmas.empty() ? nullptr : custom_sigmas.data();
     sample_params.custom_sigmas_count                 = static_cast<int>(custom_sigmas.size());
+    sample_params.extra_sample_args                   = extra_sample_args.empty() ? nullptr : extra_sample_args.c_str();
+    high_noise_sample_params.extra_sample_args        = high_noise_extra_sample_args.empty() ? nullptr : high_noise_extra_sample_args.c_str();
     cache_params.scm_mask                             = scm_mask.empty() ? nullptr : scm_mask.c_str();
 
     sd_pm_params_t pm_params = {
@@ -2161,6 +2211,8 @@ sd_vid_gen_params_t SDGenerationParams::to_sd_vid_gen_params_t() {
     high_noise_sample_params.guidance.slg.layer_count = high_noise_skip_layers.size();
     sample_params.custom_sigmas                       = custom_sigmas.empty() ? nullptr : custom_sigmas.data();
     sample_params.custom_sigmas_count                 = static_cast<int>(custom_sigmas.size());
+    sample_params.extra_sample_args                   = extra_sample_args.empty() ? nullptr : extra_sample_args.c_str();
+    high_noise_sample_params.extra_sample_args        = high_noise_extra_sample_args.empty() ? nullptr : high_noise_extra_sample_args.c_str();
     cache_params.scm_mask                             = scm_mask.empty() ? nullptr : scm_mask.c_str();
 
     params.loras                    = lora_vec.empty() ? nullptr : lora_vec.data();
@@ -2180,6 +2232,7 @@ sd_vid_gen_params_t SDGenerationParams::to_sd_vid_gen_params_t() {
     params.strength                 = strength;
     params.seed                     = seed;
     params.video_frames             = video_frames;
+    params.fps                      = fps;
     params.vace_strength            = vace_strength;
     params.vae_tiling_params        = vae_tiling_params;
     params.cache                    = cache_params;
@@ -2268,6 +2321,7 @@ std::string SDGenerationParams::to_string() const {
         << ", upscale_tile_size: " << hires_upscale_tile_size << " },\n"
         << "  vae_tiling_params: { "
         << vae_tiling_params.enabled << ", "
+        << vae_tiling_params.temporal_tiling << ", "
         << vae_tiling_params.tile_size_x << ", "
         << vae_tiling_params.tile_size_y << ", "
         << vae_tiling_params.target_overlap << ", "
@@ -2299,6 +2353,7 @@ static json build_sampling_metadata_json(const sd_sample_params_t& sample_params
         {"eta", sample_params.eta},
         {"shifted_timestep", sample_params.shifted_timestep},
         {"flow_shift", sample_params.flow_shift},
+        {"extra_sample_args", safe_json_string(sample_params.extra_sample_args)},
         {"guidance",
          {
              {"txt_cfg", sample_params.guidance.txt_cfg},
@@ -2490,6 +2545,9 @@ std::string get_image_params(const SDContextParams& ctx_params,
     }
     parameter_string += "Guidance: " + std::to_string(gen_params.sample_params.guidance.distilled_guidance) + ", ";
     parameter_string += "Eta: " + std::to_string(gen_params.sample_params.eta) + ", ";
+    if (!gen_params.extra_sample_args.empty()) {
+        parameter_string += "Extra sample args: " + gen_params.extra_sample_args + ", ";
+    }
     parameter_string += "Seed: " + std::to_string(seed) + ", ";
     parameter_string += "Size: " + std::to_string(gen_params.get_resolved_width()) + "x" + std::to_string(gen_params.get_resolved_height()) + ", ";
     parameter_string += "Model: " + sd_basename(ctx_params.model_path) + ", ";

examples/common/common.h

@@ -92,7 +92,9 @@ struct SDContextParams {
     std::string llm_vision_path;
     std::string diffusion_model_path;
     std::string high_noise_diffusion_model_path;
+    std::string embeddings_connectors_path;
     std::string vae_path;
+    std::string audio_vae_path;
     std::string taesd_path;
     std::string esrgan_path;
     std::string control_net_path;
@@ -109,14 +111,17 @@ struct SDContextParams {
     rng_type_t rng_type         = CUDA_RNG;
     rng_type_t sampler_rng_type = RNG_TYPE_COUNT;
     bool offload_params_to_cpu  = false;
-    bool enable_mmap            = false;
-    bool control_net_cpu        = false;
-    bool clip_on_cpu            = false;
-    bool vae_on_cpu             = false;
-    bool flash_attn             = false;
-    bool diffusion_flash_attn   = false;
-    bool diffusion_conv_direct  = false;
-    bool vae_conv_direct        = false;
+    float max_vram              = 0.f;
+    std::string backend;
+    std::string params_backend;
+    bool enable_mmap           = false;
+    bool control_net_cpu       = false;
+    bool clip_on_cpu           = false;
+    bool vae_on_cpu            = false;
+    bool flash_attn            = false;
+    bool diffusion_flash_attn  = false;
+    bool diffusion_conv_direct = false;
+    bool vae_conv_direct       = false;
 
     bool circular   = false;
     bool circular_x = false;
@@ -167,6 +172,8 @@ struct SDGenerationParams {
 
     sd_sample_params_t sample_params;
     sd_sample_params_t high_noise_sample_params;
+    std::string extra_sample_args;
+    std::string high_noise_extra_sample_args;
     std::vector<int> skip_layers            = {7, 8, 9};
     std::vector<int> high_noise_skip_layers = {7, 8, 9};
 
@@ -182,7 +189,7 @@ struct SDGenerationParams {
     int video_frames                     = 1;
     int fps                              = 16;
     float vace_strength                  = 1.f;
-    sd_tiling_params_t vae_tiling_params = {false, 0, 0, 0.5f, 0.0f, 0.0f};
+    sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f};
 
     std::string pm_id_images_dir;
     std::string pm_id_embed_path;

examples/common/media_io.cpp

@@ -613,6 +613,13 @@ typedef struct {
     uint32_t size;
 } avi_index_entry;
 
+typedef struct {
+    char fourcc[4];
+    uint32_t flags;
+    uint32_t offset;
+    uint32_t size;
+} avi_chunk_index_entry;
+
 void write_u32_le(FILE* f, uint32_t val) {
     fwrite(&val, 4, 1, f);
 }
@@ -647,6 +654,33 @@ void write_fourcc(std::vector<uint8_t>& data, const char* fourcc) {
     data.insert(data.end(), fourcc, fourcc + 4);
 }
 
+static std::vector<uint8_t> audio_to_pcm16_bytes(const sd_audio_t* audio) {
+    if (audio == nullptr || audio->data == nullptr || audio->sample_count == 0 || audio->channels == 0 || audio->sample_rate == 0) {
+        return {};
+    }
+
+    const size_t pcm_samples = static_cast<size_t>(audio->sample_count) * static_cast<size_t>(audio->channels);
+    std::vector<uint8_t> bytes(pcm_samples * sizeof(int16_t));
+    auto* pcm = reinterpret_cast<int16_t*>(bytes.data());
+    for (size_t i = 0; i < pcm_samples; ++i) {
+        const float sample = std::clamp(audio->data[i], -1.0f, 1.0f);
+        pcm[i]             = static_cast<int16_t>(std::lrint(sample * 32767.0f));
+    }
+    return bytes;
+}
+
+static std::pair<uint64_t, uint64_t> audio_sample_range_for_video_frame(const sd_audio_t* audio, int frame_idx, int num_frames, int fps) {
+    if (audio == nullptr || fps <= 0 || num_frames <= 0) {
+        return {0, 0};
+    }
+    const uint64_t total = audio->sample_count;
+    const uint64_t start = static_cast<uint64_t>((static_cast<long double>(frame_idx) * total) / num_frames);
+    const uint64_t end   = frame_idx + 1 == num_frames
+                               ? total
+                               : static_cast<uint64_t>((static_cast<long double>(frame_idx + 1) * total) / num_frames);
+    return {start, std::max(start, end)};
+}
+
 EncodedImageFormat encoded_image_format_from_path(const std::string& path) {
     std::string ext = fs::path(path).extension().string();
     std::transform(ext.begin(), ext.end(), ext.begin(), ::tolower);
@@ -776,7 +810,7 @@ uint8_t* load_image_from_memory(const char* image_bytes,
     return load_image_common(true, image_bytes, len, width, height, expected_width, expected_height, expected_channel);
 }
 
-std::vector<uint8_t> create_mjpg_avi_from_sd_images_to_vector(sd_image_t* images, int num_images, int fps, int quality) {
+std::vector<uint8_t> create_mjpg_avi_from_sd_images_to_vector(sd_image_t* images, int num_images, int fps, int quality, const sd_audio_t* audio) {
     if (num_images == 0) {
         fprintf(stderr, "Error: Image array is empty.\n");
         return {};
@@ -793,7 +827,13 @@ std::vector<uint8_t> create_mjpg_avi_from_sd_images_to_vector(sd_image_t* images
     // stb_image_write changes JPEG sampling behavior above quality 90.
     // MJPG AVI playback is more compatible when we keep the encoder on the
     // <= 90 path.
-    const int mjpg_quality = std::clamp(quality, 1, 90);
+    const int mjpg_quality               = std::clamp(quality, 1, 90);
+    const bool has_audio                 = audio != nullptr && audio->data != nullptr && audio->sample_count > 0 && audio->channels > 0 && audio->sample_rate > 0;
+    const std::vector<uint8_t> audio_pcm = audio_to_pcm16_bytes(audio);
+    const uint16_t audio_bits_per_sample = 16;
+    const uint16_t audio_block_align     = has_audio ? static_cast<uint16_t>(audio->channels * (audio_bits_per_sample / 8)) : 0;
+    const uint32_t audio_byte_rate       = has_audio ? static_cast<uint32_t>(audio->sample_rate * audio_block_align) : 0;
+    const uint32_t audio_data_size       = has_audio ? static_cast<uint32_t>(audio_pcm.size()) : 0;
 
     std::vector<uint8_t> avi_data;
     avi_data.reserve(static_cast<size_t>(num_images) * 1024);
@@ -804,7 +844,11 @@ std::vector<uint8_t> create_mjpg_avi_from_sd_images_to_vector(sd_image_t* images
     write_fourcc(avi_data, "AVI ");
 
     write_fourcc(avi_data, "LIST");
-    write_u32_le(avi_data, 4 + 8 + 56 + 8 + 4 + 8 + 56 + 8 + 40);
+    uint32_t hdrl_size = 4 + 8 + 56 + 8 + 4 + 8 + 56 + 8 + 40;
+    if (has_audio) {
+        hdrl_size += 8 + (4 + 8 + 56 + 8 + 16);
+    }
+    write_u32_le(avi_data, hdrl_size);
     write_fourcc(avi_data, "hdrl");
 
     write_fourcc(avi_data, "avih");
@@ -815,7 +859,7 @@ std::vector<uint8_t> create_mjpg_avi_from_sd_images_to_vector(sd_image_t* images
     write_u32_le(avi_data, 0x110);
     write_u32_le(avi_data, num_images);
     write_u32_le(avi_data, 0);
-    write_u32_le(avi_data, 1);
+    write_u32_le(avi_data, has_audio ? 2 : 1);
     write_u32_le(avi_data, width * height * 3);
     write_u32_le(avi_data, width);
     write_u32_le(avi_data, height);
@@ -862,12 +906,48 @@ std::vector<uint8_t> create_mjpg_avi_from_sd_images_to_vector(sd_image_t* images
     write_u32_le(avi_data, 0);
     write_u32_le(avi_data, 0);
 
+    if (has_audio) {
+        write_fourcc(avi_data, "LIST");
+        write_u32_le(avi_data, 4 + 8 + 56 + 8 + 16);
+        write_fourcc(avi_data, "strl");
+
+        write_fourcc(avi_data, "strh");
+        write_u32_le(avi_data, 56);
+        write_fourcc(avi_data, "auds");
+        write_u32_le(avi_data, 0);
+        write_u32_le(avi_data, 0);
+        write_u16_le(avi_data, 0);
+        write_u16_le(avi_data, 0);
+        write_u32_le(avi_data, 0);
+        write_u32_le(avi_data, audio_block_align);
+        write_u32_le(avi_data, audio_byte_rate);
+        write_u32_le(avi_data, 0);
+        write_u32_le(avi_data, static_cast<uint32_t>(audio->sample_count));
+        write_u32_le(avi_data, audio_data_size);
+        write_u32_le(avi_data, static_cast<uint32_t>(-1));
+        write_u32_le(avi_data, audio_block_align);
+        write_u16_le(avi_data, 0);
+        write_u16_le(avi_data, 0);
+        write_u16_le(avi_data, 0);
+        write_u16_le(avi_data, 0);
+
+        write_fourcc(avi_data, "strf");
+        write_u32_le(avi_data, 16);
+        write_u16_le(avi_data, 1);
+        write_u16_le(avi_data, static_cast<uint16_t>(audio->channels));
+        write_u32_le(avi_data, audio->sample_rate);
+        write_u32_le(avi_data, audio_byte_rate);
+        write_u16_le(avi_data, audio_block_align);
+        write_u16_le(avi_data, audio_bits_per_sample);
+    }
+
     write_fourcc(avi_data, "LIST");
     const size_t movi_size_pos = avi_data.size();
     write_u32_le(avi_data, 0);
     write_fourcc(avi_data, "movi");
 
-    std::vector<avi_index_entry> index(static_cast<size_t>(num_images));
+    std::vector<avi_chunk_index_entry> index;
+    index.reserve(static_cast<size_t>(num_images) + (has_audio ? 1 : 0));
     std::vector<uint8_t> jpeg_data;
 
     for (int i = 0; i < num_images; i++) {
@@ -884,27 +964,46 @@ std::vector<uint8_t> create_mjpg_avi_from_sd_images_to_vector(sd_image_t* images
             return {};
         }
 
-        index[i].offset = static_cast<uint32_t>(avi_data.size());
+        avi_chunk_index_entry video_entry = {};
+        memcpy(video_entry.fourcc, "00dc", 4);
+        video_entry.flags  = 0x10;
+        video_entry.offset = static_cast<uint32_t>(avi_data.size());
         write_fourcc(avi_data, "00dc");
         write_u32_le(avi_data, static_cast<uint32_t>(jpeg_data.size()));
-        index[i].size = (uint32_t)jpeg_data.size();
+        video_entry.size = static_cast<uint32_t>(jpeg_data.size());
         avi_data.insert(avi_data.end(), jpeg_data.begin(), jpeg_data.end());
+        index.push_back(video_entry);
 
         if (jpeg_data.size() % 2) {
             avi_data.push_back(0);
         }
     }
 
+    if (has_audio && !audio_pcm.empty()) {
+        avi_chunk_index_entry audio_entry = {};
+        memcpy(audio_entry.fourcc, "01wb", 4);
+        audio_entry.flags  = 0;
+        audio_entry.offset = static_cast<uint32_t>(avi_data.size());
+        audio_entry.size   = static_cast<uint32_t>(audio_pcm.size());
+        write_fourcc(avi_data, "01wb");
+        write_u32_le(avi_data, static_cast<uint32_t>(audio_pcm.size()));
+        avi_data.insert(avi_data.end(), audio_pcm.begin(), audio_pcm.end());
+        index.push_back(audio_entry);
+        if (audio_pcm.size() % 2 != 0) {
+            avi_data.push_back(0);
+        }
+    }
+
     const size_t movi_size = avi_data.size() - movi_size_pos - 4;
     patch_u32_le(avi_data, movi_size_pos, static_cast<uint32_t>(movi_size));
 
     write_fourcc(avi_data, "idx1");
-    write_u32_le(avi_data, num_images * 16);
-    for (int i = 0; i < num_images; i++) {
-        write_fourcc(avi_data, "00dc");
-        write_u32_le(avi_data, 0x10);
-        write_u32_le(avi_data, index[i].offset);
-        write_u32_le(avi_data, index[i].size);
+    write_u32_le(avi_data, static_cast<uint32_t>(index.size() * 16));
+    for (const auto& entry : index) {
+        write_fourcc(avi_data, entry.fourcc);
+        write_u32_le(avi_data, entry.flags);
+        write_u32_le(avi_data, entry.offset);
+        write_u32_le(avi_data, entry.size);
     }
 
     const size_t file_size = avi_data.size() - riff_size_pos - 4;
@@ -913,8 +1012,8 @@ std::vector<uint8_t> create_mjpg_avi_from_sd_images_to_vector(sd_image_t* images
     return avi_data;
 }
 
-int create_mjpg_avi_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality) {
-    std::vector<uint8_t> avi_data = create_mjpg_avi_from_sd_images_to_vector(images, num_images, fps, quality);
+int create_mjpg_avi_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality, const sd_audio_t* audio) {
+    std::vector<uint8_t> avi_data = create_mjpg_avi_from_sd_images_to_vector(images, num_images, fps, quality, audio);
     if (avi_data.empty()) {
         return -1;
     }
@@ -1044,7 +1143,7 @@ int create_animated_webp_from_sd_images(const char* filename, sd_image_t* images
 #endif
 
 #ifdef SD_USE_WEBM
-std::vector<uint8_t> create_webm_from_sd_images_to_vector(sd_image_t* images, int num_images, int fps, int quality) {
+std::vector<uint8_t> create_webm_from_sd_images_to_vector(sd_image_t* images, int num_images, int fps, int quality, const sd_audio_t* audio) {
     if (num_images == 0) {
         fprintf(stderr, "Error: Image array is empty.\n");
         return {};
@@ -1089,6 +1188,25 @@ std::vector<uint8_t> create_webm_from_sd_images_to_vector(sd_image_t* images, in
             video_track->set_display_height(static_cast<uint64_t>(height));
             video_track->set_frame_rate(static_cast<double>(fps));
         }
+
+        uint64_t audio_track_number    = 0;
+        std::vector<uint8_t> audio_pcm = audio_to_pcm16_bytes(audio);
+        if (audio != nullptr && !audio_pcm.empty()) {
+            audio_track_number = segment.AddAudioTrack(static_cast<int32_t>(audio->sample_rate), static_cast<int32_t>(audio->channels), 0);
+            if (audio_track_number == 0) {
+                fprintf(stderr, "Error: Failed to add audio track.\n");
+                return -1;
+            }
+            auto* audio_track = static_cast<mkvmuxer::AudioTrack*>(segment.GetTrackByNumber(audio_track_number));
+            if (audio_track == nullptr) {
+                fprintf(stderr, "Error: Failed to get audio track.\n");
+                return -1;
+            }
+            audio_track->set_codec_id("A_PCM/INT/LIT");
+            audio_track->set_bit_depth(16);
+            audio_track->set_sample_rate(static_cast<double>(audio->sample_rate));
+            audio_track->set_channels(audio->channels);
+        }
         segment.GetSegmentInfo()->set_writing_app("stable-diffusion.cpp");
         segment.GetSegmentInfo()->set_muxing_app("stable-diffusion.cpp");
 
@@ -1118,6 +1236,23 @@ std::vector<uint8_t> create_webm_from_sd_images_to_vector(sd_image_t* images, in
                 return -1;
             }
 
+            if (audio_track_number != 0) {
+                auto [audio_begin, audio_end] = audio_sample_range_for_video_frame(audio, i, num_images, fps);
+                const uint64_t frame_samples  = audio_end - audio_begin;
+                if (frame_samples > 0) {
+                    const uint64_t frame_bytes = frame_samples * audio->channels * sizeof(int16_t);
+                    const uint8_t* frame_ptr   = audio_pcm.data() + audio_begin * audio->channels * sizeof(int16_t);
+                    if (!segment.AddFrame(frame_ptr,
+                                          frame_bytes,
+                                          audio_track_number,
+                                          timestamp_ns,
+                                          true)) {
+                        fprintf(stderr, "Error: Failed to mux audio chunk %d into WebM.\n", i);
+                        return -1;
+                    }
+                }
+            }
+
             timestamp_ns += frame_duration_ns;
         }
 
@@ -1133,8 +1268,8 @@ std::vector<uint8_t> create_webm_from_sd_images_to_vector(sd_image_t* images, in
     return writer.data();
 }
 
-int create_webm_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality) {
-    std::vector<uint8_t> webm_data = create_webm_from_sd_images_to_vector(images, num_images, fps, quality);
+int create_webm_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality, const sd_audio_t* audio) {
+    std::vector<uint8_t> webm_data = create_webm_from_sd_images_to_vector(images, num_images, fps, quality, audio);
     if (webm_data.empty()) {
         return -1;
     }
@@ -1150,7 +1285,8 @@ std::vector<uint8_t> create_video_from_sd_images_to_vector(const std::string& ou
                                                            sd_image_t* images,
                                                            int num_images,
                                                            int fps,
-                                                           int quality) {
+                                                           int quality,
+                                                           const sd_audio_t* audio) {
     std::string format = output_format;
     std::transform(format.begin(), format.end(), format.begin(),
                    [](unsigned char c) { return static_cast<char>(tolower(c)); });
@@ -1160,7 +1296,7 @@ std::vector<uint8_t> create_video_from_sd_images_to_vector(const std::string& ou
 
 #ifdef SD_USE_WEBM
     if (format == "webm") {
-        return create_webm_from_sd_images_to_vector(images, num_images, fps, quality);
+        return create_webm_from_sd_images_to_vector(images, num_images, fps, quality, audio);
     }
 #endif
 
@@ -1170,14 +1306,14 @@ std::vector<uint8_t> create_video_from_sd_images_to_vector(const std::string& ou
     }
 #endif
 
-    return create_mjpg_avi_from_sd_images_to_vector(images, num_images, fps, quality);
+    return create_mjpg_avi_from_sd_images_to_vector(images, num_images, fps, quality, audio);
 }
 
-int create_video_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality) {
+int create_video_from_sd_images(const char* filename, sd_image_t* images, int num_images, int fps, int quality, const sd_audio_t* audio) {
     std::string path                = filename ? filename : "";
     auto pos                        = path.find_last_of('.');
     std::string ext                 = pos == std::string::npos ? "" : path.substr(pos);
-    std::vector<uint8_t> video_data = create_video_from_sd_images_to_vector(ext, images, num_images, fps, quality);
+    std::vector<uint8_t> video_data = create_video_from_sd_images_to_vector(ext, images, num_images, fps, quality, audio);
     if (video_data.empty()) {
         return -1;
     }
@@ -1187,3 +1323,54 @@ int create_video_from_sd_images(const char* filename, sd_image_t* images, int nu
     }
     return 0;
 }
+
+bool write_wav_to_file(const std::string& path,
+                       const float* interleaved_samples,
+                       uint64_t sample_count,
+                       uint32_t channels,
+                       uint32_t sample_rate) {
+    if (interleaved_samples == nullptr || sample_count == 0 || channels == 0 || sample_rate == 0) {
+        return false;
+    }
+
+    std::ofstream file(path, std::ios::binary);
+    if (!file.is_open()) {
+        return false;
+    }
+
+    uint32_t bits_per_sample  = 16;
+    uint32_t bytes_per_sample = bits_per_sample / 8;
+    uint32_t block_align      = channels * bytes_per_sample;
+    uint32_t byte_rate        = sample_rate * block_align;
+    uint32_t data_size        = static_cast<uint32_t>(sample_count * channels * bytes_per_sample);
+    uint32_t riff_size        = 36 + data_size;
+
+    file.write("RIFF", 4);
+    file.write(reinterpret_cast<const char*>(&riff_size), sizeof(riff_size));
+    file.write("WAVE", 4);
+    file.write("fmt ", 4);
+
+    uint32_t fmt_size            = 16;
+    uint16_t audio_format        = 1;
+    uint16_t wav_channels        = static_cast<uint16_t>(channels);
+    uint16_t wav_block_align     = static_cast<uint16_t>(block_align);
+    uint16_t wav_bits_per_sample = static_cast<uint16_t>(bits_per_sample);
+    file.write(reinterpret_cast<const char*>(&fmt_size), sizeof(fmt_size));
+    file.write(reinterpret_cast<const char*>(&audio_format), sizeof(audio_format));
+    file.write(reinterpret_cast<const char*>(&wav_channels), sizeof(wav_channels));
+    file.write(reinterpret_cast<const char*>(&sample_rate), sizeof(sample_rate));
+    file.write(reinterpret_cast<const char*>(&byte_rate), sizeof(byte_rate));
+    file.write(reinterpret_cast<const char*>(&wav_block_align), sizeof(wav_block_align));
+    file.write(reinterpret_cast<const char*>(&wav_bits_per_sample), sizeof(wav_bits_per_sample));
+
+    file.write("data", 4);
+    file.write(reinterpret_cast<const char*>(&data_size), sizeof(data_size));
+
+    std::vector<int16_t> pcm(sample_count * channels);
+    for (size_t i = 0; i < pcm.size(); ++i) {
+        float sample = std::max(-1.0f, std::min(1.0f, interleaved_samples[i]));
+        pcm[i]       = static_cast<int16_t>(std::lrint(sample * 32767.0f));
+    }
+    file.write(reinterpret_cast<const char*>(pcm.data()), static_cast<std::streamsize>(pcm.size() * sizeof(int16_t)));
+    return file.good();
+}

examples/common/media_io.h

@@ -57,11 +57,13 @@ int create_mjpg_avi_from_sd_images(const char* filename,
                                    sd_image_t* images,
                                    int num_images,
                                    int fps,
-                                   int quality = 90);
+                                   int quality             = 90,
+                                   const sd_audio_t* audio = nullptr);
 std::vector<uint8_t> create_mjpg_avi_from_sd_images_to_vector(sd_image_t* images,
                                                               int num_images,
                                                               int fps,
-                                                              int quality = 90);
+                                                              int quality             = 90,
+                                                              const sd_audio_t* audio = nullptr);
 
 #ifdef SD_USE_WEBP
 int create_animated_webp_from_sd_images(const char* filename,
@@ -80,22 +82,32 @@ int create_webm_from_sd_images(const char* filename,
                                sd_image_t* images,
                                int num_images,
                                int fps,
-                               int quality = 90);
+                               int quality             = 90,
+                               const sd_audio_t* audio = nullptr);
 std::vector<uint8_t> create_webm_from_sd_images_to_vector(sd_image_t* images,
                                                           int num_images,
                                                           int fps,
-                                                          int quality = 90);
+                                                          int quality             = 90,
+                                                          const sd_audio_t* audio = nullptr);
 #endif
 
 int create_video_from_sd_images(const char* filename,
                                 sd_image_t* images,
                                 int num_images,
                                 int fps,
-                                int quality = 90);
+                                int quality             = 90,
+                                const sd_audio_t* audio = nullptr);
 std::vector<uint8_t> create_video_from_sd_images_to_vector(const std::string& output_format,
                                                            sd_image_t* images,
                                                            int num_images,
                                                            int fps,
-                                                           int quality = 90);
+                                                           int quality             = 90,
+                                                           const sd_audio_t* audio = nullptr);
+
+bool write_wav_to_file(const std::string& path,
+                       const float* interleaved_samples,
+                       uint64_t sample_count,
+                       uint32_t channels,
+                       uint32_t sample_rate);
 
 #endif  // __MEDIA_IO_H__

examples/server/README.md

@@ -156,6 +156,8 @@ Context Options:
   -t, --threads <int>                      number of threads to use during computation (default: -1). If threads <= 0,
                                            then threads will be set to the number of CPU physical cores
   --chroma-t5-mask-pad <int>               t5 mask pad size of chroma
+  --max-vram <float>                       maximum VRAM budget in GiB for graph-cut segmented execution. 0 disables
+                                           graph splitting; -1 auto-detects free VRAM minus 1 GiB
   --force-sdxl-vae-conv-scale              force use of conv scale on sdxl vae
   --offload-to-cpu                         place the weights in RAM to save VRAM, and automatically load them into VRAM
                                            when needed
@@ -203,6 +205,9 @@ 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
   -H, --height <int>                       image height, in pixel space (default: 512)
   -W, --width <int>                        image width, in pixel space (default: 512)
   --steps <int>                            number of sample steps (default: 20)
@@ -262,13 +267,13 @@ Default Generation Options:
   -s, --seed                               RNG seed (default: 42, use random seed for < 0)
   --sampling-method                        sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m,
                                            dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s,
-                                           er_sde] (default: euler for Flux/SD3/Wan, euler_a otherwise)
+                                           er_sde, euler_cfg_pp, euler_a_cfg_pp] (default: euler for Flux/SD3/Wan, euler_a otherwise)
   --high-noise-sampling-method             (high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a,
                                            dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep,
-                                           res_2s, er_sde] default: euler for Flux/SD3/Wan, euler_a otherwise
+                                           res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp] default: euler for Flux/SD3/Wan, euler_a otherwise
   --scheduler                              denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits,
-                                           smoothstep, sgm_uniform, simple, kl_optimal, lcm, bong_tangent], default:
-                                           discrete
+                                           smoothstep, sgm_uniform, simple, kl_optimal, lcm, bong_tangent, ltx2], default:
+                                           model-specific
   --sigmas                                 custom sigma values for the sampler, comma-separated (e.g.,
                                            "14.61,7.8,3.5,0.0").
   --skip-layers                            layers to skip for SLG steps (default: [7,8,9])

examples/server/async_jobs.cpp

@@ -231,16 +231,21 @@ bool execute_vid_gen_job(ServerRuntime& runtime,
     sd_vid_gen_params_t params = job.vid_gen.to_sd_vid_gen_params_t();
 
     SDImageVec results;
-    int num_results = 0;
+    int num_results             = 0;
+    sd_audio_t* generated_audio = nullptr;
 
     {
         std::lock_guard<std::mutex> lock(*runtime.sd_ctx_mutex);
-        sd_image_t* raw_results = generate_video(runtime.sd_ctx, &params, &num_results);
+        sd_image_t* raw_results = nullptr;
+        if (!generate_video(runtime.sd_ctx, &params, &raw_results, &num_results, &generated_audio)) {
+            raw_results = nullptr;
+        }
         results.adopt(raw_results, num_results);
     }
 
     num_results = results.count();
     if (num_results <= 0) {
+        free_sd_audio(generated_audio);
         error_message = "generate_video returned no results";
         return false;
     }
@@ -249,7 +254,9 @@ bool execute_vid_gen_job(ServerRuntime& runtime,
                                                                              results.data(),
                                                                              num_results,
                                                                              job.vid_gen.gen_params.fps,
-                                                                             job.vid_gen.output_compression);
+                                                                             job.vid_gen.output_compression,
+                                                                             generated_audio);
+    free_sd_audio(generated_audio);
     if (video_bytes.empty()) {
         error_message = "failed to encode generated video container";
         return false;

examples/server/main.cpp

@@ -145,7 +145,7 @@ int main(int argc, const char** argv) {
     register_sdapi_endpoints(svr, runtime);
     register_sdcpp_api_endpoints(svr, runtime);
 
-    LOG_INFO("listening on: %s:%d\n", svr_params.listen_ip.c_str(), svr_params.listen_port);
+    LOG_INFO("listening on: http://%s:%d\n", svr_params.listen_ip.c_str(), svr_params.listen_port);
     svr.listen(svr_params.listen_ip, svr_params.listen_port);
 
     {

examples/server/routes_sdapi.cpp

@@ -67,6 +67,10 @@ static enum sample_method_t get_sdapi_sample_method(std::string name) {
         {"k_res_multistep", RES_MULTISTEP_SAMPLE_METHOD},
         {"res 2s", RES_2S_SAMPLE_METHOD},
         {"k_res_2s", RES_2S_SAMPLE_METHOD},
+        {"euler_cfg_pp", EULER_CFG_PP_SAMPLE_METHOD},
+        {"k_euler_cfg_pp", EULER_CFG_PP_SAMPLE_METHOD},
+        {"euler_a_cfg_pp", EULER_CFG_PP_SAMPLE_METHOD},
+        {"k_euler_a_cfg_pp", EULER_CFG_PP_SAMPLE_METHOD},
     };
     auto it = hardcoded.find(name);
     return it != hardcoded.end() ? it->second : SAMPLE_METHOD_COUNT;

include/stable-diffusion.h

@@ -51,6 +51,8 @@ enum sample_method_t {
     RES_MULTISTEP_SAMPLE_METHOD,
     RES_2S_SAMPLE_METHOD,
     ER_SDE_SAMPLE_METHOD,
+    EULER_CFG_PP_SAMPLE_METHOD,
+    EULER_A_CFG_PP_SAMPLE_METHOD,
     SAMPLE_METHOD_COUNT
 };
 
@@ -66,6 +68,7 @@ enum scheduler_t {
     KL_OPTIMAL_SCHEDULER,
     LCM_SCHEDULER,
     BONG_TANGENT_SCHEDULER,
+    LTX2_SCHEDULER,
     SCHEDULER_COUNT
 };
 
@@ -149,6 +152,7 @@ enum lora_apply_mode_t {
 
 typedef struct {
     bool enabled;
+    bool temporal_tiling;
     int tile_size_x;
     int tile_size_y;
     float target_overlap;
@@ -171,7 +175,9 @@ typedef struct {
     const char* llm_vision_path;
     const char* diffusion_model_path;
     const char* high_noise_diffusion_model_path;
+    const char* embeddings_connectors_path;
     const char* vae_path;
+    const char* audio_vae_path;
     const char* taesd_path;
     const char* control_net_path;
     const sd_embedding_t* embeddings;
@@ -203,8 +209,18 @@ typedef struct {
     bool chroma_use_t5_mask;
     int chroma_t5_mask_pad;
     bool qwen_image_zero_cond_t;
+    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;
 } sd_ctx_params_t;
 
+typedef struct {
+    uint32_t sample_rate;
+    uint32_t channels;
+    uint64_t sample_count;
+    float* data;
+} sd_audio_t;
+
 typedef struct {
     uint32_t width;
     uint32_t height;
@@ -237,6 +253,7 @@ typedef struct {
     float* custom_sigmas;
     int custom_sigmas_count;
     float flow_shift;
+    const char* extra_sample_args;
 } sd_sample_params_t;
 
 typedef struct {
@@ -359,6 +376,7 @@ typedef struct {
     float strength;
     int64_t seed;
     int video_frames;
+    int fps;
     float vace_strength;
     sd_tiling_params_t vae_tiling_params;
     sd_cache_params_t cache;
@@ -403,6 +421,7 @@ SD_API char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params);
 
 SD_API sd_ctx_t* new_sd_ctx(const sd_ctx_params_t* sd_ctx_params);
 SD_API void free_sd_ctx(sd_ctx_t* sd_ctx);
+SD_API void free_sd_audio(sd_audio_t* audio);
 
 SD_API void sd_sample_params_init(sd_sample_params_t* sample_params);
 SD_API char* sd_sample_params_to_str(const sd_sample_params_t* sample_params);
@@ -415,7 +434,11 @@ SD_API char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_para
 SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_gen_params);
 
 SD_API void sd_vid_gen_params_init(sd_vid_gen_params_t* sd_vid_gen_params);
-SD_API sd_image_t* generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* sd_vid_gen_params, int* num_frames_out);
+SD_API bool generate_video(sd_ctx_t* sd_ctx,
+                           const sd_vid_gen_params_t* sd_vid_gen_params,
+                           sd_image_t** frames_out,
+                           int* num_frames_out,
+                           sd_audio_t** audio_out);
 
 typedef struct upscaler_ctx_t upscaler_ctx_t;
 
@@ -423,7 +446,9 @@ SD_API upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path,
                                         bool offload_params_to_cpu,
                                         bool direct,
                                         int n_threads,
-                                        int tile_size);
+                                        int tile_size,
+                                        const char* backend,
+                                        const char* params_backend);
 SD_API void free_upscaler_ctx(upscaler_ctx_t* upscaler_ctx);
 
 SD_API sd_image_t upscale(upscaler_ctx_t* upscaler_ctx,

src/anima.hpp

@@ -499,9 +499,15 @@ namespace Anima {
                 encoder_hidden_states = adapted_context;
             }
 
+            sd::ggml_graph_cut::mark_graph_cut(x, "anima.prelude", "x");
+            sd::ggml_graph_cut::mark_graph_cut(embedded_timestep, "anima.prelude", "embedded_timestep");
+            sd::ggml_graph_cut::mark_graph_cut(temb, "anima.prelude", "temb");
+            sd::ggml_graph_cut::mark_graph_cut(encoder_hidden_states, "anima.prelude", "context");
+
             for (int i = 0; i < num_layers; i++) {
                 auto block = std::dynamic_pointer_cast<TransformerBlock>(blocks["blocks." + std::to_string(i)]);
                 x          = block->forward(ctx, x, encoder_hidden_states, embedded_timestep, temb, image_pe);
+                sd::ggml_graph_cut::mark_graph_cut(x, "anima.blocks." + std::to_string(i), "x");
             }
 
             x = final_layer->forward(ctx, x, embedded_timestep, temb);  // [N, h*w, ph*pw*C]
@@ -520,10 +526,10 @@ namespace Anima {
         AnimaNet net;
 
         AnimaRunner(ggml_backend_t backend,
-                    bool offload_params_to_cpu,
+                    ggml_backend_t params_backend,
                     const String2TensorStorage& tensor_storage_map = {},
                     const std::string prefix                       = "model.diffusion_model")
-            : GGMLRunner(backend, offload_params_to_cpu) {
+            : GGMLRunner(backend, params_backend) {
             int64_t num_layers    = 0;
             std::string layer_tag = prefix + ".net.blocks.";
             for (const auto& kv : tensor_storage_map) {

src/auto_encoder_kl.hpp

@@ -328,6 +328,7 @@ public:
         auto conv_out    = std::dynamic_pointer_cast<Conv2d>(blocks["conv_out"]);
 
         auto h = conv_in->forward(ctx, x);  // [N, ch, h, w]
+        // sd::ggml_graph_cut::mark_graph_cut(h, "vae.encoder.prelude", "h");
 
         // downsampling
         size_t num_resolutions = ch_mult.size();
@@ -337,12 +338,14 @@ public:
                 auto down_block  = std::dynamic_pointer_cast<ResnetBlock>(blocks[name]);
 
                 h = down_block->forward(ctx, h);
+                // sd::ggml_graph_cut::mark_graph_cut(h, "vae.encoder.down." + std::to_string(i) + ".block." + std::to_string(j), "h");
             }
             if (i != num_resolutions - 1) {
                 std::string name = "down." + std::to_string(i) + ".downsample";
                 auto down_sample = std::dynamic_pointer_cast<DownSampleBlock>(blocks[name]);
 
                 h = down_sample->forward(ctx, h);
+                // sd::ggml_graph_cut::mark_graph_cut(h, "vae.encoder.down." + std::to_string(i) + ".downsample", "h");
             }
         }
 
@@ -350,6 +353,7 @@ public:
         h = mid_block_1->forward(ctx, h);
         h = mid_attn_1->forward(ctx, h);
         h = mid_block_2->forward(ctx, h);  // [N, block_in, h, w]
+        // sd::ggml_graph_cut::mark_graph_cut(h, "vae.encoder.mid", "h");
 
         // end
         h = norm_out->forward(ctx, h);
@@ -450,6 +454,7 @@ public:
 
         // conv_in
         auto h = conv_in->forward(ctx, z);  // [N, block_in, h, w]
+        // sd::ggml_graph_cut::mark_graph_cut(h, "vae.decoder.prelude", "h");
 
         // middle
         h = mid_block_1->forward(ctx, h);
@@ -457,6 +462,7 @@ public:
 
         h = mid_attn_1->forward(ctx, h);
         h = mid_block_2->forward(ctx, h);  // [N, block_in, h, w]
+        // sd::ggml_graph_cut::mark_graph_cut(h, "vae.decoder.mid", "h");
 
         // upsampling
         int num_resolutions = static_cast<int>(ch_mult.size());
@@ -466,12 +472,14 @@ public:
                 auto up_block    = std::dynamic_pointer_cast<ResnetBlock>(blocks[name]);
 
                 h = up_block->forward(ctx, h);
+                // sd::ggml_graph_cut::mark_graph_cut(h, "vae.decoder.up." + std::to_string(i) + ".block." + std::to_string(j), "h");
             }
             if (i != 0) {
                 std::string name = "up." + std::to_string(i) + ".upsample";
                 auto up_sample   = std::dynamic_pointer_cast<UpSampleBlock>(blocks[name]);
 
                 h = up_sample->forward(ctx, h);
+                // sd::ggml_graph_cut::mark_graph_cut(h, "vae.decoder.up." + std::to_string(i) + ".upsample", "h");
             }
         }
 
@@ -599,6 +607,7 @@ public:
         if (use_quant) {
             auto post_quant_conv = std::dynamic_pointer_cast<Conv2d>(blocks["post_quant_conv"]);
             z                    = post_quant_conv->forward(ctx, z);  // [N, z_channels, h, w]
+            // sd::ggml_graph_cut::mark_graph_cut(z, "vae.decode.prelude", "z");
         }
         auto decoder = std::dynamic_pointer_cast<Decoder>(blocks["decoder"]);
 
@@ -616,6 +625,7 @@ public:
         if (use_quant) {
             auto quant_conv = std::dynamic_pointer_cast<Conv2d>(blocks["quant_conv"]);
             z               = quant_conv->forward(ctx, z);  // [N, 2*embed_dim, h/8, w/8]
+            // sd::ggml_graph_cut::mark_graph_cut(z, "vae.encode.final", "z");
         }
         if (sd_version_uses_flux2_vae(version)) {
             z = ggml_ext_chunk(ctx->ggml_ctx, z, 2, 2)[0];
@@ -654,13 +664,13 @@ struct AutoEncoderKL : public VAE {
     AutoEncoderKLModel ae;
 
     AutoEncoderKL(ggml_backend_t backend,
-                  bool offload_params_to_cpu,
+                  ggml_backend_t params_backend,
                   const String2TensorStorage& tensor_storage_map,
                   const std::string prefix,
                   bool decode_only       = false,
                   bool use_video_decoder = false,
                   SDVersion version      = VERSION_SD1)
-        : decode_only(decode_only), VAE(version, backend, offload_params_to_cpu) {
+        : decode_only(decode_only), VAE(version, backend, params_backend) {
         if (sd_version_is_sd1(version) || sd_version_is_sd2(version)) {
             scale_factor = 0.18215f;
             shift_factor = 0.f;

src/clip.hpp

@@ -95,8 +95,9 @@ public:
 
     ggml_tensor* forward(GGMLRunnerContext* ctx,
                          ggml_tensor* x,
-                         ggml_tensor* mask = nullptr,
-                         int clip_skip     = -1) {
+                         ggml_tensor* mask                   = nullptr,
+                         int clip_skip                       = -1,
+                         const std::string& graph_cut_prefix = "") {
         // x: [N, n_token, d_model]
         int layer_idx = n_layer - 1;
         // LOG_DEBUG("clip_skip %d", clip_skip);
@@ -112,6 +113,9 @@ public:
             std::string name = "layers." + std::to_string(i);
             auto layer       = std::dynamic_pointer_cast<CLIPLayer>(blocks[name]);
             x                = layer->forward(ctx, x, mask);  // [N, n_token, d_model]
+            if (!graph_cut_prefix.empty()) {
+                sd::ggml_graph_cut::mark_graph_cut(x, graph_cut_prefix + ".layers." + std::to_string(i), "x");
+            }
             // LOG_DEBUG("layer %d", i);
         }
         return x;
@@ -304,7 +308,8 @@ public:
         auto final_layer_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["final_layer_norm"]);
 
         auto x = embeddings->forward(ctx, input_ids, tkn_embeddings);  // [N, n_token, hidden_size]
-        x      = encoder->forward(ctx, x, mask, return_pooled ? -1 : clip_skip);
+        sd::ggml_graph_cut::mark_graph_cut(x, "clip_text.prelude", "x");
+        x = encoder->forward(ctx, x, mask, return_pooled ? -1 : clip_skip, "clip_text");
         if (return_pooled || with_final_ln) {
             x = final_layer_norm->forward(ctx, x);
         }
@@ -368,7 +373,8 @@ public:
 
         auto x = embeddings->forward(ctx, pixel_values);  // [N, num_positions, embed_dim]
         x      = pre_layernorm->forward(ctx, x);
-        x      = encoder->forward(ctx, x, nullptr, clip_skip);
+        sd::ggml_graph_cut::mark_graph_cut(x, "clip_vision.prelude", "x");
+        x = encoder->forward(ctx, x, nullptr, clip_skip, "clip_vision");
 
         auto last_hidden_state = x;
 
@@ -463,13 +469,13 @@ struct CLIPTextModelRunner : public GGMLRunner {
     std::vector<float> attention_mask_vec;
 
     CLIPTextModelRunner(ggml_backend_t backend,
-                        bool offload_params_to_cpu,
+                        ggml_backend_t params_backend,
                         const String2TensorStorage& tensor_storage_map,
                         const std::string prefix,
                         CLIPVersion version = OPENAI_CLIP_VIT_L_14,
                         bool with_final_ln  = true,
                         bool force_clip_f32 = false)
-        : GGMLRunner(backend, offload_params_to_cpu) {
+        : GGMLRunner(backend, params_backend) {
         bool proj_in = false;
         for (const auto& [name, tensor_storage] : tensor_storage_map) {
             if (!starts_with(name, prefix)) {

src/common_block.hpp

@@ -1,7 +1,9 @@
 #ifndef __COMMON_BLOCK_HPP__
 #define __COMMON_BLOCK_HPP__
 
+#include "ggml-backend.h"
 #include "ggml_extend.hpp"
+#include "util.h"
 
 class DownSampleBlock : public GGMLBlock {
 protected:
@@ -248,9 +250,6 @@ public:
         float scale         = 1.f;
         if (precision_fix) {
             scale = 1.f / 128.f;
-#ifdef SD_USE_VULKAN
-            force_prec_f32 = true;
-#endif
         }
         // The purpose of the scale here is to prevent NaN issues in certain situations.
         // For example, when using Vulkan without enabling force_prec_f32,
@@ -264,6 +263,9 @@ public:
 
         auto net_0 = std::dynamic_pointer_cast<UnaryBlock>(blocks["net.0"]);
         auto net_2 = std::dynamic_pointer_cast<Linear>(blocks["net.2"]);
+        if (sd_backend_is(ctx->backend, "Vulkan")) {
+            net_2->set_force_prec_f32(true);
+        }
 
         x = net_0->forward(ctx, x);  // [ne3, ne2, ne1, inner_dim]
         x = net_2->forward(ctx, x);  // [ne3, ne2, ne1, dim_out]

src/common_dit.hpp

@@ -103,6 +103,64 @@ namespace DiT {
         x         = ggml_ext_slice(ctx, x, 0, 0, W);               // [N, C, H, W]
         return x;
     }
+
+    inline ggml_tensor* patchify(ggml_context* ctx,
+                                 ggml_tensor* x,
+                                 int pt,
+                                 int ph,
+                                 int pw,
+                                 int64_t N = 1) {
+        // x: [N*C, T, H, W]
+        // return: [N, h*w, C*pt*ph*pw]
+        int64_t C     = x->ne[3] / N;
+        int64_t T     = x->ne[2];
+        int64_t H     = x->ne[1];
+        int64_t W     = x->ne[0];
+        int64_t t_len = T / pt;
+        int64_t h_len = H / ph;
+        int64_t w_len = W / pw;
+
+        GGML_ASSERT(C * N == x->ne[3]);
+        GGML_ASSERT(t_len * pt == T && h_len * ph == H && w_len * pw == W);
+
+        x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt, t_len * C * N);      // [N*C*t_len, pt, h_len*ph, w_len*pw]
+        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));          // [N*C*t_len, h_len*ph, pt, w_len*pw]
+        x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph, h_len * t_len * C * N);      // [N*C*t_len*h_len, ph, pt, w_len*pw]
+        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));          // [N*C*t_len*h_len, pt, ph, w_len*pw]
+        x = ggml_reshape_4d(ctx, x, pw, w_len, ph * pt, h_len * t_len * C * N);      // [N*C*t_len*h_len, pt*ph, w_len, pw]
+        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));          // [N*C*t_len*h_len, w_len, pt*ph, pw]
+        x = ggml_reshape_4d(ctx, x, pw * ph * pt, w_len * h_len * t_len, C, N);      // [N, C, t_len*h_len*w_len, pt*ph*pw]
+        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));          // [N, t_len*h_len*w_len, C, pt*ph*pw]
+        x = ggml_reshape_4d(ctx, x, pw * ph * pt * C, w_len * h_len * t_len, N, 1);  // [N, t_len*h_len*w_len, C*pt*ph*pw]
+        return x;
+    }
+
+    inline ggml_tensor* unpatchify(ggml_context* ctx,
+                                   ggml_tensor* x,
+                                   int64_t t_len,
+                                   int64_t h_len,
+                                   int64_t w_len,
+                                   int pt,
+                                   int ph,
+                                   int pw) {
+        // x: [N, t_len*h_len*w_len, pt*ph*pw*C]
+        // return: [N*C, t_len*pt, h_len*ph, w_len*pw]
+        int64_t N = x->ne[3];
+        int64_t C = x->ne[0] / pt / ph / pw;
+
+        GGML_ASSERT(C * pt * ph * pw == x->ne[0]);
+
+        x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N);  // [N, t_len*h_len*w_len, pt*ph*pw, C]
+        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3));      // [N, C, t_len*h_len*w_len, pt*ph*pw]
+        x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N);  // [N*C*t_len*h_len, w_len, pt*ph, pw]
+        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len*h_len, pt*ph, w_len, pw]
+        x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N);  // [N*C*t_len*h_len, pt, ph, w_len*pw]
+        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len*h_len, ph, pt, w_len*pw]
+        x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph * h_len, t_len * C * N);  // [N*C*t_len, h_len*ph, pt, w_len*pw]
+        x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len, pt, h_len*ph, w_len*pw]
+        x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt * t_len, C * N);  // [N*C, t_len*pt, h_len*ph, w_len*pw]
+        return x;
+    }
 }  // namespace DiT
 
 #endif  // __COMMON_DIT_HPP__

src/conditioner.hpp

@@ -1,6 +1,8 @@
 #ifndef __CONDITIONER_HPP__
 #define __CONDITIONER_HPP__
 
+#include <cmath>
+#include <limits>
 #include <optional>
 
 #include "clip.hpp"
@@ -14,6 +16,12 @@ struct SDCondition {
     sd::Tensor<float> c_concat;
     sd::Tensor<int32_t> c_t5_ids;
     sd::Tensor<float> c_t5_weights;
+    sd::Tensor<int32_t> c_input_ids;
+    sd::Tensor<int32_t> c_position_ids;
+    sd::Tensor<int32_t> c_token_types;
+    sd::Tensor<int32_t> c_vinput_mask;
+    std::vector<std::pair<int, sd::Tensor<float>>> c_image_embeds;
+    std::vector<sd::Tensor<float>> c_ref_images;
 
     std::vector<sd::Tensor<float>> extra_c_crossattns;
 
@@ -26,10 +34,24 @@ struct SDCondition {
 
     bool empty() const {
         if (!c_crossattn.empty() || !c_vector.empty() || !c_concat.empty() ||
-            !c_t5_ids.empty() || !c_t5_weights.empty()) {
+            !c_t5_ids.empty() || !c_t5_weights.empty() ||
+            !c_input_ids.empty() || !c_position_ids.empty() ||
+            !c_token_types.empty() || !c_vinput_mask.empty()) {
             return false;
         }
 
+        for (const auto& image_embed : c_image_embeds) {
+            if (!image_embed.second.empty()) {
+                return false;
+            }
+        }
+
+        for (const auto& tensor : c_ref_images) {
+            if (!tensor.empty()) {
+                return false;
+            }
+        }
+
         for (const auto& tensor : extra_c_crossattns) {
             if (!tensor.empty()) {
                 return false;
@@ -46,6 +68,17 @@ static inline sd::Tensor<float> apply_token_weights(sd::Tensor<float> hidden_sta
         return hidden_states;
     }
 
+    bool all_one = true;
+    for (float weight : weights) {
+        if (weight != 1.0f) {
+            all_one = false;
+            break;
+        }
+    }
+    if (all_one) {
+        return hidden_states;
+    }
+
     if (hidden_states.dim() == 1) {
         hidden_states.unsqueeze_(1);
     }
@@ -57,7 +90,7 @@ static inline sd::Tensor<float> apply_token_weights(sd::Tensor<float> hidden_sta
     chunk_weights.reshape_({1, static_cast<int64_t>(weights.size())});
     hidden_states *= chunk_weights;
     float new_mean = hidden_states.mean();
-    if (new_mean != 0.0f) {
+    if (std::isfinite(original_mean) && std::isfinite(new_mean) && new_mean != 0.0f) {
         hidden_states *= (original_mean / new_mean);
     }
 
@@ -85,7 +118,8 @@ public:
     virtual void free_params_buffer()                                                      = 0;
     virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors)           = 0;
     virtual size_t get_params_buffer_size()                                                = 0;
-    virtual void set_flash_attention_enabled(bool enabled)                                 = 0;
+    virtual void set_max_graph_vram_bytes(size_t max_vram_bytes) {}
+    virtual void set_flash_attention_enabled(bool enabled) = 0;
     virtual void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) {}
     virtual std::tuple<SDCondition, std::vector<bool>> get_learned_condition_with_trigger(int n_threads,
                                                                                           const ConditionerParams& conditioner_params) {
@@ -113,7 +147,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
     std::map<std::string, std::pair<int, int>> embedding_pos_map;
 
     FrozenCLIPEmbedderWithCustomWords(ggml_backend_t backend,
-                                      bool offload_params_to_cpu,
+                                      ggml_backend_t params_backend,
                                       const String2TensorStorage& tensor_storage_map,
                                       const std::map<std::string, std::string>& orig_embedding_map,
                                       SDVersion version = VERSION_SD1,
@@ -127,12 +161,12 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
         }
         bool force_clip_f32 = !embedding_map.empty();
         if (sd_version_is_sd1(version)) {
-            text_model = std::make_shared<CLIPTextModelRunner>(backend, offload_params_to_cpu, tensor_storage_map, "cond_stage_model.transformer.text_model", OPENAI_CLIP_VIT_L_14, true, force_clip_f32);
+            text_model = std::make_shared<CLIPTextModelRunner>(backend, params_backend, tensor_storage_map, "cond_stage_model.transformer.text_model", OPENAI_CLIP_VIT_L_14, true, force_clip_f32);
         } else if (sd_version_is_sd2(version)) {
-            text_model = std::make_shared<CLIPTextModelRunner>(backend, offload_params_to_cpu, tensor_storage_map, "cond_stage_model.transformer.text_model", OPEN_CLIP_VIT_H_14, true, force_clip_f32);
+            text_model = std::make_shared<CLIPTextModelRunner>(backend, params_backend, tensor_storage_map, "cond_stage_model.transformer.text_model", OPEN_CLIP_VIT_H_14, true, force_clip_f32);
         } else if (sd_version_is_sdxl(version)) {
-            text_model  = std::make_shared<CLIPTextModelRunner>(backend, offload_params_to_cpu, tensor_storage_map, "cond_stage_model.transformer.text_model", OPENAI_CLIP_VIT_L_14, false, force_clip_f32);
-            text_model2 = std::make_shared<CLIPTextModelRunner>(backend, offload_params_to_cpu, tensor_storage_map, "cond_stage_model.1.transformer.text_model", OPEN_CLIP_VIT_BIGG_14, false, force_clip_f32);
+            text_model  = std::make_shared<CLIPTextModelRunner>(backend, params_backend, tensor_storage_map, "cond_stage_model.transformer.text_model", OPENAI_CLIP_VIT_L_14, false, force_clip_f32);
+            text_model2 = std::make_shared<CLIPTextModelRunner>(backend, params_backend, tensor_storage_map, "cond_stage_model.1.transformer.text_model", OPEN_CLIP_VIT_BIGG_14, false, force_clip_f32);
         }
     }
 
@@ -165,6 +199,13 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
         return buffer_size;
     }
 
+    void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
+        text_model->set_max_graph_vram_bytes(max_vram_bytes);
+        if (sd_version_is_sdxl(version)) {
+            text_model2->set_max_graph_vram_bytes(max_vram_bytes);
+        }
+    }
+
     void set_flash_attention_enabled(bool enabled) override {
         text_model->set_flash_attention_enabled(enabled);
         if (sd_version_is_sdxl(version)) {
@@ -642,9 +683,9 @@ struct FrozenCLIPVisionEmbedder : public GGMLRunner {
     CLIPVisionModelProjection vision_model;
 
     FrozenCLIPVisionEmbedder(ggml_backend_t backend,
-                             bool offload_params_to_cpu,
+                             ggml_backend_t params_backend,
                              const String2TensorStorage& tensor_storage_map = {})
-        : GGMLRunner(backend, offload_params_to_cpu) {
+        : GGMLRunner(backend, params_backend) {
         std::string prefix = "cond_stage_model.transformer";
         bool proj_in       = false;
         for (const auto& [name, tensor_storage] : tensor_storage_map) {
@@ -701,7 +742,7 @@ struct SD3CLIPEmbedder : public Conditioner {
     std::shared_ptr<T5Runner> t5;
 
     SD3CLIPEmbedder(ggml_backend_t backend,
-                    bool offload_params_to_cpu,
+                    ggml_backend_t params_backend,
                     const String2TensorStorage& tensor_storage_map = {})
         : clip_g_tokenizer(0) {
         bool use_clip_l = false;
@@ -721,13 +762,13 @@ struct SD3CLIPEmbedder : public Conditioner {
             return;
         }
         if (use_clip_l) {
-            clip_l = std::make_shared<CLIPTextModelRunner>(backend, offload_params_to_cpu, tensor_storage_map, "text_encoders.clip_l.transformer.text_model", OPENAI_CLIP_VIT_L_14, false);
+            clip_l = std::make_shared<CLIPTextModelRunner>(backend, params_backend, tensor_storage_map, "text_encoders.clip_l.transformer.text_model", OPENAI_CLIP_VIT_L_14, false);
         }
         if (use_clip_g) {
-            clip_g = std::make_shared<CLIPTextModelRunner>(backend, offload_params_to_cpu, tensor_storage_map, "text_encoders.clip_g.transformer.text_model", OPEN_CLIP_VIT_BIGG_14, false);
+            clip_g = std::make_shared<CLIPTextModelRunner>(backend, params_backend, tensor_storage_map, "text_encoders.clip_g.transformer.text_model", OPEN_CLIP_VIT_BIGG_14, false);
         }
         if (use_t5) {
-            t5 = std::make_shared<T5Runner>(backend, offload_params_to_cpu, tensor_storage_map, "text_encoders.t5xxl.transformer");
+            t5 = std::make_shared<T5Runner>(backend, params_backend, tensor_storage_map, "text_encoders.t5xxl.transformer");
         }
     }
 
@@ -781,6 +822,18 @@ struct SD3CLIPEmbedder : public Conditioner {
         return buffer_size;
     }
 
+    void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
+        if (clip_l) {
+            clip_l->set_max_graph_vram_bytes(max_vram_bytes);
+        }
+        if (clip_g) {
+            clip_g->set_max_graph_vram_bytes(max_vram_bytes);
+        }
+        if (t5) {
+            t5->set_max_graph_vram_bytes(max_vram_bytes);
+        }
+    }
+
     void set_flash_attention_enabled(bool enabled) override {
         if (clip_l) {
             clip_l->set_flash_attention_enabled(enabled);
@@ -1057,7 +1110,7 @@ struct FluxCLIPEmbedder : public Conditioner {
     size_t chunk_len = 256;
 
     FluxCLIPEmbedder(ggml_backend_t backend,
-                     bool offload_params_to_cpu,
+                     ggml_backend_t params_backend,
                      const String2TensorStorage& tensor_storage_map = {}) {
         bool use_clip_l = false;
         bool use_t5     = false;
@@ -1075,12 +1128,12 @@ struct FluxCLIPEmbedder : public Conditioner {
         }
 
         if (use_clip_l) {
-            clip_l = std::make_shared<CLIPTextModelRunner>(backend, offload_params_to_cpu, tensor_storage_map, "text_encoders.clip_l.transformer.text_model", OPENAI_CLIP_VIT_L_14, true);
+            clip_l = std::make_shared<CLIPTextModelRunner>(backend, params_backend, tensor_storage_map, "text_encoders.clip_l.transformer.text_model", OPENAI_CLIP_VIT_L_14, true);
         } else {
             LOG_WARN("clip_l text encoder not found! Prompt adherence might be degraded.");
         }
         if (use_t5) {
-            t5 = std::make_shared<T5Runner>(backend, offload_params_to_cpu, tensor_storage_map, "text_encoders.t5xxl.transformer");
+            t5 = std::make_shared<T5Runner>(backend, params_backend, tensor_storage_map, "text_encoders.t5xxl.transformer");
         } else {
             LOG_WARN("t5xxl text encoder not found! Prompt adherence might be degraded.");
         }
@@ -1124,6 +1177,15 @@ struct FluxCLIPEmbedder : public Conditioner {
         return buffer_size;
     }
 
+    void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
+        if (clip_l) {
+            clip_l->set_max_graph_vram_bytes(max_vram_bytes);
+        }
+        if (t5) {
+            t5->set_max_graph_vram_bytes(max_vram_bytes);
+        }
+    }
+
     void set_flash_attention_enabled(bool enabled) override {
         if (clip_l) {
             clip_l->set_flash_attention_enabled(enabled);
@@ -1302,7 +1364,7 @@ struct T5CLIPEmbedder : public Conditioner {
     bool is_umt5     = false;
 
     T5CLIPEmbedder(ggml_backend_t backend,
-                   bool offload_params_to_cpu,
+                   ggml_backend_t params_backend,
                    const String2TensorStorage& tensor_storage_map = {},
                    bool use_mask                                  = false,
                    int mask_pad                                   = 0,
@@ -1319,7 +1381,7 @@ struct T5CLIPEmbedder : public Conditioner {
             LOG_WARN("IMPORTANT NOTICE: No text encoders provided, cannot process prompts!");
             return;
         } else {
-            t5 = std::make_shared<T5Runner>(backend, offload_params_to_cpu, tensor_storage_map, "text_encoders.t5xxl.transformer", is_umt5);
+            t5 = std::make_shared<T5Runner>(backend, params_backend, tensor_storage_map, "text_encoders.t5xxl.transformer", is_umt5);
         }
     }
 
@@ -1349,6 +1411,12 @@ struct T5CLIPEmbedder : public Conditioner {
         return buffer_size;
     }
 
+    void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
+        if (t5) {
+            t5->set_max_graph_vram_bytes(max_vram_bytes);
+        }
+    }
+
     void set_flash_attention_enabled(bool enabled) override {
         if (t5) {
             t5->set_flash_attention_enabled(enabled);
@@ -1498,12 +1566,12 @@ struct AnimaConditioner : public Conditioner {
     std::shared_ptr<LLM::LLMRunner> llm;
 
     AnimaConditioner(ggml_backend_t backend,
-                     bool offload_params_to_cpu,
+                     ggml_backend_t params_backend,
                      const String2TensorStorage& tensor_storage_map = {}) {
         qwen_tokenizer = std::make_shared<Qwen2Tokenizer>();
         llm            = std::make_shared<LLM::LLMRunner>(LLM::LLMArch::QWEN3,
                                                backend,
-                                               offload_params_to_cpu,
+                                               params_backend,
                                                tensor_storage_map,
                                                "text_encoders.llm",
                                                false);
@@ -1525,6 +1593,10 @@ struct AnimaConditioner : public Conditioner {
         return llm->get_params_buffer_size();
     }
 
+    void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
+        llm->set_max_graph_vram_bytes(max_vram_bytes);
+    }
+
     void set_flash_attention_enabled(bool enabled) override {
         llm->set_flash_attention_enabled(enabled);
     }
@@ -1566,10 +1638,11 @@ struct AnimaConditioner : public Conditioner {
         for (const auto& item : parsed_attention) {
             const std::string& curr_text = item.first;
             float curr_weight            = item.second;
-            std::vector<int> curr_tokens = t5_tokenizer.tokenize(curr_text, nullptr, true);
+            std::vector<int> curr_tokens = t5_tokenizer.encode(curr_text);
             t5_tokens.insert(t5_tokens.end(), curr_tokens.begin(), curr_tokens.end());
             t5_weights.insert(t5_weights.end(), curr_tokens.size(), curr_weight);
         }
+        t5_tokenizer.pad_tokens(t5_tokens, &t5_weights, nullptr);
 
         return {qwen_tokens, qwen_weights, t5_tokens, t5_weights};
     }
@@ -1612,7 +1685,7 @@ struct LLMEmbedder : public Conditioner {
     std::shared_ptr<LLM::LLMRunner> llm;
 
     LLMEmbedder(ggml_backend_t backend,
-                bool offload_params_to_cpu,
+                ggml_backend_t params_backend,
                 const String2TensorStorage& tensor_storage_map = {},
                 SDVersion version                              = VERSION_QWEN_IMAGE,
                 const std::string prefix                       = "",
@@ -1633,7 +1706,7 @@ struct LLMEmbedder : public Conditioner {
         }
         llm = std::make_shared<LLM::LLMRunner>(arch,
                                                backend,
-                                               offload_params_to_cpu,
+                                               params_backend,
                                                tensor_storage_map,
                                                "text_encoders.llm",
                                                enable_vision);
@@ -1657,6 +1730,10 @@ struct LLMEmbedder : public Conditioner {
         return buffer_size;
     }
 
+    void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
+        llm->set_max_graph_vram_bytes(max_vram_bytes);
+    }
+
     void set_flash_attention_enabled(bool enabled) override {
         llm->set_flash_attention_enabled(enabled);
     }
@@ -1958,4 +2035,277 @@ struct LLMEmbedder : public Conditioner {
     }
 };
 
+struct LTXAVTextProjection : public GGMLBlock {
+    static constexpr int64_t kHiddenSize = 3840;
+    static constexpr int64_t kNumStates  = 49;
+    bool dual_projection                 = false;
+
+    LTXAVTextProjection(bool dual_projection = false)
+        : dual_projection(dual_projection) {
+        if (dual_projection) {
+            blocks["video_aggregate_embed"] = std::make_shared<Linear>(kHiddenSize * kNumStates, 4096, true);
+            blocks["audio_aggregate_embed"] = std::make_shared<Linear>(kHiddenSize * kNumStates, 2048, true);
+        } else {
+            blocks["projection"] = std::make_shared<Linear>(kHiddenSize * kNumStates, kHiddenSize, false);
+        }
+    }
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+        if (!dual_projection) {
+            auto projection = std::dynamic_pointer_cast<Linear>(blocks["projection"]);
+            return projection->forward(ctx, x);
+        }
+
+        auto video_projection = std::dynamic_pointer_cast<Linear>(blocks["video_aggregate_embed"]);
+        auto audio_projection = std::dynamic_pointer_cast<Linear>(blocks["audio_aggregate_embed"]);
+        auto video_in         = ggml_ext_scale(ctx->ggml_ctx, x, std::sqrt(4096.f / static_cast<float>(kHiddenSize)));
+        auto audio_in         = ggml_ext_scale(ctx->ggml_ctx, x, std::sqrt(2048.f / static_cast<float>(kHiddenSize)));
+        auto video            = video_projection->forward(ctx, video_in);
+        auto audio            = audio_projection->forward(ctx, audio_in);
+        return ggml_concat(ctx->ggml_ctx, video, audio, 0);
+    }
+};
+
+struct LTXAVTextProjectionRunner : public GGMLRunner {
+    LTXAVTextProjection model;
+
+    LTXAVTextProjectionRunner(ggml_backend_t backend,
+                              ggml_backend_t params_backend,
+                              const String2TensorStorage& tensor_storage_map = {},
+                              const std::string& prefix                      = "")
+        : GGMLRunner(backend, params_backend),
+          model(tensor_storage_map.find(prefix + ".video_aggregate_embed.weight") != tensor_storage_map.end()) {
+        model.init(params_ctx, tensor_storage_map, prefix);
+    }
+
+    std::string get_desc() override {
+        return "ltxav_text_projection";
+    }
+
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) {
+        model.get_param_tensors(tensors, prefix);
+    }
+
+    ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor) {
+        ggml_cgraph* gf = ggml_new_graph(compute_ctx);
+        auto x          = make_input(x_tensor);
+        auto runner_ctx = get_context();
+        auto out        = model.forward(&runner_ctx, x);
+        ggml_build_forward_expand(gf, out);
+        return gf;
+    }
+
+    sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& x) {
+        auto get_graph = [&]() -> ggml_cgraph* {
+            return build_graph(x);
+        };
+        return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true));
+    }
+};
+
+struct LTXAVEmbedder : public Conditioner {
+    static constexpr int64_t kHiddenSize = 3840;
+    static constexpr int64_t kNumStates  = 49;
+    static constexpr int64_t kMinLength  = 1024;
+
+    std::shared_ptr<GemmaTokenizer> tokenizer;
+    std::shared_ptr<LLM::LLMRunner> llm;
+    std::shared_ptr<LTXAVTextProjectionRunner> projector;
+    bool dual_projection = false;
+
+    LTXAVEmbedder(ggml_backend_t backend,
+                  ggml_backend_t params_backend,
+                  const String2TensorStorage& tensor_storage_map = {},
+                  const std::string& llm_prefix                  = "text_encoders.llm",
+                  const std::string& projector_prefix            = "text_embedding_projection") {
+        tokenizer       = std::make_shared<GemmaTokenizer>();
+        llm             = std::make_shared<LLM::LLMRunner>(LLM::LLMArch::GEMMA3_12B,
+                                               backend,
+                                               params_backend,
+                                               tensor_storage_map,
+                                               llm_prefix,
+                                               false);
+        dual_projection = tensor_storage_map.find(projector_prefix + ".video_aggregate_embed.weight") != tensor_storage_map.end();
+        projector       = std::make_shared<LTXAVTextProjectionRunner>(backend,
+                                                                params_backend,
+                                                                tensor_storage_map,
+                                                                projector_prefix);
+    }
+
+    void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+        llm->get_param_tensors(tensors, "text_encoders.llm");
+        projector->get_param_tensors(tensors, "text_embedding_projection");
+    }
+
+    void alloc_params_buffer() override {
+        llm->alloc_params_buffer();
+        projector->alloc_params_buffer();
+    }
+
+    void free_params_buffer() override {
+        llm->free_params_buffer();
+        projector->free_params_buffer();
+    }
+
+    size_t get_params_buffer_size() override {
+        return llm->get_params_buffer_size() + projector->get_params_buffer_size();
+    }
+
+    void set_flash_attention_enabled(bool enabled) override {
+        llm->set_flash_attention_enabled(enabled);
+        projector->set_flash_attention_enabled(enabled);
+    }
+
+    void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
+        llm->set_weight_adapter(adapter);
+        projector->set_weight_adapter(adapter);
+    }
+
+    std::tuple<std::vector<int>, std::vector<float>, std::vector<float>> tokenize(std::string text,
+                                                                                  const std::pair<int, int>& attn_range) {
+        std::vector<std::pair<std::string, float>> parsed_attention;
+        if (attn_range.first >= 0 && attn_range.second > 0) {
+            if (attn_range.first > 0) {
+                parsed_attention.emplace_back(text.substr(0, attn_range.first), 1.f);
+            }
+            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));
+                parsed_attention.insert(parsed_attention.end(), new_parsed_attention.begin(), new_parsed_attention.end());
+            }
+            if (static_cast<size_t>(attn_range.second) < text.size()) {
+                parsed_attention.emplace_back(text.substr(attn_range.second), 1.f);
+            }
+        } else {
+            parsed_attention.emplace_back(text, 1.f);
+        }
+
+        std::vector<int> tokens;
+        std::vector<float> weights;
+        for (const auto& item : parsed_attention) {
+            auto curr_tokens = tokenizer->encode(item.first, nullptr);
+            tokens.insert(tokens.end(), curr_tokens.begin(), curr_tokens.end());
+            weights.insert(weights.end(), curr_tokens.size(), item.second);
+        }
+
+        std::vector<float> mask;
+        tokenizer->pad_tokens(tokens, &weights, &mask, kMinLength);
+        return {tokens, weights, mask};
+    }
+
+    sd::Tensor<float> encode_prompt(int n_threads,
+                                    const std::string& prompt,
+                                    const std::pair<int, int>& prompt_attn_range) {
+        auto tokens_weights_mask = tokenize(prompt, prompt_attn_range);
+        auto& tokens             = std::get<0>(tokens_weights_mask);
+        auto& weights            = std::get<1>(tokens_weights_mask);
+        auto& mask               = std::get<2>(tokens_weights_mask);
+
+        sd::Tensor<int32_t> input_ids({static_cast<int64_t>(tokens.size())}, std::vector<int32_t>(tokens.begin(), tokens.end()));
+        sd::Tensor<float> attention_mask;
+        if (!mask.empty()) {
+            const float mask_min = std::numeric_limits<float>::lowest() / 4.0f;
+            attention_mask       = sd::Tensor<float>({static_cast<int64_t>(mask.size()), static_cast<int64_t>(mask.size())});
+            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) {
+                        value += mask_min;
+                    }
+                    if (i0 > i1) {
+                        value += mask_min;
+                    }
+                    attention_mask[static_cast<int64_t>(i0 + mask.size() * i1)] = value;
+                }
+            }
+        }
+
+        auto hidden_states = llm->compute(n_threads,
+                                          input_ids,
+                                          attention_mask,
+                                          {},
+                                          {},
+                                          true);
+        GGML_ASSERT(!hidden_states.empty());
+        hidden_states = apply_token_weights(std::move(hidden_states), weights);
+
+        int64_t valid_tokens = 0;
+        for (float value : mask) {
+            valid_tokens += static_cast<int64_t>(value > 0.0f);
+        }
+        GGML_ASSERT(valid_tokens > 0);
+
+        hidden_states = sd::ops::slice(hidden_states,
+                                       1,
+                                       hidden_states.shape()[1] - valid_tokens,
+                                       hidden_states.shape()[1]);
+        hidden_states.reshape_({kHiddenSize, kNumStates, valid_tokens});
+        hidden_states = hidden_states.permute({1, 0, 2});
+
+        if (dual_projection) {
+            for (int64_t state_idx = 0; state_idx < kNumStates; ++state_idx) {
+                for (int64_t token_idx = 0; token_idx < valid_tokens; ++token_idx) {
+                    double sq_sum = 0.0;
+                    for (int64_t hidden_idx = 0; hidden_idx < kHiddenSize; ++hidden_idx) {
+                        float value = hidden_states.index(state_idx, hidden_idx, token_idx);
+                        sq_sum += static_cast<double>(value) * static_cast<double>(value);
+                    }
+
+                    float inv_rms = 1.0f / std::sqrt(static_cast<float>(sq_sum / static_cast<double>(kHiddenSize)) + 1e-6f);
+                    for (int64_t hidden_idx = 0; hidden_idx < kHiddenSize; ++hidden_idx) {
+                        hidden_states.index(state_idx, hidden_idx, token_idx) *= inv_rms;
+                    }
+                }
+            }
+        } else {
+            for (int64_t state_idx = 0; state_idx < kNumStates; ++state_idx) {
+                double sum      = 0.0;
+                float min_value = std::numeric_limits<float>::infinity();
+                float max_value = -std::numeric_limits<float>::infinity();
+                for (int64_t token_idx = 0; token_idx < valid_tokens; ++token_idx) {
+                    for (int64_t hidden_idx = 0; hidden_idx < kHiddenSize; ++hidden_idx) {
+                        float value = hidden_states.index(state_idx, hidden_idx, token_idx);
+                        sum += value;
+                        min_value = std::min(min_value, value);
+                        max_value = std::max(max_value, value);
+                    }
+                }
+
+                float mean_value  = static_cast<float>(sum / static_cast<double>(kHiddenSize * valid_tokens));
+                float denom       = max_value - min_value + 1e-6f;
+                float scale_value = 8.0f / denom;
+                for (int64_t token_idx = 0; token_idx < valid_tokens; ++token_idx) {
+                    for (int64_t hidden_idx = 0; hidden_idx < kHiddenSize; ++hidden_idx) {
+                        float value                                           = hidden_states.index(state_idx, hidden_idx, token_idx);
+                        hidden_states.index(state_idx, hidden_idx, token_idx) = (value - mean_value) * scale_value;
+                    }
+                }
+            }
+        }
+
+        hidden_states.reshape_({kNumStates * kHiddenSize, valid_tokens});
+        return projector->compute(n_threads, hidden_states);
+    }
+
+    SDCondition get_learned_condition(int n_threads,
+                                      const ConditionerParams& conditioner_params) override {
+        int64_t t0 = ggml_time_ms();
+
+        std::string prompt;
+        std::pair<int, int> prompt_attn_range;
+        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);
+        GGML_ASSERT(!hidden_states.empty());
+
+        int64_t t1 = ggml_time_ms();
+        LOG_DEBUG("computing LTXAV condition graph completed, taking %" PRId64 " ms", t1 - t0);
+
+        SDCondition result;
+        result.c_crossattn = std::move(hidden_states);
+        return result;
+    }
+};
+
 #endif

src/control.hpp

@@ -319,10 +319,10 @@ struct ControlNet : public GGMLRunner {
     bool guided_hint_cached = false;
 
     ControlNet(ggml_backend_t backend,
-               bool offload_params_to_cpu,
+               ggml_backend_t params_backend,
                const String2TensorStorage& tensor_storage_map = {},
                SDVersion version                              = VERSION_SD1)
-        : GGMLRunner(backend, offload_params_to_cpu), control_net(version) {
+        : GGMLRunner(backend, params_backend), control_net(version) {
         control_net.init(params_ctx, tensor_storage_map, "");
     }
 

src/denoiser.hpp

@@ -1,7 +1,10 @@
 #ifndef __DENOISER_HPP__
 #define __DENOISER_HPP__
 
+#include <algorithm>
+#include <cctype>
 #include <cmath>
+#include <string>
 #include <utility>
 
 #include "ggml_extend.hpp"
@@ -479,6 +482,141 @@ struct KLOptimalScheduler : SigmaScheduler {
     }
 };
 
+struct LTX2Scheduler : SigmaScheduler {
+    int token_count  = 4096;
+    float max_shift  = 2.05f;
+    float base_shift = 0.95f;
+    bool stretch     = true;
+    float terminal   = 0.1f;
+
+    explicit LTX2Scheduler(int token_count, const char* extra_sample_args = nullptr)
+        : token_count(token_count > 0 ? token_count : 4096) {
+        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;
+                }
+            };
+            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());
+                }
+            } 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;
+            }
+        }
+    }
+
+    std::vector<float> get_sigmas(uint32_t n, float /*sigma_min*/, float /*sigma_max*/, t_to_sigma_t /*t_to_sigma*/) override {
+        std::vector<float> sigmas;
+        if (n == 0) {
+            sigmas.push_back(0.0f);
+            return sigmas;
+        }
+
+        constexpr float base_shift_anchor = 1024.0f;
+        constexpr float max_shift_anchor  = 4096.0f;
+        float m                           = (max_shift - base_shift) / (max_shift_anchor - base_shift_anchor);
+        float b                           = base_shift - m * base_shift_anchor;
+        float sigma_shift                 = static_cast<float>(token_count) * m + b;
+        float exp_shift                   = std::exp(sigma_shift);
+        float target_terminal             = std::clamp(terminal, 0.0f, 0.99f);
+
+        LOG_DEBUG("LTX2 scheduler: tokens=%d, shift=%.4f, stretch=%d, terminal=%.4f", token_count, sigma_shift, stretch ? 1 : 0, target_terminal);
+
+        sigmas.reserve(n + 1);
+        for (uint32_t i = 0; i <= n; ++i) {
+            float sigma = 1.0f - static_cast<float>(i) / static_cast<float>(n);
+            if (sigma != 0.0f) {
+                sigma = exp_shift / (exp_shift + (1.0f / sigma - 1.0f));
+            }
+            sigmas.push_back(sigma);
+        }
+
+        if (stretch && sigmas.size() > 2) {
+            float one_minus_last = 1.0f - sigmas[n - 1];
+            float scale_factor   = one_minus_last / (1.0f - target_terminal);
+            if (scale_factor > 1e-8f) {
+                for (uint32_t i = 0; i < n; ++i) {
+                    sigmas[i] = 1.0f - (1.0f - sigmas[i]) / scale_factor;
+                }
+            }
+        }
+
+        sigmas[n] = 0.0f;
+        return sigmas;
+    }
+};
+
 struct Denoiser {
     virtual float sigma_min()                                                        = 0;
     virtual float sigma_max()                                                        = 0;
@@ -491,7 +629,7 @@ struct Denoiser {
     virtual sd::Tensor<float> inverse_noise_scaling(float sigma,
                                                     const sd::Tensor<float>& latent) = 0;
 
-    virtual std::vector<float> get_sigmas(uint32_t n, int /*image_seq_len*/, scheduler_t scheduler_type, SDVersion version) {
+    virtual std::vector<float> get_sigmas(uint32_t n, int image_seq_len, scheduler_t scheduler_type, SDVersion version, const char* extra_sample_args = nullptr) {
         auto bound_t_to_sigma = std::bind(&Denoiser::t_to_sigma, this, std::placeholders::_1);
         std::shared_ptr<SigmaScheduler> scheduler;
         switch (scheduler_type) {
@@ -539,6 +677,10 @@ struct Denoiser {
                 LOG_INFO("get_sigmas with LCM scheduler");
                 scheduler = std::make_shared<LCMScheduler>();
                 break;
+            case LTX2_SCHEDULER:
+                LOG_INFO("get_sigmas with LTX2 scheduler");
+                scheduler = std::make_shared<LTX2Scheduler>(image_seq_len, extra_sample_args);
+                break;
             default:
                 LOG_INFO("get_sigmas with discrete scheduler (default)");
                 scheduler = std::make_shared<DiscreteScheduler>();
@@ -744,15 +886,15 @@ struct Flux2FlowDenoiser : public FluxFlowDenoiser {
         return mu;
     }
 
-    std::vector<float> get_sigmas(uint32_t n, int image_seq_len, scheduler_t scheduler_type, SDVersion version) override {
+    std::vector<float> get_sigmas(uint32_t n, int image_seq_len, scheduler_t scheduler_type, SDVersion version, const char* extra_sample_args = nullptr) override {
         float mu = compute_empirical_mu(n, image_seq_len);
         LOG_DEBUG("Flux2FlowDenoiser: set shift to %.3f", mu);
         set_shift(mu);
-        return Denoiser::get_sigmas(n, image_seq_len, scheduler_type, version);
+        return Denoiser::get_sigmas(n, image_seq_len, scheduler_type, version, extra_sample_args);
     }
 };
 
-typedef std::function<sd::Tensor<float>(const sd::Tensor<float>&, float, int)> denoise_cb_t;
+typedef std::function<sd::Tensor<float>(const sd::Tensor<float>&, float, int, sd::Tensor<float>*)> denoise_cb_t;
 
 static std::pair<float, float> get_ancestral_step(float sigma_from,
                                                   float sigma_to,
@@ -808,48 +950,48 @@ static std::tuple<float, float, float> get_ancestral_step_flow(float sigma_from,
     return {sigma_down, sigma_up, alpha_scale};
 }
 
+static std::tuple<float, float, float> get_ancestral_step(float sigma_from,
+                                                          float sigma_to,
+                                                          float eta,
+                                                          bool is_flow_denoiser) {
+    if (is_flow_denoiser) {
+        return get_ancestral_step_flow(sigma_from, sigma_to, eta);
+    } else {
+        auto [sigma_down, sigma_up] = get_ancestral_step(sigma_from, sigma_to, eta);
+        return {sigma_down, sigma_up, 1.0f};
+    }
+}
+
 static sd::Tensor<float> sample_euler_ancestral(denoise_cb_t model,
                                                 sd::Tensor<float> x,
                                                 const std::vector<float>& sigmas,
-                                                std::shared_ptr<RNG> rng,
-                                                float eta) {
+                                                std::shared_ptr<RNG> rng = nullptr,
+                                                bool is_flow_denoiser    = false,
+                                                float eta                = 0.f) {
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
         float sigma       = sigmas[i];
-        auto denoised_opt = model(x, sigma, i + 1);
+        float sigma_to    = sigmas[i + 1];
+        auto denoised_opt = model(x, sigma, i + 1, nullptr);
         if (denoised_opt.empty()) {
             return {};
         }
-        sd::Tensor<float> denoised  = std::move(denoised_opt);
-        sd::Tensor<float> d         = (x - denoised) / sigma;
-        auto [sigma_down, sigma_up] = get_ancestral_step(sigmas[i], sigmas[i + 1], eta);
-        x += d * (sigma_down - sigmas[i]);
-        if (sigmas[i + 1] > 0) {
-            x += sd::Tensor<float>::randn_like(x, rng) * sigma_up;
-        }
-    }
-    return x;
-}
-
-static sd::Tensor<float> sample_euler_flow(denoise_cb_t model,
-                                           sd::Tensor<float> x,
-                                           const std::vector<float>& sigmas,
-                                           std::shared_ptr<RNG> rng,
-                                           float eta) {
-    int steps = static_cast<int>(sigmas.size()) - 1;
-    for (int i = 0; i < steps; i++) {
-        float sigma       = sigmas[i];
-        auto denoised_opt = model(x, sigma, i + 1);
-        if (denoised_opt.empty()) {
-            return {};
-        }
-        sd::Tensor<float> denoised               = std::move(denoised_opt);
-        auto [sigma_down, sigma_up, alpha_scale] = get_ancestral_step_flow(sigma, sigmas[i + 1], eta);
-        float sigma_ratio                        = sigma_down / sigma;
-        x                                        = sigma_ratio * x + (1.0f - sigma_ratio) * denoised;
-
-        if (sigma_up > 0.0f) {
-            x = alpha_scale * x + sd::Tensor<float>::randn_like(x, rng) * sigma_up;
+        sd::Tensor<float> denoised = std::move(denoised_opt);
+        if (sigma_to == 0.f) {
+            x = denoised;
+        } else if (eta == 0.f) {
+            float sigma_ratio = sigma_to / sigma;
+            x                 = sigma_ratio * x + (1.0 - sigma_ratio) * denoised;
+        } else {
+            auto [sigma_down, sigma_up, alpha_scale] = get_ancestral_step(sigma, sigma_to, eta, is_flow_denoiser);
+            float sigma_ratio                        = sigma_down / sigma;
+            x                                        = sigma_ratio * x + (1.0f - sigma_ratio) * denoised;
+            if (sigma_up > 0.f) {
+                if (is_flow_denoiser) {
+                    x *= alpha_scale;
+                }
+                x += sd::Tensor<float>::randn_like(x, rng) * sigma_up;
+            }
         }
     }
     return x;
@@ -861,7 +1003,7 @@ static sd::Tensor<float> sample_euler(denoise_cb_t model,
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
         float sigma       = sigmas[i];
-        auto denoised_opt = model(x, sigma, i + 1);
+        auto denoised_opt = model(x, sigma, i + 1, nullptr);
         if (denoised_opt.empty()) {
             return {};
         }
@@ -877,7 +1019,7 @@ static sd::Tensor<float> sample_heun(denoise_cb_t model,
                                      const std::vector<float>& sigmas) {
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
-        auto denoised_opt = model(x, sigmas[i], -(i + 1));
+        auto denoised_opt = model(x, sigmas[i], -(i + 1), nullptr);
         if (denoised_opt.empty()) {
             return {};
         }
@@ -888,7 +1030,7 @@ static sd::Tensor<float> sample_heun(denoise_cb_t model,
             x += d * dt;
         } else {
             sd::Tensor<float> x2 = x + d * dt;
-            auto denoised2_opt   = model(x2, sigmas[i + 1], i + 1);
+            auto denoised2_opt   = model(x2, sigmas[i + 1], i + 1, nullptr);
             if (denoised2_opt.empty()) {
                 return {};
             }
@@ -905,7 +1047,7 @@ static sd::Tensor<float> sample_dpm2(denoise_cb_t model,
                                      const std::vector<float>& sigmas) {
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
-        auto denoised_opt = model(x, sigmas[i], -(i + 1));
+        auto denoised_opt = model(x, sigmas[i], -(i + 1), nullptr);
         if (denoised_opt.empty()) {
             return {};
         }
@@ -918,7 +1060,7 @@ static sd::Tensor<float> sample_dpm2(denoise_cb_t model,
             float dt_1           = sigma_mid - sigmas[i];
             float dt_2           = sigmas[i + 1] - sigmas[i];
             sd::Tensor<float> x2 = x + d * dt_1;
-            auto denoised2_opt   = model(x2, sigma_mid, i + 1);
+            auto denoised2_opt   = model(x2, sigma_mid, i + 1, nullptr);
             if (denoised2_opt.empty()) {
                 return {};
             }
@@ -939,7 +1081,7 @@ static sd::Tensor<float> sample_dpmpp_2s_ancestral(denoise_cb_t model,
 
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
-        auto denoised_opt = model(x, sigmas[i], -(i + 1));
+        auto denoised_opt = model(x, sigmas[i], -(i + 1), nullptr);
         if (denoised_opt.empty()) {
             return {};
         }
@@ -955,7 +1097,7 @@ static sd::Tensor<float> sample_dpmpp_2s_ancestral(denoise_cb_t model,
             float s              = t + 0.5f * h;
             float sigma_s        = sigma_fn(s);
             sd::Tensor<float> x2 = (sigma_s / sigma_fn(t)) * x - (exp(-h * 0.5f) - 1) * denoised;
-            auto denoised2_opt   = model(x2, sigma_s, i + 1);
+            auto denoised2_opt   = model(x2, sigma_s, i + 1, nullptr);
             if (denoised2_opt.empty()) {
                 return {};
             }
@@ -982,7 +1124,7 @@ static sd::Tensor<float> sample_dpmpp_2s_ancestral_flow(denoise_cb_t model,
 
         bool opt_first_step = (1.0 - sigma < 1e-6);
 
-        auto denoised_opt = model(x, sigma, (opt_first_step ? 1 : -1) * (i + 1));
+        auto denoised_opt = model(x, sigma, (opt_first_step ? 1 : -1) * (i + 1), nullptr);
         if (denoised_opt.empty()) {
             return {};
         }
@@ -1011,8 +1153,8 @@ static sd::Tensor<float> sample_dpmpp_2s_ancestral_flow(denoise_cb_t model,
                 // so sigma_s = 1 = sigma, and sigma_s_i_ratio = sigma_s / sigma = 1
                 // u = (x*sigma_s_i_ratio)+(denoised*(1.0f-sigma_s_i_ratio))
                 //   = (x*1)+(denoised*0) = x
-                // so D_i = model(u, sigma_s, i + 1)
-                //        = model(x, sigma,   i + 1)
+                // so D_i = model(u, sigma_s, i + 1, nullptr)
+                //        = model(x, sigma,   i + 1, nullptr)
                 //        = denoised
                 D_i = denoised;
 
@@ -1045,7 +1187,7 @@ static sd::Tensor<float> sample_dpmpp_2s_ancestral_flow(denoise_cb_t model,
                 float sigma_s_i_ratio = sigma_s / sigma;
                 sd::Tensor<float> u   = (x * sigma_s_i_ratio) + (denoised * (1.0f - sigma_s_i_ratio));
 
-                auto denoised2_opt = model(u, sigma_s, i + 1);
+                auto denoised2_opt = model(u, sigma_s, i + 1, nullptr);
                 if (denoised2_opt.empty()) {
                     return {};
                 }
@@ -1072,7 +1214,7 @@ static sd::Tensor<float> sample_dpmpp_2m(denoise_cb_t model,
 
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
-        auto denoised_opt = model(x, sigmas[i], i + 1);
+        auto denoised_opt = model(x, sigmas[i], i + 1, nullptr);
         if (denoised_opt.empty()) {
             return {};
         }
@@ -1104,7 +1246,7 @@ static sd::Tensor<float> sample_dpmpp_2m_v2(denoise_cb_t model,
 
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
-        auto denoised_opt = model(x, sigmas[i], i + 1);
+        auto denoised_opt = model(x, sigmas[i], i + 1, nullptr);
         if (denoised_opt.empty()) {
             return {};
         }
@@ -1136,10 +1278,83 @@ static sd::Tensor<float> sample_lcm(denoise_cb_t model,
                                     sd::Tensor<float> x,
                                     const std::vector<float>& sigmas,
                                     std::shared_ptr<RNG> rng,
-                                    bool is_flow_denoiser) {
+                                    bool is_flow_denoiser,
+                                    const char* extra_sample_args = nullptr) {
+    struct LCMSampleArgs {
+        float noise_clip_std    = 0.0f;
+        float noise_scale_start = 1.0f;
+        float noise_scale_end   = 1.0f;
+    };
+
+    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);
+    };
+
+    LCMSampleArgs args;
+    if (extra_sample_args != nullptr && extra_sample_args[0] != '\0') {
+        std::string raw(extra_sample_args);
+        size_t start                   = 0;
+        bool noise_scale_end_was_set   = false;
+        bool noise_scale_start_was_set = false;
+        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 lcm extra sample arg '%s'", token.c_str());
+                return;
+            }
+
+            std::string key   = trim(token.substr(0, eq));
+            std::string value = trim(token.substr(eq + 1));
+            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'", token.c_str());
+                    return;
+                }
+            } catch (const std::exception&) {
+                LOG_WARN("ignoring invalid lcm extra sample arg '%s'", token.c_str());
+                return;
+            }
+
+            if (key == "noise_clip_std") {
+                args.noise_clip_std = parsed;
+            } else if (key == "noise_scale_start") {
+                args.noise_scale_start    = parsed;
+                noise_scale_start_was_set = true;
+            } else if (key == "noise_scale_end") {
+                args.noise_scale_end    = parsed;
+                noise_scale_end_was_set = true;
+            } else {
+                LOG_WARN("ignoring unknown lcm extra sample arg '%s'", key.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;
+            }
+        }
+        if (noise_scale_start_was_set && !noise_scale_end_was_set) {
+            args.noise_scale_end = args.noise_scale_start;
+        }
+    }
+
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
-        auto denoised_opt = model(x, sigmas[i], i + 1);
+        auto denoised_opt = model(x, sigmas[i], i + 1, nullptr);
         if (denoised_opt.empty()) {
             return {};
         }
@@ -1148,7 +1363,27 @@ static sd::Tensor<float> sample_lcm(denoise_cb_t model,
             if (is_flow_denoiser) {
                 x *= (1 - sigmas[i + 1]);
             }
-            x += sd::Tensor<float>::randn_like(x, rng) * sigmas[i + 1];
+            auto noise = sd::Tensor<float>::randn_like(x, rng);
+            if (args.noise_clip_std > 0.0f && noise.numel() > 0) {
+                double mean = 0.0;
+                for (int64_t j = 0; j < noise.numel(); ++j) {
+                    mean += static_cast<double>(noise[j]);
+                }
+                mean /= static_cast<double>(noise.numel());
+
+                double variance = 0.0;
+                for (int64_t j = 0; j < noise.numel(); ++j) {
+                    double centered = static_cast<double>(noise[j]) - mean;
+                    variance += centered * centered;
+                }
+                variance /= static_cast<double>(noise.numel());
+
+                float clip_val = args.noise_clip_std * static_cast<float>(std::sqrt(variance));
+                noise          = sd::ops::clamp(noise, -clip_val, clip_val);
+            }
+            float t           = steps > 1 ? static_cast<float>(i) / static_cast<float>(steps - 1) : 0.0f;
+            float noise_scale = args.noise_scale_start + (args.noise_scale_end - args.noise_scale_start) * t;
+            x += noise * (sigmas[i + 1] * noise_scale);
         }
     }
     return x;
@@ -1165,7 +1400,7 @@ static sd::Tensor<float> sample_ipndm(denoise_cb_t model,
         float sigma      = sigmas[i];
         float sigma_next = sigmas[i + 1];
 
-        auto denoised_opt = model(x, sigma, i + 1);
+        auto denoised_opt = model(x, sigma, i + 1, nullptr);
         if (denoised_opt.empty()) {
             return {};
         }
@@ -1209,7 +1444,7 @@ static sd::Tensor<float> sample_ipndm_v(denoise_cb_t model,
         float sigma  = sigmas[i];
         float t_next = sigmas[i + 1];
 
-        auto denoised_opt = model(x, sigma, i + 1);
+        auto denoised_opt = model(x, sigma, i + 1, nullptr);
         if (denoised_opt.empty()) {
             return {};
         }
@@ -1247,6 +1482,7 @@ static sd::Tensor<float> sample_res_multistep(denoise_cb_t model,
                                               sd::Tensor<float> x,
                                               const std::vector<float>& sigmas,
                                               std::shared_ptr<RNG> rng,
+                                              bool is_flow_denoiser,
                                               float eta) {
     sd::Tensor<float> old_denoised = x;
     bool have_old_sigma            = false;
@@ -1270,15 +1506,16 @@ static sd::Tensor<float> sample_res_multistep(denoise_cb_t model,
 
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
-        auto denoised_opt = model(x, sigmas[i], i + 1);
+        auto denoised_opt = model(x, sigmas[i], i + 1, nullptr);
         if (denoised_opt.empty()) {
             return {};
         }
         sd::Tensor<float> denoised = std::move(denoised_opt);
 
-        float sigma_from            = sigmas[i];
-        float sigma_to              = sigmas[i + 1];
-        auto [sigma_down, sigma_up] = get_ancestral_step(sigma_from, sigma_to, eta);
+        float sigma_from = sigmas[i];
+        float sigma_to   = sigmas[i + 1];
+
+        auto [sigma_down, sigma_up, alpha_scale] = get_ancestral_step(sigma_from, sigma_to, eta, is_flow_denoiser);
 
         if (sigma_down == 0.0f || !have_old_sigma) {
             x += ((x - denoised) / sigma_from) * (sigma_down - sigma_from);
@@ -1305,7 +1542,10 @@ static sd::Tensor<float> sample_res_multistep(denoise_cb_t model,
             x = sigma_fn(h) * x + h * (b1 * denoised + b2 * old_denoised);
         }
 
-        if (sigmas[i + 1] > 0 && sigma_up > 0.0f) {
+        if (sigma_to > 0.0f && sigma_up > 0.0f) {
+            if (is_flow_denoiser) {
+                x *= alpha_scale;
+            }
             x += sd::Tensor<float>::randn_like(x, rng) * sigma_up;
         }
 
@@ -1320,6 +1560,7 @@ static sd::Tensor<float> sample_res_2s(denoise_cb_t model,
                                        sd::Tensor<float> x,
                                        const std::vector<float>& sigmas,
                                        std::shared_ptr<RNG> rng,
+                                       bool is_flow_denoiser,
                                        float eta) {
     const float c2 = 0.5f;
     auto t_fn      = [](float sigma) -> float { return -logf(sigma); };
@@ -1342,13 +1583,13 @@ static sd::Tensor<float> sample_res_2s(denoise_cb_t model,
         float sigma_from = sigmas[i];
         float sigma_to   = sigmas[i + 1];
 
-        auto denoised_opt = model(x, sigma_from, -(i + 1));
+        auto denoised_opt = model(x, sigma_from, -(i + 1), nullptr);
         if (denoised_opt.empty()) {
             return {};
         }
         sd::Tensor<float> denoised = std::move(denoised_opt);
 
-        auto [sigma_down, sigma_up] = get_ancestral_step(sigma_from, sigma_to, eta);
+        auto [sigma_down, sigma_up, alpha_scale] = get_ancestral_step(sigma_from, sigma_to, eta, is_flow_denoiser);
 
         sd::Tensor<float> x0 = x;
         if (sigma_down == 0.0f || sigma_from == 0.0f) {
@@ -1368,7 +1609,7 @@ static sd::Tensor<float> sample_res_2s(denoise_cb_t model,
             sd::Tensor<float> eps1 = denoised - x0;
             sd::Tensor<float> x2   = x0 + eps1 * (h * a21);
 
-            auto denoised2_opt = model(x2, sigma_c2, i + 1);
+            auto denoised2_opt = model(x2, sigma_c2, i + 1, nullptr);
             if (denoised2_opt.empty()) {
                 return {};
             }
@@ -1377,7 +1618,10 @@ static sd::Tensor<float> sample_res_2s(denoise_cb_t model,
             x                           = x0 + h * (b1 * eps1 + b2 * eps2);
         }
 
-        if (sigmas[i + 1] > 0 && sigma_up > 0.0f) {
+        if (sigma_to > 0.0f && sigma_up > 0.0f) {
+            if (is_flow_denoiser) {
+                x *= alpha_scale;
+            }
             x += sd::Tensor<float>::randn_like(x, rng) * sigma_up;
         }
     }
@@ -1442,7 +1686,7 @@ static sd::Tensor<float> sample_er_sde(denoise_cb_t model,
 
     int steps = static_cast<int>(sigmas.size()) - 1;
     for (int i = 0; i < steps; i++) {
-        sd::Tensor<float> denoised = model(x, sigmas[i], i + 1);
+        sd::Tensor<float> denoised = model(x, sigmas[i], i + 1, nullptr);
         if (denoised.empty()) {
             return {};
         }
@@ -1518,46 +1762,6 @@ static sd::Tensor<float> sample_er_sde(denoise_cb_t model,
     return x;
 }
 
-static sd::Tensor<float> sample_ddim_trailing(denoise_cb_t model,
-                                              sd::Tensor<float> x,
-                                              const std::vector<float>& sigmas,
-                                              std::shared_ptr<RNG> rng,
-                                              float eta) {
-    int steps = static_cast<int>(sigmas.size()) - 1;
-    for (int i = 0; i < steps; i++) {
-        float sigma    = sigmas[i];
-        float sigma_to = sigmas[i + 1];
-
-        auto model_output_opt = model(x, sigma, i + 1);
-        if (model_output_opt.empty()) {
-            return {};
-        }
-        sd::Tensor<float> model_output = std::move(model_output_opt);
-        model_output                   = (x - model_output) * (1.0f / sigma);
-
-        float alpha_prod_t      = 1.0f / (sigma * sigma + 1.0f);
-        float alpha_prod_t_prev = 1.0f / (sigma_to * sigma_to + 1.0f);
-        float beta_prod_t       = 1.0f - alpha_prod_t;
-
-        sd::Tensor<float> pred_original_sample = ((x / std::sqrt(sigma * sigma + 1)) -
-                                                  std::sqrt(beta_prod_t) * model_output) *
-                                                 (1.0f / std::sqrt(alpha_prod_t));
-
-        float beta_prod_t_prev = 1.0f - alpha_prod_t_prev;
-        float variance         = (beta_prod_t_prev / beta_prod_t) *
-                         (1.0f - alpha_prod_t / alpha_prod_t_prev);
-        float std_dev_t = eta * std::sqrt(variance);
-
-        x = pred_original_sample +
-            std::sqrt((1.0f - alpha_prod_t_prev - std::pow(std_dev_t, 2)) / alpha_prod_t_prev) * model_output;
-
-        if (eta > 0) {
-            x += std_dev_t / std::sqrt(alpha_prod_t_prev) * sd::Tensor<float>::randn_like(x, rng);
-        }
-    }
-    return x;
-}
-
 static sd::Tensor<float> sample_tcd(denoise_cb_t model,
                                     sd::Tensor<float> x,
                                     const std::vector<float>& sigmas,
@@ -1600,12 +1804,12 @@ static sd::Tensor<float> sample_tcd(denoise_cb_t model,
         int timestep_s    = (int)floor((1 - eta) * prev_timestep);
         float sigma       = sigmas[i];
 
-        auto model_output_opt = model(x, sigma, i + 1);
-        if (model_output_opt.empty()) {
+        auto denoised_opt = model(x, sigma, i + 1, nullptr);
+        if (denoised_opt.empty()) {
             return {};
         }
-        sd::Tensor<float> model_output = std::move(model_output_opt);
-        model_output                   = (x - model_output) * (1.0f / sigma);
+        sd::Tensor<float> denoised = std::move(denoised_opt);
+        sd::Tensor<float> d        = (x - denoised) / sigma;
 
         float alpha_prod_t      = 1.0f / (sigma * sigma + 1.0f);
         float beta_prod_t       = 1.0f - alpha_prod_t;
@@ -1613,12 +1817,8 @@ static sd::Tensor<float> sample_tcd(denoise_cb_t model,
         float alpha_prod_s      = static_cast<float>(alphas_cumprod[timestep_s]);
         float beta_prod_s       = 1.0f - alpha_prod_s;
 
-        sd::Tensor<float> pred_original_sample = ((x / std::sqrt(sigma * sigma + 1)) -
-                                                  std::sqrt(beta_prod_t) * model_output) *
-                                                 (1.0f / std::sqrt(alpha_prod_t));
-
-        x = std::sqrt(alpha_prod_s / alpha_prod_t_prev) * pred_original_sample +
-            std::sqrt(beta_prod_s / alpha_prod_t_prev) * model_output;
+        x = std::sqrt(alpha_prod_s / alpha_prod_t_prev) * denoised +
+            std::sqrt(beta_prod_s / alpha_prod_t_prev) * d;
 
         if (eta > 0 && sigma_to > 0.0f) {
             x = std::sqrt(alpha_prod_t_prev / alpha_prod_s) * x +
@@ -1628,6 +1828,56 @@ static sd::Tensor<float> sample_tcd(denoise_cb_t model,
     return x;
 }
 
+static sd::Tensor<float> sample_euler_cfg_pp(denoise_cb_t model,
+                                             sd::Tensor<float> x,
+                                             const std::vector<float>& sigmas) {
+    int steps = static_cast<int>(sigmas.size()) - 1;
+    for (int i = 0; i < steps; i++) {
+        float sigma = sigmas[i];
+        sd::Tensor<float> uncond_denoised;
+
+        auto denoised_opt = model(x, sigma, i + 1, &uncond_denoised);
+        if (denoised_opt.empty() || uncond_denoised.empty()) {
+            return {};
+        }
+
+        sd::Tensor<float> denoised = std::move(denoised_opt);
+        sd::Tensor<float> d        = (x - uncond_denoised) / sigma;
+
+        x = denoised + d * sigmas[i + 1];
+    }
+    return x;
+}
+
+static sd::Tensor<float> sample_euler_ancestral_cfg_pp(denoise_cb_t model,
+                                                       sd::Tensor<float> x,
+                                                       const std::vector<float>& sigmas,
+                                                       std::shared_ptr<RNG> rng,
+                                                       float eta) {
+    int steps = static_cast<int>(sigmas.size()) - 1;
+    for (int i = 0; i < steps; i++) {
+        float sigma = sigmas[i];
+        sd::Tensor<float> uncond_denoised;
+
+        auto denoised_opt = model(x, sigma, i + 1, &uncond_denoised);
+        if (denoised_opt.empty() || uncond_denoised.empty()) {
+            return {};
+        }
+
+        sd::Tensor<float> denoised = std::move(denoised_opt);
+        sd::Tensor<float> d        = (x - uncond_denoised) / sigma;
+
+        auto [sigma_down, sigma_up] = get_ancestral_step(sigmas[i], sigmas[i + 1], eta);
+
+        x = denoised + d * sigma_down;
+
+        if (sigmas[i + 1] > 0) {
+            x += sd::Tensor<float>::randn_like(x, rng) * sigma_up;
+        }
+    }
+    return x;
+}
+
 // 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,
@@ -1635,13 +1885,11 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
                                             std::vector<float> sigmas,
                                             std::shared_ptr<RNG> rng,
                                             float eta,
-                                            bool is_flow_denoiser) {
+                                            bool is_flow_denoiser,
+                                            const char* extra_sample_args) {
     switch (method) {
         case EULER_A_SAMPLE_METHOD:
-            if (is_flow_denoiser)
-                return sample_euler_flow(model, std::move(x), sigmas, rng, eta);
-            else
-                return sample_euler_ancestral(model, std::move(x), sigmas, rng, eta);
+            return sample_euler_ancestral(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
         case EULER_SAMPLE_METHOD:
             return sample_euler(model, std::move(x), sigmas);
         case HEUN_SAMPLE_METHOD:
@@ -1658,21 +1906,26 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
         case DPMPP2Mv2_SAMPLE_METHOD:
             return sample_dpmpp_2m_v2(model, std::move(x), sigmas);
         case LCM_SAMPLE_METHOD:
-            return sample_lcm(model, std::move(x), sigmas, rng, is_flow_denoiser);
+            return sample_lcm(model, std::move(x), sigmas, rng, is_flow_denoiser, extra_sample_args);
         case IPNDM_SAMPLE_METHOD:
             return sample_ipndm(model, std::move(x), sigmas);
         case IPNDM_V_SAMPLE_METHOD:
             return sample_ipndm_v(model, std::move(x), sigmas);
         case RES_MULTISTEP_SAMPLE_METHOD:
-            return sample_res_multistep(model, std::move(x), sigmas, rng, eta);
+            return sample_res_multistep(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
         case RES_2S_SAMPLE_METHOD:
-            return sample_res_2s(model, std::move(x), sigmas, rng, eta);
+            return sample_res_2s(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
         case ER_SDE_SAMPLE_METHOD:
             return sample_er_sde(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
         case DDIM_TRAILING_SAMPLE_METHOD:
-            return sample_ddim_trailing(model, std::move(x), sigmas, rng, eta);
+            // DDIM is equivalent to Euler Ancestral with the Simple scheduler
+            return sample_euler_ancestral(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
         case TCD_SAMPLE_METHOD:
             return sample_tcd(model, std::move(x), sigmas, rng, eta);
+        case EULER_CFG_PP_SAMPLE_METHOD:
+            return sample_euler_cfg_pp(model, std::move(x), sigmas);
+        case EULER_A_CFG_PP_SAMPLE_METHOD:
+            return sample_euler_ancestral_cfg_pp(model, std::move(x), sigmas, rng, eta);
         default:
             return {};
     }

src/diffusion_model.hpp

@@ -5,6 +5,8 @@
 #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"
@@ -13,22 +15,33 @@
 #include "z_image.hpp"
 
 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 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;
-    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;
-    const std::vector<int>* skip_layers               = nullptr;
+    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;
+    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;
 };
 
 template <typename T>
@@ -49,6 +62,7 @@ struct DiffusionModel {
     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;
 };
 
@@ -56,10 +70,10 @@ struct UNetModel : public DiffusionModel {
     UNetModelRunner unet;
 
     UNetModel(ggml_backend_t backend,
-              bool offload_params_to_cpu,
+              ggml_backend_t params_backend,
               const String2TensorStorage& tensor_storage_map = {},
               SDVersion version                              = VERSION_SD1)
-        : unet(backend, offload_params_to_cpu, tensor_storage_map, "model.diffusion_model", version) {
+        : unet(backend, params_backend, tensor_storage_map, "model.diffusion_model", version) {
     }
 
     std::string get_desc() override {
@@ -98,6 +112,10 @@ struct UNetModel : public DiffusionModel {
         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);
     }
@@ -123,9 +141,9 @@ struct MMDiTModel : public DiffusionModel {
     MMDiTRunner mmdit;
 
     MMDiTModel(ggml_backend_t backend,
-               bool offload_params_to_cpu,
+               ggml_backend_t params_backend,
                const String2TensorStorage& tensor_storage_map = {})
-        : mmdit(backend, offload_params_to_cpu, tensor_storage_map, "model.diffusion_model") {
+        : mmdit(backend, params_backend, tensor_storage_map, "model.diffusion_model") {
     }
 
     std::string get_desc() override {
@@ -164,6 +182,10 @@ struct MMDiTModel : public DiffusionModel {
         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);
     }
@@ -186,11 +208,11 @@ struct FluxModel : public DiffusionModel {
     Flux::FluxRunner flux;
 
     FluxModel(ggml_backend_t backend,
-              bool offload_params_to_cpu,
+              ggml_backend_t params_backend,
               const String2TensorStorage& tensor_storage_map = {},
               SDVersion version                              = VERSION_FLUX,
               bool use_mask                                  = false)
-        : flux(backend, offload_params_to_cpu, tensor_storage_map, "model.diffusion_model", version, use_mask) {
+        : flux(backend, params_backend, tensor_storage_map, "model.diffusion_model", version, use_mask) {
     }
 
     std::string get_desc() override {
@@ -229,6 +251,10 @@ struct FluxModel : public DiffusionModel {
         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);
     }
@@ -257,10 +283,10 @@ struct AnimaModel : public DiffusionModel {
     Anima::AnimaRunner anima;
 
     AnimaModel(ggml_backend_t backend,
-               bool offload_params_to_cpu,
+               ggml_backend_t params_backend,
                const String2TensorStorage& tensor_storage_map = {},
                const std::string prefix                       = "model.diffusion_model")
-        : prefix(prefix), anima(backend, offload_params_to_cpu, tensor_storage_map, prefix) {
+        : prefix(prefix), anima(backend, params_backend, tensor_storage_map, prefix) {
     }
 
     std::string get_desc() override {
@@ -299,6 +325,10 @@ struct AnimaModel : public DiffusionModel {
         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);
     }
@@ -321,11 +351,11 @@ struct WanModel : public DiffusionModel {
     WAN::WanRunner wan;
 
     WanModel(ggml_backend_t backend,
-             bool offload_params_to_cpu,
+             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, offload_params_to_cpu, tensor_storage_map, prefix, version) {
+        : prefix(prefix), wan(backend, params_backend, tensor_storage_map, prefix, version) {
     }
 
     std::string get_desc() override {
@@ -364,6 +394,10 @@ struct WanModel : public DiffusionModel {
         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);
     }
@@ -389,12 +423,12 @@ struct QwenImageModel : public DiffusionModel {
     Qwen::QwenImageRunner qwen_image;
 
     QwenImageModel(ggml_backend_t backend,
-                   bool offload_params_to_cpu,
+                   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, offload_params_to_cpu, tensor_storage_map, prefix, version, zero_cond_t) {
+        : prefix(prefix), qwen_image(backend, params_backend, tensor_storage_map, prefix, version, zero_cond_t) {
     }
 
     std::string get_desc() override {
@@ -433,6 +467,10 @@ struct QwenImageModel : public DiffusionModel {
         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);
     }
@@ -451,16 +489,92 @@ struct QwenImageModel : public DiffusionModel {
     }
 };
 
+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,
-                bool offload_params_to_cpu,
+                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, offload_params_to_cpu, tensor_storage_map, prefix, version) {
+        : prefix(prefix), z_image(backend, params_backend, tensor_storage_map, prefix, version) {
     }
 
     std::string get_desc() override {
@@ -499,6 +613,10 @@ struct ZImageModel : public DiffusionModel {
         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);
     }
@@ -522,10 +640,10 @@ struct ErnieImageModel : public DiffusionModel {
     ErnieImage::ErnieImageRunner ernie_image;
 
     ErnieImageModel(ggml_backend_t backend,
-                    bool offload_params_to_cpu,
+                    ggml_backend_t params_backend,
                     const String2TensorStorage& tensor_storage_map = {},
                     const std::string prefix                       = "model.diffusion_model")
-        : prefix(prefix), ernie_image(backend, offload_params_to_cpu, tensor_storage_map, prefix) {
+        : prefix(prefix), ernie_image(backend, params_backend, tensor_storage_map, prefix) {
     }
 
     std::string get_desc() override {
@@ -564,6 +682,10 @@ struct ErnieImageModel : public DiffusionModel {
         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);
     }
@@ -579,4 +701,75 @@ struct ErnieImageModel : public DiffusionModel {
     }
 };
 
+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));
+    }
+};
+
 #endif

src/ernie_image.hpp

@@ -295,7 +295,9 @@ namespace ErnieImage {
             auto c      = time_embedding->forward(ctx, sample);  // [N, hidden_size]
 
             auto mod_params = adaLN_mod->forward(ctx, ggml_silu(ctx->ggml_ctx, c));  // [N, 6 * hidden_size]
-            auto chunks     = ggml_ext_chunk(ctx->ggml_ctx, mod_params, 6, 0);
+            sd::ggml_graph_cut::mark_graph_cut(hidden_states, "ernie_image.prelude", "hidden_states");
+            // sd::ggml_graph_cut::mark_graph_cut(mod_params, "ernie_image.prelude", "mod_params");
+            auto chunks = ggml_ext_chunk(ctx->ggml_ctx, mod_params, 6, 0);
             std::vector<ggml_tensor*> temb;
             temb.reserve(6);
             for (auto chunk : chunks) {
@@ -305,6 +307,7 @@ namespace ErnieImage {
             for (int i = 0; i < params.num_layers; i++) {
                 auto layer    = std::dynamic_pointer_cast<ErnieImageSharedAdaLNBlock>(blocks["layers." + std::to_string(i)]);
                 hidden_states = layer->forward(ctx, hidden_states, pe, temb);
+                sd::ggml_graph_cut::mark_graph_cut(hidden_states, "ernie_image.layers." + std::to_string(i), "hidden_states");
             }
 
             hidden_states = final_norm->forward(ctx, hidden_states, c);
@@ -328,10 +331,10 @@ namespace ErnieImage {
         std::vector<float> pe_vec;
 
         ErnieImageRunner(ggml_backend_t backend,
-                         bool offload_params_to_cpu,
+                         ggml_backend_t params_backend,
                          const String2TensorStorage& tensor_storage_map = {},
                          const std::string prefix                       = "")
-            : GGMLRunner(backend, offload_params_to_cpu) {
+            : GGMLRunner(backend, params_backend) {
             ernie_params.num_layers = 0;
             for (const auto& [name, tensor_storage] : tensor_storage_map) {
                 if (!starts_with(name, prefix)) {

src/esrgan.hpp

@@ -124,27 +124,33 @@ public:
         auto conv_hr    = std::dynamic_pointer_cast<Conv2d>(blocks["conv_hr"]);
         auto conv_last  = std::dynamic_pointer_cast<Conv2d>(blocks["conv_last"]);
 
-        auto feat      = conv_first->forward(ctx, x);
+        auto feat = conv_first->forward(ctx, x);
+        sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.prelude", "feat");
         auto body_feat = feat;
         for (int i = 0; i < num_block; i++) {
             std::string name = "body." + std::to_string(i);
             auto block       = std::dynamic_pointer_cast<RRDB>(blocks[name]);
 
             body_feat = block->forward(ctx, body_feat);
+            sd::ggml_graph_cut::mark_graph_cut(body_feat, "esrgan.body." + std::to_string(i), "feat");
         }
         body_feat = conv_body->forward(ctx, body_feat);
         feat      = ggml_add(ctx->ggml_ctx, feat, body_feat);
+        sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.body.out", "feat");
         // upsample
         if (scale >= 2) {
             auto conv_up1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv_up1"]);
             feat          = lrelu(ctx, conv_up1->forward(ctx, ggml_upscale(ctx->ggml_ctx, feat, 2, GGML_SCALE_MODE_NEAREST)));
+            sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.up1", "feat");
             if (scale == 4) {
                 auto conv_up2 = std::dynamic_pointer_cast<Conv2d>(blocks["conv_up2"]);
                 feat          = lrelu(ctx, conv_up2->forward(ctx, ggml_upscale(ctx->ggml_ctx, feat, 2, GGML_SCALE_MODE_NEAREST)));
+                sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.up2", "feat");
             }
         }
         // for all scales
         auto out = conv_last->forward(ctx, lrelu(ctx, conv_hr->forward(ctx, feat)));
+        sd::ggml_graph_cut::mark_graph_cut(out, "esrgan.final", "out");
         return out;
     }
 };
@@ -155,10 +161,10 @@ struct ESRGAN : public GGMLRunner {
     int tile_size = 128;  // avoid cuda OOM for 4gb VRAM
 
     ESRGAN(ggml_backend_t backend,
-           bool offload_params_to_cpu,
+           ggml_backend_t params_backend,
            int tile_size                                  = 128,
            const String2TensorStorage& tensor_storage_map = {})
-        : GGMLRunner(backend, offload_params_to_cpu) {
+        : GGMLRunner(backend, params_backend) {
         this->tile_size = tile_size;
     }
 

src/flux.hpp

@@ -928,6 +928,9 @@ namespace Flux {
             }
 
             txt = txt_in->forward(ctx, txt);
+            sd::ggml_graph_cut::mark_graph_cut(img, "flux.prelude", "img");
+            sd::ggml_graph_cut::mark_graph_cut(txt, "flux.prelude", "txt");
+            sd::ggml_graph_cut::mark_graph_cut(vec, "flux.prelude", "vec");
 
             for (int i = 0; i < params.depth; i++) {
                 if (skip_layers.size() > 0 && std::find(skip_layers.begin(), skip_layers.end(), i) != skip_layers.end()) {
@@ -939,6 +942,8 @@ namespace Flux {
                 auto img_txt = block->forward(ctx, img, txt, vec, pe, txt_img_mask, ds_img_mods, ds_txt_mods);
                 img          = img_txt.first;   // [N, n_img_token, hidden_size]
                 txt          = img_txt.second;  // [N, n_txt_token, hidden_size]
+                sd::ggml_graph_cut::mark_graph_cut(img, "flux.double_blocks." + std::to_string(i), "img");
+                sd::ggml_graph_cut::mark_graph_cut(txt, "flux.double_blocks." + std::to_string(i), "txt");
             }
 
             auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1);  // [N, n_txt_token + n_img_token, hidden_size]
@@ -949,6 +954,7 @@ namespace Flux {
                 auto block = std::dynamic_pointer_cast<SingleStreamBlock>(blocks["single_blocks." + std::to_string(i)]);
 
                 txt_img = block->forward(ctx, txt_img, vec, pe, txt_img_mask, ss_mods);
+                sd::ggml_graph_cut::mark_graph_cut(txt_img, "flux.single_blocks." + std::to_string(i), "txt_img");
             }
 
             img = ggml_view_3d(ctx->ggml_ctx,
@@ -1183,12 +1189,12 @@ namespace Flux {
         bool use_mask = false;
 
         FluxRunner(ggml_backend_t backend,
-                   bool offload_params_to_cpu,
+                   ggml_backend_t params_backend,
                    const String2TensorStorage& tensor_storage_map = {},
                    const std::string prefix                       = "",
                    SDVersion version                              = VERSION_FLUX,
                    bool use_mask                                  = false)
-            : GGMLRunner(backend, offload_params_to_cpu), version(version), use_mask(use_mask) {
+            : GGMLRunner(backend, params_backend), version(version), use_mask(use_mask) {
             flux_params.version             = version;
             flux_params.guidance_embed      = false;
             flux_params.depth               = 0;
@@ -1558,7 +1564,7 @@ namespace Flux {
             }
 
             std::shared_ptr<FluxRunner> flux = std::make_shared<FluxRunner>(backend,
-                                                                            false,
+                                                                            backend,
                                                                             tensor_storage_map,
                                                                             "model.diffusion_model",
                                                                             VERSION_FLUX2,

src/ggml_extend_backend.cpp

@@ -0,0 +1,600 @@
+#include "ggml_extend_backend.h"
+
+#include <algorithm>
+#include <cctype>
+#include <cstdlib>
+#include <mutex>
+#include <sstream>
+#include <stdexcept>
+#include <vector>
+
+#include "util.h"
+
+static std::string trim_copy(const std::string& value) {
+    size_t begin = 0;
+    while (begin < value.size() && std::isspace(static_cast<unsigned char>(value[begin]))) {
+        ++begin;
+    }
+    size_t end = value.size();
+    while (end > begin && std::isspace(static_cast<unsigned char>(value[end - 1]))) {
+        --end;
+    }
+    return value.substr(begin, end - begin);
+}
+
+static std::string lower_copy(std::string value) {
+    std::transform(value.begin(), value.end(), value.begin(), [](unsigned char c) {
+        return static_cast<char>(std::tolower(c));
+    });
+    return value;
+}
+
+static std::vector<std::string> split_copy(const std::string& value, char delimiter) {
+    std::vector<std::string> parts;
+    std::string part;
+    std::istringstream stream(value);
+    while (std::getline(stream, part, delimiter)) {
+        parts.push_back(part);
+    }
+    return parts;
+}
+
+static bool is_default_backend_token(const std::string& name) {
+    const std::string lower = lower_copy(trim_copy(name));
+    return lower.empty() || lower == "default" || lower == "auto";
+}
+
+static bool parse_backend_module(const std::string& raw_name, SDBackendModule* module) {
+    std::string name = lower_copy(trim_copy(raw_name));
+    name.erase(std::remove(name.begin(), name.end(), '-'), name.end());
+    name.erase(std::remove(name.begin(), name.end(), '_'), name.end());
+
+    if (name == "diffusion" || name == "model" || name == "unet" || name == "dit") {
+        *module = SDBackendModule::DIFFUSION;
+        return true;
+    }
+    if (name == "te" || name == "clip" || name == "text" || name == "textencoder" || name == "textencoders" || name == "conditioner" || name == "cond" || name == "llm" || name == "t5" || name == "t5xxl") {
+        *module = SDBackendModule::TE;
+        return true;
+    }
+    if (name == "clipvision" || name == "vision") {
+        *module = SDBackendModule::CLIP_VISION;
+        return true;
+    }
+    if (name == "vae" || name == "firststage" || name == "autoencoder" || name == "tae") {
+        *module = SDBackendModule::VAE;
+        return true;
+    }
+    if (name == "controlnet" || name == "control") {
+        *module = SDBackendModule::CONTROL_NET;
+        return true;
+    }
+    if (name == "photomaker" || name == "photomakerid" || name == "pmid" || name == "photo") {
+        *module = SDBackendModule::PHOTOMAKER;
+        return true;
+    }
+    if (name == "upscaler" || name == "esrgan" || name == "hires") {
+        *module = SDBackendModule::UPSCALER;
+        return true;
+    }
+    return false;
+}
+
+static std::string module_assignment_name(const SDBackendAssignment& assignment, SDBackendModule module) {
+    auto it = assignment.module_names.find(module);
+    if (it != assignment.module_names.end()) {
+        return it->second;
+    }
+    return assignment.default_name;
+}
+
+static std::string backend_cache_key(ggml_backend_t backend) {
+    if (backend == nullptr) {
+        return "";
+    }
+    ggml_backend_dev_t dev = ggml_backend_get_device(backend);
+    if (dev != nullptr) {
+        return lower_copy(ggml_backend_dev_name(dev));
+    }
+    const char* backend_name = ggml_backend_name(backend);
+    return backend_name != nullptr ? lower_copy(backend_name) : "";
+}
+
+static std::string resolve_first_device_by_type(enum ggml_backend_dev_type type) {
+    ggml_backend_dev_t dev = ggml_backend_dev_by_type(type);
+    if (dev == nullptr) {
+        return "";
+    }
+    return ggml_backend_dev_name(dev);
+}
+
+static ggml_backend_buffer_t ggml_backend_tensor_buffer(const struct ggml_tensor* tensor) {
+    if (tensor == nullptr) {
+        return nullptr;
+    }
+
+    return tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
+}
+
+static bool ggml_backend_tensor_is_host_accessible(const struct ggml_tensor* tensor) {
+    if (tensor == nullptr || tensor->data == nullptr) {
+        return false;
+    }
+
+    ggml_backend_buffer_t buffer = ggml_backend_tensor_buffer(tensor);
+    return buffer == nullptr || ggml_backend_buffer_is_host(buffer);
+}
+
+static size_t ggml_backend_tensor_offset(const struct ggml_tensor* tensor, int64_t i0, int64_t i1, int64_t i2, int64_t i3) {
+    return static_cast<size_t>(i0 * tensor->nb[0] + i1 * tensor->nb[1] + i2 * tensor->nb[2] + i3 * tensor->nb[3]);
+}
+
+template <typename T>
+static void ggml_backend_tensor_write_scalar(const struct ggml_tensor* tensor, int64_t i0, int64_t i1, int64_t i2, int64_t i3, T value) {
+    const size_t offset = ggml_backend_tensor_offset(tensor, i0, i1, i2, i3);
+
+    if (ggml_backend_tensor_is_host_accessible(tensor)) {
+        auto* dst = reinterpret_cast<T*>(reinterpret_cast<char*>(tensor->data) + offset);
+        *dst      = value;
+        return;
+    }
+
+    ggml_backend_tensor_set(const_cast<struct ggml_tensor*>(tensor), &value, offset, sizeof(T));
+}
+
+static void ggml_set_f32_nd(const struct ggml_tensor* tensor, int64_t i0, int64_t i1, int64_t i2, int64_t i3, float value) {
+    switch (tensor->type) {
+        case GGML_TYPE_I8:
+            ggml_backend_tensor_write_scalar(tensor, i0, i1, i2, i3, static_cast<int8_t>(value));
+            break;
+        case GGML_TYPE_I16:
+            ggml_backend_tensor_write_scalar(tensor, i0, i1, i2, i3, static_cast<int16_t>(value));
+            break;
+        case GGML_TYPE_I32:
+            ggml_backend_tensor_write_scalar(tensor, i0, i1, i2, i3, static_cast<int32_t>(value));
+            break;
+        case GGML_TYPE_F16:
+            ggml_backend_tensor_write_scalar(tensor, i0, i1, i2, i3, ggml_fp32_to_fp16(value));
+            break;
+        case GGML_TYPE_BF16:
+            ggml_backend_tensor_write_scalar(tensor, i0, i1, i2, i3, ggml_fp32_to_bf16(value));
+            break;
+        case GGML_TYPE_F32:
+            ggml_backend_tensor_write_scalar(tensor, i0, i1, i2, i3, value);
+            break;
+        default:
+            GGML_ABORT("fatal error");
+    }
+}
+
+void ggml_ext_im_set_f32_1d(const struct ggml_tensor* tensor, int i, float value) {
+    if (!ggml_is_contiguous(tensor)) {
+        int64_t id[4] = {0, 0, 0, 0};
+        ggml_unravel_index(tensor, i, &id[0], &id[1], &id[2], &id[3]);
+        ggml_set_f32_nd(tensor, id[0], id[1], id[2], id[3], value);
+        return;
+    }
+
+    switch (tensor->type) {
+        case GGML_TYPE_I8:
+            ggml_backend_tensor_write_scalar(tensor, i, 0, 0, 0, static_cast<int8_t>(value));
+            break;
+        case GGML_TYPE_I16:
+            ggml_backend_tensor_write_scalar(tensor, i, 0, 0, 0, static_cast<int16_t>(value));
+            break;
+        case GGML_TYPE_I32:
+            ggml_backend_tensor_write_scalar(tensor, i, 0, 0, 0, static_cast<int32_t>(value));
+            break;
+        case GGML_TYPE_F16:
+            ggml_backend_tensor_write_scalar(tensor, i, 0, 0, 0, ggml_fp32_to_fp16(value));
+            break;
+        case GGML_TYPE_BF16:
+            ggml_backend_tensor_write_scalar(tensor, i, 0, 0, 0, ggml_fp32_to_bf16(value));
+            break;
+        case GGML_TYPE_F32:
+            ggml_backend_tensor_write_scalar(tensor, i, 0, 0, 0, value);
+            break;
+        default:
+            GGML_ABORT("fatal error");
+    }
+}
+
+static void ggml_backend_load_all_once() {
+    // If the registry already has devices and the CPU backend is present,
+    // assume either static registration or explicit host-side preloading has
+    // completed and avoid rescanning the default paths.
+    if (ggml_backend_dev_count() > 0 && ggml_backend_reg_by_name("CPU") != nullptr) {
+        return;
+    }
+    // In dynamic-backend mode the backend modules are discovered at runtime,
+    // so we must load them before asking for the CPU backend or its proc table.
+    // If the host preloaded only a subset of backends, allow one default-path
+    // scan so missing modules can still be discovered.
+    static std::once_flag once;
+    std::call_once(once, []() {
+        if (ggml_backend_dev_count() > 0 && ggml_backend_reg_by_name("CPU") != nullptr) {
+            return;
+        }
+        ggml_backend_load_all();
+    });
+}
+
+bool sd_backend_is(ggml_backend_t backend, const std::string& name) {
+    if (!backend) {
+        return false;
+    }
+    ggml_backend_dev_t dev = ggml_backend_get_device(backend);
+    if (!dev) {
+        return false;
+    }
+    std::string dev_name = ggml_backend_dev_name(dev);
+    return lower_copy(dev_name).find(lower_copy(name)) != std::string::npos;
+}
+
+static std::string get_default_backend_name() {
+    ggml_backend_load_all_once();
+    // should pick the same backend preference as ggml_backend_init_best
+    std::string name = resolve_first_device_by_type(GGML_BACKEND_DEVICE_TYPE_GPU);
+    if (!name.empty()) {
+        return name;
+    }
+    name = resolve_first_device_by_type(GGML_BACKEND_DEVICE_TYPE_IGPU);
+    if (!name.empty()) {
+        return name;
+    }
+    return resolve_first_device_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
+}
+
+static std::string sd_resolve_backend_name(const std::string& name) {
+    ggml_backend_load_all_once();
+    std::string requested = trim_copy(name);
+    std::string lower     = lower_copy(requested);
+
+    if (is_default_backend_token(lower)) {
+        return get_default_backend_name();
+    }
+    if (lower == "gpu") {
+        std::string result = resolve_first_device_by_type(GGML_BACKEND_DEVICE_TYPE_GPU);
+        if (!result.empty()) {
+            return result;
+        }
+        return resolve_first_device_by_type(GGML_BACKEND_DEVICE_TYPE_IGPU);
+    }
+
+    const size_t device_count = ggml_backend_dev_count();
+    for (size_t i = 0; i < device_count; ++i) {
+        ggml_backend_dev_t dev = ggml_backend_dev_get(i);
+        std::string dev_name   = ggml_backend_dev_name(dev);
+        if (lower_copy(dev_name) == lower) {
+            return dev_name;
+        }
+    }
+
+    for (size_t i = 0; i < device_count; ++i) {
+        ggml_backend_dev_t dev = ggml_backend_dev_get(i);
+        std::string dev_name   = ggml_backend_dev_name(dev);
+        std::string dev_lower  = lower_copy(dev_name);
+        if (dev_lower.rfind(lower, 0) == 0) {
+            return dev_name;
+        }
+    }
+
+    return "";
+}
+
+static bool backend_name_exists(const std::string& name) {
+    return !sd_resolve_backend_name(name).empty();
+}
+
+static ggml_backend_t init_named_backend(const std::string& name) {
+    ggml_backend_load_all_once();
+    LOG_DEBUG("Initializing backend: %s", name.c_str());
+    if (trim_copy(name).empty()) {
+        return ggml_backend_init_best();
+    }
+
+    std::string resolved = sd_resolve_backend_name(name);
+    if (resolved.empty()) {
+        return nullptr;
+    }
+    return ggml_backend_init_by_name(resolved.c_str(), nullptr);
+}
+
+static ggml_backend_t sd_get_default_backend() {
+    ggml_backend_load_all_once();
+    static std::once_flag once;
+    std::call_once(once, []() {
+        size_t dev_count = ggml_backend_dev_count();
+        if (dev_count == 0) {
+            LOG_ERROR("No devices found!");
+        } else {
+            LOG_DEBUG("Found %zu backend devices:", dev_count);
+            for (size_t i = 0; i < dev_count; ++i) {
+                auto dev = ggml_backend_dev_get(i);
+                LOG_DEBUG("#%zu: %s", i, ggml_backend_dev_name(dev));
+            }
+        }
+    });
+
+    ggml_backend_t backend   = nullptr;
+    const char* SD_VK_DEVICE = getenv("SD_VK_DEVICE");
+    if (SD_VK_DEVICE != nullptr) {
+        std::string sd_vk_device_str = SD_VK_DEVICE;
+        try {
+            unsigned long long device  = std::stoull(sd_vk_device_str);
+            std::string vk_device_name = "Vulkan" + std::to_string(device);
+            if (backend_name_exists(vk_device_name)) {
+                LOG_INFO("Selecting %s as main device by env var SD_VK_DEVICE", vk_device_name.c_str());
+                backend = init_named_backend(vk_device_name);
+                if (!backend) {
+                    LOG_WARN("Device %s requested by SD_VK_DEVICE failed to init. Falling back to the default device.", vk_device_name.c_str());
+                }
+            } else {
+                LOG_WARN("Device %s requested by SD_VK_DEVICE was not found. Falling back to the default device.", vk_device_name.c_str());
+            }
+        } catch (const std::invalid_argument&) {
+            LOG_WARN("SD_VK_DEVICE environment variable is not a valid integer (%s). Falling back to the default device.", SD_VK_DEVICE);
+        } catch (const std::out_of_range&) {
+            LOG_WARN("SD_VK_DEVICE environment variable value is out of range for `unsigned long long` type (%s). Falling back to the default device.", SD_VK_DEVICE);
+        }
+    }
+
+    if (!backend) {
+        std::string dev_name = get_default_backend_name();
+        backend              = init_named_backend(dev_name);
+        if (!backend && !dev_name.empty()) {
+            LOG_WARN("device %s failed to init", dev_name.c_str());
+        }
+    }
+
+    if (!backend) {
+        LOG_WARN("loading CPU backend");
+        backend = ggml_backend_cpu_init();
+    }
+
+    if (ggml_backend_is_cpu(backend)) {
+        LOG_DEBUG("Using CPU backend");
+    }
+
+    return backend;
+}
+
+static bool sd_parse_backend_assignment(const std::string& spec, SDBackendAssignment* assignment, std::string* error) {
+    if (assignment == nullptr) {
+        return false;
+    }
+
+    *assignment          = {};
+    const std::string in = trim_copy(spec);
+    if (in.empty()) {
+        return true;
+    }
+
+    for (const std::string& raw_part : split_copy(in, ',')) {
+        const std::string part = trim_copy(raw_part);
+        if (part.empty()) {
+            continue;
+        }
+
+        const size_t eq = part.find('=');
+        if (eq == std::string::npos) {
+            assignment->set_default(part);
+            continue;
+        }
+
+        const std::string key   = trim_copy(part.substr(0, eq));
+        const std::string value = trim_copy(part.substr(eq + 1));
+        if (key.empty() || value.empty()) {
+            if (error != nullptr) {
+                *error = "invalid backend assignment '" + part + "'";
+            }
+            return false;
+        }
+
+        const std::string key_lower = lower_copy(key);
+        if (key_lower == "all" || key_lower == "default" || key_lower == "*") {
+            assignment->set_default(value);
+            continue;
+        }
+
+        SDBackendModule module = SDBackendModule::DIFFUSION;
+        if (!parse_backend_module(key, &module)) {
+            if (error != nullptr) {
+                *error = "unknown backend module '" + key + "'";
+            }
+            return false;
+        }
+        assignment->set_module(module, value);
+    }
+    return true;
+}
+
+bool SDBackendAssignment::empty() const {
+    return default_name.empty() && module_names.empty();
+}
+
+std::string SDBackendAssignment::get(SDBackendModule module) const {
+    return module_assignment_name(*this, module);
+}
+
+void SDBackendAssignment::set_default(const std::string& name) {
+    default_name = trim_copy(name);
+}
+
+void SDBackendAssignment::set_module(SDBackendModule module, const std::string& name) {
+    module_names[module] = trim_copy(name);
+}
+
+void SDBackendHandleDeleter::operator()(ggml_backend_t backend) const {
+    ggml_backend_free(backend);
+}
+
+SDBackendManager::~SDBackendManager() {
+    reset();
+}
+
+void SDBackendManager::reset() {
+    backends_.clear();
+    runtime_assignment_ = {};
+    params_assignment_  = {};
+}
+
+ggml_backend_t SDBackendManager::runtime_backend(SDBackendModule module) {
+    return init_cached_backend(runtime_assignment_.get(module));
+}
+
+ggml_backend_t SDBackendManager::params_backend(SDBackendModule module) {
+    std::string name = params_assignment_.get(module);
+    if (name.empty()) {
+        return runtime_backend(module);
+    }
+    return init_cached_backend(name);
+}
+
+bool SDBackendManager::runtime_backend_is_cpu(SDBackendModule module) {
+    return ggml_backend_is_cpu(runtime_backend(module));
+}
+
+bool SDBackendManager::params_backend_is_cpu(SDBackendModule module) {
+    return ggml_backend_is_cpu(params_backend(module));
+}
+
+bool SDBackendManager::runtime_backend_supports_host_buffer(SDBackendModule module) {
+    ggml_backend_t backend = runtime_backend(module);
+    if (backend == nullptr) {
+        return false;
+    }
+    if (ggml_backend_is_cpu(backend)) {
+        return true;
+    }
+    ggml_backend_dev_t dev = ggml_backend_get_device(backend);
+    if (dev == nullptr) {
+        return false;
+    }
+    ggml_backend_dev_props props;
+    ggml_backend_dev_get_props(dev, &props);
+    return props.caps.buffer_from_host_ptr;
+}
+
+bool SDBackendManager::init(const char* backend_spec,
+                            const char* params_backend_spec,
+                            bool offload_params_to_cpu,
+                            bool keep_clip_on_cpu,
+                            bool keep_vae_on_cpu,
+                            bool keep_control_net_on_cpu,
+                            std::string* error) {
+    reset();
+
+    if (!sd_parse_backend_assignment(SAFE_STR(backend_spec), &runtime_assignment_, error)) {
+        return false;
+    }
+    if (!sd_parse_backend_assignment(SAFE_STR(params_backend_spec), &params_assignment_, error)) {
+        return false;
+    }
+
+    if (runtime_assignment_.empty()) {
+        if (keep_clip_on_cpu) {
+            runtime_assignment_.set_module(SDBackendModule::TE, "cpu");
+        }
+        if (keep_vae_on_cpu) {
+            runtime_assignment_.set_module(SDBackendModule::VAE, "cpu");
+        }
+        if (keep_control_net_on_cpu) {
+            runtime_assignment_.set_module(SDBackendModule::CONTROL_NET, "cpu");
+        }
+    }
+
+    if (params_assignment_.empty() && offload_params_to_cpu) {
+        params_assignment_.set_default("cpu");
+    }
+
+    return validate(error);
+}
+
+bool SDBackendManager::validate(std::string* error) const {
+    auto validate_name = [&](const std::string& name) -> bool {
+        if (is_default_backend_token(name)) {
+            return true;
+        }
+        if (!sd_resolve_backend_name(name).empty()) {
+            return true;
+        }
+        if (error != nullptr) {
+            *error = "backend '" + name + "' was not found";
+        }
+        return false;
+    };
+
+    if (!validate_name(runtime_assignment_.default_name) ||
+        !validate_name(params_assignment_.default_name)) {
+        return false;
+    }
+    for (const auto& kv : runtime_assignment_.module_names) {
+        if (!validate_name(kv.second)) {
+            return false;
+        }
+    }
+    for (const auto& kv : params_assignment_.module_names) {
+        if (!validate_name(kv.second)) {
+            return false;
+        }
+    }
+    return true;
+}
+
+ggml_backend_t SDBackendManager::init_cached_backend(const std::string& name) {
+    std::string resolved   = sd_resolve_backend_name(name);
+    std::string key        = lower_copy(resolved);
+    ggml_backend_t backend = nullptr;
+
+    if (!key.empty()) {
+        auto it = backends_.find(key);
+        if (it != backends_.end()) {
+            return it->second.get();
+        }
+    } else if (!is_default_backend_token(name)) {
+        LOG_ERROR("backend '%s' was not found", name.c_str());
+        return nullptr;
+    }
+
+    backend = is_default_backend_token(name) ? sd_get_default_backend() : init_named_backend(resolved);
+    if (backend == nullptr) {
+        LOG_ERROR("failed to initialize backend '%s'", name.c_str());
+        return nullptr;
+    }
+
+    std::string actual_key = backend_cache_key(backend);
+    if (actual_key.empty()) {
+        actual_key = !key.empty() ? key : lower_copy(trim_copy(name));
+    }
+
+    auto it = backends_.find(actual_key);
+    if (it != backends_.end()) {
+        ggml_backend_free(backend);
+        return it->second.get();
+    }
+
+    SDBackendHandle handle(backend);
+    backends_.emplace(actual_key, std::move(handle));
+    return backend;
+}
+
+const char* sd_backend_module_name(SDBackendModule module) {
+    switch (module) {
+        case SDBackendModule::DIFFUSION:
+            return "diffusion";
+        case SDBackendModule::TE:
+            return "te";
+        case SDBackendModule::CLIP_VISION:
+            return "clip_vision";
+        case SDBackendModule::VAE:
+            return "vae";
+        case SDBackendModule::CONTROL_NET:
+            return "controlnet";
+        case SDBackendModule::PHOTOMAKER:
+            return "photomaker";
+        case SDBackendModule::UPSCALER:
+            return "upscaler";
+    }
+    return "unknown";
+}

src/ggml_extend_backend.h

@@ -0,0 +1,77 @@
+#ifndef __SD_GGML_EXTEND_BACKEND_H__
+#define __SD_GGML_EXTEND_BACKEND_H__
+
+#include <cstdint>
+#include <cstring>
+#include <memory>
+#include <string>
+#include <unordered_map>
+
+#include "ggml-backend.h"
+#include "ggml-cpu.h"
+#include "ggml.h"
+
+enum class SDBackendModule {
+    DIFFUSION,
+    TE,
+    CLIP_VISION,
+    VAE,
+    CONTROL_NET,
+    PHOTOMAKER,
+    UPSCALER,
+};
+
+struct SDBackendAssignment {
+    std::string default_name;
+    std::unordered_map<SDBackendModule, std::string> module_names;
+
+    bool empty() const;
+    std::string get(SDBackendModule module) const;
+    void set_default(const std::string& name);
+    void set_module(SDBackendModule module, const std::string& name);
+};
+
+struct SDBackendHandleDeleter {
+    void operator()(ggml_backend_t backend) const;
+};
+
+using SDBackendHandle = std::unique_ptr<struct ggml_backend, SDBackendHandleDeleter>;
+
+class SDBackendManager {
+private:
+    SDBackendAssignment runtime_assignment_;
+    SDBackendAssignment params_assignment_;
+    std::unordered_map<std::string, SDBackendHandle> backends_;
+
+public:
+    SDBackendManager() = default;
+    ~SDBackendManager();
+
+    SDBackendManager(const SDBackendManager&)            = delete;
+    SDBackendManager& operator=(const SDBackendManager&) = delete;
+
+    bool init(const char* backend_spec,
+              const char* params_backend_spec,
+              bool offload_params_to_cpu,
+              bool keep_clip_on_cpu,
+              bool keep_vae_on_cpu,
+              bool keep_control_net_on_cpu,
+              std::string* error);
+    void reset();
+
+    ggml_backend_t runtime_backend(SDBackendModule module);
+    ggml_backend_t params_backend(SDBackendModule module);
+
+    bool runtime_backend_is_cpu(SDBackendModule module);
+    bool params_backend_is_cpu(SDBackendModule module);
+    bool runtime_backend_supports_host_buffer(SDBackendModule module);
+
+private:
+    bool validate(std::string* error) const;
+    ggml_backend_t init_cached_backend(const std::string& name);
+};
+
+bool sd_backend_is(ggml_backend_t backend, const std::string& name);
+const char* sd_backend_module_name(SDBackendModule module);
+void ggml_ext_im_set_f32_1d(const struct ggml_tensor* tensor, int i, float value);
+#endif

src/ggml_graph_cut.cpp

@@ -0,0 +1,755 @@
+#include "ggml_graph_cut.h"
+
+#include <algorithm>
+#include <cstring>
+#include <map>
+#include <set>
+#include <sstream>
+#include <stack>
+#include <unordered_map>
+
+#include "ggml-alloc.h"
+#include "ggml-backend.h"
+#include "util.h"
+
+#include "../ggml/src/ggml-impl.h"
+
+namespace sd::ggml_graph_cut {
+
+    static constexpr double MAX_VRAM_BYTES_PER_GIB      = 1024.0 * 1024.0 * 1024.0;
+    static constexpr size_t MAX_VRAM_AUTO_RESERVE_BYTES = 1024ULL * 1024ULL * 1024ULL;
+
+    static std::string graph_cut_tensor_display_name(const ggml_tensor* tensor) {
+        if (tensor == nullptr) {
+            return "<null>";
+        }
+        if (tensor->name[0] != '\0') {
+            return tensor->name;
+        }
+        return sd_format("<tensor@%p>", (const void*)tensor);
+    }
+
+    static int graph_leaf_index(ggml_cgraph* gf, const ggml_tensor* tensor) {
+        GGML_ASSERT(gf != nullptr);
+        GGML_ASSERT(tensor != nullptr);
+        for (int i = 0; i < gf->n_leafs; ++i) {
+            if (gf->leafs[i] == tensor) {
+                return i;
+            }
+        }
+        return -1;
+    }
+
+    static bool is_params_tensor(const std::unordered_set<const ggml_tensor*>& params_tensor_set,
+                                 const ggml_tensor* tensor) {
+        if (tensor == nullptr) {
+            return false;
+        }
+        return params_tensor_set.find(tensor) != params_tensor_set.end();
+    }
+
+    static int graph_node_index_by_name(ggml_cgraph* gf, const char* name) {
+        GGML_ASSERT(gf != nullptr);
+        if (name == nullptr || name[0] == '\0') {
+            return -1;
+        }
+        const int n_nodes = ggml_graph_n_nodes(gf);
+        for (int i = 0; i < n_nodes; ++i) {
+            ggml_tensor* node = ggml_graph_node(gf, i);
+            if (node != nullptr && std::strcmp(node->name, name) == 0) {
+                return i;
+            }
+        }
+        return -1;
+    }
+
+    static Plan::InputShape input_shape(const ggml_tensor* tensor) {
+        Plan::InputShape shape;
+        if (tensor == nullptr) {
+            return shape;
+        }
+        shape.type = tensor->type;
+        for (int i = 0; i < GGML_MAX_DIMS; ++i) {
+            shape.ne[static_cast<size_t>(i)] = tensor->ne[i];
+        }
+        return shape;
+    }
+
+    static size_t graph_cut_segment_vram_bytes(const Segment& segment) {
+        return segment.compute_buffer_size +
+               segment.input_param_bytes +
+               segment.input_previous_cut_bytes +
+               segment.output_bytes;
+    }
+
+    size_t max_vram_gib_to_bytes(float max_vram) {
+        if (max_vram <= 0.f) {
+            return 0;
+        }
+        return static_cast<size_t>(static_cast<double>(max_vram) * MAX_VRAM_BYTES_PER_GIB);
+    }
+
+    static float max_vram_bytes_to_gib(size_t max_vram_bytes) {
+        return static_cast<float>(static_cast<double>(max_vram_bytes) / MAX_VRAM_BYTES_PER_GIB);
+    }
+
+    static size_t resolve_auto_max_vram_bytes(ggml_backend_t backend) {
+        if (backend == nullptr) {
+            LOG_WARN("--max-vram -1 requested, but no backend is available; disabling graph splitting");
+            return 0;
+        }
+
+        ggml_backend_dev_t dev = ggml_backend_get_device(backend);
+        if (dev == nullptr) {
+            LOG_WARN("--max-vram -1 requested, but no backend device is available; disabling graph splitting");
+            return 0;
+        }
+        if (ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU) {
+            LOG_WARN("--max-vram -1 requested, but the main backend is CPU; disabling graph splitting");
+            return 0;
+        }
+
+        size_t free_vram  = 0;
+        size_t total_vram = 0;
+        ggml_backend_dev_memory(dev, &free_vram, &total_vram);
+
+        if (free_vram <= MAX_VRAM_AUTO_RESERVE_BYTES) {
+            LOG_WARN("--max-vram -1 requested, but free VRAM is %.2f GiB; reserving 1.00 GiB leaves no graph budget",
+                     free_vram / MAX_VRAM_BYTES_PER_GIB);
+            return 0;
+        }
+
+        const size_t max_vram_bytes = free_vram - MAX_VRAM_AUTO_RESERVE_BYTES;
+        LOG_INFO("--max-vram -1 auto-detected %.2f GiB free VRAM (%.2f GiB total), reserving 1.00 GiB; using %.2f GiB",
+                 free_vram / MAX_VRAM_BYTES_PER_GIB,
+                 total_vram / MAX_VRAM_BYTES_PER_GIB,
+                 max_vram_bytes / MAX_VRAM_BYTES_PER_GIB);
+        return max_vram_bytes;
+    }
+
+    float resolve_max_vram_gib(float max_vram, ggml_backend_t backend) {
+        if (max_vram != -1.f) {
+            return max_vram;
+        }
+        return max_vram_bytes_to_gib(resolve_auto_max_vram_bytes(backend));
+    }
+
+    static Segment make_segment_seed(const Plan& plan,
+                                     size_t start_segment_index,
+                                     size_t end_segment_index) {
+        GGML_ASSERT(start_segment_index < plan.segments.size());
+        GGML_ASSERT(end_segment_index < plan.segments.size());
+        GGML_ASSERT(start_segment_index <= end_segment_index);
+
+        Segment seed;
+        const auto& start_segment  = plan.segments[start_segment_index];
+        const auto& target_segment = plan.segments[end_segment_index];
+        std::unordered_set<int> seen_output_node_indices;
+        for (size_t seg_idx = start_segment_index; seg_idx <= end_segment_index; ++seg_idx) {
+            for (int output_node_index : plan.segments[seg_idx].output_node_indices) {
+                if (seen_output_node_indices.insert(output_node_index).second) {
+                    seed.output_node_indices.push_back(output_node_index);
+                }
+            }
+        }
+        if (start_segment_index == end_segment_index) {
+            seed.group_name = target_segment.group_name;
+        } else {
+            seed.group_name = sd_format("%s..%s",
+                                        start_segment.group_name.c_str(),
+                                        target_segment.group_name.c_str());
+        }
+        return seed;
+    }
+
+    static void build_segment(ggml_cgraph* gf,
+                              Plan& plan,
+                              Segment& segment,
+                              const std::unordered_map<const ggml_tensor*, int>& producer_index,
+                              std::unordered_set<int>& available_cut_output_node_indices,
+                              ggml_backend_t backend,
+                              const std::unordered_set<const ggml_tensor*>& params_tensor_set,
+                              const char* log_desc) {
+        std::set<int> internal_nodes;
+        std::unordered_set<const ggml_tensor*> input_seen;
+        std::vector<Segment::InputRef> input_refs;
+
+        std::stack<ggml_tensor*> work_stack;
+        for (int output_node_index : segment.output_node_indices) {
+            ggml_tensor* output = ggml_graph_node(gf, output_node_index);
+            if (output != nullptr) {
+                work_stack.push(output);
+            }
+        }
+
+        while (!work_stack.empty()) {
+            ggml_tensor* tensor = work_stack.top();
+            work_stack.pop();
+
+            if (tensor == nullptr) {
+                continue;
+            }
+
+            auto producer_it = producer_index.find(tensor);
+            if (producer_it == producer_index.end()) {
+                if (input_seen.insert(tensor).second) {
+                    Segment::InputRef input_ref;
+                    input_ref.type         = is_params_tensor(params_tensor_set, tensor) ? Segment::INPUT_PARAM : Segment::INPUT_EXTERNAL;
+                    input_ref.display_name = graph_cut_tensor_display_name(tensor);
+                    input_ref.leaf_index   = graph_leaf_index(gf, tensor);
+                    input_refs.push_back(std::move(input_ref));
+                }
+                continue;
+            }
+
+            int node_idx = producer_it->second;
+            if (available_cut_output_node_indices.find(node_idx) != available_cut_output_node_indices.end()) {
+                if (input_seen.insert(tensor).second) {
+                    Segment::InputRef input_ref;
+                    input_ref.type         = Segment::INPUT_PREVIOUS_CUT;
+                    input_ref.display_name = graph_cut_tensor_display_name(tensor);
+                    input_ref.node_index   = node_idx;
+                    input_refs.push_back(std::move(input_ref));
+                }
+                continue;
+            }
+
+            if (!internal_nodes.insert(node_idx).second) {
+                continue;
+            }
+
+            ggml_tensor* node = ggml_graph_node(gf, node_idx);
+            for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) {
+                if (node->src[src_idx] != nullptr) {
+                    work_stack.push(node->src[src_idx]);
+                }
+            }
+        }
+
+        if (!internal_nodes.empty()) {
+            segment.internal_node_indices.assign(internal_nodes.begin(), internal_nodes.end());
+        }
+
+        std::sort(input_refs.begin(),
+                  input_refs.end(),
+                  [](const Segment::InputRef& a, const Segment::InputRef& b) {
+                      if (a.type != b.type) {
+                          return a.type < b.type;
+                      }
+                      return a.display_name < b.display_name;
+                  });
+        segment.input_refs = input_refs;
+        for (const auto& input : input_refs) {
+            ggml_tensor* current_input = input_tensor(gf, input);
+            size_t tensor_bytes        = current_input == nullptr
+                                             ? 0
+                                             : (input.type == Segment::INPUT_PREVIOUS_CUT
+                                                    ? cache_tensor_bytes(current_input)
+                                                    : ggml_nbytes(current_input));
+            switch (input.type) {
+                case Segment::INPUT_PREVIOUS_CUT:
+                    segment.input_previous_cut_bytes += tensor_bytes;
+                    break;
+                case Segment::INPUT_PARAM:
+                    segment.input_param_bytes += tensor_bytes;
+                    break;
+                case Segment::INPUT_EXTERNAL:
+                default:
+                    segment.input_external_bytes += tensor_bytes;
+                    break;
+            }
+        }
+        for (int output_node_index : segment.output_node_indices) {
+            ggml_tensor* output = ggml_graph_node(gf, output_node_index);
+            segment.output_bytes += cache_tensor_bytes(output);
+        }
+        segment.compute_buffer_size = measure_segment_compute_buffer(backend, gf, segment, log_desc);
+
+        for (int output_node_index : segment.output_node_indices) {
+            available_cut_output_node_indices.insert(output_node_index);
+        }
+        plan.segments.push_back(std::move(segment));
+    }
+
+    bool is_graph_cut_tensor(const ggml_tensor* tensor) {
+        if (tensor == nullptr || tensor->name[0] == '\0') {
+            return false;
+        }
+        return std::strncmp(tensor->name, GGML_RUNNER_CUT_PREFIX, std::strlen(GGML_RUNNER_CUT_PREFIX)) == 0;
+    }
+
+    std::string make_graph_cut_name(const std::string& group, const std::string& output) {
+        return std::string(GGML_RUNNER_CUT_PREFIX) + group + "|" + output;
+    }
+
+    void mark_graph_cut(ggml_tensor* tensor, const std::string& group, const std::string& output) {
+        if (tensor == nullptr) {
+            return;
+        }
+        auto name = make_graph_cut_name(group, output);
+        ggml_set_name(tensor, name.c_str());
+    }
+
+    int leaf_count(ggml_cgraph* gf) {
+        GGML_ASSERT(gf != nullptr);
+        return gf->n_leafs;
+    }
+
+    ggml_tensor* leaf_tensor(ggml_cgraph* gf, int leaf_index) {
+        GGML_ASSERT(gf != nullptr);
+        if (leaf_index < 0 || leaf_index >= gf->n_leafs) {
+            return nullptr;
+        }
+        return gf->leafs[leaf_index];
+    }
+
+    ggml_backend_buffer_t tensor_buffer(const ggml_tensor* tensor) {
+        if (tensor == nullptr) {
+            return nullptr;
+        }
+        return tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
+    }
+
+    ggml_tensor* cache_source_tensor(ggml_tensor* tensor) {
+        if (tensor == nullptr) {
+            return nullptr;
+        }
+        if (tensor_buffer(tensor) == nullptr && tensor->src[0] != nullptr &&
+            ggml_nelements(tensor->src[0]) == ggml_nelements(tensor) &&
+            ggml_nbytes(tensor->src[0]) == ggml_nbytes(tensor)) {
+            return cache_source_tensor(tensor->src[0]);
+        }
+        return tensor->view_src ? tensor->view_src : tensor;
+    }
+
+    size_t cache_tensor_bytes(const ggml_tensor* tensor) {
+        if (tensor == nullptr) {
+            return 0;
+        }
+        const ggml_tensor* cache_src = tensor->view_src ? tensor->view_src : tensor;
+        return ggml_nbytes(cache_src);
+    }
+
+    bool plan_matches_graph(ggml_cgraph* gf, const Plan& plan) {
+        GGML_ASSERT(gf != nullptr);
+        if (ggml_graph_n_nodes(gf) != plan.n_nodes || gf->n_leafs != plan.n_leafs) {
+            return false;
+        }
+        for (const auto& input_shape_ref : plan.input_shapes) {
+            if (input_shape_ref.leaf_index < 0 || input_shape_ref.leaf_index >= gf->n_leafs) {
+                return false;
+            }
+            ggml_tensor* leaf = gf->leafs[input_shape_ref.leaf_index];
+            if (leaf == nullptr || input_shape_ref.type != leaf->type) {
+                return false;
+            }
+            for (int d = 0; d < GGML_MAX_DIMS; ++d) {
+                if (input_shape_ref.ne[static_cast<size_t>(d)] != leaf->ne[d]) {
+                    return false;
+                }
+            }
+        }
+        return true;
+    }
+
+    ggml_tensor* output_tensor(ggml_cgraph* gf, const Segment& segment, size_t output_index) {
+        GGML_ASSERT(gf != nullptr);
+        if (output_index >= segment.output_node_indices.size()) {
+            return nullptr;
+        }
+        int node_index = segment.output_node_indices[output_index];
+        if (node_index < 0 || node_index >= ggml_graph_n_nodes(gf)) {
+            return nullptr;
+        }
+        return ggml_graph_node(gf, node_index);
+    }
+
+    ggml_tensor* input_tensor(ggml_cgraph* gf, const Segment::InputRef& input_ref) {
+        GGML_ASSERT(gf != nullptr);
+        if (input_ref.type == Segment::INPUT_PREVIOUS_CUT) {
+            if (input_ref.node_index < 0 || input_ref.node_index >= ggml_graph_n_nodes(gf)) {
+                return nullptr;
+            }
+            return ggml_graph_node(gf, input_ref.node_index);
+        }
+        if (input_ref.leaf_index < 0 || input_ref.leaf_index >= gf->n_leafs) {
+            return nullptr;
+        }
+        return leaf_tensor(gf, input_ref.leaf_index);
+    }
+
+    std::vector<ggml_tensor*> param_tensors(ggml_cgraph* gf, const Segment& segment) {
+        GGML_ASSERT(gf != nullptr);
+        std::vector<ggml_tensor*> tensors;
+        std::unordered_set<ggml_tensor*> seen_tensors;
+        tensors.reserve(segment.input_refs.size());
+        seen_tensors.reserve(segment.input_refs.size());
+        for (const auto& input_ref : segment.input_refs) {
+            if (input_ref.type != Segment::INPUT_PARAM) {
+                continue;
+            }
+            ggml_tensor* tensor = input_tensor(gf, input_ref);
+            if (tensor == nullptr) {
+                continue;
+            }
+            if (seen_tensors.insert(tensor).second) {
+                tensors.push_back(tensor);
+            }
+        }
+        return tensors;
+    }
+
+    std::vector<ggml_tensor*> runtime_param_tensors(ggml_cgraph* gf, const Segment& segment, const char* log_desc) {
+        std::vector<ggml_tensor*> tensors = param_tensors(gf, segment);
+        std::vector<ggml_tensor*> filtered_tensors;
+        filtered_tensors.reserve(tensors.size());
+        for (ggml_tensor* tensor : tensors) {
+            if (tensor_buffer(tensor) == nullptr) {
+                LOG_WARN("%s graph cut skipping param input without buffer: segment=%s tensor=%s",
+                         log_desc == nullptr ? "unknown" : log_desc,
+                         segment.group_name.c_str(),
+                         tensor->name);
+                continue;
+            }
+            filtered_tensors.push_back(tensor);
+        }
+        return filtered_tensors;
+    }
+
+    std::unordered_set<std::string> collect_future_input_names(ggml_cgraph* gf,
+                                                               const Plan& plan,
+                                                               size_t current_segment_index) {
+        GGML_ASSERT(gf != nullptr);
+        std::unordered_set<std::string> future_input_names;
+        for (size_t seg_idx = current_segment_index + 1; seg_idx < plan.segments.size(); ++seg_idx) {
+            const auto& segment = plan.segments[seg_idx];
+            for (const auto& input_ref : segment.input_refs) {
+                if (input_ref.type != Segment::INPUT_PREVIOUS_CUT) {
+                    continue;
+                }
+                ggml_tensor* current_input = input_tensor(gf, input_ref);
+                if (current_input != nullptr && current_input->name[0] != '\0') {
+                    future_input_names.insert(current_input->name);
+                }
+            }
+        }
+        return future_input_names;
+    }
+
+    ggml_cgraph* build_segment_graph(ggml_cgraph* gf,
+                                     const Segment& segment,
+                                     ggml_context** graph_ctx_out) {
+        GGML_ASSERT(gf != nullptr);
+        GGML_ASSERT(graph_ctx_out != nullptr);
+
+        const size_t graph_size = segment.internal_node_indices.size() + segment.input_refs.size() + 8;
+        ggml_init_params params = {
+            /*.mem_size   =*/ggml_graph_overhead_custom(graph_size, false) + 1024,
+            /*.mem_buffer =*/nullptr,
+            /*.no_alloc   =*/true,
+        };
+        ggml_context* graph_ctx = ggml_init(params);
+        GGML_ASSERT(graph_ctx != nullptr);
+        ggml_cgraph* segment_graph = ggml_new_graph_custom(graph_ctx, graph_size, false);
+        GGML_ASSERT(segment_graph != nullptr);
+
+        for (const auto& input : segment.input_refs) {
+            ggml_tensor* current_input = input_tensor(gf, input);
+            if (current_input == nullptr) {
+                continue;
+            }
+            GGML_ASSERT(segment_graph->n_leafs < segment_graph->size);
+            segment_graph->leafs[segment_graph->n_leafs++] = current_input;
+        }
+
+        for (int output_node_index : segment.output_node_indices) {
+            ggml_tensor* output = ggml_graph_node(gf, output_node_index);
+            if (output == nullptr) {
+                continue;
+            }
+            ggml_set_output(output);
+        }
+        for (int node_idx : segment.internal_node_indices) {
+            ggml_graph_add_node(segment_graph, ggml_graph_node(gf, node_idx));
+        }
+        *graph_ctx_out = graph_ctx;
+        return segment_graph;
+    }
+
+    size_t measure_segment_compute_buffer(ggml_backend_t backend,
+                                          ggml_cgraph* gf,
+                                          const Segment& segment,
+                                          const char* log_desc) {
+        GGML_ASSERT(backend != nullptr);
+        GGML_ASSERT(gf != nullptr);
+        if (segment.internal_node_indices.empty()) {
+            return 0;
+        }
+
+        ggml_context* graph_ctx    = nullptr;
+        ggml_cgraph* segment_graph = build_segment_graph(gf, segment, &graph_ctx);
+        ggml_gallocr_t allocr      = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend));
+
+        size_t sizes[1] = {0};
+        ggml_gallocr_reserve_n_size(
+            allocr,
+            segment_graph,
+            nullptr,
+            nullptr,
+            sizes);
+        size_t buffer_size = sizes[0];
+
+        ggml_gallocr_free(allocr);
+        ggml_free(graph_ctx);
+        return buffer_size;
+    }
+
+    Plan build_plan(ggml_backend_t backend,
+                    ggml_cgraph* gf,
+                    const std::unordered_set<const ggml_tensor*>& params_tensor_set,
+                    const char* log_desc) {
+        GGML_ASSERT(backend != nullptr);
+        GGML_ASSERT(gf != nullptr);
+        Plan plan;
+        plan.available    = true;
+        const int n_nodes = ggml_graph_n_nodes(gf);
+        if (n_nodes <= 0) {
+            return plan;
+        }
+        plan.n_nodes = n_nodes;
+        plan.n_leafs = gf->n_leafs;
+        for (int i = 0; i < gf->n_leafs; ++i) {
+            ggml_tensor* leaf = gf->leafs[i];
+            if (is_params_tensor(params_tensor_set, leaf)) {
+                continue;
+            }
+            auto shape       = input_shape(leaf);
+            shape.leaf_index = i;
+            plan.input_shapes.push_back(shape);
+        }
+
+        std::unordered_map<const ggml_tensor*, int> producer_index;
+        producer_index.reserve(static_cast<size_t>(n_nodes));
+        for (int i = 0; i < n_nodes; ++i) {
+            producer_index[ggml_graph_node(gf, i)] = i;
+        }
+
+        std::vector<Segment> grouped_segments;
+        std::unordered_map<std::string, size_t> group_to_segment;
+        for (int i = 0; i < n_nodes; ++i) {
+            ggml_tensor* node = ggml_graph_node(gf, i);
+            if (!is_graph_cut_tensor(node)) {
+                continue;
+            }
+
+            plan.has_cuts = true;
+            std::string full_name(node->name);
+            std::string payload = full_name.substr(std::strlen(GGML_RUNNER_CUT_PREFIX));
+            size_t sep          = payload.find('|');
+            std::string group   = sep == std::string::npos ? payload : payload.substr(0, sep);
+
+            auto it = group_to_segment.find(group);
+            if (it == group_to_segment.end()) {
+                Segment segment;
+                segment.group_name = group;
+                segment.output_node_indices.push_back(i);
+                group_to_segment[group] = grouped_segments.size();
+                grouped_segments.push_back(std::move(segment));
+            } else {
+                auto& segment = grouped_segments[it->second];
+                segment.output_node_indices.push_back(i);
+            }
+        }
+
+        if (!plan.has_cuts) {
+            return plan;
+        }
+
+        std::unordered_set<int> available_cut_output_node_indices;
+        available_cut_output_node_indices.reserve(static_cast<size_t>(n_nodes));
+        for (auto& segment : grouped_segments) {
+            build_segment(gf,
+                          plan,
+                          segment,
+                          producer_index,
+                          available_cut_output_node_indices,
+                          backend,
+                          params_tensor_set,
+                          log_desc);
+        }
+
+        int final_output_index = graph_node_index_by_name(gf, "ggml_runner_final_result_tensor");
+        if (final_output_index < 0) {
+            final_output_index = n_nodes - 1;
+        }
+        ggml_tensor* final_output = final_output_index >= 0 ? ggml_graph_node(gf, final_output_index) : nullptr;
+        if (final_output != nullptr && available_cut_output_node_indices.find(final_output_index) == available_cut_output_node_indices.end()) {
+            Segment final_segment;
+            final_segment.group_name = "ggml_runner.final";
+            final_segment.output_node_indices.push_back(final_output_index);
+            build_segment(gf,
+                          plan,
+                          final_segment,
+                          producer_index,
+                          available_cut_output_node_indices,
+                          backend,
+                          params_tensor_set,
+                          log_desc);
+        }
+
+        return plan;
+    }
+
+    Plan apply_max_vram_budget(ggml_cgraph* gf,
+                               const Plan& base_plan,
+                               size_t max_graph_vram_bytes,
+                               ggml_backend_t backend,
+                               const std::unordered_set<const ggml_tensor*>& params_tensor_set,
+                               const char* log_desc) {
+        GGML_ASSERT(backend != nullptr);
+        GGML_ASSERT(gf != nullptr);
+        int64_t t_budget_begin = ggml_time_ms();
+        if (max_graph_vram_bytes == 0 || !base_plan.has_cuts || base_plan.segments.size() <= 1) {
+            return base_plan;
+        }
+
+        const int n_nodes = ggml_graph_n_nodes(gf);
+        std::unordered_map<const ggml_tensor*, int> producer_index;
+        producer_index.reserve(static_cast<size_t>(n_nodes));
+        for (int i = 0; i < n_nodes; ++i) {
+            producer_index[ggml_graph_node(gf, i)] = i;
+        }
+
+        Plan merged_plan;
+        merged_plan.available = true;
+        merged_plan.has_cuts  = base_plan.has_cuts;
+        merged_plan.valid     = base_plan.valid;
+        merged_plan.n_nodes   = base_plan.n_nodes;
+        merged_plan.n_leafs   = base_plan.n_leafs;
+
+        std::unordered_set<int> available_cut_output_node_indices;
+        available_cut_output_node_indices.reserve(static_cast<size_t>(n_nodes));
+
+        size_t start_segment_index = 0;
+        while (start_segment_index < base_plan.segments.size()) {
+            Plan single_plan;
+            auto single_available_cut_output_node_indices = available_cut_output_node_indices;
+            auto single_seed                              = make_segment_seed(base_plan,
+                                                                              start_segment_index,
+                                                                              start_segment_index);
+            build_segment(gf,
+                          single_plan,
+                          single_seed,
+                          producer_index,
+                          single_available_cut_output_node_indices,
+                          backend,
+                          params_tensor_set,
+                          log_desc);
+            GGML_ASSERT(!single_plan.segments.empty());
+
+            size_t best_end_segment_index = start_segment_index;
+            bool can_merge_next_segment   = graph_cut_segment_vram_bytes(single_plan.segments.back()) <= max_graph_vram_bytes;
+
+            while (can_merge_next_segment && best_end_segment_index + 1 < base_plan.segments.size()) {
+                const size_t next_end_segment_index = best_end_segment_index + 1;
+                Plan candidate_plan;
+                auto candidate_available_cut_output_node_indices = available_cut_output_node_indices;
+                auto candidate_seed                              = make_segment_seed(base_plan,
+                                                                                     start_segment_index,
+                                                                                     next_end_segment_index);
+                build_segment(gf,
+                              candidate_plan,
+                              candidate_seed,
+                              producer_index,
+                              candidate_available_cut_output_node_indices,
+                              backend,
+                              params_tensor_set,
+                              log_desc);
+                GGML_ASSERT(!candidate_plan.segments.empty());
+
+                const auto& candidate_segment = candidate_plan.segments.back();
+                if (graph_cut_segment_vram_bytes(candidate_segment) > max_graph_vram_bytes) {
+                    break;
+                }
+
+                best_end_segment_index = next_end_segment_index;
+            }
+
+            auto best_seed = make_segment_seed(base_plan,
+                                               start_segment_index,
+                                               best_end_segment_index);
+            build_segment(gf,
+                          merged_plan,
+                          best_seed,
+                          producer_index,
+                          available_cut_output_node_indices,
+                          backend,
+                          params_tensor_set,
+                          log_desc);
+            start_segment_index = best_end_segment_index + 1;
+        }
+
+        if (log_desc != nullptr && merged_plan.segments.size() != base_plan.segments.size()) {
+            LOG_INFO("%s graph cut max_vram=%.2f MB merged %zu segments -> %zu segments",
+                     log_desc,
+                     max_graph_vram_bytes / 1024.0 / 1024.0,
+                     base_plan.segments.size(),
+                     merged_plan.segments.size());
+        }
+
+        if (log_desc != nullptr) {
+            LOG_INFO("%s graph cut max_vram budget merge took %lld ms",
+                     log_desc,
+                     ggml_time_ms() - t_budget_begin);
+        }
+
+        return merged_plan;
+    }
+
+    Plan resolve_plan(ggml_backend_t backend,
+                      ggml_cgraph* gf,
+                      PlanCache* cache,
+                      size_t max_graph_vram_bytes,
+                      const std::unordered_set<const ggml_tensor*>& params_tensor_set,
+                      const char* log_desc) {
+        GGML_ASSERT(backend != nullptr);
+        GGML_ASSERT(gf != nullptr);
+        GGML_ASSERT(cache != nullptr);
+
+        int64_t t_prepare_begin = ggml_time_ms();
+        Plan base_plan;
+        int64_t t_plan_begin = ggml_time_ms();
+        if (cache->graph_cut_plan.available && plan_matches_graph(gf, cache->graph_cut_plan)) {
+            base_plan = cache->graph_cut_plan;
+        } else {
+            base_plan                                = build_plan(backend, gf, params_tensor_set, log_desc);
+            cache->graph_cut_plan                    = base_plan;
+            cache->graph_cut_plan.available          = true;
+            cache->budgeted_graph_cut_plan.available = false;
+            if (log_desc != nullptr) {
+                LOG_INFO("%s build cached graph cut plan done (taking %lld ms)", log_desc, ggml_time_ms() - t_plan_begin);
+            }
+        }
+
+        Plan resolved_plan = base_plan;
+        if (max_graph_vram_bytes > 0 && base_plan.has_cuts) {
+            if (cache->budgeted_graph_cut_plan.available &&
+                cache->budgeted_graph_cut_plan_max_vram_bytes == max_graph_vram_bytes &&
+                plan_matches_graph(gf, cache->budgeted_graph_cut_plan)) {
+                resolved_plan = cache->budgeted_graph_cut_plan;
+            } else {
+                resolved_plan                                 = apply_max_vram_budget(gf,
+                                                                                      base_plan,
+                                                                                      max_graph_vram_bytes,
+                                                                                      backend,
+                                                                                      params_tensor_set,
+                                                                                      log_desc);
+                cache->budgeted_graph_cut_plan                = resolved_plan;
+                cache->budgeted_graph_cut_plan.available      = true;
+                cache->budgeted_graph_cut_plan_max_vram_bytes = max_graph_vram_bytes;
+            }
+        }
+        return resolved_plan;
+    }
+
+}  // namespace sd::ggml_graph_cut

src/ggml_graph_cut.h

@@ -0,0 +1,106 @@
+#ifndef __SD_GGML_GRAPH_CUT_H__
+#define __SD_GGML_GRAPH_CUT_H__
+
+#include <array>
+#include <string>
+#include <unordered_set>
+#include <vector>
+
+#include "ggml-backend.h"
+#include "ggml.h"
+
+namespace sd::ggml_graph_cut {
+
+    struct Segment {
+        enum InputType {
+            INPUT_EXTERNAL = 0,
+            INPUT_PREVIOUS_CUT,
+            INPUT_PARAM,
+        };
+
+        struct InputRef {
+            InputType type = INPUT_EXTERNAL;
+            std::string display_name;
+            int leaf_index = -1;
+            int node_index = -1;
+        };
+
+        size_t compute_buffer_size      = 0;
+        size_t output_bytes             = 0;
+        size_t input_external_bytes     = 0;
+        size_t input_previous_cut_bytes = 0;
+        size_t input_param_bytes        = 0;
+        std::string group_name;
+        std::vector<int> internal_node_indices;
+        std::vector<int> output_node_indices;
+        std::vector<InputRef> input_refs;
+    };
+
+    struct Plan {
+        struct InputShape {
+            int leaf_index                        = -1;
+            ggml_type type                        = GGML_TYPE_COUNT;
+            std::array<int64_t, GGML_MAX_DIMS> ne = {0, 0, 0, 0};
+        };
+
+        bool available = false;
+        bool has_cuts  = false;
+        bool valid     = true;
+        int n_nodes    = 0;
+        int n_leafs    = 0;
+        std::vector<InputShape> input_shapes;
+        std::vector<Segment> segments;
+    };
+
+    struct PlanCache {
+        Plan graph_cut_plan;
+        Plan budgeted_graph_cut_plan;
+        size_t budgeted_graph_cut_plan_max_vram_bytes = 0;
+    };
+
+    static constexpr const char* GGML_RUNNER_CUT_PREFIX = "ggml_runner_cut:";
+
+    bool is_graph_cut_tensor(const ggml_tensor* tensor);
+    std::string make_graph_cut_name(const std::string& group, const std::string& output);
+    void mark_graph_cut(ggml_tensor* tensor, const std::string& group, const std::string& output);
+    int leaf_count(ggml_cgraph* gf);
+    ggml_tensor* leaf_tensor(ggml_cgraph* gf, int leaf_index);
+    ggml_backend_buffer_t tensor_buffer(const ggml_tensor* tensor);
+    ggml_tensor* cache_source_tensor(ggml_tensor* tensor);
+    size_t cache_tensor_bytes(const ggml_tensor* tensor);
+    bool plan_matches_graph(ggml_cgraph* gf, const Plan& plan);
+    ggml_tensor* output_tensor(ggml_cgraph* gf, const Segment& segment, size_t output_index);
+    ggml_tensor* input_tensor(ggml_cgraph* gf, const Segment::InputRef& input_ref);
+    std::vector<ggml_tensor*> param_tensors(ggml_cgraph* gf, const Segment& segment);
+    std::vector<ggml_tensor*> runtime_param_tensors(ggml_cgraph* gf, const Segment& segment, const char* log_desc);
+    std::unordered_set<std::string> collect_future_input_names(ggml_cgraph* gf,
+                                                               const Plan& plan,
+                                                               size_t current_segment_index);
+    ggml_cgraph* build_segment_graph(ggml_cgraph* gf,
+                                     const Segment& segment,
+                                     ggml_context** graph_ctx_out);
+    size_t measure_segment_compute_buffer(ggml_backend_t backend,
+                                          ggml_cgraph* gf,
+                                          const Segment& segment,
+                                          const char* log_desc);
+    size_t max_vram_gib_to_bytes(float max_vram);
+    float resolve_max_vram_gib(float max_vram, ggml_backend_t backend);
+    Plan build_plan(ggml_backend_t backend,
+                    ggml_cgraph* gf,
+                    const std::unordered_set<const ggml_tensor*>& params_tensor_set,
+                    const char* log_desc);
+    Plan apply_max_vram_budget(ggml_cgraph* gf,
+                               const Plan& base_plan,
+                               size_t max_graph_vram_bytes,
+                               ggml_backend_t backend,
+                               const std::unordered_set<const ggml_tensor*>& params_tensor_set,
+                               const char* log_desc);
+    Plan resolve_plan(ggml_backend_t backend,
+                      ggml_cgraph* gf,
+                      PlanCache* cache,
+                      size_t max_graph_vram_bytes,
+                      const std::unordered_set<const ggml_tensor*>& params_tensor_set,
+                      const char* log_desc);
+}  // namespace sd::ggml_graph_cut
+
+#endif

src/hidream_o1.hpp

@@ -0,0 +1,653 @@
+#ifndef __SD_HIDREAM_O1_H__
+#define __SD_HIDREAM_O1_H__
+
+#include <algorithm>
+#include <array>
+#include <cmath>
+#include <cstring>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "common_dit.hpp"
+#include "conditioner.hpp"
+#include "llm.hpp"
+#include "util.h"
+
+namespace HiDreamO1 {
+    constexpr int HIDREAM_O1_GRAPH_SIZE = 32768;
+    constexpr int PATCH_SIZE            = 32;
+    constexpr int TIMESTEP_TOKEN_NUM    = 1;
+    constexpr int IMAGE_TOKEN_ID        = 151655;
+    constexpr int VISION_START_TOKEN_ID = 151652;
+
+    static inline std::string repeat_special_token(const std::string& token, int64_t count) {
+        std::string out;
+        out.reserve(static_cast<size_t>(count) * token.size());
+        for (int64_t i = 0; i < count; ++i) {
+            out += token;
+        }
+        return out;
+    }
+
+    static inline std::pair<int, int> calculate_dimensions(int max_size, double ratio) {
+        int width  = static_cast<int>(std::sqrt(max_size * max_size * ratio));
+        int height = static_cast<int>(width / ratio);
+        width      = (width / PATCH_SIZE) * PATCH_SIZE;
+        height     = (height / PATCH_SIZE) * PATCH_SIZE;
+        width      = std::max(width, PATCH_SIZE);
+        height     = std::max(height, PATCH_SIZE);
+        return {width, height};
+    }
+
+    static inline sd::Tensor<float> resize_to_area(const sd::Tensor<float>& image, int image_size) {
+        int64_t width  = image.shape()[0];
+        int64_t height = image.shape()[1];
+        int64_t s_max  = static_cast<int64_t>(image_size) * image_size;
+        double scale   = std::sqrt(static_cast<double>(s_max) / static_cast<double>(width * height));
+
+        std::vector<std::pair<int64_t, int64_t>> sizes = {
+            {(static_cast<int64_t>(std::llround(width * scale)) / PATCH_SIZE) * PATCH_SIZE, (static_cast<int64_t>(std::llround(height * scale)) / PATCH_SIZE) * PATCH_SIZE},
+            {(static_cast<int64_t>(std::llround(width * scale)) / PATCH_SIZE) * PATCH_SIZE, (static_cast<int64_t>(std::floor(height * scale)) / PATCH_SIZE) * PATCH_SIZE},
+            {(static_cast<int64_t>(std::floor(width * scale)) / PATCH_SIZE) * PATCH_SIZE, (static_cast<int64_t>(std::llround(height * scale)) / PATCH_SIZE) * PATCH_SIZE},
+            {(static_cast<int64_t>(std::floor(width * scale)) / PATCH_SIZE) * PATCH_SIZE, (static_cast<int64_t>(std::floor(height * scale)) / PATCH_SIZE) * PATCH_SIZE},
+        };
+        std::sort(sizes.begin(), sizes.end(), [](const auto& a, const auto& b) {
+            return a.first * a.second > b.first * b.second;
+        });
+
+        std::pair<int64_t, int64_t> new_size = sizes.back();
+        for (const auto& size : sizes) {
+            if (size.first > 0 && size.second > 0 && size.first * size.second <= s_max) {
+                new_size = size;
+                break;
+            }
+        }
+
+        double s1 = static_cast<double>(width) / static_cast<double>(new_size.first);
+        double s2 = static_cast<double>(height) / static_cast<double>(new_size.second);
+        sd::Tensor<float> resized;
+        if (s1 < s2) {
+            int64_t resized_h = static_cast<int64_t>(std::llround(height / s1));
+            resized           = sd::ops::interpolate(image,
+                                                     {new_size.first, resized_h, image.shape()[2], image.shape()[3]},
+                                                     sd::ops::InterpolateMode::Bicubic);
+            int64_t top       = (resized_h - new_size.second) / 2;
+            resized           = sd::ops::slice(resized, 1, top, top + new_size.second);
+        } else {
+            int64_t resized_w = static_cast<int64_t>(std::llround(width / s2));
+            resized           = sd::ops::interpolate(image,
+                                                     {resized_w, new_size.second, image.shape()[2], image.shape()[3]},
+                                                     sd::ops::InterpolateMode::Bicubic);
+            int64_t left      = (resized_w - new_size.first) / 2;
+            resized           = sd::ops::slice(resized, 0, left, left + new_size.first);
+        }
+        return resized;
+    }
+
+    static inline std::vector<int32_t> build_position_ids(const std::vector<int32_t>& input_ids,
+                                                          const std::vector<std::array<int32_t, 3>>& image_grids,
+                                                          const std::vector<int32_t>& skip_vision_start_token) {
+        std::vector<int32_t> position_ids(4 * input_ids.size(), 0);
+        int image_index = 0;
+        int st          = 0;
+        int fix_point   = 4096;
+        std::vector<int32_t> out_t;
+        std::vector<int32_t> out_h;
+        std::vector<int32_t> out_w;
+
+        while (st < static_cast<int>(input_ids.size())) {
+            int ed = st;
+            while (ed < static_cast<int>(input_ids.size()) && input_ids[ed] != IMAGE_TOKEN_ID) {
+                ed++;
+            }
+
+            if (ed >= static_cast<int>(input_ids.size())) {
+                int st_idx = out_t.empty() ? 0 : (*std::max_element(out_t.begin(), out_t.end()) + 1);
+                for (int i = 0; i < static_cast<int>(input_ids.size()) - st; ++i) {
+                    out_t.push_back(st_idx + i);
+                    out_h.push_back(st_idx + i);
+                    out_w.push_back(st_idx + i);
+                }
+                break;
+            }
+
+            int text_len = std::max(0, ed - st - skip_vision_start_token[image_index]);
+            int st_idx   = out_t.empty() ? 0 : (*std::max_element(out_t.begin(), out_t.end()) + 1);
+            for (int i = 0; i < text_len; ++i) {
+                out_t.push_back(st_idx + i);
+                out_h.push_back(st_idx + i);
+                out_w.push_back(st_idx + i);
+            }
+
+            auto grid = image_grids[image_index];
+            int base;
+            if (skip_vision_start_token[image_index]) {
+                if (fix_point > 0) {
+                    base      = fix_point;
+                    fix_point = 0;
+                } else {
+                    base = st_idx;
+                }
+            } else {
+                base = text_len + st_idx;
+            }
+            for (int32_t ti = 0; ti < grid[0]; ++ti) {
+                for (int32_t hi = 0; hi < grid[1]; ++hi) {
+                    for (int32_t wi = 0; wi < grid[2]; ++wi) {
+                        out_t.push_back(base + ti);
+                        out_h.push_back(base + hi);
+                        out_w.push_back(base + wi);
+                    }
+                }
+            }
+
+            st = ed + grid[0] * grid[1] * grid[2];
+            image_index++;
+        }
+
+        GGML_ASSERT(out_t.size() == input_ids.size());
+        for (size_t i = 0; i < input_ids.size(); ++i) {
+            // ggml IMROPE consumes 4 flattened position streams:
+            //   [t, h, w, e]
+            // llama.cpp's generic Qwen-VL fallback expands text positions as
+            // [pos, pos, pos, 0]. Keep the extra stream zeroed here too.
+            position_ids[i]                        = out_t[i];
+            position_ids[input_ids.size() + i]     = out_h[i];
+            position_ids[input_ids.size() * 2 + i] = out_w[i];
+            position_ids[input_ids.size() * 3 + i] = 0;
+        }
+        return position_ids;
+    }
+
+    struct TimestepEmbedder : public GGMLBlock {
+        int frequency_embedding_size = 256;
+
+        TimestepEmbedder(int64_t hidden_size) {
+            blocks["mlp.0"] = std::make_shared<Linear>(frequency_embedding_size, hidden_size, true);
+            blocks["mlp.2"] = std::make_shared<Linear>(hidden_size, hidden_size, 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 emb   = ggml_ext_timestep_embedding(ctx->ggml_ctx, t, frequency_embedding_size, 10000, 1000.0f);
+            emb        = mlp_0->forward(ctx, emb);
+            emb        = ggml_silu_inplace(ctx->ggml_ctx, emb);
+            emb        = mlp_2->forward(ctx, emb);
+            return emb;
+        }
+    };
+
+    struct BottleneckPatchEmbed : public GGMLBlock {
+        BottleneckPatchEmbed(int64_t in_dim, int64_t pca_dim, int64_t embed_dim) {
+            blocks["proj1"] = std::make_shared<Linear>(in_dim, pca_dim, false);
+            blocks["proj2"] = std::make_shared<Linear>(pca_dim, embed_dim, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            auto proj1 = std::dynamic_pointer_cast<Linear>(blocks["proj1"]);
+            auto proj2 = std::dynamic_pointer_cast<Linear>(blocks["proj2"]);
+            return proj2->forward(ctx, proj1->forward(ctx, x));
+        }
+    };
+
+    struct FinalLayer : public GGMLBlock {
+        FinalLayer(int64_t hidden_size, int64_t out_dim) {
+            blocks["linear"] = std::make_shared<Linear>(hidden_size, out_dim, true);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
+            return linear->forward(ctx, x);
+        }
+    };
+
+    struct HiDreamO1Params {
+        LLM::LLMParams llm;
+        int patch_size = PATCH_SIZE;
+    };
+
+    static inline HiDreamO1Params make_hidream_o1_params() {
+        HiDreamO1Params params;
+        params.llm.arch                           = LLM::LLMArch::QWEN3_VL;
+        params.llm.hidden_size                    = 4096;
+        params.llm.intermediate_size              = 12288;
+        params.llm.num_layers                     = 36;
+        params.llm.num_heads                      = 32;
+        params.llm.num_kv_heads                   = 8;
+        params.llm.head_dim                       = 128;
+        params.llm.qkv_bias                       = false;
+        params.llm.qk_norm                        = true;
+        params.llm.vocab_size                     = 151936;
+        params.llm.rms_norm_eps                   = 1e-6f;
+        params.llm.vision.arch                    = LLM::LLMVisionArch::QWEN3_VL;
+        params.llm.vision.num_layers              = 27;
+        params.llm.vision.hidden_size             = 1152;
+        params.llm.vision.intermediate_size       = 4304;
+        params.llm.vision.num_heads               = 16;
+        params.llm.vision.out_hidden_size         = 4096;
+        params.llm.vision.patch_size              = 16;
+        params.llm.vision.spatial_merge_size      = 2;
+        params.llm.vision.temporal_patch_size     = 2;
+        params.llm.vision.num_position_embeddings = 2304;
+        return params;
+    }
+
+    struct HiDreamO1Model : public GGMLBlock {
+        HiDreamO1Params params;
+
+        HiDreamO1Model() = default;
+        explicit HiDreamO1Model(HiDreamO1Params params)
+            : params(std::move(params)) {
+            blocks["language_model"] = std::make_shared<LLM::TextModel>(this->params.llm);
+            blocks["t_embedder1"]    = std::make_shared<TimestepEmbedder>(this->params.llm.hidden_size);
+            blocks["x_embedder"]     = std::make_shared<BottleneckPatchEmbed>(this->params.patch_size * this->params.patch_size * 3,
+                                                                          this->params.llm.hidden_size / 4,
+                                                                          this->params.llm.hidden_size);
+            blocks["final_layer2"]   = std::make_shared<FinalLayer>(this->params.llm.hidden_size,
+                                                                  this->params.patch_size * this->params.patch_size * 3);
+        }
+
+        std::shared_ptr<LLM::TextModel> text_model() {
+            return std::dynamic_pointer_cast<LLM::TextModel>(blocks["language_model"]);
+        }
+
+        std::shared_ptr<TimestepEmbedder> timestep_embedder() {
+            return std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder1"]);
+        }
+
+        std::shared_ptr<BottleneckPatchEmbed> patch_embedder() {
+            return std::dynamic_pointer_cast<BottleneckPatchEmbed>(blocks["x_embedder"]);
+        }
+
+        std::shared_ptr<FinalLayer> final_layer() {
+            return std::dynamic_pointer_cast<FinalLayer>(blocks["final_layer2"]);
+        }
+    };
+
+    struct HiDreamO1VisionRunner : public GGMLRunner {
+        HiDreamO1Params params;
+        std::shared_ptr<LLM::VisionModel> model;
+
+        std::vector<int> window_index_vec;
+        std::vector<int> window_inverse_index_vec;
+        std::vector<float> window_mask_vec;
+        std::vector<float> pe_vec;
+        std::array<std::vector<int32_t>, 4> pos_embed_idx_data_;
+        std::array<std::vector<float>, 4> pos_embed_weight_data_;
+
+        HiDreamO1VisionRunner(ggml_backend_t backend,
+                              ggml_backend_t params_backend,
+                              const String2TensorStorage& tensor_storage_map = {},
+                              const std::string& prefix                      = "model.visual")
+            : GGMLRunner(backend, params_backend),
+              params(make_hidream_o1_params()),
+              model(std::make_shared<LLM::VisionModel>(false, params.llm.vision)) {
+            model->init(params_ctx, tensor_storage_map, prefix);
+        }
+
+        std::string get_desc() override {
+            return "hidream_o1_vision";
+        }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix = "model.visual") {
+            model->get_param_tensors(tensors, prefix);
+        }
+
+        ggml_tensor* encode_image(GGMLRunnerContext* runner_ctx, ggml_tensor* image) {
+            return LLM::LLMRunner::encode_image_common(this,
+                                                       compute_ctx,
+                                                       runner_ctx,
+                                                       image,
+                                                       params.llm.vision,
+                                                       model,
+                                                       window_index_vec,
+                                                       window_inverse_index_vec,
+                                                       window_mask_vec,
+                                                       pe_vec,
+                                                       pos_embed_idx_data_,
+                                                       pos_embed_weight_data_);
+        }
+
+        ggml_cgraph* build_graph(const sd::Tensor<float>& image_tensor) {
+            ggml_cgraph* gf    = new_graph_custom(HIDREAM_O1_GRAPH_SIZE);
+            ggml_tensor* image = make_input(image_tensor);
+            auto runner_ctx    = get_context();
+            auto image_embeds  = encode_image(&runner_ctx, image);
+            ggml_build_forward_expand(gf, image_embeds);
+            return gf;
+        }
+
+        sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& image) {
+            auto get_graph = [&]() {
+                return build_graph(image);
+            };
+            auto output = GGMLRunner::compute<float>(get_graph, n_threads, false);
+            return output.has_value() ? std::move(output.value()) : sd::Tensor<float>();
+        }
+    };
+
+    struct HiDreamO1Runner : public GGMLRunner {
+        HiDreamO1Params params;
+        HiDreamO1Model model;
+
+        std::vector<float> attention_mask_vec;
+
+        HiDreamO1Runner(ggml_backend_t backend,
+                        ggml_backend_t params_backend,
+                        const String2TensorStorage& tensor_storage_map = {},
+                        const std::string& prefix                      = "model")
+            : GGMLRunner(backend, params_backend),
+              params(make_hidream_o1_params()) {
+            model = HiDreamO1Model(params);
+            model.init(params_ctx, tensor_storage_map, prefix);
+        }
+
+        std::string get_desc() override {
+            return "hidream_o1";
+        }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) {
+            model.get_param_tensors(tensors, prefix);
+        }
+
+        ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
+                                 const sd::Tensor<float>& timestep_tensor,
+                                 const sd::Tensor<int32_t>& input_ids_tensor,
+                                 const sd::Tensor<int32_t>& input_pos_tensor,
+                                 const sd::Tensor<int32_t>& token_types_tensor,
+                                 const sd::Tensor<int32_t>& vinput_mask_tensor,
+                                 const std::vector<std::pair<int, sd::Tensor<float>>>& image_embeds_tensor,
+                                 const std::vector<sd::Tensor<float>>& ref_images) {
+            ggml_cgraph* gf        = new_graph_custom(HIDREAM_O1_GRAPH_SIZE);
+            ggml_tensor* x         = make_input(x_tensor);
+            ggml_tensor* timestep  = make_input(timestep_tensor);
+            ggml_tensor* input_ids = make_input(input_ids_tensor);
+            ggml_tensor* input_pos = make_input(input_pos_tensor);
+
+            auto text_model   = model.text_model();
+            auto t_embedder1  = model.timestep_embedder();
+            auto x_embedder   = model.patch_embedder();
+            auto final_layer2 = model.final_layer();
+
+            std::vector<ggml_tensor*> ref_image_tensors;
+            for (const auto& image : ref_images) {
+                ref_image_tensors.push_back(make_input(image));
+            }
+
+            attention_mask_vec    = std::vector<float>(static_cast<size_t>(token_types_tensor.shape()[0] * token_types_tensor.shape()[0]), 0.0f);
+            int64_t total_seq_len = token_types_tensor.shape()[0];
+            for (int64_t query = 0; query < total_seq_len; ++query) {
+                bool is_gen = token_types_tensor.values()[static_cast<size_t>(query)] > 0;
+                for (int64_t key = 0; key < total_seq_len; ++key) {
+                    if (!is_gen && key > query) {
+                        attention_mask_vec[static_cast<size_t>(query * total_seq_len + key)] = -INFINITY;
+                    }
+                }
+            }
+            auto attention_mask = ggml_new_tensor_2d(compute_ctx, GGML_TYPE_F32, total_seq_len, total_seq_len);
+            set_backend_tensor_data(attention_mask, attention_mask_vec.data());
+
+            auto runner_ctx = get_context();
+            auto txt        = text_model->embed(&runner_ctx, input_ids);
+            std::vector<std::pair<int, ggml_tensor*>> image_embeds;
+            image_embeds.reserve(image_embeds_tensor.size());
+            for (const auto& image_embed : image_embeds_tensor) {
+                image_embeds.emplace_back(image_embed.first, make_input(image_embed.second));
+            }
+            txt = LLM::splice_image_embeds(&runner_ctx, txt, image_embeds);
+
+            auto t_emb          = t_embedder1->forward(&runner_ctx, timestep);
+            int64_t txt_seq_len = input_ids->ne[0];
+            if (txt_seq_len > 1) {
+                auto prefix = ggml_ext_slice(compute_ctx, txt, 1, 0, txt_seq_len - 1);
+                txt         = ggml_concat(compute_ctx, prefix, ggml_reshape_3d(compute_ctx, t_emb, t_emb->ne[0], 1, 1), 1);
+            } else {
+                txt = ggml_reshape_3d(compute_ctx, t_emb, t_emb->ne[0], 1, 1);
+            }
+
+            auto vinputs          = DiT::pad_and_patchify(&runner_ctx, x, PATCH_SIZE, PATCH_SIZE);
+            int64_t target_tokens = vinputs->ne[1];
+            for (ggml_tensor* ref_image : ref_image_tensors) {
+                auto ref = DiT::pad_and_patchify(&runner_ctx, ref_image, PATCH_SIZE, PATCH_SIZE);
+                vinputs  = ggml_concat(compute_ctx, vinputs, ref, 1);
+            }
+            auto vis = x_embedder->forward(&runner_ctx, vinputs);
+
+            auto inputs_embeds = ggml_concat(compute_ctx, txt, vis, 1);
+            auto hidden_states = text_model->forward_embeds(&runner_ctx, inputs_embeds, input_pos, attention_mask, {});
+            auto x_pred_all    = final_layer2->forward(&runner_ctx, hidden_states);
+
+            int64_t x_pred_start = txt_seq_len;
+            if (!vinput_mask_tensor.empty()) {
+                int64_t seq_len      = static_cast<int64_t>(vinput_mask_tensor.shape()[0]);
+                int64_t first_vinput = 0;
+                while (first_vinput < seq_len && vinput_mask_tensor.values()[static_cast<size_t>(first_vinput)] == 0) {
+                    first_vinput++;
+                }
+                x_pred_start = first_vinput;
+            }
+            auto x_pred = ggml_ext_slice(compute_ctx, x_pred_all, 1, x_pred_start, x_pred_start + target_tokens);
+            x_pred      = DiT::unpatchify_and_crop(compute_ctx, x_pred, x->ne[1], x->ne[0], PATCH_SIZE, PATCH_SIZE);
+
+            float sigma = 1.0f - timestep_tensor.values()[0];
+            sigma       = std::max(1e-6f, sigma);
+            auto out    = ggml_scale(compute_ctx, ggml_sub(compute_ctx, x, x_pred), 1.0f / sigma);
+
+            ggml_build_forward_expand(gf, out);
+            return gf;
+        }
+
+        sd::Tensor<float> compute(int n_threads,
+                                  const sd::Tensor<float>& x,
+                                  const sd::Tensor<float>& timestep,
+                                  const sd::Tensor<int32_t>& input_ids,
+                                  const sd::Tensor<int32_t>& input_pos,
+                                  const sd::Tensor<int32_t>& token_types,
+                                  const sd::Tensor<int32_t>& vinput_mask,
+                                  const std::vector<std::pair<int, sd::Tensor<float>>>& image_embeds,
+                                  const std::vector<sd::Tensor<float>>& ref_images) {
+            auto get_graph = [&]() {
+                return build_graph(x, timestep, input_ids, input_pos, token_types, vinput_mask, image_embeds, ref_images);
+            };
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), x.dim());
+        }
+    };
+
+    struct HiDreamO1Conditioner : public Conditioner {
+        Qwen2Tokenizer tokenizer;
+        std::shared_ptr<HiDreamO1VisionRunner> vision_runner;
+
+        HiDreamO1Conditioner(ggml_backend_t backend,
+                             ggml_backend_t params_backend,
+                             const String2TensorStorage& tensor_storage_map = {})
+            : vision_runner(std::make_shared<HiDreamO1VisionRunner>(backend, params_backend, tensor_storage_map)) {}
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+            vision_runner->get_param_tensors(tensors);
+        }
+
+        void alloc_params_buffer() override {
+            vision_runner->alloc_params_buffer();
+        }
+
+        void free_params_buffer() override {
+            vision_runner->free_params_buffer();
+        }
+
+        size_t get_params_buffer_size() override {
+            return vision_runner->get_params_buffer_size();
+        }
+
+        void set_max_graph_vram_bytes(size_t max_graph_vram_bytes) override {
+            vision_runner->set_max_graph_vram_bytes(max_graph_vram_bytes);
+        }
+
+        void set_flash_attention_enabled(bool enabled) override {
+            vision_runner->set_flash_attention_enabled(enabled);
+        }
+
+        void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
+            vision_runner->set_weight_adapter(adapter);
+        }
+
+        SDCondition get_learned_condition(int n_threads,
+                                          const ConditionerParams& conditioner_params) override {
+            SDCondition result;
+
+            int width                = conditioner_params.width;
+            int height               = conditioner_params.height;
+            int64_t target_image_len = static_cast<int64_t>(width / PATCH_SIZE) * static_cast<int64_t>(height / PATCH_SIZE);
+
+            std::vector<sd::Tensor<float>> ref_images;
+            if (conditioner_params.ref_images != nullptr) {
+                ref_images = *conditioner_params.ref_images;
+            }
+
+            std::vector<std::pair<int, sd::Tensor<float>>> vlm_images;
+            std::vector<std::array<int32_t, 3>> image_grids;
+            std::vector<int32_t> skip_vision_start;
+
+            std::string prompt = "<|im_start|>user\n";
+
+            if (ref_images.empty()) {
+                prompt += conditioner_params.text;
+                prompt += "<|im_end|>\n<|im_start|>assistant\n<|boi_token|><|tms_token|>";
+                auto input_ids = tokenizer.encode(prompt, nullptr);
+
+                std::vector<int32_t> input_ids_pad = input_ids;
+                input_ids_pad.push_back(VISION_START_TOKEN_ID);
+                input_ids_pad.insert(input_ids_pad.end(), target_image_len - 1, IMAGE_TOKEN_ID);
+
+                image_grids.push_back({1, static_cast<int32_t>(height / PATCH_SIZE), static_cast<int32_t>(width / PATCH_SIZE)});
+                skip_vision_start.push_back(1);
+
+                std::vector<int32_t> token_types(input_ids_pad.size(), 0);
+                int txt_seq_len = static_cast<int>(input_ids.size());
+                int bgn         = txt_seq_len - TIMESTEP_TOKEN_NUM;
+                for (int i = bgn; i < static_cast<int>(token_types.size()); ++i) {
+                    token_types[i] = 1;
+                }
+
+                auto position_ids = build_position_ids(input_ids_pad, image_grids, skip_vision_start);
+
+                std::vector<int64_t> input_shape{static_cast<int64_t>(input_ids.size())};
+                std::vector<int64_t> position_shape{static_cast<int64_t>(input_ids_pad.size() * 4)};
+                std::vector<int64_t> token_type_shape{static_cast<int64_t>(token_types.size())};
+                std::vector<int32_t> vinput_mask(token_types.size(), 0);
+                for (int64_t i = txt_seq_len; i < static_cast<int64_t>(vinput_mask.size()); ++i) {
+                    vinput_mask[static_cast<size_t>(i)] = 1;
+                }
+                std::vector<int64_t> vinput_mask_shape{static_cast<int64_t>(vinput_mask.size())};
+
+                result.c_input_ids    = sd::Tensor<int32_t>(input_shape, std::move(input_ids));
+                result.c_position_ids = sd::Tensor<int32_t>(position_shape, position_ids);
+                result.c_token_types  = sd::Tensor<int32_t>(token_type_shape, std::move(token_types));
+                result.c_vinput_mask  = sd::Tensor<int32_t>(vinput_mask_shape, std::move(vinput_mask));
+                return result;
+            }
+
+            int K = static_cast<int>(ref_images.size());
+            int max_size;
+            if (K == 1) {
+                max_size = std::max(height, width);
+            } else if (K == 2) {
+                max_size = std::max(height, width) * 48 / 64;
+            } else if (K <= 4) {
+                max_size = std::max(height, width) / 2;
+            } else if (K <= 8) {
+                max_size = std::max(height, width) * 24 / 64;
+            } else {
+                max_size = std::max(height, width) / 4;
+            }
+
+            int cond_img_size;
+            if (K <= 4) {
+                cond_img_size = 384;
+            } else if (K <= 8) {
+                cond_img_size = 384 * 48 / 64;
+            } else {
+                cond_img_size = 384 / 2;
+            }
+
+            for (const auto& ref_image : ref_images) {
+                auto resized_ref = resize_to_area(ref_image, max_size);
+                resized_ref      = sd::ops::clamp(resized_ref, 0.0f, 1.0f);
+
+                // VLM image: Qwen3-VL expects mean=[0.5]/std=[0.5] (i.e. range [-1,1]),
+                // not CLIP normalization. Resize the already-resized ref directly to
+                // (cond_w, cond_h) to match the Python pipeline's pil_r.resize().
+                auto dims                   = calculate_dimensions(cond_img_size,
+                                                                   static_cast<double>(resized_ref.shape()[0]) / static_cast<double>(resized_ref.shape()[1]));
+                sd::Tensor<float> vlm_image = sd::ops::interpolate(
+                    resized_ref,
+                    {dims.first, dims.second, resized_ref.shape()[2], resized_ref.shape()[3]});
+                vlm_image            = vlm_image * 2.0f - 1.0f;
+                int64_t image_tokens = static_cast<int64_t>(dims.first / PATCH_SIZE) * static_cast<int64_t>(dims.second / PATCH_SIZE);
+
+                auto patch_img = resized_ref * 2.0f - 1.0f;
+                result.c_ref_images.push_back(std::move(patch_img));
+                int64_t prompt_start = static_cast<int64_t>(tokenizer.encode(prompt + "<|vision_start|>", nullptr).size());
+                prompt += "<|vision_start|>";
+                prompt += repeat_special_token("<|image_pad|>", image_tokens);
+                prompt += "<|vision_end|>";
+                vlm_images.emplace_back(static_cast<int>(prompt_start), std::move(vlm_image));
+                image_grids.push_back({1, dims.second / PATCH_SIZE, dims.first / PATCH_SIZE});
+                skip_vision_start.push_back(0);
+            }
+
+            prompt += conditioner_params.text;
+            prompt += "<|im_end|>\n<|im_start|>assistant\n<|boi_token|><|tms_token|>";
+            auto input_ids = tokenizer.encode(prompt, nullptr);
+
+            std::vector<int32_t> input_ids_pad = input_ids;
+            input_ids_pad.push_back(VISION_START_TOKEN_ID);
+            input_ids_pad.insert(input_ids_pad.end(), target_image_len - 1, IMAGE_TOKEN_ID);
+            image_grids.push_back({1, static_cast<int32_t>(height / PATCH_SIZE), static_cast<int32_t>(width / PATCH_SIZE)});
+            skip_vision_start.push_back(1);
+
+            for (const auto& ref_image : result.c_ref_images) {
+                int64_t ref_len = static_cast<int64_t>(ref_image.shape()[0] / PATCH_SIZE) * static_cast<int64_t>(ref_image.shape()[1] / PATCH_SIZE);
+                input_ids_pad.push_back(VISION_START_TOKEN_ID);
+                input_ids_pad.insert(input_ids_pad.end(), ref_len - 1, IMAGE_TOKEN_ID);
+                image_grids.push_back({1, static_cast<int32_t>(ref_image.shape()[1] / PATCH_SIZE), static_cast<int32_t>(ref_image.shape()[0] / PATCH_SIZE)});
+                skip_vision_start.push_back(1);
+            }
+
+            std::vector<int32_t> token_types(input_ids_pad.size(), 0);
+            int txt_seq_len = static_cast<int>(input_ids.size());
+            int bgn         = txt_seq_len - TIMESTEP_TOKEN_NUM;
+            for (int i = bgn; i < static_cast<int>(token_types.size()); ++i) {
+                token_types[i] = 1;
+            }
+
+            std::vector<int64_t> input_shape{static_cast<int64_t>(input_ids.size())};
+            std::vector<int64_t> position_shape{static_cast<int64_t>(input_ids_pad.size() * 4)};
+            std::vector<int64_t> token_type_shape{static_cast<int64_t>(token_types.size())};
+            std::vector<int32_t> vinput_mask(token_types.size(), 0);
+            for (int i = txt_seq_len; i < static_cast<int>(vinput_mask.size()); ++i) {
+                vinput_mask[static_cast<size_t>(i)] = 1;
+            }
+            std::vector<int64_t> vinput_mask_shape{static_cast<int64_t>(vinput_mask.size())};
+
+            result.c_input_ids    = sd::Tensor<int32_t>(input_shape, std::move(input_ids));
+            result.c_position_ids = sd::Tensor<int32_t>(position_shape, build_position_ids(input_ids_pad, image_grids, skip_vision_start));
+            result.c_token_types  = sd::Tensor<int32_t>(token_type_shape, std::move(token_types));
+            result.c_vinput_mask  = sd::Tensor<int32_t>(vinput_mask_shape, std::move(vinput_mask));
+            result.c_image_embeds.reserve(vlm_images.size());
+            for (const auto& vlm_image : vlm_images) {
+                auto image_embed = vision_runner->compute(n_threads, vlm_image.second);
+                if (image_embed.empty()) {
+                    LOG_ERROR("hidream_o1 conditioner: encode VLM image failed");
+                    return SDCondition();
+                }
+                result.c_image_embeds.emplace_back(vlm_image.first, std::move(image_embed));
+            }
+            return result;
+        }
+    };
+}  // namespace HiDreamO1
+
+#endif  // __SD_HIDREAM_O1_H__

src/latent-preview.h

@@ -4,6 +4,138 @@
 #include "ggml.h"
 #include "tensor.hpp"
 
+const float ltxav_latent_rgb_proj[128][3] = {
+    {-0.0293802f, -0.0362516f, -0.0291386f},
+    {0.0117735f, 0.0223435f, 0.018856f},
+    {0.00922335f, 0.0145666f, 0.0038772f},
+    {0.0227299f, 0.0109122f, 0.0131384f},
+    {0.00192413f, 0.0024648f, 0.00689245f},
+    {-0.0105576f, -0.0135933f, -0.00873841f},
+    {-0.0310222f, -0.0396358f, -0.0408445f},
+    {0.0149737f, 0.0316323f, 0.03415f},
+    {0.0027752f, 0.00814889f, 0.0108575f},
+    {-0.000678017f, -0.00180589f, -0.0161684f},
+    {0.0153964f, 0.0159774f, 0.0186479f},
+    {-0.0222799f, -0.0202068f, -0.0181082f},
+    {0.0128696f, 0.00754416f, -0.00673279f},
+    {0.0142729f, 0.00448099f, -0.00193934f},
+    {-0.014066f, -0.0193755f, -0.0160104f},
+    {-0.0176785f, -0.015903f, -0.0152621f},
+    {0.0307381f, 0.0292082f, 0.0328668f},
+    {0.0332928f, 0.0368629f, 0.0440893f},
+    {0.0186304f, 0.0124069f, 0.0160734f},
+    {0.00477787f, -0.00315658f, -0.000145702f},
+    {0.0183099f, 0.0122593f, 0.00599732f},
+    {-0.0194551f, -0.0183924f, -0.0147465f},
+    {0.0025732f, 0.00442582f, 0.0173176f},
+    {-0.0169423f, -0.0293863f, -0.0225908f},
+    {-0.021228f, -0.0265094f, -0.0253049f},
+    {0.0327111f, 0.0187133f, 0.0266184f},
+    {-0.0226425f, -0.0313781f, -0.0414356f},
+    {-0.0163142f, -0.0146144f, -0.0171793f},
+    {0.0192183f, 0.0108411f, 0.00829186f},
+    {-0.032246f, -0.0274846f, -0.0287434f},
+    {0.00345399f, 0.0115567f, 0.015288f},
+    {0.000972292f, 0.00331303f, 0.0110501f},
+    {0.000939494f, -0.00705084f, -0.00979449f},
+    {0.0405155f, 0.0339534f, 0.0419513f},
+    {0.0198596f, 0.0186626f, 0.0213766f},
+    {-0.00982375f, -0.00880439f, -0.00470429f},
+    {-0.0313707f, -0.0258098f, -0.0211663f},
+    {0.0144159f, 0.0117896f, 0.0141573f},
+    {0.0164571f, 0.0149178f, 0.00921599f},
+    {0.0436184f, 0.0346583f, 0.0360647f},
+    {-0.00289744f, -0.000752502f, 0.000675415f},
+    {-0.00621715f, -0.000558851f, 0.0135814f},
+    {-0.00817579f, -0.0113584f, -0.00556793f},
+    {0.00965067f, 0.0178221f, 0.015821f},
+    {0.0211832f, 0.0180827f, 0.0154707f},
+    {-0.00412858f, -0.00374182f, 0.0029568f},
+    {-0.0175603f, -0.0226242f, -0.0279012f},
+    {-0.00437471f, -0.00668329f, 0.000164887f},
+    {-0.0355983f, -0.0419093f, -0.0383065f},
+    {0.0144314f, 0.0192514f, 0.0175639f},
+    {-0.0130693f, -0.00569884f, -0.00341647f},
+    {-0.00184689f, 0.00189034f, -0.00190561f},
+    {0.019457f, 0.00842282f, 0.0123738f},
+    {-0.00477146f, -0.00206932f, 0.00283336f},
+    {-0.0364544f, -0.0256141f, -0.0322336f},
+    {-0.0295634f, -0.0295048f, -0.021057f},
+    {0.0144484f, 0.0191862f, 0.0112445f},
+    {0.0536406f, 0.0582376f, 0.0570966f},
+    {0.0085178f, 0.00748455f, 0.00995162f},
+    {-0.0136637f, -0.0172914f, -0.0195978f},
+    {-0.0339128f, -0.0392692f, -0.0355216f},
+    {0.00612855f, 0.00568303f, -0.00212333f},
+    {-0.0029225f, 0.00668819f, 0.0122131f},
+    {0.00841843f, 0.000181587f, -0.00650644f},
+    {-0.00514432f, 0.0127043f, 0.0168049f},
+    {-0.00997384f, -0.00602262f, -0.0164031f},
+    {0.0233226f, 0.033254f, 0.0307266f},
+    {-0.0110201f, -0.0164169f, -0.0161829f},
+    {-0.0195952f, -0.0177943f, -0.0115377f},
+    {-0.00523918f, -0.00452043f, 0.00267397f},
+    {0.0313464f, 0.0288241f, 0.0262496f},
+    {0.0324018f, 0.0339792f, 0.0312209f},
+    {-0.0163247f, -0.0230503f, -0.0263239f},
+    {0.000420577f, -0.00535659f, -0.00663426f},
+    {-0.012897f, -0.00203767f, -0.000622678f},
+    {-0.0632956f, -0.0651325f, -0.0584479f},
+    {-0.00426634f, -0.0150098f, -0.00719348f},
+    {0.00476109f, 0.00674315f, 0.00895472f},
+    {0.0129384f, 0.0158352f, 0.00963773f},
+    {-0.0333379f, -0.0410522f, -0.0317462f},
+    {0.00344054f, 0.00275915f, 0.00355732f},
+    {0.0209062f, 0.0273453f, 0.0222967f},
+    {0.00827287f, 0.00223045f, 0.00325844f},
+    {-0.0149132f, -0.0183973f, -0.0199781f},
+    {-0.0100786f, -0.0103681f, -0.00218224f},
+    {-0.00791409f, -0.00405153f, -0.00599893f},
+    {0.0176126f, 0.00618342f, -6.6569e-05f},
+    {0.00942486f, -0.00206494f, -0.00580324f},
+    {0.00678093f, -0.00291742f, -0.000921195f},
+    {-0.0221992f, -0.00483162f, -0.000848514f},
+    {-0.0151587f, -0.0157166f, -0.0107302f},
+    {0.00909646f, 0.0171985f, 0.0169785f},
+    {0.0127224f, 0.0170612f, 0.0303428f},
+    {0.0196562f, 0.00212451f, 0.0127744f},
+    {0.0233013f, 0.0228994f, 0.0108387f},
+    {0.00520761f, 0.00992992f, 0.0066267f},
+    {-3.77736e-05f, 0.00460229f, -0.00475132f},
+    {-0.0311763f, -0.0453566f, -0.0486901f},
+    {0.0195798f, 0.0281246f, 0.0180102f},
+    {-0.0174149f, -0.0240867f, -0.0188785f},
+    {0.000104658f, 0.00659008f, 0.0144594f},
+    {-0.00311086f, -0.0241426f, -0.0244164f},
+    {0.0336462f, 0.0305173f, 0.0331101f},
+    {0.0613625f, 0.066561f, 0.0610198f},
+    {-0.0286757f, -0.0325401f, -0.0338036f},
+    {0.0141534f, 0.0188266f, 0.0253059f},
+    {-0.00548197f, -0.00170198f, 0.00561745f},
+    {-0.0117872f, -0.00763218f, -0.0145037f},
+    {-0.0253304f, -0.0245217f, -0.0144905f},
+    {-0.00393624f, 0.00350048f, 0.00765561f},
+    {0.0113625f, 0.00561576f, -0.0113672f},
+    {-0.0301278f, -0.0261472f, -0.0301903f},
+    {0.016863f, 0.0173781f, 0.0170916f},
+    {-0.00495108f, 0.00686749f, 0.00282767f},
+    {0.00125409f, -0.00378072f, -0.00264117f},
+    {-0.00264001f, -0.00529772f, -0.0113109f},
+    {-0.054888f, -0.0575461f, -0.0509146f},
+    {-0.019442f, -0.0232916f, -0.0258637f},
+    {0.0133362f, 0.0161808f, 0.00917951f},
+    {-0.0349002f, -0.0372642f, -0.0466206f},
+    {-0.00216926f, 0.00208738f, 0.00766492f},
+    {0.0268528f, 0.0301179f, 0.0228579f},
+    {0.0226176f, 0.021536f, 0.023152f},
+    {-0.0110646f, -0.00511349f, -0.0137346f},
+    {-0.0098424f, -0.00218176f, 0.00414545f},
+    {0.00200216f, 0.00441732f, -0.0136515f},
+    {0.00695946f, 0.00313109f, -0.00379435f},
+    {0.0188377f, 0.0144059f, 0.0229724f},
+};
+float ltxav_latent_rgb_bias[3] = {0.043849f, 0.0201085f, 0.0150286f};
+
 const float wan_21_latent_rgb_proj[16][3] = {
     {0.015123f, -0.148418f, 0.479828f},
     {0.003652f, -0.010680f, -0.037142f},

src/lora.hpp

@@ -22,10 +22,11 @@ struct LoraModel : public GGMLRunner {
 
     LoraModel(const std::string& lora_id,
               ggml_backend_t backend,
+              ggml_backend_t params_backend,
               const std::string& file_path = "",
               std::string prefix           = "",
               SDVersion version            = VERSION_COUNT)
-        : lora_id(lora_id), file_path(file_path), GGMLRunner(backend, false) {
+        : lora_id(lora_id), file_path(file_path), GGMLRunner(backend, params_backend) {
         prefix = "lora." + prefix;
         if (!model_loader.init_from_file_and_convert_name(file_path, prefix, version)) {
             load_failed = true;
@@ -129,7 +130,7 @@ struct LoraModel : public GGMLRunner {
         }
     }
 
-    ggml_tensor* get_lora_weight_diff(const std::string& model_tensor_name, ggml_context* ctx) {
+    ggml_tensor* get_lora_weight_diff(const std::string& model_tensor_name, ggml_context* ctx, ggml_backend_t backend) {
         ggml_tensor* updown = nullptr;
         int index           = 0;
         while (true) {
@@ -152,17 +153,17 @@ struct LoraModel : public GGMLRunner {
 
             auto iter = lora_tensors.find(lora_up_name);
             if (iter != lora_tensors.end()) {
-                lora_up = ggml_ext_cast_f32(ctx, iter->second);
+                lora_up = ggml_ext_cast_f32(ctx, backend, iter->second);
             }
 
             iter = lora_tensors.find(lora_mid_name);
             if (iter != lora_tensors.end()) {
-                lora_mid = ggml_ext_cast_f32(ctx, iter->second);
+                lora_mid = ggml_ext_cast_f32(ctx, backend, iter->second);
             }
 
             iter = lora_tensors.find(lora_down_name);
             if (iter != lora_tensors.end()) {
-                lora_down = ggml_ext_cast_f32(ctx, iter->second);
+                lora_down = ggml_ext_cast_f32(ctx, backend, iter->second);
             }
 
             if (lora_up == nullptr || lora_down == nullptr) {
@@ -208,7 +209,7 @@ struct LoraModel : public GGMLRunner {
         return updown;
     }
 
-    ggml_tensor* get_raw_weight_diff(const std::string& model_tensor_name, ggml_context* ctx) {
+    ggml_tensor* get_raw_weight_diff(const std::string& model_tensor_name, ggml_context* ctx, ggml_backend_t backend) {
         ggml_tensor* updown = nullptr;
         int index           = 0;
         while (true) {
@@ -225,7 +226,7 @@ struct LoraModel : public GGMLRunner {
 
             auto iter = lora_tensors.find(diff_name);
             if (iter != lora_tensors.end()) {
-                curr_updown = ggml_ext_cast_f32(ctx, iter->second);
+                curr_updown = ggml_ext_cast_f32(ctx, backend, iter->second);
             } else {
                 break;
             }
@@ -248,7 +249,7 @@ struct LoraModel : public GGMLRunner {
         return updown;
     }
 
-    ggml_tensor* get_loha_weight_diff(const std::string& model_tensor_name, ggml_context* ctx) {
+    ggml_tensor* get_loha_weight_diff(const std::string& model_tensor_name, ggml_context* ctx, ggml_backend_t backend) {
         ggml_tensor* updown = nullptr;
         int index           = 0;
         while (true) {
@@ -276,33 +277,33 @@ struct LoraModel : public GGMLRunner {
 
             auto iter = lora_tensors.find(hada_1_down_name);
             if (iter != lora_tensors.end()) {
-                hada_1_down = ggml_ext_cast_f32(ctx, iter->second);
+                hada_1_down = ggml_ext_cast_f32(ctx, backend, iter->second);
             }
 
             iter = lora_tensors.find(hada_1_up_name);
             if (iter != lora_tensors.end()) {
-                hada_1_up = ggml_ext_cast_f32(ctx, iter->second);
+                hada_1_up = ggml_ext_cast_f32(ctx, backend, iter->second);
             }
 
             iter = lora_tensors.find(hada_1_mid_name);
             if (iter != lora_tensors.end()) {
-                hada_1_mid = ggml_ext_cast_f32(ctx, iter->second);
+                hada_1_mid = ggml_ext_cast_f32(ctx, backend, iter->second);
                 hada_1_up  = ggml_cont(ctx, ggml_transpose(ctx, hada_1_up));
             }
 
             iter = lora_tensors.find(hada_2_down_name);
             if (iter != lora_tensors.end()) {
-                hada_2_down = ggml_ext_cast_f32(ctx, iter->second);
+                hada_2_down = ggml_ext_cast_f32(ctx, backend, iter->second);
             }
 
             iter = lora_tensors.find(hada_2_up_name);
             if (iter != lora_tensors.end()) {
-                hada_2_up = ggml_ext_cast_f32(ctx, iter->second);
+                hada_2_up = ggml_ext_cast_f32(ctx, backend, iter->second);
             }
 
             iter = lora_tensors.find(hada_2_mid_name);
             if (iter != lora_tensors.end()) {
-                hada_2_mid = ggml_ext_cast_f32(ctx, iter->second);
+                hada_2_mid = ggml_ext_cast_f32(ctx, backend, iter->second);
                 hada_2_up  = ggml_cont(ctx, ggml_transpose(ctx, hada_2_up));
             }
 
@@ -351,7 +352,7 @@ struct LoraModel : public GGMLRunner {
         return updown;
     }
 
-    ggml_tensor* get_lokr_weight_diff(const std::string& model_tensor_name, ggml_context* ctx) {
+    ggml_tensor* get_lokr_weight_diff(const std::string& model_tensor_name, ggml_context* ctx, ggml_backend_t backend) {
         ggml_tensor* updown = nullptr;
         int index           = 0;
         while (true) {
@@ -378,24 +379,24 @@ struct LoraModel : public GGMLRunner {
 
             auto iter = lora_tensors.find(lokr_w1_name);
             if (iter != lora_tensors.end()) {
-                lokr_w1 = ggml_ext_cast_f32(ctx, iter->second);
+                lokr_w1 = ggml_ext_cast_f32(ctx, backend, iter->second);
             }
 
             iter = lora_tensors.find(lokr_w2_name);
             if (iter != lora_tensors.end()) {
-                lokr_w2 = ggml_ext_cast_f32(ctx, iter->second);
+                lokr_w2 = ggml_ext_cast_f32(ctx, backend, iter->second);
             }
 
             int64_t rank = 1;
             if (lokr_w1 == nullptr) {
                 iter = lora_tensors.find(lokr_w1_a_name);
                 if (iter != lora_tensors.end()) {
-                    lokr_w1_a = ggml_ext_cast_f32(ctx, iter->second);
+                    lokr_w1_a = ggml_ext_cast_f32(ctx, backend, iter->second);
                 }
 
                 iter = lora_tensors.find(lokr_w1_b_name);
                 if (iter != lora_tensors.end()) {
-                    lokr_w1_b = ggml_ext_cast_f32(ctx, iter->second);
+                    lokr_w1_b = ggml_ext_cast_f32(ctx, backend, iter->second);
                 }
 
                 if (lokr_w1_a == nullptr || lokr_w1_b == nullptr) {
@@ -410,12 +411,12 @@ struct LoraModel : public GGMLRunner {
             if (lokr_w2 == nullptr) {
                 iter = lora_tensors.find(lokr_w2_a_name);
                 if (iter != lora_tensors.end()) {
-                    lokr_w2_a = ggml_ext_cast_f32(ctx, iter->second);
+                    lokr_w2_a = ggml_ext_cast_f32(ctx, backend, iter->second);
                 }
 
                 iter = lora_tensors.find(lokr_w2_b_name);
                 if (iter != lora_tensors.end()) {
-                    lokr_w2_b = ggml_ext_cast_f32(ctx, iter->second);
+                    lokr_w2_b = ggml_ext_cast_f32(ctx, backend, iter->second);
                 }
 
                 if (lokr_w2_a == nullptr || lokr_w2_b == nullptr) {
@@ -468,23 +469,23 @@ struct LoraModel : public GGMLRunner {
         return updown;
     }
 
-    ggml_tensor* get_weight_diff(const std::string& model_tensor_name, ggml_context* ctx, ggml_tensor* model_tensor, bool with_lora_and_lokr = true) {
+    ggml_tensor* get_weight_diff(const std::string& model_tensor_name, ggml_backend_t backend, ggml_context* ctx, ggml_tensor* model_tensor, bool with_lora_and_lokr = true) {
         // lora
         ggml_tensor* diff = nullptr;
         if (with_lora_and_lokr) {
-            diff = get_lora_weight_diff(model_tensor_name, ctx);
+            diff = get_lora_weight_diff(model_tensor_name, ctx, backend);
         }
         // diff
         if (diff == nullptr) {
-            diff = get_raw_weight_diff(model_tensor_name, ctx);
+            diff = get_raw_weight_diff(model_tensor_name, ctx, backend);
         }
         // loha
         if (diff == nullptr) {
-            diff = get_loha_weight_diff(model_tensor_name, ctx);
+            diff = get_loha_weight_diff(model_tensor_name, ctx, backend);
         }
         // lokr
         if (diff == nullptr && with_lora_and_lokr) {
-            diff = get_lokr_weight_diff(model_tensor_name, ctx);
+            diff = get_lokr_weight_diff(model_tensor_name, ctx, backend);
         }
         if (diff != nullptr) {
             if (ggml_nelements(diff) < ggml_nelements(model_tensor)) {
@@ -502,6 +503,7 @@ struct LoraModel : public GGMLRunner {
     }
 
     ggml_tensor* get_out_diff(ggml_context* ctx,
+                              ggml_backend_t backend,
                               ggml_tensor* x,
                               WeightAdapter::ForwardParams forward_params,
                               const std::string& model_tensor_name) {
@@ -590,7 +592,7 @@ struct LoraModel : public GGMLRunner {
                 }
                 scale_value *= multiplier;
 
-                auto curr_out_diff = ggml_ext_lokr_forward(ctx, x, lokr_w1, lokr_w1_a, lokr_w1_b, lokr_w2, lokr_w2_a, lokr_w2_b, is_conv2d, forward_params.conv2d, scale_value);
+                auto curr_out_diff = ggml_ext_lokr_forward(ctx, backend, x, lokr_w1, lokr_w1_a, lokr_w1_b, lokr_w2, lokr_w2_a, lokr_w2_b, is_conv2d, forward_params.conv2d, scale_value);
                 if (out_diff == nullptr) {
                     out_diff = curr_out_diff;
                 } else {
@@ -761,7 +763,7 @@ struct LoraModel : public GGMLRunner {
             ggml_tensor* model_tensor     = it.second;
 
             // lora
-            ggml_tensor* diff = get_weight_diff(model_tensor_name, compute_ctx, model_tensor);
+            ggml_tensor* diff = get_weight_diff(model_tensor_name, runtime_backend, compute_ctx, model_tensor);
             if (diff == nullptr) {
                 continue;
             }
@@ -774,7 +776,7 @@ struct LoraModel : public GGMLRunner {
 
             ggml_tensor* final_tensor;
             if (model_tensor->type != GGML_TYPE_F32 && model_tensor->type != GGML_TYPE_F16) {
-                final_tensor = ggml_ext_cast_f32(compute_ctx, model_tensor);
+                final_tensor = ggml_ext_cast_f32(compute_ctx, runtime_backend, model_tensor);
                 final_tensor = ggml_add_inplace(compute_ctx, final_tensor, diff);
                 final_tensor = ggml_cpy(compute_ctx, final_tensor, model_tensor);
             } else {
@@ -841,34 +843,35 @@ public:
         : lora_models(lora_models) {
     }
 
-    ggml_tensor* patch_weight(ggml_context* ctx, ggml_tensor* weight, const std::string& weight_name, bool with_lora_and_lokr) {
+    ggml_tensor* patch_weight(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* weight, const std::string& weight_name, bool with_lora_and_lokr) {
         for (auto& lora_model : lora_models) {
-            ggml_tensor* diff = lora_model->get_weight_diff(weight_name, ctx, weight, with_lora_and_lokr);
+            ggml_tensor* diff = lora_model->get_weight_diff(weight_name, backend, ctx, weight, with_lora_and_lokr);
             if (diff == nullptr) {
                 continue;
             }
 
             if (weight->type != GGML_TYPE_F32 && weight->type != GGML_TYPE_F16) {
-                weight = ggml_ext_cast_f32(ctx, weight);
+                weight = ggml_ext_cast_f32(ctx, backend, weight);
             }
             weight = ggml_add(ctx, weight, diff);
         }
         return weight;
     }
 
-    ggml_tensor* patch_weight(ggml_context* ctx, ggml_tensor* weight, const std::string& weight_name) override {
-        return patch_weight(ctx, weight, weight_name, true);
+    ggml_tensor* patch_weight(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* weight, const std::string& weight_name) override {
+        return patch_weight(ctx, backend, weight, weight_name, true);
     }
 
     ggml_tensor* forward_with_lora(ggml_context* ctx,
+                                   ggml_backend_t backend,
                                    ggml_tensor* x,
                                    ggml_tensor* w,
                                    ggml_tensor* b,
                                    const std::string& prefix,
                                    WeightAdapter::ForwardParams forward_params) override {
-        w = patch_weight(ctx, w, prefix + "weight", false);
+        w = patch_weight(ctx, backend, w, prefix + "weight", false);
         if (b) {
-            b = patch_weight(ctx, b, prefix + "bias", false);
+            b = patch_weight(ctx, backend, b, prefix + "bias", false);
         }
         ggml_tensor* out;
         if (forward_params.op_type == ForwardParams::op_type_t::OP_LINEAR) {
@@ -890,7 +893,7 @@ public:
                                    forward_params.conv2d.scale);
         }
         for (auto& lora_model : lora_models) {
-            ggml_tensor* out_diff = lora_model->get_out_diff(ctx, x, forward_params, prefix + "weight");
+            ggml_tensor* out_diff = lora_model->get_out_diff(ctx, backend, x, forward_params, prefix + "weight");
             if (out_diff == nullptr) {
                 continue;
             }

src/mmdit.hpp

@@ -767,6 +767,8 @@ public:
             auto context_x = block->forward(ctx, context, x, c_mod);
             context        = context_x.first;
             x              = context_x.second;
+            sd::ggml_graph_cut::mark_graph_cut(context, "mmdit.joint_blocks." + std::to_string(i), "context");
+            sd::ggml_graph_cut::mark_graph_cut(x, "mmdit.joint_blocks." + std::to_string(i), "x");
         }
 
         x = final_layer->forward(ctx, x, c_mod);  // (N, T, patch_size ** 2 * out_channels)
@@ -809,6 +811,11 @@ public:
 
             context = context_embedder->forward(ctx, context);  // [N, L, D] aka [N, L, 1536]
         }
+        sd::ggml_graph_cut::mark_graph_cut(x, "mmdit.prelude", "x");
+        sd::ggml_graph_cut::mark_graph_cut(c, "mmdit.prelude", "c");
+        if (context != nullptr) {
+            sd::ggml_graph_cut::mark_graph_cut(context, "mmdit.prelude", "context");
+        }
 
         x = forward_core_with_concat(ctx, x, c, context, skip_layers);  // (N, H*W, patch_size ** 2 * out_channels)
 
@@ -821,10 +828,10 @@ struct MMDiTRunner : public GGMLRunner {
     MMDiT mmdit;
 
     MMDiTRunner(ggml_backend_t backend,
-                bool offload_params_to_cpu,
+                ggml_backend_t params_backend,
                 const String2TensorStorage& tensor_storage_map = {},
                 const std::string prefix                       = "")
-        : GGMLRunner(backend, offload_params_to_cpu), mmdit(tensor_storage_map) {
+        : GGMLRunner(backend, params_backend), mmdit(tensor_storage_map) {
         mmdit.init(params_ctx, tensor_storage_map, prefix);
     }
 
@@ -927,7 +934,7 @@ struct MMDiTRunner : public GGMLRunner {
         // ggml_backend_t backend    = ggml_backend_cuda_init(0);
         ggml_backend_t backend             = ggml_backend_cpu_init();
         ggml_type model_data_type          = GGML_TYPE_F16;
-        std::shared_ptr<MMDiTRunner> mmdit = std::make_shared<MMDiTRunner>(backend, false);
+        std::shared_ptr<MMDiTRunner> mmdit = std::make_shared<MMDiTRunner>(backend, backend);
         {
             LOG_INFO("loading from '%s'", file_path.c_str());
 

src/model.cpp

@@ -23,24 +23,11 @@
 
 #include "ggml-alloc.h"
 #include "ggml-backend.h"
-#include "ggml-cpu.h"
 #include "ggml.h"
+#include "ggml_extend_backend.h"
 #include "zip.h"
 
 #include "name_conversion.h"
-#include "stable-diffusion.h"
-
-#ifdef SD_USE_METAL
-#include "ggml-metal.h"
-#endif
-
-#ifdef SD_USE_VULKAN
-#include "ggml-vulkan.h"
-#endif
-
-#ifdef SD_USE_OPENCL
-#include "ggml-opencl.h"
-#endif
 
 /*================================================= Preprocess ==================================================*/
 
@@ -450,6 +437,10 @@ SDVersion ModelLoader::get_sd_version() {
         if (tensor_storage.name.find("model.diffusion_model.joint_blocks.") != std::string::npos) {
             return VERSION_SD3;
         }
+        if (tensor_storage.name.find("model.x_embedder.proj1.weight") != std::string::npos &&
+            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.img_mod.1.weight") != std::string::npos) {
             return VERSION_QWEN_IMAGE;
         }
@@ -471,6 +462,9 @@ SDVersion ModelLoader::get_sd_version() {
         if (tensor_storage.name.find("model.diffusion_model.layers.0.adaLN_sa_ln.weight") != std::string::npos) {
             return VERSION_ERNIE_IMAGE;
         }
+        if (tensor_storage.name.find("model.diffusion_model.adaln_single.emb.timestep_embedder.linear_1.bias") != std::string::npos) {
+            return VERSION_LTXAV;
+        }
         if (tensor_storage.name.find("model.diffusion_model.blocks.0.cross_attn.norm_k.weight") != std::string::npos) {
             is_wan = true;
         }
@@ -743,16 +737,10 @@ void ModelLoader::set_wtype_override(ggml_type wtype, std::string tensor_type_ru
     }
 }
 
-bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_threads_p, bool enable_mmap) {
-    int64_t process_time_ms = 0;
-    std::atomic<int64_t> read_time_ms(0);
-    std::atomic<int64_t> memcpy_time_ms(0);
-    std::atomic<int64_t> copy_to_backend_time_ms(0);
-    std::atomic<int64_t> convert_time_ms(0);
-    std::atomic<uint64_t> bytes_processed(0);
-
-    int num_threads_to_use = n_threads_p > 0 ? n_threads_p : sd_get_num_physical_cores();
-    LOG_DEBUG("using %d threads for model loading", num_threads_to_use);
+void ModelLoader::process_model_files(bool enable_mmap, bool writable_mmap) {
+    if (model_files_processed) {
+        return;
+    }
 
     int64_t start_time = ggml_time_ms();
 
@@ -764,22 +752,13 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
         processed_tensor_storages.push_back(tensor_storage);
     }
 
-    process_time_ms = ggml_time_ms() - start_time;
-
-    bool success                          = true;
-    size_t total_tensors_processed        = 0;
-    const size_t total_tensors_to_process = processed_tensor_storages.size();
-    const int64_t t_start                 = ggml_time_ms();
-    int last_n_threads                    = 1;
-
     for (size_t file_index = 0; file_index < file_paths_.size(); file_index++) {
         std::string file_path = file_paths_[file_index];
-        LOG_DEBUG("loading tensors from %s", file_path.c_str());
 
-        std::vector<const TensorStorage*> file_tensors;
+        std::vector<TensorStorage> file_tensors;
         for (const auto& ts : processed_tensor_storages) {
             if (ts.file_index == file_index) {
-                file_tensors.push_back(&ts);
+                file_tensors.push_back(ts);
             }
         }
         if (file_tensors.empty()) {
@@ -788,21 +767,169 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
 
         bool is_zip = false;
         for (auto const& ts : file_tensors) {
-            if (ts->index_in_zip >= 0) {
+            if (ts.index_in_zip >= 0) {
                 is_zip = true;
                 break;
             }
         }
 
-        std::unique_ptr<MmapWrapper> mmapped;
+        ModelFileData fdata = {};
+        fdata.path          = file_path;
+        fdata.is_zip        = is_zip;
+        fdata.tensors       = std::move(file_tensors);
+
         if (enable_mmap && !is_zip) {
             LOG_DEBUG("using mmap for I/O");
-            mmapped = MmapWrapper::create(file_path);
-            if (!mmapped) {
-                LOG_WARN("failed to memory-map '%s'", file_path.c_str());
+            std::unique_ptr<MmapWrapper> mmapped = MmapWrapper::create(file_path, writable_mmap);
+            if (mmapped) {
+                uint8_t* mmap_data             = static_cast<uint8_t*>(mmapped->writable_data());
+                ggml_backend_buffer_t buf_mmap = ggml_backend_cpu_buffer_from_ptr(mmap_data, mmapped->size());
+                if (buf_mmap) {
+                    LOG_INFO("using mmap for '%s'", file_path.c_str());
+                    fdata.mmbuffer = std::shared_ptr<struct ggml_backend_buffer>(buf_mmap, ggml_backend_buffer_free);
+                } else {
+                    LOG_WARN("mmap: failed to create backend buffer for file %s", fdata.path.c_str());
+                }
+                fdata.mmapped = std::shared_ptr<MmapWrapper>(std::move(mmapped));
+            } else {
+                LOG_WARN("failed to memory-map '%s' (falling back to read())", file_path.c_str());
             }
+        } else if (!is_zip) {
+            LOG_INFO("NOT using mmap for '%s' (mmap disabled by caller)",
+                     file_path.c_str());
         }
 
+        file_data.push_back(std::move(fdata));
+    }
+
+    model_files_processed = true;
+
+    int64_t end_time        = ggml_time_ms();
+    int64_t process_time_ms = end_time - start_time;
+
+    LOG_INFO("model files processing completed in %.2fs", process_time_ms / 1000.f);
+}
+
+std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggml_tensor*>& tensors,
+                                                       std::set<std::string> ignore_tensors,
+                                                       bool writable_mmap) {
+    process_model_files(true, writable_mmap);
+
+    std::vector<MmapTensorStore> result;
+    uint64_t mapped_bytes = 0;
+    size_t mapped_tensors = 0;
+
+    LOG_DEBUG("memory-mapping tensors...");
+
+    int64_t t_start = ggml_time_ms();
+
+    for (auto& fdata : file_data) {
+        if (!fdata.mmbuffer)
+            continue;
+
+        const std::vector<TensorStorage>& file_tensors = fdata.tensors;
+
+        size_t file_mapped_bytes   = 0;
+        size_t file_mapped_tensors = 0;
+
+        for (const auto& tensor_storage : file_tensors) {
+            const std::string& name = tensor_storage.name;
+
+            bool is_ignored = false;
+            for (const auto& ignore_prefix : ignore_tensors) {
+                if (starts_with(name, ignore_prefix)) {
+                    is_ignored = true;
+                    break;
+                }
+            }
+            if (is_ignored)
+                continue;
+
+            auto it = tensors.find(name);
+            if (it == tensors.end())
+                continue;
+
+            ggml_tensor* dst_tensor = it->second;
+            if (dst_tensor == nullptr)
+                continue;
+
+            if (tensor_storage.type != dst_tensor->type)
+                continue;
+
+            size_t tensor_size   = tensor_storage.nbytes();
+            size_t tensor_offset = tensor_storage.offset;
+
+            if (tensor_storage.ne[0] != dst_tensor->ne[0] ||
+                tensor_storage.ne[1] != dst_tensor->ne[1] ||
+                tensor_storage.ne[2] != dst_tensor->ne[2] ||
+                tensor_storage.ne[3] != dst_tensor->ne[3] ||
+                tensor_size != ggml_nbytes(dst_tensor)) {
+                // let load_tensors worry about this
+                continue;
+            }
+
+            ggml_backend_buffer_t buf_mmap = fdata.mmbuffer.get();
+            uint8_t* mmap_data             = static_cast<uint8_t*>(ggml_backend_buffer_get_base(buf_mmap));
+            dst_tensor->buffer             = buf_mmap;
+            dst_tensor->data               = mmap_data + tensor_offset;
+
+            file_mapped_bytes += tensor_size;
+            file_mapped_tensors++;
+        }
+
+        if (file_mapped_bytes > 0) {
+            mapped_tensors += file_mapped_tensors;
+            mapped_bytes += file_mapped_bytes;
+            result.push_back({fdata.mmapped, fdata.mmbuffer});
+        }
+    }
+
+    int64_t t_end       = ggml_time_ms();
+    int64_t duration_ms = t_end - t_start;
+
+    LOG_INFO("memory-mapped %zu tensors in %zu files (%.2f MB), taking %.2fs",
+             mapped_tensors,
+             result.size(),
+             mapped_bytes / (1024.0 * 1024.0),
+             duration_ms / 1000.0);
+
+    return result;
+}
+
+bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_threads_p, bool enable_mmap) {
+    process_model_files(enable_mmap, false);
+
+    std::atomic<int64_t> read_time_ms(0);
+    std::atomic<int64_t> memcpy_time_ms(0);
+    std::atomic<int64_t> copy_to_backend_time_ms(0);
+    std::atomic<int64_t> convert_time_ms(0);
+    std::atomic<uint64_t> bytes_processed(0);
+
+    int num_threads_to_use = n_threads_p > 0 ? n_threads_p : sd_get_num_physical_cores();
+    LOG_DEBUG("using %d threads for model loading", num_threads_to_use);
+
+    int64_t start_time = ggml_time_ms();
+
+    size_t total_tensors_to_process = 0;
+    for (const auto& fdata : file_data) {
+        total_tensors_to_process += fdata.tensors.size();
+    }
+
+    bool success                   = true;
+    size_t total_tensors_processed = 0;
+    const int64_t t_start          = start_time;
+    int last_n_threads             = 1;
+
+    for (auto& fdata : file_data) {
+        const std::string& file_path = fdata.path;
+        LOG_DEBUG("loading tensors from %s", file_path.c_str());
+
+        const std::vector<TensorStorage>& file_tensors = fdata.tensors;
+
+        bool is_zip = fdata.is_zip;
+
+        std::shared_ptr<MmapWrapper> mmapped = fdata.mmapped;
+
         int n_threads = is_zip ? 1 : std::min(num_threads_to_use, (int)file_tensors.size());
         if (n_threads < 1) {
             n_threads = 1;
@@ -843,7 +970,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
                         break;
                     }
 
-                    const TensorStorage& tensor_storage = *file_tensors[idx];
+                    const TensorStorage& tensor_storage = file_tensors[idx];
                     ggml_tensor* dst_tensor             = nullptr;
 
                     t0 = ggml_time_ms();
@@ -860,6 +987,11 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
                         continue;
                     }
 
+                    // skip mmapped tensors
+                    if (dst_tensor->buffer != nullptr && dst_tensor->buffer == fdata.mmbuffer.get()) {
+                        continue;
+                    }
+
                     size_t nbytes_to_read = tensor_storage.nbytes_to_read();
 
                     auto read_data = [&](char* buf, size_t n) {
@@ -1003,9 +1135,8 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread
     }
 
     int64_t end_time = ggml_time_ms();
-    LOG_INFO("loading tensors completed, taking %.2fs (process: %.2fs, read: %.2fs, memcpy: %.2fs, convert: %.2fs, copy_to_backend: %.2fs)",
+    LOG_INFO("loading tensors completed, taking %.2fs (read: %.2fs, memcpy: %.2fs, convert: %.2fs, copy_to_backend: %.2fs)",
              (end_time - start_time) / 1000.f,
-             process_time_ms / 1000.f,
              (read_time_ms.load() / (float)last_n_threads) / 1000.f,
              (memcpy_time_ms.load() / (float)last_n_threads) / 1000.f,
              (convert_time_ms.load() / (float)last_n_threads) / 1000.f,

src/model.h

@@ -42,6 +42,8 @@ enum SDVersion {
     VERSION_ANIMA,
     VERSION_FLUX2,
     VERSION_FLUX2_KLEIN,
+    VERSION_LTXAV,
+    VERSION_HIDREAM_O1,
     VERSION_Z_IMAGE,
     VERSION_OVIS_IMAGE,
     VERSION_ERNIE_IMAGE,
@@ -104,6 +106,13 @@ static inline bool sd_version_is_flux2(SDVersion version) {
     return false;
 }
 
+static inline bool sd_version_is_ltxav(SDVersion version) {
+    if (version == VERSION_LTXAV) {
+        return true;
+    }
+    return false;
+}
+
 static inline bool sd_version_is_wan(SDVersion version) {
     if (version == VERSION_WAN2 || version == VERSION_WAN2_2_I2V || version == VERSION_WAN2_2_TI2V) {
         return true;
@@ -160,9 +169,11 @@ static inline bool sd_version_is_inpaint(SDVersion version) {
 static inline bool sd_version_is_dit(SDVersion version) {
     if (sd_version_is_flux(version) ||
         sd_version_is_flux2(version) ||
+        sd_version_is_ltxav(version) ||
         sd_version_is_sd3(version) ||
         sd_version_is_wan(version) ||
         sd_version_is_qwen_image(version) ||
+        version == VERSION_HIDREAM_O1 ||
         sd_version_is_anima(version) ||
         sd_version_is_z_image(version) ||
         sd_version_is_ernie_image(version)) {
@@ -193,10 +204,27 @@ using TensorTypeRules = std::vector<std::pair<std::string, ggml_type>>;
 
 TensorTypeRules parse_tensor_type_rules(const std::string& tensor_type_rules);
 
+class MmapWrapper;
+
+struct ModelFileData {
+    std::string path;
+    std::vector<TensorStorage> tensors;
+    std::shared_ptr<MmapWrapper> mmapped;
+    std::shared_ptr<struct ggml_backend_buffer> mmbuffer;
+    bool is_zip;
+};
+
+struct MmapTensorStore {
+    std::shared_ptr<MmapWrapper> mmapped;
+    std::shared_ptr<struct ggml_backend_buffer> mmbuffer;
+};
+
 class ModelLoader {
 protected:
     SDVersion version_ = VERSION_COUNT;
     std::vector<std::string> file_paths_;
+    std::vector<ModelFileData> file_data;
+    bool model_files_processed = false;
     String2TensorStorage tensor_storage_map;
 
     void add_tensor_storage(const TensorStorage& tensor_storage);
@@ -220,6 +248,10 @@ public:
     std::map<ggml_type, uint32_t> get_vae_wtype_stat();
     String2TensorStorage& get_tensor_storage_map() { return tensor_storage_map; }
     void set_wtype_override(ggml_type wtype, std::string tensor_type_rules = "");
+    void process_model_files(bool enable_mmap = false, bool writable_mmap = true);
+    std::vector<MmapTensorStore> mmap_tensors(std::map<std::string, ggml_tensor*>& tensors,
+                                              std::set<std::string> ignore_tensors = {},
+                                              bool writable                        = true);
     bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_threads = 0, bool use_mmap = false);
     bool load_tensors(std::map<std::string, ggml_tensor*>& tensors,
                       std::set<std::string> ignore_tensors = {},

src/pmid.hpp

@@ -411,13 +411,13 @@ public:
 
 public:
     PhotoMakerIDEncoder(ggml_backend_t backend,
-                        bool offload_params_to_cpu,
+                        ggml_backend_t params_backend,
                         const String2TensorStorage& tensor_storage_map,
                         const std::string prefix,
                         SDVersion version = VERSION_SDXL,
                         PMVersion pm_v    = PM_VERSION_1,
                         float sty         = 20.f)
-        : GGMLRunner(backend, offload_params_to_cpu),
+        : GGMLRunner(backend, params_backend),
           version(version),
           pm_version(pm_v),
           style_strength(sty) {
@@ -568,11 +568,11 @@ struct PhotoMakerIDEmbed : public GGMLRunner {
     bool applied     = false;
 
     PhotoMakerIDEmbed(ggml_backend_t backend,
-                      bool offload_params_to_cpu,
+                      ggml_backend_t params_backend,
                       ModelLoader* ml,
                       const std::string& file_path = "",
                       const std::string& prefix    = "")
-        : file_path(file_path), GGMLRunner(backend, offload_params_to_cpu), model_loader(ml) {
+        : file_path(file_path), GGMLRunner(backend, params_backend), model_loader(ml) {
         if (!model_loader->init_from_file_and_convert_name(file_path, prefix)) {
             load_failed = true;
         }

src/preprocessing.hpp

@@ -24,6 +24,75 @@ static inline void preprocessing_set_4d(sd::Tensor<float>& tensor, float value,
     tensor.values()[static_cast<size_t>(preprocessing_offset_4d(tensor, i0, i1, i2, i3))] = value;
 }
 
+static inline uint8_t preprocessing_float_to_u8(float value) {
+    if (value <= 0.0f) {
+        return 0;
+    }
+    if (value >= 1.0f) {
+        return 255;
+    }
+    return static_cast<uint8_t>(value * 255.0f + 0.5f);
+}
+
+static inline void preprocessing_tensor_frame_to_sd_image(const sd::Tensor<float>& tensor, int frame_index, uint8_t* image_data) {
+    const auto& shape = tensor.shape();
+    GGML_ASSERT(shape.size() == 4 || shape.size() == 5);
+    GGML_ASSERT(image_data != nullptr);
+
+    const int width     = static_cast<int>(shape[0]);
+    const int height    = static_cast<int>(shape[1]);
+    const int channel   = static_cast<int>(shape[shape.size() == 5 ? 3 : 2]);
+    const size_t pixels = static_cast<size_t>(width) * static_cast<size_t>(height);
+    const float* src    = tensor.data();
+
+    if (shape.size() == 4) {
+        GGML_ASSERT(frame_index >= 0 && frame_index < shape[3]);
+        const size_t frame_stride = pixels * static_cast<size_t>(channel);
+        const float* frame_ptr    = src + static_cast<size_t>(frame_index) * frame_stride;
+        if (channel == 3) {
+            const float* c0 = frame_ptr;
+            const float* c1 = frame_ptr + pixels;
+            const float* c2 = frame_ptr + pixels * 2;
+            for (size_t i = 0; i < pixels; ++i) {
+                image_data[i * 3 + 0] = preprocessing_float_to_u8(c0[i]);
+                image_data[i * 3 + 1] = preprocessing_float_to_u8(c1[i]);
+                image_data[i * 3 + 2] = preprocessing_float_to_u8(c2[i]);
+            }
+            return;
+        }
+
+        for (size_t i = 0; i < pixels; ++i) {
+            for (int c = 0; c < channel; ++c) {
+                image_data[i * static_cast<size_t>(channel) + static_cast<size_t>(c)] =
+                    preprocessing_float_to_u8(frame_ptr[i + pixels * static_cast<size_t>(c)]);
+            }
+        }
+        return;
+    }
+
+    GGML_ASSERT(frame_index >= 0 && frame_index < shape[2]);
+    const size_t channel_stride = pixels * static_cast<size_t>(shape[2]);
+    const float* frame_ptr      = src + static_cast<size_t>(frame_index) * pixels;
+    if (channel == 3) {
+        const float* c0 = frame_ptr;
+        const float* c1 = frame_ptr + channel_stride;
+        const float* c2 = frame_ptr + channel_stride * 2;
+        for (size_t i = 0; i < pixels; ++i) {
+            image_data[i * 3 + 0] = preprocessing_float_to_u8(c0[i]);
+            image_data[i * 3 + 1] = preprocessing_float_to_u8(c1[i]);
+            image_data[i * 3 + 2] = preprocessing_float_to_u8(c2[i]);
+        }
+        return;
+    }
+
+    for (size_t i = 0; i < pixels; ++i) {
+        for (int c = 0; c < channel; ++c) {
+            image_data[i * static_cast<size_t>(channel) + static_cast<size_t>(c)] =
+                preprocessing_float_to_u8(frame_ptr[i + channel_stride * static_cast<size_t>(c)]);
+        }
+    }
+}
+
 static inline sd::Tensor<float> sd_image_to_preprocessing_tensor(sd_image_t image) {
     sd::Tensor<float> tensor({static_cast<int64_t>(image.width), static_cast<int64_t>(image.height), static_cast<int64_t>(image.channel), 1});
     for (uint32_t y = 0; y < image.height; ++y) {
@@ -39,20 +108,7 @@ static inline sd::Tensor<float> sd_image_to_preprocessing_tensor(sd_image_t imag
 static inline void preprocessing_tensor_to_sd_image(const sd::Tensor<float>& tensor, uint8_t* image_data) {
     GGML_ASSERT(tensor.dim() == 4);
     GGML_ASSERT(tensor.shape()[3] == 1);
-    GGML_ASSERT(image_data != nullptr);
-
-    int width   = static_cast<int>(tensor.shape()[0]);
-    int height  = static_cast<int>(tensor.shape()[1]);
-    int channel = static_cast<int>(tensor.shape()[2]);
-    for (int y = 0; y < height; ++y) {
-        for (int x = 0; x < width; ++x) {
-            for (int c = 0; c < channel; ++c) {
-                float value                               = preprocessing_get_4d(tensor, x, y, c, 0);
-                value                                     = std::min(1.0f, std::max(0.0f, value));
-                image_data[(y * width + x) * channel + c] = static_cast<uint8_t>(std::round(value * 255.0f));
-            }
-        }
-    }
+    preprocessing_tensor_frame_to_sd_image(tensor, 0, image_data);
 }
 
 static inline sd::Tensor<float> gaussian_kernel_tensor(int kernel_size) {

src/qwen_image.hpp

@@ -95,9 +95,7 @@ namespace Qwen {
 
             float scale         = 1.f / 32.f;
             bool force_prec_f32 = false;
-#ifdef SD_USE_VULKAN
-            force_prec_f32 = true;
-#endif
+
             // The purpose of the scale here is to prevent NaN issues in certain situations.
             // For example when using CUDA but the weights are k-quants (not all prompts).
             blocks["to_out.0"] = std::shared_ptr<GGMLBlock>(new Linear(inner_dim, out_dim, out_bias, false, force_prec_f32, scale));
@@ -124,6 +122,10 @@ namespace Qwen {
             auto to_v     = std::dynamic_pointer_cast<Linear>(blocks["to_v"]);
             auto to_out_0 = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
 
+            if (sd_backend_is(ctx->backend, "Vulkan")) {
+                to_out_0->set_force_prec_f32(true);
+            }
+
             auto norm_added_q = std::dynamic_pointer_cast<UnaryBlock>(blocks["norm_added_q"]);
             auto norm_added_k = std::dynamic_pointer_cast<UnaryBlock>(blocks["norm_added_k"]);
 
@@ -410,6 +412,9 @@ namespace Qwen {
             auto img = img_in->forward(ctx, x);
             auto txt = txt_norm->forward(ctx, context);
             txt      = txt_in->forward(ctx, txt);
+            sd::ggml_graph_cut::mark_graph_cut(img, "qwen_image.prelude", "img");
+            sd::ggml_graph_cut::mark_graph_cut(txt, "qwen_image.prelude", "txt");
+            // sd::ggml_graph_cut::mark_graph_cut(t_emb, "qwen_image.prelude", "t_emb");
 
             for (int i = 0; i < params.num_layers; i++) {
                 auto block = std::dynamic_pointer_cast<QwenImageTransformerBlock>(blocks["transformer_blocks." + std::to_string(i)]);
@@ -417,6 +422,8 @@ namespace Qwen {
                 auto result = block->forward(ctx, img, txt, t_emb, pe, modulate_index);
                 img         = result.first;
                 txt         = result.second;
+                sd::ggml_graph_cut::mark_graph_cut(img, "qwen_image.transformer_blocks." + std::to_string(i), "img");
+                sd::ggml_graph_cut::mark_graph_cut(txt, "qwen_image.transformer_blocks." + std::to_string(i), "txt");
             }
 
             if (params.zero_cond_t) {
@@ -481,12 +488,12 @@ namespace Qwen {
         SDVersion version;
 
         QwenImageRunner(ggml_backend_t backend,
-                        bool offload_params_to_cpu,
+                        ggml_backend_t params_backend,
                         const String2TensorStorage& tensor_storage_map = {},
                         const std::string prefix                       = "",
                         SDVersion version                              = VERSION_QWEN_IMAGE,
                         bool zero_cond_t                               = false)
-            : GGMLRunner(backend, offload_params_to_cpu) {
+            : GGMLRunner(backend, params_backend) {
             qwen_image_params.num_layers  = 0;
             qwen_image_params.zero_cond_t = zero_cond_t;
             for (auto pair : tensor_storage_map) {
@@ -679,7 +686,7 @@ namespace Qwen {
             }
 
             std::shared_ptr<QwenImageRunner> qwen_image = std::make_shared<QwenImageRunner>(backend,
-                                                                                            false,
+                                                                                            backend,
                                                                                             tensor_storage_map,
                                                                                             "model.diffusion_model",
                                                                                             VERSION_QWEN_IMAGE);

src/t5.hpp

@@ -251,7 +251,8 @@ public:
                          ggml_tensor* x,
                          ggml_tensor* past_bias                = nullptr,
                          ggml_tensor* attention_mask           = nullptr,
-                         ggml_tensor* relative_position_bucket = nullptr) {
+                         ggml_tensor* relative_position_bucket = nullptr,
+                         const std::string& graph_cut_prefix   = "") {
         // x: [N, n_token, model_dim]
         for (int i = 0; i < num_layers; i++) {
             auto block = std::dynamic_pointer_cast<T5Block>(blocks["block." + std::to_string(i)]);
@@ -259,6 +260,9 @@ public:
             auto ret  = block->forward(ctx, x, past_bias, attention_mask, relative_position_bucket);
             x         = ret.first;
             past_bias = ret.second;
+            if (!graph_cut_prefix.empty()) {
+                sd::ggml_graph_cut::mark_graph_cut(x, graph_cut_prefix + ".block." + std::to_string(i), "x");
+            }
         }
 
         auto final_layer_norm = std::dynamic_pointer_cast<T5LayerNorm>(blocks["final_layer_norm"]);
@@ -305,7 +309,8 @@ public:
         auto encoder = std::dynamic_pointer_cast<T5Stack>(blocks["encoder"]);
 
         auto x = shared->forward(ctx, input_ids);
-        x      = encoder->forward(ctx, x, past_bias, attention_mask, relative_position_bucket);
+        sd::ggml_graph_cut::mark_graph_cut(x, "t5.prelude", "x");
+        x = encoder->forward(ctx, x, past_bias, attention_mask, relative_position_bucket, "t5");
         return x;
     }
 };
@@ -316,11 +321,11 @@ struct T5Runner : public GGMLRunner {
     std::vector<int> relative_position_bucket_vec;
 
     T5Runner(ggml_backend_t backend,
-             bool offload_params_to_cpu,
+             ggml_backend_t params_backend,
              const String2TensorStorage& tensor_storage_map,
              const std::string prefix,
              bool is_umt5 = false)
-        : GGMLRunner(backend, offload_params_to_cpu) {
+        : GGMLRunner(backend, params_backend) {
         if (is_umt5) {
             params.vocab_size         = 256384;
             params.relative_attention = false;
@@ -459,11 +464,11 @@ struct T5Embedder {
     T5Runner model;
 
     T5Embedder(ggml_backend_t backend,
-               bool offload_params_to_cpu,
+               ggml_backend_t params_backend,
                const String2TensorStorage& tensor_storage_map = {},
                const std::string prefix                       = "",
                bool is_umt5                                   = false)
-        : model(backend, offload_params_to_cpu, tensor_storage_map, prefix, is_umt5), tokenizer(is_umt5) {
+        : model(backend, params_backend, tensor_storage_map, prefix, is_umt5), tokenizer(is_umt5) {
     }
 
     void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
@@ -571,7 +576,7 @@ struct T5Embedder {
             }
         }
 
-        std::shared_ptr<T5Embedder> t5 = std::make_shared<T5Embedder>(backend, false, tensor_storage_map, "", true);
+        std::shared_ptr<T5Embedder> t5 = std::make_shared<T5Embedder>(backend, backend, tensor_storage_map, "", true);
 
         t5->alloc_params_buffer();
         std::map<std::string, ggml_tensor*> tensors;

src/tae.hpp

@@ -2,7 +2,6 @@
 #define __TAE_HPP__
 
 #include "ggml_extend.hpp"
-
 #include "model.h"
 
 /*
@@ -542,14 +541,14 @@ struct TinyImageAutoEncoder : public VAE {
     bool decode_only = false;
 
     TinyImageAutoEncoder(ggml_backend_t backend,
-                         bool offload_params_to_cpu,
+                         ggml_backend_t params_backend,
                          const String2TensorStorage& tensor_storage_map,
                          const std::string prefix,
                          bool decoder_only = true,
                          SDVersion version = VERSION_SD1)
         : decode_only(decoder_only),
           taesd(decoder_only, version),
-          VAE(version, backend, offload_params_to_cpu) {
+          VAE(version, backend, params_backend) {
         scale_input = false;
         taesd.init(params_ctx, tensor_storage_map, prefix);
     }
@@ -604,14 +603,14 @@ struct TinyVideoAutoEncoder : public VAE {
     bool decode_only = false;
 
     TinyVideoAutoEncoder(ggml_backend_t backend,
-                         bool offload_params_to_cpu,
+                         ggml_backend_t params_backend,
                          const String2TensorStorage& tensor_storage_map,
                          const std::string prefix,
                          bool decoder_only = true,
                          SDVersion version = VERSION_WAN2)
         : decode_only(decoder_only),
           taehv(decoder_only, version),
-          VAE(version, backend, offload_params_to_cpu) {
+          VAE(version, backend, params_backend) {
         scale_input = false;
         taehv.init(params_ctx, tensor_storage_map, prefix);
     }

src/tensor_ggml.hpp

@@ -104,7 +104,7 @@ namespace sd {
             throw std::invalid_argument("tensor file type does not match requested sd::Tensor type");
         }
 
-        std::vector<int64_t> shape(4, 1);
+        std::vector<int64_t> shape(n_dims, 1);
         for (int i = 0; i < n_dims; ++i) {
             int32_t dim = 1;
             file.read(reinterpret_cast<char*>(&dim), sizeof(dim));

src/tokenizers/bpe_tokenizer.cpp

@@ -162,13 +162,37 @@ std::vector<int> BPETokenizer::encode(const std::string& text, on_new_token_cb_t
 
             std::string token_str = token;
             std::u32string utf32_token;
-            for (int i = 0; i < static_cast<int>(token_str.length()); i++) {
-                unsigned char b = token_str[i];
-                utf32_token += byte_encoder[b];
+            if (byte_level_bpe) {
+                for (int i = 0; i < token_str.length(); i++) {
+                    unsigned char b = token_str[i];
+                    utf32_token += byte_encoder[b];
+                }
+            } else {
+                utf32_token = utf8_to_utf32(token_str);
             }
             auto bpe_strs = bpe(utf32_token);
             for (auto bpe_str : bpe_strs) {
-                bpe_tokens.push_back(encoder[bpe_str]);
+                int token_id;
+                auto iter = encoder.find(bpe_str);
+                if (iter != encoder.end()) {
+                    token_id = iter->second;
+                } else {
+                    if (byte_fallback) {
+                        auto utf8_token_str = utf32_to_utf8(bpe_str);
+                        for (int i = 0; i < utf8_token_str.length(); i++) {
+                            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));
+                            bpe_tokens.push_back(token_id);
+                            token_strs.push_back(hex_buf);
+                        }
+                        continue;
+                    } else {
+                        token_id = UNK_TOKEN_ID;
+                    }
+                }
+                bpe_tokens.push_back(token_id);
                 token_strs.push_back(utf32_to_utf8(bpe_str));
             }
         }

src/tokenizers/bpe_tokenizer.h

@@ -20,8 +20,10 @@ protected:
     std::map<std::u32string, int> encoder;
     std::map<int, std::u32string> decoder;
     std::map<std::pair<std::u32string, std::u32string>, int> bpe_ranks;
-    int encoder_len = 0;
-    int bpe_len     = 0;
+    int encoder_len     = 0;
+    int bpe_len         = 0;
+    bool byte_level_bpe = true;
+    bool byte_fallback  = false;
 
 protected:
     static std::vector<std::pair<int, std::u32string>> bytes_to_unicode();

src/tokenizers/gemma_tokenizer.cpp

@@ -0,0 +1,191 @@
+#include "gemma_tokenizer.h"
+
+#include "ggml.h"
+#include "json.hpp"
+#include "util.h"
+#include "vocab/vocab.h"
+
+std::string GemmaTokenizer::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 GemmaTokenizer::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;
+}
+
+GemmaTokenizer::GemmaTokenizer(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
+    byte_level_bpe = false;
+    byte_fallback  = true;
+    add_bos_token  = true;
+    pad_left       = 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>",
+        "[multimodal]",
+    };
+    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>",
+        "<a>",
+        "<html>",
+        "<body>",
+        "<img>",
+        "<span>",
+        "<bbox>",
+        "<ul>",
+        "<li>",
+        "<div>",
+        "<iframe>",
+        "<footer>",
+        "</a>",
+        "</html>",
+        "</body>",
+        "</img>",
+        "</span>",
+        "</bbox>",
+        "</ul>",
+        "</li>",
+        "</div>",
+        "</iframe>",
+        "</footer>",
+    };
+    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 = {
+        "<start_of_image>",
+        "<end_of_image>",
+    };
+    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 <= 6241; i++) {
+        special_tokens_after_merge.push_back("<unused" + std::to_string(i) + ">");
+    }
+    special_tokens_after_merge.push_back("<image_soft_token>");
+
+    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_gemma_merges(), load_gemma_vocab_json());
+    }
+}

src/tokenizers/gemma_tokenizer.h

@@ -0,0 +1,17 @@
+#ifndef __SD_TOKENIZERS_GEMMA_TOKENIZER_H__
+#define __SD_TOKENIZERS_GEMMA_TOKENIZER_H__
+
+#include <string>
+
+#include "bpe_tokenizer.h"
+
+class GemmaTokenizer : 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 GemmaTokenizer(const std::string& merges_utf8_str = "", const std::string& vocab_utf8_str = "");
+};
+
+#endif  // __SD_TOKENIZERS_GEMMA_TOKENIZER_H__

src/tokenizers/qwen2_tokenizer.cpp

@@ -81,6 +81,11 @@ Qwen2Tokenizer::Qwen2Tokenizer(const std::string& merges_utf8_str) {
         "</tool_response>",
         "<think>",
         "</think>",
+        "<|boi_token|>",
+        "<|bor_token|>",
+        "<|eor_token|>",
+        "<|bot_token|>",
+        "<|tms_token|>",
     };
 
     if (merges_utf8_str.size() > 0) {

src/tokenizers/vocab/vocab.cpp

@@ -1,5 +1,7 @@
 #include "vocab.h"
 #include "clip_t5.hpp"
+#include "gemma_merges.hpp"
+#include "gemma_vocab.hpp"
 #include "mistral.hpp"
 #include "qwen.hpp"
 #include "umt5.hpp"
@@ -32,4 +34,14 @@ std::string load_t5_tokenizer_json() {
 std::string load_umt5_tokenizer_json() {
     std::string json_str(reinterpret_cast<const char*>(umt5_tokenizer_json_str), sizeof(umt5_tokenizer_json_str));
     return json_str;
+}
+
+std::string load_gemma_merges() {
+    std::string merges_utf8_str(reinterpret_cast<const char*>(gemma_merges_utf8_c_str), sizeof(gemma_merges_utf8_c_str));
+    return merges_utf8_str;
+}
+
+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;
 }

src/tokenizers/vocab/vocab.h

@@ -9,5 +9,7 @@ std::string load_mistral_merges();
 std::string load_mistral_vocab_json();
 std::string load_t5_tokenizer_json();
 std::string load_umt5_tokenizer_json();
+std::string load_gemma_merges();
+std::string load_gemma_vocab_json();
 
 #endif  // __SD_TOKENIZERS_VOCAB_VOCAB_H__

src/unet.hpp

@@ -482,12 +482,14 @@ public:
 
             emb = ggml_add(ctx->ggml_ctx, emb, label_emb);  // [N, time_embed_dim]
         }
+        // sd::ggml_graph_cut::mark_graph_cut(emb, "unet.prelude", "emb");
 
         // input_blocks
         std::vector<ggml_tensor*> hs;
 
         // input block 0
         auto h = input_blocks_0_0->forward(ctx, x);
+        sd::ggml_graph_cut::mark_graph_cut(h, "unet.input_blocks.0", "h");
 
         ggml_set_name(h, "bench-start");
         hs.push_back(h);
@@ -505,6 +507,7 @@ public:
                     std::string name = "input_blocks." + std::to_string(input_block_idx) + ".1";
                     h                = attention_layer_forward(name, ctx, h, context, num_video_frames);  // [N, mult*model_channels, h, w]
                 }
+                sd::ggml_graph_cut::mark_graph_cut(h, "unet.input_blocks." + std::to_string(input_block_idx), "h");
                 hs.push_back(h);
             }
             if (tiny_unet) {
@@ -518,6 +521,7 @@ public:
                 auto block       = std::dynamic_pointer_cast<DownSampleBlock>(blocks[name]);
 
                 h = block->forward(ctx, h);  // [N, mult*model_channels, h/(2^(i+1)), w/(2^(i+1))]
+                // sd::ggml_graph_cut::mark_graph_cut(h, "unet.input_blocks." + std::to_string(input_block_idx), "h");
                 hs.push_back(h);
             }
         }
@@ -531,6 +535,7 @@ public:
                 h = resblock_forward("middle_block.2", ctx, h, emb, num_video_frames);             // [N, 4*model_channels, h/8, w/8]
             }
         }
+        sd::ggml_graph_cut::mark_graph_cut(h, "unet.middle_block", "h");
         if (controls.size() > 0) {
             auto cs = ggml_ext_scale(ctx->ggml_ctx, controls[controls.size() - 1], control_strength, true);
             h       = ggml_add(ctx->ggml_ctx, h, cs);  // middle control
@@ -581,6 +586,7 @@ public:
                 }
 
                 output_block_idx += 1;
+                sd::ggml_graph_cut::mark_graph_cut(h, "unet.output_blocks." + std::to_string(output_block_idx - 1), "h");
             }
         }
 
@@ -597,11 +603,11 @@ struct UNetModelRunner : public GGMLRunner {
     UnetModelBlock unet;
 
     UNetModelRunner(ggml_backend_t backend,
-                    bool offload_params_to_cpu,
+                    ggml_backend_t params_backend,
                     const String2TensorStorage& tensor_storage_map,
                     const std::string prefix,
                     SDVersion version = VERSION_SD1)
-        : GGMLRunner(backend, offload_params_to_cpu), unet(version, tensor_storage_map) {
+        : GGMLRunner(backend, params_backend), unet(version, tensor_storage_map) {
         unet.init(params_ctx, tensor_storage_map, prefix);
     }
 

src/upscaler.cpp

@@ -4,49 +4,78 @@
 #include "stable-diffusion.h"
 #include "util.h"
 
+#include <utility>
+
 UpscalerGGML::UpscalerGGML(int n_threads,
                            bool direct,
-                           int tile_size)
+                           int tile_size,
+                           std::string backend_spec,
+                           std::string params_backend_spec)
     : n_threads(n_threads),
       direct(direct),
-      tile_size(tile_size) {
+      tile_size(tile_size),
+      backend_spec(std::move(backend_spec)),
+      params_backend_spec(std::move(params_backend_spec)) {
+}
+
+void UpscalerGGML::set_max_graph_vram_bytes(size_t max_vram_bytes) {
+    max_graph_vram_bytes = max_vram_bytes;
+    if (esrgan_upscaler) {
+        esrgan_upscaler->set_max_graph_vram_bytes(max_vram_bytes);
+    }
 }
 
 bool UpscalerGGML::load_from_file(const std::string& esrgan_path,
                                   bool offload_params_to_cpu,
                                   int n_threads) {
     ggml_log_set(ggml_log_callback_default, nullptr);
-#ifdef SD_USE_CUDA
-    LOG_DEBUG("Using CUDA backend");
-    backend = ggml_backend_cuda_init(0);
-#endif
-#ifdef SD_USE_METAL
-    LOG_DEBUG("Using Metal backend");
-    backend = ggml_backend_metal_init();
-#endif
-#ifdef SD_USE_VULKAN
-    LOG_DEBUG("Using Vulkan backend");
-    backend = ggml_backend_vk_init(0);
-#endif
-#ifdef SD_USE_OPENCL
-    LOG_DEBUG("Using OpenCL backend");
-    backend = ggml_backend_opencl_init();
-#endif
-#ifdef SD_USE_SYCL
-    LOG_DEBUG("Using SYCL backend");
-    backend = ggml_backend_sycl_init(0);
-#endif
+
+    std::string error;
+    if (!backend_manager.init(backend_spec.c_str(),
+                              params_backend_spec.c_str(),
+                              offload_params_to_cpu,
+                              false,
+                              false,
+                              false,
+                              &error)) {
+        LOG_ERROR("upscaler backend config failed: %s", error.c_str());
+        return false;
+    }
+    auto backend_for = [&](SDBackendModule module) {
+        ggml_backend_t module_backend = backend_manager.runtime_backend(module);
+        if (module_backend == nullptr) {
+            LOG_ERROR("failed to initialize %s backend", sd_backend_module_name(module));
+        }
+        return module_backend;
+    };
+    auto params_backend_for = [&](SDBackendModule module) {
+        ggml_backend_t module_backend = backend_manager.params_backend(module);
+        if (module_backend == nullptr) {
+            LOG_ERROR("failed to initialize %s params backend", sd_backend_module_name(module));
+        }
+        return module_backend;
+    };
+    auto ensure_backend_pair = [&](SDBackendModule module) {
+        if (backend_for(module) == nullptr) {
+            return false;
+        }
+        return params_backend_for(module) != nullptr;
+    };
+    if (!ensure_backend_pair(SDBackendModule::UPSCALER)) {
+        return false;
+    }
+
     ModelLoader model_loader;
     if (!model_loader.init_from_file_and_convert_name(esrgan_path)) {
         LOG_ERROR("init model loader from file failed: '%s'", esrgan_path.c_str());
     }
     model_loader.set_wtype_override(model_data_type);
-    if (!backend) {
-        LOG_DEBUG("Using CPU backend");
-        backend = ggml_backend_cpu_init();
-    }
     LOG_INFO("Upscaler weight type: %s", ggml_type_name(model_data_type));
-    esrgan_upscaler = std::make_shared<ESRGAN>(backend, offload_params_to_cpu, tile_size, model_loader.get_tensor_storage_map());
+    esrgan_upscaler = std::make_shared<ESRGAN>(backend_for(SDBackendModule::UPSCALER),
+                                               params_backend_for(SDBackendModule::UPSCALER),
+                                               tile_size,
+                                               model_loader.get_tensor_storage_map());
+    esrgan_upscaler->set_max_graph_vram_bytes(max_graph_vram_bytes);
     if (direct) {
         esrgan_upscaler->set_conv2d_direct_enabled(true);
     }
@@ -119,14 +148,16 @@ upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path_c_str,
                                  bool offload_params_to_cpu,
                                  bool direct,
                                  int n_threads,
-                                 int tile_size) {
+                                 int tile_size,
+                                 const char* backend,
+                                 const char* params_backend) {
     upscaler_ctx_t* upscaler_ctx = (upscaler_ctx_t*)malloc(sizeof(upscaler_ctx_t));
     if (upscaler_ctx == nullptr) {
         return nullptr;
     }
     std::string esrgan_path(esrgan_path_c_str);
 
-    upscaler_ctx->upscaler = new UpscalerGGML(n_threads, direct, tile_size);
+    upscaler_ctx->upscaler = new UpscalerGGML(n_threads, direct, tile_size, SAFE_STR(backend), SAFE_STR(params_backend));
     if (upscaler_ctx->upscaler == nullptr) {
         return nullptr;
     }

src/upscaler.h

@@ -2,6 +2,7 @@
 #define __SD_UPSCALER_H__
 
 #include "esrgan.hpp"
+#include "ggml_extend_backend.h"
 #include "stable-diffusion.h"
 #include "tensor.hpp"
 
@@ -9,21 +10,27 @@
 #include <string>
 
 struct UpscalerGGML {
-    ggml_backend_t backend    = nullptr;  // general backend
+    SDBackendManager backend_manager;
     ggml_type model_data_type = GGML_TYPE_F16;
     std::shared_ptr<ESRGAN> esrgan_upscaler;
     std::string esrgan_path;
     int n_threads;
-    bool direct   = false;
-    int tile_size = 128;
+    bool direct                 = false;
+    int tile_size               = 128;
+    size_t max_graph_vram_bytes = 0;
+    std::string backend_spec;
+    std::string params_backend_spec;
 
     UpscalerGGML(int n_threads,
-                 bool direct   = false,
-                 int tile_size = 128);
+                 bool direct                     = false,
+                 int tile_size                   = 128,
+                 std::string backend_spec        = "",
+                 std::string params_backend_spec = "");
 
     bool load_from_file(const std::string& esrgan_path,
                         bool offload_params_to_cpu,
                         int n_threads);
+    void set_max_graph_vram_bytes(size_t max_vram_bytes);
     sd::Tensor<float> upscale_tensor(const sd::Tensor<float>& input_tensor);
     sd_image_t upscale(sd_image_t input_image, uint32_t upscale_factor);
 };

src/util.cpp

@@ -23,8 +23,9 @@
 #include <unistd.h>
 #endif
 
-#include "ggml-cpu.h"
+#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) {
@@ -111,7 +112,7 @@ private:
     HANDLE hmapping_;
 };
 
-std::unique_ptr<MmapWrapper> MmapWrapper::create(const std::string& filename) {
+std::unique_ptr<MmapWrapper> MmapWrapper::create(const std::string& filename, bool writable) {
     void* mapped_data = nullptr;
     size_t file_size  = 0;
 
@@ -136,14 +137,18 @@ std::unique_ptr<MmapWrapper> MmapWrapper::create(const std::string& filename) {
 
     file_size = static_cast<size_t>(size.QuadPart);
 
-    HANDLE mapping_handle = CreateFileMapping(file_handle, nullptr, PAGE_READONLY, 0, 0, nullptr);
+    DWORD page_prot = writable ? PAGE_WRITECOPY : PAGE_READONLY;
+
+    HANDLE mapping_handle = CreateFileMapping(file_handle, nullptr, page_prot, 0, 0, nullptr);
 
     if (mapping_handle == nullptr) {
         CloseHandle(file_handle);
         return nullptr;
     }
 
-    mapped_data = MapViewOfFile(mapping_handle, FILE_MAP_READ, 0, 0, file_size);
+    DWORD view_access = writable ? FILE_MAP_COPY : FILE_MAP_READ;
+
+    mapped_data = MapViewOfFile(mapping_handle, view_access, 0, 0, file_size);
 
     if (mapped_data == nullptr) {
         CloseHandle(mapping_handle);
@@ -171,28 +176,85 @@ bool is_directory(const std::string& path) {
     return (stat(path.c_str(), &buffer) == 0 && S_ISDIR(buffer.st_mode));
 }
 
-class MmapWrapperImpl : public MmapWrapper {
-public:
-    MmapWrapperImpl(void* data, size_t size)
-        : MmapWrapper(data, size) {}
-
-    ~MmapWrapperImpl() override {
-        munmap(data_, size_);
-    }
+struct MmapFlags {
+    bool sequential;
+    bool populate;
+    bool willneed;
+    bool dontneed;
 };
 
-std::unique_ptr<MmapWrapper> MmapWrapper::create(const std::string& filename) {
+static MmapFlags get_mmap_flags() {
+    MmapFlags result          = {};
+    const char* SD_MMAP_FLAGS = std::getenv("SD_MMAP_FLAGS");
+    if (SD_MMAP_FLAGS && *SD_MMAP_FLAGS) {
+        std::stringstream ss(SD_MMAP_FLAGS);
+        std::string token;
+        while (std::getline(ss, token, ',')) {
+            std::string ntoken = trim(token);
+            std::transform(ntoken.begin(), ntoken.end(), ntoken.begin(), ::tolower);
+            if (ntoken == "sequential") {
+                result.sequential = true;
+            } else if (ntoken == "populate") {
+                result.populate = true;
+            } else if (ntoken == "willneed") {
+                result.willneed = true;
+            } else if (ntoken == "dontneed") {
+                result.dontneed = true;
+            }
+        }
+    }
+    return result;
+}
+
+class MmapWrapperImpl : public MmapWrapper {
+public:
+    MmapWrapperImpl(void* data, size_t size, int fd)
+        : MmapWrapper(data, size), fd_(fd) {}
+
+    ~MmapWrapperImpl() override {
+#ifdef __linux__
+        auto cfg_flags = get_mmap_flags();
+
+        // Drop the kernel pagecache pages for this file. madvise(DONTNEED)
+        // alone only unmaps from the process address space; pagecache
+        // entries persist (`free` reports them as buff/cache and the OOM
+        // killer doesn't touch them, but they ARE counted against
+        // overcommit and can starve other allocations on tight-RAM
+        // systems). posix_fadvise(POSIX_FADV_DONTNEED) is the documented
+        // way to evict pagecache for a specific fd's pages.
+        if (cfg_flags.dontneed) {
+            madvise(data_, size_, MADV_DONTNEED);
+            posix_fadvise(fd_, 0, 0, POSIX_FADV_DONTNEED);
+        }
+#endif
+        munmap(data_, size_);
+        close(fd_);
+    }
+
+private:
+    int fd_;
+};
+
+std::unique_ptr<MmapWrapper> MmapWrapper::create(const std::string& filename, bool writable) {
     int file_descriptor = open(filename.c_str(), O_RDONLY);
     if (file_descriptor == -1) {
         return nullptr;
     }
 
+    auto cfg_flags = get_mmap_flags();
+
     int mmap_flags = MAP_PRIVATE;
 
 #ifdef __linux__
-    // performance flags used by llama.cpp
-    // posix_fadvise(file_descriptor, 0, 0, POSIX_FADV_SEQUENTIAL);
-    // mmap_flags |= MAP_POPULATE;
+    // Sequential access hint helps the kernel read-ahead efficiently and
+    // also encourages eviction of already-read pages (the kernel keeps
+    // a smaller working set when this is set).
+    if (cfg_flags.sequential) {
+        posix_fadvise(file_descriptor, 0, 0, POSIX_FADV_SEQUENTIAL);
+    }
+    if (cfg_flags.populate) {
+        mmap_flags |= MAP_POPULATE;
+    }
 #endif
 
     struct stat sb;
@@ -203,20 +265,27 @@ std::unique_ptr<MmapWrapper> MmapWrapper::create(const std::string& filename) {
 
     size_t file_size = sb.st_size;
 
-    void* mapped_data = mmap(nullptr, file_size, PROT_READ, mmap_flags, file_descriptor, 0);
+    if (file_size == 0) {
+        close(file_descriptor);
+        return nullptr;
+    }
 
-    close(file_descriptor);
+    int mmap_prot = PROT_READ | (writable ? PROT_WRITE : 0);
+
+    void* mapped_data = mmap(nullptr, file_size, mmap_prot, mmap_flags, file_descriptor, 0);
 
     if (mapped_data == MAP_FAILED) {
+        close(file_descriptor);
         return nullptr;
     }
 
 #ifdef __linux__
-    // performance flags used by llama.cpp
-    // posix_madvise(mapped_data, file_size, POSIX_MADV_WILLNEED);
+    if (cfg_flags.willneed) {
+        posix_madvise(mapped_data, file_size, POSIX_MADV_WILLNEED);
+    }
 #endif
 
-    return std::make_unique<MmapWrapperImpl>(mapped_data, file_size);
+    return std::make_unique<MmapWrapperImpl>(mapped_data, file_size, file_descriptor);
 }
 
 #endif
@@ -495,26 +564,6 @@ sd_progress_cb_t sd_get_progress_callback() {
 void* sd_get_progress_callback_data() {
     return sd_progress_cb_data;
 }
-const char* sd_get_system_info() {
-    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;
-}
 
 sd_image_t tensor_to_sd_image(const sd::Tensor<float>& tensor, int frame_index) {
     const auto& shape = tensor.shape();
@@ -524,17 +573,7 @@ sd_image_t tensor_to_sd_image(const sd::Tensor<float>& tensor, int frame_index)
     int channel   = static_cast<int>(shape[shape.size() == 5 ? 3 : 2]);
     uint8_t* data = (uint8_t*)malloc(static_cast<size_t>(width * height * channel));
     GGML_ASSERT(data != nullptr);
-
-    for (int iw = 0; iw < width; ++iw) {
-        for (int ih = 0; ih < height; ++ih) {
-            for (int ic = 0; ic < channel; ++ic) {
-                float value                            = shape.size() == 5 ? tensor.index(iw, ih, frame_index, ic, 0)
-                                                                           : tensor.index(iw, ih, ic, frame_index);
-                value                                  = std::clamp(value, 0.0f, 1.0f);
-                data[(ih * width + iw) * channel + ic] = static_cast<uint8_t>(std::round(value * 255.0f));
-            }
-        }
-    }
+    preprocessing_tensor_frame_to_sd_image(tensor, frame_index, data);
     return {
         static_cast<uint32_t>(width),
         static_cast<uint32_t>(height),
@@ -718,3 +757,30 @@ 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"
+}
+
+const char* sd_get_system_info() {
+    using namespace ggml_cpu;
+    static char buffer[1024];
+    std::stringstream ss;
+    ss << "System Info: \n";
+    ss << "    SSE3 = " << ggml_cpu_has_sse3() << " | ";
+    ss << "    AVX = " << ggml_cpu_has_avx() << " | ";
+    ss << "    AVX2 = " << ggml_cpu_has_avx2() << " | ";
+    ss << "    AVX512 = " << ggml_cpu_has_avx512() << " | ";
+    ss << "    AVX512_VBMI = " << ggml_cpu_has_avx512_vbmi() << " | ";
+    ss << "    AVX512_VNNI = " << ggml_cpu_has_avx512_vnni() << " | ";
+    ss << "    FMA = " << ggml_cpu_has_fma() << " | ";
+    ss << "    NEON = " << ggml_cpu_has_neon() << " | ";
+    ss << "    ARM_FMA = " << ggml_cpu_has_arm_fma() << " | ";
+    ss << "    F16C = " << ggml_cpu_has_f16c() << " | ";
+    ss << "    FP16_VA = " << ggml_cpu_has_fp16_va() << " | ";
+    ss << "    WASM_SIMD = " << ggml_cpu_has_wasm_simd() << " | ";
+    ss << "    VSX = " << ggml_cpu_has_vsx() << " | ";
+    snprintf(buffer, sizeof(buffer), "%s", ss.str().c_str());
+    return buffer;
+}

src/util.h

@@ -6,6 +6,7 @@
 #include <string>
 #include <vector>
 
+#include "ggml-backend.h"
 #include "stable-diffusion.h"
 #include "tensor.hpp"
 
@@ -41,7 +42,7 @@ sd::Tensor<float> clip_preprocess(const sd::Tensor<float>& image, int target_wid
 
 class MmapWrapper {
 public:
-    static std::unique_ptr<MmapWrapper> create(const std::string& filename);
+    static std::unique_ptr<MmapWrapper> create(const std::string& filename, bool writable = false);
 
     virtual ~MmapWrapper() = default;
 
@@ -51,6 +52,7 @@ public:
     MmapWrapper& operator=(MmapWrapper&&)      = delete;
 
     const uint8_t* data() const { return static_cast<uint8_t*>(data_); }
+    uint8_t* writable_data() { return static_cast<uint8_t*>(data_); }
     size_t size() const { return size_; }
     bool copy_data(void* buf, size_t n, size_t offset) const;
 
@@ -82,6 +84,9 @@ int sd_get_preview_interval();
 bool sd_should_preview_denoised();
 bool sd_should_preview_noisy();
 
+// test if the backend is a specific one, e.g. "CUDA", "ROCm", "Vulkan" etc.
+bool sd_backend_is(ggml_backend_t backend, const std::string& name);
+
 #define LOG_DEBUG(format, ...) log_printf(SD_LOG_DEBUG, __FILE__, __LINE__, format, ##__VA_ARGS__)
 #define LOG_INFO(format, ...) log_printf(SD_LOG_INFO, __FILE__, __LINE__, format, ##__VA_ARGS__)
 #define LOG_WARN(format, ...) log_printf(SD_LOG_WARN, __FILE__, __LINE__, format, ##__VA_ARGS__)

src/vae.hpp

@@ -62,16 +62,18 @@ protected:
     }
 
 public:
-    VAE(SDVersion version, ggml_backend_t backend, bool offload_params_to_cpu)
-        : version(version), GGMLRunner(backend, offload_params_to_cpu) {}
+    VAE(SDVersion version, ggml_backend_t backend, ggml_backend_t params_backend)
+        : version(version), GGMLRunner(backend, params_backend) {}
 
     int get_scale_factor() {
         int scale_factor = 8;
-        if (version == VERSION_WAN2_2_TI2V) {
+        if (version == VERSION_LTXAV) {
+            scale_factor = 32;
+        } else if (version == VERSION_WAN2_2_TI2V) {
             scale_factor = 16;
         } else if (sd_version_uses_flux2_vae(version)) {
             scale_factor = 16;
-        } else if (version == VERSION_CHROMA_RADIANCE) {
+        } else if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1) {
             scale_factor = 1;
         }
         return scale_factor;
@@ -213,11 +215,12 @@ public:
     virtual sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents)                           = 0;
     virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix)         = 0;
     virtual void set_conv2d_scale(float scale) { SD_UNUSED(scale); };
+    virtual void set_temporal_tiling_enabled(bool enabled) { SD_UNUSED(enabled); };
 };
 
 struct FakeVAE : public VAE {
-    FakeVAE(SDVersion version, ggml_backend_t backend, bool offload_params_to_cpu)
-        : VAE(version, backend, offload_params_to_cpu) {}
+    FakeVAE(SDVersion version, ggml_backend_t backend, ggml_backend_t params_backend)
+        : VAE(version, backend, params_backend) {}
 
     int get_encoder_output_channels(int input_channels) {
         return input_channels;

src/wan.hpp

@@ -692,6 +692,7 @@ namespace WAN {
             } else {
                 x = conv1->forward(ctx, x);
             }
+            // sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.encoder.prelude", "x");
 
             // downsamples
             std::vector<int64_t> dims = {dim};
@@ -717,12 +718,14 @@ namespace WAN {
                         x = layer->forward(ctx, x, b, feat_cache, feat_idx, chunk_idx);
                     }
                 }
+                // sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.encoder.down." + std::to_string(i), "x");
             }
 
             // middle
             x = middle_0->forward(ctx, x, b, feat_cache, feat_idx);
             x = middle_1->forward(ctx, x, b);
             x = middle_2->forward(ctx, x, b, feat_cache, feat_idx);
+            // sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.encoder.mid", "x");
 
             // head
             x = head_0->forward(ctx, x);
@@ -863,11 +866,13 @@ namespace WAN {
             } else {
                 x = conv1->forward(ctx, x);
             }
+            // sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.decoder.prelude", "x");
 
             // middle
             x = middle_0->forward(ctx, x, b, feat_cache, feat_idx);
             x = middle_1->forward(ctx, x, b);
             x = middle_2->forward(ctx, x, b, feat_cache, feat_idx);
+            // sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.decoder.mid", "x");
 
             // upsamples
             std::vector<int64_t> dims = {dim_mult[dim_mult.size() - 1] * dim};
@@ -893,6 +898,7 @@ namespace WAN {
                         x = layer->forward(ctx, x, b, feat_cache, feat_idx, chunk_idx);
                     }
                 }
+                // sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.decoder.up." + std::to_string(i), "x");
             }
 
             // head
@@ -966,10 +972,10 @@ namespace WAN {
             blocks["conv2"]   = std::shared_ptr<GGMLBlock>(new CausalConv3d(z_dim, z_dim, {1, 1, 1}));
         }
 
-        ggml_tensor* patchify(ggml_context* ctx,
-                              ggml_tensor* x,
-                              int64_t patch_size,
-                              int64_t b = 1) {
+        static ggml_tensor* patchify(ggml_context* ctx,
+                                     ggml_tensor* x,
+                                     int64_t patch_size,
+                                     int64_t b = 1) {
             // x: [b*c, f, h*q, w*r]
             // return: [b*c*r*q, f, h, w]
             if (patch_size == 1) {
@@ -993,10 +999,10 @@ namespace WAN {
             return x;
         }
 
-        ggml_tensor* unpatchify(ggml_context* ctx,
-                                ggml_tensor* x,
-                                int64_t patch_size,
-                                int64_t b = 1) {
+        static ggml_tensor* unpatchify(ggml_context* ctx,
+                                       ggml_tensor* x,
+                                       int64_t patch_size,
+                                       int64_t b = 1) {
             // x: [b*c*r*q, f, h, w]
             // return: [b*c, f, h*q, w*r]
             if (patch_size == 1) {
@@ -1031,6 +1037,7 @@ namespace WAN {
             if (wan2_2) {
                 x = patchify(ctx->ggml_ctx, x, 2, b);
             }
+            // sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.encode.prelude", "x");
 
             auto encoder = std::dynamic_pointer_cast<Encoder3d>(blocks["encoder"]);
             auto conv1   = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv1"]);
@@ -1051,6 +1058,7 @@ namespace WAN {
             }
             out     = conv1->forward(ctx, out);
             auto mu = ggml_ext_chunk(ctx->ggml_ctx, out, 2, 3)[0];
+            // sd::ggml_graph_cut::mark_graph_cut(mu, "wan_vae.encode.final", "mu");
             clear_cache();
             return mu;
         }
@@ -1068,6 +1076,7 @@ namespace WAN {
 
             int64_t iter_ = z->ne[2];
             auto x        = conv2->forward(ctx, z);
+            // sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.decode.prelude", "x");
             ggml_tensor* out;
             for (int i = 0; i < iter_; i++) {
                 _conv_idx = 0;
@@ -1083,6 +1092,7 @@ namespace WAN {
             if (wan2_2) {
                 out = unpatchify(ctx->ggml_ctx, out, 2, b);
             }
+            // sd::ggml_graph_cut::mark_graph_cut(out, "wan_vae.decode.final", "out");
             clear_cache();
             return out;
         }
@@ -1097,13 +1107,15 @@ namespace WAN {
             auto decoder = std::dynamic_pointer_cast<Decoder3d>(blocks["decoder"]);
             auto conv2   = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv2"]);
 
-            auto x    = conv2->forward(ctx, z);
+            auto x = conv2->forward(ctx, z);
+            // sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.decode_partial.prelude", "x");
             auto in   = ggml_ext_slice(ctx->ggml_ctx, x, 2, i, i + 1);  // [b*c, 1, h, w]
             _conv_idx = 0;
             auto out  = decoder->forward(ctx, in, b, _feat_map, _conv_idx, i);
             if (wan2_2) {
                 out = unpatchify(ctx->ggml_ctx, out, 2, b);
             }
+            // sd::ggml_graph_cut::mark_graph_cut(out, "wan_vae.decode_partial.final", "out");
             return out;
         }
     };
@@ -1114,12 +1126,12 @@ namespace WAN {
         WanVAE ae;
 
         WanVAERunner(ggml_backend_t backend,
-                     bool offload_params_to_cpu,
+                     ggml_backend_t params_backend,
                      const String2TensorStorage& tensor_storage_map = {},
                      const std::string prefix                       = "",
                      bool decode_only                               = false,
                      SDVersion version                              = VERSION_WAN2)
-            : decode_only(decode_only), ae(decode_only, version == VERSION_WAN2_2_TI2V), VAE(version, backend, offload_params_to_cpu) {
+            : decode_only(decode_only), ae(decode_only, version == VERSION_WAN2_2_TI2V), VAE(version, backend, params_backend) {
             ae.init(params_ctx, tensor_storage_map, prefix);
         }
 
@@ -1317,7 +1329,7 @@ namespace WAN {
             // ggml_backend_t backend = ggml_backend_cuda_init(0);
             ggml_backend_t backend            = ggml_backend_cpu_init();
             ggml_type model_data_type         = GGML_TYPE_F16;
-            std::shared_ptr<WanVAERunner> vae = std::make_shared<WanVAERunner>(backend, false, String2TensorStorage{}, "", false, VERSION_WAN2_2_TI2V);
+            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());
 
@@ -1984,6 +1996,13 @@ namespace WAN {
                 c = ggml_reshape_3d(ctx->ggml_ctx, c, c->ne[0] * c->ne[1] * c->ne[2], c->ne[3] / N, N);  // [N, dim, t_len*h_len*w_len]
                 c = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, c, 1, 0, 2, 3));  // [N, t_len*h_len*w_len, dim]
             }
+            sd::ggml_graph_cut::mark_graph_cut(x, "wan.prelude", "x");
+            // sd::ggml_graph_cut::mark_graph_cut(e, "wan.prelude", "e");
+            // sd::ggml_graph_cut::mark_graph_cut(e0, "wan.prelude", "e0");
+            // sd::ggml_graph_cut::mark_graph_cut(context, "wan.prelude", "context");
+            if (c != nullptr) {
+                sd::ggml_graph_cut::mark_graph_cut(c, "wan.prelude", "c");
+            }
 
             auto x_orig = x;
 
@@ -2004,6 +2023,10 @@ namespace WAN {
                     c_skip      = ggml_ext_scale(ctx->ggml_ctx, c_skip, vace_strength);
                     x           = ggml_add(ctx->ggml_ctx, x, c_skip);
                 }
+                sd::ggml_graph_cut::mark_graph_cut(x, "wan.blocks." + std::to_string(i), "x");
+                if (c != nullptr) {
+                    sd::ggml_graph_cut::mark_graph_cut(c, "wan.blocks." + std::to_string(i), "c");
+                }
             }
 
             x = head->forward(ctx, x, e);  // [N, t_len*h_len*w_len, pt*ph*pw*out_dim]
@@ -2071,11 +2094,11 @@ namespace WAN {
         SDVersion version;
 
         WanRunner(ggml_backend_t backend,
-                  bool offload_params_to_cpu,
+                  ggml_backend_t params_backend,
                   const String2TensorStorage& tensor_storage_map = {},
                   const std::string prefix                       = "",
                   SDVersion version                              = VERSION_WAN2)
-            : GGMLRunner(backend, offload_params_to_cpu) {
+            : GGMLRunner(backend, params_backend) {
             wan_params.num_layers = 0;
             for (auto pair : tensor_storage_map) {
                 std::string tensor_name = pair.first;
@@ -2323,7 +2346,7 @@ namespace WAN {
             }
 
             std::shared_ptr<WanRunner> wan = std::make_shared<WanRunner>(backend,
-                                                                         false,
+                                                                         backend,
                                                                          tensor_storage_map,
                                                                          "model.diffusion_model",
                                                                          VERSION_WAN2_2_TI2V);

src/z_image.hpp

@@ -31,10 +31,6 @@ namespace ZImage {
             : head_dim(head_dim), num_heads(num_heads), num_kv_heads(num_kv_heads), qk_norm(qk_norm) {
             blocks["qkv"] = std::make_shared<Linear>(hidden_size, (num_heads + num_kv_heads * 2) * head_dim, false);
             float scale   = 1.f;
-#if GGML_USE_HIP
-            // Prevent NaN issues with certain ROCm setups
-            scale = 1.f / 16.f;
-#endif
             blocks["out"] = std::make_shared<Linear>(num_heads * head_dim, hidden_size, false, false, false, scale);
             if (qk_norm) {
                 blocks["q_norm"] = std::make_shared<RMSNorm>(head_dim);
@@ -52,6 +48,10 @@ namespace ZImage {
             auto qkv_proj   = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
             auto out_proj   = std::dynamic_pointer_cast<Linear>(blocks["out"]);
 
+            if (sd_backend_is(ctx->backend, "ROCm")) {
+                out_proj->set_scale(1.f / 16.f);
+            }
+
             auto qkv = qkv_proj->forward(ctx, x);                                                                            // [N, n_token, (num_heads + num_kv_heads*2)*head_dim]
             qkv      = ggml_reshape_4d(ctx->ggml_ctx, qkv, head_dim, num_heads + num_kv_heads * 2, qkv->ne[1], qkv->ne[2]);  // [N, n_token, num_heads + num_kv_heads*2, head_dim]
 
@@ -115,9 +115,7 @@ namespace ZImage {
 
             bool force_prec_f32 = false;
             float scale         = 1.f / 128.f;
-#ifdef SD_USE_VULKAN
-            force_prec_f32 = true;
-#endif
+
             // The purpose of the scale here is to prevent NaN issues in certain situations.
             // For example, when using CUDA but the weights are k-quants.
             blocks["w2"] = std::make_shared<Linear>(hidden_dim, dim, false, false, force_prec_f32, scale);
@@ -129,6 +127,10 @@ namespace ZImage {
             auto w2 = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
             auto w3 = std::dynamic_pointer_cast<Linear>(blocks["w3"]);
 
+            if (sd_backend_is(ctx->backend, "Vulkan")) {
+                w2->set_force_prec_f32(true);
+            }
+
             auto x1 = w1->forward(ctx, x);
             auto x3 = w3->forward(ctx, x);
             x       = ggml_swiglu_split(ctx->ggml_ctx, x1, x3);
@@ -369,6 +371,9 @@ namespace ZImage {
 
             auto txt = cap_embedder_1->forward(ctx, cap_embedder_0->forward(ctx, context));  // [N, n_txt_token, hidden_size]
             auto img = x_embedder->forward(ctx, x);                                          // [N, n_img_token, hidden_size]
+            sd::ggml_graph_cut::mark_graph_cut(txt, "z_image.prelude", "txt");
+            sd::ggml_graph_cut::mark_graph_cut(img, "z_image.prelude", "img");
+            sd::ggml_graph_cut::mark_graph_cut(t_emb, "z_image.prelude", "t_emb");
 
             int64_t n_txt_pad_token = Rope::bound_mod(static_cast<int>(n_txt_token), SEQ_MULTI_OF);
             if (n_txt_pad_token > 0) {
@@ -391,20 +396,24 @@ namespace ZImage {
                 auto block = std::dynamic_pointer_cast<JointTransformerBlock>(blocks["context_refiner." + std::to_string(i)]);
 
                 txt = block->forward(ctx, txt, txt_pe, nullptr, nullptr);
+                sd::ggml_graph_cut::mark_graph_cut(txt, "z_image.context_refiner." + std::to_string(i), "txt");
             }
 
             for (int i = 0; i < z_image_params.num_refiner_layers; i++) {
                 auto block = std::dynamic_pointer_cast<JointTransformerBlock>(blocks["noise_refiner." + std::to_string(i)]);
 
                 img = block->forward(ctx, img, img_pe, nullptr, t_emb);
+                sd::ggml_graph_cut::mark_graph_cut(img, "z_image.noise_refiner." + std::to_string(i), "img");
             }
 
             auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1);  // [N, n_txt_token + n_txt_pad_token + n_img_token + n_img_pad_token, hidden_size]
+            sd::ggml_graph_cut::mark_graph_cut(txt_img, "z_image.prelude", "txt_img");
 
             for (int i = 0; i < z_image_params.num_layers; i++) {
                 auto block = std::dynamic_pointer_cast<JointTransformerBlock>(blocks["layers." + std::to_string(i)]);
 
                 txt_img = block->forward(ctx, txt_img, pe, nullptr, t_emb);
+                sd::ggml_graph_cut::mark_graph_cut(txt_img, "z_image.layers." + std::to_string(i), "txt_img");
             }
 
             txt_img = final_layer->forward(ctx, txt_img, t_emb);  // [N, n_txt_token + n_txt_pad_token + n_img_token + n_img_pad_token, ph*pw*C]
@@ -464,11 +473,11 @@ namespace ZImage {
         SDVersion version;
 
         ZImageRunner(ggml_backend_t backend,
-                     bool offload_params_to_cpu,
+                     ggml_backend_t params_backend,
                      const String2TensorStorage& tensor_storage_map = {},
                      const std::string prefix                       = "",
                      SDVersion version                              = VERSION_Z_IMAGE)
-            : GGMLRunner(backend, offload_params_to_cpu) {
+            : GGMLRunner(backend, params_backend) {
             z_image = ZImageModel(z_image_params);
             z_image.init(params_ctx, tensor_storage_map, prefix);
         }
@@ -611,7 +620,7 @@ namespace ZImage {
             }
 
             std::shared_ptr<ZImageRunner> z_image = std::make_shared<ZImageRunner>(backend,
-                                                                                   false,
+                                                                                   backend,
                                                                                    tensor_storage_map,
                                                                                    "model.diffusion_model",
                                                                                    VERSION_QWEN_IMAGE);