Compare commits
39 Commits
master-775
...
master
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
c6beeef355 | ||
|
|
bfbef5b7e6 | ||
|
|
ea7f0c87cf | ||
|
|
b4e67d1221 | ||
|
|
db99efdd6d | ||
|
|
eb7f35ca49 | ||
|
|
50062a4bba | ||
|
|
8457624101 | ||
|
|
10378f42db | ||
|
|
e31a86ce91 | ||
|
|
735a4ef520 | ||
|
|
af92790ffc | ||
|
|
e92e86fb11 | ||
|
|
9cfe2af8f9 | ||
|
|
2993b7fb43 | ||
|
|
53856e7ec8 | ||
|
|
22516991cb | ||
|
|
5ef4a7557d | ||
|
|
2d0385ba85 | ||
|
|
87a01773be | ||
|
|
b0f856804c | ||
|
|
78124b6454 | ||
|
|
b338b4b4b9 | ||
|
|
b8bf67615c | ||
|
|
8d377073e3 | ||
|
|
5114672c48 | ||
|
|
8a51eb9284 | ||
|
|
35fb21f3e0 | ||
|
|
cfd4cff2e6 | ||
|
|
5e4e03c8dd | ||
|
|
2961182365 | ||
|
|
ea4e566ccf | ||
|
|
b290693977 | ||
|
|
7717e82c00 | ||
|
|
fafe8e606c | ||
|
|
a8a91b24cd | ||
|
|
c00a9e956d | ||
|
|
833369da84 | ||
|
|
74bce049d0 |
263
.github/workflows/build.yml
vendored
@ -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:
|
||||
|
||||
@ -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
|
||||
|
||||
16
README.md
@ -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,20 +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)
|
||||
@ -121,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
|
||||
@ -163,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)
|
||||
@ -177,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!
|
||||
|
||||
[](https://github.com/leejet/stable-diffusion.cpp/graphs/contributors)
|
||||
|
||||
## Star History
|
||||
|
||||
[](https://star-history.com/#leejet/stable-diffusion.cpp&Date)
|
||||
|
||||
BIN
assets/animatediff/img2video_demo.gif
Normal file
|
After Width: | Height: | Size: 1.7 MiB |
BIN
assets/animatediff/v2_coast.gif
Normal file
|
After Width: | Height: | Size: 1.4 MiB |
BIN
assets/animatediff/v2_house.gif
Normal file
|
After Width: | Height: | Size: 1.8 MiB |
BIN
assets/animatediff/v2_man.gif
Normal file
|
After Width: | Height: | Size: 1.6 MiB |
BIN
assets/animatediff/v2_rabbit.gif
Normal file
|
After Width: | Height: | Size: 1002 KiB |
BIN
assets/animatediff/v3_coast.gif
Normal file
|
After Width: | Height: | Size: 1.1 MiB |
BIN
assets/animatediff/v3_house.gif
Normal file
|
After Width: | Height: | Size: 1.8 MiB |
BIN
assets/animatediff/v3_man.gif
Normal file
|
After Width: | Height: | Size: 1.8 MiB |
BIN
assets/animatediff/v3_rabbit.gif
Normal file
|
After Width: | Height: | Size: 1.2 MiB |
BIN
assets/animatediff/v3_rabbit_domain_lora.gif
Normal file
|
After Width: | Height: | Size: 1.4 MiB |
BIN
assets/huanyuan_video/hy1.5_t2v.mp4
Normal file
BIN
assets/hunyuan_video/hy1.5_t2v.mp4
Normal file
BIN
assets/mage_flow/edit_example.png
Normal file
|
After Width: | Height: | Size: 466 KiB |
BIN
assets/mage_flow/example.png
Normal file
|
After Width: | Height: | Size: 399 KiB |
BIN
assets/minimax-h3/i2av.mp4
Normal file
BIN
assets/minimax-h3/r2av.mp4
Normal file
BIN
assets/minimax-h3/t2av.mp4
Normal 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
@ -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
@ -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.
|
||||
@ -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
@ -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
@ -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
@ -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
|
||||
```
|
||||
45
docs/mage_flow.md
Normal 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
@ -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 kitten’s 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 cat’s 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 cat’s face, focusing on its cute expression and detailed fur texture. The camera slowly rotates around the cat’s 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 cat’s 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 cat’s 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.
|
||||
@ -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.
|
||||
|
||||
|
||||
@ -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
|
||||
@ -34,4 +34,4 @@ Using formats of different precisions will yield results of varying quality.
|
||||
|
||||
<p align="center">
|
||||
<img src="../assets/img2img_output.png" width="256x">
|
||||
</p>
|
||||
</p>
|
||||
|
||||
@ -1,4 +1,4 @@
|
||||
include_directories(${CMAKE_CURRENT_SOURCE_DIR})
|
||||
|
||||
add_subdirectory(cli)
|
||||
add_subdirectory(server)
|
||||
add_subdirectory(server)
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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,30 +2529,47 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
|
||||
pulid_id_weight,
|
||||
};
|
||||
|
||||
params.loras = lora_vec.empty() ? nullptr : lora_vec.data();
|
||||
params.lora_count = static_cast<uint32_t>(lora_vec.size());
|
||||
params.prompt = prompt.c_str();
|
||||
params.negative_prompt = negative_prompt.c_str();
|
||||
params.clip_skip = clip_skip;
|
||||
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.mask_image = mask_image.get();
|
||||
params.width = get_resolved_width();
|
||||
params.height = get_resolved_height();
|
||||
params.sample_params = sample_params;
|
||||
params.strength = strength;
|
||||
params.seed = seed;
|
||||
params.batch_count = batch_count;
|
||||
params.qwen_image_layers = qwen_image_layers;
|
||||
params.control_image = control_image.get();
|
||||
params.control_strength = control_strength;
|
||||
params.pm_params = pm_params;
|
||||
params.pulid_params = pulid_params;
|
||||
params.vae_tiling_params = vae_tiling_params;
|
||||
params.cache = cache_params;
|
||||
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();
|
||||
params.negative_prompt = negative_prompt.c_str();
|
||||
params.clip_skip = clip_skip;
|
||||
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.ref_image_args = ref_image_args.c_str();
|
||||
params.mask_image = mask_image.get();
|
||||
params.width = get_resolved_width();
|
||||
params.height = get_resolved_height();
|
||||
params.sample_params = sample_params;
|
||||
params.strength = strength;
|
||||
params.seed = seed;
|
||||
params.batch_count = batch_count;
|
||||
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;
|
||||
params.cache = cache_params;
|
||||
|
||||
params.hires.enabled = hires_enabled;
|
||||
params.hires.upscaler = resolved_hires_upscaler;
|
||||
@ -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"
|
||||
@ -2675,8 +2845,13 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params,
|
||||
int64_t seed,
|
||||
SDMode mode) {
|
||||
json root;
|
||||
root["schema"] = "sdcpp.image.params/v1";
|
||||
root["mode"] = mode == VID_GEN ? "vid_gen" : "img_gen";
|
||||
root["schema"] = "sdcpp.image.params/v1";
|
||||
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) + ", ";
|
||||
|
||||
@ -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;
|
||||
|
||||
|
||||
@ -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;
|
||||
}
|
||||
|
||||
@ -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__
|
||||
|
||||
@ -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 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_;
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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},
|
||||
|
||||
@ -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,
|
||||
|
||||
86
scripts/convert_yolov8_to_safetensors.py
Normal 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()
|
||||
335
scripts/merge_safetensors.py
Normal 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()
|
||||
@ -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()) {
|
||||
@ -101,11 +124,13 @@ static inline sd::Tensor<float> apply_token_weights(sd::Tensor<float> hidden_sta
|
||||
|
||||
struct ConditionerParams {
|
||||
std::string text;
|
||||
int clip_skip = -1;
|
||||
int width = -1;
|
||||
int height = -1;
|
||||
bool zero_out_masked = false;
|
||||
const std::vector<sd::Tensor<float>>* ref_images = nullptr; // for qwen image edit
|
||||
int clip_skip = -1;
|
||||
int width = -1;
|
||||
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_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)) {
|
||||
} 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,
|
||||
@ -1932,8 +2016,10 @@ struct LLMEmbedder : public Conditioner {
|
||||
const std::vector<std::pair<int, sd::Tensor<float>>>& image_embeds,
|
||||
const std::set<int>& out_layers,
|
||||
int prompt_template_encode_start_idx,
|
||||
bool spell_quotes = false,
|
||||
int max_length = 100000000) {
|
||||
bool spell_quotes = false,
|
||||
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;
|
||||
@ -2007,9 +2145,151 @@ struct LLMEmbedder : public Conditioner {
|
||||
bool spell_quotes = false;
|
||||
std::set<int> out_layers;
|
||||
|
||||
int64_t t0 = ggml_time_ms();
|
||||
int64_t t0 = ggml_time_ms();
|
||||
RefImageResizeMode resize_mode = conditioner_params.ref_image_params.vlm_resize_mode;
|
||||
|
||||
if (sd_version_is_lingbot_video(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, "
|
||||
@ -2040,28 +2320,35 @@ struct LLMEmbedder : public Conditioner {
|
||||
|
||||
for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
|
||||
const auto& image = (*conditioner_params.ref_images)[i];
|
||||
double factor = llm->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 min_pixels = static_cast<int>(4 * factor * factor);
|
||||
int max_pixels = static_cast<int>(16384 * factor * factor);
|
||||
int h_bar = std::max(static_cast<int>(factor), static_cast<int>(std::round(height / factor) * factor));
|
||||
int w_bar = std::max(static_cast<int>(factor), static_cast<int>(std::round(width / factor) * factor));
|
||||
|
||||
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);
|
||||
}
|
||||
if (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 (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 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);
|
||||
@ -2086,36 +2373,39 @@ struct LLMEmbedder : public Conditioner {
|
||||
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)) {
|
||||
} 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("QwenImageEditPlusPipeline");
|
||||
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);
|
||||
|
||||
@ -2126,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++) {
|
||||
@ -2154,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;
|
||||
|
||||
@ -2170,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);
|
||||
|
||||
@ -2221,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);
|
||||
|
||||
@ -2268,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);
|
||||
|
||||
@ -2491,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],
|
||||
@ -2512,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;
|
||||
}
|
||||
};
|
||||
@ -2622,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);
|
||||
|
||||
@ -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.
|
||||
ggml_set_input(tensor);
|
||||
}
|
||||
// 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");
|
||||
|
||||
@ -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";
|
||||
}
|
||||
|
||||
@ -20,6 +20,7 @@ enum class SDBackendModule {
|
||||
CONTROL_NET,
|
||||
PHOTOMAKER,
|
||||
UPSCALER,
|
||||
DETECTOR,
|
||||
};
|
||||
|
||||
struct SDBackendAssignment {
|
||||
|
||||
1020
src/detailer.cpp
Normal file
75
src/detailer.h
Normal 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__
|
||||
@ -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);
|
||||
|
||||
27
src/model.h
@ -38,10 +38,12 @@ enum SDVersion {
|
||||
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,
|
||||
@ -54,6 +56,7 @@ enum SDVersion {
|
||||
VERSION_IDEOGRAM4,
|
||||
VERSION_SEFI_IMAGE,
|
||||
VERSION_KREA2,
|
||||
VERSION_MAGE_FLOW,
|
||||
VERSION_ESRGAN,
|
||||
VERSION_COUNT,
|
||||
};
|
||||
@ -121,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;
|
||||
@ -142,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;
|
||||
@ -219,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;
|
||||
@ -240,6 +258,10 @@ static inline bool sd_version_uses_wan_vae(SDVersion version) {
|
||||
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 ||
|
||||
@ -255,10 +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) ||
|
||||
@ -270,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;
|
||||
|
||||
209
src/model/adapter/ip_adapter.hpp
Normal 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__
|
||||
@ -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,7 +342,9 @@ 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);
|
||||
hada_1_up = ggml_cont(ctx, ggml_transpose(ctx, hada_1_up));
|
||||
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);
|
||||
@ -354,7 +360,9 @@ 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);
|
||||
hada_2_up = ggml_cont(ctx, ggml_transpose(ctx, hada_2_up));
|
||||
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) {
|
||||
@ -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_value = ggml_ext_backend_tensor_get_f32(iter->second);
|
||||
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;
|
||||
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;
|
||||
}
|
||||
|
||||
@ -310,17 +310,33 @@ protected:
|
||||
int64_t context_dim;
|
||||
int64_t n_head;
|
||||
int64_t d_head;
|
||||
bool xtra_dim = false;
|
||||
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;
|
||||
|
||||
@ -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,
|
||||
|
||||
362
src/model/detector/yolov8.h
Normal 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__
|
||||
@ -484,10 +484,11 @@ namespace Anima {
|
||||
ggml_tensor* timestep,
|
||||
ggml_tensor* encoder_hidden_states,
|
||||
ggml_tensor* image_pe,
|
||||
ggml_tensor* t5_ids = nullptr,
|
||||
ggml_tensor* t5_weights = nullptr,
|
||||
ggml_tensor* adapter_q_pe = nullptr,
|
||||
ggml_tensor* adapter_k_pe = nullptr) {
|
||||
ggml_tensor* t5_ids = nullptr,
|
||||
ggml_tensor* t5_weights = nullptr,
|
||||
ggml_tensor* adapter_q_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"]);
|
||||
@ -502,8 +503,16 @@ namespace Anima {
|
||||
auto padding_mask = ggml_ext_zeros(ctx->ggml_ctx, x->ne[0], x->ne[1], 1, x->ne[3]);
|
||||
x = ggml_concat(ctx->ggml_ctx, x, padding_mask, 2); // [N, C + 1, H, W]
|
||||
|
||||
x = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size); // [N, h*w, (C+1)*ph*pw]
|
||||
|
||||
x = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size); // [N, h*w, (C+1)*ph*pw]
|
||||
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);
|
||||
@ -626,14 +636,20 @@ namespace Anima {
|
||||
|
||||
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<int32_t>& t5_ids_tensor = {},
|
||||
const sd::Tensor<float>& t5_weights_tensor = {}) {
|
||||
const sd::Tensor<float>& context_tensor = {},
|
||||
const sd::Tensor<int32_t>& t5_ids_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;
|
||||
@ -691,11 +709,13 @@ namespace Anima {
|
||||
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<int32_t>& t5_ids = {},
|
||||
const sd::Tensor<float>& t5_weights = {}) {
|
||||
const sd::Tensor<float>& context = {},
|
||||
const sd::Tensor<int32_t>& t5_ids = {},
|
||||
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
|
||||
|
||||
182
src/model/diffusion/animatediff.hpp
Normal 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__
|
||||
@ -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
|
||||
|
||||
@ -706,11 +706,13 @@ namespace Flux {
|
||||
LastLayer(int64_t hidden_size,
|
||||
int64_t patch_size,
|
||||
int64_t out_channels,
|
||||
bool prune_mod = false,
|
||||
bool bias = true)
|
||||
bool prune_mod = false,
|
||||
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);
|
||||
|
||||
681
src/model/diffusion/hunyuan.hpp
Normal 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__
|
||||
@ -180,9 +180,12 @@ namespace Krea2 {
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||
ggml_tensor* scale = params["scale"];
|
||||
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);
|
||||
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);
|
||||
return x;
|
||||
}
|
||||
};
|
||||
@ -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
|
||||
|
||||
@ -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);
|
||||
}
|
||||
|
||||
162
src/model/diffusion/mage_flow.hpp
Normal 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__
|
||||
1178
src/model/diffusion/minimax_h3.hpp
Normal 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,
|
||||
embed_dim,
|
||||
{patch_size, patch_size},
|
||||
{patch_size, patch_size},
|
||||
{0, 0},
|
||||
{1, 1},
|
||||
bias));
|
||||
if (is_3d) {
|
||||
blocks["proj"] = std::make_shared<Conv3d>(in_chans,
|
||||
embed_dim,
|
||||
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 = ggml_pad(ctx->ggml_ctx, x, pad_w, pad_h, pad_t, 0); // TODO: reflect pad mode
|
||||
}
|
||||
x = proj->forward(ctx, x);
|
||||
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;
|
||||
}
|
||||
|
||||
@ -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{};
|
||||
};
|
||||
|
||||
|
||||
@ -17,30 +17,38 @@ namespace Pid {
|
||||
constexpr float PID_PI = 3.14159265358979323846f;
|
||||
|
||||
struct PixelDiTConfig {
|
||||
int64_t in_channels = 3;
|
||||
int64_t hidden_size = 1536;
|
||||
int64_t num_groups = 24;
|
||||
int64_t patch_mlp_hidden_dim = 4096;
|
||||
int64_t pixel_hidden_size = 16;
|
||||
int64_t pixel_attn_hidden_size = 1152;
|
||||
int64_t pixel_num_groups = 16;
|
||||
int64_t patch_depth = 14;
|
||||
int64_t pixel_depth = 2;
|
||||
int64_t patch_size = 16;
|
||||
int64_t txt_embed_dim = 2304;
|
||||
int64_t txt_max_length = 300;
|
||||
float text_rope_theta = 10000.f;
|
||||
int64_t lq_latent_channels = 16;
|
||||
int64_t lq_hidden_dim = 512;
|
||||
int64_t lq_num_res_blocks = 4;
|
||||
int64_t lq_interval = 2;
|
||||
int64_t lq_sr_scale = 4;
|
||||
int64_t lq_latent_down_factor = 8;
|
||||
int64_t rope_ref_grid_h = 64;
|
||||
int64_t rope_ref_grid_w = 64;
|
||||
int64_t in_channels = 3;
|
||||
int64_t hidden_size = 1536;
|
||||
int64_t num_groups = 24;
|
||||
int64_t patch_mlp_hidden_dim = 4096;
|
||||
int64_t pixel_hidden_size = 16;
|
||||
int64_t pixel_attn_hidden_size = 1152;
|
||||
int64_t pixel_num_groups = 16;
|
||||
int64_t patch_depth = 14;
|
||||
int64_t pixel_depth = 2;
|
||||
int64_t patch_size = 16;
|
||||
int64_t txt_embed_dim = 2304;
|
||||
int64_t txt_max_length = 300;
|
||||
float text_rope_theta = 10000.f;
|
||||
int64_t lq_latent_channels = 16;
|
||||
int64_t lq_hidden_dim = 512;
|
||||
int64_t lq_num_res_blocks = 4;
|
||||
int64_t lq_interval = 2;
|
||||
int64_t lq_sr_scale = 4;
|
||||
int64_t lq_latent_down_factor = 8;
|
||||
int64_t 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) {
|
||||
blocks["block.0"] = std::make_shared<GroupNorm>(4, channels, 1e-5f);
|
||||
blocks["block.2"] = std::make_shared<Conv2d>(channels, channels, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, std::pair<int, int>{1, 1});
|
||||
blocks["block.3"] = std::make_shared<GroupNorm>(4, channels, 1e-5f);
|
||||
blocks["block.5"] = std::make_shared<Conv2d>(channels, channels, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, std::pair<int, int>{1, 1});
|
||||
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}, 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}, padding);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
@ -497,9 +557,15 @@ 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));
|
||||
h = conv1->forward(ctx, h);
|
||||
h = ggml_silu_inplace(ctx->ggml_ctx, norm2->forward(ctx, h));
|
||||
h = conv2->forward(ctx, h);
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
@ -543,9 +616,29 @@ namespace Pid {
|
||||
ggml_tensor* lq_latent,
|
||||
int64_t target_pH,
|
||||
int64_t target_pW) {
|
||||
auto conv0 = std::dynamic_pointer_cast<Conv2d>(blocks["latent_proj.0"]);
|
||||
auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["latent_proj.2"]);
|
||||
float z_to_patch_ratio = static_cast<float>(config.lq_sr_scale * config.lq_latent_down_factor) /
|
||||
auto conv0 = std::dynamic_pointer_cast<Conv2d>(blocks["latent_proj.0"]);
|
||||
auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["latent_proj.2"]);
|
||||
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,9 +651,15 @@ 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);
|
||||
feat = conv2->forward(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)]);
|
||||
feat = block->forward(ctx, feat);
|
||||
@ -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);
|
||||
|
||||
@ -103,9 +103,9 @@ namespace Qwen {
|
||||
}
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* timesteps,
|
||||
ggml_tensor* addition_t_cond = nullptr) {
|
||||
virtual ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* timesteps,
|
||||
ggml_tensor* addition_t_cond = nullptr) {
|
||||
// timesteps: [N,]
|
||||
// return: [N, embedding_dim]
|
||||
auto timestep_embedder = std::dynamic_pointer_cast<TimestepEmbedding>(blocks["timestep_embedder"]);
|
||||
@ -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));
|
||||
int64_t inner_dim = config.num_attention_heads * config.attention_head_dim;
|
||||
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));
|
||||
@ -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() {
|
||||
|
||||
@ -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,14 +775,17 @@ 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);
|
||||
ggml_tensor* timesteps = make_input(timesteps_tensor);
|
||||
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* x = make_input(x_tensor);
|
||||
ggml_tensor* timesteps = make_input(timesteps_tensor);
|
||||
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) {
|
||||
@ -737,7 +796,9 @@ struct UNetModelRunner : public DiffusionModelRunner {
|
||||
num_video_frames = static_cast<int>(x->ne[3]);
|
||||
}
|
||||
|
||||
auto runner_ctx = get_context();
|
||||
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() {
|
||||
|
||||
@ -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() {
|
||||
|
||||
@ -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};
|
||||
bool split_patch_embed = false;
|
||||
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,13 +981,13 @@ namespace LLM {
|
||||
return spatial_merge_size;
|
||||
}
|
||||
|
||||
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) {
|
||||
std::vector<ggml_tensor*> forward_outputs(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) {
|
||||
// pixel_values: [grid_t*(H/mh/ph)*(W/mw/pw)*mh*mw, C*pt*ph*pw]
|
||||
// window_index: [grid_t*(H/mh/ph)*(W/mw/pw)]
|
||||
// window_inverse_index: [grid_t*(H/mh/ph)*(W/mw/pw)]
|
||||
@ -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)]);
|
||||
|
||||
@ -922,8 +1015,12 @@ namespace LLM {
|
||||
if (fullatt_block_indexes.find(i) != fullatt_block_indexes.end()) {
|
||||
mask = nullptr;
|
||||
}
|
||||
x = block->forward(ctx, x, pe, mask);
|
||||
if (i == 0) {
|
||||
x = block->forward(ctx, x, pe, mask);
|
||||
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,7 +1368,9 @@ namespace LLM {
|
||||
for (int i = 0; i < num_layers; i++) {
|
||||
blocks["layers." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new TransformerBlock(config, i));
|
||||
}
|
||||
blocks["norm"] = std::shared_ptr<GGMLBlock>(new LLMRMSNorm(config.hidden_size, config.rms_norm_eps, config.rms_norm_add));
|
||||
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,
|
||||
@ -1274,9 +1385,11 @@ namespace LLM {
|
||||
ggml_tensor* input_pos,
|
||||
ggml_tensor* attention_mask,
|
||||
std::set<int> out_layers,
|
||||
ggml_tensor* sliding_attention_mask = nullptr,
|
||||
bool return_all_hidden_states = false) {
|
||||
auto norm = std::dynamic_pointer_cast<LLMRMSNorm>(blocks["norm"]);
|
||||
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 = 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 return_all_hidden_states = false,
|
||||
bool auto_free = true,
|
||||
bool free_compute_buffer = 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 {
|
||||
|
||||
@ -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));
|
||||
|
||||
28
src/model/vae/audio_vae.hpp
Normal 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__
|
||||
834
src/model/vae/hunyuan_vae.hpp
Normal 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__
|
||||
@ -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
@ -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__
|
||||
497
src/model/vae/minimax_h3_audio_vae.hpp
Normal 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__
|
||||
805
src/model/vae/minimax_h3_vae.hpp
Normal 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__
|
||||
@ -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));
|
||||
}
|
||||
|
||||
@ -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,11 +115,11 @@ public:
|
||||
tile_size_y = get_tile_size(params.tile_size_y, params.rel_size_y, latent_y);
|
||||
}
|
||||
|
||||
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) {
|
||||
virtual 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) {
|
||||
int64_t t0 = ggml_time_ms();
|
||||
sd::Tensor<float> input = x;
|
||||
sd::Tensor<float> output;
|
||||
@ -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,13 +167,13 @@ public:
|
||||
return std::move(output);
|
||||
}
|
||||
|
||||
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) {
|
||||
virtual 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) {
|
||||
int64_t t0 = ggml_time_ms();
|
||||
sd::Tensor<float> input = x;
|
||||
sd::Tensor<float> output;
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -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));
|
||||
|
||||
@ -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();
|
||||
|
||||
@ -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);
|
||||
|
||||
@ -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;
|
||||
}
|
||||
|
||||
|
||||
@ -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;
|
||||
}
|
||||
|
||||
@ -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,6 +537,10 @@ 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;
|
||||
}
|
||||
@ -1046,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]() {
|
||||
@ -1070,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;
|
||||
@ -1108,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;
|
||||
@ -1169,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);
|
||||
|
||||
@ -1207,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);
|
||||
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));
|
||||
}
|
||||
std::lock_guard<std::mutex> lock(backend_tensor_set_mutex);
|
||||
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);
|
||||
|
||||
@ -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);
|
||||
|
||||
@ -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;
|
||||
}
|
||||
|
||||
|
||||
@ -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;
|
||||
|
||||
@ -130,9 +131,10 @@ public:
|
||||
ResidencyMode residency_mode,
|
||||
ggml_backend_t compute_backend,
|
||||
ggml_backend_t params_backend,
|
||||
size_t* registered_tensor_size = nullptr,
|
||||
bool allow_split_buffer = false,
|
||||
bool params_follow_compute_backend = false);
|
||||
size_t* registered_tensor_size = nullptr,
|
||||
bool allow_split_buffer = 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);
|
||||
|
||||
@ -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.")) {
|
||||
name = convert_qwen3_vl_vision_name(std::move(name));
|
||||
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;
|
||||
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
|
||||
|
||||
@ -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,9 +1178,9 @@ struct CompVisDenoiser : public Denoiser {
|
||||
return {c_skip, c_out, c_in};
|
||||
}
|
||||
|
||||
virtual sd::Tensor<float> noise_scaling(float sigma,
|
||||
const sd::Tensor<float>& noise,
|
||||
const sd::Tensor<float>& latent) override {
|
||||
sd::Tensor<float> noise_scaling(float sigma,
|
||||
const sd::Tensor<float>& noise,
|
||||
const sd::Tensor<float>& latent) override {
|
||||
GGML_ASSERT(noise.numel() == latent.numel());
|
||||
return latent + noise * sigma;
|
||||
}
|
||||
@ -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:
|
||||
|
||||
@ -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},
|
||||
|
||||
@ -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;
|
||||
}
|
||||
|
||||
|
||||