Compare commits
40 Commits
master-777
...
master
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
bcc7e29568 | ||
|
|
487de75c96 | ||
|
|
b4f1fd6d7f | ||
|
|
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 |
263
.github/workflows/build.yml
vendored
@ -449,8 +449,8 @@ jobs:
|
|||||||
runs-on: windows-2022
|
runs-on: windows-2022
|
||||||
|
|
||||||
env:
|
env:
|
||||||
ROCM_VERSION: "7.13.0"
|
ROCM_VERSION: "7.14.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"
|
GPU_TARGETS: "gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201"
|
||||||
|
|
||||||
steps:
|
steps:
|
||||||
- uses: actions/checkout@v3
|
- uses: actions/checkout@v3
|
||||||
@ -472,34 +472,68 @@ jobs:
|
|||||||
uses: actions/cache@v4
|
uses: actions/cache@v4
|
||||||
with:
|
with:
|
||||||
path: C:\TheRock\build
|
path: C:\TheRock\build
|
||||||
key: rocm-${{ env.ROCM_VERSION }}-gfx1151-${{ runner.os }}
|
key: rocm-wheels-${{ env.ROCM_VERSION }}-${{ runner.os }}
|
||||||
|
|
||||||
- name: ccache
|
- name: ccache
|
||||||
uses: ggml-org/ccache-action@v1.2.16
|
uses: ggml-org/ccache-action@v1.2.16
|
||||||
with:
|
with:
|
||||||
key: windows-latest-rocm-${{ env.ROCM_VERSION }}-x64
|
key: windows-rocm-${{ env.ROCM_VERSION }}-x64
|
||||||
evict-old-files: 1d
|
evict-old-files: 1d
|
||||||
|
|
||||||
- name: Install ROCm
|
- name: Install ROCm with Wheels
|
||||||
if: steps.cache-rocm.outputs.cache-hit != 'true'
|
if: steps.cache-rocm.outputs.cache-hit != 'true'
|
||||||
run: |
|
run: |
|
||||||
$ErrorActionPreference = "Stop"
|
$ErrorActionPreference = "Stop"
|
||||||
write-host "Downloading AMD ROCm ${{ env.ROCM_VERSION }} tarball"
|
write-host "Setting up Python virtual environment"
|
||||||
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"
|
# Create the venv directly at the cache location to avoid relocation issues
|
||||||
mkdir C:\TheRock\build -Force
|
New-Item -Path "C:\TheRock\build" -ItemType Directory -Force | Out-Null
|
||||||
tar -xzf "${env:RUNNER_TEMP}\rocm.tar.gz" -C C:\TheRock\build --strip-components=1
|
python -m venv C:\TheRock\build\.venv
|
||||||
write-host "Completed ROCm extraction"
|
& 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
|
- name: Setup ROCm Environment
|
||||||
run: |
|
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 "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_DEVICE_LIB_PATH=$rocmPath\lib\llvm\amdgcn\bitcode" >> $env:GITHUB_ENV
|
||||||
echo "HIP_PLATFORM=amd" >> $env:GITHUB_ENV
|
echo "HIP_PLATFORM=amd" >> $env:GITHUB_ENV
|
||||||
echo "LLVM_PATH=$rocmPath\lib\llvm" >> $env:GITHUB_ENV
|
echo "LLVM_PATH=$rocmPath\lib\llvm" >> $env:GITHUB_ENV
|
||||||
echo "$rocmPath\bin" >> $env:GITHUB_PATH
|
echo "$binPath" >> $env:GITHUB_PATH
|
||||||
echo "$rocmPath\lib\llvm\bin" >> $env:GITHUB_PATH
|
|
||||||
|
# Keep venv in PATH for subsequent steps
|
||||||
|
echo "C:\TheRock\build\.venv\Scripts" >> $env:GITHUB_PATH
|
||||||
|
|
||||||
- name: Build
|
- name: Build
|
||||||
run: |
|
run: |
|
||||||
@ -527,139 +561,6 @@ jobs:
|
|||||||
- name: Pack artifacts
|
- name: Pack artifacts
|
||||||
if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
|
if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
|
||||||
run: |
|
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\*
|
7z a sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-rocm-${{ env.ROCM_VERSION }}-x64.zip .\build\bin\*
|
||||||
|
|
||||||
- name: Upload artifacts
|
- name: Upload artifacts
|
||||||
@ -679,11 +580,8 @@ jobs:
|
|||||||
strategy:
|
strategy:
|
||||||
matrix:
|
matrix:
|
||||||
include:
|
include:
|
||||||
- ROCM_VERSION: "7.2.1"
|
- ROCM_VERSION: "7.14.0"
|
||||||
gpu_targets: "gfx908;gfx90a;gfx942;gfx1030;gfx1031;gfx1032;gfx1100;gfx1101;gfx1102;gfx1151;gfx1150;gfx1200;gfx1201"
|
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'
|
|
||||||
- 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"
|
|
||||||
build: x64
|
build: x64
|
||||||
|
|
||||||
steps:
|
steps:
|
||||||
@ -702,7 +600,7 @@ jobs:
|
|||||||
- name: Dependencies
|
- name: Dependencies
|
||||||
id: depends
|
id: depends
|
||||||
run: |
|
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
|
- name: Free disk space
|
||||||
run: |
|
run: |
|
||||||
@ -723,38 +621,36 @@ jobs:
|
|||||||
sudo apt clean
|
sudo apt clean
|
||||||
df -h
|
df -h
|
||||||
|
|
||||||
- name: Setup Legacy ROCm
|
- name: Setup TheRock with Wheels
|
||||||
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'
|
|
||||||
id: therock_env
|
id: therock_env
|
||||||
run: |
|
run: |
|
||||||
wget https://repo.amd.com/rocm/tarball/therock-dist-linux-gfx1151-${{ matrix.ROCM_VERSION }}.tar.gz
|
# Create Python virtual environment
|
||||||
mkdir install
|
python3 -m venv .venv
|
||||||
tar -xf *.tar.gz -C install
|
source .venv/bin/activate
|
||||||
export ROCM_PATH=$(pwd)/install
|
|
||||||
echo ROCM_PATH=$ROCM_PATH >> $GITHUB_ENV
|
# Install ROCm wheels for build
|
||||||
echo PATH=$PATH:$ROCM_PATH/bin >> $GITHUB_ENV
|
# libraries = HIP runtime and CMake configs needed for linking
|
||||||
echo LD_LIBRARY_PATH=$ROCM_PATH/lib:$ROCM_PATH/llvm/lib:$ROCM_PATH/lib/rocprofiler-systems >> $GITHUB_ENV
|
# 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.
|
# setup-node installs into /opt/hostedtoolcache, which is removed above.
|
||||||
# Keep Node/pnpm setup after disk cleanup so the server frontend can be embedded.
|
# Keep Node/pnpm setup after disk cleanup so the server frontend can be embedded.
|
||||||
@ -839,7 +735,6 @@ jobs:
|
|||||||
- build-and-push-docker-images
|
- build-and-push-docker-images
|
||||||
- macOS-latest-cmake
|
- macOS-latest-cmake
|
||||||
- windows-latest-cmake
|
- windows-latest-cmake
|
||||||
- windows-latest-cmake-hip
|
|
||||||
- windows-latest-rocm
|
- windows-latest-rocm
|
||||||
|
|
||||||
steps:
|
steps:
|
||||||
|
|||||||
@ -11,10 +11,11 @@ endif()
|
|||||||
if (MSVC)
|
if (MSVC)
|
||||||
add_compile_definitions(_CRT_SECURE_NO_WARNINGS)
|
add_compile_definitions(_CRT_SECURE_NO_WARNINGS)
|
||||||
add_compile_definitions(_SILENCE_CXX17_CODECVT_HEADER_DEPRECATION_WARNING)
|
add_compile_definitions(_SILENCE_CXX17_CODECVT_HEADER_DEPRECATION_WARNING)
|
||||||
|
# /MP is MSVC-only: icx rejects it outright once offloading is enabled.
|
||||||
add_compile_options(
|
add_compile_options(
|
||||||
$<$<COMPILE_LANGUAGE:C>:/MP>
|
$<$<AND:$<COMPILE_LANGUAGE:C>,$<C_COMPILER_ID:MSVC>>:/MP>
|
||||||
$<$<COMPILE_LANGUAGE:C>:/utf-8>
|
$<$<COMPILE_LANGUAGE:C>:/utf-8>
|
||||||
$<$<COMPILE_LANGUAGE:CXX>:/MP>
|
$<$<AND:$<COMPILE_LANGUAGE:CXX>,$<CXX_COMPILER_ID:MSVC>>:/MP>
|
||||||
$<$<COMPILE_LANGUAGE:CXX>:/utf-8>
|
$<$<COMPILE_LANGUAGE:CXX>:/utf-8>
|
||||||
)
|
)
|
||||||
endif()
|
endif()
|
||||||
@ -312,7 +313,7 @@ set(CMAKE_POLICY_DEFAULT_CMP0077 NEW)
|
|||||||
|
|
||||||
if (NOT SD_USE_SYSTEM_GGML)
|
if (NOT SD_USE_SYSTEM_GGML)
|
||||||
# see https://github.com/ggerganov/ggml/pull/682
|
# see https://github.com/ggerganov/ggml/pull/682
|
||||||
add_definitions(-DGGML_MAX_NAME=128)
|
add_definitions(-DGGML_MAX_NAME=160)
|
||||||
endif()
|
endif()
|
||||||
|
|
||||||
# deps
|
# deps
|
||||||
|
|||||||
15
README.md
@ -15,6 +15,7 @@ API and command-line option may change frequently.***
|
|||||||
|
|
||||||
## 🔥Important News
|
## 🔥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/25** 🚀 stable-diffusion.cpp now supports **Krea2**
|
||||||
* **2026/06/04** 🚀 stable-diffusion.cpp now supports **Ideogram4**
|
* **2026/06/04** 🚀 stable-diffusion.cpp now supports **Ideogram4**
|
||||||
* **2026/05/31** 🚀 stable-diffusion.cpp now supports **PiD**
|
* **2026/05/31** 🚀 stable-diffusion.cpp now supports **PiD**
|
||||||
@ -54,6 +55,7 @@ API and command-line option may change frequently.***
|
|||||||
- [ERNIE-Image](./docs/ernie_image.md)
|
- [ERNIE-Image](./docs/ernie_image.md)
|
||||||
- [Boogu Image](./docs/boogu_image.md)
|
- [Boogu Image](./docs/boogu_image.md)
|
||||||
- [Krea2](./docs/krea2.md)
|
- [Krea2](./docs/krea2.md)
|
||||||
|
- [Mage-Flow](./docs/mage_flow.md)
|
||||||
- [SeFi-Image](./docs/sefi_image.md)
|
- [SeFi-Image](./docs/sefi_image.md)
|
||||||
- [HiDream-O1-Image](./docs/hidream_o1_image.md)
|
- [HiDream-O1-Image](./docs/hidream_o1_image.md)
|
||||||
- [Ideogram4](./docs/ideogram4.md)
|
- [Ideogram4](./docs/ideogram4.md)
|
||||||
@ -62,12 +64,17 @@ API and command-line option may change frequently.***
|
|||||||
- [Qwen Image Edit series](./docs/qwen_image_edit.md)
|
- [Qwen Image Edit series](./docs/qwen_image_edit.md)
|
||||||
- [LongCat Image Edit](./docs/longcat_image.md)
|
- [LongCat Image Edit](./docs/longcat_image.md)
|
||||||
- [Boogu Image Edit](./docs/boogu_image.md)
|
- [Boogu Image Edit](./docs/boogu_image.md)
|
||||||
|
- [Mage-Flow-Edit](./docs/mage_flow.md#image-editing)
|
||||||
- Video Models
|
- Video Models
|
||||||
- [Wan2.1/Wan2.2](./docs/wan.md)
|
- [Wan2.1/Wan2.2](./docs/wan.md)
|
||||||
|
- [MiniMax-H3](./docs/minimax_h3.md)
|
||||||
- [LTX-2.3](./docs/ltx2.md)
|
- [LTX-2.3](./docs/ltx2.md)
|
||||||
|
- [HunyuanVideo 1.5](./docs/hunyuan_video.md)
|
||||||
- [LingBot-Video](./docs/lingbot_video.md)
|
- [LingBot-Video](./docs/lingbot_video.md)
|
||||||
- [PhotoMaker](./docs/photo_maker.md) support.
|
- [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
|
- 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)
|
- LoRA support, same as [stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#lora)
|
||||||
- Latent Consistency Models support (LCM/LCM-LoRA)
|
- Latent Consistency Models support (LCM/LCM-LoRA)
|
||||||
- Faster and memory efficient latent decoding with [TAESD](./docs/taesd.md)
|
- 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
|
- Stable Diffusion v1.5 from https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5
|
||||||
|
|
||||||
```sh
|
```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
|
### Generate an image with just one command
|
||||||
@ -145,6 +152,7 @@ For runtime and parameter backend placement, see the [backend selection guide](.
|
|||||||
- [LCM/LCM-LoRA](./docs/lcm.md)
|
- [LCM/LCM-LoRA](./docs/lcm.md)
|
||||||
- [Docker](./docs/docker.md)
|
- [Docker](./docs/docker.md)
|
||||||
- [Quantization and GGUF](./docs/quantization_and_gguf.md)
|
- [Quantization and GGUF](./docs/quantization_and_gguf.md)
|
||||||
|
- [INT8 convrot safetensors](./docs/int8_convrot.md)
|
||||||
- [Inference acceleration via caching](./docs/caching.md)
|
- [Inference acceleration via caching](./docs/caching.md)
|
||||||
|
|
||||||
## Bindings
|
## Bindings
|
||||||
@ -163,6 +171,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.
|
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)
|
- [Jellybox](https://jellybox.com)
|
||||||
- [Stable Diffusion GUI](https://github.com/fszontagh/sd.cpp.gui.wx)
|
- [Stable Diffusion GUI](https://github.com/fszontagh/sd.cpp.gui.wx)
|
||||||
- [Stable Diffusion CLI-GUI](https://github.com/piallai/stable-diffusion.cpp)
|
- [Stable Diffusion CLI-GUI](https://github.com/piallai/stable-diffusion.cpp)
|
||||||
@ -177,7 +186,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!
|
Thank you to all the people who have already contributed to stable-diffusion.cpp!
|
||||||
|
|
||||||
[](https://github.com/leejet/stable-diffusion.cpp/graphs/contributors)
|
[](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
|
FROM ubuntu:$UBUNTU_VERSION AS runtime
|
||||||
|
|
||||||
RUN apt-get update && \
|
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
|
apt-get clean
|
||||||
|
|
||||||
COPY --from=build /sd.cpp/build/bin /sd.cpp/bin
|
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` |
|
| `controlnet` | ControlNet | `controlnet`, `control` |
|
||||||
| `photomaker` | PhotoMaker ID encoder and PhotoMaker LoRA | `photomaker`, `photomakerid`, `pmid`, `photo` |
|
| `photomaker` | PhotoMaker ID encoder and PhotoMaker LoRA | `photomaker`, `photomakerid`, `pmid`, `photo` |
|
||||||
| `upscaler` | ESRGAN upscaler | `upscaler`, `esrgan`, `hires` |
|
| `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.
|
`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.
|
||||||
|
|
||||||
|
|||||||
@ -16,6 +16,7 @@ Depending on the architecture, different models handle reference images differen
|
|||||||
| [**Flux.2 [Dev] / Flux.2 [Klein]**](./flux2.md) | `flux2` |
|
| [**Flux.2 [Dev] / Flux.2 [Klein]**](./flux2.md) | `flux2` |
|
||||||
| [**Boogu Image Edit**](./boogu_image.md) | `z_image_omni` |
|
| [**Boogu Image Edit**](./boogu_image.md) | `z_image_omni` |
|
||||||
| **Krea2 (Community Edit LoRAs)** | `krea2_ostris_edit` |
|
| **Krea2 (Community Edit LoRAs)** | `krea2_ostris_edit` |
|
||||||
|
| [**Mage-Flow-Edit**](./mage_flow.md#image-editing) | `mage_flow` |
|
||||||
| **Anima (Community Edit LoRAs)** | `cosmos_reference` |
|
| **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.
|
Stable-diffusion.spp also supports basic Unet-based editing models like instruct-pix2pix or CosXL-Edit. This document is not about those.
|
||||||
@ -48,6 +49,7 @@ The `--ref-image-args` argument accepts a comma-separated list of key-value pair
|
|||||||
| `qwen_layered` | Qwen Image Layered |
|
| `qwen_layered` | Qwen Image Layered |
|
||||||
| `z_image_omni` | Boogu, Z-Image Omni |
|
| `z_image_omni` | Boogu, Z-Image Omni |
|
||||||
| `krea2_ostris_edit` | Most Krea2 Community edit LoRAs (trained with Ostris script) |
|
| `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) |
|
| `krea2_edit` | Specifically for [lbouaraba/krea2edit](https://huggingface.co/conradlocke/krea2-identity-edit). (or similar) |
|
||||||
| `cosmos_reference` | For Anima |
|
| `cosmos_reference` | For Anima |
|
||||||
| `default` | Uses the automatic detection based on model architecture. |
|
| `default` | Uses the automatic detection based on model architecture. |
|
||||||
@ -83,6 +85,7 @@ For a technical overview of how each preset is configured, see the table below.
|
|||||||
| `flux2` | No | `increase` | `none` | |
|
| `flux2` | No | `increase` | `none` | |
|
||||||
| `qwen` | Yes | `increase` | `area` | |
|
| `qwen` | Yes | `increase` | `area` | |
|
||||||
| `qwen_layered` | Yes | `decrease` | `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` | |
|
| `z_image_omni` | Yes | `fixed` | `area` | |
|
||||||
| `krea2_ostris_edit`| Yes | `increase` | `area` | `force_ref_timestep_zero = true` |
|
| `krea2_ostris_edit`| Yes | `increase` | `area` | `force_ref_timestep_zero = true` |
|
||||||
| `krea2_edit` | Yes | `increase` | `longest` | `vlm_size = 768` |
|
| `krea2_edit` | Yes | `increase` | `longest` | `vlm_size = 768` |
|
||||||
|
|||||||
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>
|
||||||
102
docs/int8_convrot.md
Normal file
@ -0,0 +1,102 @@
|
|||||||
|
# INT8 Convrot Safetensors
|
||||||
|
|
||||||
|
sd.cpp can load and execute ComfyUI `int8_tensorwise` safetensors with `convrot` metadata directly. The stored INT8 weights are not converted to another weight type at load time.
|
||||||
|
|
||||||
|
## Checkpoint format
|
||||||
|
|
||||||
|
Each quantized linear module contains the following tensors:
|
||||||
|
|
||||||
|
- `<module>.weight`: an I8 weight matrix.
|
||||||
|
- `<module>.weight_scale`: one floating-point scale for each output row. ComfyUI's two-dimensional `[out_features, 1]` representation is normalized to a one-dimensional tensor while loading.
|
||||||
|
- `<module>.comfy_quant`: a U8 tensor containing the JSON quantization configuration.
|
||||||
|
|
||||||
|
A supported configuration has this form:
|
||||||
|
|
||||||
|
```json
|
||||||
|
{
|
||||||
|
"format": "int8_tensorwise",
|
||||||
|
"convrot": true,
|
||||||
|
"convrot_groupsize": 256
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
The convrot group size must be a power of four and must divide the input feature dimension. The commonly used configuration is H256, with `convrot_groupsize` set to `256`.
|
||||||
|
|
||||||
|
## How INT8 convrot works
|
||||||
|
|
||||||
|
Convrot combines an offline rotation of the weights with the same rotation of the activations at runtime. The rotation uses a normalized regular Hadamard matrix constructed recursively from
|
||||||
|
|
||||||
|
```text
|
||||||
|
[ 1 1 1 -1 ]
|
||||||
|
[ 1 1 -1 1 ]
|
||||||
|
H4 = [ 1 -1 1 1 ] / 2
|
||||||
|
[-1 1 1 1 ]
|
||||||
|
```
|
||||||
|
|
||||||
|
For a group size `G = 4^n`, the transform is the normalized Kronecker power of `H4`. It is applied independently to every contiguous group of `G` input features. The resulting block-diagonal rotation matrix `R` is orthogonal and symmetric, so `R R^T = I`.
|
||||||
|
|
||||||
|
For an original floating-point linear layer
|
||||||
|
|
||||||
|
```text
|
||||||
|
Y = X W^T + b
|
||||||
|
```
|
||||||
|
|
||||||
|
the checkpoint stores a rotated weight matrix `W_rot = W R`, quantized per output row. At runtime sd.cpp computes `X_rot = X R`. Ignoring quantization error,
|
||||||
|
|
||||||
|
```text
|
||||||
|
X_rot W_rot^T = X R (W R)^T = X R R^T W^T = X W^T
|
||||||
|
```
|
||||||
|
|
||||||
|
The rotation therefore preserves the linear operation. Its purpose is to spread isolated large values across each feature group, reducing the effect of outliers on tensorwise INT8 quantization.
|
||||||
|
|
||||||
|
### Weight quantization
|
||||||
|
|
||||||
|
The rotated weights are quantized offline with one scale per output row:
|
||||||
|
|
||||||
|
```text
|
||||||
|
s_w[o] = max_i(abs(W_rot[o, i])) / 127
|
||||||
|
Q_w[o, i] = clamp(round(W_rot[o, i] / s_w[o]), -127, 127)
|
||||||
|
```
|
||||||
|
|
||||||
|
`Q_w` is stored in `<module>.weight`, and `s_w` is stored in `<module>.weight_scale`.
|
||||||
|
|
||||||
|
### Runtime activation quantization
|
||||||
|
|
||||||
|
For every activation row, sd.cpp applies the group-wise Hadamard rotation and then calculates one dynamic scale across the entire rotated row:
|
||||||
|
|
||||||
|
```text
|
||||||
|
s_x[r] = max_i(abs(X_rot[r, i])) / 127
|
||||||
|
Q_x[r, i] = clamp(round(X_rot[r, i] / s_x[r]), -127, 127)
|
||||||
|
```
|
||||||
|
|
||||||
|
The matrix multiplication accumulates into signed 32-bit integers:
|
||||||
|
|
||||||
|
```text
|
||||||
|
A[r, o] = sum_i(Q_x[r, i] * Q_w[o, i])
|
||||||
|
```
|
||||||
|
|
||||||
|
The floating-point output is reconstructed as
|
||||||
|
|
||||||
|
```text
|
||||||
|
Y[r, o] ~= A[r, o] * s_x[r] * s_w[o] + b[o]
|
||||||
|
```
|
||||||
|
|
||||||
|
The packed runtime activation tensor contains the I8 activation rows and their floating-point row scales. Linear layers that share the same input and convrot group size reuse this packed tensor, avoiding repeated rotation and activation quantization within the graph.
|
||||||
|
|
||||||
|
## Backend support
|
||||||
|
|
||||||
|
- CPU provides the portable regular Hadamard, activation quantization, INT8 matrix multiplication, and scale restoration implementations.
|
||||||
|
- NVIDIA CUDA devices with compute capability 7.5 or newer use the native accelerated path. For H256, CUDA fuses the rotation, row-wise maximum reduction, and activation quantization. It uses cuBLAS for I8 x I8 to I32 GEMM and a CUDA kernel for scale restoration and bias addition.
|
||||||
|
- Vulkan and other GPU backends do not currently have dedicated INT8 convrot kernels. They use the backend scheduler to fall back to CPU, which is expected to be substantially slower than the CUDA path.
|
||||||
|
|
||||||
|
LoRA adapters are applied at runtime without modifying the INT8 weights. The INT8 convrot path computes the base linear output, while LoRA, LoHa, LoKr, and raw weight-difference adapters compute their output corrections from the original, unrotated activation and add them to the base output. `--lora-apply-mode auto` selects this path for models containing INT8 tensorwise weights. If `immediately` is requested, sd.cpp falls back to runtime application because merging an adapter would require dequantizing and rotating its weight update, then recalculating the per-row scales and requantizing the result.
|
||||||
|
|
||||||
|
The dedicated CUDA convrot activation path currently requires a group size of `256`; other supported group sizes use CPU execution.
|
||||||
|
|
||||||
|
## Example
|
||||||
|
|
||||||
|
ComfyUI INT8 convrot safetensors can be passed to `--diffusion-model` without conversion:
|
||||||
|
|
||||||
|
```powershell
|
||||||
|
.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\krea2_turbo_int8_convrot.safetensors --llm ..\models\text_encoders\Qwen3-VL-4B-Instruct-Q4_K_M.gguf --vae ..\models\vae\wan_2.1_vae.safetensors -p "a lovely cat holding a sign says 'krea2.cpp'" --steps 8 --cfg-scale 1 --diffusion-fa -v --offload-to-cpu
|
||||||
|
```
|
||||||
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
|
# 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
|
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.
|
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`
|
- Flux / Z-Image PiD: use the Flux VAE and pass `--vae-format flux`
|
||||||
- SD3 PiD: use the SD3 VAE and pass `--vae-format sd3`
|
- SD3 PiD: use the SD3 VAE and pass `--vae-format sd3`
|
||||||
- Flux.2 PiD: use the Flux.2 VAE and pass `--vae-format flux2`
|
- 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.
|
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
|
- 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.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 Diffusion v1.5 from https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5
|
||||||
- Stable Diffuison v2.1 from https://huggingface.co/stabilityai/stable-diffusion-2-1
|
- 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
|
- Stable Diffusion 3 2B from https://huggingface.co/stabilityai/stable-diffusion-3-medium
|
||||||
|
|
||||||
### txt2img example
|
### txt2img example
|
||||||
@ -34,4 +34,4 @@ Using formats of different precisions will yield results of varying quality.
|
|||||||
|
|
||||||
<p align="center">
|
<p align="center">
|
||||||
<img src="../assets/img2img_output.png" width="256x">
|
<img src="../assets/img2img_output.png" width="256x">
|
||||||
</p>
|
</p>
|
||||||
|
|||||||
@ -1,4 +1,4 @@
|
|||||||
include_directories(${CMAKE_CURRENT_SOURCE_DIR})
|
include_directories(${CMAKE_CURRENT_SOURCE_DIR})
|
||||||
|
|
||||||
add_subdirectory(cli)
|
add_subdirectory(cli)
|
||||||
add_subdirectory(server)
|
add_subdirectory(server)
|
||||||
|
|||||||
@ -6,6 +6,9 @@ For detailed command-line arguments, run:
|
|||||||
./bin/sd-cli -h
|
./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:
|
Metadata mode inspects PNG/JPEG container metadata without loading any model:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
|
|||||||
@ -199,7 +199,7 @@ struct SDCliParams {
|
|||||||
options.manual_options = {
|
options.manual_options = {
|
||||||
{"-M",
|
{"-M",
|
||||||
"--mode",
|
"--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},
|
on_mode_arg},
|
||||||
{"",
|
{"",
|
||||||
"--preview",
|
"--preview",
|
||||||
@ -566,6 +566,65 @@ bool save_results(const SDCliParams& cli_params,
|
|||||||
return sucessful_reults != 0;
|
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[]) {
|
int main(int argc, const char* argv[]) {
|
||||||
if (argc > 1 && std::string(argv[1]) == "--version") {
|
if (argc > 1 && std::string(argv[1]) == "--version") {
|
||||||
std::cout << version_string() << "\n";
|
std::cout << version_string() << "\n";
|
||||||
@ -598,6 +657,11 @@ int main(int argc, const char* argv[]) {
|
|||||||
return 0;
|
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) {
|
if (gen_params.video_frames > 4) {
|
||||||
size_t last_dot_pos = cli_params.preview_path.find_last_of(".");
|
size_t last_dot_pos = cli_params.preview_path.find_last_of(".");
|
||||||
std::string base_path = cli_params.preview_path;
|
std::string base_path = cli_params.preview_path;
|
||||||
@ -690,6 +754,18 @@ int main(int argc, const char* argv[]) {
|
|||||||
return true;
|
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 (gen_params.init_image_path.size() > 0) {
|
||||||
if (!load_image_and_update_size(gen_params.init_image_path, gen_params.init_image)) {
|
if (!load_image_and_update_size(gen_params.init_image_path, gen_params.init_image)) {
|
||||||
return 1;
|
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 (gen_params.mask_image_path.size() > 0) {
|
||||||
if (!load_sd_image_from_file(gen_params.mask_image.put(),
|
if (!load_sd_image_from_file(gen_params.mask_image.put(),
|
||||||
gen_params.mask_image_path.c_str(),
|
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()) {
|
if (!gen_params.control_video_path.empty()) {
|
||||||
gen_params.control_frames.clear();
|
gen_params.control_frames.clear();
|
||||||
if (!load_images_from_dir(gen_params.control_video_path,
|
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);
|
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{};
|
||||||
sd_img_gen_params_t img_gen_params = gen_params.to_sd_img_gen_params_t();
|
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;
|
sd_image_t* generated_images = nullptr;
|
||||||
if (!generate_image(sd_ctx.get(), &img_gen_params, &generated_images, &num_results)) {
|
if (!generate_image(sd_ctx.get(), &img_gen_params, &generated_images, &num_results)) {
|
||||||
generated_images = nullptr;
|
generated_images = nullptr;
|
||||||
num_results = 0;
|
num_results = 0;
|
||||||
}
|
}
|
||||||
results.adopt(generated_images, num_results);
|
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) {
|
} else if (cli_params.mode == VID_GEN) {
|
||||||
sd_vid_gen_params_t vid_gen_params = gen_params.to_sd_vid_gen_params_t();
|
sd_vid_gen_params_t vid_gen_params = gen_params.to_sd_vid_gen_params_t();
|
||||||
sd_image_t* generated_video = nullptr;
|
sd_image_t* generated_video = nullptr;
|
||||||
@ -828,6 +952,18 @@ int main(int argc, const char* argv[]) {
|
|||||||
LOG_ERROR("generate failed");
|
LOG_ERROR("generate failed");
|
||||||
return 1;
|
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
|
int upscale_factor = 4; // unused for RealESRGAN_x4plus_anime_6B.pth
|
||||||
|
|||||||
@ -30,6 +30,7 @@ namespace fs = std::filesystem;
|
|||||||
|
|
||||||
const char* const modes_str[] = {
|
const char* const modes_str[] = {
|
||||||
"img_gen",
|
"img_gen",
|
||||||
|
"adetailer",
|
||||||
"vid_gen",
|
"vid_gen",
|
||||||
"convert",
|
"convert",
|
||||||
"upscale",
|
"upscale",
|
||||||
@ -49,6 +50,9 @@ static sd_vae_format_t str_to_vae_format(const std::string& value) {
|
|||||||
if (value == "flux2") {
|
if (value == "flux2") {
|
||||||
return SD_VAE_FORMAT_FLUX2;
|
return SD_VAE_FORMAT_FLUX2;
|
||||||
}
|
}
|
||||||
|
if (value == "wan") {
|
||||||
|
return SD_VAE_FORMAT_WAN;
|
||||||
|
}
|
||||||
return SD_VAE_FORMAT_COUNT;
|
return SD_VAE_FORMAT_COUNT;
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -400,7 +404,7 @@ ArgOptions SDContextParams::get_options() {
|
|||||||
&vae_path},
|
&vae_path},
|
||||||
{"",
|
{"",
|
||||||
"--vae-format",
|
"--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,
|
0,
|
||||||
&vae_format},
|
&vae_format},
|
||||||
{"",
|
{"",
|
||||||
@ -423,6 +427,16 @@ ArgOptions SDContextParams::get_options() {
|
|||||||
"path to control net model",
|
"path to control net model",
|
||||||
0,
|
0,
|
||||||
&control_net_path},
|
&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",
|
"--embd-dir",
|
||||||
"embeddings directory",
|
"embeddings directory",
|
||||||
@ -672,7 +686,7 @@ ArgOptions SDContextParams::get_options() {
|
|||||||
}
|
}
|
||||||
|
|
||||||
void SDContextParams::build_embedding_map() {
|
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)) {
|
if (!fs::exists(embedding_dir) || !fs::is_directory(embedding_dir)) {
|
||||||
return;
|
return;
|
||||||
@ -737,7 +751,7 @@ bool SDContextParams::validate(SDMode mode) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
if (str_to_vae_format(vae_format) == SD_VAE_FORMAT_COUNT) {
|
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;
|
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.audio_vae_path = audio_vae_path.c_str();
|
||||||
sd_ctx_params.taesd_path = taesd_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.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.embeddings = embedding_vec.data();
|
||||||
sd_ctx_params.embedding_count = static_cast<uint32_t>(embedding_vec.size());
|
sd_ctx_params.embedding_count = static_cast<uint32_t>(embedding_vec.size());
|
||||||
sd_ctx_params.photo_maker_path = photo_maker_path.c_str();
|
sd_ctx_params.photo_maker_path = photo_maker_path.c_str();
|
||||||
@ -916,6 +932,26 @@ ArgOptions SDGenerationParams::get_options() {
|
|||||||
"the negative prompt (default: \"\")",
|
"the negative prompt (default: \"\")",
|
||||||
0,
|
0,
|
||||||
&negative_prompt},
|
&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",
|
{"-i",
|
||||||
"--init-img",
|
"--init-img",
|
||||||
"path to the init image",
|
"path to the init image",
|
||||||
@ -936,6 +972,11 @@ ArgOptions SDGenerationParams::get_options() {
|
|||||||
"path to control image, control net",
|
"path to control image, control net",
|
||||||
0,
|
0,
|
||||||
&control_image_path},
|
&control_image_path},
|
||||||
|
{"",
|
||||||
|
"--ip-adapter-image",
|
||||||
|
"path to the IP-Adapter reference image",
|
||||||
|
0,
|
||||||
|
&ip_adapter_image_path},
|
||||||
{"",
|
{"",
|
||||||
"--control-video",
|
"--control-video",
|
||||||
"path to control video frames, It must be a directory path. The video frames inside should be stored as images in "
|
"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},
|
&hires_upscaler},
|
||||||
{"",
|
{"",
|
||||||
"--extra-sample-args",
|
"--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)',',
|
(int)',',
|
||||||
&extra_sample_args},
|
&extra_sample_args},
|
||||||
{"",
|
{"",
|
||||||
@ -1128,6 +1169,10 @@ ArgOptions SDGenerationParams::get_options() {
|
|||||||
"--control-strength",
|
"--control-strength",
|
||||||
"strength to apply Control Net (default: 0.9). 1.0 corresponds to full destruction of information in init image",
|
"strength to apply Control Net (default: 0.9). 1.0 corresponds to full destruction of information in init image",
|
||||||
&control_strength},
|
&control_strength},
|
||||||
|
{"",
|
||||||
|
"--ip-adapter-strength",
|
||||||
|
"strength to apply IP-Adapter (default: 1.0)",
|
||||||
|
&ip_adapter_strength},
|
||||||
{"",
|
{"",
|
||||||
"--moe-boundary",
|
"--moe-boundary",
|
||||||
"timestep boundary for Wan2.2 MoE model. (default: 0.875). Only enabled if `--high-noise-steps` is set to -1",
|
"timestep boundary for Wan2.2 MoE model. (default: 0.875). Only enabled if `--high-noise-steps` is set to -1",
|
||||||
@ -1359,6 +1404,30 @@ ArgOptions SDGenerationParams::get_options() {
|
|||||||
return 1;
|
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) {
|
auto on_cache_mode_arg = [&](int argc, const char** argv, int index) {
|
||||||
if (++index >= argc) {
|
if (++index >= argc) {
|
||||||
return -1;
|
return -1;
|
||||||
@ -1493,12 +1562,12 @@ ArgOptions SDGenerationParams::get_options() {
|
|||||||
on_seed_arg},
|
on_seed_arg},
|
||||||
{"",
|
{"",
|
||||||
"--sampling-method",
|
"--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)",
|
"(default: euler for Flux/SD3/Wan, euler_a otherwise)",
|
||||||
on_sample_method_arg},
|
on_sample_method_arg},
|
||||||
{"",
|
{"",
|
||||||
"--high-noise-sampling-method",
|
"--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",
|
" default: euler for Flux/SD3/Wan, euler_a otherwise",
|
||||||
on_high_noise_sample_method_arg},
|
on_high_noise_sample_method_arg},
|
||||||
{"",
|
{"",
|
||||||
@ -1523,8 +1592,20 @@ ArgOptions SDGenerationParams::get_options() {
|
|||||||
on_high_noise_skip_layers_arg},
|
on_high_noise_skip_layers_arg},
|
||||||
{"-r",
|
{"-r",
|
||||||
"--ref-image",
|
"--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},
|
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",
|
"--cache-mode",
|
||||||
"caching method: 'easycache' (DiT), 'ucache' (UNET), 'dbcache'/'taylorseer'/'cache-dit' (DiT block-level), 'spectrum' (UNET/DiT Chebyshev+Taylor forecasting)",
|
"caching method: 'easycache' (DiT), 'ucache' (UNET), 'dbcache'/'taylorseer'/'cache-dit' (DiT block-level), 'spectrum' (UNET/DiT Chebyshev+Taylor forecasting)",
|
||||||
@ -1836,6 +1917,10 @@ bool SDGenerationParams::from_json_str(
|
|||||||
|
|
||||||
load_if_exists("prompt", prompt);
|
load_if_exists("prompt", prompt);
|
||||||
load_if_exists("negative_prompt", negative_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_mode", cache_mode);
|
||||||
load_if_exists("cache_option", cache_option);
|
load_if_exists("cache_option", cache_option);
|
||||||
load_if_exists("scm_mask", scm_mask);
|
load_if_exists("scm_mask", scm_mask);
|
||||||
@ -1852,6 +1937,7 @@ bool SDGenerationParams::from_json_str(
|
|||||||
|
|
||||||
load_if_exists("strength", strength);
|
load_if_exists("strength", strength);
|
||||||
load_if_exists("control_strength", control_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("moe_boundary", moe_boundary);
|
||||||
load_if_exists("vace_strength", vace_strength);
|
load_if_exists("vace_strength", vace_strength);
|
||||||
|
|
||||||
@ -2023,6 +2109,10 @@ bool SDGenerationParams::from_json_str(
|
|||||||
LOG_ERROR("invalid control_image");
|
LOG_ERROR("invalid control_image");
|
||||||
return false;
|
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;
|
return true;
|
||||||
}
|
}
|
||||||
@ -2032,7 +2122,7 @@ void SDGenerationParams::extract_and_remove_lora(const std::string& lora_model_d
|
|||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
static const std::regex re(R"(<lora:([^:>]+):([^>]+)>)");
|
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::smatch m;
|
||||||
|
|
||||||
std::string tmp = prompt;
|
std::string tmp = prompt;
|
||||||
@ -2312,6 +2402,16 @@ bool SDGenerationParams::validate(SDMode mode) {
|
|||||||
return false;
|
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) {
|
if (sample_params.shifted_timestep < 0 || sample_params.shifted_timestep > 1000) {
|
||||||
LOG_ERROR("error: shifted_timestep must be in range [0, 1000]");
|
LOG_ERROR("error: shifted_timestep must be in range [0, 1000]");
|
||||||
return false;
|
return false;
|
||||||
@ -2352,13 +2452,19 @@ bool SDGenerationParams::validate(SDMode mode) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
if (mode == UPSCALE) {
|
if (mode == UPSCALE || mode == ADETAILER) {
|
||||||
if (init_image_path.length() == 0) {
|
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;
|
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;
|
return true;
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -2439,29 +2545,31 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
|
|||||||
LOG_WARN("Notice: --increase-ref-index is deprecated. Use --ref-image-args \"ref_index_mode=increase\" instead.");
|
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.loras = lora_vec.empty() ? nullptr : lora_vec.data();
|
||||||
params.lora_count = static_cast<uint32_t>(lora_vec.size());
|
params.lora_count = static_cast<uint32_t>(lora_vec.size());
|
||||||
params.prompt = prompt.c_str();
|
params.prompt = prompt.c_str();
|
||||||
params.negative_prompt = negative_prompt.c_str();
|
params.negative_prompt = negative_prompt.c_str();
|
||||||
params.clip_skip = clip_skip;
|
params.clip_skip = clip_skip;
|
||||||
params.init_image = init_image.get();
|
params.init_image = init_image.get();
|
||||||
params.ref_images = ref_image_views.empty() ? nullptr : ref_image_views.data();
|
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_images_count = static_cast<int>(ref_image_views.size());
|
||||||
params.ref_image_args = ref_image_args.c_str();
|
params.ref_image_args = ref_image_args.c_str();
|
||||||
params.mask_image = mask_image.get();
|
params.mask_image = mask_image.get();
|
||||||
params.width = get_resolved_width();
|
params.width = get_resolved_width();
|
||||||
params.height = get_resolved_height();
|
params.height = get_resolved_height();
|
||||||
params.sample_params = sample_params;
|
params.sample_params = sample_params;
|
||||||
params.strength = strength;
|
params.strength = strength;
|
||||||
params.seed = seed;
|
params.seed = seed;
|
||||||
params.batch_count = batch_count;
|
params.batch_count = batch_count;
|
||||||
params.qwen_image_layers = qwen_image_layers;
|
params.qwen_image_layers = qwen_image_layers;
|
||||||
params.control_image = control_image.get();
|
params.control_image = control_image.get();
|
||||||
params.control_strength = control_strength;
|
params.control_strength = control_strength;
|
||||||
params.pm_params = pm_params;
|
params.ip_adapter_image = ip_adapter_image.get();
|
||||||
params.pulid_params = pulid_params;
|
params.ip_adapter_strength = ip_adapter_strength;
|
||||||
params.vae_tiling_params = vae_tiling_params;
|
params.pm_params = pm_params;
|
||||||
params.cache = cache_params;
|
params.pulid_params = pulid_params;
|
||||||
|
params.vae_tiling_params = vae_tiling_params;
|
||||||
|
params.cache = cache_params;
|
||||||
|
|
||||||
params.hires.enabled = hires_enabled;
|
params.hires.enabled = hires_enabled;
|
||||||
params.hires.upscaler = resolved_hires_upscaler;
|
params.hires.upscaler = resolved_hires_upscaler;
|
||||||
@ -2498,6 +2606,35 @@ sd_vid_gen_params_t SDGenerationParams::to_sd_vid_gen_params_t() {
|
|||||||
control_frame_views.push_back(frame.get());
|
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.layers = skip_layers.empty() ? nullptr : skip_layers.data();
|
||||||
sample_params.guidance.slg.layer_count = skip_layers.size();
|
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();
|
high_noise_sample_params.guidance.slg.layers = high_noise_skip_layers.empty() ? nullptr : high_noise_skip_layers.data();
|
||||||
@ -2516,6 +2653,12 @@ sd_vid_gen_params_t SDGenerationParams::to_sd_vid_gen_params_t() {
|
|||||||
params.clip_skip = clip_skip;
|
params.clip_skip = clip_skip;
|
||||||
params.init_image = init_image.get();
|
params.init_image = init_image.get();
|
||||||
params.end_image = end_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 = control_frame_views.empty() ? nullptr : control_frame_views.data();
|
||||||
params.control_frames_size = static_cast<int>(control_frame_views.size());
|
params.control_frames_size = static_cast<int>(control_frame_views.size());
|
||||||
params.width = get_resolved_width();
|
params.width = get_resolved_width();
|
||||||
@ -2581,6 +2724,10 @@ std::string SDGenerationParams::to_string() const {
|
|||||||
<< " high_noise_loras: \"" << high_noise_loras_str << "\",\n"
|
<< " high_noise_loras: \"" << high_noise_loras_str << "\",\n"
|
||||||
<< " prompt: \"" << prompt << "\",\n"
|
<< " prompt: \"" << prompt << "\",\n"
|
||||||
<< " negative_prompt: \"" << negative_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"
|
<< " clip_skip: " << clip_skip << ",\n"
|
||||||
<< " width: " << width << ",\n"
|
<< " width: " << width << ",\n"
|
||||||
<< " height: " << height << ",\n"
|
<< " height: " << height << ",\n"
|
||||||
@ -2591,6 +2738,9 @@ std::string SDGenerationParams::to_string() const {
|
|||||||
<< " mask_image_path: \"" << mask_image_path << "\",\n"
|
<< " mask_image_path: \"" << mask_image_path << "\",\n"
|
||||||
<< " control_image_path: \"" << control_image_path << "\",\n"
|
<< " control_image_path: \"" << control_image_path << "\",\n"
|
||||||
<< " ref_image_paths: " << vec_str_to_string(ref_image_paths) << ",\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"
|
<< " control_video_path: \"" << control_video_path << "\",\n"
|
||||||
<< " auto_resize_ref_image: " << (auto_resize_ref_image ? "true" : "false") << ",\n"
|
<< " auto_resize_ref_image: " << (auto_resize_ref_image ? "true" : "false") << ",\n"
|
||||||
<< " increase_ref_index: " << (increase_ref_index ? "true" : "false") << ",\n"
|
<< " increase_ref_index: " << (increase_ref_index ? "true" : "false") << ",\n"
|
||||||
@ -2695,8 +2845,13 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params,
|
|||||||
int64_t seed,
|
int64_t seed,
|
||||||
SDMode mode) {
|
SDMode mode) {
|
||||||
json root;
|
json root;
|
||||||
root["schema"] = "sdcpp.image.params/v1";
|
root["schema"] = "sdcpp.image.params/v1";
|
||||||
root["mode"] = mode == VID_GEN ? "vid_gen" : "img_gen";
|
root["mode"] = "img_gen";
|
||||||
|
if (mode == VID_GEN) {
|
||||||
|
root["mode"] = "vid_gen";
|
||||||
|
} else if (mode == ADETAILER) {
|
||||||
|
root["mode"] = "adetailer";
|
||||||
|
}
|
||||||
root["generator"] = {
|
root["generator"] = {
|
||||||
{"name", "stable-diffusion.cpp"},
|
{"name", "stable-diffusion.cpp"},
|
||||||
{"version", safe_json_string(sd_version())},
|
{"version", safe_json_string(sd_version())},
|
||||||
@ -2710,6 +2865,14 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params,
|
|||||||
{"positive", gen_params.prompt},
|
{"positive", gen_params.prompt},
|
||||||
{"negative", gen_params.negative_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,
|
root["sampling"] = build_sampling_metadata_json(gen_params.sample_params,
|
||||||
gen_params.skip_layers,
|
gen_params.skip_layers,
|
||||||
&gen_params.custom_sigmas);
|
&gen_params.custom_sigmas);
|
||||||
@ -2733,6 +2896,7 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params,
|
|||||||
root["clip_skip"] = gen_params.clip_skip;
|
root["clip_skip"] = gen_params.clip_skip;
|
||||||
root["strength"] = gen_params.strength;
|
root["strength"] = gen_params.strength;
|
||||||
root["control_strength"] = gen_params.control_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["auto_resize_ref_image"] = gen_params.auto_resize_ref_image;
|
||||||
root["increase_ref_index"] = gen_params.increase_ref_index;
|
root["increase_ref_index"] = gen_params.increase_ref_index;
|
||||||
if (mode == VID_GEN) {
|
if (mode == VID_GEN) {
|
||||||
@ -2864,6 +3028,18 @@ std::string get_image_params(const SDContextParams& ctx_params,
|
|||||||
if (!gen_params.extra_sample_args.empty()) {
|
if (!gen_params.extra_sample_args.empty()) {
|
||||||
parameter_string += "Extra sample args: " + gen_params.extra_sample_args + ", ";
|
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 += "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 += "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) + ", ";
|
parameter_string += "Model: " + sd_basename(ctx_params.model_path) + ", ";
|
||||||
|
|||||||
@ -16,10 +16,11 @@
|
|||||||
#define BOOL_STR(b) ((b) ? "true" : "false")
|
#define BOOL_STR(b) ((b) ? "true" : "false")
|
||||||
|
|
||||||
extern const char* const modes_str[];
|
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 {
|
enum SDMode {
|
||||||
IMG_GEN,
|
IMG_GEN,
|
||||||
|
ADETAILER,
|
||||||
VID_GEN,
|
VID_GEN,
|
||||||
CONVERT,
|
CONVERT,
|
||||||
UPSCALE,
|
UPSCALE,
|
||||||
@ -132,6 +133,8 @@ struct SDContextParams {
|
|||||||
std::string taesd_path;
|
std::string taesd_path;
|
||||||
std::string esrgan_path;
|
std::string esrgan_path;
|
||||||
std::string control_net_path;
|
std::string control_net_path;
|
||||||
|
std::string ip_adapter_path;
|
||||||
|
std::string motion_module_path;
|
||||||
std::string embedding_dir;
|
std::string embedding_dir;
|
||||||
std::string photo_maker_path;
|
std::string photo_maker_path;
|
||||||
std::string pulid_weights_path;
|
std::string pulid_weights_path;
|
||||||
@ -186,6 +189,10 @@ struct SDGenerationParams {
|
|||||||
// User-facing input fields.
|
// User-facing input fields.
|
||||||
std::string prompt;
|
std::string prompt;
|
||||||
std::string negative_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 clip_skip = -1; // <= 0 represents unspecified
|
||||||
int width = -1;
|
int width = -1;
|
||||||
int height = -1;
|
int height = -1;
|
||||||
@ -194,6 +201,7 @@ struct SDGenerationParams {
|
|||||||
int64_t seed = 42;
|
int64_t seed = 42;
|
||||||
float strength = 0.75f;
|
float strength = 0.75f;
|
||||||
float control_strength = 0.9f;
|
float control_strength = 0.9f;
|
||||||
|
float ip_adapter_strength = 1.0f;
|
||||||
bool auto_resize_ref_image = true;
|
bool auto_resize_ref_image = true;
|
||||||
bool increase_ref_index = false;
|
bool increase_ref_index = false;
|
||||||
bool embed_image_metadata = true;
|
bool embed_image_metadata = true;
|
||||||
@ -202,7 +210,11 @@ struct SDGenerationParams {
|
|||||||
std::string end_image_path;
|
std::string end_image_path;
|
||||||
std::string mask_image_path;
|
std::string mask_image_path;
|
||||||
std::string control_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_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;
|
std::string control_video_path;
|
||||||
|
|
||||||
sd_sample_params_t sample_params;
|
sd_sample_params_t sample_params;
|
||||||
@ -266,13 +278,20 @@ struct SDGenerationParams {
|
|||||||
SDImageOwner init_image;
|
SDImageOwner init_image;
|
||||||
SDImageOwner end_image;
|
SDImageOwner end_image;
|
||||||
std::vector<SDImageOwner> ref_images;
|
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 mask_image;
|
||||||
SDImageOwner control_image;
|
SDImageOwner control_image;
|
||||||
|
SDImageOwner ip_adapter_image;
|
||||||
std::vector<SDImageOwner> pm_id_images;
|
std::vector<SDImageOwner> pm_id_images;
|
||||||
std::vector<SDImageOwner> control_frames;
|
std::vector<SDImageOwner> control_frames;
|
||||||
|
|
||||||
// Backing storage for sd_img_gen_params_t view fields.
|
// Backing storage for sd_img_gen_params_t view fields.
|
||||||
std::vector<sd_image_t> ref_image_views;
|
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> pm_id_image_views;
|
||||||
std::vector<sd_image_t> control_frame_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)));
|
file.write(reinterpret_cast<const char*>(pcm.data()), static_cast<std::streamsize>(pcm.size() * sizeof(int16_t)));
|
||||||
return file.good();
|
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 channels,
|
||||||
uint32_t sample_rate);
|
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__
|
#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>
|
template <typename T>
|
||||||
using FreeUniquePtr = std::unique_ptr<T, FreeDeleter>;
|
using FreeUniquePtr = std::unique_ptr<T, FreeDeleter>;
|
||||||
|
|
||||||
using FilePtr = std::unique_ptr<FILE, FileCloser>;
|
using FilePtr = std::unique_ptr<FILE, FileCloser>;
|
||||||
using SDCtxPtr = std::unique_ptr<sd_ctx_t, SDCtxDeleter>;
|
using SDCtxPtr = std::unique_ptr<sd_ctx_t, SDCtxDeleter>;
|
||||||
using UpscalerCtxPtr = std::unique_ptr<upscaler_ctx_t, UpscalerCtxDeleter>;
|
using UpscalerCtxPtr = std::unique_ptr<upscaler_ctx_t, UpscalerCtxDeleter>;
|
||||||
|
using ADetailerCtxPtr = std::unique_ptr<adetailer_ctx_t, ADetailerCtxDeleter>;
|
||||||
|
|
||||||
class SDImageOwner {
|
class SDImageOwner {
|
||||||
private:
|
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 {
|
class SDImageVec {
|
||||||
private:
|
private:
|
||||||
std::vector<sd_image_t> images_;
|
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` |
|
| `auto_resize_ref_image` | `boolean` |
|
||||||
| `increase_ref_index` | `boolean` |
|
| `increase_ref_index` | `boolean` |
|
||||||
| `control_strength` | `number` |
|
| `control_strength` | `number` |
|
||||||
|
| `ip_adapter_strength` | `number` |
|
||||||
| `hires` | `object` |
|
| `hires` | `object` |
|
||||||
| `hires.enabled` | `boolean` |
|
| `hires.enabled` | `boolean` |
|
||||||
| `hires.upscaler` | `string` |
|
| `hires.upscaler` | `string` |
|
||||||
@ -567,6 +568,7 @@ Fields returned in `features_by_mode.img_gen`:
|
|||||||
- `init_image`
|
- `init_image`
|
||||||
- `mask_image`
|
- `mask_image`
|
||||||
- `control_image`
|
- `control_image`
|
||||||
|
- `ip_adapter_image`
|
||||||
- `ref_images`
|
- `ref_images`
|
||||||
- `lora`
|
- `lora`
|
||||||
- `vae_tiling`
|
- `vae_tiling`
|
||||||
@ -653,12 +655,14 @@ Example:
|
|||||||
"auto_resize_ref_image": true,
|
"auto_resize_ref_image": true,
|
||||||
"increase_ref_index": false,
|
"increase_ref_index": false,
|
||||||
"control_strength": 0.9,
|
"control_strength": 0.9,
|
||||||
|
"ip_adapter_strength": 1.0,
|
||||||
"embed_image_metadata": true,
|
"embed_image_metadata": true,
|
||||||
|
|
||||||
"init_image": null,
|
"init_image": null,
|
||||||
"ref_images": [],
|
"ref_images": [],
|
||||||
"mask_image": null,
|
"mask_image": null,
|
||||||
"control_image": null,
|
"control_image": null,
|
||||||
|
"ip_adapter_image": null,
|
||||||
|
|
||||||
"sample_params": {
|
"sample_params": {
|
||||||
"scheduler": "discrete",
|
"scheduler": "discrete",
|
||||||
@ -733,6 +737,7 @@ Channel expectations:
|
|||||||
- `init_image`: 3 channels
|
- `init_image`: 3 channels
|
||||||
- `ref_images[]`: 3 channels
|
- `ref_images[]`: 3 channels
|
||||||
- `control_image`: 3 channels
|
- `control_image`: 3 channels
|
||||||
|
- `ip_adapter_image`: 3 channels
|
||||||
- `mask_image`: 1 channel
|
- `mask_image`: 1 channel
|
||||||
|
|
||||||
If omitted or null:
|
If omitted or null:
|
||||||
@ -757,6 +762,7 @@ Top-level scalar fields:
|
|||||||
| `auto_resize_ref_image` | `boolean` |
|
| `auto_resize_ref_image` | `boolean` |
|
||||||
| `increase_ref_index` | `boolean` |
|
| `increase_ref_index` | `boolean` |
|
||||||
| `control_strength` | `number` |
|
| `control_strength` | `number` |
|
||||||
|
| `ip_adapter_strength` | `number` |
|
||||||
| `embed_image_metadata` | `boolean` |
|
| `embed_image_metadata` | `boolean` |
|
||||||
|
|
||||||
Image fields:
|
Image fields:
|
||||||
@ -767,6 +773,7 @@ Image fields:
|
|||||||
| `ref_images` | `array<string>` |
|
| `ref_images` | `array<string>` |
|
||||||
| `mask_image` | `string \| null` |
|
| `mask_image` | `string \| null` |
|
||||||
| `control_image` | `string \| null` |
|
| `control_image` | `string \| null` |
|
||||||
|
| `ip_adapter_image` | `string \| null` |
|
||||||
|
|
||||||
LoRA fields:
|
LoRA fields:
|
||||||
|
|
||||||
@ -958,7 +965,7 @@ Response fields:
|
|||||||
Compared with `img_gen`, the `vid_gen` request body:
|
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
|
- `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`
|
- `vid_gen` adds `end_image`, `control_frames`, `high_noise_sample_params`, `video_frames`, `fps`, `moe_boundary`, and `vace_strength`
|
||||||
|
|
||||||
Example:
|
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},
|
{"auto_resize_ref_image", defaults.auto_resize_ref_image},
|
||||||
{"increase_ref_index", defaults.increase_ref_index},
|
{"increase_ref_index", defaults.increase_ref_index},
|
||||||
{"control_strength", defaults.control_strength},
|
{"control_strength", defaults.control_strength},
|
||||||
|
{"ip_adapter_strength", defaults.ip_adapter_strength},
|
||||||
{"sample_params", make_sample_params_json(defaults.sample_params, defaults.skip_layers)},
|
{"sample_params", make_sample_params_json(defaults.sample_params, defaults.skip_layers)},
|
||||||
{"hires", make_hires_json(defaults)},
|
{"hires", make_hires_json(defaults)},
|
||||||
{"vae_tiling_params", make_vae_tiling_json(defaults.vae_tiling_params)},
|
{"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},
|
{"init_image", true},
|
||||||
{"mask_image", true},
|
{"mask_image", true},
|
||||||
{"control_image", true},
|
{"control_image", true},
|
||||||
|
{"ip_adapter_image", true},
|
||||||
{"ref_images", true},
|
{"ref_images", true},
|
||||||
{"lora", true},
|
{"lora", true},
|
||||||
{"vae_tiling", true},
|
{"vae_tiling", true},
|
||||||
|
|||||||
2
ggml
@ -1 +1 @@
|
|||||||
Subproject commit eced84c86f8b012c752c016f7fe789adea168e1e
|
Subproject commit 3f85508cddcba4e4a8a22414af379816f7945222
|
||||||
@ -56,6 +56,7 @@ enum sample_method_t {
|
|||||||
EULER_GE_SAMPLE_METHOD,
|
EULER_GE_SAMPLE_METHOD,
|
||||||
DPMPP2M_SDE_SAMPLE_METHOD,
|
DPMPP2M_SDE_SAMPLE_METHOD,
|
||||||
DPMPP2M_SDE_BT_SAMPLE_METHOD,
|
DPMPP2M_SDE_BT_SAMPLE_METHOD,
|
||||||
|
LMS_SAMPLE_METHOD,
|
||||||
SAMPLE_METHOD_COUNT
|
SAMPLE_METHOD_COUNT
|
||||||
};
|
};
|
||||||
|
|
||||||
@ -180,6 +181,7 @@ enum sd_vae_format_t {
|
|||||||
SD_VAE_FORMAT_FLUX,
|
SD_VAE_FORMAT_FLUX,
|
||||||
SD_VAE_FORMAT_SD3,
|
SD_VAE_FORMAT_SD3,
|
||||||
SD_VAE_FORMAT_FLUX2,
|
SD_VAE_FORMAT_FLUX2,
|
||||||
|
SD_VAE_FORMAT_WAN,
|
||||||
SD_VAE_FORMAT_COUNT,
|
SD_VAE_FORMAT_COUNT,
|
||||||
};
|
};
|
||||||
|
|
||||||
@ -199,6 +201,8 @@ typedef struct {
|
|||||||
const char* audio_vae_path;
|
const char* audio_vae_path;
|
||||||
const char* taesd_path;
|
const char* taesd_path;
|
||||||
const char* control_net_path;
|
const char* control_net_path;
|
||||||
|
const char* ip_adapter_path;
|
||||||
|
const char* motion_module_path;
|
||||||
const sd_embedding_t* embeddings;
|
const sd_embedding_t* embeddings;
|
||||||
uint32_t embedding_count;
|
uint32_t embedding_count;
|
||||||
const char* photo_maker_path;
|
const char* photo_maker_path;
|
||||||
@ -243,6 +247,13 @@ typedef struct {
|
|||||||
uint8_t* data;
|
uint8_t* data;
|
||||||
} sd_image_t;
|
} sd_image_t;
|
||||||
|
|
||||||
|
typedef struct {
|
||||||
|
sd_image_t* frames;
|
||||||
|
int frame_count;
|
||||||
|
int fps;
|
||||||
|
sd_audio_t audio;
|
||||||
|
} sd_ref_video_t;
|
||||||
|
|
||||||
typedef struct {
|
typedef struct {
|
||||||
int* layers;
|
int* layers;
|
||||||
size_t layer_count;
|
size_t layer_count;
|
||||||
@ -373,6 +384,8 @@ typedef struct {
|
|||||||
int batch_count;
|
int batch_count;
|
||||||
sd_image_t control_image;
|
sd_image_t control_image;
|
||||||
float control_strength;
|
float control_strength;
|
||||||
|
sd_image_t ip_adapter_image;
|
||||||
|
float ip_adapter_strength;
|
||||||
sd_pm_params_t pm_params;
|
sd_pm_params_t pm_params;
|
||||||
sd_pulid_params_t pulid_params;
|
sd_pulid_params_t pulid_params;
|
||||||
sd_tiling_params_t vae_tiling_params;
|
sd_tiling_params_t vae_tiling_params;
|
||||||
@ -391,6 +404,12 @@ typedef struct {
|
|||||||
int clip_skip;
|
int clip_skip;
|
||||||
sd_image_t init_image;
|
sd_image_t init_image;
|
||||||
sd_image_t end_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;
|
sd_image_t* control_frames;
|
||||||
int control_frames_size;
|
int control_frames_size;
|
||||||
int width;
|
int width;
|
||||||
@ -508,6 +527,27 @@ SD_API bool upscale(upscaler_ctx_t* upscaler_ctx,
|
|||||||
|
|
||||||
SD_API int get_upscale_factor(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,
|
SD_API bool convert(const char* input_path,
|
||||||
const char* vae_path,
|
const char* vae_path,
|
||||||
const char* output_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,8 +2,10 @@
|
|||||||
#define __SD_CONDITIONING_CONDITIONER_HPP__
|
#define __SD_CONDITIONING_CONDITIONER_HPP__
|
||||||
|
|
||||||
#include <cmath>
|
#include <cmath>
|
||||||
|
#include <iomanip>
|
||||||
#include <limits>
|
#include <limits>
|
||||||
#include <optional>
|
#include <optional>
|
||||||
|
#include <sstream>
|
||||||
|
|
||||||
#include "core/tensor_ggml.hpp"
|
#include "core/tensor_ggml.hpp"
|
||||||
#include "core/util.h"
|
#include "core/util.h"
|
||||||
@ -25,6 +27,8 @@ struct SDCondition {
|
|||||||
sd::Tensor<int32_t> c_vinput_mask;
|
sd::Tensor<int32_t> c_vinput_mask;
|
||||||
std::vector<std::pair<int, sd::Tensor<float>>> c_image_embeds;
|
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_images;
|
||||||
|
std::vector<sd::Tensor<float>> c_ref_audios;
|
||||||
|
std::vector<MiniMaxH3ReferenceBlock> c_reference_blocks;
|
||||||
|
|
||||||
std::vector<sd::Tensor<float>> extra_c_crossattns;
|
std::vector<sd::Tensor<float>> extra_c_crossattns;
|
||||||
|
|
||||||
@ -55,6 +59,12 @@ struct SDCondition {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
for (const auto& tensor : c_ref_audios) {
|
||||||
|
if (!tensor.empty()) {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
for (const auto& tensor : extra_c_crossattns) {
|
for (const auto& tensor : extra_c_crossattns) {
|
||||||
if (!tensor.empty()) {
|
if (!tensor.empty()) {
|
||||||
return false;
|
return false;
|
||||||
@ -65,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,
|
static inline sd::Tensor<float> apply_token_weights(sd::Tensor<float> hidden_states,
|
||||||
const std::vector<float>& weights) {
|
const std::vector<float>& weights) {
|
||||||
if (hidden_states.empty()) {
|
if (hidden_states.empty()) {
|
||||||
@ -102,11 +124,12 @@ static inline sd::Tensor<float> apply_token_weights(sd::Tensor<float> hidden_sta
|
|||||||
|
|
||||||
struct ConditionerParams {
|
struct ConditionerParams {
|
||||||
std::string text;
|
std::string text;
|
||||||
int clip_skip = -1;
|
int clip_skip = -1;
|
||||||
int width = -1;
|
int width = -1;
|
||||||
int height = -1;
|
int height = -1;
|
||||||
bool zero_out_masked = false;
|
bool zero_out_masked = false;
|
||||||
const std::vector<sd::Tensor<float>>* ref_images = nullptr; // for qwen image edit
|
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;
|
RefImageParams ref_image_params;
|
||||||
};
|
};
|
||||||
|
|
||||||
@ -117,6 +140,7 @@ public:
|
|||||||
virtual SDCondition get_learned_condition(int n_threads,
|
virtual SDCondition get_learned_condition(int n_threads,
|
||||||
const ConditionerParams& conditioner_params) = 0;
|
const ConditionerParams& conditioner_params) = 0;
|
||||||
virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) = 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_max_graph_vram_bytes(size_t max_vram_bytes) {}
|
||||||
virtual void set_stream_layers_enabled(bool enabled) {}
|
virtual void set_stream_layers_enabled(bool enabled) {}
|
||||||
virtual void set_runtime_backends(const std::vector<ggml_backend_t>& backends) {}
|
virtual void set_runtime_backends(const std::vector<ggml_backend_t>& backends) {}
|
||||||
@ -1664,6 +1688,10 @@ struct AnimaConditioner : public Conditioner {
|
|||||||
llm->get_param_tensors(tensors, "text_encoders.llm");
|
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 {
|
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
|
||||||
llm->set_max_graph_vram_bytes(max_vram_bytes);
|
llm->set_max_graph_vram_bytes(max_vram_bytes);
|
||||||
}
|
}
|
||||||
@ -1782,6 +1810,7 @@ struct LLMEmbedder : public Conditioner {
|
|||||||
SDVersion version;
|
SDVersion version;
|
||||||
std::shared_ptr<BPETokenizer> tokenizer;
|
std::shared_ptr<BPETokenizer> tokenizer;
|
||||||
std::shared_ptr<LLM::LLMRunner> llm;
|
std::shared_ptr<LLM::LLMRunner> llm;
|
||||||
|
std::shared_ptr<T5Runner> byt5;
|
||||||
|
|
||||||
LLMEmbedder(ggml_backend_t backend,
|
LLMEmbedder(ggml_backend_t backend,
|
||||||
const String2TensorStorage& tensor_storage_map = {},
|
const String2TensorStorage& tensor_storage_map = {},
|
||||||
@ -1799,7 +1828,13 @@ struct LLMEmbedder : public Conditioner {
|
|||||||
arch = LLM::LLMArch::GPT_OSS_20B;
|
arch = LLM::LLMArch::GPT_OSS_20B;
|
||||||
} else if (sd_version_is_pid(version)) {
|
} else if (sd_version_is_pid(version)) {
|
||||||
arch = LLM::LLMArch::GEMMA2_2B;
|
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;
|
arch = LLM::LLMArch::QWEN3_VL;
|
||||||
} else if (sd_version_is_z_image(version) || version == VERSION_OVIS_IMAGE || version == VERSION_FLUX2_KLEIN) {
|
} else if (sd_version_is_z_image(version) || version == VERSION_OVIS_IMAGE || version == VERSION_FLUX2_KLEIN) {
|
||||||
arch = LLM::LLMArch::QWEN3;
|
arch = LLM::LLMArch::QWEN3;
|
||||||
@ -1819,54 +1854,101 @@ struct LLMEmbedder : public Conditioner {
|
|||||||
"text_encoders.llm",
|
"text_encoders.llm",
|
||||||
enable_vision,
|
enable_vision,
|
||||||
weight_manager);
|
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 {
|
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
|
||||||
llm->get_param_tensors(tensors, "text_encoders.llm");
|
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 {
|
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
|
||||||
llm->set_max_graph_vram_bytes(max_vram_bytes);
|
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 {
|
void set_stream_layers_enabled(bool enabled) override {
|
||||||
llm->set_stream_layers_enabled(enabled);
|
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 {
|
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
|
||||||
llm->set_runtime_backends(backends);
|
llm->set_runtime_backends(backends);
|
||||||
|
if (byt5) {
|
||||||
|
byt5->set_runtime_backends(backends);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
void set_graph_cut_layer_split_enabled(bool enabled) override {
|
void set_graph_cut_layer_split_enabled(bool enabled) override {
|
||||||
if (llm) {
|
if (llm) {
|
||||||
llm->set_graph_cut_layer_split_enabled(enabled);
|
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 {
|
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
|
||||||
if (llm) {
|
if (llm) {
|
||||||
llm->set_graph_cut_layer_split_backend_vram_limits(limits);
|
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 {
|
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
|
||||||
llm->get_param_tensors(tensors, "text_encoders.llm");
|
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 {
|
void set_flash_attention_enabled(bool enabled) override {
|
||||||
llm->set_flash_attention_enabled(enabled);
|
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 {
|
void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
|
||||||
if (llm) {
|
if (llm) {
|
||||||
llm->set_weight_adapter(adapter);
|
llm->set_weight_adapter(adapter);
|
||||||
}
|
}
|
||||||
|
if (byt5) {
|
||||||
|
byt5->set_weight_adapter(adapter);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
void runner_done() override {
|
void runner_done() override {
|
||||||
if (llm) {
|
if (llm) {
|
||||||
llm->runner_done();
|
llm->runner_done();
|
||||||
}
|
}
|
||||||
|
if (byt5) {
|
||||||
|
byt5->runner_done();
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
std::tuple<std::vector<int>, std::vector<float>, std::vector<float>> tokenize(std::string text,
|
std::tuple<std::vector<int>, std::vector<float>, std::vector<float>> tokenize(std::string text,
|
||||||
@ -1934,8 +2016,10 @@ struct LLMEmbedder : public Conditioner {
|
|||||||
const std::vector<std::pair<int, sd::Tensor<float>>>& image_embeds,
|
const std::vector<std::pair<int, sd::Tensor<float>>>& image_embeds,
|
||||||
const std::set<int>& out_layers,
|
const std::set<int>& out_layers,
|
||||||
int prompt_template_encode_start_idx,
|
int prompt_template_encode_start_idx,
|
||||||
bool spell_quotes = false,
|
bool spell_quotes = false,
|
||||||
int max_length = 100000000) {
|
int max_length = 100000000,
|
||||||
|
const LLM::DeepStackImageEmbeds& deepstack_image_embeds = {},
|
||||||
|
const std::vector<LLM::ImageGrid>& image_grids = {}) {
|
||||||
auto tokens_weights_mask = tokenize(prompt, prompt_attn_range, min_length, max_length, spell_quotes);
|
auto tokens_weights_mask = tokenize(prompt, prompt_attn_range, min_length, max_length, spell_quotes);
|
||||||
auto& tokens = std::get<0>(tokens_weights_mask);
|
auto& tokens = std::get<0>(tokens_weights_mask);
|
||||||
auto& weights = std::get<1>(tokens_weights_mask);
|
auto& weights = std::get<1>(tokens_weights_mask);
|
||||||
@ -1968,7 +2052,9 @@ struct LLMEmbedder : public Conditioner {
|
|||||||
false,
|
false,
|
||||||
false,
|
false,
|
||||||
true,
|
true,
|
||||||
true);
|
true,
|
||||||
|
deepstack_image_embeds,
|
||||||
|
image_grids);
|
||||||
GGML_ASSERT(!hidden_states.empty());
|
GGML_ASSERT(!hidden_states.empty());
|
||||||
hidden_states = apply_token_weights(std::move(hidden_states), weights);
|
hidden_states = apply_token_weights(std::move(hidden_states), weights);
|
||||||
GGML_ASSERT(hidden_states.shape()[1] > prompt_template_encode_start_idx);
|
GGML_ASSERT(hidden_states.shape()[1] > prompt_template_encode_start_idx);
|
||||||
@ -2050,6 +2136,8 @@ struct LLMEmbedder : public Conditioner {
|
|||||||
std::vector<std::string> extra_prompts;
|
std::vector<std::string> extra_prompts;
|
||||||
std::vector<std::pair<int, int>> extra_prompts_attn_range;
|
std::vector<std::pair<int, int>> extra_prompts_attn_range;
|
||||||
std::vector<std::pair<int, sd::Tensor<float>>> image_embeds;
|
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 prompt_template_encode_start_idx = 34;
|
||||||
int min_length = 0; // pad tokens
|
int min_length = 0; // pad tokens
|
||||||
int max_length = 100000000;
|
int max_length = 100000000;
|
||||||
@ -2060,7 +2148,148 @@ struct LLMEmbedder : public Conditioner {
|
|||||||
int64_t t0 = ggml_time_ms();
|
int64_t t0 = ggml_time_ms();
|
||||||
RefImageResizeMode resize_mode = conditioner_params.ref_image_params.vlm_resize_mode;
|
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 int pad_token = 151643;
|
||||||
const std::string prompt_prefix =
|
const std::string prompt_prefix =
|
||||||
"<|im_start|>system\nGiven a user input that may include a text prompt alone, "
|
"<|im_start|>system\nGiven a user input that may include a text prompt alone, "
|
||||||
@ -2144,22 +2373,22 @@ struct LLMEmbedder : public Conditioner {
|
|||||||
prompt += conditioner_params.text;
|
prompt += conditioner_params.text;
|
||||||
prompt_attn_range = {0, 0};
|
prompt_attn_range = {0, 0};
|
||||||
prompt += "<|im_end|>\n<|im_start|>assistant\n";
|
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()) {
|
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;
|
prompt_template_encode_start_idx = 64;
|
||||||
int image_embed_idx = 64 + 6;
|
int image_embed_idx = 64 + 6;
|
||||||
|
|
||||||
int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
|
int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
|
||||||
if (min_pixels <= 0) {
|
if (min_pixels <= 0) {
|
||||||
min_pixels = 384;
|
min_pixels = sd_version_is_mage_flow(version) ? -1 : 384;
|
||||||
if (resize_mode == RefImageResizeMode::AREA) {
|
if (min_pixels > 0 && resize_mode == RefImageResizeMode::AREA) {
|
||||||
min_pixels *= min_pixels;
|
min_pixels *= min_pixels;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
|
int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
|
||||||
if (max_pixels <= 0) {
|
if (max_pixels <= 0) {
|
||||||
max_pixels = 560;
|
max_pixels = sd_version_is_mage_flow(version) ? 384 : 560;
|
||||||
if (resize_mode == RefImageResizeMode::AREA) {
|
if (resize_mode == RefImageResizeMode::AREA) {
|
||||||
max_pixels *= max_pixels;
|
max_pixels *= max_pixels;
|
||||||
}
|
}
|
||||||
@ -2187,7 +2416,7 @@ struct LLMEmbedder : public Conditioner {
|
|||||||
image_embeds.emplace_back(image_embed_idx, image_embed);
|
image_embeds.emplace_back(image_embed_idx, image_embed);
|
||||||
image_embed_idx += 1 + static_cast<int>(image_embed.shape()[1]) + 6;
|
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];
|
int64_t num_image_tokens = image_embed.shape()[1];
|
||||||
img_prompt.reserve(num_image_tokens * placeholder.size());
|
img_prompt.reserve(num_image_tokens * placeholder.size());
|
||||||
for (int j = 0; j < num_image_tokens; j++) {
|
for (int j = 0; j < num_image_tokens; j++) {
|
||||||
@ -2215,6 +2444,9 @@ struct LLMEmbedder : public Conditioner {
|
|||||||
|
|
||||||
prompt += "<|im_end|>\n<|im_start|>assistant\n";
|
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)) {
|
} else if (sd_version_is_boogu_image(version)) {
|
||||||
prompt_template_encode_start_idx = 0;
|
prompt_template_encode_start_idx = 0;
|
||||||
|
|
||||||
@ -2588,8 +2820,50 @@ struct LLMEmbedder : public Conditioner {
|
|||||||
out_layers,
|
out_layers,
|
||||||
prompt_template_encode_start_idx,
|
prompt_template_encode_start_idx,
|
||||||
spell_quotes,
|
spell_quotes,
|
||||||
max_length);
|
max_length,
|
||||||
|
deepstack_image_embeds,
|
||||||
|
image_grids);
|
||||||
std::vector<sd::Tensor<float>> extra_hidden_states_vec;
|
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++) {
|
for (int i = 0; i < extra_prompts.size(); i++) {
|
||||||
auto extra_hidden_states = encode_prompt(n_threads,
|
auto extra_hidden_states = encode_prompt(n_threads,
|
||||||
extra_prompts[i],
|
extra_prompts[i],
|
||||||
@ -2609,6 +2883,17 @@ struct LLMEmbedder : public Conditioner {
|
|||||||
SDCondition result;
|
SDCondition result;
|
||||||
result.c_crossattn = std::move(hidden_states);
|
result.c_crossattn = std::move(hidden_states);
|
||||||
result.extra_c_crossattns = std::move(extra_hidden_states_vec);
|
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;
|
return result;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@ -2719,6 +3004,10 @@ struct LTXAVEmbedder : public Conditioner {
|
|||||||
projector->get_param_tensors(tensors, "text_embedding_projection");
|
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 {
|
void set_flash_attention_enabled(bool enabled) override {
|
||||||
llm->set_flash_attention_enabled(enabled);
|
llm->set_flash_attention_enabled(enabled);
|
||||||
projector->set_flash_attention_enabled(enabled);
|
projector->set_flash_attention_enabled(enabled);
|
||||||
|
|||||||
@ -1039,6 +1039,38 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_linear(ggml_context* ctx,
|
|||||||
return x;
|
return x;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
__STATIC_INLINE__ ggml_tensor* ggml_ext_linear_i8_tensorwise(ggml_context* ctx,
|
||||||
|
ggml_tensor* x,
|
||||||
|
ggml_tensor* w,
|
||||||
|
ggml_tensor* weight_scale,
|
||||||
|
ggml_tensor* b,
|
||||||
|
int convrot_group_size,
|
||||||
|
float scale = 1.f) {
|
||||||
|
GGML_ASSERT(x->type == GGML_TYPE_F32 || (x->type == GGML_TYPE_I8 && scale == 1.f));
|
||||||
|
if (scale != 1.f) {
|
||||||
|
x = ggml_ext_scale(ctx, x, scale);
|
||||||
|
}
|
||||||
|
|
||||||
|
ggml_tensor* fused_bias = scale == 1.f ? b : nullptr;
|
||||||
|
if (x->ne[2] * x->ne[3] > 1024) {
|
||||||
|
int64_t ne2 = x->ne[2];
|
||||||
|
int64_t ne3 = x->ne[3];
|
||||||
|
x = ggml_reshape_2d(ctx, x, x->ne[0], x->ne[1] * x->ne[2] * x->ne[3]);
|
||||||
|
x = ggml_mul_mat_i8_tensorwise(ctx, w, x, weight_scale, fused_bias, convrot_group_size);
|
||||||
|
x = ggml_reshape_4d(ctx, x, x->ne[0], x->ne[1] / ne2 / ne3, ne2, ne3);
|
||||||
|
} else {
|
||||||
|
x = ggml_mul_mat_i8_tensorwise(ctx, w, x, weight_scale, fused_bias, convrot_group_size);
|
||||||
|
}
|
||||||
|
|
||||||
|
if (scale != 1.f) {
|
||||||
|
x = ggml_ext_scale(ctx, x, 1.f / scale);
|
||||||
|
if (b != nullptr) {
|
||||||
|
x = ggml_add_inplace(ctx, x, b);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return x;
|
||||||
|
}
|
||||||
|
|
||||||
__STATIC_INLINE__ ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx,
|
__STATIC_INLINE__ ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx,
|
||||||
ggml_backend_t backend,
|
ggml_backend_t backend,
|
||||||
ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
@ -1679,6 +1711,13 @@ struct WeightAdapter {
|
|||||||
ggml_tensor* b,
|
ggml_tensor* b,
|
||||||
const std::string& prefix,
|
const std::string& prefix,
|
||||||
ForwardParams forward_params) = 0;
|
ForwardParams forward_params) = 0;
|
||||||
|
virtual ggml_tensor* add_lora_to_output(ggml_context* ctx,
|
||||||
|
ggml_backend_t backend,
|
||||||
|
ggml_tensor* x,
|
||||||
|
ggml_tensor* w,
|
||||||
|
ggml_tensor* output,
|
||||||
|
const std::string& prefix,
|
||||||
|
ForwardParams forward_params) = 0;
|
||||||
virtual size_t get_extra_graph_size() = 0;
|
virtual size_t get_extra_graph_size() = 0;
|
||||||
};
|
};
|
||||||
|
|
||||||
@ -1689,11 +1728,14 @@ struct GGMLRunnerContext {
|
|||||||
bool conv2d_direct_enabled = false;
|
bool conv2d_direct_enabled = false;
|
||||||
bool circular_x_enabled = false;
|
bool circular_x_enabled = false;
|
||||||
bool circular_y_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::shared_ptr<WeightAdapter> weight_adapter = nullptr;
|
||||||
std::vector<std::pair<ggml_tensor*, std::string>>* debug_tensors = nullptr;
|
std::vector<std::pair<ggml_tensor*, std::string>>* debug_tensors = nullptr;
|
||||||
std::function<ggml_tensor*(const std::string&)> get_cache_tensor;
|
std::function<ggml_tensor*(const std::string&)> get_cache_tensor;
|
||||||
std::function<void(const std::string&, ggml_tensor*)> cache_tensor;
|
std::function<void(const std::string&, ggml_tensor*)> cache_tensor;
|
||||||
std::function<void(ggml_tensor*, const void*)> set_backend_tensor_data;
|
std::function<void(ggml_tensor*, const void*)> set_backend_tensor_data;
|
||||||
|
std::map<std::pair<ggml_tensor*, int>, ggml_tensor*> int8_convrot_cache;
|
||||||
|
|
||||||
void capture_tensor(const std::string& name, ggml_tensor* tensor) {
|
void capture_tensor(const std::string& name, ggml_tensor* tensor) {
|
||||||
if (debug_tensors == nullptr || tensor == nullptr) {
|
if (debug_tensors == nullptr || tensor == nullptr) {
|
||||||
@ -1751,8 +1793,9 @@ protected:
|
|||||||
std::vector<size_t> graph_cut_layer_split_backend_vram_limits_;
|
std::vector<size_t> graph_cut_layer_split_backend_vram_limits_;
|
||||||
|
|
||||||
std::vector<ggml_backend_t> extra_runtime_backends; // borrowed (SDBackendManager-owned)
|
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
|
size_t sched_graph_capacity = 0;
|
||||||
|
ggml_backend_t cpu_fallback_backend = nullptr; // owned, sched requires a trailing CPU backend
|
||||||
bool multi_device_eval_callback_warned = false;
|
bool multi_device_eval_callback_warned = false;
|
||||||
|
|
||||||
std::shared_ptr<WeightAdapter> weight_adapter = nullptr;
|
std::shared_ptr<WeightAdapter> weight_adapter = nullptr;
|
||||||
@ -2038,9 +2081,19 @@ protected:
|
|||||||
// Pass explicit buffer types: synthesized defaults can make CUDA devices
|
// Pass explicit buffer types: synthesized defaults can make CUDA devices
|
||||||
// report supporting each other's buffers and skip a required copy.
|
// report supporting each other's buffers and skip a required copy.
|
||||||
bool ensure_sched(ggml_cgraph* gf) {
|
bool ensure_sched(ggml_cgraph* gf) {
|
||||||
if (sched != nullptr) {
|
const size_t required_graph_size = gf != nullptr
|
||||||
|
? std::max<size_t>(1,
|
||||||
|
(size_t)ggml_graph_n_nodes(gf) +
|
||||||
|
sd::ggml_graph_cut::leaf_count(gf))
|
||||||
|
: 1;
|
||||||
|
if (sched != nullptr && sched_graph_capacity >= required_graph_size) {
|
||||||
return true;
|
return true;
|
||||||
}
|
}
|
||||||
|
if (sched != nullptr) {
|
||||||
|
ggml_backend_sched_free(sched);
|
||||||
|
sched = nullptr;
|
||||||
|
sched_graph_capacity = 0;
|
||||||
|
}
|
||||||
std::vector<ggml_backend_t> backends;
|
std::vector<ggml_backend_t> backends;
|
||||||
backends.reserve(extra_runtime_backends.size() + 2);
|
backends.reserve(extra_runtime_backends.size() + 2);
|
||||||
backends.push_back(runtime_backend);
|
backends.push_back(runtime_backend);
|
||||||
@ -2068,20 +2121,17 @@ protected:
|
|||||||
bufts.push_back(buft);
|
bufts.push_back(buft);
|
||||||
}
|
}
|
||||||
|
|
||||||
size_t graph_size = MAX_GRAPH_SIZE;
|
|
||||||
if (gf != nullptr) {
|
|
||||||
graph_size = std::max<size_t>(graph_size, (size_t)ggml_graph_n_nodes(gf));
|
|
||||||
}
|
|
||||||
sched = ggml_backend_sched_new(backends.data(),
|
sched = ggml_backend_sched_new(backends.data(),
|
||||||
bufts.data(),
|
bufts.data(),
|
||||||
(int)backends.size(),
|
(int)backends.size(),
|
||||||
graph_size,
|
required_graph_size,
|
||||||
/*parallel=*/false,
|
/*parallel=*/false,
|
||||||
/*op_offload=*/false);
|
/*op_offload=*/false);
|
||||||
if (sched == nullptr) {
|
if (sched == nullptr) {
|
||||||
LOG_ERROR("%s: failed to create backend sched", get_desc().c_str());
|
LOG_ERROR("%s: failed to create backend sched", get_desc().c_str());
|
||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
|
sched_graph_capacity = required_graph_size;
|
||||||
return true;
|
return true;
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -2145,8 +2195,22 @@ protected:
|
|||||||
return !extra_runtime_backends.empty();
|
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) {
|
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 sched replaces the gallocr. Do NOT ggml_backend_sched_reserve
|
||||||
// the graph here: reserve runs split_graph, which rewires the
|
// the graph here: reserve runs split_graph, which rewires the
|
||||||
// graph's src pointers to sched-internal copy tensors, and the
|
// graph's src pointers to sched-internal copy tensors, and the
|
||||||
@ -2154,6 +2218,10 @@ protected:
|
|||||||
// rewired graph, silently corrupting every cross-backend input. A
|
// rewired graph, silently corrupting every cross-backend input. A
|
||||||
// graph must be split at most once; the alloc in execute_graph
|
// graph must be split at most once; the alloc in execute_graph
|
||||||
// performs the real allocation.
|
// performs the real allocation.
|
||||||
|
if (compute_allocr != nullptr) {
|
||||||
|
ggml_gallocr_free(compute_allocr);
|
||||||
|
compute_allocr = nullptr;
|
||||||
|
}
|
||||||
return ensure_sched(gf);
|
return ensure_sched(gf);
|
||||||
}
|
}
|
||||||
if (compute_allocr != nullptr) {
|
if (compute_allocr != nullptr) {
|
||||||
@ -2751,7 +2819,7 @@ protected:
|
|||||||
};
|
};
|
||||||
ComputeBufferGuard compute_buffer_guard(this, free_compute_buffer);
|
ComputeBufferGuard compute_buffer_guard(this, free_compute_buffer);
|
||||||
|
|
||||||
if (is_multi_device()) {
|
if (sched != nullptr) {
|
||||||
ggml_backend_sched_reset(sched);
|
ggml_backend_sched_reset(sched);
|
||||||
pin_multi_device_nodes(gf); // reset clears the pins; re-apply before alloc
|
pin_multi_device_nodes(gf); // reset clears the pins; re-apply before alloc
|
||||||
if (!ggml_backend_sched_alloc_graph(sched, gf)) {
|
if (!ggml_backend_sched_alloc_graph(sched, gf)) {
|
||||||
@ -2772,9 +2840,9 @@ protected:
|
|||||||
}
|
}
|
||||||
|
|
||||||
ggml_status status;
|
ggml_status status;
|
||||||
if (is_multi_device()) {
|
if (sched != nullptr) {
|
||||||
if (sd_get_backend_eval_callback() != nullptr && !multi_device_eval_callback_warned) {
|
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());
|
get_desc().c_str());
|
||||||
multi_device_eval_callback_warned = true;
|
multi_device_eval_callback_warned = true;
|
||||||
}
|
}
|
||||||
@ -3010,18 +3078,16 @@ public:
|
|||||||
}
|
}
|
||||||
if (sched != nullptr) {
|
if (sched != nullptr) {
|
||||||
ggml_backend_sched_free(sched);
|
ggml_backend_sched_free(sched);
|
||||||
sched = nullptr;
|
sched = nullptr;
|
||||||
|
sched_graph_capacity = 0;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// do copy after alloc graph
|
// do copy after alloc graph
|
||||||
void set_backend_tensor_data(ggml_tensor* tensor, const void* data) {
|
void set_backend_tensor_data(ggml_tensor* tensor, const void* data) {
|
||||||
if (is_multi_device()) {
|
// The scheduler only allocates standalone data tensors when they are
|
||||||
// The sched only assigns a backend (and thus a buffer) to tensors
|
// marked as graph inputs. The flag is harmless for single-backend graphs.
|
||||||
// that participate in the graph; flag standalone data tensors as
|
ggml_set_input(tensor);
|
||||||
// inputs so they get one.
|
|
||||||
ggml_set_input(tensor);
|
|
||||||
}
|
|
||||||
backend_tensor_data_map[tensor] = data;
|
backend_tensor_data_map[tensor] = data;
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -3238,6 +3304,11 @@ protected:
|
|||||||
|
|
||||||
virtual void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") {}
|
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:
|
public:
|
||||||
void init(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") {
|
void init(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") {
|
||||||
if (prefix.size() > 0) {
|
if (prefix.size() > 0) {
|
||||||
@ -3288,6 +3359,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() {
|
virtual std::string get_desc() {
|
||||||
return "GGMLBlock";
|
return "GGMLBlock";
|
||||||
}
|
}
|
||||||
@ -3309,7 +3392,7 @@ public:
|
|||||||
|
|
||||||
class Identity : public UnaryBlock {
|
class Identity : public UnaryBlock {
|
||||||
public:
|
public:
|
||||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||||
return x;
|
return x;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@ -3322,14 +3405,18 @@ protected:
|
|||||||
bool force_f32;
|
bool force_f32;
|
||||||
bool force_prec_f32;
|
bool force_prec_f32;
|
||||||
bool allow_weight_scale;
|
bool allow_weight_scale;
|
||||||
bool has_weight_scale = false;
|
bool has_weight_scale = false;
|
||||||
|
bool int8_convrot = false;
|
||||||
|
int int8_convrot_group_size = 0;
|
||||||
float scale;
|
float scale;
|
||||||
std::string prefix;
|
std::string prefix;
|
||||||
|
|
||||||
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
||||||
this->prefix = prefix;
|
this->prefix = prefix;
|
||||||
has_weight_scale = false;
|
has_weight_scale = false;
|
||||||
enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32);
|
int8_convrot = false;
|
||||||
|
int8_convrot_group_size = 0;
|
||||||
|
enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32);
|
||||||
if (in_features % ggml_blck_size(wtype) != 0 || force_f32) {
|
if (in_features % ggml_blck_size(wtype) != 0 || force_f32) {
|
||||||
wtype = GGML_TYPE_F32;
|
wtype = GGML_TYPE_F32;
|
||||||
}
|
}
|
||||||
@ -3338,10 +3425,18 @@ protected:
|
|||||||
enum ggml_type wtype = GGML_TYPE_F32;
|
enum ggml_type wtype = GGML_TYPE_F32;
|
||||||
params["bias"] = ggml_new_tensor_1d(ctx, wtype, out_features);
|
params["bias"] = ggml_new_tensor_1d(ctx, wtype, out_features);
|
||||||
}
|
}
|
||||||
if (allow_weight_scale && tensor_storage_map.find(prefix + "weight_scale") != tensor_storage_map.end()) {
|
auto weight_storage = tensor_storage_map.find(prefix + "weight");
|
||||||
|
const bool is_int8_tensorwise = weight_storage != tensor_storage_map.end() && weight_storage->second.is_int8_tensorwise;
|
||||||
|
if ((allow_weight_scale || is_int8_tensorwise) && tensor_storage_map.find(prefix + "weight_scale") != tensor_storage_map.end()) {
|
||||||
params["weight_scale"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_features);
|
params["weight_scale"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_features);
|
||||||
has_weight_scale = true;
|
has_weight_scale = true;
|
||||||
}
|
}
|
||||||
|
if (is_int8_tensorwise) {
|
||||||
|
GGML_ASSERT(wtype == GGML_TYPE_I8);
|
||||||
|
GGML_ASSERT(has_weight_scale);
|
||||||
|
int8_convrot = weight_storage->second.int8_convrot;
|
||||||
|
int8_convrot_group_size = weight_storage->second.int8_convrot_group_size;
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
public:
|
public:
|
||||||
@ -3368,7 +3463,7 @@ public:
|
|||||||
force_prec_f32 = force_prec_f32_;
|
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* w = params["weight"];
|
||||||
ggml_tensor* b = nullptr;
|
ggml_tensor* b = nullptr;
|
||||||
if (bias) {
|
if (bias) {
|
||||||
@ -3376,6 +3471,49 @@ public:
|
|||||||
}
|
}
|
||||||
ggml_tensor* linear_bias = has_weight_scale ? nullptr : b;
|
ggml_tensor* linear_bias = has_weight_scale ? nullptr : b;
|
||||||
ggml_tensor* out = nullptr;
|
ggml_tensor* out = nullptr;
|
||||||
|
if (w->type == GGML_TYPE_I8) {
|
||||||
|
if (x->type != GGML_TYPE_F32) {
|
||||||
|
x = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, x);
|
||||||
|
}
|
||||||
|
if (!ggml_is_contiguous(x)) {
|
||||||
|
x = ggml_cont(ctx->ggml_ctx, x);
|
||||||
|
}
|
||||||
|
ggml_tensor* lora_input = x;
|
||||||
|
if (ctx->weight_adapter && b != nullptr) {
|
||||||
|
b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
|
||||||
|
}
|
||||||
|
if (int8_convrot && scale == 1.f) {
|
||||||
|
const auto cache_key = std::make_pair(x, int8_convrot_group_size);
|
||||||
|
auto cached = ctx->int8_convrot_cache.find(cache_key);
|
||||||
|
if (cached == ctx->int8_convrot_cache.end()) {
|
||||||
|
x = ggml_quantize_i8_convrot(ctx->ggml_ctx, x, int8_convrot_group_size);
|
||||||
|
ctx->int8_convrot_cache.emplace(cache_key, x);
|
||||||
|
} else {
|
||||||
|
x = cached->second;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
out = ggml_ext_linear_i8_tensorwise(ctx->ggml_ctx,
|
||||||
|
x,
|
||||||
|
w,
|
||||||
|
params["weight_scale"],
|
||||||
|
b,
|
||||||
|
int8_convrot ? int8_convrot_group_size : 0,
|
||||||
|
scale);
|
||||||
|
if (ctx->weight_adapter) {
|
||||||
|
WeightAdapter::ForwardParams forward_params;
|
||||||
|
forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR;
|
||||||
|
forward_params.linear.force_prec_f32 = force_prec_f32;
|
||||||
|
forward_params.linear.scale = scale;
|
||||||
|
out = ctx->weight_adapter->add_lora_to_output(ctx->ggml_ctx,
|
||||||
|
ctx->backend,
|
||||||
|
lora_input,
|
||||||
|
w,
|
||||||
|
out,
|
||||||
|
prefix,
|
||||||
|
forward_params);
|
||||||
|
}
|
||||||
|
return out;
|
||||||
|
}
|
||||||
if (ctx->weight_adapter) {
|
if (ctx->weight_adapter) {
|
||||||
WeightAdapter::ForwardParams forward_params;
|
WeightAdapter::ForwardParams forward_params;
|
||||||
forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR;
|
forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR;
|
||||||
@ -3415,6 +3553,10 @@ protected:
|
|||||||
params["weight"] = ggml_new_tensor_2d(ctx, wtype, embedding_dim, num_embeddings);
|
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:
|
public:
|
||||||
Embedding(int64_t num_embeddings, int64_t embedding_dim)
|
Embedding(int64_t num_embeddings, int64_t embedding_dim)
|
||||||
: embedding_dim(embedding_dim),
|
: embedding_dim(embedding_dim),
|
||||||
@ -3422,7 +3564,7 @@ public:
|
|||||||
}
|
}
|
||||||
|
|
||||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
ggml_tensor* input_ids) {
|
ggml_tensor* input_ids) override {
|
||||||
// input_ids: [N, n_token]
|
// input_ids: [N, n_token]
|
||||||
auto weight = params["weight"];
|
auto weight = params["weight"];
|
||||||
|
|
||||||
@ -3482,11 +3624,11 @@ public:
|
|||||||
scale = scale_value;
|
scale = scale_value;
|
||||||
}
|
}
|
||||||
|
|
||||||
std::string get_desc() {
|
std::string get_desc() override {
|
||||||
return "Conv2d";
|
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* w = params["weight"];
|
||||||
ggml_tensor* b = nullptr;
|
ggml_tensor* b = nullptr;
|
||||||
if (bias) {
|
if (bias) {
|
||||||
@ -3569,11 +3711,11 @@ public:
|
|||||||
scale = scale_value;
|
scale = scale_value;
|
||||||
}
|
}
|
||||||
|
|
||||||
std::string get_desc() {
|
std::string get_desc() override {
|
||||||
return "Conv2d_grouped";
|
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* w = params["weight"];
|
||||||
ggml_tensor* b = nullptr;
|
ggml_tensor* b = nullptr;
|
||||||
if (bias) {
|
if (bias) {
|
||||||
@ -3609,18 +3751,19 @@ public:
|
|||||||
if (groups == in_channels && groups == out_channels) {
|
if (groups == in_channels && groups == out_channels) {
|
||||||
ggml_tensor* res;
|
ggml_tensor* res;
|
||||||
if (ctx->conv2d_direct_enabled) {
|
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,
|
stride.second, stride.first,
|
||||||
padding.second, padding.first,
|
padding.second, padding.first,
|
||||||
dilation.second, dilation.first);
|
dilation.second, dilation.first);
|
||||||
} else {
|
} 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,
|
stride.second, stride.first,
|
||||||
padding.second, padding.first,
|
padding.second, padding.first,
|
||||||
dilation.second, dilation.first);
|
dilation.second, dilation.first);
|
||||||
}
|
}
|
||||||
if (b) {
|
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;
|
return res;
|
||||||
}
|
}
|
||||||
@ -3725,7 +3868,7 @@ public:
|
|||||||
bias(bias),
|
bias(bias),
|
||||||
force_prec_f32(force_prec_f32) {}
|
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* w = params["weight"];
|
||||||
ggml_tensor* b = nullptr;
|
ggml_tensor* b = nullptr;
|
||||||
if (ctx->weight_adapter) {
|
if (ctx->weight_adapter) {
|
||||||
@ -3778,7 +3921,7 @@ public:
|
|||||||
elementwise_affine(elementwise_affine),
|
elementwise_affine(elementwise_affine),
|
||||||
bias(bias) {}
|
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* w = nullptr;
|
||||||
ggml_tensor* b = nullptr;
|
ggml_tensor* b = nullptr;
|
||||||
|
|
||||||
@ -3865,7 +4008,7 @@ public:
|
|||||||
: hidden_size(hidden_size),
|
: hidden_size(hidden_size),
|
||||||
eps(eps) {}
|
eps(eps) {}
|
||||||
|
|
||||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||||
ggml_tensor* w = params["weight"];
|
ggml_tensor* w = params["weight"];
|
||||||
if (ctx->weight_adapter) {
|
if (ctx->weight_adapter) {
|
||||||
w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");
|
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;
|
*module = SDBackendModule::UPSCALER;
|
||||||
return true;
|
return true;
|
||||||
}
|
}
|
||||||
|
if (name == "detector" || name == "adetailer" || name == "yolo") {
|
||||||
|
*module = SDBackendModule::DETECTOR;
|
||||||
|
return true;
|
||||||
|
}
|
||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -956,6 +960,8 @@ const char* sd_backend_module_name(SDBackendModule module) {
|
|||||||
return "photomaker";
|
return "photomaker";
|
||||||
case SDBackendModule::UPSCALER:
|
case SDBackendModule::UPSCALER:
|
||||||
return "upscaler";
|
return "upscaler";
|
||||||
|
case SDBackendModule::DETECTOR:
|
||||||
|
return "detector";
|
||||||
}
|
}
|
||||||
return "unknown";
|
return "unknown";
|
||||||
}
|
}
|
||||||
|
|||||||
@ -20,6 +20,7 @@ enum class SDBackendModule {
|
|||||||
CONTROL_NET,
|
CONTROL_NET,
|
||||||
PHOTOMAKER,
|
PHOTOMAKER,
|
||||||
UPSCALER,
|
UPSCALER,
|
||||||
|
DETECTOR,
|
||||||
};
|
};
|
||||||
|
|
||||||
struct SDBackendAssignment {
|
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);
|
true);
|
||||||
std::vector<bool> class_token_mask;
|
std::vector<bool> class_token_mask;
|
||||||
for (int i = 0; i < tokens.size(); i++) {
|
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);
|
return std::make_tuple(tokens, weights, class_token_mask);
|
||||||
|
|||||||
27
src/model.h
@ -38,10 +38,12 @@ enum SDVersion {
|
|||||||
VERSION_LINGBOT_VIDEO,
|
VERSION_LINGBOT_VIDEO,
|
||||||
VERSION_QWEN_IMAGE,
|
VERSION_QWEN_IMAGE,
|
||||||
VERSION_QWEN_IMAGE_LAYERED,
|
VERSION_QWEN_IMAGE_LAYERED,
|
||||||
|
VERSION_HUNYUAN_VIDEO,
|
||||||
VERSION_ANIMA,
|
VERSION_ANIMA,
|
||||||
VERSION_FLUX2,
|
VERSION_FLUX2,
|
||||||
VERSION_FLUX2_KLEIN,
|
VERSION_FLUX2_KLEIN,
|
||||||
VERSION_LTXAV,
|
VERSION_LTXAV,
|
||||||
|
VERSION_MINIMAX_H3,
|
||||||
VERSION_HIDREAM_O1,
|
VERSION_HIDREAM_O1,
|
||||||
VERSION_Z_IMAGE,
|
VERSION_Z_IMAGE,
|
||||||
VERSION_BOOGU_IMAGE,
|
VERSION_BOOGU_IMAGE,
|
||||||
@ -54,6 +56,7 @@ enum SDVersion {
|
|||||||
VERSION_IDEOGRAM4,
|
VERSION_IDEOGRAM4,
|
||||||
VERSION_SEFI_IMAGE,
|
VERSION_SEFI_IMAGE,
|
||||||
VERSION_KREA2,
|
VERSION_KREA2,
|
||||||
|
VERSION_MAGE_FLOW,
|
||||||
VERSION_ESRGAN,
|
VERSION_ESRGAN,
|
||||||
VERSION_COUNT,
|
VERSION_COUNT,
|
||||||
};
|
};
|
||||||
@ -121,6 +124,10 @@ static inline bool sd_version_is_ltxav(SDVersion version) {
|
|||||||
return false;
|
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) {
|
static inline bool sd_version_is_wan(SDVersion version) {
|
||||||
if (version == VERSION_WAN2 || version == VERSION_WAN2_2_I2V || version == VERSION_WAN2_2_TI2V) {
|
if (version == VERSION_WAN2 || version == VERSION_WAN2_2_I2V || version == VERSION_WAN2_2_TI2V) {
|
||||||
return true;
|
return true;
|
||||||
@ -142,6 +149,13 @@ static inline bool sd_version_is_qwen_image(SDVersion version) {
|
|||||||
return false;
|
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) {
|
static inline bool sd_version_is_anima(SDVersion version) {
|
||||||
if (version == VERSION_ANIMA) {
|
if (version == VERSION_ANIMA) {
|
||||||
return true;
|
return true;
|
||||||
@ -219,6 +233,10 @@ static inline bool sd_version_is_krea2(SDVersion version) {
|
|||||||
return false;
|
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) {
|
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)) {
|
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;
|
return true;
|
||||||
@ -240,6 +258,10 @@ static inline bool sd_version_uses_wan_vae(SDVersion version) {
|
|||||||
return false;
|
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) {
|
static inline bool sd_version_is_inpaint(SDVersion version) {
|
||||||
if (version == VERSION_SD1_INPAINT ||
|
if (version == VERSION_SD1_INPAINT ||
|
||||||
version == VERSION_SD2_INPAINT ||
|
version == VERSION_SD2_INPAINT ||
|
||||||
@ -255,10 +277,12 @@ static inline bool sd_version_is_dit(SDVersion version) {
|
|||||||
if (sd_version_is_flux(version) ||
|
if (sd_version_is_flux(version) ||
|
||||||
sd_version_is_flux2(version) ||
|
sd_version_is_flux2(version) ||
|
||||||
sd_version_is_ltxav(version) ||
|
sd_version_is_ltxav(version) ||
|
||||||
|
sd_version_is_minimax_h3(version) ||
|
||||||
sd_version_is_sd3(version) ||
|
sd_version_is_sd3(version) ||
|
||||||
sd_version_is_wan(version) ||
|
sd_version_is_wan(version) ||
|
||||||
sd_version_is_lingbot_video(version) ||
|
sd_version_is_lingbot_video(version) ||
|
||||||
sd_version_is_qwen_image(version) ||
|
sd_version_is_qwen_image(version) ||
|
||||||
|
sd_version_is_hunyuan_video(version) ||
|
||||||
version == VERSION_HIDREAM_O1 ||
|
version == VERSION_HIDREAM_O1 ||
|
||||||
sd_version_is_anima(version) ||
|
sd_version_is_anima(version) ||
|
||||||
sd_version_is_z_image(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_pid(version) ||
|
||||||
sd_version_is_ideogram4(version) ||
|
sd_version_is_ideogram4(version) ||
|
||||||
sd_version_is_sefi_image(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 true;
|
||||||
}
|
}
|
||||||
return false;
|
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::unordered_map<std::string, ggml_tensor*> lora_tensors;
|
||||||
std::map<ggml_tensor*, ggml_tensor*> original_tensor_to_final_tensor;
|
std::map<ggml_tensor*, ggml_tensor*> original_tensor_to_final_tensor;
|
||||||
std::set<std::string> applied_lora_tensors;
|
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::string file_path;
|
||||||
std::shared_ptr<ModelManager> model_manager;
|
std::shared_ptr<ModelManager> model_manager;
|
||||||
ggml_backend_t params_backend = nullptr;
|
ggml_backend_t params_backend = nullptr;
|
||||||
@ -133,6 +135,8 @@ struct LoraModel : public GGMLRunner {
|
|||||||
lora_tensors.clear();
|
lora_tensors.clear();
|
||||||
original_tensor_to_final_tensor.clear();
|
original_tensor_to_final_tensor.clear();
|
||||||
applied_lora_tensors.clear();
|
applied_lora_tensors.clear();
|
||||||
|
skipped_incompatible_lora_tensors.clear();
|
||||||
|
warned_incompatible_model_tensors.clear();
|
||||||
applied = false;
|
applied = false;
|
||||||
tensor_preprocessed = false;
|
tensor_preprocessed = false;
|
||||||
}
|
}
|
||||||
@ -338,7 +342,9 @@ struct LoraModel : public GGMLRunner {
|
|||||||
iter = lora_tensors.find(hada_1_mid_name);
|
iter = lora_tensors.find(hada_1_mid_name);
|
||||||
if (iter != lora_tensors.end()) {
|
if (iter != lora_tensors.end()) {
|
||||||
hada_1_mid = ggml_ext_cast_f32(ctx, backend, iter->second);
|
hada_1_mid = ggml_ext_cast_f32(ctx, backend, iter->second);
|
||||||
hada_1_up = ggml_cont(ctx, ggml_transpose(ctx, hada_1_up));
|
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);
|
iter = lora_tensors.find(hada_2_down_name);
|
||||||
@ -354,7 +360,9 @@ struct LoraModel : public GGMLRunner {
|
|||||||
iter = lora_tensors.find(hada_2_mid_name);
|
iter = lora_tensors.find(hada_2_mid_name);
|
||||||
if (iter != lora_tensors.end()) {
|
if (iter != lora_tensors.end()) {
|
||||||
hada_2_mid = ggml_ext_cast_f32(ctx, backend, iter->second);
|
hada_2_mid = ggml_ext_cast_f32(ctx, backend, iter->second);
|
||||||
hada_2_up = ggml_cont(ctx, ggml_transpose(ctx, hada_2_up));
|
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) {
|
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);
|
diff = ggml_reshape(ctx, diff, model_tensor);
|
||||||
}
|
}
|
||||||
return diff;
|
return diff;
|
||||||
@ -555,10 +583,15 @@ struct LoraModel : public GGMLRunner {
|
|||||||
ggml_tensor* get_out_diff(ggml_context* ctx,
|
ggml_tensor* get_out_diff(ggml_context* ctx,
|
||||||
ggml_backend_t backend,
|
ggml_backend_t backend,
|
||||||
ggml_tensor* x,
|
ggml_tensor* x,
|
||||||
|
ggml_tensor* model_weight,
|
||||||
WeightAdapter::ForwardParams forward_params,
|
WeightAdapter::ForwardParams forward_params,
|
||||||
const std::string& model_tensor_name) {
|
const std::string& model_tensor_name) {
|
||||||
ggml_tensor* out_diff = nullptr;
|
ggml_tensor* out_diff = nullptr;
|
||||||
int index = 0;
|
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) {
|
while (true) {
|
||||||
std::string key;
|
std::string key;
|
||||||
if (index == 0) {
|
if (index == 0) {
|
||||||
@ -566,7 +599,6 @@ struct LoraModel : public GGMLRunner {
|
|||||||
} else {
|
} else {
|
||||||
key = model_tensor_name + "." + std::to_string(index);
|
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_name = "lora." + key + ".lokr_w1";
|
||||||
std::string lokr_w1_a_name = "lora." + key + ".lokr_w1_a";
|
std::string lokr_w1_a_name = "lora." + key + ".lokr_w1_a";
|
||||||
@ -634,7 +666,6 @@ struct LoraModel : public GGMLRunner {
|
|||||||
if (iter != lora_tensors.end()) {
|
if (iter != lora_tensors.end()) {
|
||||||
float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
|
float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
|
||||||
scale_value = alpha / rank;
|
scale_value = alpha / rank;
|
||||||
applied_lora_tensors.insert(alpha_name);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
if (rank == 1) {
|
if (rank == 1) {
|
||||||
@ -649,19 +680,27 @@ struct LoraModel : public GGMLRunner {
|
|||||||
out_diff = ggml_concat(ctx, out_diff, curr_out_diff, 0);
|
out_diff = ggml_concat(ctx, out_diff, curr_out_diff, 0);
|
||||||
}
|
}
|
||||||
|
|
||||||
if (lokr_w1)
|
if (lokr_w1) {
|
||||||
applied_lora_tensors.insert(lokr_w1_name);
|
used_tensors.push_back(lokr_w1_name);
|
||||||
if (lokr_w1_a)
|
}
|
||||||
applied_lora_tensors.insert(lokr_w1_a_name);
|
if (lokr_w1_a) {
|
||||||
if (lokr_w1_b)
|
used_tensors.push_back(lokr_w1_a_name);
|
||||||
applied_lora_tensors.insert(lokr_w1_b_name);
|
}
|
||||||
if (lokr_w2)
|
if (lokr_w1_b) {
|
||||||
applied_lora_tensors.insert(lokr_w2_name);
|
used_tensors.push_back(lokr_w1_b_name);
|
||||||
if (lokr_w2_a)
|
}
|
||||||
applied_lora_tensors.insert(lokr_w2_a_name);
|
if (lokr_w2) {
|
||||||
if (lokr_w2_b)
|
used_tensors.push_back(lokr_w2_name);
|
||||||
applied_lora_tensors.insert(lokr_w2_b_name);
|
}
|
||||||
applied_lora_tensors.insert(alpha_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++;
|
index++;
|
||||||
continue;
|
continue;
|
||||||
@ -707,27 +746,60 @@ struct LoraModel : public GGMLRunner {
|
|||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
|
|
||||||
applied_lora_tensors.insert(lora_up_name);
|
if (!is_conv2d) {
|
||||||
applied_lora_tensors.insert(lora_down_name);
|
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) {
|
bool compatible = down_in == model_weight->ne[0];
|
||||||
applied_lora_tensors.insert(lora_mid_name);
|
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;
|
float scale_value = 1.0f;
|
||||||
|
std::string scale_tensor_name;
|
||||||
|
|
||||||
int64_t rank = lora_down->ne[ggml_n_dims(lora_down) - 1];
|
int64_t rank = lora_down->ne[ggml_n_dims(lora_down) - 1];
|
||||||
iter = lora_tensors.find(scale_name);
|
iter = lora_tensors.find(scale_name);
|
||||||
if (iter != lora_tensors.end()) {
|
if (iter != lora_tensors.end()) {
|
||||||
scale_value = ggml_ext_backend_tensor_get_f32(iter->second);
|
scale_value = ggml_ext_backend_tensor_get_f32(iter->second);
|
||||||
applied_lora_tensors.insert(scale_name);
|
scale_tensor_name = scale_name;
|
||||||
} else {
|
} else {
|
||||||
iter = lora_tensors.find(alpha_name);
|
iter = lora_tensors.find(alpha_name);
|
||||||
if (iter != lora_tensors.end()) {
|
if (iter != lora_tensors.end()) {
|
||||||
float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
|
float alpha = ggml_ext_backend_tensor_get_f32(iter->second);
|
||||||
scale_value = alpha / rank;
|
scale_value = alpha / rank;
|
||||||
|
scale_tensor_name = alpha_name;
|
||||||
// LOG_DEBUG("rank %s %ld %.2f %.2f", alpha_name.c_str(), rank, alpha, scale_value);
|
// LOG_DEBUG("rank %s %ld %.2f %.2f", alpha_name.c_str(), rank, alpha, scale_value);
|
||||||
applied_lora_tensors.insert(alpha_name);
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
scale_value *= multiplier;
|
scale_value *= multiplier;
|
||||||
@ -787,15 +859,45 @@ struct LoraModel : public GGMLRunner {
|
|||||||
}
|
}
|
||||||
|
|
||||||
auto curr_out_diff = ggml_ext_scale(ctx, lx, scale_value, true);
|
auto curr_out_diff = ggml_ext_scale(ctx, lx, scale_value, true);
|
||||||
|
|
||||||
if (out_diff == nullptr) {
|
if (out_diff == nullptr) {
|
||||||
out_diff = curr_out_diff;
|
out_diff = curr_out_diff;
|
||||||
} else {
|
} 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++;
|
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;
|
return out_diff;
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -869,10 +971,13 @@ struct LoraModel : public GGMLRunner {
|
|||||||
void stat(bool at_runntime = false) {
|
void stat(bool at_runntime = false) {
|
||||||
size_t total_lora_tensors_count = 0;
|
size_t total_lora_tensors_count = 0;
|
||||||
size_t applied_lora_tensors_count = 0;
|
size_t applied_lora_tensors_count = 0;
|
||||||
|
size_t skipped_lora_tensors_count = 0;
|
||||||
|
|
||||||
for (auto& kv : lora_tensors) {
|
for (auto& kv : lora_tensors) {
|
||||||
total_lora_tensors_count++;
|
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) {
|
if (!at_runntime) {
|
||||||
LOG_WARN("unused lora tensor |%s|", kv.first.c_str());
|
LOG_WARN("unused lora tensor |%s|", kv.first.c_str());
|
||||||
print_ggml_tensor(kv.second, true);
|
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,
|
/* 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
|
* this function is called once to calculate the required buffer size
|
||||||
* and then again to actually generate a graph to be used */
|
* 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",
|
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 {
|
} else {
|
||||||
LOG_INFO("(%lu / %lu) LoRA tensors have been applied, lora_file_path = %s",
|
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);
|
forward_params.conv2d.scale);
|
||||||
}
|
}
|
||||||
for (auto& lora_model : lora_models) {
|
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) {
|
if (out_diff == nullptr) {
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
@ -962,6 +1072,34 @@ public:
|
|||||||
return out;
|
return out;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
ggml_tensor* add_lora_to_output(ggml_context* ctx,
|
||||||
|
ggml_backend_t backend,
|
||||||
|
ggml_tensor* x,
|
||||||
|
ggml_tensor* w,
|
||||||
|
ggml_tensor* output,
|
||||||
|
const std::string& prefix,
|
||||||
|
WeightAdapter::ForwardParams forward_params) override {
|
||||||
|
for (auto& lora_model : lora_models) {
|
||||||
|
ggml_tensor* weight_diff = lora_model->get_weight_diff(prefix + "weight", backend, ctx, w, false);
|
||||||
|
if (weight_diff != nullptr) {
|
||||||
|
GGML_ASSERT(forward_params.op_type == ForwardParams::op_type_t::OP_LINEAR);
|
||||||
|
ggml_tensor* out_diff = ggml_ext_linear(ctx,
|
||||||
|
x,
|
||||||
|
weight_diff,
|
||||||
|
nullptr,
|
||||||
|
forward_params.linear.force_prec_f32,
|
||||||
|
forward_params.linear.scale);
|
||||||
|
output = ggml_add_inplace(ctx, output, out_diff);
|
||||||
|
}
|
||||||
|
|
||||||
|
ggml_tensor* out_diff = lora_model->get_out_diff(ctx, backend, x, w, forward_params, prefix + "weight");
|
||||||
|
if (out_diff != nullptr) {
|
||||||
|
output = ggml_add_inplace(ctx, output, out_diff);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return output;
|
||||||
|
}
|
||||||
|
|
||||||
size_t get_extra_graph_size() override {
|
size_t get_extra_graph_size() override {
|
||||||
size_t lora_tensor_num = 0;
|
size_t lora_tensor_num = 0;
|
||||||
for (auto& lora_model : lora_models) {
|
for (auto& lora_model : lora_models) {
|
||||||
|
|||||||
@ -310,17 +310,33 @@ protected:
|
|||||||
int64_t context_dim;
|
int64_t context_dim;
|
||||||
int64_t n_head;
|
int64_t n_head;
|
||||||
int64_t d_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:
|
public:
|
||||||
CrossAttention(int64_t query_dim,
|
CrossAttention(int64_t query_dim,
|
||||||
int64_t context_dim,
|
int64_t context_dim,
|
||||||
int64_t n_head,
|
int64_t n_head,
|
||||||
int64_t d_head)
|
int64_t d_head,
|
||||||
|
bool enable_ip = false)
|
||||||
: n_head(n_head),
|
: n_head(n_head),
|
||||||
d_head(d_head),
|
d_head(d_head),
|
||||||
query_dim(query_dim),
|
query_dim(query_dim),
|
||||||
context_dim(context_dim) {
|
context_dim(context_dim),
|
||||||
|
enable_ip(enable_ip) {
|
||||||
int64_t inner_dim = d_head * n_head;
|
int64_t inner_dim = d_head * n_head;
|
||||||
if (context_dim == 320 && d_head == 320) {
|
if (context_dim == 320 && d_head == 320) {
|
||||||
// LOG_DEBUG("CrossAttention: temp set dim to 1024 for sdxs_09");
|
// 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]
|
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]
|
x = to_out_0->forward(ctx, x); // [N, n_token, query_dim]
|
||||||
return x;
|
return x;
|
||||||
}
|
}
|
||||||
@ -387,7 +412,7 @@ public:
|
|||||||
// inner_dim is always None or equal to dim
|
// inner_dim is always None or equal to dim
|
||||||
// gated_ff is always True
|
// gated_ff is always True
|
||||||
blocks["attn1"] = std::shared_ptr<GGMLBlock>(new CrossAttention(dim, dim, n_head, d_head));
|
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["ff"] = std::shared_ptr<GGMLBlock>(new FeedForward(dim, dim));
|
||||||
blocks["norm1"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
|
blocks["norm1"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
|
||||||
blocks["norm2"] = 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
|
int64_t context_dim = 768; // hidden_size, 1024 for VERSION_SD2
|
||||||
bool use_linear = false;
|
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");
|
auto iter = tensor_storage_map.find(prefix + "proj_out.weight");
|
||||||
if (iter != tensor_storage_map.end()) {
|
if (iter != tensor_storage_map.end()) {
|
||||||
int64_t inner_dim = n_head * d_head;
|
int64_t inner_dim = n_head * d_head;
|
||||||
|
|||||||
@ -535,6 +535,33 @@ namespace Rope {
|
|||||||
return vid_ids_repeated;
|
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,
|
__STATIC_INLINE__ std::vector<std::vector<float>> gen_qwen_image_ids(int t,
|
||||||
int h,
|
int h,
|
||||||
int w,
|
int w,
|
||||||
@ -627,6 +654,43 @@ namespace Rope {
|
|||||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
|
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,
|
__STATIC_INLINE__ std::vector<std::vector<float>> gen_lens_ids(int h,
|
||||||
int w,
|
int w,
|
||||||
int bs,
|
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__
|
||||||
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__
|
||||||
@ -706,11 +706,13 @@ namespace Flux {
|
|||||||
LastLayer(int64_t hidden_size,
|
LastLayer(int64_t hidden_size,
|
||||||
int64_t patch_size,
|
int64_t patch_size,
|
||||||
int64_t out_channels,
|
int64_t out_channels,
|
||||||
bool prune_mod = false,
|
bool prune_mod = false,
|
||||||
bool bias = true)
|
bool bias = true,
|
||||||
|
int64_t patch_volume = 0)
|
||||||
: prune_mod(prune_mod) {
|
: prune_mod(prune_mod) {
|
||||||
blocks["norm_final"] = std::shared_ptr<GGMLBlock>(new LayerNorm(hidden_size, 1e-06f, false));
|
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) {
|
if (!prune_mod) {
|
||||||
blocks["adaLN_modulation.1"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, 2 * hidden_size, bias));
|
blocks["adaLN_modulation.1"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, 2 * hidden_size, bias));
|
||||||
}
|
}
|
||||||
|
|||||||
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* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||||
ggml_tensor* scale = params["scale"];
|
ggml_tensor* scale = params["scale"];
|
||||||
scale = ggml_add(ctx->ggml_ctx, scale, ggml_ext_ones(ctx->ggml_ctx, scale->ne[0], 1, 1, 1));
|
if (ctx->weight_adapter) {
|
||||||
x = ggml_rms_norm(ctx->ggml_ctx, x, eps);
|
scale = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, scale, prefix + "scale.weight");
|
||||||
x = ggml_mul_inplace(ctx->ggml_ctx, x, 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);
|
||||||
return x;
|
return x;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@ -295,10 +298,11 @@ namespace Krea2 {
|
|||||||
class KreaDoubleSharedModulation : public GGMLBlock {
|
class KreaDoubleSharedModulation : public GGMLBlock {
|
||||||
protected:
|
protected:
|
||||||
int64_t dim;
|
int64_t dim;
|
||||||
|
std::string prefix;
|
||||||
|
|
||||||
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
||||||
GGML_UNUSED(tensor_storage_map);
|
GGML_UNUSED(tensor_storage_map);
|
||||||
GGML_UNUSED(prefix);
|
this->prefix = prefix;
|
||||||
params["lin"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dim * 6);
|
params["lin"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dim * 6);
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -307,7 +311,11 @@ namespace Krea2 {
|
|||||||
: dim(dim) {}
|
: dim(dim) {}
|
||||||
|
|
||||||
std::vector<ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* vec) {
|
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);
|
auto out = ggml_add(ctx->ggml_ctx, vec, lin);
|
||||||
return ggml_ext_chunk(ctx->ggml_ctx, out, 6, 0);
|
return ggml_ext_chunk(ctx->ggml_ctx, out, 6, 0);
|
||||||
}
|
}
|
||||||
@ -316,10 +324,11 @@ namespace Krea2 {
|
|||||||
class KreaFinalModulation : public GGMLBlock {
|
class KreaFinalModulation : public GGMLBlock {
|
||||||
protected:
|
protected:
|
||||||
int64_t dim;
|
int64_t dim;
|
||||||
|
std::string prefix;
|
||||||
|
|
||||||
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
||||||
GGML_UNUSED(tensor_storage_map);
|
GGML_UNUSED(tensor_storage_map);
|
||||||
GGML_UNUSED(prefix);
|
this->prefix = prefix;
|
||||||
params["lin"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, dim, 2);
|
params["lin"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, dim, 2);
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -328,7 +337,11 @@ namespace Krea2 {
|
|||||||
: dim(dim) {}
|
: dim(dim) {}
|
||||||
|
|
||||||
std::vector<ggml_tensor*> forward(GGMLRunnerContext* ctx, ggml_tensor* vec) {
|
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);
|
return ggml_ext_chunk(ctx->ggml_ctx, out, 2, 1);
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|||||||
@ -800,7 +800,7 @@ namespace LTXV {
|
|||||||
auto gate_mlp = mods[5];
|
auto gate_mlp = mods[5];
|
||||||
|
|
||||||
auto x_norm = rms_norm(ctx->ggml_ctx, x);
|
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);
|
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));
|
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 gate_q = mods[8];
|
||||||
|
|
||||||
auto q = rms_norm(ctx->ggml_ctx, x);
|
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;
|
auto context_mod = context;
|
||||||
if (prompt_timestep != nullptr) {
|
if (prompt_timestep != nullptr) {
|
||||||
auto prompt_mods = get_prompt_scale_shift_values(ctx, prompt_timestep);
|
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);
|
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);
|
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);
|
auto mlp_out = ff->forward(ctx, y);
|
||||||
x = ggml_add(ctx->ggml_ctx, x, apply_gate(ctx->ggml_ctx, mlp_out, gate_mlp));
|
x = ggml_add(ctx->ggml_ctx, x, apply_gate(ctx->ggml_ctx, mlp_out, gate_mlp));
|
||||||
return x;
|
return x;
|
||||||
@ -1177,11 +1177,11 @@ namespace LTXV {
|
|||||||
if (cross_attention_adaln) {
|
if (cross_attention_adaln) {
|
||||||
auto q_mods = get_ada_values(ctx, table, timestep, dim, 9, 6, 3);
|
auto q_mods = get_ada_values(ctx, table, timestep, dim, 9, 6, 3);
|
||||||
auto q = rms_norm(ctx->ggml_ctx, x);
|
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;
|
auto context_mod = context;
|
||||||
if (prompt_timestep != nullptr && prompt_table != nullptr) {
|
if (prompt_timestep != nullptr && prompt_table != nullptr) {
|
||||||
auto p_mods = get_ada_values(ctx, prompt_table, prompt_timestep, dim, 2);
|
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);
|
auto out = attn->forward(ctx, q, context_mod, attention_mask, nullptr, nullptr);
|
||||||
return apply_gate(ctx->ggml_ctx, out, q_mods[2]);
|
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_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);
|
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);
|
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]));
|
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,
|
auto v_txt = apply_text_cross_attention(ctx,
|
||||||
@ -1246,7 +1246,7 @@ namespace LTXV {
|
|||||||
if (run_ax) {
|
if (run_ax) {
|
||||||
auto a_mods = get_ada_values(ctx, a_table, a_timestep, a_dim, cross_attention_adaln ? 9 : 6);
|
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);
|
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);
|
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]));
|
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,
|
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_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_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 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 vx_scaled = LTXV::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 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_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_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];
|
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_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_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 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 ax_scaled = LTXV::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 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_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_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];
|
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 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);
|
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);
|
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]));
|
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 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);
|
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);
|
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]));
|
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);
|
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 = 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 = proj_out->forward(ctx, vx);
|
||||||
vx = unpatchify_video(ctx, vx, width, height, frames);
|
vx = unpatchify_video(ctx, vx, width, height, frames);
|
||||||
|
|
||||||
if (ax != nullptr && audio_time > 0) {
|
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);
|
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 = 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 = audio_proj_out->forward(ctx, ax);
|
||||||
ax = unpatchify_audio(ctx, ax, audio_time);
|
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 {
|
struct PatchEmbed : public GGMLBlock {
|
||||||
// 2D Image to Patch Embedding
|
// 2D/3D Image to Patch Embedding
|
||||||
protected:
|
protected:
|
||||||
|
bool is_3d;
|
||||||
bool flatten;
|
bool flatten;
|
||||||
bool dynamic_img_pad;
|
bool dynamic_img_pad;
|
||||||
int patch_size;
|
int patch_t;
|
||||||
|
int patch_h;
|
||||||
|
int patch_w;
|
||||||
|
int64_t embed_dim;
|
||||||
|
|
||||||
public:
|
public:
|
||||||
PatchEmbed(int64_t img_size = 224,
|
PatchEmbed(int64_t img_size = 224,
|
||||||
@ -149,42 +153,90 @@ public:
|
|||||||
int64_t embed_dim = 1536,
|
int64_t embed_dim = 1536,
|
||||||
bool bias = true,
|
bool bias = true,
|
||||||
bool flatten = true,
|
bool flatten = true,
|
||||||
bool dynamic_img_pad = true)
|
bool dynamic_img_pad = true,
|
||||||
: patch_size(patch_size),
|
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),
|
flatten(flatten),
|
||||||
dynamic_img_pad(dynamic_img_pad) {
|
dynamic_img_pad(dynamic_img_pad),
|
||||||
|
is_3d(is_3d) {
|
||||||
// img_size is always None
|
// img_size is always None
|
||||||
// patch_size is always 2
|
// patch_size is always 2
|
||||||
// in_chans is always 16
|
// in_chans is always 16
|
||||||
// norm_layer is always False
|
// norm_layer is always False
|
||||||
// strict_img_size is always true, but not used
|
// strict_img_size is always true, but not used
|
||||||
|
|
||||||
blocks["proj"] = std::shared_ptr<GGMLBlock>(new Conv2d(in_chans,
|
if (is_3d) {
|
||||||
embed_dim,
|
blocks["proj"] = std::make_shared<Conv3d>(in_chans,
|
||||||
{patch_size, patch_size},
|
embed_dim,
|
||||||
{patch_size, patch_size},
|
std::tuple{patch_size, patch_size, patch_size},
|
||||||
{0, 0},
|
std::tuple{patch_size, patch_size, patch_size},
|
||||||
{1, 1},
|
std::tuple{0, 0, 0},
|
||||||
bias));
|
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) {
|
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
// x: [N, C, H, W]
|
// x: [N, C, H, W] or [N*C, T, H, W]
|
||||||
// return: [N, H*W, embed_dim]
|
// return: [N, h_len*w_len, embed_dim] or [N, t_len*h_len*w_len, embed_dim]
|
||||||
auto proj = std::dynamic_pointer_cast<Conv2d>(blocks["proj"]);
|
auto proj = std::dynamic_pointer_cast<UnaryBlock>(blocks["proj"]);
|
||||||
|
|
||||||
if (dynamic_img_pad) {
|
if (dynamic_img_pad) {
|
||||||
int64_t W = x->ne[0];
|
int64_t W = x->ne[0];
|
||||||
int64_t H = x->ne[1];
|
int64_t H = x->ne[1];
|
||||||
int pad_h = (patch_size - H % patch_size) % patch_size;
|
int pad_t = 0;
|
||||||
int pad_w = (patch_size - W % patch_size) % patch_size;
|
int pad_h = (patch_h - static_cast<int>(H % patch_h)) % patch_h;
|
||||||
x = ggml_pad(ctx->ggml_ctx, x, pad_w, pad_h, 0, 0); // TODO: reflect pad mode
|
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) {
|
if (flatten) {
|
||||||
x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1], x->ne[2], x->ne[3]);
|
if (is_3d) {
|
||||||
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
|
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;
|
return x;
|
||||||
}
|
}
|
||||||
|
|||||||
@ -26,6 +26,7 @@ struct RefImageParams {
|
|||||||
RefImageResizeMode vlm_resize_mode = RefImageResizeMode::AREA;
|
RefImageResizeMode vlm_resize_mode = RefImageResizeMode::AREA;
|
||||||
int vlm_min_size = -1;
|
int vlm_min_size = -1;
|
||||||
int vlm_max_size = -1;
|
int vlm_max_size = -1;
|
||||||
|
bool resize_vae_to_target = false;
|
||||||
};
|
};
|
||||||
|
|
||||||
const std::unordered_map<std::string, RefImageParams> REF_IMAGE_PRESETS = {
|
const std::unordered_map<std::string, RefImageParams> REF_IMAGE_PRESETS = {
|
||||||
@ -34,6 +35,7 @@ const std::unordered_map<std::string, RefImageParams> REF_IMAGE_PRESETS = {
|
|||||||
{"flux2", {false, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::NONE, -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", {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}},
|
{"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}},
|
{"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_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}},
|
{"krea2_edit", {true, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::LONGEST_SIDE, 768, 768}},
|
||||||
@ -44,6 +46,8 @@ struct UNetDiffusionExtra {
|
|||||||
int num_video_frames = -1;
|
int num_video_frames = -1;
|
||||||
const std::vector<sd::Tensor<float>>* controls = nullptr;
|
const std::vector<sd::Tensor<float>>* controls = nullptr;
|
||||||
float control_strength = 0.f;
|
float control_strength = 0.f;
|
||||||
|
const sd::Tensor<float>* ip_context = nullptr;
|
||||||
|
float ip_scale = 1.f;
|
||||||
};
|
};
|
||||||
|
|
||||||
struct SkipLayerDiffusionExtra {
|
struct SkipLayerDiffusionExtra {
|
||||||
@ -83,10 +87,40 @@ struct LTXAVDiffusionExtra {
|
|||||||
const sd::Tensor<float>* video_positions = nullptr;
|
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 {
|
struct MiniT2IDiffusionExtra {
|
||||||
const sd::Tensor<float>* mask = nullptr;
|
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,
|
using DiffusionExtraParams = std::variant<std::monostate,
|
||||||
UNetDiffusionExtra,
|
UNetDiffusionExtra,
|
||||||
SkipLayerDiffusionExtra,
|
SkipLayerDiffusionExtra,
|
||||||
@ -95,7 +129,9 @@ using DiffusionExtraParams = std::variant<std::monostate,
|
|||||||
WanDiffusionExtra,
|
WanDiffusionExtra,
|
||||||
HiDreamO1DiffusionExtra,
|
HiDreamO1DiffusionExtra,
|
||||||
LTXAVDiffusionExtra,
|
LTXAVDiffusionExtra,
|
||||||
MiniT2IDiffusionExtra>;
|
MiniMaxH3DiffusionExtra,
|
||||||
|
MiniT2IDiffusionExtra,
|
||||||
|
HunyuanVideoDiffusionExtra>;
|
||||||
|
|
||||||
struct DiffusionParams {
|
struct DiffusionParams {
|
||||||
const sd::Tensor<float>* x = nullptr;
|
const sd::Tensor<float>* x = nullptr;
|
||||||
|
|||||||
@ -17,30 +17,38 @@ namespace Pid {
|
|||||||
constexpr float PID_PI = 3.14159265358979323846f;
|
constexpr float PID_PI = 3.14159265358979323846f;
|
||||||
|
|
||||||
struct PixelDiTConfig {
|
struct PixelDiTConfig {
|
||||||
int64_t in_channels = 3;
|
int64_t in_channels = 3;
|
||||||
int64_t hidden_size = 1536;
|
int64_t hidden_size = 1536;
|
||||||
int64_t num_groups = 24;
|
int64_t num_groups = 24;
|
||||||
int64_t patch_mlp_hidden_dim = 4096;
|
int64_t patch_mlp_hidden_dim = 4096;
|
||||||
int64_t pixel_hidden_size = 16;
|
int64_t pixel_hidden_size = 16;
|
||||||
int64_t pixel_attn_hidden_size = 1152;
|
int64_t pixel_attn_hidden_size = 1152;
|
||||||
int64_t pixel_num_groups = 16;
|
int64_t pixel_num_groups = 16;
|
||||||
int64_t patch_depth = 14;
|
int64_t patch_depth = 14;
|
||||||
int64_t pixel_depth = 2;
|
int64_t pixel_depth = 2;
|
||||||
int64_t patch_size = 16;
|
int64_t patch_size = 16;
|
||||||
int64_t txt_embed_dim = 2304;
|
int64_t txt_embed_dim = 2304;
|
||||||
int64_t txt_max_length = 300;
|
int64_t txt_max_length = 300;
|
||||||
float text_rope_theta = 10000.f;
|
float text_rope_theta = 10000.f;
|
||||||
int64_t lq_latent_channels = 16;
|
int64_t lq_latent_channels = 16;
|
||||||
int64_t lq_hidden_dim = 512;
|
int64_t lq_hidden_dim = 512;
|
||||||
int64_t lq_num_res_blocks = 4;
|
int64_t lq_num_res_blocks = 4;
|
||||||
int64_t lq_interval = 2;
|
int64_t lq_interval = 2;
|
||||||
int64_t lq_sr_scale = 4;
|
int64_t lq_sr_scale = 4;
|
||||||
int64_t lq_latent_down_factor = 8;
|
int64_t lq_latent_down_factor = 8;
|
||||||
int64_t rope_ref_grid_h = 64;
|
int64_t lq_latent_unpatchify_factor = 1;
|
||||||
int64_t rope_ref_grid_w = 64;
|
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) {
|
static PixelDiTConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
|
||||||
PixelDiTConfig config;
|
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) {
|
for (const auto& [name, tensor_storage] : tensor_storage_map) {
|
||||||
if (!starts_with(name, prefix)) {
|
if (!starts_with(name, prefix)) {
|
||||||
continue;
|
continue;
|
||||||
@ -61,20 +69,56 @@ namespace Pid {
|
|||||||
config.pixel_depth = std::max<int64_t>(config.pixel_depth, block_index + 1);
|
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) {
|
if (name == lq_prefix + "latent_proj.0.weight") {
|
||||||
config.lq_latent_channels = tensor_storage.ne[2];
|
latent_proj_in_channels = tensor_storage.ne[2];
|
||||||
config.lq_latent_down_factor = config.lq_latent_channels >= 64 ? 16 : 8;
|
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) {
|
if (name.find("patch_blocks.0.mlp_x.w1.weight") != std::string::npos) {
|
||||||
config.patch_mlp_hidden_dim = tensor_storage.ne[1];
|
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.patch_depth,
|
||||||
config.pixel_depth,
|
config.pixel_depth,
|
||||||
config.patch_mlp_hidden_dim,
|
config.patch_mlp_hidden_dim,
|
||||||
config.lq_latent_channels,
|
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;
|
return config;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@ -135,6 +179,18 @@ namespace Pid {
|
|||||||
return ggml_add(ctx, ggml_add(ctx, x, ggml_mul(ctx, x, scale)), shift);
|
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 {
|
struct PatchTokenEmbedder : public GGMLBlock {
|
||||||
bool use_rms_norm;
|
bool use_rms_norm;
|
||||||
|
|
||||||
@ -457,9 +513,9 @@ namespace Pid {
|
|||||||
struct SigmaAwareGate : public GGMLBlock {
|
struct SigmaAwareGate : public GGMLBlock {
|
||||||
int64_t dim;
|
int64_t dim;
|
||||||
|
|
||||||
SigmaAwareGate(int64_t dim)
|
SigmaAwareGate(int64_t dim, bool per_token = false)
|
||||||
: dim(dim) {
|
: 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,
|
void init_params(ggml_context* ctx,
|
||||||
@ -479,16 +535,20 @@ namespace Pid {
|
|||||||
auto alpha = ggml_exp(ctx->ggml_ctx, params["log_alpha"]);
|
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 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));
|
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 {
|
struct PiDResBlock : public GGMLBlock {
|
||||||
PiDResBlock(int64_t channels) {
|
bool replicate_padding;
|
||||||
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});
|
PiDResBlock(int64_t channels, bool replicate_padding = false)
|
||||||
blocks["block.3"] = std::make_shared<GroupNorm>(4, channels, 1e-5f);
|
: replicate_padding(replicate_padding) {
|
||||||
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});
|
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) {
|
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 norm2 = std::dynamic_pointer_cast<GroupNorm>(blocks["block.3"]);
|
||||||
auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["block.5"]);
|
auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["block.5"]);
|
||||||
auto h = ggml_silu_inplace(ctx->ggml_ctx, norm1->forward(ctx, x));
|
auto h = ggml_silu_inplace(ctx->ggml_ctx, norm1->forward(ctx, x));
|
||||||
h = conv1->forward(ctx, h);
|
if (replicate_padding) {
|
||||||
h = ggml_silu_inplace(ctx->ggml_ctx, norm2->forward(ctx, h));
|
h = replicate_pad_2d(ctx->ggml_ctx, h);
|
||||||
h = conv2->forward(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);
|
return ggml_add(ctx->ggml_ctx, x, h);
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@ -509,16 +575,23 @@ namespace Pid {
|
|||||||
|
|
||||||
LQProjection2D(const PixelDiTConfig& config)
|
LQProjection2D(const PixelDiTConfig& config)
|
||||||
: config(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});
|
int64_t unpatchify_area = config.lq_latent_unpatchify_factor * config.lq_latent_unpatchify_factor;
|
||||||
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});
|
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) {
|
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);
|
int num_outputs = static_cast<int>((config.patch_depth + config.lq_interval - 1) / config.lq_interval);
|
||||||
for (int i = 0; i < num_outputs; ++i) {
|
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["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,
|
ggml_tensor* lq_latent,
|
||||||
int64_t target_pH,
|
int64_t target_pH,
|
||||||
int64_t target_pW) {
|
int64_t target_pW) {
|
||||||
auto conv0 = std::dynamic_pointer_cast<Conv2d>(blocks["latent_proj.0"]);
|
auto conv0 = std::dynamic_pointer_cast<Conv2d>(blocks["latent_proj.0"]);
|
||||||
auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["latent_proj.2"]);
|
auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["latent_proj.2"]);
|
||||||
float z_to_patch_ratio = static_cast<float>(config.lq_sr_scale * config.lq_latent_down_factor) /
|
int64_t unpatchify_factor = config.lq_latent_unpatchify_factor;
|
||||||
|
if (unpatchify_factor > 1) {
|
||||||
|
int64_t latent_h = lq_latent->ne[1];
|
||||||
|
int64_t latent_w = lq_latent->ne[0];
|
||||||
|
lq_latent = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, lq_latent, 2, 0, 1, 3));
|
||||||
|
lq_latent = ggml_reshape_3d(ctx->ggml_ctx,
|
||||||
|
lq_latent,
|
||||||
|
lq_latent->ne[0],
|
||||||
|
lq_latent->ne[1] * lq_latent->ne[2],
|
||||||
|
lq_latent->ne[3]);
|
||||||
|
lq_latent = DiT::unpatchify(ctx->ggml_ctx,
|
||||||
|
lq_latent,
|
||||||
|
latent_h,
|
||||||
|
latent_w,
|
||||||
|
static_cast<int>(unpatchify_factor),
|
||||||
|
static_cast<int>(unpatchify_factor),
|
||||||
|
true);
|
||||||
|
}
|
||||||
|
|
||||||
|
int64_t effective_down_factor = config.lq_latent_down_factor / unpatchify_factor;
|
||||||
|
float z_to_patch_ratio = static_cast<float>(config.lq_sr_scale * effective_down_factor) /
|
||||||
static_cast<float>(config.patch_size);
|
static_cast<float>(config.patch_size);
|
||||||
GGML_ASSERT(z_to_patch_ratio >= 1.0f);
|
GGML_ASSERT(z_to_patch_ratio >= 1.0f);
|
||||||
if (lq_latent->ne[0] != target_pW || lq_latent->ne[1] != target_pH) {
|
if (lq_latent->ne[0] != target_pW || lq_latent->ne[1] != target_pH) {
|
||||||
@ -558,9 +651,15 @@ namespace Pid {
|
|||||||
GGML_SCALE_MODE_NEAREST);
|
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);
|
auto feat = conv0->forward(ctx, lq_latent);
|
||||||
feat = ggml_silu_inplace(ctx->ggml_ctx, feat);
|
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) {
|
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)]);
|
auto block = std::dynamic_pointer_cast<PiDResBlock>(blocks["latent_proj." + std::to_string(3 + i)]);
|
||||||
feat = block->forward(ctx, feat);
|
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);
|
int num_outputs = static_cast<int>((config.patch_depth + config.lq_interval - 1) / config.lq_interval);
|
||||||
std::vector<ggml_tensor*> outputs;
|
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) {
|
for (int i = 0; i < num_outputs; ++i) {
|
||||||
auto head = std::dynamic_pointer_cast<Linear>(blocks["output_heads." + std::to_string(i)]);
|
auto head = std::dynamic_pointer_cast<Linear>(blocks["output_heads." + std::to_string(i)]);
|
||||||
outputs.push_back(head->forward(ctx, tokens));
|
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;
|
return outputs;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@ -606,6 +709,9 @@ namespace Pid {
|
|||||||
}
|
}
|
||||||
blocks["final_layer"] = std::make_shared<FinalLayer>(config.pixel_hidden_size, config.in_channels);
|
blocks["final_layer"] = std::make_shared<FinalLayer>(config.pixel_hidden_size, config.in_channels);
|
||||||
blocks["lq_proj"] = std::make_shared<LQProjection2D>(config);
|
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,
|
void init_params(ggml_context* ctx,
|
||||||
@ -654,6 +760,11 @@ namespace Pid {
|
|||||||
y_emb = ggml_add(ctx->ggml_ctx, y_emb, y_pos);
|
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);
|
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);
|
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");
|
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));
|
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 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);
|
auto pixels = pixel_embedder->forward(ctx, x, config.patch_size, pixel_pos_full);
|
||||||
|
|||||||
@ -103,9 +103,9 @@ namespace Qwen {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
virtual ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||||
ggml_tensor* timesteps,
|
ggml_tensor* timesteps,
|
||||||
ggml_tensor* addition_t_cond = nullptr) {
|
ggml_tensor* addition_t_cond = nullptr) {
|
||||||
// timesteps: [N,]
|
// timesteps: [N,]
|
||||||
// return: [N, embedding_dim]
|
// return: [N, embedding_dim]
|
||||||
auto timestep_embedder = std::dynamic_pointer_cast<TimestepEmbedding>(blocks["timestep_embedder"]);
|
auto timestep_embedder = std::dynamic_pointer_cast<TimestepEmbedding>(blocks["timestep_embedder"]);
|
||||||
@ -416,10 +416,14 @@ namespace Qwen {
|
|||||||
|
|
||||||
public:
|
public:
|
||||||
QwenImageModel() {}
|
QwenImageModel() {}
|
||||||
QwenImageModel(QwenImageConfig config)
|
QwenImageModel(QwenImageConfig config,
|
||||||
|
std::shared_ptr<QwenTimestepProjEmbeddings> time_text_embed = nullptr)
|
||||||
: config(config) {
|
: config(config) {
|
||||||
int64_t inner_dim = config.num_attention_heads * config.attention_head_dim;
|
int64_t inner_dim = config.num_attention_heads * config.attention_head_dim;
|
||||||
blocks["time_text_embed"] = std::shared_ptr<GGMLBlock>(new QwenTimestepProjEmbeddings(inner_dim, config.use_additional_t_cond));
|
if (time_text_embed == nullptr) {
|
||||||
|
time_text_embed = std::make_shared<QwenTimestepProjEmbeddings>(inner_dim, config.use_additional_t_cond);
|
||||||
|
}
|
||||||
|
blocks["time_text_embed"] = std::move(time_text_embed);
|
||||||
blocks["txt_norm"] = std::shared_ptr<GGMLBlock>(new RMSNorm(config.joint_attention_dim, 1e-6f));
|
blocks["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["img_in"] = std::shared_ptr<GGMLBlock>(new Linear(config.in_channels, inner_dim));
|
||||||
blocks["txt_in"] = std::shared_ptr<GGMLBlock>(new Linear(config.joint_attention_dim, inner_dim));
|
blocks["txt_in"] = std::shared_ptr<GGMLBlock>(new Linear(config.joint_attention_dim, inner_dim));
|
||||||
|
|||||||
@ -6,6 +6,7 @@
|
|||||||
|
|
||||||
#include "model.h"
|
#include "model.h"
|
||||||
#include "model/common/block.hpp"
|
#include "model/common/block.hpp"
|
||||||
|
#include "model/diffusion/animatediff.hpp"
|
||||||
#include "model/diffusion/model.hpp"
|
#include "model/diffusion/model.hpp"
|
||||||
|
|
||||||
/*==================================================== UnetModel =====================================================*/
|
/*==================================================== UnetModel =====================================================*/
|
||||||
@ -29,6 +30,8 @@ struct UNetConfig {
|
|||||||
bool tiny_unet = false;
|
bool tiny_unet = false;
|
||||||
int model_channels = 320;
|
int model_channels = 320;
|
||||||
int adm_in_channels = 2816; // only for VERSION_SDXL/SVD
|
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,
|
static UNetConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
|
||||||
const std::string& prefix,
|
const std::string& prefix,
|
||||||
@ -84,6 +87,13 @@ struct UNetConfig {
|
|||||||
return &it->second;
|
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 (const TensorStorage* input = find_weight("input_blocks.0.0.weight")) {
|
||||||
if (input->n_dims == 4) {
|
if (input->n_dims == 4) {
|
||||||
config.in_channels = static_cast<int>(input->ne[2]);
|
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
|
blocks["out.0"] = std::shared_ptr<GGMLBlock>(new GroupNorm32(ch)); // ch == model_channels
|
||||||
// out_1 is nn.SiLU()
|
// out_1 is nn.SiLU()
|
||||||
blocks["out.2"] = std::shared_ptr<GGMLBlock>(new Conv2d(model_channels, out_channels, {3, 3}, {1, 1}, {1, 1}));
|
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,
|
ggml_tensor* resblock_forward(std::string name,
|
||||||
@ -583,6 +599,42 @@ public:
|
|||||||
|
|
||||||
ggml_set_name(h, "bench-start");
|
ggml_set_name(h, "bench-start");
|
||||||
hs.push_back(h);
|
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
|
// input block 1-11
|
||||||
size_t len_mults = channel_mult.size();
|
size_t len_mults = channel_mult.size();
|
||||||
int input_block_idx = 0;
|
int input_block_idx = 0;
|
||||||
@ -597,6 +649,7 @@ public:
|
|||||||
std::string name = "input_blocks." + std::to_string(input_block_idx) + ".1";
|
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 = 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");
|
sd::ggml_graph_cut::mark_graph_cut(h, "unet.input_blocks." + std::to_string(input_block_idx), "h");
|
||||||
hs.push_back(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 = 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 = 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");
|
sd::ggml_graph_cut::mark_graph_cut(h, "unet.middle_block", "h");
|
||||||
if (controls.size() > 0) {
|
if (controls.size() > 0) {
|
||||||
@ -660,6 +714,8 @@ public:
|
|||||||
up_sample_idx++;
|
up_sample_idx++;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
h = apply_motion_output(output_block_idx, h);
|
||||||
|
|
||||||
if (i > 0 && j == num_res_blocks) {
|
if (i > 0 && j == num_res_blocks) {
|
||||||
if (tiny_unet) {
|
if (tiny_unet) {
|
||||||
output_block_idx++;
|
output_block_idx++;
|
||||||
@ -719,14 +775,17 @@ struct UNetModelRunner : public DiffusionModelRunner {
|
|||||||
const sd::Tensor<float>& y_tensor = {},
|
const sd::Tensor<float>& y_tensor = {},
|
||||||
int num_video_frames = -1,
|
int num_video_frames = -1,
|
||||||
const std::vector<sd::Tensor<float>>& controls_tensor = {},
|
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_cgraph* gf = new_graph_custom(UNET_GRAPH_SIZE);
|
||||||
|
|
||||||
ggml_tensor* x = make_input(x_tensor);
|
ggml_tensor* x = make_input(x_tensor);
|
||||||
ggml_tensor* timesteps = make_input(timesteps_tensor);
|
ggml_tensor* timesteps = make_input(timesteps_tensor);
|
||||||
ggml_tensor* context = make_optional_input(context_tensor);
|
ggml_tensor* context = make_optional_input(context_tensor);
|
||||||
ggml_tensor* c_concat = make_optional_input(c_concat_tensor);
|
ggml_tensor* c_concat = make_optional_input(c_concat_tensor);
|
||||||
ggml_tensor* y = make_optional_input(y_tensor);
|
ggml_tensor* y = make_optional_input(y_tensor);
|
||||||
|
ggml_tensor* ip_context = make_optional_input(ip_context_tensor);
|
||||||
std::vector<ggml_tensor*> controls;
|
std::vector<ggml_tensor*> controls;
|
||||||
controls.reserve(controls_tensor.size());
|
controls.reserve(controls_tensor.size());
|
||||||
for (const auto& control_tensor : controls_tensor) {
|
for (const auto& control_tensor : controls_tensor) {
|
||||||
@ -737,7 +796,9 @@ struct UNetModelRunner : public DiffusionModelRunner {
|
|||||||
num_video_frames = static_cast<int>(x->ne[3]);
|
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,
|
ggml_tensor* out = unet.forward(&runner_ctx,
|
||||||
x,
|
x,
|
||||||
@ -762,14 +823,16 @@ struct UNetModelRunner : public DiffusionModelRunner {
|
|||||||
const sd::Tensor<float>& y = {},
|
const sd::Tensor<float>& y = {},
|
||||||
int num_video_frames = -1,
|
int num_video_frames = -1,
|
||||||
const std::vector<sd::Tensor<float>>& controls = {},
|
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]
|
// x: [N, in_channels, h, w]
|
||||||
// timesteps: [N, ]
|
// timesteps: [N, ]
|
||||||
// context: [N, max_position, hidden_size]([N, 77, 768]) or [1, max_position, hidden_size]
|
// 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]
|
// c_concat: [N, in_channels, h, w] or [1, in_channels, h, w]
|
||||||
// y: [N, adm_in_channels] or [1, adm_in_channels]
|
// y: [N, adm_in_channels] or [1, adm_in_channels]
|
||||||
auto get_graph = [&]() -> ggml_cgraph* {
|
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());
|
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),
|
tensor_or_empty(diffusion_params.y),
|
||||||
extra->num_video_frames,
|
extra->num_video_frames,
|
||||||
extra->controls ? *extra->controls : empty_controls,
|
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() {
|
void test() {
|
||||||
|
|||||||
@ -79,9 +79,20 @@ namespace LLM {
|
|||||||
int window_size = 112;
|
int window_size = 112;
|
||||||
int num_position_embeddings = 0;
|
int num_position_embeddings = 0;
|
||||||
std::set<int> fullatt_block_indexes = {7, 15, 23, 31};
|
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 {
|
struct LLMConfig {
|
||||||
LLMArch arch = LLMArch::QWEN2_5_VL;
|
LLMArch arch = LLMArch::QWEN2_5_VL;
|
||||||
int64_t num_layers = 28;
|
int64_t num_layers = 28;
|
||||||
@ -93,6 +104,7 @@ namespace LLM {
|
|||||||
bool qkv_bias = true;
|
bool qkv_bias = true;
|
||||||
bool attention_out_bias = false;
|
bool attention_out_bias = false;
|
||||||
bool qk_norm = false;
|
bool qk_norm = false;
|
||||||
|
bool final_norm = true;
|
||||||
bool rms_norm_add = false;
|
bool rms_norm_add = false;
|
||||||
bool normalize_input = false;
|
bool normalize_input = false;
|
||||||
int64_t vocab_size = 152064;
|
int64_t vocab_size = 152064;
|
||||||
@ -200,7 +212,11 @@ namespace LLM {
|
|||||||
config.vision.in_channels = tensor_storage.ne[2];
|
config.vision.in_channels = tensor_storage.ne[2];
|
||||||
config.vision.hidden_size = tensor_storage.ne[3];
|
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];
|
config.vision.hidden_size = tensor_storage.ne[0];
|
||||||
}
|
}
|
||||||
if (contains(name, "visual.pos_embed.weight")) {
|
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) {
|
if ((arch == LLMArch::QWEN3 || arch == LLMArch::QWEN3_VL) && config.num_layers == 28) {
|
||||||
config.num_heads = 16;
|
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) {
|
if (detected_vision_layers > 0) {
|
||||||
config.vision.num_layers = detected_vision_layers;
|
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,
|
LOG_DEBUG("llm: num_layers = %" PRId64 ", vocab_size = %" PRId64 ", hidden_size = %" PRId64 ", intermediate_size = %" PRId64,
|
||||||
config.num_layers,
|
config.num_layers,
|
||||||
config.vocab_size,
|
config.vocab_size,
|
||||||
@ -285,7 +312,7 @@ namespace LLM {
|
|||||||
bool add_unit_offset = false)
|
bool add_unit_offset = false)
|
||||||
: hidden_size(hidden_size), eps(eps), add_unit_offset(add_unit_offset) {}
|
: 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"];
|
ggml_tensor* w = params["weight"];
|
||||||
if (ctx->weight_adapter) {
|
if (ctx->weight_adapter) {
|
||||||
w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");
|
w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");
|
||||||
@ -537,6 +564,37 @@ namespace LLM {
|
|||||||
return input_embed;
|
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 {
|
struct VisionMLP : public GGMLBlock {
|
||||||
protected:
|
protected:
|
||||||
LLMVisionArch arch_;
|
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 {
|
struct VisionAttention : public GGMLBlock {
|
||||||
protected:
|
protected:
|
||||||
bool llama_cpp_style;
|
bool llama_cpp_style;
|
||||||
@ -840,6 +925,7 @@ namespace LLM {
|
|||||||
int spatial_merge_size;
|
int spatial_merge_size;
|
||||||
int num_grid_per_side;
|
int num_grid_per_side;
|
||||||
std::set<int> fullatt_block_indexes;
|
std::set<int> fullatt_block_indexes;
|
||||||
|
std::vector<int> deepstack_visual_indexes;
|
||||||
|
|
||||||
public:
|
public:
|
||||||
VisionModel(bool llama_cpp_style,
|
VisionModel(bool llama_cpp_style,
|
||||||
@ -849,7 +935,8 @@ namespace LLM {
|
|||||||
num_layers(vision_params.num_layers),
|
num_layers(vision_params.num_layers),
|
||||||
spatial_merge_size(vision_params.spatial_merge_size),
|
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),
|
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,
|
blocks["patch_embed"] = std::shared_ptr<GGMLBlock>(new VisionPatchEmbed(vision_params.split_patch_embed,
|
||||||
arch_,
|
arch_,
|
||||||
vision_params.patch_size,
|
vision_params.patch_size,
|
||||||
@ -871,6 +958,11 @@ namespace LLM {
|
|||||||
vision_params.out_hidden_size,
|
vision_params.out_hidden_size,
|
||||||
vision_params.hidden_size,
|
vision_params.hidden_size,
|
||||||
spatial_merge_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() {
|
std::shared_ptr<Embedding> pos_embedder() {
|
||||||
@ -889,13 +981,13 @@ namespace LLM {
|
|||||||
return spatial_merge_size;
|
return spatial_merge_size;
|
||||||
}
|
}
|
||||||
|
|
||||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
std::vector<ggml_tensor*> forward_outputs(GGMLRunnerContext* ctx,
|
||||||
ggml_tensor* pixel_values,
|
ggml_tensor* pixel_values,
|
||||||
ggml_tensor* pe,
|
ggml_tensor* pe,
|
||||||
ggml_tensor* window_index,
|
ggml_tensor* window_index,
|
||||||
ggml_tensor* window_inverse_index,
|
ggml_tensor* window_inverse_index,
|
||||||
ggml_tensor* window_mask,
|
ggml_tensor* window_mask,
|
||||||
ggml_tensor* pos_embeds = nullptr) {
|
ggml_tensor* pos_embeds = nullptr) {
|
||||||
// pixel_values: [grid_t*(H/mh/ph)*(W/mw/pw)*mh*mw, C*pt*ph*pw]
|
// 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_index: [grid_t*(H/mh/ph)*(W/mw/pw)]
|
||||||
// window_inverse_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]);
|
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++) {
|
for (int i = 0; i < num_layers; i++) {
|
||||||
auto block = std::dynamic_pointer_cast<VisionBlock>(blocks["blocks." + std::to_string(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()) {
|
if (fullatt_block_indexes.find(i) != fullatt_block_indexes.end()) {
|
||||||
mask = nullptr;
|
mask = nullptr;
|
||||||
}
|
}
|
||||||
x = block->forward(ctx, x, pe, mask);
|
x = block->forward(ctx, x, pe, mask);
|
||||||
if (i == 0) {
|
auto deepstack_it = std::find(deepstack_visual_indexes.begin(), deepstack_visual_indexes.end(), i);
|
||||||
|
if (deepstack_it != deepstack_visual_indexes.end()) {
|
||||||
|
size_t deepstack_index = static_cast<size_t>(std::distance(deepstack_visual_indexes.begin(), deepstack_it));
|
||||||
|
auto deepstack_merger = std::dynamic_pointer_cast<Qwen3VLDeepStackMerger>(blocks["deepstack_merger_list." + std::to_string(deepstack_index)]);
|
||||||
|
deepstack_outputs.push_back(deepstack_merger->forward(ctx, x));
|
||||||
}
|
}
|
||||||
sd::ggml_graph_cut::mark_graph_cut(x, "llm.vision.blocks." + std::to_string(i), "x");
|
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);
|
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++) {
|
for (int i = 0; i < num_layers; i++) {
|
||||||
blocks["layers." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new TransformerBlock(config, 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,
|
ggml_tensor* embed(GGMLRunnerContext* ctx,
|
||||||
@ -1274,9 +1385,11 @@ namespace LLM {
|
|||||||
ggml_tensor* input_pos,
|
ggml_tensor* input_pos,
|
||||||
ggml_tensor* attention_mask,
|
ggml_tensor* attention_mask,
|
||||||
std::set<int> out_layers,
|
std::set<int> out_layers,
|
||||||
ggml_tensor* sliding_attention_mask = nullptr,
|
const std::vector<std::vector<std::pair<int, ggml_tensor*>>>& deepstack_image_embeds = {},
|
||||||
bool return_all_hidden_states = false) {
|
ggml_tensor* sliding_attention_mask = nullptr,
|
||||||
auto norm = std::dynamic_pointer_cast<LLMRMSNorm>(blocks["norm"]);
|
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;
|
std::vector<ggml_tensor*> intermediate_outputs;
|
||||||
|
|
||||||
if (config.normalize_input) {
|
if (config.normalize_input) {
|
||||||
@ -1291,6 +1404,9 @@ namespace LLM {
|
|||||||
auto block = std::dynamic_pointer_cast<TransformerBlock>(blocks["layers." + std::to_string(i)]);
|
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);
|
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) {
|
if (return_all_hidden_states || out_layers.size() > 1) {
|
||||||
x = ggml_cont(ctx->ggml_ctx, x);
|
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) {
|
if (return_all_hidden_states) {
|
||||||
intermediate_outputs.push_back(normed_x);
|
intermediate_outputs.push_back(normed_x);
|
||||||
x = intermediate_outputs[0];
|
x = intermediate_outputs[0];
|
||||||
@ -1332,6 +1448,7 @@ namespace LLM {
|
|||||||
ggml_tensor* attention_mask,
|
ggml_tensor* attention_mask,
|
||||||
ggml_tensor* sliding_attention_mask,
|
ggml_tensor* sliding_attention_mask,
|
||||||
std::vector<std::pair<int, ggml_tensor*>> image_embeds,
|
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,
|
std::set<int> out_layers,
|
||||||
bool return_all_hidden_states = false) {
|
bool return_all_hidden_states = false) {
|
||||||
// input_ids: [N, n_token]
|
// input_ids: [N, n_token]
|
||||||
@ -1343,6 +1460,7 @@ namespace LLM {
|
|||||||
input_pos,
|
input_pos,
|
||||||
attention_mask,
|
attention_mask,
|
||||||
std::move(out_layers),
|
std::move(out_layers),
|
||||||
|
deepstack_image_embeds,
|
||||||
sliding_attention_mask,
|
sliding_attention_mask,
|
||||||
return_all_hidden_states);
|
return_all_hidden_states);
|
||||||
}
|
}
|
||||||
@ -1368,6 +1486,7 @@ namespace LLM {
|
|||||||
ggml_tensor* attention_mask,
|
ggml_tensor* attention_mask,
|
||||||
ggml_tensor* sliding_attention_mask,
|
ggml_tensor* sliding_attention_mask,
|
||||||
std::vector<std::pair<int, ggml_tensor*>> image_embeds,
|
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,
|
std::set<int> out_layers,
|
||||||
bool return_all_hidden_states = false) {
|
bool return_all_hidden_states = false) {
|
||||||
// input_ids: [N, n_token]
|
// input_ids: [N, n_token]
|
||||||
@ -1379,6 +1498,7 @@ namespace LLM {
|
|||||||
attention_mask,
|
attention_mask,
|
||||||
sliding_attention_mask,
|
sliding_attention_mask,
|
||||||
image_embeds,
|
image_embeds,
|
||||||
|
deepstack_image_embeds,
|
||||||
out_layers,
|
out_layers,
|
||||||
return_all_hidden_states);
|
return_all_hidden_states);
|
||||||
return x;
|
return x;
|
||||||
@ -1520,7 +1640,8 @@ namespace LLM {
|
|||||||
std::vector<float>& window_mask_vec,
|
std::vector<float>& window_mask_vec,
|
||||||
std::vector<float>& pe_vec,
|
std::vector<float>& pe_vec,
|
||||||
std::array<std::vector<int32_t>, 4>& pos_embed_idx_data,
|
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[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);
|
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);
|
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);
|
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());
|
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;
|
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);
|
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());
|
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:
|
public:
|
||||||
@ -1653,12 +1782,17 @@ namespace LLM {
|
|||||||
model.get_param_tensors(tensors, prefix);
|
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* forward(GGMLRunnerContext* ctx,
|
||||||
ggml_tensor* input_ids,
|
ggml_tensor* input_ids,
|
||||||
ggml_tensor* input_pos,
|
ggml_tensor* input_pos,
|
||||||
ggml_tensor* attention_mask,
|
ggml_tensor* attention_mask,
|
||||||
ggml_tensor* sliding_attention_mask,
|
ggml_tensor* sliding_attention_mask,
|
||||||
std::vector<std::pair<int, ggml_tensor*>> image_embeds,
|
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,
|
std::set<int> out_layers,
|
||||||
bool return_all_hidden_states = false) {
|
bool return_all_hidden_states = false) {
|
||||||
auto hidden_states = model.forward(ctx,
|
auto hidden_states = model.forward(ctx,
|
||||||
@ -1667,6 +1801,7 @@ namespace LLM {
|
|||||||
attention_mask,
|
attention_mask,
|
||||||
sliding_attention_mask,
|
sliding_attention_mask,
|
||||||
image_embeds,
|
image_embeds,
|
||||||
|
deepstack_image_embeds,
|
||||||
out_layers,
|
out_layers,
|
||||||
return_all_hidden_states); // [N, n_token, hidden_size]
|
return_all_hidden_states); // [N, n_token, hidden_size]
|
||||||
return hidden_states;
|
return hidden_states;
|
||||||
@ -1685,7 +1820,9 @@ namespace LLM {
|
|||||||
|
|
||||||
ggml_cgraph* build_graph(const sd::Tensor<int32_t>& input_ids_tensor,
|
ggml_cgraph* build_graph(const sd::Tensor<int32_t>& input_ids_tensor,
|
||||||
const sd::Tensor<float>& attention_mask_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,
|
std::set<int> out_layers,
|
||||||
bool return_all_hidden_states = false) {
|
bool return_all_hidden_states = false) {
|
||||||
ggml_cgraph* gf = new_graph_custom(LLM_GRAPH_SIZE);
|
ggml_cgraph* gf = new_graph_custom(LLM_GRAPH_SIZE);
|
||||||
@ -1696,6 +1833,13 @@ namespace LLM {
|
|||||||
ggml_tensor* embed = make_input(embed_tensor);
|
ggml_tensor* embed = make_input(embed_tensor);
|
||||||
image_embeds.emplace_back(idx, embed);
|
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];
|
int64_t n_tokens = input_ids->ne[0];
|
||||||
if (config.arch == LLMArch::MISTRAL_SMALL_3_2 ||
|
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[2 * n_tokens + i] = i;
|
||||||
input_pos_vec[3 * n_tokens + i] = 0;
|
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,
|
auto input_pos = ggml_new_tensor_1d(compute_ctx,
|
||||||
@ -1773,6 +1941,7 @@ namespace LLM {
|
|||||||
attention_mask,
|
attention_mask,
|
||||||
sliding_attention_mask,
|
sliding_attention_mask,
|
||||||
image_embeds,
|
image_embeds,
|
||||||
|
deepstack_image_embeds,
|
||||||
out_layers,
|
out_layers,
|
||||||
return_all_hidden_states);
|
return_all_hidden_states);
|
||||||
|
|
||||||
@ -1784,16 +1953,20 @@ namespace LLM {
|
|||||||
sd::Tensor<float> compute(const int n_threads,
|
sd::Tensor<float> compute(const int n_threads,
|
||||||
const sd::Tensor<int32_t>& input_ids,
|
const sd::Tensor<int32_t>& input_ids,
|
||||||
const sd::Tensor<float>& attention_mask,
|
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,
|
std::set<int> out_layers,
|
||||||
bool return_all_hidden_states = false,
|
bool return_all_hidden_states = false,
|
||||||
bool auto_free = true,
|
bool auto_free = true,
|
||||||
bool free_compute_buffer = true,
|
bool free_compute_buffer = true,
|
||||||
bool free_compute_params = true) {
|
bool free_compute_params = true,
|
||||||
|
const DeepStackImageEmbeds& deepstack_image_embeds = {},
|
||||||
|
const std::vector<ImageGrid>& image_grids = {}) {
|
||||||
auto get_graph = [&]() -> ggml_cgraph* {
|
auto get_graph = [&]() -> ggml_cgraph* {
|
||||||
return build_graph(input_ids,
|
return build_graph(input_ids,
|
||||||
attention_mask,
|
attention_mask,
|
||||||
image_embeds,
|
image_embeds,
|
||||||
|
deepstack_image_embeds,
|
||||||
|
image_grids,
|
||||||
out_layers,
|
out_layers,
|
||||||
return_all_hidden_states);
|
return_all_hidden_states);
|
||||||
};
|
};
|
||||||
@ -1843,6 +2016,24 @@ namespace LLM {
|
|||||||
pos_embed_weight_data_);
|
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* build_encode_image_graph(const sd::Tensor<float>& image_tensor) {
|
||||||
ggml_cgraph* gf = new_graph_custom(LLM_GRAPH_SIZE);
|
ggml_cgraph* gf = new_graph_custom(LLM_GRAPH_SIZE);
|
||||||
ggml_tensor* image = make_input(image_tensor);
|
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));
|
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 {
|
struct LLMEmbedder {
|
||||||
|
|||||||
@ -18,6 +18,7 @@
|
|||||||
struct T5Config {
|
struct T5Config {
|
||||||
int64_t num_layers = 24;
|
int64_t num_layers = 24;
|
||||||
int64_t model_dim = 4096;
|
int64_t model_dim = 4096;
|
||||||
|
int64_t inner_dim = 4096;
|
||||||
int64_t ff_dim = 10240;
|
int64_t ff_dim = 10240;
|
||||||
int64_t num_heads = 64;
|
int64_t num_heads = 64;
|
||||||
int64_t vocab_size = 32128;
|
int64_t vocab_size = 32128;
|
||||||
@ -53,6 +54,7 @@ struct T5Config {
|
|||||||
if (q->n_dims == 2) {
|
if (q->n_dims == 2) {
|
||||||
config.model_dim = q->ne[0];
|
config.model_dim = q->ne[0];
|
||||||
int64_t inner_dim = q->ne[1];
|
int64_t inner_dim = q->ne[1];
|
||||||
|
config.inner_dim = inner_dim;
|
||||||
// Flan-T5/T5 uses d_kv=64 for common sizes.
|
// Flan-T5/T5 uses d_kv=64 for common sizes.
|
||||||
if (inner_dim % 64 == 0) {
|
if (inner_dim % 64 == 0) {
|
||||||
config.num_heads = inner_dim / 64;
|
config.num_heads = inner_dim / 64;
|
||||||
@ -357,7 +359,7 @@ public:
|
|||||||
: config(config) {
|
: config(config) {
|
||||||
blocks["encoder"] = std::shared_ptr<GGMLBlock>(new T5Stack(config.num_layers,
|
blocks["encoder"] = std::shared_ptr<GGMLBlock>(new T5Stack(config.num_layers,
|
||||||
config.model_dim,
|
config.model_dim,
|
||||||
config.model_dim,
|
config.inner_dim,
|
||||||
config.ff_dim,
|
config.ff_dim,
|
||||||
config.num_heads,
|
config.num_heads,
|
||||||
config.relative_attention));
|
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 <vector>
|
||||||
|
|
||||||
#include "core/ggml_extend.hpp"
|
#include "core/ggml_extend.hpp"
|
||||||
|
#include "model/vae/audio_vae.hpp"
|
||||||
#include "model_loader.h"
|
#include "model_loader.h"
|
||||||
#include "model_manager.h"
|
#include "model_manager.h"
|
||||||
|
|
||||||
@ -996,7 +997,7 @@ namespace LTXV {
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
struct LTXAudioVAERunner : public GGMLRunner {
|
struct LTXAudioVAERunner : public AudioVAERunner {
|
||||||
LTXAudioVAEConfig config;
|
LTXAudioVAEConfig config;
|
||||||
LTXAudioVAE model;
|
LTXAudioVAE model;
|
||||||
std::string weight_prefix;
|
std::string weight_prefix;
|
||||||
@ -1006,7 +1007,7 @@ namespace LTXV {
|
|||||||
const String2TensorStorage& tensor_storage_map,
|
const String2TensorStorage& tensor_storage_map,
|
||||||
const std::string& prefix = "",
|
const std::string& prefix = "",
|
||||||
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
|
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
|
||||||
: GGMLRunner(backend, weight_manager),
|
: AudioVAERunner(backend, weight_manager),
|
||||||
weight_prefix(prefix),
|
weight_prefix(prefix),
|
||||||
config(LTXAudioVAEConfig::detect_from_weights(tensor_storage_map)),
|
config(LTXAudioVAEConfig::detect_from_weights(tensor_storage_map)),
|
||||||
model(config) {
|
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);
|
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();
|
return model.get_params_mem_size();
|
||||||
}
|
}
|
||||||
|
|
||||||
std::string get_desc() {
|
std::string get_desc() override {
|
||||||
return "ltx_audio_vae";
|
return "ltx_audio_vae";
|
||||||
}
|
}
|
||||||
|
|
||||||
sd::Tensor<float> decode(int n_threads,
|
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();
|
int64_t t0 = ggml_time_ms();
|
||||||
auto get_graph = [&]() -> ggml_cgraph* {
|
auto get_graph = [&]() -> ggml_cgraph* {
|
||||||
auto latent = make_input(latent_tensor);
|
auto latent = make_input(latent_tensor);
|
||||||
@ -1047,6 +1048,10 @@ namespace LTXV {
|
|||||||
return result;
|
return result;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
int output_sample_rate() const override {
|
||||||
|
return config.output_sample_rate();
|
||||||
|
}
|
||||||
|
|
||||||
void test(const std::string& input_path) {
|
void test(const std::string& input_path) {
|
||||||
auto z = sd::load_tensor_from_file_as_tensor<float>(input_path);
|
auto z = sd::load_tensor_from_file_as_tensor<float>(input_path);
|
||||||
GGML_ASSERT(!z.empty());
|
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) {
|
if (version == VERSION_WAN2_2_TI2V) {
|
||||||
z_channels = 48;
|
z_channels = 48;
|
||||||
patch = 2;
|
patch = 2;
|
||||||
|
} else if (sd_version_is_hunyuan_video(version)) {
|
||||||
|
z_channels = 32;
|
||||||
|
patch = 2;
|
||||||
} else if (sd_version_is_ltxav(version)) {
|
} else if (sd_version_is_ltxav(version)) {
|
||||||
z_channels = 128;
|
z_channels = 128;
|
||||||
patch = 4;
|
patch = 4;
|
||||||
@ -542,12 +545,12 @@ public:
|
|||||||
|
|
||||||
ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* z) {
|
ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* z) {
|
||||||
auto decoder = std::dynamic_pointer_cast<TinyVideoDecoder>(blocks["decoder"]);
|
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)
|
// (W, H, C, T) -> (W, H, T, C)
|
||||||
z = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, z, 0, 1, 3, 2));
|
z = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, z, 0, 1, 3, 2));
|
||||||
}
|
}
|
||||||
auto result = decoder->forward(ctx, z);
|
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)
|
// (W, H, T, C) -> (W, H, C, T)
|
||||||
result = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, result, 0, 1, 3, 2));
|
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) {
|
ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||||
auto encoder = std::dynamic_pointer_cast<TinyVideoEncoder>(blocks["encoder"]);
|
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)
|
// (W, H, T, C) -> (W, H, C, T)
|
||||||
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
|
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);
|
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)
|
// (W, H, C, T) -> (W, H, T, C)
|
||||||
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
|
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;
|
int scale_factor = 8;
|
||||||
if (version == VERSION_LTXAV) {
|
if (version == VERSION_LTXAV) {
|
||||||
scale_factor = 32;
|
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;
|
scale_factor = 16;
|
||||||
} else if (sd_version_uses_flux2_vae(version)) {
|
} else if (sd_version_uses_flux2_vae(version)) {
|
||||||
scale_factor = 16;
|
scale_factor = 16;
|
||||||
@ -115,11 +115,11 @@ public:
|
|||||||
tile_size_y = get_tile_size(params.tile_size_y, params.rel_size_y, latent_y);
|
tile_size_y = get_tile_size(params.tile_size_y, params.rel_size_y, latent_y);
|
||||||
}
|
}
|
||||||
|
|
||||||
sd::Tensor<float> encode(int n_threads,
|
virtual sd::Tensor<float> encode(int n_threads,
|
||||||
const sd::Tensor<float>& x,
|
const sd::Tensor<float>& x,
|
||||||
sd_tiling_params_t tiling_params,
|
sd_tiling_params_t tiling_params,
|
||||||
bool circular_x = false,
|
bool circular_x = false,
|
||||||
bool circular_y = false) {
|
bool circular_y = false) {
|
||||||
int64_t t0 = ggml_time_ms();
|
int64_t t0 = ggml_time_ms();
|
||||||
sd::Tensor<float> input = x;
|
sd::Tensor<float> input = x;
|
||||||
sd::Tensor<float> output;
|
sd::Tensor<float> output;
|
||||||
@ -136,7 +136,8 @@ public:
|
|||||||
// Image VAE encode is more sensitive to tile boundary context than decode.
|
// Image VAE encode is more sensitive to tile boundary context than decode.
|
||||||
// Keep the smaller legacy factor for video VAEs, but default image encode
|
// 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.
|
// 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);
|
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);
|
LOG_DEBUG("VAE Tile size: %dx%d", tile_size_x, tile_size_y);
|
||||||
output = tiled_compute(input,
|
output = tiled_compute(input,
|
||||||
@ -166,13 +167,13 @@ public:
|
|||||||
return std::move(output);
|
return std::move(output);
|
||||||
}
|
}
|
||||||
|
|
||||||
sd::Tensor<float> decode(int n_threads,
|
virtual sd::Tensor<float> decode(int n_threads,
|
||||||
const sd::Tensor<float>& x,
|
const sd::Tensor<float>& x,
|
||||||
sd_tiling_params_t tiling_params,
|
sd_tiling_params_t tiling_params,
|
||||||
bool decode_video = false,
|
bool decode_video = false,
|
||||||
bool circular_x = false,
|
bool circular_x = false,
|
||||||
bool circular_y = false,
|
bool circular_y = false,
|
||||||
bool silent = false) {
|
bool silent = false) {
|
||||||
int64_t t0 = ggml_time_ms();
|
int64_t t0 = ggml_time_ms();
|
||||||
sd::Tensor<float> input = x;
|
sd::Tensor<float> input = x;
|
||||||
sd::Tensor<float> output;
|
sd::Tensor<float> output;
|
||||||
|
|||||||
@ -21,7 +21,7 @@ bool write_gguf_file(const std::string& file_path,
|
|||||||
class GGUFStreamingWriter : public StreamingModelWriter {
|
class GGUFStreamingWriter : public StreamingModelWriter {
|
||||||
public:
|
public:
|
||||||
GGUFStreamingWriter() = default;
|
GGUFStreamingWriter() = default;
|
||||||
~GGUFStreamingWriter();
|
~GGUFStreamingWriter() override;
|
||||||
|
|
||||||
bool write_metadata(const std::string& file_path,
|
bool write_metadata(const std::string& file_path,
|
||||||
const std::vector<TensorWritePlan>& tensors,
|
const std::vector<TensorWritePlan>& tensors,
|
||||||
|
|||||||
@ -2,6 +2,7 @@
|
|||||||
|
|
||||||
#include <cstdlib>
|
#include <cstdlib>
|
||||||
#include <cstring>
|
#include <cstring>
|
||||||
|
#include <limits>
|
||||||
#include <string>
|
#include <string>
|
||||||
#include <unordered_map>
|
#include <unordered_map>
|
||||||
#include <utility>
|
#include <utility>
|
||||||
@ -512,8 +513,51 @@ static bool parse_storage_type(const std::string& global_name, PickleStorageInfo
|
|||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
|
|
||||||
static bool tensor_is_contiguous(const PickleTensorInfo& tensor) {
|
static bool checked_pickle_byte_count(int64_t element_count,
|
||||||
if (tensor.tensor_storage.nelements() == 0) {
|
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;
|
return true;
|
||||||
}
|
}
|
||||||
if (tensor.stride_n_dims != tensor.tensor_storage.n_dims) {
|
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)) {
|
if (storage.key.empty() || !parse_storage_type(pid.items[1].str_value, &storage)) {
|
||||||
return false;
|
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;
|
storage_nbytes[storage.key] = storage.nbytes;
|
||||||
stack.push_back(make_storage_value(storage));
|
stack.push_back(make_storage_value(storage));
|
||||||
} break;
|
} 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_f64 = args.items[0].storage.is_f64;
|
||||||
tensor.tensor_storage.is_i64 = args.items[0].storage.is_i64;
|
tensor.tensor_storage.is_i64 = args.items[0].storage.is_i64;
|
||||||
tensor.tensor_storage.storage_key = args.items[0].storage.key;
|
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) {
|
for (const auto& item : args.items[2].items) {
|
||||||
if (item.kind != PickleValue::INT || tensor.tensor_storage.n_dims >= SD_MAX_DIMS) {
|
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;
|
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;
|
return false;
|
||||||
}
|
}
|
||||||
stack.push_back(make_tensor_value(tensor));
|
stack.push_back(make_tensor_value(tensor));
|
||||||
|
|||||||
@ -7,6 +7,7 @@
|
|||||||
#include <fstream>
|
#include <fstream>
|
||||||
#include <ostream>
|
#include <ostream>
|
||||||
#include <string>
|
#include <string>
|
||||||
|
#include <unordered_map>
|
||||||
#include <unordered_set>
|
#include <unordered_set>
|
||||||
#include <vector>
|
#include <vector>
|
||||||
|
|
||||||
@ -93,14 +94,57 @@ static ggml_type safetensors_dtype_to_ggml_type(const std::string& dtype) {
|
|||||||
ttype = GGML_TYPE_I32;
|
ttype = GGML_TYPE_I32;
|
||||||
} else if (dtype == "I64") {
|
} else if (dtype == "I64") {
|
||||||
ttype = GGML_TYPE_I32;
|
ttype = GGML_TYPE_I32;
|
||||||
|
} else if (dtype == "I8") {
|
||||||
|
ttype = GGML_TYPE_I8;
|
||||||
}
|
}
|
||||||
return ttype;
|
return ttype;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
struct ComfyQuantConfig {
|
||||||
|
std::string format;
|
||||||
|
bool convrot = false;
|
||||||
|
int group_size = 0;
|
||||||
|
};
|
||||||
|
|
||||||
|
static bool read_comfy_quant_config(std::ifstream& file,
|
||||||
|
const std::string& file_path,
|
||||||
|
const std::string& tensor_name,
|
||||||
|
size_t offset,
|
||||||
|
size_t size,
|
||||||
|
ComfyQuantConfig& config,
|
||||||
|
std::string* error) {
|
||||||
|
static constexpr size_t MAX_COMFY_QUANT_CONFIG_SIZE = 64 * 1024;
|
||||||
|
if (size == 0 || size > MAX_COMFY_QUANT_CONFIG_SIZE) {
|
||||||
|
set_error(error, "invalid ComfyUI quantization metadata tensor '" + tensor_name + "' in '" + file_path + "'");
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
|
||||||
|
std::vector<char> data(size + 1, '\0');
|
||||||
|
file.clear();
|
||||||
|
file.seekg((std::streamoff)offset, std::ios::beg);
|
||||||
|
file.read(data.data(), (std::streamsize)size);
|
||||||
|
if (!file) {
|
||||||
|
set_error(error, "read ComfyUI quantization metadata tensor failed: '" + tensor_name + "'");
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
|
||||||
|
try {
|
||||||
|
const nlohmann::json json = nlohmann::json::parse(data.data(), data.data() + size);
|
||||||
|
config.format = json.value("format", "");
|
||||||
|
config.convrot = json.value("convrot", false);
|
||||||
|
config.group_size = json.value("convrot_groupsize", 0);
|
||||||
|
} catch (const std::exception&) {
|
||||||
|
set_error(error, "parsing ComfyUI quantization metadata tensor failed: '" + tensor_name + "'");
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
|
||||||
// https://huggingface.co/docs/safetensors/index
|
// https://huggingface.co/docs/safetensors/index
|
||||||
bool read_safetensors_file(const std::string& file_path,
|
bool read_safetensors_file(const std::string& file_path,
|
||||||
std::vector<TensorStorage>& tensor_storages,
|
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);
|
std::ifstream file(file_path, std::ios::binary);
|
||||||
if (!file.is_open()) {
|
if (!file.is_open()) {
|
||||||
set_error(error, "failed to open '" + file_path + "'");
|
set_error(error, "failed to open '" + file_path + "'");
|
||||||
@ -150,6 +194,45 @@ bool read_safetensors_file(const std::string& file_path,
|
|||||||
return false;
|
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>());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
std::unordered_map<std::string, ComfyQuantConfig> comfy_quant_configs;
|
||||||
|
for (const auto& item : header_.items()) {
|
||||||
|
const std::string& name = item.key();
|
||||||
|
if (name == "__metadata__" || !ends_with(name, ".comfy_quant")) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
const nlohmann::json& tensor_info = item.value();
|
||||||
|
if (tensor_info.value("dtype", "") != "U8") {
|
||||||
|
set_error(error, "invalid dtype for ComfyUI quantization metadata tensor '" + name + "'");
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
const size_t begin = tensor_info["data_offsets"][0].get<size_t>();
|
||||||
|
const size_t end = tensor_info["data_offsets"][1].get<size_t>();
|
||||||
|
if (begin > end || end > file_size_ - data_start) {
|
||||||
|
set_error(error, "data offsets out of bounds for tensor '" + name + "'");
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
|
||||||
|
ComfyQuantConfig config;
|
||||||
|
if (!read_comfy_quant_config(file, file_path, name, data_start + begin, end - begin, config, error)) {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
const std::string module_name = name.substr(0, name.size() - std::string(".comfy_quant").size());
|
||||||
|
comfy_quant_configs.emplace(module_name, std::move(config));
|
||||||
|
}
|
||||||
|
|
||||||
tensor_storages.clear();
|
tensor_storages.clear();
|
||||||
for (auto& item : header_.items()) {
|
for (auto& item : header_.items()) {
|
||||||
std::string name = item.key();
|
std::string name = item.key();
|
||||||
@ -207,6 +290,39 @@ bool read_safetensors_file(const std::string& file_path,
|
|||||||
TensorStorage tensor_storage(name, type, ne, n_dims, 0, data_start + begin);
|
TensorStorage tensor_storage(name, type, ne, n_dims, 0, data_start + begin);
|
||||||
tensor_storage.reverse_ne();
|
tensor_storage.reverse_ne();
|
||||||
|
|
||||||
|
if (ends_with(name, ".weight")) {
|
||||||
|
const std::string module_name = name.substr(0, name.size() - std::string(".weight").size());
|
||||||
|
auto config = comfy_quant_configs.find(module_name);
|
||||||
|
if (config != comfy_quant_configs.end() && config->second.format == "int8_tensorwise") {
|
||||||
|
if (type != GGML_TYPE_I8) {
|
||||||
|
set_error(error, "ComfyUI int8_tensorwise weight is not I8: '" + name + "'");
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
if (config->second.convrot) {
|
||||||
|
int group_size_remainder = config->second.group_size;
|
||||||
|
while (group_size_remainder > 1 && group_size_remainder % 4 == 0) {
|
||||||
|
group_size_remainder /= 4;
|
||||||
|
}
|
||||||
|
if (group_size_remainder != 1 || tensor_storage.ne[0] % config->second.group_size != 0) {
|
||||||
|
set_error(error, "invalid ComfyUI convrot group size for tensor '" + name + "'");
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
tensor_storage.is_int8_tensorwise = true;
|
||||||
|
tensor_storage.int8_convrot = config->second.convrot;
|
||||||
|
tensor_storage.int8_convrot_group_size = config->second.group_size;
|
||||||
|
}
|
||||||
|
} else if (ends_with(name, ".weight_scale")) {
|
||||||
|
const std::string module_name = name.substr(0, name.size() - std::string(".weight_scale").size());
|
||||||
|
auto config = comfy_quant_configs.find(module_name);
|
||||||
|
if (config != comfy_quant_configs.end() && config->second.format == "int8_tensorwise" &&
|
||||||
|
tensor_storage.n_dims == 2 && tensor_storage.ne[0] == 1) {
|
||||||
|
tensor_storage.ne[0] = tensor_storage.ne[1];
|
||||||
|
tensor_storage.ne[1] = 1;
|
||||||
|
tensor_storage.n_dims = 1;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
size_t tensor_data_size = end - begin;
|
size_t tensor_data_size = end - begin;
|
||||||
|
|
||||||
bool tensor_size_ok;
|
bool tensor_size_ok;
|
||||||
|
|||||||
@ -1,6 +1,7 @@
|
|||||||
#ifndef __SD_MODEL_IO_SAFETENSORS_IO_H__
|
#ifndef __SD_MODEL_IO_SAFETENSORS_IO_H__
|
||||||
#define __SD_MODEL_IO_SAFETENSORS_IO_H__
|
#define __SD_MODEL_IO_SAFETENSORS_IO_H__
|
||||||
|
|
||||||
|
#include <map>
|
||||||
#include <string>
|
#include <string>
|
||||||
#include <vector>
|
#include <vector>
|
||||||
|
|
||||||
@ -10,7 +11,8 @@
|
|||||||
bool is_safetensors_file(const std::string& file_path);
|
bool is_safetensors_file(const std::string& file_path);
|
||||||
bool read_safetensors_file(const std::string& file_path,
|
bool read_safetensors_file(const std::string& file_path,
|
||||||
std::vector<TensorStorage>& tensor_storages,
|
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,
|
bool read_safetensors_index_file(const std::string& file_path,
|
||||||
std::vector<std::string>& shard_paths,
|
std::vector<std::string>& shard_paths,
|
||||||
std::string* error = nullptr);
|
std::string* error = nullptr);
|
||||||
|
|||||||
@ -15,14 +15,17 @@
|
|||||||
|
|
||||||
struct TensorStorage {
|
struct TensorStorage {
|
||||||
std::string name;
|
std::string name;
|
||||||
ggml_type type = GGML_TYPE_F32;
|
ggml_type type = GGML_TYPE_F32;
|
||||||
ggml_type expected_type = GGML_TYPE_COUNT;
|
ggml_type expected_type = GGML_TYPE_COUNT;
|
||||||
bool is_f8_e4m3 = false;
|
bool is_f8_e4m3 = false;
|
||||||
bool is_f8_e5m2 = false;
|
bool is_f8_e5m2 = false;
|
||||||
bool is_f64 = false;
|
bool is_f64 = false;
|
||||||
bool is_i64 = false;
|
bool is_i64 = false;
|
||||||
int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
|
bool is_int8_tensorwise = false;
|
||||||
int n_dims = 0;
|
bool int8_convrot = false;
|
||||||
|
int int8_convrot_group_size = 0;
|
||||||
|
int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1};
|
||||||
|
int n_dims = 0;
|
||||||
|
|
||||||
std::string storage_key;
|
std::string storage_key;
|
||||||
size_t file_index = 0;
|
size_t file_index = 0;
|
||||||
|
|||||||
@ -139,11 +139,16 @@ bool read_torch_legacy_file(const std::string& file_path,
|
|||||||
if (it == legacy_storage_map.end()) {
|
if (it == legacy_storage_map.end()) {
|
||||||
return false;
|
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;
|
return false;
|
||||||
}
|
}
|
||||||
storage_offsets[storage_key] = current_offset + LEGACY_STORAGE_HEADER_SIZE;
|
uint64_t storage_offset = current_offset + LEGACY_STORAGE_HEADER_SIZE;
|
||||||
current_offset += LEGACY_STORAGE_HEADER_SIZE + it->second;
|
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) {
|
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 base_offset = it_offset->second;
|
||||||
uint64_t storage_nbytes = it_size->second;
|
uint64_t storage_nbytes = it_size->second;
|
||||||
uint64_t tensor_nbytes = tensor_storage.nbytes_to_read();
|
int64_t tensor_nbytes = tensor_storage.nbytes_to_read();
|
||||||
if (tensor_storage.offset + tensor_nbytes > storage_nbytes) {
|
if (tensor_nbytes < 0 ||
|
||||||
|
tensor_storage.offset > storage_nbytes ||
|
||||||
|
static_cast<uint64_t>(tensor_nbytes) > storage_nbytes - tensor_storage.offset) {
|
||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@ -76,8 +76,10 @@ static bool parse_zip_data_pkl(const uint8_t* buffer,
|
|||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
|
|
||||||
uint64_t tensor_nbytes = tensor_storage.nbytes_to_read();
|
int64_t tensor_nbytes = tensor_storage.nbytes_to_read();
|
||||||
if (tensor_storage.offset + tensor_nbytes > entry_size) {
|
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 + "'");
|
set_error(error, "tensor '" + tensor_storage.name + "' exceeds storage entry '" + entry_name + "'");
|
||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
|
|||||||
@ -317,7 +317,7 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const
|
|||||||
|
|
||||||
std::vector<TensorStorage> tensor_storages;
|
std::vector<TensorStorage> tensor_storages;
|
||||||
std::string error;
|
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());
|
LOG_ERROR("%s", error.c_str());
|
||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
@ -498,11 +498,18 @@ SDVersion ModelLoader::get_sd_version() {
|
|||||||
return VERSION_MINIT2I;
|
return VERSION_MINIT2I;
|
||||||
}
|
}
|
||||||
if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.img_mod.1.weight") != std::string::npos) {
|
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()) {
|
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_LAYERED;
|
||||||
}
|
}
|
||||||
return VERSION_QWEN_IMAGE;
|
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) {
|
if (tensor_storage.name.find("llm_adapter.blocks.0.cross_attn.q_proj.weight") != std::string::npos) {
|
||||||
return VERSION_ANIMA;
|
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) {
|
if (tensor_storage.name.find("model.diffusion_model.adaln_single.emb.timestep_embedder.linear_1.bias") != std::string::npos) {
|
||||||
return VERSION_LTXAV;
|
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) {
|
if (tensor_storage.name.find("model.diffusion_model.blocks.0.cross_attn.norm_k.weight") != std::string::npos) {
|
||||||
is_wan = true;
|
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<size_t> tensor_idx(0);
|
||||||
std::atomic<bool> failed(false);
|
std::atomic<bool> failed(false);
|
||||||
std::vector<std::thread> workers;
|
std::vector<std::thread> workers;
|
||||||
std::mutex rpc_backend_mutex;
|
std::mutex backend_tensor_set_mutex;
|
||||||
|
|
||||||
for (int i = 0; i < n_threads; ++i) {
|
for (int i = 0; i < n_threads; ++i) {
|
||||||
workers.emplace_back([&, file_path, is_zip]() {
|
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> read_buffer;
|
||||||
std::vector<uint8_t> convert_buffer;
|
std::vector<uint8_t> convert_buffer;
|
||||||
|
std::vector<uint8_t> zip_entry_buffer;
|
||||||
|
|
||||||
while (true) {
|
while (true) {
|
||||||
int64_t t0, t1;
|
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();
|
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) {
|
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);
|
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) {
|
if (entry_size != n) {
|
||||||
int64_t t_memcpy_start;
|
int64_t t_memcpy_start;
|
||||||
read_buffer.resize(entry_size);
|
zip_entry_buffer.resize(entry_size);
|
||||||
zip_entry_noallocread(zip, (void*)read_buffer.data(), 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();
|
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);
|
memcpy_time_ms.fetch_add(ggml_time_ms() - t_memcpy_start);
|
||||||
} else {
|
} 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);
|
zip_entry_close(zip);
|
||||||
} else if (mmapped) {
|
} else if (mmapped) {
|
||||||
if (!mmapped->copy_data(buf, n, tensor_storage.offset)) {
|
if (!mmapped->copy_data(buf, n, tensor_storage.offset)) {
|
||||||
LOG_ERROR("read tensor data failed: '%s'", file_path.c_str());
|
LOG_ERROR("read tensor data failed: '%s'", file_path.c_str());
|
||||||
failed = true;
|
return false;
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
file.seekg(tensor_storage.offset);
|
file.seekg(tensor_storage.offset);
|
||||||
file.read(buf, n);
|
file.read(buf, n);
|
||||||
if (!file) {
|
if (!file) {
|
||||||
LOG_ERROR("read tensor data failed: '%s'", file_path.c_str());
|
LOG_ERROR("read tensor data failed: '%s'", file_path.c_str());
|
||||||
failed = true;
|
return false;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
return true;
|
||||||
};
|
};
|
||||||
|
|
||||||
char* read_buf = nullptr;
|
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();
|
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();
|
t1 = ggml_time_ms();
|
||||||
read_time_ms.fetch_add(t1 - t0);
|
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)) {
|
if (dst_tensor->buffer != nullptr && !ggml_backend_buffer_is_host(dst_tensor->buffer)) {
|
||||||
t0 = ggml_time_ms();
|
t0 = ggml_time_ms();
|
||||||
|
|
||||||
// RPC backends require serialized access to prevent concurrency issues
|
std::lock_guard<std::mutex> lock(backend_tensor_set_mutex);
|
||||||
const char* buffer_type_name = ggml_backend_buft_name(ggml_backend_buffer_get_type(dst_tensor->buffer));
|
ggml_backend_tensor_set(dst_tensor, convert_buf, 0, ggml_nbytes(dst_tensor));
|
||||||
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));
|
|
||||||
}
|
|
||||||
|
|
||||||
t1 = ggml_time_ms();
|
t1 = ggml_time_ms();
|
||||||
copy_to_backend_time_ms.fetch_add(t1 - t0);
|
copy_to_backend_time_ms.fetch_add(t1 - t0);
|
||||||
@ -1476,6 +1508,9 @@ bool ModelLoader::load_tensors(std::map<std::string, ggml_tensor*>& tensors,
|
|||||||
|
|
||||||
bool ModelLoader::tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type) {
|
bool ModelLoader::tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type) {
|
||||||
const std::string& name = tensor_storage.name;
|
const std::string& name = tensor_storage.name;
|
||||||
|
if (tensor_storage.is_int8_tensorwise) {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
if (type != GGML_TYPE_COUNT) {
|
if (type != GGML_TYPE_COUNT) {
|
||||||
if (ggml_is_quantized(type) && tensor_storage.ne[0] % ggml_blck_size(type) != 0) {
|
if (ggml_is_quantized(type) && tensor_storage.ne[0] % ggml_blck_size(type) != 0) {
|
||||||
// Pass, do not convert
|
// Pass, do not convert
|
||||||
|
|||||||
@ -36,6 +36,7 @@ protected:
|
|||||||
std::vector<ModelFileData> file_data;
|
std::vector<ModelFileData> file_data;
|
||||||
bool model_files_processed = false;
|
bool model_files_processed = false;
|
||||||
String2TensorStorage tensor_storage_map;
|
String2TensorStorage tensor_storage_map;
|
||||||
|
std::map<std::string, std::string> metadata_;
|
||||||
int n_threads_;
|
int n_threads_;
|
||||||
|
|
||||||
size_t add_file_path(const std::string& file_path);
|
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();
|
std::map<ggml_type, uint32_t> get_vae_wtype_stat();
|
||||||
String2TensorStorage& get_tensor_storage_map() { return tensor_storage_map; }
|
String2TensorStorage& get_tensor_storage_map() { return tensor_storage_map; }
|
||||||
const String2TensorStorage& get_tensor_storage_map() const { 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_n_threads(int n_threads);
|
||||||
void set_wtype_override(ggml_type wtype, std::string tensor_type_rules = "");
|
void set_wtype_override(ggml_type wtype, std::string tensor_type_rules = "");
|
||||||
void process_model_files(bool enable_mmap = false, bool writable_mmap = true);
|
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;
|
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() {
|
ModelManager::~ModelManager() {
|
||||||
release_all();
|
release_all();
|
||||||
}
|
}
|
||||||
@ -135,7 +177,8 @@ bool ModelManager::register_param_tensors(const std::string& desc,
|
|||||||
ggml_backend_t params_backend,
|
ggml_backend_t params_backend,
|
||||||
size_t* registered_tensor_size,
|
size_t* registered_tensor_size,
|
||||||
bool allow_split_buffer,
|
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()) {
|
if (desc.empty()) {
|
||||||
LOG_ERROR("model manager tensor desc is empty");
|
LOG_ERROR("model manager tensor desc is empty");
|
||||||
return false;
|
return false;
|
||||||
@ -168,6 +211,12 @@ bool ModelManager::register_param_tensors(const std::string& desc,
|
|||||||
state->params_backend = params_backend;
|
state->params_backend = params_backend;
|
||||||
state->allow_split_buffer = allow_split_buffer;
|
state->allow_split_buffer = allow_split_buffer;
|
||||||
state->params_follow_compute_backend = params_follow_compute_backend;
|
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));
|
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) {
|
if (params_buft == nullptr) {
|
||||||
params_buft = ggml_backend_get_default_buffer_type(state.params_backend);
|
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;
|
return params_buft;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@ -39,6 +39,7 @@ private:
|
|||||||
bool allow_split_buffer = false;
|
bool allow_split_buffer = false;
|
||||||
bool params_follow_compute_backend = false;
|
bool params_follow_compute_backend = false;
|
||||||
bool metadata_validated = false;
|
bool metadata_validated = false;
|
||||||
|
enum ggml_op usage_op = GGML_OP_NONE;
|
||||||
|
|
||||||
int active_prepare_count = 0;
|
int active_prepare_count = 0;
|
||||||
|
|
||||||
@ -130,9 +131,10 @@ public:
|
|||||||
ResidencyMode residency_mode,
|
ResidencyMode residency_mode,
|
||||||
ggml_backend_t compute_backend,
|
ggml_backend_t compute_backend,
|
||||||
ggml_backend_t params_backend,
|
ggml_backend_t params_backend,
|
||||||
size_t* registered_tensor_size = nullptr,
|
size_t* registered_tensor_size = nullptr,
|
||||||
bool allow_split_buffer = false,
|
bool allow_split_buffer = false,
|
||||||
bool params_follow_compute_backend = false);
|
bool params_follow_compute_backend = false,
|
||||||
|
const std::map<ggml_tensor*, enum ggml_op>* tensor_ops = nullptr);
|
||||||
|
|
||||||
bool unregister_param_tensors(const std::string& desc,
|
bool unregister_param_tensors(const std::string& desc,
|
||||||
size_t* registered_tensor_size = nullptr);
|
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) {
|
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{
|
static const std::vector<std::pair<std::string, std::string>> qwen3_vl_vision_name_map{
|
||||||
{"mm.0.", "merger.linear_fc1."},
|
{"mm.0.", "merger.linear_fc1."},
|
||||||
{"mm.2.", "merger.linear_fc2."},
|
{"mm.2.", "merger.linear_fc2."},
|
||||||
@ -201,6 +215,10 @@ std::string convert_qwen3_vl_vision_name(std::string name) {
|
|||||||
{"ln1.", "norm1."},
|
{"ln1.", "norm1."},
|
||||||
{"ln2.", "norm2."},
|
{"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);
|
replace_with_name_map(name, qwen3_vl_vision_name_map);
|
||||||
return name;
|
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;
|
return result;
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -658,6 +682,72 @@ std::string convert_diffusers_dit_to_original_flux(std::string name) {
|
|||||||
return 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) {
|
std::string convert_diffusers_dit_to_original_lumina2(std::string name) {
|
||||||
int num_layers = 30;
|
int num_layers = 30;
|
||||||
int num_refiner_layers = 2;
|
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);
|
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)) {
|
} 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);
|
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)) {
|
} else if (sd_version_is_z_image(version)) {
|
||||||
name = convert_diffusers_dit_to_original_lumina2(name);
|
name = convert_diffusers_dit_to_original_lumina2(name);
|
||||||
} else if (sd_version_is_anima(version)) {
|
} 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) {
|
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)) {
|
if (sd_version_uses_wan_vae(version)) {
|
||||||
return convert_diffusers_to_original_wan_vae(name);
|
return convert_diffusers_to_original_wan_vae(name);
|
||||||
}
|
}
|
||||||
@ -1104,6 +1199,7 @@ std::string convert_sep_to_dot(std::string name) {
|
|||||||
"x_embedder",
|
"x_embedder",
|
||||||
"cross_attn",
|
"cross_attn",
|
||||||
"output_proj",
|
"output_proj",
|
||||||
|
"token_refiner",
|
||||||
};
|
};
|
||||||
|
|
||||||
// record the positions of underscores that should NOT be replaced
|
// record the positions of underscores that should NOT be replaced
|
||||||
@ -1208,11 +1304,54 @@ static std::string convert_esrgan_tensor_name(std::string name) {
|
|||||||
return 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) {
|
std::string convert_tensor_name(std::string name, SDVersion version) {
|
||||||
if (version == VERSION_ESRGAN) {
|
if (version == VERSION_ESRGAN) {
|
||||||
return convert_esrgan_tensor_name(std::move(name));
|
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_lora = false;
|
||||||
bool is_lycoris_underline = false;
|
bool is_lycoris_underline = false;
|
||||||
bool is_underline = false;
|
bool is_underline = false;
|
||||||
@ -1264,6 +1403,8 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
|
|||||||
{".lora_B.weight", ".weight.lora_up"},
|
{".lora_B.weight", ".weight.lora_up"},
|
||||||
{".lora_A.default.weight", ".weight.lora_down"},
|
{".lora_A.default.weight", ".weight.lora_down"},
|
||||||
{".lora_B.default.weight", ".weight.lora_up"},
|
{".lora_B.default.weight", ".weight.lora_up"},
|
||||||
|
{".lora_A", ".weight.lora_down"},
|
||||||
|
{".lora_B", ".weight.lora_up"},
|
||||||
{".lora_linear", ".weight.alpha"},
|
{".lora_linear", ".weight.alpha"},
|
||||||
{".alpha", ".weight.alpha"},
|
{".alpha", ".weight.alpha"},
|
||||||
{".scale", ".weight.scale"},
|
{".scale", ".weight.scale"},
|
||||||
@ -1329,27 +1470,54 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
|
|||||||
{"te2.", "cond_stage_model.1.transformer."},
|
{"te2.", "cond_stage_model.1.transformer."},
|
||||||
{"te1.", "cond_stage_model.transformer."},
|
{"te1.", "cond_stage_model.transformer."},
|
||||||
{"te3.", "text_encoders.t5xxl.transformer."},
|
{"te3.", "text_encoders.t5xxl.transformer."},
|
||||||
|
{"clip_vision.", "cond_stage_model.transformer."},
|
||||||
};
|
};
|
||||||
|
|
||||||
if (sd_version_is_flux(version)) {
|
if (sd_version_is_flux(version)) {
|
||||||
prefix_map["te1."] = "text_encoders.clip_l.transformer.";
|
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);
|
replace_with_prefix_map(name, prefix_map);
|
||||||
|
|
||||||
if ((sd_version_is_boogu_image(version) || sd_version_is_krea2(version)) && starts_with(name, "text_encoders.llm.visual.")) {
|
if (sd_version_is_boogu_image(version) || sd_version_is_krea2(version) || sd_version_is_mage_flow(version) || sd_version_is_minimax_h3(version)) {
|
||||||
name = convert_qwen3_vl_vision_name(std::move(name));
|
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
|
// diffusion model
|
||||||
{
|
{
|
||||||
|
bool matched = false;
|
||||||
for (const auto& prefix : diffuison_model_prefix_vec) {
|
for (const auto& prefix : diffuison_model_prefix_vec) {
|
||||||
if (starts_with(name, prefix)) {
|
if (starts_with(name, prefix)) {
|
||||||
name = convert_diffusion_model_name(name.substr(prefix.size()), prefix, version);
|
name = convert_diffusion_model_name(name.substr(prefix.size()), prefix, version);
|
||||||
name = prefix + name;
|
name = prefix + name;
|
||||||
|
matched = true;
|
||||||
break;
|
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
|
// cond_stage_model
|
||||||
|
|||||||
@ -306,8 +306,33 @@ struct KarrasScheduler : SigmaScheduler {
|
|||||||
};
|
};
|
||||||
|
|
||||||
struct BetaScheduler : SigmaScheduler {
|
struct BetaScheduler : SigmaScheduler {
|
||||||
static constexpr double alpha = 0.6;
|
double alpha = 0.6;
|
||||||
static constexpr double beta = 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) {
|
static double log_beta(double a, double b) {
|
||||||
return std::lgamma(a) + std::lgamma(b) - std::lgamma(a + b);
|
return std::lgamma(a) + std::lgamma(b) - std::lgamma(a + b);
|
||||||
@ -1032,7 +1057,7 @@ struct Denoiser {
|
|||||||
break;
|
break;
|
||||||
case BETA_SCHEDULER:
|
case BETA_SCHEDULER:
|
||||||
LOG_INFO("get_sigmas with Beta scheduler");
|
LOG_INFO("get_sigmas with Beta scheduler");
|
||||||
scheduler = std::make_shared<BetaScheduler>();
|
scheduler = std::make_shared<BetaScheduler>(extra_sample_args);
|
||||||
break;
|
break;
|
||||||
case EXPONENTIAL_SCHEDULER:
|
case EXPONENTIAL_SCHEDULER:
|
||||||
LOG_INFO("get_sigmas exponential scheduler");
|
LOG_INFO("get_sigmas exponential scheduler");
|
||||||
@ -1153,9 +1178,9 @@ struct CompVisDenoiser : public Denoiser {
|
|||||||
return {c_skip, c_out, c_in};
|
return {c_skip, c_out, c_in};
|
||||||
}
|
}
|
||||||
|
|
||||||
virtual sd::Tensor<float> noise_scaling(float sigma,
|
sd::Tensor<float> noise_scaling(float sigma,
|
||||||
const sd::Tensor<float>& noise,
|
const sd::Tensor<float>& noise,
|
||||||
const sd::Tensor<float>& latent) override {
|
const sd::Tensor<float>& latent) override {
|
||||||
GGML_ASSERT(noise.numel() == latent.numel());
|
GGML_ASSERT(noise.numel() == latent.numel());
|
||||||
return latent + noise * sigma;
|
return latent + noise * sigma;
|
||||||
}
|
}
|
||||||
@ -1165,7 +1190,7 @@ struct CompVisDenoiser : public Denoiser {
|
|||||||
return latent;
|
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);
|
return noise_level / (1.0f - noise_level);
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@ -1256,7 +1281,7 @@ struct DiscreteFlowDenoiser : public Denoiser {
|
|||||||
return latent * (1.0f / (1.0f - sigma));
|
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;
|
return noise_level;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
@ -1396,7 +1421,7 @@ struct MiniT2IFlowDenoiser : public Denoiser {
|
|||||||
return latent;
|
return latent;
|
||||||
}
|
}
|
||||||
|
|
||||||
float noise_level_to_sigma(float noise_level) {
|
float noise_level_to_sigma(float noise_level) override {
|
||||||
SD_UNUSED(noise_level);
|
SD_UNUSED(noise_level);
|
||||||
return 1.0f;
|
return 1.0f;
|
||||||
}
|
}
|
||||||
@ -2553,6 +2578,88 @@ static sd::Tensor<float> sample_tcd(denoise_cb_t model,
|
|||||||
return x;
|
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,
|
static sd::Tensor<float> sample_euler_cfg_pp(denoise_cb_t model,
|
||||||
sd::Tensor<float> x,
|
sd::Tensor<float> x,
|
||||||
const std::vector<float>& sigmas) {
|
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);
|
return sample_euler_ancestral(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
|
||||||
case TCD_SAMPLE_METHOD:
|
case TCD_SAMPLE_METHOD:
|
||||||
return sample_tcd(model, std::move(x), sigmas, rng, eta);
|
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:
|
case EULER_CFG_PP_SAMPLE_METHOD:
|
||||||
return sample_euler_cfg_pp(model, std::move(x), sigmas);
|
return sample_euler_cfg_pp(model, std::move(x), sigmas);
|
||||||
case EULER_A_CFG_PP_SAMPLE_METHOD:
|
case EULER_A_CFG_PP_SAMPLE_METHOD:
|
||||||
|
|||||||
@ -4,6 +4,33 @@
|
|||||||
#include "core/tensor.hpp"
|
#include "core/tensor.hpp"
|
||||||
#include "ggml.h"
|
#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] = {
|
const float ltxav_latent_rgb_proj[128][3] = {
|
||||||
{-0.0293802f, -0.0362516f, -0.0291386f},
|
{-0.0293802f, -0.0362516f, -0.0291386f},
|
||||||
{0.0117735f, 0.0223435f, 0.018856f},
|
{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 << "\"" << token << "\", ";
|
||||||
}
|
}
|
||||||
ss << "]";
|
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;
|
return bpe_tokens;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||