mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-10-02 10:24:37 +08:00
Compare commits
32
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
92dc7268fc | ||
|
|
07b2b18e70 | ||
|
|
1ceb5bd9df | ||
|
|
202c6154a2 | ||
|
|
a397e03488 | ||
|
|
72e512a0cc | ||
|
|
0baf721215 | ||
|
|
645e6e9089 | ||
|
|
cbf92191c3 | ||
|
|
8cf55a3b3b | ||
|
|
3a8788cb7d | ||
|
|
449165caf5 | ||
|
|
adaa599a3b | ||
|
|
2e3514625a | ||
|
|
47d8198b69 | ||
|
|
ef92a0027e | ||
|
|
b3374e6a71 | ||
|
|
bdd937f29a | ||
|
|
c51ec7cad9 | ||
|
|
5b0267e941 | ||
|
|
0045a72b96 | ||
|
|
99bd062546 | ||
|
|
9d8c9e4279 | ||
|
|
caa823a8c0 | ||
|
|
22c8c40b0d | ||
|
|
b706d682ad | ||
|
|
b758b7de13 | ||
|
|
f683c88a28 | ||
|
|
21fd4e6788 | ||
|
|
830804262b | ||
|
|
82e03ef137 | ||
|
|
baf7eda1e4 |
@@ -14,6 +14,8 @@ on:
|
||||
paths:
|
||||
[
|
||||
".github/workflows/**",
|
||||
".dockerignore",
|
||||
"Dockerfile*",
|
||||
"**/CMakeLists.txt",
|
||||
"**/Makefile",
|
||||
"**/*.h",
|
||||
@@ -29,6 +31,8 @@ on:
|
||||
paths:
|
||||
[
|
||||
".github/workflows/**",
|
||||
".dockerignore",
|
||||
"Dockerfile*",
|
||||
"**/CMakeLists.txt",
|
||||
"**/Makefile",
|
||||
"**/*.h",
|
||||
@@ -177,7 +181,7 @@ jobs:
|
||||
build-and-push-docker-images:
|
||||
name: Build and push container images
|
||||
if: ${{ github.event_name != 'pull_request' }}
|
||||
runs-on: ubuntu-latest
|
||||
runs-on: ${{ matrix.runner }}
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
@@ -189,6 +193,20 @@ jobs:
|
||||
strategy:
|
||||
matrix:
|
||||
variant: [musa, sycl, vulkan, cuda]
|
||||
platform: [linux/amd64]
|
||||
runner: [ubuntu-latest]
|
||||
build-args: [""]
|
||||
tag-suffix: [""]
|
||||
include:
|
||||
- variant: cuda
|
||||
platform: linux/arm64
|
||||
runner: ubuntu-24.04-arm
|
||||
tag-suffix: "-spark"
|
||||
build-args: |
|
||||
CUDA_VERSION=13.0.0
|
||||
UBUNTU_VERSION=24.04
|
||||
CUDA_ARCHITECTURES=121
|
||||
GGML_CUDA_FA_ALL_QUANTS=ON
|
||||
|
||||
env:
|
||||
REGISTRY: ghcr.io
|
||||
@@ -243,12 +261,13 @@ jobs:
|
||||
uses: docker/build-push-action@v6
|
||||
with:
|
||||
context: .
|
||||
platforms: linux/amd64
|
||||
platforms: ${{ matrix.platform }}
|
||||
push: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
|
||||
file: Dockerfile.${{ matrix.variant }}
|
||||
tags: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:${{ env.BRANCH_NAME }}-${{ matrix.variant }}
|
||||
tags: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:${{ env.BRANCH_NAME }}-${{ matrix.variant }}${{ matrix.tag-suffix }}
|
||||
labels: ${{ steps.meta.outputs.labels }}
|
||||
annotations: ${{ steps.meta.outputs.annotations }}
|
||||
build-args: ${{ matrix.build-args }}
|
||||
|
||||
macOS-latest-cmake:
|
||||
runs-on: macos-latest
|
||||
@@ -449,13 +468,23 @@ jobs:
|
||||
|
||||
env:
|
||||
ROCM_VERSION: "7.13.0"
|
||||
GPU_TARGETS: "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1102;gfx1150;gfx1151;gfx1200;gfx1201"
|
||||
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"
|
||||
|
||||
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
|
||||
@@ -516,6 +545,30 @@ 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
|
||||
@@ -648,7 +701,7 @@ jobs:
|
||||
gpu_targets: "gfx908;gfx90a;gfx942;gfx1030;gfx1031;gfx1032;gfx1100;gfx1101;gfx1102;gfx1151;gfx1150;gfx1200;gfx1201"
|
||||
build: 'x64'
|
||||
- ROCM_VERSION: "7.13.0"
|
||||
gpu_targets: "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1100;gfx1101;gfx1102;gfx1150;gfx1151;gfx1200;gfx1201"
|
||||
gpu_targets: "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1150;gfx1151;gfx1152;gfx1200;gfx1201"
|
||||
build: x64
|
||||
|
||||
steps:
|
||||
|
||||
@@ -0,0 +1,55 @@
|
||||
name: Close inactive PRs
|
||||
|
||||
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 PRs"
|
||||
required: false
|
||||
default: false
|
||||
type: boolean
|
||||
|
||||
permissions:
|
||||
issues: write
|
||||
pull-requests: write
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}
|
||||
cancel-in-progress: false
|
||||
|
||||
jobs:
|
||||
stale-prs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Mark and close inactive PRs
|
||||
uses: actions/stale@v10
|
||||
with:
|
||||
days-before-issue-stale: -1
|
||||
days-before-issue-close: -1
|
||||
|
||||
days-before-pr-stale: 365
|
||||
days-before-pr-close: 7
|
||||
|
||||
stale-pr-label: pr:inactive
|
||||
close-pr-label: pr:auto-closed
|
||||
exempt-pr-labels: pr:keep-open
|
||||
|
||||
stale-pr-message: >
|
||||
This PR has been inactive for 365 days. If there is no new activity
|
||||
within 7 days, it will be closed automatically. Comment, push new
|
||||
commits, or remove the pr:inactive label to keep it open. Add
|
||||
pr:keep-open to exempt it from future inactive PR cleanup.
|
||||
|
||||
close-pr-message: >
|
||||
Closing this PR because it has had no activity for 7 days after
|
||||
being marked inactive. If this is still useful or ready to move
|
||||
forward, feel free to reopen it with fresh context or updated
|
||||
details. Sorry for any inconvenience.
|
||||
|
||||
remove-pr-stale-when-updated: true
|
||||
delete-branch: false
|
||||
operations-per-run: 100
|
||||
debug-only: ${{ github.event_name == 'workflow_dispatch' && inputs.debug_only || false }}
|
||||
@@ -22,6 +22,31 @@ endif()
|
||||
set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin)
|
||||
set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin)
|
||||
|
||||
if(APPLE)
|
||||
function(sd_set_macos_rpaths target)
|
||||
get_target_property(target_type ${target} TYPE)
|
||||
if(target_type STREQUAL "EXECUTABLE")
|
||||
set(runtime_paths "@executable_path" "@executable_path/../lib")
|
||||
elseif(target_type STREQUAL "SHARED_LIBRARY" OR target_type STREQUAL "MODULE_LIBRARY")
|
||||
set(runtime_paths "@loader_path" "@loader_path/../lib")
|
||||
set_target_properties(${target} PROPERTIES
|
||||
MACOSX_RPATH ON
|
||||
INSTALL_NAME_DIR "@rpath"
|
||||
BUILD_WITH_INSTALL_NAME_DIR ON
|
||||
)
|
||||
else()
|
||||
return()
|
||||
endif()
|
||||
|
||||
# Release artifacts zip the build output directly, so keep macOS rpaths relocatable.
|
||||
set_target_properties(${target} PROPERTIES
|
||||
BUILD_RPATH "${runtime_paths}"
|
||||
INSTALL_RPATH "${runtime_paths}"
|
||||
BUILD_WITH_INSTALL_RPATH ON
|
||||
)
|
||||
endfunction()
|
||||
endif()
|
||||
|
||||
if(CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR)
|
||||
set(SD_STANDALONE ON)
|
||||
else()
|
||||
@@ -237,6 +262,10 @@ else()
|
||||
add_library(${SD_LIB} STATIC ${SD_LIB_SOURCES})
|
||||
endif()
|
||||
|
||||
if(APPLE)
|
||||
sd_set_macos_rpaths(${SD_LIB})
|
||||
endif()
|
||||
|
||||
if(SD_SYCL)
|
||||
message("-- Use SYCL as backend stable-diffusion")
|
||||
set(GGML_SYCL ON)
|
||||
|
||||
+7
-1
@@ -10,7 +10,13 @@ WORKDIR /sd.cpp
|
||||
COPY . .
|
||||
|
||||
ARG CUDACXX=/usr/local/cuda/bin/nvcc
|
||||
RUN cmake . -B ./build -DSD_CUDA=ON
|
||||
ARG CUDA_ARCHITECTURES=""
|
||||
ARG GGML_CUDA_FA_ALL_QUANTS=""
|
||||
|
||||
RUN cmake . -B ./build \
|
||||
-DSD_CUDA=ON \
|
||||
${CUDA_ARCHITECTURES:+-DCMAKE_CUDA_ARCHITECTURES="${CUDA_ARCHITECTURES}"} \
|
||||
${GGML_CUDA_FA_ALL_QUANTS:+-DGGML_CUDA_FA_ALL_QUANTS=${GGML_CUDA_FA_ALL_QUANTS}}
|
||||
RUN cmake --build ./build --config Release -j$(nproc)
|
||||
|
||||
FROM nvidia/cuda:${CUDA_VERSION}-cudnn-runtime-ubuntu${UBUNTU_VERSION} AS runtime
|
||||
|
||||
+2
-1
@@ -1,4 +1,5 @@
|
||||
ARG SYCL_VERSION=2025.1.0-0
|
||||
# ggml SYCL hardware detection uses BMG G31/WCL architecture enums added in oneAPI 2025.3.
|
||||
ARG SYCL_VERSION=2025.3.2-0
|
||||
|
||||
FROM intel/oneapi-basekit:${SYCL_VERSION}-devel-ubuntu24.04 AS build
|
||||
|
||||
|
||||
@@ -15,29 +15,16 @@ API and command-line option may change frequently.***
|
||||
|
||||
## 🔥Important News
|
||||
|
||||
* **2026/05/27** 🚀 stable-diffusion.cpp now supports **Lens**
|
||||
* **2026/05/17** 🚀 stable-diffusion.cpp now supports **LTX-2.3**
|
||||
* **2026/04/11** 🚀 stable-diffusion.cpp now uses a brand-new embedded web UI.
|
||||
👉 Details: [PR #1408](https://github.com/leejet/stable-diffusion.cpp/pull/1408)
|
||||
|
||||
* **2026/01/18** 🚀 stable-diffusion.cpp now supports **FLUX.2-klein**
|
||||
👉 Details: [PR #1193](https://github.com/leejet/stable-diffusion.cpp/pull/1193)
|
||||
|
||||
* **2025/12/01** 🚀 stable-diffusion.cpp now supports **Z-Image**
|
||||
👉 Details: [PR #1020](https://github.com/leejet/stable-diffusion.cpp/pull/1020)
|
||||
|
||||
* **2025/11/30** 🚀 stable-diffusion.cpp now supports **FLUX.2-dev**
|
||||
👉 Details: [PR #1016](https://github.com/leejet/stable-diffusion.cpp/pull/1016)
|
||||
|
||||
* **2025/10/13** 🚀 stable-diffusion.cpp now supports **Qwen-Image-Edit / Qwen-Image-Edit 2509**
|
||||
👉 Details: [PR #877](https://github.com/leejet/stable-diffusion.cpp/pull/877)
|
||||
|
||||
* **2025/10/12** 🚀 stable-diffusion.cpp now supports **Qwen-Image**
|
||||
👉 Details: [PR #851](https://github.com/leejet/stable-diffusion.cpp/pull/851)
|
||||
|
||||
* **2025/09/14** 🚀 stable-diffusion.cpp now supports **Wan2.1 Vace**
|
||||
👉 Details: [PR #819](https://github.com/leejet/stable-diffusion.cpp/pull/819)
|
||||
|
||||
* **2025/09/06** 🚀 stable-diffusion.cpp now supports **Wan2.1 / Wan2.2**
|
||||
👉 Details: [PR #778](https://github.com/leejet/stable-diffusion.cpp/pull/778)
|
||||
|
||||
## Features
|
||||
|
||||
@@ -51,9 +38,11 @@ API and command-line option may change frequently.***
|
||||
- [SD3/SD3.5](./docs/sd3.md)
|
||||
- [FLUX.1-dev/FLUX.1-schnell](./docs/flux.md)
|
||||
- [FLUX.2-dev/FLUX.2-klein](./docs/flux2.md)
|
||||
- [Lens](./docs/lens.md)
|
||||
- [Chroma](./docs/chroma.md)
|
||||
- [Chroma1-Radiance](./docs/chroma_radiance.md)
|
||||
- [Qwen Image](./docs/qwen_image.md)
|
||||
- [LongCat Image](./docs/longcat_image.md)
|
||||
- [Z-Image](./docs/z_image.md)
|
||||
- [Ovis-Image](./docs/ovis_image.md)
|
||||
- [Anima](./docs/anima.md)
|
||||
@@ -62,6 +51,7 @@ API and command-line option may change frequently.***
|
||||
- Image Edit Models
|
||||
- [FLUX.1-Kontext-dev](./docs/kontext.md)
|
||||
- [Qwen Image Edit series](./docs/qwen_image_edit.md)
|
||||
- [LongCat Image Edit](./docs/longcat_image.md)
|
||||
- Video Models
|
||||
- [Wan2.1/Wan2.2](./docs/wan.md)
|
||||
- [LTX-2.3](./docs/ltx2.md)
|
||||
@@ -154,6 +144,8 @@ For runtime and parameter backend placement, see the [backend selection guide](.
|
||||
- [Anima](./docs/anima.md)
|
||||
- [ERNIE-Image](./docs/ernie_image.md)
|
||||
- [HiDream-O1-Image](./docs/hidream_o1_image.md)
|
||||
- [Lens](./docs/lens.md)
|
||||
- [LongCat Image / LongCat Image Edit](./docs/longcat_image.md)
|
||||
- [LoRA](./docs/lora.md)
|
||||
- [LCM/LCM-LoRA](./docs/lcm.md)
|
||||
- [Using PhotoMaker to personalize image generation](./docs/photo_maker.md)
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 630 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 555 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 423 KiB |
Binary file not shown.
@@ -0,0 +1,32 @@
|
||||
# How to Use
|
||||
|
||||
Lens uses a Lens diffusion transformer, the FLUX.2 VAE, and GPT-OSS-20B as the LLM text encoder.
|
||||
|
||||
## Download weights
|
||||
|
||||
- Download Lens
|
||||
- safetensors: https://huggingface.co/Comfy-Org/Lens/tree/main/diffusion_models
|
||||
- Download Lens Turbo
|
||||
- safetensors: https://huggingface.co/Comfy-Org/Lens/tree/main/diffusion_models
|
||||
- Download vae
|
||||
- safetensors: https://huggingface.co/black-forest-labs/FLUX.2-dev/tree/main
|
||||
- Download GPT-OSS-20B
|
||||
- gguf: https://huggingface.co/unsloth/gpt-oss-20b-GGUF/tree/main
|
||||
|
||||
## Examples
|
||||
|
||||
### Lens
|
||||
|
||||
```
|
||||
.\bin\Release\sd-cli.exe --diffusion-model ..\..\ComfyUI\models\diffusion_models\lens_bf16.safetensors --llm "..\..\llm\gpt-oss-20b-UD-Q8_K_XL.gguf" --vae ..\..\ComfyUI\models\vae\flux2_ae.safetensors --cfg-scale 5.0 -p "A crystal dragon soaring through an aurora borealis sky, its entire body made of transparent faceted crystal refracting the green and purple aurora light into rainbow spectra, ice particles trailing from its wings, high fantasy digital art" --diffusion-fa -v
|
||||
```
|
||||
|
||||
<img width="256" alt="Lens example" src="../assets/lens/example.png" />
|
||||
|
||||
### Lens Turbo
|
||||
|
||||
```
|
||||
.\bin\Release\sd-cli.exe --diffusion-model ..\..\ComfyUI\models\diffusion_models\lens_turbo_bf16.safetensors --llm "..\..\llm\gpt-oss-20b-UD-Q8_K_XL.gguf" --vae ..\..\ComfyUI\models\vae\flux2_ae.safetensors --cfg-scale 1.0 -p "A crystal dragon soaring through an aurora borealis sky, its entire body made of transparent faceted crystal refracting the green and purple aurora light into rainbow spectra, ice particles trailing from its wings, high fantasy digital art" --diffusion-fa -v --steps 4
|
||||
```
|
||||
|
||||
<img width="256" alt="Lens Turbo example" src="../assets/lens/turbo_example.png" />
|
||||
@@ -0,0 +1,30 @@
|
||||
# How to Use
|
||||
|
||||
LongCat-Image uses a LongCat diffusion transformer, the FLUX VAE, and Qwen2.5-VL as the LLM text encoder.
|
||||
|
||||
## Download weights
|
||||
|
||||
- Download LongCat Image
|
||||
- safetensors: https://huggingface.co/Comfy-Org/LongCat-Image/tree/main/split_files/diffusion_models
|
||||
- gguf: https://huggingface.co/vantagewithai/LongCat-Image-GGUF/tree/main/comfy
|
||||
- Download LongCat Image Edit
|
||||
- LongCat Image Edit Turbo: https://huggingface.co/meituan-longcat/LongCat-Image-Edit-Turbo
|
||||
- gguf: https://huggingface.co/vantagewithai/LongCat-Image-Edit-GGUF/tree/main
|
||||
- Download vae
|
||||
- safetensors: https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/ae.safetensors
|
||||
- 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
|
||||
- For image editing with GGUF text encoders, also download the matching mmproj file and pass it with `--llm_vision`.
|
||||
|
||||
## Run
|
||||
|
||||
LongCat uses quoted text for character-level text rendering. Put target text inside single quotes, double quotes, or Chinese quotes.
|
||||
|
||||
### LongCat Image
|
||||
|
||||
```
|
||||
.\bin\Release\sd-cli.exe --diffusion-model ..\..\ComfyUI\models\diffusion_models\LongCat-Image-Q4_K_M.gguf --vae ..\..\ComfyUI\models\vae\ae.sft --llm ..\..\ComfyUI\models\text_encoders\Qwen2.5-VL-7B-Instruct-Q8_0.gguf -p "a lovely cat holding a sign says 'longcat.cpp'" --cfg-scale 5.0 --sampling-method euler --flow-shift 3 -v --offload-to-cpu --diffusion-fa
|
||||
```
|
||||
|
||||
<img alt="longcat example" src="../assets/longcat/example.png" />
|
||||
@@ -13,6 +13,8 @@
|
||||
- safetensors: https://huggingface.co/unsloth/LTX-2.3-GGUF/tree/main/vae
|
||||
- Download audio vae
|
||||
- safetensors: https://huggingface.co/unsloth/LTX-2.3-GGUF/tree/main/vae
|
||||
- Download LTX spatial latent upscaler
|
||||
- safetensors: https://huggingface.co/Lightricks/LTX-2.3/resolve/main/ltx-2.3-spatial-upscaler-x2-1.1.safetensors
|
||||
|
||||
## Examples
|
||||
|
||||
@@ -50,4 +52,26 @@
|
||||
src="../assets/ltx2/flf2v.webm"
|
||||
controls
|
||||
muted
|
||||
style="max-width: 100%; height: auto;"></video>
|
||||
|
||||
### LTX-2.3 spatial latent upscale
|
||||
|
||||
LTX spatial latent upscale runs a model-backed x2 latent upsampler between the low-resolution video pass and the high-resolution refine pass. `-W` and `-H` are the pre-upscale generation size; the spatial upsampler produces x2 latent dimensions.
|
||||
|
||||
Put `ltx-2.3-spatial-upscaler-x2-1.1.safetensors` under the directory passed to `--hires-upscalers-dir`, then use the model name without path or extension in `--hires-upscaler`.
|
||||
|
||||
```
|
||||
.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model ..\..\ComfyUI\models\diffusion_models\ltx-2.3-22b-dev-UD-Q4_K_M.gguf --vae ..\..\ComfyUI\models\vae\ltx-2.3-22b-dev_video_vae.safetensors --audio-vae ..\..\ComfyUI\models\vae\ltx-2.3-22b-dev_audio_vae.safetensors --llm ..\..\ComfyUI\models\text_encoders\gemma-3-12b-it-qat-UD-Q4_K_XL.gguf --embeddings-connectors ..\..\ComfyUI\models\text_encoders\ltx-2.3-22b-dev_embeddings_connectors.safetensors --hires-upscalers-dir ..\..\ComfyUI\models\latent_upscale_models --hires-upscaler ltx-2.3-spatial-upscaler-x2-1.1 --hires --hires-steps 4 -p "a lovely cat" --cfg-scale 6.0 --sampling-method euler -v -W 640 -H 360 --diffusion-fa --offload-to-cpu --video-frames 33 -i ..\assets\ernie_image\turbo_example.png -o hires_i2v.webm
|
||||
```
|
||||
|
||||
By default, the hires refine pass uses the main sampler and scheduler, then trims the second-pass sigma schedule by `--hires-denoising-strength` (`0.7` by default). To reproduce a ComfyUI-style explicit refine schedule, pass custom hires sigmas:
|
||||
|
||||
```
|
||||
--hires-sigmas "0.85,0.725,0.421875,0.0"
|
||||
```
|
||||
|
||||
<video
|
||||
src="../assets/ltx2/hires_i2v.webm"
|
||||
controls
|
||||
muted
|
||||
style="max-width: 100%; height: auto;"></video>
|
||||
@@ -7,6 +7,9 @@ add_executable(${TARGET}
|
||||
image_metadata.cpp
|
||||
main.cpp
|
||||
)
|
||||
if(APPLE)
|
||||
sd_set_macos_rpaths(${TARGET})
|
||||
endif()
|
||||
target_include_directories(${TARGET} PRIVATE
|
||||
"${CMAKE_CURRENT_SOURCE_DIR}/.."
|
||||
"${PROJECT_SOURCE_DIR}/src"
|
||||
|
||||
@@ -55,7 +55,8 @@ Context Options:
|
||||
then threads will be set to the number of CPU physical cores
|
||||
--chroma-t5-mask-pad <int> t5 mask pad size of chroma
|
||||
--max-vram <float> maximum VRAM budget in GiB for graph-cut segmented execution. 0 disables
|
||||
graph splitting; -1 auto-detects free VRAM minus 1 GiB
|
||||
graph splitting; a negative value auto-detects free VRAM, sparing the
|
||||
specified value (e.g. -0.5 will keep at least 0.5 GiB free)
|
||||
--force-sdxl-vae-conv-scale force use of conv scale on sdxl vae
|
||||
--offload-to-cpu place the weights in RAM to save VRAM, and automatically load them into VRAM
|
||||
when needed
|
||||
@@ -106,6 +107,8 @@ Generation Options:
|
||||
--extra-sample-args <string> extra sampler/scheduler args, key=value list. lcm supports noise_clip_std,
|
||||
noise_scale_start, noise_scale_end; ltx2 supports max_shift, base_shift,
|
||||
stretch, terminal; euler_ge supports gamma
|
||||
--extra-tiling-args <string> extra VAE tiling args, key=value list. LTX video VAE supports
|
||||
temporal_tile_frames (default: 4), temporal_tile_overlap (default: 1)
|
||||
-H, --height <int> image height, in pixel space (default: 512)
|
||||
-W, --width <int> image width, in pixel space (default: 512)
|
||||
--steps <int> number of sample steps (default: 20)
|
||||
@@ -175,6 +178,8 @@ Generation Options:
|
||||
model-specific
|
||||
--sigmas custom sigma values for the sampler, comma-separated (e.g.,
|
||||
"14.61,7.8,3.5,0.0").
|
||||
--hires-sigmas custom sigma values for the highres fix second pass, comma-separated (e.g.,
|
||||
"0.85,0.725,0.421875,0.0").
|
||||
--skip-layers layers to skip for SLG steps (default: [7,8,9])
|
||||
--high-noise-skip-layers (high noise) layers to skip for SLG steps (default: [7,8,9])
|
||||
-r, --ref-image reference image for Flux Kontext models (can be used multiple times)
|
||||
|
||||
+122
-42
@@ -413,7 +413,7 @@ ArgOptions SDContextParams::get_options() {
|
||||
options.float_options = {
|
||||
{"",
|
||||
"--max-vram",
|
||||
"maximum VRAM budget in GiB for graph-cut segmented execution. 0 disables graph splitting; -1 auto-detects free VRAM minus 1 GiB",
|
||||
"maximum VRAM budget in GiB for graph-cut segmented execution. 0 disables graph splitting; a negative value auto-detects free VRAM, sparing the specified value (e.g. -0.5 will keep at least 0.5 GiB free)",
|
||||
&max_vram},
|
||||
};
|
||||
|
||||
@@ -835,6 +835,10 @@ ArgOptions SDGenerationParams::get_options() {
|
||||
"--extra-sample-args",
|
||||
"extra sampler/scheduler args, key=value list. lcm supports noise_clip_std, noise_scale_start, noise_scale_end; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma",
|
||||
&extra_sample_args},
|
||||
{"",
|
||||
"--extra-tiling-args",
|
||||
"extra VAE tiling args, key=value list. LTX video VAE supports temporal_tile_frames (default: 4), temporal_tile_overlap (default: 1)",
|
||||
&extra_tiling_args},
|
||||
};
|
||||
|
||||
options.int_options = {
|
||||
@@ -1134,11 +1138,11 @@ ArgOptions SDGenerationParams::get_options() {
|
||||
return 1;
|
||||
};
|
||||
|
||||
auto on_sigmas_arg = [&](int argc, const char** argv, int index) {
|
||||
if (++index >= argc) {
|
||||
auto parse_sigmas_arg = [&](const char* value, std::vector<float>* target, const char* option_name) {
|
||||
if (target == nullptr || value == nullptr) {
|
||||
return -1;
|
||||
}
|
||||
std::string sigmas_str = argv[index];
|
||||
std::string sigmas_str = value;
|
||||
if (!sigmas_str.empty() && sigmas_str.front() == '[') {
|
||||
sigmas_str.erase(0, 1);
|
||||
}
|
||||
@@ -1146,6 +1150,7 @@ ArgOptions SDGenerationParams::get_options() {
|
||||
sigmas_str.pop_back();
|
||||
}
|
||||
|
||||
size_t before = target->size();
|
||||
std::stringstream ss(sigmas_str);
|
||||
std::string item;
|
||||
while (std::getline(ss, item, ',')) {
|
||||
@@ -1153,24 +1158,38 @@ ArgOptions SDGenerationParams::get_options() {
|
||||
item.erase(item.find_last_not_of(" \t\n\r\f\v") + 1);
|
||||
if (!item.empty()) {
|
||||
try {
|
||||
custom_sigmas.push_back(std::stof(item));
|
||||
target->push_back(std::stof(item));
|
||||
} catch (const std::invalid_argument&) {
|
||||
LOG_ERROR("error: invalid float value '%s' in --sigmas", item.c_str());
|
||||
LOG_ERROR("error: invalid float value '%s' in %s", item.c_str(), option_name);
|
||||
return -1;
|
||||
} catch (const std::out_of_range&) {
|
||||
LOG_ERROR("error: float value '%s' out of range in --sigmas", item.c_str());
|
||||
LOG_ERROR("error: float value '%s' out of range in %s", item.c_str(), option_name);
|
||||
return -1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (custom_sigmas.empty() && !sigmas_str.empty()) {
|
||||
LOG_ERROR("error: could not parse any sigma values from '%s'", argv[index]);
|
||||
if (target->size() == before && !sigmas_str.empty()) {
|
||||
LOG_ERROR("error: could not parse any sigma values from '%s'", value);
|
||||
return -1;
|
||||
}
|
||||
return 1;
|
||||
};
|
||||
|
||||
auto on_sigmas_arg = [&](int argc, const char** argv, int index) {
|
||||
if (++index >= argc) {
|
||||
return -1;
|
||||
}
|
||||
return parse_sigmas_arg(argv[index], &custom_sigmas, "--sigmas");
|
||||
};
|
||||
|
||||
auto on_hires_sigmas_arg = [&](int argc, const char** argv, int index) {
|
||||
if (++index >= argc) {
|
||||
return -1;
|
||||
}
|
||||
return parse_sigmas_arg(argv[index], &hires_custom_sigmas, "--hires-sigmas");
|
||||
};
|
||||
|
||||
auto on_ref_image_arg = [&](int argc, const char** argv, int index) {
|
||||
if (++index >= argc) {
|
||||
return -1;
|
||||
@@ -1293,6 +1312,10 @@ ArgOptions SDGenerationParams::get_options() {
|
||||
"--sigmas",
|
||||
"custom sigma values for the sampler, comma-separated (e.g., \"14.61,7.8,3.5,0.0\").",
|
||||
on_sigmas_arg},
|
||||
{"",
|
||||
"--hires-sigmas",
|
||||
"custom sigma values for the highres fix second pass, comma-separated (e.g., \"0.85,0.725,0.421875,0.0\").",
|
||||
on_hires_sigmas_arg},
|
||||
{"",
|
||||
"--skip-layers",
|
||||
"layers to skip for SLG steps (default: [7,8,9])",
|
||||
@@ -1525,11 +1548,31 @@ static bool resolve_model_file_from_dir(const std::string& model_name,
|
||||
LOG_ERROR("%s directory is empty", label);
|
||||
return false;
|
||||
}
|
||||
auto ends_with_valid_ext = [&]() {
|
||||
for (const auto& ext : valid_ext) {
|
||||
if (model_name.size() < ext.size()) {
|
||||
continue;
|
||||
}
|
||||
auto suffix = model_name.substr(model_name.size() - ext.size());
|
||||
std::transform(suffix.begin(), suffix.end(), suffix.begin(), [](unsigned char c) {
|
||||
return static_cast<char>(std::tolower(c));
|
||||
});
|
||||
std::string lower_ext = ext;
|
||||
std::transform(lower_ext.begin(), lower_ext.end(), lower_ext.begin(), [](unsigned char c) {
|
||||
return static_cast<char>(std::tolower(c));
|
||||
});
|
||||
if (suffix == lower_ext) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
return false;
|
||||
};
|
||||
|
||||
if (model_name.empty() ||
|
||||
model_name.find('/') != std::string::npos ||
|
||||
model_name.find('\\') != std::string::npos ||
|
||||
fs::path(model_name).has_root_path() ||
|
||||
fs::path(model_name).has_extension()) {
|
||||
ends_with_valid_ext()) {
|
||||
LOG_ERROR("%s must be a model name without path or extension: %s", label, model_name.c_str());
|
||||
return false;
|
||||
}
|
||||
@@ -1633,6 +1676,9 @@ bool SDGenerationParams::from_json_str(
|
||||
if (hires_json.contains("denoising_strength") && hires_json["denoising_strength"].is_number()) {
|
||||
hires_denoising_strength = hires_json["denoising_strength"];
|
||||
}
|
||||
if (hires_json.contains("custom_sigmas") && hires_json["custom_sigmas"].is_array()) {
|
||||
hires_custom_sigmas = hires_json["custom_sigmas"].get<std::vector<float>>();
|
||||
}
|
||||
if (hires_json.contains("upscale_tile_size") && hires_json["upscale_tile_size"].is_number_integer()) {
|
||||
hires_upscale_tile_size = hires_json["upscale_tile_size"];
|
||||
}
|
||||
@@ -1738,6 +1784,9 @@ bool SDGenerationParams::from_json_str(
|
||||
if (tiling_json.contains("rel_size_y") && tiling_json["rel_size_y"].is_number()) {
|
||||
vae_tiling_params.rel_size_y = tiling_json["rel_size_y"];
|
||||
}
|
||||
if (tiling_json.contains("extra_tiling_args") && tiling_json["extra_tiling_args"].is_string()) {
|
||||
extra_tiling_args = tiling_json["extra_tiling_args"].get<std::string>();
|
||||
}
|
||||
}
|
||||
|
||||
if (!parse_lora_json_field(j, lora_path_resolver, lora_map, high_noise_lora_map)) {
|
||||
@@ -1960,6 +2009,8 @@ bool SDGenerationParams::initialize_cache_params() {
|
||||
}
|
||||
|
||||
bool SDGenerationParams::resolve(const std::string& lora_model_dir, const std::string& hires_upscalers_dir, bool strict) {
|
||||
vae_tiling_params.extra_tiling_args = extra_tiling_args.empty() ? nullptr : extra_tiling_args.c_str();
|
||||
|
||||
if (high_noise_sample_params.sample_steps <= 0) {
|
||||
high_noise_sample_params.sample_steps = -1;
|
||||
}
|
||||
@@ -2080,6 +2131,10 @@ bool SDGenerationParams::validate(SDMode mode) {
|
||||
LOG_ERROR("error: hires denoising strength must be in (0.0, 1.0]");
|
||||
return false;
|
||||
}
|
||||
if (!hires_custom_sigmas.empty() && hires_custom_sigmas.size() < 2) {
|
||||
LOG_ERROR("error: hires custom sigmas must contain at least two values");
|
||||
return false;
|
||||
}
|
||||
if (hires_upscale_tile_size < 1) {
|
||||
LOG_ERROR("error: hires upscale tile size must be positive");
|
||||
return false;
|
||||
@@ -2142,6 +2197,7 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
|
||||
sample_params.custom_sigmas_count = static_cast<int>(custom_sigmas.size());
|
||||
sample_params.extra_sample_args = extra_sample_args.empty() ? nullptr : extra_sample_args.c_str();
|
||||
high_noise_sample_params.extra_sample_args = high_noise_extra_sample_args.empty() ? nullptr : high_noise_extra_sample_args.c_str();
|
||||
vae_tiling_params.extra_tiling_args = extra_tiling_args.empty() ? nullptr : extra_tiling_args.c_str();
|
||||
cache_params.scm_mask = scm_mask.empty() ? nullptr : scm_mask.c_str();
|
||||
|
||||
sd_pm_params_t pm_params = {
|
||||
@@ -2174,15 +2230,17 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
|
||||
params.vae_tiling_params = vae_tiling_params;
|
||||
params.cache = cache_params;
|
||||
|
||||
params.hires.enabled = hires_enabled;
|
||||
params.hires.upscaler = resolved_hires_upscaler;
|
||||
params.hires.model_path = hires_upscaler_model_path.empty() ? nullptr : hires_upscaler_model_path.c_str();
|
||||
params.hires.scale = hires_scale;
|
||||
params.hires.target_width = hires_width;
|
||||
params.hires.target_height = hires_height;
|
||||
params.hires.steps = hires_steps;
|
||||
params.hires.denoising_strength = hires_denoising_strength;
|
||||
params.hires.upscale_tile_size = hires_upscale_tile_size;
|
||||
params.hires.enabled = hires_enabled;
|
||||
params.hires.upscaler = resolved_hires_upscaler;
|
||||
params.hires.model_path = hires_upscaler_model_path.empty() ? nullptr : hires_upscaler_model_path.c_str();
|
||||
params.hires.scale = hires_scale;
|
||||
params.hires.target_width = hires_width;
|
||||
params.hires.target_height = hires_height;
|
||||
params.hires.steps = hires_steps;
|
||||
params.hires.denoising_strength = hires_denoising_strength;
|
||||
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());
|
||||
return params;
|
||||
}
|
||||
|
||||
@@ -2213,29 +2271,41 @@ sd_vid_gen_params_t SDGenerationParams::to_sd_vid_gen_params_t() {
|
||||
sample_params.custom_sigmas_count = static_cast<int>(custom_sigmas.size());
|
||||
sample_params.extra_sample_args = extra_sample_args.empty() ? nullptr : extra_sample_args.c_str();
|
||||
high_noise_sample_params.extra_sample_args = high_noise_extra_sample_args.empty() ? nullptr : high_noise_extra_sample_args.c_str();
|
||||
vae_tiling_params.extra_tiling_args = extra_tiling_args.empty() ? nullptr : extra_tiling_args.c_str();
|
||||
cache_params.scm_mask = scm_mask.empty() ? nullptr : scm_mask.c_str();
|
||||
|
||||
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.end_image = end_image.get();
|
||||
params.control_frames = control_frame_views.empty() ? nullptr : control_frame_views.data();
|
||||
params.control_frames_size = static_cast<int>(control_frame_views.size());
|
||||
params.width = get_resolved_width();
|
||||
params.height = get_resolved_height();
|
||||
params.sample_params = sample_params;
|
||||
params.high_noise_sample_params = high_noise_sample_params;
|
||||
params.moe_boundary = moe_boundary;
|
||||
params.strength = strength;
|
||||
params.seed = seed;
|
||||
params.video_frames = video_frames;
|
||||
params.fps = fps;
|
||||
params.vace_strength = vace_strength;
|
||||
params.vae_tiling_params = vae_tiling_params;
|
||||
params.cache = cache_params;
|
||||
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.end_image = end_image.get();
|
||||
params.control_frames = control_frame_views.empty() ? nullptr : control_frame_views.data();
|
||||
params.control_frames_size = static_cast<int>(control_frame_views.size());
|
||||
params.width = get_resolved_width();
|
||||
params.height = get_resolved_height();
|
||||
params.sample_params = sample_params;
|
||||
params.high_noise_sample_params = high_noise_sample_params;
|
||||
params.moe_boundary = moe_boundary;
|
||||
params.strength = strength;
|
||||
params.seed = seed;
|
||||
params.video_frames = video_frames;
|
||||
params.fps = fps;
|
||||
params.vace_strength = vace_strength;
|
||||
params.vae_tiling_params = vae_tiling_params;
|
||||
params.cache = cache_params;
|
||||
params.hires.enabled = hires_enabled;
|
||||
params.hires.upscaler = resolved_hires_upscaler;
|
||||
params.hires.model_path = hires_upscaler_model_path.empty() ? nullptr : hires_upscaler_model_path.c_str();
|
||||
params.hires.scale = hires_scale;
|
||||
params.hires.target_width = hires_width;
|
||||
params.hires.target_height = hires_height;
|
||||
params.hires.steps = hires_steps;
|
||||
params.hires.denoising_strength = hires_denoising_strength;
|
||||
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());
|
||||
return params;
|
||||
}
|
||||
|
||||
@@ -2318,6 +2388,7 @@ std::string SDGenerationParams::to_string() const {
|
||||
<< ", target_height: " << hires_height
|
||||
<< ", steps: " << hires_steps
|
||||
<< ", denoising_strength: " << hires_denoising_strength
|
||||
<< ", custom_sigmas: " << vec_to_string(hires_custom_sigmas)
|
||||
<< ", upscale_tile_size: " << hires_upscale_tile_size << " },\n"
|
||||
<< " vae_tiling_params: { "
|
||||
<< vae_tiling_params.enabled << ", "
|
||||
@@ -2326,7 +2397,8 @@ std::string SDGenerationParams::to_string() const {
|
||||
<< vae_tiling_params.tile_size_y << ", "
|
||||
<< vae_tiling_params.target_overlap << ", "
|
||||
<< vae_tiling_params.rel_size_x << ", "
|
||||
<< vae_tiling_params.rel_size_y << " },\n"
|
||||
<< vae_tiling_params.rel_size_y << ", "
|
||||
<< "\"" << extra_tiling_args << "\" },\n"
|
||||
<< "}";
|
||||
return oss.str();
|
||||
}
|
||||
@@ -2469,6 +2541,7 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params,
|
||||
{"target_height", gen_params.hires_height},
|
||||
{"steps", gen_params.hires_steps},
|
||||
{"denoising_strength", gen_params.hires_denoising_strength},
|
||||
{"custom_sigmas", gen_params.hires_custom_sigmas},
|
||||
{"upscale_tile_size", gen_params.hires_upscale_tile_size},
|
||||
};
|
||||
}
|
||||
@@ -2504,14 +2577,18 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params,
|
||||
};
|
||||
}
|
||||
|
||||
if (gen_params.vae_tiling_params.enabled) {
|
||||
if (gen_params.vae_tiling_params.enabled ||
|
||||
gen_params.vae_tiling_params.temporal_tiling ||
|
||||
!gen_params.extra_tiling_args.empty()) {
|
||||
root["vae_tiling"] = {
|
||||
{"enabled", gen_params.vae_tiling_params.enabled},
|
||||
{"temporal_tiling", gen_params.vae_tiling_params.temporal_tiling},
|
||||
{"tile_size_x", gen_params.vae_tiling_params.tile_size_x},
|
||||
{"tile_size_y", gen_params.vae_tiling_params.tile_size_y},
|
||||
{"target_overlap", gen_params.vae_tiling_params.target_overlap},
|
||||
{"rel_size_x", gen_params.vae_tiling_params.rel_size_x},
|
||||
{"rel_size_y", gen_params.vae_tiling_params.rel_size_y},
|
||||
{"extra_tiling_args", gen_params.extra_tiling_args},
|
||||
};
|
||||
}
|
||||
|
||||
@@ -2588,6 +2665,9 @@ std::string get_image_params(const SDContextParams& ctx_params,
|
||||
parameter_string += "Hires resize: " + std::to_string(gen_params.hires_width) + "x" + std::to_string(gen_params.hires_height) + ", ";
|
||||
parameter_string += "Hires steps: " + std::to_string(gen_params.hires_steps) + ", ";
|
||||
parameter_string += "Denoising strength: " + std::to_string(gen_params.hires_denoising_strength) + ", ";
|
||||
if (!gen_params.hires_custom_sigmas.empty()) {
|
||||
parameter_string += "Hires custom sigmas: " + vec_to_string(gen_params.hires_custom_sigmas) + ", ";
|
||||
}
|
||||
}
|
||||
parameter_string += "Version: stable-diffusion.cpp";
|
||||
parameter_string += ", SDCPP: " + build_sdcpp_image_metadata_json(ctx_params, gen_params, seed, mode);
|
||||
|
||||
@@ -189,7 +189,8 @@ struct SDGenerationParams {
|
||||
int video_frames = 1;
|
||||
int fps = 16;
|
||||
float vace_strength = 1.f;
|
||||
sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f};
|
||||
sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr};
|
||||
std::string extra_tiling_args;
|
||||
|
||||
std::string pm_id_images_dir;
|
||||
std::string pm_id_embed_path;
|
||||
@@ -207,6 +208,7 @@ struct SDGenerationParams {
|
||||
int hires_steps = 0;
|
||||
float hires_denoising_strength = 0.7f;
|
||||
int hires_upscale_tile_size = 128;
|
||||
std::vector<float> hires_custom_sigmas;
|
||||
|
||||
std::map<std::string, float> lora_map;
|
||||
std::map<std::string, float> high_noise_lora_map;
|
||||
|
||||
@@ -74,6 +74,9 @@ add_executable(${TARGET}
|
||||
routes_sdapi.cpp
|
||||
routes_sdcpp.cpp
|
||||
)
|
||||
if(APPLE)
|
||||
sd_set_macos_rpaths(${TARGET})
|
||||
endif()
|
||||
|
||||
if(HAVE_FRONTEND_BUILD)
|
||||
add_dependencies(${TARGET} ${TARGET}_frontend)
|
||||
|
||||
@@ -157,7 +157,8 @@ Context Options:
|
||||
then threads will be set to the number of CPU physical cores
|
||||
--chroma-t5-mask-pad <int> t5 mask pad size of chroma
|
||||
--max-vram <float> maximum VRAM budget in GiB for graph-cut segmented execution. 0 disables
|
||||
graph splitting; -1 auto-detects free VRAM minus 1 GiB
|
||||
graph splitting; a negative value auto-detects free VRAM, sparing the
|
||||
specified value (e.g. -0.5 will keep at least 0.5 GiB free)
|
||||
--force-sdxl-vae-conv-scale force use of conv scale on sdxl vae
|
||||
--offload-to-cpu place the weights in RAM to save VRAM, and automatically load them into VRAM
|
||||
when needed
|
||||
@@ -208,6 +209,8 @@ Default Generation Options:
|
||||
--extra-sample-args <string> extra sampler/scheduler args, key=value list. lcm supports noise_clip_std,
|
||||
noise_scale_start, noise_scale_end; ltx2 supports max_shift, base_shift,
|
||||
stretch, terminal; euler_ge supports gamma
|
||||
--extra-tiling-args <string> extra VAE tiling args, key=value list. LTX video VAE supports
|
||||
temporal_tile_frames (default: 4), temporal_tile_overlap (default: 1)
|
||||
-H, --height <int> image height, in pixel space (default: 512)
|
||||
-W, --width <int> image width, in pixel space (default: 512)
|
||||
--steps <int> number of sample steps (default: 20)
|
||||
@@ -263,6 +266,7 @@ Default Generation Options:
|
||||
--disable-auto-resize-ref-image disable auto resize of ref images
|
||||
--disable-image-metadata do not embed generation metadata on image files
|
||||
--vae-tiling process vae in tiles to reduce memory usage
|
||||
--temporal-tiling enable temporal tiling for LTX video VAE decode
|
||||
--hires enable highres fix
|
||||
-s, --seed RNG seed (default: 42, use random seed for < 0)
|
||||
--sampling-method sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m,
|
||||
@@ -276,6 +280,8 @@ Default Generation Options:
|
||||
model-specific
|
||||
--sigmas custom sigma values for the sampler, comma-separated (e.g.,
|
||||
"14.61,7.8,3.5,0.0").
|
||||
--hires-sigmas custom sigma values for the highres fix second pass, comma-separated (e.g.,
|
||||
"0.85,0.725,0.421875,0.0").
|
||||
--skip-layers layers to skip for SLG steps (default: [7,8,9])
|
||||
--high-noise-skip-layers (high noise) layers to skip for SLG steps (default: [7,8,9])
|
||||
-r, --ref-image reference image for Flux Kontext models (can be used multiple times)
|
||||
|
||||
+30
-3
@@ -504,11 +504,13 @@ Shared default fields used by both `img_gen` and `vid_gen`:
|
||||
| `sample_params.guidance.slg.scale` | `number` |
|
||||
| `vae_tiling_params` | `object` |
|
||||
| `vae_tiling_params.enabled` | `boolean` |
|
||||
| `vae_tiling_params.temporal_tiling` | `boolean` |
|
||||
| `vae_tiling_params.tile_size_x` | `integer` |
|
||||
| `vae_tiling_params.tile_size_y` | `integer` |
|
||||
| `vae_tiling_params.target_overlap` | `number` |
|
||||
| `vae_tiling_params.rel_size_x` | `number` |
|
||||
| `vae_tiling_params.rel_size_y` | `number` |
|
||||
| `vae_tiling_params.extra_tiling_args` | `string` |
|
||||
| `cache_mode` | `string` |
|
||||
| `cache_option` | `string` |
|
||||
| `scm_mask` | `string` |
|
||||
@@ -516,6 +518,8 @@ Shared default fields used by both `img_gen` and `vid_gen`:
|
||||
| `output_format` | `string` |
|
||||
| `output_compression` | `integer` |
|
||||
|
||||
`vae_tiling_params.extra_tiling_args` accepts a key=value list. For LTX video VAE temporal tiling, `temporal_tile_frames` defaults to `4` and `temporal_tile_overlap` defaults to `1`.
|
||||
|
||||
`img_gen`-specific default fields:
|
||||
|
||||
| Field | Type |
|
||||
@@ -532,6 +536,7 @@ Shared default fields used by both `img_gen` and `vid_gen`:
|
||||
| `hires.target_height` | `integer` |
|
||||
| `hires.steps` | `integer` |
|
||||
| `hires.denoising_strength` | `number` |
|
||||
| `hires.custom_sigmas` | `array<number>` |
|
||||
| `hires.upscale_tile_size` | `integer` |
|
||||
|
||||
`vid_gen`-specific default fields:
|
||||
@@ -685,16 +690,19 @@ Example:
|
||||
"target_height": 0,
|
||||
"steps": 0,
|
||||
"denoising_strength": 0.7,
|
||||
"custom_sigmas": [],
|
||||
"upscale_tile_size": 128
|
||||
},
|
||||
|
||||
"vae_tiling_params": {
|
||||
"enabled": false,
|
||||
"temporal_tiling": false,
|
||||
"tile_size_x": 0,
|
||||
"tile_size_y": 0,
|
||||
"target_overlap": 0.5,
|
||||
"rel_size_x": 0.0,
|
||||
"rel_size_y": 0.0
|
||||
"rel_size_y": 0.0,
|
||||
"extra_tiling_args": ""
|
||||
},
|
||||
|
||||
"cache_mode": "disabled",
|
||||
@@ -799,14 +807,23 @@ Other native fields:
|
||||
| `hires.target_height` | `integer` |
|
||||
| `hires.steps` | `integer` |
|
||||
| `hires.denoising_strength` | `number` |
|
||||
| `hires.custom_sigmas` | `array<number>` |
|
||||
| `hires.upscale_tile_size` | `integer` |
|
||||
| `vae_tiling_params` | `object` |
|
||||
| `vae_tiling_params.enabled` | `boolean` |
|
||||
| `vae_tiling_params.temporal_tiling` | `boolean` |
|
||||
| `vae_tiling_params.tile_size_x` | `integer` |
|
||||
| `vae_tiling_params.tile_size_y` | `integer` |
|
||||
| `vae_tiling_params.target_overlap` | `number` |
|
||||
| `vae_tiling_params.rel_size_x` | `number` |
|
||||
| `vae_tiling_params.rel_size_y` | `number` |
|
||||
| `vae_tiling_params.extra_tiling_args` | `string` |
|
||||
| `cache_mode` | `string` |
|
||||
| `cache_option` | `string` |
|
||||
| `scm_mask` | `string` |
|
||||
| `scm_policy_dynamic` | `boolean` |
|
||||
|
||||
For `hires.upscaler`, use `Lanczos`, `Nearest`, `Latent`, `Latent (nearest)`, `Latent (nearest-exact)`, `Latent (antialiased)`, `Latent (bicubic)`, `Latent (bicubic antialiased)`, or an `upscalers[].name` value from `GET /sdcpp/v1/capabilities`. Model-backed upscalers are resolved as `--hires-upscalers-dir / (name + ext)` and must live directly in that directory.
|
||||
For `hires.upscaler`, use `Lanczos`, `Nearest`, `Latent`, `Latent (nearest)`, `Latent (nearest-exact)`, `Latent (antialiased)`, `Latent (bicubic)`, `Latent (bicubic antialiased)`, or an `upscalers[].name` value from `GET /sdcpp/v1/capabilities`. Model-backed upscalers are resolved as `--hires-upscalers-dir / (name + ext)` and must live directly in that directory. `hires.custom_sigmas`, when present, overrides the generated second-pass hires sigma schedule; otherwise the hires schedule is trimmed by `hires.denoising_strength`.
|
||||
|
||||
HTTP-only output fields:
|
||||
|
||||
@@ -1009,11 +1026,13 @@ Example:
|
||||
|
||||
"vae_tiling_params": {
|
||||
"enabled": false,
|
||||
"temporal_tiling": false,
|
||||
"tile_size_x": 0,
|
||||
"tile_size_y": 0,
|
||||
"target_overlap": 0.5,
|
||||
"rel_size_x": 0.0,
|
||||
"rel_size_y": 0.0
|
||||
"rel_size_y": 0.0,
|
||||
"extra_tiling_args": ""
|
||||
},
|
||||
|
||||
"cache_mode": "disabled",
|
||||
@@ -1131,6 +1150,14 @@ Other native fields:
|
||||
| Field | Type |
|
||||
| --- | --- |
|
||||
| `vae_tiling_params` | `object` |
|
||||
| `vae_tiling_params.enabled` | `boolean` |
|
||||
| `vae_tiling_params.temporal_tiling` | `boolean` |
|
||||
| `vae_tiling_params.tile_size_x` | `integer` |
|
||||
| `vae_tiling_params.tile_size_y` | `integer` |
|
||||
| `vae_tiling_params.target_overlap` | `number` |
|
||||
| `vae_tiling_params.rel_size_x` | `number` |
|
||||
| `vae_tiling_params.rel_size_y` | `number` |
|
||||
| `vae_tiling_params.extra_tiling_args` | `string` |
|
||||
| `cache_mode` | `string` |
|
||||
| `cache_option` | `string` |
|
||||
| `scm_mask` | `string` |
|
||||
|
||||
@@ -56,11 +56,13 @@ static const char* capability_sample_method_name(enum sample_method_t sample_met
|
||||
static json make_vae_tiling_json(const sd_tiling_params_t& params) {
|
||||
return {
|
||||
{"enabled", params.enabled},
|
||||
{"temporal_tiling", params.temporal_tiling},
|
||||
{"tile_size_x", params.tile_size_x},
|
||||
{"tile_size_y", params.tile_size_y},
|
||||
{"target_overlap", params.target_overlap},
|
||||
{"rel_size_x", params.rel_size_x},
|
||||
{"rel_size_y", params.rel_size_y},
|
||||
{"extra_tiling_args", params.extra_tiling_args ? params.extra_tiling_args : ""},
|
||||
};
|
||||
}
|
||||
|
||||
@@ -100,6 +102,20 @@ static json make_sample_params_json(const sd_sample_params_t& sample_params, con
|
||||
};
|
||||
}
|
||||
|
||||
static json make_hires_json(const SDGenerationParams& defaults) {
|
||||
return {
|
||||
{"enabled", defaults.hires_enabled},
|
||||
{"upscaler", defaults.hires_upscaler},
|
||||
{"scale", defaults.hires_scale},
|
||||
{"target_width", defaults.hires_width},
|
||||
{"target_height", defaults.hires_height},
|
||||
{"steps", defaults.hires_steps},
|
||||
{"denoising_strength", defaults.hires_denoising_strength},
|
||||
{"custom_sigmas", defaults.hires_custom_sigmas},
|
||||
{"upscale_tile_size", defaults.hires_upscale_tile_size},
|
||||
};
|
||||
}
|
||||
|
||||
static json make_img_gen_defaults_json(const SDGenerationParams& defaults, const std::string& output_format) {
|
||||
return {
|
||||
{"prompt", defaults.prompt},
|
||||
@@ -114,17 +130,7 @@ static json make_img_gen_defaults_json(const SDGenerationParams& defaults, const
|
||||
{"increase_ref_index", defaults.increase_ref_index},
|
||||
{"control_strength", defaults.control_strength},
|
||||
{"sample_params", make_sample_params_json(defaults.sample_params, defaults.skip_layers)},
|
||||
{"hires",
|
||||
{
|
||||
{"enabled", defaults.hires_enabled},
|
||||
{"upscaler", defaults.hires_upscaler},
|
||||
{"scale", defaults.hires_scale},
|
||||
{"target_width", defaults.hires_width},
|
||||
{"target_height", defaults.hires_height},
|
||||
{"steps", defaults.hires_steps},
|
||||
{"denoising_strength", defaults.hires_denoising_strength},
|
||||
{"upscale_tile_size", defaults.hires_upscale_tile_size},
|
||||
}},
|
||||
{"hires", make_hires_json(defaults)},
|
||||
{"vae_tiling_params", make_vae_tiling_json(defaults.vae_tiling_params)},
|
||||
{"cache_mode", defaults.cache_mode},
|
||||
{"cache_option", defaults.cache_option},
|
||||
@@ -150,6 +156,7 @@ static json make_vid_gen_defaults_json(const SDGenerationParams& defaults, const
|
||||
{"vace_strength", defaults.vace_strength},
|
||||
{"sample_params", make_sample_params_json(defaults.sample_params, defaults.skip_layers)},
|
||||
{"high_noise_sample_params", make_sample_params_json(defaults.high_noise_sample_params, defaults.high_noise_skip_layers)},
|
||||
{"hires", make_hires_json(defaults)},
|
||||
{"vae_tiling_params", make_vae_tiling_json(defaults.vae_tiling_params)},
|
||||
{"cache_mode", defaults.cache_mode},
|
||||
{"cache_option", defaults.cache_option},
|
||||
|
||||
@@ -254,7 +254,7 @@ void refresh_lora_cache(ServerRuntime& rt) {
|
||||
|
||||
fs::path lora_dir = rt.ctx_params->lora_model_dir;
|
||||
if (fs::exists(lora_dir) && fs::is_directory(lora_dir)) {
|
||||
for (auto& entry : fs::recursive_directory_iterator(lora_dir)) {
|
||||
for (auto& entry : fs::recursive_directory_iterator(lora_dir, fs::directory_options::skip_permission_denied)) {
|
||||
if (!entry.is_regular_file()) {
|
||||
continue;
|
||||
}
|
||||
|
||||
+1
-1
Submodule ggml updated: 7f4ab364b2...0ce7ad348a
@@ -126,7 +126,8 @@ enum sd_type_t {
|
||||
// SD_TYPE_IQ4_NL_8_8 = 38,
|
||||
SD_TYPE_MXFP4 = 39, // MXFP4 (1 block)
|
||||
SD_TYPE_NVFP4 = 40, // NVFP4 (4 blocks, E4M3 scale)
|
||||
SD_TYPE_COUNT = 41,
|
||||
SD_TYPE_Q1_0 = 41,
|
||||
SD_TYPE_COUNT = 42,
|
||||
};
|
||||
|
||||
enum sd_log_level_t {
|
||||
@@ -159,6 +160,7 @@ typedef struct {
|
||||
float target_overlap;
|
||||
float rel_size_x;
|
||||
float rel_size_y;
|
||||
const char* extra_tiling_args;
|
||||
} sd_tiling_params_t;
|
||||
|
||||
typedef struct {
|
||||
@@ -331,6 +333,8 @@ typedef struct {
|
||||
int steps;
|
||||
float denoising_strength;
|
||||
int upscale_tile_size;
|
||||
float* custom_sigmas;
|
||||
int custom_sigmas_count;
|
||||
} sd_hires_params_t;
|
||||
|
||||
typedef struct {
|
||||
@@ -381,6 +385,7 @@ typedef struct {
|
||||
float vace_strength;
|
||||
sd_tiling_params_t vae_tiling_params;
|
||||
sd_cache_params_t cache;
|
||||
sd_hires_params_t hires;
|
||||
} sd_vid_gen_params_t;
|
||||
|
||||
typedef struct sd_ctx_t sd_ctx_t;
|
||||
|
||||
+2
-1
@@ -598,7 +598,8 @@ namespace Anima {
|
||||
{},
|
||||
empty_ref_latents,
|
||||
false,
|
||||
1.0f);
|
||||
1.0f,
|
||||
false);
|
||||
|
||||
std::vector<float> axis_thetas = {
|
||||
static_cast<float>(theta) * calc_ntk_factor(t_extrapolation_ratio, axes_dim[0]),
|
||||
|
||||
@@ -680,7 +680,7 @@ struct AutoEncoderKL : public VAE {
|
||||
} else if (sd_version_is_sd3(version)) {
|
||||
scale_factor = 1.5305f;
|
||||
shift_factor = 0.0609f;
|
||||
} else if (sd_version_is_flux(version) || sd_version_is_z_image(version)) {
|
||||
} else if (sd_version_is_flux(version) || sd_version_is_z_image(version) || sd_version_is_longcat(version)) {
|
||||
scale_factor = 0.3611f;
|
||||
shift_factor = 0.1159f;
|
||||
} else if (sd_version_uses_flux2_vae(version)) {
|
||||
|
||||
+110
-5
@@ -1696,11 +1696,15 @@ struct LLMEmbedder : public Conditioner {
|
||||
arch = LLM::LLMArch::MISTRAL_SMALL_3_2;
|
||||
} else if (sd_version_is_ernie_image(version)) {
|
||||
arch = LLM::LLMArch::MINISTRAL_3_3B;
|
||||
} else if (sd_version_is_lens(version)) {
|
||||
arch = LLM::LLMArch::GPT_OSS_20B;
|
||||
} else if (sd_version_is_z_image(version) || version == VERSION_OVIS_IMAGE || version == VERSION_FLUX2_KLEIN) {
|
||||
arch = LLM::LLMArch::QWEN3;
|
||||
}
|
||||
if (arch == LLM::LLMArch::MISTRAL_SMALL_3_2 || arch == LLM::LLMArch::MINISTRAL_3_3B) {
|
||||
tokenizer = std::make_shared<MistralTokenizer>();
|
||||
} else if (arch == LLM::LLMArch::GPT_OSS_20B) {
|
||||
tokenizer = std::make_shared<GPTOSSTokenizer>();
|
||||
} else {
|
||||
tokenizer = std::make_shared<Qwen2Tokenizer>();
|
||||
}
|
||||
@@ -1747,7 +1751,8 @@ struct LLMEmbedder : public Conditioner {
|
||||
std::tuple<std::vector<int>, std::vector<float>, std::vector<float>> tokenize(std::string text,
|
||||
const std::pair<int, int>& attn_range,
|
||||
size_t min_length = 0,
|
||||
size_t max_length = 100000000) {
|
||||
size_t max_length = 100000000,
|
||||
bool spell_quotes = false) {
|
||||
std::vector<std::pair<std::string, float>> parsed_attention;
|
||||
if (attn_range.first >= 0 && attn_range.second > 0) {
|
||||
if (attn_range.first > 0) {
|
||||
@@ -1755,6 +1760,9 @@ struct LLMEmbedder : public Conditioner {
|
||||
}
|
||||
if (attn_range.second - attn_range.first > 0) {
|
||||
auto new_parsed_attention = parse_prompt_attention(text.substr(attn_range.first, attn_range.second - attn_range.first));
|
||||
if (spell_quotes) {
|
||||
new_parsed_attention = split_quotation_attention(new_parsed_attention);
|
||||
}
|
||||
parsed_attention.insert(parsed_attention.end(),
|
||||
new_parsed_attention.begin(),
|
||||
new_parsed_attention.end());
|
||||
@@ -1804,8 +1812,10 @@ struct LLMEmbedder : public Conditioner {
|
||||
int hidden_states_min_length,
|
||||
const std::vector<std::pair<int, sd::Tensor<float>>>& image_embeds,
|
||||
const std::set<int>& out_layers,
|
||||
int prompt_template_encode_start_idx) {
|
||||
auto tokens_weights_mask = tokenize(prompt, prompt_attn_range, min_length);
|
||||
int prompt_template_encode_start_idx,
|
||||
bool spell_quotes = false,
|
||||
int max_length = 100000000) {
|
||||
auto tokens_weights_mask = tokenize(prompt, prompt_attn_range, min_length, max_length, spell_quotes);
|
||||
auto& tokens = std::get<0>(tokens_weights_mask);
|
||||
auto& weights = std::get<1>(tokens_weights_mask);
|
||||
auto& mask = std::get<2>(tokens_weights_mask);
|
||||
@@ -1865,7 +1875,9 @@ struct LLMEmbedder : public Conditioner {
|
||||
std::vector<std::pair<int, sd::Tensor<float>>> image_embeds;
|
||||
int prompt_template_encode_start_idx = 34;
|
||||
int min_length = 0; // pad tokens
|
||||
int max_length = 100000000;
|
||||
int hidden_states_min_length = 0; // zero pad hidden_states
|
||||
bool spell_quotes = false;
|
||||
std::set<int> out_layers;
|
||||
|
||||
int64_t t0 = ggml_time_ms();
|
||||
@@ -1938,6 +1950,71 @@ struct LLMEmbedder : public Conditioner {
|
||||
|
||||
prompt += "<|im_end|>\n<|im_start|>assistant\n";
|
||||
}
|
||||
} else if (sd_version_is_longcat(version)) {
|
||||
spell_quotes = true;
|
||||
|
||||
if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
|
||||
LOG_INFO("LongCatEditPipeline");
|
||||
prompt_template_encode_start_idx = 67;
|
||||
min_length = 512 + prompt_template_encode_start_idx;
|
||||
int image_embed_idx = 36 + 6;
|
||||
|
||||
int min_pixels = 384 * 384;
|
||||
int max_pixels = 560 * 560;
|
||||
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->params.vision.patch_size * llm->params.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);
|
||||
|
||||
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);
|
||||
}
|
||||
|
||||
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);
|
||||
GGML_ASSERT(!image_embed.empty());
|
||||
image_embeds.emplace_back(image_embed_idx, image_embed);
|
||||
image_embed_idx += 1 + static_cast<int>(image_embed.shape()[1]) + 6;
|
||||
|
||||
img_prompt += "<|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++) {
|
||||
img_prompt += placeholder;
|
||||
}
|
||||
img_prompt += "<|vision_end|>";
|
||||
}
|
||||
|
||||
prompt = "<|im_start|>system\nAs an image editing expert, first analyze the content and attributes of the input image(s). Then, based on the user's editing instructions, clearly and precisely determine how to modify the given image(s), ensuring that only the specified parts are altered and all other aspects remain consistent with the original(s).<|im_end|>\n<|im_start|>user\n";
|
||||
prompt += img_prompt;
|
||||
} else {
|
||||
prompt_template_encode_start_idx = 36;
|
||||
min_length = 512 + prompt_template_encode_start_idx;
|
||||
|
||||
prompt = "<|im_start|>system\nAs an image captioning expert, generate a descriptive text prompt based on an image content, suitable for input to a text-to-image model.<|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 (version == VERSION_FLUX2) {
|
||||
prompt_template_encode_start_idx = 0;
|
||||
hidden_states_min_length = 512;
|
||||
@@ -1957,6 +2034,30 @@ struct LLMEmbedder : public Conditioner {
|
||||
prompt_attn_range.first = 0;
|
||||
prompt += conditioner_params.text;
|
||||
prompt_attn_range.second = static_cast<int>(prompt.size());
|
||||
} else if (sd_version_is_lens(version)) {
|
||||
prompt_template_encode_start_idx = 97;
|
||||
min_length = 0;
|
||||
max_length = 512;
|
||||
out_layers = {6, 12, 18, 24};
|
||||
|
||||
prompt =
|
||||
"<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI.\n"
|
||||
"Knowledge cutoff: 2024-06\n"
|
||||
"Current date: 2026-05-26\n" // fix for current date
|
||||
"\n"
|
||||
"Reasoning: medium\n"
|
||||
"\n"
|
||||
"# Valid channels: analysis, commentary, final. Channel must be included for every message.<|end|><|start|>developer<|message|># Instructions\n"
|
||||
"\n"
|
||||
"Describe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background.\n"
|
||||
"\n"
|
||||
"<|end|><|start|>user<|message|>";
|
||||
|
||||
prompt_attn_range.first = static_cast<int>(prompt.size());
|
||||
prompt += conditioner_params.text;
|
||||
prompt_attn_range.second = static_cast<int>(prompt.size());
|
||||
|
||||
prompt += "<|end|><|start|>assistant<|channel|>analysis<|message|>Need to generate one image according to the description.<|end|><|start|>assistant<|channel|>final<|message|>";
|
||||
} else if (sd_version_is_z_image(version)) {
|
||||
prompt_template_encode_start_idx = 0;
|
||||
out_layers = {35}; // -2
|
||||
@@ -2012,7 +2113,9 @@ struct LLMEmbedder : public Conditioner {
|
||||
hidden_states_min_length,
|
||||
image_embeds,
|
||||
out_layers,
|
||||
prompt_template_encode_start_idx);
|
||||
prompt_template_encode_start_idx,
|
||||
spell_quotes,
|
||||
max_length);
|
||||
std::vector<sd::Tensor<float>> extra_hidden_states_vec;
|
||||
for (int i = 0; i < extra_prompts.size(); i++) {
|
||||
auto extra_hidden_states = encode_prompt(n_threads,
|
||||
@@ -2022,7 +2125,9 @@ struct LLMEmbedder : public Conditioner {
|
||||
hidden_states_min_length,
|
||||
image_embeds,
|
||||
out_layers,
|
||||
prompt_template_encode_start_idx);
|
||||
prompt_template_encode_start_idx,
|
||||
spell_quotes,
|
||||
max_length);
|
||||
extra_hidden_states_vec.push_back(std::move(extra_hidden_states));
|
||||
}
|
||||
|
||||
|
||||
+23
-135
@@ -496,84 +496,26 @@ struct LTX2Scheduler : SigmaScheduler {
|
||||
parse_extra_sample_args(extra_sample_args);
|
||||
}
|
||||
|
||||
static std::string trim(std::string value) {
|
||||
const char* whitespace = " \t\r\n";
|
||||
size_t begin = value.find_first_not_of(whitespace);
|
||||
if (begin == std::string::npos) {
|
||||
return "";
|
||||
}
|
||||
size_t end = value.find_last_not_of(whitespace);
|
||||
return value.substr(begin, end - begin + 1);
|
||||
}
|
||||
|
||||
void parse_extra_sample_args(const char* extra_sample_args) {
|
||||
if (extra_sample_args == nullptr || extra_sample_args[0] == '\0') {
|
||||
return;
|
||||
}
|
||||
|
||||
std::string raw(extra_sample_args);
|
||||
size_t start = 0;
|
||||
auto parse_arg = [&](const std::string& item) {
|
||||
std::string token = trim(item);
|
||||
if (token.empty()) {
|
||||
return;
|
||||
}
|
||||
size_t eq = token.find('=');
|
||||
if (eq == std::string::npos) {
|
||||
LOG_WARN("ignoring invalid ltx2 scheduler arg '%s'", token.c_str());
|
||||
return;
|
||||
}
|
||||
|
||||
std::string key = trim(token.substr(0, eq));
|
||||
std::string value = trim(token.substr(eq + 1));
|
||||
auto parse_float = [&](float* out) -> bool {
|
||||
try {
|
||||
size_t consumed = 0;
|
||||
float parsed = std::stof(value, &consumed);
|
||||
if (!trim(value.substr(consumed)).empty()) {
|
||||
return false;
|
||||
}
|
||||
*out = parsed;
|
||||
return true;
|
||||
} catch (const std::exception&) {
|
||||
return false;
|
||||
for (const auto& [key, value] : parse_key_value_args(extra_sample_args, "ltx2 scheduler arg")) {
|
||||
if (key == "max_shift") {
|
||||
if (!parse_strict_float(value, max_shift)) {
|
||||
LOG_WARN("ignoring invalid ltx2 scheduler arg '%s=%s'", key.c_str(), value.c_str());
|
||||
}
|
||||
};
|
||||
try {
|
||||
if (key == "max_shift") {
|
||||
if (!parse_float(&max_shift)) {
|
||||
LOG_WARN("ignoring invalid ltx2 scheduler arg '%s'", token.c_str());
|
||||
}
|
||||
} else if (key == "base_shift") {
|
||||
if (!parse_float(&base_shift)) {
|
||||
LOG_WARN("ignoring invalid ltx2 scheduler arg '%s'", token.c_str());
|
||||
}
|
||||
} else if (key == "terminal") {
|
||||
if (!parse_float(&terminal)) {
|
||||
LOG_WARN("ignoring invalid ltx2 scheduler arg '%s'", token.c_str());
|
||||
}
|
||||
} else if (key == "stretch") {
|
||||
std::string v = value;
|
||||
std::transform(v.begin(), v.end(), v.begin(), [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
|
||||
if (v == "1" || v == "true" || v == "yes" || v == "on") {
|
||||
stretch = true;
|
||||
} else if (v == "0" || v == "false" || v == "no" || v == "off") {
|
||||
stretch = false;
|
||||
} else {
|
||||
LOG_WARN("ignoring invalid ltx2 scheduler arg '%s'", token.c_str());
|
||||
}
|
||||
} else {
|
||||
LOG_WARN("ignoring unknown ltx2 scheduler arg '%s'", key.c_str());
|
||||
} else if (key == "base_shift") {
|
||||
if (!parse_strict_float(value, base_shift)) {
|
||||
LOG_WARN("ignoring invalid ltx2 scheduler arg '%s=%s'", key.c_str(), value.c_str());
|
||||
}
|
||||
} catch (const std::exception&) {
|
||||
LOG_WARN("ignoring invalid ltx2 scheduler arg '%s'", token.c_str());
|
||||
}
|
||||
};
|
||||
|
||||
for (size_t pos = 0; pos <= raw.size(); ++pos) {
|
||||
if (pos == raw.size() || raw[pos] == ',' || raw[pos] == ';') {
|
||||
parse_arg(raw.substr(start, pos - start));
|
||||
start = pos + 1;
|
||||
} else if (key == "terminal") {
|
||||
if (!parse_strict_float(value, terminal)) {
|
||||
LOG_WARN("ignoring invalid ltx2 scheduler arg '%s=%s'", key.c_str(), value.c_str());
|
||||
}
|
||||
} else if (key == "stretch") {
|
||||
if (!parse_strict_bool(value, stretch)) {
|
||||
LOG_WARN("ignoring invalid ltx2 scheduler arg '%s=%s'", key.c_str(), value.c_str());
|
||||
}
|
||||
} else {
|
||||
LOG_WARN("ignoring unknown ltx2 scheduler arg '%s'", key.c_str());
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1276,7 +1218,7 @@ static sd::Tensor<float> sample_dpmpp_2m_v2(denoise_cb_t model,
|
||||
return x;
|
||||
}
|
||||
|
||||
using SamplerExtraArgs = std::vector<std::pair<std::string, std::string>>;
|
||||
using SamplerExtraArgs = KeyValueArgs;
|
||||
|
||||
static sd::Tensor<float> sample_lcm(denoise_cb_t model,
|
||||
sd::Tensor<float> x,
|
||||
@@ -1296,15 +1238,8 @@ static sd::Tensor<float> sample_lcm(denoise_cb_t model,
|
||||
|
||||
for (const auto& [key, value] : extra_sample_args) {
|
||||
float parsed = 0.0f;
|
||||
try {
|
||||
size_t consumed = 0;
|
||||
parsed = std::stof(value, &consumed);
|
||||
if (trim(value.substr(consumed)).size() != 0) {
|
||||
LOG_WARN("ignoring invalid lcm extra sample arg '%s'", key.c_str());
|
||||
continue;
|
||||
}
|
||||
} catch (const std::exception&) {
|
||||
LOG_WARN("ignoring invalid lcm extra sample arg '%s=%s'", key.c_str());
|
||||
if (!parse_strict_float(value, parsed)) {
|
||||
LOG_WARN("ignoring invalid lcm extra sample arg '%s=%s'", key.c_str(), value.c_str());
|
||||
continue;
|
||||
}
|
||||
if (key == "noise_clip_std") {
|
||||
@@ -1861,15 +1796,8 @@ static sd::Tensor<float> sample_gradient_estimation(denoise_cb_t model,
|
||||
|
||||
for (const auto& [key, value] : extra_sample_args) {
|
||||
float parsed = 0.0f;
|
||||
try {
|
||||
size_t consumed = 0;
|
||||
parsed = std::stof(value, &consumed);
|
||||
if (trim(value.substr(consumed)).size() != 0) {
|
||||
LOG_WARN("ignoring invalid euler_ge extra sample arg '%s'", key.c_str());
|
||||
continue;
|
||||
}
|
||||
} catch (const std::exception&) {
|
||||
LOG_WARN("ignoring invalid euler_ge extra sample arg '%s'", key.c_str());
|
||||
if (!parse_strict_float(value, parsed)) {
|
||||
LOG_WARN("ignoring invalid euler_ge extra sample arg '%s=%s'", key.c_str(), value.c_str());
|
||||
continue;
|
||||
}
|
||||
if (key == "gamma") {
|
||||
@@ -1916,46 +1844,6 @@ static sd::Tensor<float> sample_gradient_estimation(denoise_cb_t model,
|
||||
return x;
|
||||
}
|
||||
|
||||
static SamplerExtraArgs parse_sampler_args(const char* extra_sample_args) {
|
||||
SamplerExtraArgs pairs;
|
||||
|
||||
if (extra_sample_args == nullptr || extra_sample_args[0] == '\0') {
|
||||
return pairs;
|
||||
}
|
||||
|
||||
auto trim = [](std::string value) -> std::string {
|
||||
const char* whitespace = " \t\r\n";
|
||||
size_t begin = value.find_first_not_of(whitespace);
|
||||
if (begin == std::string::npos) {
|
||||
return "";
|
||||
}
|
||||
size_t end = value.find_last_not_of(whitespace);
|
||||
return value.substr(begin, end - begin + 1);
|
||||
};
|
||||
|
||||
std::string raw(extra_sample_args);
|
||||
size_t start = 0;
|
||||
|
||||
for (size_t pos = 0; pos <= raw.size(); ++pos) {
|
||||
if (pos == raw.size() || raw[pos] == ',' || raw[pos] == ';') {
|
||||
std::string item = raw.substr(start, pos - start);
|
||||
std::string token = trim(item);
|
||||
|
||||
if (!token.empty()) {
|
||||
size_t eq = token.find('=');
|
||||
if (eq != std::string::npos) {
|
||||
std::string key = trim(token.substr(0, eq));
|
||||
std::string value = trim(token.substr(eq + 1));
|
||||
pairs.emplace_back(std::move(key), std::move(value));
|
||||
}
|
||||
}
|
||||
start = pos + 1;
|
||||
}
|
||||
}
|
||||
|
||||
return pairs;
|
||||
}
|
||||
|
||||
// k diffusion reverse ODE: dx = (x - D(x;\sigma)) / \sigma dt; \sigma(t) = t
|
||||
static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
|
||||
denoise_cb_t model,
|
||||
@@ -1965,7 +1853,7 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
|
||||
float eta,
|
||||
bool is_flow_denoiser,
|
||||
const char* extra_sample_args) {
|
||||
SamplerExtraArgs extra_args = parse_sampler_args(extra_sample_args);
|
||||
SamplerExtraArgs extra_args = parse_key_value_args(extra_sample_args, "extra sample arg");
|
||||
switch (method) {
|
||||
case EULER_A_SAMPLE_METHOD:
|
||||
return sample_euler_ancestral(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);
|
||||
|
||||
@@ -6,6 +6,7 @@
|
||||
#include "ernie_image.hpp"
|
||||
#include "flux.hpp"
|
||||
#include "hidream_o1.hpp"
|
||||
#include "lens.hpp"
|
||||
#include "ltxv.hpp"
|
||||
#include "mmdit.hpp"
|
||||
#include "qwen_image.hpp"
|
||||
@@ -701,6 +702,72 @@ struct ErnieImageModel : public DiffusionModel {
|
||||
}
|
||||
};
|
||||
|
||||
struct LensModel : public DiffusionModel {
|
||||
std::string prefix;
|
||||
Lens::LensRunner lens;
|
||||
|
||||
LensModel(ggml_backend_t backend,
|
||||
ggml_backend_t params_backend,
|
||||
const String2TensorStorage& tensor_storage_map = {},
|
||||
const std::string prefix = "model.diffusion_model")
|
||||
: prefix(prefix), lens(backend, params_backend, tensor_storage_map, prefix) {
|
||||
}
|
||||
|
||||
std::string get_desc() override {
|
||||
return lens.get_desc();
|
||||
}
|
||||
|
||||
void alloc_params_buffer() override {
|
||||
lens.alloc_params_buffer();
|
||||
}
|
||||
|
||||
void free_params_buffer() override {
|
||||
lens.free_params_buffer();
|
||||
}
|
||||
|
||||
void free_compute_buffer() override {
|
||||
lens.free_compute_buffer();
|
||||
}
|
||||
|
||||
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
|
||||
lens.get_param_tensors(tensors, prefix);
|
||||
}
|
||||
|
||||
size_t get_params_buffer_size() override {
|
||||
return lens.get_params_buffer_size();
|
||||
}
|
||||
|
||||
void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
|
||||
lens.set_weight_adapter(adapter);
|
||||
}
|
||||
|
||||
int64_t get_adm_in_channels() override {
|
||||
return 768;
|
||||
}
|
||||
|
||||
void set_flash_attention_enabled(bool enabled) {
|
||||
lens.set_flash_attention_enabled(enabled);
|
||||
}
|
||||
|
||||
void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
|
||||
lens.set_max_graph_vram_bytes(max_vram_bytes);
|
||||
}
|
||||
|
||||
void set_circular_axes(bool circular_x, bool circular_y) override {
|
||||
lens.set_circular_axes(circular_x, circular_y);
|
||||
}
|
||||
|
||||
sd::Tensor<float> compute(int n_threads,
|
||||
const DiffusionParams& diffusion_params) override {
|
||||
GGML_ASSERT(diffusion_params.x != nullptr);
|
||||
GGML_ASSERT(diffusion_params.timesteps != nullptr);
|
||||
return lens.compute(n_threads,
|
||||
*diffusion_params.x,
|
||||
*diffusion_params.timesteps,
|
||||
tensor_or_empty(diffusion_params.context));
|
||||
}
|
||||
};
|
||||
|
||||
struct LTXAVModel : public DiffusionModel {
|
||||
std::string prefix;
|
||||
LTXV::LTXAVRunner ltxav;
|
||||
|
||||
+5
-3
@@ -446,7 +446,6 @@ namespace Flux {
|
||||
if (use_yak_mlp || use_mlp_silu_act) {
|
||||
mlp_mult_factor = 2;
|
||||
}
|
||||
|
||||
blocks["linear1"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, hidden_size * 3 + mlp_hidden_dim * mlp_mult_factor, mlp_proj_bias));
|
||||
blocks["linear2"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size + mlp_hidden_dim, hidden_size, mlp_proj_bias));
|
||||
blocks["norm"] = std::shared_ptr<GGMLBlock>(new QKNorm(head_dim));
|
||||
@@ -1225,6 +1224,9 @@ namespace Flux {
|
||||
flux_params.share_modulation = true;
|
||||
flux_params.ref_index_scale = 10.f;
|
||||
flux_params.use_mlp_silu_act = true;
|
||||
} else if (sd_version_is_longcat(version)) {
|
||||
flux_params.context_in_dim = 3584;
|
||||
flux_params.vec_in_dim = 0;
|
||||
}
|
||||
int64_t head_dim = 0;
|
||||
int64_t actual_radiance_patch_size = -1;
|
||||
@@ -1412,7 +1414,6 @@ namespace Flux {
|
||||
} else if (version == VERSION_OVIS_IMAGE) {
|
||||
txt_arange_dims = {1, 2};
|
||||
}
|
||||
|
||||
pe_vec = Rope::gen_flux_pe(static_cast<int>(x->ne[1]),
|
||||
static_cast<int>(x->ne[0]),
|
||||
flux_params.patch_size,
|
||||
@@ -1425,7 +1426,8 @@ namespace Flux {
|
||||
flux_params.theta,
|
||||
circular_y_enabled,
|
||||
circular_x_enabled,
|
||||
flux_params.axes_dim);
|
||||
flux_params.axes_dim,
|
||||
sd_version_is_longcat(version));
|
||||
int pos_len = static_cast<int>(pe_vec.size() / flux_params.axes_dim_sum / 2);
|
||||
// LOG_DEBUG("pos_len %d", pos_len);
|
||||
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, flux_params.axes_dim_sum / 2, pos_len);
|
||||
|
||||
+183
-3
@@ -953,11 +953,17 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_group_norm_32(ggml_context* ctx,
|
||||
return ggml_group_norm(ctx, a, 32, eps);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ bool ggml_ext_is_padded_1d(const ggml_tensor* x) {
|
||||
return x->nb[0] == ggml_type_size(x->type) &&
|
||||
x->nb[2] == x->nb[1] * x->ne[1] &&
|
||||
x->nb[3] == x->nb[2] * x->ne[2];
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ ggml_tensor* ggml_ext_scale(ggml_context* ctx,
|
||||
ggml_tensor* x,
|
||||
float factor,
|
||||
bool inplace = false) {
|
||||
if (!ggml_is_contiguous(x)) {
|
||||
if (!ggml_ext_is_padded_1d(x)) {
|
||||
x = ggml_cont(ctx, x);
|
||||
}
|
||||
if (inplace) {
|
||||
@@ -1602,6 +1608,23 @@ __STATIC_INLINE__ size_t ggml_tensor_num(ggml_context* ctx) {
|
||||
return num;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ ggml_tensor* ggml_ext_vec_concat(ggml_context* ctx,
|
||||
std::vector<ggml_tensor*>& tensors,
|
||||
int dim) {
|
||||
while (tensors.size() > 1) {
|
||||
std::vector<ggml_tensor*> next_level;
|
||||
for (size_t i = 0; i < tensors.size(); i += 2) {
|
||||
if (i + 1 < tensors.size()) {
|
||||
next_level.push_back(ggml_concat(ctx, tensors[i], tensors[i + 1], dim));
|
||||
} else {
|
||||
next_level.push_back(tensors[i]);
|
||||
}
|
||||
}
|
||||
tensors = std::move(next_level);
|
||||
}
|
||||
return tensors[0];
|
||||
}
|
||||
|
||||
/* SDXL with LoRA requires more space */
|
||||
#define MAX_PARAMS_TENSOR_NUM 32768
|
||||
#define MAX_GRAPH_SIZE 327680
|
||||
@@ -3139,6 +3162,163 @@ public:
|
||||
}
|
||||
};
|
||||
|
||||
class Conv2d_grouped : public UnaryBlock {
|
||||
protected:
|
||||
int64_t in_channels;
|
||||
int64_t out_channels;
|
||||
int groups;
|
||||
std::pair<int, int> kernel_size;
|
||||
std::pair<int, int> stride;
|
||||
std::pair<int, int> padding;
|
||||
std::pair<int, int> dilation;
|
||||
bool bias;
|
||||
float scale = 1.f;
|
||||
std::string prefix;
|
||||
|
||||
void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map, const std::string prefix = "") override {
|
||||
this->prefix = prefix;
|
||||
enum ggml_type wtype = GGML_TYPE_F16;
|
||||
params["weight"] = ggml_new_tensor_4d(ctx, wtype, kernel_size.second, kernel_size.first, in_channels / groups, out_channels);
|
||||
if (bias) {
|
||||
enum ggml_type wtype = GGML_TYPE_F32;
|
||||
params["bias"] = ggml_new_tensor_1d(ctx, wtype, out_channels);
|
||||
}
|
||||
}
|
||||
|
||||
public:
|
||||
Conv2d_grouped(int64_t in_channels,
|
||||
int64_t out_channels,
|
||||
int groups,
|
||||
std::pair<int, int> kernel_size,
|
||||
std::pair<int, int> stride = {1, 1},
|
||||
std::pair<int, int> padding = {0, 0},
|
||||
std::pair<int, int> dilation = {1, 1},
|
||||
bool bias = true)
|
||||
: in_channels(in_channels),
|
||||
out_channels(out_channels),
|
||||
groups(groups),
|
||||
kernel_size(kernel_size),
|
||||
stride(stride),
|
||||
padding(padding),
|
||||
dilation(dilation),
|
||||
bias(bias) {}
|
||||
|
||||
void set_scale(float scale_value) {
|
||||
scale = scale_value;
|
||||
}
|
||||
|
||||
std::string get_desc() {
|
||||
return "Conv2d_grouped";
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
ggml_tensor* w = params["weight"];
|
||||
ggml_tensor* b = nullptr;
|
||||
if (bias) {
|
||||
b = params["bias"];
|
||||
}
|
||||
|
||||
if (groups == 1) {
|
||||
if (ctx->weight_adapter) {
|
||||
WeightAdapter::ForwardParams forward_params;
|
||||
forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_CONV2D;
|
||||
forward_params.conv2d.s0 = stride.second;
|
||||
forward_params.conv2d.s1 = stride.first;
|
||||
forward_params.conv2d.p0 = padding.second;
|
||||
forward_params.conv2d.p1 = padding.first;
|
||||
forward_params.conv2d.d0 = dilation.second;
|
||||
forward_params.conv2d.d1 = dilation.first;
|
||||
forward_params.conv2d.direct = ctx->conv2d_direct_enabled;
|
||||
forward_params.conv2d.circular_x = ctx->circular_x_enabled;
|
||||
forward_params.conv2d.circular_y = ctx->circular_y_enabled;
|
||||
forward_params.conv2d.scale = scale;
|
||||
return ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x, w, b, prefix, forward_params);
|
||||
}
|
||||
return ggml_ext_conv_2d(ctx->ggml_ctx, x, w, b,
|
||||
stride.second, stride.first,
|
||||
padding.second, padding.first,
|
||||
dilation.second, dilation.first,
|
||||
ctx->conv2d_direct_enabled,
|
||||
ctx->circular_x_enabled,
|
||||
ctx->circular_y_enabled,
|
||||
scale);
|
||||
}
|
||||
|
||||
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,
|
||||
stride.second, stride.first,
|
||||
padding.second, padding.first,
|
||||
dilation.second, dilation.first);
|
||||
} else {
|
||||
res = ggml_conv_2d_dw(ctx->ggml_ctx, x, w,
|
||||
stride.second, stride.first,
|
||||
padding.second, padding.first,
|
||||
dilation.second, dilation.first);
|
||||
}
|
||||
if (b) {
|
||||
res = ggml_add(ctx->ggml_ctx, res, b);
|
||||
}
|
||||
return res;
|
||||
}
|
||||
|
||||
int64_t ic_g = in_channels / groups;
|
||||
int64_t oc_g = out_channels / groups;
|
||||
|
||||
std::vector<ggml_tensor*> out_slices(groups);
|
||||
|
||||
for (int i = 0; i < groups; ++i) {
|
||||
size_t x_offset = i * ic_g * x->nb[2];
|
||||
ggml_tensor* x_i = ggml_view_4d(ctx->ggml_ctx, x,
|
||||
x->ne[0], x->ne[1], ic_g, x->ne[3],
|
||||
x->nb[1], x->nb[2], x->nb[3],
|
||||
x_offset);
|
||||
|
||||
size_t w_offset = i * oc_g * w->nb[3];
|
||||
ggml_tensor* w_i = ggml_view_4d(ctx->ggml_ctx, w,
|
||||
w->ne[0], w->ne[1], w->ne[2], oc_g,
|
||||
w->nb[1], w->nb[2], w->nb[3],
|
||||
w_offset);
|
||||
|
||||
ggml_tensor* b_i = nullptr;
|
||||
if (b) {
|
||||
size_t b_offset = i * oc_g * b->nb[0];
|
||||
b_i = ggml_view_1d(ctx->ggml_ctx, b, oc_g, b_offset);
|
||||
}
|
||||
|
||||
if (ctx->weight_adapter) {
|
||||
WeightAdapter::ForwardParams forward_params;
|
||||
forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_CONV2D;
|
||||
forward_params.conv2d.s0 = stride.second;
|
||||
forward_params.conv2d.s1 = stride.first;
|
||||
forward_params.conv2d.p0 = padding.second;
|
||||
forward_params.conv2d.p1 = padding.first;
|
||||
forward_params.conv2d.d0 = dilation.second;
|
||||
forward_params.conv2d.d1 = dilation.first;
|
||||
forward_params.conv2d.direct = ctx->conv2d_direct_enabled;
|
||||
forward_params.conv2d.circular_x = ctx->circular_x_enabled;
|
||||
forward_params.conv2d.circular_y = ctx->circular_y_enabled;
|
||||
forward_params.conv2d.scale = scale;
|
||||
out_slices[i] = ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x_i, w_i, b_i, prefix, forward_params);
|
||||
} else {
|
||||
out_slices[i] = ggml_ext_conv_2d(ctx->ggml_ctx, x_i, w_i, b_i,
|
||||
stride.second, stride.first,
|
||||
padding.second, padding.first,
|
||||
dilation.second, dilation.first,
|
||||
ctx->conv2d_direct_enabled,
|
||||
ctx->circular_x_enabled,
|
||||
ctx->circular_y_enabled,
|
||||
scale);
|
||||
}
|
||||
}
|
||||
|
||||
ggml_tensor* out = ggml_ext_vec_concat(ctx->ggml_ctx, out_slices, 2);
|
||||
|
||||
return out;
|
||||
}
|
||||
};
|
||||
|
||||
class Conv3d : public UnaryBlock {
|
||||
protected:
|
||||
int64_t in_channels;
|
||||
@@ -3490,7 +3670,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_lokr_forward(
|
||||
|
||||
ggml_tensor* hc = ggml_transpose(ctx, hc_t);
|
||||
ggml_tensor* out = ggml_reshape_2d(ctx, ggml_cont(ctx, hc), up * vp, batch);
|
||||
return ggml_scale(ctx, out, scale);
|
||||
return ggml_ext_scale(ctx, out, scale);
|
||||
} else {
|
||||
int batch = (int)h->ne[3];
|
||||
// 1. Reshape input: [W, H, vq*uq, batch] -> [W, H, vq, uq * batch]
|
||||
@@ -3573,7 +3753,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_lokr_forward(
|
||||
ggml_tensor* hc = ggml_transpose(ctx, hc_t);
|
||||
// ungroup
|
||||
ggml_tensor* out = ggml_reshape_4d(ctx, ggml_cont(ctx, hc), w_out, h_out, up * vp, batch);
|
||||
return ggml_scale(ctx, out, scale);
|
||||
return ggml_ext_scale(ctx, out, scale);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+14
-13
@@ -16,8 +16,7 @@
|
||||
|
||||
namespace sd::ggml_graph_cut {
|
||||
|
||||
static constexpr double MAX_VRAM_BYTES_PER_GIB = 1024.0 * 1024.0 * 1024.0;
|
||||
static constexpr size_t MAX_VRAM_AUTO_RESERVE_BYTES = 1024ULL * 1024ULL * 1024ULL;
|
||||
static constexpr double MAX_VRAM_BYTES_PER_GIB = 1024.0 * 1024.0 * 1024.0;
|
||||
|
||||
static std::string graph_cut_tensor_display_name(const ggml_tensor* tensor) {
|
||||
if (tensor == nullptr) {
|
||||
@@ -93,45 +92,47 @@ namespace sd::ggml_graph_cut {
|
||||
return static_cast<float>(static_cast<double>(max_vram_bytes) / MAX_VRAM_BYTES_PER_GIB);
|
||||
}
|
||||
|
||||
static size_t resolve_auto_max_vram_bytes(ggml_backend_t backend) {
|
||||
static size_t resolve_auto_max_vram_bytes(float spare_vram, ggml_backend_t backend) {
|
||||
if (backend == nullptr) {
|
||||
LOG_WARN("--max-vram -1 requested, but no backend is available; disabling graph splitting");
|
||||
LOG_WARN("--max-vram < 0 requested, but no backend is available; disabling graph splitting");
|
||||
return 0;
|
||||
}
|
||||
|
||||
ggml_backend_dev_t dev = ggml_backend_get_device(backend);
|
||||
if (dev == nullptr) {
|
||||
LOG_WARN("--max-vram -1 requested, but no backend device is available; disabling graph splitting");
|
||||
LOG_WARN("--max-vram < 0 requested, but no backend device is available; disabling graph splitting");
|
||||
return 0;
|
||||
}
|
||||
if (ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU) {
|
||||
LOG_WARN("--max-vram -1 requested, but the main backend is CPU; disabling graph splitting");
|
||||
LOG_WARN("--max-vram < 0 requested, but the main backend is CPU; disabling graph splitting");
|
||||
return 0;
|
||||
}
|
||||
|
||||
size_t free_vram = 0;
|
||||
size_t total_vram = 0;
|
||||
ggml_backend_dev_memory(dev, &free_vram, &total_vram);
|
||||
size_t spare_bytes = static_cast<size_t>(MAX_VRAM_BYTES_PER_GIB * spare_vram);
|
||||
|
||||
if (free_vram <= MAX_VRAM_AUTO_RESERVE_BYTES) {
|
||||
LOG_WARN("--max-vram -1 requested, but free VRAM is %.2f GiB; reserving 1.00 GiB leaves no graph budget",
|
||||
free_vram / MAX_VRAM_BYTES_PER_GIB);
|
||||
if (free_vram <= spare_bytes) {
|
||||
LOG_WARN("--max-vram < 0 requested, but free VRAM is %.2f GiB; reserving %.2f GiB leaves no graph budget",
|
||||
free_vram / MAX_VRAM_BYTES_PER_GIB, spare_vram);
|
||||
return 0;
|
||||
}
|
||||
|
||||
const size_t max_vram_bytes = free_vram - MAX_VRAM_AUTO_RESERVE_BYTES;
|
||||
LOG_INFO("--max-vram -1 auto-detected %.2f GiB free VRAM (%.2f GiB total), reserving 1.00 GiB; using %.2f GiB",
|
||||
const size_t max_vram_bytes = free_vram - spare_bytes;
|
||||
LOG_INFO("--max-vram < 0 auto-detected %.2f GiB free VRAM (%.2f GiB total), reserving %.2f GiB; using %.2f GiB",
|
||||
free_vram / MAX_VRAM_BYTES_PER_GIB,
|
||||
total_vram / MAX_VRAM_BYTES_PER_GIB,
|
||||
spare_vram,
|
||||
max_vram_bytes / MAX_VRAM_BYTES_PER_GIB);
|
||||
return max_vram_bytes;
|
||||
}
|
||||
|
||||
float resolve_max_vram_gib(float max_vram, ggml_backend_t backend) {
|
||||
if (max_vram != -1.f) {
|
||||
if (max_vram >= 0.f) {
|
||||
return max_vram;
|
||||
}
|
||||
return max_vram_bytes_to_gib(resolve_auto_max_vram_bytes(backend));
|
||||
return max_vram_bytes_to_gib(resolve_auto_max_vram_bytes(-max_vram, backend));
|
||||
}
|
||||
|
||||
static Segment make_segment_seed(const Plan& plan,
|
||||
|
||||
+408
@@ -0,0 +1,408 @@
|
||||
#ifndef __SD_LENS_HPP__
|
||||
#define __SD_LENS_HPP__
|
||||
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
#include "common_block.hpp"
|
||||
#include "flux.hpp"
|
||||
#include "qwen_image.hpp"
|
||||
#include "rope.hpp"
|
||||
|
||||
namespace Lens {
|
||||
constexpr int LENS_GRAPH_SIZE = 40960;
|
||||
|
||||
struct LensTimestepProjEmbeddings : public GGMLBlock {
|
||||
LensTimestepProjEmbeddings(int64_t embedding_dim) {
|
||||
blocks["timestep_embedder"] = std::make_shared<Qwen::TimestepEmbedding>(256, embedding_dim);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* timesteps) {
|
||||
auto timestep_embedder = std::dynamic_pointer_cast<Qwen::TimestepEmbedding>(blocks["timestep_embedder"]);
|
||||
auto timesteps_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, 256, 10000, 1000.f);
|
||||
return timestep_embedder->forward(ctx, timesteps_proj);
|
||||
}
|
||||
};
|
||||
|
||||
struct LensGateMLP : public GGMLBlock {
|
||||
LensGateMLP(int64_t dim, int64_t hidden_dim) {
|
||||
blocks["w1"] = std::make_shared<Linear>(dim, hidden_dim, false);
|
||||
blocks["w2"] = std::make_shared<Linear>(hidden_dim, dim, false);
|
||||
blocks["w3"] = std::make_shared<Linear>(dim, hidden_dim, false);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
auto w1 = std::dynamic_pointer_cast<Linear>(blocks["w1"]);
|
||||
auto w2 = std::dynamic_pointer_cast<Linear>(blocks["w2"]);
|
||||
auto w3 = std::dynamic_pointer_cast<Linear>(blocks["w3"]);
|
||||
|
||||
auto gate = ggml_silu(ctx->ggml_ctx, w1->forward(ctx, x));
|
||||
auto up = w3->forward(ctx, x);
|
||||
x = ggml_mul(ctx->ggml_ctx, gate, up);
|
||||
return w2->forward(ctx, x);
|
||||
}
|
||||
};
|
||||
|
||||
struct LensJointAttention : public GGMLBlock {
|
||||
int64_t dim_head;
|
||||
int64_t num_heads;
|
||||
|
||||
LensJointAttention(int64_t query_dim,
|
||||
int64_t dim_head,
|
||||
int64_t num_heads,
|
||||
float eps = 1e-5f)
|
||||
: dim_head(dim_head), num_heads(num_heads) {
|
||||
int64_t inner_dim = dim_head * num_heads;
|
||||
blocks["img_qkv"] = std::make_shared<Linear>(query_dim, inner_dim * 3, true);
|
||||
blocks["txt_qkv"] = std::make_shared<Linear>(query_dim, inner_dim * 3, true);
|
||||
|
||||
blocks["norm_q"] = std::make_shared<RMSNorm>(dim_head, eps);
|
||||
blocks["norm_k"] = std::make_shared<RMSNorm>(dim_head, eps);
|
||||
blocks["norm_added_q"] = std::make_shared<RMSNorm>(dim_head, eps);
|
||||
blocks["norm_added_k"] = std::make_shared<RMSNorm>(dim_head, eps);
|
||||
|
||||
blocks["to_out.0"] = std::make_shared<Linear>(inner_dim, query_dim, true);
|
||||
blocks["to_add_out"] = std::make_shared<Linear>(inner_dim, query_dim, true);
|
||||
}
|
||||
|
||||
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* img,
|
||||
ggml_tensor* txt,
|
||||
ggml_tensor* pe,
|
||||
ggml_tensor* mask = nullptr) {
|
||||
auto img_qkv = std::dynamic_pointer_cast<Linear>(blocks["img_qkv"]);
|
||||
auto txt_qkv = std::dynamic_pointer_cast<Linear>(blocks["txt_qkv"]);
|
||||
auto norm_q = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_q"]);
|
||||
auto norm_k = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_k"]);
|
||||
auto norm_add_q = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_added_q"]);
|
||||
auto norm_add_k = std::dynamic_pointer_cast<RMSNorm>(blocks["norm_added_k"]);
|
||||
auto to_out_0 = std::dynamic_pointer_cast<Linear>(blocks["to_out.0"]);
|
||||
auto to_add_out = std::dynamic_pointer_cast<Linear>(blocks["to_add_out"]);
|
||||
int64_t n_img = img->ne[1];
|
||||
int64_t n_txt = txt->ne[1];
|
||||
int64_t N = img->ne[2];
|
||||
int64_t inner = dim_head * num_heads;
|
||||
|
||||
auto img_qkv_vec = split_qkv(ctx->ggml_ctx, img_qkv->forward(ctx, img));
|
||||
auto txt_qkv_vec = split_qkv(ctx->ggml_ctx, txt_qkv->forward(ctx, txt));
|
||||
|
||||
auto img_q = ggml_reshape_4d(ctx->ggml_ctx, img_qkv_vec[0], dim_head, num_heads, n_img, N);
|
||||
auto img_k = ggml_reshape_4d(ctx->ggml_ctx, img_qkv_vec[1], dim_head, num_heads, n_img, N);
|
||||
auto img_v = ggml_reshape_4d(ctx->ggml_ctx, img_qkv_vec[2], dim_head, num_heads, n_img, N);
|
||||
|
||||
img_q = norm_q->forward(ctx, img_q);
|
||||
img_k = norm_k->forward(ctx, img_k);
|
||||
|
||||
auto txt_q = ggml_reshape_4d(ctx->ggml_ctx, txt_qkv_vec[0], dim_head, num_heads, n_txt, N);
|
||||
auto txt_k = ggml_reshape_4d(ctx->ggml_ctx, txt_qkv_vec[1], dim_head, num_heads, n_txt, N);
|
||||
auto txt_v = ggml_reshape_4d(ctx->ggml_ctx, txt_qkv_vec[2], dim_head, num_heads, n_txt, N);
|
||||
|
||||
txt_q = norm_add_q->forward(ctx, txt_q);
|
||||
txt_k = norm_add_k->forward(ctx, txt_k);
|
||||
|
||||
auto q = ggml_concat(ctx->ggml_ctx, img_q, txt_q, 2);
|
||||
auto k = ggml_concat(ctx->ggml_ctx, img_k, txt_k, 2);
|
||||
auto v = ggml_concat(ctx->ggml_ctx, img_v, txt_v, 2);
|
||||
|
||||
auto attn = Rope::attention(ctx, q, k, v, pe, mask, (1.0f / 128.f));
|
||||
|
||||
auto img_attn_out = ggml_view_3d(ctx->ggml_ctx,
|
||||
attn,
|
||||
inner,
|
||||
n_img,
|
||||
N,
|
||||
attn->nb[1],
|
||||
attn->nb[2],
|
||||
0);
|
||||
auto txt_attn_out = ggml_view_3d(ctx->ggml_ctx,
|
||||
attn,
|
||||
inner,
|
||||
n_txt,
|
||||
N,
|
||||
attn->nb[1],
|
||||
attn->nb[2],
|
||||
n_img * attn->nb[1]);
|
||||
|
||||
img_attn_out = to_out_0->forward(ctx, ggml_cont(ctx->ggml_ctx, img_attn_out));
|
||||
txt_attn_out = to_add_out->forward(ctx, ggml_cont(ctx->ggml_ctx, txt_attn_out));
|
||||
return {img_attn_out, txt_attn_out};
|
||||
}
|
||||
};
|
||||
|
||||
struct LensTransformerBlock : public GGMLBlock {
|
||||
LensTransformerBlock(int64_t dim,
|
||||
int64_t num_attention_heads,
|
||||
int64_t attention_head_dim,
|
||||
float eps = 1e-6f) {
|
||||
int64_t mlp_hidden_dim = dim / 3 * 8;
|
||||
blocks["img_mod.1"] = std::make_shared<Linear>(dim, 6 * dim, true);
|
||||
blocks["txt_mod.1"] = std::make_shared<Linear>(dim, 6 * dim, true);
|
||||
blocks["img_norm1"] = std::make_shared<RMSNorm>(dim, eps);
|
||||
blocks["img_norm2"] = std::make_shared<RMSNorm>(dim, eps);
|
||||
blocks["txt_norm1"] = std::make_shared<RMSNorm>(dim, eps);
|
||||
blocks["txt_norm2"] = std::make_shared<RMSNorm>(dim, eps);
|
||||
blocks["img_mlp"] = std::make_shared<LensGateMLP>(dim, mlp_hidden_dim);
|
||||
blocks["txt_mlp"] = std::make_shared<LensGateMLP>(dim, mlp_hidden_dim);
|
||||
blocks["attn"] = std::make_shared<LensJointAttention>(dim, attention_head_dim, num_attention_heads);
|
||||
}
|
||||
|
||||
std::pair<ggml_tensor*, ggml_tensor*> forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* img,
|
||||
ggml_tensor* txt,
|
||||
ggml_tensor* t_emb,
|
||||
ggml_tensor* pe) {
|
||||
auto img_mod_1 = std::dynamic_pointer_cast<Linear>(blocks["img_mod.1"]);
|
||||
auto txt_mod_1 = std::dynamic_pointer_cast<Linear>(blocks["txt_mod.1"]);
|
||||
auto img_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["img_norm1"]);
|
||||
auto img_norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["img_norm2"]);
|
||||
auto txt_norm1 = std::dynamic_pointer_cast<RMSNorm>(blocks["txt_norm1"]);
|
||||
auto txt_norm2 = std::dynamic_pointer_cast<RMSNorm>(blocks["txt_norm2"]);
|
||||
auto img_mlp = std::dynamic_pointer_cast<LensGateMLP>(blocks["img_mlp"]);
|
||||
auto txt_mlp = std::dynamic_pointer_cast<LensGateMLP>(blocks["txt_mlp"]);
|
||||
auto attn = std::dynamic_pointer_cast<LensJointAttention>(blocks["attn"]);
|
||||
|
||||
auto temb = ggml_silu(ctx->ggml_ctx, t_emb);
|
||||
|
||||
auto img_mod_params = img_mod_1->forward(ctx, temb);
|
||||
auto img_mod_vec = ggml_ext_chunk(ctx->ggml_ctx, img_mod_params, 6, 0);
|
||||
auto txt_mod_params = txt_mod_1->forward(ctx, temb);
|
||||
auto txt_mod_vec = ggml_ext_chunk(ctx->ggml_ctx, txt_mod_params, 6, 0);
|
||||
|
||||
auto img_normed = img_norm1->forward(ctx, img);
|
||||
auto img_modulated = Flux::modulate(ctx->ggml_ctx, img_normed, img_mod_vec[0], img_mod_vec[1]);
|
||||
auto txt_normed = txt_norm1->forward(ctx, txt);
|
||||
auto txt_modulated = Flux::modulate(ctx->ggml_ctx, txt_normed, txt_mod_vec[0], txt_mod_vec[1]);
|
||||
|
||||
auto [img_attn_output, txt_attn_output] = attn->forward(ctx, img_modulated, txt_modulated, pe);
|
||||
|
||||
img = ggml_add(ctx->ggml_ctx, img, ggml_mul(ctx->ggml_ctx, img_attn_output, img_mod_vec[2]));
|
||||
txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, txt_attn_output, txt_mod_vec[2]));
|
||||
|
||||
auto img_normed2 = img_norm2->forward(ctx, img);
|
||||
auto img_modulated2 = Flux::modulate(ctx->ggml_ctx, img_normed2, img_mod_vec[3], img_mod_vec[4]);
|
||||
auto txt_normed2 = txt_norm2->forward(ctx, txt);
|
||||
auto txt_modulated2 = Flux::modulate(ctx->ggml_ctx, txt_normed2, txt_mod_vec[3], txt_mod_vec[4]);
|
||||
|
||||
img = ggml_add(ctx->ggml_ctx, img, ggml_mul(ctx->ggml_ctx, img_mlp->forward(ctx, img_modulated2), img_mod_vec[5]));
|
||||
txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, txt_mlp->forward(ctx, txt_modulated2), txt_mod_vec[5]));
|
||||
return {img, txt};
|
||||
}
|
||||
};
|
||||
|
||||
struct LensAdaLayerNormContinuous : public GGMLBlock {
|
||||
int64_t hidden_size;
|
||||
float eps;
|
||||
|
||||
LensAdaLayerNormContinuous(int64_t hidden_size, float eps = 1e-6f)
|
||||
: hidden_size(hidden_size), eps(eps) {
|
||||
blocks["linear"] = std::make_shared<Linear>(hidden_size, hidden_size * 2, true);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* conditioning) {
|
||||
auto linear = std::dynamic_pointer_cast<Linear>(blocks["linear"]);
|
||||
auto mods = ggml_ext_chunk(ctx->ggml_ctx, linear->forward(ctx, ggml_silu(ctx->ggml_ctx, conditioning)), 2, 0);
|
||||
auto scale = mods[0];
|
||||
auto shift = mods[1];
|
||||
x = ggml_norm(ctx->ggml_ctx, x, eps);
|
||||
return Flux::modulate(ctx->ggml_ctx, x, shift, scale);
|
||||
}
|
||||
};
|
||||
|
||||
struct LensParams {
|
||||
int patch_size = 2;
|
||||
int64_t in_channels = 128;
|
||||
int64_t out_channels = 32;
|
||||
int num_layers = 48;
|
||||
int64_t attention_head_dim = 64;
|
||||
int64_t num_attention_heads = 24;
|
||||
int64_t joint_attention_dim = 2880;
|
||||
int selected_layer_count = 4;
|
||||
int theta = 10000;
|
||||
std::vector<int> axes_dim = {8, 28, 28};
|
||||
int axes_dim_sum = 64;
|
||||
};
|
||||
|
||||
class LensModel : public GGMLBlock {
|
||||
public:
|
||||
LensParams params;
|
||||
|
||||
LensModel() = default;
|
||||
LensModel(LensParams params)
|
||||
: params(params) {
|
||||
int64_t inner_dim = params.num_attention_heads * params.attention_head_dim;
|
||||
blocks["time_text_embed"] = std::make_shared<LensTimestepProjEmbeddings>(inner_dim);
|
||||
blocks["img_in"] = std::make_shared<Linear>(params.in_channels, inner_dim, true);
|
||||
blocks["txt_in"] = std::make_shared<Linear>(params.joint_attention_dim * params.selected_layer_count, inner_dim, true);
|
||||
for (int i = 0; i < params.selected_layer_count; ++i) {
|
||||
blocks["txt_norm." + std::to_string(i)] = std::make_shared<RMSNorm>(params.joint_attention_dim, 1e-5f);
|
||||
}
|
||||
for (int i = 0; i < params.num_layers; ++i) {
|
||||
blocks["transformer_blocks." + std::to_string(i)] = std::make_shared<LensTransformerBlock>(inner_dim,
|
||||
params.num_attention_heads,
|
||||
params.attention_head_dim);
|
||||
}
|
||||
blocks["norm_out"] = std::make_shared<LensAdaLayerNormContinuous>(inner_dim, 1e-6f);
|
||||
blocks["proj_out"] = std::make_shared<Linear>(inner_dim, params.patch_size * params.patch_size * params.out_channels, true);
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* timestep,
|
||||
ggml_tensor* context,
|
||||
ggml_tensor* pe) {
|
||||
GGML_ASSERT(context != nullptr);
|
||||
int64_t W = x->ne[0];
|
||||
int64_t H = x->ne[1];
|
||||
int64_t C = x->ne[2];
|
||||
int64_t N = x->ne[3];
|
||||
|
||||
auto time_text_embed = std::dynamic_pointer_cast<LensTimestepProjEmbeddings>(blocks["time_text_embed"]);
|
||||
auto img_in = std::dynamic_pointer_cast<Linear>(blocks["img_in"]);
|
||||
auto txt_in = std::dynamic_pointer_cast<Linear>(blocks["txt_in"]);
|
||||
auto norm_out = std::dynamic_pointer_cast<LensAdaLayerNormContinuous>(blocks["norm_out"]);
|
||||
auto proj_out = std::dynamic_pointer_cast<Linear>(blocks["proj_out"]);
|
||||
|
||||
auto t_emb = time_text_embed->forward(ctx, timestep);
|
||||
|
||||
auto img = ggml_reshape_3d(ctx->ggml_ctx, x, W * H, C, N);
|
||||
img = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, img, 1, 0, 2, 3));
|
||||
img = img_in->forward(ctx, img);
|
||||
|
||||
std::vector<ggml_tensor*> txt_chunks = ggml_ext_chunk(ctx->ggml_ctx, context, params.selected_layer_count, 0);
|
||||
ggml_tensor* txt = nullptr;
|
||||
for (int i = 0; i < params.selected_layer_count; ++i) {
|
||||
auto txt_norm = std::dynamic_pointer_cast<RMSNorm>(blocks["txt_norm." + std::to_string(i)]);
|
||||
auto chunk = txt_norm->forward(ctx, txt_chunks[i]);
|
||||
txt = txt == nullptr ? chunk : ggml_concat(ctx->ggml_ctx, txt, chunk, 0);
|
||||
}
|
||||
txt = txt_in->forward(ctx, txt);
|
||||
|
||||
sd::ggml_graph_cut::mark_graph_cut(img, "lens.prelude", "img");
|
||||
sd::ggml_graph_cut::mark_graph_cut(txt, "lens.prelude", "txt");
|
||||
|
||||
for (int i = 0; i < params.num_layers; ++i) {
|
||||
auto block = std::dynamic_pointer_cast<LensTransformerBlock>(blocks["transformer_blocks." + std::to_string(i)]);
|
||||
auto out = block->forward(ctx, img, txt, t_emb, pe);
|
||||
img = out.first;
|
||||
txt = out.second;
|
||||
sd::ggml_graph_cut::mark_graph_cut(img, "lens.transformer_blocks." + std::to_string(i), "img");
|
||||
sd::ggml_graph_cut::mark_graph_cut(txt, "lens.transformer_blocks." + std::to_string(i), "txt");
|
||||
}
|
||||
|
||||
img = norm_out->forward(ctx, img, t_emb);
|
||||
img = proj_out->forward(ctx, img);
|
||||
|
||||
auto out = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, img, 1, 0, 2, 3));
|
||||
out = ggml_reshape_4d(ctx->ggml_ctx, out, W, H, params.patch_size * params.patch_size * params.out_channels, N);
|
||||
return out;
|
||||
}
|
||||
};
|
||||
|
||||
struct LensRunner : public GGMLRunner {
|
||||
LensParams lens_params;
|
||||
LensModel lens;
|
||||
std::vector<float> pe_vec;
|
||||
|
||||
LensRunner(ggml_backend_t backend,
|
||||
ggml_backend_t params_backend,
|
||||
const String2TensorStorage& tensor_storage_map = {},
|
||||
const std::string prefix = "")
|
||||
: GGMLRunner(backend, params_backend) {
|
||||
lens_params.num_layers = 0;
|
||||
for (const auto& [name, tensor_storage] : tensor_storage_map) {
|
||||
if (!starts_with(name, prefix)) {
|
||||
continue;
|
||||
}
|
||||
if (ends_with(name, "img_in.weight") && tensor_storage.n_dims == 2) {
|
||||
lens_params.in_channels = tensor_storage.ne[0];
|
||||
int64_t inner_dim = tensor_storage.ne[1];
|
||||
lens_params.num_attention_heads = inner_dim / lens_params.attention_head_dim;
|
||||
} else if (ends_with(name, "txt_in.weight") && tensor_storage.n_dims == 2) {
|
||||
lens_params.selected_layer_count = static_cast<int>(tensor_storage.ne[0] / lens_params.joint_attention_dim);
|
||||
} else if (ends_with(name, "proj_out.weight") && tensor_storage.n_dims == 2) {
|
||||
lens_params.out_channels = tensor_storage.ne[1] / lens_params.patch_size / lens_params.patch_size;
|
||||
} else if (ends_with(name, "transformer_blocks.0.attn.norm_q.weight") && tensor_storage.n_dims == 1) {
|
||||
lens_params.attention_head_dim = tensor_storage.ne[0];
|
||||
}
|
||||
|
||||
size_t pos = name.find("transformer_blocks.");
|
||||
if (pos != std::string::npos) {
|
||||
std::string layer_name = name.substr(pos);
|
||||
auto items = split_string(layer_name, '.');
|
||||
if (items.size() > 1) {
|
||||
int block_index = atoi(items[1].c_str());
|
||||
if (block_index + 1 > lens_params.num_layers) {
|
||||
lens_params.num_layers = block_index + 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if (lens_params.num_layers == 0) {
|
||||
lens_params.num_layers = 48;
|
||||
}
|
||||
lens_params.axes_dim_sum = 0;
|
||||
for (int axis_dim : lens_params.axes_dim) {
|
||||
lens_params.axes_dim_sum += axis_dim;
|
||||
}
|
||||
|
||||
LOG_INFO("lens: layers = %d, in_channels = %" PRId64 ", out_channels = %" PRId64
|
||||
", heads = %" PRId64 ", head_dim = %" PRId64,
|
||||
lens_params.num_layers,
|
||||
lens_params.in_channels,
|
||||
lens_params.out_channels,
|
||||
lens_params.num_attention_heads,
|
||||
lens_params.attention_head_dim);
|
||||
|
||||
lens = LensModel(lens_params);
|
||||
lens.init(params_ctx, tensor_storage_map, prefix);
|
||||
}
|
||||
|
||||
std::string get_desc() override {
|
||||
return "lens";
|
||||
}
|
||||
|
||||
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
|
||||
lens.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(LENS_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);
|
||||
|
||||
pe_vec = Rope::gen_lens_pe(static_cast<int>(x->ne[1]),
|
||||
static_cast<int>(x->ne[0]),
|
||||
static_cast<int>(x->ne[3]),
|
||||
static_cast<int>(context->ne[1]),
|
||||
lens_params.theta,
|
||||
circular_y_enabled,
|
||||
circular_x_enabled,
|
||||
lens_params.axes_dim);
|
||||
int pos_len = static_cast<int>(pe_vec.size() / lens_params.axes_dim_sum / 2);
|
||||
auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, lens_params.axes_dim_sum / 2, pos_len);
|
||||
set_backend_tensor_data(pe, pe_vec.data());
|
||||
|
||||
auto runner_ctx = get_context();
|
||||
ggml_tensor* out = lens.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), x.dim());
|
||||
}
|
||||
};
|
||||
} // namespace Lens
|
||||
|
||||
#endif // __SD_LENS_HPP__
|
||||
+288
-17
@@ -23,11 +23,12 @@
|
||||
#include "rope.hpp"
|
||||
#include "tokenizers/bpe_tokenizer.h"
|
||||
#include "tokenizers/gemma_tokenizer.h"
|
||||
#include "tokenizers/gpt_oss_tokenizer.h"
|
||||
#include "tokenizers/mistral_tokenizer.h"
|
||||
#include "tokenizers/qwen2_tokenizer.h"
|
||||
|
||||
namespace LLM {
|
||||
constexpr int LLM_GRAPH_SIZE = 10240;
|
||||
constexpr int LLM_GRAPH_SIZE = 65536;
|
||||
|
||||
enum class LLMArch {
|
||||
QWEN2_5_VL,
|
||||
@@ -36,6 +37,7 @@ namespace LLM {
|
||||
MISTRAL_SMALL_3_2,
|
||||
MINISTRAL_3_3B,
|
||||
GEMMA3_12B,
|
||||
GPT_OSS_20B,
|
||||
ARCH_COUNT,
|
||||
};
|
||||
|
||||
@@ -46,6 +48,7 @@ namespace LLM {
|
||||
"mistral_small3.2",
|
||||
"ministral3.3b",
|
||||
"gemma3_12b",
|
||||
"gpt_oss_20b",
|
||||
};
|
||||
|
||||
enum class MLPActivation {
|
||||
@@ -83,6 +86,7 @@ namespace LLM {
|
||||
int num_kv_heads = 4;
|
||||
int head_dim = 128;
|
||||
bool qkv_bias = true;
|
||||
bool attention_out_bias = false;
|
||||
bool qk_norm = false;
|
||||
bool rms_norm_add = false;
|
||||
bool normalize_input = false;
|
||||
@@ -93,6 +97,8 @@ namespace LLM {
|
||||
std::vector<float> rope_thetas = {1000000.f};
|
||||
std::vector<float> rope_scales = {1.f};
|
||||
std::vector<int> sliding_attention;
|
||||
int64_t num_experts = 0;
|
||||
int64_t num_experts_per_tok = 0;
|
||||
LLMVisionParams vision;
|
||||
};
|
||||
|
||||
@@ -163,6 +169,170 @@ namespace LLM {
|
||||
}
|
||||
};
|
||||
|
||||
struct GPTOSSMLP : public GGMLBlock {
|
||||
protected:
|
||||
int64_t hidden_size;
|
||||
int64_t intermediate_size;
|
||||
int64_t num_experts;
|
||||
int64_t num_experts_per_tok;
|
||||
bool has_combined_gate_up = false;
|
||||
|
||||
void init_params(ggml_context* ctx,
|
||||
const String2TensorStorage& tensor_storage_map = {},
|
||||
std::string prefix = "") override {
|
||||
auto supported_type = [](ggml_type wtype, int64_t in_features) {
|
||||
if (in_features % ggml_blck_size(wtype) != 0) {
|
||||
return GGML_TYPE_F32;
|
||||
}
|
||||
return wtype;
|
||||
};
|
||||
|
||||
params["router.weight"] = ggml_new_tensor_2d(ctx,
|
||||
supported_type(get_type(prefix + "router.weight", tensor_storage_map, GGML_TYPE_F32), hidden_size),
|
||||
hidden_size,
|
||||
num_experts);
|
||||
params["router.bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_experts);
|
||||
|
||||
has_combined_gate_up = tensor_storage_map.find(prefix + "experts.gate_up_proj.weight") != tensor_storage_map.end();
|
||||
if (has_combined_gate_up) {
|
||||
ggml_type gate_up_type = supported_type(get_type(prefix + "experts.gate_up_proj.weight", tensor_storage_map, GGML_TYPE_F32), hidden_size);
|
||||
params["experts.gate_up_proj.weight"] = ggml_new_tensor_3d(ctx,
|
||||
gate_up_type,
|
||||
hidden_size,
|
||||
intermediate_size * 2,
|
||||
num_experts);
|
||||
params["experts.gate_up_proj.bias"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, intermediate_size * 2, num_experts);
|
||||
} else {
|
||||
ggml_type gate_type = supported_type(get_type(prefix + "experts.gate_proj.weight", tensor_storage_map, GGML_TYPE_F32), hidden_size);
|
||||
ggml_type up_type = supported_type(get_type(prefix + "experts.up_proj.weight", tensor_storage_map, GGML_TYPE_F32), hidden_size);
|
||||
params["experts.gate_proj.weight"] = ggml_new_tensor_3d(ctx, gate_type, hidden_size, intermediate_size, num_experts);
|
||||
params["experts.up_proj.weight"] = ggml_new_tensor_3d(ctx, up_type, hidden_size, intermediate_size, num_experts);
|
||||
params["experts.gate_proj.bias"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, intermediate_size, num_experts);
|
||||
params["experts.up_proj.bias"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, intermediate_size, num_experts);
|
||||
}
|
||||
|
||||
ggml_type down_type = supported_type(get_type(prefix + "experts.down_proj.weight", tensor_storage_map, GGML_TYPE_F32), intermediate_size);
|
||||
params["experts.down_proj.weight"] = ggml_new_tensor_3d(ctx, down_type, intermediate_size, hidden_size, num_experts);
|
||||
params["experts.down_proj.bias"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, num_experts);
|
||||
}
|
||||
|
||||
ggml_tensor* expert_linear(GGMLRunnerContext* ctx,
|
||||
const std::string& weight_name,
|
||||
const std::string& bias_name,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* selected_experts) {
|
||||
auto out = ggml_mul_mat_id(ctx->ggml_ctx, params[weight_name], x, selected_experts);
|
||||
auto it = params.find(bias_name);
|
||||
if (it != params.end()) {
|
||||
out = ggml_add_id(ctx->ggml_ctx, out, it->second, selected_experts);
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
public:
|
||||
GPTOSSMLP(const LLMParams& params)
|
||||
: hidden_size(params.hidden_size),
|
||||
intermediate_size(params.intermediate_size),
|
||||
num_experts(params.num_experts),
|
||||
num_experts_per_tok(params.num_experts_per_tok) {}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
// x: [N, n_token, hidden_size]
|
||||
GGML_ASSERT(num_experts > 0 && num_experts_per_tok > 0);
|
||||
|
||||
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_weight = params["router.weight"];
|
||||
ggml_tensor* router_bias = params["router.bias"];
|
||||
ggml_tensor* router_logits = ggml_mul_mat(ctx->ggml_ctx, router_weight, x);
|
||||
router_logits = ggml_add(ctx->ggml_ctx, router_logits, router_bias);
|
||||
router_logits = ggml_reshape_2d(ctx->ggml_ctx, router_logits, num_experts, n_token_total);
|
||||
|
||||
ggml_tensor* selected_experts = ggml_argsort_top_k(ctx->ggml_ctx, router_logits, (int)num_experts_per_tok); // [top_k, tokens]
|
||||
ggml_tensor* probs = ggml_reshape_3d(ctx->ggml_ctx, router_logits, 1, num_experts, n_token_total);
|
||||
ggml_tensor* weights = ggml_get_rows(ctx->ggml_ctx, probs, selected_experts); // [1, top_k, tokens]
|
||||
weights = ggml_reshape_2d(ctx->ggml_ctx, weights, num_experts_per_tok, n_token_total);
|
||||
weights = ggml_soft_max(ctx->ggml_ctx, weights);
|
||||
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);
|
||||
|
||||
ggml_tensor* gate = nullptr;
|
||||
ggml_tensor* up = nullptr;
|
||||
if (has_combined_gate_up) {
|
||||
auto gate_up = expert_linear(ctx,
|
||||
"experts.gate_up_proj.weight",
|
||||
"experts.gate_up_proj.bias",
|
||||
x,
|
||||
selected_experts); // [2 * intermediate, top_k, tokens]
|
||||
gate_up = ggml_reshape_4d(ctx->ggml_ctx,
|
||||
gate_up,
|
||||
2,
|
||||
intermediate_size,
|
||||
num_experts_per_tok,
|
||||
n_token_total);
|
||||
gate = ggml_view_4d(ctx->ggml_ctx,
|
||||
gate_up,
|
||||
1,
|
||||
intermediate_size,
|
||||
num_experts_per_tok,
|
||||
n_token_total,
|
||||
gate_up->nb[1],
|
||||
gate_up->nb[2],
|
||||
gate_up->nb[3],
|
||||
0);
|
||||
up = ggml_view_4d(ctx->ggml_ctx,
|
||||
gate_up,
|
||||
1,
|
||||
intermediate_size,
|
||||
num_experts_per_tok,
|
||||
n_token_total,
|
||||
gate_up->nb[1],
|
||||
gate_up->nb[2],
|
||||
gate_up->nb[3],
|
||||
gate_up->nb[0]);
|
||||
gate = ggml_reshape_3d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, gate), intermediate_size, num_experts_per_tok, n_token_total);
|
||||
up = ggml_reshape_3d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, up), intermediate_size, num_experts_per_tok, n_token_total);
|
||||
} else {
|
||||
gate = expert_linear(ctx,
|
||||
"experts.gate_proj.weight",
|
||||
"experts.gate_proj.bias",
|
||||
x,
|
||||
selected_experts);
|
||||
up = expert_linear(ctx,
|
||||
"experts.up_proj.weight",
|
||||
"experts.up_proj.bias",
|
||||
x,
|
||||
selected_experts);
|
||||
}
|
||||
|
||||
auto activated = ggml_swiglu_oai(ctx->ggml_ctx, gate, up, 1.702f, 7.0f);
|
||||
auto experts = expert_linear(ctx,
|
||||
"experts.down_proj.weight",
|
||||
"experts.down_proj.bias",
|
||||
activated,
|
||||
selected_experts);
|
||||
experts = ggml_mul(ctx->ggml_ctx, experts, weights);
|
||||
|
||||
ggml_tensor* out = nullptr;
|
||||
for (int64_t i = 0; i < num_experts_per_tok; ++i) {
|
||||
auto expert_out = ggml_view_2d(ctx->ggml_ctx,
|
||||
experts,
|
||||
hidden_size,
|
||||
n_token_total,
|
||||
experts->nb[2],
|
||||
i * experts->nb[1]);
|
||||
out = out == nullptr ? expert_out : ggml_add(ctx->ggml_ctx, out, expert_out);
|
||||
}
|
||||
if (num_experts_per_tok == 1) {
|
||||
out = ggml_cont(ctx->ggml_ctx, out);
|
||||
}
|
||||
|
||||
return ggml_reshape_3d(ctx->ggml_ctx, out, hidden_size, n_token, N);
|
||||
}
|
||||
};
|
||||
|
||||
static ggml_tensor* splice_image_embeds(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* x,
|
||||
const std::vector<std::pair<int, ggml_tensor*>>& image_embeds) {
|
||||
@@ -601,6 +771,15 @@ namespace LLM {
|
||||
int64_t max_position_embeddings;
|
||||
std::vector<float> rope_thetas;
|
||||
std::vector<float> rope_scales;
|
||||
bool has_attention_sinks;
|
||||
|
||||
void init_params(ggml_context* ctx,
|
||||
const String2TensorStorage& tensor_storage_map = {},
|
||||
std::string prefix = "") override {
|
||||
if (has_attention_sinks) {
|
||||
params["sinks"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_heads);
|
||||
}
|
||||
}
|
||||
|
||||
public:
|
||||
Attention(const LLMParams& params)
|
||||
@@ -611,11 +790,12 @@ namespace LLM {
|
||||
qk_norm(params.qk_norm),
|
||||
max_position_embeddings(params.max_position_embeddings),
|
||||
rope_thetas(params.rope_thetas),
|
||||
rope_scales(params.rope_scales) {
|
||||
rope_scales(params.rope_scales),
|
||||
has_attention_sinks(params.arch == LLMArch::GPT_OSS_20B) {
|
||||
blocks["q_proj"] = std::make_shared<Linear>(params.hidden_size, num_heads * head_dim, params.qkv_bias);
|
||||
blocks["k_proj"] = std::make_shared<Linear>(params.hidden_size, num_kv_heads * head_dim, params.qkv_bias);
|
||||
blocks["v_proj"] = std::make_shared<Linear>(params.hidden_size, num_kv_heads * head_dim, params.qkv_bias);
|
||||
blocks["o_proj"] = std::make_shared<Linear>(num_heads * head_dim, params.hidden_size, false);
|
||||
blocks["o_proj"] = std::make_shared<Linear>(num_heads * head_dim, params.hidden_size, params.attention_out_bias);
|
||||
if (params.qk_norm) {
|
||||
blocks["q_norm"] = std::make_shared<LLMRMSNorm>(head_dim, params.rms_norm_eps, params.rms_norm_add);
|
||||
blocks["k_norm"] = std::make_shared<LLMRMSNorm>(head_dim, params.rms_norm_eps, params.rms_norm_add);
|
||||
@@ -660,6 +840,36 @@ namespace LLM {
|
||||
} else if (arch == LLMArch::QWEN3) {
|
||||
q = ggml_rope_ext(ctx->ggml_ctx, q, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 40960, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
|
||||
k = ggml_rope_ext(ctx->ggml_ctx, k, input_pos, nullptr, 128, GGML_ROPE_TYPE_NEOX, 40960, 1000000.f, 1.f, 0.f, 1.f, 32.f, 1.f);
|
||||
} else if (arch == LLMArch::GPT_OSS_20B) {
|
||||
float rope_theta = rope_thetas.empty() ? 150000.f : rope_thetas[0];
|
||||
float rope_scale = rope_scales.empty() ? 32.f : rope_scales[0];
|
||||
float freq_scale = 1.f / rope_scale;
|
||||
q = ggml_rope_ext(ctx->ggml_ctx,
|
||||
q,
|
||||
input_pos,
|
||||
nullptr,
|
||||
head_dim,
|
||||
GGML_ROPE_TYPE_NEOX,
|
||||
4096,
|
||||
rope_theta,
|
||||
freq_scale,
|
||||
1.f,
|
||||
1.f,
|
||||
32.f,
|
||||
1.f);
|
||||
k = ggml_rope_ext(ctx->ggml_ctx,
|
||||
k,
|
||||
input_pos,
|
||||
nullptr,
|
||||
head_dim,
|
||||
GGML_ROPE_TYPE_NEOX,
|
||||
4096,
|
||||
rope_theta,
|
||||
freq_scale,
|
||||
1.f,
|
||||
1.f,
|
||||
32.f,
|
||||
1.f);
|
||||
} else if (arch == LLMArch::GEMMA3_12B) {
|
||||
float rope_theta = (rope_index == 1 ? 10000.0f : 1000000.0f);
|
||||
float rope_scale = (rope_index == 1 ? 1.f : 8.f);
|
||||
@@ -706,7 +916,28 @@ namespace LLM {
|
||||
k = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, k, 0, 2, 1, 3)); // [N, num_kv_heads, n_token, head_dim]
|
||||
k = ggml_reshape_3d(ctx->ggml_ctx, k, k->ne[0], k->ne[1], k->ne[2] * k->ne[3]); // [N*num_kv_heads, n_token, head_dim]
|
||||
|
||||
x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, attention_mask, true, false); // [N, n_token, hidden_size]
|
||||
if (arch == LLMArch::GPT_OSS_20B) {
|
||||
GGML_ASSERT(N == 1);
|
||||
auto v_attn = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, v, 1, 2, 0, 3)); // [N, kv_heads, head_dim, tokens]
|
||||
v_attn = ggml_reshape_3d(ctx->ggml_ctx, v_attn, n_token, head_dim, num_kv_heads * N);
|
||||
|
||||
auto kq = ggml_mul_mat(ctx->ggml_ctx, k, q);
|
||||
ggml_mul_mat_set_prec(kq, GGML_PREC_F32);
|
||||
kq = ggml_scale_inplace(ctx->ggml_ctx, kq, 1.0f / std::sqrt(static_cast<float>(head_dim)));
|
||||
if (attention_mask != nullptr) {
|
||||
kq = ggml_add_inplace(ctx->ggml_ctx, kq, attention_mask);
|
||||
}
|
||||
kq = ggml_soft_max_inplace(ctx->ggml_ctx, kq);
|
||||
ggml_soft_max_add_sinks(kq, params["sinks"]);
|
||||
|
||||
auto kqv = ggml_mul_mat(ctx->ggml_ctx, v_attn, kq);
|
||||
kqv = ggml_reshape_4d(ctx->ggml_ctx, kqv, head_dim, n_token, num_heads, N);
|
||||
kqv = ggml_permute(ctx->ggml_ctx, kqv, 0, 2, 1, 3);
|
||||
x = ggml_ext_cont(ctx->ggml_ctx, kqv);
|
||||
x = ggml_reshape_3d(ctx->ggml_ctx, x, head_dim * num_heads, n_token, N);
|
||||
} else {
|
||||
x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, attention_mask, true, false); // [N, n_token, hidden_size]
|
||||
}
|
||||
|
||||
x = out_proj->forward(ctx, x); // [N, n_token, hidden_size]
|
||||
return x;
|
||||
@@ -726,11 +957,15 @@ namespace LLM {
|
||||
sliding_attention(0),
|
||||
has_post_attention_norm(params.arch == LLMArch::GEMMA3_12B),
|
||||
has_post_ffw_norm(params.arch == LLMArch::GEMMA3_12B) {
|
||||
blocks["self_attn"] = std::make_shared<Attention>(params);
|
||||
blocks["mlp"] = std::make_shared<MLP>(params.hidden_size,
|
||||
params.intermediate_size,
|
||||
false,
|
||||
params.mlp_activation);
|
||||
blocks["self_attn"] = std::make_shared<Attention>(params);
|
||||
if (params.arch == LLMArch::GPT_OSS_20B) {
|
||||
blocks["mlp"] = std::make_shared<GPTOSSMLP>(params);
|
||||
} else {
|
||||
blocks["mlp"] = std::make_shared<MLP>(params.hidden_size,
|
||||
params.intermediate_size,
|
||||
false,
|
||||
params.mlp_activation);
|
||||
}
|
||||
blocks["input_layernorm"] = std::make_shared<LLMRMSNorm>(params.hidden_size, params.rms_norm_eps, params.rms_norm_add);
|
||||
blocks["post_attention_layernorm"] = std::make_shared<LLMRMSNorm>(params.hidden_size, params.rms_norm_eps, params.rms_norm_add);
|
||||
if (has_post_attention_norm) {
|
||||
@@ -751,7 +986,6 @@ namespace LLM {
|
||||
ggml_tensor* sliding_attention_mask = nullptr) {
|
||||
// x: [N, n_token, hidden_size]
|
||||
auto self_attn = std::dynamic_pointer_cast<Attention>(blocks["self_attn"]);
|
||||
auto mlp = std::dynamic_pointer_cast<MLP>(blocks["mlp"]);
|
||||
auto input_layernorm = std::dynamic_pointer_cast<LLMRMSNorm>(blocks["input_layernorm"]);
|
||||
auto post_attention_layernorm = std::dynamic_pointer_cast<LLMRMSNorm>(blocks["post_attention_layernorm"]);
|
||||
std::shared_ptr<LLMRMSNorm> post_attention_norm = nullptr;
|
||||
@@ -764,7 +998,7 @@ namespace LLM {
|
||||
}
|
||||
ggml_tensor* block_attention_mask = attention_mask;
|
||||
int rope_index = 0;
|
||||
if (arch == LLMArch::GEMMA3_12B && sliding_attention > 0) {
|
||||
if ((arch == LLMArch::GEMMA3_12B || arch == LLMArch::GPT_OSS_20B) && sliding_attention > 0) {
|
||||
block_attention_mask = sliding_attention_mask;
|
||||
rope_index = 1;
|
||||
}
|
||||
@@ -779,7 +1013,13 @@ namespace LLM {
|
||||
|
||||
residual = x;
|
||||
x = post_attention_layernorm->forward(ctx, x);
|
||||
x = mlp->forward(ctx, x);
|
||||
if (arch == LLMArch::GPT_OSS_20B) {
|
||||
auto mlp = std::dynamic_pointer_cast<GPTOSSMLP>(blocks["mlp"]);
|
||||
x = mlp->forward(ctx, x);
|
||||
} else {
|
||||
auto mlp = std::dynamic_pointer_cast<MLP>(blocks["mlp"]);
|
||||
x = mlp->forward(ctx, x);
|
||||
}
|
||||
if (post_ffw_norm != nullptr) {
|
||||
x = post_ffw_norm->forward(ctx, x);
|
||||
}
|
||||
@@ -1202,6 +1442,24 @@ namespace LLM {
|
||||
params.rope_thetas = {1000000.f, 10000.f};
|
||||
params.rope_scales = {8.f, 1.f};
|
||||
params.sliding_attention = {1024, 1024, 1024, 1024, 1024, 0};
|
||||
} else if (arch == LLMArch::GPT_OSS_20B) {
|
||||
params.head_dim = 64;
|
||||
params.num_heads = 64;
|
||||
params.num_kv_heads = 8;
|
||||
params.qkv_bias = true;
|
||||
params.attention_out_bias = true;
|
||||
params.qk_norm = false;
|
||||
params.rms_norm_eps = 1e-5f;
|
||||
params.hidden_size = 2880;
|
||||
params.intermediate_size = 2880;
|
||||
params.num_layers = 24;
|
||||
params.vocab_size = 201088;
|
||||
params.max_position_embeddings = 131072;
|
||||
params.rope_thetas = {150000.f};
|
||||
params.rope_scales = {32.f};
|
||||
params.sliding_attention = {128, 0};
|
||||
params.num_experts = 32;
|
||||
params.num_experts_per_tok = 4;
|
||||
}
|
||||
bool have_vision_weight = false;
|
||||
bool llama_cpp_style = false;
|
||||
@@ -1236,6 +1494,12 @@ namespace LLM {
|
||||
if (contains(tensor_name, "layers.0.mlp.gate_proj.weight")) {
|
||||
params.intermediate_size = pair.second.ne[1];
|
||||
}
|
||||
if (contains(tensor_name, "layers.0.mlp.experts.gate_up_proj.weight")) {
|
||||
params.intermediate_size = pair.second.ne[1] / 2;
|
||||
}
|
||||
if (contains(tensor_name, "layers.0.mlp.experts.gate_proj.weight")) {
|
||||
params.intermediate_size = pair.second.ne[1];
|
||||
}
|
||||
}
|
||||
if (arch == LLMArch::QWEN3 && params.num_layers == 28) { // Qwen3 2B
|
||||
params.num_heads = 16;
|
||||
@@ -1315,7 +1579,8 @@ namespace LLM {
|
||||
if (params.arch == LLMArch::MISTRAL_SMALL_3_2 ||
|
||||
params.arch == LLMArch::MINISTRAL_3_3B ||
|
||||
params.arch == LLMArch::QWEN3 ||
|
||||
params.arch == LLMArch::GEMMA3_12B) {
|
||||
params.arch == LLMArch::GEMMA3_12B ||
|
||||
params.arch == LLMArch::GPT_OSS_20B) {
|
||||
input_pos_vec.resize(n_tokens);
|
||||
for (int i = 0; i < n_tokens; ++i) {
|
||||
input_pos_vec[i] = i;
|
||||
@@ -1354,7 +1619,11 @@ namespace LLM {
|
||||
set_backend_tensor_data(attention_mask, attention_mask_vec.data());
|
||||
}
|
||||
|
||||
if (params.arch == LLMArch::GEMMA3_12B) {
|
||||
if (params.arch == LLMArch::GEMMA3_12B || params.arch == LLMArch::GPT_OSS_20B) {
|
||||
int sliding_window = 0;
|
||||
for (int window : params.sliding_attention) {
|
||||
sliding_window = std::max(sliding_window, window);
|
||||
}
|
||||
sliding_attention_mask_vec.resize(n_tokens * n_tokens);
|
||||
if (!attention_mask_tensor.empty()) {
|
||||
GGML_ASSERT(attention_mask_tensor.numel() == n_tokens * n_tokens);
|
||||
@@ -1364,8 +1633,7 @@ namespace LLM {
|
||||
}
|
||||
for (int i0 = 0; i0 < n_tokens; i0++) {
|
||||
for (int i1 = 0; i1 < n_tokens; i1++) {
|
||||
if (i0 + 1024 <= i1) {
|
||||
LOG_DEBUG("xxxxxxxxxxxxxx");
|
||||
if (sliding_window > 0 && i0 + sliding_window <= i1) {
|
||||
sliding_attention_mask_vec[i1 * n_tokens + i0] = -INFINITY;
|
||||
}
|
||||
}
|
||||
@@ -1403,7 +1671,8 @@ namespace LLM {
|
||||
out_layers,
|
||||
return_all_hidden_states);
|
||||
};
|
||||
return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true));
|
||||
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, true),
|
||||
input_ids.dim() + 1);
|
||||
}
|
||||
|
||||
int64_t get_num_image_tokens(int64_t t, int64_t h, int64_t w) {
|
||||
@@ -1484,6 +1753,8 @@ namespace LLM {
|
||||
: model(arch, backend, params_backend, tensor_storage_map, prefix, enable_vision) {
|
||||
if (arch == LLMArch::MISTRAL_SMALL_3_2 || arch == LLMArch::MINISTRAL_3_3B) {
|
||||
tokenizer = std::make_shared<MistralTokenizer>();
|
||||
} else if (arch == LLMArch::GPT_OSS_20B) {
|
||||
tokenizer = std::make_shared<GPTOSSTokenizer>();
|
||||
} else {
|
||||
tokenizer = std::make_shared<Qwen2Tokenizer>();
|
||||
}
|
||||
|
||||
+266
-288
@@ -2,6 +2,7 @@
|
||||
#define __SD_LTX_AUDIO_VAE_H__
|
||||
|
||||
#include <cmath>
|
||||
#include <limits>
|
||||
#include <numeric>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
@@ -171,90 +172,59 @@ namespace LTXV {
|
||||
}
|
||||
};
|
||||
|
||||
static sd::Tensor<float> squeeze_trailing_singleton_dims(sd::Tensor<float> tensor) {
|
||||
while (tensor.dim() > 0 && tensor.shape().back() == 1) {
|
||||
tensor = tensor.squeeze(static_cast<size_t>(tensor.dim() - 1));
|
||||
}
|
||||
return tensor;
|
||||
}
|
||||
static ggml_tensor* compute_log_mel_spectrogram(GGMLRunnerContext* runner_ctx,
|
||||
ggml_tensor* waveform,
|
||||
ggml_tensor* forward_basis,
|
||||
ggml_tensor* mel_basis,
|
||||
int hop_length) {
|
||||
auto ctx = runner_ctx->ggml_ctx;
|
||||
GGML_ASSERT(ctx != nullptr);
|
||||
GGML_ASSERT(waveform != nullptr);
|
||||
GGML_ASSERT(forward_basis != nullptr);
|
||||
GGML_ASSERT(mel_basis != nullptr);
|
||||
GGML_ASSERT(waveform->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(forward_basis->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(mel_basis->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(forward_basis->ne[1] == 1);
|
||||
|
||||
static sd::Tensor<float> normalize_waveform_for_host(sd::Tensor<float> waveform) {
|
||||
waveform = squeeze_trailing_singleton_dims(std::move(waveform));
|
||||
if (waveform.empty()) {
|
||||
return waveform;
|
||||
}
|
||||
if (waveform.dim() == 1) {
|
||||
return waveform.reshape({waveform.shape()[0], 1, 1});
|
||||
}
|
||||
if (waveform.dim() == 2) {
|
||||
return waveform.reshape({waveform.shape()[0], waveform.shape()[1], 1});
|
||||
}
|
||||
if (waveform.dim() == 3) {
|
||||
return waveform;
|
||||
}
|
||||
throw std::runtime_error("Unsupported waveform rank for host processing: rank=" + std::to_string(waveform.dim()));
|
||||
}
|
||||
const int64_t time = waveform->ne[0];
|
||||
const int64_t channels = waveform->ne[1];
|
||||
const int64_t batch = waveform->ne[2];
|
||||
const int64_t filter_len = forward_basis->ne[0];
|
||||
const int64_t stft_channels = forward_basis->ne[2];
|
||||
const int64_t n_freqs = stft_channels / 2;
|
||||
const int64_t n_mels = mel_basis->ne[1];
|
||||
const int64_t left_pad = std::max<int64_t>(0, filter_len - hop_length);
|
||||
const int64_t padded_time = time + left_pad;
|
||||
const int64_t frame_count = padded_time < filter_len ? 0 : 1 + (padded_time - filter_len) / hop_length;
|
||||
|
||||
static sd::Tensor<float> load_param_tensor_f32(ggml_tensor* tensor) {
|
||||
GGML_ASSERT(tensor != nullptr);
|
||||
return squeeze_trailing_singleton_dims(sd::make_sd_tensor_from_ggml<float>(tensor));
|
||||
}
|
||||
GGML_ASSERT(stft_channels % 2 == 0);
|
||||
GGML_ASSERT(mel_basis->ne[0] == n_freqs);
|
||||
GGML_ASSERT(waveform->ne[3] == 1);
|
||||
GGML_ASSERT(frame_count > 0);
|
||||
|
||||
static sd::Tensor<float> compute_log_mel_spectrogram(const sd::Tensor<float>& waveform_in,
|
||||
const sd::Tensor<float>& forward_basis,
|
||||
const sd::Tensor<float>& mel_basis,
|
||||
int hop_length) {
|
||||
auto waveform = normalize_waveform_for_host(waveform_in);
|
||||
GGML_ASSERT(forward_basis.dim() >= 3);
|
||||
GGML_ASSERT(mel_basis.dim() >= 2);
|
||||
|
||||
const int64_t time = waveform.shape()[0];
|
||||
const int64_t channels = waveform.shape()[1];
|
||||
const int64_t batch = waveform.shape()[2];
|
||||
const int64_t filter_len = forward_basis.shape()[0];
|
||||
const int64_t basis_freq2 = forward_basis.shape().back();
|
||||
const int64_t n_freqs = basis_freq2 / 2;
|
||||
const int64_t n_mels = mel_basis.shape()[1];
|
||||
const int64_t left_pad = std::max<int64_t>(0, filter_len - hop_length);
|
||||
const int64_t padded_time = time + left_pad;
|
||||
const int64_t frame_count = padded_time < filter_len ? 0 : 1 + (padded_time - filter_len) / hop_length;
|
||||
|
||||
sd::Tensor<float> log_mel({n_mels, frame_count, channels, batch});
|
||||
std::vector<float> padded(static_cast<size_t>(padded_time), 0.0f);
|
||||
std::vector<float> magnitude(static_cast<size_t>(n_freqs), 0.0f);
|
||||
|
||||
for (int64_t b = 0; b < batch; ++b) {
|
||||
for (int64_t c = 0; c < channels; ++c) {
|
||||
std::fill(padded.begin(), padded.end(), 0.0f);
|
||||
for (int64_t t = 0; t < time; ++t) {
|
||||
padded[static_cast<size_t>(t + left_pad)] = waveform.index(t, c, b);
|
||||
}
|
||||
|
||||
for (int64_t frame = 0; frame < frame_count; ++frame) {
|
||||
const int64_t frame_offset = frame * hop_length;
|
||||
for (int64_t f = 0; f < n_freqs; ++f) {
|
||||
double real = 0.0;
|
||||
double imag = 0.0;
|
||||
for (int64_t k = 0; k < filter_len; ++k) {
|
||||
const float sample = padded[static_cast<size_t>(frame_offset + k)];
|
||||
real += static_cast<double>(sample) * static_cast<double>(forward_basis.index(k, 0, f));
|
||||
imag += static_cast<double>(sample) * static_cast<double>(forward_basis.index(k, 0, f + n_freqs));
|
||||
}
|
||||
magnitude[static_cast<size_t>(f)] = static_cast<float>(std::sqrt(real * real + imag * imag));
|
||||
}
|
||||
|
||||
for (int64_t m = 0; m < n_mels; ++m) {
|
||||
double mel_value = 0.0;
|
||||
for (int64_t f = 0; f < n_freqs; ++f) {
|
||||
mel_value += static_cast<double>(mel_basis.index(f, m)) * static_cast<double>(magnitude[static_cast<size_t>(f)]);
|
||||
}
|
||||
log_mel.index(m, frame, c, b) = static_cast<float>(std::log(std::max(mel_value, 1e-5)));
|
||||
}
|
||||
}
|
||||
}
|
||||
auto x = ggml_reshape_3d(ctx, waveform, time, 1, channels * batch);
|
||||
if (left_pad > 0) {
|
||||
x = ggml_pad_ext(ctx, x, static_cast<int>(left_pad), 0, 0, 0, 0, 0, 0, 0);
|
||||
}
|
||||
|
||||
return log_mel;
|
||||
auto frames = ggml_conv_1d(ctx, forward_basis, x, hop_length, 0, 1);
|
||||
GGML_ASSERT(frames->ne[0] == frame_count);
|
||||
GGML_ASSERT(frames->ne[1] == stft_channels);
|
||||
GGML_ASSERT(frames->ne[2] == channels * batch);
|
||||
|
||||
auto real = ggml_ext_slice(ctx, frames, 1, 0, n_freqs);
|
||||
auto imag = ggml_ext_slice(ctx, frames, 1, n_freqs, stft_channels);
|
||||
auto magnitude = ggml_sqrt(ctx,
|
||||
ggml_add(ctx,
|
||||
ggml_sqr(ctx, real),
|
||||
ggml_sqr(ctx, imag)));
|
||||
|
||||
magnitude = ggml_cont(ctx, ggml_permute(ctx, magnitude, 1, 0, 2, 3));
|
||||
auto mel = ggml_mul_mat(ctx, mel_basis, magnitude);
|
||||
mel = ggml_log(ctx, ggml_clamp(ctx, mel, 1e-5f, std::numeric_limits<float>::max()));
|
||||
|
||||
return ggml_reshape_4d(ctx, mel, n_mels, frame_count, channels, batch);
|
||||
}
|
||||
|
||||
static std::vector<float> build_hann_resample_filter(int ratio) {
|
||||
@@ -276,8 +246,135 @@ namespace LTXV {
|
||||
return filter;
|
||||
}
|
||||
|
||||
static sd::Tensor<float> upsample_waveform_hann(const sd::Tensor<float>& waveform_in, int ratio) {
|
||||
auto waveform = normalize_waveform_for_host(waveform_in);
|
||||
static ggml_type audio_conv_weight_type(ggml_type type) {
|
||||
return type == GGML_TYPE_BF16 ? GGML_TYPE_F16 : type;
|
||||
}
|
||||
|
||||
static ggml_tensor* repeat_with_vulkan_f32_workaround(ggml_backend_t backend,
|
||||
ggml_context* ctx,
|
||||
ggml_tensor* x,
|
||||
int64_t ne0,
|
||||
int64_t ne1,
|
||||
int64_t ne2,
|
||||
int64_t ne3) {
|
||||
if (x->type != GGML_TYPE_F32 &&
|
||||
(x->type == GGML_TYPE_F16 || x->type == GGML_TYPE_BF16) &&
|
||||
sd_backend_is(backend, "vulkan")) {
|
||||
auto x_f32 = ggml_cast(ctx, x, GGML_TYPE_F32);
|
||||
auto repeated = ggml_repeat_4d(ctx,
|
||||
x_f32,
|
||||
ne0,
|
||||
ne1,
|
||||
ne2,
|
||||
ne3);
|
||||
return ggml_cast(ctx, repeated, x->type);
|
||||
}
|
||||
return ggml_repeat_4d(ctx, x, ne0, ne1, ne2, ne3);
|
||||
}
|
||||
|
||||
static ggml_tensor* repeat_1d_value(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int64_t count) {
|
||||
auto ctx = runner_ctx->ggml_ctx;
|
||||
GGML_ASSERT(x->ne[0] == 1);
|
||||
return repeat_with_vulkan_f32_workaround(runner_ctx->backend, ctx, x, count, x->ne[1], x->ne[2], x->ne[3]);
|
||||
}
|
||||
|
||||
static ggml_tensor* replicate_pad_1d(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int64_t left, int64_t right) {
|
||||
auto ctx = runner_ctx->ggml_ctx;
|
||||
if (left > 0) {
|
||||
auto first = ggml_ext_slice(ctx, x, 0, 0, 1);
|
||||
x = ggml_concat(ctx, repeat_1d_value(runner_ctx, first, left), x, 0);
|
||||
}
|
||||
if (right > 0) {
|
||||
auto last = ggml_ext_slice(ctx, x, 0, x->ne[0] - 1, x->ne[0]);
|
||||
x = ggml_concat(ctx, x, repeat_1d_value(runner_ctx, last, right), 0);
|
||||
}
|
||||
return x;
|
||||
}
|
||||
|
||||
static ggml_tensor* tile_depthwise_filter_1d(GGMLRunnerContext* runner_ctx, ggml_tensor* filter, int64_t channels) {
|
||||
auto ctx = runner_ctx->ggml_ctx;
|
||||
ggml_tensor* base = filter;
|
||||
if (ggml_n_dims(base) == 3) {
|
||||
base = ggml_reshape_4d(ctx, base, base->ne[0], 1, 1, 1);
|
||||
} else if (ggml_n_dims(base) == 1) {
|
||||
base = ggml_reshape_4d(ctx, base, base->ne[0], 1, 1, 1);
|
||||
}
|
||||
return repeat_with_vulkan_f32_workaround(runner_ctx->backend, ctx, base, base->ne[0], 1, channels, 1);
|
||||
}
|
||||
|
||||
static ggml_tensor* depthwise_conv1d(GGMLRunnerContext* runner_ctx,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* filter,
|
||||
int stride,
|
||||
int padding) {
|
||||
auto ctx = runner_ctx->ggml_ctx;
|
||||
GGML_ASSERT(x->ne[3] == 1);
|
||||
auto tiled = tile_depthwise_filter_1d(runner_ctx, filter, x->ne[1]);
|
||||
auto out = ggml_conv_1d_dw(ctx, tiled, x, stride, padding, 1);
|
||||
return ggml_reshape_4d(ctx, out, out->ne[0], out->ne[1], 1, 1);
|
||||
}
|
||||
|
||||
static ggml_tensor* reverse_1d_filter(ggml_context* ctx, ggml_tensor* filter) {
|
||||
GGML_ASSERT(ctx != nullptr);
|
||||
GGML_ASSERT(filter != nullptr);
|
||||
GGML_ASSERT(filter->ne[1] == 1);
|
||||
GGML_ASSERT(filter->ne[2] == 1);
|
||||
GGML_ASSERT(filter->ne[3] == 1);
|
||||
|
||||
ggml_tensor* reversed = nullptr;
|
||||
for (int64_t k = filter->ne[0] - 1; k >= 0; --k) {
|
||||
auto slice = ggml_ext_slice(ctx, filter, 0, k, k + 1);
|
||||
reversed = reversed == nullptr ? slice : ggml_concat(ctx, reversed, slice, 0);
|
||||
}
|
||||
return reversed;
|
||||
}
|
||||
|
||||
static ggml_tensor* depthwise_conv_transpose1d(ggml_context* ctx,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* filter,
|
||||
int stride) {
|
||||
GGML_ASSERT(x->ne[2] == 1 && x->ne[3] == 1);
|
||||
GGML_ASSERT(filter->ne[1] == 1);
|
||||
GGML_ASSERT(filter->ne[2] == 1 && filter->ne[3] == 1);
|
||||
|
||||
const int64_t time = x->ne[0];
|
||||
const int64_t channels = x->ne[1];
|
||||
const int64_t kernel_size = filter->ne[0];
|
||||
const int64_t out_time = (time - 1) * stride + kernel_size;
|
||||
|
||||
auto x_flat = ggml_reshape_3d(ctx, x, 1, time, channels);
|
||||
if (stride > 1) {
|
||||
auto zero_unit = ggml_ext_scale(ctx, x_flat, 0.0f);
|
||||
auto zero_tail = zero_unit;
|
||||
for (int i = 1; i < stride - 1; ++i) {
|
||||
zero_tail = ggml_concat(ctx, zero_tail, zero_unit, 0);
|
||||
}
|
||||
x_flat = ggml_concat(ctx, x_flat, zero_tail, 0);
|
||||
}
|
||||
x_flat = ggml_reshape_3d(ctx, x_flat, time * stride, 1, channels);
|
||||
|
||||
auto reversed_filter = reverse_1d_filter(ctx, filter);
|
||||
auto out = ggml_conv_1d(ctx, reversed_filter, x_flat, 1, static_cast<int>(kernel_size - 1), 1);
|
||||
if (out->ne[0] > out_time) {
|
||||
out = ggml_ext_slice(ctx, out, 0, 0, out_time);
|
||||
}
|
||||
GGML_ASSERT(out->ne[0] == out_time);
|
||||
GGML_ASSERT(out->ne[1] == 1);
|
||||
GGML_ASSERT(out->ne[2] == channels);
|
||||
|
||||
out = ggml_ext_scale(ctx, out, static_cast<float>(stride));
|
||||
return ggml_reshape_4d(ctx, out, out_time, channels, 1, 1);
|
||||
}
|
||||
|
||||
static ggml_tensor* upsample_waveform_hann(GGMLRunnerContext* runner_ctx,
|
||||
ggml_tensor* waveform,
|
||||
ggml_tensor* filter,
|
||||
int ratio) {
|
||||
auto ctx = runner_ctx->ggml_ctx;
|
||||
GGML_ASSERT(ctx != nullptr);
|
||||
GGML_ASSERT(waveform != nullptr);
|
||||
GGML_ASSERT(filter != nullptr);
|
||||
GGML_ASSERT(waveform->ne[3] == 1);
|
||||
if (ratio <= 1) {
|
||||
return waveform;
|
||||
}
|
||||
@@ -289,122 +386,29 @@ namespace LTXV {
|
||||
const int pad = width;
|
||||
const int pad_left = 2 * width * ratio;
|
||||
const int pad_right = kernel_size - ratio;
|
||||
const int64_t time = waveform.shape()[0];
|
||||
const int64_t channels = waveform.shape()[1];
|
||||
const int64_t batch = waveform.shape()[2];
|
||||
const int64_t padded_time = time + 2 * pad;
|
||||
const int64_t conv_out_time = (padded_time - 1) * ratio + kernel_size;
|
||||
const int64_t cropped_time = conv_out_time - pad_left - pad_right;
|
||||
auto filter = build_hann_resample_filter(ratio);
|
||||
const int64_t time = waveform->ne[0];
|
||||
const int64_t channels = waveform->ne[1];
|
||||
const int64_t batch = waveform->ne[2];
|
||||
|
||||
sd::Tensor<float> output({cropped_time, channels, batch});
|
||||
std::vector<float> padded(static_cast<size_t>(padded_time), 0.0f);
|
||||
std::vector<float> conv_out(static_cast<size_t>(conv_out_time), 0.0f);
|
||||
GGML_ASSERT(filter->ne[0] == kernel_size);
|
||||
|
||||
for (int64_t b = 0; b < batch; ++b) {
|
||||
for (int64_t c = 0; c < channels; ++c) {
|
||||
std::fill(padded.begin(), padded.end(), 0.0f);
|
||||
const float first = waveform.index(0, c, b);
|
||||
const float last = waveform.index(time - 1, c, b);
|
||||
for (int i = 0; i < pad; ++i) {
|
||||
padded[static_cast<size_t>(i)] = first;
|
||||
padded[static_cast<size_t>(pad + time + i)] = last;
|
||||
}
|
||||
for (int64_t t = 0; t < time; ++t) {
|
||||
padded[static_cast<size_t>(pad + t)] = waveform.index(t, c, b);
|
||||
}
|
||||
|
||||
std::fill(conv_out.begin(), conv_out.end(), 0.0f);
|
||||
for (int64_t t = 0; t < padded_time; ++t) {
|
||||
const double sample = static_cast<double>(padded[static_cast<size_t>(t)]) * ratio;
|
||||
const int64_t out_base = t * ratio;
|
||||
for (int k = 0; k < kernel_size; ++k) {
|
||||
conv_out[static_cast<size_t>(out_base + k)] += static_cast<float>(sample * filter[static_cast<size_t>(k)]);
|
||||
}
|
||||
}
|
||||
|
||||
for (int64_t t = 0; t < cropped_time; ++t) {
|
||||
output.index(t, c, b) = conv_out[static_cast<size_t>(t + pad_left)];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return output;
|
||||
auto x = ggml_reshape_3d(ctx, waveform, time, channels * batch, 1);
|
||||
x = replicate_pad_1d(runner_ctx, x, pad, pad);
|
||||
x = depthwise_conv_transpose1d(ctx, x, filter, ratio);
|
||||
x = ggml_ext_slice(ctx, x, 0, pad_left, x->ne[0] - pad_right);
|
||||
return ggml_reshape_3d(ctx, x, x->ne[0], channels, batch);
|
||||
}
|
||||
|
||||
static sd::Tensor<float> crop_waveform_samples(const sd::Tensor<float>& waveform_in, int64_t target_samples) {
|
||||
auto waveform = normalize_waveform_for_host(waveform_in);
|
||||
if (waveform.shape()[0] == target_samples) {
|
||||
static ggml_tensor* crop_waveform_samples(ggml_context* ctx,
|
||||
ggml_tensor* waveform,
|
||||
int64_t target_samples) {
|
||||
GGML_ASSERT(ctx != nullptr);
|
||||
GGML_ASSERT(waveform != nullptr);
|
||||
if (waveform->ne[0] == target_samples) {
|
||||
return waveform;
|
||||
}
|
||||
if (waveform.shape()[0] > target_samples) {
|
||||
return sd::ops::slice(waveform, 0, 0, target_samples);
|
||||
}
|
||||
sd::Tensor<float> output({target_samples, waveform.shape()[1], waveform.shape()[2]});
|
||||
sd::ops::slice_assign(&output, 0, 0, waveform.shape()[0], waveform);
|
||||
return output;
|
||||
}
|
||||
|
||||
static ggml_type audio_conv_weight_type(ggml_type type) {
|
||||
return type == GGML_TYPE_BF16 ? GGML_TYPE_F16 : type;
|
||||
}
|
||||
|
||||
static ggml_tensor* repeat_1d_value(ggml_context* ctx, ggml_tensor* x, int64_t count) {
|
||||
GGML_ASSERT(x->ne[0] == 1);
|
||||
ggml_tensor* target = ggml_new_tensor_4d(ctx, x->type, count, x->ne[1], x->ne[2], x->ne[3]);
|
||||
return ggml_repeat(ctx, x, target);
|
||||
}
|
||||
|
||||
static ggml_tensor* replicate_pad_1d(ggml_context* ctx, ggml_tensor* x, int64_t left, int64_t right) {
|
||||
if (left > 0) {
|
||||
auto first = ggml_ext_slice(ctx, x, 0, 0, 1);
|
||||
x = ggml_concat(ctx, repeat_1d_value(ctx, first, left), x, 0);
|
||||
}
|
||||
if (right > 0) {
|
||||
auto last = ggml_ext_slice(ctx, x, 0, x->ne[0] - 1, x->ne[0]);
|
||||
x = ggml_concat(ctx, x, repeat_1d_value(ctx, last, right), 0);
|
||||
}
|
||||
return x;
|
||||
}
|
||||
|
||||
static ggml_tensor* tile_depthwise_filter_1d(ggml_context* ctx, ggml_tensor* filter, int64_t channels) {
|
||||
ggml_tensor* base = filter;
|
||||
if (ggml_n_dims(base) == 3) {
|
||||
base = ggml_reshape_4d(ctx, base, base->ne[0], 1, 1, 1);
|
||||
} else if (ggml_n_dims(base) == 1) {
|
||||
base = ggml_reshape_4d(ctx, base, base->ne[0], 1, 1, 1);
|
||||
}
|
||||
ggml_tensor* target = ggml_new_tensor_4d(ctx, base->type, base->ne[0], 1, channels, 1);
|
||||
return ggml_repeat(ctx, base, target);
|
||||
}
|
||||
|
||||
static ggml_tensor* depthwise_conv1d(ggml_context* ctx,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* filter,
|
||||
int stride,
|
||||
int padding) {
|
||||
GGML_ASSERT(x->ne[3] == 1);
|
||||
auto tiled = tile_depthwise_filter_1d(ctx, filter, x->ne[1]);
|
||||
auto out = ggml_conv_1d_dw(ctx, tiled, x, stride, padding, 1);
|
||||
return ggml_reshape_4d(ctx, out, out->ne[0], out->ne[1], 1, 1);
|
||||
}
|
||||
|
||||
static ggml_tensor* depthwise_conv_transpose1d(ggml_context* ctx,
|
||||
ggml_tensor* x,
|
||||
ggml_tensor* filter,
|
||||
int stride) {
|
||||
GGML_ASSERT(x->ne[2] == 1 && x->ne[3] == 1);
|
||||
GGML_ASSERT(filter->ne[1] == 1);
|
||||
|
||||
ggml_tensor* out = nullptr;
|
||||
for (int64_t c = 0; c < x->ne[1]; ++c) {
|
||||
auto xi = ggml_ext_slice(ctx, x, 1, c, c + 1);
|
||||
auto yi = ggml_conv_transpose_1d(ctx, filter, xi, stride, 0, 1);
|
||||
yi = ggml_ext_scale(ctx, yi, static_cast<float>(stride));
|
||||
yi = ggml_reshape_4d(ctx, yi, yi->ne[0], 1, 1, 1);
|
||||
out = out == nullptr ? yi : ggml_concat(ctx, out, yi, 1);
|
||||
}
|
||||
return out;
|
||||
GGML_ASSERT(waveform->ne[0] > target_samples);
|
||||
return ggml_ext_slice(ctx, waveform, 0, 0, target_samples);
|
||||
}
|
||||
|
||||
struct PixelNorm2D : public UnaryBlock {
|
||||
@@ -654,7 +658,7 @@ namespace LTXV {
|
||||
int up_pad_left = up_pad * up_ratio + (up_kernel_size - up_ratio) / 2;
|
||||
int up_pad_right = up_pad * up_ratio + (up_kernel_size - up_ratio + 1) / 2;
|
||||
|
||||
x = replicate_pad_1d(ctx->ggml_ctx, x, up_pad, up_pad);
|
||||
x = replicate_pad_1d(ctx, x, up_pad, up_pad);
|
||||
x = depthwise_conv_transpose1d(ctx->ggml_ctx, x, up_filter, up_ratio);
|
||||
x = ggml_ext_slice(ctx->ggml_ctx, x, 0, up_pad_left, x->ne[0] - up_pad_right);
|
||||
|
||||
@@ -662,8 +666,8 @@ namespace LTXV {
|
||||
|
||||
int down_pad_left = down_kernel_size / 2 - (down_kernel_size % 2 == 0 ? 1 : 0);
|
||||
int down_pad_right = down_kernel_size / 2;
|
||||
x = replicate_pad_1d(ctx->ggml_ctx, x, down_pad_left, down_pad_right);
|
||||
x = depthwise_conv1d(ctx->ggml_ctx, x, down_filter, down_ratio, 0);
|
||||
x = replicate_pad_1d(ctx, x, down_pad_left, down_pad_right);
|
||||
x = depthwise_conv1d(ctx, x, down_filter, down_ratio, 0);
|
||||
return x;
|
||||
}
|
||||
};
|
||||
@@ -926,41 +930,66 @@ namespace LTXV {
|
||||
}
|
||||
}
|
||||
|
||||
ggml_tensor* decode_to_mel(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* latent,
|
||||
int target_time,
|
||||
int target_freq) {
|
||||
auto mean = params["audio_vae.per_channel_statistics.mean-of-means"];
|
||||
auto stddev = params["audio_vae.per_channel_statistics.std-of-means"];
|
||||
auto decoder = std::dynamic_pointer_cast<AudioDecoder>(blocks["audio_vae.decoder"]);
|
||||
return decoder->forward(ctx, latent, mean, stddev, target_time, target_freq);
|
||||
}
|
||||
ggml_tensor* decode(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* latent,
|
||||
ggml_tensor* bwe_skip_filter) {
|
||||
int target_time = static_cast<int>(latent->ne[1]) * config.latent_downsample_factor() -
|
||||
(config.latent_downsample_factor() - 1);
|
||||
int target_freq = config.mel_bins;
|
||||
|
||||
ggml_tensor* run_vocoder(GGMLRunnerContext* ctx, ggml_tensor* mel) {
|
||||
auto vocoder = std::dynamic_pointer_cast<Vocoder>(blocks["vocoder.vocoder"]);
|
||||
return vocoder->forward(ctx, mel);
|
||||
}
|
||||
auto decoder = std::dynamic_pointer_cast<AudioDecoder>(blocks["audio_vae.decoder"]);
|
||||
auto mean = params["audio_vae.per_channel_statistics.mean-of-means"];
|
||||
auto stddev = params["audio_vae.per_channel_statistics.std-of-means"];
|
||||
auto mel = decoder->forward(ctx, latent, mean, stddev, target_time, target_freq);
|
||||
auto vocoder = std::dynamic_pointer_cast<Vocoder>(blocks["vocoder.vocoder"]);
|
||||
auto waveform = vocoder->forward(ctx, mel);
|
||||
|
||||
ggml_tensor* run_bwe_generator(GGMLRunnerContext* ctx, ggml_tensor* mel) {
|
||||
GGML_ASSERT(config.has_bwe);
|
||||
auto bwe_generator = std::dynamic_pointer_cast<Vocoder>(blocks["vocoder.bwe_generator"]);
|
||||
return bwe_generator->forward(ctx, mel);
|
||||
}
|
||||
if (config.has_bwe) {
|
||||
GGML_ASSERT(bwe_skip_filter != nullptr);
|
||||
const int bwe_ratio = config.bwe_output_sample_rate / config.bwe_input_sample_rate;
|
||||
const int64_t low_time = waveform->ne[0];
|
||||
const int64_t out_time = low_time * bwe_ratio;
|
||||
int64_t remainder = low_time % config.bwe_hop_length;
|
||||
auto bwe_waveform = waveform;
|
||||
if (remainder != 0) {
|
||||
bwe_waveform = ggml_pad_ext(ctx->ggml_ctx,
|
||||
bwe_waveform,
|
||||
0,
|
||||
static_cast<int>(config.bwe_hop_length - remainder),
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
0);
|
||||
}
|
||||
|
||||
ggml_tensor* mel_basis_tensor() const {
|
||||
auto iter = params.find("vocoder.mel_stft.mel_basis");
|
||||
return iter == params.end() ? nullptr : iter->second;
|
||||
}
|
||||
auto mel_basis = params["vocoder.mel_stft.mel_basis"];
|
||||
auto stft_basis = params["vocoder.mel_stft.stft_fn.forward_basis"];
|
||||
GGML_ASSERT(mel_basis != nullptr && stft_basis != nullptr);
|
||||
auto bwe_mel = compute_log_mel_spectrogram(ctx, bwe_waveform, stft_basis, mel_basis, config.bwe_hop_length);
|
||||
auto bwe_generator = std::dynamic_pointer_cast<Vocoder>(blocks["vocoder.bwe_generator"]);
|
||||
auto residual = bwe_generator->forward(ctx, bwe_mel);
|
||||
|
||||
ggml_tensor* stft_forward_basis_tensor() const {
|
||||
auto iter = params.find("vocoder.mel_stft.stft_fn.forward_basis");
|
||||
return iter == params.end() ? nullptr : iter->second;
|
||||
auto skip = upsample_waveform_hann(ctx,
|
||||
bwe_waveform,
|
||||
bwe_skip_filter,
|
||||
bwe_ratio);
|
||||
waveform = ggml_clamp(ctx->ggml_ctx,
|
||||
ggml_add(ctx->ggml_ctx, residual, skip),
|
||||
-1.0f,
|
||||
1.0f);
|
||||
waveform = crop_waveform_samples(ctx->ggml_ctx, waveform, out_time);
|
||||
}
|
||||
|
||||
return waveform;
|
||||
}
|
||||
};
|
||||
|
||||
struct LTXAudioVAERunner : public GGMLRunner {
|
||||
LTXAudioVAEConfig config;
|
||||
LTXAudioVAE model;
|
||||
sd::Tensor<float> bwe_skip_filter_tensor;
|
||||
|
||||
LTXAudioVAERunner(ggml_backend_t backend,
|
||||
ggml_backend_t params_backend,
|
||||
@@ -970,6 +999,10 @@ namespace LTXV {
|
||||
config(LTXAudioVAEConfig::detect_from_weights(tensor_storage_map)),
|
||||
model(config) {
|
||||
model.init(params_ctx, tensor_storage_map, prefix);
|
||||
if (config.has_bwe) {
|
||||
const int bwe_ratio = config.bwe_output_sample_rate / config.bwe_input_sample_rate;
|
||||
bwe_skip_filter_tensor = sd::Tensor<float>::from_vector(build_hann_resample_filter(bwe_ratio));
|
||||
}
|
||||
}
|
||||
|
||||
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
|
||||
@@ -984,77 +1017,22 @@ namespace LTXV {
|
||||
return "ltx_audio_vae";
|
||||
}
|
||||
|
||||
ggml_cgraph* build_base_graph(const sd::Tensor<float>& latent_tensor) {
|
||||
auto latent = make_input(latent_tensor);
|
||||
int target_time = static_cast<int>(latent_tensor.shape()[1]) * config.latent_downsample_factor() -
|
||||
(config.latent_downsample_factor() - 1);
|
||||
int target_freq = config.mel_bins;
|
||||
|
||||
ggml_cgraph* gf = new_graph_custom(655360);
|
||||
auto runner_ctx = GGMLRunner::get_context();
|
||||
auto mel = model.decode_to_mel(&runner_ctx, latent, target_time, target_freq);
|
||||
auto waveform = model.run_vocoder(&runner_ctx, mel);
|
||||
ggml_build_forward_expand(gf, waveform);
|
||||
return gf;
|
||||
}
|
||||
|
||||
ggml_cgraph* build_bwe_graph(const sd::Tensor<float>& mel_tensor) {
|
||||
auto mel = make_input(mel_tensor);
|
||||
ggml_cgraph* gf = new_graph_custom(655360);
|
||||
auto runner_ctx = GGMLRunner::get_context();
|
||||
auto residual = model.run_bwe_generator(&runner_ctx, mel);
|
||||
ggml_build_forward_expand(gf, residual);
|
||||
return gf;
|
||||
}
|
||||
|
||||
sd::Tensor<float> compute_base_waveform(int n_threads,
|
||||
const sd::Tensor<float>& latent_tensor) {
|
||||
auto get_graph = [&]() -> ggml_cgraph* {
|
||||
return build_base_graph(latent_tensor);
|
||||
};
|
||||
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), 4);
|
||||
}
|
||||
|
||||
sd::Tensor<float> compute_bwe_residual(int n_threads,
|
||||
const sd::Tensor<float>& mel_tensor) {
|
||||
auto get_graph = [&]() -> ggml_cgraph* {
|
||||
return build_bwe_graph(mel_tensor);
|
||||
};
|
||||
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), 4);
|
||||
}
|
||||
|
||||
sd::Tensor<float> decode(int n_threads,
|
||||
const sd::Tensor<float>& latent_tensor) {
|
||||
auto waveform = compute_base_waveform(n_threads, latent_tensor);
|
||||
if (!config.has_bwe || waveform.empty()) {
|
||||
return waveform;
|
||||
}
|
||||
|
||||
auto waveform_host = normalize_waveform_for_host(waveform);
|
||||
const int64_t low_time = waveform_host.shape()[0];
|
||||
const int64_t out_time = low_time * config.bwe_output_sample_rate / config.bwe_input_sample_rate;
|
||||
int64_t remainder = low_time % config.bwe_hop_length;
|
||||
if (remainder != 0) {
|
||||
sd::Tensor<float> padded({low_time + (config.bwe_hop_length - remainder), waveform_host.shape()[1], waveform_host.shape()[2]});
|
||||
sd::ops::slice_assign(&padded, 0, 0, low_time, waveform_host);
|
||||
waveform_host = std::move(padded);
|
||||
}
|
||||
|
||||
auto mel_basis_tensor = model.mel_basis_tensor();
|
||||
auto stft_basis_tensor = model.stft_forward_basis_tensor();
|
||||
GGML_ASSERT(mel_basis_tensor != nullptr && stft_basis_tensor != nullptr);
|
||||
auto mel_basis = load_param_tensor_f32(mel_basis_tensor);
|
||||
auto forward_basis = load_param_tensor_f32(stft_basis_tensor);
|
||||
auto bwe_mel = compute_log_mel_spectrogram(waveform_host, forward_basis, mel_basis, config.bwe_hop_length);
|
||||
auto residual_raw = compute_bwe_residual(n_threads, bwe_mel);
|
||||
if (residual_raw.empty()) {
|
||||
return waveform;
|
||||
}
|
||||
auto residual = normalize_waveform_for_host(residual_raw);
|
||||
auto skip = upsample_waveform_hann(waveform_host, config.bwe_output_sample_rate / config.bwe_input_sample_rate);
|
||||
auto combined = sd::ops::clamp(residual + skip, -1.0f, 1.0f);
|
||||
auto cropped = crop_waveform_samples(combined, out_time);
|
||||
return restore_trailing_singleton_dims(cropped, 4);
|
||||
int64_t t0 = ggml_time_ms();
|
||||
auto get_graph = [&]() -> ggml_cgraph* {
|
||||
auto latent = make_input(latent_tensor);
|
||||
ggml_tensor* bwe_skip_filter = config.has_bwe ? make_input(bwe_skip_filter_tensor) : nullptr;
|
||||
ggml_cgraph* gf = new_graph_custom(655360);
|
||||
auto runner_ctx = GGMLRunner::get_context();
|
||||
auto waveform = model.decode(&runner_ctx, latent, bwe_skip_filter);
|
||||
ggml_build_forward_expand(gf, waveform);
|
||||
return gf;
|
||||
};
|
||||
auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), 4);
|
||||
int64_t t1 = ggml_time_ms();
|
||||
LOG_INFO("ltx audio vae decode completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
|
||||
return result;
|
||||
}
|
||||
|
||||
void test(const std::string& input_path) {
|
||||
|
||||
@@ -0,0 +1,551 @@
|
||||
#ifndef __SD_LTX_LATENT_UPSCALER_HPP__
|
||||
#define __SD_LTX_LATENT_UPSCALER_HPP__
|
||||
|
||||
#include <cinttypes>
|
||||
#include <cmath>
|
||||
#include <cstdlib>
|
||||
#include <map>
|
||||
#include <memory>
|
||||
#include <set>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "common_dit.hpp"
|
||||
#include "ggml_extend.hpp"
|
||||
#include "ggml_graph_cut.h"
|
||||
#include "model.h"
|
||||
#include "util.h"
|
||||
|
||||
namespace LTXVUpsampler {
|
||||
constexpr int LTX_UPSAMPLER_GRAPH_SIZE = 10240;
|
||||
|
||||
struct LatentUpsamplerConfig {
|
||||
int64_t in_channels = 128;
|
||||
int64_t mid_channels = 1024;
|
||||
int num_blocks_per_stage = 4;
|
||||
int dims = 3;
|
||||
bool spatial_upsample = true;
|
||||
bool temporal_upsample = false;
|
||||
bool rational_resampler = false;
|
||||
float spatial_scale = 2.f;
|
||||
int spatial_up_num = 2;
|
||||
int spatial_down_den = 1;
|
||||
int temporal_up_factor = 1;
|
||||
};
|
||||
|
||||
static inline bool has_tensor(const String2TensorStorage& tensor_storage_map,
|
||||
const std::string& name) {
|
||||
return tensor_storage_map.find(name) != tensor_storage_map.end();
|
||||
}
|
||||
|
||||
static inline int64_t get_tensor_ne(const String2TensorStorage& tensor_storage_map,
|
||||
const std::string& name,
|
||||
int axis,
|
||||
int64_t fallback) {
|
||||
auto it = tensor_storage_map.find(name);
|
||||
if (it == tensor_storage_map.end() || axis < 0 || axis >= GGML_MAX_DIMS) {
|
||||
return fallback;
|
||||
}
|
||||
return it->second.ne[axis];
|
||||
}
|
||||
|
||||
static inline int64_t get_tensor_ne0(const String2TensorStorage& tensor_storage_map,
|
||||
const std::string& name,
|
||||
int64_t fallback) {
|
||||
return get_tensor_ne(tensor_storage_map, name, 0, fallback);
|
||||
}
|
||||
|
||||
static inline int count_module_blocks(const String2TensorStorage& tensor_storage_map,
|
||||
const std::string& module_name) {
|
||||
int max_block = -1;
|
||||
const std::string prefix = module_name + ".";
|
||||
for (const auto& pair : tensor_storage_map) {
|
||||
const std::string& name = pair.first;
|
||||
if (name.find(prefix) != 0) {
|
||||
continue;
|
||||
}
|
||||
size_t begin = prefix.size();
|
||||
size_t end = name.find('.', begin);
|
||||
if (end == std::string::npos) {
|
||||
continue;
|
||||
}
|
||||
int index = atoi(name.substr(begin, end - begin).c_str());
|
||||
max_block = std::max(max_block, index);
|
||||
}
|
||||
return max_block + 1;
|
||||
}
|
||||
|
||||
static inline LatentUpsamplerConfig detect_config_from_weights(const String2TensorStorage& tensor_storage_map) {
|
||||
LatentUpsamplerConfig config;
|
||||
config.mid_channels = get_tensor_ne0(tensor_storage_map, "initial_norm.weight", config.mid_channels);
|
||||
config.in_channels = get_tensor_ne0(tensor_storage_map, "final_conv.bias", config.in_channels);
|
||||
int detected_blocks = count_module_blocks(tensor_storage_map, "res_blocks");
|
||||
if (detected_blocks > 0) {
|
||||
config.num_blocks_per_stage = detected_blocks;
|
||||
}
|
||||
config.rational_resampler = has_tensor(tensor_storage_map, "upsampler.conv.weight");
|
||||
int64_t upsampler_out_channels = get_tensor_ne0(tensor_storage_map, "upsampler.0.bias", 0);
|
||||
config.spatial_upsample = config.rational_resampler || upsampler_out_channels == 4 * config.mid_channels;
|
||||
config.temporal_upsample = upsampler_out_channels == 2 * config.mid_channels;
|
||||
if (config.temporal_upsample) {
|
||||
config.temporal_up_factor = 2;
|
||||
}
|
||||
if (config.rational_resampler) {
|
||||
int64_t out_channels = get_tensor_ne(tensor_storage_map,
|
||||
"upsampler.conv.weight",
|
||||
3,
|
||||
config.mid_channels * 9);
|
||||
if (config.mid_channels > 0 && out_channels % config.mid_channels == 0) {
|
||||
int64_t ratio = out_channels / config.mid_channels;
|
||||
int num = static_cast<int>(std::round(std::sqrt(static_cast<double>(ratio))));
|
||||
if (num > 0 && static_cast<int64_t>(num) * num == ratio) {
|
||||
config.spatial_up_num = num;
|
||||
}
|
||||
}
|
||||
if (config.spatial_up_num == 3) {
|
||||
config.spatial_down_den = 2;
|
||||
config.spatial_scale = 1.5f;
|
||||
} else if (config.spatial_up_num == 4) {
|
||||
config.spatial_down_den = 1;
|
||||
config.spatial_scale = 4.f;
|
||||
} else {
|
||||
config.spatial_down_den = 1;
|
||||
config.spatial_scale = static_cast<float>(config.spatial_up_num);
|
||||
}
|
||||
}
|
||||
return config;
|
||||
}
|
||||
|
||||
class VideoGroupNorm : public GGMLBlock {
|
||||
protected:
|
||||
int num_groups;
|
||||
int64_t num_channels;
|
||||
float eps;
|
||||
std::string prefix;
|
||||
|
||||
void init_params(ggml_context* ctx,
|
||||
const String2TensorStorage& tensor_storage_map = {},
|
||||
const std::string prefix = "") override {
|
||||
SD_UNUSED(tensor_storage_map);
|
||||
this->prefix = prefix;
|
||||
params["weight"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_channels);
|
||||
params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, num_channels);
|
||||
}
|
||||
|
||||
public:
|
||||
VideoGroupNorm(int num_groups, int64_t num_channels, float eps = 1e-05f)
|
||||
: num_groups(num_groups),
|
||||
num_channels(num_channels),
|
||||
eps(eps) {}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
// LTX video latent layout is [W, H, T, C]. ggml_group_norm treats ne[2]
|
||||
// as channels, so fold only H/T internally and restore the same layout.
|
||||
GGML_ASSERT(x->ne[3] == num_channels);
|
||||
const int64_t W = x->ne[0];
|
||||
const int64_t H = x->ne[1];
|
||||
const int64_t T = x->ne[2];
|
||||
x = ggml_ext_cont(ctx->ggml_ctx, x);
|
||||
x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H * T, num_channels, 1);
|
||||
x = ggml_group_norm(ctx->ggml_ctx, x, num_groups, eps);
|
||||
|
||||
ggml_tensor* weight = params["weight"];
|
||||
ggml_tensor* bias = params["bias"];
|
||||
if (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");
|
||||
}
|
||||
weight = ggml_reshape_4d(ctx->ggml_ctx, weight, 1, 1, num_channels, 1);
|
||||
bias = ggml_reshape_4d(ctx->ggml_ctx, bias, 1, 1, num_channels, 1);
|
||||
x = ggml_mul_inplace(ctx->ggml_ctx, x, weight);
|
||||
x = ggml_add_inplace(ctx->ggml_ctx, x, bias);
|
||||
return ggml_reshape_4d(ctx->ggml_ctx, x, W, H, T, num_channels);
|
||||
}
|
||||
};
|
||||
|
||||
class ResBlock : public GGMLBlock {
|
||||
public:
|
||||
ResBlock(int64_t channels, int dims = 3) {
|
||||
GGML_ASSERT(dims == 3);
|
||||
blocks["conv1"] = std::shared_ptr<GGMLBlock>(new Conv3d(channels, channels, {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
|
||||
blocks["norm1"] = std::shared_ptr<GGMLBlock>(new VideoGroupNorm(32, channels));
|
||||
blocks["conv2"] = std::shared_ptr<GGMLBlock>(new Conv3d(channels, channels, {3, 3, 3}, {1, 1, 1}, {1, 1, 1}));
|
||||
blocks["norm2"] = std::shared_ptr<GGMLBlock>(new VideoGroupNorm(32, channels));
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
auto conv1 = std::dynamic_pointer_cast<Conv3d>(blocks["conv1"]);
|
||||
auto norm1 = std::dynamic_pointer_cast<VideoGroupNorm>(blocks["norm1"]);
|
||||
auto conv2 = std::dynamic_pointer_cast<Conv3d>(blocks["conv2"]);
|
||||
auto norm2 = std::dynamic_pointer_cast<VideoGroupNorm>(blocks["norm2"]);
|
||||
|
||||
ggml_tensor* residual = x;
|
||||
|
||||
x = conv1->forward(ctx, x);
|
||||
x = norm1->forward(ctx, x);
|
||||
x = ggml_silu_inplace(ctx->ggml_ctx, x);
|
||||
x = conv2->forward(ctx, x);
|
||||
x = norm2->forward(ctx, x);
|
||||
x = ggml_add(ctx->ggml_ctx, x, residual);
|
||||
return ggml_silu(ctx->ggml_ctx, x);
|
||||
}
|
||||
};
|
||||
|
||||
class PixelShuffleND : public UnaryBlock {
|
||||
protected:
|
||||
int upscale_factor;
|
||||
|
||||
public:
|
||||
explicit PixelShuffleND(int upscale_factor)
|
||||
: upscale_factor(upscale_factor) {}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||
GGML_ASSERT(upscale_factor > 0);
|
||||
int64_t h = x->ne[1];
|
||||
int64_t w = x->ne[0];
|
||||
GGML_ASSERT(x->ne[2] % (upscale_factor * upscale_factor) == 0);
|
||||
// x: [b*f, c*p1*p2, h, w] -> [b*f, c, h*p1, w*p2]
|
||||
x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3)); // [b*f, h, w, c*p1*p2]
|
||||
x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0], x->ne[1] * x->ne[2], x->ne[3]); // [b*f, h*w, c*p1*p2]
|
||||
return DiT::unpatchify(ctx->ggml_ctx, x, h, w, upscale_factor, upscale_factor, true);
|
||||
}
|
||||
};
|
||||
|
||||
class TemporalPixelShuffleND : public UnaryBlock {
|
||||
protected:
|
||||
int upscale_factor;
|
||||
|
||||
public:
|
||||
explicit TemporalPixelShuffleND(int upscale_factor)
|
||||
: upscale_factor(upscale_factor) {}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
|
||||
GGML_ASSERT(upscale_factor > 0);
|
||||
GGML_ASSERT(x->ne[3] % upscale_factor == 0);
|
||||
const int64_t W = x->ne[0];
|
||||
const int64_t H = x->ne[1];
|
||||
const int64_t F = x->ne[2];
|
||||
const int64_t C = x->ne[3] / upscale_factor;
|
||||
|
||||
// x: [b, c*p, f, h, w] -> [b, c, f*p, h, w]
|
||||
x = ggml_ext_cont(ctx->ggml_ctx, x);
|
||||
x = ggml_reshape_4d(ctx->ggml_ctx, x, W * H, F, upscale_factor, C);
|
||||
x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));
|
||||
return ggml_reshape_4d(ctx->ggml_ctx, x, W, H, F * upscale_factor, C);
|
||||
}
|
||||
};
|
||||
|
||||
class BlurDownsample : public GGMLBlock {
|
||||
protected:
|
||||
int64_t channels;
|
||||
int stride;
|
||||
ggml_tensor* kernel = nullptr;
|
||||
std::vector<float> kernel_data;
|
||||
|
||||
void init_params(ggml_context* ctx,
|
||||
const String2TensorStorage& tensor_storage_map = {},
|
||||
const std::string prefix = "") override {
|
||||
SD_UNUSED(tensor_storage_map);
|
||||
if (stride == 1) {
|
||||
return;
|
||||
}
|
||||
kernel = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 5, 5, 1, channels);
|
||||
std::string name = prefix + "kernel";
|
||||
ggml_set_name(kernel, name.c_str());
|
||||
|
||||
static const float binomial[5] = {1.f, 4.f, 6.f, 4.f, 1.f};
|
||||
kernel_data.resize(static_cast<size_t>(5 * 5 * channels));
|
||||
for (int64_t c = 0; c < channels; ++c) {
|
||||
for (int y = 0; y < 5; ++y) {
|
||||
for (int x = 0; x < 5; ++x) {
|
||||
kernel_data[static_cast<size_t>(x + 5 * (y + 5 * c))] =
|
||||
binomial[y] * binomial[x] / 256.f;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
public:
|
||||
BlurDownsample(int64_t channels, int stride)
|
||||
: channels(channels),
|
||||
stride(stride) {
|
||||
GGML_ASSERT(stride >= 1);
|
||||
}
|
||||
|
||||
void load_fixed_tensors() {
|
||||
if (kernel == nullptr || kernel_data.empty()) {
|
||||
return;
|
||||
}
|
||||
ggml_backend_tensor_set(kernel, kernel_data.data(), 0, kernel_data.size() * sizeof(float));
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
if (stride == 1) {
|
||||
return x;
|
||||
}
|
||||
GGML_ASSERT(kernel != nullptr);
|
||||
GGML_ASSERT(x->ne[2] == channels);
|
||||
if (ctx->conv2d_direct_enabled) {
|
||||
return ggml_conv_2d_dw_direct(ctx->ggml_ctx, kernel, x, stride, stride, 2, 2, 1, 1);
|
||||
}
|
||||
return ggml_conv_2d_dw(ctx->ggml_ctx, kernel, x, stride, stride, 2, 2, 1, 1);
|
||||
}
|
||||
};
|
||||
|
||||
class SpatialRationalResampler : public GGMLBlock {
|
||||
protected:
|
||||
int64_t mid_channels;
|
||||
int num;
|
||||
int den;
|
||||
|
||||
public:
|
||||
SpatialRationalResampler(int64_t mid_channels, int num, int den)
|
||||
: mid_channels(mid_channels),
|
||||
num(num),
|
||||
den(den) {
|
||||
GGML_ASSERT(num >= 1);
|
||||
GGML_ASSERT(den >= 1);
|
||||
blocks["conv"] = std::shared_ptr<GGMLBlock>(new Conv2d(mid_channels, num * num * mid_channels, {3, 3}, {1, 1}, {1, 1}));
|
||||
blocks["pixel_shuffle"] = std::shared_ptr<GGMLBlock>(new PixelShuffleND(num));
|
||||
blocks["blur_down"] = std::shared_ptr<GGMLBlock>(new BlurDownsample(mid_channels, den));
|
||||
}
|
||||
|
||||
void load_fixed_tensors() {
|
||||
auto blur_down = std::dynamic_pointer_cast<BlurDownsample>(blocks["blur_down"]);
|
||||
blur_down->load_fixed_tensors();
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
auto conv = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
|
||||
auto pixel_shuffle = std::dynamic_pointer_cast<PixelShuffleND>(blocks["pixel_shuffle"]);
|
||||
auto blur_down = std::dynamic_pointer_cast<BlurDownsample>(blocks["blur_down"]);
|
||||
|
||||
// rearrange(x, "b c f h w -> (b f) c h w")
|
||||
x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
|
||||
x = conv->forward(ctx, x);
|
||||
x = pixel_shuffle->forward(ctx, x);
|
||||
x = blur_down->forward(ctx, x);
|
||||
return ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
|
||||
}
|
||||
};
|
||||
|
||||
class LatentUpsampler : public GGMLBlock {
|
||||
public:
|
||||
LatentUpsamplerConfig config;
|
||||
|
||||
explicit LatentUpsampler(LatentUpsamplerConfig config)
|
||||
: config(std::move(config)) {
|
||||
GGML_ASSERT(this->config.dims == 3);
|
||||
GGML_ASSERT(this->config.spatial_upsample || this->config.temporal_upsample);
|
||||
|
||||
blocks["initial_conv"] = std::shared_ptr<GGMLBlock>(new Conv3d(this->config.in_channels,
|
||||
this->config.mid_channels,
|
||||
{3, 3, 3},
|
||||
{1, 1, 1},
|
||||
{1, 1, 1}));
|
||||
blocks["initial_norm"] = std::shared_ptr<GGMLBlock>(new VideoGroupNorm(32, this->config.mid_channels));
|
||||
for (int i = 0; i < this->config.num_blocks_per_stage; ++i) {
|
||||
blocks["res_blocks." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new ResBlock(this->config.mid_channels, this->config.dims));
|
||||
}
|
||||
if (this->config.rational_resampler) {
|
||||
blocks["upsampler"] = std::shared_ptr<GGMLBlock>(new SpatialRationalResampler(this->config.mid_channels,
|
||||
this->config.spatial_up_num,
|
||||
this->config.spatial_down_den));
|
||||
} else if (this->config.temporal_upsample) {
|
||||
blocks["upsampler.0"] = std::shared_ptr<GGMLBlock>(new Conv3d(this->config.mid_channels,
|
||||
this->config.temporal_up_factor * this->config.mid_channels,
|
||||
{3, 3, 3},
|
||||
{1, 1, 1},
|
||||
{1, 1, 1}));
|
||||
blocks["upsampler.1"] = std::shared_ptr<GGMLBlock>(new TemporalPixelShuffleND(this->config.temporal_up_factor));
|
||||
} else {
|
||||
blocks["upsampler.0"] = std::shared_ptr<GGMLBlock>(new Conv2d(this->config.mid_channels,
|
||||
4 * this->config.mid_channels,
|
||||
{3, 3},
|
||||
{1, 1},
|
||||
{1, 1}));
|
||||
blocks["upsampler.1"] = std::shared_ptr<GGMLBlock>(new PixelShuffleND(2));
|
||||
}
|
||||
for (int i = 0; i < this->config.num_blocks_per_stage; ++i) {
|
||||
blocks["post_upsample_res_blocks." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new ResBlock(this->config.mid_channels, this->config.dims));
|
||||
}
|
||||
blocks["final_conv"] = std::shared_ptr<GGMLBlock>(new Conv3d(this->config.mid_channels,
|
||||
this->config.in_channels,
|
||||
{3, 3, 3},
|
||||
{1, 1, 1},
|
||||
{1, 1, 1}));
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
|
||||
// x: [b, c, f, h, w]
|
||||
// return: [b, c, scaled_f, scaled_h, scaled_w]
|
||||
auto initial_conv = std::dynamic_pointer_cast<Conv3d>(blocks["initial_conv"]);
|
||||
auto initial_norm = std::dynamic_pointer_cast<VideoGroupNorm>(blocks["initial_norm"]);
|
||||
auto final_conv = std::dynamic_pointer_cast<Conv3d>(blocks["final_conv"]);
|
||||
|
||||
x = initial_conv->forward(ctx, x);
|
||||
x = initial_norm->forward(ctx, x);
|
||||
x = ggml_silu(ctx->ggml_ctx, x);
|
||||
sd::ggml_graph_cut::mark_graph_cut(x, "ltx_latent_upsampler.initial", "x");
|
||||
|
||||
for (int i = 0; i < config.num_blocks_per_stage; ++i) {
|
||||
auto block = std::dynamic_pointer_cast<ResBlock>(blocks["res_blocks." + std::to_string(i)]);
|
||||
x = block->forward(ctx, x);
|
||||
sd::ggml_graph_cut::mark_graph_cut(x, "ltx_latent_upsampler.res_blocks." + std::to_string(i), "x");
|
||||
}
|
||||
|
||||
if (config.rational_resampler) {
|
||||
auto upsampler = std::dynamic_pointer_cast<SpatialRationalResampler>(blocks["upsampler"]);
|
||||
x = upsampler->forward(ctx, x);
|
||||
} else if (config.temporal_upsample) {
|
||||
auto upsample_conv = std::dynamic_pointer_cast<Conv3d>(blocks["upsampler.0"]);
|
||||
auto pixel_shuffle = std::dynamic_pointer_cast<TemporalPixelShuffleND>(blocks["upsampler.1"]);
|
||||
x = upsample_conv->forward(ctx, x); // [b, c*2, f, h, w]
|
||||
x = pixel_shuffle->forward(ctx, x); // [b, c, f*2, h, w]
|
||||
x = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1, x->ne[2]); // x[:, :, 1:, :, :]
|
||||
} else {
|
||||
auto upsample_conv = std::dynamic_pointer_cast<Conv2d>(blocks["upsampler.0"]);
|
||||
auto pixel_shuffle = std::dynamic_pointer_cast<PixelShuffleND>(blocks["upsampler.1"]);
|
||||
|
||||
// rearrange(x, "b c f h w -> (b f) c h w"),
|
||||
x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2)); // [b*f, c, h, w]
|
||||
x = upsample_conv->forward(ctx, x); // [b*f, c*4, h, w]
|
||||
x = pixel_shuffle->forward(ctx, x); // [b*f, c, h*2, w*2]
|
||||
x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2)); // [b*c, f, h, w]
|
||||
}
|
||||
sd::ggml_graph_cut::mark_graph_cut(x, "ltx_latent_upsampler.spatial_up", "x");
|
||||
|
||||
for (int i = 0; i < config.num_blocks_per_stage; ++i) {
|
||||
auto block = std::dynamic_pointer_cast<ResBlock>(blocks["post_upsample_res_blocks." + std::to_string(i)]);
|
||||
x = block->forward(ctx, x);
|
||||
sd::ggml_graph_cut::mark_graph_cut(x, "ltx_latent_upsampler.post_blocks." + std::to_string(i), "x");
|
||||
}
|
||||
|
||||
x = final_conv->forward(ctx, x);
|
||||
sd::ggml_graph_cut::mark_graph_cut(x, "ltx_latent_upsampler.final", "x");
|
||||
return x;
|
||||
}
|
||||
|
||||
void load_fixed_tensors() {
|
||||
if (!config.rational_resampler) {
|
||||
return;
|
||||
}
|
||||
auto upsampler = std::dynamic_pointer_cast<SpatialRationalResampler>(blocks["upsampler"]);
|
||||
upsampler->load_fixed_tensors();
|
||||
}
|
||||
};
|
||||
|
||||
struct LatentUpsamplerRunner : public GGMLRunner {
|
||||
std::unique_ptr<LatentUpsampler> model;
|
||||
|
||||
LatentUpsamplerRunner(ggml_backend_t backend,
|
||||
ggml_backend_t params_backend)
|
||||
: GGMLRunner(backend, params_backend) {}
|
||||
|
||||
std::string get_desc() override {
|
||||
return "ltx_latent_upsampler";
|
||||
}
|
||||
|
||||
bool load_from_file(const std::string& file_path, int n_threads) {
|
||||
LOG_INFO("loading LTX latent upsampler from '%s'", file_path.c_str());
|
||||
ModelLoader model_loader;
|
||||
if (!model_loader.init_from_file(file_path)) {
|
||||
LOG_ERROR("init LTX latent upsampler model loader from file failed: '%s'", file_path.c_str());
|
||||
return false;
|
||||
}
|
||||
|
||||
const auto& tensor_storage_map = model_loader.get_tensor_storage_map();
|
||||
bool has_regular_upsampler = has_tensor(tensor_storage_map, "upsampler.0.weight");
|
||||
bool has_rational_spatial = has_tensor(tensor_storage_map, "upsampler.conv.weight");
|
||||
if (!has_tensor(tensor_storage_map, "post_upsample_res_blocks.0.conv2.bias") ||
|
||||
(!has_regular_upsampler && !has_rational_spatial)) {
|
||||
LOG_ERROR("unsupported LTX latent upsampler weights: expected upsampler tensors");
|
||||
return false;
|
||||
}
|
||||
|
||||
LatentUpsamplerConfig config = detect_config_from_weights(tensor_storage_map);
|
||||
if (config.dims != 3 || (!config.spatial_upsample && !config.temporal_upsample) ||
|
||||
config.spatial_up_num < 1 || config.spatial_down_den < 1 || config.temporal_up_factor < 1) {
|
||||
LOG_ERROR("unsupported LTX latent upsampler config: dims=%d spatial=%d temporal=%d rational=%d scale=%.3f temporal_factor=%d",
|
||||
config.dims,
|
||||
config.spatial_upsample,
|
||||
config.temporal_upsample,
|
||||
config.rational_resampler,
|
||||
config.spatial_scale,
|
||||
config.temporal_up_factor);
|
||||
return false;
|
||||
}
|
||||
|
||||
model = std::make_unique<LatentUpsampler>(config);
|
||||
model->init(params_ctx, tensor_storage_map, "");
|
||||
if (!alloc_params_buffer()) {
|
||||
LOG_ERROR("LTX latent upsampler params buffer allocation failed");
|
||||
return false;
|
||||
}
|
||||
|
||||
std::map<std::string, ggml_tensor*> tensors;
|
||||
model->get_param_tensors(tensors);
|
||||
std::set<std::string> ignore_tensors;
|
||||
if (config.rational_resampler) {
|
||||
ignore_tensors.insert("upsampler.blur_down.kernel");
|
||||
}
|
||||
if (!model_loader.load_tensors(tensors, ignore_tensors, n_threads)) {
|
||||
LOG_ERROR("load LTX latent upsampler tensors failed");
|
||||
return false;
|
||||
}
|
||||
model->load_fixed_tensors();
|
||||
|
||||
LOG_INFO("LTX latent upsampler loaded: in_channels=%" PRId64 ", mid_channels=%" PRId64 ", blocks=%d, scale=%.3f, temporal_factor=%d, rational=%d",
|
||||
config.in_channels,
|
||||
config.mid_channels,
|
||||
config.num_blocks_per_stage,
|
||||
config.spatial_scale,
|
||||
config.temporal_up_factor,
|
||||
config.rational_resampler);
|
||||
return true;
|
||||
}
|
||||
|
||||
ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor) {
|
||||
if (!model) {
|
||||
return nullptr;
|
||||
}
|
||||
ggml_cgraph* gf = new_graph_custom(LTX_UPSAMPLER_GRAPH_SIZE);
|
||||
ggml_tensor* x = make_input(x_tensor);
|
||||
auto runner_ctx = get_context();
|
||||
ggml_tensor* out = model->forward(&runner_ctx, x);
|
||||
ggml_build_forward_expand(gf, out);
|
||||
return gf;
|
||||
}
|
||||
|
||||
sd::Tensor<float> compute(const int n_threads,
|
||||
const sd::Tensor<float>& x) {
|
||||
if (!model) {
|
||||
LOG_ERROR("LTX latent upsampler is not loaded");
|
||||
return {};
|
||||
}
|
||||
if (x.dim() != 4 && x.dim() != 5) {
|
||||
LOG_ERROR("LTX latent upsampler expects 4D or 5D video latent, got dim=%lld",
|
||||
(long long)x.dim());
|
||||
return {};
|
||||
}
|
||||
if (x.dim() == 5 && x.shape()[4] != 1) {
|
||||
LOG_ERROR("LTX latent upsampler currently supports batch size 1, got batch=%lld",
|
||||
(long long)x.shape()[4]);
|
||||
return {};
|
||||
}
|
||||
if (x.shape()[3] != model->config.in_channels) {
|
||||
LOG_ERROR("LTX latent upsampler expected %" PRId64 " channels, got %lld",
|
||||
model->config.in_channels,
|
||||
(long long)x.shape()[3]);
|
||||
return {};
|
||||
}
|
||||
size_t expected_dim = static_cast<size_t>(x.dim());
|
||||
auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x); };
|
||||
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), expected_dim);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace LTXVUpsampler
|
||||
|
||||
#endif // __SD_LTX_LATENT_UPSCALER_HPP__
|
||||
+280
-28
@@ -1,6 +1,7 @@
|
||||
#ifndef __SD_LTX_VAE_HPP__
|
||||
#define __SD_LTX_VAE_HPP__
|
||||
|
||||
#include <algorithm>
|
||||
#include <fstream>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
@@ -143,16 +144,25 @@ namespace LTXVAE {
|
||||
std::vector<ggml_tensor*>& feat_map,
|
||||
int& feat_idx,
|
||||
int chunk_idx,
|
||||
bool causal = true) {
|
||||
bool causal = true,
|
||||
int temporal_pad = 0) {
|
||||
auto conv = std::dynamic_pointer_cast<Conv3d>(blocks["conv"]);
|
||||
const int pad = causal ? (time_kernel_size - 1) : (time_kernel_size - 1) / 2;
|
||||
|
||||
ggml_tensor* prev = (feat_idx < (int)feat_map.size()) ? feat_map[feat_idx] : nullptr;
|
||||
|
||||
GGML_ASSERT(x->ne[2] >= temporal_pad);
|
||||
|
||||
int end_idx = (int)x->ne[2] - temporal_pad;
|
||||
int start_idx = std::max(end_idx - pad, 0);
|
||||
|
||||
// Save a contiguous copy of the last `pad` frames so the large `x`
|
||||
// tensor is not kept alive across iterations by a dangling view.
|
||||
if (feat_idx < (int)feat_map.size() && pad > 0 && x->ne[2] >= pad) {
|
||||
auto slice = ggml_ext_slice(ctx->ggml_ctx, x, 2, x->ne[2] - pad, x->ne[2]);
|
||||
if (feat_idx < (int)feat_map.size() && end_idx - start_idx > 0) {
|
||||
GGML_ASSERT(start_idx >= 0);
|
||||
GGML_ASSERT(end_idx > 0);
|
||||
|
||||
auto slice = ggml_ext_slice(ctx->ggml_ctx, x, 2, start_idx, end_idx);
|
||||
feat_map[feat_idx] = ggml_cont(ctx->ggml_ctx, slice);
|
||||
}
|
||||
feat_idx++;
|
||||
@@ -284,7 +294,8 @@ namespace LTXVAE {
|
||||
bool causal,
|
||||
std::vector<ggml_tensor*>& feat_map,
|
||||
int& feat_idx,
|
||||
int chunk_idx) {
|
||||
int chunk_idx,
|
||||
int temporal_pad = 0) {
|
||||
auto norm1 = std::dynamic_pointer_cast<PixelNorm3D>(blocks["norm1"]);
|
||||
auto conv1 = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv1"]);
|
||||
auto norm2 = std::dynamic_pointer_cast<PixelNorm3D>(blocks["norm2"]);
|
||||
@@ -311,14 +322,14 @@ namespace LTXVAE {
|
||||
h = apply_scale_shift(ctx->ggml_ctx, h, scale1, shift1);
|
||||
}
|
||||
h = ggml_silu_inplace(ctx->ggml_ctx, h);
|
||||
h = conv1->forward(ctx, h, feat_map, feat_idx, chunk_idx, causal);
|
||||
h = conv1->forward(ctx, h, feat_map, feat_idx, chunk_idx, causal, temporal_pad);
|
||||
|
||||
h = norm2->forward(ctx, h);
|
||||
if (timestep_conditioning) {
|
||||
h = apply_scale_shift(ctx->ggml_ctx, h, scale2, shift2);
|
||||
}
|
||||
h = ggml_silu_inplace(ctx->ggml_ctx, h);
|
||||
h = conv2->forward(ctx, h, feat_map, feat_idx, chunk_idx, causal);
|
||||
h = conv2->forward(ctx, h, feat_map, feat_idx, chunk_idx, causal, temporal_pad);
|
||||
|
||||
return ggml_add(ctx->ggml_ctx, h, x);
|
||||
}
|
||||
@@ -367,7 +378,8 @@ namespace LTXVAE {
|
||||
bool causal,
|
||||
std::vector<ggml_tensor*>& feat_map,
|
||||
int& feat_idx,
|
||||
int chunk_idx) {
|
||||
int chunk_idx,
|
||||
int temporal_pad = 0) {
|
||||
ggml_tensor* timestep_embed = nullptr;
|
||||
if (timestep_conditioning) {
|
||||
GGML_ASSERT(timestep != nullptr);
|
||||
@@ -376,7 +388,7 @@ namespace LTXVAE {
|
||||
}
|
||||
for (int i = 0; i < num_layers; i++) {
|
||||
auto resnet = std::dynamic_pointer_cast<ResnetBlock3D>(blocks["res_blocks." + std::to_string(i)]);
|
||||
x = resnet->forward(ctx, x, timestep_embed, causal, feat_map, feat_idx, chunk_idx);
|
||||
x = resnet->forward(ctx, x, timestep_embed, causal, feat_map, feat_idx, chunk_idx, temporal_pad);
|
||||
}
|
||||
return x;
|
||||
}
|
||||
@@ -437,7 +449,8 @@ namespace LTXVAE {
|
||||
bool causal,
|
||||
std::vector<ggml_tensor*>& feat_map,
|
||||
int& feat_idx,
|
||||
int chunk_idx) {
|
||||
int chunk_idx,
|
||||
int temporal_pad = 0) {
|
||||
auto conv = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv"]);
|
||||
|
||||
bool drop_first = (chunk_idx == 0) && (factor_t > 1);
|
||||
@@ -453,7 +466,7 @@ namespace LTXVAE {
|
||||
x_in = res;
|
||||
}
|
||||
|
||||
x = conv->forward(ctx, x, feat_map, feat_idx, chunk_idx, causal);
|
||||
x = conv->forward(ctx, x, feat_map, feat_idx, chunk_idx, causal, temporal_pad);
|
||||
x = depth_to_space_3d(ctx->ggml_ctx, x, get_output_channels(), factor_t, factor_s, drop_first);
|
||||
if (residual) {
|
||||
x = ggml_add(ctx->ggml_ctx, x, x_in);
|
||||
@@ -986,7 +999,8 @@ namespace LTXVAE {
|
||||
ggml_tensor* timestep,
|
||||
std::vector<ggml_tensor*>& feat_map,
|
||||
int& feat_idx,
|
||||
int chunk_idx) {
|
||||
int chunk_idx,
|
||||
int& temporal_pad) {
|
||||
auto conv_in = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv_in"]);
|
||||
auto conv_norm_out = std::dynamic_pointer_cast<PixelNorm3D>(blocks["conv_norm_out"]);
|
||||
auto conv_out = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv_out"]);
|
||||
@@ -998,7 +1012,7 @@ namespace LTXVAE {
|
||||
}
|
||||
|
||||
// conv_in with feat_map for left temporal context
|
||||
x = conv_in->forward(ctx, x, feat_map, feat_idx, chunk_idx, causal_decoder);
|
||||
x = conv_in->forward(ctx, x, feat_map, feat_idx, chunk_idx, causal_decoder, temporal_pad);
|
||||
|
||||
// up_blocks
|
||||
int block_idx = 0;
|
||||
@@ -1006,12 +1020,13 @@ namespace LTXVAE {
|
||||
auto mid_block = std::dynamic_pointer_cast<UNetMidBlock3D>(blocks["up_blocks." + std::to_string(block_idx)]);
|
||||
if (mid_block) {
|
||||
x = mid_block->forward(ctx, x, scaled_timestep, causal_decoder,
|
||||
feat_map, feat_idx, chunk_idx);
|
||||
feat_map, feat_idx, chunk_idx, temporal_pad);
|
||||
} else {
|
||||
auto upsample = std::dynamic_pointer_cast<DepthToSpaceUpsample>(
|
||||
blocks["up_blocks." + std::to_string(block_idx)]);
|
||||
x = upsample->forward(ctx, x, causal_decoder,
|
||||
feat_map, feat_idx, chunk_idx);
|
||||
feat_map, feat_idx, chunk_idx, temporal_pad);
|
||||
temporal_pad *= upsample->factor_t;
|
||||
}
|
||||
block_idx++;
|
||||
}
|
||||
@@ -1028,7 +1043,7 @@ namespace LTXVAE {
|
||||
x = apply_scale_shift(ctx->ggml_ctx, x, scale, shift);
|
||||
}
|
||||
x = ggml_silu_inplace(ctx->ggml_ctx, x);
|
||||
x = conv_out->forward(ctx, x, feat_map, feat_idx, chunk_idx, causal_decoder);
|
||||
x = conv_out->forward(ctx, x, feat_map, feat_idx, chunk_idx, causal_decoder, temporal_pad);
|
||||
return x;
|
||||
}
|
||||
};
|
||||
@@ -1084,7 +1099,9 @@ namespace LTXVAE {
|
||||
// tensors can be freed by GGML before the next iteration starts.
|
||||
ggml_tensor* decode_tiled(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* z,
|
||||
ggml_tensor* timestep) {
|
||||
ggml_tensor* timestep,
|
||||
int temporal_window_size = 1,
|
||||
int temporal_tile_overlap = 0) {
|
||||
auto decoder = std::dynamic_pointer_cast<Decoder>(blocks["decoder"]);
|
||||
auto processor = std::dynamic_pointer_cast<PerChannelStatistics>(blocks["per_channel_statistics"]);
|
||||
auto latents = processor->un_normalize(ctx, z);
|
||||
@@ -1099,18 +1116,69 @@ namespace LTXVAE {
|
||||
// 128 slots is generous enough for any supported decoder configuration.
|
||||
std::vector<ggml_tensor*> feat_map(128, nullptr);
|
||||
|
||||
// Ensure window size is at least 1
|
||||
int window = std::max(1, temporal_window_size);
|
||||
int overlap = std::max(0, temporal_tile_overlap);
|
||||
|
||||
if (overlap >= window) {
|
||||
LOG_WARN("temporal_tile_overlap (%d) is greater than or equal to temporal_tile_frames (%d), adjusting values to avoid empty decode windows",
|
||||
overlap, window);
|
||||
overlap = window - 1;
|
||||
}
|
||||
LOG_DEBUG("Using temporal tiling: temporal_tile_frames = %d, temporal_tile_overlap = %d, total frames = %d, resulting in %d tiles",
|
||||
window,
|
||||
overlap,
|
||||
(int)T,
|
||||
(T + window - overlap - 1) / (window - overlap));
|
||||
ggml_tensor* out = nullptr;
|
||||
for (int i = 0; i < (int)T; i++) {
|
||||
for (int i = 0; i < (int)T - overlap; i += (window - overlap)) {
|
||||
int feat_idx = 0;
|
||||
auto z_i = ggml_ext_slice(ctx->ggml_ctx, latents, 2, i, i + 1);
|
||||
auto out_i = decoder->forward_tiled_frame(ctx, z_i, timestep,
|
||||
feat_map, feat_idx, i);
|
||||
out = (out == nullptr) ? out_i : ggml_concat(ctx->ggml_ctx, out, out_i, 2);
|
||||
|
||||
// Calculate the end index for the current temporal chunk
|
||||
int end_i = std::min((int)T, i + window);
|
||||
if (end_i >= (int)T) {
|
||||
overlap = 0; // avoid overlap issues in the last chunk
|
||||
}
|
||||
|
||||
int chunk_overlap = overlap; // modified by forward_tiled_frame temporal inflation
|
||||
|
||||
auto z_chunk = ggml_ext_slice(ctx->ggml_ctx, latents, 2, i, end_i);
|
||||
|
||||
auto out_chunk = decoder->forward_tiled_frame(ctx, z_chunk, timestep,
|
||||
feat_map, feat_idx, i, chunk_overlap);
|
||||
|
||||
// discard overlap frames if it's not the final chunk
|
||||
if (overlap > 0 && end_i < (int)T) {
|
||||
out_chunk = ggml_ext_slice(ctx->ggml_ctx, out_chunk, 2, 0, out_chunk->ne[2] - chunk_overlap);
|
||||
}
|
||||
|
||||
out = (out == nullptr) ? out_chunk : ggml_concat(ctx->ggml_ctx, out, out_chunk, 2);
|
||||
}
|
||||
|
||||
return WAN::WanVAE::unpatchify(ctx->ggml_ctx, out, patch_size, 1);
|
||||
}
|
||||
|
||||
ggml_tensor* decode_tiled_chunk(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* z,
|
||||
ggml_tensor* timestep,
|
||||
std::vector<ggml_tensor*>& feat_map,
|
||||
int chunk_idx,
|
||||
int temporal_tile_overlap,
|
||||
int& feat_idx) {
|
||||
auto decoder = std::dynamic_pointer_cast<Decoder>(blocks["decoder"]);
|
||||
auto processor = std::dynamic_pointer_cast<PerChannelStatistics>(blocks["per_channel_statistics"]);
|
||||
auto latents = processor->un_normalize(ctx, z);
|
||||
|
||||
feat_idx = 0;
|
||||
int chunk_overlap = temporal_tile_overlap; // modified by forward_tiled_frame temporal inflation
|
||||
auto out_chunk = decoder->forward_tiled_frame(ctx, latents, timestep,
|
||||
feat_map, feat_idx, chunk_idx, chunk_overlap);
|
||||
if (chunk_overlap > 0) {
|
||||
out_chunk = ggml_ext_slice(ctx->ggml_ctx, out_chunk, 2, 0, out_chunk->ne[2] - chunk_overlap);
|
||||
}
|
||||
return WAN::WanVAE::unpatchify(ctx->ggml_ctx, out_chunk, patch_size, 1);
|
||||
}
|
||||
|
||||
ggml_tensor* encode(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* x) {
|
||||
GGML_ASSERT(!decode_only);
|
||||
@@ -1123,13 +1191,30 @@ namespace LTXVAE {
|
||||
mean = ggml_cont(ctx->ggml_ctx, mean);
|
||||
return processor->normalize(ctx, mean);
|
||||
}
|
||||
|
||||
ggml_tensor* normalize_latents(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* x) {
|
||||
auto processor = std::dynamic_pointer_cast<PerChannelStatistics>(blocks["per_channel_statistics"]);
|
||||
return processor->normalize(ctx, x);
|
||||
}
|
||||
|
||||
ggml_tensor* un_normalize_latents(GGMLRunnerContext* ctx,
|
||||
ggml_tensor* x) {
|
||||
auto processor = std::dynamic_pointer_cast<PerChannelStatistics>(blocks["per_channel_statistics"]);
|
||||
return processor->un_normalize(ctx, x);
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace LTXVAE
|
||||
|
||||
struct LTXVideoVAE : public VAE {
|
||||
static constexpr int DEFAULT_TEMPORAL_TILE_FRAMES = 4;
|
||||
static constexpr int DEFAULT_TEMPORAL_TILE_OVERLAP = 1;
|
||||
|
||||
bool decode_only;
|
||||
bool temporal_tiling_enabled = false;
|
||||
int temporal_tile_frames = DEFAULT_TEMPORAL_TILE_FRAMES;
|
||||
int temporal_tile_overlap = DEFAULT_TEMPORAL_TILE_OVERLAP;
|
||||
int ltx_vae_version;
|
||||
bool timestep_conditioning;
|
||||
int patch_size;
|
||||
@@ -1166,10 +1251,64 @@ struct LTXVideoVAE : public VAE {
|
||||
temporal_tiling_enabled = enabled;
|
||||
}
|
||||
|
||||
void set_tiling_params(const sd_tiling_params_t& params) override {
|
||||
temporal_tiling_enabled = params.temporal_tiling;
|
||||
temporal_tile_frames = DEFAULT_TEMPORAL_TILE_FRAMES;
|
||||
temporal_tile_overlap = DEFAULT_TEMPORAL_TILE_OVERLAP;
|
||||
|
||||
for (const auto& [key, value] : parse_key_value_args(params.extra_tiling_args, "LTX VAE extra tiling arg")) {
|
||||
int parsed = 0;
|
||||
if (!parse_strict_int(value, parsed)) {
|
||||
LOG_WARN("ignoring invalid LTX VAE extra tiling arg '%s=%s'", key.c_str(), value.c_str());
|
||||
} else if (key == "temporal_tile_frames") {
|
||||
temporal_tile_frames = std::max(1, parsed);
|
||||
} else if (key == "temporal_tile_overlap") {
|
||||
temporal_tile_overlap = std::max(0, parsed);
|
||||
} else {
|
||||
LOG_WARN("ignoring unknown LTX VAE extra tiling arg '%s'", key.c_str());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) override {
|
||||
vae.get_param_tensors(tensors, prefix);
|
||||
}
|
||||
|
||||
struct TemporalTilePlan {
|
||||
int frames = 1;
|
||||
int overlap = 0;
|
||||
int stride = 1;
|
||||
int num_tiles = 1;
|
||||
};
|
||||
|
||||
TemporalTilePlan resolve_temporal_tile_plan(int64_t total_frames) const {
|
||||
TemporalTilePlan plan;
|
||||
plan.frames = std::max(1, temporal_tile_frames);
|
||||
plan.overlap = std::max(0, temporal_tile_overlap);
|
||||
|
||||
if (plan.overlap >= plan.frames) {
|
||||
LOG_WARN("temporal_tile_overlap (%d) is greater than or equal to temporal_tile_frames (%d), adjusting values to avoid empty decode windows",
|
||||
plan.overlap,
|
||||
plan.frames);
|
||||
plan.overlap = plan.frames - 1;
|
||||
}
|
||||
if (total_frames > 1 && plan.overlap >= total_frames) {
|
||||
LOG_WARN("temporal_tile_overlap (%d) is greater than or equal to total latent frames (%lld), adjusting values to decode at least one tile",
|
||||
plan.overlap,
|
||||
(long long)total_frames);
|
||||
plan.overlap = static_cast<int>(total_frames - 1);
|
||||
}
|
||||
|
||||
plan.stride = std::max(1, plan.frames - plan.overlap);
|
||||
int64_t tiled_frames = std::max<int64_t>(1, total_frames - plan.overlap);
|
||||
plan.num_tiles = total_frames > 0 ? static_cast<int>((tiled_frames + plan.stride - 1) / plan.stride) : 0;
|
||||
return plan;
|
||||
}
|
||||
|
||||
std::string temporal_feat_cache_name(size_t feat_idx) const {
|
||||
return "ltx_vae_temporal_feat:" + std::to_string(feat_idx);
|
||||
}
|
||||
|
||||
ggml_cgraph* build_graph(const sd::Tensor<float>& z_tensor, bool decode_graph) {
|
||||
ggml_cgraph* gf = new_graph_custom(20480);
|
||||
ggml_tensor* z = make_input(z_tensor);
|
||||
@@ -1180,18 +1319,108 @@ struct LTXVideoVAE : public VAE {
|
||||
|
||||
auto runner_ctx = get_context();
|
||||
ggml_tensor* out;
|
||||
bool use_tiled = decode_graph && temporal_tiling_enabled &&
|
||||
z_tensor.dim() == 5 && z_tensor.shape()[2] > 1;
|
||||
if (use_tiled) {
|
||||
out = vae.decode_tiled(&runner_ctx, z, timestep);
|
||||
} else {
|
||||
out = decode_graph ? vae.decode(&runner_ctx, z, timestep) : vae.encode(&runner_ctx, z);
|
||||
}
|
||||
out = decode_graph ? vae.decode(&runner_ctx, z, timestep) : vae.encode(&runner_ctx, z);
|
||||
ggml_build_forward_expand(gf, out);
|
||||
|
||||
return gf;
|
||||
}
|
||||
|
||||
ggml_cgraph* build_temporal_tile_graph(const sd::Tensor<float>& z_chunk_tensor,
|
||||
int chunk_idx,
|
||||
int chunk_overlap) {
|
||||
ggml_cgraph* gf = new_graph_custom(20480);
|
||||
ggml_tensor* z = make_input(z_chunk_tensor);
|
||||
ggml_tensor* timestep = nullptr;
|
||||
if (timestep_conditioning) {
|
||||
timestep = make_input(decode_timestep_tensor);
|
||||
}
|
||||
|
||||
std::vector<ggml_tensor*> feat_map(128, nullptr);
|
||||
for (size_t feat_idx = 0; feat_idx < feat_map.size(); ++feat_idx) {
|
||||
feat_map[feat_idx] = get_cache_tensor_by_name(temporal_feat_cache_name(feat_idx));
|
||||
}
|
||||
|
||||
auto runner_ctx = get_context();
|
||||
int feat_count = 0;
|
||||
ggml_tensor* out = vae.decode_tiled_chunk(&runner_ctx,
|
||||
z,
|
||||
timestep,
|
||||
feat_map,
|
||||
chunk_idx,
|
||||
chunk_overlap,
|
||||
feat_count);
|
||||
|
||||
for (int feat_idx = 0; feat_idx < feat_count && feat_idx < static_cast<int>(feat_map.size()); ++feat_idx) {
|
||||
ggml_tensor* feat_cache = feat_map[static_cast<size_t>(feat_idx)];
|
||||
if (feat_cache != nullptr) {
|
||||
cache(temporal_feat_cache_name(static_cast<size_t>(feat_idx)), feat_cache);
|
||||
ggml_build_forward_expand(gf, feat_cache);
|
||||
}
|
||||
}
|
||||
|
||||
ggml_build_forward_expand(gf, out);
|
||||
return gf;
|
||||
}
|
||||
|
||||
sd::Tensor<float> decode_temporal_tiled_streaming(const int n_threads,
|
||||
const sd::Tensor<float>& input,
|
||||
size_t expected_dim) {
|
||||
const int64_t total_frames = input.shape()[2];
|
||||
TemporalTilePlan plan = resolve_temporal_tile_plan(total_frames);
|
||||
|
||||
LOG_DEBUG("Using streaming temporal tiling: temporal_tile_frames=%d, temporal_tile_overlap=%d, total latent frames=%lld, resulting in %d tiles",
|
||||
plan.frames,
|
||||
plan.overlap,
|
||||
(long long)total_frames,
|
||||
plan.num_tiles);
|
||||
|
||||
free_cache_ctx_and_buffer();
|
||||
cache_tensor_map.clear();
|
||||
|
||||
sd::Tensor<float> output;
|
||||
for (int64_t start = 0; start < total_frames - plan.overlap; start += plan.stride) {
|
||||
const int64_t end = std::min<int64_t>(total_frames, start + plan.frames);
|
||||
const int chunk_overlap = end < total_frames ? plan.overlap : 0;
|
||||
auto z_chunk = sd::ops::slice(input, 2, start, end);
|
||||
|
||||
LOG_DEBUG("LTX VAE temporal tile %lld/%d: latent frames [%lld, %lld), overlap=%d",
|
||||
(long long)(start / plan.stride + 1),
|
||||
plan.num_tiles,
|
||||
(long long)start,
|
||||
(long long)end,
|
||||
chunk_overlap);
|
||||
|
||||
auto get_graph = [&]() -> ggml_cgraph* {
|
||||
return build_temporal_tile_graph(z_chunk,
|
||||
static_cast<int>(start),
|
||||
chunk_overlap);
|
||||
};
|
||||
auto chunk = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, true),
|
||||
expected_dim);
|
||||
if (chunk.empty()) {
|
||||
free_cache_ctx_and_buffer();
|
||||
cache_tensor_map.clear();
|
||||
return {};
|
||||
}
|
||||
output = output.empty() ? std::move(chunk) : sd::ops::concat(output, chunk, 2);
|
||||
}
|
||||
|
||||
free_cache_ctx_and_buffer();
|
||||
cache_tensor_map.clear();
|
||||
return output;
|
||||
}
|
||||
|
||||
ggml_cgraph* build_latent_statistics_graph(const sd::Tensor<float>& z_tensor, bool normalize) {
|
||||
ggml_cgraph* gf = new_graph_custom(1024);
|
||||
ggml_tensor* z = make_input(z_tensor);
|
||||
|
||||
auto runner_ctx = get_context();
|
||||
ggml_tensor* out = normalize ? vae.normalize_latents(&runner_ctx, z)
|
||||
: vae.un_normalize_latents(&runner_ctx, z);
|
||||
ggml_build_forward_expand(gf, out);
|
||||
return gf;
|
||||
}
|
||||
|
||||
sd::Tensor<float> _compute(const int n_threads,
|
||||
const sd::Tensor<float>& z,
|
||||
bool decode_graph) override {
|
||||
@@ -1216,6 +1445,9 @@ struct LTXVideoVAE : public VAE {
|
||||
input = sd::ops::slice(input, 2, 0, cropped_t);
|
||||
}
|
||||
}
|
||||
if (decode_graph && temporal_tiling_enabled && input.dim() == 5 && input.shape()[2] > 1) {
|
||||
return decode_temporal_tiled_streaming(n_threads, input, expected_dim);
|
||||
}
|
||||
auto get_graph = [&]() -> ggml_cgraph* {
|
||||
return build_graph(input, decode_graph);
|
||||
};
|
||||
@@ -1226,6 +1458,26 @@ struct LTXVideoVAE : public VAE {
|
||||
return result;
|
||||
}
|
||||
|
||||
sd::Tensor<float> apply_latent_statistics(const int n_threads,
|
||||
const sd::Tensor<float>& z,
|
||||
bool normalize) {
|
||||
auto get_graph = [&]() -> ggml_cgraph* {
|
||||
return build_latent_statistics_graph(z, normalize);
|
||||
};
|
||||
return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false),
|
||||
static_cast<size_t>(z.dim()));
|
||||
}
|
||||
|
||||
sd::Tensor<float> normalize_latents(const int n_threads,
|
||||
const sd::Tensor<float>& z) {
|
||||
return apply_latent_statistics(n_threads, z, true);
|
||||
}
|
||||
|
||||
sd::Tensor<float> un_normalize_latents(const int n_threads,
|
||||
const sd::Tensor<float>& z) {
|
||||
return apply_latent_statistics(n_threads, z, false);
|
||||
}
|
||||
|
||||
int get_encoder_output_channels(int input_channels) override {
|
||||
SD_UNUSED(input_channels);
|
||||
return 256;
|
||||
|
||||
@@ -1487,6 +1487,9 @@ namespace LTXV {
|
||||
->forward(ctx, ggml_ext_scale(ctx->ggml_ctx, av_ca_audio_timestep, av_ca_factor))
|
||||
.first;
|
||||
|
||||
sd::ggml_graph_cut::mark_graph_cut(vx, "ltxav.prelude", "vx");
|
||||
sd::ggml_graph_cut::mark_graph_cut(ax, "ltxav.prelude", "ax");
|
||||
|
||||
for (int i = 0; i < cfg.num_layers; i++) {
|
||||
auto block = std::dynamic_pointer_cast<BasicAVTransformerBlock>(blocks["transformer_blocks." + std::to_string(i)]);
|
||||
auto out = block->forward(ctx,
|
||||
@@ -1509,6 +1512,8 @@ namespace LTXV {
|
||||
a_prompt_timestep_mod);
|
||||
vx = out.first;
|
||||
ax = out.second;
|
||||
sd::ggml_graph_cut::mark_graph_cut(vx, "ltxav.transformer_blocks." + std::to_string(i), "vx");
|
||||
sd::ggml_graph_cut::mark_graph_cut(ax, "ltxav.transformer_blocks." + std::to_string(i), "ax");
|
||||
}
|
||||
|
||||
auto v_shift_scale = get_output_scale_shift(ctx, params["scale_shift_table"], v_embedded_time, cfg.hidden_size);
|
||||
|
||||
+23
-11
@@ -410,7 +410,7 @@ bool ModelLoader::init_from_diffusers_file(const std::string& file_path, const s
|
||||
}
|
||||
|
||||
SDVersion ModelLoader::get_sd_version() {
|
||||
TensorStorage token_embedding_weight, input_block_weight;
|
||||
TensorStorage token_embedding_weight, input_block_weight, context_ebedding_weight;
|
||||
|
||||
bool has_multiple_encoders = false;
|
||||
bool is_unet = false;
|
||||
@@ -428,7 +428,8 @@ SDVersion ModelLoader::get_sd_version() {
|
||||
bool has_attn_1024 = false;
|
||||
|
||||
for (auto& [name, tensor_storage] : tensor_storage_map) {
|
||||
if (tensor_storage.name.find("model.diffusion_model.double_blocks.") != std::string::npos) {
|
||||
if (tensor_storage.name.find("model.diffusion_model.double_blocks.") != std::string::npos ||
|
||||
tensor_storage.name.find("model.diffusion_model.single_transformer_blocks.") != std::string::npos) {
|
||||
is_flux = true;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.nerf_final_layer_conv.") != std::string::npos) {
|
||||
@@ -441,6 +442,10 @@ SDVersion ModelLoader::get_sd_version() {
|
||||
tensor_storage_map.find("model.language_model.layers.0.self_attn.q_proj.weight") != tensor_storage_map.end()) {
|
||||
return VERSION_HIDREAM_O1;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.attn.norm_added_q.weight") != std::string::npos &&
|
||||
tensor_storage_map.find("model.diffusion_model.transformer_blocks.0.img_mlp.w1.weight") != tensor_storage_map.end()) {
|
||||
return VERSION_LENS;
|
||||
}
|
||||
if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.img_mod.1.weight") != std::string::npos) {
|
||||
return VERSION_QWEN_IMAGE;
|
||||
}
|
||||
@@ -522,6 +527,9 @@ SDVersion ModelLoader::get_sd_version() {
|
||||
tensor_storage.name == "unet.conv_in.weight") {
|
||||
input_block_weight = tensor_storage;
|
||||
}
|
||||
if (tensor_storage.name == "model.diffusion_model.txt_in.weight" || tensor_storage.name == "model.diffusion_model.context_embedder.weight") {
|
||||
context_ebedding_weight = tensor_storage;
|
||||
}
|
||||
}
|
||||
if (is_wan) {
|
||||
LOG_DEBUG("patch_embedding_channels %d", patch_embedding_channels);
|
||||
@@ -552,16 +560,20 @@ SDVersion ModelLoader::get_sd_version() {
|
||||
}
|
||||
|
||||
if (is_flux && !is_flux2) {
|
||||
if (input_block_weight.ne[0] == 384) {
|
||||
return VERSION_FLUX_FILL;
|
||||
if (context_ebedding_weight.ne[0] == 3584) {
|
||||
return VERSION_LONGCAT;
|
||||
} else {
|
||||
if (input_block_weight.ne[0] == 384) {
|
||||
return VERSION_FLUX_FILL;
|
||||
}
|
||||
if (input_block_weight.ne[0] == 128) {
|
||||
return VERSION_FLUX_CONTROLS;
|
||||
}
|
||||
if (input_block_weight.ne[0] == 196) {
|
||||
return VERSION_FLEX_2;
|
||||
}
|
||||
return VERSION_FLUX;
|
||||
}
|
||||
if (input_block_weight.ne[0] == 128) {
|
||||
return VERSION_FLUX_CONTROLS;
|
||||
}
|
||||
if (input_block_weight.ne[0] == 196) {
|
||||
return VERSION_FLEX_2;
|
||||
}
|
||||
return VERSION_FLUX;
|
||||
}
|
||||
|
||||
if (is_flux2) {
|
||||
|
||||
+20
-2
@@ -47,6 +47,8 @@ enum SDVersion {
|
||||
VERSION_Z_IMAGE,
|
||||
VERSION_OVIS_IMAGE,
|
||||
VERSION_ERNIE_IMAGE,
|
||||
VERSION_LENS,
|
||||
VERSION_LONGCAT,
|
||||
VERSION_COUNT,
|
||||
};
|
||||
|
||||
@@ -141,6 +143,13 @@ static inline bool sd_version_is_z_image(SDVersion version) {
|
||||
return false;
|
||||
}
|
||||
|
||||
static inline bool sd_version_is_longcat(SDVersion version) {
|
||||
if (version == VERSION_LONGCAT) {
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
static inline bool sd_version_is_ernie_image(SDVersion version) {
|
||||
if (version == VERSION_ERNIE_IMAGE) {
|
||||
return true;
|
||||
@@ -148,8 +157,15 @@ static inline bool sd_version_is_ernie_image(SDVersion version) {
|
||||
return false;
|
||||
}
|
||||
|
||||
static inline bool sd_version_is_lens(SDVersion version) {
|
||||
if (version == VERSION_LENS) {
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
static inline bool sd_version_uses_flux2_vae(SDVersion version) {
|
||||
if (sd_version_is_flux2(version) || sd_version_is_ernie_image(version)) {
|
||||
if (sd_version_is_flux2(version) || sd_version_is_ernie_image(version) || sd_version_is_lens(version)) {
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
@@ -176,7 +192,9 @@ static inline bool sd_version_is_dit(SDVersion version) {
|
||||
version == VERSION_HIDREAM_O1 ||
|
||||
sd_version_is_anima(version) ||
|
||||
sd_version_is_z_image(version) ||
|
||||
sd_version_is_ernie_image(version)) {
|
||||
sd_version_is_ernie_image(version) ||
|
||||
sd_version_is_lens(version) ||
|
||||
sd_version_is_longcat(version)) {
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
|
||||
+24
-3
@@ -128,6 +128,7 @@ std::string convert_cond_stage_model_name(std::string name, std::string prefix)
|
||||
};
|
||||
|
||||
static const std::vector<std::pair<std::string, std::string>> llm_name_map{
|
||||
{"attn_sinks.weight", "self_attn.sinks"},
|
||||
{"token_embd.", "model.embed_tokens."},
|
||||
{"blk.", "model.layers."},
|
||||
{"attn_q.", "self_attn.q_proj."},
|
||||
@@ -137,6 +138,12 @@ std::string convert_cond_stage_model_name(std::string name, std::string prefix)
|
||||
{"attn_k_norm.", "self_attn.k_norm."},
|
||||
{"attn_output.", "self_attn.o_proj."},
|
||||
{"attn_norm.", "input_layernorm."},
|
||||
{"attn_post_norm.", "post_attention_layernorm."},
|
||||
{"post_attention_norm.", "post_attention_layernorm."},
|
||||
{"ffn_gate_inp.", "mlp.router."},
|
||||
{"ffn_gate_exps.", "mlp.experts.gate_proj."},
|
||||
{"ffn_up_exps.", "mlp.experts.up_proj."},
|
||||
{"ffn_down_exps.", "mlp.experts.down_proj."},
|
||||
{"ffn_down.", "mlp.down_proj."},
|
||||
{"ffn_gate.", "mlp.gate_proj."},
|
||||
{"ffn_up.", "mlp.up_proj."},
|
||||
@@ -508,6 +515,12 @@ std::string convert_diffusers_dit_to_original_flux(std::string name) {
|
||||
static std::unordered_map<std::string, std::string> flux_name_map;
|
||||
|
||||
if (flux_name_map.empty()) {
|
||||
// --- time_embed (longcat) ---
|
||||
flux_name_map["time_embed.timestep_embedder.linear_1.weight"] = "time_in.in_layer.weight";
|
||||
flux_name_map["time_embed.timestep_embedder.linear_1.bias"] = "time_in.in_layer.bias";
|
||||
flux_name_map["time_embed.timestep_embedder.linear_2.weight"] = "time_in.out_layer.weight";
|
||||
flux_name_map["time_embed.timestep_embedder.linear_2.bias"] = "time_in.out_layer.bias";
|
||||
|
||||
// --- time_text_embed ---
|
||||
flux_name_map["time_text_embed.timestep_embedder.linear_1.weight"] = "time_in.in_layer.weight";
|
||||
flux_name_map["time_text_embed.timestep_embedder.linear_1.bias"] = "time_in.in_layer.bias";
|
||||
@@ -561,6 +574,11 @@ std::string convert_diffusers_dit_to_original_flux(std::string name) {
|
||||
flux_name_map[block_prefix + "attn.norm_k.weight"] = dst_prefix + "img_attn.norm.key_norm.scale";
|
||||
flux_name_map[block_prefix + "attn.norm_added_q.weight"] = dst_prefix + "txt_attn.norm.query_norm.scale";
|
||||
flux_name_map[block_prefix + "attn.norm_added_k.weight"] = dst_prefix + "txt_attn.norm.key_norm.scale";
|
||||
// Comfy-Org/LongCat-Image stores already-converted RMSNorm tensors as *.weight.
|
||||
flux_name_map[dst_prefix + "img_attn.norm.query_norm.weight"] = dst_prefix + "img_attn.norm.query_norm.scale";
|
||||
flux_name_map[dst_prefix + "img_attn.norm.key_norm.weight"] = dst_prefix + "img_attn.norm.key_norm.scale";
|
||||
flux_name_map[dst_prefix + "txt_attn.norm.query_norm.weight"] = dst_prefix + "txt_attn.norm.query_norm.scale";
|
||||
flux_name_map[dst_prefix + "txt_attn.norm.key_norm.weight"] = dst_prefix + "txt_attn.norm.key_norm.scale";
|
||||
|
||||
// ff
|
||||
flux_name_map[block_prefix + "ff.net.0.proj.weight"] = dst_prefix + "img_mlp.0.weight";
|
||||
@@ -599,8 +617,11 @@ std::string convert_diffusers_dit_to_original_flux(std::string name) {
|
||||
|
||||
flux_name_map[block_prefix + "attn.norm_q.weight"] = dst_prefix + "norm.query_norm.scale";
|
||||
flux_name_map[block_prefix + "attn.norm_k.weight"] = dst_prefix + "norm.key_norm.scale";
|
||||
flux_name_map[block_prefix + "proj_out.weight"] = dst_prefix + "linear2.weight";
|
||||
flux_name_map[block_prefix + "proj_out.bias"] = dst_prefix + "linear2.bias";
|
||||
// Comfy-Org/LongCat-Image stores already-converted RMSNorm tensors as *.weight.
|
||||
flux_name_map[dst_prefix + "norm.query_norm.weight"] = dst_prefix + "norm.query_norm.scale";
|
||||
flux_name_map[dst_prefix + "norm.key_norm.weight"] = dst_prefix + "norm.key_norm.scale";
|
||||
flux_name_map[block_prefix + "proj_out.weight"] = dst_prefix + "linear2.weight";
|
||||
flux_name_map[block_prefix + "proj_out.bias"] = dst_prefix + "linear2.bias";
|
||||
}
|
||||
|
||||
// --- final layers ---
|
||||
@@ -668,7 +689,7 @@ std::string convert_diffusion_model_name(std::string name, std::string prefix, S
|
||||
name = convert_diffusers_unet_to_original_sdxl(name);
|
||||
} else if (sd_version_is_sd3(version)) {
|
||||
name = convert_diffusers_dit_to_original_sd3(name);
|
||||
} else if (sd_version_is_flux(version) || sd_version_is_flux2(version)) {
|
||||
} else if (sd_version_is_flux(version) || sd_version_is_flux2(version) || sd_version_is_longcat(version)) {
|
||||
name = convert_diffusers_dit_to_original_flux(name);
|
||||
} else if (sd_version_is_z_image(version)) {
|
||||
name = convert_diffusers_dit_to_original_lumina2(name);
|
||||
|
||||
+75
-13
@@ -111,6 +111,16 @@ namespace Rope {
|
||||
return txt_ids;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_longcat_txt_ids(int bs, int context_len, int axes_dim_num) {
|
||||
auto txt_ids = std::vector<std::vector<float>>(bs * context_len, std::vector<float>(axes_dim_num, 0.0f));
|
||||
for (int i = 0; i < bs * context_len; i++) {
|
||||
float token_index = static_cast<float>(i % context_len);
|
||||
txt_ids[i][1] = token_index;
|
||||
txt_ids[i][2] = token_index;
|
||||
}
|
||||
return txt_ids;
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_flux_img_ids(int h,
|
||||
int w,
|
||||
int patch_size,
|
||||
@@ -122,7 +132,6 @@ namespace Rope {
|
||||
bool scale_rope = false) {
|
||||
int h_len = (h + (patch_size / 2)) / patch_size;
|
||||
int w_len = (w + (patch_size / 2)) / patch_size;
|
||||
|
||||
std::vector<std::vector<float>> img_ids(h_len * w_len, std::vector<float>(axes_dim_num, 0.0));
|
||||
|
||||
int h_start = h_offset;
|
||||
@@ -135,7 +144,6 @@ namespace Rope {
|
||||
|
||||
std::vector<float> row_ids = linspace<float>(1.f * h_start, 1.f * h_start + h_len - 1, h_len);
|
||||
std::vector<float> col_ids = linspace<float>(1.f * w_start, 1.f * w_start + w_len - 1, w_len);
|
||||
|
||||
for (int i = 0; i < h_len; ++i) {
|
||||
for (int j = 0; j < w_len; ++j) {
|
||||
img_ids[i * w_len + j][0] = 1.f * index;
|
||||
@@ -244,14 +252,16 @@ namespace Rope {
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_refs_ids(int patch_size,
|
||||
int bs,
|
||||
int axes_dim_num,
|
||||
int start_index,
|
||||
const std::vector<ggml_tensor*>& ref_latents,
|
||||
bool increase_ref_index,
|
||||
float ref_index_scale,
|
||||
bool scale_rope) {
|
||||
bool scale_rope,
|
||||
int base_offset = 0) {
|
||||
std::vector<std::vector<float>> ids;
|
||||
int curr_h_offset = 0;
|
||||
int curr_w_offset = 0;
|
||||
int index = 1;
|
||||
int index = start_index;
|
||||
for (ggml_tensor* ref : ref_latents) {
|
||||
int h_offset = 0;
|
||||
int w_offset = 0;
|
||||
@@ -270,8 +280,8 @@ namespace Rope {
|
||||
bs,
|
||||
axes_dim_num,
|
||||
static_cast<int>(index * ref_index_scale),
|
||||
h_offset,
|
||||
w_offset,
|
||||
h_offset + base_offset,
|
||||
w_offset + base_offset,
|
||||
scale_rope);
|
||||
ids = concat_ids(ids, ref_ids, bs);
|
||||
|
||||
@@ -294,13 +304,17 @@ namespace Rope {
|
||||
std::set<int> txt_arange_dims,
|
||||
const std::vector<ggml_tensor*>& ref_latents,
|
||||
bool increase_ref_index,
|
||||
float ref_index_scale) {
|
||||
auto txt_ids = gen_flux_txt_ids(bs, context_len, axes_dim_num, txt_arange_dims);
|
||||
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num);
|
||||
float ref_index_scale,
|
||||
bool is_longcat) {
|
||||
int x_index = is_longcat ? 1 : 0;
|
||||
|
||||
auto txt_ids = is_longcat ? gen_longcat_txt_ids(bs, context_len, axes_dim_num) : gen_flux_txt_ids(bs, context_len, axes_dim_num, txt_arange_dims);
|
||||
int offset = is_longcat ? context_len : 0;
|
||||
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, x_index, offset, offset);
|
||||
|
||||
auto ids = concat_ids(txt_ids, img_ids, bs);
|
||||
if (ref_latents.size() > 0) {
|
||||
auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, ref_latents, increase_ref_index, ref_index_scale, false);
|
||||
auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, x_index + 1, ref_latents, increase_ref_index, ref_index_scale, false, offset);
|
||||
ids = concat_ids(ids, refs_ids, bs);
|
||||
}
|
||||
return ids;
|
||||
@@ -319,7 +333,8 @@ namespace Rope {
|
||||
int theta,
|
||||
bool circular_h,
|
||||
bool circular_w,
|
||||
const std::vector<int>& axes_dim) {
|
||||
const std::vector<int>& axes_dim,
|
||||
bool is_longcat) {
|
||||
std::vector<std::vector<float>> ids = gen_flux_ids(h,
|
||||
w,
|
||||
patch_size,
|
||||
@@ -329,7 +344,8 @@ namespace Rope {
|
||||
txt_arange_dims,
|
||||
ref_latents,
|
||||
increase_ref_index,
|
||||
ref_index_scale);
|
||||
ref_index_scale,
|
||||
is_longcat);
|
||||
std::vector<std::vector<int>> wrap_dims;
|
||||
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
|
||||
int h_len = (h + (patch_size / 2)) / patch_size;
|
||||
@@ -394,7 +410,7 @@ namespace Rope {
|
||||
auto img_ids = gen_flux_img_ids(h, w, patch_size, bs, axes_dim_num, 0, 0, 0, true);
|
||||
auto ids = concat_ids(txt_ids_repeated, img_ids, bs);
|
||||
if (ref_latents.size() > 0) {
|
||||
auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, ref_latents, increase_ref_index, 1.f, true);
|
||||
auto refs_ids = gen_refs_ids(patch_size, bs, axes_dim_num, 1, ref_latents, increase_ref_index, 1.f, true);
|
||||
ids = concat_ids(ids, refs_ids, bs);
|
||||
}
|
||||
return ids;
|
||||
@@ -462,6 +478,52 @@ namespace Rope {
|
||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_lens_ids(int h,
|
||||
int w,
|
||||
int bs,
|
||||
int context_len,
|
||||
bool scale_rope = true) {
|
||||
auto img_ids_repeated = gen_flux_img_ids(h, w, 1, bs, 3, 0, 0, 0, scale_rope);
|
||||
|
||||
int txt_id_start = scale_rope ? std::max(h / 2, w / 2) : 0;
|
||||
auto txt_ids = linspace<float>(1.f * txt_id_start, 1.f * context_len + txt_id_start, context_len);
|
||||
std::vector<std::vector<float>> txt_ids_repeated(bs * context_len, std::vector<float>(3));
|
||||
for (int i = 0; i < bs; ++i) {
|
||||
for (int j = 0; j < txt_ids.size(); ++j) {
|
||||
txt_ids_repeated[i * txt_ids.size() + j] = {txt_ids[j], txt_ids[j], txt_ids[j]};
|
||||
}
|
||||
}
|
||||
|
||||
return concat_ids(img_ids_repeated, txt_ids_repeated, bs);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<float> gen_lens_pe(int h,
|
||||
int w,
|
||||
int bs,
|
||||
int context_len,
|
||||
int theta,
|
||||
bool circular_h,
|
||||
bool circular_w,
|
||||
const std::vector<int>& axes_dim) {
|
||||
std::vector<std::vector<float>> ids = gen_lens_ids(h, w, bs, context_len, true);
|
||||
std::vector<std::vector<int>> wrap_dims;
|
||||
if ((circular_h || circular_w) && bs > 0 && axes_dim.size() >= 3) {
|
||||
size_t pos_len = ids.size() / bs;
|
||||
wrap_dims.assign(axes_dim.size(), std::vector<int>(pos_len, 0));
|
||||
const size_t img_tokens = static_cast<size_t>(h) * static_cast<size_t>(w);
|
||||
for (size_t token_i = 0; token_i < img_tokens; ++token_i) {
|
||||
if (circular_h) {
|
||||
wrap_dims[1][token_i] = h;
|
||||
}
|
||||
if (circular_w) {
|
||||
wrap_dims[2][token_i] = w;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
|
||||
}
|
||||
|
||||
__STATIC_INLINE__ std::vector<std::vector<float>> gen_ernie_image_ids(int h,
|
||||
int w,
|
||||
int patch_size,
|
||||
|
||||
+710
-147
File diff suppressed because it is too large
Load Diff
+118
-28
@@ -259,10 +259,54 @@ public:
|
||||
}
|
||||
};
|
||||
|
||||
ggml_tensor* patchify(ggml_context* ctx,
|
||||
ggml_tensor* x,
|
||||
int64_t patch_size,
|
||||
int64_t b = 1) {
|
||||
class WideMemBlock : public GGMLBlock {
|
||||
bool has_skip_conv = false;
|
||||
|
||||
public:
|
||||
WideMemBlock(int channels, int out_channels)
|
||||
: has_skip_conv(channels != out_channels) {
|
||||
int groups = std::max(1, out_channels / 64);
|
||||
blocks["conv.0"] = std::shared_ptr<GGMLBlock>(new Conv2d(channels * 2, out_channels, {1, 1}, {1, 1}));
|
||||
blocks["conv.2"] = std::shared_ptr<GGMLBlock>(new Conv2d_grouped(out_channels, out_channels, groups, {3, 3}, {1, 1}, {1, 1}));
|
||||
blocks["conv.4"] = std::shared_ptr<GGMLBlock>(new Conv2d(out_channels, out_channels, {1, 1}, {1, 1}));
|
||||
blocks["conv.6"] = std::shared_ptr<GGMLBlock>(new Conv2d_grouped(out_channels, out_channels, groups, {3, 3}, {1, 1}, {1, 1}));
|
||||
if (has_skip_conv) {
|
||||
blocks["skip"] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, out_channels, {1, 1}, {1, 1}, {0, 0}, {1, 1}, false));
|
||||
}
|
||||
}
|
||||
|
||||
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* past) {
|
||||
// x: [n, channels, h, w]
|
||||
auto conv0 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.0"]);
|
||||
auto conv1 = std::dynamic_pointer_cast<Conv2d_grouped>(blocks["conv.2"]);
|
||||
auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.4"]);
|
||||
auto conv3 = std::dynamic_pointer_cast<Conv2d_grouped>(blocks["conv.6"]);
|
||||
|
||||
auto h = ggml_concat(ctx->ggml_ctx, x, past, 2);
|
||||
h = conv0->forward(ctx, h);
|
||||
h = ggml_relu_inplace(ctx->ggml_ctx, h);
|
||||
h = conv1->forward(ctx, h);
|
||||
h = ggml_relu_inplace(ctx->ggml_ctx, h);
|
||||
h = conv2->forward(ctx, h);
|
||||
h = ggml_relu_inplace(ctx->ggml_ctx, h);
|
||||
h = conv3->forward(ctx, h);
|
||||
|
||||
auto skip = x;
|
||||
if (has_skip_conv) {
|
||||
auto skip_conv = std::dynamic_pointer_cast<Conv2d>(blocks["skip"]);
|
||||
skip = skip_conv->forward(ctx, x);
|
||||
}
|
||||
h = ggml_add_inplace(ctx->ggml_ctx, h, skip);
|
||||
h = ggml_relu_inplace(ctx->ggml_ctx, h);
|
||||
return h;
|
||||
}
|
||||
};
|
||||
|
||||
ggml_tensor*
|
||||
patchify(ggml_context* ctx,
|
||||
ggml_tensor* x,
|
||||
int64_t patch_size,
|
||||
int64_t b = 1) {
|
||||
// x: [f, b*c, h*q, w*r]
|
||||
// return: [f, b*c*r*q, h, w]
|
||||
if (patch_size == 1) {
|
||||
@@ -322,13 +366,20 @@ class TinyVideoEncoder : public UnaryBlock {
|
||||
int patch_size = 1;
|
||||
|
||||
public:
|
||||
TinyVideoEncoder(int z_channels = 4, int patch_size = 1)
|
||||
int t_downscale = 1;
|
||||
TinyVideoEncoder(int z_channels = 4, int patch_size = 1, std::vector<bool> time_downscale = {true, true, false})
|
||||
: z_channels(z_channels), patch_size(patch_size) {
|
||||
t_downscale = 1;
|
||||
for (bool downscale : time_downscale) {
|
||||
if (downscale) {
|
||||
t_downscale *= 2;
|
||||
}
|
||||
}
|
||||
int index = 0;
|
||||
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(in_channels * patch_size * patch_size, hidden, {3, 3}, {1, 1}, {1, 1}));
|
||||
index++; // nn.ReLU()
|
||||
for (int i = 0; i < num_layers; i++) {
|
||||
int stride = i == num_layers - 1 ? 1 : 2;
|
||||
int stride = time_downscale[i] ? 2 : 1;
|
||||
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TPool(hidden, stride));
|
||||
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(hidden, hidden, {3, 3}, {2, 2}, {1, 1}, {1, 1}, false));
|
||||
for (int j = 0; j < num_blocks; j++) {
|
||||
@@ -375,17 +426,35 @@ class TinyVideoDecoder : public UnaryBlock {
|
||||
static const int num_layers = 3;
|
||||
int channels[num_layers + 1] = {256, 128, 64, 64};
|
||||
int patch_size = 1;
|
||||
int t_upscale = 1;
|
||||
bool is_wide = false;
|
||||
|
||||
public:
|
||||
TinyVideoDecoder(int z_channels = 4, int patch_size = 1)
|
||||
: z_channels(z_channels), patch_size(patch_size) {
|
||||
TinyVideoDecoder(int z_channels = 4, int patch_size = 1, std::vector<bool> time_upscale = {false, true, true}, bool is_wide = false)
|
||||
: z_channels(z_channels), patch_size(patch_size), is_wide(is_wide) {
|
||||
t_upscale = 1;
|
||||
if (is_wide) {
|
||||
channels[0] = 1024;
|
||||
channels[1] = 512;
|
||||
channels[2] = 256;
|
||||
}
|
||||
|
||||
for (bool upscale : time_upscale) {
|
||||
if (upscale) {
|
||||
t_upscale *= 2;
|
||||
}
|
||||
}
|
||||
int index = 1; // Clamp()
|
||||
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new Conv2d(z_channels, channels[0], {3, 3}, {1, 1}, {1, 1}));
|
||||
index++; // nn.ReLU()
|
||||
for (int i = 0; i < num_layers; i++) {
|
||||
int stride = i == 0 ? 1 : 2;
|
||||
int stride = time_upscale[i] ? 2 : 1;
|
||||
for (int j = 0; j < num_blocks; j++) {
|
||||
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new MemBlock(channels[i], channels[i]));
|
||||
if (is_wide) {
|
||||
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new WideMemBlock(channels[i], channels[i]));
|
||||
} else {
|
||||
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new MemBlock(channels[i], channels[i]));
|
||||
}
|
||||
}
|
||||
index++; // nn.Upsample()
|
||||
blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TGrow(channels[i], stride));
|
||||
@@ -410,10 +479,15 @@ public:
|
||||
int index = 3;
|
||||
for (int i = 0; i < num_layers; i++) {
|
||||
for (int j = 0; j < num_blocks; j++) {
|
||||
auto block = std::dynamic_pointer_cast<MemBlock>(blocks[std::to_string(index++)]);
|
||||
auto mem = ggml_pad_ext(ctx->ggml_ctx, h, 0, 0, 0, 0, 0, 0, 1, 0);
|
||||
mem = ggml_view_4d(ctx->ggml_ctx, mem, h->ne[0], h->ne[1], h->ne[2], h->ne[3], h->nb[1], h->nb[2], h->nb[3], 0);
|
||||
h = block->forward(ctx, h, mem);
|
||||
auto mem = ggml_pad_ext(ctx->ggml_ctx, h, 0, 0, 0, 0, 0, 0, 1, 0);
|
||||
mem = ggml_view_4d(ctx->ggml_ctx, mem, h->ne[0], h->ne[1], h->ne[2], h->ne[3], h->nb[1], h->nb[2], h->nb[3], 0);
|
||||
if (is_wide) {
|
||||
auto block = std::dynamic_pointer_cast<WideMemBlock>(blocks[std::to_string(index++)]);
|
||||
h = block->forward(ctx, h, mem);
|
||||
} else {
|
||||
auto block = std::dynamic_pointer_cast<MemBlock>(blocks[std::to_string(index++)]);
|
||||
h = block->forward(ctx, h, mem);
|
||||
}
|
||||
}
|
||||
// upsample
|
||||
index++;
|
||||
@@ -430,8 +504,8 @@ public:
|
||||
if (patch_size > 1) {
|
||||
h = unpatchify(ctx->ggml_ctx, h, patch_size, 1);
|
||||
}
|
||||
// shape(W, H, 3, 3 + T) => shape(W, H, 3, T)
|
||||
h = ggml_view_4d(ctx->ggml_ctx, h, h->ne[0], h->ne[1], h->ne[2], h->ne[3] - 3, h->nb[1], h->nb[2], h->nb[3], 3 * h->nb[3]);
|
||||
// shape(W, H, 3, (t_upscale - 1) + T) => shape(W, H, 3, T)
|
||||
h = ggml_view_4d(ctx->ggml_ctx, h, h->ne[0], h->ne[1], h->ne[2], h->ne[3] - (t_upscale - 1), h->nb[1], h->nb[2], h->nb[3], (t_upscale - 1) * h->nb[3]);
|
||||
return h;
|
||||
}
|
||||
};
|
||||
@@ -440,32 +514,40 @@ class TAEHV : public GGMLBlock {
|
||||
protected:
|
||||
bool decode_only;
|
||||
SDVersion version;
|
||||
bool is_wide;
|
||||
|
||||
public:
|
||||
int z_channels = 16;
|
||||
int z_channels = 16;
|
||||
std::vector<bool> time_downscale = {true, true, false};
|
||||
std::vector<bool> time_upscale = {false, true, true};
|
||||
|
||||
public:
|
||||
TAEHV(bool decode_only = true, SDVersion version = VERSION_WAN2)
|
||||
: decode_only(decode_only), version(version) {
|
||||
TAEHV(bool decode_only = true, SDVersion version = VERSION_WAN2, bool is_wide = false)
|
||||
: decode_only(decode_only), version(version), is_wide(is_wide) {
|
||||
int patch = 1;
|
||||
if (version == VERSION_WAN2_2_TI2V) {
|
||||
z_channels = 48;
|
||||
patch = 2;
|
||||
} else if (sd_version_is_ltxav(version)) {
|
||||
z_channels = 128;
|
||||
patch = 4;
|
||||
time_downscale = {true, true, true};
|
||||
time_upscale = {true, true, true};
|
||||
}
|
||||
blocks["decoder"] = std::shared_ptr<GGMLBlock>(new TinyVideoDecoder(z_channels, patch));
|
||||
blocks["decoder"] = std::shared_ptr<GGMLBlock>(new TinyVideoDecoder(z_channels, patch, time_upscale, is_wide));
|
||||
if (!decode_only) {
|
||||
blocks["encoder"] = std::shared_ptr<GGMLBlock>(new TinyVideoEncoder(z_channels, patch));
|
||||
blocks["encoder"] = std::shared_ptr<GGMLBlock>(new TinyVideoEncoder(z_channels, patch, time_downscale));
|
||||
}
|
||||
}
|
||||
|
||||
ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* z) {
|
||||
auto decoder = std::dynamic_pointer_cast<TinyVideoDecoder>(blocks["decoder"]);
|
||||
if (sd_version_is_wan(version)) {
|
||||
if (sd_version_is_wan(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)) {
|
||||
if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) {
|
||||
// (W, H, C, T) -> (W, H, T, C)
|
||||
result = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, result, 0, 1, 3, 2));
|
||||
}
|
||||
@@ -477,10 +559,10 @@ public:
|
||||
// (W, H, T, C) -> (W, H, C, T)
|
||||
x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
|
||||
int64_t num_frames = x->ne[3];
|
||||
if (num_frames % 4) {
|
||||
// pad to multiple of 4 at the end
|
||||
if (num_frames % encoder->t_downscale) {
|
||||
// pad to multiple of encoder->t_downscale at the end
|
||||
auto last_frame = ggml_view_4d(ctx->ggml_ctx, x, x->ne[0], x->ne[1], x->ne[2], 1, x->nb[1], x->nb[2], x->nb[3], (num_frames - 1) * x->nb[3]);
|
||||
for (int i = 0; i < 4 - num_frames % 4; i++) {
|
||||
for (int i = 0; i < encoder->t_downscale - num_frames % encoder->t_downscale; i++) {
|
||||
x = ggml_concat(ctx->ggml_ctx, x, last_frame, 3);
|
||||
}
|
||||
}
|
||||
@@ -601,6 +683,7 @@ struct TinyImageAutoEncoder : public VAE {
|
||||
struct TinyVideoAutoEncoder : public VAE {
|
||||
TAEHV taehv;
|
||||
bool decode_only = false;
|
||||
bool is_wide = false;
|
||||
|
||||
TinyVideoAutoEncoder(ggml_backend_t backend,
|
||||
ggml_backend_t params_backend,
|
||||
@@ -609,8 +692,14 @@ struct TinyVideoAutoEncoder : public VAE {
|
||||
bool decoder_only = true,
|
||||
SDVersion version = VERSION_WAN2)
|
||||
: decode_only(decoder_only),
|
||||
taehv(decoder_only, version),
|
||||
VAE(version, backend, params_backend) {
|
||||
for (auto tensor_storage : tensor_storage_map) {
|
||||
if (tensor_storage.first.find(prefix + ".3.conv.6.weight") != std::string::npos) {
|
||||
is_wide = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
taehv = TAEHV(decoder_only, version, is_wide);
|
||||
scale_input = false;
|
||||
taehv.init(params_ctx, tensor_storage_map, prefix);
|
||||
}
|
||||
@@ -641,7 +730,8 @@ struct TinyVideoAutoEncoder : public VAE {
|
||||
}
|
||||
|
||||
ggml_cgraph* build_graph(const sd::Tensor<float>& z_tensor, bool decode_graph) {
|
||||
ggml_cgraph* gf = ggml_new_graph(compute_ctx);
|
||||
ggml_cgraph* gf = decode_graph && is_wide ? ggml_new_graph_custom(compute_ctx, 4096, false)
|
||||
: ggml_new_graph(compute_ctx);
|
||||
ggml_tensor* z = make_input(z_tensor);
|
||||
auto runner_ctx = get_context();
|
||||
ggml_tensor* out = decode_graph ? taehv.decode(&runner_ctx, z) : taehv.encode(&runner_ctx, z);
|
||||
|
||||
@@ -0,0 +1,91 @@
|
||||
#include "gpt_oss_tokenizer.h"
|
||||
|
||||
#include "json.hpp"
|
||||
#include "util.h"
|
||||
#include "vocab/vocab.h"
|
||||
|
||||
void GPTOSSTokenizer::load_from_merges(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
|
||||
auto byte_unicode_pairs = bytes_to_unicode();
|
||||
byte_encoder = std::map<int, std::u32string>(byte_unicode_pairs.begin(), byte_unicode_pairs.end());
|
||||
for (auto& pair : byte_unicode_pairs) {
|
||||
byte_decoder[pair.second] = pair.first;
|
||||
}
|
||||
|
||||
nlohmann::json vocab;
|
||||
try {
|
||||
vocab = nlohmann::json::parse(vocab_utf8_str);
|
||||
} catch (const nlohmann::json::parse_error&) {
|
||||
GGML_ABORT("invalid vocab json str");
|
||||
}
|
||||
for (const auto& [key, value] : vocab.items()) {
|
||||
std::u32string token = utf8_to_utf32(key);
|
||||
int i = value;
|
||||
encoder[token] = i;
|
||||
decoder[i] = token;
|
||||
}
|
||||
encoder_len = static_cast<int>(encoder.size());
|
||||
for (auto& special_token : special_tokens) {
|
||||
auto token = utf8_to_utf32(special_token);
|
||||
encoder[token] = encoder_len;
|
||||
decoder[encoder_len] = token;
|
||||
encoder_len++;
|
||||
}
|
||||
encoder_len = static_cast<int>(encoder.size());
|
||||
LOG_DEBUG("vocab size: %d", encoder_len);
|
||||
|
||||
std::vector<std::u32string> merges = split_utf32(merges_utf8_str);
|
||||
std::vector<std::pair<std::u32string, std::u32string>> merge_pairs;
|
||||
for (const auto& merge : merges) {
|
||||
size_t space_pos = merge.find(' ');
|
||||
merge_pairs.emplace_back(merge.substr(0, space_pos), merge.substr(space_pos + 1));
|
||||
}
|
||||
LOG_DEBUG("merges size %zu", merge_pairs.size());
|
||||
|
||||
int rank = 0;
|
||||
for (const auto& merge : merge_pairs) {
|
||||
bpe_ranks[merge] = rank++;
|
||||
}
|
||||
bpe_len = rank;
|
||||
}
|
||||
|
||||
GPTOSSTokenizer::GPTOSSTokenizer(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) {
|
||||
BOS_TOKEN = "<|startoftext|>";
|
||||
UNK_TOKEN = "<|endoftext|>";
|
||||
EOS_TOKEN = "<|endoftext|>";
|
||||
PAD_TOKEN = "<|endoftext|>";
|
||||
|
||||
BOS_TOKEN_ID = 199998;
|
||||
EOS_TOKEN_ID = 199999;
|
||||
UNK_TOKEN_ID = 199999;
|
||||
PAD_TOKEN_ID = 199999;
|
||||
|
||||
special_tokens = {
|
||||
"<|startoftext|>",
|
||||
"<|endoftext|>",
|
||||
"<|reserved_200000|>",
|
||||
"<|reserved_200001|>",
|
||||
"<|return|>",
|
||||
"<|constrain|>",
|
||||
"<|reserved_200004|>",
|
||||
"<|channel|>",
|
||||
"<|start|>",
|
||||
"<|end|>",
|
||||
"<|message|>",
|
||||
"<|reserved_200009|>",
|
||||
"<|reserved_200010|>",
|
||||
"<|reserved_200011|>",
|
||||
"<|call|>",
|
||||
"<|reserved_200013|>",
|
||||
"<|reserved_200014|>",
|
||||
"<|reserved_200015|>",
|
||||
"<|reserved_200016|>",
|
||||
"<|reserved_200017|>",
|
||||
"<|endofprompt|>",
|
||||
};
|
||||
|
||||
if (merges_utf8_str.size() > 0) {
|
||||
load_from_merges(merges_utf8_str, vocab_utf8_str);
|
||||
} else {
|
||||
load_from_merges(load_gpt_oss_merges(), load_gpt_oss_vocab_json());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
#ifndef __SD_TOKENIZERS_GPT_OSS_TOKENIZER_H__
|
||||
#define __SD_TOKENIZERS_GPT_OSS_TOKENIZER_H__
|
||||
|
||||
#include <string>
|
||||
|
||||
#include "bpe_tokenizer.h"
|
||||
|
||||
class GPTOSSTokenizer : public BPETokenizer {
|
||||
protected:
|
||||
void load_from_merges(const std::string& merges_utf8_str, const std::string& vocab_utf8_str);
|
||||
|
||||
public:
|
||||
explicit GPTOSSTokenizer(const std::string& merges_utf8_str = "", const std::string& vocab_utf8_str = "");
|
||||
};
|
||||
|
||||
#endif // __SD_TOKENIZERS_GPT_OSS_TOKENIZER_H__
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because it is too large
Load Diff
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because it is too large
Load Diff
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
+3
-762304
File diff suppressed because one or more lines are too long
@@ -1,9 +1,13 @@
|
||||
#include "vocab.h"
|
||||
#include "clip_t5.hpp"
|
||||
#include "clip_merges.hpp"
|
||||
#include "gemma_merges.hpp"
|
||||
#include "gemma_vocab.hpp"
|
||||
#include "mistral.hpp"
|
||||
#include "qwen.hpp"
|
||||
#include "gpt_oss_merges.hpp"
|
||||
#include "gpt_oss_vocab.hpp"
|
||||
#include "mistral_merges.hpp"
|
||||
#include "mistral_vocab.hpp"
|
||||
#include "qwen_merges.hpp"
|
||||
#include "t5.hpp"
|
||||
#include "umt5.hpp"
|
||||
|
||||
std::string load_clip_merges() {
|
||||
@@ -44,4 +48,14 @@ std::string load_gemma_merges() {
|
||||
std::string load_gemma_vocab_json() {
|
||||
std::string json_str(reinterpret_cast<const char*>(gemma_vocab_json_utf8_c_str), sizeof(gemma_vocab_json_utf8_c_str));
|
||||
return json_str;
|
||||
}
|
||||
|
||||
std::string load_gpt_oss_merges() {
|
||||
std::string merges_utf8_str(reinterpret_cast<const char*>(gpt_oss_merges_utf8_c_str), sizeof(gpt_oss_merges_utf8_c_str));
|
||||
return merges_utf8_str;
|
||||
}
|
||||
|
||||
std::string load_gpt_oss_vocab_json() {
|
||||
std::string json_str(reinterpret_cast<const char*>(gpt_oss_vocab_json_utf8_c_str), sizeof(gpt_oss_vocab_json_utf8_c_str));
|
||||
return json_str;
|
||||
}
|
||||
@@ -11,5 +11,7 @@ std::string load_t5_tokenizer_json();
|
||||
std::string load_umt5_tokenizer_json();
|
||||
std::string load_gemma_merges();
|
||||
std::string load_gemma_vocab_json();
|
||||
std::string load_gpt_oss_merges();
|
||||
std::string load_gpt_oss_vocab_json();
|
||||
|
||||
#endif // __SD_TOKENIZERS_VOCAB_VOCAB_H__
|
||||
+217
@@ -1,8 +1,10 @@
|
||||
#include "util.h"
|
||||
#include <algorithm>
|
||||
#include <cctype>
|
||||
#include <cmath>
|
||||
#include <codecvt>
|
||||
#include <cstdarg>
|
||||
#include <exception>
|
||||
#include <fstream>
|
||||
#include <locale>
|
||||
#include <regex>
|
||||
@@ -406,6 +408,88 @@ std::vector<std::string> split_string(const std::string& str, char delimiter) {
|
||||
return result;
|
||||
}
|
||||
|
||||
KeyValueArgs parse_key_value_args(const char* args, const char* context) {
|
||||
KeyValueArgs pairs;
|
||||
|
||||
if (args == nullptr || args[0] == '\0') {
|
||||
return pairs;
|
||||
}
|
||||
|
||||
std::string raw(args);
|
||||
size_t start = 0;
|
||||
for (size_t pos = 0; pos <= raw.size(); ++pos) {
|
||||
if (pos != raw.size() && raw[pos] != ',' && raw[pos] != ';') {
|
||||
continue;
|
||||
}
|
||||
|
||||
std::string token = trim(raw.substr(start, pos - start));
|
||||
if (!token.empty()) {
|
||||
size_t eq = token.find('=');
|
||||
if (eq == std::string::npos) {
|
||||
const char* log_context = context ? context : "key=value arg";
|
||||
LOG_WARN("ignoring malformed %s '%s'", log_context, token.c_str());
|
||||
} else {
|
||||
std::string key = trim(token.substr(0, eq));
|
||||
std::string value = trim(token.substr(eq + 1));
|
||||
pairs.emplace_back(std::move(key), std::move(value));
|
||||
}
|
||||
}
|
||||
|
||||
start = pos + 1;
|
||||
}
|
||||
|
||||
return pairs;
|
||||
}
|
||||
|
||||
KeyValueArgs parse_key_value_args(const std::string& args, const char* context) {
|
||||
return parse_key_value_args(args.c_str(), context);
|
||||
}
|
||||
|
||||
bool parse_strict_float(const std::string& text, float& value) {
|
||||
try {
|
||||
size_t consumed = 0;
|
||||
float parsed = std::stof(text, &consumed);
|
||||
if (!trim(text.substr(consumed)).empty()) {
|
||||
return false;
|
||||
}
|
||||
value = parsed;
|
||||
return true;
|
||||
} catch (const std::exception&) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
bool parse_strict_int(const std::string& text, int& value) {
|
||||
try {
|
||||
size_t consumed = 0;
|
||||
int parsed = std::stoi(text, &consumed);
|
||||
if (!trim(text.substr(consumed)).empty()) {
|
||||
return false;
|
||||
}
|
||||
value = parsed;
|
||||
return true;
|
||||
} catch (const std::exception&) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
bool parse_strict_bool(const std::string& text, bool& value) {
|
||||
std::string lowered = trim(text);
|
||||
std::transform(lowered.begin(), lowered.end(), lowered.begin(), [](unsigned char c) {
|
||||
return static_cast<char>(std::tolower(c));
|
||||
});
|
||||
|
||||
if (lowered == "1" || lowered == "true" || lowered == "yes" || lowered == "on") {
|
||||
value = true;
|
||||
return true;
|
||||
}
|
||||
if (lowered == "0" || lowered == "false" || lowered == "no" || lowered == "off") {
|
||||
value = false;
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
static std::string build_progress_bar(int step, int steps) {
|
||||
std::string progress = " |";
|
||||
int max_progress = 50;
|
||||
@@ -758,6 +842,139 @@ std::vector<std::pair<std::string, float>> parse_prompt_attention(const std::str
|
||||
return res;
|
||||
}
|
||||
|
||||
static size_t get_utf8_char_len(char c) {
|
||||
unsigned char uc = static_cast<unsigned char>(c);
|
||||
if ((uc & 0x80) == 0) {
|
||||
return 1;
|
||||
}
|
||||
if ((uc & 0xE0) == 0xC0) {
|
||||
return 2;
|
||||
}
|
||||
if ((uc & 0xF0) == 0xE0) {
|
||||
return 3;
|
||||
}
|
||||
if ((uc & 0xF8) == 0xF0) {
|
||||
return 4;
|
||||
}
|
||||
return 1;
|
||||
}
|
||||
|
||||
static bool is_ascii_alpha(char c) {
|
||||
return (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z');
|
||||
}
|
||||
|
||||
static bool starts_with_at(const std::string& text, size_t pos, const std::string& needle) {
|
||||
return pos + needle.size() <= text.size() && text.compare(pos, needle.size(), needle) == 0;
|
||||
}
|
||||
|
||||
static bool is_word_internal_apostrophe(const std::string& text, size_t pos) {
|
||||
return pos > 0 && pos + 1 < text.size() &&
|
||||
is_ascii_alpha(text[pos - 1]) && is_ascii_alpha(text[pos + 1]);
|
||||
}
|
||||
|
||||
static std::vector<std::pair<std::string, bool>> split_quotation(const std::string& text) {
|
||||
static const std::vector<std::pair<std::string, std::string>> quote_pairs = {
|
||||
{"'", "'"},
|
||||
{"\"", "\""},
|
||||
{"\xE2\x80\x98", "\xE2\x80\x99"},
|
||||
{"\xE2\x80\x9C", "\xE2\x80\x9D"},
|
||||
};
|
||||
|
||||
std::vector<std::pair<std::string, bool>> result;
|
||||
size_t segment_start = 0;
|
||||
size_t i = 0;
|
||||
|
||||
auto push_segment = [&](size_t begin, size_t end, bool matched) {
|
||||
if (end > begin) {
|
||||
result.emplace_back(text.substr(begin, end - begin), matched);
|
||||
}
|
||||
};
|
||||
|
||||
while (i < text.size()) {
|
||||
bool matched_quote = false;
|
||||
for (const auto& quote_pair : quote_pairs) {
|
||||
const std::string& open_quote = quote_pair.first;
|
||||
const std::string& close_quote = quote_pair.second;
|
||||
if (!starts_with_at(text, i, open_quote)) {
|
||||
continue;
|
||||
}
|
||||
if (open_quote == "'" && is_word_internal_apostrophe(text, i)) {
|
||||
continue;
|
||||
}
|
||||
|
||||
size_t search_pos = i + open_quote.size();
|
||||
size_t close_pos = std::string::npos;
|
||||
bool invalid = false;
|
||||
while (search_pos < text.size()) {
|
||||
if (open_quote != close_quote && starts_with_at(text, search_pos, open_quote)) {
|
||||
invalid = true;
|
||||
break;
|
||||
}
|
||||
if (starts_with_at(text, search_pos, close_quote)) {
|
||||
if (close_quote == "'" && is_word_internal_apostrophe(text, search_pos)) {
|
||||
search_pos += close_quote.size();
|
||||
continue;
|
||||
}
|
||||
close_pos = search_pos;
|
||||
break;
|
||||
}
|
||||
|
||||
size_t char_len = get_utf8_char_len(text[search_pos]);
|
||||
if (search_pos + char_len > text.size()) {
|
||||
char_len = 1;
|
||||
}
|
||||
search_pos += char_len;
|
||||
}
|
||||
if (invalid || close_pos == std::string::npos) {
|
||||
continue;
|
||||
}
|
||||
|
||||
size_t quote_start = i;
|
||||
push_segment(segment_start, quote_start, false);
|
||||
i = close_pos + close_quote.size();
|
||||
push_segment(quote_start, i, true);
|
||||
segment_start = i;
|
||||
matched_quote = true;
|
||||
break;
|
||||
}
|
||||
if (!matched_quote) {
|
||||
size_t char_len = get_utf8_char_len(text[i]);
|
||||
if (i + char_len > text.size()) {
|
||||
char_len = 1;
|
||||
}
|
||||
i += char_len;
|
||||
}
|
||||
}
|
||||
|
||||
push_segment(segment_start, text.size(), false);
|
||||
return result;
|
||||
}
|
||||
|
||||
std::vector<std::pair<std::string, float>> split_quotation_attention(
|
||||
const std::vector<std::pair<std::string, float>>& parsed_attention) {
|
||||
std::vector<std::pair<std::string, float>> result;
|
||||
for (const auto& item : parsed_attention) {
|
||||
const std::string& text = item.first;
|
||||
float weight = item.second;
|
||||
for (const auto& part : split_quotation(text)) {
|
||||
if (part.second) {
|
||||
size_t i = 0;
|
||||
while (i < part.first.size()) {
|
||||
size_t char_len = get_utf8_char_len(part.first[i]);
|
||||
if (i + char_len > part.first.size()) {
|
||||
char_len = 1;
|
||||
}
|
||||
result.emplace_back(part.first.substr(i, char_len), weight);
|
||||
i += char_len;
|
||||
}
|
||||
} else {
|
||||
result.emplace_back(part.first, weight);
|
||||
}
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
// namespace is needed to avoid conflicts with ggml_backend_extend.hpp
|
||||
namespace ggml_cpu {
|
||||
#include "ggml-cpu.h"
|
||||
|
||||
+12
@@ -4,6 +4,7 @@
|
||||
#include <cstdint>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "ggml-backend.h"
|
||||
@@ -65,6 +66,15 @@ protected:
|
||||
|
||||
std::string path_join(const std::string& p1, const std::string& p2);
|
||||
std::vector<std::string> split_string(const std::string& str, char delimiter);
|
||||
|
||||
using KeyValueArgs = std::vector<std::pair<std::string, std::string>>;
|
||||
|
||||
KeyValueArgs parse_key_value_args(const char* args, const char* context = "key=value arg");
|
||||
KeyValueArgs parse_key_value_args(const std::string& args, const char* context = "key=value arg");
|
||||
bool parse_strict_float(const std::string& text, float& value);
|
||||
bool parse_strict_int(const std::string& text, int& value);
|
||||
bool parse_strict_bool(const std::string& text, bool& value);
|
||||
|
||||
void pretty_progress(int step, int steps, float time);
|
||||
void pretty_bytes_progress(int step, int steps, uint64_t bytes_processed, float elapsed_seconds);
|
||||
|
||||
@@ -73,6 +83,8 @@ void log_printf(sd_log_level_t level, const char* file, int line, const char* fo
|
||||
std::string trim(const std::string& s);
|
||||
|
||||
std::vector<std::pair<std::string, float>> parse_prompt_attention(const std::string& text);
|
||||
std::vector<std::pair<std::string, float>> split_quotation_attention(
|
||||
const std::vector<std::pair<std::string, float>>& parsed_attention);
|
||||
|
||||
sd_progress_cb_t sd_get_progress_callback();
|
||||
void* sd_get_progress_callback_data();
|
||||
|
||||
@@ -167,6 +167,7 @@ public:
|
||||
int64_t t0 = ggml_time_ms();
|
||||
sd::Tensor<float> input = x;
|
||||
sd::Tensor<float> output;
|
||||
set_tiling_params(tiling_params);
|
||||
|
||||
if (tiling_params.enabled) {
|
||||
const int scale_factor = get_scale_factor();
|
||||
@@ -216,6 +217,9 @@ public:
|
||||
virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) = 0;
|
||||
virtual void set_conv2d_scale(float scale) { SD_UNUSED(scale); };
|
||||
virtual void set_temporal_tiling_enabled(bool enabled) { SD_UNUSED(enabled); };
|
||||
virtual void set_tiling_params(const sd_tiling_params_t& params) {
|
||||
set_temporal_tiling_enabled(params.temporal_tiling);
|
||||
};
|
||||
};
|
||||
|
||||
struct FakeVAE : public VAE {
|
||||
|
||||
Reference in New Issue
Block a user