Compare commits

...

40 Commits

Author SHA1 Message Date
leejet
c6beeef355
fix: map Qwen3-VL DeepStack GGUF tensor names (#1858) 2026-08-06 01:50:55 +08:00
stduhpf
bfbef5b7e6
feat: trained Minimax VAE Latent2rgb proj (#1856) 2026-08-06 01:39:10 +08:00
leejet
ea7f0c87cf
feat: add minimax-h3 support (#1854) 2026-08-04 23:14:18 +08:00
Huang, Hong-Chang
b4e67d1221
fix(cmake): only apply /MP to the MSVC compiler, not icx (#1846) 2026-08-04 22:37:30 +08:00
leejet
db99efdd6d
refactor: extract model loader initialization (#1844) 2026-08-02 17:24:10 +08:00
vmobilis
eb7f35ca49
feat: add linear multi-step sampling method (#1843) 2026-08-02 16:19:47 +08:00
fszontagh
50062a4bba
feat: add IP-Adapter Plus (Resampler image projection) support (#1839) 2026-08-02 16:15:28 +08:00
stduhpf
8457624101
feat: support more LoRA models (Kroma-v0.1 support) (#1842) 2026-08-02 16:08:47 +08:00
stduhpf
10378f42db
fix: lora with split qkv compatibility check at runtime (#1836)
Co-authored-by: leejet <leejet714@gmail.com>
2026-08-02 16:08:08 +08:00
leejet
e31a86ce91
refactor: centralize CLIP prefix conversion (#1837) 2026-07-30 22:44:31 +08:00
akleine
735a4ef520
fix(PhotoMaker): avoid GGML_ASSERT if trigger word 'img' was not found in prompt (#1835) 2026-07-30 21:26:57 +08:00
Wagner Bruna
af92790ffc
feat: allow customizing the alpha and beta parameters of the beta scheduler (#1834) 2026-07-30 21:24:51 +08:00
leejet
e92e86fb11
fix: prevent torch checkpoint offset overflow (#1832) 2026-07-29 23:16:29 +08:00
vmobilis
9cfe2af8f9
feat: display number of tokens for SD models (#1831) 2026-07-29 22:21:54 +08:00
yzyyzyhhh
2993b7fb43
fix: make parameter loading backend-aware (#1828) 2026-07-29 22:18:44 +08:00
Nicky Mouha
53856e7ec8
fix: null pointer dereference when loading malformed LoHa file (#1826) 2026-07-29 21:53:26 +08:00
leejet
22516991cb
fix: skip incompatible LoRA weights (#1825) 2026-07-28 00:06:46 +08:00
fszontagh
5ef4a7557d
feat: expose IP-Adapter in server request schema and capabilities (#1824) 2026-07-27 23:26:34 +08:00
vmobilis
2d0385ba85
fix: add missing sampler names (#1819) 2026-07-26 22:58:05 +08:00
leejet
87a01773be
fix: add frame dimension for Hunyuan IMG2VID encoding (#1816) 2026-07-24 22:37:26 +08:00
leejet
b0f856804c
fix: correct IP-Adapter CFG conditioning and defaults (#1815) 2026-07-24 22:01:54 +08:00
Mario Limonciello
78124b6454
ci: update ROCm releases to 7.14.0 (#1802) 2026-07-24 21:25:04 +08:00
Henry Kroll III
b338b4b4b9
docs: add Gimp plugins to UIs section (#1799) 2026-07-24 21:21:17 +08:00
Wagner Bruna
b8bf67615c
fix: correct dangling pointer to empty image reference vector (#1813) 2026-07-24 21:20:33 +08:00
fszontagh
8d377073e3
feat: add IP-Adapter support for SD 1.5 and SDXL (#1803) 2026-07-24 21:19:44 +08:00
fszontagh
5114672c48
fix: detect vision patch size for unsplit (HF-format) Qwen3-VL (#1811) 2026-07-23 22:22:21 +08:00
leejet
8a51eb9284
feat: add Mage-Flow support (#1808) 2026-07-23 01:23:25 +08:00
leejet
35fb21f3e0
fix: avoid structured binding capture in Hunyuan config (#1809) 2026-07-23 01:23:03 +08:00
somewhatfrog
cfd4cff2e6
fix: Dockerfile.vulkan add missing libraries for nvidia support (#1805) 2026-07-23 00:45:48 +08:00
project516
5e4e03c8dd
docs: fix links to sd 1.5 and sd 2.1 (#1798) 2026-07-21 22:54:34 +08:00
Wagner Bruna
2961182365
chore: add missing override declarations (#1800) 2026-07-21 22:49:52 +08:00
leejet
ea4e566ccf
feat: add hunyuan video 1.5 support (#1795) 2026-07-18 22:30:10 +08:00
leejet
b290693977
feat: add PiD 1.5 support (#1790) 2026-07-17 01:35:18 +08:00
fszontagh
7717e82c00
feat(animatediff): support img2video via --init-img (#1789) 2026-07-17 01:24:12 +08:00
leejet
fafe8e606c docs: remove star history 2026-07-17 00:10:10 +08:00
leejet
a8a91b24cd
feat: add ADetailer support (#1785) 2026-07-14 23:59:11 +08:00
fszontagh
c00a9e956d
feat: AnimateDiff SD 1.5 motion modules (v2 + v3) (#1784) 2026-07-14 23:58:03 +08:00
stduhpf
833369da84
fix: protect cross_attn and output_proj tokens for Anima LoRAs (#1786) 2026-07-14 22:59:07 +08:00
stduhpf
74bce049d0
feat: add configurable reference image processing for edit models (#1780) 2026-07-14 22:58:22 +08:00
leejet
b5d812008e
feat: add lingbot video support (#1770) 2026-07-11 17:50:59 +08:00
100 changed files with 12524 additions and 808 deletions

View File

@ -449,8 +449,8 @@ jobs:
runs-on: windows-2022
env:
ROCM_VERSION: "7.13.0"
GPU_TARGETS: "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1150;gfx1151;gfx1152;gfx1200;gfx1201"
ROCM_VERSION: "7.14.0"
GPU_TARGETS: "gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201"
steps:
- uses: actions/checkout@v3
@ -472,34 +472,68 @@ jobs:
uses: actions/cache@v4
with:
path: C:\TheRock\build
key: rocm-${{ env.ROCM_VERSION }}-gfx1151-${{ runner.os }}
key: rocm-wheels-${{ env.ROCM_VERSION }}-${{ runner.os }}
- name: ccache
uses: ggml-org/ccache-action@v1.2.16
with:
key: windows-latest-rocm-${{ env.ROCM_VERSION }}-x64
key: windows-rocm-${{ env.ROCM_VERSION }}-x64
evict-old-files: 1d
- name: Install ROCm
- name: Install ROCm with Wheels
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"
write-host "Setting up Python virtual environment"
# Create the venv directly at the cache location to avoid relocation issues
New-Item -Path "C:\TheRock\build" -ItemType Directory -Force | Out-Null
python -m venv C:\TheRock\build\.venv
& C:\TheRock\build\.venv\Scripts\Activate.ps1
write-host "Upgrading pip"
python -m pip install --upgrade pip
write-host "Installing ROCm wheels for multi-arch support"
# Install ROCm wheels for multi-arch support (this may take several minutes)
python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ "rocm[libraries,devel]==${{env.ROCM_VERSION}}"
# Pre-expand the devel tree so it is included in the cache
write-host "Initializing ROCm devel tree"
rocm-sdk init
if ($LASTEXITCODE -ne 0) { throw "rocm-sdk init failed with exit code $LASTEXITCODE" }
write-host "Completed ROCm wheel installation to C:\TheRock\build"
- name: Setup ROCm Environment
run: |
$rocmPath = "C:\TheRock\build"
$ErrorActionPreference = "Stop"
# Activate venv from cache or fresh install
& C:\TheRock\build\.venv\Scripts\Activate.ps1
# Expand the devel tree (idempotent; no-op if already done during install)
rocm-sdk init
if ($LASTEXITCODE -ne 0) { throw "rocm-sdk init failed with exit code $LASTEXITCODE" }
# Get ROCm installation paths using the rocm-sdk CLI tool
$rocmPath = (rocm-sdk path --root)
if (-not $rocmPath) { throw "rocm-sdk path --root returned empty - devel package may not be installed" }
$rocmPath = $rocmPath.Trim()
$cmakePath = (rocm-sdk path --cmake).Trim()
$binPath = (rocm-sdk path --bin).Trim()
write-host "ROCm root: $rocmPath"
write-host "CMake path: $cmakePath"
write-host "Bin path: $binPath"
echo "HIP_PATH=$rocmPath" >> $env:GITHUB_ENV
echo "CMAKE_PREFIX_PATH=$cmakePath" >> $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
echo "$binPath" >> $env:GITHUB_PATH
# Keep venv in PATH for subsequent steps
echo "C:\TheRock\build\.venv\Scripts" >> $env:GITHUB_PATH
- name: Build
run: |
@ -527,139 +561,6 @@ jobs:
- name: Pack artifacts
if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
run: |
$ErrorActionPreference = "Stop"
$dst = "build\bin"
$rocmBin = Join-Path "${env:HIP_PATH}" "bin"
$requiredRocmPaths = @(
(Join-Path $rocmBin "rocblas.dll"),
(Join-Path $rocmBin "rocblas\library")
)
foreach ($path in $requiredRocmPaths) {
if (!(Test-Path $path)) {
throw "Missing ROCm runtime dependency: $path"
}
}
foreach ($pattern in @("rocblas*.dll", "hipblas*.dll", "libhipblas*.dll")) {
Copy-Item -Path (Join-Path $rocmBin $pattern) -Destination $dst -Force -ErrorAction SilentlyContinue
}
foreach ($dir in @("rocblas", "hipblaslt")) {
$src = Join-Path $rocmBin $dir
if (Test-Path $src) {
Copy-Item -Path $src -Destination $dst -Recurse -Force
}
}
7z a sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-rocm-${{ env.ROCM_VERSION }}-x64.zip .\build\bin\*
- name: Upload artifacts
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: "26.Q1"
ROCM_VERSION: "7.1.1"
GPU_TARGETS: "gfx1150;gfx1151;gfx1200;gfx1201;gfx1100;gfx1101;gfx1102;gfx1030;gfx1031;gfx1032"
steps:
- uses: actions/checkout@v3
with:
submodules: recursive
- name: Setup Node
uses: actions/setup-node@v4
with:
node-version: 20
- name: Setup pnpm
uses: pnpm/action-setup@v4
with:
version: 10.15.1
- name: Cache ROCm Installation
id: cache-rocm
uses: actions/cache@v4
with:
path: C:\Program Files\AMD\ROCm
key: rocm-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ runner.os }}
- name: ccache
uses: ggml-org/ccache-action@v1.2.16
with:
key: windows-latest-cmake-hip-${{ env.HIPSDK_INSTALLER_VERSION }}-x64
evict-old-files: 1d
- name: Install ROCm
if: steps.cache-rocm.outputs.cache-hit != 'true'
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 }}-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)
if (-not $completed) {
Write-Error "ROCm installation timed out after 10 minutes. Killing the process"
$proc.Kill()
exit 1
}
if ($proc.ExitCode -ne 0) {
Write-Error "ROCm installation failed with exit code $($proc.ExitCode)"
exit 1
}
write-host "Completed AMD HIP SDK installation"
- name: Verify ROCm
run: |
# Find and test ROCm installation
$clangPath = Get-ChildItem 'C:\Program Files\AMD\ROCm\*\bin\clang.exe' | Select-Object -First 1
if (-not $clangPath) {
Write-Error "ROCm installation not found"
exit 1
}
& $clangPath.FullName --version
# Set HIP_PATH environment variable for later steps
echo "HIP_PATH=$(Resolve-Path 'C:\Program Files\AMD\ROCm\*\bin\clang.exe' | split-path | split-path)" >> $env:GITHUB_ENV
- name: Build
run: |
mkdir build
cd build
$env:CMAKE_PREFIX_PATH="${env:HIP_PATH}"
cmake .. `
-G "Unix Makefiles" `
-DSD_HIPBLAS=ON `
-DSD_BUILD_SHARED_LIBS=ON `
-DGGML_NATIVE=OFF `
-DCMAKE_C_COMPILER=clang `
-DCMAKE_CXX_COMPILER=clang++ `
-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: prompt/actions-commit-hash@v2
- name: Pack artifacts
if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
run: |
md "build\bin\rocblas\library\"
md "build\bin\hipblaslt\library"
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-${{ env.ROCM_VERSION }}-x64.zip .\build\bin\*
- name: Upload artifacts
@ -679,11 +580,8 @@ jobs:
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;gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1150;gfx1151;gfx1152;gfx1200;gfx1201"
- ROCM_VERSION: "7.14.0"
gpu_targets: "gfx900;gfx906;gfx908;gfx90a;gfx90c;gfx942;gfx950;gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201"
build: x64
steps:
@ -702,7 +600,7 @@ jobs:
- name: Dependencies
id: depends
run: |
sudo apt install -y build-essential cmake wget zip ninja-build
sudo apt install -y build-essential git cmake wget
- name: Free disk space
run: |
@ -723,38 +621,36 @@ jobs:
sudo apt clean
df -h
- 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: Setup TheRock
if: matrix.ROCM_VERSION != '7.2.1'
- name: Setup TheRock with Wheels
id: therock_env
run: |
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
# Create Python virtual environment
python3 -m venv .venv
source .venv/bin/activate
# Install ROCm wheels for build
# libraries = HIP runtime and CMake configs needed for linking
# devel = compilers, headers, static libs
python -m pip install --upgrade pip
python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ "rocm[libraries,devel]==${{matrix.ROCM_VERSION}}"
# Get ROCm installation paths using the rocm-sdk CLI tool
ROCM_PATH=$(rocm-sdk path --root)
CMAKE_PATH=$(rocm-sdk path --cmake)
BIN_PATH=$(rocm-sdk path --bin)
echo "ROCM_PATH=$ROCM_PATH"
echo "CMAKE_PATH=$CMAKE_PATH"
echo "BIN_PATH=$BIN_PATH"
# Set environment variables
echo "ROCM_PATH=$ROCM_PATH" >> $GITHUB_ENV
echo "CMAKE_PREFIX_PATH=$CMAKE_PATH" >> $GITHUB_ENV
echo "HIP_PATH=$ROCM_PATH" >> $GITHUB_ENV
echo "PATH=$BIN_PATH:${PATH}" >> $GITHUB_ENV
echo "LD_LIBRARY_PATH=$ROCM_PATH/lib:${LD_LIBRARY_PATH:-}" >> $GITHUB_ENV
# Keep venv activated for subsequent steps
echo "$(pwd)/.venv/bin" >> $GITHUB_PATH
# setup-node installs into /opt/hostedtoolcache, which is removed above.
# Keep Node/pnpm setup after disk cleanup so the server frontend can be embedded.
@ -839,7 +735,6 @@ jobs:
- build-and-push-docker-images
- macOS-latest-cmake
- windows-latest-cmake
- windows-latest-cmake-hip
- windows-latest-rocm
steps:

View File

@ -11,10 +11,11 @@ endif()
if (MSVC)
add_compile_definitions(_CRT_SECURE_NO_WARNINGS)
add_compile_definitions(_SILENCE_CXX17_CODECVT_HEADER_DEPRECATION_WARNING)
# /MP is MSVC-only: icx rejects it outright once offloading is enabled.
add_compile_options(
$<$<COMPILE_LANGUAGE:C>:/MP>
$<$<AND:$<COMPILE_LANGUAGE:C>,$<C_COMPILER_ID:MSVC>>:/MP>
$<$<COMPILE_LANGUAGE:C>:/utf-8>
$<$<COMPILE_LANGUAGE:CXX>:/MP>
$<$<AND:$<COMPILE_LANGUAGE:CXX>,$<CXX_COMPILER_ID:MSVC>>:/MP>
$<$<COMPILE_LANGUAGE:CXX>:/utf-8>
)
endif()
@ -312,7 +313,7 @@ set(CMAKE_POLICY_DEFAULT_CMP0077 NEW)
if (NOT SD_USE_SYSTEM_GGML)
# see https://github.com/ggerganov/ggml/pull/682
add_definitions(-DGGML_MAX_NAME=128)
add_definitions(-DGGML_MAX_NAME=160)
endif()
# deps

View File

@ -15,6 +15,7 @@ API and command-line option may change frequently.***
## 🔥Important News
* **2026/08/04** 🚀 stable-diffusion.cpp adds **Day-1 support for MiniMax-H3**
* **2026/06/25** 🚀 stable-diffusion.cpp now supports **Krea2**
* **2026/06/04** 🚀 stable-diffusion.cpp now supports **Ideogram4**
* **2026/05/31** 🚀 stable-diffusion.cpp now supports **PiD**
@ -54,19 +55,26 @@ API and command-line option may change frequently.***
- [ERNIE-Image](./docs/ernie_image.md)
- [Boogu Image](./docs/boogu_image.md)
- [Krea2](./docs/krea2.md)
- [Mage-Flow](./docs/mage_flow.md)
- [SeFi-Image](./docs/sefi_image.md)
- [HiDream-O1-Image](./docs/hidream_o1_image.md)
- [Ideogram4](./docs/ideogram4.md)
- Image Edit Models
- [Image Edit Models](./docs/edit.md)
- [FLUX.1-Kontext-dev](./docs/kontext.md)
- [Qwen Image Edit series](./docs/qwen_image_edit.md)
- [LongCat Image Edit](./docs/longcat_image.md)
- [Boogu Image Edit](./docs/boogu_image.md)
- [Mage-Flow-Edit](./docs/mage_flow.md#image-editing)
- Video Models
- [Wan2.1/Wan2.2](./docs/wan.md)
- [MiniMax-H3](./docs/minimax_h3.md)
- [LTX-2.3](./docs/ltx2.md)
- [HunyuanVideo 1.5](./docs/hunyuan_video.md)
- [LingBot-Video](./docs/lingbot_video.md)
- [PhotoMaker](./docs/photo_maker.md) support.
- [IP-Adapter](./docs/ip_adapter.md) support (SD 1.5 and SDXL, including Plus)
- Control Net support with SD 1.5
- [ADetailer](./docs/adetailer.md)
- LoRA support, same as [stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#lora)
- Latent Consistency Models support (LCM/LCM-LoRA)
- Faster and memory efficient latent decoding with [TAESD](./docs/taesd.md)
@ -120,7 +128,7 @@ API and command-line option may change frequently.***
- Stable Diffusion v1.5 from https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5
```sh
curl -L -O https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors
curl -L -O https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors
```
### Generate an image with just one command
@ -162,6 +170,7 @@ These projects wrap `stable-diffusion.cpp` for easier use in other languages/fra
These projects use `stable-diffusion.cpp` as a backend for their image generation.
- [GIMP Plugins](https://github.com/themanyone/gimp-plugins)
- [Jellybox](https://jellybox.com)
- [Stable Diffusion GUI](https://github.com/fszontagh/sd.cpp.gui.wx)
- [Stable Diffusion CLI-GUI](https://github.com/piallai/stable-diffusion.cpp)
@ -176,7 +185,3 @@ These projects use `stable-diffusion.cpp` as a backend for their image generatio
Thank you to all the people who have already contributed to stable-diffusion.cpp!
[![Contributors](https://contrib.rocks/image?repo=leejet/stable-diffusion.cpp)](https://github.com/leejet/stable-diffusion.cpp/graphs/contributors)
## Star History
[![Star History Chart](https://api.star-history.com/svg?repos=leejet/stable-diffusion.cpp&type=Date)](https://star-history.com/#leejet/stable-diffusion.cpp&Date)

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.7 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.4 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.8 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.6 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1002 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.1 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.8 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.8 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.2 MiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 1.4 MiB

Binary file not shown.

Binary file not shown.

Binary file not shown.

After

Width:  |  Height:  |  Size: 466 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 399 KiB

BIN
assets/minimax-h3/i2av.mp4 Normal file

Binary file not shown.

BIN
assets/minimax-h3/r2av.mp4 Normal file

Binary file not shown.

BIN
assets/minimax-h3/t2av.mp4 Normal file

Binary file not shown.

View File

@ -33,7 +33,7 @@ RUN cmake --build ./build --config Release -j$(nproc)
FROM ubuntu:$UBUNTU_VERSION AS runtime
RUN apt-get update && \
apt-get install --yes --no-install-recommends libgomp1 libvulkan1 mesa-vulkan-drivers && \
apt-get install --yes --no-install-recommends libgomp1 libvulkan1 mesa-vulkan-drivers libglvnd0 libgl1 libglx0 libegl1 libgles2 && \
apt-get clean
COPY --from=build /sd.cpp/build/bin /sd.cpp/bin

110
docs/adetailer.md Normal file
View File

@ -0,0 +1,110 @@
# ADetailer
`sd-cli` can run a YOLOv8 object detector on an existing or newly generated
image and perform a cropped inpaint pass for every detected object. The first
implementation supports YOLOv8 detection checkpoints. YOLOv8 segmentation and
MediaPipe models are not supported yet.
## Convert a detector
Ultralytics checkpoints must be converted before use. The converter fuses
BatchNorm into convolution layers and writes a safetensors file with the weight
names expected by the native GGML implementation.
```bash
python scripts/convert_yolov8_to_safetensors.py face_yolov8n.pt face_yolov8n.safetensors
```
The converter requires Python packages `ultralytics`, `torch`, and
`safetensors`.
Only YOLOv8 detection checkpoints are accepted.
PyTorch checkpoints use pickle internally, so only convert `.pt` files from a
trusted source.
## Repair an existing image
Use the dedicated `adetailer` mode to detect and repair objects in an existing
image:
```bash
./bin/sd-cli \
-M adetailer \
-m model.safetensors \
-i input.png \
-o repaired.png \
-p "detailed portrait photo" \
--negative-prompt "deformed face" \
--steps 24 \
--cfg-scale 6 \
--strength 0.4 \
--sampling-method dpm++2m \
--scheduler karras \
--ad-model face_yolov8n.safetensors \
--extra-ad-args "confidence=0.3,inpaint_padding=32,mask_blur=4"
```
This mode reuses the normal image-generation options for the detail pass:
- `--init-img`, `--output`, `--prompt`, and `--negative-prompt`
- `--steps`, `--cfg-scale`, `--sampling-method`, and `--scheduler`
- `--strength`, `--seed`, LoRA settings, VAE tiling, and backend assignments
- `--width` and `--height`, which also resize the input when specified
`--ad-prompt` and `--ad-negative-prompt` optionally override the normal prompts.
Values provided in `--extra-ad-args`, such as `steps`, `cfg_scale`,
`denoising_strength`, or `inpaint_width`, take precedence over inherited values.
## Repair generated images
ADetailer can also run automatically after normal image generation:
```bash
./bin/sd-cli \
-m model.safetensors \
-p "portrait photo" \
--ad-model face_yolov8n.safetensors \
--ad-prompt "[PROMPT], detailed face" \
--ad-negative-prompt "" \
--extra-ad-args "confidence=0.3,denoising_strength=0.4,inpaint_width=512,inpaint_height=512"
```
An empty ADetailer prompt inherits the main prompt. `[PROMPT]` inserts the main
prompt, `[SEP]` assigns different prompts to consecutive masks, and `[SKIP]`
skips the corresponding mask.
All settings other than the detector path and prompts are passed through
`--extra-ad-args` as a comma-separated `key=value` list:
| Key | Default | Description |
| --- | ---: | --- |
| `input_size` | `640` | Square YOLO input size; must be a multiple of 32 |
| `confidence` | `0.3` | Detection confidence threshold |
| `nms` | `0.45` | NMS IoU threshold |
| `max_detections` | `100` | Maximum detections retained after NMS |
| `mask_k_largest` | `0` | Keep only the largest K masks; zero keeps all |
| `mask_min_ratio` | `0` | Minimum bbox area relative to the image |
| `mask_max_ratio` | `1` | Maximum bbox area relative to the image |
| `dilate_erode` | `4` | Positive values dilate; negative values erode |
| `x_offset`, `y_offset` | `0` | Mask offset in pixels; positive Y moves upward |
| `mask_mode` | `none` | `none`, `merge`, or `merge_invert` |
| `merge_masks`, `invert_mask` | `false` | Boolean alternatives to `mask_mode` |
| `mask_blur` | `4` | Final composite feather radius |
| `inpaint_padding` | `32` | Padding around the detected region |
| `inpaint_width`, `inpaint_height` | mode-specific | `512x512` after generation; input/output size in `adetailer` mode |
| `denoising_strength` | mode-specific | `0.4` after generation; inherits `--strength` in `adetailer` mode |
| `steps` | `0` | Detail steps; zero inherits the main generation |
| `cfg_scale` | `-1` | Detail CFG; a negative value inherits the main generation |
| `sample_method` | inherited | Detail sampler name |
| `scheduler` | inherited | Detail scheduler name |
| `sort_by` | `none` | `none`, `left_to_right`, `center_to_edge`, or `area` |
Multiple masks are processed serially. Each completed inpaint becomes the input
for the next mask, and the seed is incremented by the mask index. Use
`mask_mode=merge` to process all detections in one inpaint pass.
The detector uses the `detector` backend module. For example, keep detection on
the CPU while diffusion runs on CUDA:
```bash
--backend "diffusion=cuda0,detector=cpu"
```

171
docs/animatediff.md Normal file
View File

@ -0,0 +1,171 @@
# AnimateDiff (SD 1.5)
AnimateDiff adds motion to a frozen Stable Diffusion 1.5 checkpoint by
injecting a temporal-attention module at 20 UNet slots. The base SD 1.5
model, VAE, and text encoder are unchanged; only the motion module produces
the temporal residual that turns a batch of independent frames into a
coherent animation. Reference: Guo et al., "AnimateDiff: Animate Your
Personalized Text-to-Image Diffusion Models without Specific Tuning"
(https://arxiv.org/abs/2307.04725).
## Download weights
- Motion module (v3, recommended)
- fp16 safetensors: https://huggingface.co/conrevo/AnimateDiff-A1111/resolve/main/motion_module/mm_sd15_v3.safetensors
- original checkpoint: https://huggingface.co/guoyww/animatediff/resolve/main/v3_sd15_mm.ckpt
- SD 1.5 base model
- any SD 1.5 checkpoint works. `realisticVisionV60B1` and `toonyou_beta3`
are the ones used in guoyww's reference configs.
- Domain Adapter LoRA (optional, v3 only, sharpens the base UNet's output
toward the motion module's trained distribution)
- ckpt: https://huggingface.co/guoyww/animatediff/resolve/main/v3_sd15_adapter.ckpt
- place under your `--lora-model-dir` and reference in the prompt as
`<lora:v3_sd15_adapter:1.0>`.
The motion module is `~836 MB` and loads alongside the SD 1.5 UNet via
`--motion-module`.
## Motion module versions
Per [animatediff.net/models](https://animatediff.net/models):
| Module | Base | Native res | Character |
|---------------------|------|------------|-----------|
| `mm_sd_v14.ckpt` | 1.5 | 256x256 | earliest, more jittery |
| `mm_sd_v15.ckpt` | 1.5 | 256x256 | improved stability over v1.4 |
| `mm_sd_v15_v2.ckpt` | 1.5 | 384x384 | significantly better motion dynamics |
| `v3_sd15_mm.ckpt` | 1.5 | 512x512 | smoothest, highest quality; pairs with a Domain Adapter LoRA |
| `mm_sdxl_v10_beta` | SDXL | 512x512 | experimental, not yet supported here |
Match your `-H -W` to the module's native resolution for best results. v3 is
trained at 512x512 - going smaller (e.g. 384x384) still works but the motion
character is closer to v2.
## Examples
Generate an 8-frame animation at 512x512, seed 42, 20 steps. The sampler /
scheduler / CFG values below match what mm_sd15_v3 was trained with; using
SD 1.5 defaults (euler_a, low CFG) produces noise-like output.
```
.\bin\Release\sd-cli.exe -M vid_gen \
--model ..\models\checkpoints\realisticVisionV60B1.safetensors \
--motion-module ..\models\animatediff\mm_sd15_v3.safetensors \
--offload-to-cpu --diffusion-fa \
-p "a red apple on a wooden table" \
--cfg-scale 8.0 --sampling-method euler --scheduler discrete \
-H 512 -W 512 --video-frames 8 --fps 8 --steps 20 -s 42 \
-o out.avi
```
Generate at the motion module's native 16-frame context (recommended for
best temporal quality). Needs more VRAM at 512x512, so drop to 384x384 or
use layer streaming:
```
.\bin\Release\sd-cli.exe -M vid_gen \
--model ..\models\checkpoints\realisticVisionV60B1.safetensors \
--motion-module ..\models\animatediff\mm_sd15_v3.safetensors \
--offload-to-cpu --diffusion-fa \
-p "photo of coastline, rocks, storm weather, wind, waves, lightning" \
--cfg-scale 8.0 --sampling-method euler --scheduler discrete \
-H 384 -W 384 --video-frames 16 --fps 8 --steps 20 -s 42 \
-o out.avi
```
Low-VRAM streaming (verified with a 2 GiB cap on RTX 3060):
```
.\bin\Release\sd-cli.exe -M vid_gen \
--model ..\models\checkpoints\realisticVisionV60B1.safetensors \
--motion-module ..\models\animatediff\mm_sd15_v3.safetensors \
--max-vram 2.0 --stream-layers --diffusion-fa \
-p "photo of coastline, rocks, storm weather, wind, waves, lightning" \
--cfg-scale 8.0 --sampling-method euler --scheduler discrete \
-H 384 -W 384 --video-frames 8 --fps 8 --steps 20 -s 42 \
-o out.avi
```
## Reference-quality reproduction
Using guoyww's official reference configs on this impl (RealisticVision v6.0
base + `mm_sd15_v3` or `mm_sd_v15_v2` + native resolution + 16 frames + euler
+ 25 steps + CFG 8 + linear beta schedule) reproduces the reference
AnimateDiff output style.
### v3 (512x512, `mm_sd15_v3`)
| Prompt | Sample |
|---------------------------------------|--------|
| B&W man on stormy coastline | <img src="../assets/animatediff/v3_man.gif" width="256"/> |
| Close-up rabbit macro shot | <img src="../assets/animatediff/v3_rabbit.gif" width="256"/> |
| Coastline, storm, waves, lightning | <img src="../assets/animatediff/v3_coast.gif" width="256"/> |
| Old house, storm, forest, night | <img src="../assets/animatediff/v3_house.gif" width="256"/> |
### v2 (384x384, `mm_sd_v15_v2.ckpt`)
| Prompt | Sample |
|---------------------------------------|--------|
| B&W man on stormy coastline | <img src="../assets/animatediff/v2_man.gif" width="256"/> |
| Close-up rabbit macro shot | <img src="../assets/animatediff/v2_rabbit.gif" width="256"/> |
| Coastline, storm, waves, lightning | <img src="../assets/animatediff/v2_coast.gif" width="256"/> |
| Old house, storm, forest, night | <img src="../assets/animatediff/v2_house.gif" width="256"/> |
Motion is strong for scenes with motion cues in the prompt (storm/waves/wind)
and subtle for static subjects (close-up macro), matching reference behavior.
v2 has an additional motion module at the UNet middle block that v3 dropped;
this impl auto-detects the topology from the checkpoint.
### v3 + Domain Adapter LoRA
Attaching the `v3_sd15_adapter` LoRA sharpens the base UNet output toward
the training distribution the motion module was fine-tuned against. Same
prompt, seed, config as above:
<img src="../assets/animatediff/v3_rabbit_domain_lora.gif" width="256"/>
Individual fur strands, glowing inner-ear, and richer forest detail become
visible compared to the no-LoRA rendering.
```
sd-cli -M vid_gen --model realisticVisionV60B1.safetensors \
--motion-module mm_sd15_v3.safetensors \
--lora-model-dir ./loras \
-p "close up photo of a rabbit ...<lora:v3_sd15_adapter:1.0>" ...
```
## img2video
Pass a pre-rendered image via `-i / --init-img` to animate FROM it. All N output frames start from the encoded init latent, then per-frame noise is added at `--strength`. Character identity, composition, and quality are anchored by the init image; the motion module adds subtle motion on top.
Left: init image rendered with `-M img_gen`. Right: 8-frame vid_gen output.
<img src="../assets/animatediff/img2video_demo.gif" width="512"/>
```
sd-cli -M img_gen ... -o init.png # any high-quality still
sd-cli -M vid_gen --motion-module mm_sd15_v3.safetensors \
-i init.png --strength 0.75 \
--cfg-scale 7.0 --sampling-method euler --scheduler karras \
-H 512 -W 512 --video-frames 8 --steps 25 -s 42 \
-p "..." -o out.avi
```
`--strength` controls how far the motion module is allowed to deviate from the init image (higher = more motion, lower = more static).
## Notes
- The motion module was trained at `video_length=16`. Running with
`--video-frames 16` gives the best coherence; F=8 works but shows a shorter
motion arc. Frame counts up to 32 are supported by the positional encoding
but exceed the trained regime and produce more static output.
- At `--video-frames 1` the motion module is skipped entirely and the output
is bit-identical to `-M img_gen`. This avoids the single-token
temporal-attention degeneracy that would otherwise emit an untrained-magnitude
residual on a single-frame sample.
- The base UNet is frozen, so character identity and style follow the SD 1.5
checkpoint you pass to `--model`. LoRAs and prompt weighting attach to the
base model in the usual way.
- No mid_block motion module in v3. `mm_sdxl_v10_beta` (SDXL variant) is not
supported yet.
- Output is written as MJPEG AVI. Use `--fps` to set playback speed.

View File

@ -153,6 +153,7 @@ still runs out of memory, tiling is enabled and the decode retried once.
| `controlnet` | ControlNet | `controlnet`, `control` |
| `photomaker` | PhotoMaker ID encoder and PhotoMaker LoRA | `photomaker`, `photomakerid`, `pmid`, `photo` |
| `upscaler` | ESRGAN upscaler | `upscaler`, `esrgan`, `hires` |
| `detector` | ADetailer YOLOv8 detector | `detector`, `adetailer`, `yolo` |
`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.

96
docs/edit.md Normal file
View File

@ -0,0 +1,96 @@
# Image Editing
Image editing in `stable-diffusion.cpp` allows you to use reference images to guide the generation process, enabling tasks like identity preservation, style transfer, or layout modification.
## Supported Models
Depending on the architecture, different models handle reference images differently.
| Model | Default Preset |
| :--- | :--- |
| [**FLUX.1-Kontext-dev**](./kontext.md) | `flux_kontext` |
| [**LongCat Image Edit**](./longcat_image.md) | `longcat` |
| [**Qwen Image Edit**](./qwen_image_edit.md) | `qwen` |
| **Qwen Image LAYERED** | `qwen_layered` |
| [**Flux.2 [Dev] / Flux.2 [Klein]**](./flux2.md) | `flux2` |
| [**Boogu Image Edit**](./boogu_image.md) | `z_image_omni` |
| **Krea2 (Community Edit LoRAs)** | `krea2_ostris_edit` |
| [**Mage-Flow-Edit**](./mage_flow.md#image-editing) | `mage_flow` |
| **Anima (Community Edit LoRAs)** | `cosmos_reference` |
Stable-diffusion.spp also supports basic Unet-based editing models like instruct-pix2pix or CosXL-Edit. This document is not about those.
---
## Configuring Reference Modes (`--ref-image-args`)
Different DiT-based editing models require different configurations to process reference images correctly (e.g., whether to use a Vision Language Model (VLM) encoder or pass VAE-encoded images directly to the DiT).
To simplify this, we provide **Presets**. By default, the system automatically selects the best preset based on the model architecture. However, you can override this using the `--ref-image-args` argument.
### Usage
The `--ref-image-args` argument accepts a comma-separated list of key-value pairs:
**Using a preset:**
`--ref-image-args "preset=qwen_layered"`
**Using a preset with a specific override:**
`--ref-image-args "preset=krea2_edit,force_ref_timestep_zero=true"`
### Available Presets
| Preset | Primary Use Case |
| :--- | :--- |
| `flux_kontext` | FLUX.1 Kontext |
| `longcat` | LongCat Image Edit |
| `flux2` | FLUX.2 models |
| `qwen` | Qwen Image Edit |
| `qwen_layered` | Qwen Image Layered |
| `z_image_omni` | Boogu, Z-Image Omni |
| `krea2_ostris_edit` | Most Krea2 Community edit LoRAs (trained with Ostris script) |
| `mage_flow` | Mage-Flow-Edit |
| `krea2_edit` | Specifically for [lbouaraba/krea2edit](https://huggingface.co/conradlocke/krea2-identity-edit). (or similar) |
| `cosmos_reference` | For Anima |
| `default` | Uses the automatic detection based on model architecture. |
---
## Advanced Parameter Reference
If presets are insufficient, you can manually configure the following parameters via `--ref-image-args`:
| Key | Type | Description | Allowed Values |
| :--- | :--- | :--- | :--- |
| `preset` | string | Overrides the automatic preset. | (See the Presets table above) |
| `pass_to_vlm` | bool | Whether reference images are passed to the VLM encoder. | `true`, `false` |
| `pass_to_dit` | bool | Whether VAE-encoded references are passed directly to the DiT. | `true`, `false` |
| `ref_index_mode` | string | Behavior of the RoPE index. | `fixed`, `increase`, `decrease` |
| `force_ref_timestep_zero` | bool | Forces timestep=0 for reference tokens. | `true`, `false` (Krea2 only) |
| `resize_before_vae` | bool | Whether reference images are resized before VAE encoding. | `true`, `false` |
| `vae_input_max_pixels` | int | Maximum pixel area for VAE reference inputs. | Integer |
| `vlm_resize_mode` | string | How to resize VLM reference inputs. | `longest_side`, `area`, `none` |
| `vlm_max_size` | int | Maximum VLM input size; interpreted according to `vlm_resize_mode`. | Integer |
| `vlm_min_size` | int | Minimum VLM input size; interpreted according to `vlm_resize_mode`. | Integer |
| `vlm_size` | int | Shortcut to set both VLM min and max size to the same value. | Integer |
### Preset Default Values
For a technical overview of how each preset is configured, see the table below.
| Preset | VLM | RoPE Index | Cond Resize | Special Notes |
| :--- | :---: | :---: | :---: | :--- |
| `flux_kontext` | No | `fixed` | `none` | |
| `longcat` | Yes | `fixed` | `area` | |
| `flux2` | No | `increase` | `none` | |
| `qwen` | Yes | `increase` | `area` | |
| `qwen_layered` | Yes | `decrease` | `area` | |
| `mage_flow` | Yes | `increase` | `longest` | `vlm_max_size = 384`, VAE input resized to target |
| `z_image_omni` | Yes | `fixed` | `area` | |
| `krea2_ostris_edit`| Yes | `increase` | `area` | `force_ref_timestep_zero = true` |
| `krea2_edit` | Yes | `increase` | `longest` | `vlm_size = 768` |
| `cosmos_reference` | No | `fixed` | `none` | `resize_before_vae = false` |
**Additional Default Notes:**
- **VLM Input Sizes:** For most presets, `vlm_max_size` and `vlm_min_size` are set to `-1`, meaning the values are model-dependent and handled automatically. In `area` mode they represent pixel area; in `longest_side` mode they represent a side length in pixels.
- **VAE Input Size:** `vae_input_max_pixels` defaults to $1024 \times 1024$ pixels (`1048576`).

24
docs/hunyuan_video.md Normal file
View File

@ -0,0 +1,24 @@
# HunyuanVideo 1.5
HunyuanVideo 1.5 uses a HunyuanVideo diffusion transformer, a causal video VAE, Qwen2.5-VL 7B for the main text conditioning,
and ByT5 Small GlyphXL for glyph-aware text conditioning.
## Download weights
- Download HunyuanVideo 1.5
- safetensors: https://huggingface.co/Comfy-Org/HunyuanVideo_1.5_repackaged/tree/main/split_files/diffusion_models
- Download vae
- safetensors: https://huggingface.co/Comfy-Org/HunyuanVideo_1.5_repackaged/tree/main/split_files/vae
- Download qwen_2.5_vl 7b
- safetensors: https://huggingface.co/Comfy-Org/Qwen-Image_ComfyUI/tree/main/split_files/text_encoders
- gguf: https://huggingface.co/mradermacher/Qwen2.5-VL-7B-Instruct-GGUF/tree/main
- Download byt5 small glyphxl
- safetensros: https://huggingface.co/Comfy-Org/HunyuanVideo_1.5_repackaged/tree/main/split_files/text_encoders
## Text-to-video example
```shell
.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model ..\models\diffusion_models\hunyuanvideo1.5_720p_t2v_fp16.safetensors --vae ..\models\vae\hunyuanvideo15_vae_fp16.safetensors --llm ..\models\text_encoders\qwen_2.5_vl_7b.safetensors --t5xxl ..\models\text_encoders\byt5_small_glyphxl_fp16.safetensors -p "a lovely cat" --cfg-scale 6.0 --sampling-method euler -v -W 1280 -H 720 --offload-to-cpu --diffusion-fa --video-frames 33 --vae-tiling
```
<video src=../assets/hunyuan_video/hy1.5_t2v.mp4 controls="controls" muted="muted" type="video/mp4"></video>

87
docs/ip_adapter.md Normal file
View File

@ -0,0 +1,87 @@
# IP-Adapter
stable-diffusion.cpp supports [IP-Adapter](https://github.com/tencent-ailab/IP-Adapter)
image-prompt conditioning for SD 1.5 and SDXL. Given a reference image,
IP-Adapter transfers the subject and appearance of that image into the
generation, alongside the text prompt.
IP-Adapter encodes the reference image with a CLIP-Vision (ViT-H/14)
encoder, projects the embedding into a few image tokens, and injects them
through a decoupled cross-attention added to every attn2 layer of the
UNet. It composes with Control Net, so a reference image (appearance) and
an OpenPose hint (pose) can be combined in a single generation.
Both the classic adapters and the higher-fidelity **Plus** adapters are
supported; see [Plus variants](#plus-variants) below. The variant is
detected from the weight file, so the same options work for both.
## Required weights
1. A base SD 1.5 or SDXL model.
2. A CLIP-Vision (ViT-H/14) image encoder, passed with `--clip_vision`
(for example `clip_vision_h.safetensors`).
3. An IP-Adapter weight file, passed with `--ip-adapter`. The `vit-h`
variants reuse the same ViT-H encoder as above. From
[h94/IP-Adapter](https://huggingface.co/h94/IP-Adapter):
- SD 1.5: `models/ip-adapter_sd15.safetensors`
- SDXL: `sdxl_models/ip-adapter_sdxl_vit-h.safetensors`
- SD 1.5 Plus: `models/ip-adapter-plus_sd15.safetensors`
- SDXL Plus: `sdxl_models/ip-adapter-plus_sdxl_vit-h.safetensors`
The Plus files (`ip-adapter-plus_*`) are used exactly like the classic
ones; see [Plus variants](#plus-variants).
## Options
- `--ip-adapter <path>` path to the IP-Adapter weight file.
- `--ip-adapter-image <path>` path to the reference image.
- `--ip-adapter-strength <float>` strength of the IP-Adapter injection
(default 1.0). Lower values let the text prompt dominate; 0.6 to 0.8 is
a good starting range.
## Example (SD 1.5)
```
sd-cli -m ..\models\sd_v1.5.safetensors --clip_vision ..\models\clip_vision_h.safetensors --ip-adapter ..\models\ip-adapter_sd15.safetensors --ip-adapter-image ..\assets\reference.png --ip-adapter-strength 0.8 -p "a woman, best quality" -n "lowres, bad anatomy" --cfg-scale 7 --steps 30 --sampling-method dpm++2m --scheduler karras -W 512 -H 512
```
## Example (SDXL)
```
sd-cli -m ..\models\sdxl.safetensors --clip_vision ..\models\clip_vision_h.safetensors --ip-adapter ..\models\ip-adapter_sdxl_vit-h.safetensors --ip-adapter-image ..\assets\reference.png --ip-adapter-strength 0.8 -p "a woman, best quality" -n "lowres, bad anatomy" --cfg-scale 6 --steps 25 --sampling-method dpm++2m --scheduler karras -W 1024 -H 1024 --diffusion-fa --vae-tiling
```
The SDXL VAE decode at 1024x1024 is memory heavy; add `--vae-tiling` (and
`--offload-to-cpu`) on GPUs with limited VRAM.
## Plus variants
The Plus adapters (`ip-adapter-plus_sd15`, `ip-adapter-plus_sdxl_vit-h`)
replace the small linear image projection with a Resampler (a
Perceiver-style module with learned latent queries). Instead of pooling the
CLIP-Vision output into one vector, the Resampler attends over the full grid
of penultimate CLIP-Vision hidden states and emits more image tokens (16
instead of 4). The result transfers finer detail and layout from the
reference, at a small extra cost in the image-projection step.
No extra flags are needed. The variant is detected from the weight file (the
Resampler's `image_proj.latents` tensor), and every Resampler dimension is
read from the tensor shapes, so the same `--ip-adapter`,
`--ip-adapter-image`, and `--ip-adapter-strength` options apply. Plus
composes with Control Net in the same way as the classic adapters.
```
sd-cli -m ..\models\sd_v1.5.safetensors --clip_vision ..\models\clip_vision_h.safetensors --ip-adapter ..\models\ip-adapter-plus_sd15.safetensors --ip-adapter-image ..\assets\reference.png --ip-adapter-strength 0.8 -p "a woman, best quality" -n "lowres, bad anatomy" --cfg-scale 7 --steps 30 --sampling-method dpm++2m --scheduler karras -W 512 -H 512
```
The startup log line `IP-Adapter: 16 image tokens` (versus `4` for the
classic adapters) confirms a Plus file was loaded.
## Combining with Control Net
Add the usual Control Net options to keep the reference appearance while
controlling the pose:
```
sd-cli -m ..\models\sdxl.safetensors --clip_vision ..\models\clip_vision_h.safetensors --ip-adapter ..\models\ip-adapter_sdxl_vit-h.safetensors --ip-adapter-image ..\assets\character.png --ip-adapter-strength 0.9 --control-net ..\models\OpenPoseXL2.safetensors --control-image ..\assets\pose.png --control-strength 0.8 -p "a character, side view" --cfg-scale 6 --steps 25 -W 1024 -H 1024 --diffusion-fa --vae-tiling
```

32
docs/lingbot_video.md Normal file
View File

@ -0,0 +1,32 @@
# How to Use
Lingbot Video uses a Lingbot diffusion transformer, the Wan2.1 VAE, and Qwen3-VL 4B as the LLM text encoder.
## Download weights
- Download lingbot-video-dense-1.3b
- safetensors: https://huggingface.co/robbyant/lingbot-video-dense-1.3b/tree/main/transformer
- Download lingbot-video-moe-30b-a3b
- safetensors: https://huggingface.co/robbyant/lingbot-video-moe-30b-a3b/tree/main/transformer
- Download vae
- safetensors: https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/blob/main/split_files/vae/wan_2.1_vae.safetensors
- Download Qwen3-VL 4B
- safetensors: https://huggingface.co/Comfy-Org/Krea-2/tree/main/text_encoders
- gguf: https://huggingface.co/Qwen/Qwen3-VL-4B-Instruct-GGUF/tree/main
## Notes
- Use `-M vid_gen`.
- T2V uses the text prompt only.
- I2V uses `-i` as the first video frame. The same image is also passed to the
Qwen3-VL prompt enhancer when vision weights are available.
- Video frames are aligned to Wan-style temporal compression, so use frame counts
like 33, 49, or 81.
## Examples
### LingBot-Video T2V
```sh
.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model ..\models\diffusion_models\lingbot-video-dens-1.3b.safetensors --llm ..\models\text_encoders\Qwen3-VL-4B-Instruct-Q4_K_M.gguf --vae ..\models\vae\wan_2.1_vae.safetensors -p '{"caption":{"comprehensive_description":"A lovely cat sits comfortably on a soft cushion near a sunlit window, looking calm, gentle, and adorable. The cat has soft fluffy fur, bright expressive eyes, small rounded ears, delicate whiskers, and a relaxed posture. Warm daylight falls across the cat from one side, creating soft highlights on the fur and a cozy glow around the scene. The background is softly blurred, showing hints of a peaceful indoor home environment with warm tones and gentle natural light. The overall atmosphere is cute, tender, serene, and photorealistic, emphasizing the cat''s charming appearance, soft texture, and affectionate presence.","camera_info":{"color":"Warm","frame_size":"Close Up","shot_type_angle":"Eye level","lens_size":"Medium Lens","composition":"Centered balanced","lighting":"Soft light","lighting_type":"Daylight"},"world_knowledge":[],"prominent_elements":[{"name":"lovely cat","description":"A cute and gentle domestic cat with soft fluffy fur, expressive eyes, and a calm relaxed presence.","location":"center of the frame","relative_size":"large","shape_and_color":"Small animal body with rounded face, triangular ears, bright eyes, and soft fur in warm natural tones","texture":"soft, fluffy, silky","appearance_details":"The cat has clean well-groomed fur, delicate whiskers, small ears, a cute nose, and bright attentive eyes. Its expression appears calm, affectionate, and slightly curious.","relationship":"Acts as the main subject and emotional focal point of the scene.","orientation":"facing the camera","pose":"sitting comfortably with a relaxed posture","expression":"gentle, adorable, calm, slightly curious","clothing":"","gender":"","skin_tone_and_texture":""},{"name":"cat eyes","description":"Bright expressive eyes that give the cat a sweet and affectionate appearance.","location":"upper center of the cat face","relative_size":"small","shape_and_color":"Round almond-like eyes with glossy reflections","texture":"clear, glossy, reflective","appearance_details":"The eyes catch the soft daylight, creating small natural highlights that make the cat look vivid and alive.","relationship":"Enhance the emotional charm and cuteness of the cat.","orientation":"looking toward the camera","pose":"","expression":"soft and attentive","clothing":"","gender":"","skin_tone_and_texture":""},{"name":"soft cushion","description":"A comfortable cushion or blanket where the cat is resting.","location":"bottom portion of the frame","relative_size":"medium","shape_and_color":"Soft rounded fabric surface in light warm neutral tones","texture":"plush, fabric, cozy","appearance_details":"The cushion gently supports the cat and adds a comfortable home-like feeling to the scene.","relationship":"Provides a cozy resting place for the cat.","orientation":"horizontal","pose":"","expression":"","clothing":"","gender":"","skin_tone_and_texture":""},{"name":"sunlit indoor background","description":"A softly blurred indoor background with warm daylight and peaceful home atmosphere.","location":"behind the cat, filling the upper and side areas of the frame","relative_size":"large","shape_and_color":"Soft abstract shapes in warm beige, cream, and pale golden tones","texture":"soft, blurry, bokeh-like","appearance_details":"The background is intentionally out of focus, keeping attention on the cat while creating a cozy and serene mood.","relationship":"Provides a warm and gentle environment that supports the cute domestic scene.","orientation":"upright","pose":"","expression":"","clothing":"","gender":"","skin_tone_and_texture":"","is_cluster":true,"number_of_objects":"numerous"}]}}' -n '{"universal_negative":{"visual_quality":["low quality","worst quality","blurry","pixelated","jpeg artifacts","low resolution","unstable color","color flicker","underexposed","overexposed","invisible subject","subject hidden in darkness"],"artistic_style":["painting","illustration","drawing","cartoon","3d render","cgi","sketch","digital art"],"composition_and_content":["text","watermark","signature","logo","subtitles","pillarboxed","side bars","portrait image in landscape frame"],"temporal_and_motion_stability":["flickering","jittery","motion blur","temporal inconsistency","warping","morphing","incoherent motion","unnatural movement","static object with sudden jump","frame-to-frame inconsistency"],"material_and_structure":["plastic-like glass","unrealistic texture","deformed bottle","liquid freezing improperly","distorted reflections"]}}' --diffusion-fa --offload-to-cpu --cfg-scale 3 --video-frames 33 -v
```

45
docs/mage_flow.md Normal file
View File

@ -0,0 +1,45 @@
# Mage-Flow
[Mage-Flow](https://github.com/microsoft/Mage) uses a 4B native-resolution multimodal diffusion transformer, Qwen3-VL for text and image conditioning, and the 128-channel Mage-VAE. Both text-to-image and instruction-based image editing checkpoints are supported.
## Download weights
- Download Mage-Flow
- safetensors: https://huggingface.co/microsoft/Mage-Flow/tree/main/transformer
- Download Mage-Flow-Base
- safetensors: https://huggingface.co/microsoft/Mage-Flow-Base/tree/main/transformer
- Download Mage-Flow-Turbo
- safetensors: https://huggingface.co/microsoft/Mage-Flow-Turbo/tree/main/transformer
- Download Mage-Flow-Edit
- safetensors: https://huggingface.co/microsoft/Mage-Flow-Edit/tree/main/transformer
- Download Mage-Flow-Edit-Turbo
- safetensors: https://huggingface.co/microsoft/Mage-Flow-Edit-Turbo/tree/main/transformer
- Download Mage-Flow-Edit-Base
- safetensors: https://huggingface.co/microsoft/Mage-Flow-Edit-Base/tree/main/transformer
- Download Mage-Flow vae
- safetensors: https://huggingface.co/microsoft/Mage-Flow/tree/main/vae
- Download Qwen3-VL 4B
- safetensors: https://huggingface.co/Comfy-Org/Krea-2/tree/main/text_encoders
- gguf: https://huggingface.co/Qwen/Qwen3-VL-4B-Instruct-GGUF/tree/main
## Text-to-image
Use 30 steps for Base models and 4 steps with `--cfg-scale 1` for Turbo models. Image dimensions must be multiples of 16; the official checkpoints are trained for native resolutions from 512 to 2048 pixels.
```bash
.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\Mage-Flow-Turbo.safetensors --llm ..\models\text_encoders\Qwen3-VL-4B-Instruct-Q4_K_M.gguf --vae ..\models\vae\mage_vae.safetensors -p "a lovely cat holding a sign says 'mage.cpp'" --cfg-scale 1.0 --steps 4 --diffusion-fa -v --offload-to-cpu
```
<img width="256" alt="Mage-Flow example" src="../assets/mage_flow/example.png" />
## Image editing
Mage-Flow-Edit accepts one or more reference images. The default `mage_flow` reference preset sends each image to both Qwen3-VL and the diffusion transformer, caps the VLM copy's longest edge at 384 pixels, and keeps the VAE copy at the requested output resolution.
For the Turbo edit checkpoint, use 4 steps and `--cfg-scale 1`.
```bash
.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\Mage-Flow-Edit.safetensors --llm ..\models\text_encoders\Qwen3-VL-4B-Instruct-Q4_K_M.gguf --llm_vision ..\models\text_encoders\Qwen3-VL-4B-Instruct-mmproj-BF16.gguf --vae ..\models\vae\mage_vae.safetensors -r ..\assets\flux\flux1-dev-q8_0.png -p "change 'flux.cpp' to 'mage.cpp'" --cfg-scale 4.0 --sampling-method euler -v --diffusion-fa --offload-to-cpu
```
<img width="256" alt="Mage-Flow-Edit example" src="../assets/mage_flow/edit_example.png" />

96
docs/minimax_h3.md Normal file
View File

@ -0,0 +1,96 @@
# MiniMax-H3
MiniMax-H3 jointly generates video and stereo audio with a packed diffusion
transformer. The implementation supports text-to-audio-video (T2VA), optional
first-frame conditioning (I2VA), first/last-frame conditioning (FL2VA), and
image/video/audio reference conditioning (Ref2VA).
## Model files
Pass the four MiniMax-H3 components separately:
- `--diffusion-model`: MiniMax-H3 diffusion transformer
- `--vae`: MiniMax-H3 video VAE
- `--audio-vae`: MiniMax-H3 audio VAE
- `--llm`: the MiniMax-H3 Qwen3-VL-32B text encoder checkpoint
The text encoder must be the MiniMax-H3 variant: Qwen3-VL-32B truncated to 50
language layers and exported without the final language-model normalization.
Its Qwen3-VL vision tower, including the three DeepStack mergers, must also be
present. If the vision tower is stored separately, pass it with `--llm_vision`.
Both the original time-embedder DiT and the smaller AdaLN curve-table variant
are detected from their weights.
### Download weights
- Download minimax_h3_fl2va/minimax_h3_ref2va
- safetensors: https://huggingface.co/Comfy-Org/MiniMax-H3/tree/main/diffusion_models
- gguf: https://huggingface.co/leejet/MiniMax-H3-GGUF/tree/main
- Download qwen3vl_32b_minimax_h3
- safetensors: https://huggingface.co/Comfy-Org/MiniMax-H3/tree/main/text_encoders
- gguf: https://huggingface.co/leejet/MiniMax-H3-GGUF/tree/main
- Download vae
- safetensors: https://huggingface.co/Comfy-Org/MiniMax-H3/tree/main/vae
- Download audio vae
- safetensors: https://huggingface.co/Comfy-Org/MiniMax-H3/tree/main/vae
## Text-to-audio-video
```sh
.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model ..\models\diffusion_models\minimax_h3_fl2va-Q4_K_M.gguf --vae ..\models\vae\minimax_h3_video_vae_fp16.safetensors --audio-vae ..\models\vae\minimax_h3_audio_vae_fp32.safetensors --llm ..\models\text_encoders\qwen3vl_32b_minimax_h3-Q4_K_M.gguf -p "A cute American Shorthair silver tabby kitten surfs on a tropical ocean wave, riding a white surfboard with the clear text 'sd.cpp' on it. Cinematic tracking shot, realistic water, bright sunlight, smooth motion, and consistent character appearance. Add upbeat tropical surf-rock background music with cheerful drums and guitar, synchronized with the kittens energetic surfing." --cfg-scale 1.0 -v -W 864 -H 480 --diffusion-fa --offload-to-cpu --rng cpu --fps 24 --video-frames 56
```
<video src=../assets/minimax-h3/t2av.mp4 controls="controls" muted="muted" type="video/mp4"></video>
Omitting `--audio-vae` still runs the joint diffusion model but produces video without a
decoded audio track.
## First/last-frame conditioning
Add `--init-img` for I2VA, or both `--init-img` and `--end-img` for FL2VA:
```sh
.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model ..\models\diffusion_models\minimax_h3_fl2va-Q4_K_M.gguf --vae ..\models\vae\minimax_h3_video_vae_fp16.safetensors --audio-vae ..\models\vae\minimax_h3_audio_vae_fp32.safetensors --llm ..\models\text_encoders\qwen3vl_32b_minimax_h3-Q4_K_M.gguf -p "a lovely cat" -i ..\assets\ernie_image\turbo_example.png --cfg-scale 1.0 -v -W 864 -H 480 --diffusion-fa --offload-to-cpu --rng cpu --fps 24 --video-frames 56
```
<video src=../assets/minimax-h3/i2av.mp4 controls="controls" muted="muted" type="video/mp4"></video>
## Reference-to-audio-video conditioning
Ref2VA accepts any combination of reference images, reference videos, paired
video soundtracks, and standalone audio references:
```sh
.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model ..\models\diffusion_models\minimax_h3_ref2va_pruned-Q4_K_M.gguf --vae ..\models\vae\minimax_h3_video_vae_fp16.safetensors --audio-vae ..\models\vae\minimax_h3_audio_vae_fp32.safetensors --llm ..\models\text_encoders\qwen3vl_32b_minimax_h3-Q4_K_M.gguf -p "Use the cat from <Picture 1> as the main character. Keep the cats appearance, fur color, facial features, and identity consistent with the reference image. Create a 2-second cinematic video: start with an extreme close-up shot of the cats face, focusing on its cute expression and detailed fur texture. The camera slowly rotates around the cats head, creating a dynamic reveal. Then smoothly pull back and zoom out to reveal the full scene: the cat is standing confidently on a surfboard, riding ocean waves. Water splashes around the board, sea breeze gently moves the cats fur, and the cat maintains a cute and fearless expression while surfing. Smooth camera movement, cinematic orbit shot, seamless zoom-out transition, low-angle wide shot, realistic ocean environment, golden sunlight, dynamic waves, high-quality realistic style, natural motion, no distortion, keep the cats identity unchanged." -r ..\assets\ernie_image\turbo_example.png --cfg-scale 1.0 -v -W 864 -H 480 --diffusion-fa --offload-to-cpu --rng cpu --fps 24 --video-frames 56
```
<video src=../assets/minimax-h3/r2av.mp4 controls="controls" muted="muted" type="video/mp4"></video>
`--ref-image`, `--ref-video`, and `--ref-audio` can each be repeated. A
reference video is a directory of image frames sorted lexicographically and is
treated as 24 fps. Repeated `--ref-video-audio` WAV files are paired by index
with repeated `--ref-video` inputs. WAV PCM (8/16/24/32-bit) and 32/64-bit
floating-point samples are accepted; audio is converted to stereo 32 kHz by the
pipeline.
Reference inputs are presented to Qwen3-VL in image, video, then audio order.
Videos are sampled at 2 fps for the Qwen presentation while their full 24 fps
latents condition the diffusion transformer. Paired video and audio references
share the same timeline. Ref2VA cannot be combined with `--init-img` or
`--end-img` in one request.
Reference images keep their aspect ratio and are only downscaled when their
pixel area exceeds the requested generation canvas.
The C API exposes the same inputs through `ref_images`, `ref_videos`, and
`ref_audios` in `sd_vid_gen_params_t`. Each `sd_ref_video_t` supplies its own
frame rate and optional soundtrack; non-24-fps inputs are resampled internally.
## Shape and runtime notes
- Width and height are aligned upward to a multiple of 32.
- Frame count is aligned upward to the `17k + 5` grid, with a minimum of 5.
- MiniMax-H3 runs at 24 fps; another requested value is overridden.
- The default video flow shift is 12. The audio stream is mapped internally to
its shift of 3, so the regular samplers can operate on the packed AV latent.

View File

@ -1,7 +1,7 @@
# How to Use
PiD is NVIDIA's Pixel Diffusion Decoder. It replaces the usual VAE decode or decode-then-upscale path with a pixel-space diffusion decoder conditioned on a
source latent and text prompt.
source latent and text prompt. Both the original PiD checkpoints and PiD 1.5 are supported.
In stable-diffusion.cpp, PiD currently runs as an image edit pipeline: provide a reference image with `-r`/`--ref-image`, encode that image with a matching VAE, then let the PiD diffusion model decode/upscale directly to RGB.
@ -16,6 +16,7 @@ In stable-diffusion.cpp, PiD currently runs as an image edit pipeline: provide a
- Flux / Z-Image PiD: use the Flux VAE and pass `--vae-format flux`
- SD3 PiD: use the SD3 VAE and pass `--vae-format sd3`
- Flux.2 PiD: use the Flux.2 VAE and pass `--vae-format flux2`
- Qwen-Image PiD: use the Qwen-Image 2D VAE and pass `--vae-format wan`
The official PiD model card should be checked before use. At the time of the initial PiD release, the official weights are under the NSCLv1 non-commercial license.

View File

@ -2,8 +2,8 @@
- download original weights(.ckpt or .safetensors). For example
- Stable Diffusion v1.4 from https://huggingface.co/CompVis/stable-diffusion-v-1-4-original
- Stable Diffusion v1.5 from https://huggingface.co/runwayml/stable-diffusion-v1-5
- Stable Diffuison v2.1 from https://huggingface.co/stabilityai/stable-diffusion-2-1
- Stable Diffusion v1.5 from https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5
- Stable Diffuison v2.1 from https://huggingface.co/Manojb/stable-diffusion-2-1-base
- Stable Diffusion 3 2B from https://huggingface.co/stabilityai/stable-diffusion-3-medium
### txt2img example

View File

@ -6,6 +6,9 @@ For detailed command-line arguments, run:
./bin/sd-cli -h
```
For direct image repair or automatic post-generation YOLOv8 detection followed by cropped inpainting, see
[ADetailer](../../docs/adetailer.md).
Metadata mode inspects PNG/JPEG container metadata without loading any model:
```bash

View File

@ -199,7 +199,7 @@ struct SDCliParams {
options.manual_options = {
{"-M",
"--mode",
"run mode, one of [img_gen, vid_gen, upscale, convert, metadata], default: img_gen",
"run mode, one of [img_gen, adetailer, vid_gen, upscale, convert, metadata], default: img_gen",
on_mode_arg},
{"",
"--preview",
@ -566,6 +566,65 @@ bool save_results(const SDCliParams& cli_params,
return sucessful_reults != 0;
}
static bool apply_adetailer(sd_ctx_t* sd_ctx,
const sd_ctx_params_t& sd_ctx_params,
const SDContextParams& ctx_params,
const SDGenerationParams& gen_params,
const sd_img_gen_params_t& img_gen_params,
SDMode mode,
SDImageVec& results,
int num_results) {
if (gen_params.ad_model_path.empty()) {
return true;
}
sd_adetailer_params_t ad_params{};
ad_params.prompt = gen_params.ad_prompt.empty() ? nullptr : gen_params.ad_prompt.c_str();
ad_params.negative_prompt = gen_params.ad_negative_prompt.empty() ? nullptr : gen_params.ad_negative_prompt.c_str();
ad_params.extra_ad_args = gen_params.extra_ad_args.c_str();
ADetailerCtxPtr ad_ctx(new_adetailer_ctx(gen_params.ad_model_path.c_str(),
ctx_params.n_threads,
sd_ctx_params.backend,
sd_ctx_params.params_backend));
if (ad_ctx == nullptr) {
LOG_ERROR("new_adetailer_ctx failed");
return false;
}
for (int i = 0; i < num_results; ++i) {
if (results[i].data == nullptr) {
continue;
}
sd_img_gen_params_t ad_generation_params = img_gen_params;
ad_generation_params.seed = img_gen_params.seed + i;
if (mode == IMG_GEN) {
ad_generation_params.width = 512;
ad_generation_params.height = 512;
ad_generation_params.strength = 0.4f;
}
sd_image_t* detailed_images = nullptr;
int detailed_count = 0;
if (!adetail_image(ad_ctx.get(),
sd_ctx,
results[i],
&ad_params,
&ad_generation_params,
&detailed_images,
&detailed_count) ||
detailed_count <= 0 || detailed_images == nullptr || detailed_images[0].data == nullptr) {
free_sd_images(detailed_images, detailed_count);
LOG_ERROR("ADetailer failed for image %d", i + 1);
return false;
}
free(results[i].data);
results[i] = detailed_images[0];
detailed_images[0] = {0, 0, 0, nullptr};
free_sd_images(detailed_images, detailed_count);
}
return true;
}
int main(int argc, const char* argv[]) {
if (argc > 1 && std::string(argv[1]) == "--version") {
std::cout << version_string() << "\n";
@ -598,6 +657,11 @@ int main(int argc, const char* argv[]) {
return 0;
}
if (!gen_params.ad_model_path.empty() && cli_params.mode != IMG_GEN && cli_params.mode != ADETAILER) {
LOG_ERROR("--ad-model is only supported in image generation and adetailer modes");
return 1;
}
if (gen_params.video_frames > 4) {
size_t last_dot_pos = cli_params.preview_path.find_last_of(".");
std::string base_path = cli_params.preview_path;
@ -690,6 +754,18 @@ int main(int argc, const char* argv[]) {
return true;
};
auto load_audio = [&](const std::string& path, SDAudioOwner& audio) -> bool {
std::vector<float> samples;
uint32_t sample_rate = 0;
uint32_t channels = 0;
if (!load_wav_from_file(path, samples, sample_rate, channels)) {
LOG_ERROR("load WAV audio from '%s' failed", path.c_str());
return false;
}
audio.reset(std::move(samples), sample_rate, channels);
return true;
};
if (gen_params.init_image_path.size() > 0) {
if (!load_image_and_update_size(gen_params.init_image_path, gen_params.init_image)) {
return 1;
@ -713,6 +789,37 @@ int main(int argc, const char* argv[]) {
}
}
if (!gen_params.ref_video_paths.empty()) {
gen_params.ref_videos.clear();
gen_params.ref_videos.reserve(gen_params.ref_video_paths.size());
for (const auto& path : gen_params.ref_video_paths) {
std::vector<SDImageOwner> frames;
if (!load_images_from_dir(path, frames, 0, 0, 0, cli_params.verbose) || frames.empty()) {
LOG_ERROR("load reference video frames from '%s' failed", path.c_str());
return 1;
}
gen_params.ref_videos.push_back(std::move(frames));
}
gen_params.ref_video_audios.clear();
gen_params.ref_video_audios.resize(gen_params.ref_videos.size());
for (size_t i = 0; i < gen_params.ref_video_audio_paths.size(); ++i) {
if (!load_audio(gen_params.ref_video_audio_paths[i], gen_params.ref_video_audios[i])) {
return 1;
}
}
}
if (!gen_params.ref_audio_paths.empty()) {
gen_params.ref_audios.clear();
gen_params.ref_audios.resize(gen_params.ref_audio_paths.size());
for (size_t i = 0; i < gen_params.ref_audio_paths.size(); ++i) {
if (!load_audio(gen_params.ref_audio_paths[i], gen_params.ref_audios[i])) {
return 1;
}
}
}
if (gen_params.mask_image_path.size() > 0) {
if (!load_sd_image_from_file(gen_params.mask_image.put(),
gen_params.mask_image_path.c_str(),
@ -753,6 +860,16 @@ int main(int argc, const char* argv[]) {
}
}
if (gen_params.ip_adapter_image_path.size() > 0) {
if (!load_sd_image_from_file(gen_params.ip_adapter_image.put(),
gen_params.ip_adapter_image_path.c_str(),
0,
0)) {
LOG_ERROR("load image from '%s' failed", gen_params.ip_adapter_image_path.c_str());
return 1;
}
}
if (!gen_params.control_video_path.empty()) {
gen_params.control_frames.clear();
if (!load_images_from_dir(gen_params.control_video_path,
@ -806,15 +923,22 @@ int main(int argc, const char* argv[]) {
gen_params.sample_params.scheduler = sd_get_default_scheduler(sd_ctx.get(), gen_params.sample_params.sample_method);
}
if (cli_params.mode == IMG_GEN) {
sd_img_gen_params_t img_gen_params = gen_params.to_sd_img_gen_params_t();
sd_img_gen_params_t img_gen_params{};
const bool use_img_gen_params = cli_params.mode == IMG_GEN || cli_params.mode == ADETAILER;
if (use_img_gen_params) {
img_gen_params = gen_params.to_sd_img_gen_params_t();
}
if (cli_params.mode == IMG_GEN) {
sd_image_t* generated_images = nullptr;
if (!generate_image(sd_ctx.get(), &img_gen_params, &generated_images, &num_results)) {
generated_images = nullptr;
num_results = 0;
}
results.adopt(generated_images, num_results);
} else if (cli_params.mode == ADETAILER) {
num_results = 1;
results.push_back(gen_params.init_image.release());
} 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 = nullptr;
@ -828,6 +952,18 @@ int main(int argc, const char* argv[]) {
LOG_ERROR("generate failed");
return 1;
}
if (use_img_gen_params &&
!apply_adetailer(sd_ctx.get(),
sd_ctx_params,
ctx_params,
gen_params,
img_gen_params,
cli_params.mode,
results,
num_results)) {
return 1;
}
}
int upscale_factor = 4; // unused for RealESRGAN_x4plus_anime_6B.pth

View File

@ -30,6 +30,7 @@ namespace fs = std::filesystem;
const char* const modes_str[] = {
"img_gen",
"adetailer",
"vid_gen",
"convert",
"upscale",
@ -49,6 +50,9 @@ static sd_vae_format_t str_to_vae_format(const std::string& value) {
if (value == "flux2") {
return SD_VAE_FORMAT_FLUX2;
}
if (value == "wan") {
return SD_VAE_FORMAT_WAN;
}
return SD_VAE_FORMAT_COUNT;
}
@ -400,7 +404,7 @@ ArgOptions SDContextParams::get_options() {
&vae_path},
{"",
"--vae-format",
"VAE latent format override: auto, flux, sd3, or flux2 (default: auto)",
"VAE latent format override: auto, flux, sd3, flux2, or wan (default: auto)",
0,
&vae_format},
{"",
@ -423,6 +427,16 @@ ArgOptions SDContextParams::get_options() {
"path to control net model",
0,
&control_net_path},
{"",
"--ip-adapter",
"path to IP-Adapter model (requires --clip_vision)",
0,
&ip_adapter_path},
{"",
"--motion-module",
"path to AnimateDiff motion module (SD 1.5); enables video generation on --video-frames > 1",
0,
&motion_module_path},
{"",
"--embd-dir",
"embeddings directory",
@ -672,7 +686,7 @@ ArgOptions SDContextParams::get_options() {
}
void SDContextParams::build_embedding_map() {
static const std::vector<std::string> valid_ext = {".gguf", ".safetensors", ".pt"};
static const std::vector<std::string> valid_ext = {".gguf", ".safetensors", ".pt", ".ckpt"};
if (!fs::exists(embedding_dir) || !fs::is_directory(embedding_dir)) {
return;
@ -737,7 +751,7 @@ bool SDContextParams::validate(SDMode mode) {
}
if (str_to_vae_format(vae_format) == SD_VAE_FORMAT_COUNT) {
LOG_ERROR("error: vae_format must be 'auto', 'flux', 'sd3', or 'flux2'");
LOG_ERROR("error: vae_format must be 'auto', 'flux', 'sd3', 'flux2', or 'wan'");
return false;
}
@ -867,6 +881,8 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) {
sd_ctx_params.audio_vae_path = audio_vae_path.c_str();
sd_ctx_params.taesd_path = taesd_path.c_str();
sd_ctx_params.control_net_path = control_net_path.c_str();
sd_ctx_params.ip_adapter_path = ip_adapter_path.c_str();
sd_ctx_params.motion_module_path = motion_module_path.c_str();
sd_ctx_params.embeddings = embedding_vec.data();
sd_ctx_params.embedding_count = static_cast<uint32_t>(embedding_vec.size());
sd_ctx_params.photo_maker_path = photo_maker_path.c_str();
@ -916,6 +932,26 @@ ArgOptions SDGenerationParams::get_options() {
"the negative prompt (default: \"\")",
0,
&negative_prompt},
{"",
"--ad-model",
"path to a converted YOLOv8 detection model for ADetailer",
0,
&ad_model_path},
{"",
"--ad-prompt",
"ADetailer prompt; empty inherits the main prompt, supports [PROMPT], [SEP], and [SKIP]",
0,
&ad_prompt},
{"",
"--ad-negative-prompt",
"ADetailer negative prompt; empty inherits the main negative prompt, supports [PROMPT] and [SEP]",
0,
&ad_negative_prompt},
{"",
"--extra-ad-args",
"extra ADetailer args, key=value list. Supports input_size, confidence, nms, max_detections, mask_k_largest, mask_min_ratio, mask_max_ratio, dilate_erode, x_offset, y_offset, mask_mode, merge_masks, invert_mask, mask_blur, inpaint_padding, inpaint_width, inpaint_height, denoising_strength, steps, cfg_scale, sample_method, scheduler, sort_by",
(int)',',
&extra_ad_args},
{"-i",
"--init-img",
"path to the init image",
@ -936,6 +972,11 @@ ArgOptions SDGenerationParams::get_options() {
"path to control image, control net",
0,
&control_image_path},
{"",
"--ip-adapter-image",
"path to the IP-Adapter reference image",
0,
&ip_adapter_image_path},
{"",
"--control-video",
"path to control video frames, It must be a directory path. The video frames inside should be stored as images in "
@ -967,7 +1008,7 @@ ArgOptions SDGenerationParams::get_options() {
&hires_upscaler},
{"",
"--extra-sample-args",
"extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma;; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware",
"extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma; beta scheduler supports alpha, beta; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware; lms supports lms_divisions",
(int)',',
&extra_sample_args},
{"",
@ -975,6 +1016,11 @@ ArgOptions SDGenerationParams::get_options() {
"extra VAE tiling args, key=value list. LTX video VAE supports temporal_tile_frames (default: 4), temporal_tile_overlap (default: 1)",
(int)',',
&extra_tiling_args},
{"",
"--ref-image-args",
"Key-value list to set up the way the reference images are processed (empty = auto-detect from model weigths)",
(int)',',
&ref_image_args},
};
options.int_options = {
@ -1123,6 +1169,10 @@ ArgOptions SDGenerationParams::get_options() {
"--control-strength",
"strength to apply Control Net (default: 0.9). 1.0 corresponds to full destruction of information in init image",
&control_strength},
{"",
"--ip-adapter-strength",
"strength to apply IP-Adapter (default: 1.0)",
&ip_adapter_strength},
{"",
"--moe-boundary",
"timestep boundary for Wan2.2 MoE model. (default: 0.875). Only enabled if `--high-noise-steps` is set to -1",
@ -1354,6 +1404,30 @@ ArgOptions SDGenerationParams::get_options() {
return 1;
};
auto on_ref_video_arg = [&](int argc, const char** argv, int index) {
if (++index >= argc) {
return -1;
}
ref_video_paths.push_back(argv[index]);
return 1;
};
auto on_ref_video_audio_arg = [&](int argc, const char** argv, int index) {
if (++index >= argc) {
return -1;
}
ref_video_audio_paths.push_back(argv[index]);
return 1;
};
auto on_ref_audio_arg = [&](int argc, const char** argv, int index) {
if (++index >= argc) {
return -1;
}
ref_audio_paths.push_back(argv[index]);
return 1;
};
auto on_cache_mode_arg = [&](int argc, const char** argv, int index) {
if (++index >= argc) {
return -1;
@ -1488,12 +1562,12 @@ 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, dpm++2m_sde, dpm++2m_sde_bt, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp]"
"sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, dpm++2m_sde, dpm++2m_sde_bt, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp, lms]"
"(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, dpm++2m_sde, dpm++2m_sde_bt, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp]"
"(high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m, dpm++2mv2, dpm++2m_sde, dpm++2m_sde_bt, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp, lms]"
" default: euler for Flux/SD3/Wan, euler_a otherwise",
on_high_noise_sample_method_arg},
{"",
@ -1518,8 +1592,20 @@ ArgOptions SDGenerationParams::get_options() {
on_high_noise_skip_layers_arg},
{"-r",
"--ref-image",
"reference image for Flux Kontext models (can be used multiple times)",
"reference image for Flux Kontext or MiniMax-H3 Ref2VA (can be used multiple times)",
on_ref_image_arg},
{"",
"--ref-video",
"MiniMax-H3 Ref2VA reference video frame directory at 24 fps (can be used multiple times)",
on_ref_video_arg},
{"",
"--ref-video-audio",
"WAV soundtrack paired by index with --ref-video (can be used multiple times)",
on_ref_video_audio_arg},
{"",
"--ref-audio",
"standalone WAV reference for MiniMax-H3 Ref2VA (can be used multiple times)",
on_ref_audio_arg},
{"",
"--cache-mode",
"caching method: 'easycache' (DiT), 'ucache' (UNET), 'dbcache'/'taylorseer'/'cache-dit' (DiT block-level), 'spectrum' (UNET/DiT Chebyshev+Taylor forecasting)",
@ -1831,6 +1917,10 @@ bool SDGenerationParams::from_json_str(
load_if_exists("prompt", prompt);
load_if_exists("negative_prompt", negative_prompt);
load_if_exists("ad_model", ad_model_path);
load_if_exists("ad_prompt", ad_prompt);
load_if_exists("ad_negative_prompt", ad_negative_prompt);
load_if_exists("extra_ad_args", extra_ad_args);
load_if_exists("cache_mode", cache_mode);
load_if_exists("cache_option", cache_option);
load_if_exists("scm_mask", scm_mask);
@ -1847,6 +1937,7 @@ bool SDGenerationParams::from_json_str(
load_if_exists("strength", strength);
load_if_exists("control_strength", control_strength);
load_if_exists("ip_adapter_strength", ip_adapter_strength);
load_if_exists("moe_boundary", moe_boundary);
load_if_exists("vace_strength", vace_strength);
@ -2018,6 +2109,10 @@ bool SDGenerationParams::from_json_str(
LOG_ERROR("invalid control_image");
return false;
}
if (!parse_image_json_field(j, "ip_adapter_image", 3, width, height, ip_adapter_image)) {
LOG_ERROR("invalid ip_adapter_image");
return false;
}
return true;
}
@ -2027,7 +2122,7 @@ void SDGenerationParams::extract_and_remove_lora(const std::string& lora_model_d
return;
}
static const std::regex re(R"(<lora:([^:>]+):([^>]+)>)");
static const std::vector<std::string> valid_ext = {".gguf", ".safetensors", ".pt"};
static const std::vector<std::string> valid_ext = {".gguf", ".safetensors", ".pt", ".ckpt"};
std::smatch m;
std::string tmp = prompt;
@ -2307,6 +2402,16 @@ bool SDGenerationParams::validate(SDMode mode) {
return false;
}
if (ref_video_audio_paths.size() > ref_video_paths.size()) {
LOG_ERROR("error: each --ref-video-audio needs a corresponding --ref-video");
return false;
}
if (mode != VID_GEN && (!ref_video_paths.empty() || !ref_video_audio_paths.empty() || !ref_audio_paths.empty())) {
LOG_ERROR("error: reference video and audio inputs require vid_gen mode");
return false;
}
if (sample_params.shifted_timestep < 0 || sample_params.shifted_timestep > 1000) {
LOG_ERROR("error: shifted_timestep must be in range [0, 1000]");
return false;
@ -2347,13 +2452,19 @@ bool SDGenerationParams::validate(SDMode mode) {
}
}
if (mode == UPSCALE) {
if (mode == UPSCALE || mode == ADETAILER) {
if (init_image_path.length() == 0) {
LOG_ERROR("error: upscale mode needs an init image (--init-img)\n");
LOG_ERROR("error: %s mode needs an init image (--init-img)\n",
mode == UPSCALE ? "upscale" : "adetailer");
return false;
}
}
if (mode == ADETAILER && ad_model_path.empty()) {
LOG_ERROR("error: adetailer mode needs a detector model (--ad-model)\n");
return false;
}
return true;
}
@ -2418,6 +2529,22 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
pulid_id_weight,
};
if (!auto_resize_ref_image) {
if (!ref_image_args.empty()) {
ref_image_args += ",";
}
ref_image_args += "resize_before_vae=0";
LOG_WARN("Notice: --disable-auto-resize-ref-image is deprecated. Use --ref-image-args \"resize_before_vae=off\" instead.");
}
if (increase_ref_index) {
if (!ref_image_args.empty()) {
ref_image_args += ",";
}
ref_image_args += "ref_index_mode=increase";
LOG_WARN("Notice: --increase-ref-index is deprecated. Use --ref-image-args \"ref_index_mode=increase\" instead.");
}
params.loras = lora_vec.empty() ? nullptr : lora_vec.data();
params.lora_count = static_cast<uint32_t>(lora_vec.size());
params.prompt = prompt.c_str();
@ -2426,8 +2553,7 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
params.init_image = init_image.get();
params.ref_images = ref_image_views.empty() ? nullptr : ref_image_views.data();
params.ref_images_count = static_cast<int>(ref_image_views.size());
params.auto_resize_ref_image = auto_resize_ref_image;
params.increase_ref_index = increase_ref_index;
params.ref_image_args = ref_image_args.c_str();
params.mask_image = mask_image.get();
params.width = get_resolved_width();
params.height = get_resolved_height();
@ -2438,6 +2564,8 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
params.qwen_image_layers = qwen_image_layers;
params.control_image = control_image.get();
params.control_strength = control_strength;
params.ip_adapter_image = ip_adapter_image.get();
params.ip_adapter_strength = ip_adapter_strength;
params.pm_params = pm_params;
params.pulid_params = pulid_params;
params.vae_tiling_params = vae_tiling_params;
@ -2478,6 +2606,35 @@ sd_vid_gen_params_t SDGenerationParams::to_sd_vid_gen_params_t() {
control_frame_views.push_back(frame.get());
}
ref_image_views.clear();
ref_image_views.reserve(ref_images.size());
for (auto& image : ref_images) {
ref_image_views.push_back(image.get());
}
ref_video_frame_views.clear();
ref_video_frame_views.resize(ref_videos.size());
ref_video_views.clear();
ref_video_views.reserve(ref_videos.size());
for (size_t i = 0; i < ref_videos.size(); ++i) {
auto& frame_views = ref_video_frame_views[i];
frame_views.reserve(ref_videos[i].size());
for (auto& frame : ref_videos[i]) {
frame_views.push_back(frame.get());
}
sd_audio_t audio = i < ref_video_audios.size() ? ref_video_audios[i].get() : sd_audio_t{};
ref_video_views.push_back({frame_views.empty() ? nullptr : frame_views.data(),
static_cast<int>(frame_views.size()),
24,
audio});
}
ref_audio_views.clear();
ref_audio_views.reserve(ref_audios.size());
for (auto& audio : ref_audios) {
ref_audio_views.push_back(audio.get());
}
sample_params.guidance.slg.layers = skip_layers.empty() ? nullptr : skip_layers.data();
sample_params.guidance.slg.layer_count = skip_layers.size();
high_noise_sample_params.guidance.slg.layers = high_noise_skip_layers.empty() ? nullptr : high_noise_skip_layers.data();
@ -2496,6 +2653,12 @@ sd_vid_gen_params_t SDGenerationParams::to_sd_vid_gen_params_t() {
params.clip_skip = clip_skip;
params.init_image = init_image.get();
params.end_image = end_image.get();
params.ref_images = ref_image_views.empty() ? nullptr : ref_image_views.data();
params.ref_images_count = static_cast<int>(ref_image_views.size());
params.ref_videos = ref_video_views.empty() ? nullptr : ref_video_views.data();
params.ref_videos_count = static_cast<int>(ref_video_views.size());
params.ref_audios = ref_audio_views.empty() ? nullptr : ref_audio_views.data();
params.ref_audios_count = static_cast<int>(ref_audio_views.size());
params.control_frames = control_frame_views.empty() ? nullptr : control_frame_views.data();
params.control_frames_size = static_cast<int>(control_frame_views.size());
params.width = get_resolved_width();
@ -2561,6 +2724,10 @@ std::string SDGenerationParams::to_string() const {
<< " high_noise_loras: \"" << high_noise_loras_str << "\",\n"
<< " prompt: \"" << prompt << "\",\n"
<< " negative_prompt: \"" << negative_prompt << "\",\n"
<< " ad_model_path: \"" << ad_model_path << "\",\n"
<< " ad_prompt: \"" << ad_prompt << "\",\n"
<< " ad_negative_prompt: \"" << ad_negative_prompt << "\",\n"
<< " extra_ad_args: \"" << extra_ad_args << "\",\n"
<< " clip_skip: " << clip_skip << ",\n"
<< " width: " << width << ",\n"
<< " height: " << height << ",\n"
@ -2571,6 +2738,9 @@ std::string SDGenerationParams::to_string() const {
<< " mask_image_path: \"" << mask_image_path << "\",\n"
<< " control_image_path: \"" << control_image_path << "\",\n"
<< " ref_image_paths: " << vec_str_to_string(ref_image_paths) << ",\n"
<< " ref_video_paths: " << vec_str_to_string(ref_video_paths) << ",\n"
<< " ref_video_audio_paths: " << vec_str_to_string(ref_video_audio_paths) << ",\n"
<< " ref_audio_paths: " << vec_str_to_string(ref_audio_paths) << ",\n"
<< " control_video_path: \"" << control_video_path << "\",\n"
<< " auto_resize_ref_image: " << (auto_resize_ref_image ? "true" : "false") << ",\n"
<< " increase_ref_index: " << (increase_ref_index ? "true" : "false") << ",\n"
@ -2676,7 +2846,12 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params,
SDMode mode) {
json root;
root["schema"] = "sdcpp.image.params/v1";
root["mode"] = mode == VID_GEN ? "vid_gen" : "img_gen";
root["mode"] = "img_gen";
if (mode == VID_GEN) {
root["mode"] = "vid_gen";
} else if (mode == ADETAILER) {
root["mode"] = "adetailer";
}
root["generator"] = {
{"name", "stable-diffusion.cpp"},
{"version", safe_json_string(sd_version())},
@ -2690,6 +2865,14 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params,
{"positive", gen_params.prompt},
{"negative", gen_params.negative_prompt},
};
if (!gen_params.ad_model_path.empty()) {
root["adetailer"] = {
{"model", sd_basename(gen_params.ad_model_path)},
{"prompt", gen_params.ad_prompt},
{"negative_prompt", gen_params.ad_negative_prompt},
{"extra_args", gen_params.extra_ad_args},
};
}
root["sampling"] = build_sampling_metadata_json(gen_params.sample_params,
gen_params.skip_layers,
&gen_params.custom_sigmas);
@ -2713,6 +2896,7 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params,
root["clip_skip"] = gen_params.clip_skip;
root["strength"] = gen_params.strength;
root["control_strength"] = gen_params.control_strength;
root["ip_adapter_strength"] = gen_params.ip_adapter_strength;
root["auto_resize_ref_image"] = gen_params.auto_resize_ref_image;
root["increase_ref_index"] = gen_params.increase_ref_index;
if (mode == VID_GEN) {
@ -2844,6 +3028,18 @@ std::string get_image_params(const SDContextParams& ctx_params,
if (!gen_params.extra_sample_args.empty()) {
parameter_string += "Extra sample args: " + gen_params.extra_sample_args + ", ";
}
if (!gen_params.ad_model_path.empty()) {
parameter_string += "ADetailer model: " + sd_basename(gen_params.ad_model_path) + ", ";
if (!gen_params.ad_prompt.empty()) {
parameter_string += "ADetailer prompt: " + gen_params.ad_prompt + ", ";
}
if (!gen_params.ad_negative_prompt.empty()) {
parameter_string += "ADetailer negative prompt: " + gen_params.ad_negative_prompt + ", ";
}
if (!gen_params.extra_ad_args.empty()) {
parameter_string += "ADetailer args: " + gen_params.extra_ad_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) + ", ";

View File

@ -16,10 +16,11 @@
#define BOOL_STR(b) ((b) ? "true" : "false")
extern const char* const modes_str[];
#define SD_ALL_MODES_STR "img_gen, vid_gen, convert, upscale, metadata"
#define SD_ALL_MODES_STR "img_gen, adetailer, vid_gen, convert, upscale, metadata"
enum SDMode {
IMG_GEN,
ADETAILER,
VID_GEN,
CONVERT,
UPSCALE,
@ -132,6 +133,8 @@ struct SDContextParams {
std::string taesd_path;
std::string esrgan_path;
std::string control_net_path;
std::string ip_adapter_path;
std::string motion_module_path;
std::string embedding_dir;
std::string photo_maker_path;
std::string pulid_weights_path;
@ -186,6 +189,10 @@ struct SDGenerationParams {
// User-facing input fields.
std::string prompt;
std::string negative_prompt;
std::string ad_model_path;
std::string ad_prompt;
std::string ad_negative_prompt;
std::string extra_ad_args;
int clip_skip = -1; // <= 0 represents unspecified
int width = -1;
int height = -1;
@ -194,6 +201,7 @@ struct SDGenerationParams {
int64_t seed = 42;
float strength = 0.75f;
float control_strength = 0.9f;
float ip_adapter_strength = 1.0f;
bool auto_resize_ref_image = true;
bool increase_ref_index = false;
bool embed_image_metadata = true;
@ -202,7 +210,11 @@ struct SDGenerationParams {
std::string end_image_path;
std::string mask_image_path;
std::string control_image_path;
std::string ip_adapter_image_path;
std::vector<std::string> ref_image_paths;
std::vector<std::string> ref_video_paths;
std::vector<std::string> ref_video_audio_paths;
std::vector<std::string> ref_audio_paths;
std::string control_video_path;
sd_sample_params_t sample_params;
@ -227,6 +239,8 @@ struct SDGenerationParams {
sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr};
std::string extra_tiling_args;
std::string ref_image_args;
std::string pm_id_images_dir;
std::string pm_id_embed_path;
float pm_style_strength = 20.f;
@ -264,13 +278,20 @@ struct SDGenerationParams {
SDImageOwner init_image;
SDImageOwner end_image;
std::vector<SDImageOwner> ref_images;
std::vector<std::vector<SDImageOwner>> ref_videos;
std::vector<SDAudioOwner> ref_video_audios;
std::vector<SDAudioOwner> ref_audios;
SDImageOwner mask_image;
SDImageOwner control_image;
SDImageOwner ip_adapter_image;
std::vector<SDImageOwner> pm_id_images;
std::vector<SDImageOwner> control_frames;
// Backing storage for sd_img_gen_params_t view fields.
std::vector<sd_image_t> ref_image_views;
std::vector<std::vector<sd_image_t>> ref_video_frame_views;
std::vector<sd_ref_video_t> ref_video_views;
std::vector<sd_audio_t> ref_audio_views;
std::vector<sd_image_t> pm_id_image_views;
std::vector<sd_image_t> control_frame_views;

View File

@ -1374,3 +1374,132 @@ bool write_wav_to_file(const std::string& path,
file.write(reinterpret_cast<const char*>(pcm.data()), static_cast<std::streamsize>(pcm.size() * sizeof(int16_t)));
return file.good();
}
static uint16_t read_le16(const uint8_t* data) {
return static_cast<uint16_t>(data[0]) |
(static_cast<uint16_t>(data[1]) << 8);
}
static uint32_t read_le32(const uint8_t* data) {
return static_cast<uint32_t>(data[0]) |
(static_cast<uint32_t>(data[1]) << 8) |
(static_cast<uint32_t>(data[2]) << 16) |
(static_cast<uint32_t>(data[3]) << 24);
}
bool load_wav_from_file(const std::string& path,
std::vector<float>& interleaved_samples,
uint32_t& sample_rate,
uint32_t& channels) {
interleaved_samples.clear();
sample_rate = 0;
channels = 0;
std::ifstream file(path, std::ios::binary);
uint8_t riff_header[12];
if (!file.read(reinterpret_cast<char*>(riff_header), sizeof(riff_header)) ||
std::memcmp(riff_header, "RIFF", 4) != 0 ||
std::memcmp(riff_header + 8, "WAVE", 4) != 0) {
return false;
}
uint16_t audio_format = 0;
uint16_t bits_per_sample = 0;
uint16_t block_align = 0;
std::streampos data_pos = std::streampos(-1);
uint32_t data_size = 0;
while (file.good()) {
uint8_t chunk_header[8];
if (!file.read(reinterpret_cast<char*>(chunk_header), sizeof(chunk_header))) {
break;
}
uint32_t chunk_size = read_le32(chunk_header + 4);
std::streampos chunk_data_pos = file.tellg();
if (std::memcmp(chunk_header, "fmt ", 4) == 0) {
if (chunk_size < 16) {
return false;
}
std::vector<uint8_t> fmt(chunk_size);
if (!file.read(reinterpret_cast<char*>(fmt.data()), chunk_size)) {
return false;
}
audio_format = read_le16(fmt.data());
channels = read_le16(fmt.data() + 2);
sample_rate = read_le32(fmt.data() + 4);
block_align = read_le16(fmt.data() + 12);
bits_per_sample = read_le16(fmt.data() + 14);
if (audio_format == 0xfffe && chunk_size >= 40) {
audio_format = read_le16(fmt.data() + 24);
}
} else if (std::memcmp(chunk_header, "data", 4) == 0) {
data_pos = chunk_data_pos;
data_size = chunk_size;
file.seekg(chunk_size, std::ios::cur);
} else {
file.seekg(chunk_size, std::ios::cur);
}
if (!file.good()) {
break;
}
if ((chunk_size & 1) != 0) {
file.seekg(1, std::ios::cur);
}
}
const uint32_t bytes_per_sample = (bits_per_sample + 7) / 8;
if (data_pos == std::streampos(-1) || data_size == 0 || channels == 0 || sample_rate == 0 ||
block_align == 0 || bytes_per_sample == 0 || block_align < channels * bytes_per_sample ||
(audio_format != 1 && audio_format != 3)) {
return false;
}
const uint64_t frame_count = data_size / block_align;
if (frame_count == 0 || frame_count > SIZE_MAX / channels) {
return false;
}
std::vector<uint8_t> pcm(data_size);
file.clear();
file.seekg(data_pos);
if (!file.read(reinterpret_cast<char*>(pcm.data()), data_size)) {
return false;
}
interleaved_samples.resize(static_cast<size_t>(frame_count * channels));
for (uint64_t frame = 0; frame < frame_count; ++frame) {
const uint8_t* frame_data = pcm.data() + frame * block_align;
for (uint32_t channel = 0; channel < channels; ++channel) {
const uint8_t* sample_data = frame_data + channel * bytes_per_sample;
float sample = 0.0f;
if (audio_format == 3 && bits_per_sample == 32) {
std::memcpy(&sample, sample_data, sizeof(sample));
} else if (audio_format == 3 && bits_per_sample == 64) {
double value;
std::memcpy(&value, sample_data, sizeof(value));
sample = static_cast<float>(value);
} else if (audio_format == 1 && bits_per_sample == 8) {
sample = (static_cast<int>(sample_data[0]) - 128) / 128.0f;
} else if (audio_format == 1 && bits_per_sample == 16) {
sample = static_cast<int16_t>(read_le16(sample_data)) / 32768.0f;
} else if (audio_format == 1 && bits_per_sample == 24) {
int32_t value = static_cast<int32_t>(sample_data[0]) |
(static_cast<int32_t>(sample_data[1]) << 8) |
(static_cast<int32_t>(sample_data[2]) << 16);
if ((value & 0x800000) != 0) {
value |= ~0xffffff;
}
sample = value / 8388608.0f;
} else if (audio_format == 1 && bits_per_sample == 32) {
int32_t value = static_cast<int32_t>(read_le32(sample_data));
sample = value / 2147483648.0f;
} else {
interleaved_samples.clear();
return false;
}
interleaved_samples[static_cast<size_t>(frame * channels + channel)] = sample;
}
}
return true;
}

View File

@ -110,4 +110,9 @@ bool write_wav_to_file(const std::string& path,
uint32_t channels,
uint32_t sample_rate);
bool load_wav_from_file(const std::string& path,
std::vector<float>& interleaved_samples,
uint32_t& sample_rate,
uint32_t& channels);
#endif // __MEDIA_IO_H__

View File

@ -40,12 +40,21 @@ struct UpscalerCtxDeleter {
}
};
struct ADetailerCtxDeleter {
void operator()(adetailer_ctx_t* ctx) const {
if (ctx != nullptr) {
free_adetailer_ctx(ctx);
}
}
};
template <typename T>
using FreeUniquePtr = std::unique_ptr<T, FreeDeleter>;
using FilePtr = std::unique_ptr<FILE, FileCloser>;
using SDCtxPtr = std::unique_ptr<sd_ctx_t, SDCtxDeleter>;
using UpscalerCtxPtr = std::unique_ptr<upscaler_ctx_t, UpscalerCtxDeleter>;
using ADetailerCtxPtr = std::unique_ptr<adetailer_ctx_t, ADetailerCtxDeleter>;
class SDImageOwner {
private:
@ -132,6 +141,37 @@ public:
}
};
class SDAudioOwner {
private:
uint32_t sample_rate_ = 0;
uint32_t channels_ = 0;
std::vector<float> samples_;
public:
SDAudioOwner() = default;
void reset(std::vector<float> samples = {}, uint32_t sample_rate = 0, uint32_t channels = 0) {
samples_ = std::move(samples);
sample_rate_ = sample_rate;
channels_ = channels;
}
bool empty() const {
return samples_.empty();
}
sd_audio_t get() {
return {sample_rate_,
channels_,
channels_ == 0 ? 0 : static_cast<uint64_t>(samples_.size() / channels_),
samples_.empty() ? nullptr : samples_.data()};
}
const std::vector<float>& samples() const {
return samples_;
}
};
class SDImageVec {
private:
std::vector<sd_image_t> images_;

View File

@ -528,6 +528,7 @@ Shared default fields used by both `img_gen` and `vid_gen`:
| `auto_resize_ref_image` | `boolean` |
| `increase_ref_index` | `boolean` |
| `control_strength` | `number` |
| `ip_adapter_strength` | `number` |
| `hires` | `object` |
| `hires.enabled` | `boolean` |
| `hires.upscaler` | `string` |
@ -567,6 +568,7 @@ Fields returned in `features_by_mode.img_gen`:
- `init_image`
- `mask_image`
- `control_image`
- `ip_adapter_image`
- `ref_images`
- `lora`
- `vae_tiling`
@ -653,12 +655,14 @@ Example:
"auto_resize_ref_image": true,
"increase_ref_index": false,
"control_strength": 0.9,
"ip_adapter_strength": 1.0,
"embed_image_metadata": true,
"init_image": null,
"ref_images": [],
"mask_image": null,
"control_image": null,
"ip_adapter_image": null,
"sample_params": {
"scheduler": "discrete",
@ -733,6 +737,7 @@ Channel expectations:
- `init_image`: 3 channels
- `ref_images[]`: 3 channels
- `control_image`: 3 channels
- `ip_adapter_image`: 3 channels
- `mask_image`: 1 channel
If omitted or null:
@ -757,6 +762,7 @@ Top-level scalar fields:
| `auto_resize_ref_image` | `boolean` |
| `increase_ref_index` | `boolean` |
| `control_strength` | `number` |
| `ip_adapter_strength` | `number` |
| `embed_image_metadata` | `boolean` |
Image fields:
@ -767,6 +773,7 @@ Image fields:
| `ref_images` | `array<string>` |
| `mask_image` | `string \| null` |
| `control_image` | `string \| null` |
| `ip_adapter_image` | `string \| null` |
LoRA fields:
@ -958,7 +965,7 @@ Response fields:
Compared with `img_gen`, the `vid_gen` request body:
- `vid_gen` is a single video sequence job, so `batch_count` is not part of the request schema
- `ref_images`, `mask_image`, `control_image`, `control_strength`, and `embed_image_metadata` are not part of the request schema
- `ref_images`, `mask_image`, `control_image`, `control_strength`, `ip_adapter_image`, `ip_adapter_strength`, and `embed_image_metadata` are not part of the request schema
- `vid_gen` adds `end_image`, `control_frames`, `high_noise_sample_params`, `video_frames`, `fps`, `moe_boundary`, and `vace_strength`
Example:

View File

@ -130,6 +130,7 @@ static json make_img_gen_defaults_json(const SDGenerationParams& defaults, const
{"auto_resize_ref_image", defaults.auto_resize_ref_image},
{"increase_ref_index", defaults.increase_ref_index},
{"control_strength", defaults.control_strength},
{"ip_adapter_strength", defaults.ip_adapter_strength},
{"sample_params", make_sample_params_json(defaults.sample_params, defaults.skip_layers)},
{"hires", make_hires_json(defaults)},
{"vae_tiling_params", make_vae_tiling_json(defaults.vae_tiling_params)},
@ -173,6 +174,7 @@ static json make_img_gen_features_json() {
{"init_image", true},
{"mask_image", true},
{"control_image", true},
{"ip_adapter_image", true},
{"ref_images", true},
{"lora", true},
{"vae_tiling", true},

View File

@ -56,6 +56,7 @@ enum sample_method_t {
EULER_GE_SAMPLE_METHOD,
DPMPP2M_SDE_SAMPLE_METHOD,
DPMPP2M_SDE_BT_SAMPLE_METHOD,
LMS_SAMPLE_METHOD,
SAMPLE_METHOD_COUNT
};
@ -180,6 +181,7 @@ enum sd_vae_format_t {
SD_VAE_FORMAT_FLUX,
SD_VAE_FORMAT_SD3,
SD_VAE_FORMAT_FLUX2,
SD_VAE_FORMAT_WAN,
SD_VAE_FORMAT_COUNT,
};
@ -199,6 +201,8 @@ typedef struct {
const char* audio_vae_path;
const char* taesd_path;
const char* control_net_path;
const char* ip_adapter_path;
const char* motion_module_path;
const sd_embedding_t* embeddings;
uint32_t embedding_count;
const char* photo_maker_path;
@ -243,6 +247,13 @@ typedef struct {
uint8_t* data;
} sd_image_t;
typedef struct {
sd_image_t* frames;
int frame_count;
int fps;
sd_audio_t audio;
} sd_ref_video_t;
typedef struct {
int* layers;
size_t layer_count;
@ -363,8 +374,7 @@ typedef struct {
sd_image_t init_image;
sd_image_t* ref_images;
int ref_images_count;
bool auto_resize_ref_image;
bool increase_ref_index;
const char* ref_image_args;
sd_image_t mask_image;
int width;
int height;
@ -374,6 +384,8 @@ typedef struct {
int batch_count;
sd_image_t control_image;
float control_strength;
sd_image_t ip_adapter_image;
float ip_adapter_strength;
sd_pm_params_t pm_params;
sd_pulid_params_t pulid_params;
sd_tiling_params_t vae_tiling_params;
@ -392,6 +404,12 @@ typedef struct {
int clip_skip;
sd_image_t init_image;
sd_image_t end_image;
sd_image_t* ref_images;
int ref_images_count;
sd_ref_video_t* ref_videos;
int ref_videos_count;
sd_audio_t* ref_audios;
int ref_audios_count;
sd_image_t* control_frames;
int control_frames_size;
int width;
@ -509,6 +527,27 @@ SD_API bool upscale(upscaler_ctx_t* upscaler_ctx,
SD_API int get_upscale_factor(upscaler_ctx_t* upscaler_ctx);
typedef struct adetailer_ctx_t adetailer_ctx_t;
typedef struct {
const char* prompt;
const char* negative_prompt;
const char* extra_ad_args;
} sd_adetailer_params_t;
SD_API adetailer_ctx_t* new_adetailer_ctx(const char* detector_path,
int n_threads,
const char* backend,
const char* params_backend);
SD_API void free_adetailer_ctx(adetailer_ctx_t* adetailer_ctx);
SD_API bool adetail_image(adetailer_ctx_t* adetailer_ctx,
sd_ctx_t* sd_ctx,
sd_image_t input_image,
const sd_adetailer_params_t* adetailer_params,
const sd_img_gen_params_t* inpaint_params,
sd_image_t** images_out,
int* num_images_out);
SD_API bool convert(const char* input_path,
const char* vae_path,
const char* output_path,

View File

@ -0,0 +1,86 @@
#!/usr/bin/env python3
"""Convert an Ultralytics YOLOv8 detection checkpoint for sd.cpp ADetailer."""
import argparse
import json
from pathlib import Path
def parse_args():
parser = argparse.ArgumentParser(
description="Convert an Ultralytics YOLOv8 detection .pt checkpoint to safetensors."
)
parser.add_argument("input", type=Path, help="input YOLOv8 detection checkpoint")
parser.add_argument("output", type=Path, help="output safetensors path")
parser.add_argument(
"--input-size", type=int, default=640, help="detector input size metadata (default: 640)"
)
return parser.parse_args()
def main():
args = parse_args()
if args.input_size < 32 or args.input_size % 32 != 0:
raise ValueError("--input-size must be a positive multiple of 32")
if args.output.suffix.lower() != ".safetensors":
raise ValueError("output path must use the .safetensors extension")
try:
import torch
from safetensors.torch import save_file
from ultralytics import YOLO
from ultralytics.nn.modules.head import Detect
except ImportError as exc:
raise SystemExit("conversion requires ultralytics, torch, and safetensors") from exc
torch_load = torch.load
def load_trusted_checkpoint(*load_args, **load_kwargs):
load_kwargs.setdefault("weights_only", False)
return torch_load(*load_args, **load_kwargs)
torch.load = load_trusted_checkpoint
try:
yolo = YOLO(str(args.input))
finally:
torch.load = torch_load
network = yolo.model
if not isinstance(network.model[-1], Detect) or network.model[-1].__class__.__name__ != "Detect":
raise ValueError("only YOLOv8 detection checkpoints are supported; segmentation is not yet supported")
network.eval()
network.fuse()
state_dict = network.state_dict()
required = {
"model.0.conv.weight",
"model.22.cv2.0.2.weight",
"model.22.cv3.0.2.weight",
}
missing = sorted(required.difference(state_dict))
if missing:
raise ValueError(f"checkpoint does not match the supported YOLOv8 layout; missing {missing}")
tensors = {}
for name, tensor in state_dict.items():
if not name.startswith("model.") or ".bn." in name or name.endswith("dfl.conv.weight"):
continue
if not (name.endswith(".weight") or name.endswith(".bias")):
continue
dtype = torch.float16 if name.endswith(".weight") else torch.float32
tensors[name] = tensor.detach().to(device="cpu", dtype=dtype).contiguous()
metadata = {
"format": "pt",
"yolov8.variant": "detect",
"yolov8.input_size": str(args.input_size),
"yolov8.num_classes": str(int(network.model[-1].nc)),
"yolov8.reg_max": str(int(network.model[-1].reg_max)),
"yolov8.names": json.dumps(yolo.names, ensure_ascii=False),
}
args.output.parent.mkdir(parents=True, exist_ok=True)
save_file(tensors, str(args.output), metadata=metadata)
print(f"wrote {args.output}: {len(tensors)} tensors")
if __name__ == "__main__":
main()

View File

@ -0,0 +1,335 @@
#!/usr/bin/env python3
"""Merge selected tensors from multiple safetensors files without loading weights.
Edit ``OUTPUT_PATH`` and ``SOURCE_RULES`` below, then run:
python scripts/merge_safetensors.py
Each source rule uses regular expressions against complete tensor names.
``include`` is required and matches when any expression succeeds. ``exclude``
wins over ``include``. Expressions are evaluated with ``re.search``.
"""
import json
import os
import re
import struct
from dataclasses import dataclass
from pathlib import Path
from typing import BinaryIO
# -----------------------------------------------------------------------------
# Configuration
# -----------------------------------------------------------------------------
OUTPUT_PATH = Path(r"..\models\diffusion_models\minimax_h3_ref2va_pruned_bf16.safetensors")
SOURCE_RULES = [
{
"path": Path(r"..\models\diffusion_models\minimax_h3_ref2va_bf16.safetensors"),
"include": [r".*"],
"exclude": [r".*adaln_proj\.linear.*", r"time_embedder.*"],
},
{
"path": Path(r"..\models\diffusion_models\minimax_h3_ref2va_pruned_int8_convrot.safetensors"),
"include": [r"^.*adaln_proj\.linear.*", "adaln_t_table"],
"exclude": [],
},
]
# Safetensors metadata is optional. Set this to a dict[str, str] if needed.
OUTPUT_METADATA = None
# Refuse to replace an existing output unless explicitly enabled.
OVERWRITE_OUTPUT = False
# Only tensor headers and this fixed-size buffer are held in memory.
COPY_BUFFER_SIZE = 8 * 1024 * 1024
PROGRESS_INTERVAL = 1024 * 1024 * 1024
MAX_HEADER_SIZE = 256 * 1024 * 1024
@dataclass(frozen=True)
class TensorEntry:
name: str
source_path: Path
source_data_offset: int
source_start: int
source_end: int
dtype: str
shape: list[int]
@property
def size(self) -> int:
return self.source_end - self.source_start
def format_bytes(size: int) -> str:
value = float(size)
for unit in ("B", "KiB", "MiB", "GiB", "TiB"):
if value < 1024.0 or unit == "TiB":
return f"{value:.2f} {unit}"
value /= 1024.0
raise AssertionError("unreachable")
def read_exact(file: BinaryIO, size: int, description: str) -> bytes:
data = file.read(size)
if len(data) != size:
raise ValueError(f"truncated {description}: expected {size} bytes, got {len(data)}")
return data
def read_safetensors_header(path: Path) -> tuple[dict, int, int]:
file_size = path.stat().st_size
with path.open("rb") as file:
header_size = struct.unpack("<Q", read_exact(file, 8, f"header size in {path}"))[0]
if header_size == 0 or header_size > MAX_HEADER_SIZE:
raise ValueError(
f"invalid header size in {path}: {header_size} "
f"(limit: {MAX_HEADER_SIZE})"
)
header_bytes = read_exact(file, header_size, f"header in {path}")
try:
header = json.loads(header_bytes)
except (UnicodeDecodeError, json.JSONDecodeError) as error:
raise ValueError(f"invalid safetensors JSON header in {path}: {error}") from error
if not isinstance(header, dict):
raise ValueError(f"safetensors header in {path} is not an object")
data_offset = 8 + header_size
if data_offset > file_size:
raise ValueError(f"safetensors data offset is past end of file: {path}")
return header, data_offset, file_size
def parse_tensor_entry(
name: str,
info: object,
source_path: Path,
source_data_offset: int,
source_file_size: int,
) -> TensorEntry:
if not isinstance(info, dict):
raise ValueError(f"{source_path}: tensor {name!r} has an invalid header entry")
dtype = info.get("dtype")
shape = info.get("shape")
offsets = info.get("data_offsets")
if not isinstance(dtype, str):
raise ValueError(f"{source_path}: tensor {name!r} has an invalid dtype")
if not isinstance(shape, list) or not all(
isinstance(dimension, int) and dimension >= 0 for dimension in shape
):
raise ValueError(f"{source_path}: tensor {name!r} has an invalid shape")
if (
not isinstance(offsets, list)
or len(offsets) != 2
or not all(isinstance(offset, int) for offset in offsets)
):
raise ValueError(f"{source_path}: tensor {name!r} has invalid data offsets")
start, end = offsets
if start < 0 or end < start or source_data_offset + end > source_file_size:
raise ValueError(
f"{source_path}: tensor {name!r} byte range [{start}, {end}) "
"is outside the file"
)
return TensorEntry(
name=name,
source_path=source_path,
source_data_offset=source_data_offset,
source_start=start,
source_end=end,
dtype=dtype,
shape=list(shape),
)
def compile_patterns(rule_index: int, field: str, values: object) -> list[re.Pattern[str]]:
if not isinstance(values, list) or not all(isinstance(value, str) for value in values):
raise TypeError(f"SOURCE_RULES[{rule_index}][{field!r}] must be a list of strings")
try:
return [re.compile(value) for value in values]
except re.error as error:
raise ValueError(
f"invalid regex in SOURCE_RULES[{rule_index}][{field!r}]: {error}"
) from error
def collect_entries() -> list[TensorEntry]:
if not SOURCE_RULES:
raise ValueError("SOURCE_RULES must contain at least one source")
entries: list[TensorEntry] = []
selected_by_name: dict[str, TensorEntry] = {}
header_cache: dict[Path, tuple[dict, int, int]] = {}
for rule_index, rule in enumerate(SOURCE_RULES):
if not isinstance(rule, dict) or "path" not in rule or "include" not in rule:
raise TypeError(
f"SOURCE_RULES[{rule_index}] must contain 'path' and 'include'"
)
source_path = Path(rule["path"])
if not source_path.is_file():
raise FileNotFoundError(f"source file does not exist: {source_path}")
source_path = source_path.resolve()
include = compile_patterns(rule_index, "include", rule["include"])
exclude = compile_patterns(rule_index, "exclude", rule.get("exclude", []))
if not include:
raise ValueError(f"SOURCE_RULES[{rule_index}]['include'] must not be empty")
if source_path not in header_cache:
header_cache[source_path] = read_safetensors_header(source_path)
header, data_offset, file_size = header_cache[source_path]
matched = 0
for name, info in header.items():
if name == "__metadata__":
continue
if not any(pattern.search(name) for pattern in include):
continue
if any(pattern.search(name) for pattern in exclude):
continue
entry = parse_tensor_entry(name, info, source_path, data_offset, file_size)
previous = selected_by_name.get(name)
if previous is not None:
raise ValueError(
f"tensor {name!r} was selected more than once:\n"
f" first: {previous.source_path}\n"
f" second: {source_path}"
)
selected_by_name[name] = entry
print(f"entry {entry}")
entries.append(entry)
matched += 1
print(f"Rule {rule_index}: selected {matched} tensors from {source_path}")
if matched == 0:
raise ValueError(
f"SOURCE_RULES[{rule_index}] did not select any tensors; check its regexes"
)
if not entries:
raise ValueError("no tensors were selected")
return entries
def build_output_header(entries: list[TensorEntry]) -> tuple[bytes, int]:
header: dict[str, object] = {}
if OUTPUT_METADATA is not None:
if not isinstance(OUTPUT_METADATA, dict) or not all(
isinstance(key, str) and isinstance(value, str)
for key, value in OUTPUT_METADATA.items()
):
raise TypeError("OUTPUT_METADATA must be None or a dict[str, str]")
header["__metadata__"] = OUTPUT_METADATA
output_offset = 0
for entry in entries:
header[entry.name] = {
"dtype": entry.dtype,
"shape": entry.shape,
"data_offsets": [output_offset, output_offset + entry.size],
}
output_offset += entry.size
header_bytes = json.dumps(header, separators=(",", ":"), ensure_ascii=False).encode(
"utf-8"
)
header_bytes += b" " * (-len(header_bytes) % 8)
return header_bytes, output_offset
def copy_tensor(source: BinaryIO, output: BinaryIO, entry: TensorEntry) -> None:
source.seek(entry.source_data_offset + entry.source_start)
remaining = entry.size
while remaining:
chunk = source.read(min(COPY_BUFFER_SIZE, remaining))
if not chunk:
raise OSError(
f"unexpected end of file while copying {entry.name!r} "
f"from {entry.source_path}"
)
output.write(chunk)
remaining -= len(chunk)
def write_output(entries: list[TensorEntry]) -> None:
if COPY_BUFFER_SIZE <= 0:
raise ValueError("COPY_BUFFER_SIZE must be positive")
output_path = OUTPUT_PATH.resolve()
source_paths = {entry.source_path.resolve() for entry in entries}
if output_path in source_paths:
raise ValueError("OUTPUT_PATH must not be one of the source files")
if output_path.exists() and not OVERWRITE_OUTPUT:
raise FileExistsError(
f"output already exists: {output_path}; set OVERWRITE_OUTPUT = True to replace it"
)
output_path.parent.mkdir(parents=True, exist_ok=True)
partial_path = output_path.with_name(output_path.name + ".partial")
if partial_path.exists():
raise FileExistsError(
f"partial output already exists: {partial_path}; remove it before retrying"
)
header_bytes, tensor_bytes = build_output_header(entries)
print(
f"Writing {len(entries)} tensors ({format_bytes(tensor_bytes)}) to {output_path}"
)
current_source_path: Path | None = None
current_source: BinaryIO | None = None
copied = 0
next_progress = PROGRESS_INTERVAL
try:
with partial_path.open("xb") as output:
output.write(struct.pack("<Q", len(header_bytes)))
output.write(header_bytes)
try:
for entry in entries:
if entry.source_path != current_source_path:
if current_source is not None:
current_source.close()
current_source = entry.source_path.open("rb")
current_source_path = entry.source_path
copy_tensor(current_source, output, entry)
copied += entry.size
if PROGRESS_INTERVAL > 0 and copied >= next_progress:
print(
f" copied {format_bytes(copied)} / "
f"{format_bytes(tensor_bytes)}"
)
while next_progress <= copied:
next_progress += PROGRESS_INTERVAL
finally:
if current_source is not None:
current_source.close()
if copied != tensor_bytes:
raise OSError(f"copied {copied} tensor bytes, expected {tensor_bytes}")
os.replace(partial_path, output_path)
except BaseException:
partial_path.unlink(missing_ok=True)
raise
print(f"Done: {output_path} ({format_bytes(output_path.stat().st_size)})")
def main() -> None:
entries = collect_entries()
write_output(entries)
if __name__ == "__main__":
main()

View File

@ -2,11 +2,14 @@
#define __SD_CONDITIONING_CONDITIONER_HPP__
#include <cmath>
#include <iomanip>
#include <limits>
#include <optional>
#include <sstream>
#include "core/tensor_ggml.hpp"
#include "core/util.h"
#include "model/diffusion/model.hpp"
#include "model/te/clip.hpp"
#include "model/te/llm.hpp"
#include "model/te/t5.hpp"
@ -24,6 +27,8 @@ struct SDCondition {
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>> c_ref_audios;
std::vector<MiniMaxH3ReferenceBlock> c_reference_blocks;
std::vector<sd::Tensor<float>> extra_c_crossattns;
@ -54,6 +59,12 @@ struct SDCondition {
}
}
for (const auto& tensor : c_ref_audios) {
if (!tensor.empty()) {
return false;
}
}
for (const auto& tensor : extra_c_crossattns) {
if (!tensor.empty()) {
return false;
@ -64,6 +75,18 @@ struct SDCondition {
}
};
enum class MiniMaxH3PresentationKind {
IMAGE,
VIDEO,
AUDIO,
};
struct MiniMaxH3PresentationItem {
MiniMaxH3PresentationKind kind = MiniMaxH3PresentationKind::IMAGE;
std::vector<sd::Tensor<float>> frames;
std::vector<float> timestamps;
};
static inline sd::Tensor<float> apply_token_weights(sd::Tensor<float> hidden_states,
const std::vector<float>& weights) {
if (hidden_states.empty()) {
@ -106,6 +129,8 @@ struct ConditionerParams {
int height = -1;
bool zero_out_masked = false;
const std::vector<sd::Tensor<float>>* ref_images = nullptr; // for qwen image edit
const std::vector<MiniMaxH3PresentationItem>* minimax_h3_references = nullptr;
RefImageParams ref_image_params;
};
struct Conditioner {
@ -115,6 +140,7 @@ public:
virtual SDCondition get_learned_condition(int n_threads,
const ConditionerParams& conditioner_params) = 0;
virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) = 0;
virtual void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) {}
virtual void set_max_graph_vram_bytes(size_t max_vram_bytes) {}
virtual void set_stream_layers_enabled(bool enabled) {}
virtual void set_runtime_backends(const std::vector<ggml_backend_t>& backends) {}
@ -1662,6 +1688,10 @@ struct AnimaConditioner : public Conditioner {
llm->get_param_tensors(tensors, "text_encoders.llm");
}
void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) override {
llm->get_param_tensor_ops(tensor_ops);
}
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
llm->set_max_graph_vram_bytes(max_vram_bytes);
}
@ -1780,6 +1810,7 @@ struct LLMEmbedder : public Conditioner {
SDVersion version;
std::shared_ptr<BPETokenizer> tokenizer;
std::shared_ptr<LLM::LLMRunner> llm;
std::shared_ptr<T5Runner> byt5;
LLMEmbedder(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
@ -1797,7 +1828,13 @@ struct LLMEmbedder : public Conditioner {
arch = LLM::LLMArch::GPT_OSS_20B;
} else if (sd_version_is_pid(version)) {
arch = LLM::LLMArch::GEMMA2_2B;
} else if (sd_version_is_ideogram4(version) || sd_version_is_boogu_image(version) || sd_version_is_sefi_image(version) || sd_version_is_krea2(version)) {
} else if (sd_version_is_lingbot_video(version) ||
sd_version_is_ideogram4(version) ||
sd_version_is_boogu_image(version) ||
sd_version_is_sefi_image(version) ||
sd_version_is_krea2(version) ||
sd_version_is_minimax_h3(version) ||
sd_version_is_mage_flow(version)) {
arch = LLM::LLMArch::QWEN3_VL;
} else if (sd_version_is_z_image(version) || version == VERSION_OVIS_IMAGE || version == VERSION_FLUX2_KLEIN) {
arch = LLM::LLMArch::QWEN3;
@ -1817,54 +1854,101 @@ struct LLMEmbedder : public Conditioner {
"text_encoders.llm",
enable_vision,
weight_manager);
if (sd_version_is_hunyuan_video(version)) {
const std::string byt5_prefix = "text_encoders.t5xxl.transformer";
for (const auto& [name, _] : tensor_storage_map) {
if (starts_with(name, byt5_prefix + ".")) {
byt5 = std::make_shared<T5Runner>(backend,
tensor_storage_map,
byt5_prefix,
false,
weight_manager);
break;
}
}
}
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
llm->get_param_tensors(tensors, "text_encoders.llm");
if (byt5) {
byt5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
}
}
void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) override {
llm->get_param_tensor_ops(tensor_ops);
}
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
llm->set_max_graph_vram_bytes(max_vram_bytes);
if (byt5) {
byt5->set_max_graph_vram_bytes(max_vram_bytes);
}
}
void set_stream_layers_enabled(bool enabled) override {
llm->set_stream_layers_enabled(enabled);
if (byt5) {
byt5->set_stream_layers_enabled(enabled);
}
}
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
llm->set_runtime_backends(backends);
if (byt5) {
byt5->set_runtime_backends(backends);
}
}
void set_graph_cut_layer_split_enabled(bool enabled) override {
if (llm) {
llm->set_graph_cut_layer_split_enabled(enabled);
}
if (byt5) {
byt5->set_graph_cut_layer_split_enabled(enabled);
}
}
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
if (llm) {
llm->set_graph_cut_layer_split_backend_vram_limits(limits);
}
if (byt5) {
byt5->set_graph_cut_layer_split_backend_vram_limits(limits);
}
}
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
llm->get_param_tensors(tensors, "text_encoders.llm");
if (byt5) {
byt5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
}
}
void set_flash_attention_enabled(bool enabled) override {
llm->set_flash_attention_enabled(enabled);
if (byt5) {
byt5->set_flash_attention_enabled(enabled);
}
}
void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
if (llm) {
llm->set_weight_adapter(adapter);
}
if (byt5) {
byt5->set_weight_adapter(adapter);
}
}
void runner_done() override {
if (llm) {
llm->runner_done();
}
if (byt5) {
byt5->runner_done();
}
}
std::tuple<std::vector<int>, std::vector<float>, std::vector<float>> tokenize(std::string text,
@ -1933,7 +2017,9 @@ struct LLMEmbedder : public Conditioner {
const std::set<int>& out_layers,
int prompt_template_encode_start_idx,
bool spell_quotes = false,
int max_length = 100000000) {
int max_length = 100000000,
const LLM::DeepStackImageEmbeds& deepstack_image_embeds = {},
const std::vector<LLM::ImageGrid>& image_grids = {}) {
auto tokens_weights_mask = tokenize(prompt, prompt_attn_range, min_length, max_length, spell_quotes);
auto& tokens = std::get<0>(tokens_weights_mask);
auto& weights = std::get<1>(tokens_weights_mask);
@ -1966,7 +2052,9 @@ struct LLMEmbedder : public Conditioner {
false,
false,
true,
true);
true,
deepstack_image_embeds,
image_grids);
GGML_ASSERT(!hidden_states.empty());
hidden_states = apply_token_weights(std::move(hidden_states), weights);
GGML_ASSERT(hidden_states.shape()[1] > prompt_template_encode_start_idx);
@ -1993,6 +2081,54 @@ struct LLMEmbedder : public Conditioner {
return new_hidden_states;
}
void resize_image_dims(int height, int width, int& h_bar, int& w_bar, int factor, int min_size, int max_size, RefImageResizeMode mode) {
if (min_size > 0 && min_size == max_size) {
if (mode == RefImageResizeMode::AREA) {
double beta = std::sqrt(static_cast<double>(min_size) / (static_cast<double>(height) * width));
h_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::round(height * beta / factor)) * static_cast<int>(factor));
w_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::round(width * beta / factor)) * static_cast<int>(factor));
} else if (mode == RefImageResizeMode::LONGEST_SIDE) {
int current_max_side = std::max(height, width);
double beta = static_cast<double>(min_size) / current_max_side;
h_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::round(height * beta / factor)) * static_cast<int>(factor));
w_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::round(width * beta / factor)) * static_cast<int>(factor));
}
return;
}
if (mode == RefImageResizeMode::AREA) {
double current_area = static_cast<double>(h_bar) * w_bar;
if (max_size > 0 && current_area > max_size) {
double beta = std::sqrt((static_cast<double>(height) * width) / static_cast<double>(max_size));
h_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::floor(height / beta / factor)) * static_cast<int>(factor));
w_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::floor(width / beta / factor)) * static_cast<int>(factor));
} else if (min_size > 0 && current_area < min_size) {
double beta = std::sqrt(static_cast<double>(min_size) / (static_cast<double>(height) * width));
h_bar = static_cast<int>(std::ceil(height * beta / factor)) * static_cast<int>(factor);
w_bar = static_cast<int>(std::ceil(width * beta / factor)) * static_cast<int>(factor);
}
} else if (mode == RefImageResizeMode::LONGEST_SIDE) {
int current_max_side = std::max(height, width);
if (max_size > 0 && current_max_side > max_size) {
double beta = static_cast<double>(max_size) / current_max_side;
h_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::floor(height * beta / factor)) * static_cast<int>(factor));
w_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::floor(width * beta / factor)) * static_cast<int>(factor));
} else if (min_size > 0 && current_max_side < min_size) {
double beta = static_cast<double>(min_size) / current_max_side;
h_bar = static_cast<int>(std::ceil(height * beta / factor)) * static_cast<int>(factor);
w_bar = static_cast<int>(std::ceil(width * beta / factor)) * static_cast<int>(factor);
}
}
}
SDCondition get_learned_condition(int n_threads,
const ConditionerParams& conditioner_params) override {
std::string prompt;
@ -2000,6 +2136,8 @@ struct LLMEmbedder : public Conditioner {
std::vector<std::string> extra_prompts;
std::vector<std::pair<int, int>> extra_prompts_attn_range;
std::vector<std::pair<int, sd::Tensor<float>>> image_embeds;
LLM::DeepStackImageEmbeds deepstack_image_embeds;
std::vector<LLM::ImageGrid> image_grids;
int prompt_template_encode_start_idx = 34;
int min_length = 0; // pad tokens
int max_length = 100000000;
@ -2008,37 +2146,266 @@ struct LLMEmbedder : public Conditioner {
std::set<int> out_layers;
int64_t t0 = ggml_time_ms();
RefImageResizeMode resize_mode = conditioner_params.ref_image_params.vlm_resize_mode;
if (sd_version_is_qwen_image(version)) {
if (sd_version_is_minimax_h3(version)) {
prompt_template_encode_start_idx = 0;
out_layers = {50};
prompt_attn_range = {0, 0};
if (llm->enable_vision) {
const std::string placeholder = "<|image_pad|>";
const int patch_size = llm->config.vision.patch_size;
const int factor = patch_size * llm->config.vision.spatial_merge_size;
auto resize_for_vision = [&](const sd::Tensor<float>& image) {
int height = static_cast<int>(image.shape()[1]);
int width = static_cast<int>(image.shape()[0]);
int h_bar = std::max(factor, static_cast<int>(std::round(static_cast<double>(height) / factor)) * factor);
int w_bar = std::max(factor, static_cast<int>(std::round(static_cast<double>(width) / factor)) * factor);
resize_image_dims(height,
width,
h_bar,
w_bar,
factor,
3136,
12845056,
RefImageResizeMode::AREA);
auto resized = sd::ops::interpolate(
image,
std::vector<int64_t>{w_bar, h_bar, image.shape()[2], image.shape()[3]});
for (int64_t i = 0; i < resized.numel(); ++i) {
resized[i] = std::clamp(resized[i], 0.f, 1.f) * 2.f - 1.f;
}
return resized;
};
auto add_vision_outputs = [&](std::vector<sd::Tensor<float>> image_outputs,
int grid_h,
int grid_w) {
GGML_ASSERT(image_outputs.size() == 4);
auto image_embed = std::move(image_outputs[0]);
prompt += "<|vision_start|>";
int image_embed_idx = static_cast<int>(tokenizer->encode(prompt, nullptr).size());
image_embeds.emplace_back(image_embed_idx, image_embed);
if (deepstack_image_embeds.empty()) {
deepstack_image_embeds.resize(image_outputs.size() - 1);
}
for (size_t layer = 0; layer < deepstack_image_embeds.size(); ++layer) {
deepstack_image_embeds[layer].emplace_back(image_embed_idx, std::move(image_outputs[layer + 1]));
}
image_grids.push_back({image_embed_idx,
static_cast<int>(image_embed.shape()[1]),
grid_h,
grid_w});
for (int64_t i = 0; i < image_embed.shape()[1]; ++i) {
prompt += placeholder;
}
prompt += "<|vision_end|>";
};
const auto* references = conditioner_params.minimax_h3_references;
if (references != nullptr && !references->empty()) {
int picture_index = 0;
int video_index = 0;
int audio_index = 0;
for (const auto& item : *references) {
if (item.kind == MiniMaxH3PresentationKind::AUDIO) {
prompt += "<Audio " + std::to_string(++audio_index) + ">: ";
continue;
}
if (item.kind == MiniMaxH3PresentationKind::IMAGE) {
GGML_ASSERT(item.frames.size() == 1);
auto resized = resize_for_vision(item.frames[0]);
prompt += "<Picture " + std::to_string(++picture_index) + ">: ";
add_vision_outputs(llm->encode_image_outputs(n_threads,
resized,
false,
true,
true),
static_cast<int>(resized.shape()[1]) / patch_size,
static_cast<int>(resized.shape()[0]) / patch_size);
continue;
}
GGML_ASSERT(!item.frames.empty());
prompt += "<Video " + std::to_string(++video_index) + ">: ";
for (size_t frame = 0; frame < item.frames.size(); frame += 2) {
size_t next = std::min(frame + 1, item.frames.size() - 1);
float t0 = frame < item.timestamps.size() ? item.timestamps[frame] : frame / 2.f;
float t1 = next < item.timestamps.size() ? item.timestamps[next] : next / 2.f;
std::ostringstream timestamp;
timestamp << '<' << std::fixed << std::setprecision(1) << (t0 + t1) * 0.5f << " seconds>";
prompt += timestamp.str();
auto first = resize_for_vision(item.frames[frame]);
auto second = resize_for_vision(item.frames[next]);
if (first.shape()[0] != second.shape()[0] || first.shape()[1] != second.shape()[1]) {
second = sd::ops::interpolate(second,
std::vector<int64_t>{first.shape()[0],
first.shape()[1],
second.shape()[2],
second.shape()[3]});
}
auto pair = sd::ops::concat(first.unsqueeze(2), second.unsqueeze(2), 2);
add_vision_outputs(llm->encode_video_block_outputs(n_threads,
pair,
false,
true,
true),
static_cast<int>(first.shape()[1]) / patch_size,
static_cast<int>(first.shape()[0]) / patch_size);
}
}
} else if (conditioner_params.ref_images != nullptr) {
for (size_t i = 0; i < conditioner_params.ref_images->size(); ++i) {
auto resized = resize_for_vision((*conditioner_params.ref_images)[i]);
prompt += "<Picture " + std::to_string(i + 1) + ">: ";
add_vision_outputs(llm->encode_image_outputs(n_threads,
resized,
false,
true,
true),
static_cast<int>(resized.shape()[1]) / patch_size,
static_cast<int>(resized.shape()[0]) / patch_size);
}
}
}
prompt += conditioner_params.text;
} else if (sd_version_is_hunyuan_video(version)) {
prompt_template_encode_start_idx = 98;
out_layers = {26};
prompt =
"<|im_start|>system\nYou are a helpful assistant. Describe the video by detailing the following aspects:\n"
"1. The main content and theme of the video.\n"
"2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects.\n"
"3. Actions, events, behaviors temporal relationships, physical movement changes of the objects.\n"
"4. background environment, light, style and atmosphere.\n"
"5. camera angles, movements, and transitions used in the video.<|im_end|>\n"
"<|im_start|>user\n";
prompt_attn_range.first = static_cast<int>(prompt.size());
prompt += conditioner_params.text;
prompt_attn_range.second = static_cast<int>(prompt.size());
prompt += "<|im_end|>\n<|im_start|>assistant\n";
} else if (sd_version_is_lingbot_video(version)) {
const int pad_token = 151643;
const std::string prompt_prefix =
"<|im_start|>system\nGiven a user input that may include a text prompt alone, "
"a text prompt with an image reference, or a text prompt with a video reference "
"or a video reference alone, generate an \"Enhanced prompt\" that provides detailed "
"visual descriptions suitable for video generation. Evaluate the level of detail "
"in the user's input: if it is simple, enrich it by adding specifics about colors, "
"shapes, sizes, textures, lighting, motion dynamics, camera movement, temporal "
"progression, and spatial relationships to create vivid, concrete, and temporally "
"coherent scenes to create vivid and concrete scenes. Please generate only the "
"enhanced description for the prompt below and avoid including any additional "
"commentary or evaluations:<|im_end|>\n<|im_start|>user\n";
auto prefix_tokens = tokenizer->encode(prompt_prefix, nullptr);
prompt_template_encode_start_idx = 0;
for (int token : prefix_tokens) {
if (token != pad_token) {
prompt_template_encode_start_idx++;
}
}
LOG_DEBUG("prompt_template_encode_start_idx %d", prompt_template_encode_start_idx);
prompt = prompt_prefix;
if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
LOG_INFO("QwenImageEditPlusPipeline");
LOG_INFO("LingBotVideoI2VPipeline");
const std::string placeholder = "<|image_pad|>";
std::string img_prompt;
for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
const auto& image = (*conditioner_params.ref_images)[i];
const int factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
int height = static_cast<int>(image.shape()[1]);
int width = static_cast<int>(image.shape()[0]);
int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
if (min_pixels <= 0) {
if (resize_mode == RefImageResizeMode::AREA) {
min_pixels = static_cast<int>(4 * factor * factor);
} else {
min_pixels = static_cast<int>(2 * factor);
}
}
int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
if (max_pixels <= 0) {
if (resize_mode == RefImageResizeMode::AREA) {
max_pixels = static_cast<int>(16384 * factor * factor);
} else {
max_pixels = static_cast<int>(128 * factor);
}
}
int h_bar = std::max(factor, static_cast<int>(std::round(static_cast<double>(height) / factor) * factor));
int w_bar = std::max(factor, static_cast<int>(std::round(static_cast<double>(width) / factor) * factor));
if (std::max(height, width) > 200 * std::min(height, width)) {
LOG_WARN("LingBotVideo image aspect ratio is very large: %dx%d", width, height);
}
resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
LOG_DEBUG("resize LingBotVideo ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
auto resized_image = clip_preprocess(image, w_bar, h_bar);
auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true);
GGML_ASSERT(!image_embed.empty());
std::string image_prefix = prompt + img_prompt + "<|vision_start|>";
int image_embed_idx = static_cast<int>(tokenizer->encode(image_prefix, nullptr).size());
image_embeds.emplace_back(image_embed_idx, image_embed);
img_prompt += "<|vision_start|>";
int64_t num_image_tokens = image_embed.shape()[1];
img_prompt.reserve(img_prompt.size() + static_cast<size_t>(num_image_tokens) * placeholder.size() + 32);
for (int j = 0; j < num_image_tokens; j++) {
img_prompt += placeholder;
}
img_prompt += "<|vision_end|>";
}
prompt += img_prompt;
}
prompt += conditioner_params.text;
prompt_attn_range = {0, 0};
prompt += "<|im_end|>\n<|im_start|>assistant\n";
} else if (sd_version_is_qwen_image(version) || sd_version_is_mage_flow(version)) {
if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
LOG_INFO("%s", sd_version_is_mage_flow(version) ? "MageFlowEditPipeline" : "QwenImageEditPlusPipeline");
prompt_template_encode_start_idx = 64;
int image_embed_idx = 64 + 6;
int min_pixels = 384 * 384;
int max_pixels = 560 * 560;
int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
if (min_pixels <= 0) {
min_pixels = sd_version_is_mage_flow(version) ? -1 : 384;
if (min_pixels > 0 && resize_mode == RefImageResizeMode::AREA) {
min_pixels *= min_pixels;
}
}
int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
if (max_pixels <= 0) {
max_pixels = sd_version_is_mage_flow(version) ? 384 : 560;
if (resize_mode == RefImageResizeMode::AREA) {
max_pixels *= max_pixels;
}
}
std::string placeholder = "<|image_pad|>";
std::string img_prompt;
for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
const auto& image = (*conditioner_params.ref_images)[i];
double factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
const int factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
int height = static_cast<int>(image.shape()[1]);
int width = static_cast<int>(image.shape()[0]);
int h_bar = static_cast<int>(std::round(height / factor) * factor);
int w_bar = static_cast<int>(std::round(width / factor) * factor);
int h_bar = static_cast<int>(std::round(static_cast<double>(height) / factor) * factor);
int w_bar = static_cast<int>(std::round(static_cast<double>(width) / factor) * factor);
if (static_cast<double>(h_bar) * w_bar > max_pixels) {
double beta = std::sqrt((height * width) / static_cast<double>(max_pixels));
h_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::floor(height / beta / factor)) * static_cast<int>(factor));
w_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::floor(width / beta / factor)) * static_cast<int>(factor));
} else if (static_cast<double>(h_bar) * w_bar < min_pixels) {
double beta = std::sqrt(static_cast<double>(min_pixels) / (height * width));
h_bar = static_cast<int>(std::ceil(height * beta / factor)) * static_cast<int>(factor);
w_bar = static_cast<int>(std::ceil(width * beta / factor)) * static_cast<int>(factor);
}
resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
@ -2049,7 +2416,7 @@ struct LLMEmbedder : public Conditioner {
image_embeds.emplace_back(image_embed_idx, image_embed);
image_embed_idx += 1 + static_cast<int>(image_embed.shape()[1]) + 6;
img_prompt += "Picture " + std::to_string(i + 1) + ": <|vision_start|>"; // [24669, 220, index, 25, 220, 151652]
img_prompt += (sd_version_is_mage_flow(version) ? "Image " : "Picture ") + std::to_string(i + 1) + ": <|vision_start|>";
int64_t num_image_tokens = image_embed.shape()[1];
img_prompt.reserve(num_image_tokens * placeholder.size());
for (int j = 0; j < num_image_tokens; j++) {
@ -2077,6 +2444,9 @@ struct LLMEmbedder : public Conditioner {
prompt += "<|im_end|>\n<|im_start|>assistant\n";
}
if (sd_version_is_mage_flow(version)) {
max_length = 2048 + prompt_template_encode_start_idx;
}
} else if (sd_version_is_boogu_image(version)) {
prompt_template_encode_start_idx = 0;
@ -2093,16 +2463,33 @@ struct LLMEmbedder : public Conditioner {
std::string img_prompt;
const std::string placeholder = "<|image_pad|>";
int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
if (min_pixels <= 0) {
min_pixels = 384;
if (resize_mode == RefImageResizeMode::AREA) {
min_pixels *= min_pixels;
}
}
int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
if (max_pixels <= 0) {
max_pixels = 384;
if (resize_mode == RefImageResizeMode::AREA) {
max_pixels *= max_pixels;
}
}
for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
const auto& image = (*conditioner_params.ref_images)[i];
double factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
const int factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
int height = static_cast<int>(image.shape()[1]);
int width = static_cast<int>(image.shape()[0]);
double beta = std::sqrt((384.0 * 384.0) / (static_cast<double>(height) * static_cast<double>(width)));
int h_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::round(height * beta / factor)) * static_cast<int>(factor));
int w_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::round(width * beta / factor)) * static_cast<int>(factor));
int h_bar = std::max(factor,
static_cast<int>(std::round(static_cast<double>(height) / factor)) * factor);
int w_bar = std::max(factor,
static_cast<int>(std::round(static_cast<double>(width) / factor)) * factor);
resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
@ -2144,17 +2531,33 @@ struct LLMEmbedder : public Conditioner {
if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
std::string img_prompt = "";
const std::string placeholder = "<|image_pad|>";
int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
if (min_pixels <= 0) {
min_pixels = 384;
if (resize_mode == RefImageResizeMode::AREA) {
min_pixels *= min_pixels;
}
}
int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
if (max_pixels <= 0) {
max_pixels = 1024;
if (resize_mode == RefImageResizeMode::AREA) {
max_pixels *= max_pixels;
}
}
for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
const auto& image = (*conditioner_params.ref_images)[i];
double factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
const int factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
int height = static_cast<int>(image.shape()[1]);
int width = static_cast<int>(image.shape()[0]);
double beta = std::sqrt((384.0 * 384.0) / (static_cast<double>(height) * static_cast<double>(width)));
int h_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::round(height * beta / factor)) * static_cast<int>(factor));
int w_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::round(width * beta / factor)) * static_cast<int>(factor));
int h_bar = std::max(factor,
static_cast<int>(std::round(static_cast<double>(height) / factor)) * factor);
int w_bar = std::max(factor,
static_cast<int>(std::round(static_cast<double>(width) / factor)) * factor);
resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
@ -2191,30 +2594,33 @@ struct LLMEmbedder : public Conditioner {
min_length = 512 + prompt_template_encode_start_idx;
int image_embed_idx = 36 + 6;
int min_pixels = 384 * 384;
int max_pixels = 560 * 560;
int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
if (min_pixels <= 0) {
min_pixels = 384;
if (resize_mode == RefImageResizeMode::AREA) {
min_pixels *= min_pixels;
}
}
int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
if (max_pixels <= 0) {
max_pixels = 560;
if (resize_mode == RefImageResizeMode::AREA) {
max_pixels *= max_pixels;
}
}
std::string placeholder = "<|image_pad|>";
std::string img_prompt;
for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
const auto& image = (*conditioner_params.ref_images)[i];
double factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
const int factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
int height = static_cast<int>(image.shape()[1]);
int width = static_cast<int>(image.shape()[0]);
int h_bar = static_cast<int>(std::round(height / factor) * factor);
int w_bar = static_cast<int>(std::round(width / factor) * factor);
int h_bar = static_cast<int>(std::round(static_cast<double>(height) / factor) * factor);
int w_bar = static_cast<int>(std::round(static_cast<double>(width) / factor) * factor);
if (static_cast<double>(h_bar) * w_bar > max_pixels) {
double beta = std::sqrt((height * width) / static_cast<double>(max_pixels));
h_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::floor(height / beta / factor)) * static_cast<int>(factor));
w_bar = std::max(static_cast<int>(factor),
static_cast<int>(std::floor(width / beta / factor)) * static_cast<int>(factor));
} else if (static_cast<double>(h_bar) * w_bar < min_pixels) {
double beta = std::sqrt(static_cast<double>(min_pixels) / (height * width));
h_bar = static_cast<int>(std::ceil(height * beta / factor)) * static_cast<int>(factor);
w_bar = static_cast<int>(std::ceil(width * beta / factor)) * static_cast<int>(factor);
}
resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
@ -2414,8 +2820,50 @@ struct LLMEmbedder : public Conditioner {
out_layers,
prompt_template_encode_start_idx,
spell_quotes,
max_length);
max_length,
deepstack_image_embeds,
image_grids);
std::vector<sd::Tensor<float>> extra_hidden_states_vec;
if (sd_version_is_hunyuan_video(version) && byt5) {
std::vector<std::string> quoted_texts;
auto collect_quoted = [&](const std::string& open, const std::string& close) {
size_t begin = 0;
while ((begin = conditioner_params.text.find(open, begin)) != std::string::npos) {
size_t content_begin = begin + open.size();
size_t end = conditioner_params.text.find(close, content_begin);
if (end == std::string::npos) {
break;
}
quoted_texts.push_back(conditioner_params.text.substr(content_begin, end - content_begin));
begin = end + close.size();
}
};
collect_quoted("\"", "\"");
collect_quoted("\xE2\x80\x98", "\xE2\x80\x99");
collect_quoted("\xE2\x80\x9C", "\xE2\x80\x9D");
if (!quoted_texts.empty()) {
std::string byt5_text;
for (const auto& text : quoted_texts) {
byt5_text += "Text \"" + text + "\". ";
}
std::vector<int> tokens;
tokens.reserve(byt5_text.size() + 1);
for (unsigned char byte : byt5_text) {
tokens.push_back(static_cast<int>(byte) + 3);
}
tokens.push_back(1);
sd::Tensor<int32_t> input_ids({static_cast<int64_t>(tokens.size())}, tokens);
auto byt5_hidden_states = byt5->compute(n_threads,
input_ids,
sd::Tensor<float>(),
false,
true,
true);
GGML_ASSERT(!byt5_hidden_states.empty());
extra_hidden_states_vec.push_back(std::move(byt5_hidden_states));
}
}
for (int i = 0; i < extra_prompts.size(); i++) {
auto extra_hidden_states = encode_prompt(n_threads,
extra_prompts[i],
@ -2435,6 +2883,17 @@ struct LLMEmbedder : public Conditioner {
SDCondition result;
result.c_crossattn = std::move(hidden_states);
result.extra_c_crossattns = std::move(extra_hidden_states_vec);
if (sd_version_is_minimax_h3(version)) {
std::vector<int32_t> tags(static_cast<size_t>(result.c_crossattn.shape()[1]), 1);
for (const auto& [index, image_embed] : image_embeds) {
int64_t begin = std::max<int64_t>(0, index - 1);
int64_t end = std::min<int64_t>(static_cast<int64_t>(tags.size()),
index + image_embed.shape()[1] + 1);
std::fill(tags.begin() + begin, tags.begin() + end, 0);
}
int64_t tag_count = static_cast<int64_t>(tags.size());
result.c_token_types = sd::Tensor<int32_t>({tag_count}, std::move(tags));
}
return result;
}
};
@ -2545,6 +3004,10 @@ struct LTXAVEmbedder : public Conditioner {
projector->get_param_tensors(tensors, "text_embedding_projection");
}
void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) override {
llm->get_param_tensor_ops(tensor_ops);
}
void set_flash_attention_enabled(bool enabled) override {
llm->set_flash_attention_enabled(enabled);
projector->set_flash_attention_enabled(enabled);

View File

@ -1689,6 +1689,8 @@ struct GGMLRunnerContext {
bool conv2d_direct_enabled = false;
bool circular_x_enabled = false;
bool circular_y_enabled = false;
ggml_tensor* ip_context = nullptr;
float ip_scale = 1.0f;
std::shared_ptr<WeightAdapter> weight_adapter = nullptr;
std::vector<std::pair<ggml_tensor*, std::string>>* debug_tensors = nullptr;
std::function<ggml_tensor*(const std::string&)> get_cache_tensor;
@ -1751,7 +1753,7 @@ protected:
std::vector<size_t> graph_cut_layer_split_backend_vram_limits_;
std::vector<ggml_backend_t> extra_runtime_backends; // borrowed (SDBackendManager-owned)
ggml_backend_sched_t sched = nullptr; // owned, multi-device only
ggml_backend_sched_t sched = nullptr; // owned
ggml_backend_t cpu_fallback_backend = nullptr; // owned, sched requires a trailing CPU backend
bool multi_device_eval_callback_warned = false;
@ -2145,8 +2147,22 @@ protected:
return !extra_runtime_backends.empty();
}
bool graph_requires_backend_fallback(ggml_cgraph* gf) const {
if (gf == nullptr || sd_backend_is_cpu(runtime_backend)) {
return false;
}
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 && !ggml_backend_supports_op(runtime_backend, node)) {
return true;
}
}
return false;
}
bool alloc_compute_buffer(ggml_cgraph* gf) {
if (is_multi_device()) {
if (sched != nullptr || is_multi_device() || graph_requires_backend_fallback(gf)) {
// The sched replaces the gallocr. Do NOT ggml_backend_sched_reserve
// the graph here: reserve runs split_graph, which rewires the
// graph's src pointers to sched-internal copy tensors, and the
@ -2154,6 +2170,10 @@ protected:
// rewired graph, silently corrupting every cross-backend input. A
// graph must be split at most once; the alloc in execute_graph
// performs the real allocation.
if (compute_allocr != nullptr) {
ggml_gallocr_free(compute_allocr);
compute_allocr = nullptr;
}
return ensure_sched(gf);
}
if (compute_allocr != nullptr) {
@ -2751,7 +2771,7 @@ protected:
};
ComputeBufferGuard compute_buffer_guard(this, free_compute_buffer);
if (is_multi_device()) {
if (sched != nullptr) {
ggml_backend_sched_reset(sched);
pin_multi_device_nodes(gf); // reset clears the pins; re-apply before alloc
if (!ggml_backend_sched_alloc_graph(sched, gf)) {
@ -2772,9 +2792,9 @@ protected:
}
ggml_status status;
if (is_multi_device()) {
if (sched != nullptr) {
if (sd_get_backend_eval_callback() != nullptr && !multi_device_eval_callback_warned) {
LOG_WARN("%s: eval callback is not supported with multiple runtime backends; ignoring",
LOG_WARN("%s: eval callback is not supported with the backend scheduler; ignoring",
get_desc().c_str());
multi_device_eval_callback_warned = true;
}
@ -3016,12 +3036,9 @@ public:
// do copy after alloc graph
void set_backend_tensor_data(ggml_tensor* tensor, const void* data) {
if (is_multi_device()) {
// The sched only assigns a backend (and thus a buffer) to tensors
// that participate in the graph; flag standalone data tensors as
// inputs so they get one.
// The scheduler only allocates standalone data tensors when they are
// marked as graph inputs. The flag is harmless for single-backend graphs.
ggml_set_input(tensor);
}
backend_tensor_data_map[tensor] = data;
}
@ -3238,6 +3255,11 @@ protected:
virtual void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") {}
virtual enum ggml_op param_usage_op(const std::string& name) const {
(void)name;
return GGML_OP_NONE;
}
public:
void init(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") {
if (prefix.size() > 0) {
@ -3288,6 +3310,18 @@ public:
}
}
void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) {
for (auto& pair : blocks) {
pair.second->get_param_tensor_ops(tensor_ops);
}
for (auto& pair : params) {
enum ggml_op op = param_usage_op(pair.first);
if (op != GGML_OP_NONE) {
tensor_ops[pair.second] = op;
}
}
}
virtual std::string get_desc() {
return "GGMLBlock";
}
@ -3309,7 +3343,7 @@ public:
class Identity : public UnaryBlock {
public:
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
return x;
}
};
@ -3368,7 +3402,7 @@ public:
force_prec_f32 = force_prec_f32_;
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
ggml_tensor* w = params["weight"];
ggml_tensor* b = nullptr;
if (bias) {
@ -3415,6 +3449,10 @@ protected:
params["weight"] = ggml_new_tensor_2d(ctx, wtype, embedding_dim, num_embeddings);
}
enum ggml_op param_usage_op(const std::string& name) const override {
return name == "weight" ? GGML_OP_GET_ROWS : GGML_OP_NONE;
}
public:
Embedding(int64_t num_embeddings, int64_t embedding_dim)
: embedding_dim(embedding_dim),
@ -3422,7 +3460,7 @@ public:
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* input_ids) {
ggml_tensor* input_ids) override {
// input_ids: [N, n_token]
auto weight = params["weight"];
@ -3482,11 +3520,11 @@ public:
scale = scale_value;
}
std::string get_desc() {
std::string get_desc() override {
return "Conv2d";
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
ggml_tensor* w = params["weight"];
ggml_tensor* b = nullptr;
if (bias) {
@ -3569,11 +3607,11 @@ public:
scale = scale_value;
}
std::string get_desc() {
std::string get_desc() override {
return "Conv2d_grouped";
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
ggml_tensor* w = params["weight"];
ggml_tensor* b = nullptr;
if (bias) {
@ -3609,18 +3647,19 @@ public:
if (groups == in_channels && groups == out_channels) {
ggml_tensor* res;
if (ctx->conv2d_direct_enabled) {
res = ggml_conv_2d_dw_direct(ctx->ggml_ctx, x, w,
res = ggml_conv_2d_dw_direct(ctx->ggml_ctx, w, x,
stride.second, stride.first,
padding.second, padding.first,
dilation.second, dilation.first);
} else {
res = ggml_conv_2d_dw(ctx->ggml_ctx, x, w,
res = ggml_conv_2d_dw(ctx->ggml_ctx, w, x,
stride.second, stride.first,
padding.second, padding.first,
dilation.second, dilation.first);
}
if (b) {
res = ggml_add(ctx->ggml_ctx, res, b);
b = ggml_reshape_4d(ctx->ggml_ctx, b, 1, 1, b->ne[0], 1);
res = ggml_add_inplace(ctx->ggml_ctx, res, b);
}
return res;
}
@ -3725,7 +3764,7 @@ public:
bias(bias),
force_prec_f32(force_prec_f32) {}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
ggml_tensor* w = params["weight"];
ggml_tensor* b = nullptr;
if (ctx->weight_adapter) {
@ -3778,7 +3817,7 @@ public:
elementwise_affine(elementwise_affine),
bias(bias) {}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
ggml_tensor* w = nullptr;
ggml_tensor* b = nullptr;
@ -3865,7 +3904,7 @@ public:
: hidden_size(hidden_size),
eps(eps) {}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
ggml_tensor* w = params["weight"];
if (ctx->weight_adapter) {
w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");

View File

@ -83,6 +83,10 @@ static bool parse_backend_module(const std::string& raw_name, SDBackendModule* m
*module = SDBackendModule::UPSCALER;
return true;
}
if (name == "detector" || name == "adetailer" || name == "yolo") {
*module = SDBackendModule::DETECTOR;
return true;
}
return false;
}
@ -956,6 +960,8 @@ const char* sd_backend_module_name(SDBackendModule module) {
return "photomaker";
case SDBackendModule::UPSCALER:
return "upscaler";
case SDBackendModule::DETECTOR:
return "detector";
}
return "unknown";
}

View File

@ -20,6 +20,7 @@ enum class SDBackendModule {
CONTROL_NET,
PHOTOMAKER,
UPSCALER,
DETECTOR,
};
struct SDBackendAssignment {

1020
src/detailer.cpp Normal file

File diff suppressed because it is too large Load Diff

75
src/detailer.h Normal file
View File

@ -0,0 +1,75 @@
#ifndef __SD_DETAILER_H__
#define __SD_DETAILER_H__
#include <memory>
#include <string>
#include <vector>
#include "core/ggml_extend_backend.h"
#include "model/detector/yolov8.h"
#include "model_manager.h"
#include "stable-diffusion.h"
struct ADetailerDetection {
float x1 = 0.f;
float y1 = 0.f;
float x2 = 0.f;
float y2 = 0.f;
float confidence = 0.f;
int class_id = 0;
};
enum ADetailerSort {
ADETAILER_SORT_NONE,
ADETAILER_SORT_LEFT_TO_RIGHT,
ADETAILER_SORT_CENTER_TO_EDGE,
ADETAILER_SORT_AREA,
};
struct ADetailerParams {
const char* prompt = nullptr;
const char* negative_prompt = nullptr;
int input_size = 640;
float confidence = 0.3f;
float nms_threshold = 0.45f;
int max_detections = 100;
int mask_k_largest = 0;
float mask_min_ratio = 0.f;
float mask_max_ratio = 1.f;
int dilate_erode = 4;
int x_offset = 0;
int y_offset = 0;
bool merge_masks = false;
bool invert_mask = false;
int mask_blur = 4;
int inpaint_padding = 32;
int inpaint_width = 512;
int inpaint_height = 512;
float denoising_strength = 0.4f;
int steps = 0;
float cfg_scale = -1.f;
sample_method_t sample_method = SAMPLE_METHOD_COUNT;
scheduler_t scheduler = SCHEDULER_COUNT;
ADetailerSort sort_by = ADETAILER_SORT_NONE;
};
struct ADetailerGGML {
SDBackendManager backend_manager;
std::shared_ptr<ModelManager> model_manager;
std::shared_ptr<YOLOv8Runner> detector;
std::vector<std::string> class_names;
int n_threads = 1;
std::string backend_spec;
std::string params_backend_spec;
ADetailerGGML(int n_threads,
std::string backend_spec,
std::string params_backend_spec);
~ADetailerGGML();
bool load_from_file(const std::string& detector_path);
std::vector<ADetailerDetection> predict(sd_image_t image,
const ADetailerParams& params);
};
#endif // __SD_DETAILER_H__

View File

@ -56,7 +56,7 @@ tokenize_photomaker_trigger(FrozenCLIPEmbedderWithCustomWords& clip_conditioner,
true);
std::vector<bool> class_token_mask;
for (int i = 0; i < tokens.size(); i++) {
class_token_mask.push_back(class_idx + 1 <= i && i < class_idx + 1 + trigger_token_count);
class_token_mask.push_back(class_idx >= 0 && class_idx + 1 <= i && i < class_idx + 1 + trigger_token_count);
}
return std::make_tuple(tokens, weights, class_token_mask);

View File

@ -35,12 +35,15 @@ enum SDVersion {
VERSION_WAN2,
VERSION_WAN2_2_I2V,
VERSION_WAN2_2_TI2V,
VERSION_LINGBOT_VIDEO,
VERSION_QWEN_IMAGE,
VERSION_QWEN_IMAGE_LAYERED,
VERSION_HUNYUAN_VIDEO,
VERSION_ANIMA,
VERSION_FLUX2,
VERSION_FLUX2_KLEIN,
VERSION_LTXAV,
VERSION_MINIMAX_H3,
VERSION_HIDREAM_O1,
VERSION_Z_IMAGE,
VERSION_BOOGU_IMAGE,
@ -53,6 +56,7 @@ enum SDVersion {
VERSION_IDEOGRAM4,
VERSION_SEFI_IMAGE,
VERSION_KREA2,
VERSION_MAGE_FLOW,
VERSION_ESRGAN,
VERSION_COUNT,
};
@ -120,6 +124,10 @@ static inline bool sd_version_is_ltxav(SDVersion version) {
return false;
}
static inline bool sd_version_is_minimax_h3(SDVersion version) {
return version == VERSION_MINIMAX_H3;
}
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;
@ -127,6 +135,13 @@ static inline bool sd_version_is_wan(SDVersion version) {
return false;
}
static inline bool sd_version_is_lingbot_video(SDVersion version) {
if (version == VERSION_LINGBOT_VIDEO) {
return true;
}
return false;
}
static inline bool sd_version_is_qwen_image(SDVersion version) {
if (version == VERSION_QWEN_IMAGE || version == VERSION_QWEN_IMAGE_LAYERED) {
return true;
@ -134,6 +149,13 @@ static inline bool sd_version_is_qwen_image(SDVersion version) {
return false;
}
static inline bool sd_version_is_hunyuan_video(SDVersion version) {
if (version == VERSION_HUNYUAN_VIDEO) {
return true;
}
return false;
}
static inline bool sd_version_is_anima(SDVersion version) {
if (version == VERSION_ANIMA) {
return true;
@ -211,6 +233,10 @@ static inline bool sd_version_is_krea2(SDVersion version) {
return false;
}
static inline bool sd_version_is_mage_flow(SDVersion version) {
return version == VERSION_MAGE_FLOW;
}
static inline bool sd_version_uses_flux_vae(SDVersion version) {
if (sd_version_is_flux(version) || sd_version_is_z_image(version) || sd_version_is_boogu_image(version) || sd_version_is_longcat(version)) {
return true;
@ -226,12 +252,16 @@ static inline bool sd_version_uses_flux2_vae(SDVersion version) {
}
static inline bool sd_version_uses_wan_vae(SDVersion version) {
if (sd_version_is_wan(version) || sd_version_is_qwen_image(version) || sd_version_is_krea2(version) || sd_version_is_anima(version)) {
if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_qwen_image(version) || sd_version_is_krea2(version) || sd_version_is_anima(version)) {
return true;
}
return false;
}
static inline bool sd_version_uses_hunyuan_video_vae(SDVersion version) {
return sd_version_is_hunyuan_video(version);
}
static inline bool sd_version_is_inpaint(SDVersion version) {
if (version == VERSION_SD1_INPAINT ||
version == VERSION_SD2_INPAINT ||
@ -247,9 +277,12 @@ 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_minimax_h3(version) ||
sd_version_is_sd3(version) ||
sd_version_is_wan(version) ||
sd_version_is_lingbot_video(version) ||
sd_version_is_qwen_image(version) ||
sd_version_is_hunyuan_video(version) ||
version == VERSION_HIDREAM_O1 ||
sd_version_is_anima(version) ||
sd_version_is_z_image(version) ||
@ -261,7 +294,8 @@ static inline bool sd_version_is_dit(SDVersion version) {
sd_version_is_pid(version) ||
sd_version_is_ideogram4(version) ||
sd_version_is_sefi_image(version) ||
sd_version_is_krea2(version)) {
sd_version_is_krea2(version) ||
sd_version_is_mage_flow(version)) {
return true;
}
return false;

View File

@ -0,0 +1,209 @@
#ifndef __SD_MODEL_ADAPTER_IP_ADAPTER_HPP__
#define __SD_MODEL_ADAPTER_IP_ADAPTER_HPP__
#include "core/ggml_extend.hpp"
#include "model/common/block.hpp"
#include "model_loader.h"
namespace IPAdapter {
struct ImageProjModel : public GGMLBlock {
int64_t num_tokens = 4;
int64_t ctx_dim = 768;
int64_t clip_dim = 1024;
ImageProjModel() {}
ImageProjModel(int64_t num_tokens, int64_t ctx_dim, int64_t clip_dim)
: num_tokens(num_tokens), ctx_dim(ctx_dim), clip_dim(clip_dim) {
blocks["proj"] = std::shared_ptr<GGMLBlock>(new Linear(clip_dim, num_tokens * ctx_dim, true));
blocks["norm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(ctx_dim));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* image_embeds) {
auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
auto norm = std::dynamic_pointer_cast<LayerNorm>(blocks["norm"]);
int64_t n = image_embeds->ne[1];
auto x = proj->forward(ctx, image_embeds);
x = ggml_reshape_3d(ctx->ggml_ctx, x, ctx_dim, num_tokens, n);
x = norm->forward(ctx, x);
return x;
}
};
struct Resampler : public GGMLBlock {
int64_t dim = 1280;
int64_t depth = 4;
int64_t num_queries = 16;
int64_t embed_dim = 1280;
int64_t output_dim = 2048;
int64_t ff_inner = 5120;
int64_t dim_head = 64;
int64_t heads = 20;
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
params["latents"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, dim, num_queries, 1);
}
Resampler() {}
Resampler(int64_t dim, int64_t depth, int64_t num_queries, int64_t embed_dim, int64_t output_dim, int64_t ff_inner)
: dim(dim), depth(depth), num_queries(num_queries), embed_dim(embed_dim), output_dim(output_dim), ff_inner(ff_inner) {
heads = dim / dim_head;
blocks["proj_in"] = std::shared_ptr<GGMLBlock>(new Linear(embed_dim, dim, true));
blocks["proj_out"] = std::shared_ptr<GGMLBlock>(new Linear(dim, output_dim, true));
blocks["norm_out"] = std::shared_ptr<GGMLBlock>(new LayerNorm(output_dim));
for (int64_t i = 0; i < depth; i++) {
std::string p = "layers." + std::to_string(i);
blocks[p + ".0.norm1"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
blocks[p + ".0.norm2"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
blocks[p + ".0.to_q"] = std::shared_ptr<GGMLBlock>(new Linear(dim, dim, false));
blocks[p + ".0.to_kv"] = std::shared_ptr<GGMLBlock>(new Linear(dim, dim * 2, false));
blocks[p + ".0.to_out"] = std::shared_ptr<GGMLBlock>(new Linear(dim, dim, false));
blocks[p + ".1.0"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
blocks[p + ".1.1"] = std::shared_ptr<GGMLBlock>(new Linear(dim, ff_inner, false));
blocks[p + ".1.3"] = std::shared_ptr<GGMLBlock>(new Linear(ff_inner, dim, false));
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* image_embeds) {
int64_t N = image_embeds->ne[2];
auto proj_in = std::dynamic_pointer_cast<Linear>(blocks["proj_in"]);
auto proj_out = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
auto norm_out = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_out"]);
ggml_tensor* x = proj_in->forward(ctx, image_embeds);
ggml_tensor* latents = params["latents"];
if (N > 1) {
latents = ggml_repeat(ctx->ggml_ctx, latents, ggml_new_tensor_3d(ctx->ggml_ctx, GGML_TYPE_F32, dim, num_queries, N));
}
for (int64_t i = 0; i < depth; i++) {
std::string p = "layers." + std::to_string(i);
auto norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks[p + ".0.norm1"]);
auto norm2 = std::dynamic_pointer_cast<LayerNorm>(blocks[p + ".0.norm2"]);
auto to_q = std::dynamic_pointer_cast<Linear>(blocks[p + ".0.to_q"]);
auto to_kv = std::dynamic_pointer_cast<Linear>(blocks[p + ".0.to_kv"]);
auto to_out = std::dynamic_pointer_cast<Linear>(blocks[p + ".0.to_out"]);
ggml_tensor* xn = norm1->forward(ctx, x);
ggml_tensor* ln = norm2->forward(ctx, latents);
ggml_tensor* q = to_q->forward(ctx, ln);
ggml_tensor* kv_in = ggml_concat(ctx->ggml_ctx, xn, ln, 1);
ggml_tensor* kv = to_kv->forward(ctx, kv_in);
int64_t L = kv->ne[1];
ggml_tensor* k = ggml_cont(ctx->ggml_ctx, ggml_view_3d(ctx->ggml_ctx, kv, dim, L, N, kv->nb[1], kv->nb[2], 0));
ggml_tensor* v = ggml_cont(ctx->ggml_ctx, ggml_view_3d(ctx->ggml_ctx, kv, dim, L, N, kv->nb[1], kv->nb[2], dim * kv->nb[0]));
ggml_tensor* attn = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, heads, nullptr, false, false);
attn = to_out->forward(ctx, attn);
latents = ggml_add(ctx->ggml_ctx, latents, attn);
auto ff_norm = std::dynamic_pointer_cast<LayerNorm>(blocks[p + ".1.0"]);
auto ff_fc1 = std::dynamic_pointer_cast<Linear>(blocks[p + ".1.1"]);
auto ff_fc2 = std::dynamic_pointer_cast<Linear>(blocks[p + ".1.3"]);
ggml_tensor* h = ff_norm->forward(ctx, latents);
h = ff_fc1->forward(ctx, h);
h = ggml_gelu_erf(ctx->ggml_ctx, h);
h = ff_fc2->forward(ctx, h);
latents = ggml_add(ctx->ggml_ctx, latents, h);
}
latents = proj_out->forward(ctx, latents);
latents = norm_out->forward(ctx, latents);
return latents;
}
};
struct IPAdapterRunner : public GGMLRunner {
ImageProjModel image_proj;
Resampler resampler;
bool is_plus = false;
int64_t num_tokens = 4;
std::string prefix;
IPAdapterRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map,
const std::string prefix,
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: GGMLRunner(backend, weight_manager), prefix(prefix) {
is_plus = tensor_storage_map.find(prefix + ".image_proj.latents") != tensor_storage_map.end();
if (is_plus) {
int64_t dim = 1280;
int64_t num_queries = 16;
int64_t embed_dim = 1280;
int64_t output_dim = 2048;
int64_t ff_inner = 5120;
auto latents_iter = tensor_storage_map.find(prefix + ".image_proj.latents");
if (latents_iter != tensor_storage_map.end()) {
dim = latents_iter->second.ne[0];
num_queries = latents_iter->second.ne[1];
}
auto proj_in_iter = tensor_storage_map.find(prefix + ".image_proj.proj_in.weight");
if (proj_in_iter != tensor_storage_map.end()) {
embed_dim = proj_in_iter->second.ne[0];
}
auto proj_out_iter = tensor_storage_map.find(prefix + ".image_proj.proj_out.weight");
if (proj_out_iter != tensor_storage_map.end()) {
output_dim = proj_out_iter->second.ne[1];
}
auto ff_iter = tensor_storage_map.find(prefix + ".image_proj.layers.0.1.1.weight");
if (ff_iter != tensor_storage_map.end()) {
ff_inner = ff_iter->second.ne[1];
}
int64_t depth = 0;
while (tensor_storage_map.find(prefix + ".image_proj.layers." + std::to_string(depth) + ".0.to_q.weight") != tensor_storage_map.end()) {
depth++;
}
num_tokens = num_queries;
resampler = Resampler(dim, depth, num_queries, embed_dim, output_dim, ff_inner);
resampler.init(params_ctx, tensor_storage_map, prefix + ".image_proj");
} else {
int64_t ctx_dim = 768;
int64_t clip_dim = 1024;
int64_t out_dim = 3072;
auto norm_iter = tensor_storage_map.find(prefix + ".image_proj.norm.weight");
if (norm_iter != tensor_storage_map.end()) {
ctx_dim = norm_iter->second.ne[0];
}
auto proj_iter = tensor_storage_map.find(prefix + ".image_proj.proj.weight");
if (proj_iter != tensor_storage_map.end()) {
clip_dim = proj_iter->second.ne[0];
out_dim = proj_iter->second.ne[1];
}
num_tokens = out_dim / ctx_dim;
image_proj = ImageProjModel(num_tokens, ctx_dim, clip_dim);
image_proj.init(params_ctx, tensor_storage_map, prefix + ".image_proj");
}
}
std::string get_desc() override {
return "ip_adapter";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string = "") {
if (is_plus) {
resampler.get_param_tensors(tensors, prefix + ".image_proj");
} else {
image_proj.get_param_tensors(tensors, prefix + ".image_proj");
}
}
ggml_cgraph* build_graph(const sd::Tensor<float>& image_embeds_tensor) {
ggml_cgraph* gf = new_graph_custom(1024);
ggml_tensor* embeds = make_input(image_embeds_tensor);
auto runner_ctx = get_context();
ggml_tensor* out = is_plus ? resampler.forward(&runner_ctx, embeds) : image_proj.forward(&runner_ctx, embeds);
ggml_build_forward_expand(gf, out);
return gf;
}
sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& image_embeds) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(image_embeds);
};
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true, true, true));
}
};
} // namespace IPAdapter
#endif // __SD_MODEL_ADAPTER_IP_ADAPTER_HPP__

View File

@ -14,6 +14,8 @@ struct LoraModel : public GGMLRunner {
std::unordered_map<std::string, ggml_tensor*> lora_tensors;
std::map<ggml_tensor*, ggml_tensor*> original_tensor_to_final_tensor;
std::set<std::string> applied_lora_tensors;
std::set<std::string> skipped_incompatible_lora_tensors;
std::set<std::string> warned_incompatible_model_tensors;
std::string file_path;
std::shared_ptr<ModelManager> model_manager;
ggml_backend_t params_backend = nullptr;
@ -133,6 +135,8 @@ struct LoraModel : public GGMLRunner {
lora_tensors.clear();
original_tensor_to_final_tensor.clear();
applied_lora_tensors.clear();
skipped_incompatible_lora_tensors.clear();
warned_incompatible_model_tensors.clear();
applied = false;
tensor_preprocessed = false;
}
@ -338,8 +342,10 @@ struct LoraModel : public GGMLRunner {
iter = lora_tensors.find(hada_1_mid_name);
if (iter != lora_tensors.end()) {
hada_1_mid = ggml_ext_cast_f32(ctx, backend, iter->second);
if (hada_1_up != nullptr) {
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()) {
@ -354,8 +360,10 @@ struct LoraModel : public GGMLRunner {
iter = lora_tensors.find(hada_2_mid_name);
if (iter != lora_tensors.end()) {
hada_2_mid = ggml_ext_cast_f32(ctx, backend, iter->second);
if (hada_2_up != nullptr) {
hada_2_up = ggml_cont(ctx, ggml_transpose(ctx, hada_2_up));
}
}
if (hada_1_up == nullptr || hada_1_down == nullptr || hada_2_up == nullptr || hada_2_down == nullptr) {
break;
@ -546,7 +554,27 @@ struct LoraModel : public GGMLRunner {
}
}
GGML_ASSERT(ggml_nelements(diff) == ggml_nelements(model_tensor));
if (ggml_nelements(diff) != ggml_nelements(model_tensor)) {
const std::string lora_tensor_prefix = "lora." + model_tensor_name + ".";
for (const auto& tensor_name : applied_lora_tensors) {
if (starts_with(tensor_name, lora_tensor_prefix)) {
skipped_incompatible_lora_tensors.insert(tensor_name);
}
}
if (warned_incompatible_model_tensors.insert(model_tensor_name).second) {
LOG_WARN("skip incompatible LoRA tensor |%s|: model shape = [%lld, %lld, %lld, %lld], LoRA shape = [%lld, %lld, %lld, %lld]",
model_tensor_name.c_str(),
static_cast<long long>(model_tensor->ne[0]),
static_cast<long long>(model_tensor->ne[1]),
static_cast<long long>(model_tensor->ne[2]),
static_cast<long long>(model_tensor->ne[3]),
static_cast<long long>(diff->ne[0]),
static_cast<long long>(diff->ne[1]),
static_cast<long long>(diff->ne[2]),
static_cast<long long>(diff->ne[3]));
}
return nullptr;
}
diff = ggml_reshape(ctx, diff, model_tensor);
}
return diff;
@ -555,10 +583,15 @@ struct LoraModel : public GGMLRunner {
ggml_tensor* get_out_diff(ggml_context* ctx,
ggml_backend_t backend,
ggml_tensor* x,
ggml_tensor* model_weight,
WeightAdapter::ForwardParams forward_params,
const std::string& model_tensor_name) {
ggml_tensor* out_diff = nullptr;
int index = 0;
std::vector<std::string> used_tensors;
bool is_conv2d = forward_params.op_type == WeightAdapter::ForwardParams::op_type_t::OP_CONV2D;
while (true) {
std::string key;
if (index == 0) {
@ -566,7 +599,6 @@ struct LoraModel : public GGMLRunner {
} else {
key = model_tensor_name + "." + std::to_string(index);
}
bool is_conv2d = forward_params.op_type == WeightAdapter::ForwardParams::op_type_t::OP_CONV2D;
std::string lokr_w1_name = "lora." + key + ".lokr_w1";
std::string lokr_w1_a_name = "lora." + key + ".lokr_w1_a";
@ -634,7 +666,6 @@ struct LoraModel : public GGMLRunner {
if (iter != lora_tensors.end()) {
float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
scale_value = alpha / rank;
applied_lora_tensors.insert(alpha_name);
}
if (rank == 1) {
@ -649,19 +680,27 @@ struct LoraModel : public GGMLRunner {
out_diff = ggml_concat(ctx, out_diff, curr_out_diff, 0);
}
if (lokr_w1)
applied_lora_tensors.insert(lokr_w1_name);
if (lokr_w1_a)
applied_lora_tensors.insert(lokr_w1_a_name);
if (lokr_w1_b)
applied_lora_tensors.insert(lokr_w1_b_name);
if (lokr_w2)
applied_lora_tensors.insert(lokr_w2_name);
if (lokr_w2_a)
applied_lora_tensors.insert(lokr_w2_a_name);
if (lokr_w2_b)
applied_lora_tensors.insert(lokr_w2_b_name);
applied_lora_tensors.insert(alpha_name);
if (lokr_w1) {
used_tensors.push_back(lokr_w1_name);
}
if (lokr_w1_a) {
used_tensors.push_back(lokr_w1_a_name);
}
if (lokr_w1_b) {
used_tensors.push_back(lokr_w1_b_name);
}
if (lokr_w2) {
used_tensors.push_back(lokr_w2_name);
}
if (lokr_w2_a) {
used_tensors.push_back(lokr_w2_a_name);
}
if (lokr_w2_b) {
used_tensors.push_back(lokr_w2_b_name);
}
if (iter != lora_tensors.end()) {
used_tensors.push_back(alpha_name);
}
index++;
continue;
@ -707,27 +746,60 @@ struct LoraModel : public GGMLRunner {
break;
}
applied_lora_tensors.insert(lora_up_name);
applied_lora_tensors.insert(lora_down_name);
if (!is_conv2d) {
const int64_t down_in = lora_down->ne[0];
const int64_t down_out = lora_down->ne[1];
const int64_t up_in = lora_up->ne[0];
if (lora_mid) {
applied_lora_tensors.insert(lora_mid_name);
bool compatible = down_in == model_weight->ne[0];
if (lora_mid != nullptr) {
compatible = compatible &&
lora_mid->ne[0] == down_out &&
up_in == lora_mid->ne[1];
} else {
compatible = compatible && up_in == down_out;
}
if (!compatible) {
skipped_incompatible_lora_tensors.insert(lora_down_name);
skipped_incompatible_lora_tensors.insert(lora_up_name);
if (lora_mid != nullptr) {
skipped_incompatible_lora_tensors.insert(lora_mid_name);
}
if (lora_tensors.find(scale_name) != lora_tensors.end()) {
skipped_incompatible_lora_tensors.insert(scale_name);
} else if (lora_tensors.find(alpha_name) != lora_tensors.end()) {
skipped_incompatible_lora_tensors.insert(alpha_name);
}
if (warned_incompatible_model_tensors.insert(model_tensor_name).second) {
LOG_WARN("skip incompatible LoRA tensor |%s|: model input dim = %lld, down shape = [%lld, %lld], up shape = [%lld, %lld]",
model_tensor_name.c_str(),
static_cast<long long>(model_weight->ne[0]),
static_cast<long long>(down_in),
static_cast<long long>(down_out),
static_cast<long long>(up_in),
static_cast<long long>(lora_up->ne[1]));
}
index++;
continue;
}
}
float scale_value = 1.0f;
std::string scale_tensor_name;
int64_t rank = lora_down->ne[ggml_n_dims(lora_down) - 1];
iter = lora_tensors.find(scale_name);
if (iter != lora_tensors.end()) {
scale_value = ggml_ext_backend_tensor_get_f32(iter->second);
applied_lora_tensors.insert(scale_name);
scale_tensor_name = scale_name;
} else {
iter = lora_tensors.find(alpha_name);
if (iter != lora_tensors.end()) {
float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
scale_value = alpha / rank;
scale_tensor_name = alpha_name;
// LOG_DEBUG("rank %s %ld %.2f %.2f", alpha_name.c_str(), rank, alpha, scale_value);
applied_lora_tensors.insert(alpha_name);
}
}
scale_value *= multiplier;
@ -787,15 +859,45 @@ struct LoraModel : public GGMLRunner {
}
auto curr_out_diff = ggml_ext_scale(ctx, lx, scale_value, true);
if (out_diff == nullptr) {
out_diff = curr_out_diff;
} else {
out_diff = ggml_concat(ctx, out_diff, curr_out_diff, 0);
out_diff = ggml_concat(ctx, out_diff, curr_out_diff, is_conv2d ? 2 : 0);
}
used_tensors.push_back(lora_up_name);
used_tensors.push_back(lora_down_name);
if (lora_mid) {
used_tensors.push_back(lora_mid_name);
}
if (!scale_tensor_name.empty()) {
used_tensors.push_back(scale_tensor_name);
}
index++;
}
if (out_diff == nullptr)
return nullptr;
int64_t expected_out_dim = is_conv2d ? model_weight->ne[3] : model_weight->ne[1];
int64_t actual_out_dim = out_diff->ne[is_conv2d ? 2 : 0];
if (actual_out_dim != expected_out_dim) {
for (const auto& name : used_tensors) {
skipped_incompatible_lora_tensors.insert(name);
}
if (warned_incompatible_model_tensors.insert(model_tensor_name).second) {
LOG_WARN("skip incompatible LoRA tensors for |%s|: output dim %lld != model dim %lld",
model_tensor_name.c_str(), actual_out_dim, expected_out_dim);
}
return nullptr;
}
for (const auto& name : used_tensors) {
applied_lora_tensors.insert(name);
}
return out_diff;
}
@ -869,10 +971,13 @@ struct LoraModel : public GGMLRunner {
void stat(bool at_runntime = false) {
size_t total_lora_tensors_count = 0;
size_t applied_lora_tensors_count = 0;
size_t skipped_lora_tensors_count = 0;
for (auto& kv : lora_tensors) {
total_lora_tensors_count++;
if (applied_lora_tensors.find(kv.first) == applied_lora_tensors.end()) {
if (skipped_incompatible_lora_tensors.find(kv.first) != skipped_incompatible_lora_tensors.end()) {
skipped_lora_tensors_count++;
} else if (applied_lora_tensors.find(kv.first) == applied_lora_tensors.end()) {
if (!at_runntime) {
LOG_WARN("unused lora tensor |%s|", kv.first.c_str());
print_ggml_tensor(kv.second, true);
@ -884,12 +989,17 @@ struct LoraModel : public GGMLRunner {
/* Don't worry if this message shows up twice in the logs per LoRA,
* this function is called once to calculate the required buffer size
* and then again to actually generate a graph to be used */
if (!at_runntime && applied_lora_tensors_count != total_lora_tensors_count) {
size_t compatible_lora_tensors_count = total_lora_tensors_count - skipped_lora_tensors_count;
if (!at_runntime && applied_lora_tensors_count != compatible_lora_tensors_count) {
LOG_WARN("Only (%lu / %lu) LoRA tensors have been applied, lora_file_path = %s",
applied_lora_tensors_count, total_lora_tensors_count, file_path.c_str());
applied_lora_tensors_count, compatible_lora_tensors_count, file_path.c_str());
} else {
LOG_INFO("(%lu / %lu) LoRA tensors have been applied, lora_file_path = %s",
applied_lora_tensors_count, total_lora_tensors_count, file_path.c_str());
applied_lora_tensors_count, compatible_lora_tensors_count, file_path.c_str());
}
if (skipped_lora_tensors_count > 0) {
LOG_WARN("(%lu / %lu) incompatible LoRA tensors have been skipped, lora_file_path = %s",
skipped_lora_tensors_count, total_lora_tensors_count, file_path.c_str());
}
}
};
@ -953,7 +1063,7 @@ public:
forward_params.conv2d.scale);
}
for (auto& lora_model : lora_models) {
ggml_tensor* out_diff = lora_model->get_out_diff(ctx, backend, x, forward_params, prefix + "weight");
ggml_tensor* out_diff = lora_model->get_out_diff(ctx, backend, x, w, forward_params, prefix + "weight");
if (out_diff == nullptr) {
continue;
}

View File

@ -311,16 +311,32 @@ protected:
int64_t n_head;
int64_t d_head;
bool xtra_dim = false;
bool enable_ip = false;
bool has_ip = false;
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
GGMLBlock::init_params(ctx, tensor_storage_map, prefix);
if (enable_ip &&
tensor_storage_map.find(prefix + "to_k_ip.weight") != tensor_storage_map.end()) {
has_ip = true;
int64_t inner_dim = d_head * n_head;
int64_t ip_dim = tensor_storage_map.at(prefix + "to_k_ip.weight").ne[0];
blocks["to_k_ip"] = std::shared_ptr<GGMLBlock>(new Linear(ip_dim, inner_dim, false));
blocks["to_v_ip"] = std::shared_ptr<GGMLBlock>(new Linear(ip_dim, inner_dim, false));
}
}
public:
CrossAttention(int64_t query_dim,
int64_t context_dim,
int64_t n_head,
int64_t d_head)
int64_t d_head,
bool enable_ip = false)
: n_head(n_head),
d_head(d_head),
query_dim(query_dim),
context_dim(context_dim) {
context_dim(context_dim),
enable_ip(enable_ip) {
int64_t inner_dim = d_head * n_head;
if (context_dim == 320 && d_head == 320) {
// LOG_DEBUG("CrossAttention: temp set dim to 1024 for sdxs_09");
@ -363,6 +379,15 @@ public:
}
x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, inner_dim]
if (has_ip && ctx->ip_context != nullptr && ctx->ip_scale != 0.0f) {
auto to_k_ip = std::dynamic_pointer_cast<Linear>(blocks["to_k_ip"]);
auto to_v_ip = std::dynamic_pointer_cast<Linear>(blocks["to_v_ip"]);
auto k_ip = to_k_ip->forward(ctx, ctx->ip_context);
auto v_ip = to_v_ip->forward(ctx, ctx->ip_context);
auto x_ip = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k_ip, v_ip, n_head, nullptr, false, ctx->flash_attn_enabled);
x = ggml_add(ctx->ggml_ctx, x, ggml_scale(ctx->ggml_ctx, x_ip, ctx->ip_scale));
}
x = to_out_0->forward(ctx, x); // [N, n_token, query_dim]
return x;
}
@ -387,7 +412,7 @@ public:
// inner_dim is always None or equal to dim
// gated_ff is always True
blocks["attn1"] = std::shared_ptr<GGMLBlock>(new CrossAttention(dim, dim, n_head, d_head));
blocks["attn2"] = std::shared_ptr<GGMLBlock>(new CrossAttention(dim, context_dim, n_head, d_head));
blocks["attn2"] = std::shared_ptr<GGMLBlock>(new CrossAttention(dim, context_dim, n_head, d_head, true));
blocks["ff"] = std::shared_ptr<GGMLBlock>(new FeedForward(dim, dim));
blocks["norm1"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
blocks["norm2"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
@ -450,7 +475,7 @@ protected:
int64_t context_dim = 768; // hidden_size, 1024 for VERSION_SD2
bool use_linear = false;
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") {
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
auto iter = tensor_storage_map.find(prefix + "proj_out.weight");
if (iter != tensor_storage_map.end()) {
int64_t inner_dim = n_head * d_head;

View File

@ -18,10 +18,6 @@ namespace Rope {
DECREASE,
};
__STATIC_INLINE__ RefIndexMode ref_index_mode_from_bool(bool increase_ref_index) {
return increase_ref_index ? RefIndexMode::INCREASE : RefIndexMode::FIXED;
}
template <class T>
__STATIC_INLINE__ std::vector<T> linspace(T start, T end, int num) {
std::vector<T> result(num);
@ -539,6 +535,33 @@ namespace Rope {
return vid_ids_repeated;
}
__STATIC_INLINE__ std::vector<std::vector<float>> gen_hunyuan_video_ids(int t,
int h,
int w,
int patch_t,
int patch_h,
int patch_w,
int bs,
int context_len) {
std::vector<std::vector<float>> txt_ids(bs * context_len, std::vector<float>(3, 0.0f));
auto img_ids = gen_vid_ids(t, h, w, patch_t, patch_h, patch_w, bs);
return concat_ids(txt_ids, img_ids, bs);
}
__STATIC_INLINE__ std::vector<float> gen_hunyuan_video_pe(int t,
int h,
int w,
int patch_t,
int patch_h,
int patch_w,
int bs,
int context_len,
float theta,
const std::vector<int>& axes_dim) {
auto ids = gen_hunyuan_video_ids(t, h, w, patch_t, patch_h, patch_w, bs, context_len);
return embed_nd(ids, bs, theta, axes_dim);
}
__STATIC_INLINE__ std::vector<std::vector<float>> gen_qwen_image_ids(int t,
int h,
int w,
@ -631,6 +654,43 @@ namespace Rope {
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
}
__STATIC_INLINE__ std::vector<float> gen_mage_flow_pe(int h,
int w,
int bs,
int context_len,
const std::vector<ggml_tensor*>& ref_latents,
int theta,
const std::vector<int>& axes_dim) {
const int axes_dim_num = static_cast<int>(axes_dim.size());
auto make_image_ids = [=](int image_h, int image_w, int image_index) {
std::vector<std::vector<float>> image_ids(static_cast<size_t>(bs) * image_h * image_w,
std::vector<float>(axes_dim_num, 0.f));
int h_start = -(image_h - image_h / 2);
int w_start = -(image_w - image_w / 2);
for (int b = 0; b < bs; ++b) {
for (int y = 0; y < image_h; ++y) {
for (int x = 0; x < image_w; ++x) {
auto& id = image_ids[static_cast<size_t>(b) * image_h * image_w + y * image_w + x];
id[0] = static_cast<float>(image_index);
id[1] = static_cast<float>(h_start + y);
id[2] = static_cast<float>(w_start + x);
}
}
}
return image_ids;
};
auto ids = gen_flux_txt_ids(bs, context_len, axes_dim_num, {});
auto img_ids = make_image_ids(h, w, 0);
ids = concat_ids(ids, img_ids, bs);
for (size_t i = 0; i < ref_latents.size(); ++i) {
auto ref_ids = make_image_ids(static_cast<int>(ref_latents[i]->ne[1]),
static_cast<int>(ref_latents[i]->ne[0]),
static_cast<int>(i + 1));
ids = concat_ids(ids, ref_ids, bs);
}
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim);
}
__STATIC_INLINE__ std::vector<std::vector<float>> gen_lens_ids(int h,
int w,
int bs,
@ -759,6 +819,40 @@ namespace Rope {
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim);
}
__STATIC_INLINE__ std::vector<std::vector<float>> gen_lingbot_video_ids(int t,
int h,
int w,
int pt,
int ph,
int pw,
int bs,
int context_len) {
auto vid_ids_repeated = gen_vid_ids(t, h, w, pt, ph, pw, bs, context_len + 1);
std::vector<std::vector<float>> txt_ids(bs * context_len, std::vector<float>(3, 0.0f));
for (int i = 0; i < bs; ++i) {
for (int j = 0; j < context_len; ++j) {
txt_ids[i * context_len + j][0] = static_cast<float>(j + 1);
}
}
return concat_ids(vid_ids_repeated, txt_ids, bs);
}
__STATIC_INLINE__ std::vector<float> gen_lingbot_video_pe(int t,
int h,
int w,
int pt,
int ph,
int pw,
int bs,
int context_len,
int theta,
const std::vector<int>& axes_dim) {
std::vector<std::vector<float>> ids = gen_lingbot_video_ids(t, h, w, pt, ph, pw, bs, context_len);
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim);
}
__STATIC_INLINE__ std::vector<std::vector<float>> gen_qwen2vl_ids(int grid_h,
int grid_w,
int merge_size,

362
src/model/detector/yolov8.h Normal file
View File

@ -0,0 +1,362 @@
#ifndef __SD_MODEL_DETECTOR_YOLOV8_H__
#define __SD_MODEL_DETECTOR_YOLOV8_H__
#include <algorithm>
#include <array>
#include <map>
#include <memory>
#include <string>
#include <vector>
#include "core/ggml_extend.hpp"
#include "core/util.h"
struct YOLOv8Config {
std::array<int, 23> out_channels{};
std::map<int, int> hidden_channels;
std::map<int, int> repeats;
int detect_box_channels = 0;
int detect_cls_channels = 0;
int reg_max = 0;
int num_classes = 0;
bool valid = false;
static YOLOv8Config detect_from_weights(const String2TensorStorage& tensor_storage_map,
const std::string& prefix = "") {
YOLOv8Config config;
auto full_name = [&](const std::string& name) {
return prefix.empty() ? name : prefix + "." + name;
};
auto find_weight = [&](const std::string& name) -> const TensorStorage* {
auto iter = tensor_storage_map.find(full_name(name));
return iter == tensor_storage_map.end() ? nullptr : &iter->second;
};
auto conv_out = [&](const std::string& name) -> int {
const TensorStorage* weight = find_weight(name);
return weight != nullptr && weight->n_dims == 4 ? static_cast<int>(weight->ne[3]) : 0;
};
for (int layer : {0, 1, 3, 5, 7, 16, 19}) {
config.out_channels[layer] = conv_out("model." + std::to_string(layer) + ".conv.weight");
}
for (int layer : {2, 4, 6, 8, 12, 15, 18, 21}) {
const std::string base = "model." + std::to_string(layer);
config.out_channels[layer] = conv_out(base + ".cv2.conv.weight");
config.hidden_channels[layer] = conv_out(base + ".cv1.conv.weight") / 2;
int repeat_count = 0;
while (find_weight(base + ".m." + std::to_string(repeat_count) + ".cv1.conv.weight") != nullptr) {
++repeat_count;
}
config.repeats[layer] = repeat_count;
}
config.out_channels[9] = conv_out("model.9.cv2.conv.weight");
config.detect_box_channels = conv_out("model.22.cv2.0.0.conv.weight");
config.detect_cls_channels = conv_out("model.22.cv3.0.0.conv.weight");
const int box_outputs = conv_out("model.22.cv2.0.2.weight");
config.num_classes = conv_out("model.22.cv3.0.2.weight");
config.reg_max = box_outputs / 4;
config.valid = config.out_channels[0] > 0 && config.out_channels[9] > 0 &&
config.out_channels[15] > 0 && config.out_channels[18] > 0 &&
config.out_channels[21] > 0 && config.detect_box_channels > 0 &&
config.detect_cls_channels > 0 && box_outputs > 0 && box_outputs % 4 == 0 &&
config.num_classes > 0;
for (int layer : {2, 4, 6, 8, 12, 15, 18, 21}) {
config.valid = config.valid && config.hidden_channels[layer] > 0 && config.repeats[layer] > 0;
}
if (config.valid) {
LOG_DEBUG("yolov8: classes=%d, reg_max=%d, p3=%d, p4=%d, p5=%d",
config.num_classes,
config.reg_max,
config.out_channels[15],
config.out_channels[18],
config.out_channels[21]);
}
return config;
}
};
class YOLOConv : public UnaryBlock {
int out_channels_ = 0;
public:
YOLOConv(int in_channels, int out_channels, int kernel, int stride = 1)
: out_channels_(out_channels) {
blocks["conv"] = std::shared_ptr<GGMLBlock>(new Conv2d(in_channels,
out_channels,
{kernel, kernel},
{stride, stride},
{kernel / 2, kernel / 2},
{1, 1},
true));
}
int out_channels() const {
return out_channels_;
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
return ggml_silu_inplace(ctx->ggml_ctx, conv->forward(ctx, x));
}
};
class YOLOBottleneck : public UnaryBlock {
bool shortcut_ = false;
public:
YOLOBottleneck(int channels, bool shortcut)
: shortcut_(shortcut) {
blocks["cv1"] = std::shared_ptr<GGMLBlock>(new YOLOConv(channels, channels, 3));
blocks["cv2"] = std::shared_ptr<GGMLBlock>(new YOLOConv(channels, channels, 3));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
auto cv1 = std::dynamic_pointer_cast<YOLOConv>(blocks["cv1"]);
auto cv2 = std::dynamic_pointer_cast<YOLOConv>(blocks["cv2"]);
auto out = cv2->forward(ctx, cv1->forward(ctx, x));
return shortcut_ ? ggml_add(ctx->ggml_ctx, x, out) : out;
}
};
class YOLOC2f : public UnaryBlock {
int hidden_channels_ = 0;
int repeats_ = 0;
public:
YOLOC2f(int in_channels,
int out_channels,
int hidden_channels,
int repeats,
bool shortcut)
: hidden_channels_(hidden_channels), repeats_(repeats) {
blocks["cv1"] = std::shared_ptr<GGMLBlock>(new YOLOConv(in_channels, hidden_channels * 2, 1));
blocks["cv2"] = std::shared_ptr<GGMLBlock>(new YOLOConv(hidden_channels * (2 + repeats), out_channels, 1));
for (int i = 0; i < repeats; ++i) {
blocks["m." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new YOLOBottleneck(hidden_channels, shortcut));
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
auto cv1 = std::dynamic_pointer_cast<YOLOConv>(blocks["cv1"]);
auto cv2 = std::dynamic_pointer_cast<YOLOConv>(blocks["cv2"]);
auto split = cv1->forward(ctx, x);
// split: [N, 2*C, H, W], ggml layout [W, H, 2*C, N].
auto y0 = ggml_view_4d(ctx->ggml_ctx,
split,
split->ne[0],
split->ne[1],
hidden_channels_,
split->ne[3],
split->nb[1],
split->nb[2],
split->nb[3],
0);
auto y1 = ggml_view_4d(ctx->ggml_ctx,
split,
split->ne[0],
split->ne[1],
hidden_channels_,
split->ne[3],
split->nb[1],
split->nb[2],
split->nb[3],
static_cast<size_t>(hidden_channels_) * split->nb[2]);
auto joined = ggml_concat(ctx->ggml_ctx, y0, y1, 2);
auto last = y1;
for (int i = 0; i < repeats_; ++i) {
auto block = std::dynamic_pointer_cast<YOLOBottleneck>(blocks["m." + std::to_string(i)]);
last = block->forward(ctx, last);
joined = ggml_concat(ctx->ggml_ctx, joined, last, 2);
}
return cv2->forward(ctx, joined);
}
};
class YOLOSPPF : public UnaryBlock {
public:
YOLOSPPF(int in_channels, int out_channels) {
blocks["cv1"] = std::shared_ptr<GGMLBlock>(new YOLOConv(in_channels, in_channels / 2, 1));
blocks["cv2"] = std::shared_ptr<GGMLBlock>(new YOLOConv(in_channels * 2, out_channels, 1));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
auto cv1 = std::dynamic_pointer_cast<YOLOConv>(blocks["cv1"]);
auto cv2 = std::dynamic_pointer_cast<YOLOConv>(blocks["cv2"]);
x = cv1->forward(ctx, x);
auto y1 = ggml_pool_2d(ctx->ggml_ctx, x, GGML_OP_POOL_MAX, 5, 5, 1, 1, 2, 2);
auto y2 = ggml_pool_2d(ctx->ggml_ctx, y1, GGML_OP_POOL_MAX, 5, 5, 1, 1, 2, 2);
auto y3 = ggml_pool_2d(ctx->ggml_ctx, y2, GGML_OP_POOL_MAX, 5, 5, 1, 1, 2, 2);
auto out = ggml_concat(ctx->ggml_ctx, x, y1, 2);
out = ggml_concat(ctx->ggml_ctx, out, y2, 2);
out = ggml_concat(ctx->ggml_ctx, out, y3, 2);
return cv2->forward(ctx, out);
}
};
class YOLODetect : public GGMLBlock {
int num_classes_ = 0;
int reg_max_ = 0;
public:
YOLODetect(const std::array<int, 3>& in_channels,
int box_channels,
int cls_channels,
int reg_max,
int num_classes)
: num_classes_(num_classes), reg_max_(reg_max) {
for (int i = 0; i < 3; ++i) {
const std::string box = "cv2." + std::to_string(i);
blocks[box + ".0"] = std::shared_ptr<GGMLBlock>(new YOLOConv(in_channels[i], box_channels, 3));
blocks[box + ".1"] = std::shared_ptr<GGMLBlock>(new YOLOConv(box_channels, box_channels, 3));
blocks[box + ".2"] = std::shared_ptr<GGMLBlock>(new Conv2d(box_channels, reg_max * 4, {1, 1}, {1, 1}, {0, 0}, {1, 1}, true));
const std::string cls = "cv3." + std::to_string(i);
blocks[cls + ".0"] = std::shared_ptr<GGMLBlock>(new YOLOConv(in_channels[i], cls_channels, 3));
blocks[cls + ".1"] = std::shared_ptr<GGMLBlock>(new YOLOConv(cls_channels, cls_channels, 3));
blocks[cls + ".2"] = std::shared_ptr<GGMLBlock>(new Conv2d(cls_channels, num_classes, {1, 1}, {1, 1}, {0, 0}, {1, 1}, true));
}
}
ggml_tensor* forward_scale(GGMLRunnerContext* ctx, ggml_tensor* x, int index) {
const std::string box = "cv2." + std::to_string(index);
auto box0 = std::dynamic_pointer_cast<YOLOConv>(blocks[box + ".0"]);
auto box1 = std::dynamic_pointer_cast<YOLOConv>(blocks[box + ".1"]);
auto box2 = std::dynamic_pointer_cast<Conv2d>(blocks[box + ".2"]);
const std::string cls = "cv3." + std::to_string(index);
auto cls0 = std::dynamic_pointer_cast<YOLOConv>(blocks[cls + ".0"]);
auto cls1 = std::dynamic_pointer_cast<YOLOConv>(blocks[cls + ".1"]);
auto cls2 = std::dynamic_pointer_cast<Conv2d>(blocks[cls + ".2"]);
auto boxes = box2->forward(ctx, box1->forward(ctx, box0->forward(ctx, x)));
auto classes = cls2->forward(ctx, cls1->forward(ctx, cls0->forward(ctx, x)));
return ggml_concat(ctx->ggml_ctx, boxes, classes, 2);
}
int output_channels() const {
return reg_max_ * 4 + num_classes_;
}
};
class YOLOv8Model : public GGMLBlock {
YOLOv8Config config_;
std::shared_ptr<YOLOC2f> make_c2f(int layer, int in_channels, bool shortcut) {
return std::make_shared<YOLOC2f>(in_channels,
config_.out_channels[layer],
config_.hidden_channels.at(layer),
config_.repeats.at(layer),
shortcut);
}
public:
explicit YOLOv8Model(YOLOv8Config config)
: config_(std::move(config)) {
blocks["model.0"] = std::make_shared<YOLOConv>(3, config_.out_channels[0], 3, 2);
blocks["model.1"] = std::make_shared<YOLOConv>(config_.out_channels[0], config_.out_channels[1], 3, 2);
blocks["model.2"] = make_c2f(2, config_.out_channels[1], true);
blocks["model.3"] = std::make_shared<YOLOConv>(config_.out_channels[2], config_.out_channels[3], 3, 2);
blocks["model.4"] = make_c2f(4, config_.out_channels[3], true);
blocks["model.5"] = std::make_shared<YOLOConv>(config_.out_channels[4], config_.out_channels[5], 3, 2);
blocks["model.6"] = make_c2f(6, config_.out_channels[5], true);
blocks["model.7"] = std::make_shared<YOLOConv>(config_.out_channels[6], config_.out_channels[7], 3, 2);
blocks["model.8"] = make_c2f(8, config_.out_channels[7], true);
blocks["model.9"] = std::make_shared<YOLOSPPF>(config_.out_channels[8], config_.out_channels[9]);
blocks["model.12"] = make_c2f(12, config_.out_channels[9] + config_.out_channels[6], false);
blocks["model.15"] = make_c2f(15, config_.out_channels[12] + config_.out_channels[4], false);
blocks["model.16"] = std::make_shared<YOLOConv>(config_.out_channels[15], config_.out_channels[16], 3, 2);
blocks["model.18"] = make_c2f(18, config_.out_channels[16] + config_.out_channels[12], false);
blocks["model.19"] = std::make_shared<YOLOConv>(config_.out_channels[18], config_.out_channels[19], 3, 2);
blocks["model.21"] = make_c2f(21, config_.out_channels[19] + config_.out_channels[9], false);
blocks["model.22"] = std::make_shared<YOLODetect>(
std::array<int, 3>{config_.out_channels[15], config_.out_channels[18], config_.out_channels[21]},
config_.detect_box_channels,
config_.detect_cls_channels,
config_.reg_max,
config_.num_classes);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto run = [&](int layer, ggml_tensor* input) {
return std::dynamic_pointer_cast<UnaryBlock>(blocks["model." + std::to_string(layer)])->forward(ctx, input);
};
auto x0 = run(0, x);
auto x1 = run(1, x0);
auto x2 = run(2, x1);
auto x3 = run(3, x2);
auto x4 = run(4, x3);
auto x5 = run(5, x4);
auto x6 = run(6, x5);
auto x7 = run(7, x6);
auto x8 = run(8, x7);
auto x9 = run(9, x8);
auto x12 = run(12, ggml_concat(ctx->ggml_ctx, ggml_upscale(ctx->ggml_ctx, x9, 2, GGML_SCALE_MODE_NEAREST), x6, 2));
auto x15 = run(15, ggml_concat(ctx->ggml_ctx, ggml_upscale(ctx->ggml_ctx, x12, 2, GGML_SCALE_MODE_NEAREST), x4, 2));
auto x16 = run(16, x15);
auto x18 = run(18, ggml_concat(ctx->ggml_ctx, x16, x12, 2));
auto x19 = run(19, x18);
auto x21 = run(21, ggml_concat(ctx->ggml_ctx, x19, x9, 2));
auto detect = std::dynamic_pointer_cast<YOLODetect>(blocks["model.22"]);
auto p3 = detect->forward_scale(ctx, x15, 0);
auto p4 = detect->forward_scale(ctx, x18, 1);
auto p5 = detect->forward_scale(ctx, x21, 2);
p3 = ggml_reshape_2d(ctx->ggml_ctx, p3, p3->ne[0] * p3->ne[1], detect->output_channels());
p4 = ggml_reshape_2d(ctx->ggml_ctx, p4, p4->ne[0] * p4->ne[1], detect->output_channels());
p5 = ggml_reshape_2d(ctx->ggml_ctx, p5, p5->ne[0] * p5->ne[1], detect->output_channels());
return ggml_concat(ctx->ggml_ctx, ggml_concat(ctx->ggml_ctx, p3, p4, 0), p5, 0);
}
};
struct YOLOv8Runner : public GGMLRunner {
YOLOv8Config config;
std::unique_ptr<YOLOv8Model> model;
YOLOv8Runner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map,
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: GGMLRunner(backend, weight_manager),
config(YOLOv8Config::detect_from_weights(tensor_storage_map)) {
if (config.valid) {
model = std::make_unique<YOLOv8Model>(config);
model->init(params_ctx, tensor_storage_map, "");
}
}
std::string get_desc() override {
return "yolov8";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
if (model) {
model->get_param_tensors(tensors);
}
}
ggml_cgraph* build_graph(const sd::Tensor<float>& input) {
if (!model) {
return nullptr;
}
ggml_cgraph* graph = new_graph_custom(1 << 16);
auto x = make_input(input);
auto runner_ctx = get_context();
auto output = model->forward(&runner_ctx, x);
ggml_build_forward_expand(graph, output);
return graph;
}
sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& input) {
auto get_graph = [&]() { return build_graph(input); };
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false));
}
};
#endif // __SD_MODEL_DETECTOR_YOLOV8_H__

View File

@ -487,7 +487,8 @@ namespace Anima {
ggml_tensor* t5_ids = nullptr,
ggml_tensor* t5_weights = nullptr,
ggml_tensor* adapter_q_pe = nullptr,
ggml_tensor* adapter_k_pe = nullptr) {
ggml_tensor* adapter_k_pe = nullptr,
std::vector<ggml_tensor*> ref_latents = {}) {
GGML_ASSERT(x->ne[3] == 1);
auto x_embedder = std::dynamic_pointer_cast<XEmbedder>(blocks["x_embedder"]);
@ -503,7 +504,15 @@ namespace Anima {
x = ggml_concat(ctx->ggml_ctx, x, padding_mask, 2); // [N, C + 1, H, W]
x = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size); // [N, h*w, (C+1)*ph*pw]
int64_t img_len = x->ne[1];
if (ref_latents.size() > 0) {
for (ggml_tensor* ref : ref_latents) {
auto padding_mask = ggml_ext_zeros(ctx->ggml_ctx, ref->ne[0], ref->ne[1], 1, ref->ne[3]);
ref = ggml_concat(ctx->ggml_ctx, ref, padding_mask, 2); // [N, C + 1, H, W]
ref = DiT::pad_and_patchify(ctx, ref, config.patch_size, config.patch_size);
x = ggml_concat(ctx->ggml_ctx, x, ref, 1);
}
}
x = x_embedder->forward(ctx, x);
auto timestep_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast<int>(config.hidden_size));
@ -543,6 +552,7 @@ namespace Anima {
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 = ggml_ext_slice(ctx->ggml_ctx, x, 1, 0, img_len);
x = final_layer->forward(ctx, x, embedded_timestep, temb); // [N, h*w, ph*pw*C]
@ -602,8 +612,8 @@ namespace Anima {
const std::vector<int>& axes_dim,
float h_extrapolation_ratio,
float w_extrapolation_ratio,
float t_extrapolation_ratio) {
static const std::vector<ggml_tensor*> empty_ref_latents;
float t_extrapolation_ratio,
const std::vector<ggml_tensor*>& ref_latents) {
auto ids = Rope::gen_flux_ids(h,
w,
patch_size,
@ -611,7 +621,7 @@ namespace Anima {
static_cast<int>(axes_dim.size()),
0,
{},
empty_ref_latents,
ref_latents,
Rope::RefIndexMode::FIXED,
1.0f,
false);
@ -628,12 +638,18 @@ namespace Anima {
const sd::Tensor<float>& timesteps_tensor,
const sd::Tensor<float>& context_tensor = {},
const sd::Tensor<int32_t>& t5_ids_tensor = {},
const sd::Tensor<float>& t5_weights_tensor = {}) {
const sd::Tensor<float>& t5_weights_tensor = {},
const std::vector<sd::Tensor<float>>& ref_latents_tensor = {}) {
ggml_tensor* x = make_input(x_tensor);
ggml_tensor* timesteps = make_input(timesteps_tensor);
ggml_tensor* context = make_optional_input(context_tensor);
ggml_tensor* t5_ids = make_optional_input(t5_ids_tensor);
ggml_tensor* t5_weights = make_optional_input(t5_weights_tensor);
std::vector<ggml_tensor*> ref_latents;
ref_latents.reserve(ref_latents_tensor.size());
for (const auto& ref_latent_tensor : ref_latents_tensor) {
ref_latents.push_back(make_input(ref_latent_tensor));
}
GGML_ASSERT(x->ne[3] == 1);
ggml_cgraph* gf = new_graph_custom(ANIMA_GRAPH_SIZE);
@ -650,7 +666,8 @@ namespace Anima {
config.axes_dim,
4.0f,
4.0f,
1.0f);
1.0f,
ref_latents);
int64_t image_pos_len = static_cast<int64_t>(image_pe_vec.size()) / (2 * 2 * (config.head_dim / 2));
auto image_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, image_pos_len);
set_backend_tensor_data(image_pe, image_pe_vec.data());
@ -682,7 +699,8 @@ namespace Anima {
t5_ids,
t5_weights,
adapter_q_pe,
adapter_k_pe);
adapter_k_pe,
ref_latents);
ggml_build_forward_expand(gf, out);
return gf;
@ -693,9 +711,11 @@ namespace Anima {
const sd::Tensor<float>& timesteps,
const sd::Tensor<float>& context = {},
const sd::Tensor<int32_t>& t5_ids = {},
const sd::Tensor<float>& t5_weights = {}) {
const sd::Tensor<float>& t5_weights = {},
const std::vector<sd::Tensor<float>>& ref_latents = {},
const RefImageParams& ref_image_params = REF_IMAGE_PRESETS.at("cosmos_reference")) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, t5_ids, t5_weights);
return build_graph(x, timesteps, context, t5_ids, t5_weights, ref_latents);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
}
@ -705,12 +725,15 @@ namespace Anima {
GGML_ASSERT(diffusion_params.x != nullptr);
GGML_ASSERT(diffusion_params.timesteps != nullptr);
const auto* extra = diffusion_extra_as<AnimaDiffusionExtra>(diffusion_params);
static const std::vector<sd::Tensor<float>> empty_ref_latents;
return compute(n_threads,
*diffusion_params.x,
*diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context),
tensor_or_empty(extra->t5_ids),
tensor_or_empty(extra->t5_weights));
tensor_or_empty(extra->t5_weights),
diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents,
diffusion_params.ref_image_params);
}
};
} // namespace Anima

View File

@ -0,0 +1,182 @@
#ifndef __SD_MODEL_DIFFUSION_ANIMATEDIFF_HPP__
#define __SD_MODEL_DIFFUSION_ANIMATEDIFF_HPP__
#include "core/ggml_extend.hpp"
#include "model/common/block.hpp"
// AnimateDiff (https://arxiv.org/abs/2307.04725) SD 1.5 motion modules.
namespace AnimateDiff {
struct MotionModuleConfig {
int max_frames = 32;
int64_t num_heads = 8;
int norm_num_groups = 32;
std::vector<int64_t> down_channels = {320, 640, 1280, 1280};
std::vector<int64_t> up_channels = {1280, 1280, 640, 320};
int num_down_motion_per_block = 2;
int num_up_motion_per_block = 3;
bool enable_mid_block = false;
int64_t mid_channels = 1280;
};
class TemporalAttention : public GGMLBlock {
protected:
int64_t channels;
int64_t num_heads;
int max_frames;
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
params["pos_encoder.pe"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, channels, max_frames, 1);
}
public:
TemporalAttention(int64_t channels, int64_t num_heads, int max_frames)
: channels(channels), num_heads(num_heads), max_frames(max_frames) {
blocks["to_q"] = std::shared_ptr<GGMLBlock>(new Linear(channels, channels, false));
blocks["to_k"] = std::shared_ptr<GGMLBlock>(new Linear(channels, channels, false));
blocks["to_v"] = std::shared_ptr<GGMLBlock>(new Linear(channels, channels, false));
blocks["to_out.0"] = std::shared_ptr<GGMLBlock>(new Linear(channels, channels, true));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto to_q = std::dynamic_pointer_cast<Linear>(blocks["to_q"]);
auto to_k = std::dynamic_pointer_cast<Linear>(blocks["to_k"]);
auto to_v = std::dynamic_pointer_cast<Linear>(blocks["to_v"]);
auto to_out = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
int64_t C = x->ne[0];
int64_t F = x->ne[1];
auto pe = params["pos_encoder.pe"];
auto pe_f = (F == pe->ne[1])
? pe
: ggml_view_3d(ctx->ggml_ctx, pe, C, F, 1, pe->nb[1], pe->nb[2], 0);
auto x_pe = ggml_add(ctx->ggml_ctx, x, ggml_repeat(ctx->ggml_ctx, pe_f, x));
auto q = to_q->forward(ctx, x_pe);
auto k = to_k->forward(ctx, x_pe);
auto v = to_v->forward(ctx, x_pe);
auto a = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, (int)num_heads, nullptr, false);
return to_out->forward(ctx, a);
}
};
class TemporalTransformerBlock : public GGMLBlock {
public:
TemporalTransformerBlock(int64_t channels, int64_t num_heads, int max_frames) {
blocks["attention_blocks.0"] = std::make_shared<TemporalAttention>(channels, num_heads, max_frames);
blocks["attention_blocks.1"] = std::make_shared<TemporalAttention>(channels, num_heads, max_frames);
blocks["norms.0"] = std::shared_ptr<GGMLBlock>(new LayerNorm(channels));
blocks["norms.1"] = std::shared_ptr<GGMLBlock>(new LayerNorm(channels));
blocks["ff"] = std::make_shared<FeedForward>(channels, channels, 4, FeedForward::Activation::GEGLU);
blocks["ff_norm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(channels));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto attn0 = std::dynamic_pointer_cast<TemporalAttention>(blocks["attention_blocks.0"]);
auto attn1 = std::dynamic_pointer_cast<TemporalAttention>(blocks["attention_blocks.1"]);
auto norm0 = std::dynamic_pointer_cast<LayerNorm>(blocks["norms.0"]);
auto norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["norms.1"]);
auto ff = std::dynamic_pointer_cast<FeedForward>(blocks["ff"]);
auto ff_norm = std::dynamic_pointer_cast<LayerNorm>(blocks["ff_norm"]);
auto r = x;
x = ggml_add(ctx->ggml_ctx, attn0->forward(ctx, norm0->forward(ctx, x)), r);
r = x;
x = ggml_add(ctx->ggml_ctx, attn1->forward(ctx, norm1->forward(ctx, x)), r);
r = x;
x = ggml_add(ctx->ggml_ctx, ff->forward(ctx, ff_norm->forward(ctx, x)), r);
return x;
}
};
class TemporalTransformer : public GGMLBlock {
public:
TemporalTransformer(int64_t channels, int64_t num_heads, int norm_num_groups, int max_frames) {
blocks["norm"] = std::shared_ptr<GGMLBlock>(new GroupNorm(norm_num_groups, channels));
blocks["proj_in"] = std::shared_ptr<GGMLBlock>(new Linear(channels, channels, true));
blocks["transformer_blocks.0"] = std::make_shared<TemporalTransformerBlock>(channels, num_heads, max_frames);
blocks["proj_out"] = std::shared_ptr<GGMLBlock>(new Linear(channels, channels, true));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, int64_t num_frames) {
auto norm = std::dynamic_pointer_cast<GroupNorm>(blocks["norm"]);
auto proj_in = std::dynamic_pointer_cast<Linear>(blocks["proj_in"]);
auto tb0 = std::dynamic_pointer_cast<TemporalTransformerBlock>(blocks["transformer_blocks.0"]);
auto proj_out = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t C = x->ne[2];
GGML_ASSERT(x->ne[3] == num_frames);
auto residual = x;
auto h = norm->forward(ctx, x);
h = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, h, 2, 3, 0, 1));
h = ggml_reshape_3d(ctx->ggml_ctx, h, C, num_frames, W * H);
h = proj_in->forward(ctx, h);
h = tb0->forward(ctx, h);
h = proj_out->forward(ctx, h);
h = ggml_reshape_4d(ctx->ggml_ctx, h, C, num_frames, W, H);
h = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, h, 2, 3, 0, 1));
return ggml_add(ctx->ggml_ctx, h, residual);
}
};
class MotionModule : public GGMLBlock {
public:
MotionModule(int64_t channels, int64_t num_heads, int norm_num_groups, int max_frames) {
blocks["temporal_transformer"] = std::make_shared<TemporalTransformer>(channels, num_heads, norm_num_groups, max_frames);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, int64_t num_frames) {
auto tt = std::dynamic_pointer_cast<TemporalTransformer>(blocks["temporal_transformer"]);
return tt->forward(ctx, x, num_frames);
}
};
class AnimateDiffModel : public GGMLBlock {
public:
MotionModuleConfig config;
AnimateDiffModel(const MotionModuleConfig& cfg)
: config(cfg) {
for (int i = 0; i < static_cast<int>(cfg.down_channels.size()); ++i) {
int64_t ch = cfg.down_channels[i];
for (int j = 0; j < cfg.num_down_motion_per_block; ++j) {
blocks["down_blocks." + std::to_string(i) + ".motion_modules." + std::to_string(j)] =
std::make_shared<MotionModule>(ch, cfg.num_heads, cfg.norm_num_groups, cfg.max_frames);
}
}
for (int i = 0; i < static_cast<int>(cfg.up_channels.size()); ++i) {
int64_t ch = cfg.up_channels[i];
for (int j = 0; j < cfg.num_up_motion_per_block; ++j) {
blocks["up_blocks." + std::to_string(i) + ".motion_modules." + std::to_string(j)] =
std::make_shared<MotionModule>(ch, cfg.num_heads, cfg.norm_num_groups, cfg.max_frames);
}
}
if (cfg.enable_mid_block) {
blocks["mid_block.motion_modules.0"] =
std::make_shared<MotionModule>(cfg.mid_channels, cfg.num_heads, cfg.norm_num_groups, cfg.max_frames);
}
}
std::shared_ptr<MotionModule> motion(const std::string& key) {
auto it = blocks.find(key);
if (it == blocks.end())
return nullptr;
return std::dynamic_pointer_cast<MotionModule>(it->second);
}
};
} // namespace AnimateDiff
#endif // __SD_MODEL_DIFFUSION_ANIMATEDIFF_HPP__

View File

@ -827,7 +827,7 @@ namespace Boogu {
*diffusion_params.x,
*diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context),
diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents);
diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents);
}
};
} // namespace Boogu

View File

@ -104,14 +104,15 @@ namespace DiT {
return x;
}
inline ggml_tensor* patchify(ggml_context* ctx,
inline ggml_tensor* patchify_3d(ggml_context* ctx,
ggml_tensor* x,
int pt,
int ph,
int pw,
int64_t N = 1) {
int64_t N = 1,
bool patch_last = true) {
// x: [N*C, T, H, W]
// return: [N, h*w, C*pt*ph*pw]
// return: [N, t_len*h_len*w_len, C*pt*ph*pw] if patch_last else [N, t_len*h_len*w_len, C*pt*ph*pw] or [N, t_len*h_len*w_len, pt*ph*pw*C]
int64_t C = x->ne[3] / N;
int64_t T = x->ne[2];
int64_t H = x->ne[1];
@ -130,8 +131,13 @@ namespace DiT {
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]
if (patch_last) {
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]
} else {
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3)); // [N, t_len*h_len*w_len, pt*ph*pw, C]
}
// [N, t_len*h_len*w_len, C*pt*ph*pw] or [N, t_len*h_len*w_len, pt*ph*pw*C]
x = ggml_reshape_4d(ctx, x, pw * ph * pt * C, w_len * h_len * t_len, N, 1);
return x;
}
@ -142,16 +148,23 @@ namespace DiT {
int64_t w_len,
int pt,
int ph,
int pw) {
// x: [N, t_len*h_len*w_len, C*pt*ph*pw]
int pw,
bool patch_last = true) {
// x: [N, t_len*h_len*w_len, C*pt*ph*pw] if patch_last else [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[2];
int64_t C = x->ne[0] / pt / ph / pw;
GGML_ASSERT(C * pt * ph * pw == x->ne[0]);
if (patch_last) {
x = ggml_reshape_4d(ctx, x, pw * ph * pt, C, w_len * h_len * t_len, N); // [N, t_len*h_len*w_len, C, pt*ph*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]
} else {
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]

View File

@ -707,10 +707,12 @@ namespace Flux {
int64_t patch_size,
int64_t out_channels,
bool prune_mod = false,
bool bias = true)
bool bias = true,
int64_t patch_volume = 0)
: prune_mod(prune_mod) {
blocks["norm_final"] = std::shared_ptr<GGMLBlock>(new LayerNorm(hidden_size, 1e-06f, false));
blocks["linear"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, patch_size * patch_size * out_channels, bias));
int64_t out_dim = (patch_volume > 0 ? patch_volume : patch_size * patch_size) * out_channels;
blocks["linear"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, out_dim, bias));
if (!prune_mod) {
blocks["adaLN_modulation.1"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, 2 * hidden_size, bias));
}
@ -1642,8 +1644,8 @@ namespace Flux {
tensor_or_empty(diffusion_params.c_concat),
tensor_or_empty(diffusion_params.y),
tensor_or_empty(extra->guidance),
diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
diffusion_params.ref_index_mode,
diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents,
diffusion_params.ref_image_params.ref_index_mode,
extra->skip_layers ? *extra->skip_layers : empty_skip_layers,
tensor_or_empty(extra->pulid_id),
extra->pulid_id_weight);

View File

@ -0,0 +1,681 @@
#ifndef __SD_MODEL_DIFFUSION_HUNYUAN_HPP__
#define __SD_MODEL_DIFFUSION_HUNYUAN_HPP__
#include <memory>
#include "model/common/block.hpp"
#include "model/diffusion/flux.hpp"
#include "model/diffusion/mmdit.hpp"
#include "model/diffusion/wan.hpp"
#include "model_manager.h"
namespace Hunyuan {
constexpr int HUNYUAN_VIDEO_GRAPH_SIZE = 65536;
// Ref: https://github.com/huggingface/diffusers/pull/12696
struct IndividualTokenRefinerBlock : public GGMLBlock {
protected:
int64_t num_heads;
public:
IndividualTokenRefinerBlock(int64_t num_heads,
int64_t head_dim,
int64_t mlp_ratio = 4,
bool attn_bias = true)
: num_heads(num_heads) {
int64_t hidden_size = num_heads * head_dim;
blocks["self_attn.qkv"] = std::make_shared<Linear>(hidden_size, hidden_size * 3, attn_bias);
blocks["self_attn.proj"] = std::make_shared<Linear>(hidden_size, hidden_size, attn_bias);
blocks["norm1"] = std::make_shared<LayerNorm>(hidden_size, 1e-6f, true);
blocks["norm2"] = std::make_shared<LayerNorm>(hidden_size, 1e-6f, true);
blocks["mlp.0"] = std::make_shared<Linear>(hidden_size, hidden_size * mlp_ratio);
blocks["mlp.2"] = std::make_shared<Linear>(hidden_size * mlp_ratio, hidden_size);
// adaLN_modulation.0 is nn.SiLU()
blocks["adaLN_modulation.1"] = std::make_shared<Linear>(hidden_size, hidden_size * 2);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* txt, ggml_tensor* t_emb, ggml_tensor* mask) {
auto norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
auto norm2 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm2"]);
auto self_attn_qkv = std::dynamic_pointer_cast<Linear>(blocks["self_attn.qkv"]);
auto self_attn_proj = std::dynamic_pointer_cast<Linear>(blocks["self_attn.proj"]);
auto mlp_fc1 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
auto mlp_fc2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
auto adaLN_modulation_1 = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.1"]);
// self attn
auto qkv = self_attn_qkv->forward(ctx, norm1->forward(ctx, txt));
auto qkv_vec = split_qkv(ctx->ggml_ctx, qkv);
auto q = qkv_vec[0];
auto k = qkv_vec[1];
auto v = qkv_vec[2];
auto attn_out = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, mask, false, ctx->flash_attn_enabled);
attn_out = self_attn_proj->forward(ctx, attn_out);
// adaLN_modulation
auto emb = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, t_emb));
auto mods = ggml_ext_chunk(ctx->ggml_ctx, emb, 2, 0);
txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, attn_out, mods[0]));
// mlp
auto mlp_out = mlp_fc1->forward(ctx, norm2->forward(ctx, txt));
mlp_out = ggml_silu_inplace(ctx->ggml_ctx, mlp_out);
mlp_out = mlp_fc2->forward(ctx, mlp_out);
txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, mlp_out, mods[1]));
return txt;
}
};
struct IndividualTokenRefiner : public GGMLBlock {
protected:
int num_layers;
public:
IndividualTokenRefiner(int64_t num_heads,
int64_t head_dim,
int num_layers,
int64_t mlp_ratio = 4,
bool attn_bias = true)
: num_layers(num_layers) {
for (int i = 0; i < num_layers; i++) {
blocks["blocks." + std::to_string(i)] = std::make_shared<IndividualTokenRefinerBlock>(num_heads, head_dim, mlp_ratio, attn_bias);
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* txt, ggml_tensor* t_emb, ggml_tensor* mask) {
for (int i = 0; i < num_layers; i++) {
auto block = std::dynamic_pointer_cast<IndividualTokenRefinerBlock>(blocks["blocks." + std::to_string(i)]);
txt = block->forward(ctx, txt, t_emb, mask);
}
return txt;
}
};
struct TokenRefiner : public GGMLBlock {
public:
TokenRefiner(int64_t in_channels,
int64_t num_heads,
int64_t head_dim,
int num_layers,
int64_t mlp_ratio = 4,
bool attn_bias = true) {
int64_t hidden_size = num_heads * head_dim;
blocks["input_embedder"] = std::make_shared<Linear>(in_channels, hidden_size);
blocks["t_embedder"] = std::make_shared<Flux::MLPEmbedder>(256, hidden_size);
blocks["c_embedder"] = std::make_shared<Flux::MLPEmbedder>(in_channels, hidden_size);
blocks["individual_token_refiner"] = std::make_shared<IndividualTokenRefiner>(num_heads, head_dim, num_layers, mlp_ratio, attn_bias);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* txt, ggml_tensor* timestep, ggml_tensor* mask) {
auto input_embedder = std::dynamic_pointer_cast<Linear>(blocks["input_embedder"]);
auto t_embedder = std::dynamic_pointer_cast<Flux::MLPEmbedder>(blocks["t_embedder"]);
auto c_embedder = std::dynamic_pointer_cast<Flux::MLPEmbedder>(blocks["c_embedder"]);
auto individual_token_refiner = std::dynamic_pointer_cast<IndividualTokenRefiner>(blocks["individual_token_refiner"]);
auto t_emb = t_embedder->forward(ctx, ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, 256, 10000, 1.f));
auto h = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, txt, 1, 0, 2, 3));
auto pooled_projections = ggml_scale(ctx->ggml_ctx, ggml_sum_rows(ctx->ggml_ctx, h), 1.f / txt->ne[1]);
pooled_projections = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, pooled_projections, 1, 0, 2, 3));
auto c_emb = c_embedder->forward(ctx, pooled_projections);
t_emb = ggml_add(ctx->ggml_ctx, t_emb, c_emb);
txt = input_embedder->forward(ctx, txt);
txt = individual_token_refiner->forward(ctx, txt, t_emb, mask);
return txt;
}
};
struct ByT5Mapper : public UnaryBlock {
ByT5Mapper(int64_t in_dim, int64_t hidden_size) {
blocks["layernorm"] = std::make_shared<LayerNorm>(in_dim);
blocks["fc1"] = std::make_shared<Linear>(in_dim, 2048);
blocks["fc2"] = std::make_shared<Linear>(2048, 2048);
blocks["fc3"] = std::make_shared<Linear>(2048, hidden_size);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
auto layernorm = std::dynamic_pointer_cast<LayerNorm>(blocks["layernorm"]);
auto fc1 = std::dynamic_pointer_cast<Linear>(blocks["fc1"]);
auto fc2 = std::dynamic_pointer_cast<Linear>(blocks["fc2"]);
auto fc3 = std::dynamic_pointer_cast<Linear>(blocks["fc3"]);
x = fc1->forward(ctx, layernorm->forward(ctx, x));
x = ggml_ext_gelu(ctx->ggml_ctx, x);
x = fc2->forward(ctx, x);
x = ggml_ext_gelu(ctx->ggml_ctx, x);
return fc3->forward(ctx, x);
}
};
struct HunyuanVideoConfig {
std::tuple<int, int, int> patch_size = {1, 2, 2};
int64_t in_channels = 65;
int64_t out_channels = 32;
int64_t hidden_size = 2048;
int64_t vec_in_dim = 0;
int64_t context_in_dim = 3584;
int64_t vision_in_dim = 0;
float mlp_ratio = 4.0f;
int num_heads = 16;
int depth = 54;
int depth_single_blocks = 0;
bool qkv_bias = true;
bool guidance_embed = false;
bool use_byt5 = false;
bool use_cond_type_embedding = false;
bool use_meanflow = false;
bool use_meanflow_sum = false;
float theta = 256;
std::vector<int> axes_dim = {16, 56, 56};
int axes_dim_sum = 128;
int64_t patch_volume() const {
return static_cast<int64_t>(std::get<0>(patch_size)) * std::get<1>(patch_size) * std::get<2>(patch_size);
}
static HunyuanVideoConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
const std::string& prefix) {
HunyuanVideoConfig config;
config.depth = 0;
config.depth_single_blocks = 0;
bool inferred = false;
int64_t img_embed_dim = 0;
for (const auto& [name, storage] : tensor_storage_map) {
if (starts_with(name, prefix) && ends_with(name, "img_in.proj.bias")) {
img_embed_dim = storage.ne[0];
break;
}
}
for (const auto& entry : tensor_storage_map) {
const auto& name = entry.first;
const auto& storage = entry.second;
if (!starts_with(name, prefix)) {
continue;
}
auto update_depth = [&](const char* block_prefix, int* depth) {
size_t pos = name.find(block_prefix);
if (pos == std::string::npos) {
return;
}
pos += strlen(block_prefix);
size_t end = name.find('.', pos);
if (end != std::string::npos) {
*depth = std::max(*depth, atoi(name.substr(pos, end - pos).c_str()) + 1);
}
};
update_depth("double_blocks.", &config.depth);
update_depth("single_blocks.", &config.depth_single_blocks);
if (ends_with(name, "img_in.proj.weight") && storage.n_dims == 5) {
config.patch_size = {static_cast<int>(storage.ne[2]),
static_cast<int>(storage.ne[1]),
static_cast<int>(storage.ne[0])};
config.in_channels = storage.ne[3];
config.hidden_size = storage.ne[4];
inferred = true;
} else if (ends_with(name, "img_in.proj.weight") && storage.n_dims == 4) {
config.patch_size = {static_cast<int>(storage.ne[2]),
static_cast<int>(storage.ne[1]),
static_cast<int>(storage.ne[0])};
if (img_embed_dim > 0 && storage.ne[3] % img_embed_dim == 0) {
config.hidden_size = img_embed_dim;
config.in_channels = storage.ne[3] / img_embed_dim;
}
inferred = true;
} else if (ends_with(name, "txt_in.input_embedder.weight")) {
config.context_in_dim = storage.ne[0];
inferred = true;
} else if (ends_with(name, "vector_in.in_layer.weight")) {
config.vec_in_dim = storage.ne[0];
} else if (ends_with(name, "vision_in.proj.0.weight")) {
config.vision_in_dim = storage.ne[0];
} else if (ends_with(name, "double_blocks.0.img_attn.norm.key_norm.scale") ||
ends_with(name, "double_blocks.0.img_attn.norm.key_norm.weight")) {
config.num_heads = static_cast<int>(config.hidden_size / storage.ne[0]);
} else if (ends_with(name, "double_blocks.0.img_mlp.0.weight")) {
config.mlp_ratio = static_cast<float>(storage.ne[1]) / static_cast<float>(storage.ne[0]);
}
config.guidance_embed = config.guidance_embed || name.find("guidance_in.") != std::string::npos;
config.use_byt5 = config.use_byt5 || name.find("byt5_in.") != std::string::npos;
config.use_meanflow = config.use_meanflow || name.find("time_r_in.") != std::string::npos;
}
config.use_cond_type_embedding = tensor_storage_map.find(prefix + ".cond_type_embedding.weight") != tensor_storage_map.end();
config.use_meanflow_sum = config.vision_in_dim > 0;
auto final_iter = tensor_storage_map.find(prefix + ".final_layer.linear.weight");
if (final_iter != tensor_storage_map.end()) {
config.out_channels = final_iter->second.ne[1] / config.patch_volume();
}
config.qkv_bias = tensor_storage_map.find(prefix + ".double_blocks.0.img_attn.qkv.bias") != tensor_storage_map.end();
GGML_ASSERT(config.hidden_size % config.num_heads == 0);
GGML_ASSERT(config.hidden_size / config.num_heads == config.axes_dim_sum);
if (inferred) {
LOG_DEBUG("hunyuan video: depth = %d, single depth = %d, in_channels = %" PRId64 ", out_channels = %" PRId64 ", hidden_size = %" PRId64 ", context_in_dim = %" PRId64 ", patch_size = %dx%dx%d",
config.depth,
config.depth_single_blocks,
config.in_channels,
config.out_channels,
config.hidden_size,
config.context_in_dim,
std::get<0>(config.patch_size),
std::get<1>(config.patch_size),
std::get<2>(config.patch_size));
}
return config;
}
};
class HunyuanVideoModel : public GGMLBlock {
protected:
HunyuanVideoConfig config;
void init_params(struct ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
if (config.use_cond_type_embedding) {
ggml_type type = get_type(prefix + "cond_type_embedding.weight", tensor_storage_map, GGML_TYPE_F16);
GGMLBlock::params["cond_type_embedding.weight"] = ggml_new_tensor_2d(ctx, type, config.hidden_size, 3);
}
}
public:
HunyuanVideoModel() {}
explicit HunyuanVideoModel(HunyuanVideoConfig config)
: config(std::move(config)) {
int64_t head_dim = this->config.hidden_size / this->config.num_heads;
blocks["txt_in"] = std::make_shared<TokenRefiner>(this->config.context_in_dim, this->config.num_heads, head_dim, 2);
blocks["img_in"] = std::make_shared<PatchEmbed>(static_cast<int64_t>(224) /*Not used*/,
this->config.patch_size,
this->config.in_channels,
this->config.hidden_size);
blocks["time_in"] = std::make_shared<Flux::MLPEmbedder>(256, this->config.hidden_size);
if (this->config.vec_in_dim > 0) {
blocks["vector_in"] = std::make_shared<Flux::MLPEmbedder>(this->config.vec_in_dim, this->config.hidden_size);
}
if (this->config.vision_in_dim > 0) {
blocks["vision_in"] = std::make_shared<WAN::MLPProj>(this->config.vision_in_dim, this->config.hidden_size);
}
if (this->config.guidance_embed) {
blocks["guidance_in"] = std::make_shared<Flux::MLPEmbedder>(256, this->config.hidden_size);
}
if (this->config.use_byt5) {
blocks["byt5_in"] = std::make_shared<ByT5Mapper>(1472, this->config.hidden_size);
}
if (this->config.use_meanflow) {
blocks["time_r_in"] = std::make_shared<Flux::MLPEmbedder>(256, this->config.hidden_size);
}
for (int i = 0; i < this->config.depth; i++) {
blocks["double_blocks." + std::to_string(i)] = std::make_shared<Flux::DoubleStreamBlock>(this->config.hidden_size,
this->config.num_heads,
this->config.mlp_ratio,
i,
this->config.qkv_bias);
}
for (int i = 0; i < this->config.depth_single_blocks; i++) {
blocks["single_blocks." + std::to_string(i)] = std::make_shared<Flux::SingleStreamBlock>(this->config.hidden_size,
this->config.num_heads,
this->config.mlp_ratio,
i,
0.f);
}
blocks["final_layer"] = std::make_shared<Flux::LastLayer>(this->config.hidden_size,
std::get<2>(this->config.patch_size),
this->config.out_channels,
false,
true,
this->config.patch_volume());
}
ggml_tensor* pad_to_patch_size(struct ggml_context* ctx,
ggml_tensor* x) {
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t T = x->ne[2];
int pt = std::get<0>(config.patch_size);
int ph = std::get<1>(config.patch_size);
int pw = std::get<2>(config.patch_size);
int pad_t = (pt - static_cast<int>(T % pt)) % pt;
int pad_h = (ph - static_cast<int>(H % ph)) % ph;
int pad_w = (pw - static_cast<int>(W % pw)) % pw;
x = ggml_pad(ctx, x, pad_w, pad_h, pad_t, 0); // [N*C, T + pad_t, H + pad_h, W + pad_w]
return x;
}
ggml_tensor* unpatchify(struct ggml_context* ctx,
ggml_tensor* x,
int64_t t_len,
int64_t h_len,
int64_t w_len) {
// x: [N, t_len*h_len*w_len, C*pt*ph*pw]
// return: [N*C, t_len*pt, h_len*ph, w_len*pw]
int64_t N = x->ne[3];
int64_t pt = std::get<0>(config.patch_size);
int64_t ph = std::get<1>(config.patch_size);
int64_t pw = std::get<2>(config.patch_size);
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;
}
ggml_tensor* add_condition_type(GGMLRunnerContext* ctx, ggml_tensor* x, int type) {
if (!config.use_cond_type_embedding) {
return x;
}
auto weight = GGMLBlock::params["cond_type_embedding.weight"];
auto row = ggml_view_1d(ctx->ggml_ctx,
weight,
weight->ne[0],
static_cast<size_t>(type) * weight->nb[1]);
auto target = ggml_new_tensor_3d(ctx->ggml_ctx, row->type, config.hidden_size, x->ne[1], x->ne[2]);
auto embed = ggml_repeat(ctx->ggml_ctx, row, target);
embed = ggml_cast(ctx->ggml_ctx, embed, x->type);
return ggml_add(ctx->ggml_ctx, x, embed);
}
ggml_tensor* forward_orig(GGMLRunnerContext* ctx,
ggml_tensor* img,
ggml_tensor* txt,
ggml_tensor* timestep,
ggml_tensor* pe,
ggml_tensor* guidance = nullptr,
ggml_tensor* y = nullptr,
ggml_tensor* txt_byt5 = nullptr,
ggml_tensor* clip_fea = nullptr,
ggml_tensor* timestep_r = nullptr,
int64_t N = 1) {
// img: [N*C, T, H, W], C => in_dim
// txt: [N, L, text_dim]
// timestep: [N,] or [T]
// return: [N, t_len*h_len*w_len, out_dim*pt*ph*pw]
GGML_ASSERT(N == 1);
auto img_in = std::dynamic_pointer_cast<PatchEmbed>(blocks["img_in"]);
auto txt_in = std::dynamic_pointer_cast<TokenRefiner>(blocks["txt_in"]);
auto time_in = std::dynamic_pointer_cast<Flux::MLPEmbedder>(blocks["time_in"]);
auto final_layer = std::dynamic_pointer_cast<Flux::LastLayer>(blocks["final_layer"]);
img = img_in->forward(ctx, img); // [N*C, t_len*h_len*w_len, hidden_size]
txt = txt_in->forward(ctx, txt, timestep, nullptr); // [N, n_txt_token, hidden_size]
auto vec = time_in->forward(ctx, ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, 256, 10000, 1.f));
if (config.use_meanflow && timestep_r != nullptr) {
auto time_r_in = std::dynamic_pointer_cast<Flux::MLPEmbedder>(blocks["time_r_in"]);
auto vec_r = time_r_in->forward(ctx, ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_r, 256, 10000, 1000.f));
vec = ggml_add(ctx->ggml_ctx, vec, vec_r);
if (!config.use_meanflow_sum) {
vec = ggml_scale(ctx->ggml_ctx, vec, 0.5f);
}
}
if (config.vec_in_dim > 0 && y != nullptr) {
auto vector_in = std::dynamic_pointer_cast<Flux::MLPEmbedder>(blocks["vector_in"]);
vec = ggml_add(ctx->ggml_ctx, vec, vector_in->forward(ctx, y));
}
if (config.guidance_embed && guidance != nullptr) {
auto guidance_in = std::dynamic_pointer_cast<Flux::MLPEmbedder>(blocks["guidance_in"]);
auto guidance_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, guidance, 256, 10000, 1.f);
vec = ggml_add(ctx->ggml_ctx, vec, guidance_in->forward(ctx, guidance_emb));
}
txt = add_condition_type(ctx, txt, 0);
if (config.use_byt5 && txt_byt5 != nullptr) {
auto byt5_in = std::dynamic_pointer_cast<ByT5Mapper>(blocks["byt5_in"]);
txt_byt5 = add_condition_type(ctx, byt5_in->forward(ctx, txt_byt5), 1);
txt = config.use_cond_type_embedding ? ggml_concat(ctx->ggml_ctx, txt_byt5, txt, 1)
: ggml_concat(ctx->ggml_ctx, txt, txt_byt5, 1);
}
if (config.vision_in_dim > 0 && clip_fea != nullptr) {
auto vision_in = std::dynamic_pointer_cast<WAN::MLPProj>(blocks["vision_in"]);
clip_fea = add_condition_type(ctx, vision_in->forward(ctx, clip_fea), 2);
txt = ggml_concat(ctx->ggml_ctx, clip_fea, txt, 1);
}
for (int i = 0; i < config.depth; i++) {
auto block = std::dynamic_pointer_cast<Flux::DoubleStreamBlock>(blocks["double_blocks." + std::to_string(i)]);
auto img_txt = block->forward(ctx, img, txt, vec, pe, nullptr);
img = img_txt.first; // [N, n_img_token, hidden_size]
txt = img_txt.second; // [N, n_txt_token, hidden_size]
}
if (config.depth_single_blocks > 0) {
auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1); // [N, n_txt_token + n_img_token, hidden_size]
for (int i = 0; i < config.depth_single_blocks; i++) {
auto block = std::dynamic_pointer_cast<Flux::SingleStreamBlock>(blocks["single_blocks." + std::to_string(i)]);
txt_img = block->forward(ctx, txt_img, vec, pe, nullptr);
}
txt_img = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, txt_img, 0, 2, 1, 3));
img = ggml_view_3d(ctx->ggml_ctx,
txt_img,
txt_img->ne[0],
txt_img->ne[1],
img->ne[1],
txt_img->nb[1],
txt_img->nb[2],
txt_img->nb[2] * txt->ne[1]);
img = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, img, 0, 2, 1, 3));
}
img = final_layer->forward(ctx, img, vec); // (N, t_len*h_len*w_len, out_channels * patch_size ** 3)
return img;
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* timestep,
ggml_tensor* context,
ggml_tensor* pe,
ggml_tensor* guidance = nullptr,
ggml_tensor* y = nullptr,
ggml_tensor* txt_byt5 = nullptr,
ggml_tensor* clip_fea = nullptr,
ggml_tensor* timestep_r = nullptr,
int64_t N = 1) {
// Forward pass of DiT.
// x: [N*C, T, H, W]
// timestep: [N,]
// context: [N, L, D]
// pe: [L, d_head/2, 2, 2]
// return: [N*C, T, H, W]
GGML_ASSERT(N == 1);
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t T = x->ne[2];
x = pad_to_patch_size(ctx->ggml_ctx, x);
int64_t pt = std::get<0>(config.patch_size);
int64_t ph = std::get<1>(config.patch_size);
int64_t pw = std::get<2>(config.patch_size);
int64_t t_len = (T + pt - 1) / pt;
int64_t h_len = (H + ph - 1) / ph;
int64_t w_len = (W + pw - 1) / pw;
auto out = forward_orig(ctx, x, context, timestep, pe, guidance, y, txt_byt5, clip_fea, timestep_r, N);
out = unpatchify(ctx->ggml_ctx, out, t_len, h_len, w_len); // [N*C, (T+pad_t) + (T2+pad_t2), H + pad_h, W + pad_w]
// slice
out = ggml_ext_slice(ctx->ggml_ctx, out, 2, 0, T); // [N*C, T, H + pad_h, W + pad_w]
out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, H); // [N*C, T, H, W + pad_w]
out = ggml_ext_slice(ctx->ggml_ctx, out, 0, 0, W); // [N*C, T, H, W]
return out;
}
};
struct HunyuanVideoRunner : public DiffusionModelRunner {
public:
HunyuanVideoConfig config;
HunyuanVideoModel hunyuan_video;
std::vector<float> pe_vec;
SDVersion version;
HunyuanVideoRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "",
SDVersion version = VERSION_HUNYUAN_VIDEO,
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: DiffusionModelRunner(backend, prefix, weight_manager),
config(HunyuanVideoConfig::detect_from_weights(tensor_storage_map, prefix)),
version(version) {
LOG_INFO("HunyuanVideo blocks: %d double, %d single", config.depth, config.depth_single_blocks);
hunyuan_video = HunyuanVideoModel(config);
hunyuan_video.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "hunyuan_video";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
hunyuan_video.get_param_tensors(tensors, prefix);
}
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
const sd::Tensor<float>& timesteps_tensor,
const sd::Tensor<float>& context_tensor,
const sd::Tensor<float>& c_concat_tensor = {},
const sd::Tensor<float>& y_tensor = {},
const sd::Tensor<float>& guidance_tensor = {},
const sd::Tensor<float>& byt5_tensor = {},
const sd::Tensor<float>& vision_tensor = {},
const sd::Tensor<float>& timestep_r_tensor = {}) {
ggml_cgraph* gf = new_graph_custom(HUNYUAN_VIDEO_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor);
ggml_tensor* timesteps = make_input(timesteps_tensor);
ggml_tensor* context = make_input(context_tensor);
ggml_tensor* c_concat = make_optional_input(c_concat_tensor);
ggml_tensor* y = make_optional_input(y_tensor);
ggml_tensor* guidance = make_optional_input(guidance_tensor);
ggml_tensor* byt5 = make_optional_input(byt5_tensor);
ggml_tensor* vision = make_optional_input(vision_tensor);
ggml_tensor* timestep_r = make_optional_input(timestep_r_tensor);
GGML_ASSERT(x->ne[3] == config.out_channels);
if (c_concat != nullptr) {
x = ggml_concat(compute_ctx, x, c_concat, 3);
}
GGML_ASSERT(x->ne[3] <= config.in_channels);
if (x->ne[3] < config.in_channels) {
x = ggml_pad(compute_ctx, x, 0, 0, 0, static_cast<int>(config.in_channels - x->ne[3]));
}
int text_len = static_cast<int>(context->ne[1]);
if (byt5 != nullptr) {
text_len += static_cast<int>(byt5->ne[1]);
}
if (vision != nullptr) {
text_len += static_cast<int>(vision->ne[1]);
}
pe_vec = Rope::gen_hunyuan_video_pe(static_cast<int>(x->ne[2]),
static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
std::get<0>(config.patch_size),
std::get<1>(config.patch_size),
std::get<2>(config.patch_size),
1,
text_len,
config.theta,
config.axes_dim);
int64_t pos_len = static_cast<int64_t>(pe_vec.size() / config.axes_dim_sum / 2);
// LOG_DEBUG("pos_len %d", pos_len);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
// pe->data = pe_vec.data();
// print_ggml_tensor(pe, true, "pe");
// pe->data = nullptr;
set_backend_tensor_data(pe, pe_vec.data());
auto runner_ctx = get_context();
ggml_tensor* out = hunyuan_video.forward(&runner_ctx,
x,
timesteps,
context,
pe,
guidance,
y,
byt5,
vision,
timestep_r);
ggml_build_forward_expand(gf, out);
return gf;
}
sd::Tensor<float> compute(int n_threads,
const sd::Tensor<float>& x,
const sd::Tensor<float>& timesteps,
const sd::Tensor<float>& context,
const sd::Tensor<float>& c_concat = {},
const sd::Tensor<float>& y = {},
const sd::Tensor<float>& guidance = {},
const sd::Tensor<float>& byt5 = {},
const sd::Tensor<float>& vision = {},
const sd::Tensor<float>& timestep_r = {}) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, c_concat, y, guidance, byt5, vision, timestep_r);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
const DiffusionParams& diffusion_params) override {
GGML_ASSERT(diffusion_params.x != nullptr);
GGML_ASSERT(diffusion_params.timesteps != nullptr);
GGML_ASSERT(diffusion_params.context != nullptr);
const auto* extra = diffusion_extra_as<HunyuanVideoDiffusionExtra>(diffusion_params);
return compute(n_threads,
*diffusion_params.x,
*diffusion_params.timesteps,
*diffusion_params.context,
tensor_or_empty(diffusion_params.c_concat),
tensor_or_empty(diffusion_params.y),
tensor_or_empty(extra->guidance),
tensor_or_empty(extra->byt5),
tensor_or_empty(extra->vision),
tensor_or_empty(extra->timestep_r));
}
};
} // namespace Hunyuan
#endif // __SD_MODEL_DIFFUSION_HUNYUAN_HPP__

View File

@ -180,6 +180,9 @@ namespace Krea2 {
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
ggml_tensor* scale = params["scale"];
if (ctx->weight_adapter) {
scale = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, scale, prefix + "scale.weight");
}
scale = ggml_add(ctx->ggml_ctx, scale, ggml_ext_ones(ctx->ggml_ctx, scale->ne[0], 1, 1, 1));
x = ggml_rms_norm(ctx->ggml_ctx, x, eps);
x = ggml_mul_inplace(ctx->ggml_ctx, x, scale);
@ -295,10 +298,11 @@ namespace Krea2 {
class KreaDoubleSharedModulation : public GGMLBlock {
protected:
int64_t dim;
std::string prefix;
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
GGML_UNUSED(tensor_storage_map);
GGML_UNUSED(prefix);
this->prefix = prefix;
params["lin"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dim * 6);
}
@ -307,7 +311,11 @@ namespace Krea2 {
: dim(dim) {}
std::vector<ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* vec) {
auto lin = ggml_repeat(ctx->ggml_ctx, params["lin"], vec);
auto lin = params["lin"];
if (ctx->weight_adapter) {
lin = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, lin, prefix + "lin.weight");
}
lin = ggml_repeat(ctx->ggml_ctx, lin, vec);
auto out = ggml_add(ctx->ggml_ctx, vec, lin);
return ggml_ext_chunk(ctx->ggml_ctx, out, 6, 0);
}
@ -316,10 +324,11 @@ namespace Krea2 {
class KreaFinalModulation : public GGMLBlock {
protected:
int64_t dim;
std::string prefix;
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
GGML_UNUSED(tensor_storage_map);
GGML_UNUSED(prefix);
this->prefix = prefix;
params["lin"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, dim, 2);
}
@ -328,7 +337,11 @@ namespace Krea2 {
: dim(dim) {}
std::vector<ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* vec) {
auto out = ggml_add(ctx->ggml_ctx, params["lin"], vec);
auto lin = params["lin"];
if (ctx->weight_adapter) {
lin = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, lin, prefix + "lin.weight");
}
auto out = ggml_add(ctx->ggml_ctx, lin, vec);
return ggml_ext_chunk(ctx->ggml_ctx, out, 2, 1);
}
};
@ -615,7 +628,8 @@ namespace Krea2 {
ggml_tensor* timestep,
ggml_tensor* context,
ggml_tensor* pe,
std::vector<ggml_tensor*> ref_latents = {}) {
std::vector<ggml_tensor*> ref_latents = {},
bool zero_timestep_refs = false) {
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t N = x->ne[3];
@ -645,7 +659,7 @@ namespace Krea2 {
auto tvec = tproj->forward(ctx, t);
ggml_tensor* tvec_0 = nullptr;
if (ref_latents.size() > 0) {
if (ref_latents.size() > 0 && zero_timestep_refs) {
// "index_timestep_zero" mode: use timestep = 0 for ref latents
auto timestep_0 = ggml_scale(ctx->ggml_ctx, timestep, 0.0f);
auto t_0 = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_0, static_cast<int>(config.timestep_dim), 10000, 1000.f);
@ -719,7 +733,7 @@ namespace Krea2 {
const sd::Tensor<float>& timesteps_tensor,
const sd::Tensor<float>& context_tensor,
const std::vector<sd::Tensor<float>>& ref_latents_tensor = {},
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED) {
const RefImageParams& ref_image_params = REF_IMAGE_PRESETS.at("krea2_ostris_edit")) {
ggml_cgraph* gf = new_graph_custom(KREA2_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor);
ggml_tensor* timesteps = make_input(timesteps_tensor);
@ -741,13 +755,13 @@ namespace Krea2 {
config.theta,
config.axes_dim,
ref_latents,
ref_index_mode);
ref_image_params.ref_index_mode);
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
set_backend_tensor_data(pe, pe_vec.data());
auto runner_ctx = get_context();
ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pe, ref_latents);
ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pe, ref_latents, ref_image_params.force_ref_timestep_zero);
ggml_build_forward_expand(gf, out);
return gf;
}
@ -757,9 +771,9 @@ namespace Krea2 {
const sd::Tensor<float>& timesteps,
const sd::Tensor<float>& context,
const std::vector<sd::Tensor<float>>& ref_latents = {},
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED) {
const RefImageParams& ref_image_params = REF_IMAGE_PRESETS.at("krea2_ostris_edit")) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, ref_latents, ref_index_mode);
return build_graph(x, timesteps, context, ref_latents, ref_image_params);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
}
@ -773,8 +787,8 @@ namespace Krea2 {
*diffusion_params.x,
*diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context),
diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
diffusion_params.ref_index_mode);
diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents,
diffusion_params.ref_image_params);
}
};
} // namespace Krea2

View File

@ -0,0 +1,692 @@
#ifndef __SD_MODEL_DIFFUSION_LINGBOT_VIDEO_HPP__
#define __SD_MODEL_DIFFUSION_LINGBOT_VIDEO_HPP__
#include <algorithm>
#include <cinttypes>
#include <cmath>
#include <cstdlib>
#include <memory>
#include <set>
#include <string>
#include <vector>
#include "core/util.h"
#include "model/common/rope.hpp"
#include "model/diffusion/dit.hpp"
#include "model/diffusion/flux.hpp"
#include "model/diffusion/model.hpp"
#include "model/diffusion/qwen_image.hpp"
namespace LingBotVideo {
constexpr int LINGBOT_VIDEO_GRAPH_SIZE = 65536;
struct LingBotVideoConfig {
int patch_t = 1;
int patch_h = 2;
int patch_w = 2;
int64_t in_channels = 16;
int64_t out_channels = 16;
int64_t hidden_size = 2048;
int64_t num_attention_heads = 16;
int64_t depth = 24;
int64_t intermediate_size = 6144;
int64_t text_dim = 2560;
int64_t freq_dim = 256;
float norm_eps = 1e-6f;
int rope_theta = 256;
std::vector<int> axes_dim = {32, 48, 48};
int axes_dim_sum = 128;
bool qkv_bias = false;
bool out_bias = true;
bool patch_embed_bias = true;
bool timestep_mlp_bias = true;
int64_t num_experts = 0;
int64_t num_experts_per_tok = 8;
int64_t moe_intermediate_size = 512;
int64_t decoder_sparse_step = 1;
int64_t n_shared_experts = 0;
bool norm_topk_prob = true;
float routed_scaling_factor = 1.0f;
int64_t n_group = 0;
int64_t topk_group = 0;
std::set<int> sparse_layers;
static LingBotVideoConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
const std::string& prefix) {
LingBotVideoConfig config;
config.depth = 0;
for (const auto& [name, tensor_storage] : tensor_storage_map) {
if (!starts_with(name, prefix)) {
continue;
}
if (ends_with(name, "patch_embedder.weight") && tensor_storage.n_dims == 2) {
int64_t patch_dim = tensor_storage.ne[0];
config.hidden_size = tensor_storage.ne[1];
int64_t patch_volume = config.patch_t * config.patch_h * config.patch_w;
if (patch_dim % patch_volume == 0) {
config.in_channels = patch_dim / patch_volume;
}
} else if (ends_with(name, "text_embedder.linear_1.weight") && tensor_storage.n_dims == 2) {
config.text_dim = tensor_storage.ne[0];
} else if (ends_with(name, "time_embedder.linear_1.weight") && tensor_storage.n_dims == 2) {
config.freq_dim = tensor_storage.ne[0];
} else if (ends_with(name, "blocks.0.attn.norm_q.weight") && tensor_storage.n_dims == 1) {
int64_t head_dim = tensor_storage.ne[0];
if (head_dim > 0) {
config.num_attention_heads = config.hidden_size / head_dim;
}
} else if (name.find(".attn.to_q.bias") != std::string::npos) {
config.qkv_bias = true;
} else if (name.find(".ffn.gate_proj.weight") != std::string::npos && tensor_storage.n_dims == 2) {
config.intermediate_size = tensor_storage.ne[1];
} else if (name.find(".ffn.experts.w1") != std::string::npos && tensor_storage.n_dims == 3) {
config.num_experts = tensor_storage.ne[2];
config.moe_intermediate_size = tensor_storage.ne[1];
} else if (name.find(".ffn.shared_experts.gate_proj.weight") != std::string::npos && tensor_storage.n_dims == 2) {
if (config.moe_intermediate_size > 0) {
config.n_shared_experts = tensor_storage.ne[1] / config.moe_intermediate_size;
}
} else if (ends_with(name, "proj_out.weight") && tensor_storage.n_dims == 2) {
int64_t out_dim = tensor_storage.ne[1];
int64_t patch_volume = config.patch_t * config.patch_h * config.patch_w;
config.out_channels = patch_volume > 0 ? out_dim / patch_volume : config.out_channels;
}
size_t block_pos = name.find("blocks.");
if (block_pos != std::string::npos) {
auto items = split_string(name.substr(block_pos), '.');
if (items.size() > 1) {
int block_index = atoi(items[1].c_str());
if (block_index + 1 > config.depth) {
config.depth = block_index + 1;
}
if (name.find("blocks." + std::to_string(block_index) + ".ffn.experts.w1") != std::string::npos) {
config.sparse_layers.insert(block_index);
}
}
}
}
if (config.depth == 0) {
config.depth = 24;
}
config.axes_dim_sum = 0;
for (int axis_dim : config.axes_dim) {
config.axes_dim_sum += axis_dim;
}
if (!config.sparse_layers.empty()) {
config.num_experts = 128;
config.num_experts_per_tok = 8;
config.moe_intermediate_size = 768;
config.decoder_sparse_step = 1;
config.n_shared_experts = 1;
config.norm_topk_prob = true;
config.n_group = 4;
config.topk_group = 2;
config.routed_scaling_factor = 2.5f;
}
LOG_DEBUG("lingbot_video: depth = %" PRId64 ", hidden_size = %" PRId64 ", heads = %" PRId64 ", text_dim = %" PRId64 ", experts = %" PRId64 ", experts_per_tok = %" PRId64 ", n_group = %" PRId64 ", topk_group = %" PRId64 ", route_scale = %.2f, sparse_layers = %zu",
config.depth,
config.hidden_size,
config.num_attention_heads,
config.text_dim,
config.num_experts,
config.num_experts_per_tok,
config.n_group,
config.topk_group,
config.routed_scaling_factor,
config.sparse_layers.size());
return config;
}
};
struct LingBotVideoTextEmbedder : public GGMLBlock {
LingBotVideoTextEmbedder(int64_t text_dim,
int64_t hidden_size,
float eps = 1e-6f) {
blocks["norm"] = std::make_shared<RMSNorm>(text_dim, eps);
blocks["linear_1"] = std::make_shared<Linear>(text_dim, hidden_size, true);
blocks["linear_2"] = std::make_shared<Linear>(hidden_size, hidden_size, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["linear_1"]);
auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["linear_2"]);
x = norm->forward(ctx, x);
x = linear_1->forward(ctx, x);
x = ggml_silu_inplace(ctx->ggml_ctx, x);
x = linear_2->forward(ctx, x);
return x;
}
};
struct LingBotVideoAttention : public GGMLBlock {
int64_t num_heads;
int64_t head_dim;
LingBotVideoAttention(int64_t hidden_size,
int64_t num_heads,
bool qkv_bias = false,
bool out_bias = true,
float eps = 1e-6f)
: num_heads(num_heads),
head_dim(hidden_size / num_heads) {
int64_t inner_dim = num_heads * head_dim;
blocks["to_q"] = std::make_shared<Linear>(hidden_size, inner_dim, qkv_bias);
blocks["to_k"] = std::make_shared<Linear>(hidden_size, inner_dim, qkv_bias);
blocks["to_v"] = std::make_shared<Linear>(hidden_size, inner_dim, qkv_bias);
blocks["norm_q"] = std::make_shared<RMSNorm>(head_dim, eps);
blocks["norm_k"] = std::make_shared<RMSNorm>(head_dim, eps);
blocks["to_out"] = std::make_shared<Linear>(inner_dim, hidden_size, out_bias);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* pe,
ggml_tensor* attention_mask = nullptr) {
// x: [N, video_tokens + text_tokens, hidden_size]
auto to_q = std::dynamic_pointer_cast<Linear>(blocks["to_q"]);
auto to_k = std::dynamic_pointer_cast<Linear>(blocks["to_k"]);
auto to_v = std::dynamic_pointer_cast<Linear>(blocks["to_v"]);
auto norm_q = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_q"]);
auto norm_k = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"]);
auto to_out = std::dynamic_pointer_cast<Linear>(blocks["to_out"]);
int64_t S = x->ne[1];
int64_t N = x->ne[2];
auto q = to_q->forward(ctx, x);
auto k = to_k->forward(ctx, x);
auto v = to_v->forward(ctx, x);
q = ggml_reshape_4d(ctx->ggml_ctx, q, head_dim, num_heads, S, N);
k = ggml_reshape_4d(ctx->ggml_ctx, k, head_dim, num_heads, S, N);
v = ggml_reshape_4d(ctx->ggml_ctx, v, head_dim, num_heads, S, N);
q = norm_q->forward(ctx, q);
k = norm_k->forward(ctx, k);
x = Rope::attention(ctx, q, k, v, pe, attention_mask);
x = to_out->forward(ctx, x);
return x;
}
};
struct LingBotVideoMLP : public UnaryBlock {
LingBotVideoMLP(int64_t hidden_size,
int64_t intermediate_size) {
blocks["gate_proj"] = std::make_shared<Linear>(hidden_size, intermediate_size, false);
blocks["up_proj"] = std::make_shared<Linear>(hidden_size, intermediate_size, false);
blocks["down_proj"] = std::make_shared<Linear>(intermediate_size, hidden_size, false);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto gate_proj = std::dynamic_pointer_cast<Linear>(blocks["gate_proj"]);
auto up_proj = std::dynamic_pointer_cast<Linear>(blocks["up_proj"]);
auto down_proj = std::dynamic_pointer_cast<Linear>(blocks["down_proj"]);
auto gate = gate_proj->forward(ctx, x);
gate = ggml_silu_inplace(ctx->ggml_ctx, gate);
auto up = up_proj->forward(ctx, x);
x = ggml_mul(ctx->ggml_ctx, gate, up);
x = down_proj->forward(ctx, x);
return x;
}
};
struct LingBotVideoSparseMoeBlock : public UnaryBlock {
int64_t hidden_size;
int64_t intermediate_size;
int64_t num_experts;
int64_t num_experts_per_tok;
bool has_shared_experts;
bool norm_topk_prob;
float routed_scaling_factor;
int64_t n_group;
int64_t topk_group;
std::vector<float> group_expert_mask_vec;
bool has_correction_bias = false;
LingBotVideoSparseMoeBlock(const LingBotVideoConfig& config)
: hidden_size(config.hidden_size),
intermediate_size(config.moe_intermediate_size),
num_experts(config.num_experts),
num_experts_per_tok(config.num_experts_per_tok),
has_shared_experts(config.n_shared_experts > 0),
norm_topk_prob(config.norm_topk_prob),
routed_scaling_factor(config.routed_scaling_factor),
n_group(config.n_group),
topk_group(config.topk_group) {
if (n_group > 1) {
GGML_ASSERT(num_experts % n_group == 0);
int64_t experts_per_group = num_experts / n_group;
group_expert_mask_vec.assign(static_cast<size_t>(num_experts * n_group), 0.f);
for (int64_t group = 0; group < n_group; ++group) {
int64_t expert_begin = group * experts_per_group;
int64_t expert_end = expert_begin + experts_per_group;
for (int64_t expert = expert_begin; expert < expert_end; ++expert) {
group_expert_mask_vec[static_cast<size_t>(group * num_experts + expert)] = 1.f;
}
}
}
if (has_shared_experts) {
blocks["shared_experts"] = std::make_shared<LingBotVideoMLP>(hidden_size,
intermediate_size * config.n_shared_experts);
}
}
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
ggml_type router_type = get_type(prefix + "router.weight", tensor_storage_map, GGML_TYPE_F32);
ggml_type w1_type = get_type(prefix + "experts.w1", tensor_storage_map, GGML_TYPE_F32);
ggml_type w2_type = get_type(prefix + "experts.w2", tensor_storage_map, GGML_TYPE_F32);
ggml_type w3_type = get_type(prefix + "experts.w3", tensor_storage_map, GGML_TYPE_F32);
params["router.weight"] = ggml_new_tensor_2d(ctx, router_type, hidden_size, num_experts);
if (tensor_storage_map.find(prefix + "router.e_score_correction_bias") != tensor_storage_map.end()) {
params["router.e_score_correction_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_experts);
has_correction_bias = true;
}
params["experts.w1"] = ggml_new_tensor_3d(ctx, w1_type, hidden_size, intermediate_size, num_experts);
params["experts.w2"] = ggml_new_tensor_3d(ctx, w2_type, intermediate_size, hidden_size, num_experts);
params["experts.w3"] = ggml_new_tensor_3d(ctx, w3_type, hidden_size, intermediate_size, num_experts);
}
ggml_tensor* expert_linear(GGMLRunnerContext* ctx,
const std::string& weight_name,
ggml_tensor* x,
ggml_tensor* selected_experts) {
return ggml_mul_mat_id(ctx->ggml_ctx, params[weight_name], x, selected_experts);
}
ggml_tensor* select_experts(GGMLRunnerContext* ctx, ggml_tensor* choice_scores) {
ggml_context* gctx = ctx->ggml_ctx;
if (n_group <= 1 || topk_group <= 0) {
return ggml_argsort_top_k(gctx, choice_scores, static_cast<int>(num_experts_per_tok));
}
GGML_ASSERT(choice_scores->ne[0] == num_experts);
GGML_ASSERT(num_experts % n_group == 0);
GGML_ASSERT(topk_group > 0 && topk_group <= n_group);
GGML_ASSERT(!group_expert_mask_vec.empty());
const int64_t n_token_total = choice_scores->ne[1];
const int64_t experts_per_group = num_experts / n_group;
const int group_score_k = 2;
GGML_ASSERT(experts_per_group >= group_score_k);
ggml_tensor* grouped_scores = ggml_reshape_3d(gctx, choice_scores, experts_per_group, n_group, n_token_total);
ggml_tensor* group_top_ids = ggml_argsort_top_k(gctx, grouped_scores, group_score_k);
grouped_scores = ggml_reshape_3d(gctx, grouped_scores, 1, experts_per_group, n_group * n_token_total);
group_top_ids = ggml_cont(gctx, group_top_ids);
group_top_ids = ggml_reshape_2d(gctx, group_top_ids, group_score_k, n_group * n_token_total);
ggml_tensor* group_top_values = ggml_get_rows(gctx, grouped_scores, group_top_ids);
group_top_values = ggml_reshape_3d(gctx, group_top_values, group_score_k, n_group, n_token_total);
ggml_tensor* group_scores = nullptr;
for (int rank = 0; rank < group_score_k; ++rank) {
ggml_tensor* value = ggml_view_3d(gctx,
group_top_values,
1,
n_group,
n_token_total,
group_top_values->nb[1],
group_top_values->nb[2],
rank * group_top_values->nb[0]);
group_scores = group_scores == nullptr ? value : ggml_add(gctx, group_scores, value);
}
group_scores = ggml_reshape_2d(gctx, group_scores, n_group, n_token_total);
ggml_tensor* selected_groups = ggml_argsort_top_k(gctx, group_scores, static_cast<int>(topk_group));
selected_groups = ggml_cont(gctx, selected_groups);
ggml_tensor* group_expert_mask = ggml_new_tensor_3d(gctx, GGML_TYPE_F32, num_experts, n_group, 1);
ctx->bind_backend_tensor_data(group_expert_mask, group_expert_mask_vec.data());
ggml_tensor* group_expert_mask_template = ggml_new_tensor_3d(gctx, GGML_TYPE_F32, num_experts, n_group, n_token_total);
group_expert_mask = ggml_repeat(gctx, group_expert_mask, group_expert_mask_template);
ggml_tensor* selected_group_masks = ggml_get_rows(gctx, group_expert_mask, selected_groups);
ggml_tensor* selected_mask = nullptr;
for (int64_t rank = 0; rank < topk_group; ++rank) {
ggml_tensor* mask = ggml_view_3d(gctx,
selected_group_masks,
num_experts,
1,
n_token_total,
selected_group_masks->nb[1],
selected_group_masks->nb[2],
rank * selected_group_masks->nb[1]);
selected_mask = selected_mask == nullptr ? mask : ggml_add(gctx, selected_mask, mask);
}
selected_mask = ggml_reshape_2d(gctx, selected_mask, num_experts, n_token_total);
ggml_tensor* excluded_group_mask = ggml_sub(gctx, selected_mask, ggml_ext_ones_like(gctx, selected_mask));
ggml_tensor* masked_scores = ggml_add(gctx, choice_scores, ggml_scale(gctx, excluded_group_mask, 1.0e9f));
return ggml_argsort_top_k(gctx, masked_scores, static_cast<int>(num_experts_per_tok));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
// x: [N, tokens, hidden_size]
GGML_ASSERT(num_experts > 0);
GGML_ASSERT(num_experts_per_tok > 0 && num_experts_per_tok <= num_experts);
ggml_tensor* residual = x;
const int64_t n_token = x->ne[1];
const int64_t N = x->ne[2];
const int64_t n_token_total = n_token * N;
ggml_tensor* router_logits = ggml_mul_mat(ctx->ggml_ctx, params["router.weight"], x);
router_logits = ggml_reshape_2d(ctx->ggml_ctx, router_logits, num_experts, n_token_total);
ggml_tensor* probs = ggml_sigmoid(ctx->ggml_ctx, router_logits);
ggml_tensor* choice_scores = probs;
if (has_correction_bias) {
choice_scores = ggml_add(ctx->ggml_ctx, choice_scores, params["router.e_score_correction_bias"]);
}
ggml_tensor* selected_experts = select_experts(ctx, choice_scores);
ggml_tensor* weights = ggml_get_rows(ctx->ggml_ctx,
ggml_reshape_3d(ctx->ggml_ctx, probs, 1, num_experts, n_token_total),
selected_experts);
weights = ggml_reshape_2d(ctx->ggml_ctx, weights, num_experts_per_tok, n_token_total);
if (norm_topk_prob && num_experts_per_tok > 1) {
auto weights_sum = ggml_sum_rows(ctx->ggml_ctx, weights);
weights_sum = ggml_clamp(ctx->ggml_ctx, weights_sum, 6.103515625e-5f, INFINITY);
weights = ggml_div(ctx->ggml_ctx, weights, weights_sum);
}
if (routed_scaling_factor != 1.0f) {
weights = ggml_scale(ctx->ggml_ctx, weights, routed_scaling_factor);
}
weights = ggml_reshape_3d(ctx->ggml_ctx, weights, 1, num_experts_per_tok, n_token_total);
x = ggml_reshape_3d(ctx->ggml_ctx, x, hidden_size, 1, n_token_total);
auto gate = expert_linear(ctx, "experts.w1", x, selected_experts);
gate = ggml_silu_inplace(ctx->ggml_ctx, gate);
auto up = expert_linear(ctx, "experts.w3", x, selected_experts);
auto act = ggml_mul(ctx->ggml_ctx, gate, up);
auto out = expert_linear(ctx, "experts.w2", act, selected_experts);
out = ggml_mul(ctx->ggml_ctx, out, weights);
ggml_tensor* summed = nullptr;
for (int64_t i = 0; i < num_experts_per_tok; ++i) {
auto expert_out = ggml_view_2d(ctx->ggml_ctx,
out,
hidden_size,
n_token_total,
out->nb[2],
i * out->nb[1]);
summed = summed == nullptr ? expert_out : ggml_add(ctx->ggml_ctx, summed, expert_out);
}
if (num_experts_per_tok == 1) {
summed = ggml_cont(ctx->ggml_ctx, summed);
}
summed = ggml_reshape_3d(ctx->ggml_ctx, summed, hidden_size, n_token, N);
if (has_shared_experts) {
auto shared_experts = std::dynamic_pointer_cast<LingBotVideoMLP>(blocks["shared_experts"]);
summed = ggml_add(ctx->ggml_ctx, summed, shared_experts->forward(ctx, residual));
}
return summed;
}
};
struct LingBotVideoBlock : public GGMLBlock {
int64_t hidden_size;
LingBotVideoBlock(const LingBotVideoConfig& config,
bool sparse)
: hidden_size(config.hidden_size) {
blocks["norm1"] = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps);
blocks["attn"] = std::make_shared<LingBotVideoAttention>(config.hidden_size,
config.num_attention_heads,
config.qkv_bias,
config.out_bias,
config.norm_eps);
blocks["norm_post_attn"] = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps);
blocks["norm2"] = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps);
if (sparse) {
blocks["ffn"] = std::make_shared<LingBotVideoSparseMoeBlock>(config);
} else {
blocks["ffn"] = std::make_shared<LingBotVideoMLP>(config.hidden_size, config.intermediate_size);
}
blocks["norm_post_ffn"] = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps);
}
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
params["scale_shift_table"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size * 6, 1);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* temb6,
ggml_tensor* pe,
ggml_tensor* attention_mask = nullptr) {
// x: [N, tokens, hidden_size], temb6: [N, tokens, 6 * hidden_size]
auto norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm1"]);
auto attn = std::dynamic_pointer_cast<LingBotVideoAttention>(blocks["attn"]);
auto norm_post_attn = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_post_attn"]);
auto norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm2"]);
auto ffn = std::dynamic_pointer_cast<UnaryBlock>(blocks["ffn"]);
auto norm_post_ffn = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_post_ffn"]);
auto table = ggml_reshape_3d(ctx->ggml_ctx, params["scale_shift_table"], hidden_size * 6, 1, 1);
auto mod = ggml_add(ctx->ggml_ctx, temb6, table);
auto mods = ggml_ext_chunk(ctx->ggml_ctx, mod, 6, 0);
auto shift_msa = mods[0];
auto scale_msa = mods[1];
auto gate_msa = ggml_tanh(ctx->ggml_ctx, mods[2]);
auto shift_mlp = mods[3];
auto scale_mlp = mods[4];
auto gate_mlp = ggml_tanh(ctx->ggml_ctx, mods[5]);
auto attn_in = Flux::modulate(ctx->ggml_ctx, norm1->forward(ctx, x), shift_msa, scale_msa, true);
auto attn_out = attn->forward(ctx, attn_in, pe, attention_mask);
attn_out = norm_post_attn->forward(ctx, attn_out);
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, gate_msa, attn_out));
auto ffn_in = Flux::modulate(ctx->ggml_ctx, norm2->forward(ctx, x), shift_mlp, scale_mlp, true);
auto ffn_out = ffn->forward(ctx, ffn_in);
ffn_out = norm_post_ffn->forward(ctx, ffn_out);
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, gate_mlp, ffn_out));
return x;
}
};
class LingBotVideoModel : public GGMLBlock {
public:
LingBotVideoConfig config;
LingBotVideoModel() = default;
LingBotVideoModel(LingBotVideoConfig config)
: config(config) {
int64_t patch_dim = config.patch_t * config.patch_h * config.patch_w * config.in_channels;
blocks["patch_embedder"] = std::make_shared<Linear>(patch_dim, config.hidden_size, config.patch_embed_bias);
blocks["time_embedder"] = std::make_shared<Qwen::TimestepEmbedding>(config.freq_dim,
config.hidden_size,
config.hidden_size,
0,
config.timestep_mlp_bias);
blocks["time_modulation.1"] = std::make_shared<Linear>(config.hidden_size, 6 * config.hidden_size, true);
blocks["text_embedder"] = std::make_shared<LingBotVideoTextEmbedder>(config.text_dim,
config.hidden_size,
config.norm_eps);
for (int i = 0; i < config.depth; i++) {
bool sparse = config.sparse_layers.find(i) != config.sparse_layers.end();
blocks["blocks." + std::to_string(i)] = std::make_shared<LingBotVideoBlock>(config, sparse);
}
blocks["norm_out"] = std::make_shared<LayerNorm>(config.hidden_size, config.norm_eps, false);
blocks["norm_out_modulation.1"] = std::make_shared<Linear>(config.hidden_size, 2 * config.hidden_size, true);
blocks["proj_out"] = std::make_shared<Linear>(config.hidden_size,
config.patch_t * config.patch_h * config.patch_w * config.out_channels,
true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* timestep,
ggml_tensor* context,
ggml_tensor* pe) {
// x: [N*C, T, H, W], context: [N, text_tokens, text_dim]
GGML_ASSERT(context != nullptr);
GGML_ASSERT(x->ne[3] == config.in_channels);
GGML_ASSERT(x->ne[2] % config.patch_t == 0);
GGML_ASSERT(x->ne[1] % config.patch_h == 0);
GGML_ASSERT(x->ne[0] % config.patch_w == 0);
auto patch_embedder = std::dynamic_pointer_cast<Linear>(blocks["patch_embedder"]);
auto time_embedder = std::dynamic_pointer_cast<Qwen::TimestepEmbedding>(blocks["time_embedder"]);
auto time_modulation = std::dynamic_pointer_cast<Linear>(blocks["time_modulation.1"]);
auto text_embedder = std::dynamic_pointer_cast<LingBotVideoTextEmbedder>(blocks["text_embedder"]);
auto norm_out = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_out"]);
auto norm_out_modulation = std::dynamic_pointer_cast<Linear>(blocks["norm_out_modulation.1"]);
auto proj_out = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t T = x->ne[2];
int64_t N = 1;
int64_t t_len = T / config.patch_t;
int64_t h_len = H / config.patch_h;
int64_t w_len = W / config.patch_w;
int64_t n_img = t_len * h_len * w_len;
auto img = DiT::patchify_3d(ctx->ggml_ctx, x, config.patch_t, config.patch_h, config.patch_w, N, false);
img = patch_embedder->forward(ctx, img);
auto txt = text_embedder->forward(ctx, context);
auto hidden_states = ggml_concat(ctx->ggml_ctx, img, txt, 1);
auto timestep_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx,
timestep,
static_cast<int>(config.freq_dim),
10000,
1.0f);
auto t_emb = time_embedder->forward(ctx, timestep_proj);
auto temb_template = ggml_new_tensor_3d(ctx->ggml_ctx, t_emb->type, t_emb->ne[0], hidden_states->ne[1], t_emb->ne[1]);
auto temb_tokens = ggml_repeat(ctx->ggml_ctx,
ggml_reshape_3d(ctx->ggml_ctx, t_emb, t_emb->ne[0], 1, t_emb->ne[1]),
temb_template);
auto temb6 = time_modulation->forward(ctx, ggml_silu(ctx->ggml_ctx, temb_tokens));
sd::ggml_graph_cut::mark_graph_cut(hidden_states, "lingbot_video.prelude", "hidden_states");
for (int i = 0; i < config.depth; i++) {
auto block = std::dynamic_pointer_cast<LingBotVideoBlock>(blocks["blocks." + std::to_string(i)]);
hidden_states = block->forward(ctx, hidden_states, temb6, pe);
sd::ggml_graph_cut::mark_graph_cut(hidden_states, "lingbot_video.blocks." + std::to_string(i), "hidden_states");
}
auto final_mods = ggml_ext_chunk(ctx->ggml_ctx,
norm_out_modulation->forward(ctx, ggml_silu(ctx->ggml_ctx, temb_tokens)),
2,
0);
hidden_states = norm_out->forward(ctx, hidden_states);
hidden_states = Flux::modulate(ctx->ggml_ctx, hidden_states, final_mods[0], final_mods[1], true);
hidden_states = proj_out->forward(ctx, hidden_states);
hidden_states = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, 0, n_img);
auto out = DiT::unpatchify_3d(ctx->ggml_ctx,
hidden_states,
t_len,
h_len,
w_len,
config.patch_t,
config.patch_h,
config.patch_w,
false);
return out;
}
};
struct LingBotVideoRunner : public DiffusionModelRunner {
LingBotVideoConfig config;
LingBotVideoModel lingbot_video;
std::vector<float> pe_vec;
LingBotVideoRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
const char* model_args = nullptr)
: DiffusionModelRunner(backend, prefix, weight_manager),
config(LingBotVideoConfig::detect_from_weights(tensor_storage_map, prefix)) {
SD_UNUSED(model_args);
lingbot_video = LingBotVideoModel(config);
lingbot_video.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "lingbot_video";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
lingbot_video.get_param_tensors(tensors, prefix);
}
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
const sd::Tensor<float>& timesteps_tensor,
const sd::Tensor<float>& context_tensor) {
ggml_cgraph* gf = new_graph_custom(LINGBOT_VIDEO_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor);
ggml_tensor* timesteps = make_input(timesteps_tensor);
GGML_ASSERT(x_tensor.dim() == 5);
GGML_ASSERT(x->ne[3] == config.in_channels);
GGML_ASSERT(!context_tensor.empty());
ggml_tensor* context = make_input(context_tensor);
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t T = x->ne[2];
int64_t N = x_tensor.shape()[4];
GGML_ASSERT(N == 1);
pe_vec = Rope::gen_lingbot_video_pe(static_cast<int>(T),
static_cast<int>(H),
static_cast<int>(W),
config.patch_t,
config.patch_h,
config.patch_w,
static_cast<int>(N),
static_cast<int>(context->ne[1]),
config.rope_theta,
config.axes_dim);
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
set_backend_tensor_data(pe, pe_vec.data());
auto runner_ctx = get_context();
ggml_tensor* out = lingbot_video.forward(&runner_ctx, x, timesteps, context, pe);
ggml_build_forward_expand(gf, out);
return gf;
}
sd::Tensor<float> compute(int n_threads,
const sd::Tensor<float>& x,
const sd::Tensor<float>& timesteps,
const sd::Tensor<float>& context) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
const DiffusionParams& diffusion_params) override {
GGML_ASSERT(diffusion_params.x != nullptr);
GGML_ASSERT(diffusion_params.timesteps != nullptr);
return compute(n_threads,
*diffusion_params.x,
*diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context));
}
};
} // namespace LingBotVideo
#endif // __SD_MODEL_DIFFUSION_LINGBOT_VIDEO_HPP__

View File

@ -800,7 +800,7 @@ namespace LTXV {
auto gate_mlp = mods[5];
auto x_norm = rms_norm(ctx->ggml_ctx, x);
x_norm = modulate(ctx->ggml_ctx, x_norm, shift_msa, scale_msa);
x_norm = LTXV::modulate(ctx->ggml_ctx, x_norm, shift_msa, scale_msa);
auto msa = attn1->forward(ctx, x_norm, nullptr, self_attention_mask, pe);
x = ggml_add(ctx->ggml_ctx, x, apply_gate(ctx->ggml_ctx, msa, gate_msa));
@ -810,12 +810,12 @@ namespace LTXV {
auto gate_q = mods[8];
auto q = rms_norm(ctx->ggml_ctx, x);
q = modulate(ctx->ggml_ctx, q, shift_q, scale_q);
q = LTXV::modulate(ctx->ggml_ctx, q, shift_q, scale_q);
auto context_mod = context;
if (prompt_timestep != nullptr) {
auto prompt_mods = get_prompt_scale_shift_values(ctx, prompt_timestep);
context_mod = modulate(ctx->ggml_ctx, context_mod, prompt_mods[0], prompt_mods[1]);
context_mod = LTXV::modulate(ctx->ggml_ctx, context_mod, prompt_mods[0], prompt_mods[1]);
}
auto mca = attn2->forward(ctx, q, context_mod, attention_mask, nullptr, nullptr);
@ -826,7 +826,7 @@ namespace LTXV {
}
auto y = rms_norm(ctx->ggml_ctx, x);
y = modulate(ctx->ggml_ctx, y, shift_mlp, scale_mlp);
y = LTXV::modulate(ctx->ggml_ctx, y, shift_mlp, scale_mlp);
auto mlp_out = ff->forward(ctx, y);
x = ggml_add(ctx->ggml_ctx, x, apply_gate(ctx->ggml_ctx, mlp_out, gate_mlp));
return x;
@ -1177,11 +1177,11 @@ namespace LTXV {
if (cross_attention_adaln) {
auto q_mods = get_ada_values(ctx, table, timestep, dim, 9, 6, 3);
auto q = rms_norm(ctx->ggml_ctx, x);
q = modulate(ctx->ggml_ctx, q, q_mods[0], q_mods[1]);
q = LTXV::modulate(ctx->ggml_ctx, q, q_mods[0], q_mods[1]);
auto context_mod = context;
if (prompt_timestep != nullptr && prompt_table != nullptr) {
auto p_mods = get_ada_values(ctx, prompt_table, prompt_timestep, dim, 2);
context_mod = modulate(ctx->ggml_ctx, context_mod, p_mods[0], p_mods[1]);
context_mod = LTXV::modulate(ctx->ggml_ctx, context_mod, p_mods[0], p_mods[1]);
}
auto out = attn->forward(ctx, q, context_mod, attention_mask, nullptr, nullptr);
return apply_gate(ctx->ggml_ctx, out, q_mods[2]);
@ -1228,7 +1228,7 @@ namespace LTXV {
auto v_mods = get_ada_values(ctx, v_table, v_timestep, v_dim, cross_attention_adaln ? 9 : 6);
auto v_norm = rms_norm(ctx->ggml_ctx, vx);
v_norm = modulate(ctx->ggml_ctx, v_norm, v_mods[0], v_mods[1]);
v_norm = LTXV::modulate(ctx->ggml_ctx, v_norm, v_mods[0], v_mods[1]);
auto v_sa = attn1->forward(ctx, v_norm, nullptr, self_attention_mask, v_pe);
vx = ggml_add(ctx->ggml_ctx, vx, apply_gate(ctx->ggml_ctx, v_sa, v_mods[2]));
auto v_txt = apply_text_cross_attention(ctx,
@ -1246,7 +1246,7 @@ namespace LTXV {
if (run_ax) {
auto a_mods = get_ada_values(ctx, a_table, a_timestep, a_dim, cross_attention_adaln ? 9 : 6);
auto a_norm = rms_norm(ctx->ggml_ctx, ax);
a_norm = modulate(ctx->ggml_ctx, a_norm, a_mods[0], a_mods[1]);
a_norm = LTXV::modulate(ctx->ggml_ctx, a_norm, a_mods[0], a_mods[1]);
auto a_sa = audio_attn1->forward(ctx, a_norm, nullptr, nullptr, a_pe);
ax = ggml_add(ctx->ggml_ctx, ax, apply_gate(ctx->ggml_ctx, a_sa, a_mods[2]));
auto a_txt = apply_text_cross_attention(ctx,
@ -1269,8 +1269,8 @@ namespace LTXV {
auto a2v_video_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 0, 4);
auto a2v_audio = get_ada_values(ctx, a2v_audio_table, a_cross_scale_shift_timestep, a_dim, 4);
auto a2v_video = get_ada_values(ctx, a2v_video_table, v_cross_scale_shift_timestep, v_dim, 4);
auto vx_scaled = modulate(ctx->ggml_ctx, vx_norm3, a2v_video[1], a2v_video[0]);
auto ax_scaled = modulate(ctx->ggml_ctx, ax_norm3, a2v_audio[1], a2v_audio[0]);
auto vx_scaled = LTXV::modulate(ctx->ggml_ctx, vx_norm3, a2v_video[1], a2v_video[0]);
auto ax_scaled = LTXV::modulate(ctx->ggml_ctx, ax_norm3, a2v_audio[1], a2v_audio[0]);
auto a2v_out = audio_to_video_attn->forward(ctx, vx_scaled, ax_scaled, nullptr, v_cross_pe, a_cross_pe);
auto a2v_gate_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 4, 5);
auto a2v_gate = get_ada_values(ctx, a2v_gate_table, v_cross_gate_timestep, v_dim, 1)[0];
@ -1282,8 +1282,8 @@ namespace LTXV {
auto v2a_video_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 0, 4);
auto v2a_audio = get_ada_values(ctx, v2a_audio_table, a_cross_scale_shift_timestep, a_dim, 4);
auto v2a_video = get_ada_values(ctx, v2a_video_table, v_cross_scale_shift_timestep, v_dim, 4);
auto ax_scaled = modulate(ctx->ggml_ctx, ax_norm3, v2a_audio[3], v2a_audio[2]);
auto vx_scaled = modulate(ctx->ggml_ctx, vx_norm3, v2a_video[3], v2a_video[2]);
auto ax_scaled = LTXV::modulate(ctx->ggml_ctx, ax_norm3, v2a_audio[3], v2a_audio[2]);
auto vx_scaled = LTXV::modulate(ctx->ggml_ctx, vx_norm3, v2a_video[3], v2a_video[2]);
auto v2a_out = video_to_audio_attn->forward(ctx, ax_scaled, vx_scaled, nullptr, a_cross_pe, v_cross_pe);
auto v2a_gate_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_audio"], 1, 4, 5);
auto v2a_gate = get_ada_values(ctx, v2a_gate_table, a_cross_gate_timestep, a_dim, 1)[0];
@ -1291,14 +1291,14 @@ namespace LTXV {
}
auto a_ff_mods = get_ada_values(ctx, a_table, a_timestep, a_dim, cross_attention_adaln ? 9 : 6, 3, 3);
auto ax_scaled = rms_norm(ctx->ggml_ctx, ax);
ax_scaled = modulate(ctx->ggml_ctx, ax_scaled, a_ff_mods[0], a_ff_mods[1]);
ax_scaled = LTXV::modulate(ctx->ggml_ctx, ax_scaled, a_ff_mods[0], a_ff_mods[1]);
auto a_ff_out = audio_ff->forward(ctx, ax_scaled);
ax = ggml_add(ctx->ggml_ctx, ax, apply_gate(ctx->ggml_ctx, a_ff_out, a_ff_mods[2]));
}
auto v_ff_mods = get_ada_values(ctx, v_table, v_timestep, v_dim, cross_attention_adaln ? 9 : 6, 3, 3);
auto vx_scaled = rms_norm(ctx->ggml_ctx, vx);
vx_scaled = modulate(ctx->ggml_ctx, vx_scaled, v_ff_mods[0], v_ff_mods[1]);
vx_scaled = LTXV::modulate(ctx->ggml_ctx, vx_scaled, v_ff_mods[0], v_ff_mods[1]);
auto v_ff_out = ff->forward(ctx, vx_scaled);
vx = ggml_add(ctx->ggml_ctx, vx, apply_gate(ctx->ggml_ctx, v_ff_out, v_ff_mods[2]));
@ -1657,14 +1657,14 @@ namespace LTXV {
auto v_shift_scale = get_output_scale_shift(ctx, params["scale_shift_table"], v_embedded_time, config.hidden_size);
vx = norm_out->forward(ctx, vx);
vx = modulate(ctx->ggml_ctx, vx, v_shift_scale[0], v_shift_scale[1]);
vx = LTXV::modulate(ctx->ggml_ctx, vx, v_shift_scale[0], v_shift_scale[1]);
vx = proj_out->forward(ctx, vx);
vx = unpatchify_video(ctx, vx, width, height, frames);
if (ax != nullptr && audio_time > 0) {
auto a_shift_scale = get_output_scale_shift(ctx, params["audio_scale_shift_table"], a_embedded_time, config.audio_hidden_size);
ax = audio_norm_out->forward(ctx, ax);
ax = modulate(ctx->ggml_ctx, ax, a_shift_scale[0], a_shift_scale[1]);
ax = LTXV::modulate(ctx->ggml_ctx, ax, a_shift_scale[0], a_shift_scale[1]);
ax = audio_proj_out->forward(ctx, ax);
ax = unpatchify_audio(ctx, ax, audio_time);
}

View File

@ -0,0 +1,162 @@
#ifndef __SD_MODEL_DIFFUSION_MAGE_FLOW_HPP__
#define __SD_MODEL_DIFFUSION_MAGE_FLOW_HPP__
#include <cmath>
#include <memory>
#include "model/diffusion/qwen_image.hpp"
namespace MageFlow {
constexpr int MAGE_FLOW_GRAPH_SIZE = 20480;
// Mage-Flow was trained with BF16-rounded timestep frequencies; using Qwen's F32 projection degrades generation quality.
struct MageFlowTimestepProjEmbeddings : public Qwen::QwenTimestepProjEmbeddings {
static constexpr int TIMESTEP_DIM = 256;
static constexpr int HALF_DIM = TIMESTEP_DIM / 2;
std::vector<float> frequencies;
std::vector<float> timesteps_proj;
explicit MageFlowTimestepProjEmbeddings(int64_t embedding_dim)
: QwenTimestepProjEmbeddings(embedding_dim), frequencies(HALF_DIM) {
for (int i = 0; i < HALF_DIM; ++i) {
float frequency = std::exp(-std::log(10000.f) * static_cast<float>(i) / HALF_DIM);
frequencies[i] = ggml_bf16_to_fp32(ggml_fp32_to_bf16(frequency));
}
}
void prepare(const sd::Tensor<float>& timesteps) {
size_t num_timesteps = static_cast<size_t>(timesteps.numel());
timesteps_proj.resize(static_cast<size_t>(TIMESTEP_DIM) * num_timesteps);
for (size_t b = 0; b < num_timesteps; ++b) {
float sigma = ggml_bf16_to_fp32(ggml_fp32_to_bf16(timesteps.values()[b] / 1000.f));
for (int i = 0; i < HALF_DIM; ++i) {
float argument = sigma * frequencies[i] * 1000.f;
timesteps_proj[b * TIMESTEP_DIM + i] =
ggml_bf16_to_fp32(ggml_fp32_to_bf16(std::cos(argument)));
timesteps_proj[b * TIMESTEP_DIM + HALF_DIM + i] =
ggml_bf16_to_fp32(ggml_fp32_to_bf16(std::sin(argument)));
}
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* timesteps,
ggml_tensor* addition_t_cond = nullptr) override {
GGML_ASSERT(addition_t_cond == nullptr);
GGML_ASSERT(timesteps_proj.size() ==
static_cast<size_t>(TIMESTEP_DIM * ggml_nelements(timesteps)));
auto projection = ggml_new_tensor_2d(ctx->ggml_ctx,
GGML_TYPE_F32,
TIMESTEP_DIM,
ggml_nelements(timesteps));
ctx->bind_backend_tensor_data(projection, timesteps_proj.data());
auto timestep_embedder = std::dynamic_pointer_cast<Qwen::TimestepEmbedding>(blocks["timestep_embedder"]);
return timestep_embedder->forward(ctx, projection);
}
};
struct MageFlowRunner : public DiffusionModelRunner {
public:
Qwen::QwenImageConfig config;
Qwen::QwenImageModel mage_flow;
std::shared_ptr<MageFlowTimestepProjEmbeddings> time_text_embed;
std::vector<float> pe_vec;
MageFlowRunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: DiffusionModelRunner(backend, prefix, weight_manager) {
config.patch_size = 1;
config.in_channels = 128;
config.out_channels = 128;
config.num_layers = 12;
config.attention_head_dim = 128;
config.num_attention_heads = 24;
config.joint_attention_dim = 2560;
config.theta = 10000;
config.axes_dim = {16, 56, 56};
config.axes_dim_sum = 128;
time_text_embed = std::make_shared<MageFlowTimestepProjEmbeddings>(
config.num_attention_heads * config.attention_head_dim);
mage_flow = Qwen::QwenImageModel(config, time_text_embed);
mage_flow.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "mage_flow";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
mage_flow.get_param_tensors(tensors, prefix);
}
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
const sd::Tensor<float>& timesteps_tensor,
const sd::Tensor<float>& context_tensor,
const std::vector<sd::Tensor<float>>& ref_latents_tensor = {}) {
ggml_cgraph* gf = new_graph_custom(MAGE_FLOW_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor);
ggml_tensor* timesteps = make_input(timesteps_tensor);
GGML_ASSERT(x->ne[3] == 1);
GGML_ASSERT(!context_tensor.empty());
ggml_tensor* context = make_input(context_tensor);
std::vector<ggml_tensor*> ref_latents;
ref_latents.reserve(ref_latents_tensor.size());
for (const auto& ref_latent_tensor : ref_latents_tensor) {
ref_latents.push_back(make_input(ref_latent_tensor));
}
int batch_size = static_cast<int>(x->ne[3]);
pe_vec = Rope::gen_mage_flow_pe(static_cast<int>(x->ne[1]),
static_cast<int>(x->ne[0]),
batch_size,
static_cast<int>(context->ne[1]),
ref_latents,
config.theta,
config.axes_dim);
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
set_backend_tensor_data(pe, pe_vec.data());
time_text_embed->prepare(timesteps_tensor);
auto runner_ctx = get_context();
auto out = mage_flow.forward(&runner_ctx,
x,
timesteps,
nullptr,
context,
pe,
ref_latents);
ggml_build_forward_expand(gf, out);
return gf;
}
sd::Tensor<float> compute(int n_threads,
const sd::Tensor<float>& x,
const sd::Tensor<float>& timesteps,
const sd::Tensor<float>& context,
const std::vector<sd::Tensor<float>>& ref_latents = {}) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, ref_latents);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
}
sd::Tensor<float> compute(int n_threads,
const DiffusionParams& diffusion_params) override {
GGML_ASSERT(diffusion_params.x != nullptr);
GGML_ASSERT(diffusion_params.timesteps != nullptr);
static const std::vector<sd::Tensor<float>> empty_ref_latents;
return compute(n_threads,
*diffusion_params.x,
*diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context),
diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents);
}
};
} // namespace MageFlow
#endif // __SD_MODEL_DIFFUSION_MAGE_FLOW_HPP__

File diff suppressed because it is too large Load Diff

View File

@ -136,11 +136,15 @@ struct MMDiTConfig {
};
struct PatchEmbed : public GGMLBlock {
// 2D Image to Patch Embedding
// 2D/3D Image to Patch Embedding
protected:
bool is_3d;
bool flatten;
bool dynamic_img_pad;
int patch_size;
int patch_t;
int patch_h;
int patch_w;
int64_t embed_dim;
public:
PatchEmbed(int64_t img_size = 224,
@ -149,42 +153,90 @@ public:
int64_t embed_dim = 1536,
bool bias = true,
bool flatten = true,
bool dynamic_img_pad = true)
: patch_size(patch_size),
bool dynamic_img_pad = true,
bool is_3d = false)
: patch_t(is_3d ? patch_size : 1),
patch_h(patch_size),
patch_w(patch_size),
embed_dim(embed_dim),
flatten(flatten),
dynamic_img_pad(dynamic_img_pad) {
dynamic_img_pad(dynamic_img_pad),
is_3d(is_3d) {
// img_size is always None
// patch_size is always 2
// in_chans is always 16
// norm_layer is always False
// strict_img_size is always true, but not used
blocks["proj"] = std::shared_ptr<GGMLBlock>(new Conv2d(in_chans,
if (is_3d) {
blocks["proj"] = std::make_shared<Conv3d>(in_chans,
embed_dim,
{patch_size, patch_size},
{patch_size, patch_size},
{0, 0},
{1, 1},
bias));
std::tuple{patch_size, patch_size, patch_size},
std::tuple{patch_size, patch_size, patch_size},
std::tuple{0, 0, 0},
std::tuple{1, 1, 1},
bias);
} else {
blocks["proj"] = std::make_shared<Conv2d>(in_chans,
embed_dim,
std::pair{patch_size, patch_size},
std::pair{patch_size, patch_size},
std::pair{0, 0},
std::pair{1, 1},
bias);
}
}
PatchEmbed(int64_t img_size,
std::tuple<int, int, int> patch_size,
int64_t in_chans,
int64_t embed_dim,
bool bias = true,
bool flatten = true,
bool dynamic_img_pad = true)
: patch_t(std::get<0>(patch_size)),
patch_h(std::get<1>(patch_size)),
patch_w(std::get<2>(patch_size)),
embed_dim(embed_dim),
flatten(flatten),
dynamic_img_pad(dynamic_img_pad),
is_3d(true) {
SD_UNUSED(img_size);
blocks["proj"] = std::make_shared<Conv3d>(in_chans,
embed_dim,
patch_size,
patch_size,
std::tuple{0, 0, 0},
std::tuple{1, 1, 1},
bias);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
// x: [N, C, H, W]
// return: [N, H*W, embed_dim]
auto proj = std::dynamic_pointer_cast<Conv2d>(blocks["proj"]);
// x: [N, C, H, W] or [N*C, T, H, W]
// return: [N, h_len*w_len, embed_dim] or [N, t_len*h_len*w_len, embed_dim]
auto proj = std::dynamic_pointer_cast<UnaryBlock>(blocks["proj"]);
if (dynamic_img_pad) {
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int pad_h = (patch_size - H % patch_size) % patch_size;
int pad_w = (patch_size - W % patch_size) % patch_size;
x = ggml_pad(ctx->ggml_ctx, x, pad_w, pad_h, 0, 0); // TODO: reflect pad mode
int pad_t = 0;
int pad_h = (patch_h - static_cast<int>(H % patch_h)) % patch_h;
int pad_w = (patch_w - static_cast<int>(W % patch_w)) % patch_w;
if (is_3d) {
int64_t T = x->ne[2];
pad_t = (patch_t - static_cast<int>(T % patch_t)) % patch_t;
}
x = proj->forward(ctx, x);
x = ggml_pad(ctx->ggml_ctx, x, pad_w, pad_h, pad_t, 0); // TODO: reflect pad mode
}
x = proj->forward(ctx, x); // [N, C, h_len, w_len] or [N*C, t_len, h_len, w_len]
if (flatten) {
x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1], x->ne[2], x->ne[3]);
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
if (is_3d) {
x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1] * x->ne[2], embed_dim, x->ne[3] / embed_dim); // [N, C, t_len*h_len*w_len]
} else {
x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1], x->ne[2], x->ne[3]); // [N, C, h_len*w_len]
}
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); // [N, h_len*w_len, C]
}
return x;
}

View File

@ -10,10 +10,44 @@
#include "model/common/rope.hpp"
#include "model_manager.h"
enum class RefImageResizeMode {
NONE,
LONGEST_SIDE,
AREA,
};
struct RefImageParams {
bool pass_to_vlm = false;
bool pass_to_dit = true;
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED;
bool force_ref_timestep_zero = false;
bool resize_before_vae = true;
int vae_input_max_pixels = -1;
RefImageResizeMode vlm_resize_mode = RefImageResizeMode::AREA;
int vlm_min_size = -1;
int vlm_max_size = -1;
bool resize_vae_to_target = false;
};
const std::unordered_map<std::string, RefImageParams> REF_IMAGE_PRESETS = {
{"flux_kontext", {false, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::NONE, -1, -1}},
{"longcat", {true, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
{"flux2", {false, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::NONE, -1, -1}},
{"qwen", {true, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
{"qwen_layered", {true, true, Rope::RefIndexMode::DECREASE, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
{"mage_flow", {true, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::LONGEST_SIDE, -1, 384, true}},
{"z_image_omni", {true, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
{"krea2_ostris_edit", {true, true, Rope::RefIndexMode::INCREASE, true, true, -1, RefImageResizeMode::AREA, -1, -1}},
{"krea2_edit", {true, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::LONGEST_SIDE, 768, 768}},
{"cosmos_reference", {false, true, Rope::RefIndexMode::INCREASE, false, false, -1, RefImageResizeMode::NONE, -1, -1}},
};
struct UNetDiffusionExtra {
int num_video_frames = -1;
const std::vector<sd::Tensor<float>>* controls = nullptr;
float control_strength = 0.f;
const sd::Tensor<float>* ip_context = nullptr;
float ip_scale = 1.f;
};
struct SkipLayerDiffusionExtra {
@ -53,10 +87,40 @@ struct LTXAVDiffusionExtra {
const sd::Tensor<float>* video_positions = nullptr;
};
enum class MiniMaxH3ReferenceKind : int32_t {
IMAGE,
VIDEO,
AUDIO,
VIDEO_AUDIO,
};
struct MiniMaxH3ReferenceBlock {
MiniMaxH3ReferenceKind kind = MiniMaxH3ReferenceKind::IMAGE;
int32_t video_index = -1;
int32_t audio_index = -1;
};
struct MiniMaxH3DiffusionExtra {
const sd::Tensor<int32_t>* text_token_tags = nullptr;
const sd::Tensor<int32_t>* keyframe_indices = nullptr;
const std::vector<sd::Tensor<float>>* reference_audio_latents = nullptr;
const std::vector<MiniMaxH3ReferenceBlock>* reference_blocks = nullptr;
int audio_length = 0;
float video_sigma_shift = 12.f;
float audio_sigma_shift = 3.f;
};
struct MiniT2IDiffusionExtra {
const sd::Tensor<float>* mask = nullptr;
};
struct HunyuanVideoDiffusionExtra {
const sd::Tensor<float>* guidance = nullptr;
const sd::Tensor<float>* byt5 = nullptr;
const sd::Tensor<float>* vision = nullptr;
const sd::Tensor<float>* timestep_r = nullptr;
};
using DiffusionExtraParams = std::variant<std::monostate,
UNetDiffusionExtra,
SkipLayerDiffusionExtra,
@ -65,7 +129,9 @@ using DiffusionExtraParams = std::variant<std::monostate,
WanDiffusionExtra,
HiDreamO1DiffusionExtra,
LTXAVDiffusionExtra,
MiniT2IDiffusionExtra>;
MiniMaxH3DiffusionExtra,
MiniT2IDiffusionExtra,
HunyuanVideoDiffusionExtra>;
struct DiffusionParams {
const sd::Tensor<float>* x = nullptr;
@ -74,7 +140,7 @@ struct DiffusionParams {
const sd::Tensor<float>* c_concat = nullptr;
const sd::Tensor<float>* y = nullptr;
const std::vector<sd::Tensor<float>>* ref_latents = nullptr;
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED;
RefImageParams ref_image_params = {false, false, Rope::RefIndexMode::FIXED, false};
DiffusionExtraParams extra = std::monostate{};
};

View File

@ -36,11 +36,19 @@ namespace Pid {
int64_t lq_interval = 2;
int64_t lq_sr_scale = 4;
int64_t lq_latent_down_factor = 8;
int64_t lq_latent_unpatchify_factor = 1;
bool lq_replicate_padding = false;
bool lq_gate_per_token = false;
bool pit_lq_inject = false;
int64_t rope_ref_grid_h = 64;
int64_t rope_ref_grid_w = 64;
static PixelDiTConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
PixelDiTConfig config;
int64_t latent_proj_in_channels = config.lq_latent_channels;
int64_t num_lq_gates = 0;
const std::string lq_prefix = prefix + ".lq_proj.";
config.pit_lq_inject = tensor_storage_map.find(lq_prefix + "pit_head.weight") != tensor_storage_map.end();
for (const auto& [name, tensor_storage] : tensor_storage_map) {
if (!starts_with(name, prefix)) {
continue;
@ -61,20 +69,56 @@ namespace Pid {
config.pixel_depth = std::max<int64_t>(config.pixel_depth, block_index + 1);
}
}
if (name.find("lq_proj.latent_proj.0.weight") != std::string::npos) {
config.lq_latent_channels = tensor_storage.ne[2];
config.lq_latent_down_factor = config.lq_latent_channels >= 64 ? 16 : 8;
if (name == lq_prefix + "latent_proj.0.weight") {
latent_proj_in_channels = tensor_storage.ne[2];
config.lq_hidden_dim = tensor_storage.ne[3];
}
if (starts_with(name, lq_prefix + "gate_modules.")) {
auto items = split_string(name.substr(lq_prefix.size()), '.');
if (items.size() > 1) {
int gate_index = atoi(items[1].c_str());
num_lq_gates = std::max<int64_t>(num_lq_gates, gate_index + 1);
}
}
if (name.find("patch_blocks.0.mlp_x.w1.weight") != std::string::npos) {
config.patch_mlp_hidden_dim = tensor_storage.ne[1];
}
}
LOG_DEBUG("pid: patch_depth = %" PRId64 ", pixel_depth = %" PRId64 ", patch_mlp_hidden_dim = %" PRId64 ", lq_latent_channels = %" PRId64 ", lq_latent_down_factor = %" PRId64,
if (num_lq_gates > 0) {
config.lq_interval = (config.patch_depth + num_lq_gates - 1) / num_lq_gates;
}
if (config.pit_lq_inject) {
if (latent_proj_in_channels == 16) {
config.lq_latent_channels = 16;
config.lq_latent_down_factor = 8;
config.lq_latent_unpatchify_factor = 1;
} else {
GGML_ASSERT(latent_proj_in_channels == 32);
config.lq_latent_channels = 128;
config.lq_latent_down_factor = 16;
config.lq_latent_unpatchify_factor = 2;
}
auto gate_weight = tensor_storage_map.find(lq_prefix + "gate_modules.0.content_proj.weight");
if (gate_weight != tensor_storage_map.end()) {
config.lq_gate_per_token = gate_weight->second.ne[1] == 1;
}
config.lq_replicate_padding = true;
config.rope_ref_grid_h = 128;
config.rope_ref_grid_w = 128;
} else {
config.lq_latent_channels = latent_proj_in_channels;
config.lq_latent_down_factor = latent_proj_in_channels >= 64 ? 16 : 8;
}
LOG_DEBUG("pid: version = %s, patch_depth = %" PRId64 ", pixel_depth = %" PRId64 ", patch_mlp_hidden_dim = %" PRId64 ", lq_latent_channels = %" PRId64 ", lq_hidden_dim = %" PRId64 ", lq_latent_down_factor = %" PRId64 ", lq_latent_unpatchify_factor = %" PRId64 ", lq_interval = %" PRId64,
config.pit_lq_inject ? "1.5" : "1",
config.patch_depth,
config.pixel_depth,
config.patch_mlp_hidden_dim,
config.lq_latent_channels,
config.lq_latent_down_factor);
config.lq_hidden_dim,
config.lq_latent_down_factor,
config.lq_latent_unpatchify_factor,
config.lq_interval);
return config;
}
};
@ -135,6 +179,18 @@ namespace Pid {
return ggml_add(ctx, ggml_add(ctx, x, ggml_mul(ctx, x, scale)), shift);
}
inline ggml_tensor* replicate_pad_2d(ggml_context* ctx, ggml_tensor* x) {
auto left = ggml_ext_slice(ctx, x, 0, 0, 1);
auto right = ggml_ext_slice(ctx, x, 0, x->ne[0] - 1, x->ne[0]);
x = ggml_concat(ctx, left, x, 0);
x = ggml_concat(ctx, x, right, 0);
auto top = ggml_ext_slice(ctx, x, 1, 0, 1);
auto bottom = ggml_ext_slice(ctx, x, 1, x->ne[1] - 1, x->ne[1]);
x = ggml_concat(ctx, top, x, 1);
return ggml_concat(ctx, x, bottom, 1);
}
struct PatchTokenEmbedder : public GGMLBlock {
bool use_rms_norm;
@ -457,9 +513,9 @@ namespace Pid {
struct SigmaAwareGate : public GGMLBlock {
int64_t dim;
SigmaAwareGate(int64_t dim)
SigmaAwareGate(int64_t dim, bool per_token = false)
: dim(dim) {
blocks["content_proj"] = std::make_shared<Linear>(dim * 2, dim, true);
blocks["content_proj"] = std::make_shared<Linear>(dim * 2, per_token ? 1 : dim, true);
}
void init_params(ggml_context* ctx,
@ -479,16 +535,20 @@ namespace Pid {
auto alpha = ggml_exp(ctx->ggml_ctx, params["log_alpha"]);
auto offset = ggml_neg(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, alpha, sigma));
auto gate = ggml_sigmoid(ctx->ggml_ctx, ggml_add(ctx->ggml_ctx, content_logit, offset));
return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, gate, lq));
return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, lq, gate));
}
};
struct PiDResBlock : public GGMLBlock {
PiDResBlock(int64_t channels) {
bool replicate_padding;
PiDResBlock(int64_t channels, bool replicate_padding = false)
: replicate_padding(replicate_padding) {
std::pair<int, int> padding = replicate_padding ? std::pair<int, int>{0, 0} : std::pair<int, int>{1, 1};
blocks["block.0"] = std::make_shared<GroupNorm>(4, channels, 1e-5f);
blocks["block.2"] = std::make_shared<Conv2d>(channels, channels, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, std::pair<int, int>{1, 1});
blocks["block.2"] = std::make_shared<Conv2d>(channels, channels, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, padding);
blocks["block.3"] = std::make_shared<GroupNorm>(4, channels, 1e-5f);
blocks["block.5"] = std::make_shared<Conv2d>(channels, channels, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, std::pair<int, int>{1, 1});
blocks["block.5"] = std::make_shared<Conv2d>(channels, channels, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, padding);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
@ -497,8 +557,14 @@ namespace Pid {
auto norm2 = std::dynamic_pointer_cast<GroupNorm>(blocks["block.3"]);
auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["block.5"]);
auto h = ggml_silu_inplace(ctx->ggml_ctx, norm1->forward(ctx, x));
if (replicate_padding) {
h = replicate_pad_2d(ctx->ggml_ctx, h);
}
h = conv1->forward(ctx, h);
h = ggml_silu_inplace(ctx->ggml_ctx, norm2->forward(ctx, h));
if (replicate_padding) {
h = replicate_pad_2d(ctx->ggml_ctx, h);
}
h = conv2->forward(ctx, h);
return ggml_add(ctx->ggml_ctx, x, h);
}
@ -509,16 +575,23 @@ namespace Pid {
LQProjection2D(const PixelDiTConfig& config)
: config(config) {
blocks["latent_proj.0"] = std::make_shared<Conv2d>(config.lq_latent_channels, config.lq_hidden_dim, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, std::pair<int, int>{1, 1});
blocks["latent_proj.2"] = std::make_shared<Conv2d>(config.lq_hidden_dim, config.lq_hidden_dim, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, std::pair<int, int>{1, 1});
int64_t unpatchify_area = config.lq_latent_unpatchify_factor * config.lq_latent_unpatchify_factor;
GGML_ASSERT(config.lq_latent_channels % unpatchify_area == 0);
int64_t latent_proj_in_channels = config.lq_latent_channels / unpatchify_area;
std::pair<int, int> padding = config.lq_replicate_padding ? std::pair<int, int>{0, 0} : std::pair<int, int>{1, 1};
blocks["latent_proj.0"] = std::make_shared<Conv2d>(latent_proj_in_channels, config.lq_hidden_dim, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, padding);
blocks["latent_proj.2"] = std::make_shared<Conv2d>(config.lq_hidden_dim, config.lq_hidden_dim, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, padding);
for (int i = 0; i < config.lq_num_res_blocks; ++i) {
blocks["latent_proj." + std::to_string(3 + i)] = std::make_shared<PiDResBlock>(config.lq_hidden_dim);
blocks["latent_proj." + std::to_string(3 + i)] = std::make_shared<PiDResBlock>(config.lq_hidden_dim, config.lq_replicate_padding);
}
int num_outputs = static_cast<int>((config.patch_depth + config.lq_interval - 1) / config.lq_interval);
for (int i = 0; i < num_outputs; ++i) {
blocks["output_heads." + std::to_string(i)] = std::make_shared<Linear>(config.lq_hidden_dim, config.hidden_size, true);
blocks["gate_modules." + std::to_string(i)] = std::make_shared<SigmaAwareGate>(config.hidden_size);
blocks["gate_modules." + std::to_string(i)] = std::make_shared<SigmaAwareGate>(config.hidden_size, config.lq_gate_per_token);
}
if (config.pit_lq_inject) {
blocks["pit_head"] = std::make_shared<Linear>(config.lq_hidden_dim, config.hidden_size, true);
}
}
@ -545,7 +618,27 @@ namespace Pid {
int64_t target_pW) {
auto conv0 = std::dynamic_pointer_cast<Conv2d>(blocks["latent_proj.0"]);
auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["latent_proj.2"]);
float z_to_patch_ratio = static_cast<float>(config.lq_sr_scale * config.lq_latent_down_factor) /
int64_t unpatchify_factor = config.lq_latent_unpatchify_factor;
if (unpatchify_factor > 1) {
int64_t latent_h = lq_latent->ne[1];
int64_t latent_w = lq_latent->ne[0];
lq_latent = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, lq_latent, 2, 0, 1, 3));
lq_latent = ggml_reshape_3d(ctx->ggml_ctx,
lq_latent,
lq_latent->ne[0],
lq_latent->ne[1] * lq_latent->ne[2],
lq_latent->ne[3]);
lq_latent = DiT::unpatchify(ctx->ggml_ctx,
lq_latent,
latent_h,
latent_w,
static_cast<int>(unpatchify_factor),
static_cast<int>(unpatchify_factor),
true);
}
int64_t effective_down_factor = config.lq_latent_down_factor / unpatchify_factor;
float z_to_patch_ratio = static_cast<float>(config.lq_sr_scale * effective_down_factor) /
static_cast<float>(config.patch_size);
GGML_ASSERT(z_to_patch_ratio >= 1.0f);
if (lq_latent->ne[0] != target_pW || lq_latent->ne[1] != target_pH) {
@ -558,8 +651,14 @@ namespace Pid {
GGML_SCALE_MODE_NEAREST);
}
if (config.lq_replicate_padding) {
lq_latent = replicate_pad_2d(ctx->ggml_ctx, lq_latent);
}
auto feat = conv0->forward(ctx, lq_latent);
feat = ggml_silu_inplace(ctx->ggml_ctx, feat);
if (config.lq_replicate_padding) {
feat = replicate_pad_2d(ctx->ggml_ctx, feat);
}
feat = conv2->forward(ctx, feat);
for (int i = 0; i < config.lq_num_res_blocks; ++i) {
auto block = std::dynamic_pointer_cast<PiDResBlock>(blocks["latent_proj." + std::to_string(3 + i)]);
@ -574,11 +673,15 @@ namespace Pid {
int num_outputs = static_cast<int>((config.patch_depth + config.lq_interval - 1) / config.lq_interval);
std::vector<ggml_tensor*> outputs;
outputs.reserve(num_outputs);
outputs.reserve(num_outputs + (config.pit_lq_inject ? 1 : 0));
for (int i = 0; i < num_outputs; ++i) {
auto head = std::dynamic_pointer_cast<Linear>(blocks["output_heads." + std::to_string(i)]);
outputs.push_back(head->forward(ctx, tokens));
}
if (config.pit_lq_inject) {
auto pit_head = std::dynamic_pointer_cast<Linear>(blocks["pit_head"]);
outputs.push_back(pit_head->forward(ctx, tokens));
}
return outputs;
}
};
@ -606,6 +709,9 @@ namespace Pid {
}
blocks["final_layer"] = std::make_shared<FinalLayer>(config.pixel_hidden_size, config.in_channels);
blocks["lq_proj"] = std::make_shared<LQProjection2D>(config);
if (config.pit_lq_inject) {
blocks["pit_lq_gate"] = std::make_shared<SigmaAwareGate>(config.hidden_size, config.lq_gate_per_token);
}
}
void init_params(ggml_context* ctx,
@ -654,6 +760,11 @@ namespace Pid {
y_emb = ggml_add(ctx->ggml_ctx, y_emb, y_pos);
std::vector<ggml_tensor*> lq_features = lq_proj->forward(ctx, lq_latent, Hs, Ws);
ggml_tensor* pit_lq_feature = nullptr;
if (config.pit_lq_inject) {
pit_lq_feature = lq_features.back();
lq_features.pop_back();
}
auto s = s_embedder->forward(ctx, x_patches);
@ -677,6 +788,10 @@ namespace Pid {
sd::ggml_graph_cut::mark_graph_cut(y_emb, "pid.patch_blocks." + std::to_string(i), "y");
}
s = ggml_silu(ctx->ggml_ctx, ggml_add(ctx->ggml_ctx, s, t_emb));
if (pit_lq_feature != nullptr) {
auto pit_lq_gate = std::dynamic_pointer_cast<SigmaAwareGate>(blocks["pit_lq_gate"]);
s = pit_lq_gate->forward(ctx, s, pit_lq_feature, degrade_sigma);
}
auto s_cond = ggml_reshape_2d(ctx->ggml_ctx, s, config.hidden_size, L * B);
auto pixels = pixel_embedder->forward(ctx, x, config.patch_size, pixel_pos_full);

View File

@ -103,7 +103,7 @@ namespace Qwen {
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
virtual ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* timesteps,
ggml_tensor* addition_t_cond = nullptr) {
// timesteps: [N,]
@ -416,10 +416,14 @@ namespace Qwen {
public:
QwenImageModel() {}
QwenImageModel(QwenImageConfig config)
QwenImageModel(QwenImageConfig config,
std::shared_ptr<QwenTimestepProjEmbeddings> time_text_embed = nullptr)
: config(config) {
int64_t inner_dim = config.num_attention_heads * config.attention_head_dim;
blocks["time_text_embed"] = std::shared_ptr<GGMLBlock>(new QwenTimestepProjEmbeddings(inner_dim, config.use_additional_t_cond));
if (time_text_embed == nullptr) {
time_text_embed = std::make_shared<QwenTimestepProjEmbeddings>(inner_dim, config.use_additional_t_cond);
}
blocks["time_text_embed"] = std::move(time_text_embed);
blocks["txt_norm"] = std::shared_ptr<GGMLBlock>(new RMSNorm(config.joint_attention_dim, 1e-6f));
blocks["img_in"] = std::shared_ptr<GGMLBlock>(new Linear(config.in_channels, inner_dim));
blocks["txt_in"] = std::shared_ptr<GGMLBlock>(new Linear(config.joint_attention_dim, inner_dim));
@ -517,7 +521,7 @@ namespace Qwen {
if (input->ne[3] == 1) {
input = ggml_reshape_4d(ctx->ggml_ctx, input, input->ne[0], input->ne[1], 1, input->ne[2]);
}
return DiT::patchify(ctx->ggml_ctx, input, 1, config.patch_size, config.patch_size, N);
return DiT::patchify_3d(ctx->ggml_ctx, input, 1, config.patch_size, config.patch_size, N);
};
auto img = patchify_input(x);
@ -715,8 +719,8 @@ namespace Qwen {
*diffusion_params.x,
*diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context),
diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
diffusion_params.ref_index_mode);
diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents,
diffusion_params.ref_image_params.ref_index_mode);
}
void test() {

View File

@ -6,6 +6,7 @@
#include "model.h"
#include "model/common/block.hpp"
#include "model/diffusion/animatediff.hpp"
#include "model/diffusion/model.hpp"
/*==================================================== UnetModel =====================================================*/
@ -29,6 +30,8 @@ struct UNetConfig {
bool tiny_unet = false;
int model_channels = 320;
int adm_in_channels = 2816; // only for VERSION_SDXL/SVD
bool enable_animatediff = false;
bool animatediff_has_mid_block = false;
static UNetConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
const std::string& prefix,
@ -84,6 +87,13 @@ struct UNetConfig {
return &it->second;
};
if (find_weight("motion_module.down_blocks.0.motion_modules.0.temporal_transformer.proj_in.weight") != nullptr) {
config.enable_animatediff = true;
if (find_weight("motion_module.mid_block.motion_modules.0.temporal_transformer.proj_in.weight") != nullptr) {
config.animatediff_has_mid_block = true;
}
}
if (const TensorStorage* input = find_weight("input_blocks.0.0.weight")) {
if (input->n_dims == 4) {
config.in_channels = static_cast<int>(input->ne[2]);
@ -473,6 +483,12 @@ public:
blocks["out.0"] = std::shared_ptr<GGMLBlock>(new GroupNorm32(ch)); // ch == model_channels
// out_1 is nn.SiLU()
blocks["out.2"] = std::shared_ptr<GGMLBlock>(new Conv2d(model_channels, out_channels, {3, 3}, {1, 1}, {1, 1}));
if (this->config.enable_animatediff) {
AnimateDiff::MotionModuleConfig mm_cfg;
mm_cfg.enable_mid_block = this->config.animatediff_has_mid_block;
blocks["motion_module"] = std::make_shared<AnimateDiff::AnimateDiffModel>(mm_cfg);
}
}
ggml_tensor* resblock_forward(std::string name,
@ -583,6 +599,42 @@ public:
ggml_set_name(h, "bench-start");
hs.push_back(h);
auto motion_root = config.enable_animatediff && num_video_frames > 1
? std::dynamic_pointer_cast<AnimateDiff::AnimateDiffModel>(blocks["motion_module"])
: nullptr;
auto apply_motion_input = [&](int input_block_idx, ggml_tensor* h_in) -> ggml_tensor* {
if (!motion_root)
return h_in;
int di = (input_block_idx - 1) / 3;
int mj = (input_block_idx - 1) % 3;
if (di < 0 || di >= (int)channel_mult.size() || mj < 0 || mj >= num_res_blocks)
return h_in;
auto mm = motion_root->motion("down_blocks." + std::to_string(di) + ".motion_modules." + std::to_string(mj));
if (!mm)
return h_in;
return mm->forward(ctx, h_in, num_video_frames);
};
auto apply_motion_output = [&](int output_block_idx, ggml_tensor* h_in) -> ggml_tensor* {
if (!motion_root)
return h_in;
int ui = output_block_idx / 3;
int mj = output_block_idx % 3;
if (ui < 0 || ui >= (int)channel_mult.size() || mj < 0 || mj > num_res_blocks)
return h_in;
auto mm = motion_root->motion("up_blocks." + std::to_string(ui) + ".motion_modules." + std::to_string(mj));
if (!mm)
return h_in;
return mm->forward(ctx, h_in, num_video_frames);
};
auto apply_motion_mid = [&](ggml_tensor* h_in) -> ggml_tensor* {
if (!motion_root)
return h_in;
auto mm = motion_root->motion("mid_block.motion_modules.0");
if (!mm)
return h_in;
return mm->forward(ctx, h_in, num_video_frames);
};
// input block 1-11
size_t len_mults = channel_mult.size();
int input_block_idx = 0;
@ -597,6 +649,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]
}
h = apply_motion_input(input_block_idx, h);
sd::ggml_graph_cut::mark_graph_cut(h, "unet.input_blocks." + std::to_string(input_block_idx), "h");
hs.push_back(h);
}
@ -624,6 +677,7 @@ public:
h = attention_layer_forward("middle_block.1", ctx, h, context, num_video_frames); // [N, 4*model_channels, h/8, w/8]
h = resblock_forward("middle_block.2", ctx, h, emb, num_video_frames); // [N, 4*model_channels, h/8, w/8]
}
h = apply_motion_mid(h);
}
sd::ggml_graph_cut::mark_graph_cut(h, "unet.middle_block", "h");
if (controls.size() > 0) {
@ -660,6 +714,8 @@ public:
up_sample_idx++;
}
h = apply_motion_output(output_block_idx, h);
if (i > 0 && j == num_res_blocks) {
if (tiny_unet) {
output_block_idx++;
@ -719,7 +775,9 @@ struct UNetModelRunner : public DiffusionModelRunner {
const sd::Tensor<float>& y_tensor = {},
int num_video_frames = -1,
const std::vector<sd::Tensor<float>>& controls_tensor = {},
float control_strength = 0.f) {
float control_strength = 0.f,
const sd::Tensor<float>& ip_context_tensor = {},
float ip_scale = 1.f) {
ggml_cgraph* gf = new_graph_custom(UNET_GRAPH_SIZE);
ggml_tensor* x = make_input(x_tensor);
@ -727,6 +785,7 @@ struct UNetModelRunner : public DiffusionModelRunner {
ggml_tensor* context = make_optional_input(context_tensor);
ggml_tensor* c_concat = make_optional_input(c_concat_tensor);
ggml_tensor* y = make_optional_input(y_tensor);
ggml_tensor* ip_context = make_optional_input(ip_context_tensor);
std::vector<ggml_tensor*> controls;
controls.reserve(controls_tensor.size());
for (const auto& control_tensor : controls_tensor) {
@ -738,6 +797,8 @@ struct UNetModelRunner : public DiffusionModelRunner {
}
auto runner_ctx = get_context();
runner_ctx.ip_context = ip_context;
runner_ctx.ip_scale = ip_scale;
ggml_tensor* out = unet.forward(&runner_ctx,
x,
@ -762,14 +823,16 @@ struct UNetModelRunner : public DiffusionModelRunner {
const sd::Tensor<float>& y = {},
int num_video_frames = -1,
const std::vector<sd::Tensor<float>>& controls = {},
float control_strength = 0.f) {
float control_strength = 0.f,
const sd::Tensor<float>& ip_context = {},
float ip_scale = 1.f) {
// x: [N, in_channels, h, w]
// timesteps: [N, ]
// context: [N, max_position, hidden_size]([N, 77, 768]) or [1, max_position, hidden_size]
// c_concat: [N, in_channels, h, w] or [1, in_channels, h, w]
// y: [N, adm_in_channels] or [1, adm_in_channels]
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(x, timesteps, context, c_concat, y, num_video_frames, controls, control_strength);
return build_graph(x, timesteps, context, c_concat, y, num_video_frames, controls, control_strength, ip_context, ip_scale);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
@ -789,7 +852,9 @@ struct UNetModelRunner : public DiffusionModelRunner {
tensor_or_empty(diffusion_params.y),
extra->num_video_frames,
extra->controls ? *extra->controls : empty_controls,
extra->control_strength);
extra->control_strength,
extra->ip_context ? *extra->ip_context : sd::Tensor<float>{},
extra->ip_scale);
}
void test() {

View File

@ -648,8 +648,8 @@ namespace ZImage {
*diffusion_params.x,
*diffusion_params.timesteps,
tensor_or_empty(diffusion_params.context),
diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
diffusion_params.ref_index_mode);
diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents,
diffusion_params.ref_image_params.ref_index_mode);
}
void test() {

View File

@ -79,9 +79,20 @@ namespace LLM {
int window_size = 112;
int num_position_embeddings = 0;
std::set<int> fullatt_block_indexes = {7, 15, 23, 31};
std::vector<int> deepstack_visual_indexes;
bool split_patch_embed = false;
};
struct ImageGrid {
int index = 0;
int size = 0;
int grid_h = 0;
int grid_w = 0;
};
using ImageEmbeds = std::vector<std::pair<int, sd::Tensor<float>>>;
using DeepStackImageEmbeds = std::vector<ImageEmbeds>;
struct LLMConfig {
LLMArch arch = LLMArch::QWEN2_5_VL;
int64_t num_layers = 28;
@ -93,6 +104,7 @@ namespace LLM {
bool qkv_bias = true;
bool attention_out_bias = false;
bool qk_norm = false;
bool final_norm = true;
bool rms_norm_add = false;
bool normalize_input = false;
int64_t vocab_size = 152064;
@ -200,7 +212,11 @@ namespace LLM {
config.vision.in_channels = tensor_storage.ne[2];
config.vision.hidden_size = tensor_storage.ne[3];
}
if (contains(name, "visual.patch_embed.bias")) {
// HF-format checkpoints keep the patch embed unsplit under a single name.
if (contains(name, "visual.patch_embed.proj.weight")) {
config.vision.patch_size = static_cast<int>(tensor_storage.ne[0]);
}
if (contains(name, "visual.patch_embed.bias") || contains(name, "visual.patch_embed.proj.bias")) {
config.vision.hidden_size = tensor_storage.ne[0];
}
if (contains(name, "visual.pos_embed.weight")) {
@ -253,9 +269,20 @@ namespace LLM {
if ((arch == LLMArch::QWEN3 || arch == LLMArch::QWEN3_VL) && config.num_layers == 28) {
config.num_heads = 16;
}
if (arch == LLMArch::QWEN3_VL && config.num_layers == 50 && config.hidden_size == 5120) {
config.num_heads = 64;
config.final_norm = false;
}
if (detected_vision_layers > 0) {
config.vision.num_layers = detected_vision_layers;
}
if (arch == LLMArch::QWEN3_VL) {
if (config.vision.num_layers == 24) {
config.vision.deepstack_visual_indexes = {5, 11, 17};
} else if (config.vision.num_layers == 27) {
config.vision.deepstack_visual_indexes = {8, 16, 24};
}
}
LOG_DEBUG("llm: num_layers = %" PRId64 ", vocab_size = %" PRId64 ", hidden_size = %" PRId64 ", intermediate_size = %" PRId64,
config.num_layers,
config.vocab_size,
@ -285,7 +312,7 @@ namespace LLM {
bool add_unit_offset = false)
: hidden_size(hidden_size), eps(eps), add_unit_offset(add_unit_offset) {}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
ggml_tensor* w = params["weight"];
if (ctx->weight_adapter) {
w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");
@ -537,6 +564,37 @@ namespace LLM {
return input_embed;
}
static ggml_tensor* add_deepstack_image_embeds(GGMLRunnerContext* ctx,
ggml_tensor* x,
const std::vector<std::pair<int, ggml_tensor*>>& image_embeds) {
if (image_embeds.empty()) {
return x;
}
GGML_ASSERT(x->ne[2] == 1);
auto raw_x = ggml_cast(ctx->ggml_ctx, x, image_embeds[0].second->type);
int64_t token_start = 0;
ggml_tensor* output = nullptr;
for (const auto& [index, image_embed] : image_embeds) {
GGML_ASSERT(index >= token_start);
GGML_ASSERT(index + image_embed->ne[1] <= raw_x->ne[1]);
if (index > token_start) {
auto text_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, token_start, index);
output = output == nullptr ? text_embed : ggml_concat(ctx->ggml_ctx, output, text_embed, 1);
}
auto visual_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, index, index + image_embed->ne[1]);
visual_embed = ggml_add(ctx->ggml_ctx, visual_embed, image_embed);
output = output == nullptr ? visual_embed : ggml_concat(ctx->ggml_ctx, output, visual_embed, 1);
token_start = index + image_embed->ne[1];
}
if (token_start < raw_x->ne[1]) {
auto text_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, token_start, raw_x->ne[1]);
output = output == nullptr ? text_embed : ggml_concat(ctx->ggml_ctx, output, text_embed, 1);
}
GGML_ASSERT(output != nullptr && output->ne[1] == raw_x->ne[1]);
return output;
}
struct VisionMLP : public GGMLBlock {
protected:
LLMVisionArch arch_;
@ -719,6 +777,33 @@ namespace LLM {
}
};
struct Qwen3VLDeepStackMerger : public GGMLBlock {
protected:
int64_t merge_dim;
public:
Qwen3VLDeepStackMerger(int64_t dim,
int64_t context_dim,
int64_t spatial_merge_size)
: merge_dim(context_dim * spatial_merge_size * spatial_merge_size) {
blocks["norm"] = std::make_shared<LayerNorm>(merge_dim, 1e-6f);
blocks["linear_fc1"] = std::make_shared<Linear>(merge_dim, merge_dim, true);
blocks["linear_fc2"] = std::make_shared<Linear>(merge_dim, dim, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto norm = std::dynamic_pointer_cast<LayerNorm>(blocks["norm"]);
auto linear_fc1 = std::dynamic_pointer_cast<Linear>(blocks["linear_fc1"]);
auto linear_fc2 = std::dynamic_pointer_cast<Linear>(blocks["linear_fc2"]);
x = ggml_reshape_2d(ctx->ggml_ctx, x, merge_dim, ggml_nelements(x) / merge_dim);
x = norm->forward(ctx, x);
x = linear_fc1->forward(ctx, x);
x = ggml_gelu_erf(ctx->ggml_ctx, x);
return linear_fc2->forward(ctx, x);
}
};
struct VisionAttention : public GGMLBlock {
protected:
bool llama_cpp_style;
@ -840,6 +925,7 @@ namespace LLM {
int spatial_merge_size;
int num_grid_per_side;
std::set<int> fullatt_block_indexes;
std::vector<int> deepstack_visual_indexes;
public:
VisionModel(bool llama_cpp_style,
@ -849,7 +935,8 @@ namespace LLM {
num_layers(vision_params.num_layers),
spatial_merge_size(vision_params.spatial_merge_size),
num_grid_per_side(vision_params.num_position_embeddings > 0 ? static_cast<int>(std::sqrt(vision_params.num_position_embeddings)) : 0),
fullatt_block_indexes(vision_params.fullatt_block_indexes) {
fullatt_block_indexes(vision_params.fullatt_block_indexes),
deepstack_visual_indexes(vision_params.deepstack_visual_indexes) {
blocks["patch_embed"] = std::shared_ptr<GGMLBlock>(new VisionPatchEmbed(vision_params.split_patch_embed,
arch_,
vision_params.patch_size,
@ -871,6 +958,11 @@ namespace LLM {
vision_params.out_hidden_size,
vision_params.hidden_size,
spatial_merge_size));
for (size_t i = 0; i < deepstack_visual_indexes.size(); ++i) {
blocks["deepstack_merger_list." + std::to_string(i)] = std::make_shared<Qwen3VLDeepStackMerger>(vision_params.out_hidden_size,
vision_params.hidden_size,
spatial_merge_size);
}
}
std::shared_ptr<Embedding> pos_embedder() {
@ -889,7 +981,7 @@ namespace LLM {
return spatial_merge_size;
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
std::vector<ggml_tensor*> forward_outputs(GGMLRunnerContext* ctx,
ggml_tensor* pixel_values,
ggml_tensor* pe,
ggml_tensor* window_index,
@ -915,6 +1007,7 @@ namespace LLM {
x = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0] / spatial_merge_size / spatial_merge_size, x->ne[1] * spatial_merge_size * spatial_merge_size, x->ne[2], x->ne[3]);
}
std::vector<ggml_tensor*> deepstack_outputs;
for (int i = 0; i < num_layers; i++) {
auto block = std::dynamic_pointer_cast<VisionBlock>(blocks["blocks." + std::to_string(i)]);
@ -923,7 +1016,11 @@ namespace LLM {
mask = nullptr;
}
x = block->forward(ctx, x, pe, mask);
if (i == 0) {
auto deepstack_it = std::find(deepstack_visual_indexes.begin(), deepstack_visual_indexes.end(), i);
if (deepstack_it != deepstack_visual_indexes.end()) {
size_t deepstack_index = static_cast<size_t>(std::distance(deepstack_visual_indexes.begin(), deepstack_it));
auto deepstack_merger = std::dynamic_pointer_cast<Qwen3VLDeepStackMerger>(blocks["deepstack_merger_list." + std::to_string(deepstack_index)]);
deepstack_outputs.push_back(deepstack_merger->forward(ctx, x));
}
sd::ggml_graph_cut::mark_graph_cut(x, "llm.vision.blocks." + std::to_string(i), "x");
}
@ -935,7 +1032,19 @@ namespace LLM {
x = ggml_get_rows(ctx->ggml_ctx, x, window_inverse_index);
}
return x;
std::vector<ggml_tensor*> outputs = {x};
outputs.insert(outputs.end(), deepstack_outputs.begin(), deepstack_outputs.end());
return outputs;
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* pixel_values,
ggml_tensor* pe,
ggml_tensor* window_index,
ggml_tensor* window_inverse_index,
ggml_tensor* window_mask,
ggml_tensor* pos_embeds = nullptr) {
return forward_outputs(ctx, pixel_values, pe, window_index, window_inverse_index, window_mask, pos_embeds)[0];
}
};
@ -1259,8 +1368,10 @@ namespace LLM {
for (int i = 0; i < num_layers; i++) {
blocks["layers." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new TransformerBlock(config, i));
}
if (config.final_norm) {
blocks["norm"] = std::shared_ptr<GGMLBlock>(new LLMRMSNorm(config.hidden_size, config.rms_norm_eps, config.rms_norm_add));
}
}
ggml_tensor* embed(GGMLRunnerContext* ctx,
ggml_tensor* input_ids) {
@ -1274,9 +1385,11 @@ namespace LLM {
ggml_tensor* input_pos,
ggml_tensor* attention_mask,
std::set<int> out_layers,
const std::vector<std::vector<std::pair<int, ggml_tensor*>>>& deepstack_image_embeds = {},
ggml_tensor* sliding_attention_mask = nullptr,
bool return_all_hidden_states = false) {
auto norm = std::dynamic_pointer_cast<LLMRMSNorm>(blocks["norm"]);
auto norm = config.final_norm ? std::dynamic_pointer_cast<LLMRMSNorm>(blocks["norm"])
: nullptr;
std::vector<ggml_tensor*> intermediate_outputs;
if (config.normalize_input) {
@ -1291,6 +1404,9 @@ namespace LLM {
auto block = std::dynamic_pointer_cast<TransformerBlock>(blocks["layers." + std::to_string(i)]);
x = block->forward(ctx, x, input_pos, attention_mask, sliding_attention_mask);
if (i < static_cast<int>(deepstack_image_embeds.size())) {
x = add_deepstack_image_embeds(ctx, x, deepstack_image_embeds[static_cast<size_t>(i)]);
}
if (return_all_hidden_states || out_layers.size() > 1) {
x = ggml_cont(ctx->ggml_ctx, x);
}
@ -1304,7 +1420,7 @@ namespace LLM {
}
}
auto normed_x = norm->forward(ctx, x);
auto normed_x = norm == nullptr ? x : norm->forward(ctx, x);
if (return_all_hidden_states) {
intermediate_outputs.push_back(normed_x);
x = intermediate_outputs[0];
@ -1332,6 +1448,7 @@ namespace LLM {
ggml_tensor* attention_mask,
ggml_tensor* sliding_attention_mask,
std::vector<std::pair<int, ggml_tensor*>> image_embeds,
const std::vector<std::vector<std::pair<int, ggml_tensor*>>>& deepstack_image_embeds,
std::set<int> out_layers,
bool return_all_hidden_states = false) {
// input_ids: [N, n_token]
@ -1343,6 +1460,7 @@ namespace LLM {
input_pos,
attention_mask,
std::move(out_layers),
deepstack_image_embeds,
sliding_attention_mask,
return_all_hidden_states);
}
@ -1368,6 +1486,7 @@ namespace LLM {
ggml_tensor* attention_mask,
ggml_tensor* sliding_attention_mask,
std::vector<std::pair<int, ggml_tensor*>> image_embeds,
const std::vector<std::vector<std::pair<int, ggml_tensor*>>>& deepstack_image_embeds,
std::set<int> out_layers,
bool return_all_hidden_states = false) {
// input_ids: [N, n_token]
@ -1379,6 +1498,7 @@ namespace LLM {
attention_mask,
sliding_attention_mask,
image_embeds,
deepstack_image_embeds,
out_layers,
return_all_hidden_states);
return x;
@ -1520,7 +1640,8 @@ namespace LLM {
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) {
std::array<std::vector<float>, 4>& pos_embed_weight_data,
std::vector<ggml_tensor*>* output_tensors = nullptr) {
GGML_ASSERT(image->ne[1] % (vision_params.patch_size * vision_params.spatial_merge_size) == 0);
GGML_ASSERT(image->ne[0] % (vision_params.patch_size * vision_params.spatial_merge_size) == 0);
@ -1552,7 +1673,11 @@ namespace LLM {
int pos_len = static_cast<int>(pe_vec.size() / head_dim / 2);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, head_dim / 2, pos_len);
runner->set_backend_tensor_data(pe, pe_vec.data());
return vision_model->forward(runner_ctx, pixel_values, pe, nullptr, nullptr, nullptr, pos_embeds);
auto outputs = vision_model->forward_outputs(runner_ctx, pixel_values, pe, nullptr, nullptr, nullptr, pos_embeds);
if (output_tensors != nullptr) {
*output_tensors = outputs;
}
return outputs[0];
}
int llm_grid_h = grid_h / vision_params.spatial_merge_size;
@ -1618,7 +1743,11 @@ namespace LLM {
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, head_dim / 2, pos_len);
runner->set_backend_tensor_data(pe, pe_vec.data());
return vision_model->forward(runner_ctx, pixel_values, pe, window_index, window_inverse_index, window_mask);
auto output = vision_model->forward(runner_ctx, pixel_values, pe, window_index, window_inverse_index, window_mask);
if (output_tensors != nullptr) {
*output_tensors = {output};
}
return output;
}
public:
@ -1653,12 +1782,17 @@ namespace LLM {
model.get_param_tensors(tensors, prefix);
}
void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) {
model.get_param_tensor_ops(tensor_ops);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* input_ids,
ggml_tensor* input_pos,
ggml_tensor* attention_mask,
ggml_tensor* sliding_attention_mask,
std::vector<std::pair<int, ggml_tensor*>> image_embeds,
const std::vector<std::vector<std::pair<int, ggml_tensor*>>>& deepstack_image_embeds,
std::set<int> out_layers,
bool return_all_hidden_states = false) {
auto hidden_states = model.forward(ctx,
@ -1667,6 +1801,7 @@ namespace LLM {
attention_mask,
sliding_attention_mask,
image_embeds,
deepstack_image_embeds,
out_layers,
return_all_hidden_states); // [N, n_token, hidden_size]
return hidden_states;
@ -1685,7 +1820,9 @@ namespace LLM {
ggml_cgraph* build_graph(const sd::Tensor<int32_t>& input_ids_tensor,
const sd::Tensor<float>& attention_mask_tensor,
const std::vector<std::pair<int, sd::Tensor<float>>>& image_embeds_tensor,
const ImageEmbeds& image_embeds_tensor,
const DeepStackImageEmbeds& deepstack_image_embeds_tensor,
const std::vector<ImageGrid>& image_grids,
std::set<int> out_layers,
bool return_all_hidden_states = false) {
ggml_cgraph* gf = new_graph_custom(LLM_GRAPH_SIZE);
@ -1696,6 +1833,13 @@ namespace LLM {
ggml_tensor* embed = make_input(embed_tensor);
image_embeds.emplace_back(idx, embed);
}
std::vector<std::vector<std::pair<int, ggml_tensor*>>> deepstack_image_embeds(deepstack_image_embeds_tensor.size());
for (size_t layer = 0; layer < deepstack_image_embeds_tensor.size(); ++layer) {
deepstack_image_embeds[layer].reserve(deepstack_image_embeds_tensor[layer].size());
for (const auto& [idx, embed_tensor] : deepstack_image_embeds_tensor[layer]) {
deepstack_image_embeds[layer].emplace_back(idx, make_input(embed_tensor));
}
}
int64_t n_tokens = input_ids->ne[0];
if (config.arch == LLMArch::MISTRAL_SMALL_3_2 ||
@ -1716,6 +1860,30 @@ namespace LLM {
input_pos_vec[2 * n_tokens + i] = i;
input_pos_vec[3 * n_tokens + i] = 0;
}
if (config.arch == LLMArch::QWEN3_VL && !image_grids.empty()) {
int offset = 0;
for (const auto& grid : image_grids) {
int end = grid.index + grid.size;
int grid_h = grid.grid_h / config.vision.spatial_merge_size;
int grid_w = grid.grid_w / config.vision.spatial_merge_size;
int len_max = std::max(grid_h, grid_w);
int next_pos = grid.index + len_max + offset;
GGML_ASSERT(grid.index >= 0 && end <= n_tokens);
GGML_ASSERT(grid_h > 0 && grid_w > 0 && grid.size == grid_h * grid_w);
for (int token = end; token < n_tokens; ++token) {
int pos = next_pos + token - end;
input_pos_vec[token] = pos;
input_pos_vec[n_tokens + token] = pos;
input_pos_vec[2 * n_tokens + token] = pos;
}
for (int token = 0; token < grid.size; ++token) {
input_pos_vec[grid.index + token] = grid.index + offset;
input_pos_vec[n_tokens + grid.index + token] = grid.index + offset + token / grid_w;
input_pos_vec[2 * n_tokens + grid.index + token] = grid.index + offset + token % grid_w;
}
offset += len_max - grid.size;
}
}
}
auto input_pos = ggml_new_tensor_1d(compute_ctx,
@ -1773,6 +1941,7 @@ namespace LLM {
attention_mask,
sliding_attention_mask,
image_embeds,
deepstack_image_embeds,
out_layers,
return_all_hidden_states);
@ -1784,16 +1953,20 @@ namespace LLM {
sd::Tensor<float> compute(const int n_threads,
const sd::Tensor<int32_t>& input_ids,
const sd::Tensor<float>& attention_mask,
const std::vector<std::pair<int, sd::Tensor<float>>>& image_embeds,
const ImageEmbeds& image_embeds,
std::set<int> out_layers,
bool return_all_hidden_states = false,
bool auto_free = true,
bool free_compute_buffer = true,
bool free_compute_params = true) {
bool free_compute_params = true,
const DeepStackImageEmbeds& deepstack_image_embeds = {},
const std::vector<ImageGrid>& image_grids = {}) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(input_ids,
attention_mask,
image_embeds,
deepstack_image_embeds,
image_grids,
out_layers,
return_all_hidden_states);
};
@ -1843,6 +2016,24 @@ namespace LLM {
pos_embed_weight_data_);
}
std::vector<ggml_tensor*> encode_image_outputs(GGMLRunnerContext* runner_ctx, ggml_tensor* image) {
std::vector<ggml_tensor*> outputs;
encode_image_common(this,
compute_ctx,
runner_ctx,
image,
config.vision,
model.vision_model(),
window_index_vec,
window_inverse_index_vec,
window_mask_vec,
pe_vec,
pos_embed_idx_data_,
pos_embed_weight_data_,
&outputs);
return outputs;
}
ggml_cgraph* build_encode_image_graph(const sd::Tensor<float>& image_tensor) {
ggml_cgraph* gf = new_graph_custom(LLM_GRAPH_SIZE);
ggml_tensor* image = make_input(image_tensor);
@ -1867,6 +2058,166 @@ namespace LLM {
};
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params));
}
ggml_cgraph* build_encode_image_outputs_graph(const sd::Tensor<float>& image_tensor) {
ggml_cgraph* gf = new_graph_custom(LLM_GRAPH_SIZE);
ggml_tensor* image = make_input(image_tensor);
auto runner_ctx = get_context();
auto outputs = encode_image_outputs(&runner_ctx, image);
GGML_ASSERT(!outputs.empty());
auto combined = outputs[0];
for (size_t i = 1; i < outputs.size(); ++i) {
combined = ggml_concat(compute_ctx, combined, outputs[i], 0);
}
ggml_build_forward_expand(gf, combined);
return gf;
}
static sd::Tensor<float> process_video_block_tensor(const sd::Tensor<float>& frames,
const LLMVisionConfig& vision_params) {
GGML_ASSERT(frames.dim() == 5);
GGML_ASSERT(frames.shape()[2] == vision_params.temporal_patch_size);
GGML_ASSERT(frames.shape()[3] == vision_params.in_channels);
GGML_ASSERT(frames.shape()[4] == 1);
int64_t width = frames.shape()[0];
int64_t height = frames.shape()[1];
int64_t temporal = frames.shape()[2];
int64_t channels = frames.shape()[3];
int64_t patch = vision_params.patch_size;
int64_t merge = vision_params.spatial_merge_size;
int64_t grid_w = width / patch;
int64_t grid_h = height / patch;
int64_t feature = channels * temporal * patch * patch;
int64_t token_count = grid_h * grid_w;
sd::Tensor<float> output({feature, token_count});
int64_t token = 0;
for (int64_t block_h = 0; block_h < grid_h / merge; ++block_h) {
for (int64_t block_w = 0; block_w < grid_w / merge; ++block_w) {
for (int64_t inner_h = 0; inner_h < merge; ++inner_h) {
for (int64_t inner_w = 0; inner_w < merge; ++inner_w) {
int64_t patch_h = block_h * merge + inner_h;
int64_t patch_w = block_w * merge + inner_w;
int64_t offset = 0;
for (int64_t c = 0; c < channels; ++c) {
for (int64_t t = 0; t < temporal; ++t) {
for (int64_t y = 0; y < patch; ++y) {
for (int64_t x = 0; x < patch; ++x) {
output.index(offset++, token) =
frames.index(patch_w * patch + x,
patch_h * patch + y,
t,
c,
0);
}
}
}
}
++token;
}
}
}
}
return output;
}
ggml_cgraph* build_encode_video_block_outputs_graph(const sd::Tensor<float>& pixel_values_tensor,
int grid_h,
int grid_w) {
ggml_cgraph* gf = new_graph_custom(LLM_GRAPH_SIZE);
auto pixel_values = make_input(pixel_values_tensor);
auto runner_ctx = get_context();
auto vision = model.vision_model();
int head_dim = static_cast<int>(config.vision.hidden_size / config.vision.num_heads);
auto pos_embeds = build_patch_pos_embeds(&runner_ctx, vision, grid_h, grid_w);
window_index_vec.resize(static_cast<size_t>((grid_h / config.vision.spatial_merge_size) *
(grid_w / config.vision.spatial_merge_size)));
for (int i = 0; i < static_cast<int>(window_index_vec.size()); ++i) {
window_index_vec[static_cast<size_t>(i)] = i;
}
pe_vec = Rope::gen_qwen2vl_pe(grid_h,
grid_w,
config.vision.spatial_merge_size,
window_index_vec,
10000,
{head_dim / 2, head_dim / 2});
int pos_len = static_cast<int>(pe_vec.size() / head_dim / 2);
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, head_dim / 2, pos_len);
set_backend_tensor_data(pe, pe_vec.data());
auto outputs = vision->forward_outputs(&runner_ctx,
pixel_values,
pe,
nullptr,
nullptr,
nullptr,
pos_embeds);
GGML_ASSERT(!outputs.empty());
auto combined = outputs[0];
for (size_t i = 1; i < outputs.size(); ++i) {
combined = ggml_concat(compute_ctx, combined, outputs[i], 0);
}
ggml_build_forward_expand(gf, combined);
return gf;
}
std::vector<sd::Tensor<float>> encode_image_outputs(const int n_threads,
const sd::Tensor<float>& image,
bool auto_free = false,
bool free_compute_buffer = false,
bool free_compute_params = false) {
auto get_graph = [&]() -> ggml_cgraph* {
return build_encode_image_outputs_graph(image);
};
auto combined = take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params));
if (combined.empty()) {
return {};
}
size_t output_count = config.vision.deepstack_visual_indexes.size() + 1;
GGML_ASSERT(combined.shape()[0] == config.hidden_size * static_cast<int64_t>(output_count));
std::vector<sd::Tensor<float>> outputs;
outputs.reserve(output_count);
for (size_t i = 0; i < output_count; ++i) {
outputs.push_back(sd::ops::slice(combined,
0,
static_cast<int64_t>(i) * config.hidden_size,
static_cast<int64_t>(i + 1) * config.hidden_size));
}
return outputs;
}
std::vector<sd::Tensor<float>> encode_video_block_outputs(const int n_threads,
const sd::Tensor<float>& frames,
bool auto_free = false,
bool free_compute_buffer = false,
bool free_compute_params = false) {
int grid_h = static_cast<int>(frames.shape()[1] / config.vision.patch_size);
int grid_w = static_cast<int>(frames.shape()[0] / config.vision.patch_size);
auto pixel_values = process_video_block_tensor(frames, config.vision);
auto get_graph = [&]() -> ggml_cgraph* {
return build_encode_video_block_outputs_graph(pixel_values, grid_h, grid_w);
};
auto combined = take_or_empty(GGMLRunner::compute<float>(get_graph,
n_threads,
auto_free,
free_compute_buffer,
free_compute_params));
if (combined.empty()) {
return {};
}
size_t output_count = config.vision.deepstack_visual_indexes.size() + 1;
GGML_ASSERT(combined.shape()[0] == config.hidden_size * static_cast<int64_t>(output_count));
std::vector<sd::Tensor<float>> outputs;
outputs.reserve(output_count);
for (size_t i = 0; i < output_count; ++i) {
outputs.push_back(sd::ops::slice(combined,
0,
static_cast<int64_t>(i) * config.hidden_size,
static_cast<int64_t>(i + 1) * config.hidden_size));
}
return outputs;
}
};
struct LLMEmbedder {

View File

@ -18,6 +18,7 @@
struct T5Config {
int64_t num_layers = 24;
int64_t model_dim = 4096;
int64_t inner_dim = 4096;
int64_t ff_dim = 10240;
int64_t num_heads = 64;
int64_t vocab_size = 32128;
@ -53,6 +54,7 @@ struct T5Config {
if (q->n_dims == 2) {
config.model_dim = q->ne[0];
int64_t inner_dim = q->ne[1];
config.inner_dim = inner_dim;
// Flan-T5/T5 uses d_kv=64 for common sizes.
if (inner_dim % 64 == 0) {
config.num_heads = inner_dim / 64;
@ -357,7 +359,7 @@ public:
: config(config) {
blocks["encoder"] = std::shared_ptr<GGMLBlock>(new T5Stack(config.num_layers,
config.model_dim,
config.model_dim,
config.inner_dim,
config.ff_dim,
config.num_heads,
config.relative_attention));

View File

@ -0,0 +1,28 @@
#ifndef __SD_MODEL_VAE_AUDIO_VAE_HPP__
#define __SD_MODEL_VAE_AUDIO_VAE_HPP__
#include "core/ggml_extend.hpp"
struct AudioVAERunner : public GGMLRunner {
AudioVAERunner(ggml_backend_t backend,
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: GGMLRunner(backend, weight_manager) {}
virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) = 0;
virtual size_t get_params_mem_size() = 0;
virtual std::string get_desc() = 0;
virtual sd::Tensor<float> encode(int n_threads,
const sd::Tensor<float>& waveform) {
SD_UNUSED(n_threads);
SD_UNUSED(waveform);
return {};
}
virtual sd::Tensor<float> decode(int n_threads,
const sd::Tensor<float>& latent_tensor) = 0;
virtual int input_sample_rate() const {
return output_sample_rate();
}
virtual int output_sample_rate() const = 0;
};
#endif // __SD_MODEL_VAE_AUDIO_VAE_HPP__

View File

@ -0,0 +1,834 @@
#ifndef __SD_MODEL_VAE_HUNYUAN_VAE_HPP__
#define __SD_MODEL_VAE_HUNYUAN_VAE_HPP__
#include <algorithm>
#include <cmath>
#include <map>
#include <memory>
#include <string>
#include <tuple>
#include <utility>
#include <vector>
#include "model/vae/wan_vae.hpp"
#include "model_manager.h"
namespace Hunyuan {
constexpr int HUNYUAN_VIDEO_VAE_GRAPH_SIZE = 65536;
constexpr int HUNYUAN_VIDEO_VAE_GRAPH_SIZE_PER_LATENT_FRAME = 8192;
constexpr int HUNYUAN_VIDEO_VAE_TEMPORAL_CHUNK_SIZE = 1;
struct TemporalConvCarry {
const std::vector<ggml_tensor*>* input = nullptr;
std::vector<ggml_tensor*>* output = nullptr;
size_t input_index = 0;
bool is_continuation() const {
return input != nullptr;
}
ggml_tensor* take() {
GGML_ASSERT(input != nullptr && input_index < input->size());
return (*input)[input_index++];
}
void push(ggml_tensor* tensor) {
if (output != nullptr) {
output->push_back(tensor);
}
}
void finish() const {
GGML_ASSERT(input == nullptr || input_index == input->size());
}
};
static ggml_tensor* repeat_interleave_channels(GGMLRunnerContext* ctx,
ggml_tensor* x,
int64_t repeats,
int64_t width,
int64_t height,
int64_t frames) {
GGML_ASSERT(repeats > 0);
GGML_ASSERT(width * height * frames == x->ne[0] * x->ne[1] * x->ne[2]);
int64_t channels = x->ne[3];
if (repeats == 1) {
return ggml_reshape_4d(ctx->ggml_ctx, x, width, height, frames, channels);
}
x = ggml_reshape_3d(ctx->ggml_ctx, x, width * height * frames, 1, channels);
auto target = ggml_new_tensor_3d(ctx->ggml_ctx, x->type, width * height * frames, repeats, channels);
x = ggml_repeat(ctx->ggml_ctx, x, target);
return ggml_reshape_4d(ctx->ggml_ctx, x, width, height, frames, channels * repeats);
}
class CausalConv3d : public GGMLBlock {
protected:
std::tuple<int, int, int> kernel_size;
public:
CausalConv3d(int64_t in_channels,
int64_t out_channels,
std::tuple<int, int, int> kernel_size,
std::tuple<int, int, int> stride = {1, 1, 1},
std::tuple<int, int, int> padding = {0, 0, 0},
std::tuple<int, int, int> dilation = {1, 1, 1},
bool bias = true)
: kernel_size(kernel_size) {
blocks["conv"] = std::make_shared<Conv3d>(in_channels, out_channels, kernel_size, stride, padding, dilation, bias);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
TemporalConvCarry* carry = nullptr) {
// x: [N*IC, ID, IH, IW]
// result: x: [N*OC, OD, OH, OW]
// assert N == 1
auto conv = std::dynamic_pointer_cast<Conv3d>(blocks["conv"]);
int pad_w = std::get<2>(kernel_size) / 2;
int pad_h = std::get<1>(kernel_size) / 2;
int pad_t = std::get<0>(kernel_size) - 1;
std::vector<ggml_tensor*> temporal_frames;
temporal_frames.reserve(x->ne[2] + pad_t);
if (pad_t > 0) {
if (carry != nullptr && carry->is_continuation()) {
auto previous = carry->take();
GGML_ASSERT(previous->ne[2] <= pad_t);
for (int64_t frame = 0; frame < previous->ne[2]; frame++) {
temporal_frames.push_back(ggml_ext_slice(ctx->ggml_ctx, previous, 2, frame, frame + 1));
}
for (int64_t frame = previous->ne[2]; frame < pad_t; frame++) {
temporal_frames.push_back(ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1));
}
} else {
auto first = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1);
for (int frame = 0; frame < pad_t; frame++) {
temporal_frames.push_back(first);
}
}
}
for (int64_t frame = 0; frame < x->ne[2]; frame++) {
temporal_frames.push_back(ggml_ext_slice(ctx->ggml_ctx, x, 2, frame, frame + 1));
}
if (pad_t > 0 && carry != nullptr && carry->output != nullptr) {
ggml_tensor* next = nullptr;
for (int frame = pad_t; frame > 0; frame--) {
auto item = temporal_frames[temporal_frames.size() - frame];
next = next == nullptr ? item : ggml_concat(ctx->ggml_ctx, next, item, 2);
}
carry->push(ggml_cont(ctx->ggml_ctx, next));
}
ggml_tensor* padded = nullptr;
for (auto frame : temporal_frames) {
padded = padded == nullptr ? frame : ggml_concat(ctx->ggml_ctx, padded, frame, 2);
}
auto replicate_pad = [&](ggml_tensor* input, int dim, int left, int right) {
if (left > 0) {
auto first = ggml_ext_slice(ctx->ggml_ctx, input, dim, 0, 1);
for (int i = 0; i < left; i++) {
input = ggml_concat(ctx->ggml_ctx, first, input, dim);
}
}
if (right > 0) {
auto last = ggml_ext_slice(ctx->ggml_ctx, input, dim, input->ne[dim] - 1, input->ne[dim]);
for (int i = 0; i < right; i++) {
input = ggml_concat(ctx->ggml_ctx, input, last, dim);
}
}
return input;
};
padded = replicate_pad(padded, 0, pad_w, pad_w);
padded = replicate_pad(padded, 1, pad_h, pad_h);
return conv->forward(ctx, padded);
}
};
class AttnBlock : public UnaryBlock {
protected:
int64_t in_channels;
public:
AttnBlock(int64_t in_channels)
: in_channels(in_channels) {
blocks["norm"] = std::make_shared<WAN::RMS_norm>(in_channels);
blocks["q"] = std::make_shared<Conv3d>(in_channels, in_channels, std::tuple{1, 1, 1});
blocks["k"] = std::make_shared<Conv3d>(in_channels, in_channels, std::tuple{1, 1, 1});
blocks["v"] = std::make_shared<Conv3d>(in_channels, in_channels, std::tuple{1, 1, 1});
blocks["proj_out"] = std::make_shared<Conv3d>(in_channels, in_channels, std::tuple{1, 1, 1});
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x) override {
// x: [b*c, t, h, w]
auto norm = std::dynamic_pointer_cast<WAN::RMS_norm>(blocks["norm"]);
auto q_proj = std::dynamic_pointer_cast<UnaryBlock>(blocks["q"]);
auto k_proj = std::dynamic_pointer_cast<UnaryBlock>(blocks["k"]);
auto v_proj = std::dynamic_pointer_cast<UnaryBlock>(blocks["v"]);
auto proj_out = std::dynamic_pointer_cast<UnaryBlock>(blocks["proj_out"]);
const int64_t b = x->ne[3] / in_channels;
auto identity = x;
x = norm->forward(ctx, x);
const int64_t c = x->ne[3] / b;
const int64_t t = x->ne[2];
const int64_t h = x->ne[1];
const int64_t w = x->ne[0];
auto q = q_proj->forward(ctx, x); // [b*c, t, h, w]
auto k = k_proj->forward(ctx, x); // [b*c, t, h, w]
auto v = v_proj->forward(ctx, x); // [b*c, t, h, w]
q = ggml_reshape_3d(ctx->ggml_ctx, q, w * h * t, c, b); // [b, c, t*h*w]
q = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, q, 1, 0, 2, 3)); // [b, t*h*w, c]
k = ggml_reshape_3d(ctx->ggml_ctx, k, w * h * t, c, b); // [b, c, t*h*w]
k = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, k, 1, 0, 2, 3)); // [b, t*h*w, c]
v = ggml_reshape_3d(ctx->ggml_ctx, v, w * h * t, c, b); // [b, c, t*h*w]
v = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, v, 1, 0, 2, 3)); // [b, t*h*w, c]
x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled); // [b, t*h*w, c]
x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); // [b, c, t*h*w]
x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, t, c * b); // [b*c, t, h, w]
x = proj_out->forward(ctx, x);
x = ggml_add(ctx->ggml_ctx, x, identity);
return x;
}
};
class ResnetBlock : public UnaryBlock {
protected:
int64_t in_channels;
int64_t out_channels;
public:
ResnetBlock(int64_t in_channels,
int64_t out_channels)
: in_channels(in_channels),
out_channels(out_channels) {
blocks["norm1"] = std::make_shared<WAN::RMS_norm>(in_channels);
blocks["conv1"] = std::make_shared<CausalConv3d>(in_channels, out_channels, std::tuple{3, 3, 3});
blocks["norm2"] = std::make_shared<WAN::RMS_norm>(out_channels);
blocks["conv2"] = std::make_shared<CausalConv3d>(out_channels, out_channels, std::tuple{3, 3, 3});
if (out_channels != in_channels) {
blocks["nin_shortcut"] = std::make_shared<CausalConv3d>(in_channels, out_channels, std::tuple{1, 1, 1});
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
return forward(ctx, x, nullptr);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
TemporalConvCarry* carry) {
// x: [B*IC, IT, OH, OW]
// return: [B*OC, OT, OH, OW]
auto norm1 = std::dynamic_pointer_cast<WAN::RMS_norm>(blocks["norm1"]);
auto conv1 = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv1"]);
auto norm2 = std::dynamic_pointer_cast<WAN::RMS_norm>(blocks["norm2"]);
auto conv2 = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv2"]);
auto h = x;
h = norm1->forward(ctx, h);
h = ggml_silu_inplace(ctx->ggml_ctx, h); // swish
h = conv1->forward(ctx, h, carry);
h = norm2->forward(ctx, h);
h = ggml_silu_inplace(ctx->ggml_ctx, h); // swish
// dropout, skip for inference
h = conv2->forward(ctx, h, carry);
// skip connection
if (out_channels != in_channels) {
auto nin_shortcut = std::dynamic_pointer_cast<CausalConv3d>(blocks["nin_shortcut"]);
x = nin_shortcut->forward(ctx, x); // [B*OC, OT, OH, OW]
}
h = ggml_add(ctx->ggml_ctx, h, x);
return h; // [B*OC, OT, OH, OW]
}
};
class Upsample : public GGMLBlock {
protected:
int64_t in_channels;
int64_t out_channels;
int64_t factor_t;
int64_t factor_s;
int64_t factor;
int64_t repeats;
public:
Upsample(int64_t in_channels, int64_t out_channels, bool add_temporal_upsample)
: in_channels(in_channels), out_channels(out_channels) {
if (add_temporal_upsample) {
factor_t = 2;
} else {
factor_t = 1;
}
factor_s = 2;
factor = factor_t * factor_s * factor_s;
GGML_ASSERT(out_channels * factor % in_channels == 0);
repeats = out_channels * factor / in_channels;
blocks["conv"] = std::make_shared<CausalConv3d>(in_channels, out_channels * factor, std::tuple{3, 3, 3});
}
static ggml_tensor* _pixel_shuffle_3d(GGMLRunnerContext* ctx,
ggml_tensor* x,
int64_t factor_t,
int64_t factor_s,
int64_t B = 1) {
// x: [B*factor*C, T, H, W]
// return: [B*C, T*factor_t, H*factor_s, W*factor_s]
GGML_ASSERT(B == 1);
int64_t factor = factor_t * factor_s * factor_s;
int64_t C = x->ne[3] / factor;
int64_t T = x->ne[2];
int64_t H = x->ne[1];
int64_t W = x->ne[0];
x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H * T, C, factor); // [factor, C, T*H, W]
x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2)); // [C, factor, T*H, W]
x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H * T, factor_s, factor_s * factor_t * C); // [C*factor_t*factor_s, factor_s, T*H, W]
x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3)); // [C*factor_t*factor_s, T*H, W, factor_s]
x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W, H * T, factor_s, factor_t * C); // [C*factor_t, factor_s, T*H, W*factor_s]
x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3)); // [C*factor_t, T*H, factor_s, W*factor_s]
x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W * factor_s * H, T, factor_t, C); // [C, factor_t, T, H*factor_s*W*factor_s]
x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3)); // [C, T, factor_t, H*factor_s*W*factor_s]
x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W, factor_s * H, factor_t * T, C); // [C, T*factor_t, H*factor_s, W*factor_s]
return x;
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
TemporalConvCarry* carry = nullptr) {
// x: [B*IC, T, H, W]
// return: [B*OC, 1 + (T - 1)*factor_t, H*factor_s, W*factor_s]
const int64_t B = x->ne[3] / in_channels;
GGML_ASSERT(B == 1);
auto conv = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv"]);
const bool continuation = carry != nullptr && carry->is_continuation();
auto h = conv->forward(ctx, x, carry); // [B*factor*OC, T, H, W]
ggml_tensor* shortcut = nullptr;
if (factor_t == 2 && !continuation) {
auto h_first = ggml_ext_slice(ctx->ggml_ctx, h, 2, 0, 1); // [B*factor*OC, 1, H, W]
h_first = _pixel_shuffle_3d(ctx, h_first, 1, factor_s, B); // [B*2*OC, 1, H*factor_s, W*factor_s]
h_first = ggml_ext_slice(ctx->ggml_ctx, h_first, 3, 0, out_channels); // [B*OC, 1, H*factor_s, W*factor_s]
auto x_first = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1);
x_first = repeat_interleave_channels(ctx, x_first, repeats / 2, x->ne[0], x->ne[1], 1);
x_first = _pixel_shuffle_3d(ctx, x_first, 1, factor_s, B);
if (x->ne[2] == 1) {
return ggml_add(ctx->ggml_ctx, h_first, x_first);
}
auto h_next = ggml_ext_slice(ctx->ggml_ctx, h, 2, 1, h->ne[2]); // [B*factor*OC, T - 1, H, W]
h_next = _pixel_shuffle_3d(ctx, h_next, factor_t, factor_s, B); // [B*OC, (T - 1)*factor_t, H*factor_s, W*factor_s]
h = ggml_concat(ctx->ggml_ctx, h_first, h_next, 2); // [B*OC, 1 + (T - 1)*factor_t, H*factor_s, W*factor_s]
auto x_next = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1, x->ne[2]);
x_next = repeat_interleave_channels(ctx, x_next, repeats, x->ne[0], x->ne[1], x->ne[2] - 1);
x_next = _pixel_shuffle_3d(ctx, x_next, factor_t, factor_s, B);
shortcut = ggml_concat(ctx->ggml_ctx, x_first, x_next, 2); // [B*OC, 1 + (T - 1)*factor_t, H*factor_s, W*factor_s]
} else {
h = _pixel_shuffle_3d(ctx, h, factor_t, factor_s, B);
shortcut = repeat_interleave_channels(ctx, x, repeats, x->ne[0], x->ne[1], x->ne[2]);
shortcut = _pixel_shuffle_3d(ctx, shortcut, factor_t, factor_s, B); // [B*OC, T*factor_t, H*factor_s, W*factor_s]
}
return ggml_add(ctx->ggml_ctx, h, shortcut);
}
};
static ggml_tensor* pixel_unshuffle_3d(GGMLRunnerContext* ctx,
ggml_tensor* x,
int64_t factor_t,
int64_t factor_s) {
GGML_ASSERT(x->ne[0] % factor_s == 0);
GGML_ASSERT(x->ne[1] % factor_s == 0);
GGML_ASSERT(x->ne[2] % factor_t == 0);
int64_t W = x->ne[0] / factor_s;
int64_t H = x->ne[1] / factor_s;
int64_t T = x->ne[2] / factor_t;
int64_t C = x->ne[3];
int64_t factor = factor_t * factor_s * factor_s;
x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W * factor_s * H, factor_t, T, C);
x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));
x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W, factor_s, H * T, factor_t * C);
x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));
x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s, W, H * T, factor_s * factor_t * C);
x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 2, 0, 3));
x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H * T, factor, C);
x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
return ggml_reshape_4d(ctx->ggml_ctx, x, W, H, T, C * factor);
}
static ggml_tensor* mean_channel_groups(GGMLRunnerContext* ctx,
ggml_tensor* x,
int64_t group_size) {
GGML_ASSERT(group_size > 0);
GGML_ASSERT(x->ne[3] % group_size == 0);
if (group_size == 1) {
return x;
}
int64_t W = x->ne[0];
int64_t H = x->ne[1];
int64_t T = x->ne[2];
int64_t spatial = W * H * T;
int64_t groups = x->ne[3] / group_size;
x = ggml_reshape_3d(ctx->ggml_ctx, x, spatial, group_size, groups);
x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
x = ggml_sum_rows(ctx->ggml_ctx, x);
x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H, T, groups);
return ggml_scale(ctx->ggml_ctx, x, 1.f / static_cast<float>(group_size));
}
class Downsample : public GGMLBlock {
protected:
int64_t in_channels;
int64_t out_channels;
int64_t factor_t;
int64_t factor_s = 2;
int64_t factor;
int64_t group_size;
public:
Downsample(int64_t in_channels, int64_t out_channels, bool add_temporal_downsample)
: in_channels(in_channels),
out_channels(out_channels),
factor_t(add_temporal_downsample ? 2 : 1),
factor(factor_t * factor_s * factor_s),
group_size(factor * in_channels / out_channels) {
GGML_ASSERT(out_channels % factor == 0);
GGML_ASSERT(factor * in_channels % out_channels == 0);
blocks["conv"] = std::make_shared<CausalConv3d>(in_channels,
out_channels / factor,
std::tuple{3, 3, 3});
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto conv = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv"]);
auto h = conv->forward(ctx, x);
ggml_tensor* h_first = nullptr;
ggml_tensor* x_first = nullptr;
if (factor_t == 2) {
h_first = ggml_ext_slice(ctx->ggml_ctx, h, 2, 0, 1);
h_first = pixel_unshuffle_3d(ctx, h_first, 1, factor_s);
h_first = ggml_concat(ctx->ggml_ctx, h_first, h_first, 3);
x_first = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1);
x_first = pixel_unshuffle_3d(ctx, x_first, 1, factor_s);
x_first = mean_channel_groups(ctx, x_first, group_size / 2);
if (x->ne[2] == 1) {
return ggml_add(ctx->ggml_ctx, h_first, x_first);
}
h = ggml_ext_slice(ctx->ggml_ctx, h, 2, 1, h->ne[2]);
x = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1, x->ne[2]);
}
GGML_ASSERT(h->ne[2] % factor_t == 0);
h = pixel_unshuffle_3d(ctx, h, factor_t, factor_s);
x = pixel_unshuffle_3d(ctx, x, factor_t, factor_s);
x = mean_channel_groups(ctx, x, group_size);
if (factor_t == 2) {
h = ggml_concat(ctx->ggml_ctx, h_first, h, 2);
x = ggml_concat(ctx->ggml_ctx, x_first, x, 2);
}
return ggml_add(ctx->ggml_ctx, h, x);
}
};
class MidBlock : public UnaryBlock {
protected:
int64_t in_channels;
int num_layers;
bool add_attention;
public:
MidBlock(int64_t in_channels,
int num_layers = 1,
bool add_attention = true)
: in_channels(in_channels),
num_layers(num_layers),
add_attention(add_attention) {
blocks["block_1"] = std::make_shared<ResnetBlock>(in_channels, in_channels);
for (int i = 0; i < num_layers; i++) {
if (add_attention) {
blocks["attn_" + std::to_string(i + 1)] = std::make_shared<AttnBlock>(in_channels);
}
blocks["block_" + std::to_string(i + 2)] = std::make_shared<ResnetBlock>(in_channels, in_channels);
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
// x: [B*C, T, H, W]
// return: [B*C, T, H, W]
auto block_1 = std::dynamic_pointer_cast<ResnetBlock>(blocks["block_1"]);
x = block_1->forward(ctx, x);
for (int i = 0; i < num_layers; i++) {
if (add_attention) {
auto block = std::dynamic_pointer_cast<AttnBlock>(blocks["attn_" + std::to_string(i + 1)]);
x = block->forward(ctx, x);
}
auto block = std::dynamic_pointer_cast<ResnetBlock>(blocks["block_" + std::to_string(i + 2)]);
x = block->forward(ctx, x);
}
return x;
}
};
class UpBlock : public UnaryBlock {
protected:
int num_layers;
int64_t upsample_out_channels;
public:
UpBlock(int64_t in_channels,
int64_t out_channels,
int num_layers = 1,
int64_t upsample_out_channels = 0,
bool add_temporal_upsample = true)
: num_layers(num_layers),
upsample_out_channels(upsample_out_channels) {
for (int i = 0; i < num_layers; i++) {
int64_t IC = i == 0 ? in_channels : out_channels;
blocks["block." + std::to_string(i)] = std::make_shared<ResnetBlock>(IC, out_channels);
}
if (upsample_out_channels > 0) {
blocks["upsample"] = std::make_shared<Upsample>(out_channels, upsample_out_channels, add_temporal_upsample);
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
return forward(ctx, x, nullptr);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
TemporalConvCarry* carry) {
// x: [B*IC, T, H, W]
// return: [B*OC, T, H, W] or [B*OC, T, H*2, W*2] or [B*OC, T*2, H*2, W*2]
for (int i = 0; i < num_layers; i++) {
auto block = std::dynamic_pointer_cast<ResnetBlock>(blocks["block." + std::to_string(i)]);
x = block->forward(ctx, x, carry);
}
if (upsample_out_channels > 0) {
auto upsample = std::dynamic_pointer_cast<Upsample>(blocks["upsample"]);
x = upsample->forward(ctx, x, carry);
}
return x;
}
};
class DownBlock : public UnaryBlock {
protected:
int num_layers;
int64_t downsample_out_channels;
public:
DownBlock(int64_t in_channels,
int64_t out_channels,
int num_layers,
int64_t downsample_out_channels = 0,
bool add_temporal_downsample = false)
: num_layers(num_layers),
downsample_out_channels(downsample_out_channels) {
for (int i = 0; i < num_layers; i++) {
int64_t IC = i == 0 ? in_channels : out_channels;
blocks["block." + std::to_string(i)] = std::make_shared<ResnetBlock>(IC, out_channels);
}
if (downsample_out_channels > 0) {
blocks["downsample"] = std::make_shared<Downsample>(out_channels,
downsample_out_channels,
add_temporal_downsample);
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
for (int i = 0; i < num_layers; i++) {
auto block = std::dynamic_pointer_cast<ResnetBlock>(blocks["block." + std::to_string(i)]);
x = block->forward(ctx, x);
}
if (downsample_out_channels > 0) {
auto downsample = std::dynamic_pointer_cast<Downsample>(blocks["downsample"]);
x = downsample->forward(ctx, x);
}
return x;
}
};
class Encoder : public GGMLBlock {
protected:
int64_t z_channels;
std::vector<int64_t> block_out_channels;
public:
Encoder(int64_t in_channels = 3,
int64_t z_channels = 32,
std::vector<int64_t> block_out_channels = {128, 256, 512, 1024, 1024},
int layers_per_block = 2,
int spatial_compression_ratio = 16,
int temporal_compression_ratio = 4,
bool downsample_match_channel = true)
: z_channels(z_channels),
block_out_channels(std::move(block_out_channels)) {
blocks["conv_in"] = std::make_shared<CausalConv3d>(in_channels,
this->block_out_channels[0],
std::tuple{3, 3, 3});
int spatial_depth = static_cast<int>(std::log2(static_cast<double>(spatial_compression_ratio)));
int temporal_start = static_cast<int>(std::log2(static_cast<double>(spatial_compression_ratio / temporal_compression_ratio)));
int64_t channels = this->block_out_channels[0];
for (int i = 0; i < static_cast<int>(this->block_out_channels.size()); i++) {
int64_t out_channels = this->block_out_channels[i];
if (i < spatial_depth) {
int64_t next_channels = downsample_match_channel ? this->block_out_channels[i + 1] : out_channels;
blocks["down." + std::to_string(i)] = std::make_shared<DownBlock>(channels,
out_channels,
layers_per_block,
next_channels,
i >= temporal_start);
channels = next_channels;
} else {
blocks["down." + std::to_string(i)] = std::make_shared<DownBlock>(channels,
out_channels,
layers_per_block);
channels = out_channels;
}
}
blocks["mid"] = std::make_shared<MidBlock>(channels);
blocks["norm_out"] = std::make_shared<WAN::RMS_norm>(channels);
blocks["conv_out"] = std::make_shared<CausalConv3d>(channels,
z_channels * 2,
std::tuple{3, 3, 3});
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto conv_in = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv_in"]);
auto mid = std::dynamic_pointer_cast<MidBlock>(blocks["mid"]);
auto norm_out = std::dynamic_pointer_cast<WAN::RMS_norm>(blocks["norm_out"]);
auto conv_out = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv_out"]);
x = conv_in->forward(ctx, x);
for (int i = 0; i < static_cast<int>(block_out_channels.size()); i++) {
auto down = std::dynamic_pointer_cast<DownBlock>(blocks["down." + std::to_string(i)]);
x = down->forward(ctx, x);
}
x = mid->forward(ctx, x);
auto shortcut = mean_channel_groups(ctx, x, x->ne[3] / (z_channels * 2));
x = norm_out->forward(ctx, x);
x = ggml_silu_inplace(ctx->ggml_ctx, x);
x = conv_out->forward(ctx, x);
x = ggml_add(ctx->ggml_ctx, x, shortcut);
return ggml_ext_slice(ctx->ggml_ctx, x, 3, 0, z_channels);
}
};
class Decoder : public GGMLBlock {
protected:
int64_t repeats;
std::vector<int64_t> block_out_channels;
public:
Decoder(int64_t in_channels = 32,
int64_t out_channels = 3,
std::vector<int64_t> block_out_channels = {1024, 1024, 512, 256, 128},
int layers_per_block = 2,
int spatial_compression_ratio = 16,
int temporal_compression_ratio = 4,
bool upsample_match_channel = true)
: block_out_channels(std::move(block_out_channels)) {
repeats = this->block_out_channels[0] / in_channels;
blocks["conv_in"] = std::make_shared<CausalConv3d>(in_channels, this->block_out_channels[0], std::tuple{3, 3, 3});
blocks["mid"] = std::make_shared<MidBlock>(this->block_out_channels[0]);
int64_t IC = this->block_out_channels[0];
for (int i = 0; i < this->block_out_channels.size(); i++) {
int64_t OC = this->block_out_channels[i];
bool add_spatial_upsample = i < std::log2(static_cast<double>(spatial_compression_ratio));
bool add_temporal_upsample = i < std::log2(static_cast<double>(temporal_compression_ratio));
if (add_spatial_upsample || add_temporal_upsample) {
int64_t upsample_out_channels = upsample_match_channel ? this->block_out_channels[i + 1] : OC;
blocks["up." + std::to_string(i)] = std::make_shared<UpBlock>(IC, OC, layers_per_block + 1, upsample_out_channels, add_temporal_upsample);
IC = upsample_out_channels;
} else {
blocks["up." + std::to_string(i)] = std::make_shared<UpBlock>(IC, OC, layers_per_block + 1, 0, false);
}
}
blocks["norm_out"] = std::make_shared<WAN::RMS_norm>(this->block_out_channels.back());
blocks["conv_out"] = std::make_shared<CausalConv3d>(this->block_out_channels.back(), out_channels, std::tuple{3, 3, 3});
}
struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* z) {
auto conv_in = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv_in"]);
auto mid_block = std::dynamic_pointer_cast<MidBlock>(blocks["mid"]);
auto norm_out = std::dynamic_pointer_cast<WAN::RMS_norm>(blocks["norm_out"]);
auto conv_out = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv_out"]);
auto h = conv_in->forward(ctx, z);
auto shortcut = repeat_interleave_channels(ctx, z, repeats, z->ne[0], z->ne[1], z->ne[2]);
h = ggml_add(ctx->ggml_ctx, h, shortcut);
h = mid_block->forward(ctx, h);
ggml_tensor* output = nullptr;
std::vector<ggml_tensor*> carry_input;
const int64_t frames = h->ne[2];
for (int64_t start = 0; start < frames; start += HUNYUAN_VIDEO_VAE_TEMPORAL_CHUNK_SIZE) {
const int64_t end = std::min(start + HUNYUAN_VIDEO_VAE_TEMPORAL_CHUNK_SIZE, frames);
auto chunk = ggml_ext_slice(ctx->ggml_ctx, h, 2, start, end);
std::vector<ggml_tensor*> carry_output;
TemporalConvCarry carry{
start == 0 ? nullptr : &carry_input,
end == frames ? nullptr : &carry_output,
};
for (int i = 0; i < block_out_channels.size(); i++) {
auto up_block = std::dynamic_pointer_cast<UpBlock>(blocks["up." + std::to_string(i)]);
chunk = up_block->forward(ctx, chunk, &carry);
}
chunk = norm_out->forward(ctx, chunk);
chunk = ggml_silu_inplace(ctx->ggml_ctx, chunk); // nonlinearity/swish
chunk = conv_out->forward(ctx, chunk, &carry);
carry.finish();
output = output == nullptr ? chunk : ggml_concat(ctx->ggml_ctx, output, chunk, 2);
carry_input = std::move(carry_output);
}
return output;
}
};
class HunyuanVideoVAERunner : public VAE {
protected:
bool decode_only;
Encoder encoder;
Decoder decoder;
public:
HunyuanVideoVAERunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map,
const std::string& prefix,
bool decode_only,
SDVersion version,
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: VAE(version, backend, prefix, weight_manager),
decode_only(decode_only ||
tensor_storage_map.find(prefix + ".encoder.conv_in.conv.weight") == tensor_storage_map.end()) {
if (!this->decode_only) {
encoder.init(params_ctx, tensor_storage_map, prefix + ".encoder");
}
decoder.init(params_ctx, tensor_storage_map, prefix + ".decoder");
}
std::string get_desc() override {
return "hunyuan_video_vae";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
if (!decode_only) {
encoder.get_param_tensors(tensors, weight_prefix + ".encoder");
}
decoder.get_param_tensors(tensors, weight_prefix + ".decoder");
}
int get_encoder_output_channels(int input_channels) override {
SD_UNUSED(input_channels);
return 32;
}
sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output,
std::shared_ptr<RNG> rng) override {
SD_UNUSED(rng);
return vae_output;
}
sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) override {
return latents / 1.03682f;
}
sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) override {
return latents * 1.03682f;
}
ggml_cgraph* build_graph(const sd::Tensor<float>& input_tensor, bool decode_graph) {
size_t graph_size = HUNYUAN_VIDEO_VAE_GRAPH_SIZE;
if (decode_graph) {
graph_size = std::max(graph_size,
HUNYUAN_VIDEO_VAE_GRAPH_SIZE_PER_LATENT_FRAME *
static_cast<size_t>(input_tensor.shape()[2]));
}
ggml_cgraph* gf = new_graph_custom(graph_size);
ggml_tensor* input = make_input(input_tensor);
auto runner_ctx = get_context();
ggml_tensor* output = decode_graph ? decoder.forward(&runner_ctx, input)
: encoder.forward(&runner_ctx, input);
ggml_build_forward_expand(gf, output);
return gf;
}
sd::Tensor<float> _compute(const int n_threads,
const sd::Tensor<float>& input,
bool decode_graph) override {
if (!decode_graph && decode_only) {
LOG_ERROR("Hunyuan Video VAE encoder weights are not available");
return {};
}
sd::Tensor<float> expanded;
if (input.dim() == 4) {
expanded = input.unsqueeze(2);
}
const auto& graph_input = expanded.empty() ? input : expanded;
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(graph_input, decode_graph);
};
auto output = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph,
n_threads,
true,
true,
true),
graph_input.dim());
if (!output.empty() && input.dim() == 4) {
output.squeeze_(2);
}
return output;
}
};
} // namespace Hunyuan
#endif // __SD_MODEL_VAE_HUNYUAN_VAE_HPP__

View File

@ -8,6 +8,7 @@
#include <vector>
#include "core/ggml_extend.hpp"
#include "model/vae/audio_vae.hpp"
#include "model_loader.h"
#include "model_manager.h"
@ -996,7 +997,7 @@ namespace LTXV {
}
};
struct LTXAudioVAERunner : public GGMLRunner {
struct LTXAudioVAERunner : public AudioVAERunner {
LTXAudioVAEConfig config;
LTXAudioVAE model;
std::string weight_prefix;
@ -1006,7 +1007,7 @@ namespace LTXV {
const String2TensorStorage& tensor_storage_map,
const std::string& prefix = "",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: GGMLRunner(backend, weight_manager),
: AudioVAERunner(backend, weight_manager),
weight_prefix(prefix),
config(LTXAudioVAEConfig::detect_from_weights(tensor_storage_map)),
model(config) {
@ -1017,20 +1018,20 @@ namespace LTXV {
}
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
model.get_param_tensors(tensors, weight_prefix);
}
size_t get_params_mem_size() {
size_t get_params_mem_size() override {
return model.get_params_mem_size();
}
std::string get_desc() {
std::string get_desc() override {
return "ltx_audio_vae";
}
sd::Tensor<float> decode(int n_threads,
const sd::Tensor<float>& latent_tensor) {
const sd::Tensor<float>& latent_tensor) override {
int64_t t0 = ggml_time_ms();
auto get_graph = [&]() -> ggml_cgraph* {
auto latent = make_input(latent_tensor);
@ -1047,6 +1048,10 @@ namespace LTXV {
return result;
}
int output_sample_rate() const override {
return config.output_sample_rate();
}
void test(const std::string& input_path) {
auto z = sd::load_tensor_from_file_as_tensor<float>(input_path);
GGML_ASSERT(!z.empty());

521
src/model/vae/mage_vae.hpp Normal file
View File

@ -0,0 +1,521 @@
#ifndef __SD_MODEL_VAE_MAGE_VAE_HPP__
#define __SD_MODEL_VAE_MAGE_VAE_HPP__
#include "model/diffusion/dit.hpp"
#include "model/vae/vae.hpp"
namespace MageVAE {
constexpr int MAGE_VAE_GRAPH_SIZE = 327680;
constexpr int HIDDEN_SIZE = 384;
constexpr int LATENT_CHANNELS = 128;
constexpr int PATCH_SIZE = 16;
struct LayerNorm2d : public UnaryBlock {
int64_t channels;
bool affine;
std::string prefix;
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
this->prefix = prefix;
if (affine) {
params["weight"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels);
params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels);
}
}
LayerNorm2d(int64_t channels, bool affine = true)
: channels(channels), affine(affine) {}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
ggml_tensor* weight = affine ? params["weight"] : nullptr;
ggml_tensor* bias = affine ? params["bias"] : nullptr;
if (affine && ctx->weight_adapter) {
weight = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, weight, prefix + "weight");
bias = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, bias, prefix + "bias");
}
// [N, C, H, W] -> [N, H, W, C] so layer norm reduces over channels.
x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 2, 0, 3));
x = ggml_ext_layer_norm(ctx->ggml_ctx, x, weight, bias, 1e-6f);
return ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));
}
};
inline ggml_tensor* modulate_2d(ggml_context* ctx,
ggml_tensor* x,
ggml_tensor* shift,
ggml_tensor* scale) {
shift = ggml_reshape_4d(ctx, shift, 1, 1, shift->ne[0], shift->ne[1]);
scale = ggml_reshape_4d(ctx, scale, 1, 1, scale->ne[0], scale->ne[1]);
return ggml_add(ctx, ggml_mul(ctx, x, ggml_add(ctx, scale, ggml_ext_ones(ctx, 1, 1, 1, 1))), shift);
}
inline ggml_tensor* channel_attention(GGMLRunnerContext* ctx,
ggml_tensor* x,
Conv2d* projection) {
auto pooled = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1], x->ne[2], x->ne[3]);
pooled = ggml_mean(ctx->ggml_ctx, pooled);
pooled = ggml_reshape_4d(ctx->ggml_ctx, pooled, 1, 1, x->ne[2], x->ne[3]);
pooled = ggml_sigmoid(ctx->ggml_ctx, projection->forward(ctx, pooled));
return ggml_mul(ctx->ggml_ctx, x, pooled);
}
struct TimestepEmbedder : public GGMLBlock {
TimestepEmbedder() {
blocks["mlp.0"] = std::make_shared<Linear>(256, HIDDEN_SIZE);
blocks["mlp.2"] = std::make_shared<Linear>(HIDDEN_SIZE, HIDDEN_SIZE);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* timestep) {
auto linear_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
auto x = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, 256, 10000, 1.f);
x = linear_0->forward(ctx, x);
x = ggml_silu_inplace(ctx->ggml_ctx, x);
return linear_2->forward(ctx, x);
}
};
struct EncoderDiCoBlock : public UnaryBlock {
explicit EncoderDiCoBlock(int64_t channels) {
blocks["conv1"] = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
blocks["conv2"] = std::make_shared<Conv2d_grouped>(channels, channels, static_cast<int>(channels), std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1});
blocks["conv3"] = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
blocks["ca.1"] = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
blocks["conv4"] = std::make_shared<Conv2d>(channels, channels * 4, std::pair{1, 1});
blocks["conv5"] = std::make_shared<Conv2d>(channels * 4, channels, std::pair{1, 1});
blocks["norm1"] = std::make_shared<LayerNorm2d>(channels);
blocks["norm2"] = std::make_shared<LayerNorm2d>(channels);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input) override {
auto conv1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv1"]);
auto conv2 = std::dynamic_pointer_cast<Conv2d_grouped>(blocks["conv2"]);
auto conv3 = std::dynamic_pointer_cast<Conv2d>(blocks["conv3"]);
auto ca = std::dynamic_pointer_cast<Conv2d>(blocks["ca.1"]);
auto conv4 = std::dynamic_pointer_cast<Conv2d>(blocks["conv4"]);
auto conv5 = std::dynamic_pointer_cast<Conv2d>(blocks["conv5"]);
auto norm1 = std::dynamic_pointer_cast<LayerNorm2d>(blocks["norm1"]);
auto norm2 = std::dynamic_pointer_cast<LayerNorm2d>(blocks["norm2"]);
auto x = norm1->forward(ctx, input);
x = conv1->forward(ctx, x);
x = conv2->forward(ctx, x);
x = ggml_gelu(ctx->ggml_ctx, x);
x = channel_attention(ctx, x, ca.get());
x = conv3->forward(ctx, x);
x = ggml_add(ctx->ggml_ctx, input, x);
auto h = norm2->forward(ctx, x);
h = conv4->forward(ctx, h);
h = ggml_gelu(ctx->ggml_ctx, h);
h = conv5->forward(ctx, h);
return ggml_add(ctx->ggml_ctx, x, h);
}
};
struct DiCoBlock : public GGMLBlock {
explicit DiCoBlock(int64_t channels) {
blocks["conv1"] = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
blocks["conv2"] = std::make_shared<Conv2d_grouped>(channels, channels, static_cast<int>(channels), std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1});
blocks["conv3"] = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
blocks["ca.1"] = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
blocks["conv4"] = std::make_shared<Conv2d>(channels, channels * 4, std::pair{1, 1});
blocks["conv5"] = std::make_shared<Conv2d>(channels * 4, channels, std::pair{1, 1});
blocks["norm1"] = std::make_shared<LayerNorm2d>(channels, false);
blocks["norm2"] = std::make_shared<LayerNorm2d>(channels, false);
blocks["adaLN_modulation.1"] = std::make_shared<Linear>(channels, channels * 6);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input, ggml_tensor* condition) {
auto conv1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv1"]);
auto conv2 = std::dynamic_pointer_cast<Conv2d_grouped>(blocks["conv2"]);
auto conv3 = std::dynamic_pointer_cast<Conv2d>(blocks["conv3"]);
auto ca = std::dynamic_pointer_cast<Conv2d>(blocks["ca.1"]);
auto conv4 = std::dynamic_pointer_cast<Conv2d>(blocks["conv4"]);
auto conv5 = std::dynamic_pointer_cast<Conv2d>(blocks["conv5"]);
auto norm1 = std::dynamic_pointer_cast<LayerNorm2d>(blocks["norm1"]);
auto norm2 = std::dynamic_pointer_cast<LayerNorm2d>(blocks["norm2"]);
auto ada = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.1"]);
auto params = ada->forward(ctx, ggml_silu(ctx->ggml_ctx, condition));
auto chunks = ggml_ext_chunk(ctx->ggml_ctx, params, 6, 0);
auto x = norm1->forward(ctx, input);
x = modulate_2d(ctx->ggml_ctx, x, chunks[0], chunks[1]);
x = conv1->forward(ctx, x);
x = conv2->forward(ctx, x);
x = ggml_gelu(ctx->ggml_ctx, x);
x = channel_attention(ctx, x, ca.get());
x = conv3->forward(ctx, x);
auto gate_1 = ggml_reshape_4d(ctx->ggml_ctx, chunks[2], 1, 1, chunks[2]->ne[0], chunks[2]->ne[1]);
x = ggml_add(ctx->ggml_ctx, input, ggml_mul(ctx->ggml_ctx, x, gate_1));
auto h = norm2->forward(ctx, x);
h = modulate_2d(ctx->ggml_ctx, h, chunks[3], chunks[4]);
h = conv4->forward(ctx, h);
h = ggml_gelu(ctx->ggml_ctx, h);
h = conv5->forward(ctx, h);
auto gate_2 = ggml_reshape_4d(ctx->ggml_ctx, chunks[5], 1, 1, chunks[5]->ne[0], chunks[5]->ne[1]);
return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, h, gate_2));
}
};
struct MageResnetBlock : public UnaryBlock {
explicit MageResnetBlock(int64_t channels) {
blocks["norm1"] = std::make_shared<GroupNorm32>(channels);
blocks["conv1"] = std::make_shared<Conv2d>(channels, channels, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1});
blocks["norm2"] = std::make_shared<GroupNorm32>(channels);
blocks["conv2"] = std::make_shared<Conv2d>(channels, channels, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1});
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input) override {
auto norm1 = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm1"]);
auto conv1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv1"]);
auto norm2 = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm2"]);
auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["conv2"]);
auto x = conv1->forward(ctx, ggml_silu(ctx->ggml_ctx, norm1->forward(ctx, input)));
x = conv2->forward(ctx, ggml_silu(ctx->ggml_ctx, norm2->forward(ctx, x)));
return ggml_add(ctx->ggml_ctx, input, x);
}
};
inline ggml_tensor* replicate_pad_right_bottom(ggml_context* ctx,
ggml_tensor* x,
int pad_w,
int pad_h) {
if (pad_w > 0) {
auto edge = ggml_ext_slice(ctx, x, 0, x->ne[0] - 1, x->ne[0]);
edge = ggml_repeat_4d(ctx, edge, pad_w, x->ne[1], x->ne[2], x->ne[3]);
x = ggml_concat(ctx, x, edge, 0);
}
if (pad_h > 0) {
auto edge = ggml_ext_slice(ctx, x, 1, x->ne[1] - 1, x->ne[1]);
edge = ggml_repeat_4d(ctx, edge, x->ne[0], pad_h, x->ne[2], x->ne[3]);
x = ggml_concat(ctx, x, edge, 1);
}
return x;
}
struct MageAttnBlock : public UnaryBlock {
int64_t channels;
int patch_size;
MageAttnBlock(int64_t channels, int patch_size = 32)
: channels(channels), patch_size(patch_size) {
blocks["norm"] = std::make_shared<GroupNorm32>(channels);
blocks["q"] = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
blocks["k"] = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
blocks["v"] = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
blocks["proj_out"] = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
}
ggml_tensor* to_patches(ggml_context* ctx, ggml_tensor* x) {
x = DiT::patchify(ctx, x, patch_size, patch_size);
x = ggml_reshape_4d(ctx, x, patch_size * patch_size, channels, x->ne[1], x->ne[2]);
// [N, np, C, P] -> [N, np, P, C] for attention over P pixels.
x = ggml_ext_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3));
return ggml_reshape_3d(ctx, x, channels, patch_size * patch_size, x->ne[2] * x->ne[3]);
}
ggml_tensor* from_patches(ggml_context* ctx,
ggml_tensor* x,
int64_t patch_count,
int64_t batch_size,
int64_t h_patches,
int64_t w_patches) {
x = ggml_reshape_4d(ctx, x, channels, patch_size * patch_size, patch_count, batch_size);
// [N, np, P, C] -> [N, np, C, P] before spatial unpatchify.
x = ggml_ext_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3));
x = ggml_reshape_3d(ctx, x, patch_size * patch_size * channels, patch_count, batch_size);
return DiT::unpatchify(ctx, x, h_patches, w_patches, patch_size, patch_size);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input) override {
auto norm = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm"]);
auto q_proj = std::dynamic_pointer_cast<Conv2d>(blocks["q"]);
auto k_proj = std::dynamic_pointer_cast<Conv2d>(blocks["k"]);
auto v_proj = std::dynamic_pointer_cast<Conv2d>(blocks["v"]);
auto proj_out = std::dynamic_pointer_cast<Conv2d>(blocks["proj_out"]);
int64_t width = input->ne[0];
int64_t height = input->ne[1];
int64_t batch = input->ne[3];
int pad_w = (patch_size - static_cast<int>(width % patch_size)) % patch_size;
int pad_h = (patch_size - static_cast<int>(height % patch_size)) % patch_size;
int64_t wp = (width + pad_w) / patch_size;
int64_t hp = (height + pad_h) / patch_size;
int64_t np = wp * hp;
auto h = norm->forward(ctx, input);
auto q = replicate_pad_right_bottom(ctx->ggml_ctx, q_proj->forward(ctx, h), pad_w, pad_h);
auto k = replicate_pad_right_bottom(ctx->ggml_ctx, k_proj->forward(ctx, h), pad_w, pad_h);
auto v = replicate_pad_right_bottom(ctx->ggml_ctx, v_proj->forward(ctx, h), pad_w, pad_h);
q = to_patches(ctx->ggml_ctx, q);
k = to_patches(ctx->ggml_ctx, k);
v = to_patches(ctx->ggml_ctx, v);
h = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled);
h = from_patches(ctx->ggml_ctx, h, np, batch, hp, wp);
if (pad_h > 0) {
h = ggml_ext_slice(ctx->ggml_ctx, h, 1, 0, height);
}
if (pad_w > 0) {
h = ggml_ext_slice(ctx->ggml_ctx, h, 0, 0, width);
}
return ggml_add(ctx->ggml_ctx, input, proj_out->forward(ctx, h));
}
};
struct Decoder : public UnaryBlock {
Decoder() {
blocks["conv_in"] = std::make_shared<Conv2d>(LATENT_CHANNELS, HIDDEN_SIZE, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1});
blocks["block.0"] = std::make_shared<MageResnetBlock>(HIDDEN_SIZE);
blocks["block.1"] = std::make_shared<MageAttnBlock>(HIDDEN_SIZE);
blocks["block.2"] = std::make_shared<MageResnetBlock>(HIDDEN_SIZE);
blocks["block.3"] = std::make_shared<MageAttnBlock>(HIDDEN_SIZE);
blocks["block.4"] = std::make_shared<MageResnetBlock>(HIDDEN_SIZE);
blocks["norm_out"] = std::make_shared<GroupNorm32>(HIDDEN_SIZE);
blocks["conv_out"] = std::make_shared<Conv2d>(HIDDEN_SIZE, HIDDEN_SIZE, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1});
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
x = std::dynamic_pointer_cast<Conv2d>(blocks["conv_in"])->forward(ctx, x);
for (int i = 0; i < 5; ++i) {
x = std::dynamic_pointer_cast<UnaryBlock>(blocks["block." + std::to_string(i)])->forward(ctx, x);
}
x = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm_out"])->forward(ctx, x);
x = ggml_silu(ctx->ggml_ctx, x);
return std::dynamic_pointer_cast<Conv2d>(blocks["conv_out"])->forward(ctx, x);
}
};
struct DConvEncoder : public UnaryBlock {
DConvEncoder() {
blocks["patch_cond_embed"] = std::make_shared<Conv2d>(3, 768, std::pair{PATCH_SIZE, PATCH_SIZE}, std::pair{PATCH_SIZE, PATCH_SIZE});
for (int i = 0; i < 2; ++i) {
blocks["head_blocks." + std::to_string(i)] = std::make_shared<EncoderDiCoBlock>(768);
}
blocks["proj_down"] = std::make_shared<Conv2d>(768, HIDDEN_SIZE, std::pair{1, 1});
blocks["z_proj"] = std::make_shared<Conv2d>(LATENT_CHANNELS, HIDDEN_SIZE, std::pair{1, 1});
blocks["fuse_proj"] = std::make_shared<Conv2d>(HIDDEN_SIZE * 2, HIDDEN_SIZE, std::pair{1, 1});
blocks["t_embedder"] = std::make_shared<TimestepEmbedder>();
for (int i = 0; i < 21; ++i) {
blocks["blocks." + std::to_string(i)] = std::make_shared<DiCoBlock>(HIDDEN_SIZE);
}
blocks["norm_out"] = std::make_shared<LayerNorm2d>(HIDDEN_SIZE);
blocks["proj_out"] = std::make_shared<Conv2d>(HIDDEN_SIZE, LATENT_CHANNELS * 2, std::pair{1, 1});
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* image) override {
auto cond = std::dynamic_pointer_cast<Conv2d>(blocks["patch_cond_embed"])->forward(ctx, image);
for (int i = 0; i < 2; ++i) {
cond = std::dynamic_pointer_cast<EncoderDiCoBlock>(blocks["head_blocks." + std::to_string(i)])->forward(ctx, cond);
}
cond = std::dynamic_pointer_cast<Conv2d>(blocks["proj_down"])->forward(ctx, cond);
auto z = ggml_ext_zeros(ctx->ggml_ctx, cond->ne[0], cond->ne[1], LATENT_CHANNELS, cond->ne[3]);
z = std::dynamic_pointer_cast<Conv2d>(blocks["z_proj"])->forward(ctx, z);
z = ggml_concat(ctx->ggml_ctx, cond, z, 2);
z = std::dynamic_pointer_cast<Conv2d>(blocks["fuse_proj"])->forward(ctx, z);
auto t = ggml_ext_zeros(ctx->ggml_ctx, image->ne[3], 1, 1, 1);
t = ggml_reshape_1d(ctx->ggml_ctx, t, image->ne[3]);
auto c = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"])->forward(ctx, t);
for (int i = 0; i < 21; ++i) {
z = std::dynamic_pointer_cast<DiCoBlock>(blocks["blocks." + std::to_string(i)])->forward(ctx, z, c);
}
z = std::dynamic_pointer_cast<LayerNorm2d>(blocks["norm_out"])->forward(ctx, z);
return std::dynamic_pointer_cast<Conv2d>(blocks["proj_out"])->forward(ctx, z);
}
};
struct MLPResBlock : public GGMLBlock {
MLPResBlock() {
blocks["in_ln"] = std::make_shared<LayerNorm>(32, 1e-6f);
blocks["mlp.0"] = std::make_shared<Linear>(32, 32);
blocks["mlp.2"] = std::make_shared<Linear>(32, 32);
blocks["adaLN_modulation.1"] = std::make_shared<Linear>(32, 96);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* condition) {
auto params = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.1"])->forward(ctx, ggml_silu(ctx->ggml_ctx, condition));
auto chunks = ggml_ext_chunk(ctx->ggml_ctx, params, 3, 0);
auto h = std::dynamic_pointer_cast<LayerNorm>(blocks["in_ln"])->forward(ctx, x);
h = ggml_add(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, h, ggml_add(ctx->ggml_ctx, chunks[1], ggml_ext_ones(ctx->ggml_ctx, 1, 1, 1, 1))), chunks[0]);
h = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"])->forward(ctx, h);
h = ggml_silu(ctx->ggml_ctx, h);
h = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"])->forward(ctx, h);
return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, chunks[2], h));
}
};
struct DConvDenoiser : public GGMLBlock {
DConvDenoiser() {
blocks["t_embedder"] = std::make_shared<TimestepEmbedder>();
blocks["y_embedder_x"] = std::make_shared<Conv2d>(HIDDEN_SIZE, 32 * PATCH_SIZE * PATCH_SIZE, std::pair{1, 1});
blocks["x_embedder.embedder.0"] = std::make_shared<Linear>(3 + 32 + 64, 32);
blocks["s_embedder.proj1"] = std::make_shared<Conv2d>(3, LATENT_CHANNELS, std::pair{PATCH_SIZE, PATCH_SIZE}, std::pair{PATCH_SIZE, PATCH_SIZE}, std::pair{0, 0}, std::pair{1, 1}, false);
blocks["s_embedder.proj2"] = std::make_shared<Conv2d>(LATENT_CHANNELS + HIDDEN_SIZE, HIDDEN_SIZE, std::pair{1, 1});
for (int i = 0; i < 21; ++i) {
blocks["blocks." + std::to_string(i)] = std::make_shared<DiCoBlock>(HIDDEN_SIZE);
}
blocks["dec_net.cond_embed"] = std::make_shared<Linear>(HIDDEN_SIZE, PATCH_SIZE * PATCH_SIZE * 32);
blocks["dec_net.input_proj"] = std::make_shared<Linear>(32, 32);
for (int i = 0; i < 3; ++i) {
blocks["dec_net.res_blocks." + std::to_string(i)] = std::make_shared<MLPResBlock>();
}
blocks["final_layer.norm"] = std::make_shared<RMSNorm>(32);
blocks["final_layer.linear"] = std::make_shared<Linear>(32, 3);
blocks["y_embedder.decoder"] = std::make_shared<Decoder>();
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* latent, ggml_tensor* dct) {
auto cond = std::dynamic_pointer_cast<Decoder>(blocks["y_embedder.decoder"])->forward(ctx, latent);
int64_t w = cond->ne[0];
int64_t h = cond->ne[1];
int64_t n = cond->ne[3];
int64_t length = w * h;
auto image = ggml_ext_zeros(ctx->ggml_ctx, w * PATCH_SIZE, h * PATCH_SIZE, 3, n);
auto t = ggml_ext_zeros(ctx->ggml_ctx, n, 1, 1, 1);
t = ggml_reshape_1d(ctx->ggml_ctx, t, n);
auto c = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"])->forward(ctx, t);
auto s0 = std::dynamic_pointer_cast<Conv2d>(blocks["s_embedder.proj1"])->forward(ctx, image);
s0 = ggml_concat(ctx->ggml_ctx, s0, cond, 2);
auto s = std::dynamic_pointer_cast<Conv2d>(blocks["s_embedder.proj2"])->forward(ctx, s0);
for (int i = 0; i < 21; ++i) {
s = std::dynamic_pointer_cast<DiCoBlock>(blocks["blocks." + std::to_string(i)])->forward(ctx, s, c);
}
// [N, C, H, W] -> [N*H*W, C].
s = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, s, 1, 2, 0, 3));
s = ggml_reshape_2d(ctx->ggml_ctx, s, HIDDEN_SIZE, length * n);
auto y = std::dynamic_pointer_cast<Conv2d>(blocks["y_embedder_x"])->forward(ctx, cond);
// Split 32*P channels as [32, P], then produce [N*L, P, 32].
y = ggml_reshape_4d(ctx->ggml_ctx, y, length, PATCH_SIZE * PATCH_SIZE, 32, n);
y = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, y, 2, 1, 0, 3));
y = ggml_reshape_3d(ctx->ggml_ctx, y, 32, PATCH_SIZE * PATCH_SIZE, length * n);
auto zeros = ggml_ext_zeros(ctx->ggml_ctx, 3, PATCH_SIZE * PATCH_SIZE, length * n, 1);
dct = ggml_repeat_4d(ctx->ggml_ctx, dct, 64, PATCH_SIZE * PATCH_SIZE, length * n, 1);
auto x = ggml_concat(ctx->ggml_ctx, zeros, y, 0);
x = ggml_concat(ctx->ggml_ctx, x, dct, 0);
x = std::dynamic_pointer_cast<Linear>(blocks["x_embedder.embedder.0"])->forward(ctx, x);
x = std::dynamic_pointer_cast<Linear>(blocks["dec_net.input_proj"])->forward(ctx, x);
auto dec_cond = std::dynamic_pointer_cast<Linear>(blocks["dec_net.cond_embed"])->forward(ctx, s);
dec_cond = ggml_reshape_3d(ctx->ggml_ctx, dec_cond, 32, PATCH_SIZE * PATCH_SIZE, length * n);
for (int i = 0; i < 3; ++i) {
x = std::dynamic_pointer_cast<MLPResBlock>(blocks["dec_net.res_blocks." + std::to_string(i)])->forward(ctx, x, dec_cond);
}
x = std::dynamic_pointer_cast<RMSNorm>(blocks["final_layer.norm"])->forward(ctx, x);
x = std::dynamic_pointer_cast<Linear>(blocks["final_layer.linear"])->forward(ctx, x);
// [N*L, P, 3] -> [N, L, 3*P] for fold/unpatchify.
x = ggml_reshape_4d(ctx->ggml_ctx, x, 3, PATCH_SIZE * PATCH_SIZE, length, n);
x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
x = ggml_reshape_3d(ctx->ggml_ctx, x, 3 * PATCH_SIZE * PATCH_SIZE, length, n);
return DiT::unpatchify(ctx->ggml_ctx, x, h, w, PATCH_SIZE, PATCH_SIZE);
}
};
struct MageVAEModel : public GGMLBlock {
MageVAEModel() {
blocks["student.dconv_encoder"] = std::make_shared<DConvEncoder>();
blocks["pipeline"] = std::make_shared<DConvDenoiser>();
}
ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* image) {
return std::dynamic_pointer_cast<DConvEncoder>(blocks["student.dconv_encoder"])->forward(ctx, image);
}
ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* latent, ggml_tensor* dct) {
return std::dynamic_pointer_cast<DConvDenoiser>(blocks["pipeline"])->forward(ctx, latent, dct);
}
};
struct MageVAERunner : public VAE {
MageVAEModel model;
std::vector<float> dct_vec;
MageVAERunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map,
const std::string& prefix,
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: VAE(VERSION_MAGE_FLOW, backend, prefix, weight_manager) {
model = MageVAEModel();
model.init(params_ctx, tensor_storage_map, prefix);
dct_vec.resize(64 * PATCH_SIZE * PATCH_SIZE);
constexpr float pi = 3.14159265358979323846f;
for (int py = 0; py < PATCH_SIZE; ++py) {
float y = static_cast<float>(py) / static_cast<float>(PATCH_SIZE - 1);
for (int px = 0; px < PATCH_SIZE; ++px) {
float x = static_cast<float>(px) / static_cast<float>(PATCH_SIZE - 1);
int pos = py * PATCH_SIZE + px;
for (int fy = 0; fy < 8; ++fy) {
for (int fx = 0; fx < 8; ++fx) {
int freq = fx * 8 + fy;
float freq_x = static_cast<float>(fx) * 8.f / 7.f;
float freq_y = static_cast<float>(fy) * 8.f / 7.f;
float coeff = 1.f / (1.f + freq_x * freq_y);
dct_vec[freq + 64 * pos] = std::cos(x * freq_x * pi) *
std::cos(y * freq_y * pi) * coeff;
}
}
}
}
}
std::string get_desc() override {
return "mage_vae";
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
model.get_param_tensors(tensors, weight_prefix);
}
ggml_cgraph* build_graph(const sd::Tensor<float>& input_tensor, bool decode_graph) {
ggml_cgraph* gf = new_graph_custom(MAGE_VAE_GRAPH_SIZE);
auto input = make_input(input_tensor);
auto runner_ctx = get_context();
ggml_tensor* dct = nullptr;
if (decode_graph) {
dct = ggml_new_tensor_3d(compute_ctx, GGML_TYPE_F32, 64, PATCH_SIZE * PATCH_SIZE, 1);
set_backend_tensor_data(dct, dct_vec.data());
}
auto out = decode_graph ? model.decode(&runner_ctx, input, dct) : model.encode(&runner_ctx, input);
ggml_build_forward_expand(gf, out);
return gf;
}
sd::Tensor<float> _compute(const int n_threads,
const sd::Tensor<float>& input,
bool decode_graph) override {
auto get_graph = [&]() -> ggml_cgraph* {
return build_graph(input, decode_graph);
};
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), input.dim());
}
int get_encoder_output_channels(int input_channels) override {
SD_UNUSED(input_channels);
return LATENT_CHANNELS * 2;
}
sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) override {
const auto chunks = sd::ops::chunk(vae_output, 2, 2);
const auto& mean = chunks[0];
const auto& logvar = chunks[1];
sd::Tensor<float> stddev = sd::ops::exp(0.5f * sd::ops::clamp(logvar, -20.0f, 10.0f));
sd::Tensor<float> noise = sd::Tensor<float>::randn_like(mean, rng);
sd::Tensor<float> latents = mean + stddev * noise;
return latents;
}
sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) override {
return latents;
}
sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) override {
return latents;
}
};
} // namespace MageVAE
#endif // __SD_MODEL_VAE_MAGE_VAE_HPP__

View File

@ -0,0 +1,497 @@
#ifndef __SD_MODEL_VAE_MINIMAX_H3_AUDIO_VAE_HPP__
#define __SD_MODEL_VAE_MINIMAX_H3_AUDIO_VAE_HPP__
#include <array>
#include <string>
#include <vector>
#include "model/vae/audio_vae.hpp"
#include "model/vae/ltx_audio_vae.hpp"
namespace MiniMaxH3 {
struct AudioSnake1D : public UnaryBlock {
int64_t channels;
explicit AudioSnake1D(int64_t channels)
: channels(channels) {}
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
SD_UNUSED(tensor_storage_map);
SD_UNUSED(prefix);
params["alpha"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, channels, 1);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
auto alpha = params["alpha"];
auto oscillation = ggml_sin(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, x, alpha));
oscillation = ggml_mul(ctx->ggml_ctx, oscillation, oscillation);
auto eps = ggml_ext_scale(ctx->ggml_ctx, ggml_ext_ones(ctx->ggml_ctx, 1, 1, 1, 1), 1e-9f);
return ggml_add(ctx->ggml_ctx,
x,
ggml_div(ctx->ggml_ctx, oscillation, ggml_add(ctx->ggml_ctx, alpha, eps)));
}
};
struct AudioEncoderResidualUnit : public GGMLBlock {
int64_t channels;
AudioEncoderResidualUnit(int64_t channels, int dilation)
: channels(channels) {
blocks["block.0"] = std::make_shared<AudioSnake1D>(channels);
blocks["block.1"] = std::make_shared<LTXV::Conv1D>(channels,
channels,
7,
1,
3 * dilation,
dilation);
blocks["block.2"] = std::make_shared<AudioSnake1D>(channels);
blocks["block.3"] = std::make_shared<LTXV::Conv1D>(channels, channels, 1);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto act1 = std::dynamic_pointer_cast<AudioSnake1D>(blocks["block.0"]);
auto conv1 = std::dynamic_pointer_cast<LTXV::Conv1D>(blocks["block.1"]);
auto act2 = std::dynamic_pointer_cast<AudioSnake1D>(blocks["block.2"]);
auto conv2 = std::dynamic_pointer_cast<LTXV::Conv1D>(blocks["block.3"]);
auto h = conv2->forward(ctx, act2->forward(ctx, conv1->forward(ctx, act1->forward(ctx, x))));
if (x->ne[0] != h->ne[0]) {
int64_t pad = (x->ne[0] - h->ne[0]) / 2;
x = ggml_ext_slice(ctx->ggml_ctx, x, 0, pad, x->ne[0] - pad);
}
return ggml_add(ctx->ggml_ctx, x, h);
}
};
struct AudioEncoderBlock : public GGMLBlock {
int64_t out_channels;
AudioEncoderBlock(int64_t out_channels, int stride)
: out_channels(out_channels) {
int64_t in_channels = out_channels / 2;
blocks["block.0"] = std::make_shared<AudioEncoderResidualUnit>(in_channels, 1);
blocks["block.1"] = std::make_shared<AudioEncoderResidualUnit>(in_channels, 3);
blocks["block.2"] = std::make_shared<AudioEncoderResidualUnit>(in_channels, 9);
blocks["block.3"] = std::make_shared<AudioSnake1D>(in_channels);
blocks["block.4"] = std::make_shared<LTXV::Conv1D>(in_channels,
out_channels,
2 * stride,
stride,
static_cast<int>(std::ceil(stride / 2.f)));
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
for (int i = 0; i < 3; ++i) {
auto unit = std::dynamic_pointer_cast<AudioEncoderResidualUnit>(blocks["block." + std::to_string(i)]);
x = unit->forward(ctx, x);
}
auto act = std::dynamic_pointer_cast<AudioSnake1D>(blocks["block.3"]);
auto conv = std::dynamic_pointer_cast<LTXV::Conv1D>(blocks["block.4"]);
return conv->forward(ctx, act->forward(ctx, x));
}
};
struct AudioEncoder : public GGMLBlock {
static constexpr std::array<int, 5> strides = {2, 4, 4, 5, 5};
AudioEncoder() {
int64_t channels = 64;
blocks["block.0"] = std::make_shared<LTXV::Conv1D>(1, channels, 7, 1, 3);
for (size_t i = 0; i < strides.size(); ++i) {
channels *= 2;
blocks["block." + std::to_string(i + 1)] = std::make_shared<AudioEncoderBlock>(channels, strides[i]);
}
blocks["block.6"] = std::make_shared<AudioSnake1D>(channels);
blocks["block.7"] = std::make_shared<LTXV::Conv1D>(channels, 2048, 3, 1, 1);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto input = std::dynamic_pointer_cast<LTXV::Conv1D>(blocks["block.0"]);
x = input->forward(ctx, x);
for (size_t i = 0; i < strides.size(); ++i) {
auto block = std::dynamic_pointer_cast<AudioEncoderBlock>(blocks["block." + std::to_string(i + 1)]);
x = block->forward(ctx, x);
}
auto act = std::dynamic_pointer_cast<AudioSnake1D>(blocks["block.6"]);
auto out = std::dynamic_pointer_cast<LTXV::Conv1D>(blocks["block.7"]);
return out->forward(ctx, act->forward(ctx, x));
}
};
struct AudioGeGLUMLP : public GGMLBlock {
AudioGeGLUMLP(int64_t hidden_size, int64_t intermediate_size) {
blocks["norm"] = std::make_shared<LayerNorm>(hidden_size);
blocks["w0"] = std::make_shared<Linear>(hidden_size, intermediate_size, true);
blocks["w1"] = std::make_shared<Linear>(hidden_size, intermediate_size, true);
blocks["w2"] = std::make_shared<Linear>(intermediate_size, hidden_size, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto norm = std::dynamic_pointer_cast<LayerNorm>(blocks["norm"]);
auto w0 = std::dynamic_pointer_cast<Linear>(blocks["w0"]);
auto w1 = std::dynamic_pointer_cast<Linear>(blocks["w1"]);
auto w2 = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
x = norm->forward(ctx, x);
auto gate = ggml_ext_gelu(ctx->ggml_ctx, w0->forward(ctx, x), true);
return w2->forward(ctx, ggml_mul(ctx->ggml_ctx, gate, w1->forward(ctx, x)));
}
};
struct AudioCausalAttention : public GGMLBlock {
static constexpr int64_t in_channels = 2048;
static constexpr int64_t out_channels = 32;
static constexpr int64_t num_head = 8;
static constexpr int64_t head_dim = in_channels / num_head;
AudioCausalAttention() {
blocks["qkv"] = std::make_shared<Linear>(in_channels, in_channels * 3, false);
blocks["proj"] = std::make_shared<Linear>(out_channels, out_channels, true);
}
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
GGMLBlock::init_params(ctx, tensor_storage_map, prefix);
params["q_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
params["v_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, in_channels);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto qkv_layer = std::dynamic_pointer_cast<Linear>(blocks["qkv"]);
auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
auto qkv = ggml_ext_chunk(ctx->ggml_ctx, qkv_layer->forward(ctx, x), 3, 0);
auto bias_shape = [&](ggml_tensor* bias) {
return ggml_reshape_4d(ctx->ggml_ctx, bias, bias->ne[0], 1, 1, 1);
};
auto q = ggml_add(ctx->ggml_ctx, qkv[0], bias_shape(params["q_bias"]));
auto k = qkv[1];
auto v = ggml_add(ctx->ggml_ctx, qkv[2], bias_shape(params["v_bias"]));
int64_t sequence = x->ne[1];
auto mask = ggml_diag_mask_inf(ctx->ggml_ctx,
ggml_ext_zeros(ctx->ggml_ctx, sequence, sequence, 1, 1),
0);
auto attn_out = ggml_ext_attention_ext(ctx->ggml_ctx,
ctx->backend,
q,
k,
v,
num_head,
mask,
false,
ctx->flash_attn_enabled);
int64_t batch = attn_out->ne[2] * attn_out->ne[3];
attn_out = ggml_reshape_4d(ctx->ggml_ctx, attn_out, head_dim, num_head, sequence, batch);
attn_out = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, attn_out, 1, 0, 2, 3));
attn_out = ggml_mean(ctx->ggml_ctx, attn_out);
attn_out = ggml_reshape_3d(ctx->ggml_ctx, attn_out, head_dim, sequence, batch);
constexpr int64_t pool = head_dim / out_channels;
attn_out = ggml_reshape_4d(ctx->ggml_ctx, attn_out, pool, out_channels, sequence, batch);
attn_out = ggml_mean(ctx->ggml_ctx, attn_out);
attn_out = ggml_reshape_3d(ctx->ggml_ctx, attn_out, out_channels, sequence, batch);
return proj->forward(ctx, attn_out);
}
};
struct AudioAttentionProjection : public GGMLBlock {
AudioAttentionProjection() {
blocks["norm1"] = std::make_shared<LayerNorm>(2048);
blocks["attn"] = std::make_shared<AudioCausalAttention>();
blocks["proj"] = std::make_shared<Linear>(2048, 32, true);
blocks["norm3"] = std::make_shared<LayerNorm>(2048);
blocks["norm2"] = std::make_shared<LayerNorm>(32);
blocks["mlp"] = std::make_shared<AudioGeGLUMLP>(32, 64);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto norm1 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
auto attn = std::dynamic_pointer_cast<AudioCausalAttention>(blocks["attn"]);
auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
auto norm3 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm3"]);
auto norm2 = std::dynamic_pointer_cast<LayerNorm>(blocks["norm2"]);
auto mlp = std::dynamic_pointer_cast<AudioGeGLUMLP>(blocks["mlp"]);
x = ggml_add(ctx->ggml_ctx,
proj->forward(ctx, norm3->forward(ctx, x)),
attn->forward(ctx, norm1->forward(ctx, x)));
return ggml_add(ctx->ggml_ctx, x, mlp->forward(ctx, norm2->forward(ctx, x)));
}
};
struct AudioAMPBlock : public GGMLBlock {
int channels;
AudioAMPBlock(int channels,
int kernel_size,
const std::array<int, 3>& dilations)
: channels(channels) {
for (int i = 0; i < 3; ++i) {
blocks["activations." + std::to_string(i * 2)] =
std::make_shared<LTXV::Activation1D>(channels);
blocks["activations." + std::to_string(i * 2 + 1)] =
std::make_shared<LTXV::Activation1D>(channels);
blocks["convs1." + std::to_string(i)] =
std::make_shared<LTXV::Conv1D>(channels,
channels,
kernel_size,
1,
(kernel_size * dilations[i] - dilations[i]) / 2,
dilations[i]);
blocks["convs2." + std::to_string(i)] =
std::make_shared<LTXV::Conv1D>(channels,
channels,
kernel_size,
1,
kernel_size / 2);
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
for (int i = 0; i < 3; ++i) {
auto act1 = std::dynamic_pointer_cast<LTXV::Activation1D>(
blocks["activations." + std::to_string(i * 2)]);
auto act2 = std::dynamic_pointer_cast<LTXV::Activation1D>(
blocks["activations." + std::to_string(i * 2 + 1)]);
auto conv1 = std::dynamic_pointer_cast<LTXV::Conv1D>(
blocks["convs1." + std::to_string(i)]);
auto conv2 = std::dynamic_pointer_cast<LTXV::Conv1D>(
blocks["convs2." + std::to_string(i)]);
auto h = conv1->forward(ctx, act1->forward(ctx, x));
h = conv2->forward(ctx, act2->forward(ctx, h));
x = ggml_add(ctx->ggml_ctx, x, h);
}
return x;
}
};
struct BigVGAN : public GGMLBlock {
static constexpr int initial_channels = 1024;
static constexpr int num_kernels = 3;
static constexpr int num_upsamples = 7;
static constexpr std::array<int, num_upsamples> rates = {5, 5, 2, 2, 2, 2, 2};
static constexpr std::array<int, num_upsamples> kernels = {9, 9, 4, 4, 4, 4, 4};
static constexpr std::array<int, num_kernels> res_kernels = {3, 7, 11};
BigVGAN() {
blocks["conv_pre"] = std::make_shared<LTXV::Conv1D>(2048,
initial_channels,
7,
1,
3);
int channels = initial_channels;
for (int i = 0; i < num_upsamples; ++i) {
int next_channels = initial_channels / (1 << (i + 1));
blocks["ups." + std::to_string(i) + ".0"] =
std::make_shared<LTXV::ConvTranspose1D>(channels,
next_channels,
kernels[i],
rates[i],
(kernels[i] - rates[i]) / 2);
for (int j = 0; j < num_kernels; ++j) {
blocks["resblocks." + std::to_string(i * num_kernels + j)] =
std::make_shared<AudioAMPBlock>(next_channels,
res_kernels[j],
std::array<int, 3>{1, 3, 5});
}
channels = next_channels;
}
blocks["activation_post"] = std::make_shared<LTXV::Activation1D>(channels);
blocks["conv_post"] = std::make_shared<LTXV::Conv1D>(channels,
1,
7,
1,
3,
1,
false);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto conv_pre = std::dynamic_pointer_cast<LTXV::Conv1D>(blocks["conv_pre"]);
x = conv_pre->forward(ctx, x);
for (int i = 0; i < num_upsamples; ++i) {
auto up = std::dynamic_pointer_cast<LTXV::ConvTranspose1D>(
blocks["ups." + std::to_string(i) + ".0"]);
x = up->forward(ctx, x);
ggml_tensor* sum = nullptr;
for (int j = 0; j < num_kernels; ++j) {
auto block = std::dynamic_pointer_cast<AudioAMPBlock>(
blocks["resblocks." + std::to_string(i * num_kernels + j)]);
auto value = block->forward(ctx, x);
sum = sum == nullptr ? value : ggml_add(ctx->ggml_ctx, sum, value);
}
x = ggml_ext_scale(ctx->ggml_ctx, sum, 1.f / num_kernels);
}
auto activation = std::dynamic_pointer_cast<LTXV::Activation1D>(blocks["activation_post"]);
auto conv_post = std::dynamic_pointer_cast<LTXV::Conv1D>(blocks["conv_post"]);
return ggml_clamp(ctx->ggml_ctx,
conv_post->forward(ctx, activation->forward(ctx, x)),
-1.f,
1.f);
}
};
struct AudioVAE : public GGMLBlock {
static constexpr int kLatentChannels = 32;
AudioVAE() {
blocks["encoder"] = std::make_shared<AudioEncoder>();
blocks["pre_block"] = std::make_shared<AudioAttentionProjection>();
blocks["mean_proj"] = std::make_shared<LTXV::Conv1D>(kLatentChannels, kLatentChannels, 1);
blocks["logs_proj"] = std::make_shared<LTXV::Conv1D>(kLatentChannels, kLatentChannels, 1);
blocks["dec_in_proj"] = std::make_shared<LTXV::Conv1D>(kLatentChannels,
2048,
1);
blocks["decoder"] = std::make_shared<BigVGAN>();
}
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
SD_UNUSED(tensor_storage_map);
SD_UNUSED(prefix);
params["latents_mean"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, kLatentChannels);
params["latents_std"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, kLatentChannels);
}
ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* waveform) {
GGML_ASSERT(waveform->ne[1] == 2);
auto encoder = std::dynamic_pointer_cast<AudioEncoder>(blocks["encoder"]);
auto pre = std::dynamic_pointer_cast<AudioAttentionProjection>(blocks["pre_block"]);
auto mean_proj = std::dynamic_pointer_cast<LTXV::Conv1D>(blocks["mean_proj"]);
waveform = ggml_reshape_3d(ctx->ggml_ctx, waveform, waveform->ne[0], 1, waveform->ne[1]);
auto x = encoder->forward(ctx, waveform); // [B*S, 2048, T]
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
x = pre->forward(ctx, x);
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
auto z = mean_proj->forward(ctx, x);
auto mean = ggml_reshape_4d(ctx->ggml_ctx, params["latents_mean"], 1, kLatentChannels, 1, 1);
auto std = ggml_reshape_4d(ctx->ggml_ctx, params["latents_std"], 1, kLatentChannels, 1, 1);
z = ggml_div(ctx->ggml_ctx, ggml_sub(ctx->ggml_ctx, z, mean), std);
return ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, z, 0, 2, 1, 3));
}
ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* latent) {
GGML_ASSERT(latent->ne[1] == 2 && latent->ne[2] == kLatentChannels);
latent = ggml_cont(ctx->ggml_ctx,
ggml_permute(ctx->ggml_ctx, latent, 0, 2, 1, 3));
auto mean = ggml_reshape_4d(ctx->ggml_ctx,
params["latents_mean"],
1,
kLatentChannels,
1,
1);
auto std = ggml_reshape_4d(ctx->ggml_ctx,
params["latents_std"],
1,
kLatentChannels,
1,
1);
latent = ggml_add(ctx->ggml_ctx,
ggml_mul(ctx->ggml_ctx, latent, std),
mean);
auto dec_in = std::dynamic_pointer_cast<LTXV::Conv1D>(blocks["dec_in_proj"]);
auto decoder = std::dynamic_pointer_cast<BigVGAN>(blocks["decoder"]);
int64_t streams = latent->ne[2] * latent->ne[3];
latent = ggml_reshape_3d(ctx->ggml_ctx,
latent,
latent->ne[0],
latent->ne[1],
streams);
ggml_tensor* waveform = nullptr;
for (int64_t stream = 0; stream < streams; ++stream) {
auto stream_latent = ggml_ext_slice(ctx->ggml_ctx, latent, 2, stream, stream + 1);
auto stream_waveform = decoder->forward(ctx, dec_in->forward(ctx, stream_latent));
waveform = waveform == nullptr
? stream_waveform
: ggml_concat(ctx->ggml_ctx, waveform, stream_waveform, 2);
}
return ggml_reshape_4d(ctx->ggml_ctx,
waveform,
waveform->ne[0],
streams,
1,
1);
}
};
struct AudioVAERunner : public ::AudioVAERunner {
AudioVAE model;
std::string weight_prefix;
AudioVAERunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map,
const std::string& prefix = "",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: ::AudioVAERunner(backend, weight_manager),
weight_prefix(prefix) {
model.init(params_ctx, tensor_storage_map, prefix);
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
model.get_param_tensors(tensors, weight_prefix);
}
size_t get_params_mem_size() override {
return model.get_params_mem_size();
}
std::string get_desc() override {
return "minimax_h3_audio_vae";
}
int output_sample_rate() const override {
return 32000;
}
sd::Tensor<float> encode(int n_threads,
const sd::Tensor<float>& waveform) override {
int64_t t0 = ggml_time_ms();
auto get_graph = [&]() -> ggml_cgraph* {
auto input = make_input(waveform);
auto runner_ctx = get_context();
auto latent = model.encode(&runner_ctx, input);
auto graph = new_graph_custom(655360);
ggml_build_forward_expand(graph, latent);
return graph;
};
auto result = restore_trailing_singleton_dims(
GGMLRunner::compute<float>(get_graph, n_threads, false, false, false),
4);
int64_t t1 = ggml_time_ms();
LOG_INFO("MiniMax-H3 audio VAE encode completed, taking %.2fs",
(t1 - t0) / 1000.f);
return result;
}
sd::Tensor<float> decode(int n_threads,
const sd::Tensor<float>& latent_tensor) override {
int64_t t0 = ggml_time_ms();
auto get_graph = [&]() -> ggml_cgraph* {
auto latent = make_input(latent_tensor);
auto runner_ctx = get_context();
auto waveform = model.decode(&runner_ctx, latent);
auto graph = new_graph_custom(655360);
ggml_build_forward_expand(graph, waveform);
return graph;
};
auto result = restore_trailing_singleton_dims(
GGMLRunner::compute<float>(get_graph, n_threads, false, false, false),
4);
int64_t t1 = ggml_time_ms();
LOG_INFO("MiniMax-H3 audio VAE decode completed, taking %.2fs",
(t1 - t0) / 1000.f);
return result;
}
};
} // namespace MiniMaxH3
#endif // __SD_MODEL_VAE_MINIMAX_H3_AUDIO_VAE_HPP__

View File

@ -0,0 +1,805 @@
#ifndef __SD_MODEL_VAE_MINIMAX_H3_VAE_HPP__
#define __SD_MODEL_VAE_MINIMAX_H3_VAE_HPP__
#include <algorithm>
#include <array>
#include <cmath>
#include <memory>
#include <string>
#include <tuple>
#include <utility>
#include <vector>
#include "model/common/rope.hpp"
#include "model/diffusion/dit.hpp"
#include "model/vae/vae.hpp"
namespace MiniMaxH3VAE {
constexpr int H3_VIDEO_VAE_GRAPH_SIZE = 262144;
struct CausalConv3d : public Conv3d {
std::tuple<int, int, int> temporal_padding;
CausalConv3d(int64_t in_channels,
int64_t out_channels,
std::tuple<int, int, int> kernel_size,
std::tuple<int, int, int> stride = {1, 1, 1},
std::tuple<int, int, int> padding = {0, 0, 0})
: Conv3d(in_channels,
out_channels,
kernel_size,
stride,
{0, 0, 0}),
temporal_padding(padding) {}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
auto reflect_pad = [&](ggml_tensor* value, int dim, int amount) {
for (int i = 0; i < amount; ++i) {
GGML_ASSERT(value->ne[dim] > 1);
auto left = ggml_ext_slice(ctx->ggml_ctx, value, dim, 1, 2);
auto right = ggml_ext_slice(ctx->ggml_ctx,
value,
dim,
value->ne[dim] - 2,
value->ne[dim] - 1);
value = ggml_concat(ctx->ggml_ctx, left, value, dim);
value = ggml_concat(ctx->ggml_ctx, value, right, dim);
}
return value;
};
x = reflect_pad(x, 0, std::get<2>(temporal_padding));
x = reflect_pad(x, 1, std::get<1>(temporal_padding));
int temporal_pad = std::get<0>(temporal_padding) * 2;
if (temporal_pad > 0) {
x = ggml_ext_pad_ext(ctx->ggml_ctx,
ctx->backend,
x,
0,
0,
0,
0,
temporal_pad,
0,
0,
0);
}
return Conv3d::forward(ctx, x);
}
};
struct TemporalGroupNorm : public GroupNorm {
explicit TemporalGroupNorm(int64_t channels)
: GroupNorm(32, channels, 1e-6f, true) {}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
ggml_tensor* result = nullptr;
for (int64_t t = 0; t < x->ne[2]; ++t) {
auto frame = ggml_ext_slice(ctx->ggml_ctx, x, 2, t, t + 1);
GGML_ASSERT(frame->ne[3] % num_channels == 0);
int64_t batch_size = frame->ne[3] / num_channels;
frame = ggml_cont(ctx->ggml_ctx, frame);
frame = ggml_reshape_4d(ctx->ggml_ctx,
frame,
frame->ne[0],
frame->ne[1],
num_channels,
batch_size);
frame = GroupNorm::forward(ctx, frame);
frame = ggml_reshape_4d(ctx->ggml_ctx,
frame,
frame->ne[0],
frame->ne[1],
1,
num_channels * batch_size);
result = result == nullptr ? frame : ggml_concat(ctx->ggml_ctx, result, frame, 2);
}
return result;
}
};
struct Downsample3D : public GGMLBlock {
int spatial_stride;
Downsample3D(int64_t in_channels,
int64_t out_channels,
int temporal_stride,
int spatial_stride)
: spatial_stride(spatial_stride) {
blocks["conv"] = std::make_shared<CausalConv3d>(in_channels,
out_channels,
std::tuple{3, 3, 3},
std::tuple{temporal_stride, spatial_stride, spatial_stride},
std::tuple{1, 0, 0});
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
if (spatial_stride == 2) {
GGML_ASSERT(x->ne[0] > 1 && x->ne[1] > 1);
auto right = ggml_ext_slice(ctx->ggml_ctx, x, 0, x->ne[0] - 2, x->ne[0] - 1);
x = ggml_concat(ctx->ggml_ctx, x, right, 0);
auto bottom = ggml_ext_slice(ctx->ggml_ctx, x, 1, x->ne[1] - 2, x->ne[1] - 1);
x = ggml_concat(ctx->ggml_ctx, x, bottom, 1);
}
return std::dynamic_pointer_cast<CausalConv3d>(blocks["conv"])->forward(ctx, x);
}
};
struct ResnetBlock3D : public GGMLBlock {
int64_t in_channels;
int64_t out_channels;
ResnetBlock3D(int64_t in_channels,
int64_t out_channels)
: in_channels(in_channels), out_channels(out_channels) {
blocks["norm1"] = std::make_shared<TemporalGroupNorm>(in_channels);
blocks["norm2"] = std::make_shared<TemporalGroupNorm>(out_channels);
blocks["conv1"] = std::make_shared<CausalConv3d>(in_channels,
out_channels,
std::tuple{3, 3, 3},
std::tuple{1, 1, 1},
std::tuple{1, 1, 1});
blocks["conv2"] = std::make_shared<CausalConv3d>(out_channels,
out_channels,
std::tuple{3, 3, 3},
std::tuple{1, 1, 1},
std::tuple{1, 1, 1});
if (in_channels != out_channels) {
blocks["nin_shortcut"] = std::make_shared<CausalConv3d>(in_channels,
out_channels,
std::tuple{1, 1, 1});
}
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto norm1 = std::dynamic_pointer_cast<TemporalGroupNorm>(blocks["norm1"]);
auto norm2 = std::dynamic_pointer_cast<TemporalGroupNorm>(blocks["norm2"]);
auto conv1 = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv1"]);
auto conv2 = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv2"]);
auto h = conv1->forward(ctx, ggml_silu(ctx->ggml_ctx, norm1->forward(ctx, x)));
h = conv2->forward(ctx, ggml_silu(ctx->ggml_ctx, norm2->forward(ctx, h)));
if (in_channels != out_channels) {
x = std::dynamic_pointer_cast<CausalConv3d>(blocks["nin_shortcut"])->forward(ctx, x);
}
return ggml_add(ctx->ggml_ctx, x, h);
}
};
struct Encoder : public GGMLBlock {
static constexpr int levels = 6;
static constexpr std::array<int, levels> multipliers = {1, 2, 2, 4, 4, 8};
static constexpr std::array<int, levels> spatial_down = {2, 2, 2, 2, 1, 1};
static constexpr std::array<int, levels> temporal_down = {1, 2, 2, 1, 1, 1};
Encoder() {
constexpr int ch = 128;
blocks["conv_in"] = std::make_shared<CausalConv3d>(3,
ch,
std::tuple{3, 3, 3},
std::tuple{1, 1, 1},
std::tuple{1, 1, 1});
int64_t previous = ch;
for (int level = 0; level < levels; ++level) {
int64_t current = ch * multipliers[level];
for (int block = 0; block < 2; ++block) {
blocks["down." + std::to_string(level) + ".block." + std::to_string(block)] =
std::make_shared<ResnetBlock3D>(block == 0 ? previous : current,
current);
}
if (spatial_down[level] * temporal_down[level] > 1) {
blocks["down." + std::to_string(level) + ".downsample"] =
std::make_shared<Downsample3D>(current,
current,
temporal_down[level],
spatial_down[level]);
}
previous = current;
}
blocks["norm_out"] = std::make_shared<TemporalGroupNorm>(previous);
blocks["conv_out"] = std::make_shared<CausalConv3d>(previous,
48,
std::tuple{3, 3, 3},
std::tuple{1, 1, 1},
std::tuple{1, 1, 1});
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
x = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv_in"])->forward(ctx, x);
for (int level = 0; level < levels; ++level) {
for (int block = 0; block < 2; ++block) {
x = std::dynamic_pointer_cast<ResnetBlock3D>(
blocks["down." + std::to_string(level) + ".block." + std::to_string(block)])
->forward(ctx, x);
}
auto downsample = blocks.find("down." + std::to_string(level) + ".downsample");
if (downsample != blocks.end()) {
x = std::dynamic_pointer_cast<Downsample3D>(downsample->second)->forward(ctx, x);
}
}
auto norm = std::dynamic_pointer_cast<TemporalGroupNorm>(blocks["norm_out"]);
auto conv = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv_out"]);
return conv->forward(ctx, ggml_silu(ctx->ggml_ctx, norm->forward(ctx, x)));
}
};
static ggml_tensor* attention_layout(ggml_context* ctx, ggml_tensor* x) {
x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3));
return ggml_reshape_3d(ctx, x, x->ne[0], x->ne[1], x->ne[2] * x->ne[3]);
}
static ggml_tensor* apply_partial_rope(ggml_context* ctx,
ggml_tensor* x,
ggml_tensor* pe) {
int64_t rot_dim = pe->ne[2] * 2;
auto rotated = Rope::apply_rope(ctx,
ggml_ext_slice(ctx, x, 0, 0, rot_dim),
pe,
false);
if (rot_dim == x->ne[0]) {
return rotated;
}
auto tail = attention_layout(ctx,
ggml_ext_slice(ctx, x, 0, rot_dim, x->ne[0]));
return ggml_concat(ctx, rotated, tail, 0);
}
struct DecoderAttention : public GGMLBlock {
static constexpr int num_head = 32;
static constexpr int head_dim = 64;
static constexpr int dim = num_head * head_dim;
DecoderAttention() {
blocks["to_qkv"] = std::make_shared<Linear>(dim, dim * 3, true);
blocks["to_out"] = std::make_shared<Linear>(dim, dim, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* pe) {
auto to_qkv = std::dynamic_pointer_cast<Linear>(blocks["to_qkv"]);
auto to_out = std::dynamic_pointer_cast<Linear>(blocks["to_out"]);
auto qkv_projection = to_qkv->forward(ctx, x);
int64_t sequence = x->ne[1];
int64_t batch_size = x->ne[2] * x->ne[3];
qkv_projection = ggml_reshape_4d(ctx->ggml_ctx,
qkv_projection,
3 * head_dim,
num_head,
sequence,
batch_size);
auto qkv = ggml_ext_chunk(ctx->ggml_ctx, qkv_projection, 3, 0);
auto q = ggml_reshape_4d(ctx->ggml_ctx,
qkv[0],
head_dim,
num_head,
sequence,
batch_size);
auto k = ggml_reshape_4d(ctx->ggml_ctx,
qkv[1],
head_dim,
num_head,
sequence,
batch_size);
auto v = ggml_reshape_4d(ctx->ggml_ctx,
qkv[2],
head_dim,
num_head,
sequence,
batch_size);
q = ggml_rms_norm(ctx->ggml_ctx, q, 1e-5f);
k = ggml_rms_norm(ctx->ggml_ctx, k, 1e-5f);
q = apply_partial_rope(ctx->ggml_ctx, q, pe);
k = apply_partial_rope(ctx->ggml_ctx, k, pe);
auto out = ggml_ext_attention_ext(ctx->ggml_ctx,
ctx->backend,
q,
k,
v,
num_head,
nullptr,
true,
ctx->flash_attn_enabled);
return to_out->forward(ctx, out);
}
};
struct DecoderFeedForward : public GGMLBlock {
static constexpr int dim = 2048;
static constexpr int kInnerDim = dim * 4;
DecoderFeedForward() {
blocks["w1"] = std::make_shared<Linear>(dim, kInnerDim * 2, true);
blocks["w2"] = std::make_shared<Linear>(kInnerDim, dim, true);
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto w1 = std::dynamic_pointer_cast<Linear>(blocks["w1"]);
auto w2 = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
auto gate = ggml_ext_chunk(ctx->ggml_ctx, w1->forward(ctx, x), 2, 0);
return w2->forward(ctx,
ggml_mul(ctx->ggml_ctx,
ggml_silu(ctx->ggml_ctx, gate[0]),
gate[1]));
}
};
struct DecoderBlock : public GGMLBlock {
static constexpr int dim = 2048;
DecoderBlock() {
blocks["norm1"] = std::make_shared<RMSNorm>(dim, 1e-5f);
blocks["attn"] = std::make_shared<DecoderAttention>();
blocks["norm2"] = std::make_shared<RMSNorm>(dim, 1e-5f);
blocks["ff"] = std::make_shared<DecoderFeedForward>();
}
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
SD_UNUSED(tensor_storage_map);
SD_UNUSED(prefix);
params["scale1"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dim);
params["scale2"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dim);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* x,
ggml_tensor* pe) {
auto norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm1"]);
auto attn = std::dynamic_pointer_cast<DecoderAttention>(blocks["attn"]);
auto norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm2"]);
auto ff = std::dynamic_pointer_cast<DecoderFeedForward>(blocks["ff"]);
x = ggml_add(ctx->ggml_ctx,
x,
ggml_mul(ctx->ggml_ctx,
attn->forward(ctx, norm1->forward(ctx, x), pe),
params["scale1"]));
return ggml_add(ctx->ggml_ctx,
x,
ggml_mul(ctx->ggml_ctx,
ff->forward(ctx, norm2->forward(ctx, x)),
params["scale2"]));
}
};
struct Decoder : public GGMLBlock {
static constexpr int dim = 2048;
static constexpr int num_layers = 36;
static constexpr int num_register_tokens = 4;
static constexpr int patch_size = 16;
static constexpr int patch_size_t = 4;
Decoder() {
blocks["x_embedder"] = std::make_shared<Linear>(24, dim, true);
for (int i = 0; i < num_layers; ++i) {
blocks["transformer_blocks." + std::to_string(i)] =
std::make_shared<DecoderBlock>();
}
blocks["norm_out"] = std::make_shared<LayerNorm>(dim, 1e-5f, true, true);
blocks["proj_out"] = std::make_shared<Linear>(dim,
3 * patch_size_t * patch_size * patch_size,
true,
true);
}
void init_params(ggml_context* ctx,
const String2TensorStorage& tensor_storage_map = {},
const std::string prefix = "") override {
SD_UNUSED(tensor_storage_map);
SD_UNUSED(prefix);
params["register_tokens"] = ggml_new_tensor_2d(ctx,
GGML_TYPE_F32,
dim,
num_register_tokens);
params["mask_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dim);
}
ggml_tensor* forward(GGMLRunnerContext* ctx,
ggml_tensor* z,
ggml_tensor* pe) {
int64_t width = z->ne[0];
int64_t height = z->ne[1];
int64_t num_frames = z->ne[2];
int64_t batch_size = z->ne[3] / 24;
GGML_ASSERT(batch_size == 1);
z = ggml_cont(ctx->ggml_ctx,
ggml_ext_torch_permute(ctx->ggml_ctx, z, 3, 0, 1, 2));
z = ggml_reshape_3d(ctx->ggml_ctx,
z,
24,
width * height * num_frames,
batch_size);
auto x_embedder = std::dynamic_pointer_cast<Linear>(blocks["x_embedder"]);
auto h = x_embedder->forward(ctx, z);
int64_t num_patches = h->ne[1];
h = ggml_concat(ctx->ggml_ctx, h, params["register_tokens"], 1);
auto zero = ggml_ext_scale(ctx->ggml_ctx,
ggml_ext_slice(ctx->ggml_ctx, h, 1, 0, 1),
0.f);
h = ggml_concat(ctx->ggml_ctx, h, zero, 1);
for (int i = 0; i < num_layers; ++i) {
auto block = std::dynamic_pointer_cast<DecoderBlock>(
blocks["transformer_blocks." + std::to_string(i)]);
h = block->forward(ctx, h, pe);
sd::ggml_graph_cut::mark_graph_cut(h,
"minimax_h3_vae.decoder.blocks." + std::to_string(i),
"hidden_states");
}
auto norm_out = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_out"]);
auto proj_out = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
h = proj_out->forward(ctx, norm_out->forward(ctx, h));
h = ggml_ext_slice(ctx->ggml_ctx, h, 1, 0, num_patches);
return DiT::unpatchify_3d(ctx->ggml_ctx,
h,
num_frames,
height,
width,
patch_size_t,
patch_size,
patch_size,
true);
}
};
struct MiniMaxH3VideoVAE : public GGMLBlock {
MiniMaxH3VideoVAE() {
blocks["encoder"] = std::make_shared<Encoder>();
blocks["quant_conv"] = std::make_shared<Conv3d>(48,
48,
std::tuple{1, 1, 1});
blocks["post_quant_conv"] = std::make_shared<Conv3d>(24,
24,
std::tuple{1, 1, 1});
blocks["decoder"] = std::make_shared<Decoder>();
}
ggml_tensor* encode(GGMLRunnerContext* ctx,
ggml_tensor* pixels,
ggml_tensor* pixel_mean,
ggml_tensor* pixel_std) {
pixels = ggml_div(ctx->ggml_ctx,
ggml_sub(ctx->ggml_ctx, pixels, pixel_mean),
pixel_std);
auto encoder = std::dynamic_pointer_cast<Encoder>(blocks["encoder"]);
auto quant = std::dynamic_pointer_cast<Conv3d>(blocks["quant_conv"]);
auto moments = quant->forward(ctx, encoder->forward(ctx, pixels));
return ggml_ext_slice(ctx->ggml_ctx, moments, 3, 0, 24);
}
ggml_tensor* decode(GGMLRunnerContext* ctx,
ggml_tensor* latent,
ggml_tensor* pe,
ggml_tensor* pixel_mean,
ggml_tensor* pixel_std) {
auto post_quant = std::dynamic_pointer_cast<Conv3d>(blocks["post_quant_conv"]);
auto decoder = std::dynamic_pointer_cast<Decoder>(blocks["decoder"]);
auto pixels = decoder->forward(ctx, post_quant->forward(ctx, latent), pe);
pixels = ggml_add(ctx->ggml_ctx,
ggml_mul(ctx->ggml_ctx, pixels, pixel_std),
pixel_mean);
return ggml_clamp(ctx->ggml_ctx, pixels, 0.f, 1.f);
}
};
struct MiniMaxH3VideoVAERunner : public VAE {
MiniMaxH3VideoVAE model;
sd::Tensor<float> pixel_mean;
sd::Tensor<float> pixel_std;
sd::Tensor<float> latents_mean;
sd::Tensor<float> latents_std;
sd::Tensor<float> rope_cache;
MiniMaxH3VideoVAERunner(ggml_backend_t backend,
const String2TensorStorage& tensor_storage_map,
const std::string& prefix = "first_stage_model",
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
: VAE(VERSION_MINIMAX_H3, backend, prefix, weight_manager),
pixel_mean({1, 1, 1, 3}, {0.485f, 0.456f, 0.406f}),
pixel_std({1, 1, 1, 3}, {0.229f, 0.224f, 0.225f}),
latents_mean({1, 1, 1, 24},
{0.858090341091156f, -0.960659146308899f, 1.066164016723633f, -0.509032547473907f,
-0.272758185863495f, -1.367541432380676f, -0.255325496196747f, -0.269075542688370f,
-0.537684082984924f, -0.046409729868174f, 0.665737032890320f, 0.196901276707649f,
-0.546060800552368f, -0.403534203767776f, -0.236830249428749f, 0.259284526109695f,
-0.301339447498322f, 0.211341992020607f, -1.120684862136841f, 0.358193337917328f,
-0.042251437902451f, 0.260482996702194f, 0.228640928864479f, 0.705603182315826f}),
latents_std({1, 1, 1, 24},
{1.222377419471741f, 1.276726365089417f, 1.683177471160889f, 1.754945516586304f,
1.563621640205383f, 2.194143533706665f, 0.965313792228699f, 1.056988596916199f,
0.841948926448822f, 0.772995293140411f, 1.895593762397766f, 0.946841835975647f,
0.799680948257446f, 0.449889004230499f, 0.719739973545075f, 0.693629324436188f,
2.961095094680786f, 2.769419908523560f, 3.049618482589722f, 2.108805418014527f,
3.276226282119751f, 3.162735700607300f, 2.281681299209595f, 2.612784385681153f}) {
scale_input = false;
model.init(params_ctx, tensor_storage_map, prefix);
}
std::string get_desc() override {
return "minimax_h3_video_vae";
}
int get_encoder_output_channels(int input_channels) override {
SD_UNUSED(input_channels);
return 24;
}
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
model.get_param_tensors(tensors, weight_prefix);
}
sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output,
std::shared_ptr<RNG> rng) override {
SD_UNUSED(rng);
return vae_output;
}
sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) override {
return latents * latents_std + latents_mean;
}
sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) override {
return (latents - latents_mean) / latents_std;
}
static sd::Tensor<float> ensure_video_shape(const sd::Tensor<float>& tensor) {
if (tensor.dim() == 5) {
return tensor;
}
GGML_ASSERT(tensor.dim() == 4);
return tensor.reshape({tensor.shape()[0],
tensor.shape()[1],
1,
tensor.shape()[2],
tensor.shape()[3]});
}
static sd_tiling_params_t h3_tiling(sd_tiling_params_t params) {
params.enabled = true;
params.tile_size_x = 16;
params.tile_size_y = 16;
params.target_overlap = 0.25f;
return params;
}
static sd::Tensor<float> repeat_last_frame(const sd::Tensor<float>& input,
int64_t count) {
auto result = input;
auto last = sd::ops::slice(input, 2, input.shape()[2] - 1, input.shape()[2]);
for (int64_t i = 0; i < count; ++i) {
result = sd::ops::concat(result, last, 2);
}
return result;
}
static sd::Tensor<float> blend_temporal(const sd::Tensor<float>& previous,
const sd::Tensor<float>& current,
int64_t extent) {
auto output = current;
extent = std::min({extent, previous.shape()[2], current.shape()[2]});
int64_t previous_start = previous.shape()[2] - extent;
for (int64_t b = 0; b < current.shape()[4]; ++b) {
for (int64_t c = 0; c < current.shape()[3]; ++c) {
for (int64_t t = 0; t < extent; ++t) {
float wb = static_cast<float>(t) / extent;
float wa = 1.f - wb;
for (int64_t h = 0; h < current.shape()[1]; ++h) {
for (int64_t w = 0; w < current.shape()[0]; ++w) {
output.index(w, h, t, c, b) =
previous.index(w, h, previous_start + t, c, b) * wa +
current.index(w, h, t, c, b) * wb;
}
}
}
}
}
return output;
}
sd::Tensor<float> encode(int n_threads,
const sd::Tensor<float>& x,
sd_tiling_params_t tiling_params,
bool circular_x = false,
bool circular_y = false) override {
auto input = ensure_video_shape(x);
auto tiling = h3_tiling(tiling_params);
if (input.shape()[2] == 1) {
auto encoded = VAE::encode(n_threads, input, tiling, circular_x, circular_y);
if (!encoded.empty() && encoded.shape()[2] > 1) {
encoded = sd::ops::slice(encoded,
2,
encoded.shape()[2] - 1,
encoded.shape()[2]);
}
return encoded;
}
int64_t pad = (-input.shape()[2]) % 17;
if (pad < 0) {
pad += 17;
}
if (pad > 0) {
input = repeat_last_frame(input, pad);
}
sd::Tensor<float> result;
for (int64_t start = 0; start < input.shape()[2]; start += 17) {
auto chunk = sd::ops::slice(input, 2, start, start + 17);
auto encoded = VAE::encode(n_threads, chunk, tiling, circular_x, circular_y);
if (encoded.empty()) {
return {};
}
result = result.empty() ? std::move(encoded)
: sd::ops::concat(result, encoded, 2);
}
if (result.shape()[2] > 3) {
result = sd::ops::slice(result, 2, 0, result.shape()[2] - 3);
}
return result;
}
sd::Tensor<float> decode(int n_threads,
const sd::Tensor<float>& x,
sd_tiling_params_t tiling_params,
bool decode_video = false,
bool circular_x = false,
bool circular_y = false,
bool silent = false) override {
auto input = ensure_video_shape(x);
auto tiling = h3_tiling(tiling_params);
if (input.shape()[2] == 1) {
auto decoded = VAE::decode(n_threads,
input,
tiling,
decode_video,
circular_x,
circular_y,
silent);
if (!decoded.empty() && decoded.shape()[2] > 1) {
decoded = sd::ops::slice(decoded,
2,
decoded.shape()[2] - 1,
decoded.shape()[2]);
}
return decoded;
}
constexpr int64_t tokens_per_chunk = 5;
constexpr int64_t token_drop = 3;
constexpr int64_t token_overlap = 2;
constexpr int64_t frames_per_chunk = 20;
constexpr int64_t frame_pre_padding = 3;
constexpr int64_t frame_overlap = 5;
int64_t pseudo_tokens = input.shape()[2] + token_drop;
int64_t pad_tokens = (tokens_per_chunk - pseudo_tokens % tokens_per_chunk) % tokens_per_chunk;
pseudo_tokens += pad_tokens;
int64_t num_chunks = pseudo_tokens / tokens_per_chunk - 1;
if (num_chunks < 1) {
pad_tokens += tokens_per_chunk;
num_chunks += 1;
}
if (pad_tokens > 0) {
input = repeat_last_frame(input, pad_tokens);
}
sd::Tensor<float> result;
sd::Tensor<float> overlap;
for (int64_t i = 0; i < num_chunks; ++i) {
int64_t start = i * tokens_per_chunk;
int64_t end = std::min(start + tokens_per_chunk + token_overlap,
input.shape()[2]);
auto chunk = sd::ops::slice(input, 2, start, end);
auto decoded = VAE::decode(n_threads,
chunk,
tiling,
true,
circular_x,
circular_y,
silent);
if (decoded.empty()) {
return {};
}
int64_t first_end = std::min<int64_t>(frames_per_chunk, decoded.shape()[2]);
auto first = sd::ops::slice(decoded,
2,
std::min<int64_t>(frame_pre_padding, first_end),
first_end);
if (!overlap.empty()) {
first = blend_temporal(overlap, first, frame_overlap);
overlap = {};
}
result = result.empty() ? std::move(first)
: sd::ops::concat(result, first, 2);
if (decoded.shape()[2] > frames_per_chunk + frame_pre_padding) {
overlap = sd::ops::slice(decoded,
2,
frames_per_chunk + frame_pre_padding,
decoded.shape()[2]);
}
if (i == num_chunks - 1 && !overlap.empty()) {
result = sd::ops::concat(result, overlap, 2);
overlap = {};
}
}
int64_t expected_frames = input.shape()[2] <= 1 ? 1 : ((x.shape()[2] - 2) / 5) * 17 + 5;
expected_frames = std::max<int64_t>(1, expected_frames);
if (result.shape()[2] > expected_frames) {
result = sd::ops::slice(result, 2, 0, expected_frames);
}
return result;
}
sd::Tensor<float> build_rope(int64_t width,
int64_t height,
int64_t num_frames) {
std::vector<std::vector<float>> ids;
ids.reserve(static_cast<size_t>(width * height * num_frames + 5));
constexpr float two_pi = 6.28318530717958647692f;
for (int64_t t = 0; t < num_frames; ++t) {
float pt = (2.f * ((t + 0.5f) / num_frames) - 1.f) * two_pi;
for (int64_t h = 0; h < height; ++h) {
float ph = (2.f * ((h + 0.5f) / height) - 1.f) * two_pi;
for (int64_t w = 0; w < width; ++w) {
float pw = (2.f * ((w + 0.5f) / width) - 1.f) * two_pi;
ids.push_back({pt, ph, pw});
}
}
}
for (int i = 0; i < 5; ++i) {
ids.push_back({0.f, 0.f, 0.f});
}
auto values = Rope::embed_nd(ids,
1,
100.f,
std::vector<int>{16, 16, 16});
return sd::Tensor<float>({2,
2,
24,
static_cast<int64_t>(ids.size())},
std::move(values));
}
sd::Tensor<float> _compute(const int n_threads,
const sd::Tensor<float>& z,
bool decode_graph) override {
auto input = ensure_video_shape(z);
if (decode_graph) {
rope_cache = build_rope(input.shape()[0],
input.shape()[1],
input.shape()[2]);
}
auto get_graph = [&]() -> ggml_cgraph* {
auto value = make_input(input);
auto mean = make_input(pixel_mean);
auto std = make_input(pixel_std);
auto runner_ctx = get_context();
ggml_tensor* out = nullptr;
if (decode_graph) {
auto pe = make_input(rope_cache);
out = model.decode(&runner_ctx, value, pe, mean, std);
} else {
out = model.encode(&runner_ctx, value, mean, std);
}
auto graph = new_graph_custom(H3_VIDEO_VAE_GRAPH_SIZE);
ggml_build_forward_expand(graph, out);
return graph;
};
return restore_trailing_singleton_dims(
GGMLRunner::compute<float>(get_graph,
n_threads,
false,
false,
false),
5);
}
};
} // namespace MiniMaxH3VAE
#endif // __SD_MODEL_VAE_MINIMAX_H3_VAE_HPP__

View File

@ -528,6 +528,9 @@ public:
if (version == VERSION_WAN2_2_TI2V) {
z_channels = 48;
patch = 2;
} else if (sd_version_is_hunyuan_video(version)) {
z_channels = 32;
patch = 2;
} else if (sd_version_is_ltxav(version)) {
z_channels = 128;
patch = 4;
@ -542,12 +545,12 @@ public:
ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* z) {
auto decoder = std::dynamic_pointer_cast<TinyVideoDecoder>(blocks["decoder"]);
if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) {
if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) {
// (W, H, C, T) -> (W, H, T, C)
z = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, z, 0, 1, 3, 2));
}
auto result = decoder->forward(ctx, z);
if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) {
if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) {
// (W, H, T, C) -> (W, H, C, T)
result = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, result, 0, 1, 3, 2));
}
@ -556,7 +559,7 @@ public:
ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* x) {
auto encoder = std::dynamic_pointer_cast<TinyVideoEncoder>(blocks["encoder"]);
if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) {
if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) {
// (W, H, T, C) -> (W, H, C, T)
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
}
@ -569,7 +572,7 @@ public:
}
}
x = encoder->forward(ctx, x);
if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) {
if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) {
// (W, H, C, T) -> (W, H, T, C)
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
}

View File

@ -74,7 +74,7 @@ public:
int scale_factor = 8;
if (version == VERSION_LTXAV) {
scale_factor = 32;
} else if (version == VERSION_WAN2_2_TI2V) {
} else if (version == VERSION_WAN2_2_TI2V || sd_version_is_hunyuan_video(version) || sd_version_is_mage_flow(version) || sd_version_is_minimax_h3(version)) {
scale_factor = 16;
} else if (sd_version_uses_flux2_vae(version)) {
scale_factor = 16;
@ -115,7 +115,7 @@ public:
tile_size_y = get_tile_size(params.tile_size_y, params.rel_size_y, latent_y);
}
sd::Tensor<float> encode(int n_threads,
virtual sd::Tensor<float> encode(int n_threads,
const sd::Tensor<float>& x,
sd_tiling_params_t tiling_params,
bool circular_x = false,
@ -136,7 +136,8 @@ public:
// Image VAE encode is more sensitive to tile boundary context than decode.
// Keep the smaller legacy factor for video VAEs, but default image encode
// tiles to 64 latent pixels so a 512px SD image is encoded as one tile.
const float encode_tile_factor = (sd_version_is_wan(version) || sd_version_is_ltxav(version)) ? 1.30539f : 2.0f;
const float encode_tile_factor = sd_version_is_minimax_h3(version) ? 1.f : (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) ? 1.30539f
: 2.0f;
get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, W, H, encode_tile_factor);
LOG_DEBUG("VAE Tile size: %dx%d", tile_size_x, tile_size_y);
output = tiled_compute(input,
@ -166,7 +167,7 @@ public:
return std::move(output);
}
sd::Tensor<float> decode(int n_threads,
virtual sd::Tensor<float> decode(int n_threads,
const sd::Tensor<float>& x,
sd_tiling_params_t tiling_params,
bool decode_video = false,

View File

@ -21,7 +21,7 @@ bool write_gguf_file(const std::string& file_path,
class GGUFStreamingWriter : public StreamingModelWriter {
public:
GGUFStreamingWriter() = default;
~GGUFStreamingWriter();
~GGUFStreamingWriter() override;
bool write_metadata(const std::string& file_path,
const std::vector<TensorWritePlan>& tensors,

View File

@ -2,6 +2,7 @@
#include <cstdlib>
#include <cstring>
#include <limits>
#include <string>
#include <unordered_map>
#include <utility>
@ -512,8 +513,51 @@ static bool parse_storage_type(const std::string& global_name, PickleStorageInfo
return false;
}
static bool tensor_is_contiguous(const PickleTensorInfo& tensor) {
if (tensor.tensor_storage.nelements() == 0) {
static bool checked_pickle_byte_count(int64_t element_count,
uint64_t element_nbytes,
uint64_t* byte_count) {
if (element_count < 0 || element_nbytes == 0) {
return false;
}
uint64_t count = static_cast<uint64_t>(element_count);
if (count > std::numeric_limits<uint64_t>::max() / element_nbytes) {
return false;
}
*byte_count = count * element_nbytes;
return true;
}
static bool tensor_layout_is_valid(const PickleTensorInfo& tensor, uint64_t raw_element_nbytes) {
if (raw_element_nbytes == 0) {
return false;
}
bool has_zero_dimension = false;
uint64_t element_count = 1;
for (int i = 0; i < tensor.tensor_storage.n_dims; ++i) {
int64_t dimension = tensor.tensor_storage.ne[i];
if (dimension < 0) {
return false;
}
if (dimension == 0) {
has_zero_dimension = true;
continue;
}
uint64_t size = static_cast<uint64_t>(dimension);
if (element_count > static_cast<uint64_t>(std::numeric_limits<int64_t>::max()) / size) {
return false;
}
element_count *= size;
}
if (!has_zero_dimension &&
element_count > static_cast<uint64_t>(std::numeric_limits<int64_t>::max()) / raw_element_nbytes) {
return false;
}
if (has_zero_dimension) {
return true;
}
if (tensor.stride_n_dims != tensor.tensor_storage.n_dims) {
@ -932,7 +976,12 @@ bool parse_torch_state_dict_pickle(const uint8_t* buffer,
if (storage.key.empty() || !parse_storage_type(pid.items[1].str_value, &storage)) {
return false;
}
storage.nbytes = (uint64_t)pid.items[4].int_value * storage.raw_element_nbytes;
if (!checked_pickle_byte_count(pid.items[4].int_value,
storage.raw_element_nbytes,
&storage.nbytes)) {
set_error(error, "invalid storage size in torch pickle");
return false;
}
storage_nbytes[storage.key] = storage.nbytes;
stack.push_back(make_storage_value(storage));
} break;
@ -963,7 +1012,12 @@ bool parse_torch_state_dict_pickle(const uint8_t* buffer,
tensor.tensor_storage.is_f64 = args.items[0].storage.is_f64;
tensor.tensor_storage.is_i64 = args.items[0].storage.is_i64;
tensor.tensor_storage.storage_key = args.items[0].storage.key;
tensor.tensor_storage.offset = (uint64_t)args.items[1].int_value * args.items[0].storage.raw_element_nbytes;
if (!checked_pickle_byte_count(args.items[1].int_value,
args.items[0].storage.raw_element_nbytes,
&tensor.tensor_storage.offset)) {
set_error(error, "invalid tensor storage offset in torch pickle");
return false;
}
for (const auto& item : args.items[2].items) {
if (item.kind != PickleValue::INT || tensor.tensor_storage.n_dims >= SD_MAX_DIMS) {
@ -979,7 +1033,8 @@ bool parse_torch_state_dict_pickle(const uint8_t* buffer,
tensor.stride[tensor.stride_n_dims++] = item.int_value;
}
if (!tensor_is_contiguous(tensor)) {
if (!tensor_layout_is_valid(tensor, args.items[0].storage.raw_element_nbytes)) {
set_error(error, "invalid tensor shape or stride in torch pickle");
return false;
}
stack.push_back(make_tensor_value(tensor));

View File

@ -100,7 +100,8 @@ static ggml_type safetensors_dtype_to_ggml_type(const std::string& dtype) {
// https://huggingface.co/docs/safetensors/index
bool read_safetensors_file(const std::string& file_path,
std::vector<TensorStorage>& tensor_storages,
std::string* error) {
std::string* error,
std::map<std::string, std::string>* metadata) {
std::ifstream file(file_path, std::ios::binary);
if (!file.is_open()) {
set_error(error, "failed to open '" + file_path + "'");
@ -150,6 +151,18 @@ bool read_safetensors_file(const std::string& file_path,
return false;
}
if (metadata != nullptr) {
metadata->clear();
auto metadata_item = header_.find("__metadata__");
if (metadata_item != header_.end() && metadata_item->is_object()) {
for (const auto& item : metadata_item->items()) {
if (item.value().is_string()) {
metadata->emplace(item.key(), item.value().get<std::string>());
}
}
}
}
tensor_storages.clear();
for (auto& item : header_.items()) {
std::string name = item.key();

View File

@ -1,6 +1,7 @@
#ifndef __SD_MODEL_IO_SAFETENSORS_IO_H__
#define __SD_MODEL_IO_SAFETENSORS_IO_H__
#include <map>
#include <string>
#include <vector>
@ -10,7 +11,8 @@
bool is_safetensors_file(const std::string& file_path);
bool read_safetensors_file(const std::string& file_path,
std::vector<TensorStorage>& tensor_storages,
std::string* error = nullptr);
std::string* error = nullptr,
std::map<std::string, std::string>* metadata = nullptr);
bool read_safetensors_index_file(const std::string& file_path,
std::vector<std::string>& shard_paths,
std::string* error = nullptr);

View File

@ -139,11 +139,16 @@ bool read_torch_legacy_file(const std::string& file_path,
if (it == legacy_storage_map.end()) {
return false;
}
if (current_offset + LEGACY_STORAGE_HEADER_SIZE + it->second > file_size) {
if (current_offset > file_size ||
LEGACY_STORAGE_HEADER_SIZE > file_size - current_offset) {
return false;
}
storage_offsets[storage_key] = current_offset + LEGACY_STORAGE_HEADER_SIZE;
current_offset += LEGACY_STORAGE_HEADER_SIZE + it->second;
uint64_t storage_offset = current_offset + LEGACY_STORAGE_HEADER_SIZE;
if (it->second > file_size - storage_offset) {
return false;
}
storage_offsets[storage_key] = storage_offset;
current_offset = storage_offset + it->second;
}
for (auto& tensor_storage : tensor_storages) {
@ -159,8 +164,10 @@ bool read_torch_legacy_file(const std::string& file_path,
uint64_t base_offset = it_offset->second;
uint64_t storage_nbytes = it_size->second;
uint64_t tensor_nbytes = tensor_storage.nbytes_to_read();
if (tensor_storage.offset + tensor_nbytes > storage_nbytes) {
int64_t tensor_nbytes = tensor_storage.nbytes_to_read();
if (tensor_nbytes < 0 ||
tensor_storage.offset > storage_nbytes ||
static_cast<uint64_t>(tensor_nbytes) > storage_nbytes - tensor_storage.offset) {
return false;
}

View File

@ -76,8 +76,10 @@ static bool parse_zip_data_pkl(const uint8_t* buffer,
return false;
}
uint64_t tensor_nbytes = tensor_storage.nbytes_to_read();
if (tensor_storage.offset + tensor_nbytes > entry_size) {
int64_t tensor_nbytes = tensor_storage.nbytes_to_read();
if (tensor_nbytes < 0 ||
tensor_storage.offset > entry_size ||
static_cast<uint64_t>(tensor_nbytes) > entry_size - tensor_storage.offset) {
set_error(error, "tensor '" + tensor_storage.name + "' exceeds storage entry '" + entry_name + "'");
return false;
}

View File

@ -317,7 +317,7 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const
std::vector<TensorStorage> tensor_storages;
std::string error;
if (!read_safetensors_file(file_path, tensor_storages, &error)) {
if (!read_safetensors_file(file_path, tensor_storages, &error, &metadata_)) {
LOG_ERROR("%s", error.c_str());
return false;
}
@ -498,11 +498,18 @@ SDVersion ModelLoader::get_sd_version() {
return VERSION_MINIT2I;
}
if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.img_mod.1.weight") != std::string::npos) {
auto img_in = tensor_storage_map.find("model.diffusion_model.img_in.weight");
if (img_in != tensor_storage_map.end() && img_in->second.ne[0] == 128) {
return VERSION_MAGE_FLOW;
}
if (tensor_storage_map.find("model.diffusion_model.time_text_embed.addition_t_embedding.weight") != tensor_storage_map.end()) {
return VERSION_QWEN_IMAGE_LAYERED;
}
return VERSION_QWEN_IMAGE;
}
if (tensor_storage.name.find("model.diffusion_model.txt_in.individual_token_refiner.blocks.0.adaLN_modulation.1.weight") != std::string::npos) {
return VERSION_HUNYUAN_VIDEO;
}
if (tensor_storage.name.find("llm_adapter.blocks.0.cross_attn.q_proj.weight") != std::string::npos) {
return VERSION_ANIMA;
}
@ -530,9 +537,16 @@ SDVersion ModelLoader::get_sd_version() {
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.video_patch_proj.weight") != std::string::npos &&
tensor_storage_map.find("model.diffusion_model.audio_patch_proj.weight") != tensor_storage_map.end()) {
return VERSION_MINIMAX_H3;
}
if (tensor_storage.name.find("model.diffusion_model.blocks.0.cross_attn.norm_k.weight") != std::string::npos) {
is_wan = true;
}
if (tensor_storage.name.find("model.diffusion_model.patch_embedder.weight") != std::string::npos) {
return VERSION_LINGBOT_VIDEO;
}
if (tensor_storage.name.find("model.diffusion_model.patch_embedding.weight") != std::string::npos) {
patch_embedding_channels = tensor_storage.ne[3];
}
@ -1043,7 +1057,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
std::atomic<size_t> tensor_idx(0);
std::atomic<bool> failed(false);
std::vector<std::thread> workers;
std::mutex rpc_backend_mutex;
std::mutex backend_tensor_set_mutex;
for (int i = 0; i < n_threads; ++i) {
workers.emplace_back([&, file_path, is_zip]() {
@ -1067,6 +1081,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
std::vector<uint8_t> read_buffer;
std::vector<uint8_t> convert_buffer;
std::vector<uint8_t> zip_entry_buffer;
while (true) {
int64_t t0, t1;
@ -1105,34 +1120,60 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
size_t nbytes_to_read = tensor_storage.nbytes_to_read();
auto read_data = [&](char* buf, size_t n) {
auto read_data = [&](char* buf, size_t n) -> bool {
if (zip != nullptr) {
zip_entry_openbyindex(zip, tensor_storage.index_in_zip);
if (zip_entry_openbyindex(zip, tensor_storage.index_in_zip) != 0) {
LOG_ERROR("failed to open zip entry for tensor '%s'", tensor_storage.name.c_str());
return false;
}
size_t entry_size = zip_entry_size(zip);
if (tensor_storage.offset > entry_size) {
LOG_ERROR("tensor '%s' exceeds its zip storage entry", tensor_storage.name.c_str());
zip_entry_close(zip);
return false;
}
size_t tensor_offset = static_cast<size_t>(tensor_storage.offset);
if (n > entry_size - tensor_offset) {
LOG_ERROR("tensor '%s' exceeds its zip storage entry", tensor_storage.name.c_str());
zip_entry_close(zip);
return false;
}
if (entry_size != n) {
int64_t t_memcpy_start;
read_buffer.resize(entry_size);
zip_entry_noallocread(zip, (void*)read_buffer.data(), entry_size);
zip_entry_buffer.resize(entry_size);
auto bytes_read = zip_entry_noallocread(zip, (void*)zip_entry_buffer.data(), entry_size);
if (bytes_read < 0 || static_cast<size_t>(bytes_read) != entry_size) {
LOG_ERROR("failed to read zip entry for tensor '%s'", tensor_storage.name.c_str());
zip_entry_close(zip);
return false;
}
t_memcpy_start = ggml_time_ms();
memcpy((void*)buf, (void*)(read_buffer.data() + tensor_storage.offset), n);
memcpy((void*)buf, (void*)(zip_entry_buffer.data() + tensor_offset), n);
memcpy_time_ms.fetch_add(ggml_time_ms() - t_memcpy_start);
} else {
zip_entry_noallocread(zip, (void*)buf, n);
auto bytes_read = zip_entry_noallocread(zip, (void*)buf, n);
if (bytes_read < 0 || static_cast<size_t>(bytes_read) != n) {
LOG_ERROR("failed to read zip entry for tensor '%s'", tensor_storage.name.c_str());
zip_entry_close(zip);
return false;
}
}
zip_entry_close(zip);
} else if (mmapped) {
if (!mmapped->copy_data(buf, n, tensor_storage.offset)) {
LOG_ERROR("read tensor data failed: '%s'", file_path.c_str());
failed = true;
return false;
}
} else {
file.seekg(tensor_storage.offset);
file.read(buf, n);
if (!file) {
LOG_ERROR("read tensor data failed: '%s'", file_path.c_str());
failed = true;
return false;
}
}
return true;
};
char* read_buf = nullptr;
@ -1166,7 +1207,10 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
}
t0 = ggml_time_ms();
read_data(read_buf, nbytes_to_read);
if (!read_data(read_buf, nbytes_to_read)) {
failed = true;
break;
}
t1 = ggml_time_ms();
read_time_ms.fetch_add(t1 - t0);
@ -1204,17 +1248,8 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
if (dst_tensor->buffer != nullptr && !ggml_backend_buffer_is_host(dst_tensor->buffer)) {
t0 = ggml_time_ms();
// RPC backends require serialized access to prevent concurrency issues
const char* buffer_type_name = ggml_backend_buft_name(ggml_backend_buffer_get_type(dst_tensor->buffer));
bool is_rpc_buffer = buffer_type_name != nullptr &&
std::string(buffer_type_name).find("RPC") != std::string::npos;
if (is_rpc_buffer) {
std::lock_guard<std::mutex> lock(rpc_backend_mutex);
std::lock_guard<std::mutex> lock(backend_tensor_set_mutex);
ggml_backend_tensor_set(dst_tensor, convert_buf, 0, ggml_nbytes(dst_tensor));
} else {
ggml_backend_tensor_set(dst_tensor, convert_buf, 0, ggml_nbytes(dst_tensor));
}
t1 = ggml_time_ms();
copy_to_backend_time_ms.fetch_add(t1 - t0);

View File

@ -36,6 +36,7 @@ protected:
std::vector<ModelFileData> file_data;
bool model_files_processed = false;
String2TensorStorage tensor_storage_map;
std::map<std::string, std::string> metadata_;
int n_threads_;
size_t add_file_path(const std::string& file_path);
@ -63,6 +64,7 @@ public:
std::map<ggml_type, uint32_t> get_vae_wtype_stat();
String2TensorStorage& get_tensor_storage_map() { return tensor_storage_map; }
const String2TensorStorage& get_tensor_storage_map() const { return tensor_storage_map; }
const std::map<std::string, std::string>& get_metadata() const { return metadata_; }
void set_n_threads(int n_threads);
void set_wtype_override(ggml_type wtype, std::string tensor_type_rules = "");
void process_model_files(bool enable_mmap = false, bool writable_mmap = true);

View File

@ -53,6 +53,48 @@ static bool backend_supports_host_buffer(ggml_backend_t backend) {
return props.caps.buffer_from_host_ptr;
}
static bool device_supports_param_op(ggml_backend_dev_t device,
ggml_tensor* weight,
enum ggml_op op,
ggml_backend_buffer_type_t buft) {
if (op == GGML_OP_NONE) {
return true;
}
if (device == nullptr || weight == nullptr || buft == nullptr || weight->buffer != nullptr) {
return false;
}
ggml_init_params params;
params.mem_size = ggml_tensor_overhead() * 2;
params.mem_buffer = nullptr;
params.no_alloc = true;
ggml_context* ctx = ggml_init(params);
if (ctx == nullptr) {
return false;
}
ggml_tensor* op_tensor = nullptr;
if (op == GGML_OP_GET_ROWS) {
ggml_tensor* indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1);
op_tensor = ggml_get_rows(ctx, weight, indices);
}
if (op_tensor == nullptr) {
ggml_free(ctx);
return false;
}
weight->buffer = ggml_backend_buft_alloc_buffer(buft, 0);
if (weight->buffer == nullptr) {
ggml_free(ctx);
return false;
}
bool supported = ggml_backend_dev_supports_op(device, op_tensor);
ggml_backend_buffer_free(weight->buffer);
weight->buffer = nullptr;
ggml_free(ctx);
return supported;
}
ModelManager::~ModelManager() {
release_all();
}
@ -135,7 +177,8 @@ bool ModelManager::register_param_tensors(const std::string& desc,
ggml_backend_t params_backend,
size_t* registered_tensor_size,
bool allow_split_buffer,
bool params_follow_compute_backend) {
bool params_follow_compute_backend,
const std::map<ggml_tensor*, enum ggml_op>* tensor_ops) {
if (desc.empty()) {
LOG_ERROR("model manager tensor desc is empty");
return false;
@ -168,6 +211,12 @@ bool ModelManager::register_param_tensors(const std::string& desc,
state->params_backend = params_backend;
state->allow_split_buffer = allow_split_buffer;
state->params_follow_compute_backend = params_follow_compute_backend;
if (tensor_ops != nullptr) {
auto op_it = tensor_ops->find(tensor);
if (op_it != tensor_ops->end()) {
state->usage_op = op_it->second;
}
}
new_states.push_back(std::move(state));
}
@ -844,6 +893,22 @@ ggml_backend_buffer_type_t ModelManager::params_buffer_type_for(const TensorStat
if (params_buft == nullptr) {
params_buft = ggml_backend_get_default_buffer_type(state.params_backend);
}
if (state.usage_op != GGML_OP_NONE &&
state.compute_backend != nullptr) {
ggml_backend_dev_t compute_dev = ggml_backend_get_device(state.compute_backend);
if (device_supports_param_op(compute_dev, state.tensor, state.usage_op, params_buft)) {
return params_buft;
}
ggml_backend_dev_t cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
params_buft = cpu_dev != nullptr ? ggml_backend_dev_buffer_type(cpu_dev) : nullptr;
if (!device_supports_param_op(cpu_dev, state.tensor, state.usage_op, params_buft)) {
LOG_ERROR("model manager has no compatible buffer for tensor '%s' used by %s",
state.name.c_str(),
ggml_op_name(state.usage_op));
return nullptr;
}
}
return params_buft;
}

View File

@ -39,6 +39,7 @@ private:
bool allow_split_buffer = false;
bool params_follow_compute_backend = false;
bool metadata_validated = false;
enum ggml_op usage_op = GGML_OP_NONE;
int active_prepare_count = 0;
@ -132,7 +133,8 @@ public:
ggml_backend_t params_backend,
size_t* registered_tensor_size = nullptr,
bool allow_split_buffer = false,
bool params_follow_compute_backend = false);
bool params_follow_compute_backend = false,
const std::map<ggml_tensor*, enum ggml_op>* tensor_ops = nullptr);
bool unregister_param_tensors(const std::string& desc,
size_t* registered_tensor_size = nullptr);

View File

@ -185,6 +185,20 @@ std::string convert_cond_stage_model_name(std::string name, std::string prefix)
}
std::string convert_qwen3_vl_vision_name(std::string name) {
static const std::vector<std::pair<std::string, std::string>> qwen3_vl_deepstack_name_map{
{"v.deepstack_merger_list.", "deepstack_merger_list."},
{"v.deepstack.5.", "deepstack_merger_list.0."},
{"v.deepstack.8.", "deepstack_merger_list.0."},
{"v.deepstack.11.", "deepstack_merger_list.1."},
{"v.deepstack.16.", "deepstack_merger_list.1."},
{"v.deepstack.17.", "deepstack_merger_list.2."},
{"v.deepstack.24.", "deepstack_merger_list.2."},
{"fc1.", "linear_fc1."},
{"fc2.", "linear_fc2."},
{"ffn_up.", "linear_fc1."},
{"ffn_down.", "linear_fc2."},
{"ffn_norm.", "norm."},
};
static const std::vector<std::pair<std::string, std::string>> qwen3_vl_vision_name_map{
{"mm.0.", "merger.linear_fc1."},
{"mm.2.", "merger.linear_fc2."},
@ -201,6 +215,10 @@ std::string convert_qwen3_vl_vision_name(std::string name) {
{"ln1.", "norm1."},
{"ln2.", "norm2."},
};
if (contains(name, "v.deepstack_merger_list.") || contains(name, "v.deepstack.")) {
replace_with_name_map(name, qwen3_vl_deepstack_name_map);
return name;
}
replace_with_name_map(name, qwen3_vl_vision_name_map);
return name;
}
@ -304,6 +322,12 @@ std::string convert_diffusers_unet_to_original_sd1(std::string name) {
}
}
static const std::vector<std::pair<std::string, std::string>> name_map{
{"to_out.weight", "to_out.0.weight"},
{"to_out.bias", "to_out.0.bias"},
};
replace_with_name_map(result, name_map);
return result;
}
@ -658,6 +682,72 @@ std::string convert_diffusers_dit_to_original_flux(std::string name) {
return name;
}
std::string convert_hunyuan_video_to_original_flux(std::string name) {
int num_layers = 54;
int num_single_layers = 0;
static std::unordered_map<std::string, std::string> hy_name_map;
if (hy_name_map.empty()) {
// --- double transformer blocks ---
for (int i = 0; i < num_layers; ++i) {
std::string block_prefix = "double_blocks." + std::to_string(i) + ".";
std::string dst_prefix = "double_blocks." + std::to_string(i) + ".";
hy_name_map[block_prefix + "img_mod.linear"] = dst_prefix + "img_mod.lin";
hy_name_map[block_prefix + "txt_mod.linear"] = dst_prefix + "txt_mod.lin";
// attn
hy_name_map[block_prefix + "img_attn_qkv"] = dst_prefix + "img_attn.qkv";
hy_name_map[block_prefix + "txt_attn_qkv"] = dst_prefix + "txt_attn.qkv";
// norm
hy_name_map[block_prefix + "img_attn_q_norm.weight"] = dst_prefix + "img_attn.norm.query_norm.scale";
hy_name_map[block_prefix + "img_attn_k_norm.weight"] = dst_prefix + "img_attn.norm.key_norm.scale";
hy_name_map[block_prefix + "txt_attn_q_norm.weight"] = dst_prefix + "txt_attn.norm.query_norm.scale";
hy_name_map[block_prefix + "txt_attn_k_norm.weight"] = dst_prefix + "txt_attn.norm.key_norm.scale";
// ff
hy_name_map[block_prefix + "img_mlp.fc1"] = dst_prefix + "img_mlp.0";
hy_name_map[block_prefix + "img_mlp.fc2"] = dst_prefix + "img_mlp.2";
hy_name_map[block_prefix + "txt_mlp.fc1"] = dst_prefix + "txt_mlp.0";
hy_name_map[block_prefix + "txt_mlp.fc2"] = dst_prefix + "txt_mlp.2";
// output projections
hy_name_map[block_prefix + "img_attn_proj"] = dst_prefix + "img_attn.proj";
hy_name_map[block_prefix + "txt_attn_proj"] = dst_prefix + "txt_attn.proj";
}
}
hy_name_map["time_in.mlp.0"] = "time_in.in_layer";
hy_name_map["time_in.mlp.2"] = "time_in.out_layer";
hy_name_map["time_r_in.mlp.0"] = "time_r_in.in_layer";
hy_name_map["time_r_in.mlp.2"] = "time_r_in.out_layer";
hy_name_map["vector_in.mlp.0"] = "vector_in.in_layer";
hy_name_map["vector_in.mlp.2"] = "vector_in.out_layer";
hy_name_map["guidance_in.mlp.0"] = "guidance_in.in_layer";
hy_name_map["guidance_in.mlp.2"] = "guidance_in.out_layer";
hy_name_map["txt_in.c_embedder.linear_1"] = "txt_in.c_embedder.in_layer";
hy_name_map["txt_in.c_embedder.linear_2"] = "txt_in.c_embedder.out_layer";
hy_name_map["txt_in.t_embedder.mlp.0"] = "txt_in.t_embedder.in_layer";
hy_name_map["txt_in.t_embedder.mlp.2"] = "txt_in.t_embedder.out_layer";
replace_with_prefix_map(name, hy_name_map);
static const std::vector<std::pair<std::string, std::string>> generic_name_map = {
{"_attn_qkv.", "_attn.qkv."},
{"_attn_proj.", "_attn.proj."},
{"mlp.fc1.", "mlp.0."},
{"mlp.fc2.", "mlp.2."},
{".modulation.linear.", ".modulation.lin."},
};
replace_with_name_map(name, generic_name_map);
return name;
}
std::string convert_diffusers_dit_to_original_lumina2(std::string name) {
int num_layers = 30;
int num_refiner_layers = 2;
@ -801,6 +891,8 @@ std::string convert_diffusion_model_name(std::string name, std::string prefix, S
name = convert_diffusers_dit_to_original_sd3(name);
} else if (sd_version_is_flux(version) || sd_version_is_flux2(version) || sd_version_is_longcat(version) || sd_version_is_sefi_image(version)) {
name = convert_diffusers_dit_to_original_flux(name);
} else if (sd_version_is_hunyuan_video(version)) {
name = convert_hunyuan_video_to_original_flux(name);
} else if (sd_version_is_z_image(version)) {
name = convert_diffusers_dit_to_original_lumina2(name);
} else if (sd_version_is_anima(version)) {
@ -974,6 +1066,9 @@ std::string convert_diffusers_to_original_wan_vae(std::string name) {
}
std::string convert_first_stage_model_name(std::string name, std::string prefix, SDVersion version) {
if (sd_version_is_hunyuan_video(version) || sd_version_is_mage_flow(version) || sd_version_is_minimax_h3(version)) {
return name;
}
if (sd_version_uses_wan_vae(version)) {
return convert_diffusers_to_original_wan_vae(name);
}
@ -1102,6 +1197,8 @@ std::string convert_sep_to_dot(std::string name) {
"norm1_context",
"ff_context",
"x_embedder",
"cross_attn",
"output_proj",
};
// record the positions of underscores that should NOT be replaced
@ -1206,11 +1303,54 @@ static std::string convert_esrgan_tensor_name(std::string name) {
return name;
}
static const std::map<int, std::string>& ip_adapter_index_map(SDVersion version) {
static const std::map<int, std::string> sd15_map = {
{1, "input_blocks.1.1.transformer_blocks.0"}, {3, "input_blocks.2.1.transformer_blocks.0"}, {5, "input_blocks.4.1.transformer_blocks.0"}, {7, "input_blocks.5.1.transformer_blocks.0"}, {9, "input_blocks.7.1.transformer_blocks.0"}, {11, "input_blocks.8.1.transformer_blocks.0"}, {13, "output_blocks.3.1.transformer_blocks.0"}, {15, "output_blocks.4.1.transformer_blocks.0"}, {17, "output_blocks.5.1.transformer_blocks.0"}, {19, "output_blocks.6.1.transformer_blocks.0"}, {21, "output_blocks.7.1.transformer_blocks.0"}, {23, "output_blocks.8.1.transformer_blocks.0"}, {25, "output_blocks.9.1.transformer_blocks.0"}, {27, "output_blocks.10.1.transformer_blocks.0"}, {29, "output_blocks.11.1.transformer_blocks.0"}, {31, "middle_block.1.transformer_blocks.0"}};
static std::map<int, std::string> sdxl_map;
if (sdxl_map.empty()) {
std::vector<std::pair<std::string, int>> order = {
{"input_blocks.4.1", 2}, {"input_blocks.5.1", 2}, {"input_blocks.7.1", 10}, {"input_blocks.8.1", 10}, {"output_blocks.0.1", 10}, {"output_blocks.1.1", 10}, {"output_blocks.2.1", 10}, {"output_blocks.3.1", 2}, {"output_blocks.4.1", 2}, {"output_blocks.5.1", 2}, {"middle_block.1", 10}};
int idx = 1;
for (const auto& [block, depth] : order) {
for (int m = 0; m < depth; m++) {
sdxl_map[idx] = block + ".transformer_blocks." + std::to_string(m);
idx += 2;
}
}
}
return sd_version_is_sdxl(version) ? sdxl_map : sd15_map;
}
static std::string convert_ip_adapter_name(std::string name, SDVersion version) {
if (starts_with(name, "image_proj.")) {
return "ip_adapter." + name;
}
if (starts_with(name, "ip_adapter.")) {
auto items = split_string(name, '.');
if (items.size() < 4) {
return name;
}
int idx = atoi(items[1].c_str());
const auto& mp = ip_adapter_index_map(version);
auto blk = mp.find(idx);
if (blk == mp.end()) {
return name;
}
return "model.diffusion_model." + blk->second + ".attn2." + items[2] + "." + items[3];
}
return name;
}
std::string convert_tensor_name(std::string name, SDVersion version) {
if (version == VERSION_ESRGAN) {
return convert_esrgan_tensor_name(std::move(name));
}
if (starts_with(name, "ip_adapter.") || starts_with(name, "image_proj.")) {
return convert_ip_adapter_name(std::move(name), version);
}
bool is_lora = false;
bool is_lycoris_underline = false;
bool is_underline = false;
@ -1262,6 +1402,8 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
{".lora_B.weight", ".weight.lora_up"},
{".lora_A.default.weight", ".weight.lora_down"},
{".lora_B.default.weight", ".weight.lora_up"},
{".lora_A", ".weight.lora_down"},
{".lora_B", ".weight.lora_up"},
{".lora_linear", ".weight.alpha"},
{".alpha", ".weight.alpha"},
{".scale", ".weight.scale"},
@ -1327,27 +1469,54 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
{"te2.", "cond_stage_model.1.transformer."},
{"te1.", "cond_stage_model.transformer."},
{"te3.", "text_encoders.t5xxl.transformer."},
{"clip_vision.", "cond_stage_model.transformer."},
};
if (sd_version_is_flux(version)) {
prefix_map["te1."] = "text_encoders.clip_l.transformer.";
}
if (sd_version_is_unet(version)) {
prefix_map["clip_l."] = "cond_stage_model.transformer.";
prefix_map["clip_g."] = "cond_stage_model.1.transformer.";
} else {
prefix_map["clip_l."] = "text_encoders.clip_l.transformer.";
prefix_map["clip_g."] = "text_encoders.clip_g.transformer.";
}
replace_with_prefix_map(name, prefix_map);
if ((sd_version_is_boogu_image(version) || sd_version_is_krea2(version)) && starts_with(name, "text_encoders.llm.visual.")) {
if (sd_version_is_boogu_image(version) || sd_version_is_krea2(version) || sd_version_is_mage_flow(version) || sd_version_is_minimax_h3(version)) {
const std::string hf_vision_prefix = "text_encoders.llm.model.visual.";
if (starts_with(name, hf_vision_prefix)) {
name = "text_encoders.llm.visual." + name.substr(hf_vision_prefix.size());
}
if (starts_with(name, "text_encoders.llm.visual.")) {
name = convert_qwen3_vl_vision_name(std::move(name));
}
}
// diffusion model
{
bool matched = false;
for (const auto& prefix : diffuison_model_prefix_vec) {
if (starts_with(name, prefix)) {
name = convert_diffusion_model_name(name.substr(prefix.size()), prefix, version);
name = prefix + name;
matched = true;
break;
}
}
if (is_lora && !matched && !diffuison_model_prefix_vec.empty()) {
if (starts_with(name, "down_blocks.") || starts_with(name, "up_blocks.") ||
starts_with(name, "mid_block.") || starts_with(name, "conv_in.") ||
starts_with(name, "conv_out.") || starts_with(name, "time_embedding.") ||
starts_with(name, "conv_norm_out.")) {
const std::string& canonical_prefix = diffuison_model_prefix_vec.front();
name = convert_diffusion_model_name(name, canonical_prefix, version);
name = canonical_prefix + name;
}
}
}
// cond_stage_model

View File

@ -306,8 +306,33 @@ struct KarrasScheduler : SigmaScheduler {
};
struct BetaScheduler : SigmaScheduler {
static constexpr double alpha = 0.6;
static constexpr double beta = 0.6;
double alpha = 0.6;
double beta = 0.6;
explicit BetaScheduler(const char* extra_sample_args = nullptr) {
parse_extra_sample_args(extra_sample_args);
LOG_DEBUG("Beta scheduler: alpha=%.4f, beta=%.4f", alpha, beta);
}
void parse_extra_sample_args(const char* extra_sample_args) {
for (const auto& [key, value] : parse_key_value_args(extra_sample_args, "beta scheduler arg")) {
if (key == "alpha") {
float parsed;
if (!parse_strict_float(value, parsed) || parsed <= 0.0) {
LOG_WARN("ignoring invalid beta scheduler arg '%s=%s'", key.c_str(), value.c_str());
} else {
alpha = static_cast<double>(parsed);
}
} else if (key == "beta") {
float parsed;
if (!parse_strict_float(value, parsed) || parsed <= 0.0) {
LOG_WARN("ignoring invalid beta scheduler arg '%s=%s'", key.c_str(), value.c_str());
} else {
beta = static_cast<double>(parsed);
}
}
}
}
static double log_beta(double a, double b) {
return std::lgamma(a) + std::lgamma(b) - std::lgamma(a + b);
@ -1032,7 +1057,7 @@ struct Denoiser {
break;
case BETA_SCHEDULER:
LOG_INFO("get_sigmas with Beta scheduler");
scheduler = std::make_shared<BetaScheduler>();
scheduler = std::make_shared<BetaScheduler>(extra_sample_args);
break;
case EXPONENTIAL_SCHEDULER:
LOG_INFO("get_sigmas exponential scheduler");
@ -1153,7 +1178,7 @@ struct CompVisDenoiser : public Denoiser {
return {c_skip, c_out, c_in};
}
virtual sd::Tensor<float> noise_scaling(float sigma,
sd::Tensor<float> noise_scaling(float sigma,
const sd::Tensor<float>& noise,
const sd::Tensor<float>& latent) override {
GGML_ASSERT(noise.numel() == latent.numel());
@ -1165,7 +1190,7 @@ struct CompVisDenoiser : public Denoiser {
return latent;
}
float noise_level_to_sigma(float noise_level) {
float noise_level_to_sigma(float noise_level) override {
return noise_level / (1.0f - noise_level);
}
};
@ -1256,7 +1281,7 @@ struct DiscreteFlowDenoiser : public Denoiser {
return latent * (1.0f / (1.0f - sigma));
}
float noise_level_to_sigma(float noise_level) {
float noise_level_to_sigma(float noise_level) override {
return noise_level;
}
};
@ -1396,7 +1421,7 @@ struct MiniT2IFlowDenoiser : public Denoiser {
return latent;
}
float noise_level_to_sigma(float noise_level) {
float noise_level_to_sigma(float noise_level) override {
SD_UNUSED(noise_level);
return 1.0f;
}
@ -2553,6 +2578,88 @@ static sd::Tensor<float> sample_tcd(denoise_cb_t model,
return x;
}
static sd::Tensor<float> sample_lms(denoise_cb_t model,
sd::Tensor<float> x,
const std::vector<float>& sigmas,
const SamplerExtraArgs& extra_sample_args) {
// Linear Multi-Step from https://github.com/crowsonkb/k-diffusion
int divisions = 1000;
for (const auto& [key, value] : extra_sample_args) {
int parsed = 0;
if (key == "lms_divisions") {
if (!parse_strict_int(value, parsed)) {
LOG_WARN("ignoring invalid lms extra sample arg '%s=%s'", key.c_str(), value.c_str());
continue;
}
divisions = parsed; // std::max(1, parsed);
// values above 35M produce noise, can be fixed by double precision
// values < 1 always produce noise
}
}
LOG_DEBUG("linear multi-step sampler: integrating using %i division%s", divisions, (divisions == 1) ? "" : "s");
auto linear_multistep_coeff = [=](const int order, const int m, const int j) -> float {
if (!divisions)
return sigmas[m + 1] - sigmas[m]; // delta / 0 * 0
#define LMS_PRECISION float // double
const LMS_PRECISION a = sigmas[m], dx = (sigmas[m + 1] - a) / divisions, s = sigmas[m - j];
const LMS_PRECISION b0 = a + 0.5f * dx; // using Riemann middle integral
LMS_PRECISION sum = 0.0f;
for (int h = 0; h < divisions; h++) {
const LMS_PRECISION b = h * dx + b0;
LMS_PRECISION prod = 1.0f;
for (int k = 0; k < j; k++) {
const LMS_PRECISION t = sigmas[m - k];
prod *= (b - t) / (s - t);
}
for (int k = j + 1; k < order; k++) {
const LMS_PRECISION t = sigmas[m - k];
prod *= (b - t) / (s - t);
}
sum += prod;
}
return sum * dx;
};
const int max_order = 4;
float lms_coeff[max_order];
std::vector<sd::Tensor<float>> hist = {};
int steps = static_cast<int>(sigmas.size()) - 1;
for (int i = 0; i < steps; i++) {
const float sigma = sigmas[i];
auto denoised_opt = model(x, sigma, i + 1);
if (denoised_opt.pred.empty()) {
return {};
}
sd::Tensor<float> denoised = std::move(denoised_opt.pred);
const int order = std::min(max_order, i + 1);
for (int c = 0; c < order; c++) // computing coefficients
lms_coeff[c] = linear_multistep_coeff(order, i, c);
sd::Tensor<float> d_cur = (x - denoised) / sigma;
switch (order) {
case 4: // derivative + 3 history points
x += hist[hist.size() - 2] * lms_coeff[3];
case 3:
x += hist[hist.size() - 1] * lms_coeff[2];
case 2:
x += hist.back() * lms_coeff[1];
case 1:
x += d_cur * lms_coeff[0];
}
if (hist.size() == static_cast<size_t>(max_order - 1)) {
hist.erase(hist.begin());
}
hist.push_back(std::move(d_cur));
}
return x;
}
static sd::Tensor<float> sample_euler_cfg_pp(denoise_cb_t model,
sd::Tensor<float> x,
const std::vector<float>& sigmas) {
@ -2714,6 +2821,8 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
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 LMS_SAMPLE_METHOD:
return sample_lms(model, std::move(x), sigmas, extra_args);
case EULER_CFG_PP_SAMPLE_METHOD:
return sample_euler_cfg_pp(model, std::move(x), sigmas);
case EULER_A_CFG_PP_SAMPLE_METHOD:

View File

@ -4,6 +4,33 @@
#include "core/tensor.hpp"
#include "ggml.h"
const float minimax_latent_rgb_proj[24][3] = {
{0.19819857f, 0.11584999f, 0.07929777f},
{-0.16047224f, -0.10601170f, -0.15996324f},
{0.47391951f, 0.37602475f, 0.20267826f},
{-0.09857441f, -0.27435449f, -0.51681751f},
{-0.18930605f, -0.10512278f, -0.28571478f},
{-0.15639569f, -0.18000929f, -0.25432852f},
{-0.07176921f, -0.10901598f, -0.06654253f},
{-0.05014077f, -0.05839826f, -0.05516087f},
{-0.05201424f, -0.04351913f, -0.01507579f},
{0.24750438f, 0.13307422f, 0.17684120f},
{0.07377446f, 0.10235858f, 0.11707827f},
{0.02908304f, 0.06587022f, 0.10643690f},
{-0.00670531f, -0.03857879f, 0.01750151f},
{-0.07119107f, -0.03083323f, -0.01995450f},
{-0.08612627f, -0.07253841f, -0.01442890f},
{0.08793202f, 0.08681750f, 0.02994647f},
{0.00876893f, 0.02721868f, 0.00091178f},
{-0.03484412f, -0.02711262f, -0.00110101f},
{-0.00679772f, -0.01844275f, -0.01683359f},
{0.04287028f, 0.01601068f, 0.04037397f},
{-0.00493432f, -0.00230528f, 0.00353911f},
{0.01495088f, 0.00292306f, 0.00416671f},
{0.00495307f, 0.05066542f, 0.05210543f},
{-0.02154842f, -0.01518524f, 0.00442402f}};
float minimax_latent_rgb_bias[3] = {0.07776964f, -0.01580954f, -0.06561434f};
const float ltxav_latent_rgb_proj[128][3] = {
{-0.0293802f, -0.0362516f, -0.0291386f},
{0.0117735f, 0.0223435f, 0.018856f},

File diff suppressed because it is too large Load Diff

View File

@ -205,7 +205,7 @@ std::vector<int> BPETokenizer::encode(const std::string& text, on_new_token_cb_t
ss << "\"" << token << "\", ";
}
ss << "]";
LOG_DEBUG("split prompt \"%s\" to tokens %s", text.c_str(), ss.str().c_str());
LOG_DEBUG("split prompt \"%s\" to %zu tokens %s", text.c_str(), bpe_tokens.size(), ss.str().c_str());
return bpe_tokens;
}