Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e92e86fb11 | ||
|
|
9cfe2af8f9 | ||
|
|
2993b7fb43 | ||
|
|
53856e7ec8 | ||
|
|
22516991cb | ||
|
|
5ef4a7557d | ||
|
|
2d0385ba85 | ||
|
|
87a01773be | ||
|
|
b0f856804c | ||
|
|
78124b6454 | ||
|
|
b338b4b4b9 | ||
|
|
b8bf67615c | ||
|
|
8d377073e3 | ||
|
|
5114672c48 | ||
|
|
8a51eb9284 | ||
|
|
35fb21f3e0 | ||
|
|
cfd4cff2e6 | ||
|
|
5e4e03c8dd | ||
|
|
2961182365 | ||
|
|
ea4e566ccf | ||
|
|
b290693977 | ||
|
|
7717e82c00 | ||
|
|
fafe8e606c | ||
|
|
a8a91b24cd | ||
|
|
c00a9e956d | ||
|
|
833369da84 | ||
|
|
74bce049d0 | ||
|
|
b5d812008e | ||
|
|
c79d24bdbd | ||
|
|
1b042838d9 | ||
|
|
ead6bf521b | ||
|
|
9beb6aca69 | ||
|
|
12b6fbff28 | ||
|
|
cc73429228 | ||
|
|
6314af404a | ||
|
|
885f01a555 | ||
|
|
9ef6e7398f | ||
|
|
bb84971129 | ||
|
|
e22272ee63 | ||
|
|
4fcc6fef10 | ||
|
|
9e1055d2e3 | ||
|
|
8b135b5489 | ||
|
|
e071aa390b | ||
|
|
dff0e88265 |
@@ -449,8 +449,8 @@ jobs:
|
||||
runs-on: windows-2022
|
||||
|
||||
env:
|
||||
ROCM_VERSION: "7.13.0"
|
||||
GPU_TARGETS: "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1150;gfx1151;gfx1152;gfx1200;gfx1201"
|
||||
ROCM_VERSION: "7.14.0"
|
||||
GPU_TARGETS: "gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201"
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v3
|
||||
@@ -472,34 +472,68 @@ jobs:
|
||||
uses: actions/cache@v4
|
||||
with:
|
||||
path: C:\TheRock\build
|
||||
key: rocm-${{ env.ROCM_VERSION }}-gfx1151-${{ runner.os }}
|
||||
key: rocm-wheels-${{ env.ROCM_VERSION }}-${{ runner.os }}
|
||||
|
||||
- name: ccache
|
||||
uses: ggml-org/ccache-action@v1.2.16
|
||||
with:
|
||||
key: windows-latest-rocm-${{ env.ROCM_VERSION }}-x64
|
||||
key: windows-rocm-${{ env.ROCM_VERSION }}-x64
|
||||
evict-old-files: 1d
|
||||
|
||||
- name: Install ROCm
|
||||
- name: Install ROCm with Wheels
|
||||
if: steps.cache-rocm.outputs.cache-hit != 'true'
|
||||
run: |
|
||||
$ErrorActionPreference = "Stop"
|
||||
write-host "Downloading AMD ROCm ${{ env.ROCM_VERSION }} tarball"
|
||||
Invoke-WebRequest -Uri "https://repo.amd.com/rocm/tarball/therock-dist-windows-gfx1151-${{ env.ROCM_VERSION }}.tar.gz" -OutFile "${env:RUNNER_TEMP}\rocm.tar.gz"
|
||||
write-host "Extracting ROCm tarball"
|
||||
mkdir C:\TheRock\build -Force
|
||||
tar -xzf "${env:RUNNER_TEMP}\rocm.tar.gz" -C C:\TheRock\build --strip-components=1
|
||||
write-host "Completed ROCm extraction"
|
||||
write-host "Setting up Python virtual environment"
|
||||
|
||||
# Create the venv directly at the cache location to avoid relocation issues
|
||||
New-Item -Path "C:\TheRock\build" -ItemType Directory -Force | Out-Null
|
||||
python -m venv C:\TheRock\build\.venv
|
||||
& C:\TheRock\build\.venv\Scripts\Activate.ps1
|
||||
|
||||
write-host "Upgrading pip"
|
||||
python -m pip install --upgrade pip
|
||||
|
||||
write-host "Installing ROCm wheels for multi-arch support"
|
||||
# Install ROCm wheels for multi-arch support (this may take several minutes)
|
||||
python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ "rocm[libraries,devel]==${{env.ROCM_VERSION}}"
|
||||
|
||||
# Pre-expand the devel tree so it is included in the cache
|
||||
write-host "Initializing ROCm devel tree"
|
||||
rocm-sdk init
|
||||
if ($LASTEXITCODE -ne 0) { throw "rocm-sdk init failed with exit code $LASTEXITCODE" }
|
||||
write-host "Completed ROCm wheel installation to C:\TheRock\build"
|
||||
|
||||
- name: Setup ROCm Environment
|
||||
run: |
|
||||
$rocmPath = "C:\TheRock\build"
|
||||
$ErrorActionPreference = "Stop"
|
||||
|
||||
# Activate venv from cache or fresh install
|
||||
& C:\TheRock\build\.venv\Scripts\Activate.ps1
|
||||
|
||||
# Expand the devel tree (idempotent; no-op if already done during install)
|
||||
rocm-sdk init
|
||||
if ($LASTEXITCODE -ne 0) { throw "rocm-sdk init failed with exit code $LASTEXITCODE" }
|
||||
|
||||
# Get ROCm installation paths using the rocm-sdk CLI tool
|
||||
$rocmPath = (rocm-sdk path --root)
|
||||
if (-not $rocmPath) { throw "rocm-sdk path --root returned empty - devel package may not be installed" }
|
||||
$rocmPath = $rocmPath.Trim()
|
||||
$cmakePath = (rocm-sdk path --cmake).Trim()
|
||||
$binPath = (rocm-sdk path --bin).Trim()
|
||||
write-host "ROCm root: $rocmPath"
|
||||
write-host "CMake path: $cmakePath"
|
||||
write-host "Bin path: $binPath"
|
||||
|
||||
echo "HIP_PATH=$rocmPath" >> $env:GITHUB_ENV
|
||||
echo "CMAKE_PREFIX_PATH=$cmakePath" >> $env:GITHUB_ENV
|
||||
echo "HIP_DEVICE_LIB_PATH=$rocmPath\lib\llvm\amdgcn\bitcode" >> $env:GITHUB_ENV
|
||||
echo "HIP_PLATFORM=amd" >> $env:GITHUB_ENV
|
||||
echo "LLVM_PATH=$rocmPath\lib\llvm" >> $env:GITHUB_ENV
|
||||
echo "$rocmPath\bin" >> $env:GITHUB_PATH
|
||||
echo "$rocmPath\lib\llvm\bin" >> $env:GITHUB_PATH
|
||||
echo "$binPath" >> $env:GITHUB_PATH
|
||||
|
||||
# Keep venv in PATH for subsequent steps
|
||||
echo "C:\TheRock\build\.venv\Scripts" >> $env:GITHUB_PATH
|
||||
|
||||
- name: Build
|
||||
run: |
|
||||
@@ -527,139 +561,6 @@ jobs:
|
||||
- name: Pack artifacts
|
||||
if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
|
||||
run: |
|
||||
$ErrorActionPreference = "Stop"
|
||||
$dst = "build\bin"
|
||||
$rocmBin = Join-Path "${env:HIP_PATH}" "bin"
|
||||
$requiredRocmPaths = @(
|
||||
(Join-Path $rocmBin "rocblas.dll"),
|
||||
(Join-Path $rocmBin "rocblas\library")
|
||||
)
|
||||
foreach ($path in $requiredRocmPaths) {
|
||||
if (!(Test-Path $path)) {
|
||||
throw "Missing ROCm runtime dependency: $path"
|
||||
}
|
||||
}
|
||||
|
||||
foreach ($pattern in @("rocblas*.dll", "hipblas*.dll", "libhipblas*.dll")) {
|
||||
Copy-Item -Path (Join-Path $rocmBin $pattern) -Destination $dst -Force -ErrorAction SilentlyContinue
|
||||
}
|
||||
|
||||
foreach ($dir in @("rocblas", "hipblaslt")) {
|
||||
$src = Join-Path $rocmBin $dir
|
||||
if (Test-Path $src) {
|
||||
Copy-Item -Path $src -Destination $dst -Recurse -Force
|
||||
}
|
||||
}
|
||||
|
||||
7z a sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-rocm-${{ env.ROCM_VERSION }}-x64.zip .\build\bin\*
|
||||
|
||||
- name: Upload artifacts
|
||||
if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-rocm-${{ env.ROCM_VERSION }}-x64.zip
|
||||
path: |
|
||||
sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-rocm-${{ env.ROCM_VERSION }}-x64.zip
|
||||
|
||||
windows-latest-cmake-hip:
|
||||
runs-on: windows-2022
|
||||
|
||||
env:
|
||||
HIPSDK_INSTALLER_VERSION: "26.Q1"
|
||||
ROCM_VERSION: "7.1.1"
|
||||
GPU_TARGETS: "gfx1150;gfx1151;gfx1200;gfx1201;gfx1100;gfx1101;gfx1102;gfx1030;gfx1031;gfx1032"
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v3
|
||||
with:
|
||||
submodules: recursive
|
||||
|
||||
- name: Setup Node
|
||||
uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: 20
|
||||
|
||||
- name: Setup pnpm
|
||||
uses: pnpm/action-setup@v4
|
||||
with:
|
||||
version: 10.15.1
|
||||
|
||||
- name: Cache ROCm Installation
|
||||
id: cache-rocm
|
||||
uses: actions/cache@v4
|
||||
with:
|
||||
path: C:\Program Files\AMD\ROCm
|
||||
key: rocm-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ runner.os }}
|
||||
|
||||
- name: ccache
|
||||
uses: ggml-org/ccache-action@v1.2.16
|
||||
with:
|
||||
key: windows-latest-cmake-hip-${{ env.HIPSDK_INSTALLER_VERSION }}-x64
|
||||
evict-old-files: 1d
|
||||
|
||||
- name: Install ROCm
|
||||
if: steps.cache-rocm.outputs.cache-hit != 'true'
|
||||
run: |
|
||||
$ErrorActionPreference = "Stop"
|
||||
write-host "Downloading AMD HIP SDK Installer"
|
||||
Invoke-WebRequest -Uri "https://download.amd.com/developer/eula/rocm-hub/AMD-Software-PRO-Edition-${{ env.HIPSDK_INSTALLER_VERSION }}-Win11-For-HIP.exe" -OutFile "${env:RUNNER_TEMP}\rocm-install.exe"
|
||||
write-host "Installing AMD HIP SDK"
|
||||
$proc = Start-Process "${env:RUNNER_TEMP}\rocm-install.exe" -ArgumentList '-install' -NoNewWindow -PassThru
|
||||
$completed = $proc.WaitForExit(600000)
|
||||
if (-not $completed) {
|
||||
Write-Error "ROCm installation timed out after 10 minutes. Killing the process"
|
||||
$proc.Kill()
|
||||
exit 1
|
||||
}
|
||||
if ($proc.ExitCode -ne 0) {
|
||||
Write-Error "ROCm installation failed with exit code $($proc.ExitCode)"
|
||||
exit 1
|
||||
}
|
||||
write-host "Completed AMD HIP SDK installation"
|
||||
|
||||
- name: Verify ROCm
|
||||
run: |
|
||||
# Find and test ROCm installation
|
||||
$clangPath = Get-ChildItem 'C:\Program Files\AMD\ROCm\*\bin\clang.exe' | Select-Object -First 1
|
||||
if (-not $clangPath) {
|
||||
Write-Error "ROCm installation not found"
|
||||
exit 1
|
||||
}
|
||||
& $clangPath.FullName --version
|
||||
# Set HIP_PATH environment variable for later steps
|
||||
echo "HIP_PATH=$(Resolve-Path 'C:\Program Files\AMD\ROCm\*\bin\clang.exe' | split-path | split-path)" >> $env:GITHUB_ENV
|
||||
|
||||
- name: Build
|
||||
run: |
|
||||
mkdir build
|
||||
cd build
|
||||
$env:CMAKE_PREFIX_PATH="${env:HIP_PATH}"
|
||||
cmake .. `
|
||||
-G "Unix Makefiles" `
|
||||
-DSD_HIPBLAS=ON `
|
||||
-DSD_BUILD_SHARED_LIBS=ON `
|
||||
-DGGML_NATIVE=OFF `
|
||||
-DCMAKE_C_COMPILER=clang `
|
||||
-DCMAKE_CXX_COMPILER=clang++ `
|
||||
-DCMAKE_BUILD_TYPE=Release `
|
||||
-DGPU_TARGETS="${{ env.GPU_TARGETS }}"
|
||||
cmake --build . --config Release --parallel ${env:NUMBER_OF_PROCESSORS}
|
||||
|
||||
- name: Get commit hash
|
||||
id: commit
|
||||
if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
|
||||
uses: prompt/actions-commit-hash@v2
|
||||
|
||||
- name: Pack artifacts
|
||||
if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
|
||||
run: |
|
||||
md "build\bin\rocblas\library\"
|
||||
md "build\bin\hipblaslt\library"
|
||||
cp "${env:HIP_PATH}\bin\libhipblas.dll" "build\bin\"
|
||||
cp "${env:HIP_PATH}\bin\libhipblaslt.dll" "build\bin\"
|
||||
cp "${env:HIP_PATH}\bin\rocblas.dll" "build\bin\"
|
||||
cp "${env:HIP_PATH}\bin\rocblas\library\*" "build\bin\rocblas\library\"
|
||||
cp "${env:HIP_PATH}\bin\hipblaslt\library\*" "build\bin\hipblaslt\library\"
|
||||
7z a sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-rocm-${{ env.ROCM_VERSION }}-x64.zip .\build\bin\*
|
||||
|
||||
- name: Upload artifacts
|
||||
@@ -679,11 +580,8 @@ jobs:
|
||||
strategy:
|
||||
matrix:
|
||||
include:
|
||||
- ROCM_VERSION: "7.2.1"
|
||||
gpu_targets: "gfx908;gfx90a;gfx942;gfx1030;gfx1031;gfx1032;gfx1100;gfx1101;gfx1102;gfx1151;gfx1150;gfx1200;gfx1201"
|
||||
build: 'x64'
|
||||
- ROCM_VERSION: "7.13.0"
|
||||
gpu_targets: "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1150;gfx1151;gfx1152;gfx1200;gfx1201"
|
||||
- ROCM_VERSION: "7.14.0"
|
||||
gpu_targets: "gfx900;gfx906;gfx908;gfx90a;gfx90c;gfx942;gfx950;gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201"
|
||||
build: x64
|
||||
|
||||
steps:
|
||||
@@ -702,7 +600,7 @@ jobs:
|
||||
- name: Dependencies
|
||||
id: depends
|
||||
run: |
|
||||
sudo apt install -y build-essential cmake wget zip ninja-build
|
||||
sudo apt install -y build-essential git cmake wget
|
||||
|
||||
- name: Free disk space
|
||||
run: |
|
||||
@@ -723,38 +621,36 @@ jobs:
|
||||
sudo apt clean
|
||||
df -h
|
||||
|
||||
- name: Setup Legacy ROCm
|
||||
if: matrix.ROCM_VERSION == '7.2.1'
|
||||
id: legacy_env
|
||||
run: |
|
||||
sudo mkdir --parents --mode=0755 /etc/apt/keyrings
|
||||
wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | \
|
||||
gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null
|
||||
|
||||
sudo tee /etc/apt/sources.list.d/rocm.list << EOF
|
||||
deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/${{ matrix.ROCM_VERSION }} noble main
|
||||
EOF
|
||||
|
||||
sudo tee /etc/apt/preferences.d/rocm-pin-600 << EOF
|
||||
Package: *
|
||||
Pin: release o=repo.radeon.com
|
||||
Pin-Priority: 600
|
||||
EOF
|
||||
|
||||
sudo apt update
|
||||
sudo apt-get install -y libssl-dev rocm-hip-sdk
|
||||
|
||||
- name: Setup TheRock
|
||||
if: matrix.ROCM_VERSION != '7.2.1'
|
||||
- name: Setup TheRock with Wheels
|
||||
id: therock_env
|
||||
run: |
|
||||
wget https://repo.amd.com/rocm/tarball/therock-dist-linux-gfx1151-${{ matrix.ROCM_VERSION }}.tar.gz
|
||||
mkdir install
|
||||
tar -xf *.tar.gz -C install
|
||||
export ROCM_PATH=$(pwd)/install
|
||||
echo ROCM_PATH=$ROCM_PATH >> $GITHUB_ENV
|
||||
echo PATH=$PATH:$ROCM_PATH/bin >> $GITHUB_ENV
|
||||
echo LD_LIBRARY_PATH=$ROCM_PATH/lib:$ROCM_PATH/llvm/lib:$ROCM_PATH/lib/rocprofiler-systems >> $GITHUB_ENV
|
||||
# Create Python virtual environment
|
||||
python3 -m venv .venv
|
||||
source .venv/bin/activate
|
||||
|
||||
# Install ROCm wheels for build
|
||||
# libraries = HIP runtime and CMake configs needed for linking
|
||||
# devel = compilers, headers, static libs
|
||||
python -m pip install --upgrade pip
|
||||
python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ "rocm[libraries,devel]==${{matrix.ROCM_VERSION}}"
|
||||
|
||||
# Get ROCm installation paths using the rocm-sdk CLI tool
|
||||
ROCM_PATH=$(rocm-sdk path --root)
|
||||
CMAKE_PATH=$(rocm-sdk path --cmake)
|
||||
BIN_PATH=$(rocm-sdk path --bin)
|
||||
echo "ROCM_PATH=$ROCM_PATH"
|
||||
echo "CMAKE_PATH=$CMAKE_PATH"
|
||||
echo "BIN_PATH=$BIN_PATH"
|
||||
|
||||
# Set environment variables
|
||||
echo "ROCM_PATH=$ROCM_PATH" >> $GITHUB_ENV
|
||||
echo "CMAKE_PREFIX_PATH=$CMAKE_PATH" >> $GITHUB_ENV
|
||||
echo "HIP_PATH=$ROCM_PATH" >> $GITHUB_ENV
|
||||
echo "PATH=$BIN_PATH:${PATH}" >> $GITHUB_ENV
|
||||
echo "LD_LIBRARY_PATH=$ROCM_PATH/lib:${LD_LIBRARY_PATH:-}" >> $GITHUB_ENV
|
||||
|
||||
# Keep venv activated for subsequent steps
|
||||
echo "$(pwd)/.venv/bin" >> $GITHUB_PATH
|
||||
|
||||
# setup-node installs into /opt/hostedtoolcache, which is removed above.
|
||||
# Keep Node/pnpm setup after disk cleanup so the server frontend can be embedded.
|
||||
@@ -839,7 +735,6 @@ jobs:
|
||||
- build-and-push-docker-images
|
||||
- macOS-latest-cmake
|
||||
- windows-latest-cmake
|
||||
- windows-latest-cmake-hip
|
||||
- windows-latest-rocm
|
||||
|
||||
steps:
|
||||
|
||||
@@ -0,0 +1,48 @@
|
||||
name: Close inactive issues
|
||||
|
||||
on:
|
||||
schedule:
|
||||
# Run daily. GitHub cron schedules use UTC.
|
||||
- cron: "30 1 * * *"
|
||||
workflow_dispatch:
|
||||
inputs:
|
||||
debug_only:
|
||||
description: "Dry run: log intended actions without changing issues"
|
||||
required: false
|
||||
default: false
|
||||
type: boolean
|
||||
|
||||
permissions:
|
||||
issues: write
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}
|
||||
cancel-in-progress: false
|
||||
|
||||
jobs:
|
||||
close-inactive-issues:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Comment and close inactive issues
|
||||
uses: actions/stale@v10
|
||||
with:
|
||||
days-before-issue-stale: 365
|
||||
days-before-issue-close: 0
|
||||
|
||||
days-before-pr-stale: -1
|
||||
days-before-pr-close: -1
|
||||
|
||||
stale-issue-label: issue:inactive
|
||||
close-issue-label: issue:auto-closed
|
||||
close-issue-reason: completed
|
||||
stale-issue-message: ""
|
||||
close-issue-message: >
|
||||
This issue has had no activity for one year. The latest version of
|
||||
the code may already have fixed the problem.
|
||||
|
||||
If the issue still exists in the latest version, you can reopen
|
||||
this issue at any time with updated reproduction details.
|
||||
|
||||
remove-issue-stale-when-updated: true
|
||||
operations-per-run: 1000
|
||||
debug-only: ${{ github.event_name == 'workflow_dispatch' && inputs.debug_only || false }}
|
||||
@@ -13,3 +13,7 @@ output*.png
|
||||
models*
|
||||
*.log
|
||||
preview.png
|
||||
.claude/
|
||||
CLAUDE.local.md
|
||||
.agents/
|
||||
.codex/
|
||||
|
||||
@@ -0,0 +1,183 @@
|
||||
# Instructions for stable-diffusion.cpp
|
||||
|
||||
This document is for AI coding agents working in this repository. It should
|
||||
describe agent-specific workflow, repository routing, editing boundaries, and
|
||||
project-specific pitfalls.
|
||||
|
||||
For general contribution rules, including PR scope, commit conventions, code
|
||||
style, dependency updates, security hygiene, and AI-assisted contribution policy,
|
||||
see `CONTRIBUTING.md`.
|
||||
|
||||
---
|
||||
|
||||
## Agent Operating Rules
|
||||
|
||||
Before analyzing or modifying the repository:
|
||||
|
||||
1. Read this file.
|
||||
2. Use `rg` / `rg --files` or directory listing commands to confirm the current
|
||||
tree before relying on a path.
|
||||
3. Start from `src/` and relevant `docs/` for runtime behavior.
|
||||
4. Read the relevant code before editing.
|
||||
5. Prefer the smallest change that fits the existing architecture.
|
||||
6. Report focused verification and mention any tests not run.
|
||||
|
||||
Agents must not:
|
||||
|
||||
* Run `git push`, create PRs, or submit issue/PR comments on the user's behalf.
|
||||
* Create commits unless the user explicitly requests that specific commit.
|
||||
* Modify `ggml/`, `thirdparty/`, or `examples/server/frontend/` unless
|
||||
explicitly requested and necessary.
|
||||
* Read large local model files or tokenizer vocabulary files.
|
||||
* Rewrite unrelated code for style-only reasons.
|
||||
* Add secrets, model weights, generated binaries, local absolute paths, or
|
||||
machine-specific output.
|
||||
|
||||
When a change is large, architectural, or likely to affect public behavior,
|
||||
pause and present a short plan before editing.
|
||||
|
||||
---
|
||||
|
||||
## Repository Map and Editing Boundaries
|
||||
|
||||
This is a routing map for agents, not a full architecture document. The layout
|
||||
can change, so verify paths before using them. Do not inspect excluded
|
||||
large-data directories while checking the tree.
|
||||
|
||||
### Primary Project Code
|
||||
|
||||
Core implementation lives under `src/`.
|
||||
|
||||
Current source layout includes:
|
||||
|
||||
* `src/core/` - shared tensor, ggml integration, backend, graph, RNG, and utility
|
||||
code.
|
||||
* `src/model/` - model families and model components.
|
||||
* `src/model_io/` - model file loading, GGUF, safetensors, pickle, and related
|
||||
serialization helpers.
|
||||
* `src/runtime/` - sampling, denoising, guidance, caching, preprocessing, and
|
||||
runtime execution helpers.
|
||||
* `src/tokenizers/` - tokenizer implementations.
|
||||
* `src/conditioning/` - conditioning and prompt-related implementation.
|
||||
* `src/extensions/` - optional feature extensions.
|
||||
* top-level `src/*.cpp` and `src/*.h` files - public implementation entry
|
||||
points, model loading, conversion, versioning, and shared managers.
|
||||
|
||||
`src/tokenizers/vocab/` contains large tokenizer vocabulary data. Do not read or
|
||||
parse files in this directory; reference the path only when necessary.
|
||||
|
||||
### Public API
|
||||
|
||||
`include/` contains the C API exposed by the project. Currently the primary
|
||||
public header is `include/stable-diffusion.h`.
|
||||
|
||||
Treat public headers as stable API. Avoid breaking compatibility unless the user
|
||||
explicitly requests it. If public behavior changes, update relevant examples or
|
||||
documentation.
|
||||
|
||||
### Examples
|
||||
|
||||
`examples/` contains programs demonstrating library usage.
|
||||
|
||||
* `examples/cli/` - command line program for running models, testing features,
|
||||
and debugging.
|
||||
* `examples/common/` - shared example support code.
|
||||
* `examples/server/` - server application built on top of the library.
|
||||
* `examples/server/frontend/` - git submodule containing independent frontend
|
||||
code. Avoid modifying it unless explicitly requested.
|
||||
|
||||
### Documentation and Tooling
|
||||
|
||||
* `docs/` - documentation for supported models, build options, behavior, and
|
||||
workflows.
|
||||
* `scripts/` - development, model processing, build automation, formatting, and
|
||||
tooling scripts.
|
||||
* `cmake/` - CMake support modules.
|
||||
* `docker/` - Docker-related project files.
|
||||
* `assets/` - documentation assets; not runtime code.
|
||||
|
||||
### External, Local, and Generated State
|
||||
|
||||
* `ggml/` - git submodule for the ggml dependency.
|
||||
* `thirdparty/` - vendored third-party dependencies.
|
||||
* `models/` - local model storage. Ignore this directory and do not read model
|
||||
files.
|
||||
* `test/` - local testing scripts. Use only when relevant to the task.
|
||||
* `build/`, `build_*`, and similar directories - generated build outputs.
|
||||
Inspect them only when debugging a build result.
|
||||
|
||||
---
|
||||
|
||||
## Agent Workflow for Code Changes
|
||||
|
||||
1. Identify the relevant modules under `src/`.
|
||||
2. Check whether the change touches the public API in `include/`.
|
||||
3. Consult relevant `docs/` and examples before changing user-facing behavior.
|
||||
4. Follow existing local patterns before adding new abstractions.
|
||||
5. Keep edits scoped to the requested behavior.
|
||||
6. Run the narrowest useful build, test, or inspection command available.
|
||||
|
||||
Follow `CONTRIBUTING.md` for formatting, naming, PR expectations, dependency
|
||||
update policy, and security rules.
|
||||
|
||||
---
|
||||
|
||||
## Code Comments
|
||||
|
||||
Keep comments rare and useful.
|
||||
|
||||
Do not add comments that only describe what the code does. Add comments only
|
||||
when the code cannot fully express the logic, the logic is unusually complex, or
|
||||
there are historical reasons, invariants, constraints, compatibility concerns,
|
||||
or known pitfalls that future maintainers need to understand.
|
||||
|
||||
Do not add task-specific comments that will be meaningless after review.
|
||||
|
||||
Examples from the current codebase:
|
||||
|
||||
```cpp
|
||||
// GOOD: explains a safety constraint that is not obvious from the assignment.
|
||||
// From src/model_io/pickle_io.cpp.
|
||||
// Non-tensor checkpoint metadata can use REDUCE for arbitrary
|
||||
// Python objects. Do not execute it; keep stack shape only.
|
||||
stack.push_back(make_none_value());
|
||||
|
||||
// BAD: describes only what the next line does.
|
||||
// Set the token count to zero.
|
||||
token_count = 0;
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Text File Encoding
|
||||
|
||||
When reading or editing repository text files:
|
||||
|
||||
* Prefer UTF-8 with LF for Markdown, frontend source, JSON, and other text-first
|
||||
project files unless the file already clearly uses a different encoding.
|
||||
* Do not assume terminal output encoding matches file encoding on Windows.
|
||||
* A file that looks garbled in PowerShell output may still be valid UTF-8.
|
||||
* When inspecting UTF-8 files in PowerShell, prefer explicit UTF-8 reads such as:
|
||||
* `[Console]::OutputEncoding = [System.Text.Encoding]::UTF8`
|
||||
* `Get-Content -Encoding utf8 <path>`
|
||||
* Avoid rewriting a file purely because console output looked garbled; verify
|
||||
the actual file encoding first.
|
||||
|
||||
---
|
||||
|
||||
## Tensor and Layout Notes
|
||||
|
||||
Additional tensor/layout rules for this codebase:
|
||||
|
||||
* `sd::Tensor` shape order is not PyTorch/NumPy-style. `shape()[0]` is the
|
||||
lowest and most contiguous dimension, and higher indices are higher
|
||||
dimensions.
|
||||
* Broadcasting for `sd::Tensor` must align dimensions from low to high dimension
|
||||
indices. If one tensor has fewer dimensions, append implicit `1`s at the
|
||||
higher-dimension end.
|
||||
* `ggml_n_dims` / `ggml_n_dims(tensor)` can drop trailing singleton high
|
||||
dimensions. Do not assume a logical trailing dimension of `1` will still be
|
||||
counted in ggml metadata.
|
||||
* Internal tensor-returning interfaces use an empty `sd::Tensor` to represent
|
||||
null, absent, or failure states. Do not add `std::optional<sd::Tensor<...>>`
|
||||
for internal APIs unless a distinct semantic state is truly required.
|
||||
@@ -0,0 +1,11 @@
|
||||
@AGENTS.md
|
||||
|
||||
## Claude Code
|
||||
|
||||
Follow `AGENTS.md` as the shared repository instructions.
|
||||
|
||||
Do not duplicate contribution, style, PR, dependency, or security policy here;
|
||||
use `CONTRIBUTING.md` as the canonical source for those rules.
|
||||
|
||||
Keep Claude-specific project notes in this file only when they do not apply to
|
||||
other coding agents.
|
||||
@@ -312,7 +312,7 @@ set(CMAKE_POLICY_DEFAULT_CMP0077 NEW)
|
||||
|
||||
if (NOT SD_USE_SYSTEM_GGML)
|
||||
# see https://github.com/ggerganov/ggml/pull/682
|
||||
add_definitions(-DGGML_MAX_NAME=128)
|
||||
add_definitions(-DGGML_MAX_NAME=160)
|
||||
endif()
|
||||
|
||||
# deps
|
||||
|
||||
@@ -27,7 +27,7 @@ PRs should include:
|
||||
|
||||
Format code according to the repository style before submitting changes.
|
||||
|
||||
Formatting follows `.clang-format` (Chromium base, 4-space indent, no tabs). Run `format-code.sh` before opening a PR. Keep C++ standard at C++17-compatible patterns used in this repo.
|
||||
Formatting follows `.clang-format` (Chromium base, 4-space indent, no tabs). Run `scripts/format-code.sh` or `scripts/format-code.ps1` before opening a PR. Keep C++ standard at C++17-compatible patterns used in this repo.
|
||||
|
||||
Naming conventions:
|
||||
|
||||
|
||||
@@ -54,19 +54,25 @@ API and command-line option may change frequently.***
|
||||
- [ERNIE-Image](./docs/ernie_image.md)
|
||||
- [Boogu Image](./docs/boogu_image.md)
|
||||
- [Krea2](./docs/krea2.md)
|
||||
- [Mage-Flow](./docs/mage_flow.md)
|
||||
- [SeFi-Image](./docs/sefi_image.md)
|
||||
- [HiDream-O1-Image](./docs/hidream_o1_image.md)
|
||||
- [Ideogram4](./docs/ideogram4.md)
|
||||
- Image Edit Models
|
||||
- [Image Edit Models](./docs/edit.md)
|
||||
- [FLUX.1-Kontext-dev](./docs/kontext.md)
|
||||
- [Qwen Image Edit series](./docs/qwen_image_edit.md)
|
||||
- [LongCat Image Edit](./docs/longcat_image.md)
|
||||
- [Boogu Image Edit](./docs/boogu_image.md)
|
||||
- [Mage-Flow-Edit](./docs/mage_flow.md#image-editing)
|
||||
- Video Models
|
||||
- [Wan2.1/Wan2.2](./docs/wan.md)
|
||||
- [LTX-2.3](./docs/ltx2.md)
|
||||
- [HunyuanVideo 1.5](./docs/hunyuan_video.md)
|
||||
- [LingBot-Video](./docs/lingbot_video.md)
|
||||
- [PhotoMaker](./docs/photo_maker.md) support.
|
||||
- [IP-Adapter](./docs/ip_adapter.md) support (SD 1.5 and SDXL)
|
||||
- Control Net support with SD 1.5
|
||||
- [ADetailer](./docs/adetailer.md)
|
||||
- LoRA support, same as [stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#lora)
|
||||
- Latent Consistency Models support (LCM/LCM-LoRA)
|
||||
- Faster and memory efficient latent decoding with [TAESD](./docs/taesd.md)
|
||||
@@ -120,7 +126,7 @@ API and command-line option may change frequently.***
|
||||
- Stable Diffusion v1.5 from https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5
|
||||
|
||||
```sh
|
||||
curl -L -O https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors
|
||||
curl -L -O https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors
|
||||
```
|
||||
|
||||
### Generate an image with just one command
|
||||
@@ -162,6 +168,7 @@ These projects wrap `stable-diffusion.cpp` for easier use in other languages/fra
|
||||
|
||||
These projects use `stable-diffusion.cpp` as a backend for their image generation.
|
||||
|
||||
- [GIMP Plugins](https://github.com/themanyone/gimp-plugins)
|
||||
- [Jellybox](https://jellybox.com)
|
||||
- [Stable Diffusion GUI](https://github.com/fszontagh/sd.cpp.gui.wx)
|
||||
- [Stable Diffusion CLI-GUI](https://github.com/piallai/stable-diffusion.cpp)
|
||||
@@ -176,7 +183,3 @@ These projects use `stable-diffusion.cpp` as a backend for their image generatio
|
||||
Thank you to all the people who have already contributed to stable-diffusion.cpp!
|
||||
|
||||
[](https://github.com/leejet/stable-diffusion.cpp/graphs/contributors)
|
||||
|
||||
## Star History
|
||||
|
||||
[](https://star-history.com/#leejet/stable-diffusion.cpp&Date)
|
||||
|
||||
|
After Width: | Height: | Size: 1.7 MiB |
|
After Width: | Height: | Size: 1.4 MiB |
|
After Width: | Height: | Size: 1.8 MiB |
|
After Width: | Height: | Size: 1.6 MiB |
|
After Width: | Height: | Size: 1002 KiB |
|
After Width: | Height: | Size: 1.1 MiB |
|
After Width: | Height: | Size: 1.8 MiB |
|
After Width: | Height: | Size: 1.8 MiB |
|
After Width: | Height: | Size: 1.2 MiB |
|
After Width: | Height: | Size: 1.4 MiB |
|
After Width: | Height: | Size: 466 KiB |
|
After Width: | Height: | Size: 399 KiB |
@@ -33,7 +33,7 @@ RUN cmake --build ./build --config Release -j$(nproc)
|
||||
FROM ubuntu:$UBUNTU_VERSION AS runtime
|
||||
|
||||
RUN apt-get update && \
|
||||
apt-get install --yes --no-install-recommends libgomp1 libvulkan1 mesa-vulkan-drivers && \
|
||||
apt-get install --yes --no-install-recommends libgomp1 libvulkan1 mesa-vulkan-drivers libglvnd0 libgl1 libglx0 libegl1 libgles2 && \
|
||||
apt-get clean
|
||||
|
||||
COPY --from=build /sd.cpp/build/bin /sd.cpp/bin
|
||||
|
||||
@@ -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"
|
||||
```
|
||||
@@ -0,0 +1,171 @@
|
||||
# AnimateDiff (SD 1.5)
|
||||
|
||||
AnimateDiff adds motion to a frozen Stable Diffusion 1.5 checkpoint by
|
||||
injecting a temporal-attention module at 20 UNet slots. The base SD 1.5
|
||||
model, VAE, and text encoder are unchanged; only the motion module produces
|
||||
the temporal residual that turns a batch of independent frames into a
|
||||
coherent animation. Reference: Guo et al., "AnimateDiff: Animate Your
|
||||
Personalized Text-to-Image Diffusion Models without Specific Tuning"
|
||||
(https://arxiv.org/abs/2307.04725).
|
||||
|
||||
## Download weights
|
||||
|
||||
- Motion module (v3, recommended)
|
||||
- fp16 safetensors: https://huggingface.co/conrevo/AnimateDiff-A1111/resolve/main/motion_module/mm_sd15_v3.safetensors
|
||||
- original checkpoint: https://huggingface.co/guoyww/animatediff/resolve/main/v3_sd15_mm.ckpt
|
||||
- SD 1.5 base model
|
||||
- any SD 1.5 checkpoint works. `realisticVisionV60B1` and `toonyou_beta3`
|
||||
are the ones used in guoyww's reference configs.
|
||||
- Domain Adapter LoRA (optional, v3 only, sharpens the base UNet's output
|
||||
toward the motion module's trained distribution)
|
||||
- ckpt: https://huggingface.co/guoyww/animatediff/resolve/main/v3_sd15_adapter.ckpt
|
||||
- place under your `--lora-model-dir` and reference in the prompt as
|
||||
`<lora:v3_sd15_adapter:1.0>`.
|
||||
|
||||
The motion module is `~836 MB` and loads alongside the SD 1.5 UNet via
|
||||
`--motion-module`.
|
||||
|
||||
## Motion module versions
|
||||
|
||||
Per [animatediff.net/models](https://animatediff.net/models):
|
||||
|
||||
| Module | Base | Native res | Character |
|
||||
|---------------------|------|------------|-----------|
|
||||
| `mm_sd_v14.ckpt` | 1.5 | 256x256 | earliest, more jittery |
|
||||
| `mm_sd_v15.ckpt` | 1.5 | 256x256 | improved stability over v1.4 |
|
||||
| `mm_sd_v15_v2.ckpt` | 1.5 | 384x384 | significantly better motion dynamics |
|
||||
| `v3_sd15_mm.ckpt` | 1.5 | 512x512 | smoothest, highest quality; pairs with a Domain Adapter LoRA |
|
||||
| `mm_sdxl_v10_beta` | SDXL | 512x512 | experimental, not yet supported here |
|
||||
|
||||
Match your `-H -W` to the module's native resolution for best results. v3 is
|
||||
trained at 512x512 - going smaller (e.g. 384x384) still works but the motion
|
||||
character is closer to v2.
|
||||
|
||||
## Examples
|
||||
|
||||
Generate an 8-frame animation at 512x512, seed 42, 20 steps. The sampler /
|
||||
scheduler / CFG values below match what mm_sd15_v3 was trained with; using
|
||||
SD 1.5 defaults (euler_a, low CFG) produces noise-like output.
|
||||
|
||||
```
|
||||
.\bin\Release\sd-cli.exe -M vid_gen \
|
||||
--model ..\models\checkpoints\realisticVisionV60B1.safetensors \
|
||||
--motion-module ..\models\animatediff\mm_sd15_v3.safetensors \
|
||||
--offload-to-cpu --diffusion-fa \
|
||||
-p "a red apple on a wooden table" \
|
||||
--cfg-scale 8.0 --sampling-method euler --scheduler discrete \
|
||||
-H 512 -W 512 --video-frames 8 --fps 8 --steps 20 -s 42 \
|
||||
-o out.avi
|
||||
```
|
||||
|
||||
Generate at the motion module's native 16-frame context (recommended for
|
||||
best temporal quality). Needs more VRAM at 512x512, so drop to 384x384 or
|
||||
use layer streaming:
|
||||
|
||||
```
|
||||
.\bin\Release\sd-cli.exe -M vid_gen \
|
||||
--model ..\models\checkpoints\realisticVisionV60B1.safetensors \
|
||||
--motion-module ..\models\animatediff\mm_sd15_v3.safetensors \
|
||||
--offload-to-cpu --diffusion-fa \
|
||||
-p "photo of coastline, rocks, storm weather, wind, waves, lightning" \
|
||||
--cfg-scale 8.0 --sampling-method euler --scheduler discrete \
|
||||
-H 384 -W 384 --video-frames 16 --fps 8 --steps 20 -s 42 \
|
||||
-o out.avi
|
||||
```
|
||||
|
||||
Low-VRAM streaming (verified with a 2 GiB cap on RTX 3060):
|
||||
|
||||
```
|
||||
.\bin\Release\sd-cli.exe -M vid_gen \
|
||||
--model ..\models\checkpoints\realisticVisionV60B1.safetensors \
|
||||
--motion-module ..\models\animatediff\mm_sd15_v3.safetensors \
|
||||
--max-vram 2.0 --stream-layers --diffusion-fa \
|
||||
-p "photo of coastline, rocks, storm weather, wind, waves, lightning" \
|
||||
--cfg-scale 8.0 --sampling-method euler --scheduler discrete \
|
||||
-H 384 -W 384 --video-frames 8 --fps 8 --steps 20 -s 42 \
|
||||
-o out.avi
|
||||
```
|
||||
|
||||
## Reference-quality reproduction
|
||||
|
||||
Using guoyww's official reference configs on this impl (RealisticVision v6.0
|
||||
base + `mm_sd15_v3` or `mm_sd_v15_v2` + native resolution + 16 frames + euler
|
||||
+ 25 steps + CFG 8 + linear beta schedule) reproduces the reference
|
||||
AnimateDiff output style.
|
||||
|
||||
### v3 (512x512, `mm_sd15_v3`)
|
||||
|
||||
| Prompt | Sample |
|
||||
|---------------------------------------|--------|
|
||||
| B&W man on stormy coastline | <img src="../assets/animatediff/v3_man.gif" width="256"/> |
|
||||
| Close-up rabbit macro shot | <img src="../assets/animatediff/v3_rabbit.gif" width="256"/> |
|
||||
| Coastline, storm, waves, lightning | <img src="../assets/animatediff/v3_coast.gif" width="256"/> |
|
||||
| Old house, storm, forest, night | <img src="../assets/animatediff/v3_house.gif" width="256"/> |
|
||||
|
||||
### v2 (384x384, `mm_sd_v15_v2.ckpt`)
|
||||
|
||||
| Prompt | Sample |
|
||||
|---------------------------------------|--------|
|
||||
| B&W man on stormy coastline | <img src="../assets/animatediff/v2_man.gif" width="256"/> |
|
||||
| Close-up rabbit macro shot | <img src="../assets/animatediff/v2_rabbit.gif" width="256"/> |
|
||||
| Coastline, storm, waves, lightning | <img src="../assets/animatediff/v2_coast.gif" width="256"/> |
|
||||
| Old house, storm, forest, night | <img src="../assets/animatediff/v2_house.gif" width="256"/> |
|
||||
|
||||
Motion is strong for scenes with motion cues in the prompt (storm/waves/wind)
|
||||
and subtle for static subjects (close-up macro), matching reference behavior.
|
||||
v2 has an additional motion module at the UNet middle block that v3 dropped;
|
||||
this impl auto-detects the topology from the checkpoint.
|
||||
|
||||
### v3 + Domain Adapter LoRA
|
||||
|
||||
Attaching the `v3_sd15_adapter` LoRA sharpens the base UNet output toward
|
||||
the training distribution the motion module was fine-tuned against. Same
|
||||
prompt, seed, config as above:
|
||||
|
||||
<img src="../assets/animatediff/v3_rabbit_domain_lora.gif" width="256"/>
|
||||
|
||||
Individual fur strands, glowing inner-ear, and richer forest detail become
|
||||
visible compared to the no-LoRA rendering.
|
||||
|
||||
```
|
||||
sd-cli -M vid_gen --model realisticVisionV60B1.safetensors \
|
||||
--motion-module mm_sd15_v3.safetensors \
|
||||
--lora-model-dir ./loras \
|
||||
-p "close up photo of a rabbit ...<lora:v3_sd15_adapter:1.0>" ...
|
||||
```
|
||||
|
||||
## img2video
|
||||
|
||||
Pass a pre-rendered image via `-i / --init-img` to animate FROM it. All N output frames start from the encoded init latent, then per-frame noise is added at `--strength`. Character identity, composition, and quality are anchored by the init image; the motion module adds subtle motion on top.
|
||||
|
||||
Left: init image rendered with `-M img_gen`. Right: 8-frame vid_gen output.
|
||||
|
||||
<img src="../assets/animatediff/img2video_demo.gif" width="512"/>
|
||||
|
||||
```
|
||||
sd-cli -M img_gen ... -o init.png # any high-quality still
|
||||
sd-cli -M vid_gen --motion-module mm_sd15_v3.safetensors \
|
||||
-i init.png --strength 0.75 \
|
||||
--cfg-scale 7.0 --sampling-method euler --scheduler karras \
|
||||
-H 512 -W 512 --video-frames 8 --steps 25 -s 42 \
|
||||
-p "..." -o out.avi
|
||||
```
|
||||
|
||||
`--strength` controls how far the motion module is allowed to deviate from the init image (higher = more motion, lower = more static).
|
||||
|
||||
## Notes
|
||||
|
||||
- The motion module was trained at `video_length=16`. Running with
|
||||
`--video-frames 16` gives the best coherence; F=8 works but shows a shorter
|
||||
motion arc. Frame counts up to 32 are supported by the positional encoding
|
||||
but exceed the trained regime and produce more static output.
|
||||
- At `--video-frames 1` the motion module is skipped entirely and the output
|
||||
is bit-identical to `-M img_gen`. This avoids the single-token
|
||||
temporal-attention degeneracy that would otherwise emit an untrained-magnitude
|
||||
residual on a single-frame sample.
|
||||
- The base UNet is frozen, so character identity and style follow the SD 1.5
|
||||
checkpoint you pass to `--model`. LoRAs and prompt weighting attach to the
|
||||
base model in the usual way.
|
||||
- No mid_block motion module in v3. `mm_sdxl_v10_beta` (SDXL variant) is not
|
||||
supported yet.
|
||||
- Output is written as MJPEG AVI. Use `--fps` to set playback speed.
|
||||
@@ -153,6 +153,7 @@ still runs out of memory, tiling is enabled and the decode retried once.
|
||||
| `controlnet` | ControlNet | `controlnet`, `control` |
|
||||
| `photomaker` | PhotoMaker ID encoder and PhotoMaker LoRA | `photomaker`, `photomakerid`, `pmid`, `photo` |
|
||||
| `upscaler` | ESRGAN upscaler | `upscaler`, `esrgan`, `hires` |
|
||||
| `detector` | ADetailer YOLOv8 detector | `detector`, `adetailer`, `yolo` |
|
||||
|
||||
`te` is the preferred module name for text encoders. `clip` is kept as an accepted alias because many existing commands and model names use CLIP terminology.
|
||||
|
||||
|
||||
@@ -24,10 +24,9 @@ You can download the preconverted gguf weights from [silveroxides/Chroma-GGUF](h
|
||||
For example:
|
||||
|
||||
```
|
||||
.\bin\Release\sd-cli.exe --diffusion-model ..\models\chroma-unlocked-v40-q8_0.gguf --vae ..\models\ae.sft --t5xxl ..\models\t5xxl_fp16.safetensors -p "a lovely cat holding a sign says 'chroma.cpp'" --cfg-scale 4.0 --sampling-method euler -v --chroma-disable-dit-mask --clip-on-cpu
|
||||
.\bin\Release\sd-cli.exe --diffusion-model ..\models\chroma-unlocked-v40-q8_0.gguf --vae ..\models\ae.sft --t5xxl ..\models\t5xxl_fp16.safetensors -p "a lovely cat holding a sign says 'chroma.cpp'" --cfg-scale 4.0 --sampling-method euler -v --model-args chroma_use_dit_mask=false --clip-on-cpu
|
||||
```
|
||||
|
||||

|
||||
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,96 @@
|
||||
# Image Editing
|
||||
|
||||
Image editing in `stable-diffusion.cpp` allows you to use reference images to guide the generation process, enabling tasks like identity preservation, style transfer, or layout modification.
|
||||
|
||||
|
||||
## Supported Models
|
||||
|
||||
Depending on the architecture, different models handle reference images differently.
|
||||
|
||||
| Model | Default Preset |
|
||||
| :--- | :--- |
|
||||
| [**FLUX.1-Kontext-dev**](./kontext.md) | `flux_kontext` |
|
||||
| [**LongCat Image Edit**](./longcat_image.md) | `longcat` |
|
||||
| [**Qwen Image Edit**](./qwen_image_edit.md) | `qwen` |
|
||||
| **Qwen Image LAYERED** | `qwen_layered` |
|
||||
| [**Flux.2 [Dev] / Flux.2 [Klein]**](./flux2.md) | `flux2` |
|
||||
| [**Boogu Image Edit**](./boogu_image.md) | `z_image_omni` |
|
||||
| **Krea2 (Community Edit LoRAs)** | `krea2_ostris_edit` |
|
||||
| [**Mage-Flow-Edit**](./mage_flow.md#image-editing) | `mage_flow` |
|
||||
| **Anima (Community Edit LoRAs)** | `cosmos_reference` |
|
||||
|
||||
Stable-diffusion.spp also supports basic Unet-based editing models like instruct-pix2pix or CosXL-Edit. This document is not about those.
|
||||
|
||||
---
|
||||
|
||||
## Configuring Reference Modes (`--ref-image-args`)
|
||||
|
||||
Different DiT-based editing models require different configurations to process reference images correctly (e.g., whether to use a Vision Language Model (VLM) encoder or pass VAE-encoded images directly to the DiT).
|
||||
|
||||
To simplify this, we provide **Presets**. By default, the system automatically selects the best preset based on the model architecture. However, you can override this using the `--ref-image-args` argument.
|
||||
|
||||
### Usage
|
||||
The `--ref-image-args` argument accepts a comma-separated list of key-value pairs:
|
||||
|
||||
**Using a preset:**
|
||||
`--ref-image-args "preset=qwen_layered"`
|
||||
|
||||
**Using a preset with a specific override:**
|
||||
`--ref-image-args "preset=krea2_edit,force_ref_timestep_zero=true"`
|
||||
|
||||
### Available Presets
|
||||
|
||||
| Preset | Primary Use Case |
|
||||
| :--- | :--- |
|
||||
| `flux_kontext` | FLUX.1 Kontext |
|
||||
| `longcat` | LongCat Image Edit |
|
||||
| `flux2` | FLUX.2 models |
|
||||
| `qwen` | Qwen Image Edit |
|
||||
| `qwen_layered` | Qwen Image Layered |
|
||||
| `z_image_omni` | Boogu, Z-Image Omni |
|
||||
| `krea2_ostris_edit` | Most Krea2 Community edit LoRAs (trained with Ostris script) |
|
||||
| `mage_flow` | Mage-Flow-Edit |
|
||||
| `krea2_edit` | Specifically for [lbouaraba/krea2edit](https://huggingface.co/conradlocke/krea2-identity-edit). (or similar) |
|
||||
| `cosmos_reference` | For Anima |
|
||||
| `default` | Uses the automatic detection based on model architecture. |
|
||||
|
||||
---
|
||||
|
||||
## Advanced Parameter Reference
|
||||
|
||||
If presets are insufficient, you can manually configure the following parameters via `--ref-image-args`:
|
||||
|
||||
| Key | Type | Description | Allowed Values |
|
||||
| :--- | :--- | :--- | :--- |
|
||||
| `preset` | string | Overrides the automatic preset. | (See the Presets table above) |
|
||||
| `pass_to_vlm` | bool | Whether reference images are passed to the VLM encoder. | `true`, `false` |
|
||||
| `pass_to_dit` | bool | Whether VAE-encoded references are passed directly to the DiT. | `true`, `false` |
|
||||
| `ref_index_mode` | string | Behavior of the RoPE index. | `fixed`, `increase`, `decrease` |
|
||||
| `force_ref_timestep_zero` | bool | Forces timestep=0 for reference tokens. | `true`, `false` (Krea2 only) |
|
||||
| `resize_before_vae` | bool | Whether reference images are resized before VAE encoding. | `true`, `false` |
|
||||
| `vae_input_max_pixels` | int | Maximum pixel area for VAE reference inputs. | Integer |
|
||||
| `vlm_resize_mode` | string | How to resize VLM reference inputs. | `longest_side`, `area`, `none` |
|
||||
| `vlm_max_size` | int | Maximum VLM input size; interpreted according to `vlm_resize_mode`. | Integer |
|
||||
| `vlm_min_size` | int | Minimum VLM input size; interpreted according to `vlm_resize_mode`. | Integer |
|
||||
| `vlm_size` | int | Shortcut to set both VLM min and max size to the same value. | Integer |
|
||||
|
||||
### Preset Default Values
|
||||
|
||||
For a technical overview of how each preset is configured, see the table below.
|
||||
|
||||
| Preset | VLM | RoPE Index | Cond Resize | Special Notes |
|
||||
| :--- | :---: | :---: | :---: | :--- |
|
||||
| `flux_kontext` | No | `fixed` | `none` | |
|
||||
| `longcat` | Yes | `fixed` | `area` | |
|
||||
| `flux2` | No | `increase` | `none` | |
|
||||
| `qwen` | Yes | `increase` | `area` | |
|
||||
| `qwen_layered` | Yes | `decrease` | `area` | |
|
||||
| `mage_flow` | Yes | `increase` | `longest` | `vlm_max_size = 384`, VAE input resized to target |
|
||||
| `z_image_omni` | Yes | `fixed` | `area` | |
|
||||
| `krea2_ostris_edit`| Yes | `increase` | `area` | `force_ref_timestep_zero = true` |
|
||||
| `krea2_edit` | Yes | `increase` | `longest` | `vlm_size = 768` |
|
||||
| `cosmos_reference` | No | `fixed` | `none` | `resize_before_vae = false` |
|
||||
|
||||
**Additional Default Notes:**
|
||||
- **VLM Input Sizes:** For most presets, `vlm_max_size` and `vlm_min_size` are set to `-1`, meaning the values are model-dependent and handled automatically. In `area` mode they represent pixel area; in `longest_side` mode they represent a side length in pixels.
|
||||
- **VAE Input Size:** `vae_input_max_pixels` defaults to $1024 \times 1024$ pixels (`1048576`).
|
||||
@@ -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>
|
||||
@@ -0,0 +1,55 @@
|
||||
# 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.
|
||||
|
||||
## 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`
|
||||
|
||||
## 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.
|
||||
|
||||
## 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
|
||||
```
|
||||
@@ -0,0 +1,32 @@
|
||||
# How to Use
|
||||
|
||||
Lingbot Video uses a Lingbot diffusion transformer, the Wan2.1 VAE, and Qwen3-VL 4B as the LLM text encoder.
|
||||
|
||||
## Download weights
|
||||
|
||||
- Download lingbot-video-dense-1.3b
|
||||
- safetensors: https://huggingface.co/robbyant/lingbot-video-dense-1.3b/tree/main/transformer
|
||||
- Download lingbot-video-moe-30b-a3b
|
||||
- safetensors: https://huggingface.co/robbyant/lingbot-video-moe-30b-a3b/tree/main/transformer
|
||||
- Download vae
|
||||
- safetensors: https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/blob/main/split_files/vae/wan_2.1_vae.safetensors
|
||||
- Download Qwen3-VL 4B
|
||||
- safetensors: https://huggingface.co/Comfy-Org/Krea-2/tree/main/text_encoders
|
||||
- gguf: https://huggingface.co/Qwen/Qwen3-VL-4B-Instruct-GGUF/tree/main
|
||||
|
||||
## Notes
|
||||
|
||||
- Use `-M vid_gen`.
|
||||
- T2V uses the text prompt only.
|
||||
- I2V uses `-i` as the first video frame. The same image is also passed to the
|
||||
Qwen3-VL prompt enhancer when vision weights are available.
|
||||
- Video frames are aligned to Wan-style temporal compression, so use frame counts
|
||||
like 33, 49, or 81.
|
||||
|
||||
## Examples
|
||||
|
||||
### LingBot-Video T2V
|
||||
|
||||
```sh
|
||||
.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model ..\models\diffusion_models\lingbot-video-dens-1.3b.safetensors --llm ..\models\text_encoders\Qwen3-VL-4B-Instruct-Q4_K_M.gguf --vae ..\models\vae\wan_2.1_vae.safetensors -p '{"caption":{"comprehensive_description":"A lovely cat sits comfortably on a soft cushion near a sunlit window, looking calm, gentle, and adorable. The cat has soft fluffy fur, bright expressive eyes, small rounded ears, delicate whiskers, and a relaxed posture. Warm daylight falls across the cat from one side, creating soft highlights on the fur and a cozy glow around the scene. The background is softly blurred, showing hints of a peaceful indoor home environment with warm tones and gentle natural light. The overall atmosphere is cute, tender, serene, and photorealistic, emphasizing the cat''s charming appearance, soft texture, and affectionate presence.","camera_info":{"color":"Warm","frame_size":"Close Up","shot_type_angle":"Eye level","lens_size":"Medium Lens","composition":"Centered balanced","lighting":"Soft light","lighting_type":"Daylight"},"world_knowledge":[],"prominent_elements":[{"name":"lovely cat","description":"A cute and gentle domestic cat with soft fluffy fur, expressive eyes, and a calm relaxed presence.","location":"center of the frame","relative_size":"large","shape_and_color":"Small animal body with rounded face, triangular ears, bright eyes, and soft fur in warm natural tones","texture":"soft, fluffy, silky","appearance_details":"The cat has clean well-groomed fur, delicate whiskers, small ears, a cute nose, and bright attentive eyes. Its expression appears calm, affectionate, and slightly curious.","relationship":"Acts as the main subject and emotional focal point of the scene.","orientation":"facing the camera","pose":"sitting comfortably with a relaxed posture","expression":"gentle, adorable, calm, slightly curious","clothing":"","gender":"","skin_tone_and_texture":""},{"name":"cat eyes","description":"Bright expressive eyes that give the cat a sweet and affectionate appearance.","location":"upper center of the cat face","relative_size":"small","shape_and_color":"Round almond-like eyes with glossy reflections","texture":"clear, glossy, reflective","appearance_details":"The eyes catch the soft daylight, creating small natural highlights that make the cat look vivid and alive.","relationship":"Enhance the emotional charm and cuteness of the cat.","orientation":"looking toward the camera","pose":"","expression":"soft and attentive","clothing":"","gender":"","skin_tone_and_texture":""},{"name":"soft cushion","description":"A comfortable cushion or blanket where the cat is resting.","location":"bottom portion of the frame","relative_size":"medium","shape_and_color":"Soft rounded fabric surface in light warm neutral tones","texture":"plush, fabric, cozy","appearance_details":"The cushion gently supports the cat and adds a comfortable home-like feeling to the scene.","relationship":"Provides a cozy resting place for the cat.","orientation":"horizontal","pose":"","expression":"","clothing":"","gender":"","skin_tone_and_texture":""},{"name":"sunlit indoor background","description":"A softly blurred indoor background with warm daylight and peaceful home atmosphere.","location":"behind the cat, filling the upper and side areas of the frame","relative_size":"large","shape_and_color":"Soft abstract shapes in warm beige, cream, and pale golden tones","texture":"soft, blurry, bokeh-like","appearance_details":"The background is intentionally out of focus, keeping attention on the cat while creating a cozy and serene mood.","relationship":"Provides a warm and gentle environment that supports the cute domestic scene.","orientation":"upright","pose":"","expression":"","clothing":"","gender":"","skin_tone_and_texture":"","is_cluster":true,"number_of_objects":"numerous"}]}}' -n '{"universal_negative":{"visual_quality":["low quality","worst quality","blurry","pixelated","jpeg artifacts","low resolution","unstable color","color flicker","underexposed","overexposed","invisible subject","subject hidden in darkness"],"artistic_style":["painting","illustration","drawing","cartoon","3d render","cgi","sketch","digital art"],"composition_and_content":["text","watermark","signature","logo","subtitles","pillarboxed","side bars","portrait image in landscape frame"],"temporal_and_motion_stability":["flickering","jittery","motion blur","temporal inconsistency","warping","morphing","incoherent motion","unnatural movement","static object with sudden jump","frame-to-frame inconsistency"],"material_and_structure":["plastic-like glass","unrealistic texture","deformed bottle","liquid freezing improperly","distorted reflections"]}}' --diffusion-fa --offload-to-cpu --cfg-scale 3 --video-frames 33 -v
|
||||
```
|
||||
@@ -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" />
|
||||
@@ -53,6 +53,26 @@ Per-module assignments can target only the largest modules:
|
||||
|
||||
See [backend selection](./backend.md) for full syntax.
|
||||
|
||||
## Run models that don't fit in VRAM (CPU streaming).
|
||||
|
||||
`--offload-to-cpu` alone keeps every parameter in system RAM and stages it to the runtime backend on first use, then leaves it resident there. If the diffusion model is larger than the runtime backend's free memory (e.g. Flux dev at bf16 on an 8 GiB GPU), that residency stops fitting during the sampling loop and generation fails. Two additional flags make it fit by trading a small amount of speed for room:
|
||||
|
||||
- `--max-vram <GiB>` sets a VRAM budget the graph-cut segmenter respects. It cuts each forward pass into segments sized to fit the budget, running them in sequence and freeing intermediate activations between them. Negative values auto-detect free VRAM and spare the given amount (`--max-vram -1` uses most of the free VRAM and keeps ~1 GiB headroom), a positive value caps the budget, `0` disables segmentation.
|
||||
- `--stream-layers` streams the diffusion model's transformer blocks one at a time. Each block's parameters are copied from the CPU to the runtime backend just before it runs and evicted when the residency budget is reached. Prefetching hides most of the copy latency behind compute. This flag only takes effect when the diffusion params backend is CPU, so it must be combined with `--offload-to-cpu` (or an explicit `--params-backend diffusion=cpu`); a warning is logged and the flag is ignored otherwise.
|
||||
|
||||
The three flags stack. The recommended shape for "biggest model my card can host":
|
||||
|
||||
```shell
|
||||
sd-cli --diffusion-model flux1-dev.safetensors ... \
|
||||
--offload-to-cpu --max-vram -1 --stream-layers
|
||||
```
|
||||
|
||||
- `--offload-to-cpu`: params in RAM, staged as needed.
|
||||
- `--max-vram -1`: use most of the free VRAM as the compute budget, spare 1 GiB headroom, let the graph-cut segmenter split each forward pass to fit.
|
||||
- `--stream-layers`: on top of the segmenter, stream individual transformer blocks so their weights don't all need to be resident at once.
|
||||
|
||||
Ordered from fastest to smallest-VRAM: no flags → `--offload-to-cpu` → `--offload-to-cpu --max-vram <N>` → `--offload-to-cpu --max-vram <N> --stream-layers`. Each step down costs a few percent of throughput to buy more room; combined they can run models roughly 3-4x larger than the raw VRAM would allow.
|
||||
|
||||
## Use quantization to reduce memory usage.
|
||||
|
||||
[quantization](./quantization_and_gguf.md)
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
# How to Use
|
||||
|
||||
PiD is NVIDIA's Pixel Diffusion Decoder. It replaces the usual VAE decode or decode-then-upscale path with a pixel-space diffusion decoder conditioned on a
|
||||
source latent and text prompt.
|
||||
source latent and text prompt. Both the original PiD checkpoints and PiD 1.5 are supported.
|
||||
|
||||
In stable-diffusion.cpp, PiD currently runs as an image edit pipeline: provide a reference image with `-r`/`--ref-image`, encode that image with a matching VAE, then let the PiD diffusion model decode/upscale directly to RGB.
|
||||
|
||||
@@ -16,6 +16,7 @@ In stable-diffusion.cpp, PiD currently runs as an image edit pipeline: provide a
|
||||
- Flux / Z-Image PiD: use the Flux VAE and pass `--vae-format flux`
|
||||
- SD3 PiD: use the SD3 VAE and pass `--vae-format sd3`
|
||||
- Flux.2 PiD: use the Flux.2 VAE and pass `--vae-format flux2`
|
||||
- Qwen-Image PiD: use the Qwen-Image 2D VAE and pass `--vae-format wan`
|
||||
|
||||
The official PiD model card should be checked before use. At the time of the initial PiD release, the official weights are under the NSCLv1 non-commercial license.
|
||||
|
||||
|
||||
@@ -52,7 +52,7 @@ to a `.pulidembd` binary file (about 131 KB). Run it once per source
|
||||
person; the same file is reused for any number of generations.
|
||||
|
||||
A reference Python script is provided alongside this docs file at
|
||||
[`script/pulid_extract_id.py`](../script/pulid_extract_id.py). It
|
||||
[`scripts/pulid_extract_id.py`](../scripts/pulid_extract_id.py). It
|
||||
requires:
|
||||
- A working CUDA / CPU PyTorch stack
|
||||
- `insightface`, `facexlib`, `eva-clip`, `torchvision`, `opencv-python`,
|
||||
|
||||
@@ -39,10 +39,10 @@
|
||||
|
||||
### Qwen Image Edit 2511
|
||||
|
||||
To use the new Qwen Image Edit 2511 mode, the `--qwen-image-zero-cond-t` flag must be enabled; otherwise, image editing quality will degrade significantly.
|
||||
To use the new Qwen Image Edit 2511 mode, `--model-args qwen_image_zero_cond_t=true` must be enabled; otherwise, image editing quality will degrade significantly.
|
||||
|
||||
```
|
||||
.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\qwen-image-edit-2511-Q4_K_M.gguf --vae ..\models\vae\qwen_image_vae.safetensors --llm ..\models\text_encoders\qwen_2.5_vl_7b.safetensors --cfg-scale 2.5 --sampling-method euler -v --offload-to-cpu --diffusion-fa --flow-shift 3 -r ..\assets\flux\flux1-dev-q8_0.png -p "change 'flux.cpp' to 'edit.cpp'" --qwen-image-zero-cond-t
|
||||
.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\qwen-image-edit-2511-Q4_K_M.gguf --vae ..\models\vae\qwen_image_vae.safetensors --llm ..\models\text_encoders\qwen_2.5_vl_7b.safetensors --cfg-scale 2.5 --sampling-method euler -v --offload-to-cpu --diffusion-fa --flow-shift 3 -r ..\assets\flux\flux1-dev-q8_0.png -p "change 'flux.cpp' to 'edit.cpp'" --model-args qwen_image_zero_cond_t=true
|
||||
```
|
||||
|
||||
<img alt="qwen_image_edit_2509" src="../assets/qwen/qwen_image_edit_2511.png" />
|
||||
<img alt="qwen_image_edit_2509" src="../assets/qwen/qwen_image_edit_2511.png" />
|
||||
|
||||
@@ -2,8 +2,8 @@
|
||||
|
||||
- download original weights(.ckpt or .safetensors). For example
|
||||
- Stable Diffusion v1.4 from https://huggingface.co/CompVis/stable-diffusion-v-1-4-original
|
||||
- Stable Diffusion v1.5 from https://huggingface.co/runwayml/stable-diffusion-v1-5
|
||||
- Stable Diffuison v2.1 from https://huggingface.co/stabilityai/stable-diffusion-2-1
|
||||
- Stable Diffusion v1.5 from https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5
|
||||
- Stable Diffuison v2.1 from https://huggingface.co/Manojb/stable-diffusion-2-1-base
|
||||
- Stable Diffusion 3 2B from https://huggingface.co/stabilityai/stable-diffusion-3-medium
|
||||
|
||||
### txt2img example
|
||||
@@ -34,4 +34,4 @@ Using formats of different precisions will yield results of varying quality.
|
||||
|
||||
<p align="center">
|
||||
<img src="../assets/img2img_output.png" width="256x">
|
||||
</p>
|
||||
</p>
|
||||
|
||||
@@ -13,8 +13,8 @@ The SeFi-Image family ships in three scales (1B / 2B / 5B) and three families (B
|
||||
Convert the transformer and text encoder to sd.cpp safetensors:
|
||||
|
||||
```bash
|
||||
python3 script/convert_sefi.py <hf_repo_dir> <out_dir>/sefi_<scale>_<family>.safetensors
|
||||
python3 script/convert_qwen3_vl.py <hf_repo_dir>/Qwen3-VL-XB-Instruct <out_dir>/qwen3_vl_<X>b.safetensors
|
||||
python3 scripts/convert_sefi.py <hf_repo_dir> <out_dir>/sefi_<scale>_<family>.safetensors
|
||||
python3 scripts/convert_qwen3_vl.py <hf_repo_dir>/Qwen3-VL-XB-Instruct <out_dir>/qwen3_vl_<X>b.safetensors
|
||||
```
|
||||
|
||||
## Variant defaults
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
include_directories(${CMAKE_CURRENT_SOURCE_DIR})
|
||||
|
||||
add_subdirectory(cli)
|
||||
add_subdirectory(server)
|
||||
add_subdirectory(server)
|
||||
|
||||
@@ -6,6 +6,9 @@ For detailed command-line arguments, run:
|
||||
./bin/sd-cli -h
|
||||
```
|
||||
|
||||
For direct image repair or automatic post-generation YOLOv8 detection followed by cropped inpainting, see
|
||||
[ADetailer](../../docs/adetailer.md).
|
||||
|
||||
Metadata mode inspects PNG/JPEG container metadata without loading any model:
|
||||
|
||||
```bash
|
||||
|
||||
@@ -199,7 +199,7 @@ struct SDCliParams {
|
||||
options.manual_options = {
|
||||
{"-M",
|
||||
"--mode",
|
||||
"run mode, one of [img_gen, vid_gen, upscale, convert, metadata], default: img_gen",
|
||||
"run mode, one of [img_gen, adetailer, vid_gen, upscale, convert, metadata], default: img_gen",
|
||||
on_mode_arg},
|
||||
{"",
|
||||
"--preview",
|
||||
@@ -566,6 +566,65 @@ bool save_results(const SDCliParams& cli_params,
|
||||
return sucessful_reults != 0;
|
||||
}
|
||||
|
||||
static bool apply_adetailer(sd_ctx_t* sd_ctx,
|
||||
const sd_ctx_params_t& sd_ctx_params,
|
||||
const SDContextParams& ctx_params,
|
||||
const SDGenerationParams& gen_params,
|
||||
const sd_img_gen_params_t& img_gen_params,
|
||||
SDMode mode,
|
||||
SDImageVec& results,
|
||||
int num_results) {
|
||||
if (gen_params.ad_model_path.empty()) {
|
||||
return true;
|
||||
}
|
||||
|
||||
sd_adetailer_params_t ad_params{};
|
||||
ad_params.prompt = gen_params.ad_prompt.empty() ? nullptr : gen_params.ad_prompt.c_str();
|
||||
ad_params.negative_prompt = gen_params.ad_negative_prompt.empty() ? nullptr : gen_params.ad_negative_prompt.c_str();
|
||||
ad_params.extra_ad_args = gen_params.extra_ad_args.c_str();
|
||||
|
||||
ADetailerCtxPtr ad_ctx(new_adetailer_ctx(gen_params.ad_model_path.c_str(),
|
||||
ctx_params.n_threads,
|
||||
sd_ctx_params.backend,
|
||||
sd_ctx_params.params_backend));
|
||||
if (ad_ctx == nullptr) {
|
||||
LOG_ERROR("new_adetailer_ctx failed");
|
||||
return false;
|
||||
}
|
||||
|
||||
for (int i = 0; i < num_results; ++i) {
|
||||
if (results[i].data == nullptr) {
|
||||
continue;
|
||||
}
|
||||
sd_img_gen_params_t ad_generation_params = img_gen_params;
|
||||
ad_generation_params.seed = img_gen_params.seed + i;
|
||||
if (mode == IMG_GEN) {
|
||||
ad_generation_params.width = 512;
|
||||
ad_generation_params.height = 512;
|
||||
ad_generation_params.strength = 0.4f;
|
||||
}
|
||||
sd_image_t* detailed_images = nullptr;
|
||||
int detailed_count = 0;
|
||||
if (!adetail_image(ad_ctx.get(),
|
||||
sd_ctx,
|
||||
results[i],
|
||||
&ad_params,
|
||||
&ad_generation_params,
|
||||
&detailed_images,
|
||||
&detailed_count) ||
|
||||
detailed_count <= 0 || detailed_images == nullptr || detailed_images[0].data == nullptr) {
|
||||
free_sd_images(detailed_images, detailed_count);
|
||||
LOG_ERROR("ADetailer failed for image %d", i + 1);
|
||||
return false;
|
||||
}
|
||||
free(results[i].data);
|
||||
results[i] = detailed_images[0];
|
||||
detailed_images[0] = {0, 0, 0, nullptr};
|
||||
free_sd_images(detailed_images, detailed_count);
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
int main(int argc, const char* argv[]) {
|
||||
if (argc > 1 && std::string(argv[1]) == "--version") {
|
||||
std::cout << version_string() << "\n";
|
||||
@@ -598,6 +657,11 @@ int main(int argc, const char* argv[]) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
if (!gen_params.ad_model_path.empty() && cli_params.mode != IMG_GEN && cli_params.mode != ADETAILER) {
|
||||
LOG_ERROR("--ad-model is only supported in image generation and adetailer modes");
|
||||
return 1;
|
||||
}
|
||||
|
||||
if (gen_params.video_frames > 4) {
|
||||
size_t last_dot_pos = cli_params.preview_path.find_last_of(".");
|
||||
std::string base_path = cli_params.preview_path;
|
||||
@@ -753,6 +817,16 @@ int main(int argc, const char* argv[]) {
|
||||
}
|
||||
}
|
||||
|
||||
if (gen_params.ip_adapter_image_path.size() > 0) {
|
||||
if (!load_sd_image_from_file(gen_params.ip_adapter_image.put(),
|
||||
gen_params.ip_adapter_image_path.c_str(),
|
||||
0,
|
||||
0)) {
|
||||
LOG_ERROR("load image from '%s' failed", gen_params.ip_adapter_image_path.c_str());
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
|
||||
if (!gen_params.control_video_path.empty()) {
|
||||
gen_params.control_frames.clear();
|
||||
if (!load_images_from_dir(gen_params.control_video_path,
|
||||
@@ -806,15 +880,22 @@ int main(int argc, const char* argv[]) {
|
||||
gen_params.sample_params.scheduler = sd_get_default_scheduler(sd_ctx.get(), gen_params.sample_params.sample_method);
|
||||
}
|
||||
|
||||
if (cli_params.mode == IMG_GEN) {
|
||||
sd_img_gen_params_t img_gen_params = gen_params.to_sd_img_gen_params_t();
|
||||
sd_img_gen_params_t img_gen_params{};
|
||||
const bool use_img_gen_params = cli_params.mode == IMG_GEN || cli_params.mode == ADETAILER;
|
||||
if (use_img_gen_params) {
|
||||
img_gen_params = gen_params.to_sd_img_gen_params_t();
|
||||
}
|
||||
|
||||
if (cli_params.mode == IMG_GEN) {
|
||||
sd_image_t* generated_images = nullptr;
|
||||
if (!generate_image(sd_ctx.get(), &img_gen_params, &generated_images, &num_results)) {
|
||||
generated_images = nullptr;
|
||||
num_results = 0;
|
||||
}
|
||||
results.adopt(generated_images, num_results);
|
||||
} else if (cli_params.mode == ADETAILER) {
|
||||
num_results = 1;
|
||||
results.push_back(gen_params.init_image.release());
|
||||
} else if (cli_params.mode == VID_GEN) {
|
||||
sd_vid_gen_params_t vid_gen_params = gen_params.to_sd_vid_gen_params_t();
|
||||
sd_image_t* generated_video = nullptr;
|
||||
@@ -828,6 +909,18 @@ int main(int argc, const char* argv[]) {
|
||||
LOG_ERROR("generate failed");
|
||||
return 1;
|
||||
}
|
||||
|
||||
if (use_img_gen_params &&
|
||||
!apply_adetailer(sd_ctx.get(),
|
||||
sd_ctx_params,
|
||||
ctx_params,
|
||||
gen_params,
|
||||
img_gen_params,
|
||||
cli_params.mode,
|
||||
results,
|
||||
num_results)) {
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
|
||||
int upscale_factor = 4; // unused for RealESRGAN_x4plus_anime_6B.pth
|
||||
|
||||
@@ -30,6 +30,7 @@ namespace fs = std::filesystem;
|
||||
|
||||
const char* const modes_str[] = {
|
||||
"img_gen",
|
||||
"adetailer",
|
||||
"vid_gen",
|
||||
"convert",
|
||||
"upscale",
|
||||
@@ -49,6 +50,9 @@ static sd_vae_format_t str_to_vae_format(const std::string& value) {
|
||||
if (value == "flux2") {
|
||||
return SD_VAE_FORMAT_FLUX2;
|
||||
}
|
||||
if (value == "wan") {
|
||||
return SD_VAE_FORMAT_WAN;
|
||||
}
|
||||
return SD_VAE_FORMAT_COUNT;
|
||||
}
|
||||
|
||||
@@ -400,7 +404,7 @@ ArgOptions SDContextParams::get_options() {
|
||||
&vae_path},
|
||||
{"",
|
||||
"--vae-format",
|
||||
"VAE latent format override: auto, flux, sd3, or flux2 (default: auto)",
|
||||
"VAE latent format override: auto, flux, sd3, flux2, or wan (default: auto)",
|
||||
0,
|
||||
&vae_format},
|
||||
{"",
|
||||
@@ -423,6 +427,16 @@ ArgOptions SDContextParams::get_options() {
|
||||
"path to control net model",
|
||||
0,
|
||||
&control_net_path},
|
||||
{"",
|
||||
"--ip-adapter",
|
||||
"path to IP-Adapter model (requires --clip_vision)",
|
||||
0,
|
||||
&ip_adapter_path},
|
||||
{"",
|
||||
"--motion-module",
|
||||
"path to AnimateDiff motion module (SD 1.5); enables video generation on --video-frames > 1",
|
||||
0,
|
||||
&motion_module_path},
|
||||
{"",
|
||||
"--embd-dir",
|
||||
"embeddings directory",
|
||||
@@ -443,6 +457,12 @@ ArgOptions SDContextParams::get_options() {
|
||||
"weight type per tensor pattern (example: \"^vae\\.=f16,model\\.=q8_0\")",
|
||||
(int)',',
|
||||
&tensor_type_rules},
|
||||
{"",
|
||||
"--model-args",
|
||||
"extra model args, key=value list. Supports chroma_use_dit_mask, chroma_use_t5_mask, "
|
||||
"chroma_t5_mask_pad, qwen_image_zero_cond_t",
|
||||
(int)',',
|
||||
&model_args},
|
||||
{"",
|
||||
"--photo-maker",
|
||||
"path to PHOTOMAKER model",
|
||||
@@ -493,10 +513,6 @@ ArgOptions SDContextParams::get_options() {
|
||||
"number of threads to use during computation (default: -1). "
|
||||
"If threads <= 0, then threads will be set to the number of CPU physical cores",
|
||||
&n_threads},
|
||||
{"",
|
||||
"--chroma-t5-mask-pad",
|
||||
"t5 mask pad size of chroma",
|
||||
&chroma_t5_mask_pad},
|
||||
};
|
||||
|
||||
options.bool_options = {
|
||||
@@ -554,30 +570,6 @@ ArgOptions SDContextParams::get_options() {
|
||||
"--vae-conv-direct",
|
||||
"use ggml_conv2d_direct in the vae model",
|
||||
true, &vae_conv_direct},
|
||||
{"",
|
||||
"--circular",
|
||||
"enable circular padding for convolutions",
|
||||
true, &circular},
|
||||
{"",
|
||||
"--circularx",
|
||||
"enable circular RoPE wrapping on x-axis (width) only",
|
||||
true, &circular_x},
|
||||
{"",
|
||||
"--circulary",
|
||||
"enable circular RoPE wrapping on y-axis (height) only",
|
||||
true, &circular_y},
|
||||
{"",
|
||||
"--chroma-disable-dit-mask",
|
||||
"disable dit mask for chroma",
|
||||
false, &chroma_use_dit_mask},
|
||||
{"",
|
||||
"--qwen-image-zero-cond-t",
|
||||
"enable zero_cond_t for qwen image",
|
||||
true, &qwen_image_zero_cond_t},
|
||||
{"",
|
||||
"--chroma-enable-t5-mask",
|
||||
"enable t5 mask for chroma",
|
||||
true, &chroma_use_t5_mask},
|
||||
};
|
||||
|
||||
auto on_type_arg = [&](int argc, const char** argv, int index) {
|
||||
@@ -694,7 +686,7 @@ ArgOptions SDContextParams::get_options() {
|
||||
}
|
||||
|
||||
void SDContextParams::build_embedding_map() {
|
||||
static const std::vector<std::string> valid_ext = {".gguf", ".safetensors", ".pt"};
|
||||
static const std::vector<std::string> valid_ext = {".gguf", ".safetensors", ".pt", ".ckpt"};
|
||||
|
||||
if (!fs::exists(embedding_dir) || !fs::is_directory(embedding_dir)) {
|
||||
return;
|
||||
@@ -759,7 +751,7 @@ bool SDContextParams::validate(SDMode mode) {
|
||||
}
|
||||
|
||||
if (str_to_vae_format(vae_format) == SD_VAE_FORMAT_COUNT) {
|
||||
LOG_ERROR("error: vae_format must be 'auto', 'flux', 'sd3', or 'flux2'");
|
||||
LOG_ERROR("error: vae_format must be 'auto', 'flux', 'sd3', 'flux2', or 'wan'");
|
||||
return false;
|
||||
}
|
||||
|
||||
@@ -844,6 +836,7 @@ std::string SDContextParams::to_string() const {
|
||||
<< " backend: \"" << backend << "\",\n"
|
||||
<< " params_backend: \"" << params_backend << "\",\n"
|
||||
<< " split_mode: \"" << split_mode << "\",\n"
|
||||
<< " model_args: \"" << model_args << "\",\n"
|
||||
<< " auto_fit: " << (auto_fit ? "true" : "false") << ",\n"
|
||||
<< " enable_mmap: " << (enable_mmap ? "true" : "false") << ",\n"
|
||||
<< " control_net_cpu: " << (control_net_cpu ? "true" : "false") << ",\n"
|
||||
@@ -853,13 +846,6 @@ std::string SDContextParams::to_string() const {
|
||||
<< " diffusion_flash_attn: " << (diffusion_flash_attn ? "true" : "false") << ",\n"
|
||||
<< " diffusion_conv_direct: " << (diffusion_conv_direct ? "true" : "false") << ",\n"
|
||||
<< " vae_conv_direct: " << (vae_conv_direct ? "true" : "false") << ",\n"
|
||||
<< " circular: " << (circular ? "true" : "false") << ",\n"
|
||||
<< " circular_x: " << (circular_x ? "true" : "false") << ",\n"
|
||||
<< " circular_y: " << (circular_y ? "true" : "false") << ",\n"
|
||||
<< " chroma_use_dit_mask: " << (chroma_use_dit_mask ? "true" : "false") << ",\n"
|
||||
<< " qwen_image_zero_cond_t: " << (qwen_image_zero_cond_t ? "true" : "false") << ",\n"
|
||||
<< " chroma_use_t5_mask: " << (chroma_use_t5_mask ? "true" : "false") << ",\n"
|
||||
<< " chroma_t5_mask_pad: " << chroma_t5_mask_pad << ",\n"
|
||||
<< " prediction: " << sd_prediction_name(prediction) << ",\n"
|
||||
<< " lora_apply_mode: " << sd_lora_apply_mode_name(lora_apply_mode) << ",\n"
|
||||
<< " force_sdxl_vae_conv_scale: " << (force_sdxl_vae_conv_scale ? "true" : "false") << "\n"
|
||||
@@ -895,6 +881,8 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) {
|
||||
sd_ctx_params.audio_vae_path = audio_vae_path.c_str();
|
||||
sd_ctx_params.taesd_path = taesd_path.c_str();
|
||||
sd_ctx_params.control_net_path = control_net_path.c_str();
|
||||
sd_ctx_params.ip_adapter_path = ip_adapter_path.c_str();
|
||||
sd_ctx_params.motion_module_path = motion_module_path.c_str();
|
||||
sd_ctx_params.embeddings = embedding_vec.data();
|
||||
sd_ctx_params.embedding_count = static_cast<uint32_t>(embedding_vec.size());
|
||||
sd_ctx_params.photo_maker_path = photo_maker_path.c_str();
|
||||
@@ -912,13 +900,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) {
|
||||
sd_ctx_params.tae_preview_only = taesd_preview;
|
||||
sd_ctx_params.diffusion_conv_direct = diffusion_conv_direct;
|
||||
sd_ctx_params.vae_conv_direct = vae_conv_direct;
|
||||
sd_ctx_params.circular_x = circular || circular_x;
|
||||
sd_ctx_params.circular_y = circular || circular_y;
|
||||
sd_ctx_params.force_sdxl_vae_conv_scale = force_sdxl_vae_conv_scale;
|
||||
sd_ctx_params.chroma_use_dit_mask = chroma_use_dit_mask;
|
||||
sd_ctx_params.chroma_use_t5_mask = chroma_use_t5_mask;
|
||||
sd_ctx_params.chroma_t5_mask_pad = chroma_t5_mask_pad;
|
||||
sd_ctx_params.qwen_image_zero_cond_t = qwen_image_zero_cond_t;
|
||||
sd_ctx_params.vae_format = str_to_vae_format(vae_format);
|
||||
sd_ctx_params.max_vram = max_vram.c_str();
|
||||
sd_ctx_params.stream_layers = stream_layers;
|
||||
@@ -928,6 +910,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) {
|
||||
sd_ctx_params.split_mode = split_mode.c_str();
|
||||
sd_ctx_params.auto_fit = auto_fit;
|
||||
sd_ctx_params.rpc_servers = rpc_servers.c_str();
|
||||
sd_ctx_params.model_args = model_args.empty() ? nullptr : model_args.c_str();
|
||||
return sd_ctx_params;
|
||||
}
|
||||
|
||||
@@ -949,6 +932,26 @@ ArgOptions SDGenerationParams::get_options() {
|
||||
"the negative prompt (default: \"\")",
|
||||
0,
|
||||
&negative_prompt},
|
||||
{"",
|
||||
"--ad-model",
|
||||
"path to a converted YOLOv8 detection model for ADetailer",
|
||||
0,
|
||||
&ad_model_path},
|
||||
{"",
|
||||
"--ad-prompt",
|
||||
"ADetailer prompt; empty inherits the main prompt, supports [PROMPT], [SEP], and [SKIP]",
|
||||
0,
|
||||
&ad_prompt},
|
||||
{"",
|
||||
"--ad-negative-prompt",
|
||||
"ADetailer negative prompt; empty inherits the main negative prompt, supports [PROMPT] and [SEP]",
|
||||
0,
|
||||
&ad_negative_prompt},
|
||||
{"",
|
||||
"--extra-ad-args",
|
||||
"extra ADetailer args, key=value list. Supports input_size, confidence, nms, max_detections, mask_k_largest, mask_min_ratio, mask_max_ratio, dilate_erode, x_offset, y_offset, mask_mode, merge_masks, invert_mask, mask_blur, inpaint_padding, inpaint_width, inpaint_height, denoising_strength, steps, cfg_scale, sample_method, scheduler, sort_by",
|
||||
(int)',',
|
||||
&extra_ad_args},
|
||||
{"-i",
|
||||
"--init-img",
|
||||
"path to the init image",
|
||||
@@ -969,6 +972,11 @@ ArgOptions SDGenerationParams::get_options() {
|
||||
"path to control image, control net",
|
||||
0,
|
||||
&control_image_path},
|
||||
{"",
|
||||
"--ip-adapter-image",
|
||||
"path to the IP-Adapter reference image",
|
||||
0,
|
||||
&ip_adapter_image_path},
|
||||
{"",
|
||||
"--control-video",
|
||||
"path to control video frames, It must be a directory path. The video frames inside should be stored as images in "
|
||||
@@ -1008,6 +1016,11 @@ ArgOptions SDGenerationParams::get_options() {
|
||||
"extra VAE tiling args, key=value list. LTX video VAE supports temporal_tile_frames (default: 4), temporal_tile_overlap (default: 1)",
|
||||
(int)',',
|
||||
&extra_tiling_args},
|
||||
{"",
|
||||
"--ref-image-args",
|
||||
"Key-value list to set up the way the reference images are processed (empty = auto-detect from model weigths)",
|
||||
(int)',',
|
||||
&ref_image_args},
|
||||
};
|
||||
|
||||
options.int_options = {
|
||||
@@ -1156,6 +1169,10 @@ ArgOptions SDGenerationParams::get_options() {
|
||||
"--control-strength",
|
||||
"strength to apply Control Net (default: 0.9). 1.0 corresponds to full destruction of information in init image",
|
||||
&control_strength},
|
||||
{"",
|
||||
"--ip-adapter-strength",
|
||||
"strength to apply IP-Adapter (default: 1.0)",
|
||||
&ip_adapter_strength},
|
||||
{"",
|
||||
"--moe-boundary",
|
||||
"timestep boundary for Wan2.2 MoE model. (default: 0.875). Only enabled if `--high-noise-steps` is set to -1",
|
||||
@@ -1189,6 +1206,18 @@ ArgOptions SDGenerationParams::get_options() {
|
||||
"disable auto resize of ref images",
|
||||
false,
|
||||
&auto_resize_ref_image},
|
||||
{"",
|
||||
"--circular",
|
||||
"enable circular padding on both axes for tileable output",
|
||||
true, &circular},
|
||||
{"",
|
||||
"--circularx",
|
||||
"enable circular padding on x-axis (width) only",
|
||||
true, &circular_x},
|
||||
{"",
|
||||
"--circulary",
|
||||
"enable circular padding on y-axis (height) only",
|
||||
true, &circular_y},
|
||||
{"",
|
||||
"--disable-image-metadata",
|
||||
"do not embed generation metadata on image files",
|
||||
@@ -1852,6 +1881,10 @@ bool SDGenerationParams::from_json_str(
|
||||
|
||||
load_if_exists("prompt", prompt);
|
||||
load_if_exists("negative_prompt", negative_prompt);
|
||||
load_if_exists("ad_model", ad_model_path);
|
||||
load_if_exists("ad_prompt", ad_prompt);
|
||||
load_if_exists("ad_negative_prompt", ad_negative_prompt);
|
||||
load_if_exists("extra_ad_args", extra_ad_args);
|
||||
load_if_exists("cache_mode", cache_mode);
|
||||
load_if_exists("cache_option", cache_option);
|
||||
load_if_exists("scm_mask", scm_mask);
|
||||
@@ -1868,6 +1901,7 @@ bool SDGenerationParams::from_json_str(
|
||||
|
||||
load_if_exists("strength", strength);
|
||||
load_if_exists("control_strength", control_strength);
|
||||
load_if_exists("ip_adapter_strength", ip_adapter_strength);
|
||||
load_if_exists("moe_boundary", moe_boundary);
|
||||
load_if_exists("vace_strength", vace_strength);
|
||||
|
||||
@@ -2039,6 +2073,10 @@ bool SDGenerationParams::from_json_str(
|
||||
LOG_ERROR("invalid control_image");
|
||||
return false;
|
||||
}
|
||||
if (!parse_image_json_field(j, "ip_adapter_image", 3, width, height, ip_adapter_image)) {
|
||||
LOG_ERROR("invalid ip_adapter_image");
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
@@ -2048,7 +2086,7 @@ void SDGenerationParams::extract_and_remove_lora(const std::string& lora_model_d
|
||||
return;
|
||||
}
|
||||
static const std::regex re(R"(<lora:([^:>]+):([^>]+)>)");
|
||||
static const std::vector<std::string> valid_ext = {".gguf", ".safetensors", ".pt"};
|
||||
static const std::vector<std::string> valid_ext = {".gguf", ".safetensors", ".pt", ".ckpt"};
|
||||
std::smatch m;
|
||||
|
||||
std::string tmp = prompt;
|
||||
@@ -2368,13 +2406,19 @@ bool SDGenerationParams::validate(SDMode mode) {
|
||||
}
|
||||
}
|
||||
|
||||
if (mode == UPSCALE) {
|
||||
if (mode == UPSCALE || mode == ADETAILER) {
|
||||
if (init_image_path.length() == 0) {
|
||||
LOG_ERROR("error: upscale mode needs an init image (--init-img)\n");
|
||||
LOG_ERROR("error: %s mode needs an init image (--init-img)\n",
|
||||
mode == UPSCALE ? "upscale" : "adetailer");
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
if (mode == ADETAILER && ad_model_path.empty()) {
|
||||
LOG_ERROR("error: adetailer mode needs a detector model (--ad-model)\n");
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
@@ -2439,30 +2483,47 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
|
||||
pulid_id_weight,
|
||||
};
|
||||
|
||||
params.loras = lora_vec.empty() ? nullptr : lora_vec.data();
|
||||
params.lora_count = static_cast<uint32_t>(lora_vec.size());
|
||||
params.prompt = prompt.c_str();
|
||||
params.negative_prompt = negative_prompt.c_str();
|
||||
params.clip_skip = clip_skip;
|
||||
params.init_image = init_image.get();
|
||||
params.ref_images = ref_image_views.empty() ? nullptr : ref_image_views.data();
|
||||
params.ref_images_count = static_cast<int>(ref_image_views.size());
|
||||
params.auto_resize_ref_image = auto_resize_ref_image;
|
||||
params.increase_ref_index = increase_ref_index;
|
||||
params.mask_image = mask_image.get();
|
||||
params.width = get_resolved_width();
|
||||
params.height = get_resolved_height();
|
||||
params.sample_params = sample_params;
|
||||
params.strength = strength;
|
||||
params.seed = seed;
|
||||
params.batch_count = batch_count;
|
||||
params.qwen_image_layers = qwen_image_layers;
|
||||
params.control_image = control_image.get();
|
||||
params.control_strength = control_strength;
|
||||
params.pm_params = pm_params;
|
||||
params.pulid_params = pulid_params;
|
||||
params.vae_tiling_params = vae_tiling_params;
|
||||
params.cache = cache_params;
|
||||
if (!auto_resize_ref_image) {
|
||||
if (!ref_image_args.empty()) {
|
||||
ref_image_args += ",";
|
||||
}
|
||||
ref_image_args += "resize_before_vae=0";
|
||||
LOG_WARN("Notice: --disable-auto-resize-ref-image is deprecated. Use --ref-image-args \"resize_before_vae=off\" instead.");
|
||||
}
|
||||
|
||||
if (increase_ref_index) {
|
||||
if (!ref_image_args.empty()) {
|
||||
ref_image_args += ",";
|
||||
}
|
||||
ref_image_args += "ref_index_mode=increase";
|
||||
LOG_WARN("Notice: --increase-ref-index is deprecated. Use --ref-image-args \"ref_index_mode=increase\" instead.");
|
||||
}
|
||||
|
||||
params.loras = lora_vec.empty() ? nullptr : lora_vec.data();
|
||||
params.lora_count = static_cast<uint32_t>(lora_vec.size());
|
||||
params.prompt = prompt.c_str();
|
||||
params.negative_prompt = negative_prompt.c_str();
|
||||
params.clip_skip = clip_skip;
|
||||
params.init_image = init_image.get();
|
||||
params.ref_images = ref_image_views.empty() ? nullptr : ref_image_views.data();
|
||||
params.ref_images_count = static_cast<int>(ref_image_views.size());
|
||||
params.ref_image_args = ref_image_args.c_str();
|
||||
params.mask_image = mask_image.get();
|
||||
params.width = get_resolved_width();
|
||||
params.height = get_resolved_height();
|
||||
params.sample_params = sample_params;
|
||||
params.strength = strength;
|
||||
params.seed = seed;
|
||||
params.batch_count = batch_count;
|
||||
params.qwen_image_layers = qwen_image_layers;
|
||||
params.control_image = control_image.get();
|
||||
params.control_strength = control_strength;
|
||||
params.ip_adapter_image = ip_adapter_image.get();
|
||||
params.ip_adapter_strength = ip_adapter_strength;
|
||||
params.pm_params = pm_params;
|
||||
params.pulid_params = pulid_params;
|
||||
params.vae_tiling_params = vae_tiling_params;
|
||||
params.cache = cache_params;
|
||||
|
||||
params.hires.enabled = hires_enabled;
|
||||
params.hires.upscaler = resolved_hires_upscaler;
|
||||
@@ -2475,6 +2536,8 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
|
||||
params.hires.upscale_tile_size = hires_upscale_tile_size;
|
||||
params.hires.custom_sigmas = hires_custom_sigmas.empty() ? nullptr : hires_custom_sigmas.data();
|
||||
params.hires.custom_sigmas_count = static_cast<int>(hires_custom_sigmas.size());
|
||||
params.circular_x = circular || circular_x;
|
||||
params.circular_y = circular || circular_y;
|
||||
return params;
|
||||
}
|
||||
|
||||
@@ -2540,6 +2603,8 @@ sd_vid_gen_params_t SDGenerationParams::to_sd_vid_gen_params_t() {
|
||||
params.hires.upscale_tile_size = hires_upscale_tile_size;
|
||||
params.hires.custom_sigmas = hires_custom_sigmas.empty() ? nullptr : hires_custom_sigmas.data();
|
||||
params.hires.custom_sigmas_count = static_cast<int>(hires_custom_sigmas.size());
|
||||
params.circular_x = circular || circular_x;
|
||||
params.circular_y = circular || circular_y;
|
||||
return params;
|
||||
}
|
||||
|
||||
@@ -2578,6 +2643,10 @@ std::string SDGenerationParams::to_string() const {
|
||||
<< " high_noise_loras: \"" << high_noise_loras_str << "\",\n"
|
||||
<< " prompt: \"" << prompt << "\",\n"
|
||||
<< " negative_prompt: \"" << negative_prompt << "\",\n"
|
||||
<< " ad_model_path: \"" << ad_model_path << "\",\n"
|
||||
<< " ad_prompt: \"" << ad_prompt << "\",\n"
|
||||
<< " ad_negative_prompt: \"" << ad_negative_prompt << "\",\n"
|
||||
<< " extra_ad_args: \"" << extra_ad_args << "\",\n"
|
||||
<< " clip_skip: " << clip_skip << ",\n"
|
||||
<< " width: " << width << ",\n"
|
||||
<< " height: " << height << ",\n"
|
||||
@@ -2692,8 +2761,13 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params,
|
||||
int64_t seed,
|
||||
SDMode mode) {
|
||||
json root;
|
||||
root["schema"] = "sdcpp.image.params/v1";
|
||||
root["mode"] = mode == VID_GEN ? "vid_gen" : "img_gen";
|
||||
root["schema"] = "sdcpp.image.params/v1";
|
||||
root["mode"] = "img_gen";
|
||||
if (mode == VID_GEN) {
|
||||
root["mode"] = "vid_gen";
|
||||
} else if (mode == ADETAILER) {
|
||||
root["mode"] = "adetailer";
|
||||
}
|
||||
root["generator"] = {
|
||||
{"name", "stable-diffusion.cpp"},
|
||||
{"version", safe_json_string(sd_version())},
|
||||
@@ -2707,6 +2781,14 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params,
|
||||
{"positive", gen_params.prompt},
|
||||
{"negative", gen_params.negative_prompt},
|
||||
};
|
||||
if (!gen_params.ad_model_path.empty()) {
|
||||
root["adetailer"] = {
|
||||
{"model", sd_basename(gen_params.ad_model_path)},
|
||||
{"prompt", gen_params.ad_prompt},
|
||||
{"negative_prompt", gen_params.ad_negative_prompt},
|
||||
{"extra_args", gen_params.extra_ad_args},
|
||||
};
|
||||
}
|
||||
root["sampling"] = build_sampling_metadata_json(gen_params.sample_params,
|
||||
gen_params.skip_layers,
|
||||
&gen_params.custom_sigmas);
|
||||
@@ -2730,6 +2812,7 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params,
|
||||
root["clip_skip"] = gen_params.clip_skip;
|
||||
root["strength"] = gen_params.strength;
|
||||
root["control_strength"] = gen_params.control_strength;
|
||||
root["ip_adapter_strength"] = gen_params.ip_adapter_strength;
|
||||
root["auto_resize_ref_image"] = gen_params.auto_resize_ref_image;
|
||||
root["increase_ref_index"] = gen_params.increase_ref_index;
|
||||
if (mode == VID_GEN) {
|
||||
@@ -2861,6 +2944,18 @@ std::string get_image_params(const SDContextParams& ctx_params,
|
||||
if (!gen_params.extra_sample_args.empty()) {
|
||||
parameter_string += "Extra sample args: " + gen_params.extra_sample_args + ", ";
|
||||
}
|
||||
if (!gen_params.ad_model_path.empty()) {
|
||||
parameter_string += "ADetailer model: " + sd_basename(gen_params.ad_model_path) + ", ";
|
||||
if (!gen_params.ad_prompt.empty()) {
|
||||
parameter_string += "ADetailer prompt: " + gen_params.ad_prompt + ", ";
|
||||
}
|
||||
if (!gen_params.ad_negative_prompt.empty()) {
|
||||
parameter_string += "ADetailer negative prompt: " + gen_params.ad_negative_prompt + ", ";
|
||||
}
|
||||
if (!gen_params.extra_ad_args.empty()) {
|
||||
parameter_string += "ADetailer args: " + gen_params.extra_ad_args + ", ";
|
||||
}
|
||||
}
|
||||
parameter_string += "Seed: " + std::to_string(seed) + ", ";
|
||||
parameter_string += "Size: " + std::to_string(gen_params.get_resolved_width()) + "x" + std::to_string(gen_params.get_resolved_height()) + ", ";
|
||||
parameter_string += "Model: " + sd_basename(ctx_params.model_path) + ", ";
|
||||
|
||||
@@ -16,10 +16,11 @@
|
||||
#define BOOL_STR(b) ((b) ? "true" : "false")
|
||||
|
||||
extern const char* const modes_str[];
|
||||
#define SD_ALL_MODES_STR "img_gen, vid_gen, convert, upscale, metadata"
|
||||
#define SD_ALL_MODES_STR "img_gen, adetailer, vid_gen, convert, upscale, metadata"
|
||||
|
||||
enum SDMode {
|
||||
IMG_GEN,
|
||||
ADETAILER,
|
||||
VID_GEN,
|
||||
CONVERT,
|
||||
UPSCALE,
|
||||
@@ -132,6 +133,8 @@ struct SDContextParams {
|
||||
std::string taesd_path;
|
||||
std::string esrgan_path;
|
||||
std::string control_net_path;
|
||||
std::string ip_adapter_path;
|
||||
std::string motion_module_path;
|
||||
std::string embedding_dir;
|
||||
std::string photo_maker_path;
|
||||
std::string pulid_weights_path;
|
||||
@@ -152,6 +155,7 @@ struct SDContextParams {
|
||||
std::string backend;
|
||||
std::string params_backend;
|
||||
std::string split_mode;
|
||||
std::string model_args;
|
||||
bool auto_fit = false;
|
||||
std::string rpc_servers;
|
||||
std::string effective_backend;
|
||||
@@ -165,16 +169,6 @@ struct SDContextParams {
|
||||
bool diffusion_conv_direct = false;
|
||||
bool vae_conv_direct = false;
|
||||
|
||||
bool circular = false;
|
||||
bool circular_x = false;
|
||||
bool circular_y = false;
|
||||
|
||||
bool chroma_use_dit_mask = true;
|
||||
bool chroma_use_t5_mask = false;
|
||||
int chroma_t5_mask_pad = 1;
|
||||
|
||||
bool qwen_image_zero_cond_t = false;
|
||||
|
||||
prediction_t prediction = PREDICTION_COUNT;
|
||||
lora_apply_mode_t lora_apply_mode = LORA_APPLY_AUTO;
|
||||
|
||||
@@ -195,6 +189,10 @@ struct SDGenerationParams {
|
||||
// User-facing input fields.
|
||||
std::string prompt;
|
||||
std::string negative_prompt;
|
||||
std::string ad_model_path;
|
||||
std::string ad_prompt;
|
||||
std::string ad_negative_prompt;
|
||||
std::string extra_ad_args;
|
||||
int clip_skip = -1; // <= 0 represents unspecified
|
||||
int width = -1;
|
||||
int height = -1;
|
||||
@@ -203,6 +201,7 @@ struct SDGenerationParams {
|
||||
int64_t seed = 42;
|
||||
float strength = 0.75f;
|
||||
float control_strength = 0.9f;
|
||||
float ip_adapter_strength = 1.0f;
|
||||
bool auto_resize_ref_image = true;
|
||||
bool increase_ref_index = false;
|
||||
bool embed_image_metadata = true;
|
||||
@@ -211,6 +210,7 @@ struct SDGenerationParams {
|
||||
std::string end_image_path;
|
||||
std::string mask_image_path;
|
||||
std::string control_image_path;
|
||||
std::string ip_adapter_image_path;
|
||||
std::vector<std::string> ref_image_paths;
|
||||
std::string control_video_path;
|
||||
|
||||
@@ -236,6 +236,8 @@ struct SDGenerationParams {
|
||||
sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr};
|
||||
std::string extra_tiling_args;
|
||||
|
||||
std::string ref_image_args;
|
||||
|
||||
std::string pm_id_images_dir;
|
||||
std::string pm_id_embed_path;
|
||||
float pm_style_strength = 20.f;
|
||||
@@ -246,6 +248,10 @@ struct SDGenerationParams {
|
||||
int upscale_repeats = 1;
|
||||
int upscale_tile_size = 128;
|
||||
|
||||
bool circular = false;
|
||||
bool circular_x = false;
|
||||
bool circular_y = false;
|
||||
|
||||
bool hires_enabled = false;
|
||||
std::string hires_upscaler = "Latent";
|
||||
std::string hires_upscaler_model_path;
|
||||
@@ -271,6 +277,7 @@ struct SDGenerationParams {
|
||||
std::vector<SDImageOwner> ref_images;
|
||||
SDImageOwner mask_image;
|
||||
SDImageOwner control_image;
|
||||
SDImageOwner ip_adapter_image;
|
||||
std::vector<SDImageOwner> pm_id_images;
|
||||
std::vector<SDImageOwner> control_frames;
|
||||
|
||||
|
||||
@@ -40,12 +40,21 @@ struct UpscalerCtxDeleter {
|
||||
}
|
||||
};
|
||||
|
||||
struct ADetailerCtxDeleter {
|
||||
void operator()(adetailer_ctx_t* ctx) const {
|
||||
if (ctx != nullptr) {
|
||||
free_adetailer_ctx(ctx);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
using FreeUniquePtr = std::unique_ptr<T, FreeDeleter>;
|
||||
|
||||
using FilePtr = std::unique_ptr<FILE, FileCloser>;
|
||||
using SDCtxPtr = std::unique_ptr<sd_ctx_t, SDCtxDeleter>;
|
||||
using UpscalerCtxPtr = std::unique_ptr<upscaler_ctx_t, UpscalerCtxDeleter>;
|
||||
using FilePtr = std::unique_ptr<FILE, FileCloser>;
|
||||
using SDCtxPtr = std::unique_ptr<sd_ctx_t, SDCtxDeleter>;
|
||||
using UpscalerCtxPtr = std::unique_ptr<upscaler_ctx_t, UpscalerCtxDeleter>;
|
||||
using ADetailerCtxPtr = std::unique_ptr<adetailer_ctx_t, ADetailerCtxDeleter>;
|
||||
|
||||
class SDImageOwner {
|
||||
private:
|
||||
|
||||
@@ -528,6 +528,7 @@ Shared default fields used by both `img_gen` and `vid_gen`:
|
||||
| `auto_resize_ref_image` | `boolean` |
|
||||
| `increase_ref_index` | `boolean` |
|
||||
| `control_strength` | `number` |
|
||||
| `ip_adapter_strength` | `number` |
|
||||
| `hires` | `object` |
|
||||
| `hires.enabled` | `boolean` |
|
||||
| `hires.upscaler` | `string` |
|
||||
@@ -567,6 +568,7 @@ Fields returned in `features_by_mode.img_gen`:
|
||||
- `init_image`
|
||||
- `mask_image`
|
||||
- `control_image`
|
||||
- `ip_adapter_image`
|
||||
- `ref_images`
|
||||
- `lora`
|
||||
- `vae_tiling`
|
||||
@@ -653,12 +655,14 @@ Example:
|
||||
"auto_resize_ref_image": true,
|
||||
"increase_ref_index": false,
|
||||
"control_strength": 0.9,
|
||||
"ip_adapter_strength": 1.0,
|
||||
"embed_image_metadata": true,
|
||||
|
||||
"init_image": null,
|
||||
"ref_images": [],
|
||||
"mask_image": null,
|
||||
"control_image": null,
|
||||
"ip_adapter_image": null,
|
||||
|
||||
"sample_params": {
|
||||
"scheduler": "discrete",
|
||||
@@ -733,6 +737,7 @@ Channel expectations:
|
||||
- `init_image`: 3 channels
|
||||
- `ref_images[]`: 3 channels
|
||||
- `control_image`: 3 channels
|
||||
- `ip_adapter_image`: 3 channels
|
||||
- `mask_image`: 1 channel
|
||||
|
||||
If omitted or null:
|
||||
@@ -757,6 +762,7 @@ Top-level scalar fields:
|
||||
| `auto_resize_ref_image` | `boolean` |
|
||||
| `increase_ref_index` | `boolean` |
|
||||
| `control_strength` | `number` |
|
||||
| `ip_adapter_strength` | `number` |
|
||||
| `embed_image_metadata` | `boolean` |
|
||||
|
||||
Image fields:
|
||||
@@ -767,6 +773,7 @@ Image fields:
|
||||
| `ref_images` | `array<string>` |
|
||||
| `mask_image` | `string \| null` |
|
||||
| `control_image` | `string \| null` |
|
||||
| `ip_adapter_image` | `string \| null` |
|
||||
|
||||
LoRA fields:
|
||||
|
||||
@@ -958,7 +965,7 @@ Response fields:
|
||||
Compared with `img_gen`, the `vid_gen` request body:
|
||||
|
||||
- `vid_gen` is a single video sequence job, so `batch_count` is not part of the request schema
|
||||
- `ref_images`, `mask_image`, `control_image`, `control_strength`, and `embed_image_metadata` are not part of the request schema
|
||||
- `ref_images`, `mask_image`, `control_image`, `control_strength`, `ip_adapter_image`, `ip_adapter_strength`, and `embed_image_metadata` are not part of the request schema
|
||||
- `vid_gen` adds `end_image`, `control_frames`, `high_noise_sample_params`, `video_frames`, `fps`, `moe_boundary`, and `vace_strength`
|
||||
|
||||
Example:
|
||||
|
||||
@@ -130,6 +130,7 @@ static json make_img_gen_defaults_json(const SDGenerationParams& defaults, const
|
||||
{"auto_resize_ref_image", defaults.auto_resize_ref_image},
|
||||
{"increase_ref_index", defaults.increase_ref_index},
|
||||
{"control_strength", defaults.control_strength},
|
||||
{"ip_adapter_strength", defaults.ip_adapter_strength},
|
||||
{"sample_params", make_sample_params_json(defaults.sample_params, defaults.skip_layers)},
|
||||
{"hires", make_hires_json(defaults)},
|
||||
{"vae_tiling_params", make_vae_tiling_json(defaults.vae_tiling_params)},
|
||||
@@ -173,6 +174,7 @@ static json make_img_gen_features_json() {
|
||||
{"init_image", true},
|
||||
{"mask_image", true},
|
||||
{"control_image", true},
|
||||
{"ip_adapter_image", true},
|
||||
{"ref_images", true},
|
||||
{"lora", true},
|
||||
{"vae_tiling", true},
|
||||
|
||||
@@ -180,6 +180,7 @@ enum sd_vae_format_t {
|
||||
SD_VAE_FORMAT_FLUX,
|
||||
SD_VAE_FORMAT_SD3,
|
||||
SD_VAE_FORMAT_FLUX2,
|
||||
SD_VAE_FORMAT_WAN,
|
||||
SD_VAE_FORMAT_COUNT,
|
||||
};
|
||||
|
||||
@@ -199,6 +200,8 @@ typedef struct {
|
||||
const char* audio_vae_path;
|
||||
const char* taesd_path;
|
||||
const char* control_net_path;
|
||||
const char* ip_adapter_path;
|
||||
const char* motion_module_path;
|
||||
const sd_embedding_t* embeddings;
|
||||
uint32_t embedding_count;
|
||||
const char* photo_maker_path;
|
||||
@@ -216,13 +219,7 @@ typedef struct {
|
||||
bool tae_preview_only;
|
||||
bool diffusion_conv_direct;
|
||||
bool vae_conv_direct;
|
||||
bool circular_x;
|
||||
bool circular_y;
|
||||
bool force_sdxl_vae_conv_scale;
|
||||
bool chroma_use_dit_mask;
|
||||
bool chroma_use_t5_mask;
|
||||
int chroma_t5_mask_pad;
|
||||
bool qwen_image_zero_cond_t;
|
||||
enum sd_vae_format_t vae_format;
|
||||
const char* max_vram; // GiB budget or backend assignment spec for graph-cut segmented param offload (0 = disabled, -1 = auto)
|
||||
bool stream_layers; // Enable residency+prefetch streaming on top of --max-vram (no effect without --max-vram)
|
||||
@@ -232,6 +229,7 @@ typedef struct {
|
||||
const char* split_mode; // weight distribution for multi-device modules: layer (default) or row, or per-module assignments e.g. "diffusion=row"
|
||||
bool auto_fit;
|
||||
const char* rpc_servers;
|
||||
const char* model_args;
|
||||
} sd_ctx_params_t;
|
||||
|
||||
typedef struct {
|
||||
@@ -368,8 +366,7 @@ typedef struct {
|
||||
sd_image_t init_image;
|
||||
sd_image_t* ref_images;
|
||||
int ref_images_count;
|
||||
bool auto_resize_ref_image;
|
||||
bool increase_ref_index;
|
||||
const char* ref_image_args;
|
||||
sd_image_t mask_image;
|
||||
int width;
|
||||
int height;
|
||||
@@ -379,12 +376,16 @@ typedef struct {
|
||||
int batch_count;
|
||||
sd_image_t control_image;
|
||||
float control_strength;
|
||||
sd_image_t ip_adapter_image;
|
||||
float ip_adapter_strength;
|
||||
sd_pm_params_t pm_params;
|
||||
sd_pulid_params_t pulid_params;
|
||||
sd_tiling_params_t vae_tiling_params;
|
||||
sd_cache_params_t cache;
|
||||
sd_hires_params_t hires;
|
||||
int qwen_image_layers;
|
||||
bool circular_x;
|
||||
bool circular_y;
|
||||
} sd_img_gen_params_t;
|
||||
|
||||
typedef struct {
|
||||
@@ -410,6 +411,8 @@ typedef struct {
|
||||
sd_tiling_params_t vae_tiling_params;
|
||||
sd_cache_params_t cache;
|
||||
sd_hires_params_t hires;
|
||||
bool circular_x;
|
||||
bool circular_y;
|
||||
} sd_vid_gen_params_t;
|
||||
|
||||
typedef struct sd_ctx_t sd_ctx_t;
|
||||
@@ -429,6 +432,11 @@ SD_API const char* sd_get_system_info();
|
||||
SD_API bool sd_ctx_supports_image_generation(const sd_ctx_t* sd_ctx);
|
||||
SD_API bool sd_ctx_supports_video_generation(const sd_ctx_t* sd_ctx);
|
||||
|
||||
// ControlNet hot-swap APIs are not safe to call while generation is in flight.
|
||||
SD_API bool sd_ctx_load_control_net(sd_ctx_t* sd_ctx, const char* path);
|
||||
SD_API bool sd_ctx_unload_control_net(sd_ctx_t* sd_ctx);
|
||||
SD_API bool sd_ctx_has_control_net(const sd_ctx_t* sd_ctx);
|
||||
|
||||
SD_API const char* sd_type_name(enum sd_type_t type);
|
||||
SD_API enum sd_type_t str_to_sd_type(const char* str);
|
||||
SD_API const char* sd_rng_type_name(enum rng_type_t rng_type);
|
||||
@@ -505,6 +513,27 @@ SD_API bool upscale(upscaler_ctx_t* upscaler_ctx,
|
||||
|
||||
SD_API int get_upscale_factor(upscaler_ctx_t* upscaler_ctx);
|
||||
|
||||
typedef struct adetailer_ctx_t adetailer_ctx_t;
|
||||
|
||||
typedef struct {
|
||||
const char* prompt;
|
||||
const char* negative_prompt;
|
||||
const char* extra_ad_args;
|
||||
} sd_adetailer_params_t;
|
||||
|
||||
SD_API adetailer_ctx_t* new_adetailer_ctx(const char* detector_path,
|
||||
int n_threads,
|
||||
const char* backend,
|
||||
const char* params_backend);
|
||||
SD_API void free_adetailer_ctx(adetailer_ctx_t* adetailer_ctx);
|
||||
SD_API bool adetail_image(adetailer_ctx_t* adetailer_ctx,
|
||||
sd_ctx_t* sd_ctx,
|
||||
sd_image_t input_image,
|
||||
const sd_adetailer_params_t* adetailer_params,
|
||||
const sd_img_gen_params_t* inpaint_params,
|
||||
sd_image_t** images_out,
|
||||
int* num_images_out);
|
||||
|
||||
SD_API bool convert(const char* input_path,
|
||||
const char* vae_path,
|
||||
const char* output_path,
|
||||
|
||||
@@ -9,7 +9,7 @@ the text side; the vision side is converted by sd.cpp's own
|
||||
Operates on raw safetensors bytes so any dtype (BF16/F16/F32) is preserved.
|
||||
|
||||
Usage:
|
||||
python3 script/convert_qwen3_vl.py <hf_qwen3_vl_dir_or_safetensors> <output.safetensors>
|
||||
python3 scripts/convert_qwen3_vl.py <hf_qwen3_vl_dir_or_safetensors> <output.safetensors>
|
||||
"""
|
||||
|
||||
import argparse
|
||||
@@ -5,7 +5,7 @@ Operates on raw safetensors bytes so any dtype (BF16, F32, ...) is preserved exa
|
||||
No numpy or torch dependency required.
|
||||
|
||||
Usage:
|
||||
python3 script/convert_sefi.py <sefi_diffusers_dir> <output.safetensors>
|
||||
python3 scripts/convert_sefi.py <sefi_diffusers_dir> <output.safetensors>
|
||||
"""
|
||||
|
||||
import argparse
|
||||
@@ -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()
|
||||
@@ -31,24 +31,30 @@ $patterns = @(
|
||||
"examples/common/*.cpp"
|
||||
)
|
||||
|
||||
$root = (Get-Location).Path
|
||||
Push-Location (Join-Path $PSScriptRoot "..")
|
||||
|
||||
foreach ($pattern in $patterns) {
|
||||
$files = Get-ChildItem -Path $pattern -File -ErrorAction SilentlyContinue | Sort-Object FullName
|
||||
try {
|
||||
$root = (Get-Location).Path
|
||||
|
||||
foreach ($file in $files) {
|
||||
$relativePath = $file.FullName.Substring($root.Length).TrimStart('\', '/') -replace '\\', '/'
|
||||
foreach ($pattern in $patterns) {
|
||||
$files = Get-ChildItem -Path $pattern -File -ErrorAction SilentlyContinue | Sort-Object FullName
|
||||
|
||||
if ($relativePath -like "vocab*") {
|
||||
continue
|
||||
foreach ($file in $files) {
|
||||
$relativePath = $file.FullName.Substring($root.Length).TrimStart('\', '/') -replace '\\', '/'
|
||||
|
||||
if ($relativePath -like "vocab*") {
|
||||
continue
|
||||
}
|
||||
|
||||
Write-Host "formatting '$relativePath'"
|
||||
|
||||
# if ($relativePath -ne "stable-diffusion.h") {
|
||||
# clang-tidy -fix -p build_linux/ "$relativePath"
|
||||
# }
|
||||
|
||||
& clang-format -style=file -i $relativePath
|
||||
}
|
||||
|
||||
Write-Host "formatting '$relativePath'"
|
||||
|
||||
# if ($relativePath -ne "stable-diffusion.h") {
|
||||
# clang-tidy -fix -p build_linux/ "$relativePath"
|
||||
# }
|
||||
|
||||
& clang-format -style=file -i $relativePath
|
||||
}
|
||||
} finally {
|
||||
Pop-Location
|
||||
}
|
||||
@@ -1,3 +1,10 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
set -e
|
||||
|
||||
repo_root="$(cd "$(dirname "$0")/.." && pwd)"
|
||||
cd "$repo_root"
|
||||
|
||||
for f in src/*.cpp src/*.h src/*.hpp \
|
||||
src/conditioning/*.cpp src/conditioning/*.h src/conditioning/*.hpp \
|
||||
src/core/*.cpp src/core/*.h src/core/*.hpp \
|
||||
@@ -13,7 +13,7 @@ packaging quirks):
|
||||
- numpy, Pillow
|
||||
|
||||
Usage:
|
||||
python script/pulid_extract_id.py \\
|
||||
python scripts/pulid_extract_id.py \\
|
||||
--portrait /path/to/source-photo.jpg \\
|
||||
--pulid-weights /path/to/pulid_flux_v0.9.1.safetensors \\
|
||||
--out /path/to/source.pulidembd
|
||||
@@ -6,6 +6,8 @@
|
||||
#include <optional>
|
||||
|
||||
#include "core/tensor_ggml.hpp"
|
||||
#include "core/util.h"
|
||||
#include "model/diffusion/model.hpp"
|
||||
#include "model/te/clip.hpp"
|
||||
#include "model/te/llm.hpp"
|
||||
#include "model/te/t5.hpp"
|
||||
@@ -105,6 +107,7 @@ struct ConditionerParams {
|
||||
int height = -1;
|
||||
bool zero_out_masked = false;
|
||||
const std::vector<sd::Tensor<float>>* ref_images = nullptr; // for qwen image edit
|
||||
RefImageParams ref_image_params;
|
||||
};
|
||||
|
||||
struct Conditioner {
|
||||
@@ -114,9 +117,12 @@ public:
|
||||
virtual SDCondition get_learned_condition(int n_threads,
|
||||
const ConditionerParams& conditioner_params) = 0;
|
||||
virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) = 0;
|
||||
virtual void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) {}
|
||||
virtual void set_max_graph_vram_bytes(size_t max_vram_bytes) {}
|
||||
virtual void set_stream_layers_enabled(bool enabled) {}
|
||||
virtual void set_runtime_backends(const std::vector<ggml_backend_t>& backends) {}
|
||||
virtual void set_graph_cut_layer_split_enabled(bool enabled) {}
|
||||
virtual void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) {}
|
||||
virtual void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {}
|
||||
virtual void set_flash_attention_enabled(bool enabled) = 0;
|
||||
virtual void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) {}
|
||||
@@ -180,6 +186,27 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner {
|
||||
}
|
||||
}
|
||||
|
||||
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
|
||||
text_model->set_runtime_backends(backends);
|
||||
if (sd_version_is_sdxl(version)) {
|
||||
text_model2->set_runtime_backends(backends);
|
||||
}
|
||||
}
|
||||
|
||||
void set_graph_cut_layer_split_enabled(bool enabled) override {
|
||||
text_model->set_graph_cut_layer_split_enabled(enabled);
|
||||
if (sd_version_is_sdxl(version)) {
|
||||
text_model2->set_graph_cut_layer_split_enabled(enabled);
|
||||
}
|
||||
}
|
||||
|
||||
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
|
||||
text_model->set_graph_cut_layer_split_backend_vram_limits(limits);
|
||||
if (sd_version_is_sdxl(version)) {
|
||||
text_model2->set_graph_cut_layer_split_backend_vram_limits(limits);
|
||||
}
|
||||
}
|
||||
|
||||
void set_flash_attention_enabled(bool enabled) override {
|
||||
text_model->set_flash_attention_enabled(enabled);
|
||||
if (sd_version_is_sdxl(version)) {
|
||||
@@ -638,11 +665,41 @@ struct SD3CLIPEmbedder : public Conditioner {
|
||||
}
|
||||
|
||||
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
|
||||
if (clip_l) {
|
||||
clip_l->set_runtime_backends(backends);
|
||||
}
|
||||
if (clip_g) {
|
||||
clip_g->set_runtime_backends(backends);
|
||||
}
|
||||
if (t5) {
|
||||
t5->set_runtime_backends(backends);
|
||||
}
|
||||
}
|
||||
|
||||
void set_graph_cut_layer_split_enabled(bool enabled) override {
|
||||
if (clip_l) {
|
||||
clip_l->set_graph_cut_layer_split_enabled(enabled);
|
||||
}
|
||||
if (clip_g) {
|
||||
clip_g->set_graph_cut_layer_split_enabled(enabled);
|
||||
}
|
||||
if (t5) {
|
||||
t5->set_graph_cut_layer_split_enabled(enabled);
|
||||
}
|
||||
}
|
||||
|
||||
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
|
||||
if (clip_l) {
|
||||
clip_l->set_graph_cut_layer_split_backend_vram_limits(limits);
|
||||
}
|
||||
if (clip_g) {
|
||||
clip_g->set_graph_cut_layer_split_backend_vram_limits(limits);
|
||||
}
|
||||
if (t5) {
|
||||
t5->set_graph_cut_layer_split_backend_vram_limits(limits);
|
||||
}
|
||||
}
|
||||
|
||||
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
|
||||
if (t5) {
|
||||
t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
|
||||
@@ -1009,11 +1066,32 @@ struct FluxCLIPEmbedder : public Conditioner {
|
||||
}
|
||||
|
||||
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
|
||||
if (clip_l) {
|
||||
clip_l->set_runtime_backends(backends);
|
||||
}
|
||||
if (t5) {
|
||||
t5->set_runtime_backends(backends);
|
||||
}
|
||||
}
|
||||
|
||||
void set_graph_cut_layer_split_enabled(bool enabled) override {
|
||||
if (clip_l) {
|
||||
clip_l->set_graph_cut_layer_split_enabled(enabled);
|
||||
}
|
||||
if (t5) {
|
||||
t5->set_graph_cut_layer_split_enabled(enabled);
|
||||
}
|
||||
}
|
||||
|
||||
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
|
||||
if (clip_l) {
|
||||
clip_l->set_graph_cut_layer_split_backend_vram_limits(limits);
|
||||
}
|
||||
if (t5) {
|
||||
t5->set_graph_cut_layer_split_backend_vram_limits(limits);
|
||||
}
|
||||
}
|
||||
|
||||
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
|
||||
if (t5) {
|
||||
t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
|
||||
@@ -1217,8 +1295,27 @@ struct T5CLIPEmbedder : public Conditioner {
|
||||
bool use_mask = false,
|
||||
int mask_pad = 0,
|
||||
bool is_umt5 = false,
|
||||
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
|
||||
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
|
||||
const char* model_args = nullptr)
|
||||
: use_mask(use_mask), mask_pad(mask_pad), t5_tokenizer(is_umt5) {
|
||||
for (const auto& [key, value] : parse_key_value_args(model_args, "model arg")) {
|
||||
if (key == "chroma_use_t5_mask") {
|
||||
bool parsed = false;
|
||||
if (parse_strict_bool(value, parsed)) {
|
||||
this->use_mask = parsed;
|
||||
} else {
|
||||
LOG_WARN("ignoring invalid Chroma T5 model arg '%s=%s'", key.c_str(), value.c_str());
|
||||
}
|
||||
} else if (key == "chroma_t5_mask_pad") {
|
||||
int parsed = 0;
|
||||
if (parse_strict_int(value, parsed)) {
|
||||
this->mask_pad = parsed;
|
||||
} else {
|
||||
LOG_WARN("ignoring invalid Chroma T5 model arg '%s=%s'", key.c_str(), value.c_str());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
bool use_t5 = false;
|
||||
for (auto pair : tensor_storage_map) {
|
||||
if (pair.first.find("text_encoders.t5xxl") != std::string::npos) {
|
||||
@@ -1258,6 +1355,18 @@ struct T5CLIPEmbedder : public Conditioner {
|
||||
}
|
||||
}
|
||||
|
||||
void set_graph_cut_layer_split_enabled(bool enabled) override {
|
||||
if (t5) {
|
||||
t5->set_graph_cut_layer_split_enabled(enabled);
|
||||
}
|
||||
}
|
||||
|
||||
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
|
||||
if (t5) {
|
||||
t5->set_graph_cut_layer_split_backend_vram_limits(limits);
|
||||
}
|
||||
}
|
||||
|
||||
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
|
||||
if (t5) {
|
||||
t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
|
||||
@@ -1462,6 +1571,18 @@ struct MiniT2IConditioner : public Conditioner {
|
||||
}
|
||||
}
|
||||
|
||||
void set_graph_cut_layer_split_enabled(bool enabled) override {
|
||||
if (t5) {
|
||||
t5->set_graph_cut_layer_split_enabled(enabled);
|
||||
}
|
||||
}
|
||||
|
||||
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
|
||||
if (t5) {
|
||||
t5->set_graph_cut_layer_split_backend_vram_limits(limits);
|
||||
}
|
||||
}
|
||||
|
||||
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
|
||||
if (t5) {
|
||||
t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
|
||||
@@ -1544,6 +1665,10 @@ struct AnimaConditioner : public Conditioner {
|
||||
llm->get_param_tensors(tensors, "text_encoders.llm");
|
||||
}
|
||||
|
||||
void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) override {
|
||||
llm->get_param_tensor_ops(tensor_ops);
|
||||
}
|
||||
|
||||
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
|
||||
llm->set_max_graph_vram_bytes(max_vram_bytes);
|
||||
}
|
||||
@@ -1556,6 +1681,14 @@ struct AnimaConditioner : public Conditioner {
|
||||
llm->set_runtime_backends(backends);
|
||||
}
|
||||
|
||||
void set_graph_cut_layer_split_enabled(bool enabled) override {
|
||||
llm->set_graph_cut_layer_split_enabled(enabled);
|
||||
}
|
||||
|
||||
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
|
||||
llm->set_graph_cut_layer_split_backend_vram_limits(limits);
|
||||
}
|
||||
|
||||
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
|
||||
llm->get_param_tensors(tensors, "text_encoders.llm");
|
||||
}
|
||||
@@ -1654,6 +1787,7 @@ struct LLMEmbedder : public Conditioner {
|
||||
SDVersion version;
|
||||
std::shared_ptr<BPETokenizer> tokenizer;
|
||||
std::shared_ptr<LLM::LLMRunner> llm;
|
||||
std::shared_ptr<T5Runner> byt5;
|
||||
|
||||
LLMEmbedder(ggml_backend_t backend,
|
||||
const String2TensorStorage& tensor_storage_map = {},
|
||||
@@ -1671,7 +1805,12 @@ struct LLMEmbedder : public Conditioner {
|
||||
arch = LLM::LLMArch::GPT_OSS_20B;
|
||||
} else if (sd_version_is_pid(version)) {
|
||||
arch = LLM::LLMArch::GEMMA2_2B;
|
||||
} else if (sd_version_is_ideogram4(version) || sd_version_is_boogu_image(version) || sd_version_is_sefi_image(version) || sd_version_is_krea2(version)) {
|
||||
} else if (sd_version_is_lingbot_video(version) ||
|
||||
sd_version_is_ideogram4(version) ||
|
||||
sd_version_is_boogu_image(version) ||
|
||||
sd_version_is_sefi_image(version) ||
|
||||
sd_version_is_krea2(version) ||
|
||||
sd_version_is_mage_flow(version)) {
|
||||
arch = LLM::LLMArch::QWEN3_VL;
|
||||
} else if (sd_version_is_z_image(version) || version == VERSION_OVIS_IMAGE || version == VERSION_FLUX2_KLEIN) {
|
||||
arch = LLM::LLMArch::QWEN3;
|
||||
@@ -1691,42 +1830,101 @@ struct LLMEmbedder : public Conditioner {
|
||||
"text_encoders.llm",
|
||||
enable_vision,
|
||||
weight_manager);
|
||||
if (sd_version_is_hunyuan_video(version)) {
|
||||
const std::string byt5_prefix = "text_encoders.t5xxl.transformer";
|
||||
for (const auto& [name, _] : tensor_storage_map) {
|
||||
if (starts_with(name, byt5_prefix + ".")) {
|
||||
byt5 = std::make_shared<T5Runner>(backend,
|
||||
tensor_storage_map,
|
||||
byt5_prefix,
|
||||
false,
|
||||
weight_manager);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
|
||||
llm->get_param_tensors(tensors, "text_encoders.llm");
|
||||
if (byt5) {
|
||||
byt5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
|
||||
}
|
||||
}
|
||||
|
||||
void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) override {
|
||||
llm->get_param_tensor_ops(tensor_ops);
|
||||
}
|
||||
|
||||
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
|
||||
llm->set_max_graph_vram_bytes(max_vram_bytes);
|
||||
if (byt5) {
|
||||
byt5->set_max_graph_vram_bytes(max_vram_bytes);
|
||||
}
|
||||
}
|
||||
|
||||
void set_stream_layers_enabled(bool enabled) override {
|
||||
llm->set_stream_layers_enabled(enabled);
|
||||
if (byt5) {
|
||||
byt5->set_stream_layers_enabled(enabled);
|
||||
}
|
||||
}
|
||||
|
||||
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
|
||||
llm->set_runtime_backends(backends);
|
||||
if (byt5) {
|
||||
byt5->set_runtime_backends(backends);
|
||||
}
|
||||
}
|
||||
|
||||
void set_graph_cut_layer_split_enabled(bool enabled) override {
|
||||
if (llm) {
|
||||
llm->set_graph_cut_layer_split_enabled(enabled);
|
||||
}
|
||||
if (byt5) {
|
||||
byt5->set_graph_cut_layer_split_enabled(enabled);
|
||||
}
|
||||
}
|
||||
|
||||
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
|
||||
if (llm) {
|
||||
llm->set_graph_cut_layer_split_backend_vram_limits(limits);
|
||||
}
|
||||
if (byt5) {
|
||||
byt5->set_graph_cut_layer_split_backend_vram_limits(limits);
|
||||
}
|
||||
}
|
||||
|
||||
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
|
||||
llm->get_param_tensors(tensors, "text_encoders.llm");
|
||||
if (byt5) {
|
||||
byt5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
|
||||
}
|
||||
}
|
||||
|
||||
void set_flash_attention_enabled(bool enabled) override {
|
||||
llm->set_flash_attention_enabled(enabled);
|
||||
if (byt5) {
|
||||
byt5->set_flash_attention_enabled(enabled);
|
||||
}
|
||||
}
|
||||
|
||||
void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
|
||||
if (llm) {
|
||||
llm->set_weight_adapter(adapter);
|
||||
}
|
||||
if (byt5) {
|
||||
byt5->set_weight_adapter(adapter);
|
||||
}
|
||||
}
|
||||
|
||||
void runner_done() override {
|
||||
if (llm) {
|
||||
llm->runner_done();
|
||||
}
|
||||
if (byt5) {
|
||||
byt5->runner_done();
|
||||
}
|
||||
}
|
||||
|
||||
std::tuple<std::vector<int>, std::vector<float>, std::vector<float>> tokenize(std::string text,
|
||||
@@ -1855,6 +2053,54 @@ struct LLMEmbedder : public Conditioner {
|
||||
return new_hidden_states;
|
||||
}
|
||||
|
||||
void resize_image_dims(int height, int width, int& h_bar, int& w_bar, int factor, int min_size, int max_size, RefImageResizeMode mode) {
|
||||
if (min_size > 0 && min_size == max_size) {
|
||||
if (mode == RefImageResizeMode::AREA) {
|
||||
double beta = std::sqrt(static_cast<double>(min_size) / (static_cast<double>(height) * width));
|
||||
h_bar = std::max(static_cast<int>(factor),
|
||||
static_cast<int>(std::round(height * beta / factor)) * static_cast<int>(factor));
|
||||
w_bar = std::max(static_cast<int>(factor),
|
||||
static_cast<int>(std::round(width * beta / factor)) * static_cast<int>(factor));
|
||||
} else if (mode == RefImageResizeMode::LONGEST_SIDE) {
|
||||
int current_max_side = std::max(height, width);
|
||||
double beta = static_cast<double>(min_size) / current_max_side;
|
||||
h_bar = std::max(static_cast<int>(factor),
|
||||
static_cast<int>(std::round(height * beta / factor)) * static_cast<int>(factor));
|
||||
w_bar = std::max(static_cast<int>(factor),
|
||||
static_cast<int>(std::round(width * beta / factor)) * static_cast<int>(factor));
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
if (mode == RefImageResizeMode::AREA) {
|
||||
double current_area = static_cast<double>(h_bar) * w_bar;
|
||||
if (max_size > 0 && current_area > max_size) {
|
||||
double beta = std::sqrt((static_cast<double>(height) * width) / static_cast<double>(max_size));
|
||||
h_bar = std::max(static_cast<int>(factor),
|
||||
static_cast<int>(std::floor(height / beta / factor)) * static_cast<int>(factor));
|
||||
w_bar = std::max(static_cast<int>(factor),
|
||||
static_cast<int>(std::floor(width / beta / factor)) * static_cast<int>(factor));
|
||||
} else if (min_size > 0 && current_area < min_size) {
|
||||
double beta = std::sqrt(static_cast<double>(min_size) / (static_cast<double>(height) * width));
|
||||
h_bar = static_cast<int>(std::ceil(height * beta / factor)) * static_cast<int>(factor);
|
||||
w_bar = static_cast<int>(std::ceil(width * beta / factor)) * static_cast<int>(factor);
|
||||
}
|
||||
} else if (mode == RefImageResizeMode::LONGEST_SIDE) {
|
||||
int current_max_side = std::max(height, width);
|
||||
if (max_size > 0 && current_max_side > max_size) {
|
||||
double beta = static_cast<double>(max_size) / current_max_side;
|
||||
h_bar = std::max(static_cast<int>(factor),
|
||||
static_cast<int>(std::floor(height * beta / factor)) * static_cast<int>(factor));
|
||||
w_bar = std::max(static_cast<int>(factor),
|
||||
static_cast<int>(std::floor(width * beta / factor)) * static_cast<int>(factor));
|
||||
} else if (min_size > 0 && current_max_side < min_size) {
|
||||
double beta = static_cast<double>(min_size) / current_max_side;
|
||||
h_bar = static_cast<int>(std::ceil(height * beta / factor)) * static_cast<int>(factor);
|
||||
w_bar = static_cast<int>(std::ceil(width * beta / factor)) * static_cast<int>(factor);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
SDCondition get_learned_condition(int n_threads,
|
||||
const ConditionerParams& conditioner_params) override {
|
||||
std::string prompt;
|
||||
@@ -1869,38 +2115,143 @@ struct LLMEmbedder : public Conditioner {
|
||||
bool spell_quotes = false;
|
||||
std::set<int> out_layers;
|
||||
|
||||
int64_t t0 = ggml_time_ms();
|
||||
int64_t t0 = ggml_time_ms();
|
||||
RefImageResizeMode resize_mode = conditioner_params.ref_image_params.vlm_resize_mode;
|
||||
|
||||
if (sd_version_is_qwen_image(version)) {
|
||||
if (sd_version_is_hunyuan_video(version)) {
|
||||
prompt_template_encode_start_idx = 98;
|
||||
out_layers = {26};
|
||||
|
||||
prompt =
|
||||
"<|im_start|>system\nYou are a helpful assistant. Describe the video by detailing the following aspects:\n"
|
||||
"1. The main content and theme of the video.\n"
|
||||
"2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects.\n"
|
||||
"3. Actions, events, behaviors temporal relationships, physical movement changes of the objects.\n"
|
||||
"4. background environment, light, style and atmosphere.\n"
|
||||
"5. camera angles, movements, and transitions used in the video.<|im_end|>\n"
|
||||
"<|im_start|>user\n";
|
||||
|
||||
prompt_attn_range.first = static_cast<int>(prompt.size());
|
||||
prompt += conditioner_params.text;
|
||||
prompt_attn_range.second = static_cast<int>(prompt.size());
|
||||
prompt += "<|im_end|>\n<|im_start|>assistant\n";
|
||||
} else if (sd_version_is_lingbot_video(version)) {
|
||||
const int pad_token = 151643;
|
||||
const std::string prompt_prefix =
|
||||
"<|im_start|>system\nGiven a user input that may include a text prompt alone, "
|
||||
"a text prompt with an image reference, or a text prompt with a video reference "
|
||||
"or a video reference alone, generate an \"Enhanced prompt\" that provides detailed "
|
||||
"visual descriptions suitable for video generation. Evaluate the level of detail "
|
||||
"in the user's input: if it is simple, enrich it by adding specifics about colors, "
|
||||
"shapes, sizes, textures, lighting, motion dynamics, camera movement, temporal "
|
||||
"progression, and spatial relationships to create vivid, concrete, and temporally "
|
||||
"coherent scenes to create vivid and concrete scenes. Please generate only the "
|
||||
"enhanced description for the prompt below and avoid including any additional "
|
||||
"commentary or evaluations:<|im_end|>\n<|im_start|>user\n";
|
||||
|
||||
auto prefix_tokens = tokenizer->encode(prompt_prefix, nullptr);
|
||||
prompt_template_encode_start_idx = 0;
|
||||
for (int token : prefix_tokens) {
|
||||
if (token != pad_token) {
|
||||
prompt_template_encode_start_idx++;
|
||||
}
|
||||
}
|
||||
LOG_DEBUG("prompt_template_encode_start_idx %d", prompt_template_encode_start_idx);
|
||||
|
||||
prompt = prompt_prefix;
|
||||
if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
|
||||
LOG_INFO("QwenImageEditPlusPipeline");
|
||||
LOG_INFO("LingBotVideoI2VPipeline");
|
||||
const std::string placeholder = "<|image_pad|>";
|
||||
std::string img_prompt;
|
||||
|
||||
for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
|
||||
const auto& image = (*conditioner_params.ref_images)[i];
|
||||
const int factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
|
||||
int height = static_cast<int>(image.shape()[1]);
|
||||
int width = static_cast<int>(image.shape()[0]);
|
||||
|
||||
int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
|
||||
if (min_pixels <= 0) {
|
||||
if (resize_mode == RefImageResizeMode::AREA) {
|
||||
min_pixels = static_cast<int>(4 * factor * factor);
|
||||
} else {
|
||||
min_pixels = static_cast<int>(2 * factor);
|
||||
}
|
||||
}
|
||||
int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
|
||||
if (max_pixels <= 0) {
|
||||
if (resize_mode == RefImageResizeMode::AREA) {
|
||||
max_pixels = static_cast<int>(16384 * factor * factor);
|
||||
} else {
|
||||
max_pixels = static_cast<int>(128 * factor);
|
||||
}
|
||||
}
|
||||
|
||||
int h_bar = std::max(factor, static_cast<int>(std::round(static_cast<double>(height) / factor) * factor));
|
||||
int w_bar = std::max(factor, static_cast<int>(std::round(static_cast<double>(width) / factor) * factor));
|
||||
|
||||
if (std::max(height, width) > 200 * std::min(height, width)) {
|
||||
LOG_WARN("LingBotVideo image aspect ratio is very large: %dx%d", width, height);
|
||||
}
|
||||
|
||||
resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
|
||||
|
||||
LOG_DEBUG("resize LingBotVideo ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
|
||||
auto resized_image = clip_preprocess(image, w_bar, h_bar);
|
||||
auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true);
|
||||
GGML_ASSERT(!image_embed.empty());
|
||||
|
||||
std::string image_prefix = prompt + img_prompt + "<|vision_start|>";
|
||||
int image_embed_idx = static_cast<int>(tokenizer->encode(image_prefix, nullptr).size());
|
||||
image_embeds.emplace_back(image_embed_idx, image_embed);
|
||||
|
||||
img_prompt += "<|vision_start|>";
|
||||
int64_t num_image_tokens = image_embed.shape()[1];
|
||||
img_prompt.reserve(img_prompt.size() + static_cast<size_t>(num_image_tokens) * placeholder.size() + 32);
|
||||
for (int j = 0; j < num_image_tokens; j++) {
|
||||
img_prompt += placeholder;
|
||||
}
|
||||
img_prompt += "<|vision_end|>";
|
||||
}
|
||||
prompt += img_prompt;
|
||||
}
|
||||
|
||||
prompt += conditioner_params.text;
|
||||
prompt_attn_range = {0, 0};
|
||||
prompt += "<|im_end|>\n<|im_start|>assistant\n";
|
||||
} else if (sd_version_is_qwen_image(version) || sd_version_is_mage_flow(version)) {
|
||||
if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
|
||||
LOG_INFO("%s", sd_version_is_mage_flow(version) ? "MageFlowEditPipeline" : "QwenImageEditPlusPipeline");
|
||||
prompt_template_encode_start_idx = 64;
|
||||
int image_embed_idx = 64 + 6;
|
||||
|
||||
int min_pixels = 384 * 384;
|
||||
int max_pixels = 560 * 560;
|
||||
int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
|
||||
if (min_pixels <= 0) {
|
||||
min_pixels = sd_version_is_mage_flow(version) ? -1 : 384;
|
||||
if (min_pixels > 0 && resize_mode == RefImageResizeMode::AREA) {
|
||||
min_pixels *= min_pixels;
|
||||
}
|
||||
}
|
||||
int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
|
||||
if (max_pixels <= 0) {
|
||||
max_pixels = sd_version_is_mage_flow(version) ? 384 : 560;
|
||||
if (resize_mode == RefImageResizeMode::AREA) {
|
||||
max_pixels *= max_pixels;
|
||||
}
|
||||
}
|
||||
|
||||
std::string placeholder = "<|image_pad|>";
|
||||
std::string img_prompt;
|
||||
|
||||
for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
|
||||
const auto& image = (*conditioner_params.ref_images)[i];
|
||||
double factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
|
||||
const int factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
|
||||
int height = static_cast<int>(image.shape()[1]);
|
||||
int width = static_cast<int>(image.shape()[0]);
|
||||
int h_bar = static_cast<int>(std::round(height / factor) * factor);
|
||||
int w_bar = static_cast<int>(std::round(width / factor) * factor);
|
||||
int h_bar = static_cast<int>(std::round(static_cast<double>(height) / factor) * factor);
|
||||
int w_bar = static_cast<int>(std::round(static_cast<double>(width) / factor) * factor);
|
||||
|
||||
if (static_cast<double>(h_bar) * w_bar > max_pixels) {
|
||||
double beta = std::sqrt((height * width) / static_cast<double>(max_pixels));
|
||||
h_bar = std::max(static_cast<int>(factor),
|
||||
static_cast<int>(std::floor(height / beta / factor)) * static_cast<int>(factor));
|
||||
w_bar = std::max(static_cast<int>(factor),
|
||||
static_cast<int>(std::floor(width / beta / factor)) * static_cast<int>(factor));
|
||||
} else if (static_cast<double>(h_bar) * w_bar < min_pixels) {
|
||||
double beta = std::sqrt(static_cast<double>(min_pixels) / (height * width));
|
||||
h_bar = static_cast<int>(std::ceil(height * beta / factor)) * static_cast<int>(factor);
|
||||
w_bar = static_cast<int>(std::ceil(width * beta / factor)) * static_cast<int>(factor);
|
||||
}
|
||||
resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
|
||||
|
||||
LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
|
||||
|
||||
@@ -1911,7 +2262,7 @@ struct LLMEmbedder : public Conditioner {
|
||||
image_embeds.emplace_back(image_embed_idx, image_embed);
|
||||
image_embed_idx += 1 + static_cast<int>(image_embed.shape()[1]) + 6;
|
||||
|
||||
img_prompt += "Picture " + std::to_string(i + 1) + ": <|vision_start|>"; // [24669, 220, index, 25, 220, 151652]
|
||||
img_prompt += (sd_version_is_mage_flow(version) ? "Image " : "Picture ") + std::to_string(i + 1) + ": <|vision_start|>";
|
||||
int64_t num_image_tokens = image_embed.shape()[1];
|
||||
img_prompt.reserve(num_image_tokens * placeholder.size());
|
||||
for (int j = 0; j < num_image_tokens; j++) {
|
||||
@@ -1939,6 +2290,9 @@ struct LLMEmbedder : public Conditioner {
|
||||
|
||||
prompt += "<|im_end|>\n<|im_start|>assistant\n";
|
||||
}
|
||||
if (sd_version_is_mage_flow(version)) {
|
||||
max_length = 2048 + prompt_template_encode_start_idx;
|
||||
}
|
||||
} else if (sd_version_is_boogu_image(version)) {
|
||||
prompt_template_encode_start_idx = 0;
|
||||
|
||||
@@ -1955,16 +2309,33 @@ struct LLMEmbedder : public Conditioner {
|
||||
std::string img_prompt;
|
||||
const std::string placeholder = "<|image_pad|>";
|
||||
|
||||
int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
|
||||
if (min_pixels <= 0) {
|
||||
min_pixels = 384;
|
||||
if (resize_mode == RefImageResizeMode::AREA) {
|
||||
min_pixels *= min_pixels;
|
||||
}
|
||||
}
|
||||
int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
|
||||
if (max_pixels <= 0) {
|
||||
max_pixels = 384;
|
||||
if (resize_mode == RefImageResizeMode::AREA) {
|
||||
max_pixels *= max_pixels;
|
||||
}
|
||||
}
|
||||
|
||||
for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
|
||||
const auto& image = (*conditioner_params.ref_images)[i];
|
||||
double factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
|
||||
const int factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
|
||||
int height = static_cast<int>(image.shape()[1]);
|
||||
int width = static_cast<int>(image.shape()[0]);
|
||||
double beta = std::sqrt((384.0 * 384.0) / (static_cast<double>(height) * static_cast<double>(width)));
|
||||
int h_bar = std::max(static_cast<int>(factor),
|
||||
static_cast<int>(std::round(height * beta / factor)) * static_cast<int>(factor));
|
||||
int w_bar = std::max(static_cast<int>(factor),
|
||||
static_cast<int>(std::round(width * beta / factor)) * static_cast<int>(factor));
|
||||
|
||||
int h_bar = std::max(factor,
|
||||
static_cast<int>(std::round(static_cast<double>(height) / factor)) * factor);
|
||||
int w_bar = std::max(factor,
|
||||
static_cast<int>(std::round(static_cast<double>(width) / factor)) * factor);
|
||||
|
||||
resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
|
||||
|
||||
LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
|
||||
|
||||
@@ -2003,6 +2374,57 @@ struct LLMEmbedder : public Conditioner {
|
||||
out_layers = {2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35};
|
||||
|
||||
prompt = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n";
|
||||
if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
|
||||
std::string img_prompt = "";
|
||||
const std::string placeholder = "<|image_pad|>";
|
||||
int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
|
||||
if (min_pixels <= 0) {
|
||||
min_pixels = 384;
|
||||
if (resize_mode == RefImageResizeMode::AREA) {
|
||||
min_pixels *= min_pixels;
|
||||
}
|
||||
}
|
||||
int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
|
||||
if (max_pixels <= 0) {
|
||||
max_pixels = 1024;
|
||||
if (resize_mode == RefImageResizeMode::AREA) {
|
||||
max_pixels *= max_pixels;
|
||||
}
|
||||
}
|
||||
|
||||
for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
|
||||
const auto& image = (*conditioner_params.ref_images)[i];
|
||||
const int factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
|
||||
int height = static_cast<int>(image.shape()[1]);
|
||||
int width = static_cast<int>(image.shape()[0]);
|
||||
|
||||
int h_bar = std::max(factor,
|
||||
static_cast<int>(std::round(static_cast<double>(height) / factor)) * factor);
|
||||
int w_bar = std::max(factor,
|
||||
static_cast<int>(std::round(static_cast<double>(width) / factor)) * factor);
|
||||
|
||||
resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
|
||||
|
||||
LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
|
||||
|
||||
auto resized_image = clip_preprocess(image, w_bar, h_bar);
|
||||
auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true);
|
||||
GGML_ASSERT(!image_embed.empty());
|
||||
|
||||
std::string image_prefix = prompt + img_prompt + "Picture " + std::to_string(i + 1) + ": <|vision_start|>";
|
||||
int image_embed_idx = static_cast<int>(tokenizer->encode(image_prefix, nullptr).size());
|
||||
image_embeds.emplace_back(image_embed_idx, image_embed);
|
||||
|
||||
img_prompt += "Picture " + std::to_string(i + 1) + ": <|vision_start|>";
|
||||
int64_t num_image_tokens = image_embed.shape()[1];
|
||||
img_prompt.reserve(img_prompt.size() + static_cast<size_t>(num_image_tokens) * placeholder.size() + 32);
|
||||
for (int j = 0; j < num_image_tokens; j++) {
|
||||
img_prompt += placeholder;
|
||||
}
|
||||
img_prompt += "<|vision_end|>";
|
||||
}
|
||||
prompt += img_prompt;
|
||||
}
|
||||
|
||||
prompt_attn_range.first = static_cast<int>(prompt.size());
|
||||
prompt += conditioner_params.text;
|
||||
@@ -2018,30 +2440,33 @@ struct LLMEmbedder : public Conditioner {
|
||||
min_length = 512 + prompt_template_encode_start_idx;
|
||||
int image_embed_idx = 36 + 6;
|
||||
|
||||
int min_pixels = 384 * 384;
|
||||
int max_pixels = 560 * 560;
|
||||
int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
|
||||
if (min_pixels <= 0) {
|
||||
min_pixels = 384;
|
||||
if (resize_mode == RefImageResizeMode::AREA) {
|
||||
min_pixels *= min_pixels;
|
||||
}
|
||||
}
|
||||
int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
|
||||
if (max_pixels <= 0) {
|
||||
max_pixels = 560;
|
||||
if (resize_mode == RefImageResizeMode::AREA) {
|
||||
max_pixels *= max_pixels;
|
||||
}
|
||||
}
|
||||
|
||||
std::string placeholder = "<|image_pad|>";
|
||||
std::string img_prompt;
|
||||
|
||||
for (int i = 0; i < conditioner_params.ref_images->size(); i++) {
|
||||
const auto& image = (*conditioner_params.ref_images)[i];
|
||||
double factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
|
||||
const int factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size;
|
||||
int height = static_cast<int>(image.shape()[1]);
|
||||
int width = static_cast<int>(image.shape()[0]);
|
||||
int h_bar = static_cast<int>(std::round(height / factor) * factor);
|
||||
int w_bar = static_cast<int>(std::round(width / factor) * factor);
|
||||
int h_bar = static_cast<int>(std::round(static_cast<double>(height) / factor) * factor);
|
||||
int w_bar = static_cast<int>(std::round(static_cast<double>(width) / factor) * factor);
|
||||
|
||||
if (static_cast<double>(h_bar) * w_bar > max_pixels) {
|
||||
double beta = std::sqrt((height * width) / static_cast<double>(max_pixels));
|
||||
h_bar = std::max(static_cast<int>(factor),
|
||||
static_cast<int>(std::floor(height / beta / factor)) * static_cast<int>(factor));
|
||||
w_bar = std::max(static_cast<int>(factor),
|
||||
static_cast<int>(std::floor(width / beta / factor)) * static_cast<int>(factor));
|
||||
} else if (static_cast<double>(h_bar) * w_bar < min_pixels) {
|
||||
double beta = std::sqrt(static_cast<double>(min_pixels) / (height * width));
|
||||
h_bar = static_cast<int>(std::ceil(height * beta / factor)) * static_cast<int>(factor);
|
||||
w_bar = static_cast<int>(std::ceil(width * beta / factor)) * static_cast<int>(factor);
|
||||
}
|
||||
resize_image_dims(height, width, h_bar, w_bar, factor, min_pixels, max_pixels, resize_mode);
|
||||
|
||||
LOG_DEBUG("resize conditioner ref image %d from %dx%d to %dx%d", i, height, width, h_bar, w_bar);
|
||||
|
||||
@@ -2243,6 +2668,46 @@ struct LLMEmbedder : public Conditioner {
|
||||
spell_quotes,
|
||||
max_length);
|
||||
std::vector<sd::Tensor<float>> extra_hidden_states_vec;
|
||||
if (sd_version_is_hunyuan_video(version) && byt5) {
|
||||
std::vector<std::string> quoted_texts;
|
||||
auto collect_quoted = [&](const std::string& open, const std::string& close) {
|
||||
size_t begin = 0;
|
||||
while ((begin = conditioner_params.text.find(open, begin)) != std::string::npos) {
|
||||
size_t content_begin = begin + open.size();
|
||||
size_t end = conditioner_params.text.find(close, content_begin);
|
||||
if (end == std::string::npos) {
|
||||
break;
|
||||
}
|
||||
quoted_texts.push_back(conditioner_params.text.substr(content_begin, end - content_begin));
|
||||
begin = end + close.size();
|
||||
}
|
||||
};
|
||||
collect_quoted("\"", "\"");
|
||||
collect_quoted("\xE2\x80\x98", "\xE2\x80\x99");
|
||||
collect_quoted("\xE2\x80\x9C", "\xE2\x80\x9D");
|
||||
|
||||
if (!quoted_texts.empty()) {
|
||||
std::string byt5_text;
|
||||
for (const auto& text : quoted_texts) {
|
||||
byt5_text += "Text \"" + text + "\". ";
|
||||
}
|
||||
std::vector<int> tokens;
|
||||
tokens.reserve(byt5_text.size() + 1);
|
||||
for (unsigned char byte : byt5_text) {
|
||||
tokens.push_back(static_cast<int>(byte) + 3);
|
||||
}
|
||||
tokens.push_back(1);
|
||||
sd::Tensor<int32_t> input_ids({static_cast<int64_t>(tokens.size())}, tokens);
|
||||
auto byt5_hidden_states = byt5->compute(n_threads,
|
||||
input_ids,
|
||||
sd::Tensor<float>(),
|
||||
false,
|
||||
true,
|
||||
true);
|
||||
GGML_ASSERT(!byt5_hidden_states.empty());
|
||||
extra_hidden_states_vec.push_back(std::move(byt5_hidden_states));
|
||||
}
|
||||
}
|
||||
for (int i = 0; i < extra_prompts.size(); i++) {
|
||||
auto extra_hidden_states = encode_prompt(n_threads,
|
||||
extra_prompts[i],
|
||||
@@ -2372,6 +2837,10 @@ struct LTXAVEmbedder : public Conditioner {
|
||||
projector->get_param_tensors(tensors, "text_embedding_projection");
|
||||
}
|
||||
|
||||
void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) override {
|
||||
llm->get_param_tensor_ops(tensor_ops);
|
||||
}
|
||||
|
||||
void set_flash_attention_enabled(bool enabled) override {
|
||||
llm->set_flash_attention_enabled(enabled);
|
||||
projector->set_flash_attention_enabled(enabled);
|
||||
@@ -2386,6 +2855,14 @@ struct LTXAVEmbedder : public Conditioner {
|
||||
llm->set_runtime_backends(backends);
|
||||
}
|
||||
|
||||
void set_graph_cut_layer_split_enabled(bool enabled) override {
|
||||
llm->set_graph_cut_layer_split_enabled(enabled);
|
||||
}
|
||||
|
||||
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
|
||||
llm->set_graph_cut_layer_split_backend_vram_limits(limits);
|
||||
}
|
||||
|
||||
void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
|
||||
llm->get_param_tensors(tensors, "text_encoders.llm");
|
||||
}
|
||||
|
||||
@@ -21,10 +21,12 @@
|
||||
#include <sstream>
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
#include <unordered_set>
|
||||
#include <vector>
|
||||
|
||||
#include "core/ggml_extend_backend.h"
|
||||
#include "core/ggml_graph_cut.h"
|
||||
#include "core/layer_split_partition.h"
|
||||
#include "ggml-alloc.h"
|
||||
#include "ggml-backend.h"
|
||||
#include "ggml.h"
|
||||
@@ -1687,6 +1689,8 @@ struct GGMLRunnerContext {
|
||||
bool conv2d_direct_enabled = false;
|
||||
bool circular_x_enabled = false;
|
||||
bool circular_y_enabled = false;
|
||||
ggml_tensor* ip_context = nullptr;
|
||||
float ip_scale = 1.0f;
|
||||
std::shared_ptr<WeightAdapter> weight_adapter = nullptr;
|
||||
std::vector<std::pair<ggml_tensor*, std::string>>* debug_tensors = nullptr;
|
||||
std::function<ggml_tensor*(const std::string&)> get_cache_tensor;
|
||||
@@ -1745,9 +1749,11 @@ protected:
|
||||
size_t max_graph_vram_bytes = 0;
|
||||
bool stream_layers_enabled = false;
|
||||
size_t observed_max_effective_budget_ = 0;
|
||||
bool graph_cut_layer_split_enabled = false;
|
||||
std::vector<size_t> graph_cut_layer_split_backend_vram_limits_;
|
||||
|
||||
std::vector<ggml_backend_t> extra_runtime_backends; // borrowed (SDBackendManager-owned)
|
||||
ggml_backend_sched_t sched = nullptr; // owned, multi-device only
|
||||
ggml_backend_sched_t sched = nullptr; // owned
|
||||
ggml_backend_t cpu_fallback_backend = nullptr; // owned, sched requires a trailing CPU backend
|
||||
bool multi_device_eval_callback_warned = false;
|
||||
|
||||
@@ -1776,6 +1782,9 @@ protected:
|
||||
|
||||
sd::ggml_graph_cut::PlanCache graph_cut_plan_cache_;
|
||||
std::unordered_set<const ggml_tensor*> params_tensor_set_;
|
||||
std::unordered_map<const ggml_tensor*, ggml_backend_t> graph_cut_layer_split_assignments_;
|
||||
std::unordered_map<const ggml_tensor*, ggml_backend_t> graph_cut_layer_split_node_assignments_;
|
||||
bool graph_cut_layer_split_primary_notice_logged_ = false;
|
||||
|
||||
template <typename T>
|
||||
static sd::Tensor<T> take_or_empty(std::optional<sd::Tensor<T>> tensor) {
|
||||
@@ -1874,6 +1883,20 @@ protected:
|
||||
params_tensor_set_dirty_ = false;
|
||||
}
|
||||
|
||||
ggml_tensor* canonical_param_tensor(ggml_tensor* tensor) {
|
||||
if (tensor == nullptr) {
|
||||
return nullptr;
|
||||
}
|
||||
if (params_tensor_set_.find(tensor) != params_tensor_set_.end()) {
|
||||
return tensor;
|
||||
}
|
||||
if (tensor->view_src != nullptr &&
|
||||
params_tensor_set_.find(tensor->view_src) != params_tensor_set_.end()) {
|
||||
return tensor->view_src;
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
std::vector<ggml_tensor*> collect_used_param_tensors(ggml_cgraph* gf) {
|
||||
std::vector<ggml_tensor*> used_params;
|
||||
rebuild_params_tensor_set();
|
||||
@@ -1886,12 +1909,8 @@ protected:
|
||||
seen_params.reserve(static_cast<size_t>(n_leafs));
|
||||
for (int i = 0; i < n_leafs; ++i) {
|
||||
ggml_tensor* leaf = sd::ggml_graph_cut::leaf_tensor(gf, i);
|
||||
ggml_tensor* param_leaf = leaf;
|
||||
if (param_leaf != nullptr && params_tensor_set_.find(param_leaf) == params_tensor_set_.end()) {
|
||||
param_leaf = param_leaf->view_src;
|
||||
}
|
||||
ggml_tensor* param_leaf = canonical_param_tensor(leaf);
|
||||
if (param_leaf != nullptr &&
|
||||
params_tensor_set_.find(param_leaf) != params_tensor_set_.end() &&
|
||||
seen_params.insert(param_leaf).second) {
|
||||
used_params.push_back(param_leaf);
|
||||
}
|
||||
@@ -2101,11 +2120,17 @@ protected:
|
||||
ggml_backend_t current = runtime_backend;
|
||||
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);
|
||||
ggml_tensor* node = ggml_graph_node(gf, i);
|
||||
auto node_assignment = graph_cut_layer_split_node_assignments_.find(node);
|
||||
if (node_assignment != graph_cut_layer_split_node_assignments_.end()) {
|
||||
current = node_assignment->second;
|
||||
}
|
||||
for (int s = 0; s < GGML_MAX_SRC; s++) {
|
||||
ggml_backend_t weight_backend = backend_for_weight(node->src[s]);
|
||||
if (weight_backend != nullptr) {
|
||||
current = weight_backend;
|
||||
if (node_assignment == graph_cut_layer_split_node_assignments_.end()) {
|
||||
current = weight_backend;
|
||||
}
|
||||
}
|
||||
}
|
||||
if (node->op == GGML_OP_NONE || node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE ||
|
||||
@@ -2122,8 +2147,22 @@ protected:
|
||||
return !extra_runtime_backends.empty();
|
||||
}
|
||||
|
||||
bool graph_requires_backend_fallback(ggml_cgraph* gf) const {
|
||||
if (gf == nullptr || sd_backend_is_cpu(runtime_backend)) {
|
||||
return false;
|
||||
}
|
||||
const int n_nodes = ggml_graph_n_nodes(gf);
|
||||
for (int i = 0; i < n_nodes; ++i) {
|
||||
ggml_tensor* node = ggml_graph_node(gf, i);
|
||||
if (node != nullptr && !ggml_backend_supports_op(runtime_backend, node)) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
bool alloc_compute_buffer(ggml_cgraph* gf) {
|
||||
if (is_multi_device()) {
|
||||
if (sched != nullptr || is_multi_device() || graph_requires_backend_fallback(gf)) {
|
||||
// The sched replaces the gallocr. Do NOT ggml_backend_sched_reserve
|
||||
// the graph here: reserve runs split_graph, which rewires the
|
||||
// graph's src pointers to sched-internal copy tensors, and the
|
||||
@@ -2131,6 +2170,10 @@ protected:
|
||||
// rewired graph, silently corrupting every cross-backend input. A
|
||||
// graph must be split at most once; the alloc in execute_graph
|
||||
// performs the real allocation.
|
||||
if (compute_allocr != nullptr) {
|
||||
ggml_gallocr_free(compute_allocr);
|
||||
compute_allocr = nullptr;
|
||||
}
|
||||
return ensure_sched(gf);
|
||||
}
|
||||
if (compute_allocr != nullptr) {
|
||||
@@ -2435,6 +2478,123 @@ protected:
|
||||
return true;
|
||||
}
|
||||
|
||||
bool resolve_graph_cut_layer_split_plan(ggml_cgraph* gf,
|
||||
GraphCutPlan* plan_out) {
|
||||
GGML_ASSERT(plan_out != nullptr);
|
||||
GGML_ASSERT(gf != nullptr);
|
||||
*plan_out = sd::ggml_graph_cut::resolve_plan(runtime_backend,
|
||||
gf,
|
||||
&graph_cut_plan_cache_,
|
||||
0,
|
||||
params_tensor_set_,
|
||||
get_desc().c_str());
|
||||
return true;
|
||||
}
|
||||
|
||||
bool assign_graph_cut_layer_split_backends(ggml_cgraph* gf) {
|
||||
graph_cut_layer_split_node_assignments_.clear();
|
||||
if (!graph_cut_layer_split_enabled) {
|
||||
return true;
|
||||
}
|
||||
if (!is_multi_device()) {
|
||||
LOG_ERROR("%s graph-cut layer split requires multiple runtime backends", get_desc().c_str());
|
||||
return false;
|
||||
}
|
||||
|
||||
GraphCutPlan plan;
|
||||
if (!resolve_graph_cut_layer_split_plan(gf, &plan)) {
|
||||
return false;
|
||||
}
|
||||
if (!plan.valid || !plan.has_cuts || plan.segments.size() <= 1) {
|
||||
auto manager = weight_manager.lock();
|
||||
if (manager == nullptr) {
|
||||
LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str());
|
||||
return false;
|
||||
}
|
||||
std::vector<ggml_tensor*> graph_params = collect_used_param_tensors(gf);
|
||||
if (!graph_params.empty() &&
|
||||
!manager->assign_compute_backend(graph_params, runtime_backend)) {
|
||||
LOG_ERROR("%s graph-cut layer split failed to assign unmarked graph params to %s",
|
||||
get_desc().c_str(),
|
||||
sd::layer_split_backend_device_display_name(runtime_backend).c_str());
|
||||
return false;
|
||||
}
|
||||
for (ggml_tensor* param : graph_params) {
|
||||
if (param != nullptr) {
|
||||
graph_cut_layer_split_assignments_[param] = runtime_backend;
|
||||
}
|
||||
}
|
||||
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) {
|
||||
graph_cut_layer_split_node_assignments_[node] = runtime_backend;
|
||||
}
|
||||
}
|
||||
if (!graph_cut_layer_split_primary_notice_logged_) {
|
||||
LOG_WARN("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params",
|
||||
get_desc().c_str(),
|
||||
sd::layer_split_backend_device_display_name(runtime_backend).c_str(),
|
||||
graph_params.size());
|
||||
graph_cut_layer_split_primary_notice_logged_ = true;
|
||||
} else {
|
||||
LOG_DEBUG("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params",
|
||||
get_desc().c_str(),
|
||||
sd::layer_split_backend_device_display_name(runtime_backend).c_str(),
|
||||
graph_params.size());
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
std::vector<ggml_backend_t> split_backends;
|
||||
split_backends.reserve(extra_runtime_backends.size() + 1);
|
||||
split_backends.push_back(runtime_backend);
|
||||
for (ggml_backend_t backend : extra_runtime_backends) {
|
||||
if (backend != nullptr) {
|
||||
split_backends.push_back(backend);
|
||||
}
|
||||
}
|
||||
|
||||
auto manager = weight_manager.lock();
|
||||
if (manager == nullptr) {
|
||||
LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str());
|
||||
return false;
|
||||
}
|
||||
|
||||
sd::GraphCutLayerSplitAssignment assignment;
|
||||
auto canonicalize_param = [this](ggml_tensor* tensor) {
|
||||
return canonical_param_tensor(tensor);
|
||||
};
|
||||
if (!sd::partition_graph_cut_layer_split(get_desc().c_str(),
|
||||
gf,
|
||||
plan,
|
||||
split_backends,
|
||||
graph_cut_layer_split_backend_vram_limits_,
|
||||
max_graph_vram_bytes,
|
||||
graph_cut_layer_split_assignments_,
|
||||
canonicalize_param,
|
||||
&assignment)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
for (size_t i = 0; i < split_backends.size(); i++) {
|
||||
if (assignment.tensors_by_backend[i].empty()) {
|
||||
continue;
|
||||
}
|
||||
if (!manager->assign_compute_backend(assignment.tensors_by_backend[i], split_backends[i])) {
|
||||
LOG_ERROR("%s graph-cut layer split failed to assign params to %s",
|
||||
get_desc().c_str(),
|
||||
sd::layer_split_backend_device_display_name(split_backends[i]).c_str());
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
graph_cut_layer_split_node_assignments_ = std::move(assignment.node_assignments);
|
||||
sd::log_graph_cut_layer_split_assignment(get_desc().c_str(), split_backends, assignment);
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
struct PersistentExternalBinding {
|
||||
ggml_backend_buffer_t buffer = nullptr;
|
||||
void* data = nullptr;
|
||||
@@ -2611,7 +2771,7 @@ protected:
|
||||
};
|
||||
ComputeBufferGuard compute_buffer_guard(this, free_compute_buffer);
|
||||
|
||||
if (is_multi_device()) {
|
||||
if (sched != nullptr) {
|
||||
ggml_backend_sched_reset(sched);
|
||||
pin_multi_device_nodes(gf); // reset clears the pins; re-apply before alloc
|
||||
if (!ggml_backend_sched_alloc_graph(sched, gf)) {
|
||||
@@ -2632,9 +2792,9 @@ protected:
|
||||
}
|
||||
|
||||
ggml_status status;
|
||||
if (is_multi_device()) {
|
||||
if (sched != nullptr) {
|
||||
if (sd_get_backend_eval_callback() != nullptr && !multi_device_eval_callback_warned) {
|
||||
LOG_WARN("%s: eval callback is not supported with multiple runtime backends; ignoring",
|
||||
LOG_WARN("%s: eval callback is not supported with the backend scheduler; ignoring",
|
||||
get_desc().c_str());
|
||||
multi_device_eval_callback_warned = true;
|
||||
}
|
||||
@@ -2876,12 +3036,9 @@ public:
|
||||
|
||||
// do copy after alloc graph
|
||||
void set_backend_tensor_data(ggml_tensor* tensor, const void* data) {
|
||||
if (is_multi_device()) {
|
||||
// The sched only assigns a backend (and thus a buffer) to tensors
|
||||
// that participate in the graph; flag standalone data tensors as
|
||||
// inputs so they get one.
|
||||
ggml_set_input(tensor);
|
||||
}
|
||||
// The scheduler only allocates standalone data tensors when they are
|
||||
// marked as graph inputs. The flag is harmless for single-backend graphs.
|
||||
ggml_set_input(tensor);
|
||||
backend_tensor_data_map[tensor] = data;
|
||||
}
|
||||
|
||||
@@ -2972,6 +3129,11 @@ public:
|
||||
GGML_ASSERT(gf != nullptr);
|
||||
rebuild_params_tensor_set();
|
||||
|
||||
if (!assign_graph_cut_layer_split_backends(gf)) {
|
||||
free_compute_ctx();
|
||||
return std::nullopt;
|
||||
}
|
||||
|
||||
if (can_attempt_graph_cut_segmented_compute()) {
|
||||
GraphCutPlan plan;
|
||||
if (!resolve_graph_cut_plan(gf, &plan)) {
|
||||
@@ -3025,6 +3187,22 @@ public:
|
||||
stream_layers_enabled = enabled;
|
||||
}
|
||||
|
||||
void set_graph_cut_layer_split_enabled(bool enabled) {
|
||||
graph_cut_layer_split_enabled = enabled;
|
||||
if (!enabled) {
|
||||
graph_cut_layer_split_assignments_.clear();
|
||||
graph_cut_layer_split_node_assignments_.clear();
|
||||
graph_cut_layer_split_primary_notice_logged_ = false;
|
||||
}
|
||||
}
|
||||
|
||||
void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) {
|
||||
graph_cut_layer_split_backend_vram_limits_ = limits;
|
||||
graph_cut_layer_split_assignments_.clear();
|
||||
graph_cut_layer_split_node_assignments_.clear();
|
||||
graph_cut_layer_split_primary_notice_logged_ = false;
|
||||
}
|
||||
|
||||
void set_runtime_backends(const std::vector<ggml_backend_t>& backends) {
|
||||
extra_runtime_backends.clear();
|
||||
for (ggml_backend_t backend : backends) {
|
||||
@@ -3036,6 +3214,9 @@ public:
|
||||
extra_runtime_backends.push_back(backend);
|
||||
}
|
||||
}
|
||||
graph_cut_layer_split_assignments_.clear();
|
||||
graph_cut_layer_split_node_assignments_.clear();
|
||||
graph_cut_layer_split_primary_notice_logged_ = false;
|
||||
if (is_multi_device() && stream_layers_enabled) {
|
||||
LOG_WARN("%s: --stream-layers is not supported with multiple runtime backends; ignoring",
|
||||
get_desc().c_str());
|
||||
@@ -3074,6 +3255,11 @@ protected:
|
||||
|
||||
virtual void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") {}
|
||||
|
||||
virtual enum ggml_op param_usage_op(const std::string& name) const {
|
||||
(void)name;
|
||||
return GGML_OP_NONE;
|
||||
}
|
||||
|
||||
public:
|
||||
void init(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") {
|
||||
if (prefix.size() > 0) {
|
||||
@@ -3124,6 +3310,18 @@ public:
|
||||
}
|
||||
}
|
||||
|
||||
void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) {
|
||||
for (auto& pair : blocks) {
|
||||
pair.second->get_param_tensor_ops(tensor_ops);
|
||||
}
|
||||
for (auto& pair : params) {
|
||||
enum ggml_op op = param_usage_op(pair.first);
|
||||
if (op != GGML_OP_NONE) {
|
||||
tensor_ops[pair.second] = op;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
virtual std::string get_desc() {
|
||||
return "GGMLBlock";
|
||||
}
|
||||
@@ -3145,7 +3343,7 @@ public:
|
||||
|
||||
class Identity : public UnaryBlock {
|
||||
public:
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||
return x;
|
||||
}
|
||||
};
|
||||
@@ -3204,7 +3402,7 @@ public:
|
||||
force_prec_f32 = force_prec_f32_;
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||
ggml_tensor* w = params["weight"];
|
||||
ggml_tensor* b = nullptr;
|
||||
if (bias) {
|
||||
@@ -3251,6 +3449,10 @@ protected:
|
||||
params["weight"] = ggml_new_tensor_2d(ctx, wtype, embedding_dim, num_embeddings);
|
||||
}
|
||||
|
||||
enum ggml_op param_usage_op(const std::string& name) const override {
|
||||
return name == "weight" ? GGML_OP_GET_ROWS : GGML_OP_NONE;
|
||||
}
|
||||
|
||||
public:
|
||||
Embedding(int64_t num_embeddings, int64_t embedding_dim)
|
||||
: embedding_dim(embedding_dim),
|
||||
@@ -3258,7 +3460,7 @@ public:
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* input_ids) {
|
||||
ggml_tensor* input_ids) override {
|
||||
// input_ids: [N, n_token]
|
||||
auto weight = params["weight"];
|
||||
|
||||
@@ -3318,11 +3520,11 @@ public:
|
||||
scale = scale_value;
|
||||
}
|
||||
|
||||
std::string get_desc() {
|
||||
std::string get_desc() override {
|
||||
return "Conv2d";
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||
ggml_tensor* w = params["weight"];
|
||||
ggml_tensor* b = nullptr;
|
||||
if (bias) {
|
||||
@@ -3405,11 +3607,11 @@ public:
|
||||
scale = scale_value;
|
||||
}
|
||||
|
||||
std::string get_desc() {
|
||||
std::string get_desc() override {
|
||||
return "Conv2d_grouped";
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||
ggml_tensor* w = params["weight"];
|
||||
ggml_tensor* b = nullptr;
|
||||
if (bias) {
|
||||
@@ -3445,18 +3647,19 @@ public:
|
||||
if (groups == in_channels && groups == out_channels) {
|
||||
ggml_tensor* res;
|
||||
if (ctx->conv2d_direct_enabled) {
|
||||
res = ggml_conv_2d_dw_direct(ctx->ggml_ctx, x, w,
|
||||
res = ggml_conv_2d_dw_direct(ctx->ggml_ctx, w, x,
|
||||
stride.second, stride.first,
|
||||
padding.second, padding.first,
|
||||
dilation.second, dilation.first);
|
||||
} else {
|
||||
res = ggml_conv_2d_dw(ctx->ggml_ctx, x, w,
|
||||
res = ggml_conv_2d_dw(ctx->ggml_ctx, w, x,
|
||||
stride.second, stride.first,
|
||||
padding.second, padding.first,
|
||||
dilation.second, dilation.first);
|
||||
}
|
||||
if (b) {
|
||||
res = ggml_add(ctx->ggml_ctx, res, b);
|
||||
b = ggml_reshape_4d(ctx->ggml_ctx, b, 1, 1, b->ne[0], 1);
|
||||
res = ggml_add_inplace(ctx->ggml_ctx, res, b);
|
||||
}
|
||||
return res;
|
||||
}
|
||||
@@ -3561,7 +3764,7 @@ public:
|
||||
bias(bias),
|
||||
force_prec_f32(force_prec_f32) {}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||
ggml_tensor* w = params["weight"];
|
||||
ggml_tensor* b = nullptr;
|
||||
if (ctx->weight_adapter) {
|
||||
@@ -3614,7 +3817,7 @@ public:
|
||||
elementwise_affine(elementwise_affine),
|
||||
bias(bias) {}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||
ggml_tensor* w = nullptr;
|
||||
ggml_tensor* b = nullptr;
|
||||
|
||||
@@ -3701,7 +3904,7 @@ public:
|
||||
: hidden_size(hidden_size),
|
||||
eps(eps) {}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||
ggml_tensor* w = params["weight"];
|
||||
if (ctx->weight_adapter) {
|
||||
w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");
|
||||
|
||||
@@ -83,6 +83,10 @@ static bool parse_backend_module(const std::string& raw_name, SDBackendModule* m
|
||||
*module = SDBackendModule::UPSCALER;
|
||||
return true;
|
||||
}
|
||||
if (name == "detector" || name == "adetailer" || name == "yolo") {
|
||||
*module = SDBackendModule::DETECTOR;
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
@@ -956,6 +960,8 @@ const char* sd_backend_module_name(SDBackendModule module) {
|
||||
return "photomaker";
|
||||
case SDBackendModule::UPSCALER:
|
||||
return "upscaler";
|
||||
case SDBackendModule::DETECTOR:
|
||||
return "detector";
|
||||
}
|
||||
return "unknown";
|
||||
}
|
||||
|
||||
@@ -20,6 +20,7 @@ enum class SDBackendModule {
|
||||
CONTROL_NET,
|
||||
PHOTOMAKER,
|
||||
UPSCALER,
|
||||
DETECTOR,
|
||||
};
|
||||
|
||||
struct SDBackendAssignment {
|
||||
|
||||
@@ -1,9 +1,11 @@
|
||||
#include "core/layer_split_partition.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstdint>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <limits>
|
||||
#include <unordered_set>
|
||||
#include <utility>
|
||||
|
||||
#include "core/util.h"
|
||||
|
||||
@@ -62,160 +64,194 @@ namespace sd {
|
||||
return name != nullptr ? name : "unknown";
|
||||
}
|
||||
|
||||
static bool layer_split_backend_supports_tensor(ggml_backend_t backend, const ggml_tensor* tensor) {
|
||||
return backend != nullptr && tensor != nullptr && ggml_backend_supports_op(backend, tensor);
|
||||
static size_t graph_cut_layer_split_backend_vram_limit(const std::vector<size_t>& backend_vram_limits,
|
||||
size_t backend_index,
|
||||
size_t primary_backend_vram_limit) {
|
||||
if (backend_index < backend_vram_limits.size()) {
|
||||
return backend_vram_limits[backend_index];
|
||||
}
|
||||
return backend_index == 0 ? primary_backend_vram_limit : 0;
|
||||
}
|
||||
|
||||
static size_t layer_split_supported_target(const std::string& desc,
|
||||
const std::string& tensor_name,
|
||||
const ggml_tensor* tensor,
|
||||
const std::vector<ggml_backend_t>& backends,
|
||||
size_t preferred) {
|
||||
if (tensor == nullptr || backends.empty()) {
|
||||
return preferred;
|
||||
}
|
||||
size_t preferred_safe = std::min(preferred, backends.size() - 1);
|
||||
if (layer_split_backend_supports_tensor(backends[preferred_safe], tensor)) {
|
||||
return preferred_safe;
|
||||
}
|
||||
for (size_t i = 0; i < backends.size(); i++) {
|
||||
if (layer_split_backend_supports_tensor(backends[i], tensor)) {
|
||||
LOG_WARN("%s layer split: moving tensor '%s' from %s to %s because the preferred backend cannot run op=%s type=%s nbytes=%.2f MB",
|
||||
desc.c_str(),
|
||||
tensor_name.c_str(),
|
||||
layer_split_backend_device_display_name(backends[preferred_safe]).c_str(),
|
||||
layer_split_backend_device_display_name(backends[i]).c_str(),
|
||||
ggml_op_name(tensor->op),
|
||||
ggml_type_name(tensor->type),
|
||||
ggml_nbytes(tensor) / (1024.0 * 1024.0));
|
||||
return i;
|
||||
}
|
||||
}
|
||||
LOG_WARN("%s layer split: tensor '%s' is not supported by any split backend: op=%s type=%s nbytes=%.2f MB",
|
||||
desc.c_str(),
|
||||
tensor_name.c_str(),
|
||||
ggml_op_name(tensor->op),
|
||||
ggml_type_name(tensor->type),
|
||||
ggml_nbytes(tensor) / (1024.0 * 1024.0));
|
||||
return preferred_safe;
|
||||
}
|
||||
|
||||
std::vector<std::map<std::string, ggml_tensor*>> partition_layer_split_tensors(
|
||||
const std::string& desc,
|
||||
const std::map<std::string, ggml_tensor*>& tensors,
|
||||
const std::map<std::string, ggml_tensor*>& split_tensors,
|
||||
const std::vector<ggml_backend_t>& backends) {
|
||||
std::vector<std::map<std::string, ggml_tensor*>> partitions(backends.size());
|
||||
if (backends.empty()) {
|
||||
LOG_WARN("%s: no backend available for a layer split", desc.c_str());
|
||||
return partitions;
|
||||
}
|
||||
|
||||
std::map<int, int64_t> block_bytes;
|
||||
std::map<std::string, size_t> non_block_targets;
|
||||
std::vector<int64_t> other_bytes_by_backend(backends.size(), 0);
|
||||
int64_t total_block_bytes = 0;
|
||||
int64_t total_other_bytes = 0;
|
||||
int n_blocks = 0;
|
||||
for (const auto& kv : tensors) {
|
||||
int64_t bytes = (int64_t)ggml_nbytes(kv.second);
|
||||
int idx = split_tensors.count(kv.first) != 0 ? layer_split_tensor_block_index(kv.first) : -1;
|
||||
if (idx >= 0) {
|
||||
block_bytes[idx] += bytes;
|
||||
total_block_bytes += bytes;
|
||||
n_blocks = std::max(n_blocks, idx + 1);
|
||||
} else {
|
||||
size_t target = layer_split_supported_target(desc, kv.first, kv.second, backends, 0);
|
||||
non_block_targets[kv.first] = target;
|
||||
other_bytes_by_backend[target] += bytes;
|
||||
total_other_bytes += bytes;
|
||||
}
|
||||
}
|
||||
if (n_blocks == 0) {
|
||||
LOG_WARN("%s: no transformer blocks found for a layer split; keeping tensors on compatible backends starting from %s",
|
||||
desc.c_str(),
|
||||
layer_split_backend_device_display_name(backends[0]).c_str());
|
||||
for (const auto& kv : tensors) {
|
||||
size_t target = 0;
|
||||
auto target_it = non_block_targets.find(kv.first);
|
||||
if (target_it != non_block_targets.end()) {
|
||||
target = target_it->second;
|
||||
}
|
||||
partitions[target][kv.first] = kv.second;
|
||||
}
|
||||
return partitions;
|
||||
}
|
||||
|
||||
// Reserve compute headroom and subtract each device's actual non-block
|
||||
// bytes from its block budget.
|
||||
static std::vector<int64_t> graph_cut_layer_split_backend_capacities(const std::vector<ggml_backend_t>& backends,
|
||||
const std::vector<size_t>& backend_vram_limits,
|
||||
size_t primary_backend_vram_limit) {
|
||||
std::vector<int64_t> capacities(backends.size(), std::numeric_limits<int64_t>::max() / 4);
|
||||
constexpr int64_t compute_headroom_bytes = 2ll * 1024 * 1024 * 1024;
|
||||
std::vector<double> device_weights(backends.size(), 1.0);
|
||||
double weight_sum = 0.0;
|
||||
for (size_t i = 0; i < backends.size(); i++) {
|
||||
ggml_backend_dev_t dev = ggml_backend_get_device(backends[i]);
|
||||
size_t free_bytes = 0, total_bytes = 0;
|
||||
if (dev != nullptr) {
|
||||
ggml_backend_dev_memory(dev, &free_bytes, &total_bytes);
|
||||
}
|
||||
// Keep a small share even for tight devices instead of dropping them.
|
||||
int64_t usable_bytes = std::max<int64_t>((int64_t)free_bytes - compute_headroom_bytes,
|
||||
(int64_t)free_bytes / 8);
|
||||
device_weights[i] = usable_bytes > 0 ? (double)usable_bytes : 1.0;
|
||||
weight_sum += device_weights[i];
|
||||
if (free_bytes > 0) {
|
||||
capacities[i] = std::max<int64_t>((int64_t)free_bytes - compute_headroom_bytes, 0);
|
||||
}
|
||||
size_t limit_bytes = graph_cut_layer_split_backend_vram_limit(backend_vram_limits,
|
||||
i,
|
||||
primary_backend_vram_limit);
|
||||
if (limit_bytes > 0) {
|
||||
capacities[i] = std::min<int64_t>(capacities[i], (int64_t)limit_bytes);
|
||||
}
|
||||
}
|
||||
return capacities;
|
||||
}
|
||||
|
||||
std::vector<int64_t> block_budgets(backends.size(), 0);
|
||||
const int64_t total_bytes = total_block_bytes + total_other_bytes;
|
||||
for (size_t i = 0; i < backends.size(); i++) {
|
||||
int64_t budget = (int64_t)((double)total_bytes * device_weights[i] / weight_sum);
|
||||
budget = std::max<int64_t>(budget - other_bytes_by_backend[i], 0);
|
||||
block_budgets[i] = budget;
|
||||
}
|
||||
bool partition_graph_cut_layer_split(const char* desc,
|
||||
ggml_cgraph* gf,
|
||||
const sd::ggml_graph_cut::Plan& plan,
|
||||
const std::vector<ggml_backend_t>& split_backends,
|
||||
const std::vector<size_t>& backend_vram_limits,
|
||||
size_t primary_backend_vram_limit,
|
||||
std::unordered_map<const ggml_tensor*, ggml_backend_t>& param_assignments,
|
||||
const std::function<ggml_tensor*(ggml_tensor*)>& canonical_param_tensor,
|
||||
GraphCutLayerSplitAssignment* assignment_out) {
|
||||
GGML_ASSERT(gf != nullptr);
|
||||
GGML_ASSERT(assignment_out != nullptr);
|
||||
GGML_ASSERT(canonical_param_tensor != nullptr);
|
||||
GGML_ASSERT(!split_backends.empty());
|
||||
|
||||
GraphCutLayerSplitAssignment assignment;
|
||||
assignment.segment_count = plan.segments.size();
|
||||
assignment.tensors_by_backend.resize(split_backends.size());
|
||||
assignment.bytes_by_backend.resize(split_backends.size(), 0);
|
||||
assignment.first_segment_by_backend.resize(split_backends.size(), plan.segments.size());
|
||||
assignment.last_segment_by_backend.resize(split_backends.size(), 0);
|
||||
|
||||
std::vector<int> boundaries(backends.size(), n_blocks);
|
||||
size_t current = 0;
|
||||
int64_t used = 0;
|
||||
for (int b = 0; b < n_blocks; b++) {
|
||||
int64_t bytes = block_bytes.count(b) != 0 ? block_bytes[b] : 0;
|
||||
if (current + 1 < backends.size() && used > 0 && used + bytes > block_budgets[current]) {
|
||||
boundaries[current] = b;
|
||||
current++;
|
||||
used = 0;
|
||||
std::vector<std::vector<ggml_tensor*>> segment_params(plan.segments.size());
|
||||
std::vector<int64_t> segment_param_bytes(plan.segments.size(), 0);
|
||||
std::unordered_set<ggml_tensor*> seen_params;
|
||||
for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) {
|
||||
std::vector<ggml_tensor*> params = sd::ggml_graph_cut::param_tensors(gf, plan.segments[seg_idx]);
|
||||
for (ggml_tensor* raw_param : params) {
|
||||
ggml_tensor* param = canonical_param_tensor(raw_param);
|
||||
if (param == nullptr || !seen_params.insert(param).second) {
|
||||
continue;
|
||||
}
|
||||
segment_params[seg_idx].push_back(param);
|
||||
segment_param_bytes[seg_idx] += (int64_t)ggml_nbytes(param);
|
||||
}
|
||||
used += bytes;
|
||||
}
|
||||
|
||||
for (const auto& kv : tensors) {
|
||||
size_t target = 0;
|
||||
int idx = split_tensors.count(kv.first) != 0 ? layer_split_tensor_block_index(kv.first) : -1;
|
||||
if (idx >= 0) {
|
||||
while (target < boundaries.size() && idx >= boundaries[target]) {
|
||||
target++;
|
||||
int64_t total_param_bytes = 0;
|
||||
for (int64_t bytes : segment_param_bytes) {
|
||||
total_param_bytes += bytes;
|
||||
}
|
||||
if (total_param_bytes <= 0) {
|
||||
LOG_ERROR("%s graph-cut layer split found no graph params to assign", desc);
|
||||
return false;
|
||||
}
|
||||
|
||||
std::vector<int64_t> backend_capacities = graph_cut_layer_split_backend_capacities(split_backends,
|
||||
backend_vram_limits,
|
||||
primary_backend_vram_limit);
|
||||
|
||||
std::vector<ggml_backend_t> backend_by_segment(plan.segments.size(), split_backends[0]);
|
||||
size_t current_backend = 0;
|
||||
int64_t current_used = 0;
|
||||
for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) {
|
||||
int64_t bytes = segment_param_bytes[seg_idx];
|
||||
while (current_backend + 1 < split_backends.size() &&
|
||||
bytes > 0 &&
|
||||
current_used + bytes > backend_capacities[current_backend]) {
|
||||
current_backend++;
|
||||
current_used = 0;
|
||||
}
|
||||
if (bytes > 0 && current_used + bytes > backend_capacities[current_backend]) {
|
||||
LOG_ERROR("%s graph-cut layer split: segment %zu needs %.1f MB on %s, but only %.1f MB is available under current VRAM limits",
|
||||
desc,
|
||||
seg_idx,
|
||||
(current_used + bytes) / (1024.0 * 1024.0),
|
||||
layer_split_backend_device_display_name(split_backends[current_backend]).c_str(),
|
||||
backend_capacities[current_backend] / (1024.0 * 1024.0));
|
||||
return false;
|
||||
}
|
||||
current_used += bytes;
|
||||
backend_by_segment[seg_idx] = split_backends[current_backend];
|
||||
|
||||
for (ggml_tensor* param : segment_params[seg_idx]) {
|
||||
ggml_backend_t target_backend = split_backends[current_backend];
|
||||
auto assigned_it = param_assignments.find(param);
|
||||
if (assigned_it == param_assignments.end()) {
|
||||
param_assignments[param] = target_backend;
|
||||
assignment.has_new_param_assignment = true;
|
||||
} else {
|
||||
target_backend = assigned_it->second;
|
||||
}
|
||||
target = std::min(target, backends.size() - 1);
|
||||
target = layer_split_supported_target(desc, kv.first, kv.second, backends, target);
|
||||
|
||||
auto backend_it = std::find(split_backends.begin(), split_backends.end(), target_backend);
|
||||
if (backend_it == split_backends.end()) {
|
||||
LOG_ERROR("%s graph-cut layer split tensor '%s' is assigned to an unavailable backend",
|
||||
desc,
|
||||
ggml_get_name(param));
|
||||
return false;
|
||||
}
|
||||
size_t backend_idx = (size_t)std::distance(split_backends.begin(), backend_it);
|
||||
assignment.first_segment_by_backend[backend_idx] = std::min(assignment.first_segment_by_backend[backend_idx], seg_idx);
|
||||
assignment.last_segment_by_backend[backend_idx] = std::max(assignment.last_segment_by_backend[backend_idx], seg_idx + 1);
|
||||
assignment.tensors_by_backend[backend_idx].push_back(param);
|
||||
assignment.bytes_by_backend[backend_idx] += (int64_t)ggml_nbytes(param);
|
||||
}
|
||||
}
|
||||
|
||||
const int n_nodes = ggml_graph_n_nodes(gf);
|
||||
for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) {
|
||||
ggml_backend_t backend = backend_by_segment[seg_idx];
|
||||
const auto& segment = plan.segments[seg_idx];
|
||||
for (int node_index : segment.internal_node_indices) {
|
||||
if (node_index < 0 || node_index >= n_nodes) {
|
||||
continue;
|
||||
}
|
||||
ggml_tensor* node = ggml_graph_node(gf, node_index);
|
||||
if (node != nullptr) {
|
||||
assignment.node_assignments[node] = backend;
|
||||
}
|
||||
}
|
||||
for (int node_index : segment.output_node_indices) {
|
||||
if (node_index < 0 || node_index >= n_nodes) {
|
||||
continue;
|
||||
}
|
||||
ggml_tensor* node = ggml_graph_node(gf, node_index);
|
||||
if (node != nullptr) {
|
||||
assignment.node_assignments[node] = backend;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
*assignment_out = std::move(assignment);
|
||||
return true;
|
||||
}
|
||||
|
||||
void log_graph_cut_layer_split_assignment(const char* desc,
|
||||
const std::vector<ggml_backend_t>& split_backends,
|
||||
const GraphCutLayerSplitAssignment& assignment) {
|
||||
for (size_t i = 0; i < split_backends.size(); i++) {
|
||||
if (i >= assignment.tensors_by_backend.size() ||
|
||||
assignment.tensors_by_backend[i].empty()) {
|
||||
continue;
|
||||
}
|
||||
size_t first_segment = assignment.first_segment_by_backend[i] == assignment.segment_count
|
||||
? 0
|
||||
: assignment.first_segment_by_backend[i];
|
||||
size_t last_segment = assignment.last_segment_by_backend[i];
|
||||
if (assignment.has_new_param_assignment) {
|
||||
LOG_INFO("%s graph-cut layer split: %s <- segments [%zu, %zu), %zu tensors, %.1f MB",
|
||||
desc,
|
||||
layer_split_backend_device_display_name(split_backends[i]).c_str(),
|
||||
first_segment,
|
||||
last_segment,
|
||||
assignment.tensors_by_backend[i].size(),
|
||||
assignment.bytes_by_backend[i] / (1024.0 * 1024.0));
|
||||
} else {
|
||||
auto target_it = non_block_targets.find(kv.first);
|
||||
if (target_it != non_block_targets.end()) {
|
||||
target = target_it->second;
|
||||
}
|
||||
LOG_DEBUG("%s graph-cut layer split: %s <- segments [%zu, %zu), %zu tensors, %.1f MB",
|
||||
desc,
|
||||
layer_split_backend_device_display_name(split_backends[i]).c_str(),
|
||||
first_segment,
|
||||
last_segment,
|
||||
assignment.tensors_by_backend[i].size(),
|
||||
assignment.bytes_by_backend[i] / (1024.0 * 1024.0));
|
||||
}
|
||||
partitions[target][kv.first] = kv.second;
|
||||
}
|
||||
|
||||
int range_start = 0;
|
||||
for (size_t i = 0; i < backends.size(); i++) {
|
||||
int range_end = boundaries[i];
|
||||
const char* non_block_suffix = other_bytes_by_backend[i] > 0 ? " + non-block tensors" : "";
|
||||
LOG_INFO("%s layer split: %s <- blocks [%d, %d)%s",
|
||||
desc.c_str(),
|
||||
layer_split_backend_device_display_name(backends[i]).c_str(),
|
||||
range_start,
|
||||
range_end,
|
||||
non_block_suffix);
|
||||
range_start = range_end;
|
||||
}
|
||||
return partitions;
|
||||
}
|
||||
|
||||
} // namespace sd
|
||||
|
||||
@@ -1,23 +1,43 @@
|
||||
#ifndef __SD_CORE_LAYER_SPLIT_PARTITION_H__
|
||||
#define __SD_CORE_LAYER_SPLIT_PARTITION_H__
|
||||
|
||||
#include <map>
|
||||
#include <cstdint>
|
||||
#include <functional>
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
#include <vector>
|
||||
|
||||
#include "ggml-backend.h"
|
||||
#include "ggml.h"
|
||||
|
||||
#include "core/ggml_graph_cut.h"
|
||||
|
||||
namespace sd {
|
||||
|
||||
struct GraphCutLayerSplitAssignment {
|
||||
std::vector<std::vector<ggml_tensor*>> tensors_by_backend;
|
||||
std::vector<int64_t> bytes_by_backend;
|
||||
std::vector<size_t> first_segment_by_backend;
|
||||
std::vector<size_t> last_segment_by_backend;
|
||||
std::unordered_map<const ggml_tensor*, ggml_backend_t> node_assignments;
|
||||
size_t segment_count = 0;
|
||||
bool has_new_param_assignment = false;
|
||||
};
|
||||
|
||||
std::string layer_split_backend_device_display_name(ggml_backend_t backend);
|
||||
int layer_split_tensor_block_index(const std::string& name);
|
||||
|
||||
std::vector<std::map<std::string, ggml_tensor*>> partition_layer_split_tensors(
|
||||
const std::string& desc,
|
||||
const std::map<std::string, ggml_tensor*>& tensors,
|
||||
const std::map<std::string, ggml_tensor*>& split_tensors,
|
||||
const std::vector<ggml_backend_t>& backends);
|
||||
bool partition_graph_cut_layer_split(const char* desc,
|
||||
ggml_cgraph* gf,
|
||||
const sd::ggml_graph_cut::Plan& plan,
|
||||
const std::vector<ggml_backend_t>& split_backends,
|
||||
const std::vector<size_t>& backend_vram_limits,
|
||||
size_t primary_backend_vram_limit,
|
||||
std::unordered_map<const ggml_tensor*, ggml_backend_t>& param_assignments,
|
||||
const std::function<ggml_tensor*(ggml_tensor*)>& canonical_param_tensor,
|
||||
GraphCutLayerSplitAssignment* assignment_out);
|
||||
void log_graph_cut_layer_split_assignment(const char* desc,
|
||||
const std::vector<ggml_backend_t>& split_backends,
|
||||
const GraphCutLayerSplitAssignment& assignment);
|
||||
|
||||
} // namespace sd
|
||||
|
||||
|
||||
@@ -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__
|
||||
@@ -35,8 +35,10 @@ enum SDVersion {
|
||||
VERSION_WAN2,
|
||||
VERSION_WAN2_2_I2V,
|
||||
VERSION_WAN2_2_TI2V,
|
||||
VERSION_LINGBOT_VIDEO,
|
||||
VERSION_QWEN_IMAGE,
|
||||
VERSION_QWEN_IMAGE_LAYERED,
|
||||
VERSION_HUNYUAN_VIDEO,
|
||||
VERSION_ANIMA,
|
||||
VERSION_FLUX2,
|
||||
VERSION_FLUX2_KLEIN,
|
||||
@@ -53,6 +55,7 @@ enum SDVersion {
|
||||
VERSION_IDEOGRAM4,
|
||||
VERSION_SEFI_IMAGE,
|
||||
VERSION_KREA2,
|
||||
VERSION_MAGE_FLOW,
|
||||
VERSION_ESRGAN,
|
||||
VERSION_COUNT,
|
||||
};
|
||||
@@ -127,6 +130,13 @@ static inline bool sd_version_is_wan(SDVersion version) {
|
||||
return false;
|
||||
}
|
||||
|
||||
static inline bool sd_version_is_lingbot_video(SDVersion version) {
|
||||
if (version == VERSION_LINGBOT_VIDEO) {
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
static inline bool sd_version_is_qwen_image(SDVersion version) {
|
||||
if (version == VERSION_QWEN_IMAGE || version == VERSION_QWEN_IMAGE_LAYERED) {
|
||||
return true;
|
||||
@@ -134,6 +144,13 @@ static inline bool sd_version_is_qwen_image(SDVersion version) {
|
||||
return false;
|
||||
}
|
||||
|
||||
static inline bool sd_version_is_hunyuan_video(SDVersion version) {
|
||||
if (version == VERSION_HUNYUAN_VIDEO) {
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
static inline bool sd_version_is_anima(SDVersion version) {
|
||||
if (version == VERSION_ANIMA) {
|
||||
return true;
|
||||
@@ -211,6 +228,10 @@ static inline bool sd_version_is_krea2(SDVersion version) {
|
||||
return false;
|
||||
}
|
||||
|
||||
static inline bool sd_version_is_mage_flow(SDVersion version) {
|
||||
return version == VERSION_MAGE_FLOW;
|
||||
}
|
||||
|
||||
static inline bool sd_version_uses_flux_vae(SDVersion version) {
|
||||
if (sd_version_is_flux(version) || sd_version_is_z_image(version) || sd_version_is_boogu_image(version) || sd_version_is_longcat(version)) {
|
||||
return true;
|
||||
@@ -226,12 +247,16 @@ static inline bool sd_version_uses_flux2_vae(SDVersion version) {
|
||||
}
|
||||
|
||||
static inline bool sd_version_uses_wan_vae(SDVersion version) {
|
||||
if (sd_version_is_wan(version) || sd_version_is_qwen_image(version) || sd_version_is_krea2(version) || sd_version_is_anima(version)) {
|
||||
if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_qwen_image(version) || sd_version_is_krea2(version) || sd_version_is_anima(version)) {
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
static inline bool sd_version_uses_hunyuan_video_vae(SDVersion version) {
|
||||
return sd_version_is_hunyuan_video(version);
|
||||
}
|
||||
|
||||
static inline bool sd_version_is_inpaint(SDVersion version) {
|
||||
if (version == VERSION_SD1_INPAINT ||
|
||||
version == VERSION_SD2_INPAINT ||
|
||||
@@ -249,7 +274,9 @@ static inline bool sd_version_is_dit(SDVersion version) {
|
||||
sd_version_is_ltxav(version) ||
|
||||
sd_version_is_sd3(version) ||
|
||||
sd_version_is_wan(version) ||
|
||||
sd_version_is_lingbot_video(version) ||
|
||||
sd_version_is_qwen_image(version) ||
|
||||
sd_version_is_hunyuan_video(version) ||
|
||||
version == VERSION_HIDREAM_O1 ||
|
||||
sd_version_is_anima(version) ||
|
||||
sd_version_is_z_image(version) ||
|
||||
@@ -261,7 +288,8 @@ static inline bool sd_version_is_dit(SDVersion version) {
|
||||
sd_version_is_pid(version) ||
|
||||
sd_version_is_ideogram4(version) ||
|
||||
sd_version_is_sefi_image(version) ||
|
||||
sd_version_is_krea2(version)) {
|
||||
sd_version_is_krea2(version) ||
|
||||
sd_version_is_mage_flow(version)) {
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
|
||||
@@ -0,0 +1,88 @@
|
||||
#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 IPAdapterRunner : public GGMLRunner {
|
||||
ImageProjModel image_proj;
|
||||
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) {
|
||||
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 = "") {
|
||||
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 = image_proj.forward(&runner_ctx, embeds);
|
||||
ggml_build_forward_expand(gf, out);
|
||||
return gf;
|
||||
}
|
||||
|
||||
sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& image_embeds) {
|
||||
auto get_graph = [&]() -> ggml_cgraph* {
|
||||
return build_graph(image_embeds);
|
||||
};
|
||||
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true, true, true));
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace IPAdapter
|
||||
|
||||
#endif // __SD_MODEL_ADAPTER_IP_ADAPTER_HPP__
|
||||
@@ -14,6 +14,8 @@ struct LoraModel : public GGMLRunner {
|
||||
std::unordered_map<std::string, ggml_tensor*> lora_tensors;
|
||||
std::map<ggml_tensor*, ggml_tensor*> original_tensor_to_final_tensor;
|
||||
std::set<std::string> applied_lora_tensors;
|
||||
std::set<std::string> skipped_incompatible_lora_tensors;
|
||||
std::set<std::string> warned_incompatible_model_tensors;
|
||||
std::string file_path;
|
||||
std::shared_ptr<ModelManager> model_manager;
|
||||
ggml_backend_t params_backend = nullptr;
|
||||
@@ -133,6 +135,8 @@ struct LoraModel : public GGMLRunner {
|
||||
lora_tensors.clear();
|
||||
original_tensor_to_final_tensor.clear();
|
||||
applied_lora_tensors.clear();
|
||||
skipped_incompatible_lora_tensors.clear();
|
||||
warned_incompatible_model_tensors.clear();
|
||||
applied = false;
|
||||
tensor_preprocessed = false;
|
||||
}
|
||||
@@ -338,7 +342,9 @@ struct LoraModel : public GGMLRunner {
|
||||
iter = lora_tensors.find(hada_1_mid_name);
|
||||
if (iter != lora_tensors.end()) {
|
||||
hada_1_mid = ggml_ext_cast_f32(ctx, backend, iter->second);
|
||||
hada_1_up = ggml_cont(ctx, ggml_transpose(ctx, hada_1_up));
|
||||
if (hada_1_up != nullptr) {
|
||||
hada_1_up = ggml_cont(ctx, ggml_transpose(ctx, hada_1_up));
|
||||
}
|
||||
}
|
||||
|
||||
iter = lora_tensors.find(hada_2_down_name);
|
||||
@@ -354,7 +360,9 @@ struct LoraModel : public GGMLRunner {
|
||||
iter = lora_tensors.find(hada_2_mid_name);
|
||||
if (iter != lora_tensors.end()) {
|
||||
hada_2_mid = ggml_ext_cast_f32(ctx, backend, iter->second);
|
||||
hada_2_up = ggml_cont(ctx, ggml_transpose(ctx, hada_2_up));
|
||||
if (hada_2_up != nullptr) {
|
||||
hada_2_up = ggml_cont(ctx, ggml_transpose(ctx, hada_2_up));
|
||||
}
|
||||
}
|
||||
|
||||
if (hada_1_up == nullptr || hada_1_down == nullptr || hada_2_up == nullptr || hada_2_down == nullptr) {
|
||||
@@ -546,7 +554,27 @@ struct LoraModel : public GGMLRunner {
|
||||
}
|
||||
}
|
||||
|
||||
GGML_ASSERT(ggml_nelements(diff) == ggml_nelements(model_tensor));
|
||||
if (ggml_nelements(diff) != ggml_nelements(model_tensor)) {
|
||||
const std::string lora_tensor_prefix = "lora." + model_tensor_name + ".";
|
||||
for (const auto& tensor_name : applied_lora_tensors) {
|
||||
if (starts_with(tensor_name, lora_tensor_prefix)) {
|
||||
skipped_incompatible_lora_tensors.insert(tensor_name);
|
||||
}
|
||||
}
|
||||
if (warned_incompatible_model_tensors.insert(model_tensor_name).second) {
|
||||
LOG_WARN("skip incompatible LoRA tensor |%s|: model shape = [%lld, %lld, %lld, %lld], LoRA shape = [%lld, %lld, %lld, %lld]",
|
||||
model_tensor_name.c_str(),
|
||||
static_cast<long long>(model_tensor->ne[0]),
|
||||
static_cast<long long>(model_tensor->ne[1]),
|
||||
static_cast<long long>(model_tensor->ne[2]),
|
||||
static_cast<long long>(model_tensor->ne[3]),
|
||||
static_cast<long long>(diff->ne[0]),
|
||||
static_cast<long long>(diff->ne[1]),
|
||||
static_cast<long long>(diff->ne[2]),
|
||||
static_cast<long long>(diff->ne[3]));
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
diff = ggml_reshape(ctx, diff, model_tensor);
|
||||
}
|
||||
return diff;
|
||||
@@ -555,6 +583,7 @@ struct LoraModel : public GGMLRunner {
|
||||
ggml_tensor* get_out_diff(ggml_context* ctx,
|
||||
ggml_backend_t backend,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* model_weight,
|
||||
WeightAdapter::ForwardParams forward_params,
|
||||
const std::string& model_tensor_name) {
|
||||
ggml_tensor* out_diff = nullptr;
|
||||
@@ -707,6 +736,43 @@ struct LoraModel : public GGMLRunner {
|
||||
break;
|
||||
}
|
||||
|
||||
if (!is_conv2d) {
|
||||
const int64_t down_in = lora_down->ne[0];
|
||||
const int64_t down_out = lora_down->ne[1];
|
||||
const int64_t up_in = lora_up->ne[0];
|
||||
const int64_t up_out = lora_up->ne[1];
|
||||
|
||||
bool compatible = down_in == model_weight->ne[0] &&
|
||||
up_out == model_weight->ne[1];
|
||||
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);
|
||||
skipped_incompatible_lora_tensors.insert(lora_mid_name);
|
||||
skipped_incompatible_lora_tensors.insert(scale_name);
|
||||
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 shape = [%lld, %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>(model_weight->ne[1]),
|
||||
static_cast<long long>(down_in),
|
||||
static_cast<long long>(down_out),
|
||||
static_cast<long long>(up_in),
|
||||
static_cast<long long>(up_out));
|
||||
}
|
||||
index++;
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
applied_lora_tensors.insert(lora_up_name);
|
||||
applied_lora_tensors.insert(lora_down_name);
|
||||
|
||||
@@ -869,10 +935,13 @@ struct LoraModel : public GGMLRunner {
|
||||
void stat(bool at_runntime = false) {
|
||||
size_t total_lora_tensors_count = 0;
|
||||
size_t applied_lora_tensors_count = 0;
|
||||
size_t skipped_lora_tensors_count = 0;
|
||||
|
||||
for (auto& kv : lora_tensors) {
|
||||
total_lora_tensors_count++;
|
||||
if (applied_lora_tensors.find(kv.first) == applied_lora_tensors.end()) {
|
||||
if (skipped_incompatible_lora_tensors.find(kv.first) != skipped_incompatible_lora_tensors.end()) {
|
||||
skipped_lora_tensors_count++;
|
||||
} else if (applied_lora_tensors.find(kv.first) == applied_lora_tensors.end()) {
|
||||
if (!at_runntime) {
|
||||
LOG_WARN("unused lora tensor |%s|", kv.first.c_str());
|
||||
print_ggml_tensor(kv.second, true);
|
||||
@@ -884,12 +953,17 @@ struct LoraModel : public GGMLRunner {
|
||||
/* Don't worry if this message shows up twice in the logs per LoRA,
|
||||
* this function is called once to calculate the required buffer size
|
||||
* and then again to actually generate a graph to be used */
|
||||
if (!at_runntime && applied_lora_tensors_count != total_lora_tensors_count) {
|
||||
size_t compatible_lora_tensors_count = total_lora_tensors_count - skipped_lora_tensors_count;
|
||||
if (!at_runntime && applied_lora_tensors_count != compatible_lora_tensors_count) {
|
||||
LOG_WARN("Only (%lu / %lu) LoRA tensors have been applied, lora_file_path = %s",
|
||||
applied_lora_tensors_count, total_lora_tensors_count, file_path.c_str());
|
||||
applied_lora_tensors_count, compatible_lora_tensors_count, file_path.c_str());
|
||||
} else {
|
||||
LOG_INFO("(%lu / %lu) LoRA tensors have been applied, lora_file_path = %s",
|
||||
applied_lora_tensors_count, total_lora_tensors_count, file_path.c_str());
|
||||
applied_lora_tensors_count, compatible_lora_tensors_count, file_path.c_str());
|
||||
}
|
||||
if (skipped_lora_tensors_count > 0) {
|
||||
LOG_WARN("(%lu / %lu) incompatible LoRA tensors have been skipped, lora_file_path = %s",
|
||||
skipped_lora_tensors_count, total_lora_tensors_count, file_path.c_str());
|
||||
}
|
||||
}
|
||||
};
|
||||
@@ -953,7 +1027,7 @@ public:
|
||||
forward_params.conv2d.scale);
|
||||
}
|
||||
for (auto& lora_model : lora_models) {
|
||||
ggml_tensor* out_diff = lora_model->get_out_diff(ctx, backend, x, forward_params, prefix + "weight");
|
||||
ggml_tensor* out_diff = lora_model->get_out_diff(ctx, backend, x, w, forward_params, prefix + "weight");
|
||||
if (out_diff == nullptr) {
|
||||
continue;
|
||||
}
|
||||
|
||||
@@ -294,7 +294,7 @@ public:
|
||||
|
||||
auto net_0 = std::dynamic_pointer_cast<UnaryBlock>(blocks["net.0"]);
|
||||
auto net_2 = std::dynamic_pointer_cast<Linear>(blocks["net.2"]);
|
||||
if (sd_backend_is(ctx->backend, "Vulkan")) {
|
||||
if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
|
||||
net_2->set_force_prec_f32(true);
|
||||
}
|
||||
|
||||
@@ -310,17 +310,33 @@ protected:
|
||||
int64_t context_dim;
|
||||
int64_t n_head;
|
||||
int64_t d_head;
|
||||
bool xtra_dim = false;
|
||||
bool xtra_dim = false;
|
||||
bool enable_ip = false;
|
||||
bool has_ip = false;
|
||||
|
||||
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
||||
GGMLBlock::init_params(ctx, tensor_storage_map, prefix);
|
||||
if (enable_ip &&
|
||||
tensor_storage_map.find(prefix + "to_k_ip.weight") != tensor_storage_map.end()) {
|
||||
has_ip = true;
|
||||
int64_t inner_dim = d_head * n_head;
|
||||
int64_t ip_dim = tensor_storage_map.at(prefix + "to_k_ip.weight").ne[0];
|
||||
blocks["to_k_ip"] = std::shared_ptr<GGMLBlock>(new Linear(ip_dim, inner_dim, false));
|
||||
blocks["to_v_ip"] = std::shared_ptr<GGMLBlock>(new Linear(ip_dim, inner_dim, false));
|
||||
}
|
||||
}
|
||||
|
||||
public:
|
||||
CrossAttention(int64_t query_dim,
|
||||
int64_t context_dim,
|
||||
int64_t n_head,
|
||||
int64_t d_head)
|
||||
int64_t d_head,
|
||||
bool enable_ip = false)
|
||||
: n_head(n_head),
|
||||
d_head(d_head),
|
||||
query_dim(query_dim),
|
||||
context_dim(context_dim) {
|
||||
context_dim(context_dim),
|
||||
enable_ip(enable_ip) {
|
||||
int64_t inner_dim = d_head * n_head;
|
||||
if (context_dim == 320 && d_head == 320) {
|
||||
// LOG_DEBUG("CrossAttention: temp set dim to 1024 for sdxs_09");
|
||||
@@ -363,6 +379,15 @@ public:
|
||||
}
|
||||
x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, inner_dim]
|
||||
|
||||
if (has_ip && ctx->ip_context != nullptr && ctx->ip_scale != 0.0f) {
|
||||
auto to_k_ip = std::dynamic_pointer_cast<Linear>(blocks["to_k_ip"]);
|
||||
auto to_v_ip = std::dynamic_pointer_cast<Linear>(blocks["to_v_ip"]);
|
||||
auto k_ip = to_k_ip->forward(ctx, ctx->ip_context);
|
||||
auto v_ip = to_v_ip->forward(ctx, ctx->ip_context);
|
||||
auto x_ip = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k_ip, v_ip, n_head, nullptr, false, ctx->flash_attn_enabled);
|
||||
x = ggml_add(ctx->ggml_ctx, x, ggml_scale(ctx->ggml_ctx, x_ip, ctx->ip_scale));
|
||||
}
|
||||
|
||||
x = to_out_0->forward(ctx, x); // [N, n_token, query_dim]
|
||||
return x;
|
||||
}
|
||||
@@ -387,7 +412,7 @@ public:
|
||||
// inner_dim is always None or equal to dim
|
||||
// gated_ff is always True
|
||||
blocks["attn1"] = std::shared_ptr<GGMLBlock>(new CrossAttention(dim, dim, n_head, d_head));
|
||||
blocks["attn2"] = std::shared_ptr<GGMLBlock>(new CrossAttention(dim, context_dim, n_head, d_head));
|
||||
blocks["attn2"] = std::shared_ptr<GGMLBlock>(new CrossAttention(dim, context_dim, n_head, d_head, true));
|
||||
blocks["ff"] = std::shared_ptr<GGMLBlock>(new FeedForward(dim, dim));
|
||||
blocks["norm1"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
|
||||
blocks["norm2"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
|
||||
@@ -450,7 +475,7 @@ protected:
|
||||
int64_t context_dim = 768; // hidden_size, 1024 for VERSION_SD2
|
||||
bool use_linear = false;
|
||||
|
||||
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") {
|
||||
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
|
||||
auto iter = tensor_storage_map.find(prefix + "proj_out.weight");
|
||||
if (iter != tensor_storage_map.end()) {
|
||||
int64_t inner_dim = n_head * d_head;
|
||||
|
||||
@@ -18,10 +18,6 @@ namespace Rope {
|
||||
DECREASE,
|
||||
};
|
||||
|
||||
__STATIC_INLINE__ RefIndexMode ref_index_mode_from_bool(bool increase_ref_index) {
|
||||
return increase_ref_index ? RefIndexMode::INCREASE : RefIndexMode::FIXED;
|
||||
}
|
||||
|
||||
template <class T>
|
||||
__STATIC_INLINE__ std::vector<T> linspace(T start, T end, int num) {
|
||||
std::vector<T> result(num);
|
||||
@@ -539,6 +535,33 @@ namespace Rope {
|
||||
return vid_ids_repeated;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_hunyuan_video_ids(int t,
|
||||
int h,
|
||||
int w,
|
||||
int patch_t,
|
||||
int patch_h,
|
||||
int patch_w,
|
||||
int bs,
|
||||
int context_len) {
|
||||
std::vector<std::vector<float>> txt_ids(bs * context_len, std::vector<float>(3, 0.0f));
|
||||
auto img_ids = gen_vid_ids(t, h, w, patch_t, patch_h, patch_w, bs);
|
||||
return concat_ids(txt_ids, img_ids, bs);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<float> gen_hunyuan_video_pe(int t,
|
||||
int h,
|
||||
int w,
|
||||
int patch_t,
|
||||
int patch_h,
|
||||
int patch_w,
|
||||
int bs,
|
||||
int context_len,
|
||||
float theta,
|
||||
const std::vector<int>& axes_dim) {
|
||||
auto ids = gen_hunyuan_video_ids(t, h, w, patch_t, patch_h, patch_w, bs, context_len);
|
||||
return embed_nd(ids, bs, theta, axes_dim);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_qwen_image_ids(int t,
|
||||
int h,
|
||||
int w,
|
||||
@@ -631,6 +654,43 @@ namespace Rope {
|
||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<float> gen_mage_flow_pe(int h,
|
||||
int w,
|
||||
int bs,
|
||||
int context_len,
|
||||
const std::vector<ggml_tensor*>& ref_latents,
|
||||
int theta,
|
||||
const std::vector<int>& axes_dim) {
|
||||
const int axes_dim_num = static_cast<int>(axes_dim.size());
|
||||
auto make_image_ids = [=](int image_h, int image_w, int image_index) {
|
||||
std::vector<std::vector<float>> image_ids(static_cast<size_t>(bs) * image_h * image_w,
|
||||
std::vector<float>(axes_dim_num, 0.f));
|
||||
int h_start = -(image_h - image_h / 2);
|
||||
int w_start = -(image_w - image_w / 2);
|
||||
for (int b = 0; b < bs; ++b) {
|
||||
for (int y = 0; y < image_h; ++y) {
|
||||
for (int x = 0; x < image_w; ++x) {
|
||||
auto& id = image_ids[static_cast<size_t>(b) * image_h * image_w + y * image_w + x];
|
||||
id[0] = static_cast<float>(image_index);
|
||||
id[1] = static_cast<float>(h_start + y);
|
||||
id[2] = static_cast<float>(w_start + x);
|
||||
}
|
||||
}
|
||||
}
|
||||
return image_ids;
|
||||
};
|
||||
auto ids = gen_flux_txt_ids(bs, context_len, axes_dim_num, {});
|
||||
auto img_ids = make_image_ids(h, w, 0);
|
||||
ids = concat_ids(ids, img_ids, bs);
|
||||
for (size_t i = 0; i < ref_latents.size(); ++i) {
|
||||
auto ref_ids = make_image_ids(static_cast<int>(ref_latents[i]->ne[1]),
|
||||
static_cast<int>(ref_latents[i]->ne[0]),
|
||||
static_cast<int>(i + 1));
|
||||
ids = concat_ids(ids, ref_ids, bs);
|
||||
}
|
||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_lens_ids(int h,
|
||||
int w,
|
||||
int bs,
|
||||
@@ -759,6 +819,40 @@ namespace Rope {
|
||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_lingbot_video_ids(int t,
|
||||
int h,
|
||||
int w,
|
||||
int pt,
|
||||
int ph,
|
||||
int pw,
|
||||
int bs,
|
||||
int context_len) {
|
||||
auto vid_ids_repeated = gen_vid_ids(t, h, w, pt, ph, pw, bs, context_len + 1);
|
||||
|
||||
std::vector<std::vector<float>> txt_ids(bs * context_len, std::vector<float>(3, 0.0f));
|
||||
for (int i = 0; i < bs; ++i) {
|
||||
for (int j = 0; j < context_len; ++j) {
|
||||
txt_ids[i * context_len + j][0] = static_cast<float>(j + 1);
|
||||
}
|
||||
}
|
||||
|
||||
return concat_ids(vid_ids_repeated, txt_ids, bs);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<float> gen_lingbot_video_pe(int t,
|
||||
int h,
|
||||
int w,
|
||||
int pt,
|
||||
int ph,
|
||||
int pw,
|
||||
int bs,
|
||||
int context_len,
|
||||
int theta,
|
||||
const std::vector<int>& axes_dim) {
|
||||
std::vector<std::vector<float>> ids = gen_lingbot_video_ids(t, h, w, pt, ph, pw, bs, context_len);
|
||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_qwen2vl_ids(int grid_h,
|
||||
int grid_w,
|
||||
int merge_size,
|
||||
|
||||
@@ -0,0 +1,362 @@
|
||||
#ifndef __SD_MODEL_DETECTOR_YOLOV8_H__
|
||||
#define __SD_MODEL_DETECTOR_YOLOV8_H__
|
||||
|
||||
#include <algorithm>
|
||||
#include <array>
|
||||
#include <map>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "core/ggml_extend.hpp"
|
||||
#include "core/util.h"
|
||||
|
||||
struct YOLOv8Config {
|
||||
std::array<int, 23> out_channels{};
|
||||
std::map<int, int> hidden_channels;
|
||||
std::map<int, int> repeats;
|
||||
int detect_box_channels = 0;
|
||||
int detect_cls_channels = 0;
|
||||
int reg_max = 0;
|
||||
int num_classes = 0;
|
||||
bool valid = false;
|
||||
|
||||
static YOLOv8Config detect_from_weights(const String2TensorStorage& tensor_storage_map,
|
||||
const std::string& prefix = "") {
|
||||
YOLOv8Config config;
|
||||
auto full_name = [&](const std::string& name) {
|
||||
return prefix.empty() ? name : prefix + "." + name;
|
||||
};
|
||||
auto find_weight = [&](const std::string& name) -> const TensorStorage* {
|
||||
auto iter = tensor_storage_map.find(full_name(name));
|
||||
return iter == tensor_storage_map.end() ? nullptr : &iter->second;
|
||||
};
|
||||
auto conv_out = [&](const std::string& name) -> int {
|
||||
const TensorStorage* weight = find_weight(name);
|
||||
return weight != nullptr && weight->n_dims == 4 ? static_cast<int>(weight->ne[3]) : 0;
|
||||
};
|
||||
|
||||
for (int layer : {0, 1, 3, 5, 7, 16, 19}) {
|
||||
config.out_channels[layer] = conv_out("model." + std::to_string(layer) + ".conv.weight");
|
||||
}
|
||||
for (int layer : {2, 4, 6, 8, 12, 15, 18, 21}) {
|
||||
const std::string base = "model." + std::to_string(layer);
|
||||
config.out_channels[layer] = conv_out(base + ".cv2.conv.weight");
|
||||
config.hidden_channels[layer] = conv_out(base + ".cv1.conv.weight") / 2;
|
||||
|
||||
int repeat_count = 0;
|
||||
while (find_weight(base + ".m." + std::to_string(repeat_count) + ".cv1.conv.weight") != nullptr) {
|
||||
++repeat_count;
|
||||
}
|
||||
config.repeats[layer] = repeat_count;
|
||||
}
|
||||
config.out_channels[9] = conv_out("model.9.cv2.conv.weight");
|
||||
|
||||
config.detect_box_channels = conv_out("model.22.cv2.0.0.conv.weight");
|
||||
config.detect_cls_channels = conv_out("model.22.cv3.0.0.conv.weight");
|
||||
const int box_outputs = conv_out("model.22.cv2.0.2.weight");
|
||||
config.num_classes = conv_out("model.22.cv3.0.2.weight");
|
||||
config.reg_max = box_outputs / 4;
|
||||
|
||||
config.valid = config.out_channels[0] > 0 && config.out_channels[9] > 0 &&
|
||||
config.out_channels[15] > 0 && config.out_channels[18] > 0 &&
|
||||
config.out_channels[21] > 0 && config.detect_box_channels > 0 &&
|
||||
config.detect_cls_channels > 0 && box_outputs > 0 && box_outputs % 4 == 0 &&
|
||||
config.num_classes > 0;
|
||||
for (int layer : {2, 4, 6, 8, 12, 15, 18, 21}) {
|
||||
config.valid = config.valid && config.hidden_channels[layer] > 0 && config.repeats[layer] > 0;
|
||||
}
|
||||
|
||||
if (config.valid) {
|
||||
LOG_DEBUG("yolov8: classes=%d, reg_max=%d, p3=%d, p4=%d, p5=%d",
|
||||
config.num_classes,
|
||||
config.reg_max,
|
||||
config.out_channels[15],
|
||||
config.out_channels[18],
|
||||
config.out_channels[21]);
|
||||
}
|
||||
return config;
|
||||
}
|
||||
};
|
||||
|
||||
class YOLOConv : public UnaryBlock {
|
||||
int out_channels_ = 0;
|
||||
|
||||
public:
|
||||
YOLOConv(int in_channels, int out_channels, int kernel, int stride = 1)
|
||||
: out_channels_(out_channels) {
|
||||
blocks["conv"] = std::shared_ptr<GGMLBlock>(new Conv2d(in_channels,
|
||||
out_channels,
|
||||
{kernel, kernel},
|
||||
{stride, stride},
|
||||
{kernel / 2, kernel / 2},
|
||||
{1, 1},
|
||||
true));
|
||||
}
|
||||
|
||||
int out_channels() const {
|
||||
return out_channels_;
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||
auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
|
||||
return ggml_silu_inplace(ctx->ggml_ctx, conv->forward(ctx, x));
|
||||
}
|
||||
};
|
||||
|
||||
class YOLOBottleneck : public UnaryBlock {
|
||||
bool shortcut_ = false;
|
||||
|
||||
public:
|
||||
YOLOBottleneck(int channels, bool shortcut)
|
||||
: shortcut_(shortcut) {
|
||||
blocks["cv1"] = std::shared_ptr<GGMLBlock>(new YOLOConv(channels, channels, 3));
|
||||
blocks["cv2"] = std::shared_ptr<GGMLBlock>(new YOLOConv(channels, channels, 3));
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||
auto cv1 = std::dynamic_pointer_cast<YOLOConv>(blocks["cv1"]);
|
||||
auto cv2 = std::dynamic_pointer_cast<YOLOConv>(blocks["cv2"]);
|
||||
auto out = cv2->forward(ctx, cv1->forward(ctx, x));
|
||||
return shortcut_ ? ggml_add(ctx->ggml_ctx, x, out) : out;
|
||||
}
|
||||
};
|
||||
|
||||
class YOLOC2f : public UnaryBlock {
|
||||
int hidden_channels_ = 0;
|
||||
int repeats_ = 0;
|
||||
|
||||
public:
|
||||
YOLOC2f(int in_channels,
|
||||
int out_channels,
|
||||
int hidden_channels,
|
||||
int repeats,
|
||||
bool shortcut)
|
||||
: hidden_channels_(hidden_channels), repeats_(repeats) {
|
||||
blocks["cv1"] = std::shared_ptr<GGMLBlock>(new YOLOConv(in_channels, hidden_channels * 2, 1));
|
||||
blocks["cv2"] = std::shared_ptr<GGMLBlock>(new YOLOConv(hidden_channels * (2 + repeats), out_channels, 1));
|
||||
for (int i = 0; i < repeats; ++i) {
|
||||
blocks["m." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new YOLOBottleneck(hidden_channels, shortcut));
|
||||
}
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||
auto cv1 = std::dynamic_pointer_cast<YOLOConv>(blocks["cv1"]);
|
||||
auto cv2 = std::dynamic_pointer_cast<YOLOConv>(blocks["cv2"]);
|
||||
auto split = cv1->forward(ctx, x);
|
||||
|
||||
// split: [N, 2*C, H, W], ggml layout [W, H, 2*C, N].
|
||||
auto y0 = ggml_view_4d(ctx->ggml_ctx,
|
||||
split,
|
||||
split->ne[0],
|
||||
split->ne[1],
|
||||
hidden_channels_,
|
||||
split->ne[3],
|
||||
split->nb[1],
|
||||
split->nb[2],
|
||||
split->nb[3],
|
||||
0);
|
||||
auto y1 = ggml_view_4d(ctx->ggml_ctx,
|
||||
split,
|
||||
split->ne[0],
|
||||
split->ne[1],
|
||||
hidden_channels_,
|
||||
split->ne[3],
|
||||
split->nb[1],
|
||||
split->nb[2],
|
||||
split->nb[3],
|
||||
static_cast<size_t>(hidden_channels_) * split->nb[2]);
|
||||
auto joined = ggml_concat(ctx->ggml_ctx, y0, y1, 2);
|
||||
auto last = y1;
|
||||
for (int i = 0; i < repeats_; ++i) {
|
||||
auto block = std::dynamic_pointer_cast<YOLOBottleneck>(blocks["m." + std::to_string(i)]);
|
||||
last = block->forward(ctx, last);
|
||||
joined = ggml_concat(ctx->ggml_ctx, joined, last, 2);
|
||||
}
|
||||
return cv2->forward(ctx, joined);
|
||||
}
|
||||
};
|
||||
|
||||
class YOLOSPPF : public UnaryBlock {
|
||||
public:
|
||||
YOLOSPPF(int in_channels, int out_channels) {
|
||||
blocks["cv1"] = std::shared_ptr<GGMLBlock>(new YOLOConv(in_channels, in_channels / 2, 1));
|
||||
blocks["cv2"] = std::shared_ptr<GGMLBlock>(new YOLOConv(in_channels * 2, out_channels, 1));
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||
auto cv1 = std::dynamic_pointer_cast<YOLOConv>(blocks["cv1"]);
|
||||
auto cv2 = std::dynamic_pointer_cast<YOLOConv>(blocks["cv2"]);
|
||||
x = cv1->forward(ctx, x);
|
||||
auto y1 = ggml_pool_2d(ctx->ggml_ctx, x, GGML_OP_POOL_MAX, 5, 5, 1, 1, 2, 2);
|
||||
auto y2 = ggml_pool_2d(ctx->ggml_ctx, y1, GGML_OP_POOL_MAX, 5, 5, 1, 1, 2, 2);
|
||||
auto y3 = ggml_pool_2d(ctx->ggml_ctx, y2, GGML_OP_POOL_MAX, 5, 5, 1, 1, 2, 2);
|
||||
auto out = ggml_concat(ctx->ggml_ctx, x, y1, 2);
|
||||
out = ggml_concat(ctx->ggml_ctx, out, y2, 2);
|
||||
out = ggml_concat(ctx->ggml_ctx, out, y3, 2);
|
||||
return cv2->forward(ctx, out);
|
||||
}
|
||||
};
|
||||
|
||||
class YOLODetect : public GGMLBlock {
|
||||
int num_classes_ = 0;
|
||||
int reg_max_ = 0;
|
||||
|
||||
public:
|
||||
YOLODetect(const std::array<int, 3>& in_channels,
|
||||
int box_channels,
|
||||
int cls_channels,
|
||||
int reg_max,
|
||||
int num_classes)
|
||||
: num_classes_(num_classes), reg_max_(reg_max) {
|
||||
for (int i = 0; i < 3; ++i) {
|
||||
const std::string box = "cv2." + std::to_string(i);
|
||||
blocks[box + ".0"] = std::shared_ptr<GGMLBlock>(new YOLOConv(in_channels[i], box_channels, 3));
|
||||
blocks[box + ".1"] = std::shared_ptr<GGMLBlock>(new YOLOConv(box_channels, box_channels, 3));
|
||||
blocks[box + ".2"] = std::shared_ptr<GGMLBlock>(new Conv2d(box_channels, reg_max * 4, {1, 1}, {1, 1}, {0, 0}, {1, 1}, true));
|
||||
|
||||
const std::string cls = "cv3." + std::to_string(i);
|
||||
blocks[cls + ".0"] = std::shared_ptr<GGMLBlock>(new YOLOConv(in_channels[i], cls_channels, 3));
|
||||
blocks[cls + ".1"] = std::shared_ptr<GGMLBlock>(new YOLOConv(cls_channels, cls_channels, 3));
|
||||
blocks[cls + ".2"] = std::shared_ptr<GGMLBlock>(new Conv2d(cls_channels, num_classes, {1, 1}, {1, 1}, {0, 0}, {1, 1}, true));
|
||||
}
|
||||
}
|
||||
|
||||
ggml_tensor* forward_scale(GGMLRunnerContext* ctx, ggml_tensor* x, int index) {
|
||||
const std::string box = "cv2." + std::to_string(index);
|
||||
auto box0 = std::dynamic_pointer_cast<YOLOConv>(blocks[box + ".0"]);
|
||||
auto box1 = std::dynamic_pointer_cast<YOLOConv>(blocks[box + ".1"]);
|
||||
auto box2 = std::dynamic_pointer_cast<Conv2d>(blocks[box + ".2"]);
|
||||
|
||||
const std::string cls = "cv3." + std::to_string(index);
|
||||
auto cls0 = std::dynamic_pointer_cast<YOLOConv>(blocks[cls + ".0"]);
|
||||
auto cls1 = std::dynamic_pointer_cast<YOLOConv>(blocks[cls + ".1"]);
|
||||
auto cls2 = std::dynamic_pointer_cast<Conv2d>(blocks[cls + ".2"]);
|
||||
|
||||
auto boxes = box2->forward(ctx, box1->forward(ctx, box0->forward(ctx, x)));
|
||||
auto classes = cls2->forward(ctx, cls1->forward(ctx, cls0->forward(ctx, x)));
|
||||
return ggml_concat(ctx->ggml_ctx, boxes, classes, 2);
|
||||
}
|
||||
|
||||
int output_channels() const {
|
||||
return reg_max_ * 4 + num_classes_;
|
||||
}
|
||||
};
|
||||
|
||||
class YOLOv8Model : public GGMLBlock {
|
||||
YOLOv8Config config_;
|
||||
|
||||
std::shared_ptr<YOLOC2f> make_c2f(int layer, int in_channels, bool shortcut) {
|
||||
return std::make_shared<YOLOC2f>(in_channels,
|
||||
config_.out_channels[layer],
|
||||
config_.hidden_channels.at(layer),
|
||||
config_.repeats.at(layer),
|
||||
shortcut);
|
||||
}
|
||||
|
||||
public:
|
||||
explicit YOLOv8Model(YOLOv8Config config)
|
||||
: config_(std::move(config)) {
|
||||
blocks["model.0"] = std::make_shared<YOLOConv>(3, config_.out_channels[0], 3, 2);
|
||||
blocks["model.1"] = std::make_shared<YOLOConv>(config_.out_channels[0], config_.out_channels[1], 3, 2);
|
||||
blocks["model.2"] = make_c2f(2, config_.out_channels[1], true);
|
||||
blocks["model.3"] = std::make_shared<YOLOConv>(config_.out_channels[2], config_.out_channels[3], 3, 2);
|
||||
blocks["model.4"] = make_c2f(4, config_.out_channels[3], true);
|
||||
blocks["model.5"] = std::make_shared<YOLOConv>(config_.out_channels[4], config_.out_channels[5], 3, 2);
|
||||
blocks["model.6"] = make_c2f(6, config_.out_channels[5], true);
|
||||
blocks["model.7"] = std::make_shared<YOLOConv>(config_.out_channels[6], config_.out_channels[7], 3, 2);
|
||||
blocks["model.8"] = make_c2f(8, config_.out_channels[7], true);
|
||||
blocks["model.9"] = std::make_shared<YOLOSPPF>(config_.out_channels[8], config_.out_channels[9]);
|
||||
|
||||
blocks["model.12"] = make_c2f(12, config_.out_channels[9] + config_.out_channels[6], false);
|
||||
blocks["model.15"] = make_c2f(15, config_.out_channels[12] + config_.out_channels[4], false);
|
||||
blocks["model.16"] = std::make_shared<YOLOConv>(config_.out_channels[15], config_.out_channels[16], 3, 2);
|
||||
blocks["model.18"] = make_c2f(18, config_.out_channels[16] + config_.out_channels[12], false);
|
||||
blocks["model.19"] = std::make_shared<YOLOConv>(config_.out_channels[18], config_.out_channels[19], 3, 2);
|
||||
blocks["model.21"] = make_c2f(21, config_.out_channels[19] + config_.out_channels[9], false);
|
||||
blocks["model.22"] = std::make_shared<YOLODetect>(
|
||||
std::array<int, 3>{config_.out_channels[15], config_.out_channels[18], config_.out_channels[21]},
|
||||
config_.detect_box_channels,
|
||||
config_.detect_cls_channels,
|
||||
config_.reg_max,
|
||||
config_.num_classes);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
auto run = [&](int layer, ggml_tensor* input) {
|
||||
return std::dynamic_pointer_cast<UnaryBlock>(blocks["model." + std::to_string(layer)])->forward(ctx, input);
|
||||
};
|
||||
|
||||
auto x0 = run(0, x);
|
||||
auto x1 = run(1, x0);
|
||||
auto x2 = run(2, x1);
|
||||
auto x3 = run(3, x2);
|
||||
auto x4 = run(4, x3);
|
||||
auto x5 = run(5, x4);
|
||||
auto x6 = run(6, x5);
|
||||
auto x7 = run(7, x6);
|
||||
auto x8 = run(8, x7);
|
||||
auto x9 = run(9, x8);
|
||||
|
||||
auto x12 = run(12, ggml_concat(ctx->ggml_ctx, ggml_upscale(ctx->ggml_ctx, x9, 2, GGML_SCALE_MODE_NEAREST), x6, 2));
|
||||
auto x15 = run(15, ggml_concat(ctx->ggml_ctx, ggml_upscale(ctx->ggml_ctx, x12, 2, GGML_SCALE_MODE_NEAREST), x4, 2));
|
||||
auto x16 = run(16, x15);
|
||||
auto x18 = run(18, ggml_concat(ctx->ggml_ctx, x16, x12, 2));
|
||||
auto x19 = run(19, x18);
|
||||
auto x21 = run(21, ggml_concat(ctx->ggml_ctx, x19, x9, 2));
|
||||
|
||||
auto detect = std::dynamic_pointer_cast<YOLODetect>(blocks["model.22"]);
|
||||
auto p3 = detect->forward_scale(ctx, x15, 0);
|
||||
auto p4 = detect->forward_scale(ctx, x18, 1);
|
||||
auto p5 = detect->forward_scale(ctx, x21, 2);
|
||||
p3 = ggml_reshape_2d(ctx->ggml_ctx, p3, p3->ne[0] * p3->ne[1], detect->output_channels());
|
||||
p4 = ggml_reshape_2d(ctx->ggml_ctx, p4, p4->ne[0] * p4->ne[1], detect->output_channels());
|
||||
p5 = ggml_reshape_2d(ctx->ggml_ctx, p5, p5->ne[0] * p5->ne[1], detect->output_channels());
|
||||
return ggml_concat(ctx->ggml_ctx, ggml_concat(ctx->ggml_ctx, p3, p4, 0), p5, 0);
|
||||
}
|
||||
};
|
||||
|
||||
struct YOLOv8Runner : public GGMLRunner {
|
||||
YOLOv8Config config;
|
||||
std::unique_ptr<YOLOv8Model> model;
|
||||
|
||||
YOLOv8Runner(ggml_backend_t backend,
|
||||
const String2TensorStorage& tensor_storage_map,
|
||||
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
|
||||
: GGMLRunner(backend, weight_manager),
|
||||
config(YOLOv8Config::detect_from_weights(tensor_storage_map)) {
|
||||
if (config.valid) {
|
||||
model = std::make_unique<YOLOv8Model>(config);
|
||||
model->init(params_ctx, tensor_storage_map, "");
|
||||
}
|
||||
}
|
||||
|
||||
std::string get_desc() override {
|
||||
return "yolov8";
|
||||
}
|
||||
|
||||
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
|
||||
if (model) {
|
||||
model->get_param_tensors(tensors);
|
||||
}
|
||||
}
|
||||
|
||||
ggml_cgraph* build_graph(const sd::Tensor<float>& input) {
|
||||
if (!model) {
|
||||
return nullptr;
|
||||
}
|
||||
ggml_cgraph* graph = new_graph_custom(1 << 16);
|
||||
auto x = make_input(input);
|
||||
auto runner_ctx = get_context();
|
||||
auto output = model->forward(&runner_ctx, x);
|
||||
ggml_build_forward_expand(graph, output);
|
||||
return graph;
|
||||
}
|
||||
|
||||
sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& input) {
|
||||
auto get_graph = [&]() { return build_graph(input); };
|
||||
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false));
|
||||
}
|
||||
};
|
||||
|
||||
#endif // __SD_MODEL_DETECTOR_YOLOV8_H__
|
||||
@@ -484,10 +484,11 @@ namespace Anima {
|
||||
ggml_tensor* timestep,
|
||||
ggml_tensor* encoder_hidden_states,
|
||||
ggml_tensor* image_pe,
|
||||
ggml_tensor* t5_ids = nullptr,
|
||||
ggml_tensor* t5_weights = nullptr,
|
||||
ggml_tensor* adapter_q_pe = nullptr,
|
||||
ggml_tensor* adapter_k_pe = nullptr) {
|
||||
ggml_tensor* t5_ids = nullptr,
|
||||
ggml_tensor* t5_weights = nullptr,
|
||||
ggml_tensor* adapter_q_pe = nullptr,
|
||||
ggml_tensor* adapter_k_pe = nullptr,
|
||||
std::vector<ggml_tensor*> ref_latents = {}) {
|
||||
GGML_ASSERT(x->ne[3] == 1);
|
||||
|
||||
auto x_embedder = std::dynamic_pointer_cast<XEmbedder>(blocks["x_embedder"]);
|
||||
@@ -502,8 +503,16 @@ namespace Anima {
|
||||
auto padding_mask = ggml_ext_zeros(ctx->ggml_ctx, x->ne[0], x->ne[1], 1, x->ne[3]);
|
||||
x = ggml_concat(ctx->ggml_ctx, x, padding_mask, 2); // [N, C + 1, H, W]
|
||||
|
||||
x = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size); // [N, h*w, (C+1)*ph*pw]
|
||||
|
||||
x = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size); // [N, h*w, (C+1)*ph*pw]
|
||||
int64_t img_len = x->ne[1];
|
||||
if (ref_latents.size() > 0) {
|
||||
for (ggml_tensor* ref : ref_latents) {
|
||||
auto padding_mask = ggml_ext_zeros(ctx->ggml_ctx, ref->ne[0], ref->ne[1], 1, ref->ne[3]);
|
||||
ref = ggml_concat(ctx->ggml_ctx, ref, padding_mask, 2); // [N, C + 1, H, W]
|
||||
ref = DiT::pad_and_patchify(ctx, ref, config.patch_size, config.patch_size);
|
||||
x = ggml_concat(ctx->ggml_ctx, x, ref, 1);
|
||||
}
|
||||
}
|
||||
x = x_embedder->forward(ctx, x);
|
||||
|
||||
auto timestep_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast<int>(config.hidden_size));
|
||||
@@ -543,6 +552,7 @@ namespace Anima {
|
||||
x = block->forward(ctx, x, encoder_hidden_states, embedded_timestep, temb, image_pe);
|
||||
sd::ggml_graph_cut::mark_graph_cut(x, "anima.blocks." + std::to_string(i), "x");
|
||||
}
|
||||
x = ggml_ext_slice(ctx->ggml_ctx, x, 1, 0, img_len);
|
||||
|
||||
x = final_layer->forward(ctx, x, embedded_timestep, temb); // [N, h*w, ph*pw*C]
|
||||
|
||||
@@ -602,8 +612,8 @@ namespace Anima {
|
||||
const std::vector<int>& axes_dim,
|
||||
float h_extrapolation_ratio,
|
||||
float w_extrapolation_ratio,
|
||||
float t_extrapolation_ratio) {
|
||||
static const std::vector<ggml_tensor*> empty_ref_latents;
|
||||
float t_extrapolation_ratio,
|
||||
const std::vector<ggml_tensor*>& ref_latents) {
|
||||
auto ids = Rope::gen_flux_ids(h,
|
||||
w,
|
||||
patch_size,
|
||||
@@ -611,7 +621,7 @@ namespace Anima {
|
||||
static_cast<int>(axes_dim.size()),
|
||||
0,
|
||||
{},
|
||||
empty_ref_latents,
|
||||
ref_latents,
|
||||
Rope::RefIndexMode::FIXED,
|
||||
1.0f,
|
||||
false);
|
||||
@@ -626,14 +636,20 @@ namespace Anima {
|
||||
|
||||
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
|
||||
const sd::Tensor<float>& timesteps_tensor,
|
||||
const sd::Tensor<float>& context_tensor = {},
|
||||
const sd::Tensor<int32_t>& t5_ids_tensor = {},
|
||||
const sd::Tensor<float>& t5_weights_tensor = {}) {
|
||||
const sd::Tensor<float>& context_tensor = {},
|
||||
const sd::Tensor<int32_t>& t5_ids_tensor = {},
|
||||
const sd::Tensor<float>& t5_weights_tensor = {},
|
||||
const std::vector<sd::Tensor<float>>& ref_latents_tensor = {}) {
|
||||
ggml_tensor* x = make_input(x_tensor);
|
||||
ggml_tensor* timesteps = make_input(timesteps_tensor);
|
||||
ggml_tensor* context = make_optional_input(context_tensor);
|
||||
ggml_tensor* t5_ids = make_optional_input(t5_ids_tensor);
|
||||
ggml_tensor* t5_weights = make_optional_input(t5_weights_tensor);
|
||||
std::vector<ggml_tensor*> ref_latents;
|
||||
ref_latents.reserve(ref_latents_tensor.size());
|
||||
for (const auto& ref_latent_tensor : ref_latents_tensor) {
|
||||
ref_latents.push_back(make_input(ref_latent_tensor));
|
||||
}
|
||||
GGML_ASSERT(x->ne[3] == 1);
|
||||
ggml_cgraph* gf = new_graph_custom(ANIMA_GRAPH_SIZE);
|
||||
|
||||
@@ -650,7 +666,8 @@ namespace Anima {
|
||||
config.axes_dim,
|
||||
4.0f,
|
||||
4.0f,
|
||||
1.0f);
|
||||
1.0f,
|
||||
ref_latents);
|
||||
int64_t image_pos_len = static_cast<int64_t>(image_pe_vec.size()) / (2 * 2 * (config.head_dim / 2));
|
||||
auto image_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, image_pos_len);
|
||||
set_backend_tensor_data(image_pe, image_pe_vec.data());
|
||||
@@ -682,7 +699,8 @@ namespace Anima {
|
||||
t5_ids,
|
||||
t5_weights,
|
||||
adapter_q_pe,
|
||||
adapter_k_pe);
|
||||
adapter_k_pe,
|
||||
ref_latents);
|
||||
|
||||
ggml_build_forward_expand(gf, out);
|
||||
return gf;
|
||||
@@ -691,11 +709,13 @@ namespace Anima {
|
||||
sd::Tensor<float> compute(int n_threads,
|
||||
const sd::Tensor<float>& x,
|
||||
const sd::Tensor<float>& timesteps,
|
||||
const sd::Tensor<float>& context = {},
|
||||
const sd::Tensor<int32_t>& t5_ids = {},
|
||||
const sd::Tensor<float>& t5_weights = {}) {
|
||||
const sd::Tensor<float>& context = {},
|
||||
const sd::Tensor<int32_t>& t5_ids = {},
|
||||
const sd::Tensor<float>& t5_weights = {},
|
||||
const std::vector<sd::Tensor<float>>& ref_latents = {},
|
||||
const RefImageParams& ref_image_params = REF_IMAGE_PRESETS.at("cosmos_reference")) {
|
||||
auto get_graph = [&]() -> ggml_cgraph* {
|
||||
return build_graph(x, timesteps, context, t5_ids, t5_weights);
|
||||
return build_graph(x, timesteps, context, t5_ids, t5_weights, ref_latents);
|
||||
};
|
||||
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
|
||||
}
|
||||
@@ -705,12 +725,15 @@ namespace Anima {
|
||||
GGML_ASSERT(diffusion_params.x != nullptr);
|
||||
GGML_ASSERT(diffusion_params.timesteps != nullptr);
|
||||
const auto* extra = diffusion_extra_as<AnimaDiffusionExtra>(diffusion_params);
|
||||
static const std::vector<sd::Tensor<float>> empty_ref_latents;
|
||||
return compute(n_threads,
|
||||
*diffusion_params.x,
|
||||
*diffusion_params.timesteps,
|
||||
tensor_or_empty(diffusion_params.context),
|
||||
tensor_or_empty(extra->t5_ids),
|
||||
tensor_or_empty(extra->t5_weights));
|
||||
tensor_or_empty(extra->t5_weights),
|
||||
diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents,
|
||||
diffusion_params.ref_image_params);
|
||||
}
|
||||
};
|
||||
} // namespace Anima
|
||||
|
||||
@@ -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__
|
||||
@@ -199,7 +199,7 @@ namespace Boogu {
|
||||
auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["linear_2"]);
|
||||
auto linear_3 = std::dynamic_pointer_cast<Linear>(blocks["linear_3"]);
|
||||
|
||||
if (sd_backend_is(ctx->backend, "Vulkan")) {
|
||||
if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
|
||||
linear_2->set_force_prec_f32(true);
|
||||
}
|
||||
|
||||
@@ -259,7 +259,7 @@ namespace Boogu {
|
||||
auto norm_k = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"]);
|
||||
auto to_out_0 = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
|
||||
|
||||
if (sd_backend_is(ctx->backend, "Vulkan")) {
|
||||
if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
|
||||
to_out_0->set_force_prec_f32(true);
|
||||
}
|
||||
|
||||
@@ -383,7 +383,7 @@ namespace Boogu {
|
||||
auto instruct_out = std::dynamic_pointer_cast<Linear>(blocks["processor.instruct_out"]);
|
||||
auto img_out = std::dynamic_pointer_cast<Linear>(blocks["processor.img_out"]);
|
||||
|
||||
if (sd_backend_is(ctx->backend, "Vulkan")) {
|
||||
if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
|
||||
to_out_0->set_force_prec_f32(true);
|
||||
}
|
||||
|
||||
@@ -827,7 +827,7 @@ namespace Boogu {
|
||||
*diffusion_params.x,
|
||||
*diffusion_params.timesteps,
|
||||
tensor_or_empty(diffusion_params.context),
|
||||
diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents);
|
||||
diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents);
|
||||
}
|
||||
};
|
||||
} // namespace Boogu
|
||||
|
||||
@@ -5,7 +5,8 @@
|
||||
#include "model_loader.h"
|
||||
#include "model_manager.h"
|
||||
|
||||
#define CONTROL_NET_GRAPH_SIZE 1536
|
||||
// Match main UNet's MAX_GRAPH_SIZE so SDXL ControlNet (transformer_depth={1,2,10}) fits.
|
||||
#define CONTROL_NET_GRAPH_SIZE MAX_GRAPH_SIZE
|
||||
|
||||
/*
|
||||
=================================== ControlNet ===================================
|
||||
|
||||
@@ -104,14 +104,15 @@ namespace DiT {
|
||||
return x;
|
||||
}
|
||||
|
||||
inline ggml_tensor* patchify(ggml_context* ctx,
|
||||
ggml_tensor* x,
|
||||
int pt,
|
||||
int ph,
|
||||
int pw,
|
||||
int64_t N = 1) {
|
||||
inline ggml_tensor* patchify_3d(ggml_context* ctx,
|
||||
ggml_tensor* x,
|
||||
int pt,
|
||||
int ph,
|
||||
int pw,
|
||||
int64_t N = 1,
|
||||
bool patch_last = true) {
|
||||
// x: [N*C, T, H, W]
|
||||
// return: [N, h*w, C*pt*ph*pw]
|
||||
// return: [N, t_len*h_len*w_len, C*pt*ph*pw] if patch_last else [N, t_len*h_len*w_len, C*pt*ph*pw] or [N, t_len*h_len*w_len, pt*ph*pw*C]
|
||||
int64_t C = x->ne[3] / N;
|
||||
int64_t T = x->ne[2];
|
||||
int64_t H = x->ne[1];
|
||||
@@ -123,15 +124,20 @@ namespace DiT {
|
||||
GGML_ASSERT(C * N == x->ne[3]);
|
||||
GGML_ASSERT(t_len * pt == T && h_len * ph == H && w_len * pw == W);
|
||||
|
||||
x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt, t_len * C * N); // [N*C*t_len, pt, h_len*ph, w_len*pw]
|
||||
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len, h_len*ph, pt, w_len*pw]
|
||||
x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph, h_len * t_len * C * N); // [N*C*t_len*h_len, ph, pt, w_len*pw]
|
||||
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, pt, ph, w_len*pw]
|
||||
x = ggml_reshape_4d(ctx, x, pw, w_len, ph * pt, h_len * t_len * C * N); // [N*C*t_len*h_len, pt*ph, w_len, pw]
|
||||
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, w_len, pt*ph, pw]
|
||||
x = ggml_reshape_4d(ctx, x, pw * ph * pt, w_len * h_len * t_len, C, N); // [N, C, t_len*h_len*w_len, pt*ph*pw]
|
||||
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N, t_len*h_len*w_len, C, pt*ph*pw]
|
||||
x = ggml_reshape_4d(ctx, x, pw * ph * pt * C, w_len * h_len * t_len, N, 1); // [N, t_len*h_len*w_len, C*pt*ph*pw]
|
||||
x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt, t_len * C * N); // [N*C*t_len, pt, h_len*ph, w_len*pw]
|
||||
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len, h_len*ph, pt, w_len*pw]
|
||||
x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph, h_len * t_len * C * N); // [N*C*t_len*h_len, ph, pt, w_len*pw]
|
||||
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, pt, ph, w_len*pw]
|
||||
x = ggml_reshape_4d(ctx, x, pw, w_len, ph * pt, h_len * t_len * C * N); // [N*C*t_len*h_len, pt*ph, w_len, pw]
|
||||
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, w_len, pt*ph, pw]
|
||||
x = ggml_reshape_4d(ctx, x, pw * ph * pt, w_len * h_len * t_len, C, N); // [N, C, t_len*h_len*w_len, pt*ph*pw]
|
||||
if (patch_last) {
|
||||
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N, t_len*h_len*w_len, C, pt*ph*pw]
|
||||
} else {
|
||||
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3)); // [N, t_len*h_len*w_len, pt*ph*pw, C]
|
||||
}
|
||||
// [N, t_len*h_len*w_len, C*pt*ph*pw] or [N, t_len*h_len*w_len, pt*ph*pw*C]
|
||||
x = ggml_reshape_4d(ctx, x, pw * ph * pt * C, w_len * h_len * t_len, N, 1);
|
||||
return x;
|
||||
}
|
||||
|
||||
@@ -142,16 +148,23 @@ namespace DiT {
|
||||
int64_t w_len,
|
||||
int pt,
|
||||
int ph,
|
||||
int pw) {
|
||||
// x: [N, t_len*h_len*w_len, C*pt*ph*pw]
|
||||
int pw,
|
||||
bool patch_last = true) {
|
||||
// x: [N, t_len*h_len*w_len, C*pt*ph*pw] if patch_last else [N, t_len*h_len*w_len, pt*ph*pw*C]
|
||||
// return: [N*C, t_len*pt, h_len*ph, w_len*pw]
|
||||
int64_t N = x->ne[2];
|
||||
int64_t C = x->ne[0] / pt / ph / pw;
|
||||
|
||||
GGML_ASSERT(C * pt * ph * pw == x->ne[0]);
|
||||
|
||||
x = ggml_reshape_4d(ctx, x, pw * ph * pt, C, w_len * h_len * t_len, N); // [N, t_len*h_len*w_len, C, pt*ph*pw]
|
||||
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N, C, t_len*h_len*w_len, pt*ph*pw]
|
||||
if (patch_last) {
|
||||
x = ggml_reshape_4d(ctx, x, pw * ph * pt, C, w_len * h_len * t_len, N); // [N, t_len*h_len*w_len, C, pt*ph*pw]
|
||||
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N, C, t_len*h_len*w_len, pt*ph*pw]
|
||||
} else {
|
||||
x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N); // [N, t_len*h_len*w_len, pt*ph*pw, C]
|
||||
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3)); // [N, C, t_len*h_len*w_len, pt*ph*pw]
|
||||
}
|
||||
|
||||
x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N); // [N*C*t_len*h_len, w_len, pt*ph, pw]
|
||||
x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, pt*ph, w_len, pw]
|
||||
x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N); // [N*C*t_len*h_len, pt, ph, w_len*pw]
|
||||
|
||||
@@ -4,6 +4,7 @@
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
#include "core/util.h"
|
||||
#include "model/adapter/pulid.hpp"
|
||||
#include "model/common/rope.hpp"
|
||||
#include "model/diffusion/dit.hpp"
|
||||
@@ -705,11 +706,13 @@ namespace Flux {
|
||||
LastLayer(int64_t hidden_size,
|
||||
int64_t patch_size,
|
||||
int64_t out_channels,
|
||||
bool prune_mod = false,
|
||||
bool bias = true)
|
||||
bool prune_mod = false,
|
||||
bool bias = true,
|
||||
int64_t patch_volume = 0)
|
||||
: prune_mod(prune_mod) {
|
||||
blocks["norm_final"] = std::shared_ptr<GGMLBlock>(new LayerNorm(hidden_size, 1e-06f, false));
|
||||
blocks["linear"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, patch_size * patch_size * out_channels, bias));
|
||||
int64_t out_dim = (patch_volume > 0 ? patch_volume : patch_size * patch_size) * out_channels;
|
||||
blocks["linear"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, out_dim, bias));
|
||||
if (!prune_mod) {
|
||||
blocks["adaLN_modulation.1"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, 2 * hidden_size, bias));
|
||||
}
|
||||
@@ -1400,18 +1403,28 @@ namespace Flux {
|
||||
std::vector<float> dct_vec;
|
||||
sd::Tensor<float> guidance_tensor;
|
||||
SDVersion version;
|
||||
bool use_mask = false;
|
||||
bool use_mask = true;
|
||||
|
||||
FluxRunner(ggml_backend_t backend,
|
||||
const String2TensorStorage& tensor_storage_map = {},
|
||||
const std::string prefix = "",
|
||||
SDVersion version = VERSION_FLUX,
|
||||
bool use_mask = false,
|
||||
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
|
||||
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
|
||||
const char* model_args = nullptr)
|
||||
: DiffusionModelRunner(backend, prefix, weight_manager),
|
||||
config(FluxConfig::detect_from_weights(tensor_storage_map, prefix, version)),
|
||||
version(version),
|
||||
use_mask(use_mask) {
|
||||
version(version) {
|
||||
for (const auto& [key, value] : parse_key_value_args(model_args, "model arg")) {
|
||||
if (key == "chroma_use_dit_mask") {
|
||||
bool parsed = true;
|
||||
if (parse_strict_bool(value, parsed)) {
|
||||
use_mask = parsed;
|
||||
} else {
|
||||
LOG_WARN("ignoring invalid Chroma DiT model arg '%s=%s'", key.c_str(), value.c_str());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (config.is_chroma) {
|
||||
LOG_INFO("Using pruned modulation (Chroma)");
|
||||
}
|
||||
@@ -1631,8 +1644,8 @@ namespace Flux {
|
||||
tensor_or_empty(diffusion_params.c_concat),
|
||||
tensor_or_empty(diffusion_params.y),
|
||||
tensor_or_empty(extra->guidance),
|
||||
diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
|
||||
diffusion_params.ref_index_mode,
|
||||
diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents,
|
||||
diffusion_params.ref_image_params.ref_index_mode,
|
||||
extra->skip_layers ? *extra->skip_layers : empty_skip_layers,
|
||||
tensor_or_empty(extra->pulid_id),
|
||||
extra->pulid_id_weight);
|
||||
@@ -1718,7 +1731,6 @@ namespace Flux {
|
||||
tensor_storage_map,
|
||||
"model.diffusion_model",
|
||||
VERSION_FLUX2,
|
||||
false,
|
||||
model_manager);
|
||||
|
||||
if (!model_manager->register_runner_params("Flux test",
|
||||
|
||||
@@ -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__
|
||||
@@ -267,7 +267,7 @@ namespace Krea2 {
|
||||
auto knorm = std::dynamic_pointer_cast<KreaRMSNorm>(blocks["qknorm.knorm"]);
|
||||
auto wo = std::dynamic_pointer_cast<Linear>(blocks["wo"]);
|
||||
|
||||
if (sd_backend_is(ctx->backend, "Vulkan")) {
|
||||
if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
|
||||
wo->set_force_prec_f32(true);
|
||||
}
|
||||
|
||||
@@ -421,29 +421,88 @@ namespace Krea2 {
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* vec,
|
||||
ggml_tensor* pe) {
|
||||
ggml_tensor* pe,
|
||||
ggml_tensor* vec_refs = nullptr,
|
||||
int64_t ref_start = -1) {
|
||||
auto mod = std::dynamic_pointer_cast<KreaDoubleSharedModulation>(blocks["mod"]);
|
||||
auto prenorm = std::dynamic_pointer_cast<KreaRMSNorm>(blocks["prenorm"]);
|
||||
auto postnorm = std::dynamic_pointer_cast<KreaRMSNorm>(blocks["postnorm"]);
|
||||
auto attn = std::dynamic_pointer_cast<KreaAttention>(blocks["attn"]);
|
||||
auto mlp = std::dynamic_pointer_cast<KreaSwiGLU>(blocks["mlp"]);
|
||||
|
||||
auto mods = mod->forward(ctx, vec);
|
||||
auto attn_input = Flux::modulate(ctx->ggml_ctx,
|
||||
prenorm->forward(ctx, x),
|
||||
mods[1],
|
||||
mods[0],
|
||||
true);
|
||||
auto attn_out = attn->forward(ctx, attn_input, pe);
|
||||
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out, mods[2]));
|
||||
if (ref_start >= 0 && vec_refs) {
|
||||
// same as normal, but since vec is different for refs and the rest, needs a lot of views and concats
|
||||
auto mods_main = mod->forward(ctx, vec);
|
||||
auto mods_refs = mod->forward(ctx, vec_refs);
|
||||
|
||||
auto mlp_input = Flux::modulate(ctx->ggml_ctx,
|
||||
postnorm->forward(ctx, x),
|
||||
mods[4],
|
||||
mods[3],
|
||||
true);
|
||||
auto mlp_out = mlp->forward(ctx, mlp_input);
|
||||
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, mlp_out, mods[5]));
|
||||
int64_t D = x->ne[0];
|
||||
int64_t N = x->ne[1];
|
||||
int64_t B = x->ne[2];
|
||||
size_t nb1 = x->nb[1];
|
||||
size_t nb2 = x->nb[2];
|
||||
|
||||
int64_t len_main = ref_start;
|
||||
int64_t len_refs = N - ref_start;
|
||||
|
||||
auto pre_x = prenorm->forward(ctx, x);
|
||||
|
||||
auto pre_x_main = ggml_view_3d(ctx->ggml_ctx, pre_x, D, len_main, B, nb1, nb2, 0);
|
||||
auto pre_x_refs = ggml_view_3d(ctx->ggml_ctx, pre_x, D, len_refs, B, nb1, nb2, len_main * nb1);
|
||||
|
||||
auto attn_in_main = Flux::modulate(ctx->ggml_ctx, pre_x_main, mods_main[1], mods_main[0], true);
|
||||
auto attn_in_refs = Flux::modulate(ctx->ggml_ctx, pre_x_refs, mods_refs[1], mods_refs[0], true);
|
||||
|
||||
auto attn_input = ggml_concat(ctx->ggml_ctx, attn_in_main, attn_in_refs, 1);
|
||||
|
||||
auto attn_out = attn->forward(ctx, attn_input, pe);
|
||||
|
||||
auto attn_out_main = ggml_view_3d(ctx->ggml_ctx, attn_out, D, len_main, B, attn_out->nb[1], attn_out->nb[2], 0);
|
||||
auto attn_out_refs = ggml_view_3d(ctx->ggml_ctx, attn_out, D, len_refs, B, attn_out->nb[1], attn_out->nb[2], len_main * attn_out->nb[1]);
|
||||
|
||||
auto res_main = ggml_mul(ctx->ggml_ctx, attn_out_main, mods_main[2]);
|
||||
auto res_refs = ggml_mul(ctx->ggml_ctx, attn_out_refs, mods_refs[2]);
|
||||
|
||||
auto attn_res = ggml_concat(ctx->ggml_ctx, res_main, res_refs, 1);
|
||||
|
||||
x = ggml_add(ctx->ggml_ctx, x, attn_res);
|
||||
|
||||
auto post_x = postnorm->forward(ctx, x);
|
||||
|
||||
auto post_x_main = ggml_view_3d(ctx->ggml_ctx, post_x, D, len_main, B, post_x->nb[1], post_x->nb[2], 0);
|
||||
auto post_x_refs = ggml_view_3d(ctx->ggml_ctx, post_x, D, len_refs, B, post_x->nb[1], post_x->nb[2], len_main * post_x->nb[1]);
|
||||
|
||||
auto mlp_in_main = Flux::modulate(ctx->ggml_ctx, post_x_main, mods_main[4], mods_main[3], true);
|
||||
auto mlp_in_refs = Flux::modulate(ctx->ggml_ctx, post_x_refs, mods_refs[4], mods_refs[3], true);
|
||||
|
||||
auto mlp_input = ggml_concat(ctx->ggml_ctx, mlp_in_main, mlp_in_refs, 1);
|
||||
auto mlp_out = mlp->forward(ctx, mlp_input);
|
||||
|
||||
auto mlp_out_main = ggml_view_3d(ctx->ggml_ctx, mlp_out, D, len_main, B, mlp_out->nb[1], mlp_out->nb[2], 0);
|
||||
auto mlp_out_refs = ggml_view_3d(ctx->ggml_ctx, mlp_out, D, len_refs, B, mlp_out->nb[1], mlp_out->nb[2], len_main * mlp_out->nb[1]);
|
||||
|
||||
auto mlp_res_main = ggml_mul(ctx->ggml_ctx, mlp_out_main, mods_main[5]);
|
||||
auto mlp_res_refs = ggml_mul(ctx->ggml_ctx, mlp_out_refs, mods_refs[5]);
|
||||
|
||||
auto mlp_res = ggml_concat(ctx->ggml_ctx, mlp_res_main, mlp_res_refs, 1);
|
||||
x = ggml_add(ctx->ggml_ctx, x, mlp_res);
|
||||
} else {
|
||||
auto mods = mod->forward(ctx, vec);
|
||||
auto attn_input = Flux::modulate(ctx->ggml_ctx,
|
||||
prenorm->forward(ctx, x),
|
||||
mods[1],
|
||||
mods[0],
|
||||
true);
|
||||
auto attn_out = attn->forward(ctx, attn_input, pe);
|
||||
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out, mods[2]));
|
||||
|
||||
auto mlp_input = Flux::modulate(ctx->ggml_ctx,
|
||||
postnorm->forward(ctx, x),
|
||||
mods[4],
|
||||
mods[3],
|
||||
true);
|
||||
auto mlp_out = mlp->forward(ctx, mlp_input);
|
||||
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, mlp_out, mods[5]));
|
||||
}
|
||||
return x;
|
||||
}
|
||||
|
||||
@@ -555,7 +614,9 @@ namespace Krea2 {
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* timestep,
|
||||
ggml_tensor* context,
|
||||
ggml_tensor* pe) {
|
||||
ggml_tensor* pe,
|
||||
std::vector<ggml_tensor*> ref_latents = {},
|
||||
bool zero_timestep_refs = false) {
|
||||
int64_t W = x->ne[0];
|
||||
int64_t H = x->ne[1];
|
||||
int64_t N = x->ne[3];
|
||||
@@ -570,6 +631,13 @@ namespace Krea2 {
|
||||
|
||||
auto img = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size, true);
|
||||
int64_t img_len = img->ne[1];
|
||||
if (ref_latents.size() > 0) {
|
||||
for (ggml_tensor* ref : ref_latents) {
|
||||
ref = DiT::pad_and_patchify(ctx, ref, config.patch_size, config.patch_size, true);
|
||||
img = ggml_concat(ctx->ggml_ctx, img, ref, 1);
|
||||
}
|
||||
}
|
||||
int64_t ref_len = img->ne[1] - img_len;
|
||||
img = first->forward(ctx, img);
|
||||
|
||||
auto t = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast<int>(config.timestep_dim), 10000, 1000.f);
|
||||
@@ -577,19 +645,30 @@ namespace Krea2 {
|
||||
t = ggml_reshape_3d(ctx->ggml_ctx, t, t->ne[0], 1, t->ne[1]);
|
||||
auto tvec = tproj->forward(ctx, t);
|
||||
|
||||
ggml_tensor* tvec_0 = nullptr;
|
||||
if (ref_latents.size() > 0 && zero_timestep_refs) {
|
||||
// "index_timestep_zero" mode: use timestep = 0 for ref latents
|
||||
auto timestep_0 = ggml_scale(ctx->ggml_ctx, timestep, 0.0f);
|
||||
auto t_0 = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_0, static_cast<int>(config.timestep_dim), 10000, 1000.f);
|
||||
t_0 = tmlp->forward(ctx, t_0);
|
||||
t_0 = ggml_reshape_3d(ctx->ggml_ctx, t_0, t_0->ne[0], 1, t_0->ne[1]);
|
||||
tvec_0 = tproj->forward(ctx, t_0);
|
||||
}
|
||||
|
||||
auto txt = txtfusion->forward(ctx, context);
|
||||
txt = txtmlp->forward(ctx, txt);
|
||||
int64_t txt_len = txt->ne[1];
|
||||
|
||||
auto hidden_states = ggml_concat(ctx->ggml_ctx, txt, img, 1);
|
||||
int64_t ref_start = hidden_states->ne[1] - ref_len;
|
||||
for (int i = 0; i < config.layers; ++i) {
|
||||
auto block = std::dynamic_pointer_cast<KreaSingleStreamBlock>(blocks["blocks." + std::to_string(i)]);
|
||||
hidden_states = block->forward(ctx, hidden_states, tvec, pe);
|
||||
hidden_states = block->forward(ctx, hidden_states, tvec, pe, tvec_0, ref_start);
|
||||
sd::ggml_graph_cut::mark_graph_cut(hidden_states, "krea2.blocks." + std::to_string(i), "hidden_states");
|
||||
}
|
||||
|
||||
hidden_states = last->forward(ctx, hidden_states, t);
|
||||
hidden_states = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, txt_len, txt_len + img_len);
|
||||
hidden_states = last->forward(ctx, hidden_states, t);
|
||||
hidden_states = DiT::unpatchify_and_crop(ctx->ggml_ctx, hidden_states, H, W, config.patch_size, config.patch_size, true);
|
||||
return hidden_states;
|
||||
}
|
||||
@@ -601,10 +680,16 @@ namespace Krea2 {
|
||||
int bs,
|
||||
int context_len,
|
||||
float theta,
|
||||
const std::vector<int>& axes_dim) {
|
||||
const std::vector<int>& axes_dim,
|
||||
const std::vector<ggml_tensor*>& ref_latents,
|
||||
Rope::RefIndexMode ref_index_mode) {
|
||||
auto txt_ids = Rope::gen_flux_txt_ids(bs, context_len, 3, {});
|
||||
auto img_ids = Rope::gen_flux_img_ids(h, w, patch_size, bs, 3, 0, 0, 0, false);
|
||||
auto ids = Rope::concat_ids(txt_ids, img_ids, bs);
|
||||
if (ref_latents.size() > 0) {
|
||||
auto refs_ids = Rope::gen_refs_ids(patch_size, bs, 3, 1, ref_latents, ref_index_mode, 1.0f, false, 0);
|
||||
ids = Rope::concat_ids(ids, refs_ids, bs);
|
||||
}
|
||||
return Rope::embed_nd(ids, bs, theta, axes_dim);
|
||||
}
|
||||
|
||||
@@ -633,7 +718,9 @@ namespace Krea2 {
|
||||
|
||||
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>& context_tensor,
|
||||
const std::vector<sd::Tensor<float>>& ref_latents_tensor = {},
|
||||
const RefImageParams& ref_image_params = REF_IMAGE_PRESETS.at("krea2_ostris_edit")) {
|
||||
ggml_cgraph* gf = new_graph_custom(KREA2_GRAPH_SIZE);
|
||||
ggml_tensor* x = make_input(x_tensor);
|
||||
ggml_tensor* timesteps = make_input(timesteps_tensor);
|
||||
@@ -641,19 +728,27 @@ namespace Krea2 {
|
||||
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));
|
||||
}
|
||||
|
||||
pe_vec = gen_krea2_pe(static_cast<int>(x->ne[1]),
|
||||
static_cast<int>(x->ne[0]),
|
||||
config.patch_size,
|
||||
static_cast<int>(x->ne[3]),
|
||||
static_cast<int>(context->ne[1]),
|
||||
config.theta,
|
||||
config.axes_dim);
|
||||
config.axes_dim,
|
||||
ref_latents,
|
||||
ref_image_params.ref_index_mode);
|
||||
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
|
||||
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
|
||||
set_backend_tensor_data(pe, pe_vec.data());
|
||||
|
||||
auto runner_ctx = get_context();
|
||||
ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pe);
|
||||
ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pe, ref_latents, ref_image_params.force_ref_timestep_zero);
|
||||
ggml_build_forward_expand(gf, out);
|
||||
return gf;
|
||||
}
|
||||
@@ -661,9 +756,11 @@ namespace Krea2 {
|
||||
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>& context,
|
||||
const std::vector<sd::Tensor<float>>& ref_latents = {},
|
||||
const RefImageParams& ref_image_params = REF_IMAGE_PRESETS.at("krea2_ostris_edit")) {
|
||||
auto get_graph = [&]() -> ggml_cgraph* {
|
||||
return build_graph(x, timesteps, context);
|
||||
return build_graph(x, timesteps, context, ref_latents, ref_image_params);
|
||||
};
|
||||
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
|
||||
}
|
||||
@@ -672,10 +769,13 @@ namespace Krea2 {
|
||||
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));
|
||||
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,
|
||||
diffusion_params.ref_image_params);
|
||||
}
|
||||
};
|
||||
} // namespace Krea2
|
||||
|
||||
@@ -0,0 +1,692 @@
|
||||
#ifndef __SD_MODEL_DIFFUSION_LINGBOT_VIDEO_HPP__
|
||||
#define __SD_MODEL_DIFFUSION_LINGBOT_VIDEO_HPP__
|
||||
|
||||
#include <algorithm>
|
||||
#include <cinttypes>
|
||||
#include <cmath>
|
||||
#include <cstdlib>
|
||||
#include <memory>
|
||||
#include <set>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include "core/util.h"
|
||||
#include "model/common/rope.hpp"
|
||||
#include "model/diffusion/dit.hpp"
|
||||
#include "model/diffusion/flux.hpp"
|
||||
#include "model/diffusion/model.hpp"
|
||||
#include "model/diffusion/qwen_image.hpp"
|
||||
|
||||
namespace LingBotVideo {
|
||||
constexpr int LINGBOT_VIDEO_GRAPH_SIZE = 65536;
|
||||
|
||||
struct LingBotVideoConfig {
|
||||
int patch_t = 1;
|
||||
int patch_h = 2;
|
||||
int patch_w = 2;
|
||||
int64_t in_channels = 16;
|
||||
int64_t out_channels = 16;
|
||||
int64_t hidden_size = 2048;
|
||||
int64_t num_attention_heads = 16;
|
||||
int64_t depth = 24;
|
||||
int64_t intermediate_size = 6144;
|
||||
int64_t text_dim = 2560;
|
||||
int64_t freq_dim = 256;
|
||||
float norm_eps = 1e-6f;
|
||||
int rope_theta = 256;
|
||||
std::vector<int> axes_dim = {32, 48, 48};
|
||||
int axes_dim_sum = 128;
|
||||
bool qkv_bias = false;
|
||||
bool out_bias = true;
|
||||
bool patch_embed_bias = true;
|
||||
bool timestep_mlp_bias = true;
|
||||
int64_t num_experts = 0;
|
||||
int64_t num_experts_per_tok = 8;
|
||||
int64_t moe_intermediate_size = 512;
|
||||
int64_t decoder_sparse_step = 1;
|
||||
int64_t n_shared_experts = 0;
|
||||
bool norm_topk_prob = true;
|
||||
float routed_scaling_factor = 1.0f;
|
||||
int64_t n_group = 0;
|
||||
int64_t topk_group = 0;
|
||||
std::set<int> sparse_layers;
|
||||
|
||||
static LingBotVideoConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
|
||||
const std::string& prefix) {
|
||||
LingBotVideoConfig config;
|
||||
config.depth = 0;
|
||||
|
||||
for (const auto& [name, tensor_storage] : tensor_storage_map) {
|
||||
if (!starts_with(name, prefix)) {
|
||||
continue;
|
||||
}
|
||||
|
||||
if (ends_with(name, "patch_embedder.weight") && tensor_storage.n_dims == 2) {
|
||||
int64_t patch_dim = tensor_storage.ne[0];
|
||||
config.hidden_size = tensor_storage.ne[1];
|
||||
int64_t patch_volume = config.patch_t * config.patch_h * config.patch_w;
|
||||
if (patch_dim % patch_volume == 0) {
|
||||
config.in_channels = patch_dim / patch_volume;
|
||||
}
|
||||
} else if (ends_with(name, "text_embedder.linear_1.weight") && tensor_storage.n_dims == 2) {
|
||||
config.text_dim = tensor_storage.ne[0];
|
||||
} else if (ends_with(name, "time_embedder.linear_1.weight") && tensor_storage.n_dims == 2) {
|
||||
config.freq_dim = tensor_storage.ne[0];
|
||||
} else if (ends_with(name, "blocks.0.attn.norm_q.weight") && tensor_storage.n_dims == 1) {
|
||||
int64_t head_dim = tensor_storage.ne[0];
|
||||
if (head_dim > 0) {
|
||||
config.num_attention_heads = config.hidden_size / head_dim;
|
||||
}
|
||||
} else if (name.find(".attn.to_q.bias") != std::string::npos) {
|
||||
config.qkv_bias = true;
|
||||
} else if (name.find(".ffn.gate_proj.weight") != std::string::npos && tensor_storage.n_dims == 2) {
|
||||
config.intermediate_size = tensor_storage.ne[1];
|
||||
} else if (name.find(".ffn.experts.w1") != std::string::npos && tensor_storage.n_dims == 3) {
|
||||
config.num_experts = tensor_storage.ne[2];
|
||||
config.moe_intermediate_size = tensor_storage.ne[1];
|
||||
} else if (name.find(".ffn.shared_experts.gate_proj.weight") != std::string::npos && tensor_storage.n_dims == 2) {
|
||||
if (config.moe_intermediate_size > 0) {
|
||||
config.n_shared_experts = tensor_storage.ne[1] / config.moe_intermediate_size;
|
||||
}
|
||||
} else if (ends_with(name, "proj_out.weight") && tensor_storage.n_dims == 2) {
|
||||
int64_t out_dim = tensor_storage.ne[1];
|
||||
int64_t patch_volume = config.patch_t * config.patch_h * config.patch_w;
|
||||
config.out_channels = patch_volume > 0 ? out_dim / patch_volume : config.out_channels;
|
||||
}
|
||||
|
||||
size_t block_pos = name.find("blocks.");
|
||||
if (block_pos != std::string::npos) {
|
||||
auto items = split_string(name.substr(block_pos), '.');
|
||||
if (items.size() > 1) {
|
||||
int block_index = atoi(items[1].c_str());
|
||||
if (block_index + 1 > config.depth) {
|
||||
config.depth = block_index + 1;
|
||||
}
|
||||
if (name.find("blocks." + std::to_string(block_index) + ".ffn.experts.w1") != std::string::npos) {
|
||||
config.sparse_layers.insert(block_index);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (config.depth == 0) {
|
||||
config.depth = 24;
|
||||
}
|
||||
config.axes_dim_sum = 0;
|
||||
for (int axis_dim : config.axes_dim) {
|
||||
config.axes_dim_sum += axis_dim;
|
||||
}
|
||||
if (!config.sparse_layers.empty()) {
|
||||
config.num_experts = 128;
|
||||
config.num_experts_per_tok = 8;
|
||||
config.moe_intermediate_size = 768;
|
||||
config.decoder_sparse_step = 1;
|
||||
config.n_shared_experts = 1;
|
||||
config.norm_topk_prob = true;
|
||||
config.n_group = 4;
|
||||
config.topk_group = 2;
|
||||
config.routed_scaling_factor = 2.5f;
|
||||
}
|
||||
LOG_DEBUG("lingbot_video: depth = %" PRId64 ", hidden_size = %" PRId64 ", heads = %" PRId64 ", text_dim = %" PRId64 ", experts = %" PRId64 ", experts_per_tok = %" PRId64 ", n_group = %" PRId64 ", topk_group = %" PRId64 ", route_scale = %.2f, sparse_layers = %zu",
|
||||
config.depth,
|
||||
config.hidden_size,
|
||||
config.num_attention_heads,
|
||||
config.text_dim,
|
||||
config.num_experts,
|
||||
config.num_experts_per_tok,
|
||||
config.n_group,
|
||||
config.topk_group,
|
||||
config.routed_scaling_factor,
|
||||
config.sparse_layers.size());
|
||||
return config;
|
||||
}
|
||||
};
|
||||
|
||||
struct LingBotVideoTextEmbedder : public GGMLBlock {
|
||||
LingBotVideoTextEmbedder(int64_t text_dim,
|
||||
int64_t hidden_size,
|
||||
float eps = 1e-6f) {
|
||||
blocks["norm"] = std::make_shared<RMSNorm>(text_dim, eps);
|
||||
blocks["linear_1"] = std::make_shared<Linear>(text_dim, hidden_size, true);
|
||||
blocks["linear_2"] = std::make_shared<Linear>(hidden_size, hidden_size, true);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
auto norm = std::dynamic_pointer_cast<RMSNorm>(blocks["norm"]);
|
||||
auto linear_1 = std::dynamic_pointer_cast<Linear>(blocks["linear_1"]);
|
||||
auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["linear_2"]);
|
||||
|
||||
x = norm->forward(ctx, x);
|
||||
x = linear_1->forward(ctx, x);
|
||||
x = ggml_silu_inplace(ctx->ggml_ctx, x);
|
||||
x = linear_2->forward(ctx, x);
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
struct LingBotVideoAttention : public GGMLBlock {
|
||||
int64_t num_heads;
|
||||
int64_t head_dim;
|
||||
|
||||
LingBotVideoAttention(int64_t hidden_size,
|
||||
int64_t num_heads,
|
||||
bool qkv_bias = false,
|
||||
bool out_bias = true,
|
||||
float eps = 1e-6f)
|
||||
: num_heads(num_heads),
|
||||
head_dim(hidden_size / num_heads) {
|
||||
int64_t inner_dim = num_heads * head_dim;
|
||||
blocks["to_q"] = std::make_shared<Linear>(hidden_size, inner_dim, qkv_bias);
|
||||
blocks["to_k"] = std::make_shared<Linear>(hidden_size, inner_dim, qkv_bias);
|
||||
blocks["to_v"] = std::make_shared<Linear>(hidden_size, inner_dim, qkv_bias);
|
||||
blocks["norm_q"] = std::make_shared<RMSNorm>(head_dim, eps);
|
||||
blocks["norm_k"] = std::make_shared<RMSNorm>(head_dim, eps);
|
||||
blocks["to_out"] = std::make_shared<Linear>(inner_dim, hidden_size, out_bias);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* pe,
|
||||
ggml_tensor* attention_mask = nullptr) {
|
||||
// x: [N, video_tokens + text_tokens, hidden_size]
|
||||
auto to_q = std::dynamic_pointer_cast<Linear>(blocks["to_q"]);
|
||||
auto to_k = std::dynamic_pointer_cast<Linear>(blocks["to_k"]);
|
||||
auto to_v = std::dynamic_pointer_cast<Linear>(blocks["to_v"]);
|
||||
auto norm_q = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_q"]);
|
||||
auto norm_k = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"]);
|
||||
auto to_out = std::dynamic_pointer_cast<Linear>(blocks["to_out"]);
|
||||
|
||||
int64_t S = x->ne[1];
|
||||
int64_t N = x->ne[2];
|
||||
|
||||
auto q = to_q->forward(ctx, x);
|
||||
auto k = to_k->forward(ctx, x);
|
||||
auto v = to_v->forward(ctx, x);
|
||||
|
||||
q = ggml_reshape_4d(ctx->ggml_ctx, q, head_dim, num_heads, S, N);
|
||||
k = ggml_reshape_4d(ctx->ggml_ctx, k, head_dim, num_heads, S, N);
|
||||
v = ggml_reshape_4d(ctx->ggml_ctx, v, head_dim, num_heads, S, N);
|
||||
|
||||
q = norm_q->forward(ctx, q);
|
||||
k = norm_k->forward(ctx, k);
|
||||
|
||||
x = Rope::attention(ctx, q, k, v, pe, attention_mask);
|
||||
x = to_out->forward(ctx, x);
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
struct LingBotVideoMLP : public UnaryBlock {
|
||||
LingBotVideoMLP(int64_t hidden_size,
|
||||
int64_t intermediate_size) {
|
||||
blocks["gate_proj"] = std::make_shared<Linear>(hidden_size, intermediate_size, false);
|
||||
blocks["up_proj"] = std::make_shared<Linear>(hidden_size, intermediate_size, false);
|
||||
blocks["down_proj"] = std::make_shared<Linear>(intermediate_size, hidden_size, false);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
auto gate_proj = std::dynamic_pointer_cast<Linear>(blocks["gate_proj"]);
|
||||
auto up_proj = std::dynamic_pointer_cast<Linear>(blocks["up_proj"]);
|
||||
auto down_proj = std::dynamic_pointer_cast<Linear>(blocks["down_proj"]);
|
||||
|
||||
auto gate = gate_proj->forward(ctx, x);
|
||||
gate = ggml_silu_inplace(ctx->ggml_ctx, gate);
|
||||
auto up = up_proj->forward(ctx, x);
|
||||
x = ggml_mul(ctx->ggml_ctx, gate, up);
|
||||
x = down_proj->forward(ctx, x);
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
struct LingBotVideoSparseMoeBlock : public UnaryBlock {
|
||||
int64_t hidden_size;
|
||||
int64_t intermediate_size;
|
||||
int64_t num_experts;
|
||||
int64_t num_experts_per_tok;
|
||||
bool has_shared_experts;
|
||||
bool norm_topk_prob;
|
||||
float routed_scaling_factor;
|
||||
int64_t n_group;
|
||||
int64_t topk_group;
|
||||
std::vector<float> group_expert_mask_vec;
|
||||
bool has_correction_bias = false;
|
||||
|
||||
LingBotVideoSparseMoeBlock(const LingBotVideoConfig& config)
|
||||
: hidden_size(config.hidden_size),
|
||||
intermediate_size(config.moe_intermediate_size),
|
||||
num_experts(config.num_experts),
|
||||
num_experts_per_tok(config.num_experts_per_tok),
|
||||
has_shared_experts(config.n_shared_experts > 0),
|
||||
norm_topk_prob(config.norm_topk_prob),
|
||||
routed_scaling_factor(config.routed_scaling_factor),
|
||||
n_group(config.n_group),
|
||||
topk_group(config.topk_group) {
|
||||
if (n_group > 1) {
|
||||
GGML_ASSERT(num_experts % n_group == 0);
|
||||
int64_t experts_per_group = num_experts / n_group;
|
||||
group_expert_mask_vec.assign(static_cast<size_t>(num_experts * n_group), 0.f);
|
||||
for (int64_t group = 0; group < n_group; ++group) {
|
||||
int64_t expert_begin = group * experts_per_group;
|
||||
int64_t expert_end = expert_begin + experts_per_group;
|
||||
for (int64_t expert = expert_begin; expert < expert_end; ++expert) {
|
||||
group_expert_mask_vec[static_cast<size_t>(group * num_experts + expert)] = 1.f;
|
||||
}
|
||||
}
|
||||
}
|
||||
if (has_shared_experts) {
|
||||
blocks["shared_experts"] = std::make_shared<LingBotVideoMLP>(hidden_size,
|
||||
intermediate_size * config.n_shared_experts);
|
||||
}
|
||||
}
|
||||
|
||||
void init_params(ggml_context* ctx,
|
||||
const String2TensorStorage& tensor_storage_map = {},
|
||||
const std::string prefix = "") override {
|
||||
ggml_type router_type = get_type(prefix + "router.weight", tensor_storage_map, GGML_TYPE_F32);
|
||||
ggml_type w1_type = get_type(prefix + "experts.w1", tensor_storage_map, GGML_TYPE_F32);
|
||||
ggml_type w2_type = get_type(prefix + "experts.w2", tensor_storage_map, GGML_TYPE_F32);
|
||||
ggml_type w3_type = get_type(prefix + "experts.w3", tensor_storage_map, GGML_TYPE_F32);
|
||||
params["router.weight"] = ggml_new_tensor_2d(ctx, router_type, hidden_size, num_experts);
|
||||
if (tensor_storage_map.find(prefix + "router.e_score_correction_bias") != tensor_storage_map.end()) {
|
||||
params["router.e_score_correction_bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_experts);
|
||||
has_correction_bias = true;
|
||||
}
|
||||
params["experts.w1"] = ggml_new_tensor_3d(ctx, w1_type, hidden_size, intermediate_size, num_experts);
|
||||
params["experts.w2"] = ggml_new_tensor_3d(ctx, w2_type, intermediate_size, hidden_size, num_experts);
|
||||
params["experts.w3"] = ggml_new_tensor_3d(ctx, w3_type, hidden_size, intermediate_size, num_experts);
|
||||
}
|
||||
|
||||
ggml_tensor* expert_linear(GGMLRunnerContext* ctx,
|
||||
const std::string& weight_name,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* selected_experts) {
|
||||
return ggml_mul_mat_id(ctx->ggml_ctx, params[weight_name], x, selected_experts);
|
||||
}
|
||||
|
||||
ggml_tensor* select_experts(GGMLRunnerContext* ctx, ggml_tensor* choice_scores) {
|
||||
ggml_context* gctx = ctx->ggml_ctx;
|
||||
if (n_group <= 1 || topk_group <= 0) {
|
||||
return ggml_argsort_top_k(gctx, choice_scores, static_cast<int>(num_experts_per_tok));
|
||||
}
|
||||
|
||||
GGML_ASSERT(choice_scores->ne[0] == num_experts);
|
||||
GGML_ASSERT(num_experts % n_group == 0);
|
||||
GGML_ASSERT(topk_group > 0 && topk_group <= n_group);
|
||||
GGML_ASSERT(!group_expert_mask_vec.empty());
|
||||
|
||||
const int64_t n_token_total = choice_scores->ne[1];
|
||||
const int64_t experts_per_group = num_experts / n_group;
|
||||
const int group_score_k = 2;
|
||||
GGML_ASSERT(experts_per_group >= group_score_k);
|
||||
|
||||
ggml_tensor* grouped_scores = ggml_reshape_3d(gctx, choice_scores, experts_per_group, n_group, n_token_total);
|
||||
ggml_tensor* group_top_ids = ggml_argsort_top_k(gctx, grouped_scores, group_score_k);
|
||||
grouped_scores = ggml_reshape_3d(gctx, grouped_scores, 1, experts_per_group, n_group * n_token_total);
|
||||
group_top_ids = ggml_cont(gctx, group_top_ids);
|
||||
group_top_ids = ggml_reshape_2d(gctx, group_top_ids, group_score_k, n_group * n_token_total);
|
||||
|
||||
ggml_tensor* group_top_values = ggml_get_rows(gctx, grouped_scores, group_top_ids);
|
||||
group_top_values = ggml_reshape_3d(gctx, group_top_values, group_score_k, n_group, n_token_total);
|
||||
|
||||
ggml_tensor* group_scores = nullptr;
|
||||
for (int rank = 0; rank < group_score_k; ++rank) {
|
||||
ggml_tensor* value = ggml_view_3d(gctx,
|
||||
group_top_values,
|
||||
1,
|
||||
n_group,
|
||||
n_token_total,
|
||||
group_top_values->nb[1],
|
||||
group_top_values->nb[2],
|
||||
rank * group_top_values->nb[0]);
|
||||
group_scores = group_scores == nullptr ? value : ggml_add(gctx, group_scores, value);
|
||||
}
|
||||
group_scores = ggml_reshape_2d(gctx, group_scores, n_group, n_token_total);
|
||||
|
||||
ggml_tensor* selected_groups = ggml_argsort_top_k(gctx, group_scores, static_cast<int>(topk_group));
|
||||
selected_groups = ggml_cont(gctx, selected_groups);
|
||||
|
||||
ggml_tensor* group_expert_mask = ggml_new_tensor_3d(gctx, GGML_TYPE_F32, num_experts, n_group, 1);
|
||||
ctx->bind_backend_tensor_data(group_expert_mask, group_expert_mask_vec.data());
|
||||
ggml_tensor* group_expert_mask_template = ggml_new_tensor_3d(gctx, GGML_TYPE_F32, num_experts, n_group, n_token_total);
|
||||
group_expert_mask = ggml_repeat(gctx, group_expert_mask, group_expert_mask_template);
|
||||
|
||||
ggml_tensor* selected_group_masks = ggml_get_rows(gctx, group_expert_mask, selected_groups);
|
||||
ggml_tensor* selected_mask = nullptr;
|
||||
for (int64_t rank = 0; rank < topk_group; ++rank) {
|
||||
ggml_tensor* mask = ggml_view_3d(gctx,
|
||||
selected_group_masks,
|
||||
num_experts,
|
||||
1,
|
||||
n_token_total,
|
||||
selected_group_masks->nb[1],
|
||||
selected_group_masks->nb[2],
|
||||
rank * selected_group_masks->nb[1]);
|
||||
selected_mask = selected_mask == nullptr ? mask : ggml_add(gctx, selected_mask, mask);
|
||||
}
|
||||
selected_mask = ggml_reshape_2d(gctx, selected_mask, num_experts, n_token_total);
|
||||
|
||||
ggml_tensor* excluded_group_mask = ggml_sub(gctx, selected_mask, ggml_ext_ones_like(gctx, selected_mask));
|
||||
ggml_tensor* masked_scores = ggml_add(gctx, choice_scores, ggml_scale(gctx, excluded_group_mask, 1.0e9f));
|
||||
return ggml_argsort_top_k(gctx, masked_scores, static_cast<int>(num_experts_per_tok));
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
// x: [N, tokens, hidden_size]
|
||||
GGML_ASSERT(num_experts > 0);
|
||||
GGML_ASSERT(num_experts_per_tok > 0 && num_experts_per_tok <= num_experts);
|
||||
|
||||
ggml_tensor* residual = x;
|
||||
const int64_t n_token = x->ne[1];
|
||||
const int64_t N = x->ne[2];
|
||||
const int64_t n_token_total = n_token * N;
|
||||
|
||||
ggml_tensor* router_logits = ggml_mul_mat(ctx->ggml_ctx, params["router.weight"], x);
|
||||
router_logits = ggml_reshape_2d(ctx->ggml_ctx, router_logits, num_experts, n_token_total);
|
||||
ggml_tensor* probs = ggml_sigmoid(ctx->ggml_ctx, router_logits);
|
||||
ggml_tensor* choice_scores = probs;
|
||||
if (has_correction_bias) {
|
||||
choice_scores = ggml_add(ctx->ggml_ctx, choice_scores, params["router.e_score_correction_bias"]);
|
||||
}
|
||||
|
||||
ggml_tensor* selected_experts = select_experts(ctx, choice_scores);
|
||||
ggml_tensor* weights = ggml_get_rows(ctx->ggml_ctx,
|
||||
ggml_reshape_3d(ctx->ggml_ctx, probs, 1, num_experts, n_token_total),
|
||||
selected_experts);
|
||||
weights = ggml_reshape_2d(ctx->ggml_ctx, weights, num_experts_per_tok, n_token_total);
|
||||
if (norm_topk_prob && num_experts_per_tok > 1) {
|
||||
auto weights_sum = ggml_sum_rows(ctx->ggml_ctx, weights);
|
||||
weights_sum = ggml_clamp(ctx->ggml_ctx, weights_sum, 6.103515625e-5f, INFINITY);
|
||||
weights = ggml_div(ctx->ggml_ctx, weights, weights_sum);
|
||||
}
|
||||
if (routed_scaling_factor != 1.0f) {
|
||||
weights = ggml_scale(ctx->ggml_ctx, weights, routed_scaling_factor);
|
||||
}
|
||||
weights = ggml_reshape_3d(ctx->ggml_ctx, weights, 1, num_experts_per_tok, n_token_total);
|
||||
|
||||
x = ggml_reshape_3d(ctx->ggml_ctx, x, hidden_size, 1, n_token_total);
|
||||
auto gate = expert_linear(ctx, "experts.w1", x, selected_experts);
|
||||
gate = ggml_silu_inplace(ctx->ggml_ctx, gate);
|
||||
auto up = expert_linear(ctx, "experts.w3", x, selected_experts);
|
||||
auto act = ggml_mul(ctx->ggml_ctx, gate, up);
|
||||
auto out = expert_linear(ctx, "experts.w2", act, selected_experts);
|
||||
out = ggml_mul(ctx->ggml_ctx, out, weights);
|
||||
|
||||
ggml_tensor* summed = nullptr;
|
||||
for (int64_t i = 0; i < num_experts_per_tok; ++i) {
|
||||
auto expert_out = ggml_view_2d(ctx->ggml_ctx,
|
||||
out,
|
||||
hidden_size,
|
||||
n_token_total,
|
||||
out->nb[2],
|
||||
i * out->nb[1]);
|
||||
summed = summed == nullptr ? expert_out : ggml_add(ctx->ggml_ctx, summed, expert_out);
|
||||
}
|
||||
if (num_experts_per_tok == 1) {
|
||||
summed = ggml_cont(ctx->ggml_ctx, summed);
|
||||
}
|
||||
summed = ggml_reshape_3d(ctx->ggml_ctx, summed, hidden_size, n_token, N);
|
||||
|
||||
if (has_shared_experts) {
|
||||
auto shared_experts = std::dynamic_pointer_cast<LingBotVideoMLP>(blocks["shared_experts"]);
|
||||
summed = ggml_add(ctx->ggml_ctx, summed, shared_experts->forward(ctx, residual));
|
||||
}
|
||||
return summed;
|
||||
}
|
||||
};
|
||||
|
||||
struct LingBotVideoBlock : public GGMLBlock {
|
||||
int64_t hidden_size;
|
||||
|
||||
LingBotVideoBlock(const LingBotVideoConfig& config,
|
||||
bool sparse)
|
||||
: hidden_size(config.hidden_size) {
|
||||
blocks["norm1"] = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps);
|
||||
blocks["attn"] = std::make_shared<LingBotVideoAttention>(config.hidden_size,
|
||||
config.num_attention_heads,
|
||||
config.qkv_bias,
|
||||
config.out_bias,
|
||||
config.norm_eps);
|
||||
blocks["norm_post_attn"] = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps);
|
||||
blocks["norm2"] = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps);
|
||||
if (sparse) {
|
||||
blocks["ffn"] = std::make_shared<LingBotVideoSparseMoeBlock>(config);
|
||||
} else {
|
||||
blocks["ffn"] = std::make_shared<LingBotVideoMLP>(config.hidden_size, config.intermediate_size);
|
||||
}
|
||||
blocks["norm_post_ffn"] = std::make_shared<RMSNorm>(config.hidden_size, config.norm_eps);
|
||||
}
|
||||
|
||||
void init_params(ggml_context* ctx,
|
||||
const String2TensorStorage& tensor_storage_map = {},
|
||||
const std::string prefix = "") override {
|
||||
params["scale_shift_table"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size * 6, 1);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* temb6,
|
||||
ggml_tensor* pe,
|
||||
ggml_tensor* attention_mask = nullptr) {
|
||||
// x: [N, tokens, hidden_size], temb6: [N, tokens, 6 * hidden_size]
|
||||
auto norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm1"]);
|
||||
auto attn = std::dynamic_pointer_cast<LingBotVideoAttention>(blocks["attn"]);
|
||||
auto norm_post_attn = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_post_attn"]);
|
||||
auto norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["norm2"]);
|
||||
auto ffn = std::dynamic_pointer_cast<UnaryBlock>(blocks["ffn"]);
|
||||
auto norm_post_ffn = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_post_ffn"]);
|
||||
|
||||
auto table = ggml_reshape_3d(ctx->ggml_ctx, params["scale_shift_table"], hidden_size * 6, 1, 1);
|
||||
auto mod = ggml_add(ctx->ggml_ctx, temb6, table);
|
||||
auto mods = ggml_ext_chunk(ctx->ggml_ctx, mod, 6, 0);
|
||||
|
||||
auto shift_msa = mods[0];
|
||||
auto scale_msa = mods[1];
|
||||
auto gate_msa = ggml_tanh(ctx->ggml_ctx, mods[2]);
|
||||
auto shift_mlp = mods[3];
|
||||
auto scale_mlp = mods[4];
|
||||
auto gate_mlp = ggml_tanh(ctx->ggml_ctx, mods[5]);
|
||||
|
||||
auto attn_in = Flux::modulate(ctx->ggml_ctx, norm1->forward(ctx, x), shift_msa, scale_msa, true);
|
||||
auto attn_out = attn->forward(ctx, attn_in, pe, attention_mask);
|
||||
attn_out = norm_post_attn->forward(ctx, attn_out);
|
||||
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, gate_msa, attn_out));
|
||||
|
||||
auto ffn_in = Flux::modulate(ctx->ggml_ctx, norm2->forward(ctx, x), shift_mlp, scale_mlp, true);
|
||||
auto ffn_out = ffn->forward(ctx, ffn_in);
|
||||
ffn_out = norm_post_ffn->forward(ctx, ffn_out);
|
||||
x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, gate_mlp, ffn_out));
|
||||
return x;
|
||||
}
|
||||
};
|
||||
|
||||
class LingBotVideoModel : public GGMLBlock {
|
||||
public:
|
||||
LingBotVideoConfig config;
|
||||
|
||||
LingBotVideoModel() = default;
|
||||
LingBotVideoModel(LingBotVideoConfig config)
|
||||
: config(config) {
|
||||
int64_t patch_dim = config.patch_t * config.patch_h * config.patch_w * config.in_channels;
|
||||
blocks["patch_embedder"] = std::make_shared<Linear>(patch_dim, config.hidden_size, config.patch_embed_bias);
|
||||
blocks["time_embedder"] = std::make_shared<Qwen::TimestepEmbedding>(config.freq_dim,
|
||||
config.hidden_size,
|
||||
config.hidden_size,
|
||||
0,
|
||||
config.timestep_mlp_bias);
|
||||
blocks["time_modulation.1"] = std::make_shared<Linear>(config.hidden_size, 6 * config.hidden_size, true);
|
||||
blocks["text_embedder"] = std::make_shared<LingBotVideoTextEmbedder>(config.text_dim,
|
||||
config.hidden_size,
|
||||
config.norm_eps);
|
||||
for (int i = 0; i < config.depth; i++) {
|
||||
bool sparse = config.sparse_layers.find(i) != config.sparse_layers.end();
|
||||
blocks["blocks." + std::to_string(i)] = std::make_shared<LingBotVideoBlock>(config, sparse);
|
||||
}
|
||||
blocks["norm_out"] = std::make_shared<LayerNorm>(config.hidden_size, config.norm_eps, false);
|
||||
blocks["norm_out_modulation.1"] = std::make_shared<Linear>(config.hidden_size, 2 * config.hidden_size, true);
|
||||
blocks["proj_out"] = std::make_shared<Linear>(config.hidden_size,
|
||||
config.patch_t * config.patch_h * config.patch_w * config.out_channels,
|
||||
true);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* timestep,
|
||||
ggml_tensor* context,
|
||||
ggml_tensor* pe) {
|
||||
// x: [N*C, T, H, W], context: [N, text_tokens, text_dim]
|
||||
GGML_ASSERT(context != nullptr);
|
||||
GGML_ASSERT(x->ne[3] == config.in_channels);
|
||||
GGML_ASSERT(x->ne[2] % config.patch_t == 0);
|
||||
GGML_ASSERT(x->ne[1] % config.patch_h == 0);
|
||||
GGML_ASSERT(x->ne[0] % config.patch_w == 0);
|
||||
|
||||
auto patch_embedder = std::dynamic_pointer_cast<Linear>(blocks["patch_embedder"]);
|
||||
auto time_embedder = std::dynamic_pointer_cast<Qwen::TimestepEmbedding>(blocks["time_embedder"]);
|
||||
auto time_modulation = std::dynamic_pointer_cast<Linear>(blocks["time_modulation.1"]);
|
||||
auto text_embedder = std::dynamic_pointer_cast<LingBotVideoTextEmbedder>(blocks["text_embedder"]);
|
||||
auto norm_out = std::dynamic_pointer_cast<LayerNorm>(blocks["norm_out"]);
|
||||
auto norm_out_modulation = std::dynamic_pointer_cast<Linear>(blocks["norm_out_modulation.1"]);
|
||||
auto proj_out = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
|
||||
|
||||
int64_t W = x->ne[0];
|
||||
int64_t H = x->ne[1];
|
||||
int64_t T = x->ne[2];
|
||||
int64_t N = 1;
|
||||
int64_t t_len = T / config.patch_t;
|
||||
int64_t h_len = H / config.patch_h;
|
||||
int64_t w_len = W / config.patch_w;
|
||||
int64_t n_img = t_len * h_len * w_len;
|
||||
|
||||
auto img = DiT::patchify_3d(ctx->ggml_ctx, x, config.patch_t, config.patch_h, config.patch_w, N, false);
|
||||
img = patch_embedder->forward(ctx, img);
|
||||
|
||||
auto txt = text_embedder->forward(ctx, context);
|
||||
auto hidden_states = ggml_concat(ctx->ggml_ctx, img, txt, 1);
|
||||
|
||||
auto timestep_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx,
|
||||
timestep,
|
||||
static_cast<int>(config.freq_dim),
|
||||
10000,
|
||||
1.0f);
|
||||
auto t_emb = time_embedder->forward(ctx, timestep_proj);
|
||||
auto temb_template = ggml_new_tensor_3d(ctx->ggml_ctx, t_emb->type, t_emb->ne[0], hidden_states->ne[1], t_emb->ne[1]);
|
||||
auto temb_tokens = ggml_repeat(ctx->ggml_ctx,
|
||||
ggml_reshape_3d(ctx->ggml_ctx, t_emb, t_emb->ne[0], 1, t_emb->ne[1]),
|
||||
temb_template);
|
||||
auto temb6 = time_modulation->forward(ctx, ggml_silu(ctx->ggml_ctx, temb_tokens));
|
||||
|
||||
sd::ggml_graph_cut::mark_graph_cut(hidden_states, "lingbot_video.prelude", "hidden_states");
|
||||
|
||||
for (int i = 0; i < config.depth; i++) {
|
||||
auto block = std::dynamic_pointer_cast<LingBotVideoBlock>(blocks["blocks." + std::to_string(i)]);
|
||||
hidden_states = block->forward(ctx, hidden_states, temb6, pe);
|
||||
sd::ggml_graph_cut::mark_graph_cut(hidden_states, "lingbot_video.blocks." + std::to_string(i), "hidden_states");
|
||||
}
|
||||
|
||||
auto final_mods = ggml_ext_chunk(ctx->ggml_ctx,
|
||||
norm_out_modulation->forward(ctx, ggml_silu(ctx->ggml_ctx, temb_tokens)),
|
||||
2,
|
||||
0);
|
||||
hidden_states = norm_out->forward(ctx, hidden_states);
|
||||
hidden_states = Flux::modulate(ctx->ggml_ctx, hidden_states, final_mods[0], final_mods[1], true);
|
||||
hidden_states = proj_out->forward(ctx, hidden_states);
|
||||
hidden_states = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, 0, n_img);
|
||||
|
||||
auto out = DiT::unpatchify_3d(ctx->ggml_ctx,
|
||||
hidden_states,
|
||||
t_len,
|
||||
h_len,
|
||||
w_len,
|
||||
config.patch_t,
|
||||
config.patch_h,
|
||||
config.patch_w,
|
||||
false);
|
||||
return out;
|
||||
}
|
||||
};
|
||||
|
||||
struct LingBotVideoRunner : public DiffusionModelRunner {
|
||||
LingBotVideoConfig config;
|
||||
LingBotVideoModel lingbot_video;
|
||||
std::vector<float> pe_vec;
|
||||
|
||||
LingBotVideoRunner(ggml_backend_t backend,
|
||||
const String2TensorStorage& tensor_storage_map = {},
|
||||
const std::string prefix = "",
|
||||
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
|
||||
const char* model_args = nullptr)
|
||||
: DiffusionModelRunner(backend, prefix, weight_manager),
|
||||
config(LingBotVideoConfig::detect_from_weights(tensor_storage_map, prefix)) {
|
||||
SD_UNUSED(model_args);
|
||||
|
||||
lingbot_video = LingBotVideoModel(config);
|
||||
lingbot_video.init(params_ctx, tensor_storage_map, prefix);
|
||||
}
|
||||
|
||||
std::string get_desc() override {
|
||||
return "lingbot_video";
|
||||
}
|
||||
|
||||
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
|
||||
lingbot_video.get_param_tensors(tensors, prefix);
|
||||
}
|
||||
|
||||
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
|
||||
const sd::Tensor<float>& timesteps_tensor,
|
||||
const sd::Tensor<float>& context_tensor) {
|
||||
ggml_cgraph* gf = new_graph_custom(LINGBOT_VIDEO_GRAPH_SIZE);
|
||||
ggml_tensor* x = make_input(x_tensor);
|
||||
ggml_tensor* timesteps = make_input(timesteps_tensor);
|
||||
GGML_ASSERT(x_tensor.dim() == 5);
|
||||
GGML_ASSERT(x->ne[3] == config.in_channels);
|
||||
GGML_ASSERT(!context_tensor.empty());
|
||||
ggml_tensor* context = make_input(context_tensor);
|
||||
|
||||
int64_t W = x->ne[0];
|
||||
int64_t H = x->ne[1];
|
||||
int64_t T = x->ne[2];
|
||||
int64_t N = x_tensor.shape()[4];
|
||||
GGML_ASSERT(N == 1);
|
||||
pe_vec = Rope::gen_lingbot_video_pe(static_cast<int>(T),
|
||||
static_cast<int>(H),
|
||||
static_cast<int>(W),
|
||||
config.patch_t,
|
||||
config.patch_h,
|
||||
config.patch_w,
|
||||
static_cast<int>(N),
|
||||
static_cast<int>(context->ne[1]),
|
||||
config.rope_theta,
|
||||
config.axes_dim);
|
||||
int pos_len = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
|
||||
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
|
||||
set_backend_tensor_data(pe, pe_vec.data());
|
||||
|
||||
auto runner_ctx = get_context();
|
||||
ggml_tensor* out = lingbot_video.forward(&runner_ctx, x, timesteps, context, pe);
|
||||
ggml_build_forward_expand(gf, out);
|
||||
return gf;
|
||||
}
|
||||
|
||||
sd::Tensor<float> compute(int n_threads,
|
||||
const sd::Tensor<float>& x,
|
||||
const sd::Tensor<float>& timesteps,
|
||||
const sd::Tensor<float>& context) {
|
||||
auto get_graph = [&]() -> ggml_cgraph* {
|
||||
return build_graph(x, timesteps, context);
|
||||
};
|
||||
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
|
||||
}
|
||||
|
||||
sd::Tensor<float> compute(int n_threads,
|
||||
const DiffusionParams& diffusion_params) override {
|
||||
GGML_ASSERT(diffusion_params.x != nullptr);
|
||||
GGML_ASSERT(diffusion_params.timesteps != nullptr);
|
||||
return compute(n_threads,
|
||||
*diffusion_params.x,
|
||||
*diffusion_params.timesteps,
|
||||
tensor_or_empty(diffusion_params.context));
|
||||
}
|
||||
};
|
||||
} // namespace LingBotVideo
|
||||
|
||||
#endif // __SD_MODEL_DIFFUSION_LINGBOT_VIDEO_HPP__
|
||||
@@ -800,7 +800,7 @@ namespace LTXV {
|
||||
auto gate_mlp = mods[5];
|
||||
|
||||
auto x_norm = rms_norm(ctx->ggml_ctx, x);
|
||||
x_norm = modulate(ctx->ggml_ctx, x_norm, shift_msa, scale_msa);
|
||||
x_norm = LTXV::modulate(ctx->ggml_ctx, x_norm, shift_msa, scale_msa);
|
||||
auto msa = attn1->forward(ctx, x_norm, nullptr, self_attention_mask, pe);
|
||||
x = ggml_add(ctx->ggml_ctx, x, apply_gate(ctx->ggml_ctx, msa, gate_msa));
|
||||
|
||||
@@ -810,12 +810,12 @@ namespace LTXV {
|
||||
auto gate_q = mods[8];
|
||||
|
||||
auto q = rms_norm(ctx->ggml_ctx, x);
|
||||
q = modulate(ctx->ggml_ctx, q, shift_q, scale_q);
|
||||
q = LTXV::modulate(ctx->ggml_ctx, q, shift_q, scale_q);
|
||||
|
||||
auto context_mod = context;
|
||||
if (prompt_timestep != nullptr) {
|
||||
auto prompt_mods = get_prompt_scale_shift_values(ctx, prompt_timestep);
|
||||
context_mod = modulate(ctx->ggml_ctx, context_mod, prompt_mods[0], prompt_mods[1]);
|
||||
context_mod = LTXV::modulate(ctx->ggml_ctx, context_mod, prompt_mods[0], prompt_mods[1]);
|
||||
}
|
||||
|
||||
auto mca = attn2->forward(ctx, q, context_mod, attention_mask, nullptr, nullptr);
|
||||
@@ -826,7 +826,7 @@ namespace LTXV {
|
||||
}
|
||||
|
||||
auto y = rms_norm(ctx->ggml_ctx, x);
|
||||
y = modulate(ctx->ggml_ctx, y, shift_mlp, scale_mlp);
|
||||
y = LTXV::modulate(ctx->ggml_ctx, y, shift_mlp, scale_mlp);
|
||||
auto mlp_out = ff->forward(ctx, y);
|
||||
x = ggml_add(ctx->ggml_ctx, x, apply_gate(ctx->ggml_ctx, mlp_out, gate_mlp));
|
||||
return x;
|
||||
@@ -1177,11 +1177,11 @@ namespace LTXV {
|
||||
if (cross_attention_adaln) {
|
||||
auto q_mods = get_ada_values(ctx, table, timestep, dim, 9, 6, 3);
|
||||
auto q = rms_norm(ctx->ggml_ctx, x);
|
||||
q = modulate(ctx->ggml_ctx, q, q_mods[0], q_mods[1]);
|
||||
q = LTXV::modulate(ctx->ggml_ctx, q, q_mods[0], q_mods[1]);
|
||||
auto context_mod = context;
|
||||
if (prompt_timestep != nullptr && prompt_table != nullptr) {
|
||||
auto p_mods = get_ada_values(ctx, prompt_table, prompt_timestep, dim, 2);
|
||||
context_mod = modulate(ctx->ggml_ctx, context_mod, p_mods[0], p_mods[1]);
|
||||
context_mod = LTXV::modulate(ctx->ggml_ctx, context_mod, p_mods[0], p_mods[1]);
|
||||
}
|
||||
auto out = attn->forward(ctx, q, context_mod, attention_mask, nullptr, nullptr);
|
||||
return apply_gate(ctx->ggml_ctx, out, q_mods[2]);
|
||||
@@ -1228,7 +1228,7 @@ namespace LTXV {
|
||||
|
||||
auto v_mods = get_ada_values(ctx, v_table, v_timestep, v_dim, cross_attention_adaln ? 9 : 6);
|
||||
auto v_norm = rms_norm(ctx->ggml_ctx, vx);
|
||||
v_norm = modulate(ctx->ggml_ctx, v_norm, v_mods[0], v_mods[1]);
|
||||
v_norm = LTXV::modulate(ctx->ggml_ctx, v_norm, v_mods[0], v_mods[1]);
|
||||
auto v_sa = attn1->forward(ctx, v_norm, nullptr, self_attention_mask, v_pe);
|
||||
vx = ggml_add(ctx->ggml_ctx, vx, apply_gate(ctx->ggml_ctx, v_sa, v_mods[2]));
|
||||
auto v_txt = apply_text_cross_attention(ctx,
|
||||
@@ -1246,7 +1246,7 @@ namespace LTXV {
|
||||
if (run_ax) {
|
||||
auto a_mods = get_ada_values(ctx, a_table, a_timestep, a_dim, cross_attention_adaln ? 9 : 6);
|
||||
auto a_norm = rms_norm(ctx->ggml_ctx, ax);
|
||||
a_norm = modulate(ctx->ggml_ctx, a_norm, a_mods[0], a_mods[1]);
|
||||
a_norm = LTXV::modulate(ctx->ggml_ctx, a_norm, a_mods[0], a_mods[1]);
|
||||
auto a_sa = audio_attn1->forward(ctx, a_norm, nullptr, nullptr, a_pe);
|
||||
ax = ggml_add(ctx->ggml_ctx, ax, apply_gate(ctx->ggml_ctx, a_sa, a_mods[2]));
|
||||
auto a_txt = apply_text_cross_attention(ctx,
|
||||
@@ -1269,8 +1269,8 @@ namespace LTXV {
|
||||
auto a2v_video_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 0, 4);
|
||||
auto a2v_audio = get_ada_values(ctx, a2v_audio_table, a_cross_scale_shift_timestep, a_dim, 4);
|
||||
auto a2v_video = get_ada_values(ctx, a2v_video_table, v_cross_scale_shift_timestep, v_dim, 4);
|
||||
auto vx_scaled = modulate(ctx->ggml_ctx, vx_norm3, a2v_video[1], a2v_video[0]);
|
||||
auto ax_scaled = modulate(ctx->ggml_ctx, ax_norm3, a2v_audio[1], a2v_audio[0]);
|
||||
auto vx_scaled = LTXV::modulate(ctx->ggml_ctx, vx_norm3, a2v_video[1], a2v_video[0]);
|
||||
auto ax_scaled = LTXV::modulate(ctx->ggml_ctx, ax_norm3, a2v_audio[1], a2v_audio[0]);
|
||||
auto a2v_out = audio_to_video_attn->forward(ctx, vx_scaled, ax_scaled, nullptr, v_cross_pe, a_cross_pe);
|
||||
auto a2v_gate_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 4, 5);
|
||||
auto a2v_gate = get_ada_values(ctx, a2v_gate_table, v_cross_gate_timestep, v_dim, 1)[0];
|
||||
@@ -1282,8 +1282,8 @@ namespace LTXV {
|
||||
auto v2a_video_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 0, 4);
|
||||
auto v2a_audio = get_ada_values(ctx, v2a_audio_table, a_cross_scale_shift_timestep, a_dim, 4);
|
||||
auto v2a_video = get_ada_values(ctx, v2a_video_table, v_cross_scale_shift_timestep, v_dim, 4);
|
||||
auto ax_scaled = modulate(ctx->ggml_ctx, ax_norm3, v2a_audio[3], v2a_audio[2]);
|
||||
auto vx_scaled = modulate(ctx->ggml_ctx, vx_norm3, v2a_video[3], v2a_video[2]);
|
||||
auto ax_scaled = LTXV::modulate(ctx->ggml_ctx, ax_norm3, v2a_audio[3], v2a_audio[2]);
|
||||
auto vx_scaled = LTXV::modulate(ctx->ggml_ctx, vx_norm3, v2a_video[3], v2a_video[2]);
|
||||
auto v2a_out = video_to_audio_attn->forward(ctx, ax_scaled, vx_scaled, nullptr, a_cross_pe, v_cross_pe);
|
||||
auto v2a_gate_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_audio"], 1, 4, 5);
|
||||
auto v2a_gate = get_ada_values(ctx, v2a_gate_table, a_cross_gate_timestep, a_dim, 1)[0];
|
||||
@@ -1291,14 +1291,14 @@ namespace LTXV {
|
||||
}
|
||||
auto a_ff_mods = get_ada_values(ctx, a_table, a_timestep, a_dim, cross_attention_adaln ? 9 : 6, 3, 3);
|
||||
auto ax_scaled = rms_norm(ctx->ggml_ctx, ax);
|
||||
ax_scaled = modulate(ctx->ggml_ctx, ax_scaled, a_ff_mods[0], a_ff_mods[1]);
|
||||
ax_scaled = LTXV::modulate(ctx->ggml_ctx, ax_scaled, a_ff_mods[0], a_ff_mods[1]);
|
||||
auto a_ff_out = audio_ff->forward(ctx, ax_scaled);
|
||||
ax = ggml_add(ctx->ggml_ctx, ax, apply_gate(ctx->ggml_ctx, a_ff_out, a_ff_mods[2]));
|
||||
}
|
||||
|
||||
auto v_ff_mods = get_ada_values(ctx, v_table, v_timestep, v_dim, cross_attention_adaln ? 9 : 6, 3, 3);
|
||||
auto vx_scaled = rms_norm(ctx->ggml_ctx, vx);
|
||||
vx_scaled = modulate(ctx->ggml_ctx, vx_scaled, v_ff_mods[0], v_ff_mods[1]);
|
||||
vx_scaled = LTXV::modulate(ctx->ggml_ctx, vx_scaled, v_ff_mods[0], v_ff_mods[1]);
|
||||
auto v_ff_out = ff->forward(ctx, vx_scaled);
|
||||
vx = ggml_add(ctx->ggml_ctx, vx, apply_gate(ctx->ggml_ctx, v_ff_out, v_ff_mods[2]));
|
||||
|
||||
@@ -1657,14 +1657,14 @@ namespace LTXV {
|
||||
|
||||
auto v_shift_scale = get_output_scale_shift(ctx, params["scale_shift_table"], v_embedded_time, config.hidden_size);
|
||||
vx = norm_out->forward(ctx, vx);
|
||||
vx = modulate(ctx->ggml_ctx, vx, v_shift_scale[0], v_shift_scale[1]);
|
||||
vx = LTXV::modulate(ctx->ggml_ctx, vx, v_shift_scale[0], v_shift_scale[1]);
|
||||
vx = proj_out->forward(ctx, vx);
|
||||
vx = unpatchify_video(ctx, vx, width, height, frames);
|
||||
|
||||
if (ax != nullptr && audio_time > 0) {
|
||||
auto a_shift_scale = get_output_scale_shift(ctx, params["audio_scale_shift_table"], a_embedded_time, config.audio_hidden_size);
|
||||
ax = audio_norm_out->forward(ctx, ax);
|
||||
ax = modulate(ctx->ggml_ctx, ax, a_shift_scale[0], a_shift_scale[1]);
|
||||
ax = LTXV::modulate(ctx->ggml_ctx, ax, a_shift_scale[0], a_shift_scale[1]);
|
||||
ax = audio_proj_out->forward(ctx, ax);
|
||||
ax = unpatchify_audio(ctx, ax, audio_time);
|
||||
}
|
||||
|
||||
@@ -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__
|
||||
@@ -136,11 +136,15 @@ struct MMDiTConfig {
|
||||
};
|
||||
|
||||
struct PatchEmbed : public GGMLBlock {
|
||||
// 2D Image to Patch Embedding
|
||||
// 2D/3D Image to Patch Embedding
|
||||
protected:
|
||||
bool is_3d;
|
||||
bool flatten;
|
||||
bool dynamic_img_pad;
|
||||
int patch_size;
|
||||
int patch_t;
|
||||
int patch_h;
|
||||
int patch_w;
|
||||
int64_t embed_dim;
|
||||
|
||||
public:
|
||||
PatchEmbed(int64_t img_size = 224,
|
||||
@@ -149,42 +153,90 @@ public:
|
||||
int64_t embed_dim = 1536,
|
||||
bool bias = true,
|
||||
bool flatten = true,
|
||||
bool dynamic_img_pad = true)
|
||||
: patch_size(patch_size),
|
||||
bool dynamic_img_pad = true,
|
||||
bool is_3d = false)
|
||||
: patch_t(is_3d ? patch_size : 1),
|
||||
patch_h(patch_size),
|
||||
patch_w(patch_size),
|
||||
embed_dim(embed_dim),
|
||||
flatten(flatten),
|
||||
dynamic_img_pad(dynamic_img_pad) {
|
||||
dynamic_img_pad(dynamic_img_pad),
|
||||
is_3d(is_3d) {
|
||||
// img_size is always None
|
||||
// patch_size is always 2
|
||||
// in_chans is always 16
|
||||
// norm_layer is always False
|
||||
// strict_img_size is always true, but not used
|
||||
|
||||
blocks["proj"] = std::shared_ptr<GGMLBlock>(new Conv2d(in_chans,
|
||||
embed_dim,
|
||||
{patch_size, patch_size},
|
||||
{patch_size, patch_size},
|
||||
{0, 0},
|
||||
{1, 1},
|
||||
bias));
|
||||
if (is_3d) {
|
||||
blocks["proj"] = std::make_shared<Conv3d>(in_chans,
|
||||
embed_dim,
|
||||
std::tuple{patch_size, patch_size, patch_size},
|
||||
std::tuple{patch_size, patch_size, patch_size},
|
||||
std::tuple{0, 0, 0},
|
||||
std::tuple{1, 1, 1},
|
||||
bias);
|
||||
} else {
|
||||
blocks["proj"] = std::make_shared<Conv2d>(in_chans,
|
||||
embed_dim,
|
||||
std::pair{patch_size, patch_size},
|
||||
std::pair{patch_size, patch_size},
|
||||
std::pair{0, 0},
|
||||
std::pair{1, 1},
|
||||
bias);
|
||||
}
|
||||
}
|
||||
|
||||
PatchEmbed(int64_t img_size,
|
||||
std::tuple<int, int, int> patch_size,
|
||||
int64_t in_chans,
|
||||
int64_t embed_dim,
|
||||
bool bias = true,
|
||||
bool flatten = true,
|
||||
bool dynamic_img_pad = true)
|
||||
: patch_t(std::get<0>(patch_size)),
|
||||
patch_h(std::get<1>(patch_size)),
|
||||
patch_w(std::get<2>(patch_size)),
|
||||
embed_dim(embed_dim),
|
||||
flatten(flatten),
|
||||
dynamic_img_pad(dynamic_img_pad),
|
||||
is_3d(true) {
|
||||
SD_UNUSED(img_size);
|
||||
blocks["proj"] = std::make_shared<Conv3d>(in_chans,
|
||||
embed_dim,
|
||||
patch_size,
|
||||
patch_size,
|
||||
std::tuple{0, 0, 0},
|
||||
std::tuple{1, 1, 1},
|
||||
bias);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
// x: [N, C, H, W]
|
||||
// return: [N, H*W, embed_dim]
|
||||
auto proj = std::dynamic_pointer_cast<Conv2d>(blocks["proj"]);
|
||||
// x: [N, C, H, W] or [N*C, T, H, W]
|
||||
// return: [N, h_len*w_len, embed_dim] or [N, t_len*h_len*w_len, embed_dim]
|
||||
auto proj = std::dynamic_pointer_cast<UnaryBlock>(blocks["proj"]);
|
||||
|
||||
if (dynamic_img_pad) {
|
||||
int64_t W = x->ne[0];
|
||||
int64_t H = x->ne[1];
|
||||
int pad_h = (patch_size - H % patch_size) % patch_size;
|
||||
int pad_w = (patch_size - W % patch_size) % patch_size;
|
||||
x = ggml_pad(ctx->ggml_ctx, x, pad_w, pad_h, 0, 0); // TODO: reflect pad mode
|
||||
int pad_t = 0;
|
||||
int pad_h = (patch_h - static_cast<int>(H % patch_h)) % patch_h;
|
||||
int pad_w = (patch_w - static_cast<int>(W % patch_w)) % patch_w;
|
||||
if (is_3d) {
|
||||
int64_t T = x->ne[2];
|
||||
pad_t = (patch_t - static_cast<int>(T % patch_t)) % patch_t;
|
||||
}
|
||||
x = ggml_pad(ctx->ggml_ctx, x, pad_w, pad_h, pad_t, 0); // TODO: reflect pad mode
|
||||
}
|
||||
x = proj->forward(ctx, x);
|
||||
x = proj->forward(ctx, x); // [N, C, h_len, w_len] or [N*C, t_len, h_len, w_len]
|
||||
|
||||
if (flatten) {
|
||||
x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1], x->ne[2], x->ne[3]);
|
||||
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
|
||||
if (is_3d) {
|
||||
x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1] * x->ne[2], embed_dim, x->ne[3] / embed_dim); // [N, C, t_len*h_len*w_len]
|
||||
} else {
|
||||
x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1], x->ne[2], x->ne[3]); // [N, C, h_len*w_len]
|
||||
}
|
||||
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); // [N, h_len*w_len, C]
|
||||
}
|
||||
return x;
|
||||
}
|
||||
|
||||
@@ -10,10 +10,44 @@
|
||||
#include "model/common/rope.hpp"
|
||||
#include "model_manager.h"
|
||||
|
||||
enum class RefImageResizeMode {
|
||||
NONE,
|
||||
LONGEST_SIDE,
|
||||
AREA,
|
||||
};
|
||||
|
||||
struct RefImageParams {
|
||||
bool pass_to_vlm = false;
|
||||
bool pass_to_dit = true;
|
||||
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED;
|
||||
bool force_ref_timestep_zero = false;
|
||||
bool resize_before_vae = true;
|
||||
int vae_input_max_pixels = -1;
|
||||
RefImageResizeMode vlm_resize_mode = RefImageResizeMode::AREA;
|
||||
int vlm_min_size = -1;
|
||||
int vlm_max_size = -1;
|
||||
bool resize_vae_to_target = false;
|
||||
};
|
||||
|
||||
const std::unordered_map<std::string, RefImageParams> REF_IMAGE_PRESETS = {
|
||||
{"flux_kontext", {false, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::NONE, -1, -1}},
|
||||
{"longcat", {true, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
|
||||
{"flux2", {false, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::NONE, -1, -1}},
|
||||
{"qwen", {true, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
|
||||
{"qwen_layered", {true, true, Rope::RefIndexMode::DECREASE, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
|
||||
{"mage_flow", {true, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::LONGEST_SIDE, -1, 384, true}},
|
||||
{"z_image_omni", {true, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
|
||||
{"krea2_ostris_edit", {true, true, Rope::RefIndexMode::INCREASE, true, true, -1, RefImageResizeMode::AREA, -1, -1}},
|
||||
{"krea2_edit", {true, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::LONGEST_SIDE, 768, 768}},
|
||||
{"cosmos_reference", {false, true, Rope::RefIndexMode::INCREASE, false, false, -1, RefImageResizeMode::NONE, -1, -1}},
|
||||
};
|
||||
|
||||
struct UNetDiffusionExtra {
|
||||
int num_video_frames = -1;
|
||||
const std::vector<sd::Tensor<float>>* controls = nullptr;
|
||||
float control_strength = 0.f;
|
||||
const sd::Tensor<float>* ip_context = nullptr;
|
||||
float ip_scale = 1.f;
|
||||
};
|
||||
|
||||
struct SkipLayerDiffusionExtra {
|
||||
@@ -57,6 +91,13 @@ struct MiniT2IDiffusionExtra {
|
||||
const sd::Tensor<float>* mask = nullptr;
|
||||
};
|
||||
|
||||
struct HunyuanVideoDiffusionExtra {
|
||||
const sd::Tensor<float>* guidance = nullptr;
|
||||
const sd::Tensor<float>* byt5 = nullptr;
|
||||
const sd::Tensor<float>* vision = nullptr;
|
||||
const sd::Tensor<float>* timestep_r = nullptr;
|
||||
};
|
||||
|
||||
using DiffusionExtraParams = std::variant<std::monostate,
|
||||
UNetDiffusionExtra,
|
||||
SkipLayerDiffusionExtra,
|
||||
@@ -65,7 +106,8 @@ using DiffusionExtraParams = std::variant<std::monostate,
|
||||
WanDiffusionExtra,
|
||||
HiDreamO1DiffusionExtra,
|
||||
LTXAVDiffusionExtra,
|
||||
MiniT2IDiffusionExtra>;
|
||||
MiniT2IDiffusionExtra,
|
||||
HunyuanVideoDiffusionExtra>;
|
||||
|
||||
struct DiffusionParams {
|
||||
const sd::Tensor<float>* x = nullptr;
|
||||
@@ -74,7 +116,7 @@ struct DiffusionParams {
|
||||
const sd::Tensor<float>* c_concat = nullptr;
|
||||
const sd::Tensor<float>* y = nullptr;
|
||||
const std::vector<sd::Tensor<float>>* ref_latents = nullptr;
|
||||
Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED;
|
||||
RefImageParams ref_image_params = {false, false, Rope::RefIndexMode::FIXED, false};
|
||||
DiffusionExtraParams extra = std::monostate{};
|
||||
};
|
||||
|
||||
|
||||
@@ -17,30 +17,38 @@ namespace Pid {
|
||||
constexpr float PID_PI = 3.14159265358979323846f;
|
||||
|
||||
struct PixelDiTConfig {
|
||||
int64_t in_channels = 3;
|
||||
int64_t hidden_size = 1536;
|
||||
int64_t num_groups = 24;
|
||||
int64_t patch_mlp_hidden_dim = 4096;
|
||||
int64_t pixel_hidden_size = 16;
|
||||
int64_t pixel_attn_hidden_size = 1152;
|
||||
int64_t pixel_num_groups = 16;
|
||||
int64_t patch_depth = 14;
|
||||
int64_t pixel_depth = 2;
|
||||
int64_t patch_size = 16;
|
||||
int64_t txt_embed_dim = 2304;
|
||||
int64_t txt_max_length = 300;
|
||||
float text_rope_theta = 10000.f;
|
||||
int64_t lq_latent_channels = 16;
|
||||
int64_t lq_hidden_dim = 512;
|
||||
int64_t lq_num_res_blocks = 4;
|
||||
int64_t lq_interval = 2;
|
||||
int64_t lq_sr_scale = 4;
|
||||
int64_t lq_latent_down_factor = 8;
|
||||
int64_t rope_ref_grid_h = 64;
|
||||
int64_t rope_ref_grid_w = 64;
|
||||
int64_t in_channels = 3;
|
||||
int64_t hidden_size = 1536;
|
||||
int64_t num_groups = 24;
|
||||
int64_t patch_mlp_hidden_dim = 4096;
|
||||
int64_t pixel_hidden_size = 16;
|
||||
int64_t pixel_attn_hidden_size = 1152;
|
||||
int64_t pixel_num_groups = 16;
|
||||
int64_t patch_depth = 14;
|
||||
int64_t pixel_depth = 2;
|
||||
int64_t patch_size = 16;
|
||||
int64_t txt_embed_dim = 2304;
|
||||
int64_t txt_max_length = 300;
|
||||
float text_rope_theta = 10000.f;
|
||||
int64_t lq_latent_channels = 16;
|
||||
int64_t lq_hidden_dim = 512;
|
||||
int64_t lq_num_res_blocks = 4;
|
||||
int64_t lq_interval = 2;
|
||||
int64_t lq_sr_scale = 4;
|
||||
int64_t lq_latent_down_factor = 8;
|
||||
int64_t lq_latent_unpatchify_factor = 1;
|
||||
bool lq_replicate_padding = false;
|
||||
bool lq_gate_per_token = false;
|
||||
bool pit_lq_inject = false;
|
||||
int64_t rope_ref_grid_h = 64;
|
||||
int64_t rope_ref_grid_w = 64;
|
||||
|
||||
static PixelDiTConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
|
||||
PixelDiTConfig config;
|
||||
int64_t latent_proj_in_channels = config.lq_latent_channels;
|
||||
int64_t num_lq_gates = 0;
|
||||
const std::string lq_prefix = prefix + ".lq_proj.";
|
||||
config.pit_lq_inject = tensor_storage_map.find(lq_prefix + "pit_head.weight") != tensor_storage_map.end();
|
||||
for (const auto& [name, tensor_storage] : tensor_storage_map) {
|
||||
if (!starts_with(name, prefix)) {
|
||||
continue;
|
||||
@@ -61,20 +69,56 @@ namespace Pid {
|
||||
config.pixel_depth = std::max<int64_t>(config.pixel_depth, block_index + 1);
|
||||
}
|
||||
}
|
||||
if (name.find("lq_proj.latent_proj.0.weight") != std::string::npos) {
|
||||
config.lq_latent_channels = tensor_storage.ne[2];
|
||||
config.lq_latent_down_factor = config.lq_latent_channels >= 64 ? 16 : 8;
|
||||
if (name == lq_prefix + "latent_proj.0.weight") {
|
||||
latent_proj_in_channels = tensor_storage.ne[2];
|
||||
config.lq_hidden_dim = tensor_storage.ne[3];
|
||||
}
|
||||
if (starts_with(name, lq_prefix + "gate_modules.")) {
|
||||
auto items = split_string(name.substr(lq_prefix.size()), '.');
|
||||
if (items.size() > 1) {
|
||||
int gate_index = atoi(items[1].c_str());
|
||||
num_lq_gates = std::max<int64_t>(num_lq_gates, gate_index + 1);
|
||||
}
|
||||
}
|
||||
if (name.find("patch_blocks.0.mlp_x.w1.weight") != std::string::npos) {
|
||||
config.patch_mlp_hidden_dim = tensor_storage.ne[1];
|
||||
}
|
||||
}
|
||||
LOG_DEBUG("pid: patch_depth = %" PRId64 ", pixel_depth = %" PRId64 ", patch_mlp_hidden_dim = %" PRId64 ", lq_latent_channels = %" PRId64 ", lq_latent_down_factor = %" PRId64,
|
||||
if (num_lq_gates > 0) {
|
||||
config.lq_interval = (config.patch_depth + num_lq_gates - 1) / num_lq_gates;
|
||||
}
|
||||
if (config.pit_lq_inject) {
|
||||
if (latent_proj_in_channels == 16) {
|
||||
config.lq_latent_channels = 16;
|
||||
config.lq_latent_down_factor = 8;
|
||||
config.lq_latent_unpatchify_factor = 1;
|
||||
} else {
|
||||
GGML_ASSERT(latent_proj_in_channels == 32);
|
||||
config.lq_latent_channels = 128;
|
||||
config.lq_latent_down_factor = 16;
|
||||
config.lq_latent_unpatchify_factor = 2;
|
||||
}
|
||||
auto gate_weight = tensor_storage_map.find(lq_prefix + "gate_modules.0.content_proj.weight");
|
||||
if (gate_weight != tensor_storage_map.end()) {
|
||||
config.lq_gate_per_token = gate_weight->second.ne[1] == 1;
|
||||
}
|
||||
config.lq_replicate_padding = true;
|
||||
config.rope_ref_grid_h = 128;
|
||||
config.rope_ref_grid_w = 128;
|
||||
} else {
|
||||
config.lq_latent_channels = latent_proj_in_channels;
|
||||
config.lq_latent_down_factor = latent_proj_in_channels >= 64 ? 16 : 8;
|
||||
}
|
||||
LOG_DEBUG("pid: version = %s, patch_depth = %" PRId64 ", pixel_depth = %" PRId64 ", patch_mlp_hidden_dim = %" PRId64 ", lq_latent_channels = %" PRId64 ", lq_hidden_dim = %" PRId64 ", lq_latent_down_factor = %" PRId64 ", lq_latent_unpatchify_factor = %" PRId64 ", lq_interval = %" PRId64,
|
||||
config.pit_lq_inject ? "1.5" : "1",
|
||||
config.patch_depth,
|
||||
config.pixel_depth,
|
||||
config.patch_mlp_hidden_dim,
|
||||
config.lq_latent_channels,
|
||||
config.lq_latent_down_factor);
|
||||
config.lq_hidden_dim,
|
||||
config.lq_latent_down_factor,
|
||||
config.lq_latent_unpatchify_factor,
|
||||
config.lq_interval);
|
||||
return config;
|
||||
}
|
||||
};
|
||||
@@ -135,6 +179,18 @@ namespace Pid {
|
||||
return ggml_add(ctx, ggml_add(ctx, x, ggml_mul(ctx, x, scale)), shift);
|
||||
}
|
||||
|
||||
inline ggml_tensor* replicate_pad_2d(ggml_context* ctx, ggml_tensor* x) {
|
||||
auto left = ggml_ext_slice(ctx, x, 0, 0, 1);
|
||||
auto right = ggml_ext_slice(ctx, x, 0, x->ne[0] - 1, x->ne[0]);
|
||||
x = ggml_concat(ctx, left, x, 0);
|
||||
x = ggml_concat(ctx, x, right, 0);
|
||||
|
||||
auto top = ggml_ext_slice(ctx, x, 1, 0, 1);
|
||||
auto bottom = ggml_ext_slice(ctx, x, 1, x->ne[1] - 1, x->ne[1]);
|
||||
x = ggml_concat(ctx, top, x, 1);
|
||||
return ggml_concat(ctx, x, bottom, 1);
|
||||
}
|
||||
|
||||
struct PatchTokenEmbedder : public GGMLBlock {
|
||||
bool use_rms_norm;
|
||||
|
||||
@@ -457,9 +513,9 @@ namespace Pid {
|
||||
struct SigmaAwareGate : public GGMLBlock {
|
||||
int64_t dim;
|
||||
|
||||
SigmaAwareGate(int64_t dim)
|
||||
SigmaAwareGate(int64_t dim, bool per_token = false)
|
||||
: dim(dim) {
|
||||
blocks["content_proj"] = std::make_shared<Linear>(dim * 2, dim, true);
|
||||
blocks["content_proj"] = std::make_shared<Linear>(dim * 2, per_token ? 1 : dim, true);
|
||||
}
|
||||
|
||||
void init_params(ggml_context* ctx,
|
||||
@@ -479,16 +535,20 @@ namespace Pid {
|
||||
auto alpha = ggml_exp(ctx->ggml_ctx, params["log_alpha"]);
|
||||
auto offset = ggml_neg(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, alpha, sigma));
|
||||
auto gate = ggml_sigmoid(ctx->ggml_ctx, ggml_add(ctx->ggml_ctx, content_logit, offset));
|
||||
return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, gate, lq));
|
||||
return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, lq, gate));
|
||||
}
|
||||
};
|
||||
|
||||
struct PiDResBlock : public GGMLBlock {
|
||||
PiDResBlock(int64_t channels) {
|
||||
blocks["block.0"] = std::make_shared<GroupNorm>(4, channels, 1e-5f);
|
||||
blocks["block.2"] = std::make_shared<Conv2d>(channels, channels, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, std::pair<int, int>{1, 1});
|
||||
blocks["block.3"] = std::make_shared<GroupNorm>(4, channels, 1e-5f);
|
||||
blocks["block.5"] = std::make_shared<Conv2d>(channels, channels, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, std::pair<int, int>{1, 1});
|
||||
bool replicate_padding;
|
||||
|
||||
PiDResBlock(int64_t channels, bool replicate_padding = false)
|
||||
: replicate_padding(replicate_padding) {
|
||||
std::pair<int, int> padding = replicate_padding ? std::pair<int, int>{0, 0} : std::pair<int, int>{1, 1};
|
||||
blocks["block.0"] = std::make_shared<GroupNorm>(4, channels, 1e-5f);
|
||||
blocks["block.2"] = std::make_shared<Conv2d>(channels, channels, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, padding);
|
||||
blocks["block.3"] = std::make_shared<GroupNorm>(4, channels, 1e-5f);
|
||||
blocks["block.5"] = std::make_shared<Conv2d>(channels, channels, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, padding);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
@@ -497,9 +557,15 @@ namespace Pid {
|
||||
auto norm2 = std::dynamic_pointer_cast<GroupNorm>(blocks["block.3"]);
|
||||
auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["block.5"]);
|
||||
auto h = ggml_silu_inplace(ctx->ggml_ctx, norm1->forward(ctx, x));
|
||||
h = conv1->forward(ctx, h);
|
||||
h = ggml_silu_inplace(ctx->ggml_ctx, norm2->forward(ctx, h));
|
||||
h = conv2->forward(ctx, h);
|
||||
if (replicate_padding) {
|
||||
h = replicate_pad_2d(ctx->ggml_ctx, h);
|
||||
}
|
||||
h = conv1->forward(ctx, h);
|
||||
h = ggml_silu_inplace(ctx->ggml_ctx, norm2->forward(ctx, h));
|
||||
if (replicate_padding) {
|
||||
h = replicate_pad_2d(ctx->ggml_ctx, h);
|
||||
}
|
||||
h = conv2->forward(ctx, h);
|
||||
return ggml_add(ctx->ggml_ctx, x, h);
|
||||
}
|
||||
};
|
||||
@@ -509,16 +575,23 @@ namespace Pid {
|
||||
|
||||
LQProjection2D(const PixelDiTConfig& config)
|
||||
: config(config) {
|
||||
blocks["latent_proj.0"] = std::make_shared<Conv2d>(config.lq_latent_channels, config.lq_hidden_dim, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, std::pair<int, int>{1, 1});
|
||||
blocks["latent_proj.2"] = std::make_shared<Conv2d>(config.lq_hidden_dim, config.lq_hidden_dim, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, std::pair<int, int>{1, 1});
|
||||
int64_t unpatchify_area = config.lq_latent_unpatchify_factor * config.lq_latent_unpatchify_factor;
|
||||
GGML_ASSERT(config.lq_latent_channels % unpatchify_area == 0);
|
||||
int64_t latent_proj_in_channels = config.lq_latent_channels / unpatchify_area;
|
||||
std::pair<int, int> padding = config.lq_replicate_padding ? std::pair<int, int>{0, 0} : std::pair<int, int>{1, 1};
|
||||
blocks["latent_proj.0"] = std::make_shared<Conv2d>(latent_proj_in_channels, config.lq_hidden_dim, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, padding);
|
||||
blocks["latent_proj.2"] = std::make_shared<Conv2d>(config.lq_hidden_dim, config.lq_hidden_dim, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, padding);
|
||||
for (int i = 0; i < config.lq_num_res_blocks; ++i) {
|
||||
blocks["latent_proj." + std::to_string(3 + i)] = std::make_shared<PiDResBlock>(config.lq_hidden_dim);
|
||||
blocks["latent_proj." + std::to_string(3 + i)] = std::make_shared<PiDResBlock>(config.lq_hidden_dim, config.lq_replicate_padding);
|
||||
}
|
||||
|
||||
int num_outputs = static_cast<int>((config.patch_depth + config.lq_interval - 1) / config.lq_interval);
|
||||
for (int i = 0; i < num_outputs; ++i) {
|
||||
blocks["output_heads." + std::to_string(i)] = std::make_shared<Linear>(config.lq_hidden_dim, config.hidden_size, true);
|
||||
blocks["gate_modules." + std::to_string(i)] = std::make_shared<SigmaAwareGate>(config.hidden_size);
|
||||
blocks["gate_modules." + std::to_string(i)] = std::make_shared<SigmaAwareGate>(config.hidden_size, config.lq_gate_per_token);
|
||||
}
|
||||
if (config.pit_lq_inject) {
|
||||
blocks["pit_head"] = std::make_shared<Linear>(config.lq_hidden_dim, config.hidden_size, true);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -543,9 +616,29 @@ namespace Pid {
|
||||
ggml_tensor* lq_latent,
|
||||
int64_t target_pH,
|
||||
int64_t target_pW) {
|
||||
auto conv0 = std::dynamic_pointer_cast<Conv2d>(blocks["latent_proj.0"]);
|
||||
auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["latent_proj.2"]);
|
||||
float z_to_patch_ratio = static_cast<float>(config.lq_sr_scale * config.lq_latent_down_factor) /
|
||||
auto conv0 = std::dynamic_pointer_cast<Conv2d>(blocks["latent_proj.0"]);
|
||||
auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["latent_proj.2"]);
|
||||
int64_t unpatchify_factor = config.lq_latent_unpatchify_factor;
|
||||
if (unpatchify_factor > 1) {
|
||||
int64_t latent_h = lq_latent->ne[1];
|
||||
int64_t latent_w = lq_latent->ne[0];
|
||||
lq_latent = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, lq_latent, 2, 0, 1, 3));
|
||||
lq_latent = ggml_reshape_3d(ctx->ggml_ctx,
|
||||
lq_latent,
|
||||
lq_latent->ne[0],
|
||||
lq_latent->ne[1] * lq_latent->ne[2],
|
||||
lq_latent->ne[3]);
|
||||
lq_latent = DiT::unpatchify(ctx->ggml_ctx,
|
||||
lq_latent,
|
||||
latent_h,
|
||||
latent_w,
|
||||
static_cast<int>(unpatchify_factor),
|
||||
static_cast<int>(unpatchify_factor),
|
||||
true);
|
||||
}
|
||||
|
||||
int64_t effective_down_factor = config.lq_latent_down_factor / unpatchify_factor;
|
||||
float z_to_patch_ratio = static_cast<float>(config.lq_sr_scale * effective_down_factor) /
|
||||
static_cast<float>(config.patch_size);
|
||||
GGML_ASSERT(z_to_patch_ratio >= 1.0f);
|
||||
if (lq_latent->ne[0] != target_pW || lq_latent->ne[1] != target_pH) {
|
||||
@@ -558,9 +651,15 @@ namespace Pid {
|
||||
GGML_SCALE_MODE_NEAREST);
|
||||
}
|
||||
|
||||
if (config.lq_replicate_padding) {
|
||||
lq_latent = replicate_pad_2d(ctx->ggml_ctx, lq_latent);
|
||||
}
|
||||
auto feat = conv0->forward(ctx, lq_latent);
|
||||
feat = ggml_silu_inplace(ctx->ggml_ctx, feat);
|
||||
feat = conv2->forward(ctx, feat);
|
||||
if (config.lq_replicate_padding) {
|
||||
feat = replicate_pad_2d(ctx->ggml_ctx, feat);
|
||||
}
|
||||
feat = conv2->forward(ctx, feat);
|
||||
for (int i = 0; i < config.lq_num_res_blocks; ++i) {
|
||||
auto block = std::dynamic_pointer_cast<PiDResBlock>(blocks["latent_proj." + std::to_string(3 + i)]);
|
||||
feat = block->forward(ctx, feat);
|
||||
@@ -574,11 +673,15 @@ namespace Pid {
|
||||
|
||||
int num_outputs = static_cast<int>((config.patch_depth + config.lq_interval - 1) / config.lq_interval);
|
||||
std::vector<ggml_tensor*> outputs;
|
||||
outputs.reserve(num_outputs);
|
||||
outputs.reserve(num_outputs + (config.pit_lq_inject ? 1 : 0));
|
||||
for (int i = 0; i < num_outputs; ++i) {
|
||||
auto head = std::dynamic_pointer_cast<Linear>(blocks["output_heads." + std::to_string(i)]);
|
||||
outputs.push_back(head->forward(ctx, tokens));
|
||||
}
|
||||
if (config.pit_lq_inject) {
|
||||
auto pit_head = std::dynamic_pointer_cast<Linear>(blocks["pit_head"]);
|
||||
outputs.push_back(pit_head->forward(ctx, tokens));
|
||||
}
|
||||
return outputs;
|
||||
}
|
||||
};
|
||||
@@ -606,6 +709,9 @@ namespace Pid {
|
||||
}
|
||||
blocks["final_layer"] = std::make_shared<FinalLayer>(config.pixel_hidden_size, config.in_channels);
|
||||
blocks["lq_proj"] = std::make_shared<LQProjection2D>(config);
|
||||
if (config.pit_lq_inject) {
|
||||
blocks["pit_lq_gate"] = std::make_shared<SigmaAwareGate>(config.hidden_size, config.lq_gate_per_token);
|
||||
}
|
||||
}
|
||||
|
||||
void init_params(ggml_context* ctx,
|
||||
@@ -654,6 +760,11 @@ namespace Pid {
|
||||
y_emb = ggml_add(ctx->ggml_ctx, y_emb, y_pos);
|
||||
|
||||
std::vector<ggml_tensor*> lq_features = lq_proj->forward(ctx, lq_latent, Hs, Ws);
|
||||
ggml_tensor* pit_lq_feature = nullptr;
|
||||
if (config.pit_lq_inject) {
|
||||
pit_lq_feature = lq_features.back();
|
||||
lq_features.pop_back();
|
||||
}
|
||||
|
||||
auto s = s_embedder->forward(ctx, x_patches);
|
||||
|
||||
@@ -677,6 +788,10 @@ namespace Pid {
|
||||
sd::ggml_graph_cut::mark_graph_cut(y_emb, "pid.patch_blocks." + std::to_string(i), "y");
|
||||
}
|
||||
s = ggml_silu(ctx->ggml_ctx, ggml_add(ctx->ggml_ctx, s, t_emb));
|
||||
if (pit_lq_feature != nullptr) {
|
||||
auto pit_lq_gate = std::dynamic_pointer_cast<SigmaAwareGate>(blocks["pit_lq_gate"]);
|
||||
s = pit_lq_gate->forward(ctx, s, pit_lq_feature, degrade_sigma);
|
||||
}
|
||||
|
||||
auto s_cond = ggml_reshape_2d(ctx->ggml_ctx, s, config.hidden_size, L * B);
|
||||
auto pixels = pixel_embedder->forward(ctx, x, config.patch_size, pixel_pos_full);
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
|
||||
#include <memory>
|
||||
|
||||
#include "core/util.h"
|
||||
#include "model/common/block.hpp"
|
||||
#include "model/diffusion/dit.hpp"
|
||||
#include "model/diffusion/flux.hpp"
|
||||
@@ -102,9 +103,9 @@ namespace Qwen {
|
||||
}
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* timesteps,
|
||||
ggml_tensor* addition_t_cond = nullptr) {
|
||||
virtual ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* timesteps,
|
||||
ggml_tensor* addition_t_cond = nullptr) {
|
||||
// timesteps: [N,]
|
||||
// return: [N, embedding_dim]
|
||||
auto timestep_embedder = std::dynamic_pointer_cast<TimestepEmbedding>(blocks["timestep_embedder"]);
|
||||
@@ -182,7 +183,7 @@ namespace Qwen {
|
||||
auto to_v = std::dynamic_pointer_cast<Linear>(blocks["to_v"]);
|
||||
auto to_out_0 = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
|
||||
|
||||
if (sd_backend_is(ctx->backend, "Vulkan")) {
|
||||
if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
|
||||
to_out_0->set_force_prec_f32(true);
|
||||
}
|
||||
|
||||
@@ -415,10 +416,14 @@ namespace Qwen {
|
||||
|
||||
public:
|
||||
QwenImageModel() {}
|
||||
QwenImageModel(QwenImageConfig config)
|
||||
QwenImageModel(QwenImageConfig config,
|
||||
std::shared_ptr<QwenTimestepProjEmbeddings> time_text_embed = nullptr)
|
||||
: config(config) {
|
||||
int64_t inner_dim = config.num_attention_heads * config.attention_head_dim;
|
||||
blocks["time_text_embed"] = std::shared_ptr<GGMLBlock>(new QwenTimestepProjEmbeddings(inner_dim, config.use_additional_t_cond));
|
||||
int64_t inner_dim = config.num_attention_heads * config.attention_head_dim;
|
||||
if (time_text_embed == nullptr) {
|
||||
time_text_embed = std::make_shared<QwenTimestepProjEmbeddings>(inner_dim, config.use_additional_t_cond);
|
||||
}
|
||||
blocks["time_text_embed"] = std::move(time_text_embed);
|
||||
blocks["txt_norm"] = std::shared_ptr<GGMLBlock>(new RMSNorm(config.joint_attention_dim, 1e-6f));
|
||||
blocks["img_in"] = std::shared_ptr<GGMLBlock>(new Linear(config.in_channels, inner_dim));
|
||||
blocks["txt_in"] = std::shared_ptr<GGMLBlock>(new Linear(config.joint_attention_dim, inner_dim));
|
||||
@@ -516,7 +521,7 @@ namespace Qwen {
|
||||
if (input->ne[3] == 1) {
|
||||
input = ggml_reshape_4d(ctx->ggml_ctx, input, input->ne[0], input->ne[1], 1, input->ne[2]);
|
||||
}
|
||||
return DiT::patchify(ctx->ggml_ctx, input, 1, config.patch_size, config.patch_size, N);
|
||||
return DiT::patchify_3d(ctx->ggml_ctx, input, 1, config.patch_size, config.patch_size, N);
|
||||
};
|
||||
|
||||
auto img = patchify_input(x);
|
||||
@@ -566,12 +571,21 @@ namespace Qwen {
|
||||
const String2TensorStorage& tensor_storage_map = {},
|
||||
const std::string prefix = "",
|
||||
SDVersion version = VERSION_QWEN_IMAGE,
|
||||
bool zero_cond_t = false,
|
||||
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
|
||||
std::shared_ptr<RunnerWeightManager> weight_manager = nullptr,
|
||||
const char* model_args = nullptr)
|
||||
: DiffusionModelRunner(backend, prefix, weight_manager),
|
||||
config(QwenImageConfig::detect_from_weights(tensor_storage_map, prefix)),
|
||||
version(version) {
|
||||
config.zero_cond_t = config.zero_cond_t || zero_cond_t;
|
||||
for (const auto& [key, value] : parse_key_value_args(model_args, "model arg")) {
|
||||
if (key == "qwen_image_zero_cond_t") {
|
||||
bool parsed = false;
|
||||
if (parse_strict_bool(value, parsed)) {
|
||||
config.zero_cond_t = config.zero_cond_t || parsed;
|
||||
} else {
|
||||
LOG_WARN("ignoring invalid Qwen Image model arg '%s=%s'", key.c_str(), value.c_str());
|
||||
}
|
||||
}
|
||||
}
|
||||
if (version == VERSION_QWEN_IMAGE_LAYERED) {
|
||||
config.use_additional_t_cond = true;
|
||||
}
|
||||
@@ -705,8 +719,8 @@ namespace Qwen {
|
||||
*diffusion_params.x,
|
||||
*diffusion_params.timesteps,
|
||||
tensor_or_empty(diffusion_params.context),
|
||||
diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
|
||||
diffusion_params.ref_index_mode);
|
||||
diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents,
|
||||
diffusion_params.ref_image_params.ref_index_mode);
|
||||
}
|
||||
|
||||
void test() {
|
||||
@@ -775,7 +789,6 @@ namespace Qwen {
|
||||
tensor_storage_map,
|
||||
"model.diffusion_model",
|
||||
VERSION_QWEN_IMAGE,
|
||||
false,
|
||||
model_manager);
|
||||
|
||||
if (!model_manager->register_runner_params("Qwen image test",
|
||||
|
||||
@@ -6,6 +6,7 @@
|
||||
|
||||
#include "model.h"
|
||||
#include "model/common/block.hpp"
|
||||
#include "model/diffusion/animatediff.hpp"
|
||||
#include "model/diffusion/model.hpp"
|
||||
|
||||
/*==================================================== UnetModel =====================================================*/
|
||||
@@ -29,6 +30,8 @@ struct UNetConfig {
|
||||
bool tiny_unet = false;
|
||||
int model_channels = 320;
|
||||
int adm_in_channels = 2816; // only for VERSION_SDXL/SVD
|
||||
bool enable_animatediff = false;
|
||||
bool animatediff_has_mid_block = false;
|
||||
|
||||
static UNetConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
|
||||
const std::string& prefix,
|
||||
@@ -84,6 +87,13 @@ struct UNetConfig {
|
||||
return &it->second;
|
||||
};
|
||||
|
||||
if (find_weight("motion_module.down_blocks.0.motion_modules.0.temporal_transformer.proj_in.weight") != nullptr) {
|
||||
config.enable_animatediff = true;
|
||||
if (find_weight("motion_module.mid_block.motion_modules.0.temporal_transformer.proj_in.weight") != nullptr) {
|
||||
config.animatediff_has_mid_block = true;
|
||||
}
|
||||
}
|
||||
|
||||
if (const TensorStorage* input = find_weight("input_blocks.0.0.weight")) {
|
||||
if (input->n_dims == 4) {
|
||||
config.in_channels = static_cast<int>(input->ne[2]);
|
||||
@@ -473,6 +483,12 @@ public:
|
||||
blocks["out.0"] = std::shared_ptr<GGMLBlock>(new GroupNorm32(ch)); // ch == model_channels
|
||||
// out_1 is nn.SiLU()
|
||||
blocks["out.2"] = std::shared_ptr<GGMLBlock>(new Conv2d(model_channels, out_channels, {3, 3}, {1, 1}, {1, 1}));
|
||||
|
||||
if (this->config.enable_animatediff) {
|
||||
AnimateDiff::MotionModuleConfig mm_cfg;
|
||||
mm_cfg.enable_mid_block = this->config.animatediff_has_mid_block;
|
||||
blocks["motion_module"] = std::make_shared<AnimateDiff::AnimateDiffModel>(mm_cfg);
|
||||
}
|
||||
}
|
||||
|
||||
ggml_tensor* resblock_forward(std::string name,
|
||||
@@ -583,6 +599,42 @@ public:
|
||||
|
||||
ggml_set_name(h, "bench-start");
|
||||
hs.push_back(h);
|
||||
|
||||
auto motion_root = config.enable_animatediff && num_video_frames > 1
|
||||
? std::dynamic_pointer_cast<AnimateDiff::AnimateDiffModel>(blocks["motion_module"])
|
||||
: nullptr;
|
||||
auto apply_motion_input = [&](int input_block_idx, ggml_tensor* h_in) -> ggml_tensor* {
|
||||
if (!motion_root)
|
||||
return h_in;
|
||||
int di = (input_block_idx - 1) / 3;
|
||||
int mj = (input_block_idx - 1) % 3;
|
||||
if (di < 0 || di >= (int)channel_mult.size() || mj < 0 || mj >= num_res_blocks)
|
||||
return h_in;
|
||||
auto mm = motion_root->motion("down_blocks." + std::to_string(di) + ".motion_modules." + std::to_string(mj));
|
||||
if (!mm)
|
||||
return h_in;
|
||||
return mm->forward(ctx, h_in, num_video_frames);
|
||||
};
|
||||
auto apply_motion_output = [&](int output_block_idx, ggml_tensor* h_in) -> ggml_tensor* {
|
||||
if (!motion_root)
|
||||
return h_in;
|
||||
int ui = output_block_idx / 3;
|
||||
int mj = output_block_idx % 3;
|
||||
if (ui < 0 || ui >= (int)channel_mult.size() || mj < 0 || mj > num_res_blocks)
|
||||
return h_in;
|
||||
auto mm = motion_root->motion("up_blocks." + std::to_string(ui) + ".motion_modules." + std::to_string(mj));
|
||||
if (!mm)
|
||||
return h_in;
|
||||
return mm->forward(ctx, h_in, num_video_frames);
|
||||
};
|
||||
auto apply_motion_mid = [&](ggml_tensor* h_in) -> ggml_tensor* {
|
||||
if (!motion_root)
|
||||
return h_in;
|
||||
auto mm = motion_root->motion("mid_block.motion_modules.0");
|
||||
if (!mm)
|
||||
return h_in;
|
||||
return mm->forward(ctx, h_in, num_video_frames);
|
||||
};
|
||||
// input block 1-11
|
||||
size_t len_mults = channel_mult.size();
|
||||
int input_block_idx = 0;
|
||||
@@ -597,6 +649,7 @@ public:
|
||||
std::string name = "input_blocks." + std::to_string(input_block_idx) + ".1";
|
||||
h = attention_layer_forward(name, ctx, h, context, num_video_frames); // [N, mult*model_channels, h, w]
|
||||
}
|
||||
h = apply_motion_input(input_block_idx, h);
|
||||
sd::ggml_graph_cut::mark_graph_cut(h, "unet.input_blocks." + std::to_string(input_block_idx), "h");
|
||||
hs.push_back(h);
|
||||
}
|
||||
@@ -624,6 +677,7 @@ public:
|
||||
h = attention_layer_forward("middle_block.1", ctx, h, context, num_video_frames); // [N, 4*model_channels, h/8, w/8]
|
||||
h = resblock_forward("middle_block.2", ctx, h, emb, num_video_frames); // [N, 4*model_channels, h/8, w/8]
|
||||
}
|
||||
h = apply_motion_mid(h);
|
||||
}
|
||||
sd::ggml_graph_cut::mark_graph_cut(h, "unet.middle_block", "h");
|
||||
if (controls.size() > 0) {
|
||||
@@ -660,6 +714,8 @@ public:
|
||||
up_sample_idx++;
|
||||
}
|
||||
|
||||
h = apply_motion_output(output_block_idx, h);
|
||||
|
||||
if (i > 0 && j == num_res_blocks) {
|
||||
if (tiny_unet) {
|
||||
output_block_idx++;
|
||||
@@ -719,14 +775,17 @@ struct UNetModelRunner : public DiffusionModelRunner {
|
||||
const sd::Tensor<float>& y_tensor = {},
|
||||
int num_video_frames = -1,
|
||||
const std::vector<sd::Tensor<float>>& controls_tensor = {},
|
||||
float control_strength = 0.f) {
|
||||
float control_strength = 0.f,
|
||||
const sd::Tensor<float>& ip_context_tensor = {},
|
||||
float ip_scale = 1.f) {
|
||||
ggml_cgraph* gf = new_graph_custom(UNET_GRAPH_SIZE);
|
||||
|
||||
ggml_tensor* x = make_input(x_tensor);
|
||||
ggml_tensor* timesteps = make_input(timesteps_tensor);
|
||||
ggml_tensor* context = make_optional_input(context_tensor);
|
||||
ggml_tensor* c_concat = make_optional_input(c_concat_tensor);
|
||||
ggml_tensor* y = make_optional_input(y_tensor);
|
||||
ggml_tensor* x = make_input(x_tensor);
|
||||
ggml_tensor* timesteps = make_input(timesteps_tensor);
|
||||
ggml_tensor* context = make_optional_input(context_tensor);
|
||||
ggml_tensor* c_concat = make_optional_input(c_concat_tensor);
|
||||
ggml_tensor* y = make_optional_input(y_tensor);
|
||||
ggml_tensor* ip_context = make_optional_input(ip_context_tensor);
|
||||
std::vector<ggml_tensor*> controls;
|
||||
controls.reserve(controls_tensor.size());
|
||||
for (const auto& control_tensor : controls_tensor) {
|
||||
@@ -737,7 +796,9 @@ struct UNetModelRunner : public DiffusionModelRunner {
|
||||
num_video_frames = static_cast<int>(x->ne[3]);
|
||||
}
|
||||
|
||||
auto runner_ctx = get_context();
|
||||
auto runner_ctx = get_context();
|
||||
runner_ctx.ip_context = ip_context;
|
||||
runner_ctx.ip_scale = ip_scale;
|
||||
|
||||
ggml_tensor* out = unet.forward(&runner_ctx,
|
||||
x,
|
||||
@@ -762,14 +823,16 @@ struct UNetModelRunner : public DiffusionModelRunner {
|
||||
const sd::Tensor<float>& y = {},
|
||||
int num_video_frames = -1,
|
||||
const std::vector<sd::Tensor<float>>& controls = {},
|
||||
float control_strength = 0.f) {
|
||||
float control_strength = 0.f,
|
||||
const sd::Tensor<float>& ip_context = {},
|
||||
float ip_scale = 1.f) {
|
||||
// x: [N, in_channels, h, w]
|
||||
// timesteps: [N, ]
|
||||
// context: [N, max_position, hidden_size]([N, 77, 768]) or [1, max_position, hidden_size]
|
||||
// c_concat: [N, in_channels, h, w] or [1, in_channels, h, w]
|
||||
// y: [N, adm_in_channels] or [1, adm_in_channels]
|
||||
auto get_graph = [&]() -> ggml_cgraph* {
|
||||
return build_graph(x, timesteps, context, c_concat, y, num_video_frames, controls, control_strength);
|
||||
return build_graph(x, timesteps, context, c_concat, y, num_video_frames, controls, control_strength, ip_context, ip_scale);
|
||||
};
|
||||
|
||||
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
|
||||
@@ -789,7 +852,9 @@ struct UNetModelRunner : public DiffusionModelRunner {
|
||||
tensor_or_empty(diffusion_params.y),
|
||||
extra->num_video_frames,
|
||||
extra->controls ? *extra->controls : empty_controls,
|
||||
extra->control_strength);
|
||||
extra->control_strength,
|
||||
extra->ip_context ? *extra->ip_context : sd::Tensor<float>{},
|
||||
extra->ip_scale);
|
||||
}
|
||||
|
||||
void test() {
|
||||
|
||||
@@ -150,6 +150,8 @@ namespace ZImage {
|
||||
|
||||
if (sd_backend_is(ctx->backend, "ROCm")) {
|
||||
out_proj->set_scale(1.f / 16.f);
|
||||
out_proj->set_force_prec_f32(true);
|
||||
qkv_proj->set_force_prec_f32(true);
|
||||
}
|
||||
|
||||
auto qkv = qkv_proj->forward(ctx, x); // [N, n_token, (num_heads + num_kv_heads*2)*head_dim]
|
||||
@@ -227,7 +229,7 @@ namespace ZImage {
|
||||
auto w2 = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
|
||||
auto w3 = std::dynamic_pointer_cast<Linear>(blocks["w3"]);
|
||||
|
||||
if (sd_backend_is(ctx->backend, "Vulkan")) {
|
||||
if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) {
|
||||
w2->set_force_prec_f32(true);
|
||||
}
|
||||
|
||||
@@ -646,8 +648,8 @@ namespace ZImage {
|
||||
*diffusion_params.x,
|
||||
*diffusion_params.timesteps,
|
||||
tensor_or_empty(diffusion_params.context),
|
||||
diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents,
|
||||
diffusion_params.ref_index_mode);
|
||||
diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents,
|
||||
diffusion_params.ref_image_params.ref_index_mode);
|
||||
}
|
||||
|
||||
void test() {
|
||||
|
||||
@@ -200,7 +200,11 @@ namespace LLM {
|
||||
config.vision.in_channels = tensor_storage.ne[2];
|
||||
config.vision.hidden_size = tensor_storage.ne[3];
|
||||
}
|
||||
if (contains(name, "visual.patch_embed.bias")) {
|
||||
// HF-format checkpoints keep the patch embed unsplit under a single name.
|
||||
if (contains(name, "visual.patch_embed.proj.weight")) {
|
||||
config.vision.patch_size = static_cast<int>(tensor_storage.ne[0]);
|
||||
}
|
||||
if (contains(name, "visual.patch_embed.bias") || contains(name, "visual.patch_embed.proj.bias")) {
|
||||
config.vision.hidden_size = tensor_storage.ne[0];
|
||||
}
|
||||
if (contains(name, "visual.pos_embed.weight")) {
|
||||
@@ -285,7 +289,7 @@ namespace LLM {
|
||||
bool add_unit_offset = false)
|
||||
: hidden_size(hidden_size), eps(eps), add_unit_offset(add_unit_offset) {}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||
ggml_tensor* w = params["weight"];
|
||||
if (ctx->weight_adapter) {
|
||||
w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");
|
||||
@@ -1653,6 +1657,10 @@ namespace LLM {
|
||||
model.get_param_tensors(tensors, prefix);
|
||||
}
|
||||
|
||||
void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) {
|
||||
model.get_param_tensor_ops(tensor_ops);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* input_ids,
|
||||
ggml_tensor* input_pos,
|
||||
|
||||
@@ -18,6 +18,7 @@
|
||||
struct T5Config {
|
||||
int64_t num_layers = 24;
|
||||
int64_t model_dim = 4096;
|
||||
int64_t inner_dim = 4096;
|
||||
int64_t ff_dim = 10240;
|
||||
int64_t num_heads = 64;
|
||||
int64_t vocab_size = 32128;
|
||||
@@ -53,6 +54,7 @@ struct T5Config {
|
||||
if (q->n_dims == 2) {
|
||||
config.model_dim = q->ne[0];
|
||||
int64_t inner_dim = q->ne[1];
|
||||
config.inner_dim = inner_dim;
|
||||
// Flan-T5/T5 uses d_kv=64 for common sizes.
|
||||
if (inner_dim % 64 == 0) {
|
||||
config.num_heads = inner_dim / 64;
|
||||
@@ -357,7 +359,7 @@ public:
|
||||
: config(config) {
|
||||
blocks["encoder"] = std::shared_ptr<GGMLBlock>(new T5Stack(config.num_layers,
|
||||
config.model_dim,
|
||||
config.model_dim,
|
||||
config.inner_dim,
|
||||
config.ff_dim,
|
||||
config.num_heads,
|
||||
config.relative_attention));
|
||||
|
||||
@@ -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__
|
||||
@@ -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__
|
||||
@@ -528,6 +528,9 @@ public:
|
||||
if (version == VERSION_WAN2_2_TI2V) {
|
||||
z_channels = 48;
|
||||
patch = 2;
|
||||
} else if (sd_version_is_hunyuan_video(version)) {
|
||||
z_channels = 32;
|
||||
patch = 2;
|
||||
} else if (sd_version_is_ltxav(version)) {
|
||||
z_channels = 128;
|
||||
patch = 4;
|
||||
@@ -542,12 +545,12 @@ public:
|
||||
|
||||
ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* z) {
|
||||
auto decoder = std::dynamic_pointer_cast<TinyVideoDecoder>(blocks["decoder"]);
|
||||
if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) {
|
||||
if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) {
|
||||
// (W, H, C, T) -> (W, H, T, C)
|
||||
z = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, z, 0, 1, 3, 2));
|
||||
}
|
||||
auto result = decoder->forward(ctx, z);
|
||||
if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) {
|
||||
if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) {
|
||||
// (W, H, T, C) -> (W, H, C, T)
|
||||
result = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, result, 0, 1, 3, 2));
|
||||
}
|
||||
@@ -556,7 +559,7 @@ public:
|
||||
|
||||
ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
auto encoder = std::dynamic_pointer_cast<TinyVideoEncoder>(blocks["encoder"]);
|
||||
if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) {
|
||||
if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) {
|
||||
// (W, H, T, C) -> (W, H, C, T)
|
||||
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
|
||||
}
|
||||
@@ -569,7 +572,7 @@ public:
|
||||
}
|
||||
}
|
||||
x = encoder->forward(ctx, x);
|
||||
if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) {
|
||||
if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) {
|
||||
// (W, H, C, T) -> (W, H, T, C)
|
||||
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
|
||||
}
|
||||
|
||||
@@ -74,7 +74,7 @@ public:
|
||||
int scale_factor = 8;
|
||||
if (version == VERSION_LTXAV) {
|
||||
scale_factor = 32;
|
||||
} else if (version == VERSION_WAN2_2_TI2V) {
|
||||
} else if (version == VERSION_WAN2_2_TI2V || sd_version_is_hunyuan_video(version) || sd_version_is_mage_flow(version)) {
|
||||
scale_factor = 16;
|
||||
} else if (sd_version_uses_flux2_vae(version)) {
|
||||
scale_factor = 16;
|
||||
@@ -136,7 +136,7 @@ public:
|
||||
// Image VAE encode is more sensitive to tile boundary context than decode.
|
||||
// Keep the smaller legacy factor for video VAEs, but default image encode
|
||||
// tiles to 64 latent pixels so a 512px SD image is encoded as one tile.
|
||||
const float encode_tile_factor = (sd_version_is_wan(version) || sd_version_is_ltxav(version)) ? 1.30539f : 2.0f;
|
||||
const float encode_tile_factor = (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) ? 1.30539f : 2.0f;
|
||||
get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, W, H, encode_tile_factor);
|
||||
LOG_DEBUG("VAE Tile size: %dx%d", tile_size_x, tile_size_y);
|
||||
output = tiled_compute(input,
|
||||
|
||||
@@ -21,7 +21,7 @@ bool write_gguf_file(const std::string& file_path,
|
||||
class GGUFStreamingWriter : public StreamingModelWriter {
|
||||
public:
|
||||
GGUFStreamingWriter() = default;
|
||||
~GGUFStreamingWriter();
|
||||
~GGUFStreamingWriter() override;
|
||||
|
||||
bool write_metadata(const std::string& file_path,
|
||||
const std::vector<TensorWritePlan>& tensors,
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <limits>
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
#include <utility>
|
||||
@@ -512,8 +513,51 @@ static bool parse_storage_type(const std::string& global_name, PickleStorageInfo
|
||||
return false;
|
||||
}
|
||||
|
||||
static bool tensor_is_contiguous(const PickleTensorInfo& tensor) {
|
||||
if (tensor.tensor_storage.nelements() == 0) {
|
||||
static bool checked_pickle_byte_count(int64_t element_count,
|
||||
uint64_t element_nbytes,
|
||||
uint64_t* byte_count) {
|
||||
if (element_count < 0 || element_nbytes == 0) {
|
||||
return false;
|
||||
}
|
||||
|
||||
uint64_t count = static_cast<uint64_t>(element_count);
|
||||
if (count > std::numeric_limits<uint64_t>::max() / element_nbytes) {
|
||||
return false;
|
||||
}
|
||||
|
||||
*byte_count = count * element_nbytes;
|
||||
return true;
|
||||
}
|
||||
|
||||
static bool tensor_layout_is_valid(const PickleTensorInfo& tensor, uint64_t raw_element_nbytes) {
|
||||
if (raw_element_nbytes == 0) {
|
||||
return false;
|
||||
}
|
||||
|
||||
bool has_zero_dimension = false;
|
||||
uint64_t element_count = 1;
|
||||
for (int i = 0; i < tensor.tensor_storage.n_dims; ++i) {
|
||||
int64_t dimension = tensor.tensor_storage.ne[i];
|
||||
if (dimension < 0) {
|
||||
return false;
|
||||
}
|
||||
if (dimension == 0) {
|
||||
has_zero_dimension = true;
|
||||
continue;
|
||||
}
|
||||
|
||||
uint64_t size = static_cast<uint64_t>(dimension);
|
||||
if (element_count > static_cast<uint64_t>(std::numeric_limits<int64_t>::max()) / size) {
|
||||
return false;
|
||||
}
|
||||
element_count *= size;
|
||||
}
|
||||
|
||||
if (!has_zero_dimension &&
|
||||
element_count > static_cast<uint64_t>(std::numeric_limits<int64_t>::max()) / raw_element_nbytes) {
|
||||
return false;
|
||||
}
|
||||
if (has_zero_dimension) {
|
||||
return true;
|
||||
}
|
||||
if (tensor.stride_n_dims != tensor.tensor_storage.n_dims) {
|
||||
@@ -932,7 +976,12 @@ bool parse_torch_state_dict_pickle(const uint8_t* buffer,
|
||||
if (storage.key.empty() || !parse_storage_type(pid.items[1].str_value, &storage)) {
|
||||
return false;
|
||||
}
|
||||
storage.nbytes = (uint64_t)pid.items[4].int_value * storage.raw_element_nbytes;
|
||||
if (!checked_pickle_byte_count(pid.items[4].int_value,
|
||||
storage.raw_element_nbytes,
|
||||
&storage.nbytes)) {
|
||||
set_error(error, "invalid storage size in torch pickle");
|
||||
return false;
|
||||
}
|
||||
storage_nbytes[storage.key] = storage.nbytes;
|
||||
stack.push_back(make_storage_value(storage));
|
||||
} break;
|
||||
@@ -963,7 +1012,12 @@ bool parse_torch_state_dict_pickle(const uint8_t* buffer,
|
||||
tensor.tensor_storage.is_f64 = args.items[0].storage.is_f64;
|
||||
tensor.tensor_storage.is_i64 = args.items[0].storage.is_i64;
|
||||
tensor.tensor_storage.storage_key = args.items[0].storage.key;
|
||||
tensor.tensor_storage.offset = (uint64_t)args.items[1].int_value * args.items[0].storage.raw_element_nbytes;
|
||||
if (!checked_pickle_byte_count(args.items[1].int_value,
|
||||
args.items[0].storage.raw_element_nbytes,
|
||||
&tensor.tensor_storage.offset)) {
|
||||
set_error(error, "invalid tensor storage offset in torch pickle");
|
||||
return false;
|
||||
}
|
||||
|
||||
for (const auto& item : args.items[2].items) {
|
||||
if (item.kind != PickleValue::INT || tensor.tensor_storage.n_dims >= SD_MAX_DIMS) {
|
||||
@@ -979,7 +1033,8 @@ bool parse_torch_state_dict_pickle(const uint8_t* buffer,
|
||||
tensor.stride[tensor.stride_n_dims++] = item.int_value;
|
||||
}
|
||||
|
||||
if (!tensor_is_contiguous(tensor)) {
|
||||
if (!tensor_layout_is_valid(tensor, args.items[0].storage.raw_element_nbytes)) {
|
||||
set_error(error, "invalid tensor shape or stride in torch pickle");
|
||||
return false;
|
||||
}
|
||||
stack.push_back(make_tensor_value(tensor));
|
||||
|
||||
@@ -3,15 +3,19 @@
|
||||
#include <algorithm>
|
||||
#include <cstdint>
|
||||
#include <exception>
|
||||
#include <filesystem>
|
||||
#include <fstream>
|
||||
#include <ostream>
|
||||
#include <string>
|
||||
#include <unordered_set>
|
||||
#include <vector>
|
||||
|
||||
#include "binary_io.h"
|
||||
#include "core/util.h"
|
||||
#include "json.hpp"
|
||||
|
||||
namespace fs = std::filesystem;
|
||||
|
||||
static constexpr size_t ST_HEADER_SIZE_LEN = 8;
|
||||
|
||||
static void set_error(std::string* error, const std::string& message) {
|
||||
@@ -20,6 +24,14 @@ static void set_error(std::string* error, const std::string& message) {
|
||||
}
|
||||
}
|
||||
|
||||
static std::string resolve_index_shard_path(const std::string& index_path, const std::string& shard_path) {
|
||||
fs::path shard_fs_path(shard_path);
|
||||
if (shard_fs_path.is_absolute()) {
|
||||
return shard_fs_path.lexically_normal().string();
|
||||
}
|
||||
return (fs::path(index_path).parent_path() / shard_fs_path).lexically_normal().string();
|
||||
}
|
||||
|
||||
bool is_safetensors_file(const std::string& file_path) {
|
||||
std::ifstream file(file_path, std::ios::binary);
|
||||
if (!file.is_open()) {
|
||||
@@ -43,7 +55,7 @@ bool is_safetensors_file(const std::string& file_path) {
|
||||
}
|
||||
|
||||
size_t header_size_ = model_io::read_u64(header_size_buf);
|
||||
if (header_size_ >= file_size_ || header_size_ <= 2) {
|
||||
if (header_size_ > file_size_ - ST_HEADER_SIZE_LEN || header_size_ <= 2) {
|
||||
return false;
|
||||
}
|
||||
|
||||
@@ -88,7 +100,8 @@ static ggml_type safetensors_dtype_to_ggml_type(const std::string& dtype) {
|
||||
// https://huggingface.co/docs/safetensors/index
|
||||
bool read_safetensors_file(const std::string& file_path,
|
||||
std::vector<TensorStorage>& tensor_storages,
|
||||
std::string* error) {
|
||||
std::string* error,
|
||||
std::map<std::string, std::string>* metadata) {
|
||||
std::ifstream file(file_path, std::ios::binary);
|
||||
if (!file.is_open()) {
|
||||
set_error(error, "failed to open '" + file_path + "'");
|
||||
@@ -114,10 +127,11 @@ bool read_safetensors_file(const std::string& file_path,
|
||||
}
|
||||
|
||||
size_t header_size_ = model_io::read_u64(header_size_buf);
|
||||
if (header_size_ >= file_size_) {
|
||||
if (header_size_ > file_size_ - ST_HEADER_SIZE_LEN) {
|
||||
set_error(error, "invalid safetensor file '" + file_path + "'");
|
||||
return false;
|
||||
}
|
||||
const size_t data_start = ST_HEADER_SIZE_LEN + header_size_;
|
||||
|
||||
// read header
|
||||
std::vector<char> header_buf;
|
||||
@@ -137,6 +151,18 @@ bool read_safetensors_file(const std::string& file_path,
|
||||
return false;
|
||||
}
|
||||
|
||||
if (metadata != nullptr) {
|
||||
metadata->clear();
|
||||
auto metadata_item = header_.find("__metadata__");
|
||||
if (metadata_item != header_.end() && metadata_item->is_object()) {
|
||||
for (const auto& item : metadata_item->items()) {
|
||||
if (item.value().is_string()) {
|
||||
metadata->emplace(item.key(), item.value().get<std::string>());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
tensor_storages.clear();
|
||||
for (auto& item : header_.items()) {
|
||||
std::string name = item.key();
|
||||
@@ -156,6 +182,10 @@ bool read_safetensors_file(const std::string& file_path,
|
||||
|
||||
size_t begin = tensor_info["data_offsets"][0].get<size_t>();
|
||||
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;
|
||||
}
|
||||
|
||||
ggml_type type = safetensors_dtype_to_ggml_type(dtype);
|
||||
if (type == GGML_TYPE_COUNT) {
|
||||
@@ -187,7 +217,7 @@ bool read_safetensors_file(const std::string& file_path,
|
||||
n_dims = 1;
|
||||
}
|
||||
|
||||
TensorStorage tensor_storage(name, type, ne, n_dims, 0, ST_HEADER_SIZE_LEN + header_size_ + begin);
|
||||
TensorStorage tensor_storage(name, type, ne, n_dims, 0, data_start + begin);
|
||||
tensor_storage.reverse_ne();
|
||||
|
||||
size_t tensor_data_size = end - begin;
|
||||
@@ -225,6 +255,52 @@ bool read_safetensors_file(const std::string& file_path,
|
||||
return true;
|
||||
}
|
||||
|
||||
bool read_safetensors_index_file(const std::string& file_path,
|
||||
std::vector<std::string>& shard_paths,
|
||||
std::string* error) {
|
||||
shard_paths.clear();
|
||||
|
||||
std::ifstream file(file_path);
|
||||
if (!file.is_open()) {
|
||||
set_error(error, "failed to open '" + file_path + "'");
|
||||
return false;
|
||||
}
|
||||
|
||||
nlohmann::json index;
|
||||
try {
|
||||
index = nlohmann::json::parse(file);
|
||||
} catch (const std::exception&) {
|
||||
set_error(error, "parsing safetensors index failed: '" + file_path + "'");
|
||||
return false;
|
||||
}
|
||||
|
||||
if (!index.is_object() || !index.contains("weight_map") || !index["weight_map"].is_object()) {
|
||||
set_error(error, "invalid safetensors index '" + file_path + "'");
|
||||
return false;
|
||||
}
|
||||
|
||||
std::unordered_set<std::string> seen_shard_paths;
|
||||
for (const auto& item : index["weight_map"].items()) {
|
||||
if (!item.value().is_string()) {
|
||||
set_error(error, "invalid shard path for tensor '" + item.key() + "'");
|
||||
return false;
|
||||
}
|
||||
|
||||
std::string shard_path = resolve_index_shard_path(file_path,
|
||||
item.value().get<std::string>());
|
||||
if (seen_shard_paths.insert(shard_path).second) {
|
||||
shard_paths.push_back(std::move(shard_path));
|
||||
}
|
||||
}
|
||||
|
||||
if (shard_paths.empty()) {
|
||||
set_error(error, "safetensors index has no tensors: '" + file_path + "'");
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
static bool ggml_type_to_safetensors_dtype(ggml_type type, std::string* dtype) {
|
||||
switch (type) {
|
||||
case GGML_TYPE_F16:
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
#ifndef __SD_MODEL_IO_SAFETENSORS_IO_H__
|
||||
#define __SD_MODEL_IO_SAFETENSORS_IO_H__
|
||||
|
||||
#include <map>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
@@ -10,7 +11,11 @@
|
||||
bool is_safetensors_file(const std::string& file_path);
|
||||
bool read_safetensors_file(const std::string& file_path,
|
||||
std::vector<TensorStorage>& tensor_storages,
|
||||
std::string* error = nullptr);
|
||||
std::string* error = nullptr,
|
||||
std::map<std::string, std::string>* metadata = nullptr);
|
||||
bool read_safetensors_index_file(const std::string& file_path,
|
||||
std::vector<std::string>& shard_paths,
|
||||
std::string* error = nullptr);
|
||||
bool write_safetensors_file(const std::string& file_path,
|
||||
const std::vector<TensorWriteInfo>& tensors,
|
||||
std::string* error = nullptr);
|
||||
|
||||
@@ -139,11 +139,16 @@ bool read_torch_legacy_file(const std::string& file_path,
|
||||
if (it == legacy_storage_map.end()) {
|
||||
return false;
|
||||
}
|
||||
if (current_offset + LEGACY_STORAGE_HEADER_SIZE + it->second > file_size) {
|
||||
if (current_offset > file_size ||
|
||||
LEGACY_STORAGE_HEADER_SIZE > file_size - current_offset) {
|
||||
return false;
|
||||
}
|
||||
storage_offsets[storage_key] = current_offset + LEGACY_STORAGE_HEADER_SIZE;
|
||||
current_offset += LEGACY_STORAGE_HEADER_SIZE + it->second;
|
||||
uint64_t storage_offset = current_offset + LEGACY_STORAGE_HEADER_SIZE;
|
||||
if (it->second > file_size - storage_offset) {
|
||||
return false;
|
||||
}
|
||||
storage_offsets[storage_key] = storage_offset;
|
||||
current_offset = storage_offset + it->second;
|
||||
}
|
||||
|
||||
for (auto& tensor_storage : tensor_storages) {
|
||||
@@ -159,8 +164,10 @@ bool read_torch_legacy_file(const std::string& file_path,
|
||||
|
||||
uint64_t base_offset = it_offset->second;
|
||||
uint64_t storage_nbytes = it_size->second;
|
||||
uint64_t tensor_nbytes = tensor_storage.nbytes_to_read();
|
||||
if (tensor_storage.offset + tensor_nbytes > storage_nbytes) {
|
||||
int64_t tensor_nbytes = tensor_storage.nbytes_to_read();
|
||||
if (tensor_nbytes < 0 ||
|
||||
tensor_storage.offset > storage_nbytes ||
|
||||
static_cast<uint64_t>(tensor_nbytes) > storage_nbytes - tensor_storage.offset) {
|
||||
return false;
|
||||
}
|
||||
|
||||
|
||||
@@ -76,8 +76,10 @@ static bool parse_zip_data_pkl(const uint8_t* buffer,
|
||||
return false;
|
||||
}
|
||||
|
||||
uint64_t tensor_nbytes = tensor_storage.nbytes_to_read();
|
||||
if (tensor_storage.offset + tensor_nbytes > entry_size) {
|
||||
int64_t tensor_nbytes = tensor_storage.nbytes_to_read();
|
||||
if (tensor_nbytes < 0 ||
|
||||
tensor_storage.offset > entry_size ||
|
||||
static_cast<uint64_t>(tensor_nbytes) > entry_size - tensor_storage.offset) {
|
||||
set_error(error, "tensor '" + tensor_storage.name + "' exceeds storage entry '" + entry_name + "'");
|
||||
return false;
|
||||
}
|
||||
|
||||
@@ -235,6 +235,9 @@ bool ModelLoader::init_from_file(const std::string& file_path, const std::string
|
||||
} else if (is_gguf_file(file_path)) {
|
||||
LOG_INFO("load %s using gguf format", file_path.c_str());
|
||||
return init_from_gguf_file(file_path, prefix);
|
||||
} else if (ends_with(file_path, ".json")) {
|
||||
LOG_INFO("load %s using safetensors index format", file_path.c_str());
|
||||
return init_from_safetensors_index_file(file_path, prefix);
|
||||
} else if (is_safetensors_file(file_path)) {
|
||||
LOG_INFO("load %s using safetensors format", file_path.c_str());
|
||||
return init_from_safetensors_file(file_path, prefix);
|
||||
@@ -314,7 +317,7 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const
|
||||
|
||||
std::vector<TensorStorage> tensor_storages;
|
||||
std::string error;
|
||||
if (!read_safetensors_file(file_path, tensor_storages, &error)) {
|
||||
if (!read_safetensors_file(file_path, tensor_storages, &error, &metadata_)) {
|
||||
LOG_ERROR("%s", error.c_str());
|
||||
return false;
|
||||
}
|
||||
@@ -339,6 +342,25 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const
|
||||
return true;
|
||||
}
|
||||
|
||||
bool ModelLoader::init_from_safetensors_index_file(const std::string& file_path, const std::string& prefix) {
|
||||
LOG_DEBUG("init from safetensors index '%s', prefix = '%s'", file_path.c_str(), prefix.c_str());
|
||||
|
||||
std::vector<std::string> shard_paths;
|
||||
std::string error;
|
||||
if (!read_safetensors_index_file(file_path, shard_paths, &error)) {
|
||||
LOG_ERROR("%s", error.c_str());
|
||||
return false;
|
||||
}
|
||||
|
||||
for (const std::string& shard_path : shard_paths) {
|
||||
if (!init_from_file(shard_path, prefix)) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
/*================================================= TorchLegacyModelLoader ==================================================*/
|
||||
|
||||
bool ModelLoader::init_from_torch_legacy_file(const std::string& file_path, const std::string& prefix) {
|
||||
@@ -476,11 +498,18 @@ SDVersion ModelLoader::get_sd_version() {
|
||||
return VERSION_MINIT2I;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.img_mod.1.weight") != std::string::npos) {
|
||||
auto img_in = tensor_storage_map.find("model.diffusion_model.img_in.weight");
|
||||
if (img_in != tensor_storage_map.end() && img_in->second.ne[0] == 128) {
|
||||
return VERSION_MAGE_FLOW;
|
||||
}
|
||||
if (tensor_storage_map.find("model.diffusion_model.time_text_embed.addition_t_embedding.weight") != tensor_storage_map.end()) {
|
||||
return VERSION_QWEN_IMAGE_LAYERED;
|
||||
}
|
||||
return VERSION_QWEN_IMAGE;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.txt_in.individual_token_refiner.blocks.0.adaLN_modulation.1.weight") != std::string::npos) {
|
||||
return VERSION_HUNYUAN_VIDEO;
|
||||
}
|
||||
if (tensor_storage.name.find("llm_adapter.blocks.0.cross_attn.q_proj.weight") != std::string::npos) {
|
||||
return VERSION_ANIMA;
|
||||
}
|
||||
@@ -511,6 +540,9 @@ SDVersion ModelLoader::get_sd_version() {
|
||||
if (tensor_storage.name.find("model.diffusion_model.blocks.0.cross_attn.norm_k.weight") != std::string::npos) {
|
||||
is_wan = true;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.patch_embedder.weight") != std::string::npos) {
|
||||
return VERSION_LINGBOT_VIDEO;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.patch_embedding.weight") != std::string::npos) {
|
||||
patch_embedding_channels = tensor_storage.ne[3];
|
||||
}
|
||||
@@ -1021,7 +1053,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
|
||||
std::atomic<size_t> tensor_idx(0);
|
||||
std::atomic<bool> failed(false);
|
||||
std::vector<std::thread> workers;
|
||||
std::mutex rpc_backend_mutex;
|
||||
std::mutex backend_tensor_set_mutex;
|
||||
|
||||
for (int i = 0; i < n_threads; ++i) {
|
||||
workers.emplace_back([&, file_path, is_zip]() {
|
||||
@@ -1045,6 +1077,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
|
||||
|
||||
std::vector<uint8_t> read_buffer;
|
||||
std::vector<uint8_t> convert_buffer;
|
||||
std::vector<uint8_t> zip_entry_buffer;
|
||||
|
||||
while (true) {
|
||||
int64_t t0, t1;
|
||||
@@ -1083,34 +1116,60 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
|
||||
|
||||
size_t nbytes_to_read = tensor_storage.nbytes_to_read();
|
||||
|
||||
auto read_data = [&](char* buf, size_t n) {
|
||||
auto read_data = [&](char* buf, size_t n) -> bool {
|
||||
if (zip != nullptr) {
|
||||
zip_entry_openbyindex(zip, tensor_storage.index_in_zip);
|
||||
if (zip_entry_openbyindex(zip, tensor_storage.index_in_zip) != 0) {
|
||||
LOG_ERROR("failed to open zip entry for tensor '%s'", tensor_storage.name.c_str());
|
||||
return false;
|
||||
}
|
||||
size_t entry_size = zip_entry_size(zip);
|
||||
if (tensor_storage.offset > entry_size) {
|
||||
LOG_ERROR("tensor '%s' exceeds its zip storage entry", tensor_storage.name.c_str());
|
||||
zip_entry_close(zip);
|
||||
return false;
|
||||
}
|
||||
size_t tensor_offset = static_cast<size_t>(tensor_storage.offset);
|
||||
if (n > entry_size - tensor_offset) {
|
||||
LOG_ERROR("tensor '%s' exceeds its zip storage entry", tensor_storage.name.c_str());
|
||||
zip_entry_close(zip);
|
||||
return false;
|
||||
}
|
||||
|
||||
if (entry_size != n) {
|
||||
int64_t t_memcpy_start;
|
||||
read_buffer.resize(entry_size);
|
||||
zip_entry_noallocread(zip, (void*)read_buffer.data(), entry_size);
|
||||
zip_entry_buffer.resize(entry_size);
|
||||
auto bytes_read = zip_entry_noallocread(zip, (void*)zip_entry_buffer.data(), entry_size);
|
||||
if (bytes_read < 0 || static_cast<size_t>(bytes_read) != entry_size) {
|
||||
LOG_ERROR("failed to read zip entry for tensor '%s'", tensor_storage.name.c_str());
|
||||
zip_entry_close(zip);
|
||||
return false;
|
||||
}
|
||||
t_memcpy_start = ggml_time_ms();
|
||||
memcpy((void*)buf, (void*)(read_buffer.data() + tensor_storage.offset), n);
|
||||
memcpy((void*)buf, (void*)(zip_entry_buffer.data() + tensor_offset), n);
|
||||
memcpy_time_ms.fetch_add(ggml_time_ms() - t_memcpy_start);
|
||||
} else {
|
||||
zip_entry_noallocread(zip, (void*)buf, n);
|
||||
auto bytes_read = zip_entry_noallocread(zip, (void*)buf, n);
|
||||
if (bytes_read < 0 || static_cast<size_t>(bytes_read) != n) {
|
||||
LOG_ERROR("failed to read zip entry for tensor '%s'", tensor_storage.name.c_str());
|
||||
zip_entry_close(zip);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
zip_entry_close(zip);
|
||||
} else if (mmapped) {
|
||||
if (!mmapped->copy_data(buf, n, tensor_storage.offset)) {
|
||||
LOG_ERROR("read tensor data failed: '%s'", file_path.c_str());
|
||||
failed = true;
|
||||
return false;
|
||||
}
|
||||
} else {
|
||||
file.seekg(tensor_storage.offset);
|
||||
file.read(buf, n);
|
||||
if (!file) {
|
||||
LOG_ERROR("read tensor data failed: '%s'", file_path.c_str());
|
||||
failed = true;
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
};
|
||||
|
||||
char* read_buf = nullptr;
|
||||
@@ -1144,7 +1203,10 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
|
||||
}
|
||||
|
||||
t0 = ggml_time_ms();
|
||||
read_data(read_buf, nbytes_to_read);
|
||||
if (!read_data(read_buf, nbytes_to_read)) {
|
||||
failed = true;
|
||||
break;
|
||||
}
|
||||
t1 = ggml_time_ms();
|
||||
read_time_ms.fetch_add(t1 - t0);
|
||||
|
||||
@@ -1182,17 +1244,8 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
|
||||
if (dst_tensor->buffer != nullptr && !ggml_backend_buffer_is_host(dst_tensor->buffer)) {
|
||||
t0 = ggml_time_ms();
|
||||
|
||||
// RPC backends require serialized access to prevent concurrency issues
|
||||
const char* buffer_type_name = ggml_backend_buft_name(ggml_backend_buffer_get_type(dst_tensor->buffer));
|
||||
bool is_rpc_buffer = buffer_type_name != nullptr &&
|
||||
std::string(buffer_type_name).find("RPC") != std::string::npos;
|
||||
|
||||
if (is_rpc_buffer) {
|
||||
std::lock_guard<std::mutex> lock(rpc_backend_mutex);
|
||||
ggml_backend_tensor_set(dst_tensor, convert_buf, 0, ggml_nbytes(dst_tensor));
|
||||
} else {
|
||||
ggml_backend_tensor_set(dst_tensor, convert_buf, 0, ggml_nbytes(dst_tensor));
|
||||
}
|
||||
std::lock_guard<std::mutex> lock(backend_tensor_set_mutex);
|
||||
ggml_backend_tensor_set(dst_tensor, convert_buf, 0, ggml_nbytes(dst_tensor));
|
||||
|
||||
t1 = ggml_time_ms();
|
||||
copy_to_backend_time_ms.fetch_add(t1 - t0);
|
||||
|
||||