diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index df4c07d4..393559ca 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -83,7 +83,7 @@ jobs: run: | mkdir build cd build - cmake .. -DGGML_AVX2=ON -DSD_BUILD_SHARED_LIBS=ON + cmake .. -DSD_BUILD_SHARED_LIBS=ON -DGGML_NATIVE=OFF -DSD_BUILD_SHARED_GGML_LIB=ON -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON -DCMAKE_INSTALL_RPATH='$ORIGIN' cmake --build . --config Release - name: Get commit hash @@ -146,7 +146,7 @@ jobs: run: | mkdir build cd build - cmake .. -DSD_BUILD_SHARED_LIBS=ON -DSD_VULKAN=ON + cmake .. -DSD_BUILD_SHARED_LIBS=ON -DSD_VULKAN=ON -DGGML_NATIVE=OFF -DSD_BUILD_SHARED_GGML_LIB=ON -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON -DCMAKE_INSTALL_RPATH='$ORIGIN' cmake --build . --config Release - name: Get commit hash @@ -207,6 +207,7 @@ jobs: UBUNTU_VERSION=24.04 CUDA_ARCHITECTURES=121 GGML_CUDA_FA_ALL_QUANTS=ON + GGML_CUDA_ENABLE_DYNAMIC_CPU_BACKENDS=OFF env: REGISTRY: ghcr.io @@ -341,18 +342,12 @@ jobs: strategy: matrix: include: - - build: "noavx" - defines: "-DGGML_NATIVE=OFF -DGGML_AVX=OFF -DGGML_AVX2=OFF -DGGML_FMA=OFF -DSD_BUILD_SHARED_LIBS=ON" - - build: "avx2" - defines: "-DGGML_NATIVE=OFF -DGGML_AVX2=ON -DSD_BUILD_SHARED_LIBS=ON" - - build: "avx" - defines: "-DGGML_NATIVE=OFF -DGGML_AVX=ON -DGGML_AVX2=OFF -DSD_BUILD_SHARED_LIBS=ON" - - build: "avx512" - defines: "-DGGML_NATIVE=OFF -DGGML_AVX512=ON -DGGML_AVX=ON -DGGML_AVX2=ON -DSD_BUILD_SHARED_LIBS=ON" + - build: "cpu" + defines: "-DGGML_NATIVE=OFF -DSD_BUILD_SHARED_LIBS=ON -DSD_BUILD_SHARED_GGML_LIB=ON -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON" - build: "cuda12" - defines: "-DSD_CUDA=ON -DSD_BUILD_SHARED_LIBS=ON -DCMAKE_CUDA_ARCHITECTURES='61;70;75;80;86;89;90;100;120' -DCMAKE_CUDA_FLAGS='-Xcudafe \"--diag_suppress=177\" -Xcudafe \"--diag_suppress=550\"'" + defines: "-DSD_CUDA=ON -DSD_BUILD_SHARED_LIBS=ON -DCMAKE_CUDA_ARCHITECTURES='61;70;75;80;86;89;90;100;120' -DCMAKE_CUDA_FLAGS='-Xcudafe \"--diag_suppress=177\" -Xcudafe \"--diag_suppress=550\"' -DGGML_NATIVE=OFF -DSD_BUILD_SHARED_GGML_LIB=ON -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON" - build: "vulkan" - defines: "-DSD_VULKAN=ON -DSD_BUILD_SHARED_LIBS=ON" + defines: "-DSD_VULKAN=ON -DSD_BUILD_SHARED_LIBS=ON -DGGML_NATIVE=OFF -DSD_BUILD_SHARED_GGML_LIB=ON -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON" steps: - name: Clone id: checkout @@ -400,19 +395,6 @@ jobs: cmake .. -DCMAKE_CXX_FLAGS='/bigobj' -G Ninja -DCMAKE_C_COMPILER=cl.exe -DCMAKE_CXX_COMPILER=cl.exe -DCMAKE_BUILD_TYPE=Release ${{ matrix.defines }} cmake --build . - - name: Check AVX512F support - id: check_avx512f - if: ${{ matrix.build == 'avx512' }} - continue-on-error: true - run: | - cd build - $vcdir = $(vswhere -latest -products * -requires Microsoft.VisualStudio.Component.VC.Tools.x86.x64 -property installationPath) - $msvc = $(join-path $vcdir $('VC\Tools\MSVC\'+$(gc -raw $(join-path $vcdir 'VC\Auxiliary\Build\Microsoft.VCToolsVersion.default.txt')).Trim())) - $cl = $(join-path $msvc 'bin\Hostx64\x64\cl.exe') - echo 'int main(void){unsigned int a[4];__cpuid(a,7);return !(a[1]&65536);}' >> avx512f.c - & $cl /O2 /GS- /kernel avx512f.c /link /nodefaultlib /entry:main - .\avx512f.exe && echo "AVX512F: YES" && ( echo HAS_AVX512F=1 >> $env:GITHUB_ENV ) || echo "AVX512F: NO" - - name: Get commit hash id: commit if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }} @@ -722,6 +704,25 @@ jobs: run: | sudo apt install -y build-essential cmake wget zip ninja-build + - name: Free disk space + run: | + df -h + + # Remove preinstalled SDKs and caches not needed for this job before + # installing ROCm. The legacy ROCm apt packages are large enough to + # exhaust ubuntu-latest if cleanup runs after installation. + sudo rm -rf /usr/share/dotnet || true + sudo rm -rf /usr/local/lib/android || true + sudo rm -rf /opt/ghc || true + sudo rm -rf /usr/local/.ghcup || true + sudo rm -rf /opt/hostedtoolcache || true + sudo rm -rf /usr/share/swift || true + sudo rm -rf /usr/local/share/boost || true + docker system prune -af || true + + sudo apt clean + df -h + - name: Setup Legacy ROCm if: matrix.ROCM_VERSION == '7.2.1' id: legacy_env @@ -743,19 +744,6 @@ jobs: sudo apt update sudo apt-get install -y libssl-dev rocm-hip-sdk - - name: Free disk space - run: | - # Remove preinstalled SDKs and caches not needed for this job - sudo rm -rf /usr/share/dotnet || true - sudo rm -rf /usr/local/lib/android || true - sudo rm -rf /opt/ghc || true - sudo rm -rf /usr/local/.ghcup || true - sudo rm -rf /opt/hostedtoolcache || true - - # Remove old package lists and caches - sudo rm -rf /var/lib/apt/lists/* || true - sudo apt clean - - name: Setup TheRock if: matrix.ROCM_VERSION != '7.2.1' id: therock_env @@ -794,6 +782,11 @@ jobs: -DGPU_TARGETS="${{ matrix.gpu_targets }}" \ -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \ -DCMAKE_POSITION_INDEPENDENT_CODE=ON \ + -DCMAKE_INSTALL_RPATH='$ORIGIN' \ + -DGGML_NATIVE=OFF \ + -DSD_BUILD_SHARED_GGML_LIB=ON \ + -DGGML_BACKEND_DL=ON \ + -DGGML_CPU_ALL_VARIANTS=ON \ -DSD_BUILD_SHARED_LIBS=ON cmake --build . --config Release diff --git a/CMakeLists.txt b/CMakeLists.txt index 2804bad4..f15fdbf2 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -204,12 +204,33 @@ if(SD_WEBM) endif() endif() +if (SD_RPC) + message("-- Use RPC as backend stable-diffusion") + set(GGML_RPC ON) + add_definitions(-DSD_USE_RPC) +endif () + set(SD_LIB stable-diffusion) file(GLOB SD_LIB_SOURCES CONFIGURE_DEPENDS "src/*.h" "src/*.cpp" "src/*.hpp" + "src/conditioning/*.h" + "src/conditioning/*.cpp" + "src/conditioning/*.hpp" + "src/core/*.h" + "src/core/*.cpp" + "src/core/*.hpp" + "src/extensions/*.h" + "src/extensions/*.cpp" + "src/extensions/*.hpp" + "src/model/*/*.h" + "src/model/*/*.cpp" + "src/model/*/*.hpp" + "src/runtime/*.h" + "src/runtime/*.cpp" + "src/runtime/*.hpp" "src/model_io/*.h" "src/model_io/*.cpp" "src/tokenizers/*.h" @@ -310,8 +331,10 @@ endif() add_subdirectory(thirdparty) -target_link_libraries(${SD_LIB} PUBLIC ggml zip) +target_sources(${SD_LIB} PRIVATE $) +target_link_libraries(${SD_LIB} PUBLIC ggml) target_include_directories(${SD_LIB} PUBLIC . src include) +target_include_directories(${SD_LIB} PRIVATE src/core) target_include_directories(${SD_LIB} PUBLIC . thirdparty) target_compile_features(${SD_LIB} PUBLIC c_std_11 cxx_std_17) @@ -320,7 +343,58 @@ if (SD_BUILD_EXAMPLES) add_subdirectory(examples) endif() -set(SD_PUBLIC_HEADERS include/stable-diffusion.h) -set_target_properties(${SD_LIB} PROPERTIES PUBLIC_HEADER "${SD_PUBLIC_HEADERS}") -install(TARGETS ${SD_LIB} LIBRARY PUBLIC_HEADER) + +# +# install +# + +include(CMakePackageConfigHelpers) +include(GNUInstallDirs) + +set(SD_INSTALL_VERSION "${SDCPP_BUILD_VERSION}") +set(SD_INSTALL_COMMIT "${SDCPP_BUILD_COMMIT}") +set(SD_SHARED_LIB ${SD_BUILD_SHARED_LIBS}) + +set(SD_INCLUDE_INSTALL_DIR ${CMAKE_INSTALL_INCLUDEDIR} CACHE PATH "Location of header files") +set(SD_LIB_INSTALL_DIR ${CMAKE_INSTALL_LIBDIR} CACHE PATH "Location of library files") +set(SD_BIN_INSTALL_DIR ${CMAKE_INSTALL_BINDIR} CACHE PATH "Location of binary files") + +set(SD_PUBLIC_HEADERS + ${CMAKE_CURRENT_SOURCE_DIR}/include/stable-diffusion.h) + +set_target_properties(${SD_LIB} + PROPERTIES + PUBLIC_HEADER "${SD_PUBLIC_HEADERS}") + + +install(TARGETS ${SD_LIB} + ARCHIVE + LIBRARY + RUNTIME + PUBLIC_HEADER) + + +configure_package_config_file( + ${CMAKE_CURRENT_SOURCE_DIR}/cmake/stable-diffusion-config.cmake.in + ${CMAKE_CURRENT_BINARY_DIR}/stable-diffusion-config.cmake + INSTALL_DESTINATION ${CMAKE_INSTALL_LIBDIR}/cmake/stable-diffusion + PATH_VARS SD_INCLUDE_INSTALL_DIR + SD_LIB_INSTALL_DIR + SD_BIN_INSTALL_DIR ) + +write_basic_package_version_file( + ${CMAKE_CURRENT_BINARY_DIR}/stable-diffusion-version.cmake + VERSION ${SD_INSTALL_VERSION} + COMPATIBILITY SameMajorVersion) + +install(FILES ${CMAKE_CURRENT_BINARY_DIR}/stable-diffusion-config.cmake + ${CMAKE_CURRENT_BINARY_DIR}/stable-diffusion-version.cmake + DESTINATION ${CMAKE_INSTALL_LIBDIR}/cmake/stable-diffusion) + +configure_file(cmake/stable-diffusion.pc.in + "${CMAKE_CURRENT_BINARY_DIR}/stable-diffusion.pc" + @ONLY) + +install(FILES "${CMAKE_CURRENT_BINARY_DIR}/stable-diffusion.pc" + DESTINATION ${CMAKE_INSTALL_LIBDIR}/pkgconfig) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 9ba9177a..f94e3904 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -44,6 +44,8 @@ Naming conventions: Some older code in the project may not fully follow the current conventions. Please do not submit PRs that only rewrite existing code to match style rules. +When adding or modifying model implementations, follow the model config and weight detection conventions in [docs/model_config.md](docs/model_config.md). + ## AI-Assisted Contributions AI tools may be used to assist development, but contributors are responsible for the quality and correctness of the submitted code. diff --git a/Dockerfile b/Dockerfile index ba27d66b..0ef3d56a 100644 --- a/Dockerfile +++ b/Dockerfile @@ -19,7 +19,14 @@ WORKDIR /sd.cpp COPY . . -RUN cmake . -B ./build +RUN cmake . -B ./build \ + -DSD_BUILD_SHARED_LIBS=ON \ + -DGGML_NATIVE=OFF \ + -DSD_BUILD_SHARED_GGML_LIB=ON \ + -DGGML_BACKEND_DL=ON \ + -DGGML_CPU_ALL_VARIANTS=ON \ + -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \ + -DCMAKE_INSTALL_RPATH='$ORIGIN' RUN cmake --build ./build --config Release --parallel FROM ubuntu:$UBUNTU_VERSION AS runtime @@ -28,7 +35,9 @@ RUN apt-get update && \ apt-get install --yes --no-install-recommends libgomp1 && \ apt-get clean -COPY --from=build /sd.cpp/build/bin/sd-cli /sd-cli -COPY --from=build /sd.cpp/build/bin/sd-server /sd-server +COPY --from=build /sd.cpp/build/bin /sd.cpp/bin +RUN printf '#!/bin/sh\nexec /sd.cpp/bin/sd-cli "$@"\n' > /sd-cli && \ + printf '#!/bin/sh\nexec /sd.cpp/bin/sd-server "$@"\n' > /sd-server && \ + chmod +x /sd-cli /sd-server ENTRYPOINT [ "/sd-cli" ] diff --git a/Dockerfile.cuda b/Dockerfile.cuda index 42b8f4fa..3e51bda6 100644 --- a/Dockerfile.cuda +++ b/Dockerfile.cuda @@ -23,11 +23,27 @@ COPY . . ARG CUDACXX=/usr/local/cuda/bin/nvcc ARG CUDA_ARCHITECTURES="" ARG GGML_CUDA_FA_ALL_QUANTS="" +ARG GGML_CUDA_ENABLE_DYNAMIC_CPU_BACKENDS=ON -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 set -- \ + -DSD_CUDA=ON; \ + if [ "${GGML_CUDA_ENABLE_DYNAMIC_CPU_BACKENDS}" = "ON" ]; then \ + set -- "$@" \ + -DSD_BUILD_SHARED_LIBS=ON \ + -DGGML_NATIVE=OFF \ + -DSD_BUILD_SHARED_GGML_LIB=ON \ + -DGGML_BACKEND_DL=ON \ + -DGGML_CPU_ALL_VARIANTS=ON \ + -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \ + "-DCMAKE_INSTALL_RPATH=\$ORIGIN"; \ + fi; \ + if [ -n "${CUDA_ARCHITECTURES}" ]; then \ + set -- "$@" "-DCMAKE_CUDA_ARCHITECTURES=${CUDA_ARCHITECTURES}"; \ + fi; \ + if [ -n "${GGML_CUDA_FA_ALL_QUANTS}" ]; then \ + set -- "$@" "-DGGML_CUDA_FA_ALL_QUANTS=${GGML_CUDA_FA_ALL_QUANTS}"; \ + fi; \ + cmake . -B ./build "$@" RUN cmake --build ./build --config Release -j$(nproc) FROM nvidia/cuda:${CUDA_VERSION}-cudnn-runtime-ubuntu${UBUNTU_VERSION} AS runtime @@ -36,7 +52,9 @@ RUN apt-get update && \ apt-get install --yes --no-install-recommends libgomp1 && \ apt-get clean -COPY --from=build /sd.cpp/build/bin/sd-cli /sd-cli -COPY --from=build /sd.cpp/build/bin/sd-server /sd-server +COPY --from=build /sd.cpp/build/bin /sd.cpp/bin +RUN printf '#!/bin/sh\nexec /sd.cpp/bin/sd-cli "$@"\n' > /sd-cli && \ + printf '#!/bin/sh\nexec /sd.cpp/bin/sd-server "$@"\n' > /sd-server && \ + chmod +x /sd-cli /sd-server ENTRYPOINT [ "/sd-cli" ] diff --git a/Dockerfile.musa b/Dockerfile.musa index 8e23318d..1c1d4d28 100644 --- a/Dockerfile.musa +++ b/Dockerfile.musa @@ -24,12 +24,22 @@ RUN mkdir build && cd build && \ cmake .. -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ \ -DCMAKE_C_FLAGS="${CMAKE_C_FLAGS} -fopenmp -I/usr/lib/llvm-14/lib/clang/14.0.0/include -L/usr/lib/llvm-14/lib" \ -DCMAKE_CXX_FLAGS="${CMAKE_CXX_FLAGS} -fopenmp -I/usr/lib/llvm-14/lib/clang/14.0.0/include -L/usr/lib/llvm-14/lib" \ - -DSD_MUSA=ON -DCMAKE_BUILD_TYPE=Release && \ + -DSD_MUSA=ON \ + -DSD_BUILD_SHARED_LIBS=ON \ + -DGGML_NATIVE=OFF \ + -DSD_BUILD_SHARED_GGML_LIB=ON \ + -DGGML_BACKEND_DL=ON \ + -DGGML_CPU_ALL_VARIANTS=ON \ + -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \ + -DCMAKE_INSTALL_RPATH='$ORIGIN' \ + -DCMAKE_BUILD_TYPE=Release && \ cmake --build . --config Release FROM mthreads/musa:${MUSA_VERSION}-runtime-ubuntu${UBUNTU_VERSION}-amd64 as runtime -COPY --from=build /sd.cpp/build/bin/sd-cli /sd-cli -COPY --from=build /sd.cpp/build/bin/sd-server /sd-server +COPY --from=build /sd.cpp/build/bin /sd.cpp/bin +RUN printf '#!/bin/sh\nexec /sd.cpp/bin/sd-cli "$@"\n' > /sd-cli && \ + printf '#!/bin/sh\nexec /sd.cpp/bin/sd-server "$@"\n' > /sd-server && \ + chmod +x /sd-cli /sd-server ENTRYPOINT [ "/sd-cli" ] diff --git a/Dockerfile.sycl b/Dockerfile.sycl index 88794be0..9927ae7e 100644 --- a/Dockerfile.sycl +++ b/Dockerfile.sycl @@ -29,4 +29,4 @@ FROM intel/oneapi-basekit:${SYCL_VERSION}-devel-ubuntu24.04 AS runtime COPY --from=build /sd.cpp/build/bin/sd-cli /sd-cli COPY --from=build /sd.cpp/build/bin/sd-server /sd-server -ENTRYPOINT [ "/sd-cli" ] +ENTRYPOINT [ "/sd-cli" ] \ No newline at end of file diff --git a/Dockerfile.vulkan b/Dockerfile.vulkan index a8f1b872..51848260 100644 --- a/Dockerfile.vulkan +++ b/Dockerfile.vulkan @@ -19,8 +19,16 @@ WORKDIR /sd.cpp COPY . . -RUN cmake . -B ./build -DSD_VULKAN=ON -RUN cmake --build ./build --config Release --parallel +RUN cmake . -B ./build \ + -DSD_VULKAN=ON \ + -DSD_BUILD_SHARED_LIBS=ON \ + -DGGML_NATIVE=OFF \ + -DSD_BUILD_SHARED_GGML_LIB=ON \ + -DGGML_BACKEND_DL=ON \ + -DGGML_CPU_ALL_VARIANTS=ON \ + -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \ + -DCMAKE_INSTALL_RPATH='$ORIGIN' +RUN cmake --build ./build --config Release -j$(nproc) FROM ubuntu:$UBUNTU_VERSION AS runtime @@ -28,7 +36,9 @@ RUN apt-get update && \ apt-get install --yes --no-install-recommends libgomp1 libvulkan1 mesa-vulkan-drivers && \ apt-get clean -COPY --from=build /sd.cpp/build/bin/sd-cli /sd-cli -COPY --from=build /sd.cpp/build/bin/sd-server /sd-server +COPY --from=build /sd.cpp/build/bin /sd.cpp/bin +RUN printf '#!/bin/sh\nexec /sd.cpp/bin/sd-cli "$@"\n' > /sd-cli && \ + printf '#!/bin/sh\nexec /sd.cpp/bin/sd-server "$@"\n' > /sd-server && \ + chmod +x /sd-cli /sd-server ENTRYPOINT [ "/sd-cli" ] diff --git a/README.md b/README.md index 6f8e62d5..d7127bb1 100644 --- a/README.md +++ b/README.md @@ -15,6 +15,8 @@ API and command-line option may change frequently.*** ## 🔥Important News +* **2026/06/25** 🚀 stable-diffusion.cpp now supports **Krea2** +* **2026/06/04** 🚀 stable-diffusion.cpp now supports **Ideogram4** * **2026/05/31** 🚀 stable-diffusion.cpp now supports **PiD** * **2026/05/27** 🚀 stable-diffusion.cpp now supports **Lens** * **2026/05/17** 🚀 stable-diffusion.cpp now supports **LTX-2.3** @@ -33,8 +35,8 @@ API and command-line option may change frequently.*** - Super lightweight and without external dependencies - Supported models - Image Models - - SD1.x, SD2.x, [SD-Turbo](https://huggingface.co/stabilityai/sd-turbo) - - SDXL, [SDXL-Turbo](https://huggingface.co/stabilityai/sdxl-turbo) + - [SD1.x, SD2.x, SD-Turbo](./docs/sd.md) + - [SDXL, SDXL-Turbo](./docs/sd.md) - [Some SD1.x and SDXL distilled models](./docs/distilled_sd.md) - [SD3/SD3.5](./docs/sd3.md) - [FLUX.1-dev/FLUX.1-schnell](./docs/flux.md) @@ -46,23 +48,29 @@ API and command-line option may change frequently.*** - [PiD](./docs/pid.md) - [LongCat Image](./docs/longcat_image.md) - [Z-Image](./docs/z_image.md) + - [MiniT2I](./docs/minit2i.md) - [Ovis-Image](./docs/ovis_image.md) - [Anima](./docs/anima.md) - [ERNIE-Image](./docs/ernie_image.md) + - [Boogu Image](./docs/boogu_image.md) + - [Krea2](./docs/krea2.md) + - [SeFi-Image](./docs/sefi_image.md) - [HiDream-O1-Image](./docs/hidream_o1_image.md) + - [Ideogram4](./docs/ideogram4.md) - 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) + - [Boogu Image Edit](./docs/boogu_image.md) - Video Models - [Wan2.1/Wan2.2](./docs/wan.md) - [LTX-2.3](./docs/ltx2.md) - - [PhotoMaker](https://github.com/TencentARC/PhotoMaker) support. + - [PhotoMaker](./docs/photo_maker.md) support. - Control Net support with SD 1.5 - LoRA support, same as [stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#lora) - Latent Consistency Models support (LCM/LCM-LoRA) - - Faster and memory efficient latent decoding with [TAESD](https://github.com/madebyollin/taesd) - - Upscale images generated with [ESRGAN](https://github.com/xinntao/Real-ESRGAN) + - Faster and memory efficient latent decoding with [TAESD](./docs/taesd.md) + - Upscale images generated with [ESRGAN](./docs/esrgan.md) - Supported backends - CPU (AVX, AVX2 and AVX512 support for x86 architectures) - CUDA @@ -131,28 +139,9 @@ For runtime and parameter backend placement, see the [backend selection guide](. ## More Guides - [Backend selection](./docs/backend.md) -- [SD1.x/SD2.x/SDXL](./docs/sd.md) -- [SD3/SD3.5](./docs/sd3.md) -- [FLUX.1-dev/FLUX.1-schnell](./docs/flux.md) -- [FLUX.2-dev/FLUX.2-klein](./docs/flux2.md) -- [FLUX.1-Kontext-dev](./docs/kontext.md) -- [Chroma](./docs/chroma.md) -- [🔥Qwen Image](./docs/qwen_image.md) -- [🔥Qwen Image Edit series](./docs/qwen_image_edit.md) -- [🔥Wan2.1/Wan2.2](./docs/wan.md) -- [🔥LTX-2.3](./docs/ltx2.md) -- [🔥Z-Image](./docs/z_image.md) -- [Ovis-Image](./docs/ovis_image.md) -- [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) +- [RPC](./docs/rpc.md) - [LoRA](./docs/lora.md) - [LCM/LCM-LoRA](./docs/lcm.md) -- [Using PhotoMaker to personalize image generation](./docs/photo_maker.md) -- [Using ESRGAN to upscale results](./docs/esrgan.md) -- [Using TAESD to faster decoding](./docs/taesd.md) - [Docker](./docs/docker.md) - [Quantization and GGUF](./docs/quantization_and_gguf.md) - [Inference acceleration via caching](./docs/caching.md) diff --git a/assets/boogu/edit_example.png b/assets/boogu/edit_example.png new file mode 100644 index 00000000..d9e30b09 Binary files /dev/null and b/assets/boogu/edit_example.png differ diff --git a/assets/boogu/example.png b/assets/boogu/example.png new file mode 100644 index 00000000..d33eeb10 Binary files /dev/null and b/assets/boogu/example.png differ diff --git a/assets/ideogram4/example.png b/assets/ideogram4/example.png new file mode 100644 index 00000000..f140c54b Binary files /dev/null and b/assets/ideogram4/example.png differ diff --git a/assets/krea2/example.png b/assets/krea2/example.png new file mode 100644 index 00000000..c665e1e7 Binary files /dev/null and b/assets/krea2/example.png differ diff --git a/assets/sefi_image/example.png b/assets/sefi_image/example.png new file mode 100644 index 00000000..14c53b81 Binary files /dev/null and b/assets/sefi_image/example.png differ diff --git a/cmake/stable-diffusion-config.cmake.in b/cmake/stable-diffusion-config.cmake.in new file mode 100644 index 00000000..c3d7433c --- /dev/null +++ b/cmake/stable-diffusion-config.cmake.in @@ -0,0 +1,37 @@ +set(SD_VERSION "@SD_INSTALL_VERSION@") +set(SD_BUILD_COMMIT "@SD_INSTALL_COMMIT@") +set(SD_SHARED_LIB @SD_SHARED_LIB@) + +@PACKAGE_INIT@ + +set_and_check(SD_INCLUDE_DIR "@PACKAGE_SD_INCLUDE_INSTALL_DIR@") +set_and_check(SD_LIB_DIR "@PACKAGE_SD_LIB_INSTALL_DIR@") +set(SD_BIN_DIR "@PACKAGE_SD_BIN_INSTALL_DIR@") + +include(CMakeFindDependencyMacro) +find_dependency(ggml REQUIRED HINTS "${SD_LIB_DIR}/cmake") + +if(NOT TARGET stable-diffusion) + find_library(stable-diffusion_LIBRARY stable-diffusion + REQUIRED + HINTS "${SD_LIB_DIR}" + NO_CMAKE_FIND_ROOT_PATH + ) + + add_library(stable-diffusion UNKNOWN IMPORTED) + set_target_properties(stable-diffusion + PROPERTIES + INTERFACE_INCLUDE_DIRECTORIES "${SD_INCLUDE_DIR}" + INTERFACE_LINK_LIBRARIES "ggml::ggml" + IMPORTED_LINK_INTERFACE_LANGUAGES "CXX" + IMPORTED_LOCATION "${stable-diffusion_LIBRARY}" + INTERFACE_COMPILE_FEATURES "c_std_11;cxx_std_17" + POSITION_INDEPENDENT_CODE ON) + + if(SD_SHARED_LIB) + target_compile_definitions(stable-diffusion + INTERFACE SD_BUILD_SHARED_LIB) + endif() +endif() + +check_required_components(stable-diffusion) diff --git a/cmake/stable-diffusion.pc.in b/cmake/stable-diffusion.pc.in new file mode 100644 index 00000000..dad257b9 --- /dev/null +++ b/cmake/stable-diffusion.pc.in @@ -0,0 +1,11 @@ +prefix=@CMAKE_INSTALL_PREFIX@ +exec_prefix=${prefix} +libdir=@CMAKE_INSTALL_FULL_LIBDIR@ +includedir=@CMAKE_INSTALL_FULL_INCLUDEDIR@ + +Name: stable-diffusion +Description: Diffusion model(SD,Flux,Wan,Qwen Image,Z-Image,...) inference in pure C/C++ +Version: @SDCPP_BUILD_VERSION@ +Libs: -L${libdir} -lstable-diffusion +Libs.private: -lggml -lggml-base +Cflags: -I${includedir} diff --git a/docs/backend.md b/docs/backend.md index 53088b0e..29ac8031 100644 --- a/docs/backend.md +++ b/docs/backend.md @@ -3,7 +3,7 @@ `stable-diffusion.cpp` has two backend assignments: - `--backend` selects the runtime backend used to execute model graphs. -- `--params-backend` selects the backend used to allocate model parameters. +- `--params-backend` selects where model parameters are kept. If `--params-backend` is not set, parameters use the same backend as their module runtime backend. @@ -29,6 +29,20 @@ The same syntax is used for parameter placement: sd-cli -m model.safetensors -p "a cat" --backend cuda0 --params-backend te=cpu,vae=cpu ``` +`--params-backend` also accepts the special value `disk`: + +```shell +sd-cli -m model.safetensors -p "a cat" --backend cuda0 --params-backend disk +``` + +`--max-vram` can target resolved backend/device names: + +```shell +sd-cli -m model.safetensors -p "a cat" --backend diffusion=cuda0,vae=vulkan0 --max-vram cuda0=6,vulkan0=2 +``` + +The budget applies to every module running on that backend. + Module names are case-insensitive. Hyphens and underscores in module names are ignored, so `clip_vision`, `clip-vision`, and `clipvision` are equivalent. `all=`, `default=`, and `*=` can be used to set the default backend inside a mixed assignment: @@ -64,9 +78,11 @@ The special values `auto`, `default`, and an empty backend name select the defau The special value `gpu` selects the first GPU backend, falling back to the first integrated GPU backend. +The special value `disk` is accepted only by `--params-backend`. `--backend disk` is invalid because `disk` is a parameter residency mode, not a runtime compute backend. + ## Runtime backend vs. parameter backend -The runtime backend controls where graph execution runs. The parameter backend controls where model weights are allocated. +The runtime backend controls where graph execution runs. The parameter backend controls where model weights are allocated or whether they are reloaded from disk on demand. For example: @@ -76,6 +92,16 @@ sd-cli -m model.safetensors -p "a cat" --backend cuda0 --params-backend cpu This runs all modules on `cuda0`, but stores parameters in CPU RAM. During execution, parameters are moved to the runtime backend as needed. +For example: + +```shell +sd-cli -m model.safetensors -p "a cat" --backend cuda0 --params-backend disk +``` + +This runs all modules on `cuda0`, reloads parameters from the model file as needed, and releases those parameter buffers after use. + +`disk` is never selected implicitly. If `--params-backend` is not set, parameters use the runtime backend. + Per-module assignments can be mixed: ```shell @@ -100,23 +126,27 @@ uses one shared CPU backend for both `te` and `vae` runtime execution. Runtime and parameter assignments also share the same backend cache. If `--backend diffusion=cuda0` and `--params-backend diffusion=cuda0` resolve to the same device, both use the same backend instance. +`--params-backend disk` does not create a separate backend instance. Parameters are loaded lazily using the module runtime backend. + `SDBackendManager` owns the backend instances and frees them when the context or upscaler is destroyed. Model runners receive non-owning runtime and parameter backend pointers and do not free them. ## Compatibility flags -The older CPU placement flags are still supported: +The example CLI/server still accepts these older CPU placement flags as compatibility aliases: - `--clip-on-cpu` - `--vae-on-cpu` - `--control-net-cpu` - `--offload-to-cpu` -`--clip-on-cpu`, `--vae-on-cpu`, and `--control-net-cpu` affect runtime backend assignment only when `--backend` is not set. They map to `te=cpu`, `vae=cpu`, and `controlnet=cpu`. +`--clip-on-cpu`, `--vae-on-cpu`, and `--control-net-cpu` are deprecated. The example argument layer prepends `te=cpu`, `vae=cpu`, and `controlnet=cpu` to `--backend` before creating the context. -`--offload-to-cpu` affects parameter backend assignment only when `--params-backend` is not set. It is equivalent to: +`--offload-to-cpu` prepends a CPU default to the parameter assignment in the caller before creating the context: ```shell ---params-backend cpu +--params-backend '*=cpu' ``` -Explicit `--backend` and `--params-backend` assignments are preferred for new commands. +Because this default is inserted first, later explicit `--params-backend` entries can still override it, for example `--offload-to-cpu --params-backend te=disk` keeps non-TE parameters on CPU and reloads TE parameters from disk. + +Library callers should set `backend` and `params_backend` directly. The old CPU/offload fields are no longer part of the C API. Explicit `--backend` and `--params-backend` assignments are preferred for new commands. diff --git a/docs/boogu_image.md b/docs/boogu_image.md new file mode 100644 index 00000000..670a46e4 --- /dev/null +++ b/docs/boogu_image.md @@ -0,0 +1,31 @@ +# How to Use + +Boogu Image uses a Boogu diffusion transformer, the FLUX VAE, and Qwen3-VL as the LLM text and vision encoder. + +## Download weights + +- Download Boogu Image + - safetensors: https://huggingface.co/Comfy-Org/Boogu-Image/tree/main/diffusion_models +- Download vae + - safetensors: https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/ae.safetensors +- Download Qwen3-VL 8B + - gguf: https://huggingface.co/unsloth/Qwen3-VL-8B-Instruct-GGUF/tree/main + - For image editing with GGUF text encoders, also download the matching mmproj file and pass it with `--llm_vision`. + +## Examples + +### Boogu Image Base + +``` +.\bin\Release\sd-cli.exe --diffusion-model ..\..\ComfyUI\models\diffusion_models\boogu_image_base_bf16.safetensors --llm ..\..\llm\Qwen3VL-8B-Instruct-Q4_K_M.gguf --vae ..\..\ComfyUI\models\vae\ae.sft -p "a lovely cat" --diffusion-fa -v --offload-to-cpu +``` + +Boogu Image Base example + +### Boogu Image Edit + +``` +.\bin\Release\sd-cli.exe --diffusion-model ..\..\ComfyUI\models\diffusion_models\boogu_image_edit_bf16.safetensors --llm ..\..\llm\Qwen3VL-8B-Instruct-Q4_K_M.gguf --llm_vision ..\..\llm\mmproj-Qwen3VL-8B-Instruct-F16.gguf --vae ..\..\ComfyUI\models\vae\ae.sft --diffusion-fa -v --offload-to-cpu -r ..\assets\flux\flux1-dev-q8_0.png -p "change 'flux.cpp' to 'boogu.cpp'" +``` + +Boogu Image Edit example \ No newline at end of file diff --git a/docs/ideogram4.md b/docs/ideogram4.md new file mode 100644 index 00000000..04864f27 --- /dev/null +++ b/docs/ideogram4.md @@ -0,0 +1,40 @@ +# How to Use + +## Download weights + +- Download Ideogram4 + - safetensors: https://huggingface.co/ideogram-ai/ideogram-4-fp8/tree/main/transformer +- Download Ideogram4 uncond + - safetensors: https://huggingface.co/ideogram-ai/ideogram-4-fp8/tree/main/unconditional_transformer +- Download vae + - safetensors: https://huggingface.co/black-forest-labs/FLUX.2-dev/tree/main +- Download Qwen3-VL-8B-Instruct + - gguf: https://huggingface.co/unsloth/Qwen3-VL-8B-Instruct-GGUF/tree/main + +## Convert weights + +fp8 scale -> bf16 + +``` +python .\convert_fp8_scale_to_bf16.py --input .\ideogram4_fp8.safetensors --output ideogram4_bf16.safetensors +python .\convert_fp8_scale_to_bf16.py --input .\ideogram4_uncond_fp8.safetensors --output ideogram4_uncond_bf16.safetensors +``` + +bf16 -> q8 + +``` +.\bin\Release\sd-cli.exe -M convert -m ideogram4_bf16.safetensors -o ideogram4-Q8_0.gguf --tensor-type-rules "^layers.*adaln_modulation.*weight=q8_0,layers.*attention.o.*weight=q8_0,layers.*attention.qkv.*weight=q8_0,layers.*feed_forward.*weight=q8_0" -v + +.\bin\Release\sd-cli.exe -M convert -m ideogram4_uncond_bf16.safetensors -o ideogram4_uncond-Q8_0.gguf --tensor-type-rules "^layers.*adaln_modulation.*weight=q8_0,layers.*attention.o.*weight=q8_0,layers.*attention.qkv.*weight=q8_0,layers.*feed_forward.*weight=q8_0" -v +``` + +If you want lower VRAM usage, you can change the quantization from q8_0 to a lower-level quantization, such as q4_0. + + +## Examples + +```sh +.\bin\Release\sd-cli.exe --diffusion-model ideogram4-Q8_0.gguf --uncond-diffusion-model ideogram4_uncond-Q8_0.gguf --llm ..\..\llm\Qwen3VL-8B-Instruct-Q4_K_M.gguf --vae ..\..\ComfyUI\models\vae\flux2_ae.safetensors -p '{"high_level_description":"A square 1024 x 1024 luxury fashion magazine cover featuring exactly one short chubby fluffy cat as the main model. The cat sits on a soft ivory studio floor, facing the viewer with a stylish calm expression, wearing tiny black sunglasses, a red silk scarf, and a small gold collar charm. In front of the cat on the floor is a wide horizontal luxury nameplate that clearly reads ideogram4.cpp. The whole design feels premium, fashionable, clean, and editorial.","style_description":{"aesthetics":"luxury fashion magazine cover, high-end pet couture campaign, minimalist editorial design, elegant studio photography, soft paper texture, refined typography, fashionable and polished","lighting":"Soft diffused studio lighting, gentle spotlight on the cat, subtle floor shadow, warm ivory highlights, clean separation between subject and background","photo":"high-resolution fashion editorial photography look, front-facing cat portrait, crisp fur details, glossy sunglasses, clear readable nameplate text, shallow depth of field","medium":"mixed media fashion photography and premium editorial graphic design","color_palette":["#F4EFE7","#111111","#D8B56D","#B73A3A","#FFFFFF","#8A7A6A"]},"compositional_deconstruction":{"canvas":"Square 1024 x 1024 canvas with a normal upright orientation. Do not rotate the poster or any text. Use a clean fashion magazine cover layout.","background":"Warm ivory studio backdrop with subtle paper grain, a soft spotlight gradient, faint floor shadow, and a few minimal gold editorial lines. The background is spacious, premium, and uncluttered.","layout":"Top center has a small elegant headline. Center area features one cat as the main fashion model. Lower foreground has a wide horizontal luxury nameplate placed on the floor in front of the cat. Bottom center has a small footer. All text is horizontal, upright, and readable left to right.","elements":[{"type":"text","desc":"Top center headline reading LOOK WHAT I FOUND in a refined high-fashion serif font. The headline is horizontal, centered, elegant, and secondary to the nameplate text."},{"type":"obj","desc":"Exactly one short chubby fluffy cat sitting in the center like a luxury fashion model. The cat has a large round head, compact body, short legs, soft detailed fur, expressive eyes, and a calm confident pose. The cat is cute and rounded, not tall, not stretched, not duplicated."},{"type":"obj","desc":"Tiny glossy black sunglasses worn naturally by the cat, slightly oversized but still showing the cat face clearly. The sunglasses add a chic fashion-editorial attitude."},{"type":"obj","desc":"A red silk scarf tied neatly around the cat neck, with soft folds and a couture feeling. The scarf must not cover the cat face or the nameplate."},{"type":"obj","desc":"A small gold collar charm or fashion accessory under the scarf, subtle and premium, adding a luxury campaign detail."},{"type":"obj","desc":"In the lower foreground, place a wide horizontal luxury nameplate on the floor in front of the cat. The nameplate is low, flat, landscape-oriented, much wider than tall, like a fashion show seat card or premium display plaque. It is centered, front-facing, level, and fully visible. It must not become vertical, tall, standing, rotated, or side-facing."},{"type":"text","desc":"Print the exact text ideogram4.cpp only on the wide horizontal nameplate. Use clean bold black lettering, perfectly spelled, lowercase, with the number 4 and .cpp extension. The text must fit completely inside the nameplate, stay horizontal, and be readable from left to right."},{"type":"obj","desc":"Add sparse premium editorial accents around the edges: thin gold lines, small code brackets, tiny cursor marks, subtle dots, and minimal geometric details. No extra cats, no stickers, no animal faces, no busy decorations."},{"type":"text","desc":"Bottom center footer reading tiny paws, big compile energy in a small refined monospace or editorial font. The footer is horizontal, centered, understated, and much smaller than the nameplate text."}]}}' --diffusion-fa -v --offload-to-cpu -H 1024 -W 1024 +``` + +ideogram4 image example diff --git a/docs/imatrix.md b/docs/imatrix.md new file mode 100644 index 00000000..df877a00 --- /dev/null +++ b/docs/imatrix.md @@ -0,0 +1,59 @@ +# Importance Matrix (imatrix) Quantization + +## What is an Importance Matrix? + +Quantization reduces the precision of a model's weights, decreasing its size and computational requirements. However, this can lead to a loss of quality. An importance matrix helps mitigate this by identifying which weights are *most* important for the model's performance. During quantization, these important weights are preserved with higher precision, while less important weights are quantized more aggressively. This allows for better overall quality at a given quantization level. + +This originates from work done with language models in [llama.cpp](https://github.com/ggml-org/llama.cpp/blob/master/tools/imatrix/README.md). + +## Usage + +The imatrix feature involves two main steps: *training* the matrix and *using* it during quantization. + +### Training the Importance Matrix + +To generate an imatrix, run stable-diffusion.cpp with the `--imat-out` flag, specifying the output filename. This process runs alongside normal image generation. + +```bash +sd.exe [same exact parameters as normal generation] --imat-out imatrix.dat +``` + +* **`[same exact parameters as normal generation]`**: Use the same command-line arguments you would normally use for image generation (e.g., prompt, dimensions, sampling method, etc.). +* **`--imat-out imatrix.dat`**: Specifies the output file for the generated imatrix. + +You can generate multiple images at once using the `-b` flag to speed up the training process. + +### Continuing Training an Existing Matrix + +If you want to refine an existing imatrix, use the `--imat-in` flag *in addition* to `--imat-out`. This will load the existing matrix and continue training it. + +```bash +sd.exe [same exact parameters as normal generation] --imat-out imatrix.dat --imat-in imatrix.dat +``` +With that, you can train and refine the imatrix while generating images like you'd normally do. + +### Using Multiple Matrices + +You can load and merge multiple imatrices together: + +```bash +sd.exe [same exact parameters as normal generation] --imat-out imatrix.dat --imat-in imatrix.dat --imat-in imatrix2.dat +``` + +### Quantizing with an Importance Matrix + +To quantize a model using a trained imatrix, use the `-M convert` option (or equivalent quantization command) and the `--imat-in` flag, specifying the imatrix file. + +```bash +sd.exe -M convert [same exact parameters as normal quantization] --imat-in imatrix.dat +``` + +* **`[same exact parameters as normal quantization]`**: Use the same command-line arguments you would normally use for quantization (e.g., target quantization method, input/output filenames). +* **`--imat-in imatrix.dat`**: Specifies the imatrix file to use during quantization. You can specify multiple `--imat-in` flags to combine multiple matrices. + +## Important Considerations + +* The quality of the imatrix depends on the prompts and settings used during training. Use prompts and settings representative of the types of images you intend to generate for the best results. +* Experiment with different training parameters (e.g., number of images, prompt variations) to optimize the imatrix for your specific use case. +* The performance impact of training an imatrix during image generation or using an imatrix for quantization is negligible. +* Using already quantized models to train the imatrix seems to be working fine. \ No newline at end of file diff --git a/docs/krea2.md b/docs/krea2.md new file mode 100644 index 00000000..b47a0354 --- /dev/null +++ b/docs/krea2.md @@ -0,0 +1,27 @@ +# How to Use + +Krea2 uses a Krea2 diffusion transformer, the Wan2.1 VAE, and Qwen3-VL 4B as the LLM text encoder. + +## Download weights + +- Download Krea2 Raw + - safetensors: https://huggingface.co/krea/Krea-2-Raw/tree/main + - gguf: https://huggingface.co/realrebelai/KREA-2_GGUFs/tree/main/BASE +- Download Krea2 Turbo + - safetensors: https://huggingface.co/krea/Krea-2-Turbo/tree/main + - gguf: https://huggingface.co/realrebelai/KREA-2_GGUFs/tree/main/TURBO +- Download vae + - safetensors: https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/blob/main/split_files/vae/wan_2.1_vae.safetensors +- Download Qwen3-VL 4B + - safetensors: https://huggingface.co/Comfy-Org/Krea-2/tree/main/text_encoders + - gguf: https://huggingface.co/Qwen/Qwen3-VL-4B-Instruct-GGUF/tree/main + +## Examples + +### Krea2 + +``` +.\bin\Release\sd-cli.exe --diffusion-model ..\..\ComfyUI\models\diffusion_models\Krea-2-Raw-Q8_0.gguf --llm ..\..\ComfyUI\models\text_encoders\Qwen3-VL-4B-Instruct-Q4_K_M.gguf --vae ..\..\ComfyUI\models\vae\wan_2.1_vae.safetensors -p "a lovely cat holding a sign says 'krea2.cpp'" --diffusion-fa -v --offload-to-cpu +``` + +Krea2 Raw example diff --git a/docs/minit2i.md b/docs/minit2i.md new file mode 100644 index 00000000..7e120b38 --- /dev/null +++ b/docs/minit2i.md @@ -0,0 +1,48 @@ +# How to Use + +MiniT2I uses a MiniT2I diffusion transformer and `google/flan-t5-large` as the text encoder. + +## Download weights + +- Download MiniT2I diffusion model + - safetensors: https://huggingface.co/MiniT2I/MiniT2I/tree/main/minit2i-b-16/transformer (`diffusion_pytorch_model.safetensors`) +- Download flan-t5-large text encoder + - safetensors: https://huggingface.co/google/flan-t5-large/tree/main (`model.safetensors`) + +## Examples + +### Mac Metal + +``` +./bin/sd-cli \ + --backend metal \ + --diffusion-model ../models/minit2i/diffusion_pytorch_model.safetensors \ + --t5xxl ../models/flan-t5-large/model.safetensors \ + --prompt "a cat" \ + --steps 100 \ + --cfg-scale 6 \ + --width 512 \ + --height 512 \ + --seed 42 \ + --sampling-method euler \ + --rng cpu \ + --output minit2i_metal.png \ + --threads 8 +``` + +### CUDA with diffusion flash attention + +``` +./bin/sd-cli \ + --diffusion-model ../models/minit2i/diffusion_pytorch_model.safetensors \ + --t5xxl ../models/flan-t5-large/model.safetensors \ + --prompt "a cat" \ + --steps 100 \ + --cfg-scale 6 \ + --width 512 \ + --height 512 \ + --seed 42 \ + --sampling-method euler \ + --diffusion-fa \ + --output minit2i_cuda.png +``` diff --git a/docs/model_config.md b/docs/model_config.md new file mode 100644 index 00000000..8c562fff --- /dev/null +++ b/docs/model_config.md @@ -0,0 +1,118 @@ +# Model Configuration Conventions + +This document describes the conventions for model configuration structs and +weight-based configuration detection. + +## Config Types + +Model configuration should live in a model-specific `*Config` struct. + +Examples: + +- `ZImageConfig` +- `UNetConfig` +- `MMDiTConfig` +- `LLMConfig` + +Preserve established acronym casing in type names, such as `UNet`, `MMDiT`, +`LLM`, `VAE`, and `T5`. + +Place the config struct near the top of the model header, before the main model +blocks and runner types that consume it. + +## Config Variables + +Variables and members that hold a config should be named `config`. + +Examples: + +```cpp +UNetConfig config; +UnetModelBlock unet; + +MMDiTRunner(...) + : DiffusionModelRunner(backend, params_backend, prefix), + config(MMDiTConfig::detect_from_weights(tensor_storage_map, prefix)), + mmdit(config) { +} +``` + +Avoid alternate names such as `params`, `params_cfg`, `model_params`, or +model-specific aliases unless an existing public API requires them. + +## Weight Detection + +If a model can derive configuration from loaded weight metadata, expose that +logic as a static method on the config type: + +```cpp +static XxxConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, + const std::string& prefix); +``` + +Additional selector arguments are allowed when required by an existing model +family, for example `SDVersion version` or an architecture enum: + +```cpp +static UNetConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, + const std::string& prefix, + SDVersion version = VERSION_SD1); +``` + +Use `TensorStorage` metadata, especially `n_dims` and `ne`, to infer shapes. +Do not load or parse tensor data for config detection. + +Detection should respect `prefix`. For nested weights, construct full names from +`prefix + "." + suffix` or filter entries with `starts_with(name, prefix)`. + +Do not add persistent config fields such as `inferred_from_weights` only to +record whether detection happened. If the function needs to decide whether to +print a debug line, keep that as local control flow inside `detect_from_weights`. + +## Logging + +When config values are inferred from weights, print one `LOG_DEBUG` line at the +end of `detect_from_weights`. + +Example: + +```cpp +LOG_DEBUG("llm: num_layers = %" PRId64 ", vocab_size = %" PRId64 ", hidden_size = %" PRId64 ", intermediate_size = %" PRId64, + config.num_layers, + config.vocab_size, + config.hidden_size, + config.intermediate_size); +``` + +Only print the config detection log when the function actually inferred values +from weights. Do not duplicate the same config summary in runner constructors or +model loading code. + +Use the correct format specifiers for field types, such as `%" PRId64 "` for +`int64_t` and `%d` for `int`. + +## Runner And Model Responsibilities + +Runners should detect the config once and pass it into the model block: + +```cpp +struct XxxRunner : public DiffusionModelRunner { + XxxConfig config; + XxxModel model; + + XxxRunner(..., const String2TensorStorage& tensor_storage_map, const std::string prefix) + : DiffusionModelRunner(backend, params_backend, prefix), + config(XxxConfig::detect_from_weights(tensor_storage_map, prefix)), + model(config) { + model.init(params_ctx, tensor_storage_map, prefix); + } +}; +``` + +Model blocks should consume `config` directly instead of re-scanning weights in +their constructors. Keep config-derived behavior centralized in the config +struct. + +If a model has no weight-derived config today, it may still provide +`detect_from_weights` for API consistency, but it should not print a config +detection log unless it actually derives values from weights. diff --git a/docs/performance.md b/docs/performance.md index 0c4735e0..ed86a4f7 100644 --- a/docs/performance.md +++ b/docs/performance.md @@ -21,6 +21,38 @@ and the compute buffer shrink in the debug log: Using `--offload-to-cpu` allows you to offload weights to the CPU, saving VRAM without reducing generation speed. +## Use params backend to reduce VRAM or RAM usage. + +`--params-backend` controls where model parameters are kept. If it is not set, parameters use the same backend as `--backend`, so a GPU runtime backend also keeps parameters in VRAM. + +Use CPU params to reduce VRAM usage: + +```shell +--backend cuda0 --params-backend cpu +``` + +This keeps model weights in system RAM and moves them to the runtime backend when needed. In the example CLI/server, `--offload-to-cpu` is a compatibility shortcut that prepends `*=cpu` to `--params-backend` before creating the context, so explicit module assignments can still override it: + +```shell +--offload-to-cpu --params-backend te=disk +``` + +Use disk params to reduce both VRAM and RAM usage: + +```shell +--backend cuda0 --params-backend disk +``` + +This reloads parameters from the model file on demand and releases them after use. It has the lowest memory residency, but can be slower because weights must be read again. `disk` is never selected implicitly; set it explicitly when RAM usage matters more than reload cost. + +Per-module assignments can target only the largest modules: + +```shell +--backend cuda0 --params-backend diffusion=disk,te=cpu,vae=cpu +``` + +See [backend selection](./backend.md) for full syntax. + ## Use quantization to reduce memory usage. -[quantization](./quantization_and_gguf.md) \ No newline at end of file +[quantization](./quantization_and_gguf.md) diff --git a/docs/pulid.md b/docs/pulid.md new file mode 100644 index 00000000..b7eec37e --- /dev/null +++ b/docs/pulid.md @@ -0,0 +1,196 @@ +# PuLID-Flux face-identity preservation + +stable-diffusion.cpp supports the [PuLID-Flux](https://github.com/ToTheBeginning/PuLID) +identity-injection technique on top of Flux.1 (schnell or dev) models. +Given a single source portrait, PuLID-Flux produces new generations that +preserve the source person's face across arbitrary scenes, poses, and +prompts. + +Unlike PhotoMaker (which extracts the identity inside the inference +process from a directory of images), PuLID-Flux's identity extractor is +a heavy stack (insightface ArcFace + EVA-CLIP-L + IDFormer encoder) that +is impractical to port to C++/ggml. To keep this implementation small and +cross-vendor, **stable-diffusion.cpp consumes a precomputed identity +embedding** produced by an external Python tool that runs once per source +portrait. Everything downstream of that one-shot extraction is C++ and +runs on any backend (Vulkan, CUDA, Metal, ROCm, CPU). + +## Architecture summary + +The PuLID-Flux contribution to the Flux denoise loop is a stack of 20 +small cross-attention modules (`PerceiverAttentionCA`) inserted between +the Flux transformer blocks: + +- After every 2nd of the 19 double-stream blocks (10 hook points) +- After every 4th of the 38 single-stream blocks (10 hook points) + +Each cross-attention layer takes the current image tokens as query, the +32-token / 2048-dim identity embedding as key+value, and adds its output +(scaled by `id_weight`, typically 1.0) back to the image tokens. + +## Required weights + +Three files in addition to the standard Flux weight set: + +1. **Flux base** (transformer + VAE + clip_l + t5xxl) -- exactly as + [docs/flux.md](flux.md) describes. +2. **PuLID weights** -- download from + [guozinan/PuLID](https://huggingface.co/guozinan/PuLID): + - `pulid_flux_v0.9.0.safetensors` or `pulid_flux_v0.9.1.safetensors` + (recommended; this implementation is verified against v0.9.1) + - **v1.1 (`pulid_v1.1.safetensors`) is NOT yet supported** -- it uses + renamed keys (`id_adapter_attn_layers.*` instead of `pulid_ca.*`) + and possibly different module structure. Future PR. +3. **Identity embedding (.pulidembd)** -- produced by the precompute + tool below. + +## Precompute the identity embedding + +The precompute tool runs the PyTorch identity-extraction stack on a +single portrait image and writes the resulting `(32, 2048)` embedding +to a `.pulidembd` binary file (about 131 KB). Run it once per source +person; the same file is reused for any number of generations. + +A reference Python script is provided alongside this docs file at +[`script/pulid_extract_id.py`](../script/pulid_extract_id.py). It +requires: +- A working CUDA / CPU PyTorch stack +- `insightface`, `facexlib`, `eva-clip`, `torchvision`, `opencv-python`, + `huggingface_hub`, `gguf` +- The PuLID weights file (same one stable-diffusion.cpp will load below) +- The ToTheBeginning/PuLID repo's `pulid/` package (including + `pulid/pipeline_flux.py`) and `eva_clip/` package on `PYTHONPATH`; `flux/` + is not needed for embedding extraction + +Run it as: + +``` +python pulid_extract_id.py \ + --portrait /path/to/source-photo.jpg \ + --pulid-weights /path/to/pulid_flux_v0.9.1.safetensors \ + --out /path/to/source.pulidembd +``` + +## Format (gguf) + +The embedding is a standard **gguf** container holding a single tensor: + +``` +tensor name : "pulid_id" +shape : [token_dim, num_tokens] (ggml order; typically [2048, 32]) +type : F16 (also accepts F32 / BF16) +metadata : general.architecture = "pulid", pulid.version = 1 +``` + +stable-diffusion.cpp loads it with the normal gguf reader +(`gguf_init_from_file`) and converts to fp32 at load time -- no bespoke +parser. Total file size for the typical (32, 2048, fp16) case is ~131 KB. + +## Command-line usage + +``` +.\bin\Release\sd-cli.exe \ + --diffusion-model models\flux1-schnell-Q4_K_S.gguf \ + --vae models\ae.safetensors \ + --clip_l models\clip_l.safetensors \ + --t5xxl models\t5xxl_fp16.safetensors \ + --pulid-weights models\pulid_flux_v0.9.1.safetensors \ + --pulid-id-embedding source.pulidembd \ + --pulid-id-weight 1.0 \ + -p "candid photograph of a young woman on a beach at sunset" \ + --cfg-scale 1.0 --sampling-method euler --steps 4 -W 512 -H 512 \ + --seed 42 --clip-on-cpu \ + -o out.png +``` + +For Flux Dev (instead of Schnell), add `--guidance 3.5` and `--steps 20`. + +## Flags + +| Flag | Purpose | +|----------------------------|-------------------------------------------------------------------| +| `--pulid-weights ` | Path to `pulid_flux_v0.9.x.safetensors`. Loaded with the model. | +| `--pulid-id-embedding

` | Path to a `.pulidembd` binary produced by the precompute tool. | +| `--pulid-id-weight ` | Identity-injection strength. Typical 0.7-1.2; default 1.0. | + +All three flags must be set together to activate PuLID. Setting only +`--pulid-weights` (no embedding) loads the weights but disables injection +at runtime. Setting `--pulid-id-weight 0` zeros out the contribution +(useful for falsification testing: outputs should be byte-identical to +a no-PuLID run with the same seed). + +## Memory budget + +At 512x512, 4 steps (Schnell), the 20 cross-attention layers add roughly +10% to denoise time and almost nothing to peak VRAM. Tested on a 12 GB +consumer card alongside Flux Schnell Q4 GGUF + CPU-offloaded clip_l and +t5xxl + GPU-resident VAE. + +At 1024x1024 with Flux Dev Q4 + 20 steps + PuLID, the VAE decode compute +buffer doesn't fit on a 12 GB card even with `--vae-on-cpu`. Workaround: +explicitly route VAE to the CPU backend instead of the offload flag: + +``` +--backend "diffusion=vulkan0,vae=cpu" +``` + +The `--vae-on-cpu` flag offloads VAE weights but leaves the compute graph +on the default backend; this is existing stable-diffusion.cpp behavior, +not a PuLID-specific issue. Documented here because anyone running PuLID +at 1024 will hit it. + +## Backend selection + +The standard `--backend` flag works as documented. Common patterns: + +``` +# AMD Vulkan +--backend "diffusion=vulkan0,vae=cpu" + +# NVIDIA Vulkan +--backend "diffusion=vulkan1,vae=cpu" + +# CUDA +--backend "diffusion=cuda0,vae=cpu" +``` + +The PuLID cross-attention layers run on the same backend as the main +diffusion model. They have not yet been independently profiled on every +backend; only Vulkan and CPU have been tested by the original contributor. + +## Verification + +A three-way SHA-256 check is the recommended sanity test when bringing up +a new combination of model + backend + hardware: + +| Run | Expected hash relation | +|----------------------------------------------|------------------------------------| +| A: no `--pulid-*` flags | baseline | +| B: PuLID flags, `--pulid-id-weight 0.0` | **byte-identical to A** | +| C: PuLID flags, `--pulid-id-weight 1.0` | **different from A,B**, preserves source identity | + +If A and C differ but A and B differ too, the injection is allocating +or computing something even at zero weight -- likely a bug. + +## Limitations / not yet supported + +- **`--skip-layers` (skip-layer-guidance / SLG) combined with PuLID** is not + supported. The `pulid_ca` index advances per non-skipped block, so a + skipped block silently misaligns the cross-attention weight assignment + vs. the trained intervals. The reference PyTorch implementation does + not have SLG either, so there is no well-defined behavior to emulate. + Use either feature alone. +- **PuLID v1.1 weights** (`pulid_v1.1.safetensors`, renamed key layout). +- **Multiple ID images.** The reference PyTorch implementation can fuse + several portraits into one embedding for stronger identity. This + implementation accepts a single embedding produced from one or more + images by the external precompute tool. +- **Negative-prompt branch of CFG.** PuLID only injects on the positive + conditioning path in the published reference, and the implementation + here follows that. Flux's distilled guidance doesn't run a separate + uncond branch in normal use, so this matters only for `--true-cfg` + workflows that aren't standard for Flux. +- **Backends other than Vulkan and CPU** are untested by the original + contributor. The implementation is pure-ggml and should work on CUDA, + ROCm, and Metal, but verification by users on those backends is + welcomed. diff --git a/docs/rpc.md b/docs/rpc.md new file mode 100644 index 00000000..617a8b5a --- /dev/null +++ b/docs/rpc.md @@ -0,0 +1,220 @@ +# Building and Using the RPC Server with `stable-diffusion.cpp` + +This guide covers how to build a version of [the RPC server from `llama.cpp`](https://github.com/ggml-org/llama.cpp/blob/master/tools/rpc/README.md) that is compatible with your version of `stable-diffusion.cpp` to manage multi-backends setups. RPC allows you to offload specific model components to a remote server. + +> **Note on Model Location:** The model files (e.g., `.safetensors` or `.gguf`) remain on the **Client** machine. The client parses the file and transmits the necessary tensor data and computational graphs to the server. The server does not need to store the model files locally. + +## 1. Building `stable-diffusion.cpp` with RPC client + +First, you should build the client application from source. It requires `SD_RPC=ON` to include the RPC backend to your client. + +```bash +mkdir build +cd build +cmake .. \ + -DSD_RPC=ON \ + # Add other build flags here (e.g., -DSD_VULKAN=ON) +cmake --build . --config Release -j $(nproc) +``` + +> **Note:** Ensure you add the other flags you would normally use (e.g., `-DSD_VULKAN=ON`, `-DSD_CUDA=ON`, `-DSD_HIPBLAS=ON`, or `-DGGML_METAL=ON`), for more information about building `stable-diffusion.cpp` from source, please refer to the [build.md](build.md) documentation. + +## 2. Ensure `llama.cpp` is at the correct commit + +`stable-diffusion.cpp`'s RPC client is designed to work with a specific version of `llama.cpp` (compatible with the `ggml` submodule) to ensure API compatibility. The commit hash for `llama.cpp` is stored in `ggml/scripts/sync-llama.last`. + +> **Start from Root:** Perform these steps from the root of your `stable-diffusion.cpp` directory. + +1. Read the target commit hash from the submodule tracker: + + ```bash + # Linux / WSL / MacOS + HASH=$(cat ggml/scripts/sync-llama.last) + + # Windows (PowerShell) + $HASH = Get-Content -Path "ggml\scripts\sync-llama.last" + ``` + +2. Clone `llama.cpp` at the target commit . + ```bash + git clone https://github.com/ggml-org/llama.cpp.git + cd llama.cpp + git checkout $HASH + ``` + To save on download time and storage, you can use a shallow clone to download only the target commit: + ```bash + mkdir -p llama.cpp + cd llama.cpp + git init + git remote add origin https://github.com/ggml-org/llama.cpp.git + git fetch --depth 1 origin $HASH + git checkout FETCH_HEAD + ``` + +## 3. Build `llama.cpp` (RPC Server) + +The RPC server acts as the worker. You must explicitly enable the **backend** (the hardware interface, such as CUDA for Nvidia, Metal for Apple Silicon, or Vulkan) when building, otherwise the server will default to using only the CPU. + +To find the correct flags for your system, refer to the official documentation for the [`llama.cpp`](https://github.com/ggml-org/llama.cpp/blob/master/docs/build.md) repository. + +> **Crucial:** You must include the compiler flags required to satisfy the API compatibility with `stable-diffusion.cpp` (`-DGGML_MAX_NAME=128`). Without this flag, `GGML_MAX_NAME` will default to `64` for the server, and data transfers between the client and server will fail. Of course, `-DGGML_RPC` must also be enabled. +> +> I recommend disabling the `LLAMA_CURL` flag to avoid unnecessary dependencies, and disabling shared library builds to avoid potential conflicts. + +> **Build Target:** We are specifically building the `rpc-server` target. This prevents the build system from compiling the entire `llama.cpp` suite (like `llama-server`), making the build significantly faster. + +### Linux / WSL (Vulkan) + +```bash +mkdir build +cd build +cmake .. -DGGML_RPC=ON \ + -DGGML_VULKAN=ON \ # Ensure backend is enabled + -DGGML_BUILD_SHARED_LIBS=OFF \ + -DLLAMA_CURL=OFF \ + -DCMAKE_C_FLAGS=-DGGML_MAX_NAME=128 \ + -DCMAKE_CXX_FLAGS=-DGGML_MAX_NAME=128 +cmake --build . --config Release --target rpc-server -j $(nproc) +``` + +### macOS (Metal) + +```bash +mkdir build +cd build +cmake .. -DGGML_RPC=ON \ + -DGGML_METAL=ON \ + -DGGML_BUILD_SHARED_LIBS=OFF \ + -DLLAMA_CURL=OFF \ + -DCMAKE_C_FLAGS=-DGGML_MAX_NAME=128 \ + -DCMAKE_CXX_FLAGS=-DGGML_MAX_NAME=128 +cmake --build . --config Release --target rpc-server +``` + +### Windows (Visual Studio 2022, Vulkan) + +```powershell +mkdir build +cd build +cmake .. -G "Visual Studio 17 2022" -A x64 ` + -DGGML_RPC=ON ` + -DGGML_VULKAN=ON ` + -DGGML_BUILD_SHARED_LIBS=OFF ` + -DLLAMA_CURL=OFF ` + -DCMAKE_C_FLAGS=-DGGML_MAX_NAME=128 ` + -DCMAKE_CXX_FLAGS=-DGGML_MAX_NAME=128 +cmake --build . --config Release --target rpc-server +``` + +## 4. Usage + +Once both applications are built, you can run the server and the client to manage your GPU allocation. + +### Step A: Run the RPC Server + +Start the server. It listens for connections on the default address (usually `localhost:50052`). If your server is on a different machine, ensure the server binds to the correct interface and your firewall allows the connection. + +**On the Server :** +If running on the same machine, you can use the default address: + +```bash +./rpc-server +``` + +If you want to allow connections from other machines on the network: + +```bash +./rpc-server --host 0.0.0.0 +``` + +> **Security Warning:** The RPC server does not currently support authentication or encryption. **Only run the server on trusted local networks**. Never expose the RPC server directly to the open internet. + +> **Drivers & Hardware:** Ensure the Server machine has the necessary drivers installed and functional (e.g., Nvidia Drivers for CUDA, Vulkan SDK, or Metal). If no devices are found, the server will simply fallback to CPU usage. + + + +### Step B: Run with RPC device + +If everything is working correctly, you can now run the client while offloading some or all of the work to the RPC server. + +Example: Setting the main backend to the RPC0 device for doing all the work on the server. + +```bash +./sd-cli -m models/sd1.5.safetensors -p "A cat" --rpc-servers localhost:50052 --backend RPC0 +``` + +--- + +## 5. Scaling: Multiple RPC Servers + +You can connect the client to multiple RPC servers simultaneously to scale out your hardware usage. + +Example: A main machine (192.168.1.10) with 3 GPUs, with one GPU running CUDA and the other two running Vulkan, and a second machine (192.168.1.11) only one GPU. + +**On the first machine (Running two server instances):** + +**Terminal 1 (CUDA):** + +```bash +# Linux / WSL +export CUDA_VISIBLE_DEVICES=0 +cd ./build_cuda/bin/Release +./rpc-server --host 0.0.0.0 + +# Windows PowerShell +$env:CUDA_VISIBLE_DEVICES="0" +cd .\build_cuda\bin\Release +./rpc-server --host 0.0.0.0 +``` + +**Terminal 2 (Vulkan):** + +```bash +cd ./build_vulkan/bin/Release +# ignore the first GPU (used by CUDA server) +./rpc-server --host 0.0.0.0 --port 50053 -d Vulkan1,Vulkan2 +``` + +**On the second machine:** + +```bash +cd ./build/bin/Release +./rpc-server --host 0.0.0.0 +``` + +**On the Client:** +Pass multiple server addresses separated by commas. + +```bash +./sd-cli --rpc-servers 192.168.1.10:50052,192.168.1.10:50053,192.168.1.11:50052 [...] +``` + +The client will map these servers to sequential device IDs (e.g., RPC0 from the first server, RPC2, RPC3 from the second, and RPC4 from the third). With this setup, you could for example use RPC0 for the main backend, RPC1 and RPC2 for the text encoders, and RPC3 for the VAE. + +--- + +## 6. Performance Considerations + +RPC performance is heavily dependent on network bandwidth, as large weights and activations must be transferred back and forth over the network, especially for large models, or when using high resolutions. For best results, ensure your network connection is stable and has sufficient bandwidth (>1Gbps recommended). This shoumd not be a concern if you are running the server and client on the same machine, as the data transfer will happen over the loopback interface. \ No newline at end of file diff --git a/docs/sefi_image.md b/docs/sefi_image.md new file mode 100644 index 00000000..e0e12fa8 --- /dev/null +++ b/docs/sefi_image.md @@ -0,0 +1,50 @@ +# How to Use + +SeFi-Image uses a Flux2-style dual-time transformer (semantic + texture streams), the standard Flux2 VAE, and Qwen3-VL as the LLM text encoder. Tech report: [arXiv:2606.22568](https://arxiv.org/abs/2606.22568). + +## Download weights + +The SeFi-Image family ships in three scales (1B / 2B / 5B) and three families (Base / RL / turbo), all gated on Hugging Face under https://huggingface.co/SeFi-Image. + +- 1B and 2B variants pair with Qwen3-VL-2B-Instruct. +- 5B variants pair with Qwen3-VL-4B-Instruct. +- All variants use the standard Flux2 VAE (`flux2_ae.safetensors` from https://huggingface.co/black-forest-labs/FLUX.2-dev). + +Convert the transformer and text encoder to sd.cpp safetensors: + +```bash +python3 script/convert_sefi.py /sefi__.safetensors +python3 script/convert_qwen3_vl.py /Qwen3-VL-XB-Instruct /qwen3_vl_b.safetensors +``` + +## Variant defaults + +| Family | timestep_shift_alpha | steps | cfg-scale | +|---|---|---|---| +| Base | 0.3 | 50 | 4.0 | +| RL | 0.3 | 50 | 4.0 | +| turbo | 1.0 | 4 | 1.0 | + +The dispatcher picks `alpha` from the filename (`turbo` substring => 1.0, otherwise 0.3). Override via `--extra-sample-args sefi_alpha=` or `sefi_delta_t=`. + +## Examples + +### 1B / 2B turbo + +``` +./build/bin/sd-cli --diffusion-model /path/to/sefi_1b_turbo.safetensors --vae /path/to/flux2_ae.safetensors --llm /path/to/qwen3_vl_2b.safetensors -p "a photograph of an orange tabby cat sitting on a couch" --cfg-scale 1.0 --steps 4 -W 1024 -H 1024 -s 42 --diffusion-fa --offload-to-cpu -o out.png +``` + +### 1B / 2B base + +``` +./build/bin/sd-cli --diffusion-model /path/to/sefi_1b_base.safetensors --vae /path/to/flux2_ae.safetensors --llm /path/to/qwen3_vl_2b.safetensors -p "a photograph of an orange tabby cat sitting on a couch" --cfg-scale 4.0 --steps 50 -W 1024 -H 1024 -s 42 --diffusion-fa --offload-to-cpu -o out.png +``` + +### 5B (needs streaming on 12 GiB VRAM) + +``` +./build/bin/sd-cli --diffusion-model /path/to/sefi_5b_turbo.safetensors --vae /path/to/flux2_ae.safetensors --llm /path/to/qwen3_vl_4b.safetensors -p "a photograph of an orange tabby cat sitting on a couch" --cfg-scale 1.0 --steps 4 -W 1024 -H 1024 -s 42 --diffusion-fa --max-vram 8 --stream-layers --offload-to-cpu -o out.png +``` + +SeFi-Image 5B turbo example diff --git a/examples/cli/README.md b/examples/cli/README.md index 4dcfb89b..e8a14098 100644 --- a/examples/cli/README.md +++ b/examples/cli/README.md @@ -1,202 +1,9 @@ -# Run +# Usage -``` -usage: ./bin/sd-cli [options] +For detailed command-line arguments, run: -CLI Options: - -o, --output path to write result image to. you can use printf-style %d format specifiers for image - sequences (default: ./output.png) (eg. output_%03d.png). Single-file video outputs - support .avi, .webm, and animated .webp - --image path to the image to inspect (for metadata mode) - --metadata-format metadata output format, one of [text, json] (default: text) - --preview-path path to write preview image to (default: ./preview.png). Multi-frame previews support - .avi, .webm, and animated .webp - --preview-interval interval in denoising steps between consecutive updates of the image preview file - (default is 1, meaning updating at every step) - --output-begin-idx starting index for output image sequence, must be non-negative (default 0 if specified - %d in output path, 1 otherwise) - --canny apply canny preprocessor (edge detection) - --convert-name convert tensor name (for convert mode) - -v, --verbose print extra info - --color colors the logging tags according to level - --taesd-preview-only prevents usage of taesd for decoding the final image. (for use with --preview tae) - --preview-noisy enables previewing noisy inputs of the models rather than the denoised outputs - --metadata-raw include raw hex previews for unparsed metadata payloads - --metadata-brief truncate long metadata text values in text output - --metadata-all include structural/container entries such as IHDR, IDAT, and non-metadata JPEG segments - -M, --mode run mode, one of [img_gen, vid_gen, upscale, convert, metadata], default: img_gen - --preview preview method. must be one of the following [none, proj, tae, vae] (default is none) - -h, --help show this help message and exit - -Context Options: - -m, --model path to full model - --clip_l path to the clip-l text encoder - --clip_g path to the clip-g text encoder - --clip_vision path to the clip-vision encoder - --t5xxl path to the t5xxl text encoder - --llm path to the llm text encoder. For example: (qwenvl2.5 for qwen-image, - mistral-small3.2 for flux2, ...) - --llm_vision path to the llm vit - --qwen2vl alias of --llm. Deprecated. - --qwen2vl_vision alias of --llm_vision. Deprecated. - --diffusion-model path to the standalone diffusion model - --high-noise-diffusion-model path to the standalone high noise diffusion model - --vae path to standalone vae model - --taesd path to taesd. Using Tiny AutoEncoder for fast decoding (low quality) - --tae alias of --taesd - --control-net path to control net model - --embd-dir embeddings directory - --lora-model-dir lora model directory - --hires-upscalers-dir highres fix upscaler model directory - --tensor-type-rules weight type per tensor pattern (example: "^vae\.=f16,model\.=q8_0") - --photo-maker path to PHOTOMAKER model - --upscale-model path to esrgan model. - -t, --threads number of threads to use during computation (default: -1). If threads <= 0, - then threads will be set to the number of CPU physical cores - --chroma-t5-mask-pad t5 mask pad size of chroma - --max-vram 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) - --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 - --mmap whether to memory-map model - --control-net-cpu keep controlnet in cpu (for low vram) - --clip-on-cpu keep clip in cpu (for low vram) - --vae-on-cpu keep vae in cpu (for low vram) - --fa use flash attention - --diffusion-fa use flash attention in the diffusion model only - --diffusion-conv-direct use ggml_conv2d_direct in the diffusion model - --vae-conv-direct use ggml_conv2d_direct in the vae model - --circular enable circular padding for convolutions - --circularx enable circular RoPE wrapping on x-axis (width) only - --circulary enable circular RoPE wrapping on y-axis (height) only - --chroma-disable-dit-mask disable dit mask for chroma - --qwen-image-zero-cond-t enable zero_cond_t for qwen image - --chroma-enable-t5-mask enable t5 mask for chroma - --type weight type (examples: f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0, q2_K, q3_K, - q4_K). If not specified, the default is the type of the weight file - --rng RNG, one of [std_default, cuda, cpu], default: cuda(sd-webui), cpu(comfyui) - --sampler-rng sampler RNG, one of [std_default, cuda, cpu]. If not specified, use --rng - --prediction prediction type override, one of [eps, v, edm_v, sd3_flow, flux_flow, - flux2_flow] - --lora-apply-mode the way to apply LoRA, one of [auto, immediately, at_runtime], default is - auto. In auto mode, if the model weights contain any quantized parameters, - the at_runtime mode will be used; otherwise, immediately will be used.The - immediately mode may have precision and compatibility issues with quantized - parameters, but it usually offers faster inference speed and, in some cases, - lower memory usage. The at_runtime mode, on the other hand, is exactly the - opposite. - -Generation Options: - -p, --prompt the prompt to render - -n, --negative-prompt the negative prompt (default: "") - -i, --init-img path to the init image - --end-img path to the end image, required by flf2v - --mask path to the mask image - --control-image path to control image, control net - --control-video path to control video frames, It must be a directory path. The video frames - inside should be stored as images in lexicographical (character) order. For - example, if the control video path is `frames`, the directory contain images - such as 00.png, 01.png, ... etc. - --pm-id-images-dir path to PHOTOMAKER input id images dir - --pm-id-embed-path path to PHOTOMAKER v2 id embed - --hires-upscaler highres fix upscaler, Lanczos, Nearest, Latent, Latent (nearest), Latent - (nearest-exact), Latent (antialiased), Latent (bicubic), Latent (bicubic - antialiased), or a model name under --hires-upscalers-dir (default: Latent) - --extra-sample-args extra sampler/scheduler/guidance args, key=value list. APG supports apg_eta, - apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports - slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; - ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma - --extra-tiling-args extra VAE tiling args, key=value list. LTX video VAE supports - temporal_tile_frames (default: 4), temporal_tile_overlap (default: 1) - -H, --height image height, in pixel space (default: 512) - -W, --width image width, in pixel space (default: 512) - --steps number of sample steps (default: 20) - --high-noise-steps (high noise) number of sample steps (default: -1 = auto) - --clip-skip ignore last layers of CLIP network; 1 ignores none, 2 ignores one layer - (default: -1). <= 0 represents unspecified, will be 1 for SD1.x, 2 for SD2.x - -b, --batch-count batch count - --video-frames video frames (default: 1) - --fps fps (default: 24) - --timestep-shift shift timestep for NitroFusion models (default: 0). recommended N for - NitroSD-Realism around 250 and 500 for NitroSD-Vibrant - --upscale-repeats Run the ESRGAN upscaler this many times (default: 1) - --upscale-tile-size tile size for ESRGAN upscaling (default: 128) - --hires-width highres fix target width, 0 to use --hires-scale (default: 0) - --hires-height highres fix target height, 0 to use --hires-scale (default: 0) - --hires-steps highres fix second pass sample steps, 0 to reuse --steps (default: 0) - --hires-upscale-tile-size highres fix upscaler tile size, reserved for model-backed upscalers (default: - 128) - --cfg-scale unconditional guidance scale: (default: 7.0) - --img-cfg-scale image guidance scale for inpaint or image edit models: (default: same as - --cfg-scale) - --guidance distilled guidance scale for models with guidance input (default: 3.5) - --slg-scale skip layer guidance (SLG) scale, only for DiT models: (default: 0). 0 means - disabled, a value of 2.5 is nice for sd3.5 medium - --skip-layer-start SLG enabling point (default: 0.01) - --skip-layer-end SLG disabling point (default: 0.2) - --eta noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and - res_2s; 1 for euler_a, er_sde and dpm++2s_a) - --flow-shift shift value for Flow models like SD3.x or WAN (default: auto) - --high-noise-cfg-scale (high noise) unconditional guidance scale: (default: 7.0) - --high-noise-img-cfg-scale (high noise) image guidance scale for inpaint or image edit models (default: - same as --cfg-scale) - --high-noise-guidance (high noise) distilled guidance scale for models with guidance input - (default: 3.5) - --high-noise-slg-scale (high noise) skip layer guidance (SLG) scale, only for DiT models: (default: - 0) - --high-noise-skip-layer-start (high noise) SLG enabling point (default: 0.01) - --high-noise-skip-layer-end (high noise) SLG disabling point (default: 0.2) - --high-noise-eta (high noise) noise multiplier (default: 0 for ddim_trailing, tcd, - res_multistep and res_2s; 1 for euler_a, er_sde and dpm++2s_a) - --strength strength for noising/unnoising (default: 0.75) - --pm-style-strength - --control-strength strength to apply Control Net (default: 0.9). 1.0 corresponds to full - destruction of information in init image - --moe-boundary timestep boundary for Wan2.2 MoE model. (default: 0.875). Only enabled if - `--high-noise-steps` is set to -1 - --vace-strength wan vace strength - --vae-tile-overlap tile overlap for vae tiling, in fraction of tile size (default: 0.5) - --hires-scale highres fix scale when target size is not set (default: 2.0) - --hires-denoising-strength highres fix second pass denoising strength (default: 0.7) - --increase-ref-index automatically increase the indices of references images based on the order - they are listed (starting with 1). - --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, - dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, - er_sde, euler_cfg_pp, euler_a_cfg_pp] (default: euler for Flux/SD3/Wan, euler_a otherwise) - --high-noise-sampling-method (high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, - dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, - res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp] default: euler for Flux/SD3/Wan, euler_a otherwise - --scheduler denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, - smoothstep, sgm_uniform, simple, kl_optimal, lcm, bong_tangent, ltx2], default: - 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) - --cache-mode caching method: 'easycache' (DiT), 'ucache' (UNET), - 'dbcache'/'taylorseer'/'cache-dit' (DiT block-level), 'spectrum' (UNET/DiT - Chebyshev+Taylor forecasting) - --cache-option named cache params (key=value format, comma-separated). easycache/ucache: - threshold=,start=,end=,decay=,relative=,reset=; dbcache/taylorseer/cache-dit: - Fn=,Bn=,threshold=,warmup=; spectrum: w=,m=,lam=,window=,flex=,warmup=,stop=. - Examples: "threshold=0.25" or "threshold=1.5,reset=0" - --scm-mask SCM steps mask for cache-dit: comma-separated 0/1 (e.g., - "1,1,1,0,0,1,0,0,1,0") - 1=compute, 0=can cache - --scm-policy SCM policy: 'dynamic' (default) or 'static' - --vae-tile-size tile size for vae tiling, format [X]x[Y] (default: 32x32) - --vae-relative-tile-size relative tile size for vae tiling, format [X]x[Y], in fraction of image size - if < 1, in number of tiles per dim if >=1 (overrides --vae-tile-size) +```bash +./bin/sd-cli -h ``` Metadata mode inspects PNG/JPEG container metadata without loading any model: diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index 2bb71730..7892d521 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -54,6 +54,9 @@ struct SDCliParams { bool metadata_brief = false; bool metadata_all = false; + std::string imatrix_out; + std::vector imatrix_in; + bool normal_exit = false; ArgOptions get_options() { @@ -63,19 +66,28 @@ struct SDCliParams { {"-o", "--output", "path to write result image to. you can use printf-style %d format specifiers for image sequences (default: ./output.png) (eg. output_%03d.png). Single-file video outputs support .avi, .webm, and animated .webp", + 0, &output_path}, {"", "--image", "path to the image to inspect (for metadata mode)", + 0, &image_path}, {"", "--metadata-format", "metadata output format, one of [text, json] (default: text)", + 0, &metadata_format}, {"", "--preview-path", "path to write preview image to (default: ./preview.png). Multi-frame previews support .avi, .webm, and animated .webp", + 0, &preview_path}, + {"", + "--imat-out", + "compute the imatrix for this run and save it to the provided path", + 0, + &imatrix_out}, }; options.int_options = { @@ -170,11 +182,20 @@ struct SDCliParams { return 1; }; - auto on_help_arg = [&](int argc, const char** argv, int index) { + auto on_help_arg = [&](int argc, const char** argv, int index, bool& valid) { normal_exit = true; + valid = true; return -1; }; + auto on_imatrix_in_arg = [&](int argc, const char** argv, int index) { + if (++index >= argc) { + return -1; + } + imatrix_in.push_back(argv[index]); + return 1; + }; + options.manual_options = { {"-M", "--mode", @@ -188,6 +209,10 @@ struct SDCliParams { "--help", "show this help message and exit", on_help_arg}, + {"", + "--imat-in", + "load an imatrix file for quantization or continued collection; can be specified multiple times", + on_imatrix_in_arg}, }; return options; @@ -249,6 +274,7 @@ struct SDCliParams { << " preview_fps: " << preview_fps << ",\n" << " taesd_preview: " << (taesd_preview ? "true" : "false") << ",\n" << " preview_noisy: " << (preview_noisy ? "true" : "false") << ",\n" + << " imatrix_out: \"" << imatrix_out << "\",\n" << " metadata_raw: " << (metadata_raw ? "true" : "false") << ",\n" << " metadata_brief: " << (metadata_brief ? "true" : "false") << ",\n" << " metadata_all: " << (metadata_all ? "true" : "false") << "\n" @@ -602,13 +628,32 @@ int main(int argc, const char* argv[]) { LOG_DEBUG("%s", ctx_params.to_string().c_str()); LOG_DEBUG("%s", gen_params.to_string().c_str()); + if (!cli_params.imatrix_out.empty()) { + if (fs::exists(cli_params.imatrix_out) && + std::find(cli_params.imatrix_in.begin(), cli_params.imatrix_in.end(), cli_params.imatrix_out) == cli_params.imatrix_in.end()) { + LOG_WARN("imatrix file '%s' already exists and will be overwritten", cli_params.imatrix_out.c_str()); + } + enable_imatrix_collection(); + } + + for (const auto& in_file : cli_params.imatrix_in) { + LOG_INFO("loading imatrix from '%s'", in_file.c_str()); + if (!load_imatrix(in_file.c_str())) { + LOG_WARN("failed to load imatrix from '%s'", in_file.c_str()); + } + } + if (cli_params.mode == CONVERT) { - bool success = convert(ctx_params.model_path.c_str(), - ctx_params.vae_path.c_str(), - cli_params.output_path.c_str(), - ctx_params.wtype, - ctx_params.tensor_type_rules.c_str(), - cli_params.convert_name); + bool success = convert_with_components(ctx_params.model_path.c_str(), + ctx_params.clip_l_path.c_str(), + ctx_params.clip_g_path.c_str(), + ctx_params.t5xxl_path.c_str(), + ctx_params.diffusion_model_path.c_str(), + ctx_params.vae_path.c_str(), + cli_params.output_path.c_str(), + ctx_params.wtype, + ctx_params.tensor_type_rules.c_str(), + cli_params.convert_name); if (!success) { LOG_ERROR("convert '%s'/'%s' to '%s' failed", ctx_params.model_path.c_str(), @@ -624,8 +669,6 @@ int main(int argc, const char* argv[]) { } } - bool vae_decode_only = true; - auto load_image_and_update_size = [&](const std::string& path, SDImageOwner& image, bool resize_image = true, @@ -647,21 +690,18 @@ int main(int argc, const char* argv[]) { }; if (gen_params.init_image_path.size() > 0) { - vae_decode_only = false; if (!load_image_and_update_size(gen_params.init_image_path, gen_params.init_image)) { return 1; } } if (gen_params.end_image_path.size() > 0) { - vae_decode_only = false; if (!load_image_and_update_size(gen_params.end_image_path, gen_params.end_image)) { return 1; } } if (gen_params.ref_image_paths.size() > 0) { - vae_decode_only = false; gen_params.ref_images.clear(); for (auto& path : gen_params.ref_image_paths) { SDImageOwner ref_image({0, 0, 3, nullptr}); @@ -736,18 +776,7 @@ int main(int argc, const char* argv[]) { } } - if (cli_params.mode == VID_GEN) { - vae_decode_only = false; - } - - if (gen_params.hires_enabled && - (gen_params.resolved_hires_upscaler == SD_HIRES_UPSCALER_MODEL || - gen_params.resolved_hires_upscaler == SD_HIRES_UPSCALER_LANCZOS || - gen_params.resolved_hires_upscaler == SD_HIRES_UPSCALER_NEAREST)) { - vae_decode_only = false; - } - - sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(vae_decode_only, true, cli_params.taesd_preview); + sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(cli_params.taesd_preview); SDImageVec results; int num_results = 0; @@ -779,8 +808,12 @@ int main(int argc, const char* argv[]) { if (cli_params.mode == IMG_GEN) { sd_img_gen_params_t img_gen_params = gen_params.to_sd_img_gen_params_t(); - num_results = sd_get_image_result_count(sd_ctx.get(), &img_gen_params); - results.adopt(generate_image(sd_ctx.get(), &img_gen_params), num_results); + sd_image_t* generated_images = nullptr; + if (!generate_image(sd_ctx.get(), &img_gen_params, &generated_images, &num_results)) { + generated_images = nullptr; + num_results = 0; + } + results.adopt(generated_images, num_results); } else if (cli_params.mode == VID_GEN) { sd_vid_gen_params_t vid_gen_params = gen_params.to_sd_vid_gen_params_t(); sd_image_t* generated_video = nullptr; @@ -799,12 +832,11 @@ int main(int argc, const char* argv[]) { int upscale_factor = 4; // unused for RealESRGAN_x4plus_anime_6B.pth if (ctx_params.esrgan_path.size() > 0 && gen_params.upscale_repeats > 0) { UpscalerCtxPtr upscaler_ctx(new_upscaler_ctx(ctx_params.esrgan_path.c_str(), - ctx_params.offload_params_to_cpu, ctx_params.diffusion_conv_direct, ctx_params.n_threads, gen_params.upscale_tile_size, - ctx_params.backend.c_str(), - ctx_params.params_backend.c_str())); + sd_ctx_params.backend, + sd_ctx_params.params_backend)); if (upscaler_ctx == nullptr) { LOG_ERROR("new_upscaler_ctx failed"); @@ -816,12 +848,22 @@ int main(int argc, const char* argv[]) { SDImageOwner current_image(results[i]); results[i] = {0, 0, 0, nullptr}; for (int u = 0; u < gen_params.upscale_repeats; ++u) { - SDImageOwner upscaled_image(upscale(upscaler_ctx.get(), current_image.get(), upscale_factor)); - if (upscaled_image.get().data == nullptr) { + sd_image_t* upscaled_images = nullptr; + int upscaled_count = 0; + bool upscale_ok = upscale(upscaler_ctx.get(), + current_image.get(), + upscale_factor, + &upscaled_images, + &upscaled_count); + if (!upscale_ok || upscaled_count <= 0 || upscaled_images[0].data == nullptr) { + free_sd_images(upscaled_images, upscaled_count); LOG_ERROR("upscale failed"); break; } - current_image = std::move(upscaled_image); + sd_image_t upscaled_image = upscaled_images[0]; + upscaled_images[0] = {0, 0, 0, nullptr}; + free_sd_images(upscaled_images, upscaled_count); + current_image.reset(upscaled_image); } results[i] = current_image.release(); // Set the final upscaled image as the result } @@ -833,6 +875,11 @@ int main(int argc, const char* argv[]) { return 1; } + if (!cli_params.imatrix_out.empty()) { + LOG_INFO("saving imatrix to '%s'", cli_params.imatrix_out.c_str()); + save_imatrix(cli_params.imatrix_out.c_str()); + } + free_sd_audio(generated_audio); return 0; diff --git a/examples/common/common.cpp b/examples/common/common.cpp index 0ecc72dc..524b8940 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -6,6 +6,7 @@ #include #include #include +#include #include #include #include @@ -51,6 +52,10 @@ static sd_vae_format_t str_to_vae_format(const std::string& value) { return SD_VAE_FORMAT_COUNT; } +static void prepend_backend_assignment(std::string& spec, const char* assignment) { + spec = spec.empty() ? assignment : std::string(assignment) + "," + spec; +} + #if defined(_WIN32) static std::string utf16_to_utf8(const std::wstring& wstr) { if (wstr.empty()) @@ -245,6 +250,7 @@ bool parse_options(int argc, const char** argv, const std::vector& o return false; }; + bool valid = false; for (int i = 1; i < argc; i++) { arg = argv[i]; bool found_arg = false; @@ -255,8 +261,15 @@ bool parse_options(int argc, const char** argv, const std::vector& o invalid_arg = true; return; } - *option.target = argv_to_utf8(i, argv); - found_arg = true; + if (option.concat && !option.target->empty()) { + if (option.concat > 0 && option.concat <= 0xff) { + *option.target += static_cast(option.concat); + } + *option.target += argv_to_utf8(i, argv); + } else { + *option.target = argv_to_utf8(i, argv); + } + found_arg = true; })) break; @@ -287,7 +300,7 @@ bool parse_options(int argc, const char** argv, const std::vector& o break; if (match_and_apply(options.manual_options, [&](auto& option) { - int ret = option.cb(argc, argv, i); + int ret = option.cb(argc, argv, i, valid); if (ret < 0) { invalid_arg = true; return; @@ -299,7 +312,9 @@ bool parse_options(int argc, const char** argv, const std::vector& o } if (invalid_arg) { - LOG_ERROR("error: invalid parameter for argument: %s", arg.c_str()); + if (!valid) { + LOG_ERROR("error: invalid parameter for argument: %s", arg.c_str()); + } return false; } if (!found_arg) { @@ -317,105 +332,152 @@ ArgOptions SDContextParams::get_options() { {"-m", "--model", "path to full model", + 0, &model_path}, {"", "--clip_l", - "path to the clip-l text encoder", &clip_l_path}, + "path to the clip-l text encoder", + 0, + &clip_l_path}, {"", "--clip_g", "path to the clip-g text encoder", + 0, &clip_g_path}, {"", "--clip_vision", "path to the clip-vision encoder", + 0, &clip_vision_path}, {"", "--t5xxl", "path to the t5xxl text encoder", + 0, &t5xxl_path}, {"", "--llm", "path to the llm text encoder. For example: (qwenvl2.5 for qwen-image, mistral-small3.2 for flux2, ...)", + 0, &llm_path}, {"", "--llm_vision", "path to the llm vit", + 0, &llm_vision_path}, {"", "--qwen2vl", "alias of --llm. Deprecated.", + 0, &llm_path}, {"", "--qwen2vl_vision", "alias of --llm_vision. Deprecated.", + 0, &llm_vision_path}, {"", "--diffusion-model", "path to the standalone diffusion model", + 0, &diffusion_model_path}, {"", "--high-noise-diffusion-model", "path to the standalone high noise diffusion model", + 0, &high_noise_diffusion_model_path}, + {"", + "--uncond-diffusion-model", + "path to the standalone unconditional diffusion model, currently used by Ideogram4 CFG", + 0, + &uncond_diffusion_model_path}, {"", "--embeddings-connectors", "path to LTXAV embeddings connectors", + 0, &embeddings_connectors_path}, {"", "--vae", "path to standalone vae model", + 0, &vae_path}, {"", "--vae-format", "VAE latent format override: auto, flux, sd3, or flux2 (default: auto)", + 0, &vae_format}, {"", "--audio-vae", "path to standalone LTX audio vae model", + 0, &audio_vae_path}, {"", "--taesd", "path to taesd. Using Tiny AutoEncoder for fast decoding (low quality)", + 0, &taesd_path}, {"", "--tae", "alias of --taesd", + 0, &taesd_path}, {"", "--control-net", "path to control net model", + 0, &control_net_path}, {"", "--embd-dir", "embeddings directory", + 0, &embedding_dir}, {"", "--lora-model-dir", "lora model directory", + 0, &lora_model_dir}, {"", "--hires-upscalers-dir", "highres fix upscaler model directory", + 0, &hires_upscalers_dir}, {"", "--tensor-type-rules", "weight type per tensor pattern (example: \"^vae\\.=f16,model\\.=q8_0\")", + (int)',', &tensor_type_rules}, {"", "--photo-maker", "path to PHOTOMAKER model", + 0, &photo_maker_path}, + {"", + "--pulid-weights", + "path to PuLID Flux weights", + 0, + &pulid_weights_path}, {"", "--upscale-model", "path to esrgan model.", + 0, &esrgan_path}, {"", "--backend", "runtime backend assignment, e.g. cpu or clip=cpu,vae=cuda0,diffusion=vulkan0", + (int)',', &backend}, {"", "--params-backend", - "parameter backend assignment, e.g. cpu or diffusion=cpu,clip=cpu", + "parameter backend assignment, e.g. disk, cpu, or diffusion=disk,clip=cpu", + (int)',', ¶ms_backend}, + {"", + "--rpc-servers", + "comma-separated list of RPC servers to connect to for offloading, in the format host:port, e.g. localhost:50052,192.168.1.3:50052", + (int)',', + &rpc_servers}, + {"", + "--max-vram", + "maximum VRAM budget in GiB for graph-cut segmented execution. Accepts a single value or assignments by backend/device, e.g. 6 or cuda0=6,vulkan0=4. 0 disables graph splitting; a negative value auto-detects free VRAM, sparing the specified value", + 0, + &max_vram}, }; options.int_options = { @@ -430,18 +492,15 @@ ArgOptions SDContextParams::get_options() { &chroma_t5_mask_pad}, }; - options.float_options = { - {"", - "--max-vram", - "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}, - }; - options.bool_options = { {"", "--stream-layers", "enable residency+prefetch streaming on top of --max-vram (no effect without --max-vram; defaults to false)", true, &stream_layers}, + {"", + "--eager-load", + "load all params into the params backend at model-load time instead of lazily on first use (defaults to false)", + true, &eager_load}, {"", "--force-sdxl-vae-conv-scale", "force use of conv scale on sdxl vae", @@ -456,15 +515,15 @@ ArgOptions SDContextParams::get_options() { true, &enable_mmap}, {"", "--control-net-cpu", - "keep controlnet in cpu (for low vram)", + "deprecated; use --backend controlnet=cpu", true, &control_net_cpu}, {"", "--clip-on-cpu", - "keep clip in cpu (for low vram)", + "deprecated; use --backend te=cpu", true, &clip_on_cpu}, {"", "--vae-on-cpu", - "keep vae in cpu (for low vram)", + "deprecated; use --backend vae=cpu", true, &vae_on_cpu}, {"", "--fa", @@ -594,7 +653,7 @@ ArgOptions SDContextParams::get_options() { on_sampler_rng_arg}, {"", "--prediction", - "prediction type override, one of [eps, v, edm_v, sd3_flow, flux_flow, flux2_flow]", + "prediction type override, one of [eps, v, edm_v, sd3_flow, flux_flow, sefi_flow]", on_prediction_arg}, {"", "--lora-apply-mode", @@ -651,7 +710,18 @@ bool SDContextParams::resolve(SDMode mode) { } bool SDContextParams::validate(SDMode mode) { - if (mode != UPSCALE && mode != METADATA && model_path.length() == 0 && diffusion_model_path.length() == 0) { + if (mode == CONVERT) { + const bool has_convert_input = model_path.length() != 0 || + clip_l_path.length() != 0 || + clip_g_path.length() != 0 || + t5xxl_path.length() != 0 || + diffusion_model_path.length() != 0 || + vae_path.length() != 0; + if (!has_convert_input) { + LOG_ERROR("error: convert mode needs at least one model input path\n"); + return false; + } + } else if (mode != UPSCALE && mode != METADATA && model_path.length() == 0 && diffusion_model_path.length() == 0) { LOG_ERROR("error: the following arguments are required: model_path/diffusion_model\n"); return false; } @@ -681,6 +751,25 @@ bool SDContextParams::resolve_and_validate(SDMode mode) { return true; } +void SDContextParams::prepare_backend_assignments() { + effective_backend = backend; + effective_params_backend = params_backend; + + if (offload_params_to_cpu) { + prepend_backend_assignment(effective_params_backend, "*=cpu"); + } + + if (clip_on_cpu) { + prepend_backend_assignment(effective_backend, "te=cpu"); + } + if (vae_on_cpu) { + prepend_backend_assignment(effective_backend, "vae=cpu"); + } + if (control_net_cpu) { + prepend_backend_assignment(effective_backend, "controlnet=cpu"); + } +} + std::string SDContextParams::to_string() const { std::ostringstream emb_ss; emb_ss << "{\n"; @@ -706,6 +795,7 @@ std::string SDContextParams::to_string() const { << " llm_vision_path: \"" << llm_vision_path << "\",\n" << " diffusion_model_path: \"" << diffusion_model_path << "\",\n" << " high_noise_diffusion_model_path: \"" << high_noise_diffusion_model_path << "\",\n" + << " uncond_diffusion_model_path: \"" << uncond_diffusion_model_path << "\",\n" << " embeddings_connectors_path: \"" << embeddings_connectors_path << "\",\n" << " vae_path: \"" << vae_path << "\",\n" << " vae_format: \"" << vae_format << "\",\n" @@ -723,8 +813,9 @@ std::string SDContextParams::to_string() const { << " rng_type: " << sd_rng_type_name(rng_type) << ",\n" << " sampler_rng_type: " << sd_rng_type_name(sampler_rng_type) << ",\n" << " offload_params_to_cpu: " << (offload_params_to_cpu ? "true" : "false") << ",\n" - << " max_vram: " << max_vram << ",\n" + << " max_vram: \"" << max_vram << "\",\n" << " stream_layers: " << (stream_layers ? "true" : "false") << ",\n" + << " eager_load: " << (eager_load ? "true" : "false") << ",\n" << " backend: \"" << backend << "\",\n" << " params_backend: \"" << params_backend << "\",\n" << " enable_mmap: " << (enable_mmap ? "true" : "false") << ",\n" @@ -749,7 +840,8 @@ std::string SDContextParams::to_string() const { return oss.str(); } -sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool vae_decode_only, bool free_params_immediately, bool taesd_preview) { +sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) { + prepare_backend_assignments(); embedding_vec.clear(); embedding_vec.reserve(embedding_map.size()); for (const auto& kv : embedding_map) { @@ -759,56 +851,54 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool vae_decode_only, bool f embedding_vec.emplace_back(item); } - sd_ctx_params_t sd_ctx_params = { - model_path.c_str(), - clip_l_path.c_str(), - clip_g_path.c_str(), - clip_vision_path.c_str(), - t5xxl_path.c_str(), - llm_path.c_str(), - llm_vision_path.c_str(), - diffusion_model_path.c_str(), - high_noise_diffusion_model_path.c_str(), - embeddings_connectors_path.c_str(), - vae_path.c_str(), - audio_vae_path.c_str(), - taesd_path.c_str(), - control_net_path.c_str(), - embedding_vec.data(), - static_cast(embedding_vec.size()), - photo_maker_path.c_str(), - tensor_type_rules.c_str(), - vae_decode_only, - free_params_immediately, - n_threads, - wtype, - rng_type, - sampler_rng_type, - prediction, - lora_apply_mode, - offload_params_to_cpu, - enable_mmap, - clip_on_cpu, - control_net_cpu, - vae_on_cpu, - flash_attn, - diffusion_flash_attn, - taesd_preview, - diffusion_conv_direct, - vae_conv_direct, - circular || circular_x, - circular || circular_y, - force_sdxl_vae_conv_scale, - chroma_use_dit_mask, - chroma_use_t5_mask, - chroma_t5_mask_pad, - qwen_image_zero_cond_t, - str_to_vae_format(vae_format), - max_vram, - stream_layers, - backend.c_str(), - params_backend.c_str(), - }; + sd_ctx_params_t sd_ctx_params; + sd_ctx_params_init(&sd_ctx_params); + sd_ctx_params.model_path = model_path.c_str(); + sd_ctx_params.clip_l_path = clip_l_path.c_str(); + sd_ctx_params.clip_g_path = clip_g_path.c_str(); + sd_ctx_params.clip_vision_path = clip_vision_path.c_str(); + sd_ctx_params.t5xxl_path = t5xxl_path.c_str(); + sd_ctx_params.llm_path = llm_path.c_str(); + sd_ctx_params.llm_vision_path = llm_vision_path.c_str(); + sd_ctx_params.diffusion_model_path = diffusion_model_path.c_str(); + sd_ctx_params.high_noise_diffusion_model_path = high_noise_diffusion_model_path.c_str(); + sd_ctx_params.uncond_diffusion_model_path = uncond_diffusion_model_path.c_str(); + sd_ctx_params.embeddings_connectors_path = embeddings_connectors_path.c_str(); + sd_ctx_params.vae_path = vae_path.c_str(); + sd_ctx_params.audio_vae_path = audio_vae_path.c_str(); + sd_ctx_params.taesd_path = taesd_path.c_str(); + sd_ctx_params.control_net_path = control_net_path.c_str(); + sd_ctx_params.embeddings = embedding_vec.data(); + sd_ctx_params.embedding_count = static_cast(embedding_vec.size()); + sd_ctx_params.photo_maker_path = photo_maker_path.c_str(); + sd_ctx_params.pulid_weights_path = pulid_weights_path.c_str(); + sd_ctx_params.tensor_type_rules = tensor_type_rules.c_str(); + sd_ctx_params.n_threads = n_threads; + sd_ctx_params.wtype = wtype; + sd_ctx_params.rng_type = rng_type; + sd_ctx_params.sampler_rng_type = sampler_rng_type; + sd_ctx_params.prediction = prediction; + sd_ctx_params.lora_apply_mode = lora_apply_mode; + sd_ctx_params.enable_mmap = enable_mmap; + sd_ctx_params.flash_attn = flash_attn; + sd_ctx_params.diffusion_flash_attn = diffusion_flash_attn; + sd_ctx_params.tae_preview_only = taesd_preview; + sd_ctx_params.diffusion_conv_direct = diffusion_conv_direct; + sd_ctx_params.vae_conv_direct = vae_conv_direct; + sd_ctx_params.circular_x = circular || circular_x; + sd_ctx_params.circular_y = circular || circular_y; + sd_ctx_params.force_sdxl_vae_conv_scale = force_sdxl_vae_conv_scale; + sd_ctx_params.chroma_use_dit_mask = chroma_use_dit_mask; + sd_ctx_params.chroma_use_t5_mask = chroma_use_t5_mask; + sd_ctx_params.chroma_t5_mask_pad = chroma_t5_mask_pad; + sd_ctx_params.qwen_image_zero_cond_t = qwen_image_zero_cond_t; + sd_ctx_params.vae_format = str_to_vae_format(vae_format); + sd_ctx_params.max_vram = max_vram.c_str(); + sd_ctx_params.stream_layers = stream_layers; + sd_ctx_params.eager_load = eager_load; + sd_ctx_params.backend = effective_backend.c_str(); + sd_ctx_params.params_backend = effective_params_backend.c_str(); + sd_ctx_params.rpc_servers = rpc_servers.c_str(); return sd_ctx_params; } @@ -823,54 +913,71 @@ ArgOptions SDGenerationParams::get_options() { {"-p", "--prompt", "the prompt to render", + 0, &prompt}, {"-n", "--negative-prompt", "the negative prompt (default: \"\")", + 0, &negative_prompt}, {"-i", "--init-img", "path to the init image", + 0, &init_image_path}, {"", "--end-img", "path to the end image, required by flf2v", + 0, &end_image_path}, {"", "--mask", "path to the mask image", + 0, &mask_image_path}, {"", "--control-image", "path to control image, control net", + 0, &control_image_path}, {"", "--control-video", "path to control video frames, It must be a directory path. The video frames inside should be stored as images in " "lexicographical (character) order. For example, if the control video path is `frames`, the directory contain images " "such as 00.png, 01.png, ... etc.", + 0, &control_video_path}, {"", "--pm-id-images-dir", "path to PHOTOMAKER input id images dir", + 0, &pm_id_images_dir}, {"", "--pm-id-embed-path", "path to PHOTOMAKER v2 id embed", + 0, &pm_id_embed_path}, + {"", + "--pulid-id-embedding", + "path to PuLID id embedding", + 0, + &pulid_id_embedding_path}, {"", "--hires-upscaler", "highres fix upscaler, Lanczos, Nearest, Latent, Latent (nearest), Latent (nearest-exact), " "Latent (antialiased), Latent (bicubic), Latent (bicubic antialiased), or a model name " "under --hires-upscalers-dir (default: Latent)", + 0, &hires_upscaler}, {"", "--extra-sample-args", - "extra sampler/scheduler/guidance args, key=value list. APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma", + "extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma;; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware", + (int)',', &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)", + (int)',', &extra_tiling_args}, }; @@ -1008,6 +1115,10 @@ ArgOptions SDGenerationParams::get_options() { "--pm-style-strength", "", &pm_style_strength}, + {"", + "--pulid-id-weight", + "strength of PuLID identity injection", + &pulid_id_weight}, {"", "--control-strength", "strength to apply Control Net (default: 0.9). 1.0 corresponds to full destruction of information in init image", @@ -1322,6 +1433,42 @@ ArgOptions SDGenerationParams::get_options() { return 1; }; + auto on_prompt_file_arg = [&](int argc, const char** argv, int index) { + if (++index >= argc) { + return -1; + } + const char* arg = argv[index]; + std::ifstream f(arg, std::ios::binary); + try { + prompt = std::string(std::istreambuf_iterator{f}, {}); + } catch (const std::ios_base::failure&) { + f.setstate(std::ios_base::failbit); + } + if (f.fail()) { + LOG_ERROR("error: failed to read prompt file '%s'\n", arg); + return -1; + } + return 1; + }; + + auto on_negative_prompt_file_arg = [&](int argc, const char** argv, int index) { + if (++index >= argc) { + return -1; + } + const char* arg = argv[index]; + std::ifstream f(arg, std::ios::binary); + try { + negative_prompt = std::string(std::istreambuf_iterator{f}, {}); + } catch (const std::ios_base::failure&) { + f.setstate(std::ios_base::failbit); + } + if (f.fail()) { + LOG_ERROR("error: failed to read negative prompt file '%s'\n", arg); + return -1; + } + return 1; + }; + options.manual_options = { {"-s", "--seed", @@ -1339,7 +1486,7 @@ ArgOptions SDGenerationParams::get_options() { on_high_noise_sample_method_arg}, {"", "--scheduler", - "denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple, kl_optimal, lcm, bong_tangent, ltx2], default: model-specific", + "denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, smoothstep, sgm_uniform, simple, kl_optimal, lcm, bong_tangent, ltx2, logit_normal, flux2, flux, beta], alias: normal=discrete, default: model-specific", on_scheduler_arg}, {"", "--sigmas", @@ -1385,6 +1532,14 @@ ArgOptions SDGenerationParams::get_options() { "--vae-relative-tile-size", "relative tile size for vae tiling, format [X]x[Y], in fraction of image size if < 1, in number of tiles per dim if >=1 (overrides --vae-tile-size)", on_relative_tile_size_arg}, + {"", + "--prompt-file", + "path to the file containing the prompt to render", + on_prompt_file_arg}, + {"", + "--negative-prompt-file", + "path to the file containing the negative prompt", + on_negative_prompt_file_arg}, }; @@ -2240,6 +2395,11 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() { pm_style_strength, }; + sd_pulid_params_t pulid_params = { + pulid_id_embedding_path.empty() ? nullptr : pulid_id_embedding_path.c_str(), + pulid_id_weight, + }; + params.loras = lora_vec.empty() ? nullptr : lora_vec.data(); params.lora_count = static_cast(lora_vec.size()); params.prompt = prompt.c_str(); @@ -2260,6 +2420,7 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() { params.control_image = control_image.get(); params.control_strength = control_strength; params.pm_params = pm_params; + params.pulid_params = pulid_params; params.vae_tiling_params = vae_tiling_params; params.cache = cache_params; @@ -2519,6 +2680,7 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params, set_json_basename_if_not_empty(models, "llm_vision", ctx_params.llm_vision_path); set_json_basename_if_not_empty(models, "diffusion_model", ctx_params.diffusion_model_path); set_json_basename_if_not_empty(models, "high_noise_diffusion_model", ctx_params.high_noise_diffusion_model_path); + set_json_basename_if_not_empty(models, "uncond_diffusion_model", ctx_params.uncond_diffusion_model_path); set_json_basename_if_not_empty(models, "vae", ctx_params.vae_path); set_json_basename_if_not_empty(models, "taesd", ctx_params.taesd_path); set_json_basename_if_not_empty(models, "control_net", ctx_params.control_net_path); @@ -2686,6 +2848,9 @@ std::string get_image_params(const SDContextParams& ctx_params, if (!ctx_params.diffusion_model_path.empty()) { parameter_string += "Unet: " + sd_basename(ctx_params.diffusion_model_path) + ", "; } + if (!ctx_params.uncond_diffusion_model_path.empty()) { + parameter_string += "Uncond Unet: " + sd_basename(ctx_params.uncond_diffusion_model_path) + ", "; + } if (!ctx_params.vae_path.empty()) { parameter_string += "VAE: " + sd_basename(ctx_params.vae_path) + ", "; } diff --git a/examples/common/common.h b/examples/common/common.h index cd02d212..e7c25015 100644 --- a/examples/common/common.h +++ b/examples/common/common.h @@ -31,6 +31,7 @@ struct StringOption { std::string short_name; std::string long_name; std::string desc; + int concat; std::string* target; }; @@ -56,11 +57,42 @@ struct BoolOption { bool* target; }; +struct ManualFunction { + std::function _func; + + ManualFunction() = default; + + ManualFunction(std::function func) + : _func(std::move(func)) { + } + + template + ManualFunction(F func) + : _func(make_function(func)) { + } + + int operator()(int argc, const char** argv, int index, bool& valid) const { + return _func(argc, argv, index, valid); + } + +private: + template + static std::function make_function(F func) { + if constexpr (std::is_invocable_v) { + return func; + } else { + return [func](int argc, const char** argv, int index, bool&) { + return func(argc, argv, index); + }; + } + } +}; + struct ManualOption { std::string short_name; std::string long_name; std::string desc; - std::function cb; + ManualFunction cb; }; struct ArgOptions { @@ -92,6 +124,7 @@ struct SDContextParams { std::string llm_vision_path; std::string diffusion_model_path; std::string high_noise_diffusion_model_path; + std::string uncond_diffusion_model_path; std::string embeddings_connectors_path; std::string vae_path; std::string vae_format = "auto"; @@ -101,6 +134,7 @@ struct SDContextParams { std::string control_net_path; std::string embedding_dir; std::string photo_maker_path; + std::string pulid_weights_path; sd_type_t wtype = SD_TYPE_COUNT; std::string tensor_type_rules; std::string lora_model_dir = "."; @@ -112,10 +146,14 @@ struct SDContextParams { rng_type_t rng_type = CUDA_RNG; rng_type_t sampler_rng_type = RNG_TYPE_COUNT; bool offload_params_to_cpu = false; - float max_vram = 0.f; + std::string max_vram = "0"; bool stream_layers = false; + bool eager_load = false; std::string backend; std::string params_backend; + std::string rpc_servers; + std::string effective_backend; + std::string effective_params_backend; bool enable_mmap = false; bool control_net_cpu = false; bool clip_on_cpu = false; @@ -143,11 +181,12 @@ struct SDContextParams { float flow_shift = INFINITY; ArgOptions get_options(); void build_embedding_map(); + void prepare_backend_assignments(); bool resolve(SDMode mode); bool validate(SDMode mode); bool resolve_and_validate(SDMode mode); std::string to_string() const; - sd_ctx_params_t to_sd_ctx_params_t(bool vae_decode_only, bool free_params_immediately, bool taesd_preview); + sd_ctx_params_t to_sd_ctx_params_t(bool taesd_preview); }; struct SDGenerationParams { @@ -198,6 +237,9 @@ struct SDGenerationParams { std::string pm_id_embed_path; float pm_style_strength = 20.f; + std::string pulid_id_embedding_path; + float pulid_id_weight = 1.0f; + int upscale_repeats = 1; int upscale_tile_size = 128; diff --git a/examples/server/README.md b/examples/server/README.md index d971f5fe..c24ed083 100644 --- a/examples/server/README.md +++ b/examples/server/README.md @@ -117,186 +117,10 @@ In this case, the server will load and serve the specified `index.html` file ins * using a custom UI * avoiding rebuilding the binary after frontend modifications -# Run +# Usage -``` -usage: ./bin/sd-server [options] - -Svr Options: - -l, --listen-ip server listen ip (default: 127.0.0.1) - --serve-html-path path to HTML file to serve at root (optional) - --listen-port server listen port (default: 1234) - -v, --verbose print extra info - --color colors the logging tags according to level - -h, --help show this help message and exit - -Context Options: - -m, --model path to full model - --clip_l path to the clip-l text encoder - --clip_g path to the clip-g text encoder - --clip_vision path to the clip-vision encoder - --t5xxl path to the t5xxl text encoder - --llm path to the llm text encoder. For example: (qwenvl2.5 for qwen-image, - mistral-small3.2 for flux2, ...) - --llm_vision path to the llm vit - --qwen2vl alias of --llm. Deprecated. - --qwen2vl_vision alias of --llm_vision. Deprecated. - --diffusion-model path to the standalone diffusion model - --high-noise-diffusion-model path to the standalone high noise diffusion model - --vae path to standalone vae model - --taesd path to taesd. Using Tiny AutoEncoder for fast decoding (low quality) - --tae alias of --taesd - --control-net path to control net model - --embd-dir embeddings directory - --lora-model-dir lora model directory - --hires-upscalers-dir highres fix upscaler model directory - --tensor-type-rules weight type per tensor pattern (example: "^vae\.=f16,model\.=q8_0") - --photo-maker path to PHOTOMAKER model - --upscale-model path to esrgan model. - -t, --threads number of threads to use during computation (default: -1). If threads <= 0, - then threads will be set to the number of CPU physical cores - --chroma-t5-mask-pad t5 mask pad size of chroma - --max-vram 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) - --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 - --mmap whether to memory-map model - --control-net-cpu keep controlnet in cpu (for low vram) - --clip-on-cpu keep clip in cpu (for low vram) - --vae-on-cpu keep vae in cpu (for low vram) - --fa use flash attention - --diffusion-fa use flash attention in the diffusion model only - --diffusion-conv-direct use ggml_conv2d_direct in the diffusion model - --vae-conv-direct use ggml_conv2d_direct in the vae model - --circular enable circular padding for convolutions - --circularx enable circular RoPE wrapping on x-axis (width) only - --circulary enable circular RoPE wrapping on y-axis (height) only - --chroma-disable-dit-mask disable dit mask for chroma - --qwen-image-zero-cond-t enable zero_cond_t for qwen image - --chroma-enable-t5-mask enable t5 mask for chroma - --type weight type (examples: f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0, q2_K, q3_K, - q4_K). If not specified, the default is the type of the weight file - --rng RNG, one of [std_default, cuda, cpu], default: cuda(sd-webui), cpu(comfyui) - --sampler-rng sampler RNG, one of [std_default, cuda, cpu]. If not specified, use --rng - --prediction prediction type override, one of [eps, v, edm_v, sd3_flow, flux_flow, - flux2_flow] - --lora-apply-mode the way to apply LoRA, one of [auto, immediately, at_runtime], default is - auto. In auto mode, if the model weights contain any quantized parameters, - the at_runtime mode will be used; otherwise, immediately will be used.The - immediately mode may have precision and compatibility issues with quantized - parameters, but it usually offers faster inference speed and, in some cases, - lower memory usage. The at_runtime mode, on the other hand, is exactly the - opposite. - -Default Generation Options: - -p, --prompt the prompt to render - -n, --negative-prompt the negative prompt (default: "") - -i, --init-img path to the init image - --end-img path to the end image, required by flf2v - --mask path to the mask image - --control-image path to control image, control net - --control-video path to control video frames, It must be a directory path. The video frames - inside should be stored as images in lexicographical (character) order. For - example, if the control video path is `frames`, the directory contain images - such as 00.png, 01.png, ... etc. - --pm-id-images-dir path to PHOTOMAKER input id images dir - --pm-id-embed-path path to PHOTOMAKER v2 id embed - --hires-upscaler highres fix upscaler, Lanczos, Nearest, Latent, Latent (nearest), Latent - (nearest-exact), Latent (antialiased), Latent (bicubic), Latent (bicubic - antialiased), or a model name under --hires-upscalers-dir (default: Latent) - --extra-sample-args extra sampler/scheduler/guidance args, key=value list. APG supports apg_eta, - apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports - slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; - ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma - --extra-tiling-args extra VAE tiling args, key=value list. LTX video VAE supports - temporal_tile_frames (default: 4), temporal_tile_overlap (default: 1) - -H, --height image height, in pixel space (default: 512) - -W, --width image width, in pixel space (default: 512) - --steps number of sample steps (default: 20) - --high-noise-steps (high noise) number of sample steps (default: -1 = auto) - --clip-skip ignore last layers of CLIP network; 1 ignores none, 2 ignores one layer - (default: -1). <= 0 represents unspecified, will be 1 for SD1.x, 2 for SD2.x - -b, --batch-count batch count - --video-frames video frames (default: 1) - --fps fps (default: 24) - --timestep-shift shift timestep for NitroFusion models (default: 0). recommended N for - NitroSD-Realism around 250 and 500 for NitroSD-Vibrant - --upscale-repeats Run the ESRGAN upscaler this many times (default: 1) - --upscale-tile-size tile size for ESRGAN upscaling (default: 128) - --hires-width highres fix target width, 0 to use --hires-scale (default: 0) - --hires-height highres fix target height, 0 to use --hires-scale (default: 0) - --hires-steps highres fix second pass sample steps, 0 to reuse --steps (default: 0) - --hires-upscale-tile-size highres fix upscaler tile size, reserved for model-backed upscalers (default: - 128) - --cfg-scale unconditional guidance scale: (default: 7.0) - --img-cfg-scale image guidance scale for inpaint or image edit models: (default: same as - --cfg-scale) - --guidance distilled guidance scale for models with guidance input (default: 3.5) - --slg-scale skip layer guidance (SLG) scale, only for DiT models: (default: 0). 0 means - disabled, a value of 2.5 is nice for sd3.5 medium - --skip-layer-start SLG enabling point (default: 0.01) - --skip-layer-end SLG disabling point (default: 0.2) - --eta noise multiplier (default: 0 for ddim_trailing, tcd, res_multistep and - res_2s; 1 for euler_a, er_sde and dpm++2s_a) - --flow-shift shift value for Flow models like SD3.x or WAN (default: auto) - --high-noise-cfg-scale (high noise) unconditional guidance scale: (default: 7.0) - --high-noise-img-cfg-scale (high noise) image guidance scale for inpaint or image edit models (default: - same as --cfg-scale) - --high-noise-guidance (high noise) distilled guidance scale for models with guidance input - (default: 3.5) - --high-noise-slg-scale (high noise) skip layer guidance (SLG) scale, only for DiT models: (default: - 0) - --high-noise-skip-layer-start (high noise) SLG enabling point (default: 0.01) - --high-noise-skip-layer-end (high noise) SLG disabling point (default: 0.2) - --high-noise-eta (high noise) noise multiplier (default: 0 for ddim_trailing, tcd, - res_multistep and res_2s; 1 for euler_a, er_sde and dpm++2s_a) - --strength strength for noising/unnoising (default: 0.75) - --pm-style-strength - --control-strength strength to apply Control Net (default: 0.9). 1.0 corresponds to full - destruction of information in init image - --moe-boundary timestep boundary for Wan2.2 MoE model. (default: 0.875). Only enabled if - `--high-noise-steps` is set to -1 - --vace-strength wan vace strength - --vae-tile-overlap tile overlap for vae tiling, in fraction of tile size (default: 0.5) - --hires-scale highres fix scale when target size is not set (default: 2.0) - --hires-denoising-strength highres fix second pass denoising strength (default: 0.7) - --increase-ref-index automatically increase the indices of references images based on the order - they are listed (starting with 1). - --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, - dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, res_2s, - er_sde, euler_cfg_pp, euler_a_cfg_pp] (default: euler for Flux/SD3/Wan, euler_a otherwise) - --high-noise-sampling-method (high noise) sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, - dpm++2m, dpm++2mv2, ipndm, ipndm_v, lcm, ddim_trailing, tcd, res_multistep, - res_2s, er_sde, euler_cfg_pp, euler_a_cfg_pp] default: euler for Flux/SD3/Wan, euler_a otherwise - --scheduler denoiser sigma scheduler, one of [discrete, karras, exponential, ays, gits, - smoothstep, sgm_uniform, simple, kl_optimal, lcm, bong_tangent, ltx2], default: - 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) - --cache-mode caching method: 'easycache' (DiT), 'ucache' (UNET), - 'dbcache'/'taylorseer'/'cache-dit' (DiT block-level), 'spectrum' (UNET/DiT - Chebyshev+Taylor forecasting) - --cache-option named cache params (key=value format, comma-separated). easycache/ucache: - threshold=,start=,end=,decay=,relative=,reset=; dbcache/taylorseer/cache-dit: - Fn=,Bn=,threshold=,warmup=; spectrum: w=,m=,lam=,window=,flex=,warmup=,stop=. - Examples: "threshold=0.25" or "threshold=1.5,reset=0" - --scm-mask SCM steps mask for cache-dit: comma-separated 0/1 (e.g., - "1,1,1,0,0,1,0,0,1,0") - 1=compute, 0=can cache - --scm-policy SCM policy: 'dynamic' (default) or 'static' - --vae-tile-size tile size for vae tiling, format [X]x[Y] (default: 32x32) - --vae-relative-tile-size relative tile size for vae tiling, format [X]x[Y], in fraction of image size - if < 1, in number of tiles per dim if >=1 (overrides --vae-tile-size) +For detailed command-line arguments, run: + +```bash +./bin/sd-server -h ``` diff --git a/examples/server/async_jobs.cpp b/examples/server/async_jobs.cpp index 5083ce01..1933dec1 100644 --- a/examples/server/async_jobs.cpp +++ b/examples/server/async_jobs.cpp @@ -174,9 +174,13 @@ bool execute_img_gen_job(ServerRuntime& runtime, { std::lock_guard lock(*runtime.sd_ctx_mutex); - int result_count = sd_get_image_result_count(runtime.sd_ctx, ¶ms); - sd_image_t* raw_results = generate_image(runtime.sd_ctx, ¶ms); - results.adopt(raw_results, result_count); + sd_image_t* raw_results = nullptr; + int num_results = 0; + if (!generate_image(runtime.sd_ctx, ¶ms, &raw_results, &num_results)) { + raw_results = nullptr; + num_results = 0; + } + results.adopt(raw_results, num_results); } const int num_results = results.count(); diff --git a/examples/server/frontend b/examples/server/frontend index 797ccf80..c4bce3d6 160000 --- a/examples/server/frontend +++ b/examples/server/frontend @@ -1 +1 @@ -Subproject commit 797ccf80825cc035508ba9b599b2a21953e7f835 +Subproject commit c4bce3d6b3f236614cca21014f076083b7270ba8 diff --git a/examples/server/main.cpp b/examples/server/main.cpp index 32d570d6..dce35c11 100644 --- a/examples/server/main.cpp +++ b/examples/server/main.cpp @@ -85,7 +85,7 @@ int main(int argc, const char** argv) { LOG_DEBUG("%s", ctx_params.to_string().c_str()); LOG_DEBUG("%s", default_gen_params.to_string().c_str()); - sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(false, false, false); + sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(false); SDCtxPtr sd_ctx(new_sd_ctx(&sd_ctx_params)); if (sd_ctx == nullptr) { diff --git a/examples/server/routes_openai.cpp b/examples/server/routes_openai.cpp index 554c9428..7386363c 100644 --- a/examples/server/routes_openai.cpp +++ b/examples/server/routes_openai.cpp @@ -229,8 +229,11 @@ static bool execute_sync_img_gen_request(ServerRuntime& runtime, { std::lock_guard lock(*runtime.sd_ctx_mutex); - sd_image_t* raw_results = generate_image(runtime.sd_ctx, &img_gen_params); - num_results = sd_get_image_result_count(runtime.sd_ctx, &img_gen_params); + sd_image_t* raw_results = nullptr; + if (!generate_image(runtime.sd_ctx, &img_gen_params, &raw_results, &num_results)) { + raw_results = nullptr; + num_results = 0; + } results.adopt(raw_results, num_results); } diff --git a/examples/server/routes_sdapi.cpp b/examples/server/routes_sdapi.cpp index 9a172d8c..cdc76f00 100644 --- a/examples/server/routes_sdapi.cpp +++ b/examples/server/routes_sdapi.cpp @@ -292,8 +292,11 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) { { std::lock_guard lock(*runtime->sd_ctx_mutex); - sd_image_t* raw_results = generate_image(runtime->sd_ctx, &img_gen_params); - num_results = sd_get_image_result_count(runtime->sd_ctx, &img_gen_params); + sd_image_t* raw_results = nullptr; + if (!generate_image(runtime->sd_ctx, &img_gen_params, &raw_results, &num_results)) { + raw_results = nullptr; + num_results = 0; + } results.adopt(raw_results, num_results); } @@ -439,6 +442,9 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) { scheduler_names.push_back("default"); for (int i = 0; i < SCHEDULER_COUNT; i++) { scheduler_names.push_back(sd_scheduler_name((scheduler_t)i)); + if (i == DISCRETE_SCHEDULER) { + scheduler_names.push_back("normal"); + } } json r = json::array(); for (auto name : scheduler_names) { diff --git a/examples/server/routes_sdcpp.cpp b/examples/server/routes_sdcpp.cpp index c60e9da6..cbd47a88 100644 --- a/examples/server/routes_sdcpp.cpp +++ b/examples/server/routes_sdcpp.cpp @@ -219,6 +219,9 @@ static json make_capabilities_json(ServerRuntime& runtime) { for (int i = 0; i < SCHEDULER_COUNT; ++i) { schedulers.push_back(sd_scheduler_name((scheduler_t)i)); + if (i == DISCRETE_SCHEDULER) { + schedulers.push_back("normal"); + } } { diff --git a/examples/server/runtime.cpp b/examples/server/runtime.cpp index 1127ab73..1fb41c71 100644 --- a/examples/server/runtime.cpp +++ b/examples/server/runtime.cpp @@ -190,8 +190,8 @@ ArgOptions SDSvrParams::get_options() { ArgOptions options; options.string_options = { - {"-l", "--listen-ip", "server listen ip (default: 127.0.0.1)", &listen_ip}, - {"", "--serve-html-path", "path to HTML file to serve at root (optional)", &serve_html_path}, + {"-l", "--listen-ip", "server listen ip (default: 127.0.0.1)", 0, &listen_ip}, + {"", "--serve-html-path", "path to HTML file to serve at root (optional)", 0, &serve_html_path}, }; options.int_options = { @@ -203,8 +203,9 @@ ArgOptions SDSvrParams::get_options() { {"", "--color", "colors the logging tags according to level", true, &color}, }; - auto on_help_arg = [&](int, const char**, int) { + auto on_help_arg = [&](int, const char**, int, bool& valid) { normal_exit = true; + valid = true; return -1; }; diff --git a/format-code.ps1 b/format-code.ps1 new file mode 100644 index 00000000..9e55cb2e --- /dev/null +++ b/format-code.ps1 @@ -0,0 +1,54 @@ +$patterns = @( + "src/*.cpp" + "src/*.h" + "src/*.hpp" + "src/conditioning/*.cpp" + "src/conditioning/*.h" + "src/conditioning/*.hpp" + "src/core/*.cpp" + "src/core/*.h" + "src/core/*.hpp" + "src/extensions/*.cpp" + "src/extensions/*.h" + "src/extensions/*.hpp" + "src/runtime/*.cpp" + "src/runtime/*.h" + "src/runtime/*.hpp" + "src/model/*/*.cpp" + "src/model/*/*.h" + "src/model/*/*.hpp" + "src/tokenizers/*.h" + "src/tokenizers/*.cpp" + "src/tokenizers/vocab/*.h" + "src/tokenizers/vocab/*.cpp" + "src/model_io/*.h" + "src/model_io/*.cpp" + "examples/cli/*.cpp" + "examples/cli/*.h" + "examples/server/*.cpp" + "examples/common/*.hpp" + "examples/common/*.h" + "examples/common/*.cpp" +) + +$root = (Get-Location).Path + +foreach ($pattern in $patterns) { + $files = Get-ChildItem -Path $pattern -File -ErrorAction SilentlyContinue | Sort-Object FullName + + foreach ($file in $files) { + $relativePath = $file.FullName.Substring($root.Length).TrimStart('\', '/') -replace '\\', '/' + + if ($relativePath -like "vocab*") { + continue + } + + Write-Host "formatting '$relativePath'" + + # if ($relativePath -ne "stable-diffusion.h") { + # clang-tidy -fix -p build_linux/ "$relativePath" + # } + + & clang-format -style=file -i $relativePath + } +} diff --git a/format-code.sh b/format-code.sh index 8aa422bc..1ed77a5b 100644 --- a/format-code.sh +++ b/format-code.sh @@ -1,10 +1,17 @@ -for f in src/*.cpp src/*.h src/*.hpp src/tokenizers/*.h src/tokenizers/*.cpp src/tokenizers/vocab/*.h src/tokenizers/vocab/*.cpp \ +for f in src/*.cpp src/*.h src/*.hpp \ + src/conditioning/*.cpp src/conditioning/*.h src/conditioning/*.hpp \ + src/core/*.cpp src/core/*.h src/core/*.hpp \ + src/extensions/*.cpp src/extensions/*.h src/extensions/*.hpp \ + src/runtime/*.cpp src/runtime/*.h src/runtime/*.hpp \ + src/model/*/*.cpp src/model/*/*.h src/model/*/*.hpp \ + src/tokenizers/*.h src/tokenizers/*.cpp src/tokenizers/vocab/*.h src/tokenizers/vocab/*.cpp \ src/model_io/*.h src/model_io/*.cpp examples/cli/*.cpp examples/cli/*.h examples/server/*.cpp \ examples/common/*.hpp examples/common/*.h examples/common/*.cpp; do + [[ -e "$f" ]] || continue [[ "$f" == vocab* ]] && continue echo "formatting '$f'" # if [ "$f" != "stable-diffusion.h" ]; then # clang-tidy -fix -p build_linux/ "$f" # fi clang-format -style=file -i "$f" -done \ No newline at end of file +done diff --git a/ggml b/ggml index 0ce7ad34..eced84c8 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit 0ce7ad348a3151e1da9f65d962044546bcaad421 +Subproject commit eced84c86f8b012c752c016f7fe789adea168e1e diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index 917a994d..54231485 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -70,6 +70,10 @@ enum scheduler_t { LCM_SCHEDULER, BONG_TANGENT_SCHEDULER, LTX2_SCHEDULER, + LOGIT_NORMAL_SCHEDULER, + FLUX2_SCHEDULER, + FLUX_SCHEDULER, + BETA_SCHEDULER, SCHEDULER_COUNT }; @@ -79,7 +83,8 @@ enum prediction_t { EDM_V_PRED, FLOW_PRED, FLUX_FLOW_PRED, - FLUX2_FLOW_PRED, + SEFI_FLOW_PRED, + MINIT2I_FLOW_PRED, PREDICTION_COUNT }; @@ -186,6 +191,7 @@ typedef struct { const char* llm_vision_path; const char* diffusion_model_path; const char* high_noise_diffusion_model_path; + const char* uncond_diffusion_model_path; const char* embeddings_connectors_path; const char* vae_path; const char* audio_vae_path; @@ -194,20 +200,15 @@ typedef struct { const sd_embedding_t* embeddings; uint32_t embedding_count; const char* photo_maker_path; + const char* pulid_weights_path; const char* tensor_type_rules; - bool vae_decode_only; - bool free_params_immediately; int n_threads; enum sd_type_t wtype; enum rng_type_t rng_type; enum rng_type_t sampler_rng_type; enum prediction_t prediction; enum lora_apply_mode_t lora_apply_mode; - bool offload_params_to_cpu; bool enable_mmap; - bool keep_clip_on_cpu; - bool keep_control_net_on_cpu; - bool keep_vae_on_cpu; bool flash_attn; bool diffusion_flash_attn; bool tae_preview_only; @@ -221,10 +222,12 @@ typedef struct { int chroma_t5_mask_pad; bool qwen_image_zero_cond_t; enum sd_vae_format_t vae_format; - float max_vram; // GiB budget for graph-cut segmented param offload (0 = disabled, -1 = auto free VRAM minus 1 GiB) + const char* max_vram; // GiB budget or backend assignment spec for graph-cut segmented param offload (0 = disabled, -1 = auto) bool stream_layers; // Enable residency+prefetch streaming on top of --max-vram (no effect without --max-vram) + bool eager_load; // Load all params into the params backend at model-load time instead of lazily on first use const char* backend; const char* params_backend; + const char* rpc_servers; } sd_ctx_params_t; typedef struct { @@ -276,6 +279,11 @@ typedef struct { float style_strength; } sd_pm_params_t; // photo maker +typedef struct { + const char* id_embedding_path; + float id_weight; +} sd_pulid_params_t; + enum sd_cache_mode_t { SD_CACHE_DISABLED = 0, SD_CACHE_EASYCACHE, @@ -368,6 +376,7 @@ typedef struct { sd_image_t control_image; float control_strength; sd_pm_params_t pm_params; + sd_pulid_params_t pulid_params; sd_tiling_params_t vae_tiling_params; sd_cache_params_t cache; sd_hires_params_t hires; @@ -399,14 +408,17 @@ typedef struct { } sd_vid_gen_params_t; typedef struct sd_ctx_t sd_ctx_t; +struct ggml_tensor; typedef void (*sd_log_cb_t)(enum sd_log_level_t level, const char* text, void* data); typedef void (*sd_progress_cb_t)(int step, int steps, float time, void* data); typedef void (*sd_preview_cb_t)(int step, int frame_count, sd_image_t* frames, bool is_noisy, void* data); +typedef bool (*sd_graph_eval_callback_t)(struct ggml_tensor* t, bool ask, void* user_data); SD_API void sd_set_log_callback(sd_log_cb_t sd_log_cb, void* data); SD_API void sd_set_progress_callback(sd_progress_cb_t cb, void* data); SD_API void sd_set_preview_callback(sd_preview_cb_t cb, enum preview_t mode, int interval, bool denoised, bool noisy, void* data); +SD_API void sd_set_backend_eval_callback(sd_graph_eval_callback_t cb, void* data); SD_API int32_t sd_get_num_physical_cores(); SD_API const char* sd_get_system_info(); SD_API bool sd_ctx_supports_image_generation(const sd_ctx_t* sd_ctx); @@ -447,8 +459,21 @@ SD_API enum scheduler_t sd_get_default_scheduler(const sd_ctx_t* sd_ctx, enum sa SD_API void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params); SD_API char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params); -SD_API int32_t sd_get_image_result_count(const sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_gen_params); -SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_gen_params); +SD_API bool generate_image(sd_ctx_t* sd_ctx, + const sd_img_gen_params_t* sd_img_gen_params, + sd_image_t** images_out, + int* num_images_out); + +enum sd_cancel_mode_t { + // Stop the current generation as soon as possible. + SD_CANCEL_ALL, + // Finish the current image sample, then skip additional batch latents and return completed images. + SD_CANCEL_NEW_LATENTS, + // Clear a pending cancellation request. + SD_CANCEL_RESET +}; + +SD_API void sd_cancel_generation(sd_ctx_t* sd_ctx, enum sd_cancel_mode_t mode); SD_API void sd_vid_gen_params_init(sd_vid_gen_params_t* sd_vid_gen_params); SD_API bool generate_video(sd_ctx_t* sd_ctx, @@ -460,7 +485,6 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, typedef struct upscaler_ctx_t upscaler_ctx_t; SD_API upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path, - bool offload_params_to_cpu, bool direct, int n_threads, int tile_size, @@ -468,9 +492,11 @@ SD_API upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path, const char* params_backend); SD_API void free_upscaler_ctx(upscaler_ctx_t* upscaler_ctx); -SD_API sd_image_t upscale(upscaler_ctx_t* upscaler_ctx, - sd_image_t input_image, - uint32_t upscale_factor); +SD_API bool upscale(upscaler_ctx_t* upscaler_ctx, + sd_image_t input_image, + uint32_t upscale_factor, + sd_image_t** images_out, + int* num_images_out); SD_API int get_upscale_factor(upscaler_ctx_t* upscaler_ctx); @@ -481,6 +507,17 @@ SD_API bool convert(const char* input_path, const char* tensor_type_rules, bool convert_name); +SD_API bool convert_with_components(const char* model_path, + const char* clip_l_path, + const char* clip_g_path, + const char* t5xxl_path, + const char* diffusion_model_path, + const char* vae_path, + const char* output_path, + enum sd_type_t output_type, + const char* tensor_type_rules, + bool convert_name); + SD_API bool preprocess_canny(sd_image_t image, float high_threshold, float low_threshold, @@ -488,9 +525,18 @@ SD_API bool preprocess_canny(sd_image_t image, float strong, bool inverse); +SD_API bool load_imatrix(const char* imatrix_path); +SD_API void save_imatrix(const char* imatrix_path); +SD_API void enable_imatrix_collection(void); +SD_API void disable_imatrix_collection(void); + SD_API const char* sd_commit(void); SD_API const char* sd_version(void); +// for C API, caller needs to call free_sd_images to free the memory after use +// This helps avoid CRT problems on Windows when memory is allocated in the library but freed in the caller, which may use a different CRT. +SD_API void free_sd_images(sd_image_t* result_images, int num_images); + #ifdef __cplusplus } #endif diff --git a/script/convert_fp8_scale_to_bf16.py b/script/convert_fp8_scale_to_bf16.py new file mode 100644 index 00000000..a3eb2acc --- /dev/null +++ b/script/convert_fp8_scale_to_bf16.py @@ -0,0 +1,283 @@ +#!/usr/bin/env python +import argparse +import json +import math +import os +import struct +from collections import Counter +from pathlib import Path + +import torch +from safetensors import safe_open + + +FLOAT_DTYPES = { + "BF16", + "F16", + "F32", + "F64", + "F8_E4M3", + "F8_E4M3FN", + "F8_E5M2", +} + +FP8_DTYPES = { + "F8_E4M3", + "F8_E4M3FN", + "F8_E5M2", +} + +DTYPE_SIZES = { + "BOOL": 1, + "U8": 1, + "I8": 1, + "F8_E4M3": 1, + "F8_E4M3FN": 1, + "F8_E5M2": 1, + "U16": 2, + "I16": 2, + "F16": 2, + "BF16": 2, + "U32": 4, + "I32": 4, + "F32": 4, + "U64": 8, + "I64": 8, + "F64": 8, +} + + +def read_safetensors_header(path: Path): + with path.open("rb") as f: + header_len = struct.unpack(" 0 and scale.ndim == 1: + if first_dim_end is not None and scale.shape[0] >= first_dim_end: + scale = scale[first_dim_start:first_dim_end] + if scale.shape[0] == chunk.shape[0]: + return scale.reshape((scale.shape[0],) + (1,) * (chunk.ndim - 1)) + + return scale + + +def write_scaled_fp8_weight(out, weight, scale, chunk_rows): + if weight.ndim == 0: + result = weight.to(torch.float32) * scale_view_for_chunk(scale, weight) + write_tensor_bytes(out, result.to(torch.bfloat16)) + return + + rows = weight.shape[0] + for start in range(0, rows, chunk_rows): + end = min(start + chunk_rows, rows) + chunk = weight[start:end].to(torch.float32) + scale_view = scale_view_for_chunk(scale, chunk, start, end) + result = chunk * scale_view + write_tensor_bytes(out, result.to(torch.bfloat16)) + + +def write_float_as_bf16(out, tensor, chunk_rows): + if tensor.dtype == torch.bfloat16: + write_tensor_bytes(out, tensor) + return + + if tensor.ndim == 0: + write_tensor_bytes(out, tensor.to(torch.bfloat16)) + return + + rows = tensor.shape[0] + for start in range(0, rows, chunk_rows): + end = min(start + chunk_rows, rows) + write_tensor_bytes(out, tensor[start:end].to(torch.bfloat16)) + + +def convert(input_path: Path, output_path: Path, chunk_rows: int, dry_run: bool): + header = read_safetensors_header(input_path) + plan, output_header, data_size = build_output_plan(header) + + source_counts = Counter(item["source_dtype"] for item in plan) + output_counts = Counter(item["output_dtype"] for item in plan) + scaled_count = sum(item["mode"] == "fp8_scaled_weight" for item in plan) + dropped_scales = sum(item["mode"] == "fp8_scaled_weight" for item in plan) + header_bytes = json.dumps(output_header, separators=(",", ":")).encode("utf-8") + expected_size = 8 + len(header_bytes) + data_size + + print(f"input: {input_path}") + print(f"output: {output_path}") + print(f"tensors written: {len(plan)}") + print(f"scaled fp8 weights dequantized: {scaled_count}") + print(f"weight_scale tensors dropped: {dropped_scales}") + print(f"source dtypes: {dict(sorted(source_counts.items()))}") + print(f"output dtypes: {dict(sorted(output_counts.items()))}") + print(f"expected output size: {expected_size / (1024 ** 3):.2f} GiB") + + if dry_run: + return + + if output_path.exists(): + raise FileExistsError(f"{output_path} already exists; pass --overwrite to replace it") + + tmp_path = output_path.with_suffix(output_path.suffix + ".tmp") + if tmp_path.exists(): + raise FileExistsError(f"{tmp_path} already exists; remove it or choose another output") + + with safe_open(str(input_path), framework="pt", device="cpu") as sf, tmp_path.open("wb") as out: + out.write(struct.pack(" {item['output_dtype']}") + + tensor = sf.get_tensor(name) + if item["mode"] == "fp8_scaled_weight": + scale = sf.get_tensor(item["scale_key"]) + write_scaled_fp8_weight(out, tensor, scale, chunk_rows) + elif item["mode"] == "float_to_bf16": + write_float_as_bf16(out, tensor, chunk_rows) + else: + write_tensor_bytes(out, tensor) + + actual_size = out.tell() + + if actual_size != expected_size: + tmp_path.unlink(missing_ok=True) + raise RuntimeError(f"wrote {actual_size} bytes, expected {expected_size} bytes") + + tmp_path.replace(output_path) + print("done") + + +def main(): + parser = argparse.ArgumentParser( + description="Convert an fp8 safetensors checkpoint with weight_scale tensors to bf16." + ) + parser.add_argument("--input", default="ideogram4_fp8.safetensors", type=Path) + parser.add_argument("--output", default="ideogram4_bf16.safetensors", type=Path) + parser.add_argument("--chunk-rows", default=1024, type=int) + parser.add_argument("--dry-run", action="store_true") + parser.add_argument("--overwrite", action="store_true") + args = parser.parse_args() + + input_path = args.input.resolve() + output_path = args.output.resolve() + + if args.chunk_rows < 1: + raise ValueError("--chunk-rows must be >= 1") + if not input_path.exists(): + raise FileNotFoundError(input_path) + if args.overwrite and output_path.exists(): + output_path.unlink() + + convert(input_path, output_path, args.chunk_rows, args.dry_run) + + +if __name__ == "__main__": + main() diff --git a/script/convert_qwen3_vl.py b/script/convert_qwen3_vl.py new file mode 100644 index 00000000..d34aeb5f --- /dev/null +++ b/script/convert_qwen3_vl.py @@ -0,0 +1,112 @@ +#!/usr/bin/env python3 +"""Convert a Qwen3-VL HF safetensors checkpoint into a sd.cpp-loadable form. + +The HF dump prefixes text-tower keys with ``model.language_model.`` and +vision-tower keys with ``model.visual.``. sd.cpp expects ``model.`` for +the text side; the vision side is converted by sd.cpp's own +``convert_qwen3_vl_vision_name`` and is left as-is here. + +Operates on raw safetensors bytes so any dtype (BF16/F16/F32) is preserved. + +Usage: + python3 script/convert_qwen3_vl.py +""" + +import argparse +import json +import os +import struct +import sys + + +def rewrite_key(key: str) -> str: + if key.startswith("model.language_model."): + return "model." + key[len("model.language_model."):] + return key + + +def read_safetensors_header(path: str): + with open(path, "rb") as f: + hdr_len = struct.unpack(" 0: + chunk = src.read(min(8 * 1024 * 1024, remaining)) + if not chunk: + raise IOError(f"Truncated tensor in {shard_path}") + out.write(chunk) + remaining -= len(chunk) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("input", help="HF Qwen3-VL directory or single safetensors file") + parser.add_argument("output", help="Output single safetensors path") + args = parser.parse_args() + + entries = stage_tensors(args.input) + print(f"Tensors: {len(entries)}") + print(f"Writing -> {args.output}") + os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True) + write_consolidated(args.output, entries) + print(f"Done. Output size: {os.path.getsize(args.output) / 1e9:.2f} GB") + + +if __name__ == "__main__": + main() diff --git a/script/convert_sefi.py b/script/convert_sefi.py new file mode 100644 index 00000000..5dbfac49 --- /dev/null +++ b/script/convert_sefi.py @@ -0,0 +1,279 @@ +#!/usr/bin/env python3 +"""Convert a SeFi-Image diffusers checkpoint into a single sd.cpp-compatible safetensors. + +Operates on raw safetensors bytes so any dtype (BF16, F32, ...) is preserved exactly. +No numpy or torch dependency required. + +Usage: + python3 script/convert_sefi.py +""" + +import argparse +import json +import os +import re +import struct +import sys + + +_LINEAR_TO_LIN = re.compile(r"\.linear\.") +_SHARED_MOD_PREFIXES = ( + "double_stream_modulation_img", + "double_stream_modulation_txt", + "single_stream_modulation", +) + + +def rewrite_transformer_key(key: str) -> str: + if key.startswith("backbone."): + key = key[len("backbone."):] + elif key.startswith("dual_time_embed."): + return key + + if any(key.startswith(prefix + ".") for prefix in _SHARED_MOD_PREFIXES): + key = _LINEAR_TO_LIN.sub(".lin.", key, count=1) + + if key == "context_embedder.weight": + return "txt_in.weight" + if key == "context_embedder.bias": + return "txt_in.bias" + if key == "x_embedder.weight": + return "img_in.weight" + if key == "x_embedder.bias": + return "img_in.bias" + + if key == "proj_out.weight": + return "final_layer.linear.weight" + if key == "proj_out.bias": + return "final_layer.linear.bias" + if key == "norm_out.linear.weight": + return "final_layer.adaLN_modulation.1.weight" + if key == "norm_out.linear.bias": + return "final_layer.adaLN_modulation.1.bias" + + m = re.match(r"transformer_blocks\.(\d+)\.(.*)$", key) + if m: + return _rewrite_double_stream(m.group(1), m.group(2)) + m = re.match(r"single_transformer_blocks\.(\d+)\.(.*)$", key) + if m: + return _rewrite_single_stream(m.group(1), m.group(2)) + + return key + + +def _rewrite_double_stream(idx: str, tail: str) -> str: + dst = f"double_blocks.{idx}." + mapping = { + "norm1.linear.weight": "img_mod.lin.weight", + "norm1_context.linear.weight": "txt_mod.lin.weight", + "attn.norm_q.weight": "img_attn.norm.query_norm.scale", + "attn.norm_k.weight": "img_attn.norm.key_norm.scale", + "attn.norm_added_q.weight": "txt_attn.norm.query_norm.scale", + "attn.norm_added_k.weight": "txt_attn.norm.key_norm.scale", + "attn.to_out.0.weight": "img_attn.proj.weight", + "attn.to_add_out.weight": "txt_attn.proj.weight", + "ff.net.0.proj.weight": "img_mlp.0.weight", + "ff.net.2.weight": "img_mlp.2.weight", + "ff_context.net.0.proj.weight": "txt_mlp.0.weight", + "ff_context.net.2.weight": "txt_mlp.2.weight", + "ff.linear_in.weight": "img_mlp.0.weight", + "ff.linear_out.weight": "img_mlp.2.weight", + "ff_context.linear_in.weight": "txt_mlp.0.weight", + "ff_context.linear_out.weight": "txt_mlp.2.weight", + } + return dst + mapping.get(tail, tail) + + +# QKV triplets to fuse on output: source tails -> target fused tail. +# Each tuple is (q_tail, k_tail, v_tail, fused_target_tail). +QKV_DOUBLE_TRIPLETS = [ + ("attn.to_q.weight", "attn.to_k.weight", "attn.to_v.weight", "img_attn.qkv.weight"), + ("attn.add_q_proj.weight", "attn.add_k_proj.weight", "attn.add_v_proj.weight", "txt_attn.qkv.weight"), +] + + +def _rewrite_single_stream(idx: str, tail: str) -> str: + dst = f"single_blocks.{idx}." + mapping = { + "norm.linear.weight": "modulation.lin.weight", + "attn.norm_q.weight": "norm.query_norm.scale", + "attn.norm_k.weight": "norm.key_norm.scale", + "attn.to_qkv_mlp_proj.weight": "linear1.weight", + "attn.to_out.weight": "linear2.weight", + } + return dst + mapping.get(tail, tail) + + + + +def read_safetensors_header(path: str): + """Return (header dict, data start byte offset).""" + with open(path, "rb") as f: + hdr_len = struct.unpack(" {tail: (key, shard_path, data_off, info)} + raw_others = [] + for shard_path in collect_shard_paths(section_dir, "diffusion_pytorch_model"): + hdr, data_off = read_safetensors_header(shard_path) + for key, info in hdr.items(): + if key == "__metadata__": + continue + m = re.match(r"backbone\.transformer_blocks\.(\d+)\.(.*)$", key) + if m and any(m.group(2) in trip[:3] for trip in QKV_DOUBLE_TRIPLETS): + idx = m.group(1) + raw_by_block.setdefault(idx, {})[m.group(2)] = (key, shard_path, data_off, info) + else: + raw_others.append((key, shard_path, data_off, info)) + + for key, shard_path, data_off, info in raw_others: + new_key = rewrite_fn(key) + # Swap the (scale, shift) halves to (shift, scale) at conversion time so + # the on-disk weight matches BFL flux ordering and the runtime stays + # version-agnostic. norm_out.linear weight shape is [2*dim, dim] and bias + # is [2*dim]; both split along axis 0 (outermost == row-major outer). + if new_key in ("final_layer.adaLN_modulation.1.weight", + "final_layer.adaLN_modulation.1.bias"): + info = dict(info) + info["_chunk_swap_halves"] = True + entries.append((new_key, shard_path, data_off, info)) + + for block_idx, tails in raw_by_block.items(): + for q_tail, k_tail, v_tail, fused_tail in QKV_DOUBLE_TRIPLETS: + if q_tail in tails and k_tail in tails and v_tail in tails: + q = tails[q_tail]; k = tails[k_tail]; v = tails[v_tail] + # Validate shapes match. + q_shape = q[3]["shape"]; k_shape = k[3]["shape"]; v_shape = v[3]["shape"] + if q_shape != k_shape or q_shape != v_shape: + raise ValueError(f"qkv shape mismatch at block {block_idx} {q_tail}: q={q_shape} k={k_shape} v={v_shape}") + fused_shape = [q_shape[0] * 3] + list(q_shape[1:]) + fused_info = { + "dtype": q[3]["dtype"], + "shape": fused_shape, + "_qkv_sources": [q, k, v], # pseudo field consumed by writer + } + entries.append((f"double_blocks.{block_idx}.{fused_tail}", + None, None, fused_info)) + del tails[q_tail]; del tails[k_tail]; del tails[v_tail] + # Anything left in tails was an unmatched single - pass through. + for tail, payload in tails.items(): + entries.append((rewrite_fn(payload[0]),) + payload[1:]) + return entries + + +_DTYPE_BYTES = { + "BF16": 2, "F16": 2, "F32": 4, "F64": 8, + "U8": 1, "I8": 1, "I16": 2, "I32": 4, "I64": 8, + "BOOL": 1, +} + + +def _total_bytes(info: dict) -> int: + if "_qkv_sources" in info: + elems = 1 + for d in info["shape"]: + elems *= d + return elems * _DTYPE_BYTES[info["dtype"]] + start, end = info["data_offsets"] + return end - start + + +def write_consolidated(out_path: str, entries): + """Write a single safetensors file by streaming raw bytes from each shard. + + For qkv-fused entries, q/k/v are concatenated along axis 0 (row-major), so a + simple byte-level concatenation produces the correct fused layout for any + standard dtype. + """ + entries = sorted(entries, key=lambda e: e[0]) + + new_header = {} + cur_offset = 0 + for new_key, shard_path, data_off, info in entries: + size = _total_bytes(info) + new_header[new_key] = { + "dtype": info["dtype"], + "shape": info["shape"], + "data_offsets": [cur_offset, cur_offset + size], + } + cur_offset += size + + header_json = json.dumps(new_header, separators=(",", ":")).encode("utf-8") + pad = (-len(header_json)) % 8 + header_json = header_json + (b" " * pad) + + def copy_range(src_path, src_data_off, src_info, out, byte_range=None): + start, end = src_info["data_offsets"] + if byte_range is not None: + sub_start, sub_end = byte_range + start, end = start + sub_start, start + sub_end + with open(src_path, "rb") as src: + src.seek(src_data_off + start) + remaining = end - start + while remaining > 0: + chunk = src.read(min(8 * 1024 * 1024, remaining)) + if not chunk: + raise IOError(f"Truncated tensor in {src_path}") + out.write(chunk) + remaining -= len(chunk) + + with open(out_path, "wb") as out: + out.write(struct.pack(" {args.output}") + os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True) + write_consolidated(args.output, transformer_entries) + print(f"Done. Output size: {os.path.getsize(args.output) / 1e9:.2f} GB") + + +if __name__ == "__main__": + main() diff --git a/script/pulid_extract_id.py b/script/pulid_extract_id.py new file mode 100644 index 00000000..aca52e24 --- /dev/null +++ b/script/pulid_extract_id.py @@ -0,0 +1,134 @@ +""" +Precompute a PuLID-Flux identity embedding from a single source portrait. + +Writes a gguf file (a single tensor `pulid_id`) that stable-diffusion.cpp's +`--pulid-id-embedding` flag consumes. + +Dependencies (recommended: vendor rather than pip-install due to upstream +packaging quirks): + - torch + safetensors + - The ToTheBeginning/PuLID repository's `pulid/` package and `eva_clip/`. + Put them on PYTHONPATH or sys.path before running this script. + - insightface, facexlib, torchvision, opencv-python, huggingface_hub, gguf + - numpy, Pillow + +Usage: + python script/pulid_extract_id.py \\ + --portrait /path/to/source-photo.jpg \\ + --pulid-weights /path/to/pulid_flux_v0.9.1.safetensors \\ + --out /path/to/source.pulidembd + +The portrait must contain a clearly visible face. insightface's antelopev2 +detector will be auto-downloaded on first run. +""" + +from __future__ import annotations + +import argparse +import os +import sys +from types import SimpleNamespace + + +def extract(portrait_path: str, pulid_weights: str) -> "torch.Tensor": + import numpy as np + import torch + from PIL import Image + from pulid.pipeline_flux import PuLIDPipeline + + if torch.cuda.is_available(): + device, onnx_provider = "cuda", "gpu" + else: + device, onnx_provider = "cpu", "cpu" + + print(f"device={device}", flush=True) + + # PuLIDPipeline only attaches pulid_ca attributes to `dit` during + # construction; get_id_embedding() never runs Flux, so a dummy object is + # enough and avoids importing/building a Flux skeleton. + print("instantiating PuLIDPipeline with a dummy Flux object", flush=True) + dit = SimpleNamespace() + pulid = PuLIDPipeline(dit=dit, + device=device, + weight_dtype=torch.bfloat16, + onnx_provider=onnx_provider) + + print(f"loading PuLID weights from {pulid_weights}", flush=True) + pulid.load_pretrain(pretrain_path=pulid_weights, version="v0.9.1") + + print(f"extracting ID embedding from {portrait_path}", flush=True) + face_img = np.array(Image.open(portrait_path).convert("RGB")) + id_embedding, _ = pulid.get_id_embedding(face_img) + print(f"id embedding shape={tuple(id_embedding.shape)} dtype={id_embedding.dtype}", + flush=True) + + if id_embedding.ndim == 3 and id_embedding.shape[0] == 1: + id_embedding = id_embedding[0] + return id_embedding + + +def write_embd(tensor, out_path: str, dtype_choice: str) -> None: + import gguf + import torch + + if tensor.ndim != 2: + raise ValueError(f"expected (num_tokens, token_dim); got {tuple(tensor.shape)}") + num_tokens, token_dim = tensor.shape + + os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True) + + writer = gguf.GGUFWriter(out_path, arch="pulid") + writer.add_uint32("pulid.version", 1) + + if dtype_choice == "fp16": + arr = tensor.to(torch.float16).contiguous().cpu().numpy() + writer.add_tensor("pulid_id", arr) + elif dtype_choice == "fp32": + arr = tensor.to(torch.float32).contiguous().cpu().numpy() + writer.add_tensor("pulid_id", arr) + elif dtype_choice == "bf16": + raw = tensor.to(torch.bfloat16).contiguous().view(torch.uint16).cpu().numpy() + writer.add_tensor("pulid_id", raw, + raw_shape=(int(num_tokens), int(token_dim)), + raw_dtype=gguf.GGMLQuantizationType.BF16) + else: + raise ValueError(f"unknown --dtype {dtype_choice}") + + writer.write_header_to_file() + writer.write_kv_data_to_file() + writer.write_tensors_to_file() + writer.close() + + print(f"wrote {out_path}: gguf, tensor pulid_id [{token_dim}, {num_tokens}] {dtype_choice}", + flush=True) + + +def main() -> int: + ap = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--portrait", required=True, + help="Path to the source portrait image (JPG/PNG).") + ap.add_argument("--pulid-weights", required=True, + help="Path to pulid_flux_v0.9.x.safetensors.") + ap.add_argument("--out", required=True, + help="Output path for the .pulidembd binary.") + ap.add_argument("--dtype", default="fp16", + choices=["fp16", "bf16", "fp32"], + help="Storage dtype (default fp16; produces ~131 KB).") + args = ap.parse_args() + + if not os.path.exists(args.portrait): + print(f"ERROR: portrait not found at {args.portrait}", file=sys.stderr) + return 2 + if not os.path.exists(args.pulid_weights): + print(f"ERROR: PuLID weights not found at {args.pulid_weights}", file=sys.stderr) + return 3 + + embedding = extract(args.portrait, args.pulid_weights) + write_embd(embedding, args.out, args.dtype) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/script/remove_utf8_bom.py b/script/remove_utf8_bom.py new file mode 100644 index 00000000..ccb07b0d --- /dev/null +++ b/script/remove_utf8_bom.py @@ -0,0 +1,234 @@ +#!/usr/bin/env python3 +"""Remove UTF-8 BOMs from files under a directory. + +By default this scans the current working directory recursively and skips +repository areas that should not be touched by ordinary maintenance scripts. +Only files whose first three bytes are the UTF-8 BOM are rewritten. +""" + +import argparse +import os +import shutil +import sys +import tempfile +from pathlib import Path + + +UTF8_BOM = b"\xef\xbb\xbf" + +DEFAULT_EXCLUDED_DIR_NAMES = { + ".git", + ".hg", + ".svn", + ".mypy_cache", + ".pytest_cache", + "__pycache__", + "test", +} + +DEFAULT_EXCLUDED_DIR_PREFIXES = { + "build", +} + +DEFAULT_EXCLUDED_REL_DIRS = { + "examples/server/frontend", + "ggml", + "models", + "src/vocab", + "thirdparty", +} + + +def rel_posix(path: Path, root: Path) -> str: + try: + return path.relative_to(root).as_posix() + except ValueError: + return path.as_posix() + + +def should_skip_dir( + path: Path, + root: Path, + excluded_rel_dirs: set[str], + excluded_names: set[str], + excluded_prefixes: set[str], +) -> bool: + rel = rel_posix(path, root) + return ( + path.name in excluded_names + or rel in excluded_rel_dirs + or any(path.name.startswith(prefix) for prefix in excluded_prefixes) + ) + + +def iter_files( + root: Path, + recursive: bool, + excluded_rel_dirs: set[str], + excluded_names: set[str], + excluded_prefixes: set[str], + follow_symlinks: bool, +): + if recursive: + for dirpath, dirnames, filenames in os.walk(root, followlinks=follow_symlinks): + current_dir = Path(dirpath) + dirnames[:] = [ + name + for name in dirnames + if not should_skip_dir( + current_dir / name, + root, + excluded_rel_dirs, + excluded_names, + excluded_prefixes, + ) + ] + for filename in filenames: + path = current_dir / filename + if path.is_symlink() and not follow_symlinks: + continue + yield path + else: + for path in root.iterdir(): + if path.is_file() and (follow_symlinks or not path.is_symlink()): + yield path + + +def has_utf8_bom(path: Path) -> bool: + with path.open("rb") as f: + return f.read(len(UTF8_BOM)) == UTF8_BOM + + +def strip_utf8_bom(path: Path) -> None: + tmp_path = None + try: + with path.open("rb") as src: + if src.read(len(UTF8_BOM)) != UTF8_BOM: + return + + fd, tmp_name = tempfile.mkstemp( + prefix=f".{path.name}.", + suffix=".tmp", + dir=str(path.parent), + ) + tmp_path = Path(tmp_name) + with os.fdopen(fd, "wb") as dst: + shutil.copyfileobj(src, dst, length=1024 * 1024) + + shutil.copystat(path, tmp_path, follow_symlinks=False) + os.replace(tmp_path, path) + tmp_path = None + finally: + if tmp_path is not None: + try: + tmp_path.unlink() + except FileNotFoundError: + pass + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Scan files and convert UTF-8 BOM files to UTF-8 without BOM.", + ) + parser.add_argument( + "root", + nargs="?", + default=".", + help="Directory to scan. Defaults to the current directory.", + ) + parser.add_argument( + "-n", + "--dry-run", + action="store_true", + help="Only list files that would be converted.", + ) + parser.add_argument( + "--no-recursive", + action="store_true", + help="Only scan files directly under root.", + ) + parser.add_argument( + "--include-repo-excluded", + action="store_true", + help="Do not skip default repository excluded directories.", + ) + parser.add_argument( + "--exclude-dir", + action="append", + default=[], + metavar="DIR", + help="Additional directory name or root-relative path to skip. Can be used multiple times.", + ) + parser.add_argument( + "--follow-symlinks", + action="store_true", + help="Follow symlinked directories and files.", + ) + parser.add_argument( + "-q", + "--quiet", + action="store_true", + help="Only print the final summary.", + ) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + root = Path(args.root).resolve() + + if not root.is_dir(): + print(f"error: not a directory: {root}", file=sys.stderr) + return 2 + + excluded_names = set() + excluded_rel_dirs = set() + excluded_prefixes = set() + if not args.include_repo_excluded: + excluded_names.update(DEFAULT_EXCLUDED_DIR_NAMES) + excluded_rel_dirs.update(DEFAULT_EXCLUDED_REL_DIRS) + excluded_prefixes.update(DEFAULT_EXCLUDED_DIR_PREFIXES) + + for item in args.exclude_dir: + normalized = Path(item).as_posix().strip("/") + if "/" in normalized: + excluded_rel_dirs.add(normalized) + else: + excluded_names.add(normalized) + + scanned = 0 + converted = 0 + errors = 0 + + for path in iter_files( + root=root, + recursive=not args.no_recursive, + excluded_rel_dirs=excluded_rel_dirs, + excluded_names=excluded_names, + excluded_prefixes=excluded_prefixes, + follow_symlinks=args.follow_symlinks, + ): + scanned += 1 + try: + if not has_utf8_bom(path): + continue + converted += 1 + rel = rel_posix(path, root) + if args.dry_run: + if not args.quiet: + print(f"would convert: {rel}") + else: + strip_utf8_bom(path) + if not args.quiet: + print(f"converted: {rel}") + except OSError as exc: + errors += 1 + print(f"error: {rel_posix(path, root)}: {exc}", file=sys.stderr) + + action = "would convert" if args.dry_run else "converted" + print(f"scanned {scanned} file(s), {action} {converted}, errors {errors}") + return 1 if errors else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/conditioner.hpp b/src/conditioning/conditioner.hpp similarity index 84% rename from src/conditioner.hpp rename to src/conditioning/conditioner.hpp index 157d3906..d63303a8 100644 --- a/src/conditioner.hpp +++ b/src/conditioning/conditioner.hpp @@ -1,14 +1,15 @@ -#ifndef __CONDITIONER_HPP__ -#define __CONDITIONER_HPP__ +#ifndef __SD_CONDITIONING_CONDITIONER_HPP__ +#define __SD_CONDITIONING_CONDITIONER_HPP__ #include #include #include -#include "clip.hpp" -#include "llm.hpp" -#include "t5.hpp" -#include "tensor_ggml.hpp" +#include "core/tensor_ggml.hpp" +#include "model/te/clip.hpp" +#include "model/te/llm.hpp" +#include "model/te/t5.hpp" +#include "model_loader.h" struct SDCondition { sd::Tensor c_crossattn; @@ -103,7 +104,6 @@ struct ConditionerParams { int width = -1; int height = -1; bool zero_out_masked = false; - int num_input_imgs = 0; // for photomaker const std::vector>* ref_images = nullptr; // for qwen image edit }; @@ -113,33 +113,22 @@ struct Conditioner { public: virtual SDCondition get_learned_condition(int n_threads, const ConditionerParams& conditioner_params) = 0; - virtual bool alloc_params_buffer() = 0; - virtual void free_params_buffer() = 0; virtual void get_param_tensors(std::map& tensors) = 0; - virtual size_t get_params_buffer_size() = 0; virtual void set_max_graph_vram_bytes(size_t max_vram_bytes) {} virtual void set_stream_layers_enabled(bool enabled) {} virtual void set_flash_attention_enabled(bool enabled) = 0; virtual void set_weight_adapter(const std::shared_ptr& adapter) {} - virtual std::tuple> get_learned_condition_with_trigger(int n_threads, - const ConditionerParams& conditioner_params) { - GGML_ABORT("Not implemented yet!"); - } - virtual std::string remove_trigger_from_prompt(const std::string& prompt) { - GGML_ABORT("Not implemented yet!"); - } + virtual void runner_done() {} }; // ldm.modules.encoders.modules.FrozenCLIPEmbedder // Ref: https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/cad87bf4e3e0b0a759afa94e933527c3123d59bc/modules/sd_hijack_clip.py#L283 struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { - SDVersion version = VERSION_SD1; - PMVersion pm_version = PM_VERSION_1; + SDVersion version = VERSION_SD1; CLIPTokenizer tokenizer; std::shared_ptr text_model; std::shared_ptr text_model2; - std::string trigger_word = "img"; // should be user settable std::map embedding_map; int32_t num_custom_embeddings = 0; int32_t num_custom_embeddings_2 = 0; @@ -147,26 +136,24 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { std::map> embedding_pos_map; FrozenCLIPEmbedderWithCustomWords(ggml_backend_t backend, - ggml_backend_t params_backend, const String2TensorStorage& tensor_storage_map, const std::map& orig_embedding_map, - SDVersion version = VERSION_SD1, - PMVersion pv = PM_VERSION_1) - : version(version), pm_version(pv), tokenizer(sd_version_is_sd2(version) ? 0 : 49407) { + SDVersion version = VERSION_SD1, + std::shared_ptr weight_manager = nullptr) + : version(version), tokenizer(sd_version_is_sd2(version) ? 0 : 49407) { for (const auto& kv : orig_embedding_map) { - std::string name = kv.first; - std::transform(name.begin(), name.end(), name.begin(), [](unsigned char c) { return std::tolower(c); }); + std::string name = normalize_embedding_name(kv.first); embedding_map[name] = kv.second; tokenizer.add_special_token(name); } bool force_clip_f32 = !embedding_map.empty(); if (sd_version_is_sd1(version)) { - text_model = std::make_shared(backend, params_backend, tensor_storage_map, "cond_stage_model.transformer.text_model", OPENAI_CLIP_VIT_L_14, true, force_clip_f32); + text_model = std::make_shared(backend, tensor_storage_map, "cond_stage_model.transformer.text_model", OPENAI_CLIP_VIT_L_14, true, force_clip_f32, weight_manager); } else if (sd_version_is_sd2(version)) { - text_model = std::make_shared(backend, params_backend, tensor_storage_map, "cond_stage_model.transformer.text_model", OPEN_CLIP_VIT_H_14, true, force_clip_f32); + text_model = std::make_shared(backend, tensor_storage_map, "cond_stage_model.transformer.text_model", OPEN_CLIP_VIT_H_14, true, force_clip_f32, weight_manager); } else if (sd_version_is_sdxl(version)) { - text_model = std::make_shared(backend, params_backend, tensor_storage_map, "cond_stage_model.transformer.text_model", OPENAI_CLIP_VIT_L_14, false, force_clip_f32); - text_model2 = std::make_shared(backend, params_backend, tensor_storage_map, "cond_stage_model.1.transformer.text_model", OPEN_CLIP_VIT_BIGG_14, false, force_clip_f32); + text_model = std::make_shared(backend, tensor_storage_map, "cond_stage_model.transformer.text_model", OPENAI_CLIP_VIT_L_14, false, force_clip_f32, weight_manager); + text_model2 = std::make_shared(backend, tensor_storage_map, "cond_stage_model.1.transformer.text_model", OPEN_CLIP_VIT_BIGG_14, false, force_clip_f32, weight_manager); } } @@ -177,33 +164,6 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { } } - bool alloc_params_buffer() override { - if (!text_model->alloc_params_buffer()) { - return false; - } - if (sd_version_is_sdxl(version)) { - if (!text_model2->alloc_params_buffer()) { - return false; - } - } - return true; - } - - void free_params_buffer() override { - text_model->free_params_buffer(); - if (sd_version_is_sdxl(version)) { - text_model2->free_params_buffer(); - } - } - - size_t get_params_buffer_size() override { - size_t buffer_size = text_model->get_params_buffer_size(); - if (sd_version_is_sdxl(version)) { - buffer_size += text_model2->get_params_buffer_size(); - } - return buffer_size; - } - void set_max_graph_vram_bytes(size_t max_vram_bytes) override { text_model->set_max_graph_vram_bytes(max_vram_bytes); if (sd_version_is_sdxl(version)) { @@ -232,6 +192,13 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { } } + void runner_done() override { + text_model->runner_done(); + if (sd_version_is_sdxl(version)) { + text_model2->runner_done(); + } + } + bool load_embedding(std::string embd_name, std::string embd_path, std::vector& bpe_tokens) { ModelLoader model_loader; if (!model_loader.init_from_file_and_convert_name(embd_path)) { @@ -273,7 +240,8 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { } return true; }; - model_loader.load_tensors(on_load, 1); + model_loader.set_n_threads(1); + model_loader.load_tensors(on_load); int pos_start = num_custom_embeddings; if (embd) { int64_t hidden_size = text_model->model.hidden_size; @@ -309,17 +277,23 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { return true; } + static std::string normalize_embedding_name(std::string name) { + std::transform(name.begin(), name.end(), name.begin(), [](unsigned char c) { return std::tolower(c); }); + return name; + } + + bool append_embedding_tokens(std::string str, std::vector& bpe_tokens) { + std::string name = normalize_embedding_name(std::move(str)); + auto iter = embedding_map.find(name); + if (iter == embedding_map.end()) { + return false; + } + return load_embedding(name, iter->second, bpe_tokens); + } + std::vector convert_token_to_id(std::string text) { auto on_new_token_cb = [&](std::string& str, std::vector& bpe_tokens) -> bool { - auto iter = embedding_map.find(str); - if (iter == embedding_map.end()) { - return false; - } - std::string embedding_path = iter->second; - if (load_embedding(str, embedding_path, bpe_tokens)) { - return true; - } - return false; + return append_embedding_tokens(str, bpe_tokens); }; std::vector curr_tokens = tokenizer.encode(text, on_new_token_cb); return curr_tokens; @@ -329,121 +303,6 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { return tokenizer.decode(tokens); } - std::tuple, std::vector, std::vector> - tokenize_with_trigger_token(std::string text, - int num_input_imgs, - int32_t image_token) { - auto parsed_attention = parse_prompt_attention(text); - - { - std::stringstream ss; - ss << "["; - for (const auto& item : parsed_attention) { - ss << "['" << item.first << "', " << item.second << "], "; - } - ss << "]"; - LOG_DEBUG("parse '%s' to %s", text.c_str(), ss.str().c_str()); - } - - auto on_new_token_cb = [&](std::string& str, std::vector& bpe_tokens) -> bool { - auto iter = embedding_map.find(str); - if (iter == embedding_map.end()) { - return false; - } - std::string embedding_path = iter->second; - if (load_embedding(str, embedding_path, bpe_tokens)) { - return true; - } - return false; - }; - - std::vector tokens; - std::vector weights; - std::vector class_token_mask; - int32_t class_idx = -1, tokens_acc = 0; - for (const auto& item : parsed_attention) { - std::vector class_token_index; - std::vector clean_input_ids; - const std::string& curr_text = item.first; - float curr_weight = item.second; - // printf(" %s: %f \n", curr_text.c_str(), curr_weight); - int32_t clean_index = 0; - if (curr_text == "BREAK" && curr_weight == -1.0f) { - // Pad token array up to chunk size at this point. - // TODO: This is a hardcoded chunk_len, like in stable-diffusion.cpp, make it a parameter for the future? - // Also, this is 75 instead of 77 to leave room for BOS and EOS tokens. - int padding_size = 75 - (tokens_acc % 75); - for (int j = 0; j < padding_size; j++) { - clean_input_ids.push_back(tokenizer.EOS_TOKEN_ID); - clean_index++; - } - - // After padding, continue to the next iteration to process the following text as a new segment - tokens.insert(tokens.end(), clean_input_ids.begin(), clean_input_ids.end()); - weights.insert(weights.end(), padding_size, curr_weight); - continue; - } - - // Regular token, process normally - std::vector curr_tokens = tokenizer.encode(curr_text, on_new_token_cb); - for (uint32_t i = 0; i < curr_tokens.size(); i++) { - int token_id = curr_tokens[i]; - if (token_id == image_token) { - class_token_index.push_back(clean_index - 1); - } else { - clean_input_ids.push_back(token_id); - clean_index++; - } - } - // GGML_ASSERT(class_token_index.size() == 1); // PhotoMaker currently does not support multiple - // trigger words in a single prompt. - if (class_token_index.size() == 1) { - // Expand the class word token and corresponding mask - int class_token = clean_input_ids[class_token_index[0]]; - class_idx = tokens_acc + class_token_index[0]; - std::vector clean_input_ids_tmp; - for (int i = 0; i < class_token_index[0]; i++) - clean_input_ids_tmp.push_back(clean_input_ids[i]); - for (int i = 0; i < (pm_version == PM_VERSION_2 ? 2 * num_input_imgs : num_input_imgs); i++) - clean_input_ids_tmp.push_back(class_token); - for (int i = class_token_index[0] + 1; i < clean_input_ids.size(); i++) - clean_input_ids_tmp.push_back(clean_input_ids[i]); - clean_input_ids.clear(); - clean_input_ids = clean_input_ids_tmp; - } - tokens_acc += clean_index; - tokens.insert(tokens.end(), clean_input_ids.begin(), clean_input_ids.end()); - weights.insert(weights.end(), clean_input_ids.size(), curr_weight); - } - // BUG!! double couting, pad_tokens will add BOS at the beginning - // tokens.insert(tokens.begin(), tokenizer.BOS_TOKEN_ID); - // weights.insert(weights.begin(), 1.0); - - tokenizer.pad_tokens(tokens, &weights, nullptr, text_model->model.n_token, text_model->model.n_token, true); - int offset = pm_version == PM_VERSION_2 ? 2 * num_input_imgs : num_input_imgs; - for (int i = 0; i < tokens.size(); i++) { - // if (class_idx + 1 <= i && i < class_idx + 1 + 2*num_input_imgs) // photomaker V2 has num_tokens(=2)*num_input_imgs - if (class_idx + 1 <= i && i < class_idx + 1 + offset) // photomaker V2 has num_tokens(=2)*num_input_imgs - // hardcode for now - class_token_mask.push_back(true); - else - class_token_mask.push_back(false); - } - - // printf("["); - // for (int i = 0; i < tokens.size(); i++) { - // printf("%d, ", class_token_mask[i] ? 1 : 0); - // } - // printf("]\n"); - - // for (int i = 0; i < tokens.size(); i++) { - // std::cout << tokens[i] << ":" << weights[i] << ", "; - // } - // std::cout << std::endl; - - return std::make_tuple(tokens, weights, class_token_mask); - } - std::pair, std::vector> tokenize(std::string text, size_t min_length = 0, size_t max_length = 0, @@ -461,15 +320,7 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { } auto on_new_token_cb = [&](std::string& str, std::vector& bpe_tokens) -> bool { - auto iter = embedding_map.find(str); - if (iter == embedding_map.end()) { - return false; - } - std::string embedding_path = iter->second; - if (load_embedding(str, embedding_path, bpe_tokens)) { - return true; - } - return false; + return append_embedding_tokens(str, bpe_tokens); }; std::vector tokens; @@ -557,7 +408,10 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { token_embed_custom.data(), max_token_idx, false, - clip_skip); + clip_skip, + false, + true, + true); GGML_ASSERT(!chunk_hidden_states.empty()); if (sd_version_is_sdxl(version)) { auto chunk_hidden_states2 = text_model2->compute(n_threads, @@ -566,7 +420,10 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { token_embed_custom.data(), max_token_idx, false, - clip_skip); + clip_skip, + false, + true, + true); GGML_ASSERT(!chunk_hidden_states2.empty()); chunk_hidden_states = sd::ops::concat(chunk_hidden_states, chunk_hidden_states2, 0); @@ -577,7 +434,10 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { token_embed_custom.data(), max_token_idx, true, - clip_skip); + clip_skip, + false, + true, + true); GGML_ASSERT(!pooled.empty()); } } @@ -631,49 +491,6 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { return result; } - std::tuple> - get_learned_condition_with_trigger(int n_threads, - const ConditionerParams& conditioner_params) override { - auto image_tokens = convert_token_to_id(trigger_word); - // if(image_tokens.size() == 1){ - // printf(" image token id is: %d \n", image_tokens[0]); - // } - GGML_ASSERT(image_tokens.size() == 1); - auto tokens_and_weights = tokenize_with_trigger_token(conditioner_params.text, - conditioner_params.num_input_imgs, - image_tokens[0]); - std::vector& tokens = std::get<0>(tokens_and_weights); - std::vector& weights = std::get<1>(tokens_and_weights); - std::vector& clsm = std::get<2>(tokens_and_weights); - // printf("tokens: \n"); - // for(int i = 0; i < tokens.size(); ++i) - // printf("%d ", tokens[i]); - // printf("\n"); - // printf("clsm: \n"); - // for(int i = 0; i < clsm.size(); ++i) - // printf("%d ", clsm[i]?1:0); - // printf("\n"); - auto cond = get_learned_condition_common(n_threads, - tokens, - weights, - conditioner_params.clip_skip, - conditioner_params.width, - conditioner_params.height, - conditioner_params.zero_out_masked); - return std::make_tuple(cond, clsm); - } - - std::string remove_trigger_from_prompt(const std::string& prompt) override { - auto image_tokens = convert_token_to_id(trigger_word); - GGML_ASSERT(image_tokens.size() == 1); - auto tokens_and_weights = tokenize(prompt); - std::vector& tokens = tokens_and_weights.first; - auto it = std::find(tokens.begin(), tokens.end(), image_tokens[0]); - GGML_ASSERT(it != tokens.end()); // prompt must have trigger word - tokens.erase(it); - return decode(tokens); - } - SDCondition get_learned_condition(int n_threads, const ConditionerParams& conditioner_params) override { auto tokens_and_weights = tokenize(conditioner_params.text, text_model->model.n_token, text_model->model.n_token, true); @@ -691,15 +508,15 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { struct FrozenCLIPVisionEmbedder : public GGMLRunner { CLIPVisionModelProjection vision_model; + std::string weight_prefix = "cond_stage_model.transformer"; FrozenCLIPVisionEmbedder(ggml_backend_t backend, - ggml_backend_t params_backend, - const String2TensorStorage& tensor_storage_map = {}) - : GGMLRunner(backend, params_backend) { - std::string prefix = "cond_stage_model.transformer"; - bool proj_in = false; + const String2TensorStorage& tensor_storage_map = {}, + std::shared_ptr weight_manager = nullptr) + : GGMLRunner(backend, weight_manager) { + bool proj_in = false; for (const auto& [name, tensor_storage] : tensor_storage_map) { - if (!starts_with(name, prefix)) { + if (!starts_with(name, weight_prefix)) { continue; } if (contains(name, "self_attn.in_proj")) { @@ -708,7 +525,7 @@ struct FrozenCLIPVisionEmbedder : public GGMLRunner { } } vision_model = CLIPVisionModelProjection(OPEN_CLIP_VIT_H_14, false, proj_in); - vision_model.init(params_ctx, tensor_storage_map, prefix); + vision_model.init(params_ctx, tensor_storage_map, weight_prefix); } std::string get_desc() override { @@ -716,7 +533,7 @@ struct FrozenCLIPVisionEmbedder : public GGMLRunner { } void get_param_tensors(std::map& tensors) { - vision_model.get_param_tensors(tensors, "cond_stage_model.transformer"); + vision_model.get_param_tensors(tensors, weight_prefix); } ggml_cgraph* build_graph(const sd::Tensor& pixel_values_tensor, bool return_pooled, int clip_skip) { @@ -739,7 +556,7 @@ struct FrozenCLIPVisionEmbedder : public GGMLRunner { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(pixel_values, return_pooled, clip_skip); }; - return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true)); + return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true, true, true)); } }; @@ -752,8 +569,8 @@ struct SD3CLIPEmbedder : public Conditioner { std::shared_ptr t5; SD3CLIPEmbedder(ggml_backend_t backend, - ggml_backend_t params_backend, - const String2TensorStorage& tensor_storage_map = {}) + const String2TensorStorage& tensor_storage_map = {}, + std::shared_ptr weight_manager = nullptr) : clip_g_tokenizer(0) { bool use_clip_l = false; bool use_clip_g = false; @@ -772,13 +589,13 @@ struct SD3CLIPEmbedder : public Conditioner { return; } if (use_clip_l) { - clip_l = std::make_shared(backend, params_backend, tensor_storage_map, "text_encoders.clip_l.transformer.text_model", OPENAI_CLIP_VIT_L_14, false); + clip_l = std::make_shared(backend, tensor_storage_map, "text_encoders.clip_l.transformer.text_model", OPENAI_CLIP_VIT_L_14, false, false, weight_manager); } if (use_clip_g) { - clip_g = std::make_shared(backend, params_backend, tensor_storage_map, "text_encoders.clip_g.transformer.text_model", OPEN_CLIP_VIT_BIGG_14, false); + clip_g = std::make_shared(backend, tensor_storage_map, "text_encoders.clip_g.transformer.text_model", OPEN_CLIP_VIT_BIGG_14, false, false, weight_manager); } if (use_t5) { - t5 = std::make_shared(backend, params_backend, tensor_storage_map, "text_encoders.t5xxl.transformer"); + t5 = std::make_shared(backend, tensor_storage_map, "text_encoders.t5xxl.transformer", false, weight_manager); } } @@ -794,51 +611,6 @@ struct SD3CLIPEmbedder : public Conditioner { } } - bool alloc_params_buffer() override { - if (clip_l) { - if (!clip_l->alloc_params_buffer()) { - return false; - } - } - if (clip_g) { - if (!clip_g->alloc_params_buffer()) { - return false; - } - } - if (t5) { - if (!t5->alloc_params_buffer()) { - return false; - } - } - return true; - } - - void free_params_buffer() override { - if (clip_l) { - clip_l->free_params_buffer(); - } - if (clip_g) { - clip_g->free_params_buffer(); - } - if (t5) { - t5->free_params_buffer(); - } - } - - size_t get_params_buffer_size() override { - size_t buffer_size = 0; - if (clip_l) { - buffer_size += clip_l->get_params_buffer_size(); - } - if (clip_g) { - buffer_size += clip_g->get_params_buffer_size(); - } - if (t5) { - buffer_size += t5->get_params_buffer_size(); - } - return buffer_size; - } - void set_max_graph_vram_bytes(size_t max_vram_bytes) override { if (clip_l) { clip_l->set_max_graph_vram_bytes(max_vram_bytes); @@ -887,6 +659,18 @@ struct SD3CLIPEmbedder : public Conditioner { } } + void runner_done() override { + if (clip_l) { + clip_l->runner_done(); + } + if (clip_g) { + clip_g->runner_done(); + } + if (t5) { + t5->runner_done(); + } + } + std::vector, std::vector>> tokenize(std::string text, size_t min_length = 0, size_t max_length = 0, @@ -1002,7 +786,10 @@ struct SD3CLIPEmbedder : public Conditioner { nullptr, max_token_idx, false, - clip_skip); + clip_skip, + false, + true, + true); GGML_ASSERT(!chunk_hidden_states_l.empty()); chunk_hidden_states_l = ::apply_token_weights(std::move(chunk_hidden_states_l), chunk_weights); @@ -1015,13 +802,16 @@ struct SD3CLIPEmbedder : public Conditioner { nullptr, max_token_idx, true, - clip_skip); + clip_skip, + false, + true, + true); GGML_ASSERT(!pooled_l.empty()); } } else { chunk_hidden_states_l = sd::Tensor::zeros({768, static_cast(chunk_len), 1}); if (chunk_idx == 0) { - pooled = sd::Tensor::zeros({768, 1}); + pooled_l = sd::Tensor::zeros({768, 1}); } } @@ -1043,7 +833,10 @@ struct SD3CLIPEmbedder : public Conditioner { nullptr, max_token_idx, false, - clip_skip); + clip_skip, + false, + true, + true); GGML_ASSERT(!chunk_hidden_states_g.empty()); chunk_hidden_states_g = ::apply_token_weights(std::move(chunk_hidden_states_g), chunk_weights); @@ -1056,7 +849,10 @@ struct SD3CLIPEmbedder : public Conditioner { nullptr, max_token_idx, true, - clip_skip); + clip_skip, + false, + true, + true); GGML_ASSERT(!pooled_g.empty()); } } else { @@ -1078,7 +874,10 @@ struct SD3CLIPEmbedder : public Conditioner { chunk_hidden_states_t5 = t5->compute(n_threads, input_ids, - sd::Tensor()); + sd::Tensor(), + false, + true, + true); GGML_ASSERT(!chunk_hidden_states_t5.empty()); chunk_hidden_states_t5 = ::apply_token_weights(std::move(chunk_hidden_states_t5), chunk_weights); } else { @@ -1139,8 +938,8 @@ struct FluxCLIPEmbedder : public Conditioner { size_t chunk_len = 256; FluxCLIPEmbedder(ggml_backend_t backend, - ggml_backend_t params_backend, - const String2TensorStorage& tensor_storage_map = {}) { + const String2TensorStorage& tensor_storage_map = {}, + std::shared_ptr weight_manager = nullptr) { bool use_clip_l = false; bool use_t5 = false; for (auto pair : tensor_storage_map) { @@ -1157,12 +956,12 @@ struct FluxCLIPEmbedder : public Conditioner { } if (use_clip_l) { - clip_l = std::make_shared(backend, params_backend, tensor_storage_map, "text_encoders.clip_l.transformer.text_model", OPENAI_CLIP_VIT_L_14, true); + clip_l = std::make_shared(backend, tensor_storage_map, "text_encoders.clip_l.transformer.text_model", OPENAI_CLIP_VIT_L_14, true, false, weight_manager); } else { LOG_WARN("clip_l text encoder not found! Prompt adherence might be degraded."); } if (use_t5) { - t5 = std::make_shared(backend, params_backend, tensor_storage_map, "text_encoders.t5xxl.transformer"); + t5 = std::make_shared(backend, tensor_storage_map, "text_encoders.t5xxl.transformer", false, weight_manager); } else { LOG_WARN("t5xxl text encoder not found! Prompt adherence might be degraded."); } @@ -1177,40 +976,6 @@ struct FluxCLIPEmbedder : public Conditioner { } } - bool alloc_params_buffer() override { - if (clip_l) { - if (!clip_l->alloc_params_buffer()) { - return false; - } - } - if (t5) { - if (!t5->alloc_params_buffer()) { - return false; - } - } - return true; - } - - void free_params_buffer() override { - if (clip_l) { - clip_l->free_params_buffer(); - } - if (t5) { - t5->free_params_buffer(); - } - } - - size_t get_params_buffer_size() override { - size_t buffer_size = 0; - if (clip_l) { - buffer_size += clip_l->get_params_buffer_size(); - } - if (t5) { - buffer_size += t5->get_params_buffer_size(); - } - return buffer_size; - } - void set_max_graph_vram_bytes(size_t max_vram_bytes) override { if (clip_l) { clip_l->set_max_graph_vram_bytes(max_vram_bytes); @@ -1238,7 +1003,7 @@ struct FluxCLIPEmbedder : public Conditioner { } } - void set_weight_adapter(const std::shared_ptr& adapter) { + void set_weight_adapter(const std::shared_ptr& adapter) override { if (clip_l) { clip_l->set_weight_adapter(adapter); } @@ -1247,6 +1012,15 @@ struct FluxCLIPEmbedder : public Conditioner { } } + void runner_done() override { + if (clip_l) { + clip_l->runner_done(); + } + if (t5) { + t5->runner_done(); + } + } + std::vector, std::vector>> tokenize(std::string text, size_t min_length = 0, size_t max_length = 0) { @@ -1345,7 +1119,10 @@ struct FluxCLIPEmbedder : public Conditioner { nullptr, max_token_idx, true, - clip_skip); + clip_skip, + false, + true, + true); GGML_ASSERT(!pooled.empty()); } else { pooled = sd::Tensor::zeros({768}); @@ -1363,7 +1140,10 @@ struct FluxCLIPEmbedder : public Conditioner { sd::Tensor input_ids({static_cast(chunk_tokens.size())}, chunk_tokens); chunk_hidden_states = t5->compute(n_threads, input_ids, - sd::Tensor()); + sd::Tensor(), + false, + true, + true); GGML_ASSERT(!chunk_hidden_states.empty()); chunk_hidden_states = ::apply_token_weights(std::move(chunk_hidden_states), chunk_weights); if (zero_out_masked) { @@ -1407,11 +1187,11 @@ struct T5CLIPEmbedder : public Conditioner { bool is_umt5 = false; T5CLIPEmbedder(ggml_backend_t backend, - ggml_backend_t params_backend, - const String2TensorStorage& tensor_storage_map = {}, - bool use_mask = false, - int mask_pad = 0, - bool is_umt5 = false) + const String2TensorStorage& tensor_storage_map = {}, + bool use_mask = false, + int mask_pad = 0, + bool is_umt5 = false, + std::shared_ptr weight_manager = nullptr) : use_mask(use_mask), mask_pad(mask_pad), t5_tokenizer(is_umt5) { bool use_t5 = false; for (auto pair : tensor_storage_map) { @@ -1424,7 +1204,7 @@ struct T5CLIPEmbedder : public Conditioner { LOG_WARN("IMPORTANT NOTICE: No text encoders provided, cannot process prompts!"); return; } else { - t5 = std::make_shared(backend, params_backend, tensor_storage_map, "text_encoders.t5xxl.transformer", is_umt5); + t5 = std::make_shared(backend, tensor_storage_map, "text_encoders.t5xxl.transformer", is_umt5, weight_manager); } } @@ -1434,29 +1214,6 @@ struct T5CLIPEmbedder : public Conditioner { } } - bool alloc_params_buffer() override { - if (t5) { - if (!t5->alloc_params_buffer()) { - return false; - } - } - return true; - } - - void free_params_buffer() override { - if (t5) { - t5->free_params_buffer(); - } - } - - size_t get_params_buffer_size() override { - size_t buffer_size = 0; - if (t5) { - buffer_size += t5->get_params_buffer_size(); - } - return buffer_size; - } - void set_max_graph_vram_bytes(size_t max_vram_bytes) override { if (t5) { t5->set_max_graph_vram_bytes(max_vram_bytes); @@ -1481,6 +1238,12 @@ struct T5CLIPEmbedder : public Conditioner { } } + void runner_done() override { + if (t5) { + t5->runner_done(); + } + } + std::tuple, std::vector, std::vector> tokenize(std::string text, size_t min_length = 0, size_t max_length = 0) { @@ -1574,7 +1337,10 @@ struct T5CLIPEmbedder : public Conditioner { auto chunk_hidden_states = t5->compute(n_threads, input_ids, - t5_attn_mask_chunk); + t5_attn_mask_chunk, + false, + true, + true); GGML_ASSERT(!chunk_hidden_states.empty()); chunk_hidden_states = apply_token_weights(std::move(chunk_hidden_states), chunk_weights); @@ -1612,42 +1378,122 @@ struct T5CLIPEmbedder : public Conditioner { } }; +struct MiniT2IConditioner : public Conditioner { + T5UniGramTokenizer tokenizer; + std::shared_ptr t5; + size_t prompt_length = 256; + + MiniT2IConditioner(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map = {}, + std::shared_ptr weight_manager = nullptr) { + bool use_t5 = false; + for (const auto& pair : tensor_storage_map) { + if (pair.first.find("text_encoders.t5xxl") != std::string::npos) { + use_t5 = true; + break; + } + } + if (!use_t5) { + LOG_WARN("IMPORTANT NOTICE: No MiniT2I T5 text encoder provided, cannot process prompts!"); + return; + } + t5 = std::make_shared(backend, tensor_storage_map, "text_encoders.t5xxl.transformer", false, weight_manager); + } + + void get_param_tensors(std::map& tensors) override { + if (t5) { + t5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer"); + } + } + + void set_max_graph_vram_bytes(size_t max_vram_bytes) override { + if (t5) { + t5->set_max_graph_vram_bytes(max_vram_bytes); + } + } + + void set_stream_layers_enabled(bool enabled) override { + if (t5) { + t5->set_stream_layers_enabled(enabled); + } + } + + void set_flash_attention_enabled(bool enabled) override { + if (t5) { + t5->set_flash_attention_enabled(enabled); + } + } + + void set_weight_adapter(const std::shared_ptr& adapter) override { + if (t5) { + t5->set_weight_adapter(adapter); + } + } + + void runner_done() override { + if (t5) { + t5->runner_done(); + } + } + + SDCondition get_learned_condition(int n_threads, + const ConditionerParams& conditioner_params) override { + SDCondition result; + if (!t5) { + result.c_crossattn = sd::Tensor::zeros({1024, static_cast(prompt_length)}); + result.c_vector = sd::Tensor::zeros({static_cast(prompt_length)}); + return result; + } + + std::vector tokens = tokenizer.encode(conditioner_params.text); + if (tokens.size() > prompt_length) { + tokens.resize(prompt_length); + } + std::vector mask(tokens.size(), 1.0f); + while (tokens.size() < prompt_length) { + tokens.push_back(tokenizer.PAD_TOKEN_ID); + mask.push_back(0.0f); + } + + sd::Tensor input_ids({static_cast(tokens.size())}, tokens); + std::vector t5_mask(mask.size(), 0.0f); + for (size_t i = 0; i < mask.size(); ++i) { + t5_mask[i] = mask[i] > 0.0f ? 0.0f : -HUGE_VALF; + } + sd::Tensor hidden_states = t5->compute(n_threads, + input_ids, + sd::Tensor::from_vector(t5_mask), + false, + true, + true); + GGML_ASSERT(!hidden_states.empty()); + result.c_crossattn = std::move(hidden_states); + result.c_vector = sd::Tensor::from_vector(mask); + return result; + } +}; + struct AnimaConditioner : public Conditioner { std::shared_ptr qwen_tokenizer; T5UniGramTokenizer t5_tokenizer; std::shared_ptr llm; AnimaConditioner(ggml_backend_t backend, - ggml_backend_t params_backend, - const String2TensorStorage& tensor_storage_map = {}) { + const String2TensorStorage& tensor_storage_map = {}, + std::shared_ptr weight_manager = nullptr) { qwen_tokenizer = std::make_shared(); llm = std::make_shared(LLM::LLMArch::QWEN3, backend, - params_backend, tensor_storage_map, "text_encoders.llm", - false); + false, + weight_manager); } void get_param_tensors(std::map& tensors) override { llm->get_param_tensors(tensors, "text_encoders.llm"); } - bool alloc_params_buffer() override { - if (!llm->alloc_params_buffer()) { - return false; - } - return true; - } - - void free_params_buffer() override { - llm->free_params_buffer(); - } - - size_t get_params_buffer_size() override { - return llm->get_params_buffer_size(); - } - void set_max_graph_vram_bytes(size_t max_vram_bytes) override { llm->set_max_graph_vram_bytes(max_vram_bytes); } @@ -1664,6 +1510,10 @@ struct AnimaConditioner : public Conditioner { llm->set_weight_adapter(adapter); } + void runner_done() override { + llm->runner_done(); + } + std::tuple, std::vector, std::vector, std::vector> tokenize(std::string text) { auto parsed_attention = parse_prompt_attention(text); @@ -1721,7 +1571,11 @@ struct AnimaConditioner : public Conditioner { input_ids, sd::Tensor(), {}, - {}); + {}, + false, + false, + true, + true); GGML_ASSERT(!hidden_states.empty()); hidden_states = apply_token_weights(std::move(hidden_states), qwen_weights); auto t5_ids_tensor = sd::Tensor::from_vector(t5_tokens); @@ -1744,11 +1598,11 @@ struct LLMEmbedder : public Conditioner { std::shared_ptr llm; LLMEmbedder(ggml_backend_t backend, - ggml_backend_t params_backend, - const String2TensorStorage& tensor_storage_map = {}, - SDVersion version = VERSION_QWEN_IMAGE, - const std::string prefix = "", - bool enable_vision = false) + const String2TensorStorage& tensor_storage_map = {}, + SDVersion version = VERSION_QWEN_IMAGE, + const std::string prefix = "", + bool enable_vision = false, + std::shared_ptr weight_manager = nullptr) : version(version) { LLM::LLMArch arch = LLM::LLMArch::QWEN2_5_VL; if (version == VERSION_FLUX2) { @@ -1759,6 +1613,8 @@ struct LLMEmbedder : public Conditioner { arch = LLM::LLMArch::GPT_OSS_20B; } else if (sd_version_is_pid(version)) { arch = LLM::LLMArch::GEMMA2_2B; + } else if (sd_version_is_ideogram4(version) || sd_version_is_boogu_image(version) || sd_version_is_sefi_image(version) || sd_version_is_krea2(version)) { + arch = LLM::LLMArch::QWEN3_VL; } else if (sd_version_is_z_image(version) || version == VERSION_OVIS_IMAGE || version == VERSION_FLUX2_KLEIN) { arch = LLM::LLMArch::QWEN3; } @@ -1773,33 +1629,16 @@ struct LLMEmbedder : public Conditioner { } llm = std::make_shared(arch, backend, - params_backend, tensor_storage_map, "text_encoders.llm", - enable_vision); + enable_vision, + weight_manager); } void get_param_tensors(std::map& tensors) override { llm->get_param_tensors(tensors, "text_encoders.llm"); } - bool alloc_params_buffer() override { - if (!llm->alloc_params_buffer()) { - return false; - } - return true; - } - - void free_params_buffer() override { - llm->free_params_buffer(); - } - - size_t get_params_buffer_size() override { - size_t buffer_size = 0; - buffer_size += llm->get_params_buffer_size(); - return buffer_size; - } - void set_max_graph_vram_bytes(size_t max_vram_bytes) override { llm->set_max_graph_vram_bytes(max_vram_bytes); } @@ -1818,6 +1657,12 @@ struct LLMEmbedder : public Conditioner { } } + void runner_done() override { + if (llm) { + llm->runner_done(); + } + } + std::tuple, std::vector, std::vector> tokenize(std::string text, const std::pair& attn_range, size_t min_length = 0, @@ -1913,7 +1758,11 @@ struct LLMEmbedder : public Conditioner { input_ids, attention_mask, image_embeds, - out_layers); + out_layers, + false, + false, + true, + true); GGML_ASSERT(!hidden_states.empty()); hidden_states = apply_token_weights(std::move(hidden_states), weights); GGML_ASSERT(hidden_states.shape()[1] > prompt_template_encode_start_idx); @@ -1969,7 +1818,7 @@ struct LLMEmbedder : public Conditioner { for (int i = 0; i < conditioner_params.ref_images->size(); i++) { const auto& image = (*conditioner_params.ref_images)[i]; - double factor = llm->params.vision.patch_size * llm->params.vision.spatial_merge_size; + double factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size; int height = static_cast(image.shape()[1]); int width = static_cast(image.shape()[0]); int h_bar = static_cast(std::round(height / factor) * factor); @@ -1991,7 +1840,7 @@ struct LLMEmbedder : public Conditioner { auto resized_image = clip_preprocess(image, w_bar, h_bar); - auto image_embed = llm->encode_image(n_threads, resized_image); + auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true); GGML_ASSERT(!image_embed.empty()); image_embeds.emplace_back(image_embed_idx, image_embed); image_embed_idx += 1 + static_cast(image_embed.shape()[1]) + 6; @@ -2024,6 +1873,76 @@ struct LLMEmbedder : public Conditioner { prompt += "<|im_end|>\n<|im_start|>assistant\n"; } + } else if (sd_version_is_boogu_image(version)) { + prompt_template_encode_start_idx = 0; + + const std::string t2i_system_prompt = + "You are a helpful assistant that generates high-quality images based on user instructions. The instructions are as follows."; + const std::string edit_system_prompt = + "Describe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate."; + const bool has_ref_images = llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty(); + const bool text_empty = conditioner_params.text.find_first_not_of(" \t\r\n") == std::string::npos; + + if (has_ref_images) { + LOG_INFO("BooguImageEditPipeline"); + const std::string prompt_prefix = "<|im_start|>system\n" + edit_system_prompt + "<|im_end|>\n<|im_start|>user\n"; + std::string img_prompt; + const std::string placeholder = "<|image_pad|>"; + + for (int i = 0; i < conditioner_params.ref_images->size(); i++) { + const auto& image = (*conditioner_params.ref_images)[i]; + double factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size; + int height = static_cast(image.shape()[1]); + int width = static_cast(image.shape()[0]); + double beta = std::sqrt((384.0 * 384.0) / (static_cast(height) * static_cast(width))); + int h_bar = std::max(static_cast(factor), + static_cast(std::round(height * beta / factor)) * static_cast(factor)); + int w_bar = std::max(static_cast(factor), + static_cast(std::round(width * beta / factor)) * static_cast(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, false, true, true); + GGML_ASSERT(!image_embed.empty()); + + std::string image_prefix = prompt_prefix + img_prompt + "<|vision_start|>"; + int image_embed_idx = static_cast(tokenizer->encode(image_prefix, nullptr).size()); + image_embeds.emplace_back(image_embed_idx, image_embed); + + img_prompt += "<|vision_start|>"; + int64_t num_image_tokens = image_embed.shape()[1]; + img_prompt.reserve(img_prompt.size() + static_cast(num_image_tokens) * placeholder.size() + 32); + for (int j = 0; j < num_image_tokens; j++) { + img_prompt += placeholder; + } + img_prompt += "<|vision_end|>"; + } + + prompt = prompt_prefix + img_prompt; + prompt_attn_range.first = static_cast(prompt.size()); + prompt += conditioner_params.text; + prompt_attn_range.second = static_cast(prompt.size()); + prompt += "<|im_end|>\n"; + } else { + const std::string& system_prompt = text_empty ? edit_system_prompt : t2i_system_prompt; + prompt = "<|im_start|>system\n" + system_prompt + "<|im_end|>\n<|im_start|>user\n"; + prompt_attn_range.first = static_cast(prompt.size()); + prompt += conditioner_params.text; + prompt_attn_range.second = static_cast(prompt.size()); + prompt += "<|im_end|>\n"; + } + } else if (sd_version_is_krea2(version)) { + prompt_template_encode_start_idx = 34; + out_layers = {2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35}; + + prompt = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n"; + + prompt_attn_range.first = static_cast(prompt.size()); + prompt += conditioner_params.text; + prompt_attn_range.second = static_cast(prompt.size()); + + prompt += "<|im_end|>\n<|im_start|>assistant\n"; } else if (sd_version_is_longcat(version)) { spell_quotes = true; @@ -2040,7 +1959,7 @@ struct LLMEmbedder : public Conditioner { for (int i = 0; i < conditioner_params.ref_images->size(); i++) { const auto& image = (*conditioner_params.ref_images)[i]; - double factor = llm->params.vision.patch_size * llm->params.vision.spatial_merge_size; + double factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size; int height = static_cast(image.shape()[1]); int width = static_cast(image.shape()[0]); int h_bar = static_cast(std::round(height / factor) * factor); @@ -2061,7 +1980,7 @@ struct LLMEmbedder : public Conditioner { 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); + auto image_embed = llm->encode_image(n_threads, resized_image, false, true, true); GGML_ASSERT(!image_embed.empty()); image_embeds.emplace_back(image_embed_idx, image_embed); image_embed_idx += 1 + static_cast(image_embed.shape()[1]) + 6; @@ -2101,6 +2020,14 @@ struct LLMEmbedder : public Conditioner { prompt_attn_range.second = static_cast(prompt.size()); prompt += "[/INST]"; + } else if (sd_version_is_ideogram4(version)) { + prompt_template_encode_start_idx = 0; + out_layers = {1, 4, 7, 10, 13, 16, 19, 22, 25, 28, 31, 34, 36}; + + prompt = "<|im_start|>user\n"; + prompt += conditioner_params.text; + prompt += "<|im_end|>\n<|im_start|>assistant\n"; + prompt_attn_range = {0, 0}; } else if (sd_version_is_ernie_image(version)) { prompt_template_encode_start_idx = 0; out_layers = {25}; // -2 @@ -2165,6 +2092,18 @@ struct LLMEmbedder : public Conditioner { prompt_attn_range.second = static_cast(prompt.size()); prompt += "<|im_end|>\n<|im_start|>assistant\n\n\n\n\n"; + } else if (sd_version_is_sefi_image(version)) { + prompt_template_encode_start_idx = 0; + min_length = 1024; + out_layers = {9, 18, 27}; + + prompt = "<|im_start|>user\n"; + + prompt_attn_range.first = static_cast(prompt.size()); + prompt += conditioner_params.text; + prompt_attn_range.second = static_cast(prompt.size()); + + prompt += "<|im_end|>\n<|im_start|>assistant\n"; } else if (version == VERSION_OVIS_IMAGE) { prompt_template_encode_start_idx = 28; min_length = prompt_template_encode_start_idx + 256; @@ -2296,10 +2235,10 @@ struct LTXAVTextProjectionRunner : public GGMLRunner { LTXAVTextProjection model; LTXAVTextProjectionRunner(ggml_backend_t backend, - ggml_backend_t params_backend, - const String2TensorStorage& tensor_storage_map = {}, - const std::string& prefix = "") - : GGMLRunner(backend, params_backend), + const String2TensorStorage& tensor_storage_map = {}, + const std::string& prefix = "", + std::shared_ptr weight_manager = nullptr) + : GGMLRunner(backend, weight_manager), model(tensor_storage_map.find(prefix + ".video_aggregate_embed.weight") != tensor_storage_map.end()) { model.init(params_ctx, tensor_storage_map, prefix); } @@ -2321,11 +2260,15 @@ struct LTXAVTextProjectionRunner : public GGMLRunner { return gf; } - sd::Tensor compute(int n_threads, const sd::Tensor& x) { + sd::Tensor compute(int n_threads, + const sd::Tensor& x, + bool auto_free = true, + bool free_compute_buffer = true, + bool free_compute_params = true) { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x); }; - return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true)); + return take_or_empty(GGMLRunner::compute(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params)); } }; @@ -2340,22 +2283,22 @@ struct LTXAVEmbedder : public Conditioner { bool dual_projection = false; LTXAVEmbedder(ggml_backend_t backend, - ggml_backend_t params_backend, - const String2TensorStorage& tensor_storage_map = {}, - const std::string& llm_prefix = "text_encoders.llm", - const std::string& projector_prefix = "text_embedding_projection") { + const String2TensorStorage& tensor_storage_map = {}, + const std::string& llm_prefix = "text_encoders.llm", + const std::string& projector_prefix = "text_embedding_projection", + std::shared_ptr weight_manager = nullptr) { tokenizer = std::make_shared(); llm = std::make_shared(LLM::LLMArch::GEMMA3_12B, backend, - params_backend, tensor_storage_map, llm_prefix, - false); + false, + weight_manager); dual_projection = tensor_storage_map.find(projector_prefix + ".video_aggregate_embed.weight") != tensor_storage_map.end(); projector = std::make_shared(backend, - params_backend, tensor_storage_map, - projector_prefix); + projector_prefix, + weight_manager); } void get_param_tensors(std::map& tensors) override { @@ -2363,25 +2306,6 @@ struct LTXAVEmbedder : public Conditioner { projector->get_param_tensors(tensors, "text_embedding_projection"); } - bool alloc_params_buffer() override { - if (!llm->alloc_params_buffer()) { - return false; - } - if (!projector->alloc_params_buffer()) { - return false; - } - return true; - } - - void free_params_buffer() override { - llm->free_params_buffer(); - projector->free_params_buffer(); - } - - size_t get_params_buffer_size() override { - return llm->get_params_buffer_size() + projector->get_params_buffer_size(); - } - void set_flash_attention_enabled(bool enabled) override { llm->set_flash_attention_enabled(enabled); projector->set_flash_attention_enabled(enabled); @@ -2397,6 +2321,11 @@ struct LTXAVEmbedder : public Conditioner { projector->set_weight_adapter(adapter); } + void runner_done() override { + llm->runner_done(); + projector->runner_done(); + } + std::tuple, std::vector, std::vector> tokenize(std::string text, const std::pair& attn_range) { std::vector> parsed_attention; @@ -2460,6 +2389,9 @@ struct LTXAVEmbedder : public Conditioner { attention_mask, {}, {}, + true, + false, + true, true); GGML_ASSERT(!hidden_states.empty()); hidden_states = apply_token_weights(std::move(hidden_states), weights); @@ -2519,7 +2451,7 @@ struct LTXAVEmbedder : public Conditioner { } hidden_states.reshape_({kNumStates * kHiddenSize, valid_tokens}); - return projector->compute(n_threads, hidden_states); + return projector->compute(n_threads, hidden_states, false, true, true); } SDCondition get_learned_condition(int n_threads, @@ -2544,4 +2476,4 @@ struct LTXAVEmbedder : public Conditioner { } }; -#endif +#endif // __SD_CONDITIONING_CONDITIONER_HPP__ diff --git a/src/convert.cpp b/src/convert.cpp index cc1cdd7e..0b7fe2cf 100644 --- a/src/convert.cpp +++ b/src/convert.cpp @@ -3,9 +3,9 @@ #include #include -#include "model.h" #include "model_io/gguf_io.h" #include "model_io/safetensors_io.h" +#include "model_loader.h" #include "util.h" #include "ggml_extend_backend.h" @@ -76,30 +76,23 @@ static bool load_tensors_for_export(ModelLoader& model_loader, return success; } -bool convert(const char* input_path, - const char* vae_path, - const char* output_path, - sd_type_t output_type, - const char* tensor_type_rules, - bool convert_name) { - ModelLoader model_loader; - - if (!model_loader.init_from_file(input_path)) { - LOG_ERROR("init model loader from file failed: '%s'", input_path); +static bool init_convert_path(ModelLoader& model_loader, const char* path, const char* prefix, bool& loaded_any) { + if (path == nullptr || strlen(path) == 0) { + return true; + } + if (!model_loader.init_from_file(path, prefix)) { + LOG_ERROR("init model loader from file failed: '%s'", path); return false; } + loaded_any = true; + return true; +} - if (vae_path != nullptr && strlen(vae_path) > 0) { - if (!model_loader.init_from_file(vae_path, "vae.")) { - LOG_ERROR("init model loader from file failed: '%s'", vae_path); - return false; - } - } - if (convert_name) { - model_loader.convert_tensors_name(); - } - - ggml_type type = (ggml_type)output_type; +static bool export_loaded_model(ModelLoader& model_loader, + const char* output_path, + sd_type_t output_type, + const char* tensor_type_rules) { + ggml_type type = sd_type_to_ggml_type(output_type); bool output_is_safetensors = ends_with(output_path, ".safetensors"); TensorTypeRules type_rules = parse_tensor_type_rules(tensor_type_rules); @@ -136,3 +129,55 @@ bool convert(const char* input_path, ggml_free(ggml_ctx); return success; } + +bool convert_with_components(const char* model_path, + const char* clip_l_path, + const char* clip_g_path, + const char* t5xxl_path, + const char* diffusion_model_path, + const char* vae_path, + const char* output_path, + sd_type_t output_type, + const char* tensor_type_rules, + bool convert_name) { + ModelLoader model_loader; + bool loaded_any = false; + + if (!init_convert_path(model_loader, model_path, "", loaded_any) || + !init_convert_path(model_loader, clip_l_path, "text_encoders.clip_l.transformer.", loaded_any) || + !init_convert_path(model_loader, clip_g_path, "text_encoders.clip_g.transformer.", loaded_any) || + !init_convert_path(model_loader, t5xxl_path, "text_encoders.t5xxl.transformer.", loaded_any) || + !init_convert_path(model_loader, diffusion_model_path, "model.diffusion_model.", loaded_any) || + !init_convert_path(model_loader, vae_path, "vae.", loaded_any)) { + return false; + } + + if (!loaded_any) { + LOG_ERROR("no input model path provided for convert"); + return false; + } + + if (convert_name) { + model_loader.convert_tensors_name(); + } + + return export_loaded_model(model_loader, output_path, output_type, tensor_type_rules); +} + +bool convert(const char* input_path, + const char* vae_path, + const char* output_path, + sd_type_t output_type, + const char* tensor_type_rules, + bool convert_name) { + return convert_with_components(input_path, + nullptr, + nullptr, + nullptr, + nullptr, + vae_path, + output_path, + output_type, + tensor_type_rules, + convert_name); +} diff --git a/src/ggml_extend.hpp b/src/core/ggml_extend.hpp similarity index 81% rename from src/ggml_extend.hpp rename to src/core/ggml_extend.hpp index 1ad31a17..776c18d3 100644 --- a/src/ggml_extend.hpp +++ b/src/core/ggml_extend.hpp @@ -1,5 +1,5 @@ -#ifndef __GGML_EXTEND_HPP__ -#define __GGML_EXTEND_HPP__ +#ifndef __SD_CORE_GGML_EXTEND_HPP__ +#define __SD_CORE_GGML_EXTEND_HPP__ #include #include @@ -23,19 +23,19 @@ #include #include +#include "core/ggml_extend_backend.h" +#include "core/ggml_graph_cut.h" #include "ggml-alloc.h" #include "ggml-backend.h" #include "ggml.h" -#include "ggml_extend_backend.h" -#include "ggml_graph_cut.h" -#include "layer_registry.h" +#include "core/tensor.hpp" #include "model.h" -#include "tensor.hpp" -#include "rng.hpp" -#include "tensor_ggml.hpp" -#include "util.h" +#include "core/rng.hpp" +#include "core/tensor_ggml.hpp" +#include "core/util.h" +#include "weight_manager.h" #define EPS 1e-05f @@ -1346,10 +1346,18 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, v_in = ggml_cast(ctx, v_in, GGML_TYPE_F16); if (mask_in != nullptr) { - mask_in = ggml_transpose(ctx, mask_in); - } - - if (mask_in != nullptr) { + // ggml_flash_attn_ext expects the mask as a contiguous F16 tensor shaped + // [n_kv, n_q, (heads), (batch)] (ne0 = key length, ne1 = query length) and, + // unlike the manual-attention path, does not broadcast the query dimension. + // Some callers (e.g. Chroma/T5) pass a per-key padding mask broadcast over + // queries ([n_kv, 1, ...]); materialize the query dimension to L_q so the + // kernel indexes it correctly. (A bare ggml_transpose here produced a + // [1, n_kv, ...] mask that the kernel silently misreads, yielding NaN/blank + // output for masked flash attention.) + if (mask_in->ne[1] != L_q) { + mask_in = ggml_repeat(ctx, mask_in, + ggml_new_tensor_4d(ctx, mask_in->type, mask_in->ne[0], L_q, mask_in->ne[2], mask_in->ne[3])); + } mask_in = ggml_cast(ctx, mask_in, GGML_TYPE_F16); } @@ -1374,7 +1382,16 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, if (!ggml_backend_supports_op(backend, kqv)) { kqv = nullptr; } else { - kqv = ggml_view_3d(ctx, kqv, d_head, n_head, L_q, kqv->nb[1], kqv->nb[2], 0); + kqv = ggml_view_4d(ctx, + kqv, + d_head, + n_head, + L_q, + N, + kqv->nb[1], + kqv->nb[2], + kqv->nb[1] * n_head, + 0); } } } @@ -1647,6 +1664,7 @@ struct GGMLRunnerContext { std::vector>* debug_tensors = nullptr; std::function get_cache_tensor; std::function cache_tensor; + std::function set_backend_tensor_data; void capture_tensor(const std::string& name, ggml_tensor* tensor) { if (debug_tensors == nullptr || tensor == nullptr) { @@ -1672,6 +1690,13 @@ struct GGMLRunnerContext { } cache_tensor(name, tensor); } + + void bind_backend_tensor_data(ggml_tensor* tensor, const void* data) const { + if (!set_backend_tensor_data || tensor == nullptr || data == nullptr) { + return; + } + set_backend_tensor_data(tensor, data); + } }; struct GGMLRunner { @@ -1680,14 +1705,9 @@ protected: using GraphCutSegment = sd::ggml_graph_cut::Segment; using GraphCutPlan = sd::ggml_graph_cut::Plan; - ggml_backend_t params_backend = nullptr; ggml_backend_t runtime_backend = nullptr; - ggml_context* params_ctx = nullptr; - ggml_backend_buffer_t params_buffer = nullptr; - ggml_context* offload_ctx = nullptr; - ggml_backend_buffer_t runtime_params_buffer = nullptr; - bool params_on_runtime_backend = false; + ggml_context* params_ctx = nullptr; ggml_context* cache_ctx = nullptr; ggml_backend_buffer_t cache_buffer = nullptr; @@ -1695,23 +1715,16 @@ protected: ggml_context* compute_ctx = nullptr; ggml_gallocr* compute_allocr = nullptr; - ggml_context* partial_offload_ctx = nullptr; - ggml_backend_buffer_t partial_runtime_params_buffer = nullptr; - std::vector> partial_offload_pairs; - - // Params kept on the runtime backend across streaming segments. - ggml_context* resident_offload_ctx = nullptr; - std::vector> resident_offload_pairs; - ggml_backend_buffer_t resident_runtime_params_buffer = nullptr; - std::unordered_set resident_param_set; - uint64_t resident_state_token = 0; - - size_t max_graph_vram_bytes = 0; - bool stream_layers_enabled = false; - - sd::layer_registry::LayerRegistry layer_registry_; + size_t max_graph_vram_bytes = 0; + bool stream_layers_enabled = false; + size_t observed_max_effective_budget_ = 0; std::shared_ptr weight_adapter = nullptr; + std::weak_ptr weight_manager; + std::unordered_set kept_compute_param_tensor_set; + std::vector runner_param_tensors; + std::unordered_set runner_param_tensor_set; + bool params_tensor_set_dirty_ = true; std::vector one_vec = {1.f}; ggml_tensor* one_tensor = nullptr; @@ -1767,10 +1780,7 @@ protected: params_ctx = ggml_init(params); GGML_ASSERT(params_ctx != nullptr); params_tensor_set_.clear(); - if (params_backend != runtime_backend) { - offload_ctx = ggml_init(params); - GGML_ASSERT(offload_ctx != nullptr); - } + params_tensor_set_dirty_ = true; } void free_params_ctx() { @@ -1779,14 +1789,7 @@ protected: params_ctx = nullptr; } params_tensor_set_.clear(); - if (offload_ctx != nullptr) { - ggml_free(offload_ctx); - offload_ctx = nullptr; - } - if (partial_offload_ctx != nullptr) { - ggml_free(partial_offload_ctx); - partial_offload_ctx = nullptr; - } + params_tensor_set_dirty_ = true; } void alloc_cache_ctx() { @@ -1826,6 +1829,9 @@ protected: } void rebuild_params_tensor_set() { + if (!params_tensor_set_dirty_) { + return; + } params_tensor_set_.clear(); if (params_ctx == nullptr) { return; @@ -1833,6 +1839,93 @@ protected: for (ggml_tensor* t = ggml_get_first_tensor(params_ctx); t != nullptr; t = ggml_get_next_tensor(params_ctx, t)) { params_tensor_set_.insert(t); } + params_tensor_set_dirty_ = false; + } + + std::vector collect_used_param_tensors(ggml_cgraph* gf) { + std::vector used_params; + rebuild_params_tensor_set(); + if (gf == nullptr || params_tensor_set_.empty()) { + return used_params; + } + + std::unordered_set seen_params; + const int n_leafs = sd::ggml_graph_cut::leaf_count(gf); + seen_params.reserve(static_cast(n_leafs)); + for (int i = 0; i < n_leafs; ++i) { + ggml_tensor* leaf = sd::ggml_graph_cut::leaf_tensor(gf, i); + ggml_tensor* param_leaf = leaf; + if (param_leaf != nullptr && params_tensor_set_.find(param_leaf) == params_tensor_set_.end()) { + param_leaf = param_leaf->view_src; + } + if (param_leaf != nullptr && + params_tensor_set_.find(param_leaf) != params_tensor_set_.end() && + seen_params.insert(param_leaf).second) { + used_params.push_back(param_leaf); + } + } + return used_params; + } + + bool prepare_execute_graph_weights(ggml_cgraph* gf, + std::vector& graph_param_tensors, + std::vector& params_to_prepare, + bool keep_compute_params) { + graph_param_tensors = collect_used_param_tensors(gf); + params_to_prepare.clear(); + params_to_prepare.reserve(graph_param_tensors.size()); + for (ggml_tensor* param : graph_param_tensors) { + if (param == nullptr) { + continue; + } + if (keep_compute_params && + kept_compute_param_tensor_set.find(param) != kept_compute_param_tensor_set.end()) { + continue; + } + params_to_prepare.push_back(param); + } + auto manager = weight_manager.lock(); + if (manager == nullptr) { + if (!params_to_prepare.empty()) { + LOG_ERROR("%s weight manager is not set for graph params", get_desc().c_str()); + return false; + } + return true; + } + + if (!manager->prepare_params(params_to_prepare)) { + LOG_ERROR("%s prepare graph weights failed", get_desc().c_str()); + return false; + } + for (ggml_tensor* param : params_to_prepare) { + if (param == nullptr) { + continue; + } + if (runner_param_tensor_set.insert(param).second) { + runner_param_tensors.push_back(param); + } + } + return true; + } + + void free_compute_backend_param_tensors(const std::vector& tensors) { + if (tensors.empty()) { + return; + } + auto manager = weight_manager.lock(); + if (manager != nullptr) { + manager->release_compute_backend_params(tensors); + } + } + + void free_params_backend_param_tensors(const std::vector& tensors) { + if (tensors.empty()) { + return; + } + auto manager = weight_manager.lock(); + if (manager != nullptr) { + manager->release_params_backend_params(tensors); + } } void prepare_build_in_tensor_before() { @@ -1923,6 +2016,10 @@ protected: } bool copy_cache_tensors_to_cache_buffer(const std::unordered_set* cache_keep_names = nullptr) { + if (cache_tensor_map.empty() && cache_keep_names == nullptr) { + return true; + } + ggml_context* old_cache_ctx = cache_ctx; ggml_backend_buffer_t old_cache_buffer = cache_buffer; cache_ctx = nullptr; @@ -2100,328 +2197,16 @@ protected: } } - bool offload_all_params() { - restore_partial_params(); - if (params_backend == runtime_backend) { - return true; - } - if (params_on_runtime_backend) { - return true; - } - GGML_ASSERT(runtime_params_buffer == nullptr); - int64_t t0 = ggml_time_ms(); - size_t num_tensors = ggml_tensor_num(offload_ctx); - if (num_tensors == 0) { - for (ggml_tensor* t = ggml_get_first_tensor(params_ctx); t != nullptr; t = ggml_get_next_tensor(params_ctx, t)) { - GGML_ASSERT(t->view_src == nullptr); - ggml_dup_tensor(offload_ctx, t); - } - } - num_tensors = ggml_tensor_num(offload_ctx); - GGML_ASSERT(num_tensors == ggml_tensor_num(params_ctx)); - - runtime_params_buffer = ggml_backend_alloc_ctx_tensors(offload_ctx, runtime_backend); - - if (runtime_params_buffer == nullptr) { - LOG_ERROR("%s alloc runtime params backend buffer failed, num_tensors = %i", - get_desc().c_str(), - num_tensors); - return false; - } - ggml_backend_buffer_set_usage(runtime_params_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); - - ggml_tensor* t = ggml_get_first_tensor(params_ctx); - ggml_tensor* offload_t = ggml_get_first_tensor(offload_ctx); - - while (t != nullptr && offload_t != nullptr) { - ggml_backend_tensor_copy(t, offload_t); - std::swap(t->buffer, offload_t->buffer); - std::swap(t->data, offload_t->data); - std::swap(t->extra, offload_t->extra); - - t = ggml_get_next_tensor(params_ctx, t); - offload_t = ggml_get_next_tensor(offload_ctx, offload_t); - } - - int64_t t1 = ggml_time_ms(); - - size_t params_buffer_size = ggml_backend_buffer_get_size(runtime_params_buffer); - LOG_INFO("%s offload params (%6.2f MB, %i tensors) to runtime backend (%s), taking %.2fs", - get_desc().c_str(), - params_buffer_size / (1024.f * 1024.f), - num_tensors, - ggml_backend_name(runtime_backend), - (t1 - t0) * 1.0f / 1000); - - params_on_runtime_backend = true; - - return true; - } - - bool offload_partial_params(const std::vector& tensors) { - restore_partial_params(); - if (params_backend == runtime_backend) { - return true; - } - if (tensors.empty()) { - return true; - } - GGML_ASSERT(!params_on_runtime_backend); - GGML_ASSERT(partial_runtime_params_buffer == nullptr); - - std::vector unique_tensors; - std::unordered_set seen_tensors; - unique_tensors.reserve(tensors.size()); - seen_tensors.reserve(tensors.size()); - for (ggml_tensor* tensor : tensors) { - if (tensor == nullptr) { - continue; - } - if (resident_param_set.find(tensor) != resident_param_set.end()) { - continue; - } - if (seen_tensors.insert(tensor).second) { - unique_tensors.push_back(tensor); - } - } - if (unique_tensors.empty()) { - return true; - } - - ggml_init_params params; - params.mem_size = std::max(1, unique_tensors.size()) * ggml_tensor_overhead(); - params.mem_buffer = nullptr; - params.no_alloc = true; - - partial_offload_ctx = ggml_init(params); - GGML_ASSERT(partial_offload_ctx != nullptr); - - partial_offload_pairs.clear(); - partial_offload_pairs.reserve(unique_tensors.size()); - - for (ggml_tensor* tensor : unique_tensors) { - GGML_ASSERT(tensor->view_src == nullptr); - ggml_tensor* offload_tensor = ggml_dup_tensor(partial_offload_ctx, tensor); - ggml_set_name(offload_tensor, tensor->name); - partial_offload_pairs.push_back({tensor, offload_tensor}); - } - - partial_runtime_params_buffer = ggml_backend_alloc_ctx_tensors(partial_offload_ctx, runtime_backend); - if (partial_runtime_params_buffer == nullptr) { - LOG_ERROR("%s alloc partial runtime params backend buffer failed, num_tensors = %zu", - get_desc().c_str(), - partial_offload_pairs.size()); - ggml_free(partial_offload_ctx); - partial_offload_ctx = nullptr; - partial_offload_pairs.clear(); - return false; - } - ggml_backend_buffer_set_usage(partial_runtime_params_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); - - for (auto& pair : partial_offload_pairs) { - ggml_tensor* tensor = pair.first; - ggml_tensor* offload_tensor = pair.second; - - ggml_backend_tensor_copy(tensor, offload_tensor); - std::swap(tensor->buffer, offload_tensor->buffer); - std::swap(tensor->data, offload_tensor->data); - std::swap(tensor->extra, offload_tensor->extra); - } - - size_t params_buffer_size = ggml_backend_buffer_get_size(partial_runtime_params_buffer); - LOG_DEBUG("%s offload partial params (%6.2f MB, %zu tensors) to runtime backend (%s)", - get_desc().c_str(), - params_buffer_size / (1024.f * 1024.f), - partial_offload_pairs.size(), - ggml_backend_name(runtime_backend)); - - return true; - } - - void restore_all_params() { - restore_partial_params(); - if (!params_on_runtime_backend) { - return; - } - ggml_tensor* t = ggml_get_first_tensor(params_ctx); - ggml_tensor* offload_t = ggml_get_first_tensor(offload_ctx); - - while (t != nullptr && offload_t != nullptr) { - t->buffer = offload_t->buffer; - t->data = offload_t->data; - t->extra = offload_t->extra; - offload_t->buffer = nullptr; - offload_t->data = nullptr; - offload_t->extra = nullptr; - - t = ggml_get_next_tensor(params_ctx, t); - offload_t = ggml_get_next_tensor(offload_ctx, offload_t); - } - - if (runtime_params_buffer != nullptr) { - ggml_backend_buffer_free(runtime_params_buffer); - runtime_params_buffer = nullptr; - } - params_on_runtime_backend = false; - } - - void restore_partial_params() { - if (partial_offload_pairs.empty()) { - if (partial_runtime_params_buffer != nullptr) { - ggml_backend_buffer_free(partial_runtime_params_buffer); - partial_runtime_params_buffer = nullptr; - } - if (partial_offload_ctx != nullptr) { - ggml_free(partial_offload_ctx); - partial_offload_ctx = nullptr; - } - return; - } - - for (auto& pair : partial_offload_pairs) { - ggml_tensor* tensor = pair.first; - ggml_tensor* offload_tensor = pair.second; - - tensor->buffer = offload_tensor->buffer; - tensor->data = offload_tensor->data; - tensor->extra = offload_tensor->extra; - offload_tensor->buffer = nullptr; - offload_tensor->data = nullptr; - offload_tensor->extra = nullptr; - } - - if (partial_runtime_params_buffer != nullptr) { - ggml_backend_buffer_free(partial_runtime_params_buffer); - partial_runtime_params_buffer = nullptr; - } - partial_offload_pairs.clear(); - - if (partial_offload_ctx != nullptr) { - ggml_free(partial_offload_ctx); - partial_offload_ctx = nullptr; - } - } - - bool offload_resident_params(const std::vector& tensors) { - if (params_backend == runtime_backend) { - return true; - } - if (tensors.empty()) { - return true; - } - GGML_ASSERT(resident_runtime_params_buffer == nullptr); - GGML_ASSERT(resident_offload_ctx == nullptr); - GGML_ASSERT(resident_offload_pairs.empty()); - GGML_ASSERT(resident_param_set.empty()); - - std::vector unique_tensors; - std::unordered_set seen; - unique_tensors.reserve(tensors.size()); - seen.reserve(tensors.size()); - for (ggml_tensor* t : tensors) { - if (t == nullptr) - continue; - if (seen.insert(t).second) - unique_tensors.push_back(t); - } - if (unique_tensors.empty()) - return true; - - ggml_init_params init = {}; - init.mem_size = std::max(1, unique_tensors.size()) * ggml_tensor_overhead(); - init.mem_buffer = nullptr; - init.no_alloc = true; - resident_offload_ctx = ggml_init(init); - GGML_ASSERT(resident_offload_ctx != nullptr); - - resident_offload_pairs.reserve(unique_tensors.size()); - for (ggml_tensor* t : unique_tensors) { - GGML_ASSERT(t->view_src == nullptr); - ggml_tensor* twin = ggml_dup_tensor(resident_offload_ctx, t); - ggml_set_name(twin, t->name); - resident_offload_pairs.push_back({t, twin}); - } - - resident_runtime_params_buffer = ggml_backend_alloc_ctx_tensors(resident_offload_ctx, runtime_backend); - if (resident_runtime_params_buffer == nullptr) { - LOG_ERROR("%s alloc resident runtime params backend buffer failed, num_tensors = %zu", - get_desc().c_str(), resident_offload_pairs.size()); - ggml_free(resident_offload_ctx); - resident_offload_ctx = nullptr; - resident_offload_pairs.clear(); - return false; - } - ggml_backend_buffer_set_usage(resident_runtime_params_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); - - for (auto& pair : resident_offload_pairs) { - ggml_tensor* t = pair.first; - ggml_tensor* twin = pair.second; - ggml_backend_tensor_copy(t, twin); - std::swap(t->buffer, twin->buffer); - std::swap(t->data, twin->data); - std::swap(t->extra, twin->extra); - resident_param_set.insert(t); - } - ggml_backend_synchronize(runtime_backend); - - size_t sz = ggml_backend_buffer_get_size(resident_runtime_params_buffer); - LOG_INFO("%s offload resident params (%6.2f MB, %zu tensors) to runtime backend (%s)", - get_desc().c_str(), - sz / (1024.f * 1024.f), - resident_offload_pairs.size(), - ggml_backend_name(runtime_backend)); - return true; - } - - void restore_resident_params() { - if (resident_offload_pairs.empty()) { - if (resident_runtime_params_buffer != nullptr) { - ggml_backend_buffer_free(resident_runtime_params_buffer); - resident_runtime_params_buffer = nullptr; - } - if (resident_offload_ctx != nullptr) { - ggml_free(resident_offload_ctx); - resident_offload_ctx = nullptr; - } - resident_param_set.clear(); - resident_state_token = 0; - return; - } - for (auto& pair : resident_offload_pairs) { - ggml_tensor* t = pair.first; - ggml_tensor* twin = pair.second; - t->buffer = twin->buffer; - t->data = twin->data; - t->extra = twin->extra; - twin->buffer = nullptr; - twin->data = nullptr; - twin->extra = nullptr; - } - if (resident_runtime_params_buffer != nullptr) { - ggml_backend_buffer_free(resident_runtime_params_buffer); - resident_runtime_params_buffer = nullptr; - } - resident_offload_pairs.clear(); - if (resident_offload_ctx != nullptr) { - ggml_free(resident_offload_ctx); - resident_offload_ctx = nullptr; - } - resident_param_set.clear(); - resident_state_token = 0; - } - bool should_use_graph_cut_segmented_compute(const GraphCutPlan& plan) { return plan.has_cuts && plan.valid && max_graph_vram_bytes > 0 && plan.segments.size() > 1 && - params_backend != runtime_backend && !sd_backend_is_cpu(runtime_backend); } bool can_attempt_graph_cut_segmented_compute() const { return max_graph_vram_bytes > 0 && - params_backend != runtime_backend && !sd_backend_is_cpu(runtime_backend); } @@ -2431,44 +2216,73 @@ protected: GGML_ASSERT(plan_out != nullptr); GGML_ASSERT(gf != nullptr); - // Keep the plan and resident params under the same live-VRAM cap. - // Add back our own resident buffer so we don't see chunk-K's - // allocation as "taken" VRAM and shrink the budget on every step. size_t effective_budget = max_graph_vram_bytes; if (stream_layers_enabled && max_graph_vram_bytes > 0 && runtime_backend != nullptr) { ggml_backend_dev_t dev = ggml_backend_get_device(runtime_backend); if (dev != nullptr && ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_CPU) { size_t free_vram = 0, total_vram = 0; ggml_backend_dev_memory(dev, &free_vram, &total_vram); - if (resident_runtime_params_buffer != nullptr) { - free_vram += ggml_backend_buffer_get_size(resident_runtime_params_buffer); - } constexpr size_t safety_margin = 512ull * 1024 * 1024; size_t free_clamp = (free_vram > safety_margin) ? (free_vram - safety_margin) : 0; if (free_clamp < effective_budget) { - LOG_INFO("%s clamping streaming budget: actual free VRAM %.2f MB < user cap %.2f MB", - get_desc().c_str(), - free_clamp / (1024.0 * 1024.0), - effective_budget / (1024.0 * 1024.0)); + LOG_DEBUG("%s clamping streaming budget: actual free VRAM %.2f MB < user cap %.2f MB", + get_desc().c_str(), + free_clamp / (1024.0 * 1024.0), + effective_budget / (1024.0 * 1024.0)); effective_budget = free_clamp; } } } + bool budget_increased = false; + if (stream_layers_enabled) { + if (effective_budget > observed_max_effective_budget_) { + observed_max_effective_budget_ = effective_budget; + budget_increased = true; + } else { + effective_budget = observed_max_effective_budget_; + } + } + if (effective_budget_out != nullptr) { *effective_budget_out = effective_budget; } + // When streaming and the model dwarfs the budget, cap the planner at + // a quarter so it builds smaller merged segments and chunk-K can fit + // alongside. Without streaming the cap only adds dispatch overhead. + size_t planner_budget = effective_budget; + if (stream_layers_enabled) { + size_t total_params_bytes = 0; + for (const ggml_tensor* t : params_tensor_set_) { + if (t != nullptr) { + total_params_bytes += ggml_nbytes(t); + } + } + if (total_params_bytes * 4 > effective_budget * 3) { + planner_budget = effective_budget / 4; + } + } + *plan_out = sd::ggml_graph_cut::resolve_plan(runtime_backend, gf, &graph_cut_plan_cache_, - effective_budget, + planner_budget, params_tensor_set_, get_desc().c_str()); if (stream_layers_enabled) { - LOG_INFO("%s streaming budget = %.2f MB", - get_desc().c_str(), - effective_budget / (1024.0 * 1024.0)); + sd::ggml_graph_cut::annotate_residency(*plan_out, effective_budget); + } + if (stream_layers_enabled) { + if (budget_increased) { + LOG_INFO("%s streaming budget = %.2f MB", + get_desc().c_str(), + effective_budget / (1024.0 * 1024.0)); + } else { + LOG_DEBUG("%s streaming budget = %.2f MB", + get_desc().c_str(), + effective_budget / (1024.0 * 1024.0)); + } } return true; } @@ -2594,310 +2408,176 @@ protected: template std::optional> execute_graph(ggml_cgraph* gf, int n_threads, - bool free_compute_buffer_immediately, - const std::vector& runtime_param_tensors, + bool free_compute_buffer, + bool free_compute_params, bool preserve_backend_tensor_data_map, bool no_return = false, const std::unordered_set* cache_keep_names = nullptr) { - int64_t t_execute_begin = ggml_time_ms(); - const bool use_partial_param_offload = !runtime_param_tensors.empty(); - int64_t t_offload_begin = ggml_time_ms(); - if (use_partial_param_offload) { - if (!offload_partial_params(runtime_param_tensors)) { - LOG_ERROR("%s offload partial params to runtime backend failed", get_desc().c_str()); - return std::nullopt; - } - } else { - if (!offload_all_params()) { - LOG_ERROR("%s offload params to runtime backend failed", get_desc().c_str()); - return std::nullopt; - } - } - int64_t t_offload_end = ggml_time_ms(); - - int64_t t_alloc_begin = ggml_time_ms(); - if (!alloc_compute_buffer(gf)) { - LOG_ERROR("%s alloc compute buffer failed", get_desc().c_str()); - if (use_partial_param_offload) { - restore_partial_params(); - } + std::vector graph_param_tensors; + std::vector params_to_prepare; + if (!prepare_execute_graph_weights(gf, graph_param_tensors, params_to_prepare, !free_compute_params)) { return std::nullopt; } + struct GraphWeightDoneGuard { + GraphWeightDoneGuard(GGMLRunner* runner, const std::vector* tensors) + : runner(runner), + tensors(tensors) {} + + GGMLRunner* runner = nullptr; + const std::vector* tensors = nullptr; + bool enabled = true; + + ~GraphWeightDoneGuard() { + if (enabled && runner != nullptr && tensors != nullptr) { + runner->free_compute_backend_param_tensors(*tensors); + } + } + + void dismiss() { enabled = false; } + + GraphWeightDoneGuard(const GraphWeightDoneGuard&) = delete; + GraphWeightDoneGuard& operator=(const GraphWeightDoneGuard&) = delete; + }; + GraphWeightDoneGuard graph_weight_done_guard(this, ¶ms_to_prepare); + + if (!alloc_compute_buffer(gf)) { + LOG_ERROR("%s alloc compute buffer failed", get_desc().c_str()); + return std::nullopt; + } + struct ComputeBufferGuard { + ComputeBufferGuard(GGMLRunner* runner, bool enabled) + : runner(runner), + enabled(enabled) {} + + GGMLRunner* runner = nullptr; + bool enabled = false; + + ~ComputeBufferGuard() { + if (enabled && runner != nullptr) { + runner->free_compute_buffer(); + } + } + + ComputeBufferGuard(const ComputeBufferGuard&) = delete; + ComputeBufferGuard& operator=(const ComputeBufferGuard&) = delete; + }; + ComputeBufferGuard compute_buffer_guard(this, free_compute_buffer); if (!ggml_gallocr_alloc_graph(compute_allocr, gf)) { LOG_ERROR("%s alloc compute graph failed", get_desc().c_str()); - if (free_compute_buffer_immediately) { - free_compute_buffer(); - } else if (use_partial_param_offload) { - restore_partial_params(); - } return std::nullopt; } - int64_t t_alloc_end = ggml_time_ms(); - int64_t t_copy_begin = ggml_time_ms(); copy_data_to_backend_tensor(gf, !preserve_backend_tensor_data_map); - int64_t t_copy_end = ggml_time_ms(); if (sd_backend_is_cpu(runtime_backend)) { sd_backend_cpu_set_n_threads(runtime_backend, n_threads); } - int64_t t_compute_begin = ggml_time_ms(); - ggml_status status = ggml_backend_graph_compute(runtime_backend, gf); - int64_t t_compute_end = ggml_time_ms(); + ggml_status status = sd_backend_graph_compute_with_eval_callback(runtime_backend, + gf, + sd_get_backend_eval_callback(), + sd_get_backend_eval_callback_data()); if (status != GGML_STATUS_SUCCESS) { LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status)); - if (free_compute_buffer_immediately) { - free_compute_buffer(); - } else if (use_partial_param_offload) { - restore_partial_params(); - } return std::nullopt; } - std::unordered_set debug_graph_tensor_set; - const int n_debug_leafs = sd::ggml_graph_cut::leaf_count(gf); - const int n_debug_nodes = ggml_graph_n_nodes(gf); - debug_graph_tensor_set.reserve(static_cast(n_debug_leafs + n_debug_nodes)); - for (int i = 0; i < n_debug_leafs; ++i) { - debug_graph_tensor_set.insert(sd::ggml_graph_cut::leaf_tensor(gf, i)); - } - for (int i = 0; i < n_debug_nodes; ++i) { - debug_graph_tensor_set.insert(ggml_graph_node(gf, i)); + if (!debug_tensors.empty()) { + std::unordered_set debug_graph_tensor_set; + const int n_debug_leafs = sd::ggml_graph_cut::leaf_count(gf); + const int n_debug_nodes = ggml_graph_n_nodes(gf); + debug_graph_tensor_set.reserve(static_cast(n_debug_leafs + n_debug_nodes)); + for (int i = 0; i < n_debug_leafs; ++i) { + debug_graph_tensor_set.insert(sd::ggml_graph_cut::leaf_tensor(gf, i)); + } + for (int i = 0; i < n_debug_nodes; ++i) { + debug_graph_tensor_set.insert(ggml_graph_node(gf, i)); + } + + for (const auto& entry : debug_tensors) { + auto tensor = entry.first; + if (tensor == nullptr) { + continue; + } + if (debug_graph_tensor_set.find(tensor) == debug_graph_tensor_set.end()) { + continue; + } + ggml_backend_buffer_t tensor_buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; + if (tensor_buf == nullptr) { + LOG_WARN("%s skip debug tensor '%s': tensor buffer not set", + get_desc().c_str(), + entry.second.c_str()); + continue; + } + if (tensor->type != GGML_TYPE_F32) { + LOG_WARN("%s skip debug tensor '%s': only GGML_TYPE_F32 is supported, got %s", + get_desc().c_str(), + entry.second.c_str(), + ggml_type_name(tensor->type)); + continue; + } + auto debug_tensor = sd::make_sd_tensor_from_ggml(tensor); + print_sd_tensor(debug_tensor, false, entry.second.c_str()); + } } - for (const auto& entry : debug_tensors) { - auto tensor = entry.first; - if (tensor == nullptr) { - continue; - } - if (debug_graph_tensor_set.find(tensor) == debug_graph_tensor_set.end()) { - continue; - } - ggml_backend_buffer_t tensor_buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; - if (tensor_buf == nullptr) { - LOG_WARN("%s skip debug tensor '%s': tensor buffer not set", - get_desc().c_str(), - entry.second.c_str()); - continue; - } - if (tensor->type != GGML_TYPE_F32) { - LOG_WARN("%s skip debug tensor '%s': only GGML_TYPE_F32 is supported, got %s", - get_desc().c_str(), - entry.second.c_str(), - ggml_type_name(tensor->type)); - continue; - } - auto debug_tensor = sd::make_sd_tensor_from_ggml(tensor); - print_sd_tensor(debug_tensor, false, entry.second.c_str()); - } - - int64_t t_cache_begin = ggml_time_ms(); if (!copy_cache_tensors_to_cache_buffer(cache_keep_names)) { - if (free_compute_buffer_immediately) { - free_compute_buffer(); - } else if (use_partial_param_offload) { - restore_partial_params(); - } return std::nullopt; } - int64_t t_cache_end = ggml_time_ms(); - auto result = ggml_get_tensor(compute_ctx, final_result_name.c_str()); + auto result = ggml_get_tensor(compute_ctx, final_result_name.c_str()); std::optional> output; if (!no_return) { output = read_graph_tensor(result, "output"); if (!output.has_value()) { - if (free_compute_buffer_immediately) { - free_compute_buffer(); - } else if (use_partial_param_offload) { - restore_partial_params(); - } return std::nullopt; } } else { output = sd::Tensor(); } - if (free_compute_buffer_immediately) { - free_compute_buffer(); - } else if (use_partial_param_offload) { - restore_partial_params(); - } - if (use_partial_param_offload) { - LOG_DEBUG("%s execute_graph timing: offload=%lld ms alloc=%lld ms copy_in=%lld ms compute=%lld ms cache=%lld ms total=%lld ms", - get_desc().c_str(), - t_offload_end - t_offload_begin, - t_alloc_end - t_alloc_begin, - t_copy_end - t_copy_begin, - t_compute_end - t_compute_begin, - t_cache_end - t_cache_begin, - ggml_time_ms() - t_execute_begin); - } - return output; - } - - template - std::optional> compute_with_graph_cuts(ggml_cgraph* gf, - const GraphCutPlan& plan, - int n_threads, - bool free_compute_buffer_immediately, - bool no_return = false) { - GGML_ASSERT(gf != nullptr); - - free_compute_buffer(); - free_cache_ctx_and_buffer(); - - std::unordered_map persistent_externals; - snapshot_persistent_externals(plan, gf, persistent_externals); - - std::optional> output = sd::Tensor(); - for (size_t seg_idx = 0; seg_idx < plan.segments.size(); ++seg_idx) { - int64_t t_segment_begin = ggml_time_ms(); - const auto& segment = plan.segments[seg_idx]; - auto future_cut_names = sd::ggml_graph_cut::collect_future_input_names(gf, plan, seg_idx); - LOG_DEBUG("%s graph cut executing segment %zu/%zu: %s", - get_desc().c_str(), - seg_idx + 1, - plan.segments.size(), - segment.group_name.c_str()); - - reset_segment_runtime_tensors(segment, gf, &persistent_externals); - if (!bind_segment_cached_inputs(gf, segment)) { - free_cache_ctx_and_buffer(); - free_compute_buffer(); - free_compute_ctx(); - return std::nullopt; - } - - const bool is_last_segment = seg_idx + 1 == plan.segments.size(); - if (!is_last_segment) { - for (size_t output_idx = 0; output_idx < segment.output_node_indices.size(); ++output_idx) { - ggml_tensor* output_tensor = sd::ggml_graph_cut::output_tensor(gf, segment, output_idx); - if (output_tensor != nullptr && - sd::ggml_graph_cut::is_graph_cut_tensor(output_tensor) && - future_cut_names.find(output_tensor->name) != future_cut_names.end()) { - cache(output_tensor->name, output_tensor); - } + if (!free_compute_params) { + for (ggml_tensor* param : params_to_prepare) { + if (param == nullptr) { + continue; } + kept_compute_param_tensor_set.insert(param); } - - ggml_context* segment_graph_ctx = nullptr; - ggml_cgraph* segment_graph = sd::ggml_graph_cut::build_segment_graph(gf, segment, &segment_graph_ctx); - auto segment_output = execute_graph(segment_graph, - n_threads, - true, - sd::ggml_graph_cut::runtime_param_tensors(gf, segment, get_desc().c_str()), - true, - !is_last_segment || no_return, - &future_cut_names); - ggml_free(segment_graph_ctx); - if (!segment_output.has_value()) { - free_cache_ctx_and_buffer(); - free_compute_buffer(); - free_compute_ctx(); - return std::nullopt; - } - output = std::move(segment_output); + graph_weight_done_guard.dismiss(); } - - backend_tensor_data_map.clear(); - free_cache_ctx_and_buffer(); - free_compute_ctx(); return output; } -public: - void release_streaming_residency() { - restore_resident_params(); - } - template - std::optional> compute_streaming_segments(ggml_cgraph* gf, + std::optional> compute_graph_cut_segments(ggml_cgraph* gf, const GraphCutPlan& plan, - size_t residency_budget_bytes, int n_threads, - bool free_compute_buffer_immediately, + bool log_residency, bool no_return = false) { GGML_ASSERT(gf != nullptr); - // Runtime LoRA composes `weight + diff` in the compute graph via - // ggml_add; the resident weight tensor's data is never mutated, so - // chunk-K residency stays valid across sampling steps. - // Reserve room for the worst merged segment so chunk-K can't grow - // large enough to starve later partial-param allocations. - size_t worst_merged_segment_footprint = 0; - for (const auto& seg : plan.segments) { - const size_t fp = seg.input_param_bytes + - seg.compute_buffer_size + - seg.output_bytes + - seg.input_previous_cut_bytes + - seg.input_external_bytes; - if (fp > worst_merged_segment_footprint) { - worst_merged_segment_footprint = fp; - } - } - const size_t residency_budget_for_annotate = - residency_budget_bytes > worst_merged_segment_footprint - ? residency_budget_bytes - worst_merged_segment_footprint - : 0; - - sd::ggml_graph_cut::Plan& base_plan = graph_cut_plan_cache_.graph_cut_plan; - if (base_plan.available) { - sd::ggml_graph_cut::annotate_residency(base_plan, residency_budget_for_annotate); - - std::vector resident_params; - uint64_t token = 0; - for (const auto& segment : base_plan.segments) { - if (segment.residency != sd::ggml_graph_cut::SegmentResidency::RESIDENT) { - continue; - } - auto seg_params = sd::ggml_graph_cut::param_tensors(gf, segment); - for (ggml_tensor* t : seg_params) { - if (t == nullptr) - continue; - resident_params.push_back(t); - token ^= reinterpret_cast(t) * 0x9E3779B97F4A7C15ull; - } - } - if (token != resident_state_token) { - restore_resident_params(); - if (!resident_params.empty()) { - if (offload_resident_params(resident_params)) { - resident_state_token = token; - } else { - LOG_ERROR("%s chunk-K: resident offload failed; continuing with per-segment streaming", - get_desc().c_str()); - restore_resident_params(); - } - } - } - } - free_compute_buffer(); free_cache_ctx_and_buffer(); - layer_registry_.move_layer_to_gpu("_global"); - std::unordered_map persistent_externals; snapshot_persistent_externals(plan, gf, persistent_externals); std::optional> output = sd::Tensor(); for (size_t seg_idx = 0; seg_idx < plan.segments.size(); ++seg_idx) { - int64_t t_segment_begin = ggml_time_ms(); - const auto& segment = plan.segments[seg_idx]; - const bool is_last = seg_idx + 1 == plan.segments.size(); - auto future_cut_names = sd::ggml_graph_cut::collect_future_input_names(gf, plan, seg_idx); - - LOG_DEBUG("%s streaming-cut executing segment %zu/%zu: %s (residency=%s)", - get_desc().c_str(), - seg_idx + 1, - plan.segments.size(), - segment.group_name.c_str(), - segment.residency == sd::ggml_graph_cut::SegmentResidency::RESIDENT ? "RESIDENT" : "STREAMED"); - - if (!layer_registry_.move_layer_to_gpu(segment.group_name)) { - LOG_DEBUG("%s streaming: no registry entry for group '%s' (using upstream offload path)", + const auto& segment = plan.segments[seg_idx]; + const bool is_last = seg_idx + 1 == plan.segments.size(); + auto future_cut_names = sd::ggml_graph_cut::collect_future_input_names(gf, plan, seg_idx); + if (log_residency) { + LOG_DEBUG("%s graph cut executing segment %zu/%zu: %s (residency=%s)", get_desc().c_str(), + seg_idx + 1, + plan.segments.size(), + segment.group_name.c_str(), + segment.residency == sd::ggml_graph_cut::SegmentResidency::RESIDENT ? "RESIDENT" : "STREAMED"); + } else { + LOG_DEBUG("%s graph cut executing segment %zu/%zu: %s", + get_desc().c_str(), + seg_idx + 1, + plan.segments.size(), segment.group_name.c_str()); } @@ -2911,23 +2591,24 @@ public: if (!is_last) { for (size_t output_idx = 0; output_idx < segment.output_node_indices.size(); ++output_idx) { - ggml_tensor* out_tensor = sd::ggml_graph_cut::output_tensor(gf, segment, output_idx); - if (out_tensor != nullptr && - sd::ggml_graph_cut::is_graph_cut_tensor(out_tensor) && - future_cut_names.find(out_tensor->name) != future_cut_names.end()) { - cache(out_tensor->name, out_tensor); + ggml_tensor* output_tensor = sd::ggml_graph_cut::output_tensor(gf, segment, output_idx); + if (output_tensor != nullptr && + sd::ggml_graph_cut::is_graph_cut_tensor(output_tensor) && + future_cut_names.find(output_tensor->name) != future_cut_names.end()) { + cache(output_tensor->name, output_tensor); } } } ggml_context* segment_graph_ctx = nullptr; ggml_cgraph* segment_graph = sd::ggml_graph_cut::build_segment_graph(gf, segment, &segment_graph_ctx); + const bool keep_segment_params = segment.residency == sd::ggml_graph_cut::SegmentResidency::RESIDENT; auto segment_output = execute_graph(segment_graph, n_threads, - /*free_compute_buffer_immediately=*/true, - sd::ggml_graph_cut::runtime_param_tensors(gf, segment, get_desc().c_str()), - /*preserve_backend_tensor_data_map=*/true, - /*no_return=*/!is_last || no_return, + true, + !keep_segment_params, + true, + !is_last || no_return, &future_cut_names); ggml_free(segment_graph_ctx); if (!segment_output.has_value()) { @@ -2937,11 +2618,6 @@ public: return std::nullopt; } output = std::move(segment_output); - - if (segment.residency == sd::ggml_graph_cut::SegmentResidency::STREAMED) { - layer_registry_.move_layer_to_cpu(segment.group_name); - } - (void)t_segment_begin; } backend_tensor_data_map.clear(); @@ -2950,21 +2626,29 @@ public: return output; } +public: + void runner_done() { + free_compute_buffer(); + std::vector tensors_to_release = std::move(this->runner_param_tensors); + this->runner_param_tensors.clear(); + runner_param_tensor_set.clear(); + kept_compute_param_tensor_set.clear(); + free_compute_backend_param_tensors(tensors_to_release); + free_params_backend_param_tensors(tensors_to_release); + } + public: virtual std::string get_desc() = 0; - GGMLRunner(ggml_backend_t backend, ggml_backend_t params_backend) - : params_backend(params_backend), - runtime_backend(backend) { + GGMLRunner(ggml_backend_t backend, + std::shared_ptr manager = nullptr) + : runtime_backend(backend), + weight_manager(manager) { GGML_ASSERT(runtime_backend != nullptr); - GGML_ASSERT(params_backend != nullptr); alloc_params_ctx(); - layer_registry_.set_backends(runtime_backend, params_backend); } virtual ~GGMLRunner() { - restore_resident_params(); - free_params_buffer(); free_compute_buffer(); free_params_ctx(); free_compute_ctx(); @@ -2987,6 +2671,9 @@ public: runner_ctx.cache_tensor = [this](const std::string& name, ggml_tensor* tensor) { this->cache(name, tensor); }; + runner_ctx.set_backend_tensor_data = [this](ggml_tensor* tensor, const void* data) { + this->set_backend_tensor_data(tensor, data); + }; return runner_ctx; } @@ -2995,62 +2682,7 @@ public: alloc_compute_ctx(); } - bool alloc_params_buffer() { - size_t num_tensors = ggml_tensor_num(params_ctx); - if (num_tensors > 0) { - // ggml_backend_alloc_ctx_tensors fails when all tensors are already allocated - // (typical for memory-mapped weights). See ggml-alloc.c n_buffers==0 branch. - bool all_have_data = true; - for (ggml_tensor* t = ggml_get_first_tensor(params_ctx); t != nullptr; t = ggml_get_next_tensor(params_ctx, t)) { - if (t->data == nullptr) { - all_have_data = false; - break; - } - } - if (all_have_data) { - LOG_DEBUG("%s all params already mmap-allocated (no separate buffer needed)", get_desc().c_str()); - params_buffer = nullptr; - rebuild_params_tensor_set(); - return true; - } - } else { - LOG_DEBUG("%s skipping params allocation (no tensors)", get_desc().c_str()); - return true; - } - params_buffer = ggml_backend_alloc_ctx_tensors(params_ctx, params_backend); - if (params_buffer == nullptr) { - LOG_ERROR("%s alloc params backend buffer failed, num_tensors = %i", - get_desc().c_str(), - num_tensors); - return false; - } - rebuild_params_tensor_set(); - ggml_backend_buffer_set_usage(params_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); - size_t params_buffer_size = ggml_backend_buffer_get_size(params_buffer); - LOG_DEBUG("%s params backend buffer size = % 6.2f MB(%s) (%i tensors)", - get_desc().c_str(), - params_buffer_size / (1024.f * 1024.f), - sd_backend_is_cpu(params_backend) ? "RAM" : "VRAM", - num_tensors); - return true; - } - - void free_params_buffer() { - // Restore swapped resident params before freeing their backing buffer. - restore_resident_params(); - if (params_buffer != nullptr) { - ggml_backend_buffer_free(params_buffer); - params_buffer = nullptr; - } - } - - size_t get_params_buffer_size() { - if (params_buffer != nullptr) { - return ggml_backend_buffer_get_size(params_buffer); - } - return 0; - } - +public: void free_cache_ctx_and_buffer() { free_cache_buffer(); free_cache_ctx(); @@ -3061,8 +2693,6 @@ public: ggml_gallocr_free(compute_allocr); compute_allocr = nullptr; } - restore_partial_params(); - restore_all_params(); } // do copy after alloc graph @@ -3127,47 +2757,57 @@ public: template std::optional> compute(get_graph_cb_t get_graph, int n_threads, - bool free_compute_buffer_immediately, - bool no_return = false) { + bool auto_free = true, + bool free_compute_buffer = true, + bool free_compute_params = true, + bool no_return = false) { + struct RunnerDoneGuard { + RunnerDoneGuard(GGMLRunner* runner, bool enabled) + : runner(runner), + enabled(enabled) {} + + ~RunnerDoneGuard() { + if (enabled && runner != nullptr) { + runner->runner_done(); + } + } + + RunnerDoneGuard(const RunnerDoneGuard&) = delete; + RunnerDoneGuard& operator=(const RunnerDoneGuard&) = delete; + + GGMLRunner* runner = nullptr; + bool enabled = false; + }; + RunnerDoneGuard runner_done_guard(this, auto_free); + ggml_cgraph* gf = nullptr; if (!prepare_compute_graph(get_graph, &gf)) { return std::nullopt; } GGML_ASSERT(gf != nullptr); + rebuild_params_tensor_set(); if (can_attempt_graph_cut_segmented_compute()) { GraphCutPlan plan; - size_t effective_graph_vram_bytes = 0; - if (!resolve_graph_cut_plan(gf, &plan, &effective_graph_vram_bytes)) { + if (!resolve_graph_cut_plan(gf, &plan)) { free_compute_ctx(); return std::nullopt; } if (should_use_graph_cut_segmented_compute(plan)) { - if (stream_layers_enabled) { - return compute_streaming_segments(gf, - plan, - effective_graph_vram_bytes, - n_threads, - free_compute_buffer_immediately, - no_return); - } - return compute_with_graph_cuts(gf, - plan, - n_threads, - free_compute_buffer_immediately, - no_return); + return compute_graph_cut_segments(gf, + plan, + n_threads, + stream_layers_enabled, + no_return); } } - if (!alloc_compute_buffer(gf)) { - LOG_ERROR("%s alloc compute buffer failed", get_desc().c_str()); - return std::nullopt; - } return execute_graph(gf, n_threads, - free_compute_buffer_immediately, - {}, + free_compute_buffer, + free_compute_params, false, - no_return); + no_return, + nullptr); } void set_flash_attention_enabled(bool enabled) { @@ -3194,16 +2834,6 @@ public: void set_stream_layers_enabled(bool enabled) { stream_layers_enabled = enabled; } - - sd::layer_registry::LayerRegistry& get_layer_registry() { return layer_registry_; } - - ggml_backend_t get_runtime_backend() { - return runtime_backend; - } - - ggml_backend_t get_params_backend() { - return params_backend; - } }; class GGMLBlock { @@ -3282,6 +2912,7 @@ public: for (auto& pair : params) { ggml_tensor* param = pair.second; tensors[prefix + pair.first] = pair.second; + ggml_set_name(param, (prefix + pair.first).c_str()); } } @@ -3318,11 +2949,14 @@ protected: bool bias; bool force_f32; bool force_prec_f32; + bool allow_weight_scale; + bool has_weight_scale = false; float scale; std::string prefix; void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { this->prefix = prefix; + has_weight_scale = false; enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32); if (in_features % ggml_blck_size(wtype) != 0 || force_f32) { wtype = GGML_TYPE_F32; @@ -3332,20 +2966,26 @@ protected: enum ggml_type wtype = GGML_TYPE_F32; params["bias"] = ggml_new_tensor_1d(ctx, wtype, out_features); } + if (allow_weight_scale && tensor_storage_map.find(prefix + "weight_scale") != tensor_storage_map.end()) { + params["weight_scale"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_features); + has_weight_scale = true; + } } public: Linear(int64_t in_features, int64_t out_features, - bool bias = true, - bool force_f32 = false, - bool force_prec_f32 = false, - float scale = 1.f / 256.f) + bool bias = true, + bool force_f32 = false, + bool force_prec_f32 = false, + float scale = 1.f, + bool allow_weight_scale = false) : in_features(in_features), out_features(out_features), bias(bias), force_f32(force_f32), force_prec_f32(force_prec_f32), + allow_weight_scale(allow_weight_scale), scale(scale) {} void set_scale(float scale_) { @@ -3362,14 +3002,24 @@ public: if (bias) { b = params["bias"]; } + ggml_tensor* linear_bias = has_weight_scale ? nullptr : b; + ggml_tensor* out = nullptr; if (ctx->weight_adapter) { WeightAdapter::ForwardParams forward_params; forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR; forward_params.linear.force_prec_f32 = force_prec_f32; forward_params.linear.scale = scale; - return ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x, w, b, prefix, forward_params); + out = ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x, w, linear_bias, prefix, forward_params); + } else { + out = ggml_ext_linear(ctx->ggml_ctx, x, w, linear_bias, force_prec_f32, scale); } - return ggml_ext_linear(ctx->ggml_ctx, x, w, b, force_prec_f32, scale); + if (has_weight_scale) { + out = ggml_mul(ctx->ggml_ctx, out, params["weight_scale"]); + if (b != nullptr) { + out = ggml_add_inplace(ctx->ggml_ctx, out, b); + } + } + return out; } }; @@ -4097,4 +3747,4 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_lokr_forward( } } -#endif // __GGML_EXTEND__HPP__ +#endif // __SD_CORE_GGML_EXTEND_HPP__ diff --git a/src/ggml_extend_backend.cpp b/src/core/ggml_extend_backend.cpp similarity index 73% rename from src/ggml_extend_backend.cpp rename to src/core/ggml_extend_backend.cpp index 5e95eae5..f29bdb69 100644 --- a/src/ggml_extend_backend.cpp +++ b/src/core/ggml_extend_backend.cpp @@ -1,4 +1,4 @@ -#include "ggml_extend_backend.h" +#include "core/ggml_extend_backend.h" #include #include @@ -8,8 +8,9 @@ #include #include +#include "core/util.h" +#include "ggml/src/ggml-impl.h" #include "stable-diffusion.h" -#include "util.h" static std::string trim_copy(const std::string& value) { size_t begin = 0; @@ -45,6 +46,10 @@ static bool is_default_backend_token(const std::string& name) { return lower.empty() || lower == "default" || lower == "auto"; } +static bool is_disk_backend_token(const std::string& name) { + return lower_copy(trim_copy(name)) == "disk"; +} + static bool parse_backend_module(const std::string& raw_name, SDBackendModule* module) { std::string name = lower_copy(trim_copy(raw_name)); name.erase(std::remove(name.begin(), name.end(), '-'), name.end()); @@ -106,7 +111,67 @@ static std::string resolve_first_device_by_type(enum ggml_backend_dev_type type) if (dev == nullptr) { return ""; } - return ggml_backend_dev_name(dev); + const char* dev_name = ggml_backend_dev_name(dev); + if (dev_name != nullptr && dev_name[0] != '\0') { + return dev_name; + } + ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev); + const char* reg_name = reg != nullptr ? ggml_backend_reg_name(reg) : nullptr; + return reg_name != nullptr ? reg_name : ""; +} + +static ggml_backend_dev_t resolve_first_device_by_registry_name(const std::string& name) { + std::string lower = lower_copy(trim_copy(name)); + if (lower == "metal") { + lower = "mtl"; + } + if (lower.empty()) { + return nullptr; + } + + const size_t device_count = ggml_backend_dev_count(); + for (size_t i = 0; i < device_count; ++i) { + ggml_backend_dev_t dev = ggml_backend_dev_get(i); + ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev); + if (reg == nullptr) { + continue; + } + const char* reg_name = ggml_backend_reg_name(reg); + if (reg_name != nullptr && lower_copy(reg_name) == lower) { + return dev; + } + } + return nullptr; +} + +static ggml_backend_dev_t resolve_device_by_name(const std::string& name) { + const std::string lower = lower_copy(trim_copy(name)); + if (lower.empty()) { + return nullptr; + } + + const size_t device_count = ggml_backend_dev_count(); + for (size_t i = 0; i < device_count; ++i) { + ggml_backend_dev_t dev = ggml_backend_dev_get(i); + const char* dev_name = ggml_backend_dev_name(dev); + if (dev_name != nullptr && lower_copy(dev_name) == lower) { + return dev; + } + } + return nullptr; +} + +static std::string backend_device_name(ggml_backend_dev_t dev) { + if (dev == nullptr) { + return ""; + } + const char* name = ggml_backend_dev_name(dev); + if (name != nullptr && name[0] != '\0') { + return name; + } + ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev); + const char* reg_name = reg != nullptr ? ggml_backend_reg_name(reg) : nullptr; + return reg_name != nullptr ? reg_name : ""; } static ggml_backend_buffer_t ggml_backend_tensor_buffer(const struct ggml_tensor* tensor) { @@ -200,6 +265,36 @@ void ggml_ext_im_set_f32_1d(const struct ggml_tensor* tensor, int i, float value } } +bool add_rpc_devices(const std::string& servers) { + const std::string in = trim_copy(servers); + if (in.empty()) { + return true; + } + auto rpc_servers = split_copy(in, ','); + if (rpc_servers.empty()) { + LOG_ERROR("invalid RPC servers specification: '%s'", servers.c_str()); + return false; + } + ggml_backend_reg_t rpc_reg = ggml_backend_reg_by_name("RPC"); + if (!rpc_reg) { + LOG_ERROR("RPC backend not found, cannot add RPC servers"); + return false; + } + typedef ggml_backend_reg_t (*ggml_backend_rpc_add_server_t)(const char* endpoint); + ggml_backend_rpc_add_server_t ggml_backend_rpc_add_server_fn = (ggml_backend_rpc_add_server_t)ggml_backend_reg_get_proc_address(rpc_reg, "ggml_backend_rpc_add_server"); + if (!ggml_backend_rpc_add_server_fn) { + LOG_ERROR("RPC backend does not have ggml_backend_rpc_add_server function, cannot add RPC servers"); + return false; + } + for (const auto& server : rpc_servers) { + LOG_INFO("Adding RPC server: %s", server.c_str()); + auto reg = ggml_backend_rpc_add_server_fn(server.c_str()); + // no return value to check for success but should print errors from the RPC backend if it fails to add the server + ggml_backend_register(reg); + } + return true; +} + static void ggml_backend_load_all_once() { // If the registry already has devices and the CPU backend is present, // assume either static registration or explicit host-side preloading has @@ -246,7 +341,7 @@ static std::string get_default_backend_name() { return resolve_first_device_by_type(GGML_BACKEND_DEVICE_TYPE_CPU); } -static std::string sd_resolve_backend_name(const std::string& name) { +std::string sd_backend_resolve_name(const std::string& name) { ggml_backend_load_all_once(); std::string requested = trim_copy(name); std::string lower = lower_copy(requested); @@ -262,6 +357,10 @@ static std::string sd_resolve_backend_name(const std::string& name) { return resolve_first_device_by_type(GGML_BACKEND_DEVICE_TYPE_IGPU); } + if (ggml_backend_dev_t dev = resolve_first_device_by_registry_name(requested)) { + return backend_device_name(dev); + } + const size_t device_count = ggml_backend_dev_count(); for (size_t i = 0; i < device_count; ++i) { ggml_backend_dev_t dev = ggml_backend_dev_get(i); @@ -284,7 +383,7 @@ static std::string sd_resolve_backend_name(const std::string& name) { } static bool backend_name_exists(const std::string& name) { - return !sd_resolve_backend_name(name).empty(); + return !sd_backend_resolve_name(name).empty(); } static ggml_backend_t init_named_backend(const std::string& name) { @@ -294,7 +393,20 @@ static ggml_backend_t init_named_backend(const std::string& name) { return ggml_backend_init_best(); } - std::string resolved = sd_resolve_backend_name(name); + if (ggml_backend_dev_t dev = resolve_device_by_name(name)) { + return ggml_backend_dev_init(dev, nullptr); + } + if (ggml_backend_dev_t dev = resolve_first_device_by_registry_name(name)) { + return ggml_backend_dev_init(dev, nullptr); + } + + std::string resolved = sd_backend_resolve_name(name); + if (ggml_backend_dev_t dev = resolve_device_by_name(resolved)) { + return ggml_backend_dev_init(dev, nullptr); + } + if (ggml_backend_dev_t dev = resolve_first_device_by_registry_name(resolved)) { + return ggml_backend_dev_init(dev, nullptr); + } if (resolved.empty()) { return nullptr; } @@ -330,6 +442,68 @@ bool sd_backend_cpu_set_n_threads(ggml_backend_t backend, int n_threads) { return false; } +static ggml_cgraph sd_ggml_graph_view(ggml_cgraph* cgraph0, int i0, int i1) { + ggml_cgraph cgraph = { + /*.size =*/0, + /*.n_nodes =*/i1 - i0, + /*.n_leafs =*/0, + /*.nodes =*/cgraph0->nodes + i0, + /*.grads =*/nullptr, + /*.grad_accs =*/nullptr, + /*.leafs =*/nullptr, + /*.use_counts =*/cgraph0->use_counts, + /*.visited_hash_set =*/cgraph0->visited_hash_set, + /*.order =*/cgraph0->order, + /*.uid =*/0, + }; + return cgraph; +} + +ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend, + ggml_cgraph* gf, + sd_graph_eval_callback_t callback_eval, + void* callback_eval_user_data) { + if (callback_eval == nullptr) { + return ggml_backend_graph_compute(backend, gf); + } + + ggml_status status = GGML_STATUS_SUCCESS; + const int n_nodes = ggml_graph_n_nodes(gf); + bool stopped = false; + + for (int j0 = 0; j0 < n_nodes; ++j0) { + ggml_tensor* t = ggml_graph_node(gf, j0); + bool need = callback_eval(t, true, callback_eval_user_data); + int j1 = j0; + + while (!need && j1 < n_nodes - 1) { + t = ggml_graph_node(gf, ++j1); + need = callback_eval(t, true, callback_eval_user_data); + } + + ggml_cgraph gv = sd_ggml_graph_view(gf, j0, j1 + 1); + status = ggml_backend_graph_compute_async(backend, &gv); + if (status != GGML_STATUS_SUCCESS) { + break; + } + + ggml_backend_synchronize(backend); + + if (need && !callback_eval(t, false, callback_eval_user_data)) { + stopped = true; + break; + } + + j0 = j1; + } + + ggml_backend_synchronize(backend); + if (stopped && status == GGML_STATUS_SUCCESS) { + status = GGML_STATUS_ABORTED; + } + return status; +} + const char* sd_get_system_info() { static std::string cache_info = []() -> std::string { ggml_backend_load_all_once(); @@ -504,6 +678,9 @@ ggml_backend_t SDBackendManager::params_backend(SDBackendModule module) { if (name.empty()) { return runtime_backend(module); } + if (is_disk_backend_token(name)) { + return runtime_backend(module); + } return init_cached_backend(name); } @@ -515,6 +692,10 @@ bool SDBackendManager::params_backend_is_cpu(SDBackendModule module) { return sd_backend_is_cpu(params_backend(module)); } +bool SDBackendManager::params_backend_is_disk(SDBackendModule module) const { + return is_disk_backend_token(params_assignment_.get(module)); +} + bool SDBackendManager::runtime_backend_supports_host_buffer(SDBackendModule module) { ggml_backend_t backend = runtime_backend(module); if (backend == nullptr) { @@ -534,10 +715,6 @@ bool SDBackendManager::runtime_backend_supports_host_buffer(SDBackendModule modu bool SDBackendManager::init(const char* backend_spec, const char* params_backend_spec, - bool offload_params_to_cpu, - bool keep_clip_on_cpu, - bool keep_vae_on_cpu, - bool keep_control_net_on_cpu, std::string* error) { reset(); @@ -548,31 +725,21 @@ bool SDBackendManager::init(const char* backend_spec, return false; } - if (runtime_assignment_.empty()) { - if (keep_clip_on_cpu) { - runtime_assignment_.set_module(SDBackendModule::TE, "cpu"); - } - if (keep_vae_on_cpu) { - runtime_assignment_.set_module(SDBackendModule::VAE, "cpu"); - } - if (keep_control_net_on_cpu) { - runtime_assignment_.set_module(SDBackendModule::CONTROL_NET, "cpu"); - } - } - - if (params_assignment_.empty() && offload_params_to_cpu) { - params_assignment_.set_default("cpu"); - } - return validate(error); } bool SDBackendManager::validate(std::string* error) const { - auto validate_name = [&](const std::string& name) -> bool { + auto validate_runtime_name = [&](const std::string& name) -> bool { if (is_default_backend_token(name)) { return true; } - if (!sd_resolve_backend_name(name).empty()) { + if (is_disk_backend_token(name)) { + if (error != nullptr) { + *error = "backend 'disk' is only supported by params_backend"; + } + return false; + } + if (!sd_backend_resolve_name(name).empty() || resolve_first_device_by_registry_name(name) != nullptr) { return true; } if (error != nullptr) { @@ -580,18 +747,24 @@ bool SDBackendManager::validate(std::string* error) const { } return false; }; + auto validate_params_name = [&](const std::string& name) -> bool { + if (is_disk_backend_token(name)) { + return true; + } + return validate_runtime_name(name); + }; - if (!validate_name(runtime_assignment_.default_name) || - !validate_name(params_assignment_.default_name)) { + if (!validate_runtime_name(runtime_assignment_.default_name) || + !validate_params_name(params_assignment_.default_name)) { return false; } for (const auto& kv : runtime_assignment_.module_names) { - if (!validate_name(kv.second)) { + if (!validate_runtime_name(kv.second)) { return false; } } for (const auto& kv : params_assignment_.module_names) { - if (!validate_name(kv.second)) { + if (!validate_params_name(kv.second)) { return false; } } @@ -599,7 +772,7 @@ bool SDBackendManager::validate(std::string* error) const { } ggml_backend_t SDBackendManager::init_cached_backend(const std::string& name) { - std::string resolved = sd_resolve_backend_name(name); + std::string resolved = sd_backend_resolve_name(name); std::string key = lower_copy(resolved); ggml_backend_t backend = nullptr; diff --git a/src/ggml_extend_backend.h b/src/core/ggml_extend_backend.h similarity index 75% rename from src/ggml_extend_backend.h rename to src/core/ggml_extend_backend.h index 972fbee7..19b71d43 100644 --- a/src/ggml_extend_backend.h +++ b/src/core/ggml_extend_backend.h @@ -1,5 +1,5 @@ -#ifndef __SD_GGML_EXTEND_BACKEND_H__ -#define __SD_GGML_EXTEND_BACKEND_H__ +#ifndef __SD_CORE_GGML_EXTEND_BACKEND_H__ +#define __SD_CORE_GGML_EXTEND_BACKEND_H__ #include #include @@ -9,6 +9,7 @@ #include "ggml-backend.h" #include "ggml.h" +#include "stable-diffusion.h" enum class SDBackendModule { DIFFUSION, @@ -51,10 +52,6 @@ public: bool init(const char* backend_spec, const char* params_backend_spec, - bool offload_params_to_cpu, - bool keep_clip_on_cpu, - bool keep_vae_on_cpu, - bool keep_control_net_on_cpu, std::string* error); void reset(); @@ -63,6 +60,7 @@ public: bool runtime_backend_is_cpu(SDBackendModule module); bool params_backend_is_cpu(SDBackendModule module); + bool params_backend_is_disk(SDBackendModule module) const; bool runtime_backend_supports_host_buffer(SDBackendModule module); private: @@ -74,6 +72,12 @@ bool sd_backend_is(ggml_backend_t backend, const std::string& name); bool sd_backend_is_cpu(ggml_backend_t backend); ggml_backend_t sd_backend_cpu_init(); bool sd_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_threads); +ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend, + ggml_cgraph* gf, + sd_graph_eval_callback_t callback_eval, + void* callback_eval_user_data); +std::string sd_backend_resolve_name(const std::string& name); const char* sd_backend_module_name(SDBackendModule module); void ggml_ext_im_set_f32_1d(const struct ggml_tensor* tensor, int i, float value); -#endif +bool add_rpc_devices(const std::string& servers); +#endif // __SD_CORE_GGML_EXTEND_BACKEND_H__ diff --git a/src/ggml_graph_cut.cpp b/src/core/ggml_graph_cut.cpp similarity index 78% rename from src/ggml_graph_cut.cpp rename to src/core/ggml_graph_cut.cpp index 61234eaf..d4874b05 100644 --- a/src/ggml_graph_cut.cpp +++ b/src/core/ggml_graph_cut.cpp @@ -1,6 +1,8 @@ -#include "ggml_graph_cut.h" +#include "core/ggml_graph_cut.h" #include +#include +#include #include #include #include @@ -8,11 +10,12 @@ #include #include +#include "core/ggml_extend_backend.h" +#include "core/util.h" #include "ggml-alloc.h" #include "ggml-backend.h" -#include "util.h" -#include "../ggml/src/ggml-impl.h" +#include "ggml/src/ggml-impl.h" namespace sd::ggml_graph_cut { @@ -44,7 +47,9 @@ namespace sd::ggml_graph_cut { if (tensor == nullptr) { return false; } - return params_tensor_set.find(tensor) != params_tensor_set.end(); + return params_tensor_set.find(tensor) != params_tensor_set.end() || + (tensor->view_src != nullptr && + params_tensor_set.find(tensor->view_src) != params_tensor_set.end()); } static int graph_node_index_by_name(ggml_cgraph* gf, const char* name) { @@ -81,6 +86,157 @@ namespace sd::ggml_graph_cut { segment.output_bytes; } + static std::string lower_ascii_copy(std::string value) { + std::transform(value.begin(), value.end(), value.begin(), [](unsigned char c) { + return static_cast(std::tolower(c)); + }); + return value; + } + + static std::string normalize_backend_budget_key(const std::string& value) { + return lower_ascii_copy(trim(value)); + } + + static bool is_default_max_vram_key(const std::string& key) { + std::string normalized = normalize_backend_budget_key(key); + return normalized == "all" || normalized == "default" || normalized == "*"; + } + + static bool parse_max_vram_budget_value(const std::string& text, float* value, std::string* error) { + float parsed = 0.f; + if (!parse_strict_float(text, parsed) || !std::isfinite(parsed)) { + if (error != nullptr) { + *error = "invalid --max-vram value '" + text + "'"; + } + return false; + } + *value = parsed; + return true; + } + + static std::vector backend_budget_keys(ggml_backend_t backend) { + std::vector keys; + if (backend == nullptr) { + return keys; + } + + ggml_backend_dev_t dev = ggml_backend_get_device(backend); + if (dev != nullptr) { + keys.push_back(normalize_backend_budget_key(ggml_backend_dev_name(dev))); + } + const char* backend_name = ggml_backend_name(backend); + if (backend_name != nullptr) { + keys.push_back(normalize_backend_budget_key(backend_name)); + } + return keys; + } + + void MaxVramAssignment::reset(float fallback_gib) { + default_gib = fallback_gib; + backend_gib.clear(); + resolved_backend_bytes.clear(); + } + + bool MaxVramAssignment::parse(const std::string& raw_spec, std::string* error) { + const std::string in = trim(raw_spec); + if (in.empty()) { + return true; + } + + for (const std::string& raw_part : split_string(in, ',')) { + const std::string part = trim(raw_part); + if (part.empty()) { + continue; + } + + const size_t eq = part.find('='); + if (eq == std::string::npos) { + float value = 0.f; + if (!parse_max_vram_budget_value(part, &value, error)) { + return false; + } + default_gib = value; + continue; + } + + const std::string key = trim(part.substr(0, eq)); + const std::string value_text = trim(part.substr(eq + 1)); + if (key.empty() || value_text.empty()) { + if (error != nullptr) { + *error = "invalid --max-vram assignment '" + part + "'"; + } + return false; + } + + float value = 0.f; + if (!parse_max_vram_budget_value(value_text, &value, error)) { + return false; + } + + if (is_default_max_vram_key(key)) { + default_gib = value; + continue; + } + + const std::string backend_key = trim(key); + if (backend_key.empty()) { + if (error != nullptr) { + *error = "invalid --max-vram backend key in '" + part + "'"; + } + return false; + } + backend_gib[backend_key] = value; + } + resolved_backend_bytes.clear(); + return true; + } + + bool MaxVramAssignment::canonicalize_backend_keys(std::string* error) { + if (backend_gib.empty()) { + return true; + } + + std::unordered_map normalized; + for (const auto& kv : backend_gib) { + std::string resolved = sd_backend_resolve_name(kv.first); + if (resolved.empty()) { + if (error != nullptr) { + *error = "unknown --max-vram backend '" + kv.first + "'"; + } + return false; + } + normalized[normalize_backend_budget_key(resolved)] = kv.second; + } + backend_gib = std::move(normalized); + resolved_backend_bytes.clear(); + return true; + } + + size_t MaxVramAssignment::bytes_for_backend(ggml_backend_t backend) { + std::vector keys = backend_budget_keys(backend); + const std::string cache_key = keys.empty() ? std::string("") : keys.front(); + auto cached = resolved_backend_bytes.find(cache_key); + if (cached != resolved_backend_bytes.end()) { + return cached->second; + } + + float budget_gib = default_gib; + if (!backend_gib.empty()) { + for (const std::string& key : keys) { + auto backend_it = backend_gib.find(key); + if (backend_it != backend_gib.end()) { + budget_gib = backend_it->second; + break; + } + } + } + + const float resolved_gib = resolve_max_vram_gib(budget_gib, backend); + const size_t bytes = max_vram_gib_to_bytes(resolved_gib); + resolved_backend_bytes[cache_key] = bytes; + return bytes; + } + size_t max_vram_gib_to_bytes(float max_vram) { if (max_vram <= 0.f) { return 0; @@ -135,6 +291,24 @@ namespace sd::ggml_graph_cut { return max_vram_bytes_to_gib(resolve_auto_max_vram_bytes(-max_vram, backend)); } + static bool is_segment_output_needed_after(const Plan& plan, + size_t end_segment_index, + int output_node_index) { + if (end_segment_index + 1 >= plan.segments.size()) { + return false; + } + for (size_t seg_idx = end_segment_index + 1; seg_idx < plan.segments.size(); ++seg_idx) { + const auto& segment = plan.segments[seg_idx]; + for (const auto& input_ref : segment.input_refs) { + if (input_ref.type == Segment::INPUT_PREVIOUS_CUT && + input_ref.node_index == output_node_index) { + return true; + } + } + } + return false; + } + static Segment make_segment_seed(const Plan& plan, size_t start_segment_index, size_t end_segment_index) { @@ -147,8 +321,11 @@ namespace sd::ggml_graph_cut { const auto& target_segment = plan.segments[end_segment_index]; std::unordered_set seen_output_node_indices; for (size_t seg_idx = start_segment_index; seg_idx <= end_segment_index; ++seg_idx) { + const bool is_boundary_segment = seg_idx == end_segment_index; for (int output_node_index : plan.segments[seg_idx].output_node_indices) { - if (seen_output_node_indices.insert(output_node_index).second) { + if ((is_boundary_segment || + is_segment_output_needed_after(plan, end_segment_index, output_node_index)) && + seen_output_node_indices.insert(output_node_index).second) { seed.output_node_indices.push_back(output_node_index); } } @@ -400,23 +577,6 @@ namespace sd::ggml_graph_cut { return tensors; } - std::vector runtime_param_tensors(ggml_cgraph* gf, const Segment& segment, const char* log_desc) { - std::vector tensors = param_tensors(gf, segment); - std::vector filtered_tensors; - filtered_tensors.reserve(tensors.size()); - for (ggml_tensor* tensor : tensors) { - if (tensor_buffer(tensor) == nullptr) { - LOG_WARN("%s graph cut skipping param input without buffer: segment=%s tensor=%s", - log_desc == nullptr ? "unknown" : log_desc, - segment.group_name.c_str(), - tensor->name); - continue; - } - filtered_tensors.push_back(tensor); - } - return filtered_tensors; - } - std::unordered_set collect_future_input_names(ggml_cgraph* gf, const Plan& plan, size_t current_segment_index) { @@ -487,6 +647,44 @@ namespace sd::ggml_graph_cut { return 0; } + struct TensorRuntimeBinding { + ggml_backend_buffer_t buffer = nullptr; + void* data = nullptr; + void* extra = nullptr; + }; + std::unordered_map saved_bindings; + auto mark_measurement_external = [&](ggml_tensor* tensor) { + if (tensor == nullptr) { + return; + } + auto save_tensor = [&](ggml_tensor* t) { + if (t == nullptr || saved_bindings.find(t) != saved_bindings.end()) { + return; + } + saved_bindings[t] = {t->buffer, t->data, t->extra}; + // During real execution params and previous-cut inputs already + // have backend/cache buffers, so gallocr must not reserve them. + t->data = reinterpret_cast(static_cast(1)); + }; + save_tensor(tensor); + save_tensor(tensor->view_src); + }; + for (const auto& input : segment.input_refs) { + if (input.type != Segment::INPUT_PARAM && + input.type != Segment::INPUT_PREVIOUS_CUT) { + continue; + } + mark_measurement_external(input_tensor(gf, input)); + } + + std::unordered_map saved_output_flags; + for (int output_node_index : segment.output_node_indices) { + ggml_tensor* output = ggml_graph_node(gf, output_node_index); + if (output != nullptr && saved_output_flags.find(output) == saved_output_flags.end()) { + saved_output_flags[output] = output->flags; + } + } + ggml_context* graph_ctx = nullptr; ggml_cgraph* segment_graph = build_segment_graph(gf, segment, &graph_ctx); ggml_gallocr_t allocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); @@ -502,6 +700,14 @@ namespace sd::ggml_graph_cut { ggml_gallocr_free(allocr); ggml_free(graph_ctx); + for (const auto& kv : saved_output_flags) { + kv.first->flags = kv.second; + } + for (const auto& kv : saved_bindings) { + kv.first->buffer = kv.second.buffer; + kv.first->data = kv.second.data; + kv.first->extra = kv.second.extra; + } return buffer_size; } @@ -669,7 +875,8 @@ namespace sd::ggml_graph_cut { GGML_ASSERT(!candidate_plan.segments.empty()); const auto& candidate_segment = candidate_plan.segments.back(); - if (graph_cut_segment_vram_bytes(candidate_segment) > max_graph_vram_bytes) { + const size_t candidate_bytes = graph_cut_segment_vram_bytes(candidate_segment); + if (candidate_bytes > max_graph_vram_bytes) { break; } @@ -699,9 +906,9 @@ namespace sd::ggml_graph_cut { } if (log_desc != nullptr) { - LOG_INFO("%s graph cut max_vram budget merge took %lld ms", - log_desc, - ggml_time_ms() - t_budget_begin); + LOG_DEBUG("%s graph cut max_vram budget merge took %lld ms", + log_desc, + ggml_time_ms() - t_budget_begin); } return merged_plan; diff --git a/src/ggml_graph_cut.h b/src/core/ggml_graph_cut.h similarity index 88% rename from src/ggml_graph_cut.h rename to src/core/ggml_graph_cut.h index 9668bc38..17f2f1d7 100644 --- a/src/ggml_graph_cut.h +++ b/src/core/ggml_graph_cut.h @@ -1,9 +1,10 @@ -#ifndef __SD_GGML_GRAPH_CUT_H__ -#define __SD_GGML_GRAPH_CUT_H__ +#ifndef __SD_CORE_GGML_GRAPH_CUT_H__ +#define __SD_CORE_GGML_GRAPH_CUT_H__ #include #include #include +#include #include #include @@ -68,6 +69,17 @@ namespace sd::ggml_graph_cut { static constexpr const char* GGML_RUNNER_CUT_PREFIX = "ggml_runner_cut:"; + struct MaxVramAssignment { + float default_gib = 0.f; + std::unordered_map backend_gib; + std::unordered_map resolved_backend_bytes; + + void reset(float fallback_gib); + bool parse(const std::string& raw_spec, std::string* error); + bool canonicalize_backend_keys(std::string* error); + size_t bytes_for_backend(ggml_backend_t backend); + }; + bool is_graph_cut_tensor(const ggml_tensor* tensor); std::string make_graph_cut_name(const std::string& group, const std::string& output); void mark_graph_cut(ggml_tensor* tensor, const std::string& group, const std::string& output); @@ -80,7 +92,6 @@ namespace sd::ggml_graph_cut { ggml_tensor* output_tensor(ggml_cgraph* gf, const Segment& segment, size_t output_index); ggml_tensor* input_tensor(ggml_cgraph* gf, const Segment::InputRef& input_ref); std::vector param_tensors(ggml_cgraph* gf, const Segment& segment); - std::vector runtime_param_tensors(ggml_cgraph* gf, const Segment& segment, const char* log_desc); std::unordered_set collect_future_input_names(ggml_cgraph* gf, const Plan& plan, size_t current_segment_index); @@ -114,4 +125,4 @@ namespace sd::ggml_graph_cut { void annotate_residency(Plan& plan, size_t max_graph_vram_bytes); } // namespace sd::ggml_graph_cut -#endif +#endif // __SD_CORE_GGML_GRAPH_CUT_H__ diff --git a/src/ordered_map.hpp b/src/core/ordered_map.hpp similarity index 97% rename from src/ordered_map.hpp rename to src/core/ordered_map.hpp index 3fbdca5d..fc4755b8 100644 --- a/src/ordered_map.hpp +++ b/src/core/ordered_map.hpp @@ -1,5 +1,5 @@ -#ifndef __ORDERED_MAP_HPP__ -#define __ORDERED_MAP_HPP__ +#ifndef __SD_CORE_ORDERED_MAP_HPP__ +#define __SD_CORE_ORDERED_MAP_HPP__ #include #include @@ -174,4 +174,4 @@ public: } }; -#endif // __ORDERED_MAP_HPP__ \ No newline at end of file +#endif // __SD_CORE_ORDERED_MAP_HPP__ \ No newline at end of file diff --git a/src/rng.hpp b/src/core/rng.hpp similarity index 89% rename from src/rng.hpp rename to src/core/rng.hpp index accc4088..20ffec7d 100644 --- a/src/rng.hpp +++ b/src/core/rng.hpp @@ -1,5 +1,5 @@ -#ifndef __RNG_H__ -#define __RNG_H__ +#ifndef __SD_CORE_RNG_HPP__ +#define __SD_CORE_RNG_HPP__ #include #include @@ -32,4 +32,4 @@ public: } }; -#endif // __RNG_H__ \ No newline at end of file +#endif // __SD_CORE_RNG_HPP__ \ No newline at end of file diff --git a/src/rng_mt19937.hpp b/src/core/rng_mt19937.hpp similarity index 97% rename from src/rng_mt19937.hpp rename to src/core/rng_mt19937.hpp index 734554bf..85da05cf 100644 --- a/src/rng_mt19937.hpp +++ b/src/core/rng_mt19937.hpp @@ -1,10 +1,10 @@ -#ifndef __RNG_MT19937_HPP__ -#define __RNG_MT19937_HPP__ +#ifndef __SD_CORE_RNG_MT19937_HPP__ +#define __SD_CORE_RNG_MT19937_HPP__ #include #include -#include "rng.hpp" +#include "core/rng.hpp" // RNG imitiating torch cpu randn on CPU. // Port from pytorch, original license: https://github.com/pytorch/pytorch/blob/d01a7b0241ed1c4cded7e7ca097249feb343f072/LICENSE @@ -144,4 +144,4 @@ public: } }; -#endif // __RNG_MT19937_HPP__ \ No newline at end of file +#endif // __SD_CORE_RNG_MT19937_HPP__ \ No newline at end of file diff --git a/src/rng_philox.hpp b/src/core/rng_philox.hpp similarity index 97% rename from src/rng_philox.hpp rename to src/core/rng_philox.hpp index 58da0703..8258ed95 100644 --- a/src/rng_philox.hpp +++ b/src/core/rng_philox.hpp @@ -1,10 +1,10 @@ -#ifndef __RNG_PHILOX_H__ -#define __RNG_PHILOX_H__ +#ifndef __SD_CORE_RNG_PHILOX_HPP__ +#define __SD_CORE_RNG_PHILOX_HPP__ #include #include -#include "rng.hpp" +#include "core/rng.hpp" // RNG imitiating torch cuda randn on CPU. // Port from: https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/5ef669de080814067961f28357256e8fe27544f4/modules/rng_philox.py @@ -122,4 +122,4 @@ public: } }; -#endif // __RNG_PHILOX_H__ \ No newline at end of file +#endif // __SD_CORE_RNG_PHILOX_HPP__ \ No newline at end of file diff --git a/src/tensor.hpp b/src/core/tensor.hpp similarity index 99% rename from src/tensor.hpp rename to src/core/tensor.hpp index 9d054353..ba5dc137 100644 --- a/src/tensor.hpp +++ b/src/core/tensor.hpp @@ -1,5 +1,5 @@ -#ifndef __SD_TENSOR_HPP__ -#define __SD_TENSOR_HPP__ +#ifndef __SD_CORE_TENSOR_HPP__ +#define __SD_CORE_TENSOR_HPP__ #include #include @@ -16,7 +16,7 @@ #include #include -#include "rng.hpp" +#include "core/rng.hpp" namespace sd { @@ -1661,4 +1661,4 @@ namespace sd { } // namespace sd -#endif +#endif // __SD_CORE_TENSOR_HPP__ diff --git a/src/tensor_ggml.hpp b/src/core/tensor_ggml.hpp similarity index 96% rename from src/tensor_ggml.hpp rename to src/core/tensor_ggml.hpp index c6e9d4ac..774574f7 100644 --- a/src/tensor_ggml.hpp +++ b/src/core/tensor_ggml.hpp @@ -1,5 +1,5 @@ -#ifndef __SD_TENSOR_GGML_HPP__ -#define __SD_TENSOR_GGML_HPP__ +#ifndef __SD_CORE_TENSOR_GGML_HPP__ +#define __SD_CORE_TENSOR_GGML_HPP__ #include #include @@ -8,8 +8,8 @@ #include #include +#include "core/tensor.hpp" #include "ggml.h" -#include "tensor.hpp" namespace sd { @@ -124,4 +124,4 @@ namespace sd { } // namespace sd -#endif +#endif // __SD_CORE_TENSOR_GGML_HPP__ diff --git a/src/util.cpp b/src/core/util.cpp similarity index 96% rename from src/util.cpp rename to src/core/util.cpp index 6f5f1990..6d2479f9 100644 --- a/src/util.cpp +++ b/src/core/util.cpp @@ -1,4 +1,4 @@ -#include "util.h" +#include "core/util.h" #include #include #include @@ -13,7 +13,7 @@ #include #include #include -#include "preprocessing.hpp" +#include "runtime/preprocessing.hpp" #if defined(__APPLE__) && defined(__MACH__) #include @@ -346,6 +346,9 @@ int sd_preview_interval = 1; bool sd_preview_denoised = true; bool sd_preview_noisy = false; +static sd_graph_eval_callback_t sd_backend_eval_cb = nullptr; +static void* sd_backend_eval_cb_data = nullptr; + std::u32string utf8_to_utf32(const std::string& utf8_str) { std::wstring_convert, char32_t> converter; return converter.from_bytes(utf8_str); @@ -406,6 +409,15 @@ std::vector split_string(const std::string& str, char delimiter) { return result; } +ggml_type sd_type_to_ggml_type(sd_type_t sdtype) { + const int type_value = static_cast(sdtype); + if (type_value < std::min(SD_TYPE_COUNT, GGML_TYPE_COUNT)) { + return static_cast(type_value); + } else { + return GGML_TYPE_COUNT; + } +} + KeyValueArgs parse_key_value_args(const char* args, const char* context) { KeyValueArgs pairs; @@ -488,7 +500,7 @@ bool parse_strict_bool(const std::string& text, bool& value) { return false; } -static std::string build_progress_bar(int step, int steps) { +static std::string build_progress_bar(int step, int steps, char progress_char = '=', bool show_head = true) { std::string progress = " |"; int max_progress = 50; int32_t current = 0; @@ -498,21 +510,21 @@ static std::string build_progress_bar(int step, int steps) { for (int i = 0; i < 50; i++) { if (i > current) { progress += " "; - } else if (i == current && i != max_progress - 1) { + } else if (show_head && i == current && i != max_progress - 1) { progress += ">"; } else { - progress += "="; + progress += progress_char; } } progress += "|"; return progress; } -static void print_progress_line(int step, int steps, const std::string& speed_text) { +static void print_progress_line(int step, int steps, const std::string& speed_text, char progress_char = '=', bool show_head = true) { if (step == 0) { return; } - std::string progress = build_progress_bar(step, steps); + std::string progress = build_progress_bar(step, steps, progress_char, show_head); const char* lf = (step == steps ? "\n" : ""); printf("\r%s %i/%i - %s\033[K%s", progress.c_str(), step, steps, speed_text.c_str(), lf); fflush(stdout); // for linux @@ -552,9 +564,9 @@ void pretty_bytes_progress(int step, int steps, uint64_t bytes_processed, float double speed_mb = bytes_per_second / (1024.0 * 1024.0); if (speed_mb >= 1024.0) { - print_progress_line(step, steps, sd_format("%.2fGB/s", speed_mb / 1024.0)); + print_progress_line(step, steps, sd_format("%.2fGB/s", speed_mb / 1024.0), '#', false); } else { - print_progress_line(step, steps, sd_format("%.2fMB/s", speed_mb)); + print_progress_line(step, steps, sd_format("%.2fMB/s", speed_mb), '#', false); } } @@ -620,6 +632,11 @@ void sd_set_preview_callback(sd_preview_cb_t cb, preview_t mode, int interval, b sd_preview_noisy = noisy; } +void sd_set_backend_eval_callback(sd_graph_eval_callback_t cb, void* data) { + sd_backend_eval_cb = cb; + sd_backend_eval_cb_data = data; +} + sd_preview_cb_t sd_get_preview_callback() { return sd_preview_cb; } @@ -640,6 +657,14 @@ bool sd_should_preview_noisy() { return sd_preview_noisy; } +sd_graph_eval_callback_t sd_get_backend_eval_callback() { + return sd_backend_eval_cb; +} + +void* sd_get_backend_eval_callback_data() { + return sd_backend_eval_cb_data; +} + sd_progress_cb_t sd_get_progress_callback() { return sd_progress_cb; } diff --git a/src/util.h b/src/core/util.h similarity index 93% rename from src/util.h rename to src/core/util.h index c3b06b1d..35b52061 100644 --- a/src/util.h +++ b/src/core/util.h @@ -1,5 +1,5 @@ -#ifndef __UTIL_H__ -#define __UTIL_H__ +#ifndef __SD_CORE_UTIL_H__ +#define __SD_CORE_UTIL_H__ #include #include @@ -7,9 +7,9 @@ #include #include +#include "core/tensor.hpp" #include "ggml-backend.h" #include "stable-diffusion.h" -#include "tensor.hpp" #define SAFE_STR(s) ((s) ? (s) : "") #define BOOL_STR(b) ((b) ? "true" : "false") @@ -80,6 +80,8 @@ void pretty_bytes_progress(int step, int steps, uint64_t bytes_processed, float void log_printf(sd_log_level_t level, const char* file, int line, const char* format, ...); +ggml_type sd_type_to_ggml_type(sd_type_t sdtype); + std::string trim(const std::string& s); std::vector> parse_prompt_attention(const std::string& text); @@ -96,6 +98,9 @@ int sd_get_preview_interval(); bool sd_should_preview_denoised(); bool sd_should_preview_noisy(); +sd_graph_eval_callback_t sd_get_backend_eval_callback(); +void* sd_get_backend_eval_callback_data(); + // test if the backend is a specific one, e.g. "CUDA", "ROCm", "Vulkan" etc. bool sd_backend_is(ggml_backend_t backend, const std::string& name); @@ -103,4 +108,4 @@ bool sd_backend_is(ggml_backend_t backend, const std::string& name); #define LOG_INFO(format, ...) log_printf(SD_LOG_INFO, __FILE__, __LINE__, format, ##__VA_ARGS__) #define LOG_WARN(format, ...) log_printf(SD_LOG_WARN, __FILE__, __LINE__, format, ##__VA_ARGS__) #define LOG_ERROR(format, ...) log_printf(SD_LOG_ERROR, __FILE__, __LINE__, format, ##__VA_ARGS__) -#endif // __UTIL_H__ +#endif // __SD_CORE_UTIL_H__ diff --git a/src/esrgan.hpp b/src/esrgan.hpp deleted file mode 100644 index a651007d..00000000 --- a/src/esrgan.hpp +++ /dev/null @@ -1,375 +0,0 @@ -#ifndef __ESRGAN_HPP__ -#define __ESRGAN_HPP__ - -#include "ggml_extend.hpp" -#include "model.h" - -/* - =================================== ESRGAN =================================== - References: - https://github.com/xinntao/Real-ESRGAN/blob/master/inference_realesrgan.py - https://github.com/XPixelGroup/BasicSR/blob/v1.4.2/basicsr/archs/rrdbnet_arch.py - -*/ - -class ResidualDenseBlock : public GGMLBlock { -protected: - int num_feat; - int num_grow_ch; - -public: - ResidualDenseBlock(int num_feat = 64, int num_grow_ch = 32) - : num_feat(num_feat), num_grow_ch(num_grow_ch) { - blocks["conv1"] = std::shared_ptr(new Conv2d(num_feat, num_grow_ch, {3, 3}, {1, 1}, {1, 1})); - blocks["conv2"] = std::shared_ptr(new Conv2d(num_feat + num_grow_ch, num_grow_ch, {3, 3}, {1, 1}, {1, 1})); - blocks["conv3"] = std::shared_ptr(new Conv2d(num_feat + 2 * num_grow_ch, num_grow_ch, {3, 3}, {1, 1}, {1, 1})); - blocks["conv4"] = std::shared_ptr(new Conv2d(num_feat + 3 * num_grow_ch, num_grow_ch, {3, 3}, {1, 1}, {1, 1})); - blocks["conv5"] = std::shared_ptr(new Conv2d(num_feat + 4 * num_grow_ch, num_feat, {3, 3}, {1, 1}, {1, 1})); - } - - ggml_tensor* lrelu(GGMLRunnerContext* ctx, ggml_tensor* x) { - return ggml_leaky_relu(ctx->ggml_ctx, x, 0.2f, true); - } - - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { - // x: [n, num_feat, h, w] - // return: [n, num_feat, h, w] - - auto conv1 = std::dynamic_pointer_cast(blocks["conv1"]); - auto conv2 = std::dynamic_pointer_cast(blocks["conv2"]); - auto conv3 = std::dynamic_pointer_cast(blocks["conv3"]); - auto conv4 = std::dynamic_pointer_cast(blocks["conv4"]); - auto conv5 = std::dynamic_pointer_cast(blocks["conv5"]); - - auto x1 = lrelu(ctx, conv1->forward(ctx, x)); - auto x_cat = ggml_concat(ctx->ggml_ctx, x, x1, 2); - auto x2 = lrelu(ctx, conv2->forward(ctx, x_cat)); - x_cat = ggml_concat(ctx->ggml_ctx, x_cat, x2, 2); - auto x3 = lrelu(ctx, conv3->forward(ctx, x_cat)); - x_cat = ggml_concat(ctx->ggml_ctx, x_cat, x3, 2); - auto x4 = lrelu(ctx, conv4->forward(ctx, x_cat)); - x_cat = ggml_concat(ctx->ggml_ctx, x_cat, x4, 2); - auto x5 = conv5->forward(ctx, x_cat); - - x5 = ggml_add(ctx->ggml_ctx, ggml_ext_scale(ctx->ggml_ctx, x5, 0.2f), x); - return x5; - } -}; - -class RRDB : public GGMLBlock { -public: - RRDB(int num_feat, int num_grow_ch = 32) { - blocks["rdb1"] = std::shared_ptr(new ResidualDenseBlock(num_feat, num_grow_ch)); - blocks["rdb2"] = std::shared_ptr(new ResidualDenseBlock(num_feat, num_grow_ch)); - blocks["rdb3"] = std::shared_ptr(new ResidualDenseBlock(num_feat, num_grow_ch)); - } - - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { - // x: [n, num_feat, h, w] - // return: [n, num_feat, h, w] - - auto rdb1 = std::dynamic_pointer_cast(blocks["rdb1"]); - auto rdb2 = std::dynamic_pointer_cast(blocks["rdb2"]); - auto rdb3 = std::dynamic_pointer_cast(blocks["rdb3"]); - - auto out = rdb1->forward(ctx, x); - out = rdb2->forward(ctx, out); - out = rdb3->forward(ctx, out); - - out = ggml_add(ctx->ggml_ctx, ggml_ext_scale(ctx->ggml_ctx, out, 0.2f), x); - return out; - } -}; - -class RRDBNet : public GGMLBlock { -protected: - int scale = 4; - int num_block = 23; - int num_in_ch = 3; - int num_out_ch = 3; - int num_feat = 64; - int num_grow_ch = 32; - -public: - RRDBNet(int scale, int num_block, int num_in_ch, int num_out_ch, int num_feat, int num_grow_ch) - : scale(scale), num_block(num_block), num_in_ch(num_in_ch), num_out_ch(num_out_ch), num_feat(num_feat), num_grow_ch(num_grow_ch) { - blocks["conv_first"] = std::shared_ptr(new Conv2d(num_in_ch, num_feat, {3, 3}, {1, 1}, {1, 1})); - for (int i = 0; i < num_block; i++) { - std::string name = "body." + std::to_string(i); - blocks[name] = std::shared_ptr(new RRDB(num_feat, num_grow_ch)); - } - blocks["conv_body"] = std::shared_ptr(new Conv2d(num_feat, num_feat, {3, 3}, {1, 1}, {1, 1})); - if (scale >= 2) { - blocks["conv_up1"] = std::shared_ptr(new Conv2d(num_feat, num_feat, {3, 3}, {1, 1}, {1, 1})); - } - if (scale == 4) { - blocks["conv_up2"] = std::shared_ptr(new Conv2d(num_feat, num_feat, {3, 3}, {1, 1}, {1, 1})); - } - blocks["conv_hr"] = std::shared_ptr(new Conv2d(num_feat, num_feat, {3, 3}, {1, 1}, {1, 1})); - blocks["conv_last"] = std::shared_ptr(new Conv2d(num_feat, num_out_ch, {3, 3}, {1, 1}, {1, 1})); - } - - int get_scale() { return scale; } - int get_num_block() { return num_block; } - - ggml_tensor* lrelu(GGMLRunnerContext* ctx, ggml_tensor* x) { - return ggml_leaky_relu(ctx->ggml_ctx, x, 0.2f, true); - } - - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { - // x: [n, num_in_ch, h, w] - // return: [n, num_out_ch, h*scale, w*scale] - auto conv_first = std::dynamic_pointer_cast(blocks["conv_first"]); - auto conv_body = std::dynamic_pointer_cast(blocks["conv_body"]); - auto conv_hr = std::dynamic_pointer_cast(blocks["conv_hr"]); - auto conv_last = std::dynamic_pointer_cast(blocks["conv_last"]); - - auto feat = conv_first->forward(ctx, x); - sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.prelude", "feat"); - auto body_feat = feat; - for (int i = 0; i < num_block; i++) { - std::string name = "body." + std::to_string(i); - auto block = std::dynamic_pointer_cast(blocks[name]); - - body_feat = block->forward(ctx, body_feat); - sd::ggml_graph_cut::mark_graph_cut(body_feat, "esrgan.body." + std::to_string(i), "feat"); - } - body_feat = conv_body->forward(ctx, body_feat); - feat = ggml_add(ctx->ggml_ctx, feat, body_feat); - sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.body.out", "feat"); - // upsample - if (scale >= 2) { - auto conv_up1 = std::dynamic_pointer_cast(blocks["conv_up1"]); - feat = lrelu(ctx, conv_up1->forward(ctx, ggml_upscale(ctx->ggml_ctx, feat, 2, GGML_SCALE_MODE_NEAREST))); - sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.up1", "feat"); - if (scale == 4) { - auto conv_up2 = std::dynamic_pointer_cast(blocks["conv_up2"]); - feat = lrelu(ctx, conv_up2->forward(ctx, ggml_upscale(ctx->ggml_ctx, feat, 2, GGML_SCALE_MODE_NEAREST))); - sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.up2", "feat"); - } - } - // for all scales - auto out = conv_last->forward(ctx, lrelu(ctx, conv_hr->forward(ctx, feat))); - sd::ggml_graph_cut::mark_graph_cut(out, "esrgan.final", "out"); - return out; - } -}; - -struct ESRGAN : public GGMLRunner { - std::unique_ptr rrdb_net; - int scale = 4; - int tile_size = 128; // avoid cuda OOM for 4gb VRAM - - ESRGAN(ggml_backend_t backend, - ggml_backend_t params_backend, - int tile_size = 128, - const String2TensorStorage& tensor_storage_map = {}) - : GGMLRunner(backend, params_backend) { - this->tile_size = tile_size; - } - - std::string get_desc() override { - return "esrgan"; - } - - bool load_from_file(const std::string& file_path, int n_threads) { - LOG_INFO("loading esrgan from '%s'", file_path.c_str()); - - ModelLoader model_loader; - if (!model_loader.init_from_file_and_convert_name(file_path)) { - LOG_ERROR("init esrgan model loader from file failed: '%s'", file_path.c_str()); - return false; - } - - // Get tensor names - auto tensor_names = model_loader.get_tensor_names(); - - // Detect if it's ESRGAN format - bool is_ESRGAN = std::find(tensor_names.begin(), tensor_names.end(), "model.0.weight") != tensor_names.end(); - - // Detect parameters from tensor names - int detected_num_block = 0; - if (is_ESRGAN) { - for (const auto& name : tensor_names) { - if (name.find("model.1.sub.") == 0) { - size_t first_dot = name.find('.', 12); - if (first_dot != std::string::npos) { - size_t second_dot = name.find('.', first_dot + 1); - if (second_dot != std::string::npos && name.substr(first_dot + 1, 3) == "RDB") { - try { - int idx = std::stoi(name.substr(12, first_dot - 12)); - detected_num_block = std::max(detected_num_block, idx + 1); - } catch (...) { - } - } - } - } - } - } else { - // Original format - for (const auto& name : tensor_names) { - if (name.find("body.") == 0) { - size_t pos = name.find('.', 5); - if (pos != std::string::npos) { - try { - int idx = std::stoi(name.substr(5, pos - 5)); - detected_num_block = std::max(detected_num_block, idx + 1); - } catch (...) { - } - } - } - } - } - - int detected_scale = 4; // default - if (is_ESRGAN) { - // For ESRGAN format, detect scale by highest model number - int max_model_num = 0; - for (const auto& name : tensor_names) { - if (name.find("model.") == 0) { - size_t dot_pos = name.find('.', 6); - if (dot_pos != std::string::npos) { - try { - int num = std::stoi(name.substr(6, dot_pos - 6)); - max_model_num = std::max(max_model_num, num); - } catch (...) { - } - } - } - } - if (max_model_num <= 4) { - detected_scale = 1; - } else if (max_model_num <= 7) { - detected_scale = 2; - } else { - detected_scale = 4; - } - } else { - // Original format - bool has_conv_up2 = std::any_of(tensor_names.begin(), tensor_names.end(), [](const std::string& name) { - return name == "conv_up2.weight"; - }); - bool has_conv_up1 = std::any_of(tensor_names.begin(), tensor_names.end(), [](const std::string& name) { - return name == "conv_up1.weight"; - }); - if (has_conv_up2) { - detected_scale = 4; - } else if (has_conv_up1) { - detected_scale = 2; - } else { - detected_scale = 1; - } - } - - int detected_num_in_ch = 3; - int detected_num_out_ch = 3; - int detected_num_feat = 64; - int detected_num_grow_ch = 32; - - // Create RRDBNet with detected parameters - rrdb_net = std::make_unique(detected_scale, detected_num_block, detected_num_in_ch, detected_num_out_ch, detected_num_feat, detected_num_grow_ch); - rrdb_net->init(params_ctx, {}, ""); - - if (!alloc_params_buffer()) { - LOG_ERROR("esrgan model buffer allocation failed"); - return false; - } - - std::map esrgan_tensors; - rrdb_net->get_param_tensors(esrgan_tensors); - - bool success; - if (is_ESRGAN) { - // Build name mapping for ESRGAN format - std::map expected_to_model; - expected_to_model["conv_first.weight"] = "model.0.weight"; - expected_to_model["conv_first.bias"] = "model.0.bias"; - - for (int i = 0; i < detected_num_block; i++) { - for (int j = 1; j <= 3; j++) { - for (int k = 1; k <= 5; k++) { - std::string expected_weight = "body." + std::to_string(i) + ".rdb" + std::to_string(j) + ".conv" + std::to_string(k) + ".weight"; - std::string model_weight = "model.1.sub." + std::to_string(i) + ".RDB" + std::to_string(j) + ".conv" + std::to_string(k) + ".0.weight"; - expected_to_model[expected_weight] = model_weight; - - std::string expected_bias = "body." + std::to_string(i) + ".rdb" + std::to_string(j) + ".conv" + std::to_string(k) + ".bias"; - std::string model_bias = "model.1.sub." + std::to_string(i) + ".RDB" + std::to_string(j) + ".conv" + std::to_string(k) + ".0.bias"; - expected_to_model[expected_bias] = model_bias; - } - } - } - - if (detected_scale == 1) { - expected_to_model["conv_body.weight"] = "model.1.sub." + std::to_string(detected_num_block) + ".weight"; - expected_to_model["conv_body.bias"] = "model.1.sub." + std::to_string(detected_num_block) + ".bias"; - expected_to_model["conv_hr.weight"] = "model.2.weight"; - expected_to_model["conv_hr.bias"] = "model.2.bias"; - expected_to_model["conv_last.weight"] = "model.4.weight"; - expected_to_model["conv_last.bias"] = "model.4.bias"; - } else { - expected_to_model["conv_body.weight"] = "model.1.sub." + std::to_string(detected_num_block) + ".weight"; - expected_to_model["conv_body.bias"] = "model.1.sub." + std::to_string(detected_num_block) + ".bias"; - if (detected_scale >= 2) { - expected_to_model["conv_up1.weight"] = "model.3.weight"; - expected_to_model["conv_up1.bias"] = "model.3.bias"; - } - if (detected_scale == 4) { - expected_to_model["conv_up2.weight"] = "model.6.weight"; - expected_to_model["conv_up2.bias"] = "model.6.bias"; - expected_to_model["conv_hr.weight"] = "model.8.weight"; - expected_to_model["conv_hr.bias"] = "model.8.bias"; - expected_to_model["conv_last.weight"] = "model.10.weight"; - expected_to_model["conv_last.bias"] = "model.10.bias"; - } else if (detected_scale == 2) { - expected_to_model["conv_hr.weight"] = "model.5.weight"; - expected_to_model["conv_hr.bias"] = "model.5.bias"; - expected_to_model["conv_last.weight"] = "model.7.weight"; - expected_to_model["conv_last.bias"] = "model.7.bias"; - } - } - - std::map model_tensors; - for (auto& p : esrgan_tensors) { - auto it = expected_to_model.find(p.first); - if (it != expected_to_model.end()) { - model_tensors[it->second] = p.second; - } - } - - success = model_loader.load_tensors(model_tensors, {}, n_threads); - } else { - success = model_loader.load_tensors(esrgan_tensors, {}, n_threads); - } - - if (!success) { - LOG_ERROR("load esrgan tensors from model loader failed"); - return false; - } - - scale = rrdb_net->get_scale(); - LOG_INFO("esrgan model loaded with scale=%d, num_block=%d", scale, detected_num_block); - return success; - } - - ggml_cgraph* build_graph(const sd::Tensor& x_tensor) { - if (!rrdb_net) - return nullptr; - constexpr int kGraphNodes = 1 << 16; // 65k - ggml_cgraph* gf = new_graph_custom(kGraphNodes); - ggml_tensor* x = make_input(x_tensor); - - auto runner_ctx = get_context(); - ggml_tensor* out = rrdb_net->forward(&runner_ctx, x); - ggml_build_forward_expand(gf, out); - return gf; - } - - sd::Tensor compute(const int n_threads, - const sd::Tensor& x) { - auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x); }; - auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); - return result; - } -}; - -#endif // __ESRGAN_HPP__ diff --git a/src/extensions/generation_extension.h b/src/extensions/generation_extension.h new file mode 100644 index 00000000..67085c15 --- /dev/null +++ b/src/extensions/generation_extension.h @@ -0,0 +1,77 @@ +#ifndef __SD_EXTENSIONS_GENERATION_EXTENSION_H__ +#define __SD_EXTENSIONS_GENERATION_EXTENSION_H__ + +#include +#include +#include +#include +#include +#include + +#include "conditioning/conditioner.hpp" +#include "core/ggml_extend_backend.h" +#include "model/diffusion/model.hpp" +#include "model_loader.h" +#include "model_manager.h" +#include "stable-diffusion.h" + +struct GenerationExtensionInitContext { + const sd_ctx_params_t* params; + SDVersion version; + const String2TensorStorage& tensor_storage_map; + ModelLoader& model_loader; + std::shared_ptr model_manager; + int n_threads; + std::function ensure_backend_pair; + std::function backend_for; + std::function params_backend_for; +}; + +struct GenerationExtensionConditionContext { + Conditioner* conditioner; + ConditionerParams& condition_params; + const sd_pm_params_t& pm_params; + const sd_pulid_params_t& pulid_params; + int n_threads; + int total_steps; +}; + +struct GenerationExtension { + virtual ~GenerationExtension() = default; + + virtual const char* name() const = 0; + virtual bool is_enabled() const { + return false; + } + virtual bool init(const GenerationExtensionInitContext&) { + return true; + } + virtual void get_param_tensors(std::map&) {} + virtual void collect_loras(std::vector&) {} + virtual void add_ignore_tensors(std::set&) const {} + virtual void runner_done() {} + virtual void reset_runtime_condition() {} + virtual bool prepare_condition(GenerationExtensionConditionContext&) { + return false; + } + virtual const SDCondition& before_condition(int step, + const SDCondition& condition) const { + return condition; + } + + // Called in the denoise loop for each enabled extension, after the per-step + // DiffusionParams (including its version-specific `extra`) has been built, + // but before diffusion_model->compute(). Lets an extension feed data into + // the diffusion forward that the conditioning-side hooks can't reach -- it + // can set/override fields on `params` (typically the architecture-specific + // `params.extra`, e.g. a guidance tensor, control payload, or an identity + // embedding for an adapter that injects inside the model's blocks). The + // extension targets whichever `extra` variant matches the active model. + // Mutates `params` only, never the extension. Default no-op. + virtual void before_diffusion(DiffusionParams& /*params*/, int /*step*/) const {} +}; + +std::shared_ptr create_photomaker_extension(); +std::shared_ptr create_pulid_extension(); + +#endif diff --git a/src/extensions/photomaker_extension.cpp b/src/extensions/photomaker_extension.cpp new file mode 100644 index 00000000..78c5cdb9 --- /dev/null +++ b/src/extensions/photomaker_extension.cpp @@ -0,0 +1,292 @@ +#include "extensions/generation_extension.h" + +#include +#include +#include +#include + +#include "core/tensor_ggml.hpp" +#include "core/util.h" +#include "model/adapter/pmid.hpp" + +static std::tuple, std::vector, std::vector> +tokenize_photomaker_trigger(FrozenCLIPEmbedderWithCustomWords& clip_conditioner, + const std::string& text, + int trigger_token_count, + int32_t image_token) { + auto tokens_and_weights = clip_conditioner.tokenize(text); + std::vector source_tokens = std::move(tokens_and_weights.first); + std::vector source_weights = std::move(tokens_and_weights.second); + + if (!source_tokens.empty() && source_tokens.front() == clip_conditioner.tokenizer.BOS_TOKEN_ID) { + source_tokens.erase(source_tokens.begin()); + source_weights.erase(source_weights.begin()); + } + if (!source_tokens.empty() && source_tokens.back() == clip_conditioner.tokenizer.EOS_TOKEN_ID) { + source_tokens.pop_back(); + source_weights.pop_back(); + } + + std::vector tokens; + std::vector weights; + int32_t class_idx = -1; + for (size_t i = 0; i < source_tokens.size(); i++) { + int token = source_tokens[i]; + if (token == image_token) { + if (!tokens.empty()) { + class_idx = static_cast(tokens.size()) - 1; + int class_token = tokens.back(); + float class_weight = weights.back(); + for (int j = 1; j < trigger_token_count; j++) { + tokens.push_back(class_token); + weights.push_back(class_weight); + } + } + continue; + } + tokens.push_back(token); + weights.push_back(source_weights[i]); + } + + clip_conditioner.tokenizer.pad_tokens(tokens, + &weights, + nullptr, + clip_conditioner.text_model->model.n_token, + clip_conditioner.text_model->model.n_token, + true); + std::vector class_token_mask; + for (int i = 0; i < tokens.size(); i++) { + class_token_mask.push_back(class_idx + 1 <= i && i < class_idx + 1 + trigger_token_count); + } + + return std::make_tuple(tokens, weights, class_token_mask); +} + +static std::tuple> +get_photomaker_condition_with_trigger(FrozenCLIPEmbedderWithCustomWords& clip_conditioner, + int n_threads, + const ConditionerParams& conditioner_params, + const std::string& trigger_word, + int trigger_token_count) { + auto image_tokens = clip_conditioner.convert_token_to_id(trigger_word); + GGML_ASSERT(image_tokens.size() == 1); + auto tokens_and_weights = tokenize_photomaker_trigger(clip_conditioner, + conditioner_params.text, + trigger_token_count, + image_tokens[0]); + std::vector& tokens = std::get<0>(tokens_and_weights); + std::vector& weights = std::get<1>(tokens_and_weights); + std::vector& trigger_mask = std::get<2>(tokens_and_weights); + auto cond = clip_conditioner.get_learned_condition_common(n_threads, + tokens, + weights, + conditioner_params.clip_skip, + conditioner_params.width, + conditioner_params.height, + conditioner_params.zero_out_masked); + return std::make_tuple(std::move(cond), trigger_mask); +} + +static std::string remove_photomaker_trigger_from_prompt(FrozenCLIPEmbedderWithCustomWords& clip_conditioner, + const std::string& prompt, + const std::string& trigger_word) { + auto image_tokens = clip_conditioner.convert_token_to_id(trigger_word); + GGML_ASSERT(image_tokens.size() == 1); + auto tokens_and_weights = clip_conditioner.tokenize(prompt); + std::vector& tokens = tokens_and_weights.first; + auto it = std::find(tokens.begin(), tokens.end(), image_tokens[0]); + GGML_ASSERT(it != tokens.end()); + tokens.erase(it); + return clip_conditioner.decode(tokens); +} + +struct PhotoMakerExtension : public GenerationExtension { + std::shared_ptr pmid_model; + bool enabled = false; + std::string model_path; + std::string trigger_word = "img"; + SDCondition id_condition; + int start_merge_step = -1; + + const char* name() const override { + return "photomaker"; + } + + bool is_enabled() const override { + return enabled; + } + + bool init(const GenerationExtensionInitContext& ctx) override { + model_path = SAFE_STR(ctx.params->photo_maker_path); + if (model_path.empty()) { + return true; + } + + if (!ctx.ensure_backend_pair(SDBackendModule::PHOTOMAKER)) { + return false; + } + + PMVersion pm_version = std::strstr(model_path.c_str(), "v2") != nullptr ? PM_VERSION_2 : PM_VERSION_1; + LOG_INFO("loading stacked ID embedding (PHOTOMAKER) model file from '%s'", model_path.c_str()); + if (!ctx.model_loader.init_from_file_and_convert_name(model_path, "pmid.")) { + LOG_WARN("loading stacked ID embedding from '%s' failed", model_path.c_str()); + return true; + } + + pmid_model = std::make_shared(ctx.backend_for(SDBackendModule::PHOTOMAKER), + ctx.tensor_storage_map, + "pmid", + ctx.version, + pm_version, + 20.f, + ctx.model_manager); + if (pm_version == PM_VERSION_2) { + LOG_INFO("using PhotoMaker Version 2"); + } + + enabled = true; + return true; + } + + void get_param_tensors(std::map& tensors) override { + if (!enabled || pmid_model == nullptr) { + return; + } + + pmid_model->get_param_tensors(tensors, "pmid"); + } + + void collect_loras(std::vector& loras) override { + if (!enabled || model_path.empty()) { + return; + } + ModelManager::LoraSpec lora; + lora.path = model_path; + lora.multiplier = 1.0f; + lora.tensor_name_prefix_filter = "lora.model"; + lora.required = true; + loras.push_back(std::move(lora)); + } + + void add_ignore_tensors(std::set& ignore_tensors) const override { + if (!enabled) { + return; + } + ignore_tensors.insert("pmid.unet."); + } + + void runner_done() override { + if (pmid_model != nullptr) { + pmid_model->runner_done(); + } + } + + void reset_runtime_condition() override { + id_condition = {}; + start_merge_step = -1; + } + + bool prepare_condition(GenerationExtensionConditionContext& ctx) override { + reset_runtime_condition(); + if (!enabled || pmid_model == nullptr) { + return false; + } + + bool pmv2 = pmid_model->get_version() == PM_VERSION_2; + if (ctx.pm_params.id_images_count <= 0 || ctx.pm_params.id_images == nullptr) { + LOG_WARN("Provided PhotoMaker model file, but NO input ID images"); + LOG_WARN("Turn off PhotoMaker for this request"); + return false; + } + auto* clip_conditioner = dynamic_cast(ctx.conditioner); + if (clip_conditioner == nullptr) { + LOG_WARN("PhotoMaker requires FrozenCLIPEmbedderWithCustomWords conditioner"); + LOG_WARN("Turn off PhotoMaker for this request"); + return false; + } + + int clip_image_size = 224; + pmid_model->style_strength = ctx.pm_params.style_strength; + sd::Tensor id_image_tensor; + for (int i = 0; i < ctx.pm_params.id_images_count; i++) { + auto id_image = sd_image_to_tensor(ctx.pm_params.id_images[i]); + auto processed_id_image = clip_preprocess(id_image, clip_image_size, clip_image_size); + if (id_image_tensor.empty()) { + id_image_tensor = processed_id_image; + } else { + id_image_tensor = sd::ops::concat(id_image_tensor, processed_id_image, 3); + } + } + + int64_t t0 = ggml_time_ms(); + int trigger_token_count = pmv2 ? 2 * ctx.pm_params.id_images_count : ctx.pm_params.id_images_count; + auto cond_tup = get_photomaker_condition_with_trigger(*clip_conditioner, + ctx.n_threads, + ctx.condition_params, + trigger_word, + trigger_token_count); + SDCondition prepared_id_condition = std::get<0>(cond_tup); + auto class_tokens_mask = std::get<1>(cond_tup); + if (std::find(class_tokens_mask.begin(), class_tokens_mask.end(), true) == class_tokens_mask.end()) { + LOG_WARN("PhotoMaker trigger word '%s' was not found in prompt", trigger_word.c_str()); + LOG_WARN("Turn off PhotoMaker for this request"); + return false; + } + + sd::Tensor id_embeds; + if (pmv2 && ctx.pm_params.id_embed_path != nullptr) { + try { + id_embeds = sd::load_tensor_from_file_as_tensor(ctx.pm_params.id_embed_path); + } catch (const std::exception&) { + id_embeds = {}; + } + } + if (pmv2 && id_embeds.empty()) { + LOG_WARN("Provided PhotoMaker images, but NO valid ID embeds file for PM v2"); + LOG_WARN("Turn off PhotoMaker for this request"); + return false; + } + if (pmv2 && ctx.pm_params.id_images_count != id_embeds.shape()[1]) { + LOG_WARN("PhotoMaker image count (%d) does NOT match ID embeds (%d). You should run face_detect.py again.", + ctx.pm_params.id_images_count, + static_cast(id_embeds.shape()[1])); + LOG_WARN("Turn off PhotoMaker for this request"); + return false; + } + + auto res = pmid_model->compute(ctx.n_threads, + id_image_tensor, + prepared_id_condition.c_crossattn, + id_embeds, + class_tokens_mask); + if (res.empty()) { + LOG_ERROR("Photomaker ID Stacking failed"); + LOG_WARN("Turn off PhotoMaker for this request"); + return false; + } + + prepared_id_condition.c_crossattn = std::move(res); + int64_t t1 = ggml_time_ms(); + id_condition = std::move(prepared_id_condition); + start_merge_step = int(ctx.pm_params.style_strength / 100.f * ctx.total_steps); + ctx.condition_params.text = remove_photomaker_trigger_from_prompt(*clip_conditioner, + ctx.condition_params.text, + trigger_word); + LOG_INFO("Photomaker ID Stacking, taking %" PRId64 " ms", t1 - t0); + LOG_INFO("PHOTOMAKER: start_merge_step: %d", start_merge_step); + + return true; + } + + const SDCondition& before_condition(int step, + const SDCondition& condition) const override { + if (!id_condition.empty() && start_merge_step != -1 && step > start_merge_step) { + return id_condition; + } + return condition; + } +}; + +std::shared_ptr create_photomaker_extension() { + return std::make_shared(); +} diff --git a/src/extensions/pulid_extension.cpp b/src/extensions/pulid_extension.cpp new file mode 100644 index 00000000..d529e571 --- /dev/null +++ b/src/extensions/pulid_extension.cpp @@ -0,0 +1,123 @@ +#include "extensions/generation_extension.h" + +#include +#include + +#include "core/tensor_ggml.hpp" +#include "core/util.h" +#include "gguf.h" + +static sd::Tensor load_pulid_id_embedding(const char* path) { + sd::Tensor empty; + if (path == nullptr || strlen(path) == 0) { + return empty; + } + + struct ggml_context* ctx_data = nullptr; + struct gguf_init_params gp = {/*.no_alloc =*/false, /*.ctx =*/&ctx_data}; + struct gguf_context* gguf_ctx = gguf_init_from_file(path, gp); + if (gguf_ctx == nullptr || ctx_data == nullptr) { + LOG_WARN("PuLID id-embedding: cannot read gguf '%s'", path); + if (gguf_ctx != nullptr) + gguf_free(gguf_ctx); + if (ctx_data != nullptr) + ggml_free(ctx_data); + return empty; + } + + struct ggml_tensor* t = ggml_get_tensor(ctx_data, "pulid_id"); + if (t == nullptr) { + LOG_WARN("PuLID id-embedding: no 'pulid_id' tensor in '%s'", path); + gguf_free(gguf_ctx); + ggml_free(ctx_data); + return empty; + } + + const int64_t token_dim = t->ne[0]; + const int64_t num_tokens = t->ne[1]; + if (token_dim <= 0 || num_tokens <= 0 || token_dim > 65536 || num_tokens > 1024 || + t->ne[2] != 1 || t->ne[3] != 1) { + LOG_WARN("PuLID id-embedding: implausible shape [%lld, %lld] in '%s'", + (long long)token_dim, (long long)num_tokens, path); + gguf_free(gguf_ctx); + ggml_free(ctx_data); + return empty; + } + + const size_t n_elem = (size_t)token_dim * (size_t)num_tokens; + sd::Tensor out({token_dim, num_tokens, 1}); + float* dst = out.data(); + if (t->type == GGML_TYPE_F32) { + memcpy(dst, t->data, n_elem * sizeof(float)); + } else if (t->type == GGML_TYPE_F16) { + const ggml_fp16_t* src = reinterpret_cast(t->data); + for (size_t i = 0; i < n_elem; i++) { + dst[i] = ggml_fp16_to_fp32(src[i]); + } + } else if (t->type == GGML_TYPE_BF16) { + const ggml_bf16_t* src = reinterpret_cast(t->data); + for (size_t i = 0; i < n_elem; i++) { + dst[i] = ggml_bf16_to_fp32(src[i]); + } + } else { + LOG_WARN("PuLID id-embedding: unsupported tensor type %s in '%s'", + ggml_type_name(t->type), path); + gguf_free(gguf_ctx); + ggml_free(ctx_data); + return empty; + } + + LOG_INFO("PuLID id-embedding: loaded [%lld, %lld] type=%s from '%s'", + (long long)token_dim, (long long)num_tokens, ggml_type_name(t->type), path); + gguf_free(gguf_ctx); + ggml_free(ctx_data); + return out; +} + +struct PuLIDExtension : public GenerationExtension { + bool enabled = false; + sd::Tensor id_embedding; + float id_weight = 1.0f; + + const char* name() const override { + return "pulid"; + } + + bool is_enabled() const override { + return enabled; + } + + bool init(const GenerationExtensionInitContext& ctx) override { + enabled = strlen(SAFE_STR(ctx.params->pulid_weights_path)) > 0; + return true; + } + + void reset_runtime_condition() override { + id_embedding = {}; + id_weight = 1.0f; + } + + bool prepare_condition(GenerationExtensionConditionContext& ctx) override { + reset_runtime_condition(); + if (!enabled) { + return false; + } + id_embedding = load_pulid_id_embedding(ctx.pulid_params.id_embedding_path); + id_weight = ctx.pulid_params.id_weight; + return false; // PuLID does not modify the conditioning + } + + void before_diffusion(DiffusionParams& params, int /*step*/) const override { + if (!enabled || id_embedding.empty()) { + return; + } + if (auto* flux_extra = std::get_if(¶ms.extra)) { + flux_extra->pulid_id = &id_embedding; + flux_extra->pulid_id_weight = id_weight; + } + } +}; + +std::shared_ptr create_pulid_extension() { + return std::make_shared(); +} diff --git a/src/gits_noise.inl b/src/gits_noise.inl deleted file mode 100644 index 7a10ff76..00000000 --- a/src/gits_noise.inl +++ /dev/null @@ -1,349 +0,0 @@ -#ifndef GITS_NOISE_INL -#define GITS_NOISE_INL - -const std::vector> GITS_NOISE_0_80 = { - { 14.61464119f, 7.49001646f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 6.77309084f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 7.49001646f, 3.07277966f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 2.05039096f, 0.02916753f }, - { 14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.85520077f, 2.05039096f, 0.02916753f }, - { 14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.85520077f, 3.07277966f, 1.56271636f, 0.02916753f }, - { 14.61464119f, 12.96784878f, 11.54541874f, 8.75849152f, 7.49001646f, 5.85520077f, 3.07277966f, 1.56271636f, 0.02916753f }, - { 14.61464119f, 13.76078796f, 12.23089790f, 10.90732002f, 8.75849152f, 7.49001646f, 5.85520077f, 3.07277966f, 1.56271636f, 0.02916753f }, - { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 10.90732002f, 8.75849152f, 7.49001646f, 5.85520077f, 3.07277966f, 1.56271636f, 0.02916753f }, - { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 5.85520077f, 3.07277966f, 1.56271636f, 0.02916753f }, - { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.07277966f, 1.56271636f, 0.02916753f }, - { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.07277966f, 1.56271636f, 0.02916753f }, - { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.07277966f, 1.56271636f, 0.02916753f }, - { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.07277966f, 1.56271636f, 0.02916753f }, - { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.75859547f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.19567990f, 1.98035145f, 0.86115354f, 0.02916753f }, - { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.75859547f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.19567990f, 1.98035145f, 0.86115354f, 0.02916753f }, - { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.75859547f, 9.24142551f, 8.75849152f, 8.30717278f, 7.88507891f, 7.49001646f, 6.77309084f, 5.85520077f, 4.65472794f, 3.07277966f, 1.84880662f, 0.83188516f, 0.02916753f } -}; - -const std::vector> GITS_NOISE_0_85 = { - { 14.61464119f, 7.49001646f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 1.84880662f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 6.77309084f, 1.56271636f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 7.11996698f, 3.07277966f, 1.24153244f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 7.49001646f, 5.09240818f, 2.84484982f, 0.95350921f, 0.02916753f }, - { 14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.09240818f, 2.84484982f, 0.95350921f, 0.02916753f }, - { 14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.58536053f, 3.19567990f, 1.84880662f, 0.803307f, 0.02916753f }, - { 14.61464119f, 12.96784878f, 11.54541874f, 8.75849152f, 7.49001646f, 5.58536053f, 3.19567990f, 1.84880662f, 0.803307f, 0.02916753f }, - { 14.61464119f, 12.96784878f, 11.54541874f, 8.75849152f, 7.49001646f, 6.14220476f, 4.65472794f, 3.07277966f, 1.84880662f, 0.803307f, 0.02916753f }, - { 14.61464119f, 13.76078796f, 12.23089790f, 10.90732002f, 8.75849152f, 7.49001646f, 6.14220476f, 4.65472794f, 3.07277966f, 1.84880662f, 0.803307f, 0.02916753f }, - { 14.61464119f, 13.76078796f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.65472794f, 3.07277966f, 1.84880662f, 0.803307f, 0.02916753f }, - { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.65472794f, 3.07277966f, 1.84880662f, 0.803307f, 0.02916753f }, - { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.65472794f, 3.07277966f, 1.84880662f, 0.803307f, 0.02916753f }, - { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.60512662f, 2.63833880f, 1.56271636f, 0.72133851f, 0.02916753f }, - { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 4.65472794f, 3.46139455f, 2.45070267f, 1.56271636f, 0.72133851f, 0.02916753f }, - { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 4.65472794f, 3.46139455f, 2.45070267f, 1.56271636f, 0.72133851f, 0.02916753f }, - { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 4.65472794f, 3.46139455f, 2.45070267f, 1.56271636f, 0.72133851f, 0.02916753f }, - { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.75859547f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 4.65472794f, 3.46139455f, 2.45070267f, 1.56271636f, 0.72133851f, 0.02916753f }, - { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.75859547f, 9.24142551f, 8.75849152f, 8.30717278f, 7.88507891f, 7.49001646f, 6.77309084f, 5.85520077f, 4.65472794f, 3.46139455f, 2.45070267f, 1.56271636f, 0.72133851f, 0.02916753f } -}; - -const std::vector> GITS_NOISE_0_90 = { - { 14.61464119f, 6.77309084f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 1.56271636f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 3.07277966f, 0.95350921f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 4.86714602f, 2.54230714f, 0.89115214f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 7.49001646f, 4.86714602f, 2.54230714f, 0.89115214f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 7.49001646f, 5.09240818f, 3.07277966f, 1.61558151f, 0.69515091f, 0.02916753f }, - { 14.61464119f, 12.23089790f, 8.75849152f, 7.11996698f, 4.86714602f, 3.07277966f, 1.61558151f, 0.69515091f, 0.02916753f }, - { 14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.85520077f, 4.45427561f, 2.95596409f, 1.61558151f, 0.69515091f, 0.02916753f }, - { 14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.19988537f, 1.24153244f, 0.57119018f, 0.02916753f }, - { 14.61464119f, 12.96784878f, 10.90732002f, 8.75849152f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.19988537f, 1.24153244f, 0.57119018f, 0.02916753f }, - { 14.61464119f, 12.96784878f, 11.54541874f, 9.24142551f, 8.30717278f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.19988537f, 1.24153244f, 0.57119018f, 0.02916753f }, - { 14.61464119f, 12.96784878f, 11.54541874f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.75677586f, 2.84484982f, 1.84880662f, 1.08895338f, 0.52423614f, 0.02916753f }, - { 14.61464119f, 13.76078796f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.75677586f, 2.84484982f, 1.84880662f, 1.08895338f, 0.52423614f, 0.02916753f }, - { 14.61464119f, 13.76078796f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.44769001f, 5.58536053f, 4.45427561f, 3.32507086f, 2.45070267f, 1.61558151f, 0.95350921f, 0.45573691f, 0.02916753f }, - { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.44769001f, 5.58536053f, 4.45427561f, 3.32507086f, 2.45070267f, 1.61558151f, 0.95350921f, 0.45573691f, 0.02916753f }, - { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 4.86714602f, 3.91689563f, 3.07277966f, 2.27973175f, 1.56271636f, 0.95350921f, 0.45573691f, 0.02916753f }, - { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 4.86714602f, 3.91689563f, 3.07277966f, 2.27973175f, 1.56271636f, 0.95350921f, 0.45573691f, 0.02916753f }, - { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 4.86714602f, 3.91689563f, 3.07277966f, 2.27973175f, 1.56271636f, 0.95350921f, 0.45573691f, 0.02916753f }, - { 14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 5.09240818f, 4.45427561f, 3.60512662f, 2.95596409f, 2.19988537f, 1.51179266f, 0.89115214f, 0.43325692f, 0.02916753f } -}; - -const std::vector> GITS_NOISE_0_95 = { - { 14.61464119f, 6.77309084f, 0.02916753f }, - { 14.61464119f, 6.77309084f, 1.56271636f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 2.84484982f, 0.89115214f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 4.86714602f, 2.36326075f, 0.803307f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 4.86714602f, 2.95596409f, 1.56271636f, 0.64427125f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 7.49001646f, 4.86714602f, 2.95596409f, 1.56271636f, 0.64427125f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 7.49001646f, 4.86714602f, 3.07277966f, 1.91321158f, 1.08895338f, 0.50118381f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.07277966f, 1.91321158f, 1.08895338f, 0.50118381f, 0.02916753f }, - { 14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.85520077f, 4.45427561f, 3.07277966f, 1.91321158f, 1.08895338f, 0.50118381f, 0.02916753f }, - { 14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.19988537f, 1.41535246f, 0.803307f, 0.38853383f, 0.02916753f }, - { 14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.85520077f, 4.65472794f, 3.46139455f, 2.63833880f, 1.84880662f, 1.24153244f, 0.72133851f, 0.34370604f, 0.02916753f }, - { 14.61464119f, 12.96784878f, 10.90732002f, 8.75849152f, 7.49001646f, 5.85520077f, 4.65472794f, 3.46139455f, 2.63833880f, 1.84880662f, 1.24153244f, 0.72133851f, 0.34370604f, 0.02916753f }, - { 14.61464119f, 12.96784878f, 10.90732002f, 8.75849152f, 7.49001646f, 6.14220476f, 4.86714602f, 3.75677586f, 2.95596409f, 2.19988537f, 1.56271636f, 1.05362725f, 0.64427125f, 0.32104823f, 0.02916753f }, - { 14.61464119f, 12.96784878f, 10.90732002f, 8.75849152f, 7.49001646f, 6.44769001f, 5.58536053f, 4.65472794f, 3.60512662f, 2.95596409f, 2.19988537f, 1.56271636f, 1.05362725f, 0.64427125f, 0.32104823f, 0.02916753f }, - { 14.61464119f, 12.96784878f, 11.54541874f, 9.24142551f, 8.30717278f, 7.49001646f, 6.44769001f, 5.58536053f, 4.65472794f, 3.60512662f, 2.95596409f, 2.19988537f, 1.56271636f, 1.05362725f, 0.64427125f, 0.32104823f, 0.02916753f }, - { 14.61464119f, 12.96784878f, 11.54541874f, 9.24142551f, 8.30717278f, 7.49001646f, 6.44769001f, 5.58536053f, 4.65472794f, 3.75677586f, 3.07277966f, 2.45070267f, 1.78698075f, 1.24153244f, 0.83188516f, 0.50118381f, 0.22545385f, 0.02916753f }, - { 14.61464119f, 12.96784878f, 11.54541874f, 9.24142551f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 5.09240818f, 4.45427561f, 3.60512662f, 2.95596409f, 2.36326075f, 1.72759056f, 1.24153244f, 0.83188516f, 0.50118381f, 0.22545385f, 0.02916753f }, - { 14.61464119f, 13.76078796f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 5.09240818f, 4.45427561f, 3.60512662f, 2.95596409f, 2.36326075f, 1.72759056f, 1.24153244f, 0.83188516f, 0.50118381f, 0.22545385f, 0.02916753f }, - { 14.61464119f, 13.76078796f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 5.09240818f, 4.45427561f, 3.75677586f, 3.07277966f, 2.45070267f, 1.91321158f, 1.46270394f, 1.05362725f, 0.72133851f, 0.43325692f, 0.19894916f, 0.02916753f } -}; - -const std::vector> GITS_NOISE_1_00 = { - { 14.61464119f, 1.56271636f, 0.02916753f }, - { 14.61464119f, 6.77309084f, 0.95350921f, 0.02916753f }, - { 14.61464119f, 6.77309084f, 2.36326075f, 0.803307f, 0.02916753f }, - { 14.61464119f, 7.11996698f, 3.07277966f, 1.56271636f, 0.59516323f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 4.86714602f, 2.84484982f, 1.41535246f, 0.57119018f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 4.86714602f, 2.84484982f, 1.61558151f, 0.86115354f, 0.38853383f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 7.49001646f, 4.86714602f, 2.84484982f, 1.61558151f, 0.86115354f, 0.38853383f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 7.49001646f, 4.86714602f, 3.07277966f, 1.98035145f, 1.24153244f, 0.72133851f, 0.34370604f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.07277966f, 1.98035145f, 1.24153244f, 0.72133851f, 0.34370604f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.27973175f, 1.51179266f, 0.95350921f, 0.54755926f, 0.25053367f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.36326075f, 1.61558151f, 1.08895338f, 0.72133851f, 0.41087446f, 0.17026083f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 8.75849152f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.36326075f, 1.61558151f, 1.08895338f, 0.72133851f, 0.41087446f, 0.17026083f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 8.75849152f, 7.49001646f, 5.85520077f, 4.65472794f, 3.60512662f, 2.84484982f, 2.12350607f, 1.56271636f, 1.08895338f, 0.72133851f, 0.41087446f, 0.17026083f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 8.75849152f, 7.49001646f, 5.85520077f, 4.65472794f, 3.60512662f, 2.84484982f, 2.19988537f, 1.61558151f, 1.162866f, 0.803307f, 0.50118381f, 0.27464288f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 8.75849152f, 7.49001646f, 5.85520077f, 4.65472794f, 3.75677586f, 3.07277966f, 2.45070267f, 1.84880662f, 1.36964464f, 1.01931262f, 0.72133851f, 0.45573691f, 0.25053367f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 8.75849152f, 7.49001646f, 6.14220476f, 5.09240818f, 4.26497746f, 3.46139455f, 2.84484982f, 2.19988537f, 1.67050016f, 1.24153244f, 0.92192322f, 0.64427125f, 0.43325692f, 0.25053367f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 8.75849152f, 7.49001646f, 6.14220476f, 5.09240818f, 4.26497746f, 3.60512662f, 2.95596409f, 2.45070267f, 1.91321158f, 1.51179266f, 1.12534678f, 0.83188516f, 0.59516323f, 0.38853383f, 0.22545385f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 12.23089790f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 5.09240818f, 4.26497746f, 3.60512662f, 2.95596409f, 2.45070267f, 1.91321158f, 1.51179266f, 1.12534678f, 0.83188516f, 0.59516323f, 0.38853383f, 0.22545385f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 12.23089790f, 9.24142551f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 5.09240818f, 4.26497746f, 3.60512662f, 2.95596409f, 2.45070267f, 1.91321158f, 1.51179266f, 1.12534678f, 0.83188516f, 0.59516323f, 0.38853383f, 0.22545385f, 0.09824532f, 0.02916753f } -}; - -const std::vector> GITS_NOISE_1_05 = { - { 14.61464119f, 0.95350921f, 0.02916753f }, - { 14.61464119f, 6.77309084f, 0.89115214f, 0.02916753f }, - { 14.61464119f, 6.77309084f, 2.05039096f, 0.72133851f, 0.02916753f }, - { 14.61464119f, 6.77309084f, 2.84484982f, 1.28281462f, 0.52423614f, 0.02916753f }, - { 14.61464119f, 6.77309084f, 3.07277966f, 1.61558151f, 0.803307f, 0.34370604f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 4.86714602f, 2.84484982f, 1.56271636f, 0.803307f, 0.34370604f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 4.86714602f, 2.84484982f, 1.61558151f, 0.95350921f, 0.52423614f, 0.22545385f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 4.86714602f, 3.07277966f, 1.98035145f, 1.24153244f, 0.74807048f, 0.41087446f, 0.17026083f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 4.86714602f, 3.19567990f, 2.27973175f, 1.51179266f, 0.95350921f, 0.59516323f, 0.34370604f, 0.13792117f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 5.09240818f, 3.46139455f, 2.45070267f, 1.61558151f, 1.08895338f, 0.72133851f, 0.45573691f, 0.25053367f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 7.49001646f, 5.09240818f, 3.46139455f, 2.45070267f, 1.61558151f, 1.08895338f, 0.72133851f, 0.45573691f, 0.25053367f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.36326075f, 1.61558151f, 1.08895338f, 0.72133851f, 0.45573691f, 0.25053367f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.45070267f, 1.72759056f, 1.24153244f, 0.86115354f, 0.59516323f, 0.38853383f, 0.22545385f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.65472794f, 3.60512662f, 2.84484982f, 2.19988537f, 1.61558151f, 1.162866f, 0.83188516f, 0.59516323f, 0.38853383f, 0.22545385f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.65472794f, 3.60512662f, 2.84484982f, 2.19988537f, 1.67050016f, 1.28281462f, 0.95350921f, 0.72133851f, 0.52423614f, 0.34370604f, 0.19894916f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.65472794f, 3.60512662f, 2.95596409f, 2.36326075f, 1.84880662f, 1.41535246f, 1.08895338f, 0.83188516f, 0.61951244f, 0.45573691f, 0.32104823f, 0.19894916f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.65472794f, 3.60512662f, 2.95596409f, 2.45070267f, 1.91321158f, 1.51179266f, 1.20157266f, 0.95350921f, 0.74807048f, 0.57119018f, 0.43325692f, 0.29807833f, 0.19894916f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 8.30717278f, 7.11996698f, 5.85520077f, 4.65472794f, 3.60512662f, 2.95596409f, 2.45070267f, 1.91321158f, 1.51179266f, 1.20157266f, 0.95350921f, 0.74807048f, 0.57119018f, 0.43325692f, 0.29807833f, 0.19894916f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 8.30717278f, 7.11996698f, 5.85520077f, 4.65472794f, 3.60512662f, 2.95596409f, 2.45070267f, 1.98035145f, 1.61558151f, 1.32549286f, 1.08895338f, 0.86115354f, 0.69515091f, 0.54755926f, 0.41087446f, 0.29807833f, 0.19894916f, 0.09824532f, 0.02916753f } -}; - -const std::vector> GITS_NOISE_1_10 = { - { 14.61464119f, 0.89115214f, 0.02916753f }, - { 14.61464119f, 2.36326075f, 0.72133851f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 1.61558151f, 0.57119018f, 0.02916753f }, - { 14.61464119f, 6.77309084f, 2.45070267f, 1.08895338f, 0.45573691f, 0.02916753f }, - { 14.61464119f, 6.77309084f, 2.95596409f, 1.56271636f, 0.803307f, 0.34370604f, 0.02916753f }, - { 14.61464119f, 6.77309084f, 3.07277966f, 1.61558151f, 0.89115214f, 0.4783645f, 0.19894916f, 0.02916753f }, - { 14.61464119f, 6.77309084f, 3.07277966f, 1.84880662f, 1.08895338f, 0.64427125f, 0.34370604f, 0.13792117f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 4.86714602f, 2.84484982f, 1.61558151f, 0.95350921f, 0.54755926f, 0.27464288f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 4.86714602f, 2.95596409f, 1.91321158f, 1.24153244f, 0.803307f, 0.4783645f, 0.25053367f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 4.86714602f, 3.07277966f, 2.05039096f, 1.41535246f, 0.95350921f, 0.64427125f, 0.41087446f, 0.22545385f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 4.86714602f, 3.19567990f, 2.27973175f, 1.61558151f, 1.12534678f, 0.803307f, 0.54755926f, 0.36617002f, 0.22545385f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 4.86714602f, 3.32507086f, 2.45070267f, 1.72759056f, 1.24153244f, 0.89115214f, 0.64427125f, 0.45573691f, 0.32104823f, 0.19894916f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 5.09240818f, 3.60512662f, 2.84484982f, 2.05039096f, 1.51179266f, 1.08895338f, 0.803307f, 0.59516323f, 0.43325692f, 0.29807833f, 0.19894916f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 5.09240818f, 3.60512662f, 2.84484982f, 2.12350607f, 1.61558151f, 1.24153244f, 0.95350921f, 0.72133851f, 0.54755926f, 0.41087446f, 0.29807833f, 0.19894916f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.45070267f, 1.84880662f, 1.41535246f, 1.08895338f, 0.83188516f, 0.64427125f, 0.50118381f, 0.36617002f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.45070267f, 1.91321158f, 1.51179266f, 1.20157266f, 0.95350921f, 0.74807048f, 0.59516323f, 0.45573691f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 5.85520077f, 4.45427561f, 3.46139455f, 2.84484982f, 2.19988537f, 1.72759056f, 1.36964464f, 1.08895338f, 0.86115354f, 0.69515091f, 0.54755926f, 0.43325692f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.46139455f, 2.84484982f, 2.19988537f, 1.72759056f, 1.36964464f, 1.08895338f, 0.86115354f, 0.69515091f, 0.54755926f, 0.43325692f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.46139455f, 2.84484982f, 2.19988537f, 1.72759056f, 1.36964464f, 1.08895338f, 0.89115214f, 0.72133851f, 0.59516323f, 0.4783645f, 0.38853383f, 0.29807833f, 0.22545385f, 0.17026083f, 0.09824532f, 0.02916753f } -}; - -const std::vector> GITS_NOISE_1_15 = { - { 14.61464119f, 0.83188516f, 0.02916753f }, - { 14.61464119f, 1.84880662f, 0.59516323f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 1.56271636f, 0.52423614f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 1.91321158f, 0.83188516f, 0.34370604f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.45070267f, 1.24153244f, 0.59516323f, 0.25053367f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.84484982f, 1.51179266f, 0.803307f, 0.41087446f, 0.17026083f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.84484982f, 1.56271636f, 0.89115214f, 0.50118381f, 0.25053367f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 6.77309084f, 3.07277966f, 1.84880662f, 1.12534678f, 0.72133851f, 0.43325692f, 0.22545385f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 6.77309084f, 3.07277966f, 1.91321158f, 1.24153244f, 0.803307f, 0.52423614f, 0.34370604f, 0.19894916f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 4.86714602f, 2.95596409f, 1.91321158f, 1.24153244f, 0.803307f, 0.52423614f, 0.34370604f, 0.19894916f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 4.86714602f, 3.07277966f, 2.05039096f, 1.36964464f, 0.95350921f, 0.69515091f, 0.4783645f, 0.32104823f, 0.19894916f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 4.86714602f, 3.07277966f, 2.12350607f, 1.51179266f, 1.08895338f, 0.803307f, 0.59516323f, 0.43325692f, 0.29807833f, 0.19894916f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 4.86714602f, 3.07277966f, 2.12350607f, 1.51179266f, 1.08895338f, 0.803307f, 0.59516323f, 0.45573691f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 4.86714602f, 3.07277966f, 2.19988537f, 1.61558151f, 1.24153244f, 0.95350921f, 0.74807048f, 0.59516323f, 0.45573691f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 4.86714602f, 3.19567990f, 2.45070267f, 1.78698075f, 1.32549286f, 1.01931262f, 0.803307f, 0.64427125f, 0.50118381f, 0.38853383f, 0.29807833f, 0.22545385f, 0.17026083f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 4.86714602f, 3.19567990f, 2.45070267f, 1.78698075f, 1.32549286f, 1.01931262f, 0.803307f, 0.64427125f, 0.52423614f, 0.41087446f, 0.32104823f, 0.25053367f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 4.86714602f, 3.19567990f, 2.45070267f, 1.84880662f, 1.41535246f, 1.12534678f, 0.89115214f, 0.72133851f, 0.59516323f, 0.4783645f, 0.38853383f, 0.32104823f, 0.25053367f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 4.86714602f, 3.19567990f, 2.45070267f, 1.84880662f, 1.41535246f, 1.12534678f, 0.89115214f, 0.72133851f, 0.59516323f, 0.50118381f, 0.41087446f, 0.34370604f, 0.29807833f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f } -}; - -const std::vector> GITS_NOISE_1_20 = { - { 14.61464119f, 0.803307f, 0.02916753f }, - { 14.61464119f, 1.56271636f, 0.52423614f, 0.02916753f }, - { 14.61464119f, 2.36326075f, 0.92192322f, 0.36617002f, 0.02916753f }, - { 14.61464119f, 2.84484982f, 1.24153244f, 0.59516323f, 0.25053367f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.05039096f, 0.95350921f, 0.45573691f, 0.17026083f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.45070267f, 1.24153244f, 0.64427125f, 0.29807833f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.45070267f, 1.36964464f, 0.803307f, 0.45573691f, 0.25053367f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.84484982f, 1.61558151f, 0.95350921f, 0.59516323f, 0.36617002f, 0.19894916f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.84484982f, 1.67050016f, 1.08895338f, 0.74807048f, 0.50118381f, 0.32104823f, 0.19894916f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.95596409f, 1.84880662f, 1.24153244f, 0.83188516f, 0.59516323f, 0.41087446f, 0.27464288f, 0.17026083f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 3.07277966f, 1.98035145f, 1.36964464f, 0.95350921f, 0.69515091f, 0.50118381f, 0.36617002f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 6.77309084f, 3.46139455f, 2.36326075f, 1.56271636f, 1.08895338f, 0.803307f, 0.59516323f, 0.45573691f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 6.77309084f, 3.46139455f, 2.45070267f, 1.61558151f, 1.162866f, 0.86115354f, 0.64427125f, 0.50118381f, 0.38853383f, 0.29807833f, 0.22545385f, 0.17026083f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 4.65472794f, 3.07277966f, 2.12350607f, 1.51179266f, 1.08895338f, 0.83188516f, 0.64427125f, 0.50118381f, 0.38853383f, 0.29807833f, 0.22545385f, 0.17026083f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 4.65472794f, 3.07277966f, 2.12350607f, 1.51179266f, 1.08895338f, 0.83188516f, 0.64427125f, 0.50118381f, 0.41087446f, 0.32104823f, 0.25053367f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 4.65472794f, 3.07277966f, 2.12350607f, 1.51179266f, 1.08895338f, 0.83188516f, 0.64427125f, 0.50118381f, 0.41087446f, 0.34370604f, 0.27464288f, 0.22545385f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 4.65472794f, 3.07277966f, 2.19988537f, 1.61558151f, 1.20157266f, 0.92192322f, 0.72133851f, 0.57119018f, 0.45573691f, 0.36617002f, 0.29807833f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 4.65472794f, 3.07277966f, 2.19988537f, 1.61558151f, 1.24153244f, 0.95350921f, 0.74807048f, 0.59516323f, 0.4783645f, 0.38853383f, 0.32104823f, 0.27464288f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 7.49001646f, 4.65472794f, 3.07277966f, 2.19988537f, 1.61558151f, 1.24153244f, 0.95350921f, 0.74807048f, 0.59516323f, 0.50118381f, 0.41087446f, 0.34370604f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f } -}; - -const std::vector> GITS_NOISE_1_25 = { - { 14.61464119f, 0.72133851f, 0.02916753f }, - { 14.61464119f, 1.56271636f, 0.50118381f, 0.02916753f }, - { 14.61464119f, 2.05039096f, 0.803307f, 0.32104823f, 0.02916753f }, - { 14.61464119f, 2.36326075f, 0.95350921f, 0.43325692f, 0.17026083f, 0.02916753f }, - { 14.61464119f, 2.84484982f, 1.24153244f, 0.59516323f, 0.27464288f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 3.07277966f, 1.51179266f, 0.803307f, 0.43325692f, 0.22545385f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.36326075f, 1.24153244f, 0.72133851f, 0.41087446f, 0.22545385f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.45070267f, 1.36964464f, 0.83188516f, 0.52423614f, 0.34370604f, 0.19894916f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.84484982f, 1.61558151f, 0.98595673f, 0.64427125f, 0.43325692f, 0.27464288f, 0.17026083f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.84484982f, 1.67050016f, 1.08895338f, 0.74807048f, 0.52423614f, 0.36617002f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.84484982f, 1.72759056f, 1.162866f, 0.803307f, 0.59516323f, 0.45573691f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.95596409f, 1.84880662f, 1.24153244f, 0.86115354f, 0.64427125f, 0.4783645f, 0.36617002f, 0.27464288f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.95596409f, 1.84880662f, 1.28281462f, 0.92192322f, 0.69515091f, 0.52423614f, 0.41087446f, 0.32104823f, 0.25053367f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.95596409f, 1.91321158f, 1.32549286f, 0.95350921f, 0.72133851f, 0.54755926f, 0.43325692f, 0.34370604f, 0.27464288f, 0.22545385f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.95596409f, 1.91321158f, 1.32549286f, 0.95350921f, 0.72133851f, 0.57119018f, 0.45573691f, 0.36617002f, 0.29807833f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.95596409f, 1.91321158f, 1.32549286f, 0.95350921f, 0.74807048f, 0.59516323f, 0.4783645f, 0.38853383f, 0.32104823f, 0.27464288f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 3.07277966f, 2.05039096f, 1.41535246f, 1.05362725f, 0.803307f, 0.61951244f, 0.50118381f, 0.41087446f, 0.34370604f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 3.07277966f, 2.05039096f, 1.41535246f, 1.05362725f, 0.803307f, 0.64427125f, 0.52423614f, 0.43325692f, 0.36617002f, 0.32104823f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 3.07277966f, 2.05039096f, 1.46270394f, 1.08895338f, 0.83188516f, 0.66947293f, 0.54755926f, 0.45573691f, 0.38853383f, 0.34370604f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f } -}; - -const std::vector> GITS_NOISE_1_30 = { - { 14.61464119f, 0.72133851f, 0.02916753f }, - { 14.61464119f, 1.24153244f, 0.43325692f, 0.02916753f }, - { 14.61464119f, 1.56271636f, 0.59516323f, 0.22545385f, 0.02916753f }, - { 14.61464119f, 1.84880662f, 0.803307f, 0.36617002f, 0.13792117f, 0.02916753f }, - { 14.61464119f, 2.36326075f, 1.01931262f, 0.52423614f, 0.25053367f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.84484982f, 1.36964464f, 0.74807048f, 0.41087446f, 0.22545385f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 3.07277966f, 1.56271636f, 0.89115214f, 0.54755926f, 0.34370604f, 0.19894916f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 3.07277966f, 1.61558151f, 0.95350921f, 0.61951244f, 0.41087446f, 0.27464288f, 0.17026083f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.45070267f, 1.36964464f, 0.83188516f, 0.54755926f, 0.36617002f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.45070267f, 1.41535246f, 0.92192322f, 0.64427125f, 0.45573691f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.6383388f, 1.56271636f, 1.01931262f, 0.72133851f, 0.50118381f, 0.36617002f, 0.27464288f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.84484982f, 1.61558151f, 1.05362725f, 0.74807048f, 0.54755926f, 0.41087446f, 0.32104823f, 0.25053367f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.84484982f, 1.61558151f, 1.08895338f, 0.77538133f, 0.57119018f, 0.43325692f, 0.34370604f, 0.27464288f, 0.22545385f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.84484982f, 1.61558151f, 1.08895338f, 0.803307f, 0.59516323f, 0.45573691f, 0.36617002f, 0.29807833f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.84484982f, 1.61558151f, 1.08895338f, 0.803307f, 0.59516323f, 0.4783645f, 0.38853383f, 0.32104823f, 0.27464288f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.84484982f, 1.72759056f, 1.162866f, 0.83188516f, 0.64427125f, 0.50118381f, 0.41087446f, 0.34370604f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.84484982f, 1.72759056f, 1.162866f, 0.83188516f, 0.64427125f, 0.52423614f, 0.43325692f, 0.36617002f, 0.32104823f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.84484982f, 1.78698075f, 1.24153244f, 0.92192322f, 0.72133851f, 0.57119018f, 0.45573691f, 0.38853383f, 0.34370604f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.84484982f, 1.78698075f, 1.24153244f, 0.92192322f, 0.72133851f, 0.57119018f, 0.4783645f, 0.41087446f, 0.36617002f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f } -}; - -const std::vector> GITS_NOISE_1_35 = { - { 14.61464119f, 0.69515091f, 0.02916753f }, - { 14.61464119f, 0.95350921f, 0.34370604f, 0.02916753f }, - { 14.61464119f, 1.56271636f, 0.57119018f, 0.19894916f, 0.02916753f }, - { 14.61464119f, 1.61558151f, 0.69515091f, 0.29807833f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 1.84880662f, 0.83188516f, 0.43325692f, 0.22545385f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.45070267f, 1.162866f, 0.64427125f, 0.36617002f, 0.19894916f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.84484982f, 1.36964464f, 0.803307f, 0.50118381f, 0.32104823f, 0.19894916f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.84484982f, 1.41535246f, 0.83188516f, 0.54755926f, 0.36617002f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.84484982f, 1.56271636f, 0.95350921f, 0.64427125f, 0.45573691f, 0.32104823f, 0.22545385f, 0.17026083f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.84484982f, 1.56271636f, 0.95350921f, 0.64427125f, 0.45573691f, 0.34370604f, 0.25053367f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 3.07277966f, 1.61558151f, 1.01931262f, 0.72133851f, 0.52423614f, 0.38853383f, 0.29807833f, 0.22545385f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 3.07277966f, 1.61558151f, 1.01931262f, 0.72133851f, 0.52423614f, 0.41087446f, 0.32104823f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 3.07277966f, 1.61558151f, 1.05362725f, 0.74807048f, 0.54755926f, 0.43325692f, 0.34370604f, 0.27464288f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 3.07277966f, 1.72759056f, 1.12534678f, 0.803307f, 0.59516323f, 0.45573691f, 0.36617002f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 3.07277966f, 1.72759056f, 1.12534678f, 0.803307f, 0.59516323f, 0.4783645f, 0.38853383f, 0.32104823f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.45070267f, 1.51179266f, 1.01931262f, 0.74807048f, 0.57119018f, 0.45573691f, 0.36617002f, 0.32104823f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.6383388f, 1.61558151f, 1.08895338f, 0.803307f, 0.61951244f, 0.50118381f, 0.41087446f, 0.34370604f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.6383388f, 1.61558151f, 1.08895338f, 0.803307f, 0.64427125f, 0.52423614f, 0.43325692f, 0.36617002f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 5.85520077f, 2.6383388f, 1.61558151f, 1.08895338f, 0.803307f, 0.64427125f, 0.52423614f, 0.45573691f, 0.38853383f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f } -}; - -const std::vector> GITS_NOISE_1_40 = { - { 14.61464119f, 0.59516323f, 0.02916753f }, - { 14.61464119f, 0.95350921f, 0.34370604f, 0.02916753f }, - { 14.61464119f, 1.08895338f, 0.43325692f, 0.13792117f, 0.02916753f }, - { 14.61464119f, 1.56271636f, 0.64427125f, 0.27464288f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 1.61558151f, 0.803307f, 0.43325692f, 0.22545385f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.05039096f, 0.95350921f, 0.54755926f, 0.34370604f, 0.19894916f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.45070267f, 1.24153244f, 0.72133851f, 0.43325692f, 0.27464288f, 0.17026083f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.45070267f, 1.24153244f, 0.74807048f, 0.50118381f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.45070267f, 1.28281462f, 0.803307f, 0.52423614f, 0.36617002f, 0.27464288f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.45070267f, 1.28281462f, 0.803307f, 0.54755926f, 0.38853383f, 0.29807833f, 0.22545385f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.84484982f, 1.41535246f, 0.86115354f, 0.59516323f, 0.43325692f, 0.32104823f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.84484982f, 1.51179266f, 0.95350921f, 0.64427125f, 0.45573691f, 0.34370604f, 0.27464288f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.84484982f, 1.51179266f, 0.95350921f, 0.64427125f, 0.4783645f, 0.36617002f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.84484982f, 1.56271636f, 0.98595673f, 0.69515091f, 0.52423614f, 0.41087446f, 0.34370604f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.84484982f, 1.56271636f, 1.01931262f, 0.72133851f, 0.54755926f, 0.43325692f, 0.36617002f, 0.32104823f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.84484982f, 1.61558151f, 1.05362725f, 0.74807048f, 0.57119018f, 0.45573691f, 0.38853383f, 0.34370604f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.84484982f, 1.61558151f, 1.08895338f, 0.803307f, 0.61951244f, 0.50118381f, 0.41087446f, 0.36617002f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.84484982f, 1.61558151f, 1.08895338f, 0.803307f, 0.61951244f, 0.50118381f, 0.43325692f, 0.38853383f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.84484982f, 1.61558151f, 1.08895338f, 0.803307f, 0.64427125f, 0.52423614f, 0.45573691f, 0.41087446f, 0.36617002f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f } -}; - -const std::vector> GITS_NOISE_1_45 = { - { 14.61464119f, 0.59516323f, 0.02916753f }, - { 14.61464119f, 0.803307f, 0.25053367f, 0.02916753f }, - { 14.61464119f, 0.95350921f, 0.34370604f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 1.24153244f, 0.54755926f, 0.25053367f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 1.56271636f, 0.72133851f, 0.36617002f, 0.19894916f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 1.61558151f, 0.803307f, 0.45573691f, 0.27464288f, 0.17026083f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 1.91321158f, 0.95350921f, 0.57119018f, 0.36617002f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.19988537f, 1.08895338f, 0.64427125f, 0.41087446f, 0.27464288f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.45070267f, 1.24153244f, 0.74807048f, 0.50118381f, 0.34370604f, 0.25053367f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.45070267f, 1.24153244f, 0.74807048f, 0.50118381f, 0.36617002f, 0.27464288f, 0.22545385f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.45070267f, 1.28281462f, 0.803307f, 0.54755926f, 0.41087446f, 0.32104823f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.45070267f, 1.28281462f, 0.803307f, 0.57119018f, 0.43325692f, 0.34370604f, 0.27464288f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.45070267f, 1.28281462f, 0.83188516f, 0.59516323f, 0.45573691f, 0.36617002f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.45070267f, 1.28281462f, 0.83188516f, 0.59516323f, 0.45573691f, 0.36617002f, 0.32104823f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.84484982f, 1.51179266f, 0.95350921f, 0.69515091f, 0.52423614f, 0.41087446f, 0.34370604f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.84484982f, 1.51179266f, 0.95350921f, 0.69515091f, 0.52423614f, 0.43325692f, 0.36617002f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.84484982f, 1.56271636f, 0.98595673f, 0.72133851f, 0.54755926f, 0.45573691f, 0.38853383f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.84484982f, 1.56271636f, 1.01931262f, 0.74807048f, 0.57119018f, 0.4783645f, 0.41087446f, 0.36617002f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.84484982f, 1.56271636f, 1.01931262f, 0.74807048f, 0.59516323f, 0.50118381f, 0.43325692f, 0.38853383f, 0.36617002f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f } -}; - -const std::vector> GITS_NOISE_1_50 = { - { 14.61464119f, 0.54755926f, 0.02916753f }, - { 14.61464119f, 0.803307f, 0.25053367f, 0.02916753f }, - { 14.61464119f, 0.86115354f, 0.32104823f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 1.24153244f, 0.54755926f, 0.25053367f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 1.56271636f, 0.72133851f, 0.36617002f, 0.19894916f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 1.61558151f, 0.803307f, 0.45573691f, 0.27464288f, 0.17026083f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 1.61558151f, 0.83188516f, 0.52423614f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 1.84880662f, 0.95350921f, 0.59516323f, 0.38853383f, 0.27464288f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 1.84880662f, 0.95350921f, 0.59516323f, 0.41087446f, 0.29807833f, 0.22545385f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 1.84880662f, 0.95350921f, 0.61951244f, 0.43325692f, 0.32104823f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.19988537f, 1.12534678f, 0.72133851f, 0.50118381f, 0.36617002f, 0.27464288f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.19988537f, 1.12534678f, 0.72133851f, 0.50118381f, 0.36617002f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.36326075f, 1.24153244f, 0.803307f, 0.57119018f, 0.43325692f, 0.34370604f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.36326075f, 1.24153244f, 0.803307f, 0.57119018f, 0.43325692f, 0.34370604f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.36326075f, 1.24153244f, 0.803307f, 0.59516323f, 0.45573691f, 0.36617002f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.36326075f, 1.24153244f, 0.803307f, 0.59516323f, 0.45573691f, 0.38853383f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.45070267f, 1.32549286f, 0.86115354f, 0.64427125f, 0.50118381f, 0.41087446f, 0.36617002f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.45070267f, 1.36964464f, 0.92192322f, 0.69515091f, 0.54755926f, 0.45573691f, 0.41087446f, 0.36617002f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f }, - { 14.61464119f, 2.45070267f, 1.41535246f, 0.95350921f, 0.72133851f, 0.57119018f, 0.4783645f, 0.43325692f, 0.38853383f, 0.36617002f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f } -}; - -const std::vector>*> GITS_NOISE = { - &GITS_NOISE_0_80, - &GITS_NOISE_0_85, - &GITS_NOISE_0_90, - &GITS_NOISE_0_95, - &GITS_NOISE_1_00, - &GITS_NOISE_1_05, - &GITS_NOISE_1_10, - &GITS_NOISE_1_15, - &GITS_NOISE_1_20, - &GITS_NOISE_1_25, - &GITS_NOISE_1_30, - &GITS_NOISE_1_35, - &GITS_NOISE_1_40, - &GITS_NOISE_1_45, - &GITS_NOISE_1_50 -}; - -#endif // GITS_NOISE_INL diff --git a/src/layer_registry.cpp b/src/layer_registry.cpp deleted file mode 100644 index c4ff1881..00000000 --- a/src/layer_registry.cpp +++ /dev/null @@ -1,132 +0,0 @@ -#include "layer_registry.h" - -#include - -#include "util.h" - -namespace sd::layer_registry { - - void LayerRegistry::register_layer(const std::string& name, ggml_tensor* tensor) { - auto& info = layers_[name]; - info.tensors.push_back(tensor); - info.bytes += ggml_nbytes(tensor); - } - - bool LayerRegistry::move_layer_to_gpu(const std::string& name) { - auto it = layers_.find(name); - if (it == layers_.end()) - return false; - - LayerInfo& info = it->second; - if (info.on_gpu) - return true; - if (gpu_backend_ == nullptr || cpu_backend_ == nullptr) { - LOG_ERROR("layer_registry: backends not set; cannot move '%s' to GPU", - name.c_str()); - return false; - } - if (info.tensors.empty()) { - info.on_gpu = true; - return true; - } - - // 1. Build a no_alloc context big enough to hold one twin tensor per CPU - // tensor, plus a little overhead. - const size_t ctx_size = info.tensors.size() * ggml_tensor_overhead() + 1024; - ggml_init_params ctx_params{ctx_size, /*mem_buffer=*/nullptr, /*no_alloc=*/true}; - ggml_context* twin_ctx = ggml_init(ctx_params); - if (twin_ctx == nullptr) { - LOG_ERROR("layer_registry: failed to allocate twin context for '%s'", - name.c_str()); - return false; - } - - // 2. Create one GPU twin per CPU tensor. The twin shares the original - // name so any name-based lookup keeps working. - std::vector gpu_twins; - gpu_twins.reserve(info.tensors.size()); - for (ggml_tensor* cpu_t : info.tensors) { - ggml_tensor* twin = ggml_dup_tensor(twin_ctx, cpu_t); - if (cpu_t->name[0] != '\0') { - ggml_set_name(twin, cpu_t->name); - } - gpu_twins.push_back(twin); - } - - // 3. Back the twins with a GPU buffer in one alloc call. - ggml_backend_buffer_t gpu_buffer = ggml_backend_alloc_ctx_tensors(twin_ctx, gpu_backend_); - if (gpu_buffer == nullptr) { - LOG_ERROR("layer_registry: failed to allocate GPU buffer for '%s'", - name.c_str()); - ggml_free(twin_ctx); - return false; - } - - // 4. H2D copy + sync. - for (size_t i = 0; i < info.tensors.size(); ++i) { - ggml_backend_tensor_copy(info.tensors[i], gpu_twins[i]); - } - ggml_backend_synchronize(gpu_backend_); - - // 5. Swap buffer/data/extra so the originals now point at GPU memory. - for (size_t i = 0; i < info.tensors.size(); ++i) { - std::swap(info.tensors[i]->buffer, gpu_twins[i]->buffer); - std::swap(info.tensors[i]->data, gpu_twins[i]->data); - std::swap(info.tensors[i]->extra, gpu_twins[i]->extra); - } - - info.gpu_twins = std::move(gpu_twins); - info.twin_ctx = twin_ctx; - info.gpu_buffer = gpu_buffer; - info.on_gpu = true; - return true; - } - - bool LayerRegistry::move_layer_to_cpu(const std::string& name) { - auto it = layers_.find(name); - if (it == layers_.end()) - return false; - - LayerInfo& info = it->second; - if (!info.on_gpu) - return true; - if (info.tensors.size() != info.gpu_twins.size()) { - LOG_ERROR("layer_registry: twin/tensor count mismatch for '%s'", - name.c_str()); - return false; - } - - // 1. Swap back: originals point at CPU memory again. - for (size_t i = 0; i < info.tensors.size(); ++i) { - if (info.gpu_twins[i] == nullptr) - continue; - std::swap(info.tensors[i]->buffer, info.gpu_twins[i]->buffer); - std::swap(info.tensors[i]->data, info.gpu_twins[i]->data); - std::swap(info.tensors[i]->extra, info.gpu_twins[i]->extra); - } - - // 2. Free the GPU buffer + twin context. - if (info.gpu_buffer != nullptr) { - ggml_backend_buffer_free(info.gpu_buffer); - info.gpu_buffer = nullptr; - } - if (info.twin_ctx != nullptr) { - ggml_free(info.twin_ctx); - info.twin_ctx = nullptr; - } - info.gpu_twins.clear(); - info.on_gpu = false; - return true; - } - - bool LayerRegistry::is_layer_on_gpu(const std::string& name) const { - auto it = layers_.find(name); - return it != layers_.end() && it->second.on_gpu; - } - - size_t LayerRegistry::get_layer_size(const std::string& name) const { - auto it = layers_.find(name); - return it != layers_.end() ? it->second.bytes : 0; - } - -} // namespace sd::layer_registry diff --git a/src/layer_registry.h b/src/layer_registry.h deleted file mode 100644 index 2dc75534..00000000 --- a/src/layer_registry.h +++ /dev/null @@ -1,50 +0,0 @@ -#ifndef __LAYER_REGISTRY_H__ -#define __LAYER_REGISTRY_H__ - -#include -#include -#include -#include - -#include "ggml-backend.h" -#include "ggml.h" - -namespace sd::layer_registry { - - struct LayerInfo { - std::vector tensors; - std::vector gpu_twins; - ggml_context* twin_ctx = nullptr; - ggml_backend_buffer_t gpu_buffer = nullptr; - bool on_gpu = false; - size_t bytes = 0; - }; - - class LayerRegistry { - public: - LayerRegistry() = default; - LayerRegistry(ggml_backend_t gpu_backend, ggml_backend_t cpu_backend) - : gpu_backend_(gpu_backend), cpu_backend_(cpu_backend) {} - - void set_backends(ggml_backend_t gpu_backend, ggml_backend_t cpu_backend) { - gpu_backend_ = gpu_backend; - cpu_backend_ = cpu_backend; - } - void register_layer(const std::string& name, ggml_tensor* tensor); - bool move_layer_to_gpu(const std::string& name); - bool move_layer_to_cpu(const std::string& name); - bool is_layer_on_gpu(const std::string& name) const; - size_t get_layer_size(const std::string& name) const; - size_t get_layer_count() const { return layers_.size(); } - - const std::map& layers() const { return layers_; } - - private: - ggml_backend_t gpu_backend_ = nullptr; - ggml_backend_t cpu_backend_ = nullptr; - std::map layers_; - }; - -} // namespace sd::layer_registry - -#endif diff --git a/src/model.h b/src/model.h index 1a765a14..2de82006 100644 --- a/src/model.h +++ b/src/model.h @@ -1,17 +1,14 @@ #ifndef __MODEL_H__ #define __MODEL_H__ -#include -#include -#include -#include #include +#include #include +#include "core/ordered_map.hpp" #include "ggml-backend.h" #include "ggml.h" #include "model_io/tensor_storage.h" -#include "ordered_map.hpp" enum SDVersion { VERSION_SD1, @@ -46,11 +43,17 @@ enum SDVersion { VERSION_LTXAV, VERSION_HIDREAM_O1, VERSION_Z_IMAGE, + VERSION_BOOGU_IMAGE, VERSION_OVIS_IMAGE, VERSION_ERNIE_IMAGE, VERSION_LENS, + VERSION_MINIT2I, VERSION_LONGCAT, VERSION_PID, + VERSION_IDEOGRAM4, + VERSION_SEFI_IMAGE, + VERSION_KREA2, + VERSION_ESRGAN, VERSION_COUNT, }; @@ -145,6 +148,13 @@ static inline bool sd_version_is_z_image(SDVersion version) { return false; } +static inline bool sd_version_is_boogu_image(SDVersion version) { + if (version == VERSION_BOOGU_IMAGE) { + return true; + } + return false; +} + static inline bool sd_version_is_longcat(SDVersion version) { if (version == VERSION_LONGCAT) { return true; @@ -166,6 +176,13 @@ static inline bool sd_version_is_lens(SDVersion version) { return false; } +static inline bool sd_version_is_minit2i(SDVersion version) { + if (version == VERSION_MINIT2I) { + return true; + } + return false; +} + static inline bool sd_version_is_pid(SDVersion version) { if (version == VERSION_PID) { return true; @@ -173,8 +190,43 @@ static inline bool sd_version_is_pid(SDVersion version) { return false; } +static inline bool sd_version_is_ideogram4(SDVersion version) { + if (version == VERSION_IDEOGRAM4) { + return true; + } + return false; +} + +static inline bool sd_version_is_sefi_image(SDVersion version) { + if (version == VERSION_SEFI_IMAGE) { + return true; + } + return false; +} + +static inline bool sd_version_is_krea2(SDVersion version) { + if (version == VERSION_KREA2) { + return true; + } + return false; +} + +static inline bool sd_version_uses_flux_vae(SDVersion version) { + if (sd_version_is_flux(version) || sd_version_is_z_image(version) || sd_version_is_boogu_image(version) || sd_version_is_longcat(version)) { + return true; + } + return false; +} + static inline bool sd_version_uses_flux2_vae(SDVersion version) { - if (sd_version_is_flux2(version) || sd_version_is_ernie_image(version) || sd_version_is_lens(version)) { + if (sd_version_is_flux2(version) || sd_version_is_ernie_image(version) || sd_version_is_lens(version) || sd_version_is_ideogram4(version) || sd_version_is_sefi_image(version)) { + return true; + } + return false; +} + +static inline bool sd_version_uses_wan_vae(SDVersion version) { + if (sd_version_is_wan(version) || sd_version_is_qwen_image(version) || sd_version_is_krea2(version) || sd_version_is_anima(version)) { return true; } return false; @@ -201,10 +253,15 @@ 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_boogu_image(version) || sd_version_is_ernie_image(version) || sd_version_is_lens(version) || + sd_version_is_minit2i(version) || sd_version_is_longcat(version) || - sd_version_is_pid(version)) { + sd_version_is_pid(version) || + sd_version_is_ideogram4(version) || + sd_version_is_sefi_image(version) || + sd_version_is_krea2(version)) { return true; } return false; @@ -230,73 +287,4 @@ enum PMVersion { typedef OrderedMap String2TensorStorage; using TensorTypeRules = std::vector>; -TensorTypeRules parse_tensor_type_rules(const std::string& tensor_type_rules); - -class MmapWrapper; - -struct ModelFileData { - std::string path; - std::vector tensors; - std::shared_ptr mmapped; - std::shared_ptr mmbuffer; - bool is_zip; -}; - -struct MmapTensorStore { - std::shared_ptr mmapped; - std::shared_ptr mmbuffer; -}; - -class ModelLoader { -protected: - SDVersion version_ = VERSION_COUNT; - std::vector file_paths_; - std::vector file_data; - bool model_files_processed = false; - String2TensorStorage tensor_storage_map; - - void add_tensor_storage(const TensorStorage& tensor_storage); - - bool init_from_gguf_file(const std::string& file_path, const std::string& prefix = ""); - bool init_from_safetensors_file(const std::string& file_path, const std::string& prefix = ""); - bool init_from_torch_zip_file(const std::string& file_path, const std::string& prefix = ""); - bool init_from_torch_legacy_file(const std::string& file_path, const std::string& prefix = ""); - bool init_from_diffusers_file(const std::string& file_path, const std::string& prefix = ""); - -public: - bool init_from_file(const std::string& file_path, const std::string& prefix = ""); - void convert_tensors_name(); - bool init_from_file_and_convert_name(const std::string& file_path, - const std::string& prefix = "", - SDVersion version = VERSION_COUNT); - SDVersion get_sd_version(); - std::map get_wtype_stat(); - std::map get_conditioner_wtype_stat(); - std::map get_diffusion_model_wtype_stat(); - std::map get_vae_wtype_stat(); - String2TensorStorage& get_tensor_storage_map() { return tensor_storage_map; } - void set_wtype_override(ggml_type wtype, std::string tensor_type_rules = ""); - void process_model_files(bool enable_mmap = false, bool writable_mmap = true); - std::vector mmap_tensors(std::map& tensors, - std::set ignore_tensors = {}, - bool writable = true); - bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_threads = 0, bool use_mmap = false); - bool load_tensors(std::map& tensors, - std::set ignore_tensors = {}, - int n_threads = 0, - bool use_mmap = false); - - std::vector get_tensor_names() const { - std::vector names; - for (const auto& [name, tensor_storage] : tensor_storage_map) { - names.push_back(name); - } - return names; - } - - bool tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type); - int64_t get_params_mem_size(ggml_backend_t backend, ggml_type type = GGML_TYPE_COUNT); - ~ModelLoader() = default; -}; - #endif // __MODEL_H__ diff --git a/src/lora.hpp b/src/model/adapter/lora.hpp similarity index 89% rename from src/lora.hpp rename to src/model/adapter/lora.hpp index d87b494c..0b759175 100644 --- a/src/lora.hpp +++ b/src/model/adapter/lora.hpp @@ -1,8 +1,10 @@ -#ifndef __LORA_HPP__ -#define __LORA_HPP__ +#ifndef __SD_MODEL_ADAPTER_LORA_HPP__ +#define __SD_MODEL_ADAPTER_LORA_HPP__ #include -#include "ggml_extend.hpp" +#include "core/ggml_extend.hpp" +#include "model_loader.h" +#include "model_manager.h" #define LORA_GRAPH_BASE_SIZE 10240 @@ -13,22 +15,24 @@ struct LoraModel : public GGMLRunner { std::map original_tensor_to_final_tensor; std::set applied_lora_tensors; std::string file_path; - ModelLoader model_loader; - bool load_failed = false; - bool applied = false; - bool tensor_preprocessed = false; + std::shared_ptr model_manager; + ggml_backend_t params_backend = nullptr; + bool load_failed = false; + bool applied = false; + bool tensor_preprocessed = false; typedef std::function filter_t; LoraModel(const std::string& lora_id, ggml_backend_t backend, - ggml_backend_t params_backend, - const std::string& file_path = "", - std::string prefix = "", - SDVersion version = VERSION_COUNT) - : lora_id(lora_id), file_path(file_path), GGMLRunner(backend, params_backend) { + ggml_backend_t params_backend_, + const std::string& file_path = "", + std::string prefix = "", + SDVersion version = VERSION_COUNT, + std::shared_ptr manager = std::make_shared()) + : GGMLRunner(backend, manager), lora_id(lora_id), file_path(file_path), model_manager(std::move(manager)), params_backend(params_backend_) { prefix = "lora." + prefix; - if (!model_loader.init_from_file_and_convert_name(file_path, prefix, version)) { + if (model_manager == nullptr || !model_manager->loader().init_from_file_and_convert_name(file_path, prefix, version)) { load_failed = true; } } @@ -70,7 +74,11 @@ struct LoraModel : public GGMLRunner { return true; }; - model_loader.load_tensors(on_new_tensor_cb, n_threads); + if (model_manager != nullptr) { + model_manager->set_n_threads(n_threads); + } + ModelLoader& model_loader = model_manager->loader(); + model_loader.load_tensors(on_new_tensor_cb); if (tensors_to_create.empty()) { return true; @@ -86,25 +94,64 @@ struct LoraModel : public GGMLRunner { lora_tensors[name] = real; } - if (!alloc_params_buffer()) { - LOG_ERROR("lora model buffer allocation failed"); + std::map tensors; + for (const auto& pair : lora_tensors) { + tensors[pair.first] = pair.second; + } + if (model_manager == nullptr || + !model_manager->register_param_tensors("LoRA", + std::move(tensors), + ModelManager::ResidencyMode::ParamBackend, + runtime_backend, + params_backend) || + !model_manager->validate_registered_tensors()) { + LOG_ERROR("lora model manager registration failed"); + return false; + } + std::vector lora_params; + lora_params.reserve(lora_tensors.size()); + for (const auto& pair : lora_tensors) { + lora_params.push_back(pair.second); + } + if (!model_manager->prepare_params(lora_params)) { + LOG_ERROR("lora model manager prepare params failed"); return false; } - - dry_run = false; - model_loader.load_tensors(on_new_tensor_cb, n_threads); LOG_DEBUG("finished loaded lora"); return true; } - void preprocess_lora_tensors(const std::map& model_tensors) { + void release_loaded_tensors() { + runner_done(); + free_compute_buffer(); + model_manager.reset(); + free_params_ctx(); + alloc_params_ctx(); + model_manager = std::make_shared(); + weight_manager = model_manager; + lora_tensors.clear(); + original_tensor_to_final_tensor.clear(); + applied_lora_tensors.clear(); + applied = false; + tensor_preprocessed = false; + } + + static std::set tensor_names(const std::map& model_tensors) { + std::set names; + for (const auto& item : model_tensors) { + names.insert(item.first); + } + return names; + } + + void preprocess_lora_tensors(const std::set& model_tensor_names) { if (tensor_preprocessed) { return; } tensor_preprocessed = true; // I really hate these hardcoded processes. - if (model_tensors.find("cond_stage_model.1.transformer.text_model.encoder.layers.0.self_attn.in_proj.weight") != model_tensors.end()) { + if (model_tensor_names.find("cond_stage_model.1.transformer.text_model.encoder.layers.0.self_attn.in_proj.weight") != model_tensor_names.end()) { std::unordered_map new_lora_tensors; for (auto& [old_name, tensor] : lora_tensors) { std::string new_name = old_name; @@ -611,7 +658,7 @@ struct LoraModel : public GGMLRunner { if (lokr_w2) applied_lora_tensors.insert(lokr_w2_name); if (lokr_w2_a) - applied_lora_tensors.insert(lokr_w2_name); + applied_lora_tensors.insert(lokr_w2_a_name); if (lokr_w2_b) applied_lora_tensors.insert(lokr_w2_b_name); applied_lora_tensors.insert(alpha_name); @@ -752,11 +799,13 @@ struct LoraModel : public GGMLRunner { return out_diff; } - ggml_cgraph* build_lora_graph(const std::map& model_tensors, SDVersion version) { + ggml_cgraph* build_lora_graph(const std::map& model_tensors, + const std::set& model_tensor_names, + SDVersion version) { size_t lora_graph_size = LORA_GRAPH_BASE_SIZE + lora_tensors.size() * 10; ggml_cgraph* gf = ggml_new_graph_custom(compute_ctx, lora_graph_size, false); - preprocess_lora_tensors(model_tensors); + preprocess_lora_tensors(model_tensor_names); original_tensor_to_final_tensor.clear(); applied_lora_tensors.clear(); @@ -793,12 +842,16 @@ struct LoraModel : public GGMLRunner { return gf; } - void apply(std::map model_tensors, SDVersion version, int n_threads) { + void apply(std::map model_tensors, + const std::set& model_tensor_names, + SDVersion version, + int n_threads, + bool warn_unused = true) { auto get_graph = [&]() -> ggml_cgraph* { - return build_lora_graph(model_tensors, version); + return build_lora_graph(model_tensors, model_tensor_names, version); }; - GGMLRunner::compute(get_graph, n_threads, false, true); - stat(); + GGMLRunner::compute(get_graph, n_threads, false, false, false, true); + stat(!warn_unused); for (auto item : original_tensor_to_final_tensor) { ggml_tensor* original_tensor = item.first; ggml_tensor* final_tensor = item.second; @@ -809,6 +862,10 @@ struct LoraModel : public GGMLRunner { GGMLRunner::free_compute_buffer(); } + void apply(std::map model_tensors, SDVersion version, int n_threads, bool warn_unused = true) { + apply(model_tensors, tensor_names(model_tensors), version, n_threads, warn_unused); + } + void stat(bool at_runntime = false) { size_t total_lora_tensors_count = 0; size_t applied_lora_tensors_count = 0; @@ -914,4 +971,4 @@ public: } }; -#endif // __LORA_HPP__ +#endif // __SD_MODEL_ADAPTER_LORA_HPP__ diff --git a/src/pmid.hpp b/src/model/adapter/pmid.hpp similarity index 92% rename from src/pmid.hpp rename to src/model/adapter/pmid.hpp index 6e3d5e61..8f7d4dbd 100644 --- a/src/pmid.hpp +++ b/src/model/adapter/pmid.hpp @@ -1,10 +1,12 @@ -#ifndef __PMI_HPP__ -#define __PMI_HPP__ +#ifndef __SD_MODEL_ADAPTER_PMID_HPP__ +#define __SD_MODEL_ADAPTER_PMID_HPP__ -#include "ggml_extend.hpp" +#include "core/ggml_extend.hpp" -#include "clip.hpp" -#include "lora.hpp" +#include "model/adapter/lora.hpp" +#include "model/common/block.hpp" +#include "model/te/clip.hpp" +#include "model_loader.h" struct FuseBlock : public GGMLBlock { // network hparams @@ -411,13 +413,13 @@ public: public: PhotoMakerIDEncoder(ggml_backend_t backend, - ggml_backend_t params_backend, const String2TensorStorage& tensor_storage_map, const std::string prefix, - SDVersion version = VERSION_SDXL, - PMVersion pm_v = PM_VERSION_1, - float sty = 20.f) - : GGMLRunner(backend, params_backend), + SDVersion version = VERSION_SDXL, + PMVersion pm_v = PM_VERSION_1, + float sty = 20.f, + std::shared_ptr weight_manager = nullptr) + : GGMLRunner(backend, weight_manager), version(version), pm_version(pm_v), style_strength(sty) { @@ -556,24 +558,25 @@ public: return build_graph(id_pixel_values, prompt_embeds, class_tokens_mask, id_embeds); }; - return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true)); + return take_or_empty(GGMLRunner::compute(get_graph, n_threads, true, true, true)); } }; struct PhotoMakerIDEmbed : public GGMLRunner { std::map tensors; std::string file_path; - ModelLoader* model_loader; - bool load_failed = false; - bool applied = false; + std::shared_ptr model_manager; + ggml_backend_t params_backend = nullptr; + bool load_failed = false; + bool applied = false; PhotoMakerIDEmbed(ggml_backend_t backend, - ggml_backend_t params_backend, - ModelLoader* ml, - const std::string& file_path = "", - const std::string& prefix = "") - : file_path(file_path), GGMLRunner(backend, params_backend), model_loader(ml) { - if (!model_loader->init_from_file_and_convert_name(file_path, prefix)) { + ggml_backend_t params_backend_, + std::shared_ptr manager = std::make_shared(), + const std::string& file_path = "", + const std::string& prefix = "") + : GGMLRunner(backend, manager), file_path(file_path), model_manager(std::move(manager)), params_backend(params_backend_) { + if (model_manager == nullptr || !model_manager->loader().init_from_file_and_convert_name(file_path, prefix)) { load_failed = true; } } @@ -614,14 +617,27 @@ struct PhotoMakerIDEmbed : public GGMLRunner { return true; }; - model_loader->load_tensors(on_new_tensor_cb, n_threads); - if (!alloc_params_buffer()) { - LOG_ERROR("PhotoMaker ID embeds buffer allocation failed"); + model_manager->set_n_threads(n_threads); + ModelLoader& model_loader = model_manager->loader(); + model_loader.load_tensors(on_new_tensor_cb); + if (!model_manager->register_param_tensors("PhotoMaker ID embeds", + tensors, + ModelManager::ResidencyMode::ParamBackend, + runtime_backend, + params_backend) || + !model_manager->validate_registered_tensors()) { + LOG_ERROR("PhotoMaker ID embeds model manager registration failed"); + return false; + } + std::vector id_embed_params; + id_embed_params.reserve(tensors.size()); + for (const auto& pair : tensors) { + id_embed_params.push_back(pair.second); + } + if (!model_manager->prepare_params(id_embed_params)) { + LOG_ERROR("PhotoMaker ID embeds model manager prepare params failed"); return false; } - - dry_run = false; - model_loader->load_tensors(on_new_tensor_cb, n_threads); LOG_DEBUG("finished loading PhotoMaker ID Embeds "); return true; @@ -636,4 +652,4 @@ struct PhotoMakerIDEmbed : public GGMLRunner { } }; -#endif // __PMI_HPP__ +#endif // __SD_MODEL_ADAPTER_PMID_HPP__ diff --git a/src/model/adapter/pulid.hpp b/src/model/adapter/pulid.hpp new file mode 100644 index 00000000..442c5b8b --- /dev/null +++ b/src/model/adapter/pulid.hpp @@ -0,0 +1,76 @@ +#ifndef __PULID_HPP__ +#define __PULID_HPP__ + +#include "core/ggml_extend.hpp" +#include "model/common/block.hpp" + +class PuLIDPerceiverAttentionCA : public GGMLBlock { +public: + static constexpr int64_t DEFAULT_DIM = 3072; // Flux hidden size + static constexpr int64_t DEFAULT_DIM_HEAD = 128; + static constexpr int64_t DEFAULT_HEADS = 16; + static constexpr int64_t DEFAULT_KV_DIM = 2048; // PuLID ID-embedding dim + +protected: + int64_t dim; + int64_t dim_head; + int64_t heads; + int64_t kv_dim; + int64_t inner_dim; + +public: + PuLIDPerceiverAttentionCA(int64_t dim = DEFAULT_DIM, + int64_t dim_head = DEFAULT_DIM_HEAD, + int64_t heads = DEFAULT_HEADS, + int64_t kv_dim = DEFAULT_KV_DIM) + : dim(dim), + dim_head(dim_head), + heads(heads), + kv_dim(kv_dim), + inner_dim(dim_head * heads) { + blocks["norm1"] = std::shared_ptr(new LayerNorm(kv_dim)); + blocks["norm2"] = std::shared_ptr(new LayerNorm(dim)); + blocks["to_q"] = std::shared_ptr(new Linear(dim, inner_dim, /*bias=*/false)); + blocks["to_kv"] = std::shared_ptr(new Linear(kv_dim, inner_dim * 2, /*bias=*/false)); + blocks["to_out"] = std::shared_ptr(new Linear(inner_dim, dim, /*bias=*/false)); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* id_embedding, + ggml_tensor* image_tokens) { + auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); + auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); + auto to_q = std::dynamic_pointer_cast(blocks["to_q"]); + auto to_kv = std::dynamic_pointer_cast(blocks["to_kv"]); + auto to_out = std::dynamic_pointer_cast(blocks["to_out"]); + + ggml_tensor* x_normed = norm1->forward(ctx, id_embedding); + ggml_tensor* lat_normed = norm2->forward(ctx, image_tokens); + + ggml_tensor* q = to_q->forward(ctx, lat_normed); // [N, T_img, 2048] + ggml_tensor* kv = to_kv->forward(ctx, x_normed); // [N, T_img, 3072] + + ggml_tensor* k = ggml_view_3d(ctx->ggml_ctx, kv, + inner_dim, kv->ne[1], kv->ne[2], + kv->nb[1], kv->nb[2], + /*offset=*/0); + ggml_tensor* v = ggml_view_3d(ctx->ggml_ctx, kv, + inner_dim, kv->ne[1], kv->ne[2], + kv->nb[1], kv->nb[2], + /*offset=*/inner_dim * ggml_element_size(kv)); + k = ggml_cont(ctx->ggml_ctx, k); + v = ggml_cont(ctx->ggml_ctx, v); + + ggml_tensor* attn_out = ggml_ext_attention_ext( + ctx->ggml_ctx, ctx->backend, + q, k, v, + heads, + /*mask=*/nullptr, + /*diag_mask_inf=*/false); + + ggml_tensor* out = to_out->forward(ctx, attn_out); + return out; + } +}; + +#endif // __PULID_HPP__ diff --git a/src/common_block.hpp b/src/model/common/block.hpp similarity index 93% rename from src/common_block.hpp rename to src/model/common/block.hpp index e6c0b06b..15bfa376 100644 --- a/src/common_block.hpp +++ b/src/model/common/block.hpp @@ -1,9 +1,9 @@ -#ifndef __COMMON_BLOCK_HPP__ -#define __COMMON_BLOCK_HPP__ +#ifndef __SD_MODEL_COMMON_BLOCK_HPP__ +#define __SD_MODEL_COMMON_BLOCK_HPP__ +#include "core/ggml_extend.hpp" +#include "core/util.h" #include "ggml-backend.h" -#include "ggml_extend.hpp" -#include "util.h" class DownSampleBlock : public GGMLBlock { protected: @@ -227,6 +227,37 @@ public: } }; +struct Mlp : public GGMLBlock { +public: + Mlp(int64_t in_features, + int64_t hidden_features = -1, + int64_t out_features = -1, + bool bias = true) { + // act_layer is always lambda: nn.GELU(approximate="tanh") + // norm_layer is always None + // use_conv is always False + if (hidden_features == -1) { + hidden_features = in_features; + } + if (out_features == -1) { + out_features = in_features; + } + blocks["fc1"] = std::shared_ptr(new Linear(in_features, hidden_features, bias)); + blocks["fc2"] = std::shared_ptr(new Linear(hidden_features, out_features, bias)); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + // x: [N, n_token, in_features] + auto fc1 = std::dynamic_pointer_cast(blocks["fc1"]); + auto fc2 = std::dynamic_pointer_cast(blocks["fc2"]); + + x = fc1->forward(ctx, x); + x = ggml_ext_gelu(ctx->ggml_ctx, x, true); + x = fc2->forward(ctx, x); + return x; + } +}; + class FeedForward : public GGMLBlock { public: enum class Activation { @@ -529,11 +560,11 @@ protected: params["mix_factor"] = ggml_new_tensor_1d(ctx, wtype, 1); } - float get_alpha() { + ggml_tensor* get_alpha(GGMLRunnerContext* ctx) { // image_only_indicator is always tensor([0.]) and since mix_factor.shape is [1,] // so learned_with_images is same as learned - float alpha = ggml_ext_backend_tensor_get_f32(params["mix_factor"]); - return sigmoid(alpha); + auto mix_factor = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, params["mix_factor"]); + return ggml_sigmoid(ctx->ggml_ctx, mix_factor); } public: @@ -547,11 +578,12 @@ public: ggml_tensor* x_spatial, ggml_tensor* x_temporal) { // image_only_indicator is always tensor([0.]) - float alpha = get_alpha(); - auto x = ggml_add(ctx->ggml_ctx, - ggml_ext_scale(ctx->ggml_ctx, x_spatial, alpha), - ggml_ext_scale(ctx->ggml_ctx, x_temporal, 1.0f - alpha)); - return x; + auto alpha = get_alpha(ctx); + return ggml_add(ctx->ggml_ctx, + x_temporal, + ggml_mul(ctx->ggml_ctx, + ggml_sub(ctx->ggml_ctx, x_spatial, x_temporal), + alpha)); } }; @@ -603,4 +635,4 @@ public: } }; -#endif // __COMMON_BLOCK_HPP__ +#endif // __SD_MODEL_COMMON_BLOCK_HPP__ diff --git a/src/rope.hpp b/src/model/common/rope.hpp similarity index 95% rename from src/rope.hpp rename to src/model/common/rope.hpp index eea1fc4a..af66dd90 100644 --- a/src/rope.hpp +++ b/src/model/common/rope.hpp @@ -1,10 +1,10 @@ -#ifndef __ROPE_HPP__ -#define __ROPE_HPP__ +#ifndef __SD_MODEL_COMMON_ROPE_HPP__ +#define __SD_MODEL_COMMON_ROPE_HPP__ #include #include #include -#include "ggml_extend.hpp" +#include "core/ggml_extend.hpp" namespace Rope { enum class EmbedNDLayout { @@ -259,6 +259,45 @@ namespace Rope { return embed_nd(ids, bs, axis_thetas, axes_dim, wrap_dims, layout); } + __STATIC_INLINE__ std::vector embed_interleaved_mrope(const std::vector>& ids, + int bs, + float theta, + int head_dim, + const std::vector& mrope_section, + const std::vector>& axis_wrap_dims = {}) { + GGML_ASSERT(bs > 0); + GGML_ASSERT(head_dim % 2 == 0); + GGML_ASSERT(mrope_section.size() >= 3); + + std::vector> trans_ids = transpose(ids); + size_t pos_len = ids.size() / bs; + int half_dim = head_dim / 2; + + std::vector>> axis_embs; + axis_embs.reserve(3); + for (int axis = 0; axis < 3; ++axis) { + std::vector axis_wrap; + if (axis < static_cast(axis_wrap_dims.size())) { + axis_wrap = axis_wrap_dims[axis]; + } + axis_embs.push_back(rope(trans_ids[axis], head_dim, theta, axis_wrap)); + } + + std::vector> emb = axis_embs[0]; + for (int axis = 1; axis < 3; ++axis) { + int length = std::min(mrope_section[axis] * 3, half_dim); + for (int freq_idx = axis; freq_idx < length; freq_idx += 3) { + for (size_t pos_idx = 0; pos_idx < bs * pos_len; ++pos_idx) { + for (int k = 0; k < 4; ++k) { + emb[pos_idx][4 * freq_idx + k] = axis_embs[axis][pos_idx][4 * freq_idx + k]; + } + } + } + } + + return flatten(emb); + } + __STATIC_INLINE__ std::vector embed_2d_interleaved(int height, int width, int dim, @@ -876,12 +915,14 @@ namespace Rope { // q,k,v: [N, L, n_head, d_head] // pe: [L, d_head/2, 2, 2] // return: [N, L, n_head*d_head] + int64_t n_head = q->ne[1]; + q = apply_rope(ctx->ggml_ctx, q, pe, rope_interleaved); // [N*n_head, L, d_head] k = apply_rope(ctx->ggml_ctx, k, pe, rope_interleaved); // [N*n_head, L, d_head] - auto x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, v->ne[1], mask, true, ctx->flash_attn_enabled, kv_scale); // [N, L, n_head*d_head] + auto x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, true, ctx->flash_attn_enabled, kv_scale); // [N, L, n_head*d_head] return x; } }; // namespace Rope -#endif // __ROPE_HPP__ +#endif // __SD_MODEL_COMMON_ROPE_HPP__ diff --git a/src/anima.hpp b/src/model/diffusion/anima.hpp similarity index 91% rename from src/anima.hpp rename to src/model/diffusion/anima.hpp index d004f9e7..9bcf076b 100644 --- a/src/anima.hpp +++ b/src/model/diffusion/anima.hpp @@ -1,19 +1,61 @@ -#ifndef __ANIMA_HPP__ -#define __ANIMA_HPP__ +#ifndef __SD_MODEL_DIFFUSION_ANIMA_HPP__ +#define __SD_MODEL_DIFFUSION_ANIMA_HPP__ +#include #include #include #include #include -#include "common_block.hpp" -#include "diffusion_model.hpp" -#include "flux.hpp" -#include "rope.hpp" +#include "model/common/block.hpp" +#include "model/common/rope.hpp" +#include "model/diffusion/flux.hpp" +#include "model/diffusion/model.hpp" namespace Anima { constexpr int ANIMA_GRAPH_SIZE = 65536; + struct AnimaConfig { + int64_t in_channels = 16; + int64_t out_channels = 16; + int64_t hidden_size = 2048; + int64_t text_embed_dim = 1024; + int64_t num_heads = 16; + int64_t head_dim = 128; + int patch_size = 2; + int64_t num_layers = 28; + std::vector axes_dim = {44, 42, 42}; + int theta = 10000; + + static AnimaConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) { + AnimaConfig config; + int64_t detected_layers = 0; + std::string layer_tag = prefix.empty() ? "blocks." : prefix + ".blocks."; + for (const auto& [name, _] : tensor_storage_map) { + size_t pos = name.find(layer_tag); + if (pos == std::string::npos) { + continue; + } + size_t start = pos + layer_tag.size(); + size_t end = name.find('.', start); + if (end == std::string::npos) { + continue; + } + int64_t layer_id = atoll(name.substr(start, end - start).c_str()); + detected_layers = std::max(detected_layers, layer_id + 1); + } + if (detected_layers > 0) { + config.num_layers = detected_layers; + LOG_DEBUG("anima: num_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %" PRId64 ", head_dim = %" PRId64, + config.num_layers, + config.hidden_size, + config.num_heads, + config.head_dim); + } + return config; + } + }; + __STATIC_INLINE__ ggml_tensor* apply_gate(ggml_context* ctx, ggml_tensor* x, ggml_tensor* gate) { @@ -418,31 +460,22 @@ namespace Anima { struct AnimaNet : public GGMLBlock { public: - int64_t in_channels = 16; - int64_t out_channels = 16; - int64_t hidden_size = 2048; - int64_t text_embed_dim = 1024; - int64_t num_heads = 16; - int64_t head_dim = 128; - int patch_size = 2; - int64_t num_layers = 28; - std::vector axes_dim = {44, 42, 42}; - int theta = 10000; + AnimaConfig config; public: AnimaNet() = default; - explicit AnimaNet(int64_t num_layers) - : num_layers(num_layers) { - blocks["x_embedder"] = std::make_shared((in_channels + 1) * patch_size * patch_size, hidden_size); - blocks["t_embedder"] = std::make_shared(hidden_size, hidden_size * 3); - blocks["t_embedding_norm"] = std::make_shared(hidden_size, 1e-6f); - for (int i = 0; i < num_layers; i++) { - blocks["blocks." + std::to_string(i)] = std::make_shared(hidden_size, - text_embed_dim, - num_heads, - head_dim); + explicit AnimaNet(AnimaConfig config) + : config(config) { + blocks["x_embedder"] = std::make_shared((config.in_channels + 1) * config.patch_size * config.patch_size, config.hidden_size); + blocks["t_embedder"] = std::make_shared(config.hidden_size, config.hidden_size * 3); + blocks["t_embedding_norm"] = std::make_shared(config.hidden_size, 1e-6f); + for (int i = 0; i < config.num_layers; i++) { + blocks["blocks." + std::to_string(i)] = std::make_shared(config.hidden_size, + config.text_embed_dim, + config.num_heads, + config.head_dim); } - blocks["final_layer"] = std::make_shared(hidden_size, patch_size, out_channels); + blocks["final_layer"] = std::make_shared(config.hidden_size, config.patch_size, config.out_channels); blocks["llm_adapter"] = std::make_shared(1024, 1024, 1024, 6, 16); } @@ -469,11 +502,11 @@ namespace Anima { auto padding_mask = ggml_ext_zeros(ctx->ggml_ctx, x->ne[0], x->ne[1], 1, x->ne[3]); x = ggml_concat(ctx->ggml_ctx, x, padding_mask, 2); // [N, C + 1, H, W] - x = DiT::pad_and_patchify(ctx, x, patch_size, patch_size); // [N, h*w, (C+1)*ph*pw] + x = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size); // [N, h*w, (C+1)*ph*pw] x = x_embedder->forward(ctx, x); - auto timestep_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast(hidden_size)); + auto timestep_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast(config.hidden_size)); auto temb = t_embedder->forward(ctx, timestep_proj); auto embedded_timestep = t_embedding_norm->forward(ctx, timestep_proj); @@ -505,7 +538,7 @@ namespace Anima { sd::ggml_graph_cut::mark_graph_cut(temb, "anima.prelude", "temb"); sd::ggml_graph_cut::mark_graph_cut(encoder_hidden_states, "anima.prelude", "context"); - for (int i = 0; i < num_layers; i++) { + for (int i = 0; i < config.num_layers; i++) { auto block = std::dynamic_pointer_cast(blocks["blocks." + std::to_string(i)]); x = block->forward(ctx, x, encoder_hidden_states, embedded_timestep, temb, image_pe); sd::ggml_graph_cut::mark_graph_cut(x, "anima.blocks." + std::to_string(i), "x"); @@ -513,7 +546,7 @@ namespace Anima { x = final_layer->forward(ctx, x, embedded_timestep, temb); // [N, h*w, ph*pw*C] - x = DiT::unpatchify_and_crop(ctx->ggml_ctx, x, H, W, patch_size, patch_size, false); // [N, C, H, W] + x = DiT::unpatchify_and_crop(ctx->ggml_ctx, x, H, W, config.patch_size, config.patch_size, false); // [N, C, H, W] return x; } @@ -524,35 +557,16 @@ namespace Anima { std::vector image_pe_vec; std::vector adapter_q_pe_vec; std::vector adapter_k_pe_vec; + AnimaConfig config; AnimaNet net; AnimaRunner(ggml_backend_t backend, - ggml_backend_t params_backend, - const String2TensorStorage& tensor_storage_map = {}, - const std::string prefix = "model.diffusion_model") - : DiffusionModelRunner(backend, params_backend, prefix) { - int64_t num_layers = 0; - std::string layer_tag = prefix + ".net.blocks."; - for (const auto& kv : tensor_storage_map) { - const std::string& tensor_name = kv.first; - size_t pos = tensor_name.find(layer_tag); - if (pos == std::string::npos) { - continue; - } - size_t start = pos + layer_tag.size(); - size_t end = tensor_name.find('.', start); - if (end == std::string::npos) { - continue; - } - int64_t layer_id = atoll(tensor_name.substr(start, end - start).c_str()); - num_layers = std::max(num_layers, layer_id + 1); - } - if (num_layers <= 0) { - num_layers = 28; - } - LOG_INFO("anima net layers: %" PRId64, num_layers); - - net = AnimaNet(num_layers); + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "model.diffusion_model", + std::shared_ptr weight_manager = nullptr) + : DiffusionModelRunner(backend, prefix, weight_manager), + config(AnimaConfig::detect_from_weights(tensor_storage_map, prefix + ".net")) { + net = AnimaNet(config); net.init(params_ctx, tensor_storage_map, prefix + ".net"); } @@ -623,22 +637,22 @@ namespace Anima { GGML_ASSERT(x->ne[3] == 1); ggml_cgraph* gf = new_graph_custom(ANIMA_GRAPH_SIZE); - int64_t pad_h = (net.patch_size - x->ne[1] % net.patch_size) % net.patch_size; - int64_t pad_w = (net.patch_size - x->ne[0] % net.patch_size) % net.patch_size; + int64_t pad_h = (config.patch_size - x->ne[1] % config.patch_size) % config.patch_size; + int64_t pad_w = (config.patch_size - x->ne[0] % config.patch_size) % config.patch_size; int64_t h_pad = x->ne[1] + pad_h; int64_t w_pad = x->ne[0] + pad_w; image_pe_vec = gen_anima_image_pe_vec(1, static_cast(h_pad), static_cast(w_pad), - static_cast(net.patch_size), - net.theta, - net.axes_dim, + static_cast(config.patch_size), + config.theta, + config.axes_dim, 4.0f, 4.0f, 1.0f); - int64_t image_pos_len = static_cast(image_pe_vec.size()) / (2 * 2 * (net.head_dim / 2)); - auto image_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, net.head_dim / 2, image_pos_len); + int64_t image_pos_len = static_cast(image_pe_vec.size()) / (2 * 2 * (config.head_dim / 2)); + auto image_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, image_pos_len); set_backend_tensor_data(image_pe, image_pe_vec.data()); ggml_tensor* adapter_q_pe = nullptr; @@ -683,7 +697,7 @@ namespace Anima { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context, t5_ids, t5_weights); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); } sd::Tensor compute(int n_threads, @@ -701,4 +715,4 @@ namespace Anima { }; } // namespace Anima -#endif // __ANIMA_HPP__ +#endif // __SD_MODEL_DIFFUSION_ANIMA_HPP__ diff --git a/src/model/diffusion/boogu.hpp b/src/model/diffusion/boogu.hpp new file mode 100644 index 00000000..27e13aeb --- /dev/null +++ b/src/model/diffusion/boogu.hpp @@ -0,0 +1,835 @@ +#ifndef __SD_MODEL_DIFFUSION_BOOGU_HPP__ +#define __SD_MODEL_DIFFUSION_BOOGU_HPP__ + +#include +#include +#include +#include + +#include "core/ggml_extend.hpp" +#include "model/common/rope.hpp" +#include "model/diffusion/dit.hpp" +#include "model/diffusion/model.hpp" +#include "model/diffusion/qwen_image.hpp" +#include "model_loader.h" + +namespace Boogu { + constexpr int BOOGU_GRAPH_SIZE = 65536; + + struct BooguConfig { + int patch_size = 2; + int64_t in_channels = 16; + int64_t out_channels = 16; + int64_t hidden_size = 3360; + int64_t num_layers = 32; + int64_t num_double_stream_layers = 8; + int64_t num_refiner_layers = 2; + int64_t num_attention_heads = 28; + int64_t num_kv_heads = 7; + int64_t head_dim = 120; + int64_t multiple_of = 256; + int64_t instruction_feat_dim = 4096; + int64_t timestep_embed_dim = 1024; + int theta = 10000; + float timestep_scale = 1000.0f; + float norm_eps = 1e-5f; + std::vector axes_dim = {40, 40, 40}; + int64_t axes_dim_sum = 120; + + static int64_t count_blocks(const String2TensorStorage& tensor_storage_map, + const std::string& prefix, + const std::string& block_prefix) { + int64_t count = 0; + for (const auto& [name, _] : tensor_storage_map) { + if (!starts_with(name, prefix)) { + continue; + } + size_t pos = name.find(block_prefix); + if (pos == std::string::npos) { + continue; + } + auto items = split_string(name.substr(pos), '.'); + if (items.size() > 1) { + count = std::max(count, atoi(items[1].c_str()) + 1); + } + } + return count; + } + + static BooguConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) { + BooguConfig config; + int64_t detected_head_dim = 0; + int64_t detected_kv_dim = 0; + + for (const auto& [name, tensor_storage] : tensor_storage_map) { + if (!starts_with(name, prefix)) { + continue; + } + if (ends_with(name, "x_embedder.weight") && tensor_storage.n_dims == 2) { + int64_t patch_area = config.patch_size * config.patch_size; + config.in_channels = tensor_storage.ne[0] / patch_area; + config.hidden_size = tensor_storage.ne[1]; + } else if (ends_with(name, "time_caption_embed.caption_embedder.1.weight") && tensor_storage.n_dims == 2) { + config.instruction_feat_dim = tensor_storage.ne[0]; + config.hidden_size = tensor_storage.ne[1]; + } else if (ends_with(name, "single_stream_layers.0.attn.norm_q.weight") && tensor_storage.n_dims == 1) { + detected_head_dim = tensor_storage.ne[0]; + } else if (ends_with(name, "double_stream_layers.0.img_self_attn.norm_q.weight") && tensor_storage.n_dims == 1) { + detected_head_dim = tensor_storage.ne[0]; + } else if (ends_with(name, "single_stream_layers.0.attn.to_k.weight") && tensor_storage.n_dims == 2) { + detected_kv_dim = tensor_storage.ne[1]; + } else if (ends_with(name, "double_stream_layers.0.img_instruct_attn.processor.img_to_k.weight") && tensor_storage.n_dims == 2) { + detected_kv_dim = tensor_storage.ne[1]; + } else if (ends_with(name, "norm_out.linear_2.weight") && tensor_storage.n_dims == 2) { + int64_t patch_area = config.patch_size * config.patch_size; + config.out_channels = tensor_storage.ne[1] / patch_area; + } + } + + config.num_layers = std::max(1, count_blocks(tensor_storage_map, prefix, "single_stream_layers.")); + config.num_double_stream_layers = std::max(0, count_blocks(tensor_storage_map, prefix, "double_stream_layers.")); + int64_t noise_refiner_layers = count_blocks(tensor_storage_map, prefix, "noise_refiner."); + int64_t ref_refiner_layers = count_blocks(tensor_storage_map, prefix, "ref_image_refiner."); + int64_t context_refiner_layers = count_blocks(tensor_storage_map, prefix, "context_refiner."); + config.num_refiner_layers = std::max(1, std::max(noise_refiner_layers, std::max(ref_refiner_layers, context_refiner_layers))); + + if (detected_head_dim > 0) { + config.head_dim = detected_head_dim; + config.num_attention_heads = config.hidden_size / config.head_dim; + config.axes_dim_sum = config.head_dim; + if (detected_kv_dim > 0) { + config.num_kv_heads = detected_kv_dim / config.head_dim; + } + if (config.axes_dim_sum == 120) { + config.axes_dim = {40, 40, 40}; + } else if (config.axes_dim_sum % 3 == 0) { + int axis = static_cast(config.axes_dim_sum / 3); + config.axes_dim = {axis, axis, axis}; + } + } + config.timestep_embed_dim = std::min(config.hidden_size, 1024); + + LOG_DEBUG("boogu_image: layers=%" PRId64 ", double_stream_layers=%" PRId64 ", refiner_layers=%" PRId64 ", hidden=%" PRId64 ", heads=%" PRId64 ", kv_heads=%" PRId64 ", head_dim=%" PRId64 ", in_channels=%" PRId64 ", out_channels=%" PRId64, + config.num_layers, + config.num_double_stream_layers, + config.num_refiner_layers, + config.hidden_size, + config.num_attention_heads, + config.num_kv_heads, + config.head_dim, + config.in_channels, + config.out_channels); + return config; + } + }; + + __STATIC_INLINE__ ggml_tensor* scale_modulate(ggml_context* ctx, ggml_tensor* x, ggml_tensor* scale) { + scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]); + return ggml_add(ctx, x, ggml_mul(ctx, x, scale)); + } + + __STATIC_INLINE__ ggml_tensor* gate_residual(ggml_context* ctx, ggml_tensor* residual, ggml_tensor* x, ggml_tensor* gate) { + gate = ggml_tanh(ctx, gate); + gate = ggml_reshape_3d(ctx, gate, gate->ne[0], 1, gate->ne[1]); + x = ggml_mul(ctx, x, gate); + return ggml_add(ctx, residual, x); + } + + struct LuminaCombinedTimestepCaptionEmbedding : public GGMLBlock { + int64_t frequency_embedding_size; + float timestep_scale; + + LuminaCombinedTimestepCaptionEmbedding(int64_t hidden_size, + int64_t instruction_feat_dim, + int64_t frequency_embedding_size, + float norm_eps, + float timestep_scale) + : frequency_embedding_size(frequency_embedding_size), + timestep_scale(timestep_scale) { + blocks["timestep_embedder"] = std::make_shared(frequency_embedding_size, std::min(hidden_size, 1024)); + blocks["caption_embedder.0"] = std::make_shared(instruction_feat_dim, norm_eps); + blocks["caption_embedder.1"] = std::make_shared(instruction_feat_dim, hidden_size, true); + } + + std::pair forward(GGMLRunnerContext* ctx, ggml_tensor* timestep, ggml_tensor* text_hidden_states) { + auto timestep_embedder = std::dynamic_pointer_cast(blocks["timestep_embedder"]); + auto caption_embedder_0 = std::dynamic_pointer_cast(blocks["caption_embedder.0"]); + auto caption_embedder_1 = std::dynamic_pointer_cast(blocks["caption_embedder.1"]); + + auto timestep_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast(frequency_embedding_size), 10000, timestep_scale); + auto time_embed = timestep_embedder->forward(ctx, timestep_proj); + auto caption_embed = caption_embedder_1->forward(ctx, caption_embedder_0->forward(ctx, text_hidden_states)); + return {time_embed, caption_embed}; + } + }; + + struct LuminaRMSNormZero : public GGMLBlock { + LuminaRMSNormZero(int64_t embedding_dim, int64_t conditioning_embedding_dim, float norm_eps) { + blocks["linear"] = std::make_shared(conditioning_embedding_dim, 4 * embedding_dim, true); + blocks["norm"] = std::make_shared(embedding_dim, norm_eps); + } + + std::tuple forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* emb) { + auto linear = std::dynamic_pointer_cast(blocks["linear"]); + auto norm = std::dynamic_pointer_cast(blocks["norm"]); + + emb = linear->forward(ctx, ggml_silu(ctx->ggml_ctx, emb)); + auto mods = ggml_ext_chunk(ctx->ggml_ctx, emb, 4, 0); + + auto scale_msa = mods[0]; + auto gate_msa = mods[1]; + auto scale_mlp = mods[2]; + auto gate_mlp = mods[3]; + + x = scale_modulate(ctx->ggml_ctx, norm->forward(ctx, x), scale_msa); + return {x, gate_msa, scale_mlp, gate_mlp}; + } + }; + + struct LuminaFeedForward : public GGMLBlock { + LuminaFeedForward(int64_t dim, int64_t inner_dim, int64_t multiple_of) { + inner_dim = multiple_of * ((inner_dim + multiple_of - 1) / multiple_of); + blocks["linear_1"] = std::make_shared(dim, inner_dim, false); + blocks["linear_2"] = std::make_shared(inner_dim, dim, false); + blocks["linear_3"] = std::make_shared(dim, inner_dim, false); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto linear_1 = std::dynamic_pointer_cast(blocks["linear_1"]); + auto linear_2 = std::dynamic_pointer_cast(blocks["linear_2"]); + auto linear_3 = std::dynamic_pointer_cast(blocks["linear_3"]); + + if (sd_backend_is(ctx->backend, "Vulkan")) { + linear_2->set_force_prec_f32(true); + } + + auto h1 = linear_1->forward(ctx, x); + auto h2 = linear_3->forward(ctx, x); + x = ggml_swiglu_split(ctx->ggml_ctx, h1, h2); + x = linear_2->forward(ctx, x); + return x; + } + }; + + struct LuminaLayerNormContinuous : public GGMLBlock { + LuminaLayerNormContinuous(int64_t embedding_dim, + int64_t conditioning_embedding_dim, + int64_t out_dim) { + blocks["linear_1"] = std::make_shared(conditioning_embedding_dim, embedding_dim, true); + blocks["norm"] = std::make_shared(embedding_dim, 1e-6f, false); + blocks["linear_2"] = std::make_shared(embedding_dim, out_dim, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* conditioning_embedding) { + auto linear_1 = std::dynamic_pointer_cast(blocks["linear_1"]); + auto norm = std::dynamic_pointer_cast(blocks["norm"]); + auto linear_2 = std::dynamic_pointer_cast(blocks["linear_2"]); + + auto emb = linear_1->forward(ctx, ggml_silu(ctx->ggml_ctx, conditioning_embedding)); + x = scale_modulate(ctx->ggml_ctx, norm->forward(ctx, x), emb); + x = linear_2->forward(ctx, x); + return x; + } + }; + + struct Attention : public GGMLBlock { + int64_t dim_head; + int64_t heads; + int64_t kv_heads; + + Attention(int64_t query_dim, int64_t dim_head, int64_t heads, int64_t kv_heads, float eps = 1e-5f) + : dim_head(dim_head), heads(heads), kv_heads(kv_heads) { + blocks["to_q"] = std::make_shared(query_dim, heads * dim_head, false); + blocks["to_k"] = std::make_shared(query_dim, kv_heads * dim_head, false); + blocks["to_v"] = std::make_shared(query_dim, kv_heads * dim_head, false); + blocks["norm_q"] = std::make_shared(dim_head, eps); + blocks["norm_k"] = std::make_shared(dim_head, eps); + blocks["to_out.0"] = std::make_shared(heads * dim_head, query_dim, false); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* hidden_states, + ggml_tensor* encoder_hidden_states, + ggml_tensor* rotary_emb, + ggml_tensor* attention_mask = nullptr) { + auto to_q = std::dynamic_pointer_cast(blocks["to_q"]); + auto to_k = std::dynamic_pointer_cast(blocks["to_k"]); + auto to_v = std::dynamic_pointer_cast(blocks["to_v"]); + auto norm_q = std::dynamic_pointer_cast(blocks["norm_q"]); + auto norm_k = std::dynamic_pointer_cast(blocks["norm_k"]); + auto to_out_0 = std::dynamic_pointer_cast(blocks["to_out.0"]); + + if (sd_backend_is(ctx->backend, "Vulkan")) { + to_out_0->set_force_prec_f32(true); + } + + int64_t N = hidden_states->ne[2]; + int64_t Lq = hidden_states->ne[1]; + int64_t Lk = encoder_hidden_states->ne[1]; + + auto q = to_q->forward(ctx, hidden_states); + q = ggml_reshape_4d(ctx->ggml_ctx, q, dim_head, heads, Lq, N); + auto k = to_k->forward(ctx, encoder_hidden_states); + k = ggml_reshape_4d(ctx->ggml_ctx, k, dim_head, kv_heads, Lk, N); + auto v = to_v->forward(ctx, encoder_hidden_states); + v = ggml_reshape_4d(ctx->ggml_ctx, v, dim_head, kv_heads, Lk, N); + + q = norm_q->forward(ctx, q); + k = norm_k->forward(ctx, k); + + auto out = Rope::attention(ctx, q, k, v, rotary_emb, attention_mask); + out = to_out_0->forward(ctx, out); + return out; + } + }; + + struct BooguImageTransformerBlock : public GGMLBlock { + bool modulation; + + BooguImageTransformerBlock(int64_t dim, + int64_t num_attention_heads, + int64_t num_kv_heads, + int64_t multiple_of, + float norm_eps, + bool modulation) + : modulation(modulation) { + int64_t head_dim = dim / num_attention_heads; + blocks["attn"] = std::make_shared(dim, head_dim, num_attention_heads, num_kv_heads, 1e-5f); + blocks["feed_forward"] = std::make_shared(dim, 4 * dim, multiple_of); + if (modulation) { + blocks["norm1"] = std::make_shared(dim, std::min(dim, 1024), norm_eps); + } else { + blocks["norm1"] = std::make_shared(dim, norm_eps); + } + blocks["ffn_norm1"] = std::make_shared(dim, norm_eps); + blocks["norm2"] = std::make_shared(dim, norm_eps); + blocks["ffn_norm2"] = std::make_shared(dim, norm_eps); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* hidden_states, + ggml_tensor* rotary_emb, + ggml_tensor* temb = nullptr, + ggml_tensor* attention_mask = nullptr) { + auto attn = std::dynamic_pointer_cast(blocks["attn"]); + auto feed_forward = std::dynamic_pointer_cast(blocks["feed_forward"]); + auto ffn_norm1 = std::dynamic_pointer_cast(blocks["ffn_norm1"]); + auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); + auto ffn_norm2 = std::dynamic_pointer_cast(blocks["ffn_norm2"]); + + if (modulation) { + auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); + auto mods = norm1->forward(ctx, hidden_states, temb); + + auto norm_hidden_states = std::get<0>(mods); + auto gate_msa = std::get<1>(mods); + auto scale_mlp = std::get<2>(mods); + auto gate_mlp = std::get<3>(mods); + + auto attn_output = attn->forward(ctx, norm_hidden_states, norm_hidden_states, rotary_emb, attention_mask); + hidden_states = gate_residual(ctx->ggml_ctx, hidden_states, norm2->forward(ctx, attn_output), gate_msa); + + auto mlp_input = scale_modulate(ctx->ggml_ctx, ffn_norm1->forward(ctx, hidden_states), scale_mlp); + auto mlp_output = feed_forward->forward(ctx, mlp_input); + hidden_states = gate_residual(ctx->ggml_ctx, hidden_states, ffn_norm2->forward(ctx, mlp_output), gate_mlp); + } else { + auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); + + auto norm_hidden_states = norm1->forward(ctx, hidden_states); + auto attn_output = attn->forward(ctx, norm_hidden_states, norm_hidden_states, rotary_emb, attention_mask); + hidden_states = ggml_add(ctx->ggml_ctx, hidden_states, norm2->forward(ctx, attn_output)); + + auto mlp_output = feed_forward->forward(ctx, ffn_norm1->forward(ctx, hidden_states)); + hidden_states = ggml_add(ctx->ggml_ctx, hidden_states, ffn_norm2->forward(ctx, mlp_output)); + } + return hidden_states; + } + }; + + struct BooguImageJointAttention : public GGMLBlock { + int64_t dim_head; + int64_t heads; + int64_t kv_heads; + + BooguImageJointAttention(int64_t dim, int64_t dim_head, int64_t heads, int64_t kv_heads) + : dim_head(dim_head), heads(heads), kv_heads(kv_heads) { + blocks["norm_q"] = std::make_shared(dim_head, 1e-5f); + blocks["norm_k"] = std::make_shared(dim_head, 1e-5f); + blocks["to_out.0"] = std::make_shared(heads * dim_head, dim, false); + blocks["processor.img_to_q"] = std::make_shared(dim, heads * dim_head, false); + blocks["processor.img_to_k"] = std::make_shared(dim, kv_heads * dim_head, false); + blocks["processor.img_to_v"] = std::make_shared(dim, kv_heads * dim_head, false); + blocks["processor.instruct_to_q"] = std::make_shared(dim, heads * dim_head, false); + blocks["processor.instruct_to_k"] = std::make_shared(dim, kv_heads * dim_head, false); + blocks["processor.instruct_to_v"] = std::make_shared(dim, kv_heads * dim_head, false); + blocks["processor.instruct_out"] = std::make_shared(heads * dim_head, dim, false); + blocks["processor.img_out"] = std::make_shared(heads * dim_head, dim, false); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* img_hidden_states, + ggml_tensor* instruct_hidden_states, + ggml_tensor* rotary_emb, + ggml_tensor* attention_mask = nullptr) { + auto norm_q = std::dynamic_pointer_cast(blocks["norm_q"]); + auto norm_k = std::dynamic_pointer_cast(blocks["norm_k"]); + auto to_out_0 = std::dynamic_pointer_cast(blocks["to_out.0"]); + auto img_to_q = std::dynamic_pointer_cast(blocks["processor.img_to_q"]); + auto img_to_k = std::dynamic_pointer_cast(blocks["processor.img_to_k"]); + auto img_to_v = std::dynamic_pointer_cast(blocks["processor.img_to_v"]); + auto instruct_to_q = std::dynamic_pointer_cast(blocks["processor.instruct_to_q"]); + auto instruct_to_k = std::dynamic_pointer_cast(blocks["processor.instruct_to_k"]); + auto instruct_to_v = std::dynamic_pointer_cast(blocks["processor.instruct_to_v"]); + auto instruct_out = std::dynamic_pointer_cast(blocks["processor.instruct_out"]); + auto img_out = std::dynamic_pointer_cast(blocks["processor.img_out"]); + + if (sd_backend_is(ctx->backend, "Vulkan")) { + to_out_0->set_force_prec_f32(true); + } + + int64_t N = img_hidden_states->ne[2]; + int64_t L_img = img_hidden_states->ne[1]; + int64_t L_instruct = instruct_hidden_states->ne[1]; + + auto img_q = img_to_q->forward(ctx, img_hidden_states); + img_q = ggml_reshape_4d(ctx->ggml_ctx, img_q, dim_head, heads, L_img, N); + auto img_k = img_to_k->forward(ctx, img_hidden_states); + img_k = ggml_reshape_4d(ctx->ggml_ctx, img_k, dim_head, kv_heads, L_img, N); + auto img_v = img_to_v->forward(ctx, img_hidden_states); + img_v = ggml_reshape_4d(ctx->ggml_ctx, img_v, dim_head, kv_heads, L_img, N); + + auto instruct_q = instruct_to_q->forward(ctx, instruct_hidden_states); + instruct_q = ggml_reshape_4d(ctx->ggml_ctx, instruct_q, dim_head, heads, L_instruct, N); + auto instruct_k = instruct_to_k->forward(ctx, instruct_hidden_states); + instruct_k = ggml_reshape_4d(ctx->ggml_ctx, instruct_k, dim_head, kv_heads, L_instruct, N); + auto instruct_v = instruct_to_v->forward(ctx, instruct_hidden_states); + instruct_v = ggml_reshape_4d(ctx->ggml_ctx, instruct_v, dim_head, kv_heads, L_instruct, N); + + auto q = ggml_concat(ctx->ggml_ctx, instruct_q, img_q, 2); + auto k = ggml_concat(ctx->ggml_ctx, instruct_k, img_k, 2); + auto v = ggml_concat(ctx->ggml_ctx, instruct_v, img_v, 2); + q = norm_q->forward(ctx, q); + k = norm_k->forward(ctx, k); + + auto hidden_states = Rope::attention(ctx, q, k, v, rotary_emb, attention_mask); + auto instruct_attn = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, 0, L_instruct); + auto img_attn = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, L_instruct, L_instruct + L_img); + + instruct_attn = instruct_out->forward(ctx, instruct_attn); + img_attn = img_out->forward(ctx, img_attn); + hidden_states = ggml_concat(ctx->ggml_ctx, instruct_attn, img_attn, 1); + hidden_states = to_out_0->forward(ctx, hidden_states); + return hidden_states; + } + }; + + struct BooguImageDoubleStreamBlock : public GGMLBlock { + BooguImageDoubleStreamBlock(int64_t dim, + int64_t num_attention_heads, + int64_t num_kv_heads, + int64_t multiple_of, + float norm_eps) { + int64_t head_dim = dim / num_attention_heads; + blocks["img_instruct_attn"] = std::make_shared(dim, head_dim, num_attention_heads, num_kv_heads); + blocks["img_self_attn"] = std::make_shared(dim, head_dim, num_attention_heads, num_kv_heads, 1e-5f); + blocks["img_feed_forward"] = std::make_shared(dim, 4 * dim, multiple_of); + blocks["instruct_feed_forward"] = std::make_shared(dim, 4 * dim, multiple_of); + blocks["img_norm1"] = std::make_shared(dim, std::min(dim, 1024), norm_eps); + blocks["img_norm2"] = std::make_shared(dim, std::min(dim, 1024), norm_eps); + blocks["img_norm3"] = std::make_shared(dim, std::min(dim, 1024), norm_eps); + blocks["instruct_norm1"] = std::make_shared(dim, std::min(dim, 1024), norm_eps); + blocks["instruct_norm2"] = std::make_shared(dim, std::min(dim, 1024), norm_eps); + blocks["img_attn_norm"] = std::make_shared(dim, norm_eps); + blocks["img_self_attn_norm"] = std::make_shared(dim, norm_eps); + blocks["img_ffn_norm1"] = std::make_shared(dim, norm_eps); + blocks["img_ffn_norm2"] = std::make_shared(dim, norm_eps); + blocks["instruct_attn_norm"] = std::make_shared(dim, norm_eps); + blocks["instruct_ffn_norm1"] = std::make_shared(dim, norm_eps); + blocks["instruct_ffn_norm2"] = std::make_shared(dim, norm_eps); + } + + std::pair forward(GGMLRunnerContext* ctx, + ggml_tensor* img_hidden_states, + ggml_tensor* instruct_hidden_states, + ggml_tensor* joint_rotary_emb, + ggml_tensor* img_rotary_emb, + ggml_tensor* temb) { + auto img_instruct_attn = std::dynamic_pointer_cast(blocks["img_instruct_attn"]); + auto img_self_attn = std::dynamic_pointer_cast(blocks["img_self_attn"]); + auto img_feed_forward = std::dynamic_pointer_cast(blocks["img_feed_forward"]); + auto instruct_feed_forward = std::dynamic_pointer_cast(blocks["instruct_feed_forward"]); + auto img_norm1 = std::dynamic_pointer_cast(blocks["img_norm1"]); + auto img_norm2 = std::dynamic_pointer_cast(blocks["img_norm2"]); + auto img_norm3 = std::dynamic_pointer_cast(blocks["img_norm3"]); + auto instruct_norm1 = std::dynamic_pointer_cast(blocks["instruct_norm1"]); + auto instruct_norm2 = std::dynamic_pointer_cast(blocks["instruct_norm2"]); + auto img_attn_norm = std::dynamic_pointer_cast(blocks["img_attn_norm"]); + auto img_self_attn_norm = std::dynamic_pointer_cast(blocks["img_self_attn_norm"]); + auto img_ffn_norm1 = std::dynamic_pointer_cast(blocks["img_ffn_norm1"]); + auto img_ffn_norm2 = std::dynamic_pointer_cast(blocks["img_ffn_norm2"]); + auto instruct_attn_norm = std::dynamic_pointer_cast(blocks["instruct_attn_norm"]); + auto instruct_ffn_norm1 = std::dynamic_pointer_cast(blocks["instruct_ffn_norm1"]); + auto instruct_ffn_norm2 = std::dynamic_pointer_cast(blocks["instruct_ffn_norm2"]); + + int64_t L_instruct = instruct_hidden_states->ne[1]; + + auto img_norm1_out_vec = img_norm1->forward(ctx, img_hidden_states, temb); + auto img_norm2_out_vec = img_norm2->forward(ctx, img_hidden_states, temb); + auto img_norm3_out_vec = img_norm3->forward(ctx, img_hidden_states, temb); + auto instruct_norm1_out_vec = instruct_norm1->forward(ctx, instruct_hidden_states, temb); + auto instruct_norm2_out_vec = instruct_norm2->forward(ctx, instruct_hidden_states, temb); + + auto img_norm1_out = std::get<0>(img_norm1_out_vec); + auto img_gate_msa = std::get<1>(img_norm1_out_vec); + auto img_scale_mlp = std::get<2>(img_norm1_out_vec); + auto img_gate_mlp = std::get<3>(img_norm1_out_vec); + + auto img_norm2_out = std::get<0>(img_norm2_out_vec); + auto img_shift_mlp = std::get<1>(img_norm2_out_vec); + + auto img_norm3_out = std::get<0>(img_norm3_out_vec); + auto img_gate_self = std::get<1>(img_norm3_out_vec); + + auto instruct_norm1_out = std::get<0>(instruct_norm1_out_vec); + auto instruct_gate_msa = std::get<1>(instruct_norm1_out_vec); + auto instruct_scale_mlp = std::get<2>(instruct_norm1_out_vec); + auto instruct_gate_mlp = std::get<3>(instruct_norm1_out_vec); + + auto instruct_norm2_out = std::get<0>(instruct_norm2_out_vec); + auto instruct_shift_mlp = std::get<1>(instruct_norm2_out_vec); + + auto joint_attn_out = img_instruct_attn->forward(ctx, img_norm1_out, instruct_norm1_out, joint_rotary_emb); + auto instruct_attn_out = ggml_ext_slice(ctx->ggml_ctx, joint_attn_out, 1, 0, L_instruct); + auto img_attn_out = ggml_ext_slice(ctx->ggml_ctx, joint_attn_out, 1, L_instruct, joint_attn_out->ne[1]); + + auto img_self_attn_out = img_self_attn->forward(ctx, img_norm3_out, img_norm3_out, img_rotary_emb); + + img_hidden_states = gate_residual(ctx->ggml_ctx, img_hidden_states, img_attn_norm->forward(ctx, img_attn_out), img_gate_msa); + img_hidden_states = gate_residual(ctx->ggml_ctx, img_hidden_states, img_self_attn_norm->forward(ctx, img_self_attn_out), img_gate_self); + + auto img_mlp_input = scale_modulate(ctx->ggml_ctx, img_norm2_out, img_scale_mlp); + img_shift_mlp = ggml_reshape_3d(ctx->ggml_ctx, img_shift_mlp, img_shift_mlp->ne[0], 1, img_shift_mlp->ne[1]); + img_mlp_input = ggml_add(ctx->ggml_ctx, img_mlp_input, img_shift_mlp); + auto img_mlp_out = img_feed_forward->forward(ctx, img_ffn_norm1->forward(ctx, img_mlp_input)); + img_hidden_states = gate_residual(ctx->ggml_ctx, img_hidden_states, img_ffn_norm2->forward(ctx, img_mlp_out), img_gate_mlp); + + instruct_hidden_states = gate_residual(ctx->ggml_ctx, instruct_hidden_states, instruct_attn_norm->forward(ctx, instruct_attn_out), instruct_gate_msa); + auto instruct_mlp_input = scale_modulate(ctx->ggml_ctx, instruct_norm2_out, instruct_scale_mlp); + instruct_shift_mlp = ggml_reshape_3d(ctx->ggml_ctx, instruct_shift_mlp, instruct_shift_mlp->ne[0], 1, instruct_shift_mlp->ne[1]); + instruct_mlp_input = ggml_add(ctx->ggml_ctx, instruct_mlp_input, instruct_shift_mlp); + auto instruct_mlp_out = instruct_feed_forward->forward(ctx, instruct_ffn_norm1->forward(ctx, instruct_mlp_input)); + instruct_hidden_states = gate_residual(ctx->ggml_ctx, instruct_hidden_states, instruct_ffn_norm2->forward(ctx, instruct_mlp_out), instruct_gate_mlp); + + return {img_hidden_states, instruct_hidden_states}; + } + }; + + struct BooguImageModel : public GGMLBlock { + BooguConfig config; + + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { + GGML_UNUSED(tensor_storage_map); + GGML_UNUSED(prefix); + params["image_index_embedding"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config.hidden_size, 5); + } + + BooguImageModel() = default; + BooguImageModel(BooguConfig config) + : config(std::move(config)) { + blocks["x_embedder"] = std::make_shared(this->config.patch_size * this->config.patch_size * this->config.in_channels, this->config.hidden_size, true); + blocks["ref_image_patch_embedder"] = std::make_shared(this->config.patch_size * this->config.patch_size * this->config.in_channels, this->config.hidden_size, true); + blocks["time_caption_embed"] = std::make_shared(this->config.hidden_size, + this->config.instruction_feat_dim, + 256, + this->config.norm_eps, + this->config.timestep_scale); + + for (int i = 0; i < this->config.num_refiner_layers; i++) { + blocks["noise_refiner." + std::to_string(i)] = std::make_shared(this->config.hidden_size, + this->config.num_attention_heads, + this->config.num_kv_heads, + this->config.multiple_of, + this->config.norm_eps, + true); + blocks["ref_image_refiner." + std::to_string(i)] = std::make_shared(this->config.hidden_size, + this->config.num_attention_heads, + this->config.num_kv_heads, + this->config.multiple_of, + this->config.norm_eps, + true); + blocks["context_refiner." + std::to_string(i)] = std::make_shared(this->config.hidden_size, + this->config.num_attention_heads, + this->config.num_kv_heads, + this->config.multiple_of, + this->config.norm_eps, + false); + } + + for (int i = 0; i < this->config.num_double_stream_layers; i++) { + blocks["double_stream_layers." + std::to_string(i)] = std::make_shared(this->config.hidden_size, + this->config.num_attention_heads, + this->config.num_kv_heads, + this->config.multiple_of, + this->config.norm_eps); + } + + for (int i = 0; i < this->config.num_layers; i++) { + blocks["single_stream_layers." + std::to_string(i)] = std::make_shared(this->config.hidden_size, + this->config.num_attention_heads, + this->config.num_kv_heads, + this->config.multiple_of, + this->config.norm_eps, + true); + } + + blocks["norm_out"] = std::make_shared(this->config.hidden_size, + this->config.timestep_embed_dim, + this->config.patch_size * this->config.patch_size * this->config.out_channels); + } + + ggml_tensor* image_index_embedding(GGMLRunnerContext* ctx, int index) { + GGML_ASSERT(index >= 0 && index < 5); + auto embedding = params["image_index_embedding"]; + auto out = ggml_view_1d(ctx->ggml_ctx, + embedding, + config.hidden_size, + index * config.hidden_size * ggml_element_size(embedding)); + out = ggml_reshape_3d(ctx->ggml_ctx, out, config.hidden_size, 1, 1); + return out; + } + + ggml_tensor* embed_refs(GGMLRunnerContext* ctx, const std::vector& ref_latents) { + if (ref_latents.empty()) { + return nullptr; + } + auto ref_image_patch_embedder = std::dynamic_pointer_cast(blocks["ref_image_patch_embedder"]); + + ggml_tensor* ref_img = nullptr; + for (int i = 0; i < static_cast(ref_latents.size()); i++) { + auto ref = DiT::pad_and_patchify(ctx, ref_latents[i], config.patch_size, config.patch_size, false); + ref = ref_image_patch_embedder->forward(ctx, ref); + ref = ggml_add(ctx->ggml_ctx, ref, image_index_embedding(ctx, std::min(i, 4))); + ref_img = ref_img == nullptr ? ref : ggml_concat(ctx->ggml_ctx, ref_img, ref, 1); + } + return ref_img; + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* timesteps, + ggml_tensor* context, + ggml_tensor* pe, + std::vector ref_latents = {}) { + int64_t W = x->ne[0]; + int64_t H = x->ne[1]; + int64_t N = x->ne[3]; + GGML_ASSERT(N == 1); + + auto x_embedder = std::dynamic_pointer_cast(blocks["x_embedder"]); + auto time_caption_embed = std::dynamic_pointer_cast(blocks["time_caption_embed"]); + auto norm_out = std::dynamic_pointer_cast(blocks["norm_out"]); + + auto timestep = ggml_sub(ctx->ggml_ctx, ggml_ext_ones_like(ctx->ggml_ctx, timesteps), timesteps); + auto embeds = time_caption_embed->forward(ctx, timestep, context); + auto temb = embeds.first; + auto txt = embeds.second; + + auto img = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size, false); + int64_t img_len = img->ne[1]; + img = x_embedder->forward(ctx, img); + auto ref_img = embed_refs(ctx, ref_latents); + int64_t ref_len = ref_img != nullptr ? ref_img->ne[1] : 0; + int64_t txt_len = txt->ne[1]; + + GGML_ASSERT(pe->ne[3] == txt_len + ref_len + img_len); + auto txt_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, 0, txt_len); + auto noise_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt_len + ref_len, txt_len + ref_len + img_len); + + for (int i = 0; i < config.num_refiner_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["context_refiner." + std::to_string(i)]); + txt = block->forward(ctx, txt, txt_pe); + sd::ggml_graph_cut::mark_graph_cut(txt, "boogu.context_refiner." + std::to_string(i), "txt"); + } + + for (int i = 0; i < config.num_refiner_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["noise_refiner." + std::to_string(i)]); + img = block->forward(ctx, img, noise_pe, temb); + sd::ggml_graph_cut::mark_graph_cut(img, "boogu.noise_refiner." + std::to_string(i), "img"); + } + + ggml_tensor* combined_img = img; + if (ref_img != nullptr) { + auto ref_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt_len, txt_len + ref_len); + for (int i = 0; i < config.num_refiner_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["ref_image_refiner." + std::to_string(i)]); + ref_img = block->forward(ctx, ref_img, ref_pe, temb); + sd::ggml_graph_cut::mark_graph_cut(ref_img, "boogu.ref_image_refiner." + std::to_string(i), "ref_img"); + } + combined_img = ggml_concat(ctx->ggml_ctx, ref_img, img, 1); + } + + auto img_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt_len, txt_len + combined_img->ne[1]); + for (int i = 0; i < config.num_double_stream_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["double_stream_layers." + std::to_string(i)]); + auto result = block->forward(ctx, combined_img, txt, pe, img_pe, temb); + combined_img = result.first; + txt = result.second; + sd::ggml_graph_cut::mark_graph_cut(combined_img, "boogu.double_stream_layers." + std::to_string(i), "img"); + sd::ggml_graph_cut::mark_graph_cut(txt, "boogu.double_stream_layers." + std::to_string(i), "txt"); + } + + auto hidden_states = ggml_concat(ctx->ggml_ctx, txt, combined_img, 1); + for (int i = 0; i < config.num_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["single_stream_layers." + std::to_string(i)]); + hidden_states = block->forward(ctx, hidden_states, pe, temb); + sd::ggml_graph_cut::mark_graph_cut(hidden_states, "boogu.single_stream_layers." + std::to_string(i), "hidden_states"); + } + + hidden_states = norm_out->forward(ctx, hidden_states, temb); + hidden_states = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, hidden_states->ne[1] - img_len, hidden_states->ne[1]); + hidden_states = DiT::unpatchify_and_crop(ctx->ggml_ctx, hidden_states, H, W, config.patch_size, config.patch_size, false); + hidden_states = ggml_ext_scale(ctx->ggml_ctx, hidden_states, -1.f); + return hidden_states; + } + }; + + __STATIC_INLINE__ int patched_token_count(int64_t size, int patch_size) { + int pad = (patch_size - (static_cast(size) % patch_size)) % patch_size; + return (static_cast(size) + pad) / patch_size; + } + + __STATIC_INLINE__ void append_spatial_ids(std::vector>& ids, + int bs, + int pe_shift, + int h_tokens, + int w_tokens) { + std::vector> image_ids(h_tokens * w_tokens, std::vector(3, 0.0f)); + for (int h = 0; h < h_tokens; h++) { + for (int w = 0; w < w_tokens; w++) { + image_ids[h * w_tokens + w][0] = static_cast(pe_shift); + image_ids[h * w_tokens + w][1] = static_cast(h); + image_ids[h * w_tokens + w][2] = static_cast(w); + } + } + for (int b = 0; b < bs; b++) { + ids.insert(ids.end(), image_ids.begin(), image_ids.end()); + } + } + + __STATIC_INLINE__ std::vector gen_boogu_pe(int h, + int w, + int patch_size, + int bs, + int context_len, + const std::vector& ref_latents, + int theta, + const std::vector& axes_dim) { + std::vector> ids; + ids.reserve(static_cast(bs) * context_len); + for (int b = 0; b < bs; b++) { + for (int i = 0; i < context_len; i++) { + float pos = static_cast(i); + ids.push_back({pos, pos, pos}); + } + } + + int pe_shift = context_len; + for (ggml_tensor* ref : ref_latents) { + int ref_h_tokens = patched_token_count(ref->ne[1], patch_size); + int ref_w_tokens = patched_token_count(ref->ne[0], patch_size); + append_spatial_ids(ids, bs, pe_shift, ref_h_tokens, ref_w_tokens); + pe_shift += std::max(ref_h_tokens, ref_w_tokens); + } + + int h_tokens = patched_token_count(h, patch_size); + int w_tokens = patched_token_count(w, patch_size); + append_spatial_ids(ids, bs, pe_shift, h_tokens, w_tokens); + + return Rope::embed_nd(ids, bs, static_cast(theta), axes_dim); + } + + struct BooguImageRunner : public DiffusionModelRunner { + BooguConfig config; + BooguImageModel boogu; + std::vector pe_vec; + + BooguImageRunner(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "", + SDVersion version = VERSION_BOOGU_IMAGE, + std::shared_ptr weight_manager = nullptr) + : DiffusionModelRunner(backend, prefix, weight_manager), + config(BooguConfig::detect_from_weights(tensor_storage_map, prefix)) { + boogu = BooguImageModel(config); + boogu.init(params_ctx, tensor_storage_map, prefix); + } + + std::string get_desc() override { + return "boogu_image"; + } + + void get_param_tensors(std::map& tensors, const std::string& prefix) override { + boogu.get_param_tensors(tensors, prefix); + } + + ggml_cgraph* build_graph(const sd::Tensor& x_tensor, + const sd::Tensor& timesteps_tensor, + const sd::Tensor& context_tensor, + const std::vector>& ref_latents_tensor = {}) { + ggml_cgraph* gf = new_graph_custom(BOOGU_GRAPH_SIZE); + ggml_tensor* x = make_input(x_tensor); + ggml_tensor* timesteps = make_input(timesteps_tensor); + GGML_ASSERT(x->ne[3] == 1); + GGML_ASSERT(!context_tensor.empty()); + ggml_tensor* context = make_input(context_tensor); + + std::vector ref_latents; + ref_latents.reserve(ref_latents_tensor.size()); + for (const auto& ref_latent_tensor : ref_latents_tensor) { + ref_latents.push_back(make_input(ref_latent_tensor)); + } + + pe_vec = gen_boogu_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + ref_latents, + config.theta, + config.axes_dim); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); + set_backend_tensor_data(pe, pe_vec.data()); + + auto runner_ctx = get_context(); + ggml_tensor* out = boogu.forward(&runner_ctx, x, timesteps, context, pe, ref_latents); + ggml_build_forward_expand(gf, out); + return gf; + } + + sd::Tensor compute(int n_threads, + const sd::Tensor& x, + const sd::Tensor& timesteps, + const sd::Tensor& context, + const std::vector>& ref_latents = {}) { + auto get_graph = [&]() -> ggml_cgraph* { + return build_graph(x, timesteps, context, ref_latents); + }; + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + } + + sd::Tensor compute(int n_threads, + const DiffusionParams& diffusion_params) override { + GGML_ASSERT(diffusion_params.x != nullptr); + GGML_ASSERT(diffusion_params.timesteps != nullptr); + static const std::vector> empty_ref_latents; + return compute(n_threads, + *diffusion_params.x, + *diffusion_params.timesteps, + tensor_or_empty(diffusion_params.context), + diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents); + } + }; +} // namespace Boogu + +#endif // __SD_MODEL_DIFFUSION_BOOGU_HPP__ diff --git a/src/control.hpp b/src/model/diffusion/control.hpp similarity index 85% rename from src/control.hpp rename to src/model/diffusion/control.hpp index a38f5d14..eeb8f510 100644 --- a/src/control.hpp +++ b/src/model/diffusion/control.hpp @@ -1,8 +1,9 @@ -#ifndef __CONTROL_HPP__ -#define __CONTROL_HPP__ +#ifndef __SD_MODEL_DIFFUSION_CONTROL_HPP__ +#define __SD_MODEL_DIFFUSION_CONTROL_HPP__ -#include "common_block.hpp" -#include "model.h" +#include "model/common/block.hpp" +#include "model_loader.h" +#include "model_manager.h" #define CONTROL_NET_GRAPH_SIZE 1536 @@ -309,73 +310,47 @@ public: struct ControlNet : public GGMLRunner { SDVersion version = VERSION_SD1; ControlNetBlock control_net; + std::string weight_prefix; - ggml_backend_buffer_t control_buffer = nullptr; - ggml_context* control_ctx = nullptr; std::vector control_outputs_ggml; ggml_tensor* guided_hint_output_ggml = nullptr; std::vector> controls; - sd::Tensor guided_hint; bool guided_hint_cached = false; + std::shared_ptr owned_model_manager; + ggml_backend_t params_backend = nullptr; + + static const char* guided_hint_cache_name() { + return "controlnet.guided_hint"; + } ControlNet(ggml_backend_t backend, - ggml_backend_t params_backend, - const String2TensorStorage& tensor_storage_map = {}, - SDVersion version = VERSION_SD1) - : GGMLRunner(backend, params_backend), control_net(version) { - control_net.init(params_ctx, tensor_storage_map, ""); + ggml_backend_t params_backend_, + const String2TensorStorage& tensor_storage_map = {}, + SDVersion version = VERSION_SD1, + const std::string& prefix = "", + std::shared_ptr weight_manager = nullptr) + : GGMLRunner(backend, weight_manager), version(version), control_net(version), weight_prefix(prefix), params_backend(params_backend_) { + control_net.init(params_ctx, tensor_storage_map, prefix); } ~ControlNet() override { free_control_ctx(); } - void alloc_control_ctx(std::vector outs) { - ggml_init_params params; - params.mem_size = static_cast(outs.size() * ggml_tensor_overhead()) + 1024 * 1024; - params.mem_buffer = nullptr; - params.no_alloc = true; - control_ctx = ggml_init(params); - - control_outputs_ggml.resize(outs.size() - 1); - - size_t control_buffer_size = 0; - - guided_hint_output_ggml = ggml_dup_tensor(control_ctx, outs[0]); - control_buffer_size += ggml_nbytes(guided_hint_output_ggml); - - for (int i = 0; i < outs.size() - 1; i++) { - control_outputs_ggml[i] = ggml_dup_tensor(control_ctx, outs[i + 1]); - control_buffer_size += ggml_nbytes(control_outputs_ggml[i]); - } - - control_buffer = ggml_backend_alloc_ctx_tensors(control_ctx, runtime_backend); - - LOG_DEBUG("control buffer size %.2fMB", control_buffer_size * 1.f / 1024.f / 1024.f); - } - void free_control_ctx() { - if (control_buffer != nullptr) { - ggml_backend_buffer_free(control_buffer); - control_buffer = nullptr; - } - if (control_ctx != nullptr) { - ggml_free(control_ctx); - control_ctx = nullptr; - } guided_hint_output_ggml = nullptr; guided_hint_cached = false; - guided_hint = {}; control_outputs_ggml.clear(); controls.clear(); + free_cache_ctx_and_buffer(); } std::string get_desc() override { return "control_net"; } - void get_param_tensors(std::map& tensors, const std::string prefix) { - control_net.get_param_tensors(tensors, prefix); + void get_param_tensors(std::map& tensors) { + control_net.get_param_tensors(tensors, weight_prefix); } ggml_cgraph* build_graph(const sd::Tensor& x_tensor, @@ -391,11 +366,17 @@ struct ControlNet : public GGMLRunner { ggml_tensor* context = make_optional_input(context_tensor); ggml_tensor* y = make_optional_input(y_tensor); + guided_hint_output_ggml = nullptr; + control_outputs_ggml.clear(); + ggml_tensor* guided_hint_input = nullptr; - if (guided_hint_cached && !guided_hint.empty()) { - guided_hint_input = make_input(guided_hint); - hint = nullptr; - } else { + if (guided_hint_cached) { + guided_hint_input = get_cache_tensor_by_name(guided_hint_cache_name()); + if (guided_hint_input == nullptr) { + guided_hint_cached = false; + } + } + if (guided_hint_input == nullptr) { hint = make_input(hint_tensor); } @@ -409,13 +390,19 @@ struct ControlNet : public GGMLRunner { context, y); - if (control_ctx == nullptr) { - alloc_control_ctx(outs); + if (guided_hint_input == nullptr && !outs.empty()) { + guided_hint_output_ggml = outs[0]; + ggml_set_output(guided_hint_output_ggml); + cache(guided_hint_cache_name(), guided_hint_output_ggml); + ggml_build_forward_expand(gf, guided_hint_output_ggml); } - ggml_build_forward_expand(gf, ggml_cpy(compute_ctx, outs[0], guided_hint_output_ggml)); - for (int i = 0; i < outs.size() - 1; i++) { - ggml_build_forward_expand(gf, ggml_cpy(compute_ctx, outs[i + 1], control_outputs_ggml[i])); + control_outputs_ggml.reserve(outs.size() > 0 ? outs.size() - 1 : 0); + for (size_t i = 1; i < outs.size(); i++) { + ggml_tensor* control_output = outs[i]; + ggml_set_output(control_output); + ggml_build_forward_expand(gf, control_output); + control_outputs_ggml.push_back(control_output); } return gf; @@ -435,15 +422,12 @@ struct ControlNet : public GGMLRunner { return build_graph(x, hint, timesteps, context, y); }; - auto compute_result = GGMLRunner::compute(get_graph, n_threads, false); + auto compute_result = GGMLRunner::compute(get_graph, n_threads, false, false, false, true); if (!compute_result.has_value()) { return std::nullopt; } - if (guided_hint_output_ggml != nullptr) { - guided_hint = restore_trailing_singleton_dims(sd::make_sd_tensor_from_ggml(guided_hint_output_ggml), - 4); - } + guided_hint_cached = get_cache_tensor_by_name(guided_hint_cache_name()) != nullptr; controls.clear(); controls.reserve(control_outputs_ggml.size()); for (ggml_tensor* control : control_outputs_ggml) { @@ -451,37 +435,41 @@ struct ControlNet : public GGMLRunner { GGML_ASSERT(!control_host.empty()); controls.push_back(std::move(control_host)); } - guided_hint_cached = true; return controls; } bool load_from_file(const std::string& file_path, int n_threads) { LOG_INFO("loading control net from '%s'", file_path.c_str()); - if (!alloc_params_buffer()) { - LOG_ERROR("control net model buffer allocation failed"); - return false; - } - std::map tensors; control_net.get_param_tensors(tensors); - std::set ignore_tensors; - ModelLoader model_loader; + auto manager = std::dynamic_pointer_cast(weight_manager.lock()); + if (manager == nullptr) { + owned_model_manager = std::make_shared(); + weight_manager = owned_model_manager; + manager = owned_model_manager; + } + + ModelLoader& model_loader = manager->loader(); if (!model_loader.init_from_file_and_convert_name(file_path)) { LOG_ERROR("init control net model loader from file failed: '%s'", file_path.c_str()); return false; } - bool success = model_loader.load_tensors(tensors, ignore_tensors, n_threads); - - if (!success) { - LOG_ERROR("load control net tensors from model loader failed"); + manager->set_n_threads(n_threads); + if (!manager->register_param_tensors("ControlNet", + std::move(tensors), + ModelManager::ResidencyMode::ParamBackend, + runtime_backend, + params_backend) || + !manager->validate_registered_tensors()) { + LOG_ERROR("register control net tensors with model manager failed"); return false; } LOG_INFO("control net model loaded"); - return success; + return true; } }; -#endif // __CONTROL_HPP__ +#endif // __SD_MODEL_DIFFUSION_CONTROL_HPP__ diff --git a/src/common_dit.hpp b/src/model/diffusion/dit.hpp similarity index 98% rename from src/common_dit.hpp rename to src/model/diffusion/dit.hpp index c25302ff..c1851856 100644 --- a/src/common_dit.hpp +++ b/src/model/diffusion/dit.hpp @@ -1,7 +1,7 @@ -#ifndef __COMMON_DIT_HPP__ -#define __COMMON_DIT_HPP__ +#ifndef __SD_MODEL_DIFFUSION_DIT_HPP__ +#define __SD_MODEL_DIFFUSION_DIT_HPP__ -#include "ggml_extend.hpp" +#include "core/ggml_extend.hpp" namespace DiT { inline ggml_tensor* patchify(ggml_context* ctx, @@ -163,4 +163,4 @@ namespace DiT { } } // namespace DiT -#endif // __COMMON_DIT_HPP__ +#endif // __SD_MODEL_DIFFUSION_DIT_HPP__ diff --git a/src/ernie_image.hpp b/src/model/diffusion/ernie_image.hpp similarity index 81% rename from src/ernie_image.hpp rename to src/model/diffusion/ernie_image.hpp index 35546895..12fcada5 100644 --- a/src/ernie_image.hpp +++ b/src/model/diffusion/ernie_image.hpp @@ -1,18 +1,88 @@ -#ifndef __SD_ERNIE_IMAGE_HPP__ -#define __SD_ERNIE_IMAGE_HPP__ +#ifndef __SD_MODEL_DIFFUSION_ERNIE_IMAGE_HPP__ +#define __SD_MODEL_DIFFUSION_ERNIE_IMAGE_HPP__ #include #include -#include "common_dit.hpp" -#include "diffusion_model.hpp" -#include "flux.hpp" -#include "qwen_image.hpp" -#include "rope.hpp" +#include "model/common/rope.hpp" +#include "model/diffusion/dit.hpp" +#include "model/diffusion/flux.hpp" +#include "model/diffusion/model.hpp" +#include "model/diffusion/qwen_image.hpp" namespace ErnieImage { constexpr int ERNIE_IMAGE_GRAPH_SIZE = 40960; + struct ErnieImageConfig { + int64_t hidden_size = 4096; + int64_t num_heads = 32; + int64_t num_layers = 36; + int64_t ffn_hidden_size = 12288; + int64_t in_channels = 128; + int64_t out_channels = 128; + int patch_size = 1; + int64_t text_in_dim = 3072; + int theta = 256; + std::vector axes_dim = {32, 48, 48}; + int axes_dim_sum = 128; + float eps = 1e-6f; + + static ErnieImageConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) { + ErnieImageConfig config; + config.num_layers = 0; + int64_t detected_head_dim = 0; + for (const auto& [name, tensor_storage] : tensor_storage_map) { + if (!starts_with(name, prefix)) { + continue; + } + if (ends_with(name, "x_embedder.proj.weight") && tensor_storage.n_dims == 4) { + config.patch_size = static_cast(tensor_storage.ne[0]); + config.in_channels = tensor_storage.ne[2]; + config.hidden_size = tensor_storage.ne[3]; + } else if (ends_with(name, "text_proj.weight") && tensor_storage.n_dims == 2) { + config.text_in_dim = tensor_storage.ne[0]; + } else if (ends_with(name, "layers.0.self_attention.norm_q.weight")) { + detected_head_dim = tensor_storage.ne[0]; + } else if (ends_with(name, "layers.0.mlp.gate_proj.weight") && tensor_storage.n_dims == 2) { + config.ffn_hidden_size = tensor_storage.ne[1]; + } else if (ends_with(name, "final_linear.weight") && tensor_storage.n_dims == 2) { + int64_t out_dim = tensor_storage.ne[1]; + int64_t patch_area = config.patch_size * config.patch_size; + config.out_channels = out_dim / patch_area; + } + + size_t pos = name.find("layers."); + if (pos != std::string::npos) { + auto items = split_string(name.substr(pos), '.'); + if (items.size() > 1) { + int block_index = atoi(items[1].c_str()); + if (block_index + 1 > config.num_layers) { + config.num_layers = block_index + 1; + } + } + } + } + if (config.num_layers == 0) { + config.num_layers = 36; + } + if (detected_head_dim > 0) { + config.num_heads = config.hidden_size / detected_head_dim; + } + config.axes_dim_sum = 0; + for (int axis_dim : config.axes_dim) { + config.axes_dim_sum += axis_dim; + } + LOG_DEBUG("ernie_image: num_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %" PRId64 ", ffn_hidden_size = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64, + config.num_layers, + config.hidden_size, + config.num_heads, + config.ffn_hidden_size, + config.in_channels, + config.out_channels); + return config; + } + }; + __STATIC_INLINE__ ggml_tensor* timestep_embedding_sin_cos(ggml_context* ctx, ggml_tensor* timesteps, int dim, @@ -208,51 +278,36 @@ namespace ErnieImage { } }; - struct ErnieImageParams { - int64_t hidden_size = 4096; - int64_t num_heads = 32; - int64_t num_layers = 36; - int64_t ffn_hidden_size = 12288; - int64_t in_channels = 128; - int64_t out_channels = 128; - int patch_size = 1; - int64_t text_in_dim = 3072; - int theta = 256; - std::vector axes_dim = {32, 48, 48}; - int axes_dim_sum = 128; - float eps = 1e-6f; - }; - class ErnieImageModel : public GGMLBlock { public: - ErnieImageParams params; + ErnieImageConfig config; ErnieImageModel() = default; - ErnieImageModel(ErnieImageParams params) - : params(params) { - blocks["x_embedder.proj"] = std::make_shared(params.in_channels, - params.hidden_size, - std::pair{params.patch_size, params.patch_size}, - std::pair{params.patch_size, params.patch_size}, + ErnieImageModel(ErnieImageConfig config) + : config(config) { + blocks["x_embedder.proj"] = std::make_shared(config.in_channels, + config.hidden_size, + std::pair{config.patch_size, config.patch_size}, + std::pair{config.patch_size, config.patch_size}, std::pair{0, 0}, std::pair{1, 1}, true); - if (params.text_in_dim != params.hidden_size) { - blocks["text_proj"] = std::make_shared(params.text_in_dim, params.hidden_size, false); + if (config.text_in_dim != config.hidden_size) { + blocks["text_proj"] = std::make_shared(config.text_in_dim, config.hidden_size, false); } - blocks["time_embedding"] = std::make_shared(params.hidden_size, params.hidden_size); - blocks["adaLN_modulation.1"] = std::make_shared(params.hidden_size, 6 * params.hidden_size, true); + blocks["time_embedding"] = std::make_shared(config.hidden_size, config.hidden_size); + blocks["adaLN_modulation.1"] = std::make_shared(config.hidden_size, 6 * config.hidden_size, true); - for (int i = 0; i < params.num_layers; i++) { - blocks["layers." + std::to_string(i)] = std::make_shared(params.hidden_size, - params.num_heads, - params.ffn_hidden_size, - params.eps); + for (int i = 0; i < config.num_layers; i++) { + blocks["layers." + std::to_string(i)] = std::make_shared(config.hidden_size, + config.num_heads, + config.ffn_hidden_size, + config.eps); } - blocks["final_norm"] = std::make_shared(params.hidden_size, params.eps); - blocks["final_linear"] = std::make_shared(params.hidden_size, - params.patch_size * params.patch_size * params.out_channels, + blocks["final_norm"] = std::make_shared(config.hidden_size, config.eps); + blocks["final_linear"] = std::make_shared(config.hidden_size, + config.patch_size * config.patch_size * config.out_channels, true); } @@ -265,12 +320,12 @@ namespace ErnieImage { // context: [N, text_tokens, 3072] // pe: [image_tokens + text_tokens, head_dim/2, 2, 2] GGML_ASSERT(context != nullptr); - GGML_ASSERT(x->ne[1] % params.patch_size == 0 && x->ne[0] % params.patch_size == 0); + GGML_ASSERT(x->ne[1] % config.patch_size == 0 && x->ne[0] % config.patch_size == 0); int64_t W = x->ne[0]; int64_t H = x->ne[1]; - int64_t Hp = H / params.patch_size; - int64_t Wp = W / params.patch_size; + int64_t Hp = H / config.patch_size; + int64_t Wp = W / config.patch_size; int64_t n_img = Hp * Wp; int64_t N = x->ne[3]; @@ -292,7 +347,7 @@ namespace ErnieImage { auto hidden_states = ggml_concat(ctx->ggml_ctx, img, txt, 1); // [N, image_tokens + text_tokens, hidden_size] - auto sample = timestep_embedding_sin_cos(ctx->ggml_ctx, timestep, static_cast(params.hidden_size)); + auto sample = timestep_embedding_sin_cos(ctx->ggml_ctx, timestep, static_cast(config.hidden_size)); auto c = time_embedding->forward(ctx, sample); // [N, hidden_size] auto mod_params = adaLN_mod->forward(ctx, ggml_silu(ctx->ggml_ctx, c)); // [N, 6 * hidden_size] @@ -305,7 +360,7 @@ namespace ErnieImage { temb.push_back(ggml_reshape_3d(ctx->ggml_ctx, chunk, chunk->ne[0], 1, chunk->ne[1])); // [N, 1, hidden_size] } - for (int i = 0; i < params.num_layers; i++) { + for (int i = 0; i < config.num_layers; i++) { auto layer = std::dynamic_pointer_cast(blocks["layers." + std::to_string(i)]); hidden_states = layer->forward(ctx, hidden_states, pe, temb); sd::ggml_graph_cut::mark_graph_cut(hidden_states, "ernie_image.layers." + std::to_string(i), "hidden_states"); @@ -319,74 +374,25 @@ namespace ErnieImage { patches, Hp, Wp, - params.patch_size, - params.patch_size, + config.patch_size, + config.patch_size, false); // [N, out_channels, H, W] return out; } }; struct ErnieImageRunner : public DiffusionModelRunner { - ErnieImageParams ernie_params; + ErnieImageConfig config; ErnieImageModel ernie_image; std::vector pe_vec; ErnieImageRunner(ggml_backend_t backend, - ggml_backend_t params_backend, - const String2TensorStorage& tensor_storage_map = {}, - const std::string prefix = "") - : DiffusionModelRunner(backend, params_backend, prefix) { - ernie_params.num_layers = 0; - for (const auto& [name, tensor_storage] : tensor_storage_map) { - if (!starts_with(name, prefix)) { - continue; - } - if (ends_with(name, "x_embedder.proj.weight") && tensor_storage.n_dims == 4) { - ernie_params.patch_size = static_cast(tensor_storage.ne[0]); - ernie_params.in_channels = tensor_storage.ne[2]; - ernie_params.hidden_size = tensor_storage.ne[3]; - } else if (ends_with(name, "text_proj.weight") && tensor_storage.n_dims == 2) { - ernie_params.text_in_dim = tensor_storage.ne[0]; - } else if (ends_with(name, "layers.0.self_attention.norm_q.weight")) { - int64_t head_dim = tensor_storage.ne[0]; - ernie_params.num_heads = ernie_params.hidden_size / head_dim; - } else if (ends_with(name, "layers.0.mlp.gate_proj.weight") && tensor_storage.n_dims == 2) { - ernie_params.ffn_hidden_size = tensor_storage.ne[1]; - } else if (ends_with(name, "final_linear.weight") && tensor_storage.n_dims == 2) { - int64_t out_dim = tensor_storage.ne[1]; - ernie_params.out_channels = out_dim / ernie_params.patch_size / ernie_params.patch_size; - } - - size_t pos = name.find("layers."); - 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 > ernie_params.num_layers) { - ernie_params.num_layers = block_index + 1; - } - } - } - } - if (ernie_params.num_layers == 0) { - ernie_params.num_layers = 36; - } - ernie_params.axes_dim_sum = 0; - for (int axis_dim : ernie_params.axes_dim) { - ernie_params.axes_dim_sum += axis_dim; - } - - LOG_INFO("ernie_image: layers = %" PRId64 ", hidden_size = %" PRId64 ", heads = %" PRId64 - ", ffn_hidden_size = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64, - ernie_params.num_layers, - ernie_params.hidden_size, - ernie_params.num_heads, - ernie_params.ffn_hidden_size, - ernie_params.in_channels, - ernie_params.out_channels); - - ernie_image = ErnieImageModel(ernie_params); + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "", + std::shared_ptr weight_manager = nullptr) + : DiffusionModelRunner(backend, prefix, weight_manager), + config(ErnieImageConfig::detect_from_weights(tensor_storage_map, prefix)) { + ernie_image = ErnieImageModel(config); ernie_image.init(params_ctx, tensor_storage_map, prefix); } @@ -410,15 +416,15 @@ namespace ErnieImage { pe_vec = Rope::gen_ernie_image_pe(static_cast(x->ne[1]), static_cast(x->ne[0]), - ernie_params.patch_size, + config.patch_size, static_cast(x->ne[3]), static_cast(context->ne[1]), - ernie_params.theta, + config.theta, circular_y_enabled, circular_x_enabled, - ernie_params.axes_dim); - int pos_len = static_cast(pe_vec.size() / ernie_params.axes_dim_sum / 2); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, ernie_params.axes_dim_sum, 1, pos_len, 2); + config.axes_dim); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, config.axes_dim_sum, 1, pos_len, 2); set_backend_tensor_data(pe, pe_vec.data()); auto runner_ctx = get_context(); @@ -434,7 +440,7 @@ namespace ErnieImage { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); } sd::Tensor compute(int n_threads, @@ -449,4 +455,4 @@ namespace ErnieImage { }; } // namespace ErnieImage -#endif // __SD_ERNIE_IMAGE_HPP__ +#endif // __SD_MODEL_DIFFUSION_ERNIE_IMAGE_HPP__ diff --git a/src/flux.hpp b/src/model/diffusion/flux.hpp similarity index 78% rename from src/flux.hpp rename to src/model/diffusion/flux.hpp index 82832cea..53f9ac31 100644 --- a/src/flux.hpp +++ b/src/model/diffusion/flux.hpp @@ -1,18 +1,194 @@ -#ifndef __FLUX_HPP__ -#define __FLUX_HPP__ +#ifndef __SD_MODEL_DIFFUSION_FLUX_HPP__ +#define __SD_MODEL_DIFFUSION_FLUX_HPP__ #include #include -#include "common_dit.hpp" -#include "diffusion_model.hpp" -#include "model.h" -#include "rope.hpp" +#include "model/adapter/pulid.hpp" +#include "model/common/rope.hpp" +#include "model/diffusion/dit.hpp" +#include "model/diffusion/model.hpp" +#include "model/diffusion/sefi_image.hpp" +#include "model_loader.h" #define FLUX_GRAPH_SIZE 10240 namespace Flux { + struct ChromaRadianceConfig { + int64_t nerf_hidden_size = 64; + int nerf_mlp_ratio = 4; + int nerf_depth = 4; + int nerf_max_freqs = 8; + bool use_x0 = false; + bool fake_patch_size_x2 = false; + }; + + struct FluxConfig { + SDVersion version = VERSION_FLUX; + bool is_chroma = false; + bool is_sefi = false; + int64_t semantic_channels = 0; + float sefi_delta_t = 0.1f; + int patch_size = 2; + int64_t in_channels = 64; + int64_t out_channels = 64; + int64_t vec_in_dim = 768; + int64_t context_in_dim = 4096; + int64_t hidden_size = 3072; + float mlp_ratio = 4.0f; + int num_heads = 24; + int depth = 19; + int depth_single_blocks = 38; + std::vector axes_dim = {16, 56, 56}; + int axes_dim_sum = 128; + int theta = 10000; + bool qkv_bias = true; + bool guidance_embed = true; + int64_t in_dim = 64; + bool disable_bias = false; + bool share_modulation = false; + bool semantic_txt_norm = false; + bool use_yak_mlp = false; + bool use_mlp_silu_act = false; + float ref_index_scale = 1.f; + ChromaRadianceConfig chroma_radiance_params; + + bool pulid_enabled = false; + int pulid_double_interval = 2; + int pulid_single_interval = 4; + + static FluxConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, + const std::string& prefix, + SDVersion version = VERSION_FLUX) { + FluxConfig config; + config.version = version; + config.guidance_embed = false; + config.depth = 0; + config.depth_single_blocks = 0; + if (version == VERSION_FLUX_FILL) { + config.in_channels = 384; + } else if (version == VERSION_FLUX_CONTROLS) { + config.in_channels = 128; + } else if (version == VERSION_FLEX_2) { + config.in_channels = 196; + } else if (version == VERSION_CHROMA_RADIANCE) { + config.in_channels = 3; + config.patch_size = 16; + } else if (version == VERSION_OVIS_IMAGE) { + config.semantic_txt_norm = true; + config.use_yak_mlp = true; + config.vec_in_dim = 0; + } else if (sd_version_is_flux2(version)) { + config.in_channels = 128; + config.patch_size = 1; + config.out_channels = 128; + config.mlp_ratio = 3.f; + config.theta = 2000; + config.axes_dim = {32, 32, 32, 32}; + config.vec_in_dim = 0; + config.qkv_bias = false; + config.disable_bias = true; + config.share_modulation = true; + config.ref_index_scale = 10.f; + config.use_mlp_silu_act = true; + } else if (sd_version_is_sefi_image(version)) { + config.is_sefi = true; + config.semantic_channels = 16; + config.in_channels = 128 + config.semantic_channels; + config.patch_size = 1; + config.out_channels = 128 + config.semantic_channels; + config.mlp_ratio = 3.f; + config.theta = 2000; + config.axes_dim = {32, 32, 32, 32}; + config.vec_in_dim = 0; + config.qkv_bias = false; + config.disable_bias = true; + config.share_modulation = true; + config.ref_index_scale = 10.f; + config.use_mlp_silu_act = true; + } else if (sd_version_is_longcat(version)) { + config.context_in_dim = 3584; + config.vec_in_dim = 0; + } + + int64_t head_dim = 0; + int64_t actual_radiance_patch_size = -1; + for (const auto& [name, tensor_storage] : tensor_storage_map) { + if (!starts_with(name, prefix)) { + continue; + } + if (name.find("guidance_in.in_layer.weight") != std::string::npos) { + config.guidance_embed = true; + } + if (name.find("__x0__") != std::string::npos) { + LOG_DEBUG("using x0 prediction"); + config.chroma_radiance_params.use_x0 = true; + } + if (name.find("__32x32__") != std::string::npos) { + LOG_DEBUG("using patch size 32"); + config.patch_size = 32; + } + if (name.find("img_in_patch.weight") != std::string::npos) { + actual_radiance_patch_size = tensor_storage.ne[0]; + LOG_DEBUG("actual radiance patch size: %" PRId64, actual_radiance_patch_size); + } + if (name.find("distilled_guidance_layer.in_proj.weight") != std::string::npos) { + config.is_chroma = true; + } + size_t db = name.find("double_blocks."); + if (db != std::string::npos) { + std::string block_name = name.substr(db); + int block_depth = atoi(block_name.substr(14, block_name.find(".", 14)).c_str()); + if (block_depth + 1 > config.depth) { + config.depth = block_depth + 1; + } + } + size_t sb = name.find("single_blocks."); + if (sb != std::string::npos) { + std::string block_name = name.substr(sb); + int block_depth = atoi(block_name.substr(14, block_name.find(".", 14)).c_str()); + if (block_depth + 1 > config.depth_single_blocks) { + config.depth_single_blocks = block_depth + 1; + } + } + if (ends_with(name, "txt_in.weight")) { + config.context_in_dim = tensor_storage.ne[0]; + config.hidden_size = tensor_storage.ne[1]; + } + if (ends_with(name, "single_blocks.0.norm.key_norm.scale")) { + head_dim = tensor_storage.ne[0]; + } + if (ends_with(name, "double_blocks.0.txt_attn.norm.key_norm.scale")) { + head_dim = tensor_storage.ne[0]; + } + if (name.find("pulid_ca.") != std::string::npos) { + config.pulid_enabled = true; + } + } + if (actual_radiance_patch_size > 0 && actual_radiance_patch_size != config.patch_size) { + GGML_ASSERT(config.patch_size == 2 * actual_radiance_patch_size); + LOG_DEBUG("using fake x2 patch size"); + config.chroma_radiance_params.fake_patch_size_x2 = true; + } + if (head_dim > 0) { + config.num_heads = static_cast(config.hidden_size / head_dim); + } + config.axes_dim_sum = 0; + for (int axis_dim : config.axes_dim) { + config.axes_dim_sum += axis_dim; + } + LOG_DEBUG("flux: depth = %d, depth_single_blocks = %d, guidance_embed = %s, context_in_dim = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %d", + config.depth, + config.depth_single_blocks, + config.guidance_embed ? "true" : "false", + config.context_in_dim, + config.hidden_size, + config.num_heads); + return config; + } + }; + struct MLPEmbedder : public UnaryBlock { public: MLPEmbedder(int64_t in_dim, int64_t hidden_dim, bool bias = true) { @@ -566,8 +742,8 @@ namespace Flux { auto m = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, c)); // [N, 2 * hidden_size] auto m_vec = ggml_ext_chunk(ctx->ggml_ctx, m, 2, 0); - shift = m_vec[0]; // [N, hidden_size] - scale = m_vec[1]; // [N, hidden_size] + shift = m_vec[0]; + scale = m_vec[1]; } x = Flux::modulate(ctx->ggml_ctx, norm_final->forward(ctx, x), shift, scale); @@ -723,127 +899,106 @@ namespace Flux { } }; - struct ChromaRadianceParams { - int64_t nerf_hidden_size = 64; - int nerf_mlp_ratio = 4; - int nerf_depth = 4; - int nerf_max_freqs = 8; - bool use_x0 = false; - bool fake_patch_size_x2 = false; - }; - - struct FluxParams { - SDVersion version = VERSION_FLUX; - bool is_chroma = false; - int patch_size = 2; - int64_t in_channels = 64; - int64_t out_channels = 64; - int64_t vec_in_dim = 768; - int64_t context_in_dim = 4096; - int64_t hidden_size = 3072; - float mlp_ratio = 4.0f; - int num_heads = 24; - int depth = 19; - int depth_single_blocks = 38; - std::vector axes_dim = {16, 56, 56}; - int axes_dim_sum = 128; - int theta = 10000; - bool qkv_bias = true; - bool guidance_embed = true; - int64_t in_dim = 64; - bool disable_bias = false; - bool share_modulation = false; - bool semantic_txt_norm = false; - bool use_yak_mlp = false; - bool use_mlp_silu_act = false; - float ref_index_scale = 1.f; - ChromaRadianceParams chroma_radiance_params; - }; - struct Flux : public GGMLBlock { public: - FluxParams params; + FluxConfig config; Flux() {} - Flux(FluxParams params) - : params(params) { - if (params.version == VERSION_CHROMA_RADIANCE) { - std::pair kernel_size = {params.patch_size, params.patch_size}; - if (params.chroma_radiance_params.fake_patch_size_x2) { - kernel_size = {params.patch_size / 2, params.patch_size / 2}; + Flux(FluxConfig config) + : config(config) { + if (config.version == VERSION_CHROMA_RADIANCE) { + std::pair kernel_size = {config.patch_size, config.patch_size}; + if (config.chroma_radiance_params.fake_patch_size_x2) { + kernel_size = {config.patch_size / 2, config.patch_size / 2}; } std::pair stride = kernel_size; - blocks["img_in_patch"] = std::make_shared(params.in_channels, - params.hidden_size, + blocks["img_in_patch"] = std::make_shared(config.in_channels, + config.hidden_size, kernel_size, stride); } else { - blocks["img_in"] = std::make_shared(params.in_channels, params.hidden_size, !params.disable_bias); + blocks["img_in"] = std::make_shared(config.in_channels, config.hidden_size, !config.disable_bias); } - if (params.is_chroma) { - blocks["distilled_guidance_layer"] = std::make_shared(params.in_dim, params.hidden_size); + if (config.is_chroma) { + blocks["distilled_guidance_layer"] = std::make_shared(config.in_dim, config.hidden_size); + } else if (config.is_sefi) { + blocks["dual_time_embed"] = std::make_shared(256, config.hidden_size); } else { - blocks["time_in"] = std::make_shared(256, params.hidden_size, !params.disable_bias); - if (params.vec_in_dim > 0) { - blocks["vector_in"] = std::make_shared(params.vec_in_dim, params.hidden_size, !params.disable_bias); + blocks["time_in"] = std::make_shared(256, config.hidden_size, !config.disable_bias); + if (config.vec_in_dim > 0) { + blocks["vector_in"] = std::make_shared(config.vec_in_dim, config.hidden_size, !config.disable_bias); } - if (params.guidance_embed) { - blocks["guidance_in"] = std::make_shared(256, params.hidden_size, !params.disable_bias); + if (config.guidance_embed) { + blocks["guidance_in"] = std::make_shared(256, config.hidden_size, !config.disable_bias); } } - if (params.semantic_txt_norm) { - blocks["txt_norm"] = std::make_shared(params.context_in_dim); + if (config.semantic_txt_norm) { + blocks["txt_norm"] = std::make_shared(config.context_in_dim); } - blocks["txt_in"] = std::make_shared(params.context_in_dim, params.hidden_size, !params.disable_bias); + blocks["txt_in"] = std::make_shared(config.context_in_dim, config.hidden_size, !config.disable_bias); - for (int i = 0; i < params.depth; i++) { - blocks["double_blocks." + std::to_string(i)] = std::make_shared(params.hidden_size, - params.num_heads, - params.mlp_ratio, + for (int i = 0; i < config.depth; i++) { + blocks["double_blocks." + std::to_string(i)] = std::make_shared(config.hidden_size, + config.num_heads, + config.mlp_ratio, i, - params.qkv_bias, - params.is_chroma, - params.share_modulation, - !params.disable_bias, - params.use_yak_mlp, - params.use_mlp_silu_act); + config.qkv_bias, + config.is_chroma, + config.share_modulation, + !config.disable_bias, + config.use_yak_mlp, + config.use_mlp_silu_act); } - for (int i = 0; i < params.depth_single_blocks; i++) { - blocks["single_blocks." + std::to_string(i)] = std::make_shared(params.hidden_size, - params.num_heads, - params.mlp_ratio, + for (int i = 0; i < config.depth_single_blocks; i++) { + blocks["single_blocks." + std::to_string(i)] = std::make_shared(config.hidden_size, + config.num_heads, + config.mlp_ratio, i, 0.f, - params.is_chroma, - params.share_modulation, - !params.disable_bias, - params.use_yak_mlp, - params.use_mlp_silu_act); + config.is_chroma, + config.share_modulation, + !config.disable_bias, + config.use_yak_mlp, + config.use_mlp_silu_act); } - if (params.version == VERSION_CHROMA_RADIANCE) { - blocks["nerf_image_embedder"] = std::make_shared(params.in_channels, - params.chroma_radiance_params.nerf_hidden_size, - params.chroma_radiance_params.nerf_max_freqs); + if (config.version == VERSION_CHROMA_RADIANCE) { + blocks["nerf_image_embedder"] = std::make_shared(config.in_channels, + config.chroma_radiance_params.nerf_hidden_size, + config.chroma_radiance_params.nerf_max_freqs); - for (int i = 0; i < params.chroma_radiance_params.nerf_depth; i++) { - blocks["nerf_blocks." + std::to_string(i)] = std::make_shared(params.hidden_size, - params.chroma_radiance_params.nerf_hidden_size, - params.chroma_radiance_params.nerf_mlp_ratio); + for (int i = 0; i < config.chroma_radiance_params.nerf_depth; i++) { + blocks["nerf_blocks." + std::to_string(i)] = std::make_shared(config.hidden_size, + config.chroma_radiance_params.nerf_hidden_size, + config.chroma_radiance_params.nerf_mlp_ratio); } - blocks["nerf_final_layer_conv"] = std::make_shared(params.chroma_radiance_params.nerf_hidden_size, - params.in_channels); + blocks["nerf_final_layer_conv"] = std::make_shared(config.chroma_radiance_params.nerf_hidden_size, + config.in_channels); } else { - blocks["final_layer"] = std::make_shared(params.hidden_size, 1, params.out_channels, params.is_chroma, !params.disable_bias); + blocks["final_layer"] = std::make_shared(config.hidden_size, 1, config.out_channels, config.is_chroma, !config.disable_bias); } - if (params.share_modulation) { - blocks["double_stream_modulation_img"] = std::make_shared(params.hidden_size, true, !params.disable_bias); - blocks["double_stream_modulation_txt"] = std::make_shared(params.hidden_size, true, !params.disable_bias); - blocks["single_stream_modulation"] = std::make_shared(params.hidden_size, false, !params.disable_bias); + if (config.share_modulation) { + blocks["double_stream_modulation_img"] = std::make_shared(config.hidden_size, true, !config.disable_bias); + blocks["double_stream_modulation_txt"] = std::make_shared(config.hidden_size, true, !config.disable_bias); + blocks["single_stream_modulation"] = std::make_shared(config.hidden_size, false, !config.disable_bias); + } + + if (config.pulid_enabled) { + int num_double_ca = (config.depth + config.pulid_double_interval - 1) / config.pulid_double_interval; + int num_single_ca = (config.depth_single_blocks + config.pulid_single_interval - 1) / config.pulid_single_interval; + int num_ca = num_double_ca + num_single_ca; + for (int i = 0; i < num_ca; i++) { + blocks["pulid_ca." + std::to_string(i)] = + std::shared_ptr(new PuLIDPerceiverAttentionCA( + /*dim=*/config.hidden_size, + /*dim_head=*/PuLIDPerceiverAttentionCA::DEFAULT_DIM_HEAD, + /*heads=*/PuLIDPerceiverAttentionCA::DEFAULT_HEADS, + /*kv_dim=*/PuLIDPerceiverAttentionCA::DEFAULT_KV_DIM)); + } } } @@ -855,7 +1010,9 @@ namespace Flux { ggml_tensor* guidance, ggml_tensor* pe, ggml_tensor* mod_index_arange = nullptr, - std::vector skip_layers = {}) { + std::vector skip_layers = {}, + ggml_tensor* pulid_id = nullptr, + float pulid_id_weight = 1.0f) { auto img_in = std::dynamic_pointer_cast(blocks["img_in"]); auto txt_in = std::dynamic_pointer_cast(blocks["txt_in"]); auto final_layer = std::dynamic_pointer_cast(blocks["final_layer"]); @@ -866,7 +1023,7 @@ namespace Flux { ggml_tensor* vec; ggml_tensor* txt_img_mask = nullptr; - if (params.is_chroma) { + if (config.is_chroma) { int64_t mod_index_length = 344; auto approx = std::dynamic_pointer_cast(blocks["distilled_guidance_layer"]); auto distill_timestep = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, 16, 10000, 1000.f); @@ -891,10 +1048,15 @@ namespace Flux { if (y != nullptr) { txt_img_mask = ggml_pad(ctx->ggml_ctx, y, static_cast(img->ne[1]), 0, 0, 0); } + } else if (config.is_sefi) { + auto dual_time_embed = std::dynamic_pointer_cast(blocks["dual_time_embed"]); + auto timestep_sem = ggml_view_1d(ctx->ggml_ctx, timesteps, 1, 0); + auto timestep_tex = ggml_view_1d(ctx->ggml_ctx, timesteps, 1, ggml_element_size(timesteps)); + vec = dual_time_embed->forward(ctx, timestep_sem, timestep_tex); } else { auto time_in = std::dynamic_pointer_cast(blocks["time_in"]); vec = time_in->forward(ctx, ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, 256, 10000, 1000.f)); - if (params.guidance_embed) { + if (config.guidance_embed) { GGML_ASSERT(guidance != nullptr); auto guidance_in = std::dynamic_pointer_cast(blocks["guidance_in"]); // bf16 and fp16 result is different @@ -902,7 +1064,7 @@ namespace Flux { vec = ggml_add(ctx->ggml_ctx, vec, guidance_in->forward(ctx, g_in)); } - if (params.vec_in_dim > 0) { + if (config.vec_in_dim > 0) { auto vector_in = std::dynamic_pointer_cast(blocks["vector_in"]); vec = ggml_add(ctx->ggml_ctx, vec, vector_in->forward(ctx, y)); } @@ -911,7 +1073,7 @@ namespace Flux { std::vector ds_img_mods; std::vector ds_txt_mods; std::vector ss_mods; - if (params.share_modulation) { + if (config.share_modulation) { auto double_stream_modulation_img = std::dynamic_pointer_cast(blocks["double_stream_modulation_img"]); auto double_stream_modulation_txt = std::dynamic_pointer_cast(blocks["double_stream_modulation_txt"]); auto single_stream_modulation = std::dynamic_pointer_cast(blocks["single_stream_modulation"]); @@ -921,7 +1083,7 @@ namespace Flux { ss_mods = single_stream_modulation->forward(ctx, vec); } - if (params.semantic_txt_norm) { + if (config.semantic_txt_norm) { auto semantic_txt_norm = std::dynamic_pointer_cast(blocks["txt_norm"]); txt = semantic_txt_norm->forward(ctx, txt); @@ -932,7 +1094,14 @@ namespace Flux { sd::ggml_graph_cut::mark_graph_cut(txt, "flux.prelude", "txt"); sd::ggml_graph_cut::mark_graph_cut(vec, "flux.prelude", "vec"); - for (int i = 0; i < params.depth; i++) { + const bool pulid_active = config.pulid_enabled && pulid_id != nullptr; + if (pulid_active && !skip_layers.empty()) { + LOG_WARN("PuLID + skip_layers is not supported; disabling PuLID for this generation."); + } + const bool pulid_run = pulid_active && skip_layers.empty(); + int ca_idx = 0; + + for (int i = 0; i < config.depth; i++) { if (skip_layers.size() > 0 && std::find(skip_layers.begin(), skip_layers.end(), i) != skip_layers.end()) { continue; } @@ -944,17 +1113,50 @@ namespace Flux { txt = img_txt.second; // [N, n_txt_token, hidden_size] sd::ggml_graph_cut::mark_graph_cut(img, "flux.double_blocks." + std::to_string(i), "img"); sd::ggml_graph_cut::mark_graph_cut(txt, "flux.double_blocks." + std::to_string(i), "txt"); + + if (pulid_run && (i % config.pulid_double_interval == 0)) { + auto pulid_ca = std::dynamic_pointer_cast( + blocks["pulid_ca." + std::to_string(ca_idx)]); + ggml_tensor* ca_out = pulid_ca->forward(ctx, pulid_id, img); // [N, n_img_token, hidden_size] + img = ggml_add(ctx->ggml_ctx, img, ggml_scale(ctx->ggml_ctx, ca_out, pulid_id_weight)); + sd::ggml_graph_cut::mark_graph_cut(img, "flux.pulid_ca." + std::to_string(ca_idx), "img"); + ca_idx++; + } } - auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1); // [N, n_txt_token + n_img_token, hidden_size] - for (int i = 0; i < params.depth_single_blocks; i++) { - if (skip_layers.size() > 0 && std::find(skip_layers.begin(), skip_layers.end(), i + params.depth) != skip_layers.end()) { + auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1); // [N, n_txt_token + n_img_token, hidden_size] + const int64_t n_txt_tok = txt->ne[1]; + for (int i = 0; i < config.depth_single_blocks; i++) { + if (skip_layers.size() > 0 && std::find(skip_layers.begin(), skip_layers.end(), i + config.depth) != skip_layers.end()) { continue; } auto block = std::dynamic_pointer_cast(blocks["single_blocks." + std::to_string(i)]); txt_img = block->forward(ctx, txt_img, vec, pe, txt_img_mask, ss_mods); sd::ggml_graph_cut::mark_graph_cut(txt_img, "flux.single_blocks." + std::to_string(i), "txt_img"); + + if (pulid_run && (i % config.pulid_single_interval == 0)) { + auto pulid_ca = std::dynamic_pointer_cast( + blocks["pulid_ca." + std::to_string(ca_idx)]); + ggml_tensor* txt_part = ggml_view_3d(ctx->ggml_ctx, txt_img, + txt_img->ne[0], n_txt_tok, txt_img->ne[2], + txt_img->nb[1], txt_img->nb[2], + 0); + ggml_tensor* img_part = ggml_view_3d(ctx->ggml_ctx, txt_img, + txt_img->ne[0], + txt_img->ne[1] - n_txt_tok, + txt_img->ne[2], + txt_img->nb[1], + txt_img->nb[2], + n_txt_tok * txt_img->nb[1]); + txt_part = ggml_cont(ctx->ggml_ctx, txt_part); + img_part = ggml_cont(ctx->ggml_ctx, img_part); + ggml_tensor* ca_out = pulid_ca->forward(ctx, pulid_id, img_part); + img_part = ggml_add(ctx->ggml_ctx, img_part, ggml_scale(ctx->ggml_ctx, ca_out, pulid_id_weight)); + txt_img = ggml_concat(ctx->ggml_ctx, txt_part, img_part, 1); + sd::ggml_graph_cut::mark_graph_cut(txt_img, "flux.pulid_ca." + std::to_string(ca_idx), "txt_img"); + ca_idx++; + } } img = ggml_view_3d(ctx->ggml_ctx, @@ -993,20 +1195,22 @@ namespace Flux { ggml_tensor* mod_index_arange = nullptr, ggml_tensor* dct = nullptr, std::vector ref_latents = {}, - std::vector skip_layers = {}) { + std::vector skip_layers = {}, + ggml_tensor* pulid_id = nullptr, + float pulid_id_weight = 1.0f) { GGML_ASSERT(x->ne[3] == 1); int64_t W = x->ne[0]; int64_t H = x->ne[1]; int64_t C = x->ne[2]; - int patch_size = params.patch_size; + int patch_size = config.patch_size; int pad_h = (patch_size - H % patch_size) % patch_size; int pad_w = (patch_size - W % patch_size) % patch_size; - auto img = DiT::pad_to_patch_size(ctx, x, params.patch_size, params.patch_size); + auto img = DiT::pad_to_patch_size(ctx, x, config.patch_size, config.patch_size); auto orig_img = img; - if (params.chroma_radiance_params.fake_patch_size_x2) { + if (config.chroma_radiance_params.fake_patch_size_x2) { // It's supposed to be using GGML_SCALE_MODE_NEAREST, but this seems more stable // Maybe the implementation of nearest-neighbor interpolation in ggml behaves differently than the one in PyTorch? // img = F.interpolate(img, size=(H//2, W//2), mode="nearest") @@ -1019,7 +1223,8 @@ namespace Flux { img = ggml_reshape_3d(ctx->ggml_ctx, img, img->ne[0] * img->ne[1], img->ne[2], img->ne[3]); // [N, hidden_size, H/patch_size*W/patch_size] img = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, img, 1, 0, 2, 3)); // [N, H/patch_size*W/patch_size, hidden_size] - auto out = forward_orig(ctx, img, context, timestep, y, guidance, pe, mod_index_arange, skip_layers); // [N, n_img_token, hidden_size] + auto out = forward_orig(ctx, img, context, timestep, y, guidance, pe, mod_index_arange, skip_layers, + pulid_id, pulid_id_weight); // [N, n_img_token, hidden_size] // nerf decode auto nerf_image_embedder = std::dynamic_pointer_cast(blocks["nerf_image_embedder"]); @@ -1037,7 +1242,7 @@ namespace Flux { auto nerf_hidden = ggml_reshape_2d(ctx->ggml_ctx, out, out->ne[0], out->ne[1] * out->ne[2]); // [N*num_patches, hidden_size] auto img_dct = nerf_image_embedder->forward(ctx, nerf_pixels, dct); // [N*num_patches, patch_size*patch_size, nerf_hidden_size] - for (int i = 0; i < params.chroma_radiance_params.nerf_depth; i++) { + for (int i = 0; i < config.chroma_radiance_params.nerf_depth; i++) { auto block = std::dynamic_pointer_cast(blocks["nerf_blocks." + std::to_string(i)]); img_dct = block->forward(ctx, img_dct, nerf_hidden); @@ -1049,7 +1254,7 @@ namespace Flux { out = nerf_final_layer_conv->forward(ctx, img_dct); // [N, C, H, W] - if (params.chroma_radiance_params.use_x0) { + if (config.chroma_radiance_params.use_x0) { out = _apply_x0_residual(ctx, out, orig_img, timestep); } @@ -1067,20 +1272,22 @@ namespace Flux { ggml_tensor* mod_index_arange = nullptr, ggml_tensor* dct = nullptr, std::vector ref_latents = {}, - std::vector skip_layers = {}) { + std::vector skip_layers = {}, + ggml_tensor* pulid_id = nullptr, + float pulid_id_weight = 1.0f) { GGML_ASSERT(x->ne[3] == 1); int64_t W = x->ne[0]; int64_t H = x->ne[1]; int64_t C = x->ne[2]; - int patch_size = params.patch_size; + int patch_size = config.patch_size; int pad_h = (patch_size - H % patch_size) % patch_size; int pad_w = (patch_size - W % patch_size) % patch_size; auto img = DiT::pad_and_patchify(ctx, x, patch_size, patch_size); int64_t img_tokens = img->ne[1]; - if (params.version == VERSION_FLUX_FILL) { + if (config.version == VERSION_FLUX_FILL) { GGML_ASSERT(c_concat != nullptr); ggml_tensor* masked = ggml_view_4d(ctx->ggml_ctx, c_concat, c_concat->ne[0], c_concat->ne[1], C, 1, c_concat->nb[1], c_concat->nb[2], c_concat->nb[3], 0); ggml_tensor* mask = ggml_view_4d(ctx->ggml_ctx, c_concat, c_concat->ne[0], c_concat->ne[1], 8 * 8, 1, c_concat->nb[1], c_concat->nb[2], c_concat->nb[3], c_concat->nb[2] * C); @@ -1089,7 +1296,7 @@ namespace Flux { mask = DiT::pad_and_patchify(ctx, mask, patch_size, patch_size); img = ggml_concat(ctx->ggml_ctx, img, ggml_concat(ctx->ggml_ctx, masked, mask, 0), 0); - } else if (params.version == VERSION_FLEX_2) { + } else if (config.version == VERSION_FLEX_2) { GGML_ASSERT(c_concat != nullptr); ggml_tensor* masked = ggml_view_4d(ctx->ggml_ctx, c_concat, c_concat->ne[0], c_concat->ne[1], C, 1, c_concat->nb[1], c_concat->nb[2], c_concat->nb[3], 0); ggml_tensor* mask = ggml_view_4d(ctx->ggml_ctx, c_concat, c_concat->ne[0], c_concat->ne[1], 1, 1, c_concat->nb[1], c_concat->nb[2], c_concat->nb[3], c_concat->nb[2] * C); @@ -1100,7 +1307,7 @@ namespace Flux { control = DiT::pad_and_patchify(ctx, control, patch_size, patch_size); img = ggml_concat(ctx->ggml_ctx, img, ggml_concat(ctx->ggml_ctx, ggml_concat(ctx->ggml_ctx, masked, mask, 0), control, 0), 0); - } else if (params.version == VERSION_FLUX_CONTROLS) { + } else if (config.version == VERSION_FLUX_CONTROLS) { GGML_ASSERT(c_concat != nullptr); auto control = DiT::pad_and_patchify(ctx, c_concat, patch_size, patch_size); @@ -1114,7 +1321,8 @@ namespace Flux { } } - auto out = forward_orig(ctx, img, context, timestep, y, guidance, pe, mod_index_arange, skip_layers); // [N, num_tokens, C * patch_size * patch_size] + auto out = forward_orig(ctx, img, context, timestep, y, guidance, pe, mod_index_arange, skip_layers, + pulid_id, pulid_id_weight); // [N, num_tokens, C * patch_size * patch_size] if (out->ne[1] > img_tokens) { out = ggml_view_3d(ctx->ggml_ctx, out, out->ne[0], img_tokens, out->ne[2], out->nb[1], out->nb[2], 0); @@ -1136,7 +1344,9 @@ namespace Flux { ggml_tensor* mod_index_arange = nullptr, ggml_tensor* dct = nullptr, std::vector ref_latents = {}, - std::vector skip_layers = {}) { + std::vector skip_layers = {}, + ggml_tensor* pulid_id = nullptr, + float pulid_id_weight = 1.0f) { // Forward pass of DiT. // x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images) // timestep: (N,) tensor of diffusion timesteps @@ -1147,7 +1357,7 @@ namespace Flux { // pe: (L, d_head/2, 2, 2) // return: (N, C, H, W) - if (params.version == VERSION_CHROMA_RADIANCE) { + if (config.version == VERSION_CHROMA_RADIANCE) { return forward_chroma_radiance(ctx, x, timestep, @@ -1159,7 +1369,9 @@ namespace Flux { mod_index_arange, dct, ref_latents, - skip_layers); + skip_layers, + pulid_id, + pulid_id_weight); } else { return forward_flux_chroma(ctx, x, @@ -1172,14 +1384,16 @@ namespace Flux { mod_index_arange, dct, ref_latents, - skip_layers); + skip_layers, + pulid_id, + pulid_id_weight); } } }; struct FluxRunner : public DiffusionModelRunner { public: - FluxParams flux_params; + FluxConfig config; Flux flux; std::vector pe_vec; std::vector mod_index_arange_vec; @@ -1189,119 +1403,20 @@ namespace Flux { bool use_mask = false; FluxRunner(ggml_backend_t backend, - ggml_backend_t params_backend, - const String2TensorStorage& tensor_storage_map = {}, - const std::string prefix = "", - SDVersion version = VERSION_FLUX, - bool use_mask = false) - : DiffusionModelRunner(backend, params_backend, prefix), version(version), use_mask(use_mask) { - flux_params.version = version; - flux_params.guidance_embed = false; - flux_params.depth = 0; - flux_params.depth_single_blocks = 0; - if (version == VERSION_FLUX_FILL) { - flux_params.in_channels = 384; - } else if (version == VERSION_FLUX_CONTROLS) { - flux_params.in_channels = 128; - } else if (version == VERSION_FLEX_2) { - flux_params.in_channels = 196; - } else if (version == VERSION_CHROMA_RADIANCE) { - flux_params.in_channels = 3; - flux_params.patch_size = 16; - } else if (version == VERSION_OVIS_IMAGE) { - flux_params.semantic_txt_norm = true; - flux_params.use_yak_mlp = true; - flux_params.vec_in_dim = 0; - } else if (sd_version_is_flux2(version)) { - flux_params.in_channels = 128; - flux_params.patch_size = 1; - flux_params.out_channels = 128; - flux_params.mlp_ratio = 3.f; - flux_params.theta = 2000; - flux_params.axes_dim = {32, 32, 32, 32}; - flux_params.vec_in_dim = 0; - flux_params.qkv_bias = false; - flux_params.disable_bias = true; - 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; - for (auto pair : tensor_storage_map) { - std::string tensor_name = pair.first; - if (!starts_with(tensor_name, prefix)) - continue; - if (tensor_name.find("guidance_in.in_layer.weight") != std::string::npos) { - flux_params.guidance_embed = true; - } - if (tensor_name.find("__x0__") != std::string::npos) { - LOG_DEBUG("using x0 prediction"); - flux_params.chroma_radiance_params.use_x0 = true; - } - if (tensor_name.find("__32x32__") != std::string::npos) { - LOG_DEBUG("using patch size 32"); - flux_params.patch_size = 32; - } - if (tensor_name.find("img_in_patch.weight") != std::string::npos) { - actual_radiance_patch_size = pair.second.ne[0]; - LOG_DEBUG("actual radiance patch size: %d", actual_radiance_patch_size); - } - if (tensor_name.find("distilled_guidance_layer.in_proj.weight") != std::string::npos) { - // Chroma - flux_params.is_chroma = true; - } - size_t db = tensor_name.find("double_blocks."); - if (db != std::string::npos) { - tensor_name = tensor_name.substr(db); // remove prefix - int block_depth = atoi(tensor_name.substr(14, tensor_name.find(".", 14)).c_str()); - if (block_depth + 1 > flux_params.depth) { - flux_params.depth = block_depth + 1; - } - } - size_t sb = tensor_name.find("single_blocks."); - if (sb != std::string::npos) { - tensor_name = tensor_name.substr(sb); // remove prefix - int block_depth = atoi(tensor_name.substr(14, tensor_name.find(".", 14)).c_str()); - if (block_depth + 1 > flux_params.depth_single_blocks) { - flux_params.depth_single_blocks = block_depth + 1; - } - } - if (ends_with(tensor_name, "txt_in.weight")) { - flux_params.context_in_dim = pair.second.ne[0]; - flux_params.hidden_size = pair.second.ne[1]; - } - if (ends_with(tensor_name, "single_blocks.0.norm.key_norm.scale")) { - head_dim = pair.second.ne[0]; - } - if (ends_with(tensor_name, "double_blocks.0.txt_attn.norm.key_norm.scale")) { - head_dim = pair.second.ne[0]; - } - } - if (actual_radiance_patch_size > 0 && actual_radiance_patch_size != flux_params.patch_size) { - GGML_ASSERT(flux_params.patch_size == 2 * actual_radiance_patch_size); - LOG_DEBUG("using fake x2 patch size"); - flux_params.chroma_radiance_params.fake_patch_size_x2 = true; - } - - flux_params.num_heads = static_cast(flux_params.hidden_size / head_dim); - - LOG_INFO("flux: depth = %d, depth_single_blocks = %d, guidance_embed = %s, context_in_dim = %" PRId64 - ", hidden_size = %" PRId64 ", num_heads = %d", - flux_params.depth, - flux_params.depth_single_blocks, - flux_params.guidance_embed ? "true" : "false", - flux_params.context_in_dim, - flux_params.hidden_size, - flux_params.num_heads); - if (flux_params.is_chroma) { + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "", + SDVersion version = VERSION_FLUX, + bool use_mask = false, + std::shared_ptr weight_manager = nullptr) + : DiffusionModelRunner(backend, prefix, weight_manager), + config(FluxConfig::detect_from_weights(tensor_storage_map, prefix, version)), + version(version), + use_mask(use_mask) { + if (config.is_chroma) { LOG_INFO("Using pruned modulation (Chroma)"); } - flux = Flux(flux_params); + flux = Flux(config); flux.init(params_ctx, tensor_storage_map, prefix); } @@ -1368,19 +1483,21 @@ namespace Flux { const sd::Tensor& context_tensor = {}, const sd::Tensor& c_concat_tensor = {}, const sd::Tensor& y_tensor = {}, - const sd::Tensor& guidance_tensor = {}, - const std::vector>& ref_latents_tensor = {}, - Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED, - std::vector skip_layers = {}) { + const sd::Tensor& guidance_tensor = {}, + const std::vector>& ref_latents_tensor = {}, + Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED, + std::vector skip_layers = {}, + const sd::Tensor& pulid_id_tensor = {}, + float pulid_id_weight = 1.0f) { ggml_tensor* x = make_input(x_tensor); ggml_tensor* timesteps = make_input(timesteps_tensor); ggml_tensor* context = make_optional_input(context_tensor); ggml_tensor* c_concat = make_optional_input(c_concat_tensor); ggml_tensor* y = make_optional_input(y_tensor); - if (flux_params.guidance_embed || flux_params.is_chroma) { + if (config.guidance_embed || config.is_chroma) { if (!guidance_tensor.empty()) { this->guidance_tensor = guidance_tensor; - if (flux_params.is_chroma) { + if (config.is_chroma) { this->guidance_tensor.fill_(0.f); } } @@ -1398,7 +1515,7 @@ namespace Flux { ggml_tensor* mod_index_arange = nullptr; ggml_tensor* dct = nullptr; // for chroma radiance - if (flux_params.is_chroma) { + if (config.is_chroma) { if (!use_mask) { y = nullptr; } @@ -1409,7 +1526,7 @@ namespace Flux { set_backend_tensor_data(mod_index_arange, mod_index_arange_vec.data()); } std::set txt_arange_dims; - if (sd_version_is_flux2(version)) { + if (sd_version_is_flux2(version) || sd_version_is_sefi_image(version)) { txt_arange_dims = {3}; ref_index_mode = Rope::RefIndexMode::INCREASE; } else if (version == VERSION_OVIS_IMAGE) { @@ -1417,29 +1534,29 @@ namespace Flux { } pe_vec = Rope::gen_flux_pe(static_cast(x->ne[1]), static_cast(x->ne[0]), - flux_params.patch_size, + config.patch_size, static_cast(x->ne[3]), static_cast(context->ne[1]), txt_arange_dims, ref_latents, ref_index_mode, - flux_params.ref_index_scale, - flux_params.theta, + config.ref_index_scale, + config.theta, circular_y_enabled, circular_x_enabled, - flux_params.axes_dim, + config.axes_dim, sd_version_is_longcat(version)); - int pos_len = static_cast(pe_vec.size() / flux_params.axes_dim_sum / 2); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); // LOG_DEBUG("pos_len %d", pos_len); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, flux_params.axes_dim_sum / 2, pos_len); + auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); // pe->data = pe_vec.data(); // print_ggml_tensor(pe); // pe->data = nullptr; set_backend_tensor_data(pe, pe_vec.data()); if (version == VERSION_CHROMA_RADIANCE) { - int patch_size = flux_params.patch_size; - int nerf_max_freqs = flux_params.chroma_radiance_params.nerf_max_freqs; + int patch_size = config.patch_size; + int nerf_max_freqs = config.chroma_radiance_params.nerf_max_freqs; dct_vec = fetch_dct_pos(patch_size, nerf_max_freqs); dct = ggml_new_tensor_2d(compute_ctx, GGML_TYPE_F32, nerf_max_freqs * nerf_max_freqs, patch_size * patch_size); // dct->data = dct_vec.data(); @@ -1448,6 +1565,10 @@ namespace Flux { set_backend_tensor_data(dct, dct_vec.data()); } + ggml_tensor* pulid_id = pulid_id_tensor.empty() + ? nullptr + : make_input(pulid_id_tensor); + auto runner_ctx = get_context(); ggml_tensor* out = flux.forward(&runner_ctx, @@ -1461,7 +1582,9 @@ namespace Flux { mod_index_arange, dct, ref_latents, - skip_layers); + skip_layers, + pulid_id, + pulid_id_weight); ggml_build_forward_expand(gf, out); @@ -1477,17 +1600,20 @@ namespace Flux { const sd::Tensor& guidance = {}, const std::vector>& ref_latents = {}, Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED, - std::vector skip_layers = std::vector()) { + std::vector skip_layers = std::vector(), + const sd::Tensor& pulid_id = {}, + float pulid_id_weight = 1.0f) { // x: [N, in_channels, h, w] // timesteps: [N, ] // context: [N, max_position, hidden_size] // y: [N, adm_in_channels] or [1, adm_in_channels] // guidance: [N, ] + // pulid_id: empty (no injection) or [N, num_id_tokens=32, kv_dim=2048] auto get_graph = [&]() -> ggml_cgraph* { - return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, ref_index_mode, skip_layers); + return build_graph(x, timesteps, context, c_concat, y, guidance, ref_latents, ref_index_mode, skip_layers, pulid_id, pulid_id_weight); }; - auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); + auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); return result; } @@ -1507,7 +1633,9 @@ namespace Flux { tensor_or_empty(extra->guidance), diffusion_params.ref_latents ? *diffusion_params.ref_latents : empty_ref_latents, diffusion_params.ref_index_mode, - extra->skip_layers ? *extra->skip_layers : empty_skip_layers); + extra->skip_layers ? *extra->skip_layers : empty_skip_layers, + tensor_or_empty(extra->pulid_id), + extra->pulid_id_weight); } void test() { @@ -1570,7 +1698,8 @@ namespace Flux { ggml_backend_t backend = sd_backend_cpu_init(); ggml_type model_data_type = GGML_TYPE_COUNT; - ModelLoader model_loader; + auto model_manager = std::make_shared(); + ModelLoader& model_loader = model_manager->loader(); if (!model_loader.init_from_file_and_convert_name(file_path, "model.diffusion_model.")) { LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str()); return; @@ -1586,24 +1715,20 @@ namespace Flux { } std::shared_ptr flux = std::make_shared(backend, - backend, tensor_storage_map, "model.diffusion_model", VERSION_FLUX2, - false); + false, + model_manager); - if (!flux->alloc_params_buffer()) { - LOG_ERROR("flux model allocation failed"); - return; - } - - std::map tensors; - flux->get_param_tensors(tensors, "model.diffusion_model"); - - bool success = model_loader.load_tensors(tensors); - - if (!success) { - LOG_ERROR("load tensors from model loader failed"); + if (!model_manager->register_runner_params("Flux test", + *flux, + "model.diffusion_model", + ModelManager::ResidencyMode::ParamBackend, + backend, + backend) || + !model_manager->validate_registered_tensors()) { + LOG_ERROR("register flux tensors with model manager failed"); return; } @@ -1614,4 +1739,4 @@ namespace Flux { } // namespace Flux -#endif // __FLUX_HPP__ +#endif // __SD_MODEL_DIFFUSION_FLUX_HPP__ diff --git a/src/hidream_o1.hpp b/src/model/diffusion/hidream_o1.hpp similarity index 88% rename from src/hidream_o1.hpp rename to src/model/diffusion/hidream_o1.hpp index c85e04b9..9d3df039 100644 --- a/src/hidream_o1.hpp +++ b/src/model/diffusion/hidream_o1.hpp @@ -1,5 +1,5 @@ -#ifndef __SD_HIDREAM_O1_H__ -#define __SD_HIDREAM_O1_H__ +#ifndef __SD_MODEL_DIFFUSION_HIDREAM_O1_HPP__ +#define __SD_MODEL_DIFFUSION_HIDREAM_O1_HPP__ #include #include @@ -10,11 +10,11 @@ #include #include -#include "common_dit.hpp" -#include "conditioner.hpp" -#include "diffusion_model.hpp" -#include "llm.hpp" -#include "util.h" +#include "conditioning/conditioner.hpp" +#include "core/util.h" +#include "model/diffusion/dit.hpp" +#include "model/diffusion/model.hpp" +#include "model/te/llm.hpp" namespace HiDreamO1 { constexpr int HIDREAM_O1_GRAPH_SIZE = 32768; @@ -23,6 +23,39 @@ namespace HiDreamO1 { constexpr int IMAGE_TOKEN_ID = 151655; constexpr int VISION_START_TOKEN_ID = 151652; + struct HiDreamO1Config { + LLM::LLMConfig llm; + int patch_size = PATCH_SIZE; + + static HiDreamO1Config detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) { + (void)tensor_storage_map; + (void)prefix; + HiDreamO1Config config; + config.llm.arch = LLM::LLMArch::QWEN3_VL; + config.llm.hidden_size = 4096; + config.llm.intermediate_size = 12288; + config.llm.num_layers = 36; + config.llm.num_heads = 32; + config.llm.num_kv_heads = 8; + config.llm.head_dim = 128; + config.llm.qkv_bias = false; + config.llm.qk_norm = true; + config.llm.vocab_size = 151936; + config.llm.rms_norm_eps = 1e-6f; + config.llm.vision.arch = LLM::LLMVisionArch::QWEN3_VL; + config.llm.vision.num_layers = 27; + config.llm.vision.hidden_size = 1152; + config.llm.vision.intermediate_size = 4304; + config.llm.vision.num_heads = 16; + config.llm.vision.out_hidden_size = 4096; + config.llm.vision.patch_size = 16; + config.llm.vision.spatial_merge_size = 2; + config.llm.vision.temporal_patch_size = 2; + config.llm.vision.num_position_embeddings = 2304; + return config; + } + }; + static inline std::string repeat_special_token(const std::string& token, int64_t count) { std::string out; out.reserve(static_cast(count) * token.size()); @@ -205,50 +238,19 @@ namespace HiDreamO1 { } }; - struct HiDreamO1Params { - LLM::LLMParams llm; - int patch_size = PATCH_SIZE; - }; - - static inline HiDreamO1Params make_hidream_o1_params() { - HiDreamO1Params params; - params.llm.arch = LLM::LLMArch::QWEN3_VL; - params.llm.hidden_size = 4096; - params.llm.intermediate_size = 12288; - params.llm.num_layers = 36; - params.llm.num_heads = 32; - params.llm.num_kv_heads = 8; - params.llm.head_dim = 128; - params.llm.qkv_bias = false; - params.llm.qk_norm = true; - params.llm.vocab_size = 151936; - params.llm.rms_norm_eps = 1e-6f; - params.llm.vision.arch = LLM::LLMVisionArch::QWEN3_VL; - params.llm.vision.num_layers = 27; - params.llm.vision.hidden_size = 1152; - params.llm.vision.intermediate_size = 4304; - params.llm.vision.num_heads = 16; - params.llm.vision.out_hidden_size = 4096; - params.llm.vision.patch_size = 16; - params.llm.vision.spatial_merge_size = 2; - params.llm.vision.temporal_patch_size = 2; - params.llm.vision.num_position_embeddings = 2304; - return params; - } - struct HiDreamO1Model : public GGMLBlock { - HiDreamO1Params params; + HiDreamO1Config config; HiDreamO1Model() = default; - explicit HiDreamO1Model(HiDreamO1Params params) - : params(std::move(params)) { - blocks["language_model"] = std::make_shared(this->params.llm); - blocks["t_embedder1"] = std::make_shared(this->params.llm.hidden_size); - blocks["x_embedder"] = std::make_shared(this->params.patch_size * this->params.patch_size * 3, - this->params.llm.hidden_size / 4, - this->params.llm.hidden_size); - blocks["final_layer2"] = std::make_shared(this->params.llm.hidden_size, - this->params.patch_size * this->params.patch_size * 3); + explicit HiDreamO1Model(HiDreamO1Config config) + : config(std::move(config)) { + blocks["language_model"] = std::make_shared(this->config.llm); + blocks["t_embedder1"] = std::make_shared(this->config.llm.hidden_size); + blocks["x_embedder"] = std::make_shared(this->config.patch_size * this->config.patch_size * 3, + this->config.llm.hidden_size / 4, + this->config.llm.hidden_size); + blocks["final_layer2"] = std::make_shared(this->config.llm.hidden_size, + this->config.patch_size * this->config.patch_size * 3); } std::shared_ptr text_model() { @@ -269,7 +271,7 @@ namespace HiDreamO1 { }; struct HiDreamO1VisionRunner : public GGMLRunner { - HiDreamO1Params params; + HiDreamO1Config config; std::shared_ptr model; std::vector window_index_vec; @@ -280,12 +282,12 @@ namespace HiDreamO1 { std::array, 4> pos_embed_weight_data_; HiDreamO1VisionRunner(ggml_backend_t backend, - ggml_backend_t params_backend, - const String2TensorStorage& tensor_storage_map = {}, - const std::string& prefix = "model.visual") - : GGMLRunner(backend, params_backend), - params(make_hidream_o1_params()), - model(std::make_shared(false, params.llm.vision)) { + const String2TensorStorage& tensor_storage_map = {}, + const std::string& prefix = "model.visual", + std::shared_ptr weight_manager = nullptr) + : GGMLRunner(backend, weight_manager), + config(HiDreamO1Config::detect_from_weights(tensor_storage_map, prefix)), + model(std::make_shared(false, config.llm.vision)) { model->init(params_ctx, tensor_storage_map, prefix); } @@ -302,7 +304,7 @@ namespace HiDreamO1 { compute_ctx, runner_ctx, image, - params.llm.vision, + config.llm.vision, model, window_index_vec, window_inverse_index_vec, @@ -321,28 +323,32 @@ namespace HiDreamO1 { return gf; } - sd::Tensor compute(int n_threads, const sd::Tensor& image) { + sd::Tensor compute(int n_threads, + const sd::Tensor& image, + bool auto_free = true, + bool free_compute_buffer = true, + bool free_compute_params = true) { auto get_graph = [&]() { return build_graph(image); }; - auto output = GGMLRunner::compute(get_graph, n_threads, false); + auto output = GGMLRunner::compute(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params); return output.has_value() ? std::move(output.value()) : sd::Tensor(); } }; struct HiDreamO1Runner : public DiffusionModelRunner { - HiDreamO1Params params; + HiDreamO1Config config; HiDreamO1Model model; std::vector attention_mask_vec; HiDreamO1Runner(ggml_backend_t backend, - ggml_backend_t params_backend, - const String2TensorStorage& tensor_storage_map = {}, - const std::string& prefix = "model") - : DiffusionModelRunner(backend, params_backend, prefix), - params(make_hidream_o1_params()) { - model = HiDreamO1Model(params); + const String2TensorStorage& tensor_storage_map = {}, + const std::string& prefix = "model", + std::shared_ptr weight_manager = nullptr) + : DiffusionModelRunner(backend, prefix, weight_manager), + config(HiDreamO1Config::detect_from_weights(tensor_storage_map, prefix)) { + model = HiDreamO1Model(config); model.init(params_ctx, tensor_storage_map, prefix); } @@ -453,7 +459,7 @@ namespace HiDreamO1 { auto get_graph = [&]() { return build_graph(x, timestep, input_ids, input_pos, token_types, vinput_mask, image_embeds, ref_images); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); } sd::Tensor compute(int n_threads, @@ -484,29 +490,14 @@ namespace HiDreamO1 { std::shared_ptr vision_runner; HiDreamO1Conditioner(ggml_backend_t backend, - ggml_backend_t params_backend, - const String2TensorStorage& tensor_storage_map = {}) - : vision_runner(std::make_shared(backend, params_backend, tensor_storage_map)) {} + const String2TensorStorage& tensor_storage_map = {}, + std::shared_ptr weight_manager = nullptr) + : vision_runner(std::make_shared(backend, tensor_storage_map, "model.visual", weight_manager)) {} void get_param_tensors(std::map& tensors) override { vision_runner->get_param_tensors(tensors); } - bool alloc_params_buffer() override { - if (!vision_runner->alloc_params_buffer()) { - return false; - } - return true; - } - - void free_params_buffer() override { - vision_runner->free_params_buffer(); - } - - size_t get_params_buffer_size() override { - return vision_runner->get_params_buffer_size(); - } - void set_max_graph_vram_bytes(size_t max_graph_vram_bytes) override { vision_runner->set_max_graph_vram_bytes(max_graph_vram_bytes); } @@ -519,6 +510,10 @@ namespace HiDreamO1 { vision_runner->set_weight_adapter(adapter); } + void runner_done() override { + vision_runner->runner_done(); + } + SDCondition get_learned_condition(int n_threads, const ConditionerParams& conditioner_params) override { SDCondition result; @@ -664,7 +659,7 @@ namespace HiDreamO1 { result.c_vinput_mask = sd::Tensor(vinput_mask_shape, std::move(vinput_mask)); result.c_image_embeds.reserve(vlm_images.size()); for (const auto& vlm_image : vlm_images) { - auto image_embed = vision_runner->compute(n_threads, vlm_image.second); + auto image_embed = vision_runner->compute(n_threads, vlm_image.second, false, true, true); if (image_embed.empty()) { LOG_ERROR("hidream_o1 conditioner: encode VLM image failed"); return SDCondition(); @@ -676,4 +671,4 @@ namespace HiDreamO1 { }; } // namespace HiDreamO1 -#endif // __SD_HIDREAM_O1_H__ +#endif // __SD_MODEL_DIFFUSION_HIDREAM_O1_HPP__ diff --git a/src/model/diffusion/ideogram4.hpp b/src/model/diffusion/ideogram4.hpp new file mode 100644 index 00000000..bfa2f86a --- /dev/null +++ b/src/model/diffusion/ideogram4.hpp @@ -0,0 +1,557 @@ +#ifndef __SD_MODEL_DIFFUSION_IDEOGRAM4_HPP__ +#define __SD_MODEL_DIFFUSION_IDEOGRAM4_HPP__ + +#include +#include +#include +#include +#include +#include + +#include "core/ggml_extend.hpp" +#include "core/ggml_graph_cut.h" +#include "model/common/rope.hpp" +#include "model/diffusion/model.hpp" + +namespace Ideogram4 { + constexpr int IDEOGRAM4_GRAPH_SIZE = 65536; + constexpr int OUTPUT_IMAGE_INDICATOR = 2; + constexpr int IMAGE_POSITION_OFFSET = 65536; + constexpr int DEFAULT_MROPE_SECTION_T = 24; + constexpr int DEFAULT_MROPE_SECTION_H = 20; + constexpr int DEFAULT_MROPE_SECTION_W = 20; + constexpr int TIMESTEP_MAX_PERIOD = 10000; + constexpr int LLM_HIDDEN_STATE_LAYERS = 13; + + struct Ideogram4Config { + int64_t emb_dim = 4608; + int64_t num_layers = 34; + int64_t num_heads = 18; + int64_t intermediate_size = 12288; + int64_t adanln_dim = 512; + int64_t in_channels = 128; + int64_t llm_features_dim = 53248; + int64_t rope_theta = 5000000; + float norm_eps = 1e-5f; + int patch_size = 2; + int ae_channels = 32; + std::vector mrope_section = {DEFAULT_MROPE_SECTION_T, + DEFAULT_MROPE_SECTION_H, + DEFAULT_MROPE_SECTION_W}; + + static Ideogram4Config detect_from_weights(const String2TensorStorage& tensor_storage_map, + const std::string& prefix) { + Ideogram4Config config; + int64_t detected_layers = 0; + std::string layer_prefix = prefix.empty() ? "layers." : prefix + ".layers."; + for (const auto& [name, _] : tensor_storage_map) { + if (name.find(layer_prefix) != 0) { + continue; + } + std::string tail = name.substr(layer_prefix.size()); + size_t dot = tail.find('.'); + if (dot == std::string::npos) { + continue; + } + int layer_idx = std::atoi(tail.substr(0, dot).c_str()); + detected_layers = std::max(detected_layers, layer_idx + 1); + } + if (detected_layers > 0) { + config.num_layers = detected_layers; + LOG_DEBUG("ideogram4: num_layers = %" PRId64 ", emb_dim = %" PRId64 ", num_heads = %" PRId64 ", intermediate_size = %" PRId64, + config.num_layers, + config.emb_dim, + config.num_heads, + config.intermediate_size); + } + return config; + } + }; + + __STATIC_INLINE__ ggml_tensor* timestep_embedding_sin_cos(ggml_context* ctx, + ggml_tensor* timesteps, + int dim) { + GGML_ASSERT(dim % 2 == 0); + auto embedding = ggml_ext_timestep_embedding(ctx, timesteps, dim, TIMESTEP_MAX_PERIOD, 10.f); + auto chunks = ggml_ext_chunk(ctx, embedding, 2, 0); + return ggml_concat(ctx, chunks[1], chunks[0], 0); + } + + __STATIC_INLINE__ ggml_tensor* to_token_modulation(ggml_context* ctx, ggml_tensor* x) { + // [N, C] -> [N, 1, C] in PyTorch layout. + if (ggml_n_dims(x) < 3 || x->ne[1] != 1) { + x = ggml_reshape_3d(ctx, x, x->ne[0], 1, x->ne[1]); + } + return x; + } + + __STATIC_INLINE__ ggml_tensor* interleave_hidden_state_layers(ggml_context* ctx, ggml_tensor* x) { + // Match upstream stack(...).permute(1, 2, 3, 0).reshape(...): + // [layers * hidden, tokens, batch] -> [hidden * layers, tokens, batch]. + GGML_ASSERT(x->ne[0] % LLM_HIDDEN_STATE_LAYERS == 0); + const int64_t hidden_size = x->ne[0] / LLM_HIDDEN_STATE_LAYERS; + const int64_t token_count = x->ne[1]; + const int64_t batch_count = x->ne[2]; + + x = ggml_reshape_4d(ctx, x, hidden_size, LLM_HIDDEN_STATE_LAYERS, token_count, batch_count); + x = ggml_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3)); + return ggml_reshape_3d(ctx, x, hidden_size * LLM_HIDDEN_STATE_LAYERS, token_count, batch_count); + } + + __STATIC_INLINE__ ggml_tensor* modulate(ggml_context* ctx, ggml_tensor* x, ggml_tensor* scale) { + scale = to_token_modulation(ctx, scale); + return ggml_add(ctx, x, ggml_mul(ctx, x, scale)); + } + + __STATIC_INLINE__ ggml_tensor* patchify(ggml_context* ctx, ggml_tensor* x, const Ideogram4Config& config) { + // x: [N, 128, H, W] with channel order [ae, ph, pw]. + // return: [N, H*W, 128] with token channel order [ph, pw, ae]. + const int64_t W = x->ne[0]; + const int64_t H = x->ne[1]; + const int64_t C = x->ne[2]; + const int64_t N = x->ne[3]; + + GGML_ASSERT(N == 1); + GGML_ASSERT(C == config.ae_channels * config.patch_size * config.patch_size); + + x = ggml_cont(ctx, x); + x = ggml_reshape_4d(ctx, x, W * H, config.patch_size, config.patch_size, config.ae_channels); + x = ggml_cont(ctx, ggml_permute(ctx, x, 3, 1, 2, 0)); + x = ggml_reshape_3d(ctx, x, C, W * H, N); + return x; + } + + __STATIC_INLINE__ ggml_tensor* unpatchify(ggml_context* ctx, + ggml_tensor* x, + int64_t H, + int64_t W, + const Ideogram4Config& config) { + const int64_t C = x->ne[0]; + const int64_t N = x->ne[2]; + + GGML_ASSERT(N == 1); + GGML_ASSERT(C == config.ae_channels * config.patch_size * config.patch_size); + GGML_ASSERT(x->ne[1] == H * W); + + x = ggml_reshape_4d(ctx, x, config.ae_channels, config.patch_size, config.patch_size, H * W); + x = ggml_cont(ctx, ggml_permute(ctx, x, 3, 1, 2, 0)); + x = ggml_reshape_4d(ctx, x, W, H, C, N); + return x; + } + + __STATIC_INLINE__ std::shared_ptr make_linear(int64_t in_features, + int64_t out_features, + bool bias = true) { + return std::make_shared(in_features, out_features, bias, false, false, 1.f, true); + } + + __STATIC_INLINE__ std::vector gen_ideogram4_pe(int grid_h, + int grid_w, + int bs, + int context_len, + int head_dim, + int rope_theta, + const std::vector& mrope_section, + bool circular_x = false, + bool circular_y = false) { + GGML_ASSERT(bs == 1); + std::vector> ids(static_cast(bs) * (context_len + grid_h * grid_w), + std::vector(3, 0.f)); + + for (int i = 0; i < context_len; ++i) { + ids[i] = {static_cast(i), static_cast(i), static_cast(i)}; + } + + int cursor = context_len; + for (int y = 0; y < grid_h; ++y) { + for (int x = 0; x < grid_w; ++x) { + ids[cursor++] = {static_cast(IMAGE_POSITION_OFFSET), + static_cast(IMAGE_POSITION_OFFSET + y), + static_cast(IMAGE_POSITION_OFFSET + x)}; + } + } + + std::vector> axis_wrap_dims(3); + if (circular_y || circular_x) { + size_t total_len = static_cast(bs) * (context_len + grid_h * grid_w); + axis_wrap_dims[1].assign(total_len, 0); + axis_wrap_dims[2].assign(total_len, 0); + if (circular_y) { + for (size_t idx = static_cast(context_len); idx < total_len; ++idx) { + axis_wrap_dims[1][idx] = grid_h; + } + } + if (circular_x) { + for (size_t idx = static_cast(context_len); idx < total_len; ++idx) { + axis_wrap_dims[2][idx] = grid_w; + } + } + } + + return Rope::embed_interleaved_mrope(ids, + bs, + static_cast(rope_theta), + head_dim, + mrope_section, + axis_wrap_dims); + } + + class Ideogram4Attention : public GGMLBlock { + protected: + int64_t hidden_size; + int64_t num_heads; + int64_t head_dim; + + public: + Ideogram4Attention(int64_t hidden_size, int64_t num_heads, float eps) + : hidden_size(hidden_size), num_heads(num_heads), head_dim(hidden_size / num_heads) { + GGML_ASSERT(hidden_size % num_heads == 0); + blocks["qkv"] = make_linear(hidden_size, hidden_size * 3, false); + blocks["norm_q"] = std::make_shared(head_dim, eps); + blocks["norm_k"] = std::make_shared(head_dim, eps); + blocks["o"] = make_linear(hidden_size, hidden_size, false); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* pe, + ggml_tensor* mask = nullptr) { + int64_t n_token = x->ne[1]; + int64_t N = x->ne[2]; + + auto qkv_proj = std::dynamic_pointer_cast(blocks["qkv"]); + auto norm_q = std::dynamic_pointer_cast(blocks["norm_q"]); + auto norm_k = std::dynamic_pointer_cast(blocks["norm_k"]); + auto out_proj = std::dynamic_pointer_cast(blocks["o"]); + + auto qkv = qkv_proj->forward(ctx, x); + auto qkv_vec = split_qkv(ctx->ggml_ctx, qkv); + auto q = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[0], head_dim, num_heads, n_token, N); + auto k = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[1], head_dim, num_heads, n_token, N); + auto v = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[2], head_dim, num_heads, n_token, N); + + q = norm_q->forward(ctx, q); + k = norm_k->forward(ctx, k); + + x = Rope::attention(ctx, q, k, v, pe, mask, 1.f / 128.f, false); + x = out_proj->forward(ctx, x); + return x; + } + }; + + class Ideogram4MLP : public GGMLBlock { + public: + Ideogram4MLP(int64_t dim, int64_t hidden_dim) { + blocks["w1"] = make_linear(dim, hidden_dim, false); + blocks["w2"] = make_linear(hidden_dim, dim, false); + blocks["w3"] = make_linear(dim, hidden_dim, false); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto w1 = std::dynamic_pointer_cast(blocks["w1"]); + auto w2 = std::dynamic_pointer_cast(blocks["w2"]); + auto w3 = std::dynamic_pointer_cast(blocks["w3"]); + + auto x1 = ggml_silu(ctx->ggml_ctx, w1->forward(ctx, x)); + auto x3 = w3->forward(ctx, x); + x = ggml_mul(ctx->ggml_ctx, x1, x3); + x = w2->forward(ctx, x); + return x; + } + }; + + class Ideogram4TransformerBlock : public GGMLBlock { + public: + Ideogram4TransformerBlock(const Ideogram4Config& config) { + blocks["attention"] = std::make_shared(config.emb_dim, config.num_heads, config.norm_eps); + blocks["feed_forward"] = std::make_shared(config.emb_dim, config.intermediate_size); + blocks["attention_norm1"] = std::make_shared(config.emb_dim, config.norm_eps); + blocks["ffn_norm1"] = std::make_shared(config.emb_dim, config.norm_eps); + blocks["attention_norm2"] = std::make_shared(config.emb_dim, config.norm_eps); + blocks["ffn_norm2"] = std::make_shared(config.emb_dim, config.norm_eps); + blocks["adaln_modulation"] = make_linear(config.adanln_dim, 4 * config.emb_dim, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* pe, + ggml_tensor* adaln_input, + ggml_tensor* mask = nullptr) { + auto attention = std::dynamic_pointer_cast(blocks["attention"]); + auto feed_forward = std::dynamic_pointer_cast(blocks["feed_forward"]); + auto attention_norm1 = std::dynamic_pointer_cast(blocks["attention_norm1"]); + auto ffn_norm1 = std::dynamic_pointer_cast(blocks["ffn_norm1"]); + auto attention_norm2 = std::dynamic_pointer_cast(blocks["attention_norm2"]); + auto ffn_norm2 = std::dynamic_pointer_cast(blocks["ffn_norm2"]); + auto adaln_modulation = std::dynamic_pointer_cast(blocks["adaln_modulation"]); + + auto mod = adaln_modulation->forward(ctx, adaln_input); + auto mods = ggml_ext_chunk(ctx->ggml_ctx, mod, 4, 0); + auto scale_msa = mods[0]; + auto gate_msa = to_token_modulation(ctx->ggml_ctx, ggml_tanh(ctx->ggml_ctx, mods[1])); + auto scale_mlp = mods[2]; + auto gate_mlp = to_token_modulation(ctx->ggml_ctx, ggml_tanh(ctx->ggml_ctx, mods[3])); + + auto attn_out = attention_norm1->forward(ctx, x); + attn_out = modulate(ctx->ggml_ctx, attn_out, scale_msa); + attn_out = attention->forward(ctx, attn_out, pe, mask); + attn_out = attention_norm2->forward(ctx, attn_out); + x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out, gate_msa)); + + auto ffn_out = ffn_norm1->forward(ctx, x); + ffn_out = modulate(ctx->ggml_ctx, ffn_out, scale_mlp); + ffn_out = feed_forward->forward(ctx, ffn_out); + ffn_out = ffn_norm2->forward(ctx, ffn_out); + x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, ffn_out, gate_mlp)); + + return x; + } + }; + + class Ideogram4EmbedScalar : public GGMLBlock { + protected: + int64_t dim; + + public: + Ideogram4EmbedScalar(int64_t dim) + : dim(dim) { + blocks["mlp_in"] = make_linear(dim, dim, true); + blocks["mlp_out"] = make_linear(dim, dim, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto mlp_in = std::dynamic_pointer_cast(blocks["mlp_in"]); + auto mlp_out = std::dynamic_pointer_cast(blocks["mlp_out"]); + + x = timestep_embedding_sin_cos(ctx->ggml_ctx, x, static_cast(dim)); + x = ggml_silu(ctx->ggml_ctx, mlp_in->forward(ctx, x)); + x = mlp_out->forward(ctx, x); + return x; + } + }; + + class Ideogram4FinalLayer : public GGMLBlock { + public: + Ideogram4FinalLayer(const Ideogram4Config& config) { + blocks["norm_final"] = std::make_shared(config.emb_dim, 1e-6f, false); + blocks["linear"] = make_linear(config.emb_dim, config.in_channels, true); + blocks["adaln_modulation"] = make_linear(config.adanln_dim, config.emb_dim, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* c) { + auto norm_final = std::dynamic_pointer_cast(blocks["norm_final"]); + auto linear = std::dynamic_pointer_cast(blocks["linear"]); + auto adaln_modulation = std::dynamic_pointer_cast(blocks["adaln_modulation"]); + + auto scale = adaln_modulation->forward(ctx, ggml_silu(ctx->ggml_ctx, c)); + x = norm_final->forward(ctx, x); + x = modulate(ctx->ggml_ctx, x, scale); + x = linear->forward(ctx, x); + return x; + } + }; + + class Ideogram4Transformer : public GGMLBlock { + protected: + Ideogram4Config config; + + public: + Ideogram4Transformer() = default; + explicit Ideogram4Transformer(Ideogram4Config config) + : config(std::move(config)) { + blocks["input_proj"] = make_linear(this->config.in_channels, this->config.emb_dim, true); + blocks["llm_cond_norm"] = std::make_shared(this->config.llm_features_dim, 1e-6f); + blocks["llm_cond_proj"] = make_linear(this->config.llm_features_dim, this->config.emb_dim, true); + blocks["t_embedding"] = std::make_shared(this->config.emb_dim); + blocks["adaln_proj"] = make_linear(this->config.emb_dim, this->config.adanln_dim, true); + blocks["embed_image_indicator"] = std::make_shared(2, this->config.emb_dim); + + for (int i = 0; i < this->config.num_layers; ++i) { + blocks["layers." + std::to_string(i)] = std::make_shared(this->config); + } + blocks["final_layer"] = std::make_shared(this->config); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* timestep, + ggml_tensor* context, + ggml_tensor* pe, + ggml_tensor* image_indicator_ids) { + int64_t W = x->ne[0]; + int64_t H = x->ne[1]; + int64_t N = x->ne[3]; + GGML_ASSERT(N == 1); + + auto input_proj = std::dynamic_pointer_cast(blocks["input_proj"]); + auto llm_cond_norm = std::dynamic_pointer_cast(blocks["llm_cond_norm"]); + auto llm_cond_proj = std::dynamic_pointer_cast(blocks["llm_cond_proj"]); + auto t_embedding = std::dynamic_pointer_cast(blocks["t_embedding"]); + auto adaln_proj = std::dynamic_pointer_cast(blocks["adaln_proj"]); + auto embed_image_indicator = std::dynamic_pointer_cast(blocks["embed_image_indicator"]); + auto final_layer = std::dynamic_pointer_cast(blocks["final_layer"]); + + auto img = patchify(ctx->ggml_ctx, x, config); + img = input_proj->forward(ctx, img); + + ggml_tensor* h = img; + int64_t context_len = 0; + if (context != nullptr) { + if (ggml_n_dims(context) < 3) { + context = ggml_reshape_3d(ctx->ggml_ctx, context, context->ne[0], context->ne[1], 1); + } + context = interleave_hidden_state_layers(ctx->ggml_ctx, context); + context_len = context->ne[1]; + auto txt = llm_cond_norm->forward(ctx, context); + txt = llm_cond_proj->forward(ctx, txt); + h = ggml_concat(ctx->ggml_ctx, txt, img, 1); + } + + auto indicator_embedding = embed_image_indicator->forward(ctx, image_indicator_ids); + h = ggml_add(ctx->ggml_ctx, h, indicator_embedding); + + auto t_cond = t_embedding->forward(ctx, timestep); + auto adaln_input = ggml_silu(ctx->ggml_ctx, adaln_proj->forward(ctx, t_cond)); + + for (int i = 0; i < config.num_layers; ++i) { + auto block = std::dynamic_pointer_cast(blocks["layers." + std::to_string(i)]); + h = block->forward(ctx, h, pe, adaln_input, nullptr); + sd::ggml_graph_cut::mark_graph_cut(h, "ideogram4.layers." + std::to_string(i), "hidden"); + } + + h = final_layer->forward(ctx, h, adaln_input); + if (context_len > 0) { + h = ggml_ext_slice(ctx->ggml_ctx, h, 1, context_len, h->ne[1]); + } + + h = unpatchify(ctx->ggml_ctx, h, H, W, config); + h = ggml_ext_scale(ctx->ggml_ctx, h, -1.f); + return h; + } + }; + + class Ideogram4Runner : public DiffusionModelRunner { + protected: + bool should_use_uncond_model(const DiffusionParams& diffusion_params) const { + return has_uncond_model && + diffusion_params.context == nullptr && + diffusion_params.y != nullptr && + !diffusion_params.y->empty(); + } + + public: + Ideogram4Config config; + Ideogram4Transformer model; + Ideogram4Transformer uncond_model; + bool has_uncond_model = false; + std::string uncond_prefix; + std::vector pe_vec; + std::vector image_indicator_vec; + + Ideogram4Runner(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "", + std::shared_ptr weight_manager = nullptr) + : DiffusionModelRunner(backend, prefix, weight_manager), + config(Ideogram4Config::detect_from_weights(tensor_storage_map, prefix)), + uncond_prefix(prefix + ".uncond") { + model = Ideogram4Transformer(config); + model.init(params_ctx, tensor_storage_map, prefix); + for (const auto& pair : tensor_storage_map) { + const std::string& name = pair.first; + if (starts_with(name, uncond_prefix)) { + has_uncond_model = true; + break; + } + } + if (has_uncond_model) { + LOG_DEBUG("using uncond model"); + uncond_model = Ideogram4Transformer(config); + uncond_model.init(params_ctx, tensor_storage_map, uncond_prefix); + } + } + + std::string get_desc() override { + return "ideogram4"; + } + + void get_param_tensors(std::map& tensors, const std::string& prefix) override { + model.get_param_tensors(tensors, prefix); + if (has_uncond_model) { + uncond_model.get_param_tensors(tensors, this->uncond_prefix); + } + } + + ggml_cgraph* build_graph(const sd::Tensor& x_tensor, + const sd::Tensor& timesteps_tensor, + const sd::Tensor& context_tensor, + bool use_uncond_model = false) { + ggml_cgraph* gf = new_graph_custom(IDEOGRAM4_GRAPH_SIZE); + ggml_tensor* x = make_input(x_tensor); + ggml_tensor* timesteps = make_input(timesteps_tensor); + GGML_ASSERT(x->ne[3] == 1); + Ideogram4Transformer& active_model = use_uncond_model ? uncond_model : model; + + ggml_tensor* context = nullptr; + int64_t context_len = 0; + if (!context_tensor.empty()) { + context = make_input(context_tensor); + context_len = context->ne[1]; + } + + int64_t grid_w = x->ne[0]; + int64_t grid_h = x->ne[1]; + int64_t pos_len = context_len + grid_h * grid_w; + int64_t head_dim = config.emb_dim / config.num_heads; + + auto runner_ctx = get_context(); + pe_vec = gen_ideogram4_pe(static_cast(grid_h), + static_cast(grid_w), + static_cast(x->ne[3]), + static_cast(context_len), + static_cast(head_dim), + static_cast(config.rope_theta), + config.mrope_section, + runner_ctx.circular_x_enabled, + runner_ctx.circular_y_enabled); + auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, head_dim / 2, pos_len); + set_backend_tensor_data(pe, pe_vec.data()); + + image_indicator_vec.assign(static_cast(pos_len), 1); + for (int64_t i = 0; i < context_len; ++i) { + image_indicator_vec[static_cast(i)] = 0; + } + auto indicator = ggml_new_tensor_2d(compute_ctx, GGML_TYPE_I32, pos_len, x->ne[3]); + set_backend_tensor_data(indicator, image_indicator_vec.data()); + + ggml_tensor* out = active_model.forward(&runner_ctx, x, timesteps, context, pe, indicator); + ggml_build_forward_expand(gf, out); + return gf; + } + + sd::Tensor compute(int n_threads, + const sd::Tensor& x, + const sd::Tensor& timesteps, + const sd::Tensor& context, + bool use_uncond_model = false) { + auto get_graph = [&]() -> ggml_cgraph* { + return build_graph(x, timesteps, context, use_uncond_model); + }; + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + } + + sd::Tensor compute(int n_threads, + const DiffusionParams& diffusion_params) override { + GGML_ASSERT(diffusion_params.x != nullptr); + GGML_ASSERT(diffusion_params.timesteps != nullptr); + bool use_uncond_model = should_use_uncond_model(diffusion_params); + return compute(n_threads, + *diffusion_params.x, + *diffusion_params.timesteps, + tensor_or_empty(diffusion_params.context), + use_uncond_model); + } + }; +} // namespace Ideogram4 + +#endif // __SD_MODEL_DIFFUSION_IDEOGRAM4_HPP__ diff --git a/src/model/diffusion/krea2.hpp b/src/model/diffusion/krea2.hpp new file mode 100644 index 00000000..02e65559 --- /dev/null +++ b/src/model/diffusion/krea2.hpp @@ -0,0 +1,683 @@ +#ifndef __SD_MODEL_DIFFUSION_KREA2_HPP__ +#define __SD_MODEL_DIFFUSION_KREA2_HPP__ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "core/ggml_extend.hpp" +#include "core/ggml_graph_cut.h" +#include "model/common/rope.hpp" +#include "model/diffusion/dit.hpp" +#include "model/diffusion/flux.hpp" +#include "model/diffusion/model.hpp" +#include "model_loader.h" + +namespace Krea2 { + constexpr int KREA2_GRAPH_SIZE = 65536; + + struct Krea2Config { + int patch_size = 2; + int64_t in_channels = 16; + int64_t out_channels = 16; + int64_t features = 6144; + int64_t timestep_dim = 256; + int64_t text_dim = 2560; + int64_t text_layers = 12; + int64_t layers = 28; + int64_t heads = 48; + int64_t kv_heads = 12; + int64_t text_heads = 20; + int64_t text_kv_heads = 20; + int64_t mlp_multiplier = 4; + float theta = 1000.f; + float norm_eps = 1e-5f; + std::vector axes_dim = {32, 48, 48}; + int axes_dim_sum = 128; + + int64_t head_dim() const { + return features / heads; + } + + static int64_t count_blocks(const String2TensorStorage& tensor_storage_map, + const std::string& prefix, + const std::string& block_prefix) { + int64_t count = 0; + std::string full_prefix = prefix.empty() ? block_prefix : prefix + "." + block_prefix; + for (const auto& [name, _] : tensor_storage_map) { + if (!starts_with(name, full_prefix)) { + continue; + } + std::string tail = name.substr(full_prefix.size()); + size_t dot = tail.find('.'); + if (dot == std::string::npos) { + continue; + } + int block_index = std::atoi(tail.substr(0, dot).c_str()); + count = std::max(count, block_index + 1); + } + return count; + } + + void update_axes_dim() { + int64_t dim_head = head_dim(); + int64_t unit = dim_head / 16; + axes_dim = { + static_cast(dim_head - 12 * unit), + static_cast(6 * unit), + static_cast(6 * unit), + }; + axes_dim_sum = axes_dim[0] + axes_dim[1] + axes_dim[2]; + } + + static Krea2Config detect_from_weights(const String2TensorStorage& tensor_storage_map, + const std::string& prefix) { + Krea2Config config; + int64_t detected_head_dim = 0; + int64_t detected_text_head_dim = 0; + + for (const auto& [name, tensor_storage] : tensor_storage_map) { + if (!starts_with(name, prefix)) { + continue; + } + if (ends_with(name, "first.weight") && tensor_storage.n_dims == 2) { + config.in_channels = tensor_storage.ne[0] / (config.patch_size * config.patch_size); + config.out_channels = config.in_channels; + config.features = tensor_storage.ne[1]; + } else if (ends_with(name, "blocks.0.attn.qknorm.qnorm.scale") && tensor_storage.n_dims == 1) { + detected_head_dim = tensor_storage.ne[0]; + } else if (ends_with(name, "blocks.0.attn.wq.weight") && tensor_storage.n_dims == 2) { + if (detected_head_dim > 0) { + config.heads = tensor_storage.ne[1] / detected_head_dim; + } + } else if (ends_with(name, "blocks.0.attn.wk.weight") && tensor_storage.n_dims == 2) { + if (detected_head_dim > 0) { + config.kv_heads = tensor_storage.ne[1] / detected_head_dim; + } + } else if (ends_with(name, "txtfusion.projector.weight") && tensor_storage.n_dims == 2) { + config.text_layers = tensor_storage.ne[0]; + } else if (ends_with(name, "txtfusion.layerwise_blocks.0.prenorm.scale") && tensor_storage.n_dims == 1) { + config.text_dim = tensor_storage.ne[0]; + } else if (ends_with(name, "txtfusion.layerwise_blocks.0.attn.qknorm.qnorm.scale") && tensor_storage.n_dims == 1) { + detected_text_head_dim = tensor_storage.ne[0]; + } else if (ends_with(name, "txtfusion.layerwise_blocks.0.attn.wq.weight") && tensor_storage.n_dims == 2) { + if (detected_text_head_dim > 0) { + config.text_heads = tensor_storage.ne[1] / detected_text_head_dim; + } + } else if (ends_with(name, "txtfusion.layerwise_blocks.0.attn.wk.weight") && tensor_storage.n_dims == 2) { + if (detected_text_head_dim > 0) { + config.text_kv_heads = tensor_storage.ne[1] / detected_text_head_dim; + } + } else if (ends_with(name, "last.linear.weight") && tensor_storage.n_dims == 2) { + config.out_channels = tensor_storage.ne[1] / (config.patch_size * config.patch_size); + } + } + + config.layers = std::max(1, count_blocks(tensor_storage_map, prefix, "blocks.")); + if (detected_head_dim > 0 && config.features > 0) { + config.heads = config.features / detected_head_dim; + } + if (detected_head_dim > 0) { + std::string wk_name = prefix.empty() ? "blocks.0.attn.wk.weight" : prefix + ".blocks.0.attn.wk.weight"; + auto it = tensor_storage_map.find(wk_name); + if (it != tensor_storage_map.end() && it->second.n_dims == 2) { + config.kv_heads = it->second.ne[1] / detected_head_dim; + } + } + if (detected_text_head_dim > 0 && config.text_dim > 0) { + config.text_heads = config.text_dim / detected_text_head_dim; + } + if (detected_text_head_dim > 0) { + std::string wk_name = prefix.empty() ? "txtfusion.layerwise_blocks.0.attn.wk.weight" : prefix + ".txtfusion.layerwise_blocks.0.attn.wk.weight"; + auto it = tensor_storage_map.find(wk_name); + if (it != tensor_storage_map.end() && it->second.n_dims == 2) { + config.text_kv_heads = it->second.ne[1] / detected_text_head_dim; + } + } + config.update_axes_dim(); + + LOG_DEBUG("krea2: layers=%" PRId64 ", features=%" PRId64 ", heads=%" PRId64 ", kv_heads=%" PRId64 ", text_dim=%" PRId64 ", text_layers=%" PRId64 ", text_heads=%" PRId64 ", text_kv_heads=%" PRId64 ", channels=%" PRId64, + config.layers, + config.features, + config.heads, + config.kv_heads, + config.text_dim, + config.text_layers, + config.text_heads, + config.text_kv_heads, + config.in_channels); + return config; + } + }; + + __STATIC_INLINE__ int64_t ceil_to_multiple(int64_t value, int64_t multiple) { + return ((value + multiple - 1) / multiple) * multiple; + } + + class KreaRMSNorm : public UnaryBlock { + protected: + int64_t hidden_size; + float eps; + std::string prefix; + + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { + GGML_UNUSED(tensor_storage_map); + this->prefix = prefix; + params["scale"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size); + } + + public: + KreaRMSNorm(int64_t hidden_size, float eps = 1e-5f) + : hidden_size(hidden_size), + eps(eps) {} + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + ggml_tensor* scale = params["scale"]; + scale = ggml_add(ctx->ggml_ctx, scale, ggml_ext_ones(ctx->ggml_ctx, scale->ne[0], 1, 1, 1)); + x = ggml_rms_norm(ctx->ggml_ctx, x, eps); + x = ggml_mul_inplace(ctx->ggml_ctx, x, scale); + return x; + } + }; + + class KreaSwiGLU : public UnaryBlock { + public: + KreaSwiGLU(int64_t features, int64_t multiplier) { + int64_t mlp_dim = ceil_to_multiple(((2 * features) / 3) * multiplier, 128); + blocks["gate"] = std::make_shared(features, mlp_dim, false); + blocks["up"] = std::make_shared(features, mlp_dim, false); + blocks["down"] = std::make_shared(mlp_dim, features, false); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto gate = std::dynamic_pointer_cast(blocks["gate"]); + auto up = std::dynamic_pointer_cast(blocks["up"]); + auto down = std::dynamic_pointer_cast(blocks["down"]); + + auto gated = ggml_silu(ctx->ggml_ctx, gate->forward(ctx, x)); + auto up_x = up->forward(ctx, x); + x = ggml_mul(ctx->ggml_ctx, gated, up_x); + return down->forward(ctx, x); + } + }; + + class KreaAttention : public GGMLBlock { + protected: + int64_t features; + int64_t heads; + int64_t kv_heads; + int64_t head_dim_; + + ggml_tensor* attention_no_rope(GGMLRunnerContext* ctx, + ggml_tensor* q, + ggml_tensor* k, + ggml_tensor* v, + ggml_tensor* mask) { + int64_t Lq = q->ne[2]; + int64_t Lk = k->ne[2]; + int64_t N = q->ne[3]; + q = ggml_reshape_3d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, q), head_dim_ * heads, Lq, N); + k = ggml_reshape_3d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, k), head_dim_ * kv_heads, Lk, N); + v = ggml_reshape_3d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, v), head_dim_ * kv_heads, Lk, N); + return ggml_ext_attention_ext(ctx->ggml_ctx, + ctx->backend, + q, + k, + v, + heads, + mask, + false, + ctx->flash_attn_enabled); + } + + public: + KreaAttention(int64_t features, + int64_t heads, + int64_t kv_heads, + float eps = 1e-5f) + : features(features), + heads(heads), + kv_heads(kv_heads), + head_dim_(features / heads) { + blocks["wq"] = std::make_shared(features, heads * head_dim_, false); + blocks["wk"] = std::make_shared(features, kv_heads * head_dim_, false); + blocks["wv"] = std::make_shared(features, kv_heads * head_dim_, false); + blocks["gate"] = std::make_shared(features, features, false); + blocks["qknorm.qnorm"] = std::make_shared(head_dim_, eps); + blocks["qknorm.knorm"] = std::make_shared(head_dim_, eps); + blocks["wo"] = std::make_shared(features, features, false); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* pe = nullptr, + ggml_tensor* mask = nullptr) { + auto wq = std::dynamic_pointer_cast(blocks["wq"]); + auto wk = std::dynamic_pointer_cast(blocks["wk"]); + auto wv = std::dynamic_pointer_cast(blocks["wv"]); + auto gate = std::dynamic_pointer_cast(blocks["gate"]); + auto qnorm = std::dynamic_pointer_cast(blocks["qknorm.qnorm"]); + auto knorm = std::dynamic_pointer_cast(blocks["qknorm.knorm"]); + auto wo = std::dynamic_pointer_cast(blocks["wo"]); + + if (sd_backend_is(ctx->backend, "Vulkan")) { + wo->set_force_prec_f32(true); + } + + int64_t L = x->ne[1]; + int64_t N = x->ne[2]; + + auto q = wq->forward(ctx, x); + q = ggml_reshape_4d(ctx->ggml_ctx, q, head_dim_, heads, L, N); + auto k = wk->forward(ctx, x); + k = ggml_reshape_4d(ctx->ggml_ctx, k, head_dim_, kv_heads, L, N); + auto v = wv->forward(ctx, x); + v = ggml_reshape_4d(ctx->ggml_ctx, v, head_dim_, kv_heads, L, N); + + q = qnorm->forward(ctx, q); + k = knorm->forward(ctx, k); + + auto out = pe != nullptr ? Rope::attention(ctx, q, k, v, pe, mask) + : attention_no_rope(ctx, q, k, v, mask); + out = ggml_mul(ctx->ggml_ctx, out, ggml_sigmoid(ctx->ggml_ctx, gate->forward(ctx, x))); + out = wo->forward(ctx, out); + return out; + } + }; + + class KreaDoubleSharedModulation : public GGMLBlock { + protected: + int64_t dim; + + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { + GGML_UNUSED(tensor_storage_map); + GGML_UNUSED(prefix); + params["lin"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dim * 6); + } + + public: + KreaDoubleSharedModulation(int64_t dim) + : dim(dim) {} + + std::vector forward(GGMLRunnerContext* ctx, ggml_tensor* vec) { + auto lin = ggml_repeat(ctx->ggml_ctx, params["lin"], vec); + auto out = ggml_add(ctx->ggml_ctx, vec, lin); + return ggml_ext_chunk(ctx->ggml_ctx, out, 6, 0); + } + }; + + class KreaFinalModulation : public GGMLBlock { + protected: + int64_t dim; + + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { + GGML_UNUSED(tensor_storage_map); + GGML_UNUSED(prefix); + params["lin"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, dim, 2); + } + + public: + KreaFinalModulation(int64_t dim) + : dim(dim) {} + + std::vector forward(GGMLRunnerContext* ctx, ggml_tensor* vec) { + auto out = ggml_add(ctx->ggml_ctx, params["lin"], vec); + return ggml_ext_chunk(ctx->ggml_ctx, out, 2, 1); + } + }; + + class KreaTextFusionBlock : public UnaryBlock { + public: + KreaTextFusionBlock(int64_t dim, + int64_t heads, + int64_t kv_heads, + int64_t multiplier, + float eps) { + blocks["prenorm"] = std::make_shared(dim, eps); + blocks["postnorm"] = std::make_shared(dim, eps); + blocks["attn"] = std::make_shared(dim, heads, kv_heads, eps); + blocks["mlp"] = std::make_shared(dim, multiplier); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto prenorm = std::dynamic_pointer_cast(blocks["prenorm"]); + auto postnorm = std::dynamic_pointer_cast(blocks["postnorm"]); + auto attn = std::dynamic_pointer_cast(blocks["attn"]); + auto mlp = std::dynamic_pointer_cast(blocks["mlp"]); + + x = ggml_add(ctx->ggml_ctx, x, attn->forward(ctx, prenorm->forward(ctx, x))); + x = ggml_add(ctx->ggml_ctx, x, mlp->forward(ctx, postnorm->forward(ctx, x))); + return x; + } + }; + + class KreaTextFusionTransformer : public UnaryBlock { + protected: + Krea2Config config; + + public: + explicit KreaTextFusionTransformer(Krea2Config config) + : config(std::move(config)) { + for (int i = 0; i < 2; ++i) { + blocks["layerwise_blocks." + std::to_string(i)] = std::make_shared(this->config.text_dim, + this->config.text_heads, + this->config.text_kv_heads, + this->config.mlp_multiplier, + this->config.norm_eps); + blocks["refiner_blocks." + std::to_string(i)] = std::make_shared(this->config.text_dim, + this->config.text_heads, + this->config.text_kv_heads, + this->config.mlp_multiplier, + this->config.norm_eps); + } + blocks["projector"] = std::make_shared(this->config.text_layers, 1, false); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* context) override { + int64_t text_tokens = context->ne[1]; + int64_t batch = context->ne[2]; + + context = ggml_reshape_3d(ctx->ggml_ctx, + context, + config.text_dim, + config.text_layers, + text_tokens * batch); + + for (int i = 0; i < 2; ++i) { + auto block = std::dynamic_pointer_cast(blocks["layerwise_blocks." + std::to_string(i)]); + context = block->forward(ctx, context); + } + + context = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, context, 1, 0, 2, 3)); + auto projector = std::dynamic_pointer_cast(blocks["projector"]); + context = projector->forward(ctx, context); + context = ggml_reshape_3d(ctx->ggml_ctx, context, config.text_dim, text_tokens, batch); + + for (int i = 0; i < 2; ++i) { + auto block = std::dynamic_pointer_cast(blocks["refiner_blocks." + std::to_string(i)]); + context = block->forward(ctx, context); + } + return context; + } + }; + + class KreaSingleStreamBlock : public UnaryBlock { + public: + explicit KreaSingleStreamBlock(Krea2Config config) { + blocks["mod"] = std::make_shared(config.features); + blocks["prenorm"] = std::make_shared(config.features, config.norm_eps); + blocks["postnorm"] = std::make_shared(config.features, config.norm_eps); + blocks["attn"] = std::make_shared(config.features, config.heads, config.kv_heads, config.norm_eps); + blocks["mlp"] = std::make_shared(config.features, config.mlp_multiplier); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* vec, + ggml_tensor* pe) { + auto mod = std::dynamic_pointer_cast(blocks["mod"]); + auto prenorm = std::dynamic_pointer_cast(blocks["prenorm"]); + auto postnorm = std::dynamic_pointer_cast(blocks["postnorm"]); + auto attn = std::dynamic_pointer_cast(blocks["attn"]); + auto mlp = std::dynamic_pointer_cast(blocks["mlp"]); + + auto mods = mod->forward(ctx, vec); + auto attn_input = Flux::modulate(ctx->ggml_ctx, + prenorm->forward(ctx, x), + mods[1], + mods[0], + true); + auto attn_out = attn->forward(ctx, attn_input, pe); + x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out, mods[2])); + + auto mlp_input = Flux::modulate(ctx->ggml_ctx, + postnorm->forward(ctx, x), + mods[4], + mods[3], + true); + auto mlp_out = mlp->forward(ctx, mlp_input); + x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, mlp_out, mods[5])); + return x; + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + GGML_UNUSED(ctx); + GGML_UNUSED(x); + GGML_ABORT("KreaSingleStreamBlock requires conditioning"); + return nullptr; + } + }; + + class KreaTimeMLP : public UnaryBlock { + public: + explicit KreaTimeMLP(Krea2Config config) { + blocks["0"] = std::make_shared(config.timestep_dim, config.features, true); + blocks["2"] = std::make_shared(config.features, config.features, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto linear_0 = std::dynamic_pointer_cast(blocks["0"]); + auto linear_2 = std::dynamic_pointer_cast(blocks["2"]); + x = linear_0->forward(ctx, x); + x = ggml_ext_gelu(ctx->ggml_ctx, x, false); + x = linear_2->forward(ctx, x); + return x; + } + }; + + class KreaTProj : public UnaryBlock { + public: + explicit KreaTProj(Krea2Config config) { + blocks["1"] = std::make_shared(config.features, config.features * 6, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto linear_1 = std::dynamic_pointer_cast(blocks["1"]); + x = ggml_ext_gelu(ctx->ggml_ctx, x, false); + x = linear_1->forward(ctx, x); + return x; + } + }; + + class KreaTextMLP : public UnaryBlock { + public: + explicit KreaTextMLP(Krea2Config config) { + blocks["0"] = std::make_shared(config.text_dim, config.norm_eps); + blocks["1"] = std::make_shared(config.text_dim, config.features, true); + blocks["3"] = std::make_shared(config.features, config.features, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto norm = std::dynamic_pointer_cast(blocks["0"]); + auto linear_1 = std::dynamic_pointer_cast(blocks["1"]); + auto linear_3 = std::dynamic_pointer_cast(blocks["3"]); + x = norm->forward(ctx, x); + x = linear_1->forward(ctx, x); + x = ggml_ext_gelu(ctx->ggml_ctx, x, true); + x = linear_3->forward(ctx, x); + return x; + } + }; + + class KreaLastLayer : public GGMLBlock { + public: + explicit KreaLastLayer(Krea2Config config) { + blocks["norm"] = std::make_shared(config.features, config.norm_eps); + blocks["linear"] = std::make_shared(config.features, config.patch_size * config.patch_size * config.out_channels, true); + blocks["modulation"] = std::make_shared(config.features); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* vec) { + auto norm = std::dynamic_pointer_cast(blocks["norm"]); + auto linear = std::dynamic_pointer_cast(blocks["linear"]); + auto modulation = std::dynamic_pointer_cast(blocks["modulation"]); + + auto mods = modulation->forward(ctx, vec); + x = Flux::modulate(ctx->ggml_ctx, + norm->forward(ctx, x), + mods[1], + mods[0], + true); + x = linear->forward(ctx, x); + return x; + } + }; + + class Krea2Model : public GGMLBlock { + protected: + Krea2Config config; + + public: + Krea2Model() = default; + explicit Krea2Model(Krea2Config config) + : config(std::move(config)) { + blocks["first"] = std::make_shared(this->config.patch_size * this->config.patch_size * this->config.in_channels, + this->config.features, + true); + blocks["tmlp"] = std::make_shared(this->config); + blocks["txtfusion"] = std::make_shared(this->config); + blocks["txtmlp"] = std::make_shared(this->config); + blocks["tproj"] = std::make_shared(this->config); + for (int i = 0; i < this->config.layers; ++i) { + blocks["blocks." + std::to_string(i)] = std::make_shared(this->config); + } + blocks["last"] = std::make_shared(this->config); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* timestep, + ggml_tensor* context, + ggml_tensor* pe) { + int64_t W = x->ne[0]; + int64_t H = x->ne[1]; + int64_t N = x->ne[3]; + GGML_ASSERT(N == 1); + + auto first = std::dynamic_pointer_cast(blocks["first"]); + auto tmlp = std::dynamic_pointer_cast(blocks["tmlp"]); + auto txtfusion = std::dynamic_pointer_cast(blocks["txtfusion"]); + auto txtmlp = std::dynamic_pointer_cast(blocks["txtmlp"]); + auto tproj = std::dynamic_pointer_cast(blocks["tproj"]); + auto last = std::dynamic_pointer_cast(blocks["last"]); + + auto img = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size, true); + int64_t img_len = img->ne[1]; + img = first->forward(ctx, img); + + auto t = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast(config.timestep_dim), 10000, 1000.f); + t = tmlp->forward(ctx, t); + t = ggml_reshape_3d(ctx->ggml_ctx, t, t->ne[0], 1, t->ne[1]); + auto tvec = tproj->forward(ctx, t); + + auto txt = txtfusion->forward(ctx, context); + txt = txtmlp->forward(ctx, txt); + int64_t txt_len = txt->ne[1]; + + auto hidden_states = ggml_concat(ctx->ggml_ctx, txt, img, 1); + for (int i = 0; i < config.layers; ++i) { + auto block = std::dynamic_pointer_cast(blocks["blocks." + std::to_string(i)]); + hidden_states = block->forward(ctx, hidden_states, tvec, pe); + sd::ggml_graph_cut::mark_graph_cut(hidden_states, "krea2.blocks." + std::to_string(i), "hidden_states"); + } + + hidden_states = last->forward(ctx, hidden_states, t); + hidden_states = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, txt_len, txt_len + img_len); + hidden_states = DiT::unpatchify_and_crop(ctx->ggml_ctx, hidden_states, H, W, config.patch_size, config.patch_size, true); + return hidden_states; + } + }; + + __STATIC_INLINE__ std::vector gen_krea2_pe(int h, + int w, + int patch_size, + int bs, + int context_len, + float theta, + const std::vector& axes_dim) { + auto txt_ids = Rope::gen_flux_txt_ids(bs, context_len, 3, {}); + auto img_ids = Rope::gen_flux_img_ids(h, w, patch_size, bs, 3, 0, 0, 0, false); + auto ids = Rope::concat_ids(txt_ids, img_ids, bs); + return Rope::embed_nd(ids, bs, theta, axes_dim); + } + + struct Krea2Runner : public DiffusionModelRunner { + Krea2Config config; + Krea2Model model; + std::vector pe_vec; + + Krea2Runner(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "", + std::shared_ptr weight_manager = nullptr) + : DiffusionModelRunner(backend, prefix, weight_manager), + config(Krea2Config::detect_from_weights(tensor_storage_map, prefix)) { + model = Krea2Model(config); + model.init(params_ctx, tensor_storage_map, prefix); + } + + std::string get_desc() override { + return "krea2"; + } + + void get_param_tensors(std::map& tensors, const std::string& prefix) override { + model.get_param_tensors(tensors, prefix); + } + + ggml_cgraph* build_graph(const sd::Tensor& x_tensor, + const sd::Tensor& timesteps_tensor, + const sd::Tensor& context_tensor) { + ggml_cgraph* gf = new_graph_custom(KREA2_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 = gen_krea2_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + config.theta, + config.axes_dim); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); + set_backend_tensor_data(pe, pe_vec.data()); + + auto runner_ctx = get_context(); + ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pe); + ggml_build_forward_expand(gf, out); + return gf; + } + + sd::Tensor compute(int n_threads, + const sd::Tensor& x, + const sd::Tensor& timesteps, + const sd::Tensor& context) { + auto get_graph = [&]() -> ggml_cgraph* { + return build_graph(x, timesteps, context); + }; + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + } + + sd::Tensor compute(int n_threads, + const DiffusionParams& diffusion_params) override { + GGML_ASSERT(diffusion_params.x != nullptr); + GGML_ASSERT(diffusion_params.timesteps != nullptr); + return compute(n_threads, + *diffusion_params.x, + *diffusion_params.timesteps, + tensor_or_empty(diffusion_params.context)); + } + }; +} // namespace Krea2 + +#endif // __SD_MODEL_DIFFUSION_KREA2_HPP__ diff --git a/src/lens.hpp b/src/model/diffusion/lens.hpp similarity index 83% rename from src/lens.hpp rename to src/model/diffusion/lens.hpp index b5ff0683..931a8527 100644 --- a/src/lens.hpp +++ b/src/model/diffusion/lens.hpp @@ -1,18 +1,83 @@ -#ifndef __SD_LENS_HPP__ -#define __SD_LENS_HPP__ +#ifndef __SD_MODEL_DIFFUSION_LENS_HPP__ +#define __SD_MODEL_DIFFUSION_LENS_HPP__ #include #include -#include "common_block.hpp" -#include "diffusion_model.hpp" -#include "flux.hpp" -#include "qwen_image.hpp" -#include "rope.hpp" +#include "model/common/block.hpp" +#include "model/common/rope.hpp" +#include "model/diffusion/flux.hpp" +#include "model/diffusion/model.hpp" +#include "model/diffusion/qwen_image.hpp" namespace Lens { constexpr int LENS_GRAPH_SIZE = 40960; + struct LensConfig { + 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 axes_dim = {8, 28, 28}; + int axes_dim_sum = 64; + + static LensConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) { + LensConfig config; + config.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) { + config.in_channels = tensor_storage.ne[0]; + int64_t inner_dim = tensor_storage.ne[1]; + if (config.attention_head_dim > 0) { + config.num_attention_heads = inner_dim / config.attention_head_dim; + } + } else if (ends_with(name, "txt_in.weight") && tensor_storage.n_dims == 2) { + config.selected_layer_count = static_cast(tensor_storage.ne[0] / config.joint_attention_dim); + } else if (ends_with(name, "proj_out.weight") && tensor_storage.n_dims == 2) { + int64_t patch_area = config.patch_size * config.patch_size; + config.out_channels = tensor_storage.ne[1] / patch_area; + } else if (ends_with(name, "transformer_blocks.0.attn.norm_q.weight") && tensor_storage.n_dims == 1) { + config.attention_head_dim = tensor_storage.ne[0]; + } + + size_t pos = name.find("transformer_blocks."); + if (pos != std::string::npos) { + auto items = split_string(name.substr(pos), '.'); + if (items.size() > 1) { + int block_index = atoi(items[1].c_str()); + if (block_index + 1 > config.num_layers) { + config.num_layers = block_index + 1; + } + } + } + } + if (config.num_layers == 0) { + config.num_layers = 48; + } + config.axes_dim_sum = 0; + for (int axis_dim : config.axes_dim) { + config.axes_dim_sum += axis_dim; + } + LOG_DEBUG("lens: num_layers = %d, selected_layer_count = %d, hidden_size = %" PRId64 ", num_attention_heads = %" PRId64 ", attention_head_dim = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64, + config.num_layers, + config.selected_layer_count, + config.num_attention_heads * config.attention_head_dim, + config.num_attention_heads, + config.attention_head_dim, + config.in_channels, + config.out_channels); + return config; + } + }; + struct LensTimestepProjEmbeddings : public GGMLBlock { LensTimestepProjEmbeddings(int64_t embedding_dim) { blocks["timestep_embedder"] = std::make_shared(256, embedding_dim); @@ -209,41 +274,27 @@ namespace Lens { } }; - 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 axes_dim = {8, 28, 28}; - int axes_dim_sum = 64; - }; - class LensModel : public GGMLBlock { public: - LensParams params; + LensConfig config; LensModel() = default; - LensModel(LensParams params) - : params(params) { - int64_t inner_dim = params.num_attention_heads * params.attention_head_dim; + LensModel(LensConfig config) + : config(config) { + int64_t inner_dim = config.num_attention_heads * config.attention_head_dim; blocks["time_text_embed"] = std::make_shared(inner_dim); - blocks["img_in"] = std::make_shared(params.in_channels, inner_dim, true); - blocks["txt_in"] = std::make_shared(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(params.joint_attention_dim, 1e-5f); + blocks["img_in"] = std::make_shared(config.in_channels, inner_dim, true); + blocks["txt_in"] = std::make_shared(config.joint_attention_dim * config.selected_layer_count, inner_dim, true); + for (int i = 0; i < config.selected_layer_count; ++i) { + blocks["txt_norm." + std::to_string(i)] = std::make_shared(config.joint_attention_dim, 1e-5f); } - for (int i = 0; i < params.num_layers; ++i) { + for (int i = 0; i < config.num_layers; ++i) { blocks["transformer_blocks." + std::to_string(i)] = std::make_shared(inner_dim, - params.num_attention_heads, - params.attention_head_dim); + config.num_attention_heads, + config.attention_head_dim); } blocks["norm_out"] = std::make_shared(inner_dim, 1e-6f); - blocks["proj_out"] = std::make_shared(inner_dim, params.patch_size * params.patch_size * params.out_channels, true); + blocks["proj_out"] = std::make_shared(inner_dim, config.patch_size * config.patch_size * config.out_channels, true); } ggml_tensor* forward(GGMLRunnerContext* ctx, @@ -269,9 +320,9 @@ namespace Lens { 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 txt_chunks = ggml_ext_chunk(ctx->ggml_ctx, context, params.selected_layer_count, 0); + std::vector txt_chunks = ggml_ext_chunk(ctx->ggml_ctx, context, config.selected_layer_count, 0); ggml_tensor* txt = nullptr; - for (int i = 0; i < params.selected_layer_count; ++i) { + for (int i = 0; i < config.selected_layer_count; ++i) { auto txt_norm = std::dynamic_pointer_cast(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); @@ -281,7 +332,7 @@ namespace Lens { 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) { + for (int i = 0; i < config.num_layers; ++i) { auto block = std::dynamic_pointer_cast(blocks["transformer_blocks." + std::to_string(i)]); auto out = block->forward(ctx, img, txt, t_emb, pe); img = out.first; @@ -294,67 +345,23 @@ namespace Lens { 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); + out = ggml_reshape_4d(ctx->ggml_ctx, out, W, H, config.patch_size * config.patch_size * config.out_channels, N); return out; } }; struct LensRunner : public DiffusionModelRunner { - LensParams lens_params; + LensConfig config; LensModel lens; std::vector pe_vec; LensRunner(ggml_backend_t backend, - ggml_backend_t params_backend, - const String2TensorStorage& tensor_storage_map = {}, - const std::string prefix = "") - : DiffusionModelRunner(backend, params_backend, prefix) { - 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(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); + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "", + std::shared_ptr weight_manager = nullptr) + : DiffusionModelRunner(backend, prefix, weight_manager), + config(LensConfig::detect_from_weights(tensor_storage_map, prefix)) { + lens = LensModel(config); lens.init(params_ctx, tensor_storage_map, prefix); } @@ -380,12 +387,12 @@ namespace Lens { static_cast(x->ne[0]), static_cast(x->ne[3]), static_cast(context->ne[1]), - lens_params.theta, + config.theta, circular_y_enabled, circular_x_enabled, - lens_params.axes_dim); - int pos_len = static_cast(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); + config.axes_dim); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); set_backend_tensor_data(pe, pe_vec.data()); auto runner_ctx = get_context(); @@ -401,7 +408,7 @@ namespace Lens { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); } sd::Tensor compute(int n_threads, @@ -416,4 +423,4 @@ namespace Lens { }; } // namespace Lens -#endif // __SD_LENS_HPP__ +#endif // __SD_MODEL_DIFFUSION_LENS_HPP__ diff --git a/src/ltxv.hpp b/src/model/diffusion/ltxv.hpp similarity index 87% rename from src/ltxv.hpp rename to src/model/diffusion/ltxv.hpp index a7d3fb04..b89ff32c 100644 --- a/src/ltxv.hpp +++ b/src/model/diffusion/ltxv.hpp @@ -1,5 +1,5 @@ -#ifndef __SD_LTXV_HPP__ -#define __SD_LTXV_HPP__ +#ifndef __SD_MODEL_DIFFUSION_LTXV_HPP__ +#define __SD_MODEL_DIFFUSION_LTXV_HPP__ #include #include @@ -9,10 +9,11 @@ #include #include -#include "common_block.hpp" -#include "diffusion_model.hpp" -#include "flux.hpp" -#include "rope.hpp" +#include "model/common/block.hpp" +#include "model/common/rope.hpp" +#include "model/diffusion/flux.hpp" +#include "model/diffusion/model.hpp" +#include "model_loader.h" namespace LTXV { @@ -72,6 +73,200 @@ namespace LTXV { return max_block + 1; } + struct LTXAVConfig { + int64_t in_channels = 128; + int64_t out_channels = 128; + int64_t hidden_size = 3840; + int64_t cross_attention_dim = 4096; + int64_t caption_channels = 3840; + int64_t num_attention_heads = 30; + int64_t attention_head_dim = 128; + int64_t num_layers = 28; + float positional_embedding_theta = 10000.f; + std::vector positional_embedding_max_pos = {20, 2048, 2048}; + std::tuple vae_scale_factors = {8, 32, 32}; + bool causal_temporal_positioning = true; + float timestep_scale_multiplier = 1000.f; + + int64_t audio_in_channels = 128; + int64_t audio_out_channels = 128; + int64_t audio_hidden_size = 2048; + int64_t audio_cross_attention_dim = 2048; + int64_t audio_num_attention_heads = 32; + int64_t audio_attention_head_dim = 64; + std::vector audio_positional_embedding_max_pos = {20}; + float av_ca_timestep_scale_multiplier = 1000.f; + int64_t num_audio_channels = 8; + int64_t audio_frequency_bins = 16; + + bool use_connector = false; + int64_t connector_hidden_size = 3840; + int64_t connector_num_heads = 30; + int64_t connector_head_dim = 128; + int64_t connector_num_layers = 2; + int64_t connector_num_registers = 128; + bool connector_rope_interleaved = false; + bool connector_apply_gated_attention = false; + + bool use_audio_connector = false; + int64_t audio_connector_hidden_size = 2048; + int64_t audio_connector_num_heads = 32; + int64_t audio_connector_head_dim = 64; + int64_t audio_connector_num_layers = 2; + int64_t audio_connector_num_registers = 128; + bool audio_connector_rope_interleaved = false; + bool audio_connector_apply_gated_attention = false; + + bool video_rope_interleaved = false; + bool use_middle_indices_grid = true; + bool cross_attention_adaln = false; + + bool use_caption_projection = true; + bool use_audio_caption_projection = true; + bool caption_proj_before_connector = true; + bool caption_projection_first_linear = false; + + bool self_attention_gated = false; + bool cross_attention_gated = false; + + static std::pair infer_attention_layout(int64_t hidden_size, + int64_t preferred_heads = -1) { + if (preferred_heads > 0 && hidden_size % preferred_heads == 0) { + return {preferred_heads, hidden_size / preferred_heads}; + } + const int candidates[] = {128, 96, 80, 64, 48, 40, 32}; + for (int head_dim : candidates) { + if (hidden_size % head_dim == 0) { + int64_t heads = hidden_size / head_dim; + if (heads >= 8 && heads <= 64) { + return {heads, head_dim}; + } + } + } + return {32, hidden_size / 32}; + } + + static int64_t infer_gate_heads(const String2TensorStorage& tensor_storage_map, + const std::string& bias_name, + int64_t fallback_heads) { + auto it = tensor_storage_map.find(bias_name); + if (it != tensor_storage_map.end()) { + return it->second.ne[0]; + } + return fallback_heads; + } + + static LTXAVConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) { + LTXAVConfig config; + auto patchify_proj_iter = tensor_storage_map.find(prefix + ".patchify_proj.weight"); + if (patchify_proj_iter != tensor_storage_map.end()) { + config.in_channels = patchify_proj_iter->second.ne[0]; + config.hidden_size = patchify_proj_iter->second.ne[1]; + int64_t video_heads = infer_gate_heads(tensor_storage_map, prefix + ".transformer_blocks.0.attn1.to_gate_logits.bias", 32); + auto attn_layout = infer_attention_layout(config.hidden_size, video_heads); + config.num_attention_heads = attn_layout.first; + config.attention_head_dim = attn_layout.second; + } + + auto audio_patchify_proj_iter = tensor_storage_map.find(prefix + ".audio_patchify_proj.weight"); + if (audio_patchify_proj_iter != tensor_storage_map.end()) { + config.audio_in_channels = audio_patchify_proj_iter->second.ne[0]; + config.audio_hidden_size = audio_patchify_proj_iter->second.ne[1]; + config.audio_out_channels = config.audio_in_channels; + int64_t audio_heads = infer_gate_heads(tensor_storage_map, prefix + ".transformer_blocks.0.audio_attn1.to_gate_logits.bias", 32); + auto audio_attn_layout = infer_attention_layout(config.audio_hidden_size, audio_heads); + config.audio_num_attention_heads = audio_attn_layout.first; + config.audio_attention_head_dim = audio_attn_layout.second; + } + + auto proj_out_iter = tensor_storage_map.find(prefix + ".proj_out.weight"); + if (proj_out_iter != tensor_storage_map.end()) { + config.out_channels = proj_out_iter->second.ne[1]; + } + auto audio_proj_out_iter = tensor_storage_map.find(prefix + ".audio_proj_out.weight"); + if (audio_proj_out_iter != tensor_storage_map.end()) { + config.audio_out_channels = audio_proj_out_iter->second.ne[1]; + } + + auto attn2_iter = tensor_storage_map.find(prefix + ".transformer_blocks.0.attn2.to_k.weight"); + if (attn2_iter != tensor_storage_map.end()) { + config.cross_attention_dim = attn2_iter->second.ne[0]; + } + auto audio_attn2_iter = tensor_storage_map.find(prefix + ".transformer_blocks.0.audio_attn2.to_k.weight"); + if (audio_attn2_iter != tensor_storage_map.end()) { + config.audio_cross_attention_dim = audio_attn2_iter->second.ne[0]; + } + if (tensor_storage_map.find(prefix + ".transformer_blocks.0.prompt_scale_shift_table") != tensor_storage_map.end()) { + config.cross_attention_adaln = true; + } + if (tensor_storage_map.find(prefix + ".transformer_blocks.0.attn1.to_gate_logits.weight") != tensor_storage_map.end() || + tensor_storage_map.find(prefix + ".transformer_blocks.0.audio_attn1.to_gate_logits.weight") != tensor_storage_map.end()) { + config.self_attention_gated = true; + } + if (tensor_storage_map.find(prefix + ".transformer_blocks.0.attn2.to_gate_logits.weight") != tensor_storage_map.end() || + tensor_storage_map.find(prefix + ".transformer_blocks.0.audio_attn2.to_gate_logits.weight") != tensor_storage_map.end()) { + config.cross_attention_gated = true; + } + if (tensor_storage_map.find(prefix + ".caption_projection.linear_1.weight") == tensor_storage_map.end() && + tensor_storage_map.find(prefix + ".caption_projection.linear_2.weight") == tensor_storage_map.end()) { + config.use_caption_projection = false; + } + if (tensor_storage_map.find(prefix + ".audio_caption_projection.linear_1.weight") == tensor_storage_map.end() && + tensor_storage_map.find(prefix + ".audio_caption_projection.linear_2.weight") == tensor_storage_map.end()) { + config.use_audio_caption_projection = false; + } + + config.num_layers = count_prefix_blocks(tensor_storage_map, prefix + ".", "transformer_blocks."); + + auto connector_iter = tensor_storage_map.find(prefix + ".video_embeddings_connector.transformer_1d_blocks.0.attn1.to_q.weight"); + if (connector_iter != tensor_storage_map.end()) { + config.use_connector = true; + config.connector_hidden_size = connector_iter->second.ne[1]; + int64_t connector_heads = infer_gate_heads(tensor_storage_map, + prefix + ".video_embeddings_connector.transformer_1d_blocks.0.attn1.to_gate_logits.bias", + 32); + auto connector_layout = infer_attention_layout(config.connector_hidden_size, connector_heads); + config.connector_num_heads = connector_layout.first; + config.connector_head_dim = connector_layout.second; + config.connector_num_layers = count_prefix_blocks(tensor_storage_map, prefix + ".video_embeddings_connector.", "transformer_1d_blocks."); + auto register_iter = tensor_storage_map.find(prefix + ".video_embeddings_connector.learnable_registers"); + if (register_iter != tensor_storage_map.end()) { + config.connector_num_registers = register_iter->second.ne[1]; + } + if (tensor_storage_map.find(prefix + ".video_embeddings_connector.transformer_1d_blocks.0.attn1.to_gate_logits.weight") != tensor_storage_map.end()) { + config.connector_apply_gated_attention = true; + } + } + + auto audio_connector_iter = tensor_storage_map.find(prefix + ".audio_embeddings_connector.transformer_1d_blocks.0.attn1.to_q.weight"); + if (audio_connector_iter != tensor_storage_map.end()) { + config.use_audio_connector = true; + config.audio_connector_hidden_size = audio_connector_iter->second.ne[1]; + int64_t connector_heads = infer_gate_heads(tensor_storage_map, + prefix + ".audio_embeddings_connector.transformer_1d_blocks.0.attn1.to_gate_logits.bias", + 32); + auto connector_layout = infer_attention_layout(config.audio_connector_hidden_size, connector_heads); + config.audio_connector_num_heads = connector_layout.first; + config.audio_connector_head_dim = connector_layout.second; + config.audio_connector_num_layers = count_prefix_blocks(tensor_storage_map, prefix + ".audio_embeddings_connector.", "transformer_1d_blocks."); + auto register_iter = tensor_storage_map.find(prefix + ".audio_embeddings_connector.learnable_registers"); + if (register_iter != tensor_storage_map.end()) { + config.audio_connector_num_registers = register_iter->second.ne[1]; + } + if (tensor_storage_map.find(prefix + ".audio_embeddings_connector.transformer_1d_blocks.0.attn1.to_gate_logits.weight") != tensor_storage_map.end()) { + config.audio_connector_apply_gated_attention = true; + } + } + LOG_DEBUG("ltxav: num_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_attention_heads = %" PRId64 ", audio_hidden_size = %" PRId64 ", audio_num_attention_heads = %" PRId64, + config.num_layers, + config.hidden_size, + config.num_attention_heads, + config.audio_hidden_size, + config.audio_num_attention_heads); + return config; + } + }; + __STATIC_INLINE__ std::vector generate_freq_grid(float theta, int positional_dims, int dim) { @@ -749,63 +944,6 @@ namespace LTXV { } }; - struct LTXAVParams { - int64_t in_channels = 128; - int64_t out_channels = 128; - int64_t hidden_size = 3840; - int64_t cross_attention_dim = 4096; - int64_t caption_channels = 3840; - int64_t num_attention_heads = 30; - int64_t attention_head_dim = 128; - int64_t num_layers = 28; - float positional_embedding_theta = 10000.f; - std::vector positional_embedding_max_pos = {20, 2048, 2048}; - std::tuple vae_scale_factors = {8, 32, 32}; - bool causal_temporal_positioning = true; - float timestep_scale_multiplier = 1000.f; - - int64_t audio_in_channels = 128; - int64_t audio_out_channels = 128; - int64_t audio_hidden_size = 2048; - int64_t audio_cross_attention_dim = 2048; - int64_t audio_num_attention_heads = 32; - int64_t audio_attention_head_dim = 64; - std::vector audio_positional_embedding_max_pos = {20}; - float av_ca_timestep_scale_multiplier = 1000.f; - int64_t num_audio_channels = 8; - int64_t audio_frequency_bins = 16; - - bool use_connector = false; - int64_t connector_hidden_size = 3840; - int64_t connector_num_heads = 30; - int64_t connector_head_dim = 128; - int64_t connector_num_layers = 2; - int64_t connector_num_registers = 128; - bool connector_rope_interleaved = false; - bool connector_apply_gated_attention = false; - - bool use_audio_connector = false; - int64_t audio_connector_hidden_size = 2048; - int64_t audio_connector_num_heads = 32; - int64_t audio_connector_head_dim = 64; - int64_t audio_connector_num_layers = 2; - int64_t audio_connector_num_registers = 128; - bool audio_connector_rope_interleaved = false; - bool audio_connector_apply_gated_attention = false; - - bool video_rope_interleaved = false; - bool use_middle_indices_grid = true; - bool cross_attention_adaln = false; - - bool use_caption_projection = true; - bool use_audio_caption_projection = true; - bool caption_proj_before_connector = true; - bool caption_projection_first_linear = false; - - bool self_attention_gated = false; - bool cross_attention_gated = false; - }; - __STATIC_INLINE__ std::pair infer_attention_layout(int64_t hidden_size, int64_t preferred_heads = -1) { if (preferred_heads > 0 && hidden_size % preferred_heads == 0) { @@ -1169,92 +1307,92 @@ namespace LTXV { }; struct LTXAVModelBlock : public GGMLBlock { - LTXAVParams cfg; + LTXAVConfig config; void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { params["scale_shift_table"] = ggml_new_tensor_2d(ctx, get_type(prefix + "scale_shift_table", tensor_storage_map, GGML_TYPE_F32), - cfg.hidden_size, + config.hidden_size, 2); params["audio_scale_shift_table"] = ggml_new_tensor_2d(ctx, get_type(prefix + "audio_scale_shift_table", tensor_storage_map, GGML_TYPE_F32), - cfg.audio_hidden_size, + config.audio_hidden_size, 2); } - LTXAVModelBlock(const LTXAVParams& params) - : cfg(params) { - blocks["patchify_proj"] = std::make_shared(cfg.in_channels, cfg.hidden_size, true, true); - blocks["audio_patchify_proj"] = std::make_shared(cfg.audio_in_channels, cfg.audio_hidden_size, true, true); - blocks["adaln_single"] = std::make_shared(cfg.hidden_size, cfg.cross_attention_adaln ? 9 : 6); - blocks["audio_adaln_single"] = std::make_shared(cfg.audio_hidden_size, cfg.cross_attention_adaln ? 9 : 6); - if (cfg.cross_attention_adaln) { - blocks["prompt_adaln_single"] = std::make_shared(cfg.hidden_size, 2); - blocks["audio_prompt_adaln_single"] = std::make_shared(cfg.audio_hidden_size, 2); + LTXAVModelBlock(const LTXAVConfig& config) + : config(config) { + blocks["patchify_proj"] = std::make_shared(config.in_channels, config.hidden_size, true, true); + blocks["audio_patchify_proj"] = std::make_shared(config.audio_in_channels, config.audio_hidden_size, true, true); + blocks["adaln_single"] = std::make_shared(config.hidden_size, config.cross_attention_adaln ? 9 : 6); + blocks["audio_adaln_single"] = std::make_shared(config.audio_hidden_size, config.cross_attention_adaln ? 9 : 6); + if (config.cross_attention_adaln) { + blocks["prompt_adaln_single"] = std::make_shared(config.hidden_size, 2); + blocks["audio_prompt_adaln_single"] = std::make_shared(config.audio_hidden_size, 2); } - blocks["av_ca_video_scale_shift_adaln_single"] = std::make_shared(cfg.hidden_size, 4); - blocks["av_ca_a2v_gate_adaln_single"] = std::make_shared(cfg.hidden_size, 1); - blocks["av_ca_audio_scale_shift_adaln_single"] = std::make_shared(cfg.audio_hidden_size, 4); - blocks["av_ca_v2a_gate_adaln_single"] = std::make_shared(cfg.audio_hidden_size, 1); + blocks["av_ca_video_scale_shift_adaln_single"] = std::make_shared(config.hidden_size, 4); + blocks["av_ca_a2v_gate_adaln_single"] = std::make_shared(config.hidden_size, 1); + blocks["av_ca_audio_scale_shift_adaln_single"] = std::make_shared(config.audio_hidden_size, 4); + blocks["av_ca_v2a_gate_adaln_single"] = std::make_shared(config.audio_hidden_size, 1); - if (cfg.use_caption_projection) { - if (cfg.caption_proj_before_connector) { - if (cfg.caption_projection_first_linear) { - blocks["caption_projection"] = std::make_shared(cfg.caption_channels, cfg.hidden_size); + if (config.use_caption_projection) { + if (config.caption_proj_before_connector) { + if (config.caption_projection_first_linear) { + blocks["caption_projection"] = std::make_shared(config.caption_channels, config.hidden_size); } } else { - blocks["caption_projection"] = std::make_shared(cfg.caption_channels, cfg.hidden_size, cfg.hidden_size); + blocks["caption_projection"] = std::make_shared(config.caption_channels, config.hidden_size, config.hidden_size); } } - if (cfg.use_audio_caption_projection) { - if (cfg.caption_proj_before_connector) { - if (cfg.caption_projection_first_linear) { - blocks["audio_caption_projection"] = std::make_shared(cfg.caption_channels, cfg.audio_hidden_size); + if (config.use_audio_caption_projection) { + if (config.caption_proj_before_connector) { + if (config.caption_projection_first_linear) { + blocks["audio_caption_projection"] = std::make_shared(config.caption_channels, config.audio_hidden_size); } } else { - blocks["audio_caption_projection"] = std::make_shared(cfg.caption_channels, cfg.audio_hidden_size, cfg.audio_hidden_size); + blocks["audio_caption_projection"] = std::make_shared(config.caption_channels, config.audio_hidden_size, config.audio_hidden_size); } } - if (cfg.use_connector) { - blocks["video_embeddings_connector"] = std::make_shared(cfg.connector_hidden_size, - cfg.connector_num_heads, - cfg.connector_head_dim, - cfg.connector_num_layers, - cfg.connector_num_registers, - cfg.connector_rope_interleaved, - cfg.connector_apply_gated_attention); + if (config.use_connector) { + blocks["video_embeddings_connector"] = std::make_shared(config.connector_hidden_size, + config.connector_num_heads, + config.connector_head_dim, + config.connector_num_layers, + config.connector_num_registers, + config.connector_rope_interleaved, + config.connector_apply_gated_attention); } - if (cfg.use_audio_connector) { - blocks["audio_embeddings_connector"] = std::make_shared(cfg.audio_connector_hidden_size, - cfg.audio_connector_num_heads, - cfg.audio_connector_head_dim, - cfg.audio_connector_num_layers, - cfg.audio_connector_num_registers, - cfg.audio_connector_rope_interleaved, - cfg.audio_connector_apply_gated_attention); + if (config.use_audio_connector) { + blocks["audio_embeddings_connector"] = std::make_shared(config.audio_connector_hidden_size, + config.audio_connector_num_heads, + config.audio_connector_head_dim, + config.audio_connector_num_layers, + config.audio_connector_num_registers, + config.audio_connector_rope_interleaved, + config.audio_connector_apply_gated_attention); } - for (int i = 0; i < cfg.num_layers; i++) { - blocks["transformer_blocks." + std::to_string(i)] = std::make_shared(cfg.hidden_size, - cfg.audio_hidden_size, - cfg.num_attention_heads, - cfg.audio_num_attention_heads, - cfg.attention_head_dim, - cfg.audio_attention_head_dim, - cfg.cross_attention_dim, - cfg.audio_cross_attention_dim, - cfg.self_attention_gated || cfg.cross_attention_gated, - cfg.cross_attention_adaln, - cfg.video_rope_interleaved); + for (int i = 0; i < config.num_layers; i++) { + blocks["transformer_blocks." + std::to_string(i)] = std::make_shared(config.hidden_size, + config.audio_hidden_size, + config.num_attention_heads, + config.audio_num_attention_heads, + config.attention_head_dim, + config.audio_attention_head_dim, + config.cross_attention_dim, + config.audio_cross_attention_dim, + config.self_attention_gated || config.cross_attention_gated, + config.cross_attention_adaln, + config.video_rope_interleaved); } - blocks["norm_out"] = std::make_shared(cfg.hidden_size, 1e-6f, false); - blocks["proj_out"] = std::make_shared(cfg.hidden_size, cfg.out_channels, true, true); - blocks["audio_norm_out"] = std::make_shared(cfg.audio_hidden_size, 1e-6f, false); - blocks["audio_proj_out"] = std::make_shared(cfg.audio_hidden_size, cfg.audio_out_channels, true, true); + blocks["norm_out"] = std::make_shared(config.hidden_size, 1e-6f, false); + blocks["proj_out"] = std::make_shared(config.hidden_size, config.out_channels, true, true); + blocks["audio_norm_out"] = std::make_shared(config.audio_hidden_size, 1e-6f, false); + blocks["audio_proj_out"] = std::make_shared(config.audio_hidden_size, config.audio_out_channels, true, true); } ggml_tensor* patchify_video(GGMLRunnerContext* ctx, ggml_tensor* x, int64_t n) { @@ -1293,8 +1431,8 @@ namespace LTXV { if (ax == nullptr) { return nullptr; } - ax = ggml_reshape_4d(ctx->ggml_ctx, ax, cfg.audio_frequency_bins, cfg.num_audio_channels, audio_length, ax->ne[2]); // [b, t, c, f] - ax = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, ax, 0, 2, 1, 3)); // [b, c, t, f] + ax = ggml_reshape_4d(ctx->ggml_ctx, ax, config.audio_frequency_bins, config.num_audio_channels, audio_length, ax->ne[2]); // [b, t, c, f] + ax = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, ax, 0, 2, 1, 3)); // [b, c, t, f] return ax; } @@ -1308,17 +1446,17 @@ namespace LTXV { } bool is_fully_processed_context = - context->ne[0] == cfg.cross_attention_dim + cfg.audio_cross_attention_dim && + context->ne[0] == config.cross_attention_dim + config.audio_cross_attention_dim && context->ne[1] >= 1024; bool is_unprocessed_dual_context = - context->ne[0] == cfg.cross_attention_dim + cfg.audio_cross_attention_dim && + context->ne[0] == config.cross_attention_dim + config.audio_cross_attention_dim && context->ne[1] < 1024; if (is_fully_processed_context) { - auto v_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, 0, cfg.cross_attention_dim); + auto v_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, 0, config.cross_attention_dim); ggml_tensor* a_context = nullptr; if (process_audio_context) { - a_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, cfg.cross_attention_dim, cfg.cross_attention_dim + cfg.audio_cross_attention_dim); + a_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, config.cross_attention_dim, config.cross_attention_dim + config.audio_cross_attention_dim); } return {v_context, a_context}; } @@ -1326,32 +1464,32 @@ namespace LTXV { ggml_tensor* v_context = context; ggml_tensor* a_context = process_audio_context ? context : nullptr; if (is_unprocessed_dual_context) { - v_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, 0, cfg.cross_attention_dim); + v_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, 0, config.cross_attention_dim); if (process_audio_context) { - a_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, cfg.cross_attention_dim, cfg.cross_attention_dim + cfg.audio_cross_attention_dim); + a_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, config.cross_attention_dim, config.cross_attention_dim + config.audio_cross_attention_dim); } - } else if (context->ne[0] == cfg.caption_channels * 2) { - v_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, 0, cfg.caption_channels); + } else if (context->ne[0] == config.caption_channels * 2) { + v_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, 0, config.caption_channels); if (process_audio_context) { - a_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, cfg.caption_channels, cfg.caption_channels * 2); + a_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, config.caption_channels, config.caption_channels * 2); } } - if (cfg.caption_proj_before_connector) { - if (cfg.use_caption_projection && + if (config.caption_proj_before_connector) { + if (config.use_caption_projection && blocks.count("caption_projection") > 0 && v_context != nullptr && - v_context->ne[0] == cfg.caption_channels) { + v_context->ne[0] == config.caption_channels) { auto caption_projection = std::dynamic_pointer_cast(blocks["caption_projection"]); if (caption_projection != nullptr) { v_context = caption_projection->forward(ctx, v_context); } } if (process_audio_context && - cfg.use_audio_caption_projection && + config.use_audio_caption_projection && blocks.count("audio_caption_projection") > 0 && a_context != nullptr && - a_context->ne[0] == cfg.caption_channels) { + a_context->ne[0] == config.caption_channels) { auto caption_projection = std::dynamic_pointer_cast(blocks["audio_caption_projection"]); if (caption_projection != nullptr) { a_context = caption_projection->forward(ctx, a_context); @@ -1359,34 +1497,34 @@ namespace LTXV { } } - if (cfg.use_connector && v_context != nullptr && v_context->ne[0] == cfg.connector_hidden_size) { + if (config.use_connector && v_context != nullptr && v_context->ne[0] == config.connector_hidden_size) { auto connector = std::dynamic_pointer_cast(blocks["video_embeddings_connector"]); v_context = connector->forward(ctx, v_context, video_connector_pe); } if (process_audio_context && - cfg.use_audio_connector && + config.use_audio_connector && a_context != nullptr && - a_context->ne[0] == cfg.audio_connector_hidden_size) { + a_context->ne[0] == config.audio_connector_hidden_size) { auto connector = std::dynamic_pointer_cast(blocks["audio_embeddings_connector"]); a_context = connector->forward(ctx, a_context, audio_connector_pe); } - if (!cfg.caption_proj_before_connector && - cfg.use_caption_projection && + if (!config.caption_proj_before_connector && + config.use_caption_projection && blocks.count("caption_projection") > 0 && v_context != nullptr && - v_context->ne[0] == cfg.caption_channels) { + v_context->ne[0] == config.caption_channels) { auto caption_projection = std::dynamic_pointer_cast(blocks["caption_projection"]); if (caption_projection != nullptr) { v_context = caption_projection->forward(ctx, v_context); } } if (process_audio_context && - !cfg.caption_proj_before_connector && - cfg.use_audio_caption_projection && + !config.caption_proj_before_connector && + config.use_audio_caption_projection && blocks.count("audio_caption_projection") > 0 && a_context != nullptr && - a_context->ne[0] == cfg.caption_channels) { + a_context->ne[0] == config.caption_channels) { auto caption_projection = std::dynamic_pointer_cast(blocks["audio_caption_projection"]); if (caption_projection != nullptr) { a_context = caption_projection->forward(ctx, a_context); @@ -1428,8 +1566,8 @@ namespace LTXV { auto audio_norm_out = std::dynamic_pointer_cast(blocks["audio_norm_out"]); auto audio_proj_out = std::dynamic_pointer_cast(blocks["audio_proj_out"]); - GGML_ASSERT(vx->ne[3] % cfg.in_channels == 0); - int64_t n = vx->ne[3] / cfg.in_channels; + GGML_ASSERT(vx->ne[3] % config.in_channels == 0); + int64_t n = vx->ne[3] / config.in_channels; int64_t width = vx->ne[0]; int64_t height = vx->ne[1]; int64_t frames = vx->ne[2]; @@ -1452,20 +1590,20 @@ namespace LTXV { a_context = ggml_cont(ctx->ggml_ctx, a_context); } - auto v_timestep_scaled = ggml_ext_scale(ctx->ggml_ctx, timestep, cfg.timestep_scale_multiplier); + auto v_timestep_scaled = ggml_ext_scale(ctx->ggml_ctx, timestep, config.timestep_scale_multiplier); auto v_pair = adaln_single->forward(ctx, v_timestep_scaled); auto v_timestep_mod = v_pair.first; auto v_embedded_time = v_pair.second; ggml_tensor* effective_audio_timestep = audio_timestep != nullptr ? audio_timestep : timestep; - auto a_timestep_scaled = ggml_ext_scale(ctx->ggml_ctx, effective_audio_timestep, cfg.timestep_scale_multiplier); + auto a_timestep_scaled = ggml_ext_scale(ctx->ggml_ctx, effective_audio_timestep, config.timestep_scale_multiplier); auto a_pair = audio_adaln_single->forward(ctx, a_timestep_scaled); auto a_timestep_mod = a_pair.first; auto a_embedded_time = a_pair.second; ggml_tensor* v_prompt_timestep_mod = nullptr; ggml_tensor* a_prompt_timestep_mod = nullptr; - if (cfg.cross_attention_adaln) { + if (config.cross_attention_adaln) { auto prompt_adaln_single = std::dynamic_pointer_cast(blocks["prompt_adaln_single"]); auto audio_prompt_adaln_single = std::dynamic_pointer_cast(blocks["audio_prompt_adaln_single"]); v_prompt_timestep_mod = prompt_adaln_single->forward(ctx, a_timestep_scaled).first; @@ -1474,7 +1612,7 @@ namespace LTXV { auto av_ca_video_timestep = repeat_scalar_timestep_like(ctx, effective_audio_timestep, timestep); auto av_ca_audio_timestep = effective_audio_timestep; - auto av_ca_factor = cfg.av_ca_timestep_scale_multiplier / cfg.timestep_scale_multiplier; + auto av_ca_factor = config.av_ca_timestep_scale_multiplier / config.timestep_scale_multiplier; auto av_ca_video_scale_shift_timestep = std::dynamic_pointer_cast(blocks["av_ca_video_scale_shift_adaln_single"])->forward(ctx, av_ca_video_timestep).first; auto av_ca_a2v_gate_noise_timestep = @@ -1491,7 +1629,7 @@ namespace LTXV { 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++) { + for (int i = 0; i < config.num_layers; i++) { auto block = std::dynamic_pointer_cast(blocks["transformer_blocks." + std::to_string(i)]); auto out = block->forward(ctx, vx, @@ -1517,14 +1655,14 @@ namespace LTXV { 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); + auto v_shift_scale = get_output_scale_shift(ctx, params["scale_shift_table"], v_embedded_time, config.hidden_size); vx = norm_out->forward(ctx, vx); vx = modulate(ctx->ggml_ctx, vx, v_shift_scale[0], v_shift_scale[1]); vx = proj_out->forward(ctx, vx); vx = unpatchify_video(ctx, vx, width, height, frames); if (ax != nullptr && audio_time > 0) { - auto a_shift_scale = get_output_scale_shift(ctx, params["audio_scale_shift_table"], a_embedded_time, cfg.audio_hidden_size); + auto a_shift_scale = get_output_scale_shift(ctx, params["audio_scale_shift_table"], a_embedded_time, config.audio_hidden_size); ax = audio_norm_out->forward(ctx, ax); ax = modulate(ctx->ggml_ctx, ax, a_shift_scale[0], a_shift_scale[1]); ax = audio_proj_out->forward(ctx, ax); @@ -1536,7 +1674,7 @@ namespace LTXV { }; struct LTXAVRunner : public DiffusionModelRunner { - LTXAVParams params; + LTXAVConfig config; LTXAVModelBlock model; std::vector video_pe_vec; std::vector audio_pe_vec; @@ -1547,124 +1685,13 @@ namespace LTXV { sd::Tensor vx_input_cache; sd::Tensor ax_input_cache; - static int64_t infer_gate_heads(const String2TensorStorage& tensor_storage_map, - const std::string& bias_name, - int64_t fallback_heads) { - auto it = tensor_storage_map.find(bias_name); - if (it != tensor_storage_map.end()) { - return it->second.ne[0]; - } - return fallback_heads; - } - LTXAVRunner(ggml_backend_t backend, - ggml_backend_t params_backend, - const String2TensorStorage& tensor_storage_map = {}, - const std::string& prefix = "model.diffusion_model") - : DiffusionModelRunner(backend, params_backend, prefix), - params(), - model(params) { - auto patchify_proj_iter = tensor_storage_map.find(prefix + ".patchify_proj.weight"); - if (patchify_proj_iter != tensor_storage_map.end()) { - params.in_channels = patchify_proj_iter->second.ne[0]; - params.hidden_size = patchify_proj_iter->second.ne[1]; - int64_t video_heads = infer_gate_heads(tensor_storage_map, prefix + ".transformer_blocks.0.attn1.to_gate_logits.bias", 32); - auto attn_layout = infer_attention_layout(params.hidden_size, video_heads); - params.num_attention_heads = attn_layout.first; - params.attention_head_dim = attn_layout.second; - } - - auto audio_patchify_proj_iter = tensor_storage_map.find(prefix + ".audio_patchify_proj.weight"); - if (audio_patchify_proj_iter != tensor_storage_map.end()) { - params.audio_in_channels = audio_patchify_proj_iter->second.ne[0]; - params.audio_hidden_size = audio_patchify_proj_iter->second.ne[1]; - params.audio_out_channels = params.audio_in_channels; - int64_t audio_heads = infer_gate_heads(tensor_storage_map, prefix + ".transformer_blocks.0.audio_attn1.to_gate_logits.bias", 32); - auto audio_attn_layout = infer_attention_layout(params.audio_hidden_size, audio_heads); - params.audio_num_attention_heads = audio_attn_layout.first; - params.audio_attention_head_dim = audio_attn_layout.second; - } - - auto proj_out_iter = tensor_storage_map.find(prefix + ".proj_out.weight"); - if (proj_out_iter != tensor_storage_map.end()) { - params.out_channels = proj_out_iter->second.ne[1]; - } - auto audio_proj_out_iter = tensor_storage_map.find(prefix + ".audio_proj_out.weight"); - if (audio_proj_out_iter != tensor_storage_map.end()) { - params.audio_out_channels = audio_proj_out_iter->second.ne[1]; - } - - auto attn2_iter = tensor_storage_map.find(prefix + ".transformer_blocks.0.attn2.to_k.weight"); - if (attn2_iter != tensor_storage_map.end()) { - params.cross_attention_dim = attn2_iter->second.ne[0]; - } - auto audio_attn2_iter = tensor_storage_map.find(prefix + ".transformer_blocks.0.audio_attn2.to_k.weight"); - if (audio_attn2_iter != tensor_storage_map.end()) { - params.audio_cross_attention_dim = audio_attn2_iter->second.ne[0]; - } - if (tensor_storage_map.find(prefix + ".transformer_blocks.0.prompt_scale_shift_table") != tensor_storage_map.end()) { - params.cross_attention_adaln = true; - } - if (tensor_storage_map.find(prefix + ".transformer_blocks.0.attn1.to_gate_logits.weight") != tensor_storage_map.end() || - tensor_storage_map.find(prefix + ".transformer_blocks.0.audio_attn1.to_gate_logits.weight") != tensor_storage_map.end()) { - params.self_attention_gated = true; - } - if (tensor_storage_map.find(prefix + ".transformer_blocks.0.attn2.to_gate_logits.weight") != tensor_storage_map.end() || - tensor_storage_map.find(prefix + ".transformer_blocks.0.audio_attn2.to_gate_logits.weight") != tensor_storage_map.end()) { - params.cross_attention_gated = true; - } - if (tensor_storage_map.find(prefix + ".caption_projection.linear_1.weight") == tensor_storage_map.end() && - tensor_storage_map.find(prefix + ".caption_projection.linear_2.weight") == tensor_storage_map.end()) { - params.use_caption_projection = false; - } - if (tensor_storage_map.find(prefix + ".audio_caption_projection.linear_1.weight") == tensor_storage_map.end() && - tensor_storage_map.find(prefix + ".audio_caption_projection.linear_2.weight") == tensor_storage_map.end()) { - params.use_audio_caption_projection = false; - } - - params.num_layers = count_prefix_blocks(tensor_storage_map, prefix + ".", "transformer_blocks."); - - auto connector_iter = tensor_storage_map.find(prefix + ".video_embeddings_connector.transformer_1d_blocks.0.attn1.to_q.weight"); - if (connector_iter != tensor_storage_map.end()) { - params.use_connector = true; - params.connector_hidden_size = connector_iter->second.ne[1]; - int64_t connector_heads = infer_gate_heads(tensor_storage_map, - prefix + ".video_embeddings_connector.transformer_1d_blocks.0.attn1.to_gate_logits.bias", - 32); - auto connector_layout = infer_attention_layout(params.connector_hidden_size, connector_heads); - params.connector_num_heads = connector_layout.first; - params.connector_head_dim = connector_layout.second; - params.connector_num_layers = count_prefix_blocks(tensor_storage_map, prefix + ".video_embeddings_connector.", "transformer_1d_blocks."); - auto register_iter = tensor_storage_map.find(prefix + ".video_embeddings_connector.learnable_registers"); - if (register_iter != tensor_storage_map.end()) { - params.connector_num_registers = register_iter->second.ne[1]; - } - if (tensor_storage_map.find(prefix + ".video_embeddings_connector.transformer_1d_blocks.0.attn1.to_gate_logits.weight") != tensor_storage_map.end()) { - params.connector_apply_gated_attention = true; - } - } - - auto audio_connector_iter = tensor_storage_map.find(prefix + ".audio_embeddings_connector.transformer_1d_blocks.0.attn1.to_q.weight"); - if (audio_connector_iter != tensor_storage_map.end()) { - params.use_audio_connector = true; - params.audio_connector_hidden_size = audio_connector_iter->second.ne[1]; - int64_t connector_heads = infer_gate_heads(tensor_storage_map, - prefix + ".audio_embeddings_connector.transformer_1d_blocks.0.attn1.to_gate_logits.bias", - 32); - auto connector_layout = infer_attention_layout(params.audio_connector_hidden_size, connector_heads); - params.audio_connector_num_heads = connector_layout.first; - params.audio_connector_head_dim = connector_layout.second; - params.audio_connector_num_layers = count_prefix_blocks(tensor_storage_map, prefix + ".audio_embeddings_connector.", "transformer_1d_blocks."); - auto register_iter = tensor_storage_map.find(prefix + ".audio_embeddings_connector.learnable_registers"); - if (register_iter != tensor_storage_map.end()) { - params.audio_connector_num_registers = register_iter->second.ne[1]; - } - if (tensor_storage_map.find(prefix + ".audio_embeddings_connector.transformer_1d_blocks.0.attn1.to_gate_logits.weight") != tensor_storage_map.end()) { - params.audio_connector_apply_gated_attention = true; - } - } - - model = LTXAVModelBlock(params); + const String2TensorStorage& tensor_storage_map = {}, + const std::string& prefix = "model.diffusion_model", + std::shared_ptr weight_manager = nullptr) + : DiffusionModelRunner(backend, prefix, weight_manager), + config(LTXAVConfig::detect_from_weights(tensor_storage_map, prefix)), + model(config) { model.init(params_ctx, tensor_storage_map, prefix); } @@ -1692,21 +1719,21 @@ namespace LTXV { int64_t total_channels = x_tensor.shape()[3]; int64_t spatial_size = width * height * frames; - GGML_ASSERT(total_channels >= params.in_channels); + GGML_ASSERT(total_channels >= config.in_channels); - sd::Tensor vx({width, height, frames, params.in_channels}); - size_t video_values = static_cast(params.in_channels * spatial_size); + sd::Tensor vx({width, height, frames, config.in_channels}); + size_t video_values = static_cast(config.in_channels * spatial_size); std::copy_n(x_tensor.data(), video_values, vx.data()); - if (audio_length <= 0 || total_channels == params.in_channels) { + if (audio_length <= 0 || total_channels == config.in_channels) { return {vx, {}}; } - int64_t needed_audio_values = static_cast(audio_length) * params.num_audio_channels * params.audio_frequency_bins; - int64_t packed_audio_values = (total_channels - params.in_channels) * spatial_size; + int64_t needed_audio_values = static_cast(audio_length) * config.num_audio_channels * config.audio_frequency_bins; + int64_t packed_audio_values = (total_channels - config.in_channels) * spatial_size; GGML_ASSERT(packed_audio_values >= needed_audio_values); - sd::Tensor ax({params.audio_frequency_bins, audio_length, params.num_audio_channels, 1}); + sd::Tensor ax({config.audio_frequency_bins, audio_length, config.num_audio_channels, 1}); const float* audio_src = x_tensor.data() + video_values; std::copy_n(audio_src, static_cast(needed_audio_values), ax.data()); return {vx, ax}; @@ -1767,25 +1794,25 @@ namespace LTXV { if (has_video_positions) { GGML_ASSERT(video_positions_tensor.shape()[2] == video_token_count); video_pe_vec = build_video_rope_matrix_from_positions(video_positions_tensor, - static_cast(params.hidden_size), - static_cast(params.num_attention_heads), - params.positional_embedding_theta, - params.positional_embedding_max_pos, - params.use_middle_indices_grid); + static_cast(config.hidden_size), + static_cast(config.num_attention_heads), + config.positional_embedding_theta, + config.positional_embedding_max_pos, + config.use_middle_indices_grid); } else { video_pe_vec = build_video_rope_matrix(vx->ne[0], vx->ne[1], vx->ne[2], - static_cast(params.hidden_size), - static_cast(params.num_attention_heads), + static_cast(config.hidden_size), + static_cast(config.num_attention_heads), video_frame_rate, - params.positional_embedding_theta, - params.positional_embedding_max_pos, - params.vae_scale_factors, - params.causal_temporal_positioning, - params.use_middle_indices_grid); + config.positional_embedding_theta, + config.positional_embedding_max_pos, + config.vae_scale_factors, + config.causal_temporal_positioning, + config.use_middle_indices_grid); } - auto video_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, params.attention_head_dim / 2, video_token_count * params.num_attention_heads); + auto video_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.attention_head_dim / 2, video_token_count * config.num_attention_heads); ggml_set_name(video_pe, "ltxav_video_pe"); set_backend_tensor_data(video_pe, video_pe_vec.data()); @@ -1794,66 +1821,66 @@ namespace LTXV { ggml_tensor* audio_cross_pe = nullptr; if (ax != nullptr && ggml_nelements(ax) > 0 && ax->ne[1] > 0) { audio_pe_vec = build_audio_rope_matrix(ax->ne[1], - static_cast(params.audio_hidden_size), - static_cast(params.audio_num_attention_heads), - params.positional_embedding_theta, - params.audio_positional_embedding_max_pos[0], - params.use_middle_indices_grid); - audio_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, params.audio_attention_head_dim / 2, ax->ne[1] * params.audio_num_attention_heads); + static_cast(config.audio_hidden_size), + static_cast(config.audio_num_attention_heads), + config.positional_embedding_theta, + config.audio_positional_embedding_max_pos[0], + config.use_middle_indices_grid); + audio_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.audio_attention_head_dim / 2, ax->ne[1] * config.audio_num_attention_heads); ggml_set_name(audio_pe, "ltxav_audio_pe"); set_backend_tensor_data(audio_pe, audio_pe_vec.data()); - int temporal_max_pos = std::max(params.positional_embedding_max_pos[0], params.audio_positional_embedding_max_pos[0]); + int temporal_max_pos = std::max(config.positional_embedding_max_pos[0], config.audio_positional_embedding_max_pos[0]); if (has_video_positions) { video_cross_pe_vec = build_video_temporal_rope_matrix_from_positions(video_positions_tensor, - static_cast(params.audio_cross_attention_dim), - static_cast(params.audio_num_attention_heads), - params.positional_embedding_theta, + static_cast(config.audio_cross_attention_dim), + static_cast(config.audio_num_attention_heads), + config.positional_embedding_theta, temporal_max_pos, true); } else { video_cross_pe_vec = build_video_temporal_rope_matrix(vx->ne[0], vx->ne[1], vx->ne[2], - static_cast(params.audio_cross_attention_dim), - static_cast(params.audio_num_attention_heads), + static_cast(config.audio_cross_attention_dim), + static_cast(config.audio_num_attention_heads), video_frame_rate, - params.positional_embedding_theta, + config.positional_embedding_theta, temporal_max_pos, - std::get<0>(params.vae_scale_factors), - params.causal_temporal_positioning, + std::get<0>(config.vae_scale_factors), + config.causal_temporal_positioning, true); } - video_cross_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, params.audio_attention_head_dim / 2, video_token_count * params.audio_num_attention_heads); + video_cross_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.audio_attention_head_dim / 2, video_token_count * config.audio_num_attention_heads); ggml_set_name(video_cross_pe, "ltxav_video_cross_pe"); set_backend_tensor_data(video_cross_pe, video_cross_pe_vec.data()); audio_cross_pe_vec = build_audio_rope_matrix(ax->ne[1], - static_cast(params.audio_cross_attention_dim), - static_cast(params.audio_num_attention_heads), - params.positional_embedding_theta, + static_cast(config.audio_cross_attention_dim), + static_cast(config.audio_num_attention_heads), + config.positional_embedding_theta, temporal_max_pos, true); - audio_cross_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, params.audio_attention_head_dim / 2, ax->ne[1] * params.audio_num_attention_heads); + audio_cross_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.audio_attention_head_dim / 2, ax->ne[1] * config.audio_num_attention_heads); ggml_set_name(audio_cross_pe, "ltxav_audio_cross_pe"); set_backend_tensor_data(audio_cross_pe, audio_cross_pe_vec.data()); } bool needs_video_connector_pe = - params.use_connector && + config.use_connector && context != nullptr && - (context->ne[0] == params.connector_hidden_size || - ((context->ne[0] == params.cross_attention_dim + params.audio_cross_attention_dim || - context->ne[0] == params.caption_channels * 2) && + (context->ne[0] == config.connector_hidden_size || + ((context->ne[0] == config.cross_attention_dim + config.audio_cross_attention_dim || + context->ne[0] == config.caption_channels * 2) && context->ne[1] < 1024)); ggml_tensor* video_connector_pe = nullptr; if (needs_video_connector_pe) { int64_t seq_len = context->ne[1]; int64_t target_len = std::max(1024, seq_len); - int64_t duplications = (target_len + params.connector_num_registers - 1) / params.connector_num_registers; - int64_t full_len = seq_len + duplications * params.connector_num_registers - seq_len; - connector_pe_vec = build_1d_rope_matrix(full_len, static_cast(params.connector_hidden_size), static_cast(params.connector_num_heads), 10000.f, 4096.f, true); - video_connector_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, params.connector_head_dim / 2, full_len * params.connector_num_heads); + int64_t duplications = (target_len + config.connector_num_registers - 1) / config.connector_num_registers; + int64_t full_len = seq_len + duplications * config.connector_num_registers - seq_len; + connector_pe_vec = build_1d_rope_matrix(full_len, static_cast(config.connector_hidden_size), static_cast(config.connector_num_heads), 10000.f, 4096.f, true); + video_connector_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.connector_head_dim / 2, full_len * config.connector_num_heads); ggml_set_name(video_connector_pe, "ltxav_video_connector_pe"); set_backend_tensor_data(video_connector_pe, connector_pe_vec.data()); } @@ -1864,20 +1891,20 @@ namespace LTXV { ax->ne[1] > 0; bool needs_audio_connector_pe = run_audio_context && - params.use_audio_connector && + config.use_audio_connector && context != nullptr && - (context->ne[0] == params.audio_connector_hidden_size || - ((context->ne[0] == params.cross_attention_dim + params.audio_cross_attention_dim || - context->ne[0] == params.caption_channels * 2) && + (context->ne[0] == config.audio_connector_hidden_size || + ((context->ne[0] == config.cross_attention_dim + config.audio_cross_attention_dim || + context->ne[0] == config.caption_channels * 2) && context->ne[1] < 1024)); ggml_tensor* audio_connector_pe = nullptr; if (needs_audio_connector_pe) { int64_t seq_len = context->ne[1]; int64_t target_len = std::max(1024, seq_len); - int64_t duplications = (target_len + params.audio_connector_num_registers - 1) / params.audio_connector_num_registers; - int64_t full_len = seq_len + duplications * params.audio_connector_num_registers - seq_len; - audio_connector_pe_vec = build_1d_rope_matrix(full_len, static_cast(params.audio_connector_hidden_size), static_cast(params.audio_connector_num_heads), 10000.f, 4096.f, true); - audio_connector_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, params.audio_connector_head_dim / 2, full_len * params.audio_connector_num_heads); + int64_t duplications = (target_len + config.audio_connector_num_registers - 1) / config.audio_connector_num_registers; + int64_t full_len = seq_len + duplications * config.audio_connector_num_registers - seq_len; + audio_connector_pe_vec = build_1d_rope_matrix(full_len, static_cast(config.audio_connector_hidden_size), static_cast(config.audio_connector_num_heads), 10000.f, 4096.f, true); + audio_connector_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.audio_connector_head_dim / 2, full_len * config.audio_connector_num_heads); ggml_set_name(audio_connector_pe, "ltxav_audio_connector_pe"); set_backend_tensor_data(audio_connector_pe, audio_connector_pe_vec.data()); } @@ -1912,7 +1939,7 @@ namespace LTXV { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context, audio_x, audio_timesteps, audio_length, frame_rate, video_positions); }; - auto out = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); + auto out = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); return out; } @@ -1998,7 +2025,8 @@ namespace LTXV { ggml_backend_t backend = sd_backend_cpu_init(); LOG_INFO("loading ltxav from '%s'", model_path.c_str()); - ModelLoader model_loader; + auto model_manager = std::make_shared(); + ModelLoader& model_loader = model_manager->loader(); if (!model_loader.init_from_file_and_convert_name(model_path, "model.diffusion_model.")) { LOG_ERROR("init model loader from file failed: '%s'", model_path.c_str()); return; @@ -2013,19 +2041,18 @@ namespace LTXV { auto& tensor_storage_map = model_loader.get_tensor_storage_map(); std::shared_ptr ltxav = std::make_shared(backend, - backend, tensor_storage_map, - "model.diffusion_model"); + "model.diffusion_model", + model_manager); - if (!ltxav->alloc_params_buffer()) { - LOG_ERROR("ltxav buffer allocation failed"); - return; - } - std::map tensors; - ltxav->get_param_tensors(tensors, "model.diffusion_model"); - - if (!model_loader.load_tensors(tensors)) { - LOG_ERROR("load tensors from model loader failed"); + if (!model_manager->register_runner_params("LTXAV test", + *ltxav, + "model.diffusion_model", + ModelManager::ResidencyMode::ParamBackend, + backend, + backend) || + !model_manager->validate_registered_tensors()) { + LOG_ERROR("register ltxav tensors with model manager failed"); return; } @@ -2036,4 +2063,4 @@ namespace LTXV { }; // namespace LTXV -#endif +#endif // __SD_MODEL_DIFFUSION_LTXV_HPP__ diff --git a/src/model/diffusion/minit2i.hpp b/src/model/diffusion/minit2i.hpp new file mode 100644 index 00000000..28466105 --- /dev/null +++ b/src/model/diffusion/minit2i.hpp @@ -0,0 +1,611 @@ +#ifndef __SD_MODEL_DIFFUSION_MINIT2I_HPP__ +#define __SD_MODEL_DIFFUSION_MINIT2I_HPP__ + +#include +#include +#include +#include +#include +#include +#include + +#include "core/ggml_extend.hpp" +#include "model/common/rope.hpp" +#include "model/diffusion/dit.hpp" +#include "model/diffusion/model.hpp" +#include "model_loader.h" + +namespace MiniT2I { + constexpr int MINIT2I_GRAPH_SIZE = 196608; + + struct MiniT2IConfig { + int64_t image_size = 512; + int64_t patch_size = 16; + int64_t in_channels = 3; + int64_t txt_input_size = 1024; + int64_t hidden_size = 768; + int64_t txt_hidden_size = 768; + int64_t cond_vec_size = 768; + int64_t depth_double = 17; + int64_t txt_preamble_depth = 2; + int64_t num_heads = 12; + int64_t head_dim = 64; + float mlp_ratio = 2.6667f; + int64_t pca_channels = 128; + int64_t prompt_length = 256; + int64_t n_T = 100; + float cfg_interval_start = 0.0f; + float cfg_interval_end = 1.0f; + + static MiniT2IConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) { + MiniT2IConfig config; + config.depth_double = 0; + config.txt_preamble_depth = 0; + + for (const auto& [name, tensor_storage] : tensor_storage_map) { + if (!starts_with(name, prefix)) { + continue; + } + if (ends_with(name, "img_embedder.proj1.weight") && tensor_storage.n_dims == 4) { + config.patch_size = tensor_storage.ne[0]; + config.in_channels = tensor_storage.ne[2]; + config.pca_channels = tensor_storage.ne[3]; + } else if (ends_with(name, "img_embedder.proj2.weight") && tensor_storage.n_dims == 4) { + config.pca_channels = tensor_storage.ne[2]; + config.hidden_size = tensor_storage.ne[3]; + } else if (ends_with(name, "txt_embedder.weight") && tensor_storage.n_dims == 2) { + config.txt_input_size = tensor_storage.ne[0]; + config.txt_hidden_size = tensor_storage.ne[1]; + } else if (ends_with(name, "pooled_embedder.weight") && tensor_storage.n_dims == 2) { + config.cond_vec_size = tensor_storage.ne[1]; + } else if (ends_with(name, "double_blocks.0.img_qkv.weight") && tensor_storage.n_dims == 2) { + int64_t inner3 = tensor_storage.ne[1]; + int64_t inner = inner3 / 3; + config.hidden_size = tensor_storage.ne[0]; + if (config.hidden_size == 768) { + config.num_heads = 12; + config.head_dim = 64; + } else if (config.hidden_size == 1248) { + config.num_heads = 24; + config.head_dim = 52; + } else if (inner > 0) { + config.head_dim = 64; + config.num_heads = std::max(1, inner / config.head_dim); + } + } else if (ends_with(name, "final_layer.linear.weight") && tensor_storage.n_dims == 2) { + int64_t patch_area = config.patch_size * config.patch_size; + config.hidden_size = tensor_storage.ne[0]; + config.in_channels = patch_area > 0 ? tensor_storage.ne[1] / patch_area : config.in_channels; + } else if (ends_with(name, "mask_token") && tensor_storage.n_dims >= 2) { + config.prompt_length = tensor_storage.ne[1]; + } + + size_t pos = name.find("double_blocks."); + if (pos != std::string::npos) { + auto items = split_string(name.substr(pos), '.'); + if (items.size() > 1) { + int64_t idx = atoi(items[1].c_str()); + config.depth_double = std::max(config.depth_double, idx + 1); + } + } + pos = name.find("txt_preamble_blocks."); + if (pos != std::string::npos) { + auto items = split_string(name.substr(pos), '.'); + if (items.size() > 1) { + int64_t idx = atoi(items[1].c_str()); + config.txt_preamble_depth = std::max(config.txt_preamble_depth, idx + 1); + } + } + } + + if (config.depth_double <= 0) { + config.depth_double = config.hidden_size == 1248 ? 23 : 17; + } + if (config.txt_preamble_depth <= 0) { + config.txt_preamble_depth = 2; + } + if (config.head_dim <= 0 || config.num_heads <= 0) { + config.head_dim = config.hidden_size == 1248 ? 52 : 64; + config.num_heads = config.hidden_size / config.head_dim; + } + LOG_DEBUG("minit2i: hidden_size=%" PRId64 ", txt_hidden_size=%" PRId64 ", heads=%" PRId64 ", head_dim=%" PRId64 ", double_blocks=%" PRId64 ", txt_blocks=%" PRId64 ", patch=%" PRId64 ", in_channels=%" PRId64, + config.hidden_size, + config.txt_hidden_size, + config.num_heads, + config.head_dim, + config.depth_double, + config.txt_preamble_depth, + config.patch_size, + config.in_channels); + return config; + } + }; + + inline std::vector make_2d_sincos_pos_embed(int grid_size, int dim) { + GGML_ASSERT(dim % 4 == 0); + int half_dim = dim / 2; + int quarter = half_dim / 2; + std::vector out(static_cast(grid_size) * grid_size * dim); + std::vector omega(quarter); + for (int i = 0; i < quarter; ++i) { + omega[i] = 1.0f / std::pow(10000.0f, static_cast(i) / static_cast(quarter)); + } + for (int y = 0; y < grid_size; ++y) { + for (int x = 0; x < grid_size; ++x) { + size_t base = static_cast(y * grid_size + x) * dim; + for (int i = 0; i < quarter; ++i) { + float ay = y * omega[i]; + float ax = x * omega[i]; + out[base + i] = std::sin(ax); + out[base + quarter + i] = std::cos(ax); + out[base + half_dim + i] = std::sin(ay); + out[base + half_dim + quarter + i] = std::cos(ay); + } + } + } + return out; + } + + inline std::vector make_text_rope(int length, int head_dim) { + return Rope::flatten(Rope::rope(Rope::linspace(0.f, static_cast(length - 1), length), head_dim, 10000.f)); + } + + inline std::vector make_vision_rope(int side, int head_dim) { + GGML_ASSERT(head_dim % 4 == 0); + int dim = head_dim / 2; + int quarter = dim / 2; + int length = side * side; + std::vector out(static_cast(length) * (head_dim / 2) * 4); + std::vector freqs(quarter); + for (int i = 0; i < quarter; ++i) { + freqs[i] = 1.0f / std::pow(10000.0f, static_cast(2 * i) / static_cast(dim)); + } + for (int y = 0; y < side; ++y) { + for (int x = 0; x < side; ++x) { + int pos = y * side + x; + size_t base = static_cast(pos) * (head_dim / 2) * 4; + for (int i = 0; i < quarter; ++i) { + float ay = y * freqs[i]; + float ax = x * freqs[i]; + float angles[2] = {ay, ax}; + for (int axis = 0; axis < 2; ++axis) { + int j = axis * quarter + i; + out[base + 4 * j] = std::cos(angles[axis]); + out[base + 4 * j + 1] = -std::sin(angles[axis]); + out[base + 4 * j + 2] = std::sin(angles[axis]); + out[base + 4 * j + 3] = std::cos(angles[axis]); + } + } + } + } + return out; + } + + struct SwiGLUMlp : public GGMLBlock { + SwiGLUMlp(int64_t in_features, int64_t hidden_features) { + int64_t hidden_dim = ((hidden_features + 7) / 8) * 8; + blocks["w1"] = std::make_shared(in_features, hidden_dim, false); + blocks["w3"] = std::make_shared(in_features, hidden_dim, false); + blocks["w2"] = std::make_shared(hidden_dim, in_features, false); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto w1 = std::dynamic_pointer_cast(blocks["w1"]); + auto w3 = std::dynamic_pointer_cast(blocks["w3"]); + auto w2 = std::dynamic_pointer_cast(blocks["w2"]); + auto gate = ggml_silu(ctx->ggml_ctx, w1->forward(ctx, x)); + auto up = w3->forward(ctx, x); + return w2->forward(ctx, ggml_mul(ctx->ggml_ctx, gate, up)); + } + }; + + struct BottleneckPatchEmbed : public GGMLBlock { + int64_t patch_size; + + BottleneckPatchEmbed(int64_t patch_size, int64_t in_channels, int64_t pca_channels, int64_t hidden_size) + : patch_size(patch_size) { + blocks["proj1"] = std::make_shared(in_channels, + pca_channels, + std::pair{static_cast(patch_size), static_cast(patch_size)}, + std::pair{static_cast(patch_size), static_cast(patch_size)}, + std::pair{0, 0}, + std::pair{1, 1}, + false); + blocks["proj2"] = std::make_shared(pca_channels, + hidden_size, + std::pair{1, 1}, + std::pair{1, 1}, + std::pair{0, 0}, + std::pair{1, 1}, + true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto proj1 = std::dynamic_pointer_cast(blocks["proj1"]); + auto proj2 = std::dynamic_pointer_cast(blocks["proj2"]); + x = proj1->forward(ctx, x); + x = proj2->forward(ctx, x); + x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1], x->ne[2], x->ne[3]); + x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); + return x; + } + }; + + struct TimestepEmbedder : public GGMLBlock { + int frequency_embedding_size; + + TimestepEmbedder(int64_t hidden_size, int frequency_embedding_size = 256) + : frequency_embedding_size(frequency_embedding_size) { + blocks["mlp.0"] = std::make_shared(frequency_embedding_size, hidden_size, true, true); + blocks["mlp.2"] = std::make_shared(hidden_size, hidden_size, true, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* t) { + auto mlp_0 = std::dynamic_pointer_cast(blocks["mlp.0"]); + auto mlp_2 = std::dynamic_pointer_cast(blocks["mlp.2"]); + auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, t, frequency_embedding_size, 10000, 1.0f); + t_emb = mlp_0->forward(ctx, t_emb); + t_emb = ggml_silu_inplace(ctx->ggml_ctx, t_emb); + return mlp_2->forward(ctx, t_emb); + } + }; + + inline std::vector split_qkv(ggml_context* ctx, ggml_tensor* qkv, int64_t num_heads, int64_t head_dim) { + int64_t N = qkv->ne[2]; + int64_t L = qkv->ne[1]; + auto q = ggml_view_4d(ctx, qkv, head_dim, num_heads, L, N, + qkv->nb[0] * head_dim, qkv->nb[1], qkv->nb[2], 0); + auto k = ggml_view_4d(ctx, qkv, head_dim, num_heads, L, N, + qkv->nb[0] * head_dim, qkv->nb[1], qkv->nb[2], qkv->nb[0] * head_dim * num_heads); + auto v = ggml_view_4d(ctx, qkv, head_dim, num_heads, L, N, + qkv->nb[0] * head_dim, qkv->nb[1], qkv->nb[2], qkv->nb[0] * head_dim * num_heads * 2); + return {q, k, v}; + } + + struct PlainTextTransformerBlock : public GGMLBlock { + int64_t num_heads; + int64_t head_dim; + + PlainTextTransformerBlock(int64_t hidden_size, int64_t num_heads, int64_t head_dim, float mlp_ratio) + : num_heads(num_heads), head_dim(head_dim) { + int64_t inner_dim = num_heads * head_dim; + blocks["norm1"] = std::make_shared(hidden_size, 1e-6f); + blocks["norm2"] = std::make_shared(hidden_size, 1e-6f); + blocks["qkv"] = std::make_shared(hidden_size, inner_dim * 3, true); + blocks["attn_proj"] = std::make_shared(inner_dim, hidden_size, true); + blocks["mlp"] = std::make_shared(hidden_size, static_cast(hidden_size * mlp_ratio)); + blocks["q_norm"] = std::make_shared(head_dim, 1e-6f); + blocks["k_norm"] = std::make_shared(head_dim, 1e-6f); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* txt, ggml_tensor* pe) { + auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); + auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); + auto qkv_proj = std::dynamic_pointer_cast(blocks["qkv"]); + auto attn_proj = std::dynamic_pointer_cast(blocks["attn_proj"]); + auto mlp = std::dynamic_pointer_cast(blocks["mlp"]); + auto q_norm = std::dynamic_pointer_cast(blocks["q_norm"]); + auto k_norm = std::dynamic_pointer_cast(blocks["k_norm"]); + + auto qkv = split_qkv(ctx->ggml_ctx, qkv_proj->forward(ctx, norm1->forward(ctx, txt)), num_heads, head_dim); + auto q = q_norm->forward(ctx, qkv[0]); + auto k = k_norm->forward(ctx, qkv[1]); + auto v = qkv[2]; + auto out = Rope::attention(ctx, q, k, v, pe, nullptr, 1.0f, false); + txt = ggml_add(ctx->ggml_ctx, txt, attn_proj->forward(ctx, out)); + txt = ggml_add(ctx->ggml_ctx, txt, mlp->forward(ctx, norm2->forward(ctx, txt))); + return txt; + } + }; + + struct DoubleStreamDiTBlock : public GGMLBlock { + int64_t num_heads; + int64_t head_dim; + + DoubleStreamDiTBlock(int64_t hidden_size, int64_t txt_hidden_size, int64_t num_heads, int64_t head_dim, float mlp_ratio) + : num_heads(num_heads), head_dim(head_dim) { + int64_t inner_dim = num_heads * head_dim; + blocks["img_norm1"] = std::make_shared(hidden_size, 1e-6f); + blocks["img_norm2"] = std::make_shared(hidden_size, 1e-6f); + blocks["txt_norm1"] = std::make_shared(txt_hidden_size, 1e-6f); + blocks["txt_norm2"] = std::make_shared(txt_hidden_size, 1e-6f); + blocks["img_qkv"] = std::make_shared(hidden_size, inner_dim * 3, true); + blocks["txt_qkv"] = std::make_shared(txt_hidden_size, inner_dim * 3, true); + blocks["q_norm"] = std::make_shared(head_dim, 1e-6f); + blocks["k_norm"] = std::make_shared(head_dim, 1e-6f); + blocks["img_attn_proj"] = std::make_shared(inner_dim, hidden_size, true); + blocks["txt_attn_proj"] = std::make_shared(inner_dim, txt_hidden_size, true); + blocks["img_mlp"] = std::make_shared(hidden_size, static_cast(hidden_size * mlp_ratio)); + blocks["txt_mlp"] = std::make_shared(txt_hidden_size, static_cast(txt_hidden_size * mlp_ratio)); + } + + std::pair forward(GGMLRunnerContext* ctx, + ggml_tensor* img, + ggml_tensor* txt, + ggml_tensor* pe) { + auto img_norm1 = std::dynamic_pointer_cast(blocks["img_norm1"]); + auto img_norm2 = std::dynamic_pointer_cast(blocks["img_norm2"]); + auto txt_norm1 = std::dynamic_pointer_cast(blocks["txt_norm1"]); + auto txt_norm2 = std::dynamic_pointer_cast(blocks["txt_norm2"]); + auto img_qkv_p = std::dynamic_pointer_cast(blocks["img_qkv"]); + auto txt_qkv_p = std::dynamic_pointer_cast(blocks["txt_qkv"]); + auto q_norm = std::dynamic_pointer_cast(blocks["q_norm"]); + auto k_norm = std::dynamic_pointer_cast(blocks["k_norm"]); + auto img_proj = std::dynamic_pointer_cast(blocks["img_attn_proj"]); + auto txt_proj = std::dynamic_pointer_cast(blocks["txt_attn_proj"]); + auto img_mlp = std::dynamic_pointer_cast(blocks["img_mlp"]); + auto txt_mlp = std::dynamic_pointer_cast(blocks["txt_mlp"]); + + int64_t li = img->ne[1]; + int64_t lt = txt->ne[1]; + + auto img_qkv = split_qkv(ctx->ggml_ctx, img_qkv_p->forward(ctx, img_norm1->forward(ctx, img)), num_heads, head_dim); + auto txt_qkv = split_qkv(ctx->ggml_ctx, txt_qkv_p->forward(ctx, txt_norm1->forward(ctx, txt)), num_heads, head_dim); + + auto q = ggml_concat(ctx->ggml_ctx, q_norm->forward(ctx, txt_qkv[0]), q_norm->forward(ctx, img_qkv[0]), 2); + auto k = ggml_concat(ctx->ggml_ctx, k_norm->forward(ctx, txt_qkv[1]), k_norm->forward(ctx, img_qkv[1]), 2); + auto v = ggml_concat(ctx->ggml_ctx, txt_qkv[2], img_qkv[2], 2); + + auto out = Rope::attention(ctx, q, k, v, pe, nullptr, 1.0f, false); + auto out_txt = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, lt); + auto out_img = ggml_ext_slice(ctx->ggml_ctx, out, 1, lt, lt + li); + + img = ggml_add(ctx->ggml_ctx, img, img_proj->forward(ctx, out_img)); + txt = ggml_add(ctx->ggml_ctx, txt, txt_proj->forward(ctx, out_txt)); + img = ggml_add(ctx->ggml_ctx, img, img_mlp->forward(ctx, img_norm2->forward(ctx, img))); + txt = ggml_add(ctx->ggml_ctx, txt, txt_mlp->forward(ctx, txt_norm2->forward(ctx, txt))); + return {img, txt}; + } + }; + + struct FinalLayer : public GGMLBlock { + FinalLayer(int64_t hidden_size, int64_t patch_size, int64_t out_channels) { + blocks["norm_final"] = std::make_shared(hidden_size, 1e-6f); + blocks["linear"] = std::make_shared(hidden_size, patch_size * patch_size * out_channels, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto norm_final = std::dynamic_pointer_cast(blocks["norm_final"]); + auto linear = std::dynamic_pointer_cast(blocks["linear"]); + return linear->forward(ctx, norm_final->forward(ctx, x)); + } + }; + + struct MMJiT : public GGMLBlock { + MiniT2IConfig config; + + MMJiT(const MiniT2IConfig& config) + : config(config) { + blocks["img_embedder"] = std::make_shared(config.patch_size, config.in_channels, config.pca_channels, config.hidden_size); + blocks["txt_embedder"] = std::make_shared(config.txt_input_size, config.txt_hidden_size, false); + blocks["t_embedder"] = std::make_shared(config.cond_vec_size); + blocks["pooled_embedder"] = std::make_shared(config.txt_input_size, config.cond_vec_size, false); + for (int64_t i = 0; i < config.txt_preamble_depth; ++i) { + blocks["txt_preamble_blocks." + std::to_string(i)] = std::make_shared(config.txt_hidden_size, config.num_heads, config.head_dim, config.mlp_ratio); + } + for (int64_t i = 0; i < config.depth_double; ++i) { + blocks["double_blocks." + std::to_string(i)] = std::make_shared(config.hidden_size, config.txt_hidden_size, config.num_heads, config.head_dim, config.mlp_ratio); + } + blocks["final_layer"] = std::make_shared(config.hidden_size, config.patch_size, config.in_channels); + } + + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { + GGMLBlock::init_params(ctx, tensor_storage_map, prefix); + enum ggml_type wtype = get_type(prefix + "mask_token", tensor_storage_map, GGML_TYPE_F32); + params["mask_token"] = ggml_new_tensor_3d(ctx, wtype, config.txt_input_size, 1, 1); + } + + ggml_tensor* apply_text_mask(GGMLRunnerContext* ctx, ggml_tensor* context, ggml_tensor* mask) { + if (mask == nullptr) { + return context; + } + mask = ggml_reshape_3d(ctx->ggml_ctx, mask, 1, mask->ne[0], mask->ne[1]); + mask = ggml_repeat(ctx->ggml_ctx, mask, context); + auto keep = ggml_mul(ctx->ggml_ctx, context, mask); + auto inv = ggml_sub(ctx->ggml_ctx, ggml_ext_ones_like(ctx->ggml_ctx, mask), mask); + auto mask_token = ggml_repeat(ctx->ggml_ctx, params["mask_token"], context); + return ggml_add(ctx->ggml_ctx, keep, ggml_mul(ctx->ggml_ctx, mask_token, inv)); + } + + ggml_tensor* pool_context(GGMLRunnerContext* ctx, ggml_tensor* context) { + int64_t dim = context->ne[0]; + int64_t len = context->ne[1]; + int64_t N = context->ne[2]; + auto x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, context, 1, 0, 2, 3)); + x = ggml_reshape_3d(ctx->ggml_ctx, x, len, dim, N); + x = ggml_mean(ctx->ggml_ctx, x); + x = ggml_reshape_2d(ctx->ggml_ctx, x, dim, N); + return x; + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* img, + ggml_tensor* context, + ggml_tensor* mask, + ggml_tensor* pos_embed, + ggml_tensor* txt_pe, + ggml_tensor* joint_pe) { + auto img_embedder = std::dynamic_pointer_cast(blocks["img_embedder"]); + auto txt_embedder = std::dynamic_pointer_cast(blocks["txt_embedder"]); + auto final_layer = std::dynamic_pointer_cast(blocks["final_layer"]); + + int64_t W = img->ne[0]; + int64_t H = img->ne[1]; + int64_t hp = H / config.patch_size; + int64_t wp = W / config.patch_size; + + context = apply_text_mask(ctx, context, mask); + auto x = img_embedder->forward(ctx, img); + x = ggml_add(ctx->ggml_ctx, x, pos_embed); + + auto txt = txt_embedder->forward(ctx, context); + for (int64_t i = 0; i < config.txt_preamble_depth; ++i) { + auto block = std::dynamic_pointer_cast(blocks["txt_preamble_blocks." + std::to_string(i)]); + txt = block->forward(ctx, txt, txt_pe); + sd::ggml_graph_cut::mark_graph_cut(txt, "minit2i.txt_preamble_blocks." + std::to_string(i), "txt"); + } + for (int64_t i = 0; i < config.depth_double; ++i) { + auto block = std::dynamic_pointer_cast(blocks["double_blocks." + std::to_string(i)]); + auto out = block->forward(ctx, x, txt, joint_pe); + x = out.first; + txt = out.second; + sd::ggml_graph_cut::mark_graph_cut(x, "minit2i.double_blocks." + std::to_string(i), "x"); + sd::ggml_graph_cut::mark_graph_cut(txt, "minit2i.double_blocks." + std::to_string(i), "txt"); + } + auto combined = ggml_concat(ctx->ggml_ctx, txt, x, 1); + auto out = final_layer->forward(ctx, combined); + auto img_out = ggml_ext_slice(ctx->ggml_ctx, out, 1, txt->ne[1], txt->ne[1] + x->ne[1]); + return DiT::unpatchify(ctx->ggml_ctx, img_out, hp, wp, static_cast(config.patch_size), static_cast(config.patch_size), false); + } + }; + + struct MiniT2IRunner : public DiffusionModelRunner { + MiniT2IConfig config; + MMJiT model; + ggml_context* position_cache_ctx = nullptr; + ggml_backend_buffer_t position_cache_buffer = nullptr; + ggml_tensor* cached_pos_embed = nullptr; + ggml_tensor* cached_txt_pe = nullptr; + ggml_tensor* cached_joint_pe = nullptr; + int64_t cached_img_side = -1; + int64_t cached_txt_len = -1; + int64_t cached_hidden_size = -1; + int64_t cached_head_dim = -1; + + MiniT2IRunner(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "", + std::shared_ptr weight_manager = nullptr) + : DiffusionModelRunner(backend, prefix, weight_manager), + config(MiniT2IConfig::detect_from_weights(tensor_storage_map, this->prefix)), + model(config) { + model.init(params_ctx, tensor_storage_map, this->prefix); + } + + ~MiniT2IRunner() override { + free_position_cache(); + } + + std::string get_desc() override { + return "MiniT2I"; + } + + void get_param_tensors(std::map& tensors, const std::string& prefix) override { + model.get_param_tensors(tensors, prefix); + } + + void free_position_cache() { + if (position_cache_buffer != nullptr) { + ggml_backend_buffer_free(position_cache_buffer); + position_cache_buffer = nullptr; + } + if (position_cache_ctx != nullptr) { + ggml_free(position_cache_ctx); + position_cache_ctx = nullptr; + } + cached_pos_embed = nullptr; + cached_txt_pe = nullptr; + cached_joint_pe = nullptr; + cached_img_side = -1; + cached_txt_len = -1; + cached_hidden_size = -1; + cached_head_dim = -1; + } + + void ensure_position_cache(int64_t img_side, int64_t txt_len) { + if (cached_img_side == img_side && + cached_txt_len == txt_len && + cached_hidden_size == config.hidden_size && + cached_head_dim == config.head_dim && + cached_pos_embed != nullptr && + cached_txt_pe != nullptr && + cached_joint_pe != nullptr) { + return; + } + + free_position_cache(); + + auto pos_embed_vec = make_2d_sincos_pos_embed(static_cast(img_side), static_cast(config.hidden_size)); + auto txt_pe_vec = make_text_rope(static_cast(txt_len), static_cast(config.head_dim)); + auto img_pe_vec = make_vision_rope(static_cast(img_side), static_cast(config.head_dim)); + auto joint_pe_vec = txt_pe_vec; + joint_pe_vec.insert(joint_pe_vec.end(), img_pe_vec.begin(), img_pe_vec.end()); + + ggml_init_params params; + params.mem_size = static_cast(3 * ggml_tensor_overhead()); + params.mem_buffer = nullptr; + params.no_alloc = true; + position_cache_ctx = ggml_init(params); + GGML_ASSERT(position_cache_ctx != nullptr); + + cached_pos_embed = ggml_new_tensor_3d(position_cache_ctx, GGML_TYPE_F32, config.hidden_size, img_side * img_side, 1); + ggml_set_name(cached_pos_embed, "minit2i.pos_embed"); + cached_txt_pe = ggml_new_tensor_4d(position_cache_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, txt_len); + ggml_set_name(cached_txt_pe, "minit2i.txt_pe"); + cached_joint_pe = ggml_new_tensor_4d(position_cache_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, txt_len + img_side * img_side); + ggml_set_name(cached_joint_pe, "minit2i.joint_pe"); + + position_cache_buffer = ggml_backend_alloc_ctx_tensors(position_cache_ctx, runtime_backend); + GGML_ASSERT(position_cache_buffer != nullptr); + ggml_backend_buffer_set_usage(position_cache_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + ggml_backend_tensor_set(cached_pos_embed, pos_embed_vec.data(), 0, ggml_nbytes(cached_pos_embed)); + ggml_backend_tensor_set(cached_txt_pe, txt_pe_vec.data(), 0, ggml_nbytes(cached_txt_pe)); + ggml_backend_tensor_set(cached_joint_pe, joint_pe_vec.data(), 0, ggml_nbytes(cached_joint_pe)); + ggml_backend_synchronize(runtime_backend); + + cached_img_side = img_side; + cached_txt_len = txt_len; + cached_hidden_size = config.hidden_size; + cached_head_dim = config.head_dim; + } + + ggml_cgraph* build_graph(const sd::Tensor& x_tensor, + const sd::Tensor& timesteps_tensor, + const sd::Tensor& context_tensor, + const sd::Tensor& mask_tensor) { + ggml_cgraph* gf = new_graph_custom(MINIT2I_GRAPH_SIZE); + ggml_tensor* x = make_input(x_tensor); + ggml_tensor* context = make_input(context_tensor); + ggml_tensor* mask = make_input(mask_tensor); + SD_UNUSED(timesteps_tensor); + + int64_t W = x->ne[0]; + int64_t H = x->ne[1]; + int64_t img_side = H / config.patch_size; + int64_t txt_len = context->ne[1]; + ensure_position_cache(img_side, txt_len); + + auto runner_ctx = get_context(); + auto out = model.forward(&runner_ctx, x, context, mask, cached_pos_embed, cached_txt_pe, cached_joint_pe); + ggml_build_forward_expand(gf, out); + return gf; + } + + sd::Tensor compute(int n_threads, + const sd::Tensor& x, + const sd::Tensor& timesteps, + const sd::Tensor& context, + const sd::Tensor& mask) { + auto get_graph = [&]() -> ggml_cgraph* { + return build_graph(x, timesteps, context, mask); + }; + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + } + + sd::Tensor compute(int n_threads, + const DiffusionParams& diffusion_params) override { + GGML_ASSERT(diffusion_params.x != nullptr); + GGML_ASSERT(diffusion_params.timesteps != nullptr); + GGML_ASSERT(diffusion_params.context != nullptr); + const auto* extra = diffusion_extra_as(diffusion_params); + GGML_ASSERT(extra->mask != nullptr); + return compute(n_threads, + *diffusion_params.x, + *diffusion_params.timesteps, + *diffusion_params.context, + *extra->mask); + } + }; +} // namespace MiniT2I + +#endif // __SD_MODEL_DIFFUSION_MINIT2I_HPP__ diff --git a/src/mmdit.hpp b/src/model/diffusion/mmdit.hpp similarity index 81% rename from src/mmdit.hpp rename to src/model/diffusion/mmdit.hpp index 45bc2d91..d8e76dfb 100644 --- a/src/mmdit.hpp +++ b/src/model/diffusion/mmdit.hpp @@ -1,42 +1,137 @@ -#ifndef __MMDIT_HPP__ -#define __MMDIT_HPP__ +#ifndef __SD_MODEL_DIFFUSION_MMDIT_HPP__ +#define __SD_MODEL_DIFFUSION_MMDIT_HPP__ +#include #include +#include +#include -#include "diffusion_model.hpp" -#include "ggml_extend.hpp" -#include "model.h" +#include "core/ggml_extend.hpp" +#include "model/common/block.hpp" +#include "model/diffusion/model.hpp" +#include "model_loader.h" #define MMDIT_GRAPH_SIZE 10240 -struct Mlp : public GGMLBlock { -public: - Mlp(int64_t in_features, - int64_t hidden_features = -1, - int64_t out_features = -1, - bool bias = true) { - // act_layer is always lambda: nn.GELU(approximate="tanh") - // norm_layer is always None - // use_conv is always False - if (hidden_features == -1) { - hidden_features = in_features; - } - if (out_features == -1) { - out_features = in_features; - } - blocks["fc1"] = std::shared_ptr(new Linear(in_features, hidden_features, bias)); - blocks["fc2"] = std::shared_ptr(new Linear(hidden_features, out_features, bias)); - } +struct MMDiTConfig { + int64_t input_size = -1; + int patch_size = 2; + int64_t in_channels = 16; + int64_t d_self = -1; // >=0 for MMdiT-X + int64_t depth = 24; + float mlp_ratio = 4.0f; + int64_t adm_in_channels = 2048; + int64_t out_channels = 16; + int64_t pos_embed_max_size = 192; + int64_t num_patches = 36864; // 192 * 192 + int64_t context_size = 4096; + int64_t context_embedder_out_dim = 1536; + int64_t hidden_size = 1536; + std::string qk_norm; - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { - // x: [N, n_token, in_features] - auto fc1 = std::dynamic_pointer_cast(blocks["fc1"]); - auto fc2 = std::dynamic_pointer_cast(blocks["fc2"]); + static MMDiTConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) { + MMDiTConfig config; + bool has_weight_config = false; + bool has_pos_embed = false; + bool has_hidden_size = false; + bool has_context_embed = false; - x = fc1->forward(ctx, x); - x = ggml_ext_gelu(ctx->ggml_ctx, x, true); - x = fc2->forward(ctx, x); - return x; + for (const auto& [name, tensor_storage] : tensor_storage_map) { + if (!starts_with(name, prefix)) { + continue; + } + + if (name.find("x_embedder.proj.weight") != std::string::npos && tensor_storage.n_dims == 4) { + has_weight_config = true; + has_hidden_size = true; + config.patch_size = static_cast(tensor_storage.ne[0]); + config.in_channels = tensor_storage.ne[2]; + config.hidden_size = tensor_storage.ne[3]; + } else if (name.find("t_embedder.mlp.0.weight") != std::string::npos && tensor_storage.n_dims == 2) { + has_weight_config = true; + has_hidden_size = true; + config.hidden_size = tensor_storage.ne[1]; + } else if (name.find("y_embedder.mlp.0.weight") != std::string::npos && tensor_storage.n_dims == 2) { + has_weight_config = true; + has_hidden_size = true; + config.adm_in_channels = tensor_storage.ne[0]; + config.hidden_size = tensor_storage.ne[1]; + } else if (name.find("context_embedder.weight") != std::string::npos && tensor_storage.n_dims == 2) { + has_weight_config = true; + has_context_embed = true; + config.context_size = tensor_storage.ne[0]; + config.context_embedder_out_dim = tensor_storage.ne[1]; + } else if (name.find("final_layer.linear.weight") != std::string::npos && tensor_storage.n_dims == 2) { + has_weight_config = true; + has_hidden_size = true; + config.hidden_size = tensor_storage.ne[0]; + int64_t patch_area = static_cast(config.patch_size) * config.patch_size; + if (patch_area > 0) { + config.out_channels = tensor_storage.ne[1] / patch_area; + } + } else if (name.find("pos_embed") != std::string::npos && tensor_storage.n_dims == 3) { + has_weight_config = true; + has_pos_embed = true; + has_hidden_size = true; + config.hidden_size = tensor_storage.ne[0]; + config.num_patches = tensor_storage.ne[1]; + for (int64_t size = 1; size * size <= config.num_patches; size++) { + if (size * size == config.num_patches) { + config.pos_embed_max_size = size; + break; + } + } + } + + size_t jb = name.find("joint_blocks."); + if (jb == std::string::npos) { + continue; + } + + has_weight_config = true; + std::string block_name = name.substr(jb); + int64_t block_depth = atoi(block_name.substr(13, block_name.find(".", 13)).c_str()); + if (block_depth + 1 > config.depth) { + config.depth = block_depth + 1; + } + if (block_name.find("attn.ln") != std::string::npos) { + if (block_name.find(".bias") != std::string::npos) { + config.qk_norm = "ln"; + } else { + config.qk_norm = "rms"; + } + } + if (block_name.find("attn2") != std::string::npos) { + if (block_depth > config.d_self) { + config.d_self = block_depth; + } + } + } + + if (!has_pos_embed && config.d_self >= 0) { + config.pos_embed_max_size *= 2; + config.num_patches *= 4; + } + if (!has_hidden_size || config.hidden_size <= 0) { + config.hidden_size = 64 * config.depth; + } + if (!has_context_embed || config.context_embedder_out_dim <= 0) { + config.context_embedder_out_dim = config.hidden_size; + } + + if (has_weight_config) { + LOG_DEBUG("mmdit: num_layers = %" PRId64 ", num_mmdit_x_layers = %" PRId64 ", hidden_size = %" PRId64 ", patch_size = %d, in_channels = %" PRId64 ", out_channels = %" PRId64 ", context_size = %" PRId64 ", adm_in_channels = %" PRId64 ", qk_norm = %s", + config.depth, + config.d_self + 1, + config.hidden_size, + config.patch_size, + config.in_channels, + config.out_channels, + config.context_size, + config.adm_in_channels, + config.qk_norm.empty() ? "none" : config.qk_norm.c_str()); + } + return config; } }; @@ -612,28 +707,16 @@ public: struct MMDiT : public GGMLBlock { // Diffusion model with a Transformer backbone. protected: - int64_t input_size = -1; - int patch_size = 2; - int64_t in_channels = 16; - int64_t d_self = -1; // >=0 for MMdiT-X - int64_t depth = 24; - float mlp_ratio = 4.0f; - int64_t adm_in_channels = 2048; - int64_t out_channels = 16; - int64_t pos_embed_max_size = 192; - int64_t num_patchs = 36864; // 192 * 192 - int64_t context_size = 4096; - int64_t context_embedder_out_dim = 1536; - int64_t hidden_size; - std::string qk_norm; - void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") override { enum ggml_type wtype = GGML_TYPE_F32; - params["pos_embed"] = ggml_new_tensor_3d(ctx, wtype, hidden_size, num_patchs, 1); + params["pos_embed"] = ggml_new_tensor_3d(ctx, wtype, config.hidden_size, config.num_patches, 1); } public: - MMDiT(const String2TensorStorage& tensor_storage_map = {}) { + MMDiTConfig config; + + explicit MMDiT(MMDiTConfig config = {}) + : config(config) { // input_size is always None // learn_sigma is always False // register_length is alwalys 0 @@ -646,64 +729,30 @@ public: // pos_embed_offset is not used // context_embedder_config is always {'target': 'torch.nn.Linear', 'params': {'in_features': 4096, 'out_features': 1536}} - for (auto pair : tensor_storage_map) { - std::string tensor_name = pair.first; - if (tensor_name.find("model.diffusion_model.") == std::string::npos) - continue; - size_t jb = tensor_name.find("joint_blocks."); - if (jb != std::string::npos) { - tensor_name = tensor_name.substr(jb); // remove prefix - int block_depth = atoi(tensor_name.substr(13, tensor_name.find(".", 13)).c_str()); - if (block_depth + 1 > depth) { - depth = block_depth + 1; - } - if (tensor_name.find("attn.ln") != std::string::npos) { - if (tensor_name.find(".bias") != std::string::npos) { - qk_norm = "ln"; - } else { - qk_norm = "rms"; - } - } - if (tensor_name.find("attn2") != std::string::npos) { - if (block_depth > d_self) { - d_self = block_depth; - } - } - } + blocks["x_embedder"] = std::shared_ptr(new PatchEmbed(config.input_size, + config.patch_size, + config.in_channels, + config.hidden_size, + true)); + blocks["t_embedder"] = std::shared_ptr(new TimestepEmbedder(config.hidden_size)); + + if (config.adm_in_channels != -1) { + blocks["y_embedder"] = std::shared_ptr(new VectorEmbedder(config.adm_in_channels, config.hidden_size)); } - if (d_self >= 0) { - pos_embed_max_size *= 2; - num_patchs *= 4; - } + blocks["context_embedder"] = std::shared_ptr(new Linear(config.context_size, config.context_embedder_out_dim, true, true)); - LOG_INFO("MMDiT layers: %d (including %d MMDiT-x layers)", depth, d_self + 1); - - int64_t default_out_channels = in_channels; - hidden_size = 64 * depth; - context_embedder_out_dim = 64 * depth; - int64_t num_heads = depth; - - blocks["x_embedder"] = std::shared_ptr(new PatchEmbed(input_size, patch_size, in_channels, hidden_size, true)); - blocks["t_embedder"] = std::shared_ptr(new TimestepEmbedder(hidden_size)); - - if (adm_in_channels != -1) { - blocks["y_embedder"] = std::shared_ptr(new VectorEmbedder(adm_in_channels, hidden_size)); - } - - blocks["context_embedder"] = std::shared_ptr(new Linear(4096, context_embedder_out_dim, true, true)); - - for (int i = 0; i < depth; i++) { - blocks["joint_blocks." + std::to_string(i)] = std::shared_ptr(new JointBlock(hidden_size, - num_heads, - mlp_ratio, - qk_norm, + for (int i = 0; i < config.depth; i++) { + blocks["joint_blocks." + std::to_string(i)] = std::shared_ptr(new JointBlock(config.hidden_size, + config.depth, + config.mlp_ratio, + config.qk_norm, true, - i == depth - 1, - i <= d_self)); + i == config.depth - 1, + i <= config.d_self)); } - blocks["final_layer"] = std::shared_ptr(new FinalLayer(hidden_size, patch_size, out_channels)); + blocks["final_layer"] = std::shared_ptr(new FinalLayer(config.hidden_size, config.patch_size, config.out_channels)); } ggml_tensor* @@ -712,22 +761,22 @@ public: int64_t w) { auto pos_embed = params["pos_embed"]; - h = (h + 1) / patch_size; - w = (w + 1) / patch_size; + h = (h + 1) / config.patch_size; + w = (w + 1) / config.patch_size; - GGML_ASSERT(h <= pos_embed_max_size && h > 0); - GGML_ASSERT(w <= pos_embed_max_size && w > 0); + GGML_ASSERT(h <= config.pos_embed_max_size && h > 0); + GGML_ASSERT(w <= config.pos_embed_max_size && w > 0); - int64_t top = (pos_embed_max_size - h) / 2; - int64_t left = (pos_embed_max_size - w) / 2; + int64_t top = (config.pos_embed_max_size - h) / 2; + int64_t left = (config.pos_embed_max_size - w) / 2; - auto spatial_pos_embed = ggml_reshape_3d(ctx, pos_embed, hidden_size, pos_embed_max_size, pos_embed_max_size); + auto spatial_pos_embed = ggml_reshape_3d(ctx, pos_embed, config.hidden_size, config.pos_embed_max_size, config.pos_embed_max_size); // spatial_pos_embed = spatial_pos_embed[:, top : top + h, left : left + w, :] spatial_pos_embed = ggml_view_3d(ctx, spatial_pos_embed, - hidden_size, - pos_embed_max_size, + config.hidden_size, + config.pos_embed_max_size, h, spatial_pos_embed->nb[1], spatial_pos_embed->nb[2], @@ -735,14 +784,14 @@ public: spatial_pos_embed = ggml_cont(ctx, ggml_permute(ctx, spatial_pos_embed, 0, 2, 1, 3)); // [pos_embed_max_size, h, hidden_size] spatial_pos_embed = ggml_view_3d(ctx, spatial_pos_embed, - hidden_size, + config.hidden_size, h, w, spatial_pos_embed->nb[1], spatial_pos_embed->nb[2], - spatial_pos_embed->nb[2] * left); // [w, h, hidden_size] - spatial_pos_embed = ggml_cont(ctx, ggml_permute(ctx, spatial_pos_embed, 0, 2, 1, 3)); // [h, w, hidden_size] - spatial_pos_embed = ggml_reshape_3d(ctx, spatial_pos_embed, hidden_size, h * w, 1); // [1, h*w, hidden_size] + spatial_pos_embed->nb[2] * left); // [w, h, hidden_size] + spatial_pos_embed = ggml_cont(ctx, ggml_permute(ctx, spatial_pos_embed, 0, 2, 1, 3)); // [h, w, hidden_size] + spatial_pos_embed = ggml_reshape_3d(ctx, spatial_pos_embed, config.hidden_size, h * w, 1); // [1, h*w, hidden_size] return spatial_pos_embed; } @@ -757,7 +806,7 @@ public: // return: [N, N*W, patch_size * patch_size * out_channels] auto final_layer = std::dynamic_pointer_cast(blocks["final_layer"]); - for (int i = 0; i < depth; i++) { + for (int i = 0; i < config.depth; i++) { // skip iteration if i is in skip_layers if (skip_layers.size() > 0 && std::find(skip_layers.begin(), skip_layers.end(), i) != skip_layers.end()) { continue; @@ -800,7 +849,7 @@ public: x = ggml_add(ctx->ggml_ctx, patch_embed, pos_embed); // [N, H*W, hidden_size] auto c = t_embedder->forward(ctx, t); // [N, hidden_size] - if (y != nullptr && adm_in_channels != -1) { + if (y != nullptr && config.adm_in_channels != -1) { auto y_embedder = std::dynamic_pointer_cast(blocks["y_embedder"]); y = y_embedder->forward(ctx, y); // [N, hidden_size] @@ -820,19 +869,22 @@ public: x = forward_core_with_concat(ctx, x, c, context, skip_layers); // (N, H*W, patch_size ** 2 * out_channels) - x = DiT::unpatchify_and_crop(ctx->ggml_ctx, x, H, W, patch_size, patch_size, /*patch_last*/ false); // [N, C, H, W] + x = DiT::unpatchify_and_crop(ctx->ggml_ctx, x, H, W, config.patch_size, config.patch_size, /*patch_last*/ false); // [N, C, H, W] return x; } }; struct MMDiTRunner : public DiffusionModelRunner { + MMDiTConfig config; MMDiT mmdit; MMDiTRunner(ggml_backend_t backend, - ggml_backend_t params_backend, - const String2TensorStorage& tensor_storage_map = {}, - const std::string prefix = "") - : DiffusionModelRunner(backend, params_backend, prefix), mmdit(tensor_storage_map) { + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "", + std::shared_ptr weight_manager = nullptr) + : DiffusionModelRunner(backend, prefix, weight_manager), + config(MMDiTConfig::detect_from_weights(tensor_storage_map, prefix)), + mmdit(config) { mmdit.init(params_ctx, tensor_storage_map, prefix); } @@ -883,7 +935,7 @@ struct MMDiTRunner : public DiffusionModelRunner { return build_graph(x, timesteps, context, y, skip_layers); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); } sd::Tensor compute(int n_threads, @@ -949,28 +1001,25 @@ struct MMDiTRunner : public DiffusionModelRunner { // ggml_backend_t backend = ggml_backend_cuda_init(0); ggml_backend_t backend = sd_backend_cpu_init(); ggml_type model_data_type = GGML_TYPE_F16; - std::shared_ptr mmdit = std::make_shared(backend, backend); + auto model_manager = std::make_shared(); + std::shared_ptr mmdit = std::make_shared(backend, String2TensorStorage{}, "", model_manager); { LOG_INFO("loading from '%s'", file_path.c_str()); - if (!mmdit->alloc_params_buffer()) { - LOG_ERROR("mmdit embeds buffer allocation failed"); - return; - } - - std::map tensors; - mmdit->get_param_tensors(tensors, "model.diffusion_model"); - - ModelLoader model_loader; + ModelLoader& model_loader = model_manager->loader(); if (!model_loader.init_from_file_and_convert_name(file_path)) { LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str()); return; } - bool success = model_loader.load_tensors(tensors); - - if (!success) { - LOG_ERROR("load tensors from model loader failed"); + if (!model_manager->register_runner_params("MMDiT test", + *mmdit, + "model.diffusion_model", + ModelManager::ResidencyMode::ParamBackend, + backend, + backend) || + !model_manager->validate_registered_tensors()) { + LOG_ERROR("register mmdit tensors with model manager failed"); return; } @@ -980,4 +1029,4 @@ struct MMDiTRunner : public DiffusionModelRunner { } }; -#endif +#endif // __SD_MODEL_DIFFUSION_MMDIT_HPP__ diff --git a/src/diffusion_model.hpp b/src/model/diffusion/model.hpp similarity index 82% rename from src/diffusion_model.hpp rename to src/model/diffusion/model.hpp index bb761082..8ef00023 100644 --- a/src/diffusion_model.hpp +++ b/src/model/diffusion/model.hpp @@ -1,13 +1,14 @@ -#ifndef __DIFFUSION_MODEL_H__ -#define __DIFFUSION_MODEL_H__ +#ifndef __SD_MODEL_DIFFUSION_MODEL_HPP__ +#define __SD_MODEL_DIFFUSION_MODEL_HPP__ #include #include #include -#include "ggml_extend.hpp" -#include "rope.hpp" -#include "tensor_ggml.hpp" +#include "core/ggml_extend.hpp" +#include "core/tensor_ggml.hpp" +#include "model/common/rope.hpp" +#include "model_manager.h" struct UNetDiffusionExtra { int num_video_frames = -1; @@ -22,6 +23,8 @@ struct SkipLayerDiffusionExtra { struct FluxDiffusionExtra { const sd::Tensor* guidance = nullptr; const std::vector* skip_layers = nullptr; + const sd::Tensor* pulid_id = nullptr; + float pulid_id_weight = 1.0f; }; struct AnimaDiffusionExtra { @@ -50,6 +53,10 @@ struct LTXAVDiffusionExtra { const sd::Tensor* video_positions = nullptr; }; +struct MiniT2IDiffusionExtra { + const sd::Tensor* mask = nullptr; +}; + using DiffusionExtraParams = std::variant; + LTXAVDiffusionExtra, + MiniT2IDiffusionExtra>; struct DiffusionParams { const sd::Tensor* x = nullptr; @@ -89,9 +97,9 @@ protected: public: DiffusionModelRunner(ggml_backend_t backend, - ggml_backend_t params_backend, - const std::string& prefix) - : GGMLRunner(backend, params_backend), + const std::string& prefix, + std::shared_ptr weight_manager = nullptr) + : GGMLRunner(backend, weight_manager), prefix(prefix) {} virtual sd::Tensor compute(int n_threads, @@ -105,4 +113,4 @@ public: const std::string& prefix) = 0; }; -#endif +#endif // __SD_MODEL_DIFFUSION_MODEL_HPP__ diff --git a/src/pid.hpp b/src/model/diffusion/pid.hpp similarity index 84% rename from src/pid.hpp rename to src/model/diffusion/pid.hpp index c29c207f..68dca00f 100644 --- a/src/pid.hpp +++ b/src/model/diffusion/pid.hpp @@ -1,5 +1,5 @@ -#ifndef __SD_PID_HPP__ -#define __SD_PID_HPP__ +#ifndef __SD_MODEL_DIFFUSION_PID_HPP__ +#define __SD_MODEL_DIFFUSION_PID_HPP__ #include #include @@ -7,16 +7,16 @@ #include #include -#include "common_dit.hpp" -#include "ggml_extend.hpp" -#include "mmdit.hpp" -#include "rope.hpp" +#include "core/ggml_extend.hpp" +#include "model/common/rope.hpp" +#include "model/diffusion/dit.hpp" +#include "model/diffusion/mmdit.hpp" namespace Pid { constexpr int PID_GRAPH_SIZE = 196608; constexpr float PID_PI = 3.14159265358979323846f; - struct PixelDiTParams { + struct PixelDiTConfig { int64_t in_channels = 3; int64_t hidden_size = 1536; int64_t num_groups = 24; @@ -38,6 +38,45 @@ namespace Pid { int64_t lq_latent_down_factor = 8; int64_t rope_ref_grid_h = 64; int64_t rope_ref_grid_w = 64; + + static PixelDiTConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) { + PixelDiTConfig config; + for (const auto& [name, tensor_storage] : tensor_storage_map) { + if (!starts_with(name, prefix)) { + continue; + } + size_t pos = name.find("patch_blocks."); + if (pos != std::string::npos) { + auto items = split_string(name.substr(pos), '.'); + if (items.size() > 1) { + int block_index = atoi(items[1].c_str()); + config.patch_depth = std::max(config.patch_depth, block_index + 1); + } + } + pos = name.find("pixel_blocks."); + if (pos != std::string::npos) { + auto items = split_string(name.substr(pos), '.'); + if (items.size() > 1) { + int block_index = atoi(items[1].c_str()); + config.pixel_depth = std::max(config.pixel_depth, block_index + 1); + } + } + if (name.find("lq_proj.latent_proj.0.weight") != std::string::npos) { + config.lq_latent_channels = tensor_storage.ne[2]; + config.lq_latent_down_factor = config.lq_latent_channels >= 64 ? 16 : 8; + } + if (name.find("patch_blocks.0.mlp_x.w1.weight") != std::string::npos) { + config.patch_mlp_hidden_dim = tensor_storage.ne[1]; + } + } + LOG_DEBUG("pid: patch_depth = %" PRId64 ", pixel_depth = %" PRId64 ", patch_mlp_hidden_dim = %" PRId64 ", lq_latent_channels = %" PRId64 ", lq_latent_down_factor = %" PRId64, + config.patch_depth, + config.pixel_depth, + config.patch_mlp_hidden_dim, + config.lq_latent_channels, + config.lq_latent_down_factor); + return config; + } }; inline std::vector make_rope_1d(int length, @@ -466,29 +505,29 @@ namespace Pid { }; struct LQProjection2D : public GGMLBlock { - PixelDiTParams params_cfg; + PixelDiTConfig config; - LQProjection2D(const PixelDiTParams& params_cfg) - : params_cfg(params_cfg) { - blocks["latent_proj.0"] = std::make_shared(params_cfg.lq_latent_channels, params_cfg.lq_hidden_dim, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1}); - blocks["latent_proj.2"] = std::make_shared(params_cfg.lq_hidden_dim, params_cfg.lq_hidden_dim, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1}); - for (int i = 0; i < params_cfg.lq_num_res_blocks; ++i) { - blocks["latent_proj." + std::to_string(3 + i)] = std::make_shared(params_cfg.lq_hidden_dim); + LQProjection2D(const PixelDiTConfig& config) + : config(config) { + blocks["latent_proj.0"] = std::make_shared(config.lq_latent_channels, config.lq_hidden_dim, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1}); + blocks["latent_proj.2"] = std::make_shared(config.lq_hidden_dim, config.lq_hidden_dim, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1}); + for (int i = 0; i < config.lq_num_res_blocks; ++i) { + blocks["latent_proj." + std::to_string(3 + i)] = std::make_shared(config.lq_hidden_dim); } - int num_outputs = static_cast((params_cfg.patch_depth + params_cfg.lq_interval - 1) / params_cfg.lq_interval); + int num_outputs = static_cast((config.patch_depth + config.lq_interval - 1) / config.lq_interval); for (int i = 0; i < num_outputs; ++i) { - blocks["output_heads." + std::to_string(i)] = std::make_shared(params_cfg.lq_hidden_dim, params_cfg.hidden_size, true); - blocks["gate_modules." + std::to_string(i)] = std::make_shared(params_cfg.hidden_size); + blocks["output_heads." + std::to_string(i)] = std::make_shared(config.lq_hidden_dim, config.hidden_size, true); + blocks["gate_modules." + std::to_string(i)] = std::make_shared(config.hidden_size); } } bool is_gate_active(int block_idx) const { - return block_idx % params_cfg.lq_interval == 0; + return block_idx % config.lq_interval == 0; } int get_output_index(int block_idx) const { - return block_idx / static_cast(params_cfg.lq_interval); + return block_idx / static_cast(config.lq_interval); } ggml_tensor* gate(GGMLRunnerContext* ctx, @@ -506,8 +545,8 @@ namespace Pid { int64_t target_pW) { auto conv0 = std::dynamic_pointer_cast(blocks["latent_proj.0"]); auto conv2 = std::dynamic_pointer_cast(blocks["latent_proj.2"]); - float z_to_patch_ratio = static_cast(params_cfg.lq_sr_scale * params_cfg.lq_latent_down_factor) / - static_cast(params_cfg.patch_size); + float z_to_patch_ratio = static_cast(config.lq_sr_scale * config.lq_latent_down_factor) / + static_cast(config.patch_size); GGML_ASSERT(z_to_patch_ratio >= 1.0f); if (lq_latent->ne[0] != target_pW || lq_latent->ne[1] != target_pH) { lq_latent = ggml_interpolate(ctx->ggml_ctx, @@ -522,7 +561,7 @@ namespace Pid { auto feat = conv0->forward(ctx, lq_latent); feat = ggml_silu_inplace(ctx->ggml_ctx, feat); feat = conv2->forward(ctx, feat); - for (int i = 0; i < params_cfg.lq_num_res_blocks; ++i) { + for (int i = 0; i < config.lq_num_res_blocks; ++i) { auto block = std::dynamic_pointer_cast(blocks["latent_proj." + std::to_string(3 + i)]); feat = block->forward(ctx, feat); } @@ -533,7 +572,7 @@ namespace Pid { auto tokens = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, feat, 2, 0, 1, 3)); tokens = ggml_reshape_3d(ctx->ggml_ctx, tokens, C, L, B); - int num_outputs = static_cast((params_cfg.patch_depth + params_cfg.lq_interval - 1) / params_cfg.lq_interval); + int num_outputs = static_cast((config.patch_depth + config.lq_interval - 1) / config.lq_interval); std::vector outputs; outputs.reserve(num_outputs); for (int i = 0; i < num_outputs; ++i) { @@ -545,34 +584,34 @@ namespace Pid { }; struct PixelDiT : public GGMLBlock { - PixelDiTParams params_cfg; + PixelDiTConfig config; PixelDiT() = default; - PixelDiT(const PixelDiTParams& params_cfg) - : params_cfg(params_cfg) { - blocks["pixel_embedder"] = std::make_shared(params_cfg.in_channels, params_cfg.pixel_hidden_size); - blocks["s_embedder"] = std::make_shared(params_cfg.in_channels * params_cfg.patch_size * params_cfg.patch_size, params_cfg.hidden_size, false, true); - blocks["t_embedder"] = std::make_shared(params_cfg.hidden_size); - blocks["y_embedder"] = std::make_shared(params_cfg.txt_embed_dim, params_cfg.hidden_size, true, true); - for (int i = 0; i < params_cfg.patch_depth; ++i) { - blocks["patch_blocks." + std::to_string(i)] = std::make_shared(params_cfg.hidden_size, params_cfg.num_groups, params_cfg.patch_mlp_hidden_dim); + PixelDiT(const PixelDiTConfig& config) + : config(config) { + blocks["pixel_embedder"] = std::make_shared(config.in_channels, config.pixel_hidden_size); + blocks["s_embedder"] = std::make_shared(config.in_channels * config.patch_size * config.patch_size, config.hidden_size, false, true); + blocks["t_embedder"] = std::make_shared(config.hidden_size); + blocks["y_embedder"] = std::make_shared(config.txt_embed_dim, config.hidden_size, true, true); + for (int i = 0; i < config.patch_depth; ++i) { + blocks["patch_blocks." + std::to_string(i)] = std::make_shared(config.hidden_size, config.num_groups, config.patch_mlp_hidden_dim); } - for (int i = 0; i < params_cfg.pixel_depth; ++i) { - blocks["pixel_blocks." + std::to_string(i)] = std::make_shared(params_cfg.pixel_hidden_size, - params_cfg.hidden_size, - params_cfg.patch_size, - params_cfg.pixel_attn_hidden_size, - params_cfg.pixel_num_groups); + for (int i = 0; i < config.pixel_depth; ++i) { + blocks["pixel_blocks." + std::to_string(i)] = std::make_shared(config.pixel_hidden_size, + config.hidden_size, + config.patch_size, + config.pixel_attn_hidden_size, + config.pixel_num_groups); } - blocks["final_layer"] = std::make_shared(params_cfg.pixel_hidden_size, params_cfg.in_channels); - blocks["lq_proj"] = std::make_shared(params_cfg); + blocks["final_layer"] = std::make_shared(config.pixel_hidden_size, config.in_channels); + blocks["lq_proj"] = std::make_shared(config); } void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") override { - params["y_pos_embedding"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, params_cfg.hidden_size, params_cfg.txt_max_length, 1); + params["y_pos_embedding"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, config.hidden_size, config.txt_max_length, 1); } ggml_tensor* forward(GGMLRunnerContext* ctx, @@ -594,21 +633,21 @@ namespace Pid { int64_t W_orig = x->ne[0]; int64_t H_orig = x->ne[1]; - x = DiT::pad_to_patch_size(ctx, x, static_cast(params_cfg.patch_size), static_cast(params_cfg.patch_size)); + x = DiT::pad_to_patch_size(ctx, x, static_cast(config.patch_size), static_cast(config.patch_size)); int64_t W = x->ne[0]; int64_t H = x->ne[1]; int64_t B = x->ne[3]; - int64_t Hs = H / params_cfg.patch_size; - int64_t Ws = W / params_cfg.patch_size; + int64_t Hs = H / config.patch_size; + int64_t Ws = W / config.patch_size; int64_t L = Hs * Ws; - int64_t P2 = params_cfg.patch_size * params_cfg.patch_size; + int64_t P2 = config.patch_size * config.patch_size; - auto x_patches = DiT::patchify(ctx->ggml_ctx, x, static_cast(params_cfg.patch_size), static_cast(params_cfg.patch_size), true); + auto x_patches = DiT::patchify(ctx->ggml_ctx, x, static_cast(config.patch_size), static_cast(config.patch_size), true); auto t_emb = t_embedder->forward(ctx, timesteps); auto condition = ggml_silu(ctx->ggml_ctx, t_emb); GGML_ASSERT(context != nullptr); - int64_t Ltxt = std::min(context->ne[1], params_cfg.txt_max_length); + int64_t Ltxt = std::min(context->ne[1], config.txt_max_length); auto y = ggml_ext_slice(ctx->ggml_ctx, context, 1, 0, Ltxt); auto y_emb = y_embedder->forward(ctx, y); auto y_pos = ggml_ext_slice(ctx->ggml_ctx, params["y_pos_embedding"], 1, 0, Ltxt); @@ -618,7 +657,7 @@ namespace Pid { auto s = s_embedder->forward(ctx, x_patches); - for (int i = 0; i < params_cfg.patch_depth; ++i) { + for (int i = 0; i < config.patch_depth; ++i) { if (lq_proj->is_gate_active(i)) { int out_idx = lq_proj->get_output_index(i); if (out_idx < static_cast(lq_features.size())) { @@ -639,22 +678,22 @@ namespace Pid { } s = ggml_silu(ctx->ggml_ctx, ggml_add(ctx->ggml_ctx, s, t_emb)); - auto s_cond = ggml_reshape_2d(ctx->ggml_ctx, s, params_cfg.hidden_size, L * B); - auto pixels = pixel_embedder->forward(ctx, x, params_cfg.patch_size, pixel_pos_full); - for (int i = 0; i < params_cfg.pixel_depth; ++i) { + auto s_cond = ggml_reshape_2d(ctx->ggml_ctx, s, config.hidden_size, L * B); + auto pixels = pixel_embedder->forward(ctx, x, config.patch_size, pixel_pos_full); + for (int i = 0; i < config.pixel_depth; ++i) { auto block = std::dynamic_pointer_cast(blocks["pixel_blocks." + std::to_string(i)]); pixels = block->forward(ctx, pixels, s_cond, H, W, pixel_pos_comp); sd::ggml_graph_cut::mark_graph_cut(pixels, "pid.pixel_blocks." + std::to_string(i), "pixels"); } pixels = final_layer->forward(ctx, pixels); - pixels = ggml_reshape_3d(ctx->ggml_ctx, pixels, params_cfg.in_channels * P2, L, B); + pixels = ggml_reshape_3d(ctx->ggml_ctx, pixels, config.in_channels * P2, L, B); auto out = DiT::unpatchify(ctx->ggml_ctx, pixels, Hs, Ws, - static_cast(params_cfg.patch_size), - static_cast(params_cfg.patch_size), + static_cast(config.patch_size), + static_cast(config.patch_size), false); out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, H_orig); out = ggml_ext_slice(ctx->ggml_ctx, out, 0, 0, W_orig); @@ -663,7 +702,7 @@ namespace Pid { }; struct PiDRunner : public DiffusionModelRunner { - PixelDiTParams params_cfg; + PixelDiTConfig config; PixelDiT model; std::vector pos_img_vec; std::vector pos_txt_vec; @@ -671,46 +710,12 @@ namespace Pid { std::vector pixel_pos_comp_vec; PiDRunner(ggml_backend_t backend, - ggml_backend_t params_backend, const String2TensorStorage& tensor_storage_map, - const std::string prefix = "model.diffusion_model") - : DiffusionModelRunner(backend, params_backend, prefix) { - for (const auto& pair : tensor_storage_map) { - const std::string& tensor_name = pair.first; - if (tensor_name.find(prefix) == std::string::npos) { - continue; - } - size_t pos = tensor_name.find("patch_blocks."); - if (pos != std::string::npos) { - auto items = split_string(tensor_name.substr(pos), '.'); - if (items.size() > 1) { - int block_index = atoi(items[1].c_str()); - params_cfg.patch_depth = std::max(params_cfg.patch_depth, block_index + 1); - } - } - pos = tensor_name.find("pixel_blocks."); - if (pos != std::string::npos) { - auto items = split_string(tensor_name.substr(pos), '.'); - if (items.size() > 1) { - int block_index = atoi(items[1].c_str()); - params_cfg.pixel_depth = std::max(params_cfg.pixel_depth, block_index + 1); - } - } - if (tensor_name.find("lq_proj.latent_proj.0.weight") != std::string::npos) { - params_cfg.lq_latent_channels = pair.second.ne[2]; - params_cfg.lq_latent_down_factor = params_cfg.lq_latent_channels >= 64 ? 16 : 8; - } - if (tensor_name.find("patch_blocks.0.mlp_x.w1.weight") != std::string::npos) { - params_cfg.patch_mlp_hidden_dim = pair.second.ne[1]; - } - } - LOG_INFO("PiD params: patch_depth=%" PRId64 ", pixel_depth=%" PRId64 ", patch_mlp_hidden_dim=%" PRId64 ", lq_latent_channels=%" PRId64 ", lq_latent_down_factor=%" PRId64, - params_cfg.patch_depth, - params_cfg.pixel_depth, - params_cfg.patch_mlp_hidden_dim, - params_cfg.lq_latent_channels, - params_cfg.lq_latent_down_factor); - model = PixelDiT(params_cfg); + const std::string prefix = "model.diffusion_model", + std::shared_ptr weight_manager = nullptr) + : DiffusionModelRunner(backend, prefix, weight_manager), + config(PixelDiTConfig::detect_from_weights(tensor_storage_map, prefix)) { + model = PixelDiT(config); model.init(params_ctx, tensor_storage_map, prefix); } @@ -737,60 +742,60 @@ namespace Pid { int64_t W = x->ne[0]; int64_t H = x->ne[1]; int64_t B = x->ne[3]; - int64_t Wp = align_up(static_cast(W), static_cast(params_cfg.patch_size)); - int64_t Hp = align_up(static_cast(H), static_cast(params_cfg.patch_size)); - int64_t Hs = Hp / params_cfg.patch_size; - int64_t Ws = Wp / params_cfg.patch_size; + int64_t Wp = align_up(static_cast(W), static_cast(config.patch_size)); + int64_t Hp = align_up(static_cast(H), static_cast(config.patch_size)); + int64_t Hs = Hp / config.patch_size; + int64_t Ws = Wp / config.patch_size; pos_img_vec = make_rope_2d(static_cast(Hs), static_cast(Ws), - static_cast(params_cfg.hidden_size / params_cfg.num_groups), + static_cast(config.hidden_size / config.num_groups), 10000.f, 16.f, - static_cast(params_cfg.rope_ref_grid_h), - static_cast(params_cfg.rope_ref_grid_w)); + static_cast(config.rope_ref_grid_h), + static_cast(config.rope_ref_grid_w)); auto pos_img = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, - params_cfg.hidden_size / params_cfg.num_groups / 2, + config.hidden_size / config.num_groups / 2, Hs * Ws); set_backend_tensor_data(pos_img, pos_img_vec.data()); - int64_t Ltxt = std::min(context->ne[1], params_cfg.txt_max_length); + int64_t Ltxt = std::min(context->ne[1], config.txt_max_length); pos_txt_vec = make_rope_1d(static_cast(Ltxt), - static_cast(params_cfg.hidden_size / params_cfg.num_groups), - params_cfg.text_rope_theta); + static_cast(config.hidden_size / config.num_groups), + config.text_rope_theta); auto pos_txt = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, - params_cfg.hidden_size / params_cfg.num_groups / 2, + config.hidden_size / config.num_groups / 2, Ltxt); set_backend_tensor_data(pos_txt, pos_txt_vec.data()); pixel_pos_vec = make_pixel_abs_pos(static_cast(Hp), static_cast(Wp), - static_cast(params_cfg.pixel_hidden_size)); + static_cast(config.pixel_hidden_size)); auto pixel_pos = ggml_new_tensor_3d(compute_ctx, GGML_TYPE_F32, - params_cfg.pixel_hidden_size, + config.pixel_hidden_size, Wp * Hp, 1); set_backend_tensor_data(pixel_pos, pixel_pos_vec.data()); pixel_pos_comp_vec = make_rope_2d(static_cast(Hs), static_cast(Ws), - static_cast(params_cfg.pixel_attn_hidden_size / params_cfg.pixel_num_groups), + static_cast(config.pixel_attn_hidden_size / config.pixel_num_groups), 10000.f, 16.f, - static_cast(params_cfg.rope_ref_grid_h), - static_cast(params_cfg.rope_ref_grid_w)); + static_cast(config.rope_ref_grid_h), + static_cast(config.rope_ref_grid_w)); auto pixel_pos_comp = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, - params_cfg.pixel_attn_hidden_size / params_cfg.pixel_num_groups / 2, + config.pixel_attn_hidden_size / config.pixel_num_groups / 2, Hs * Ws); set_backend_tensor_data(pixel_pos_comp, pixel_pos_comp_vec.data()); @@ -818,7 +823,7 @@ namespace Pid { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, timesteps, context, lq_latent, degrade_sigma); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); } sd::Tensor compute(int n_threads, @@ -839,4 +844,4 @@ namespace Pid { }; } // namespace Pid -#endif // __SD_PID_HPP__ +#endif // __SD_MODEL_DIFFUSION_PID_HPP__ diff --git a/src/qwen_image.hpp b/src/model/diffusion/qwen_image.hpp similarity index 87% rename from src/qwen_image.hpp rename to src/model/diffusion/qwen_image.hpp index 9e731b37..e9414c6d 100644 --- a/src/qwen_image.hpp +++ b/src/model/diffusion/qwen_image.hpp @@ -1,16 +1,60 @@ -#ifndef __QWEN_IMAGE_HPP__ -#define __QWEN_IMAGE_HPP__ +#ifndef __SD_MODEL_DIFFUSION_QWEN_IMAGE_HPP__ +#define __SD_MODEL_DIFFUSION_QWEN_IMAGE_HPP__ #include -#include "common_block.hpp" -#include "common_dit.hpp" -#include "diffusion_model.hpp" -#include "flux.hpp" +#include "model/common/block.hpp" +#include "model/diffusion/dit.hpp" +#include "model/diffusion/flux.hpp" +#include "model/diffusion/model.hpp" +#include "model_loader.h" namespace Qwen { constexpr int QWEN_IMAGE_GRAPH_SIZE = 20480; + struct QwenImageConfig { + int patch_size = 2; + int64_t in_channels = 64; + int64_t out_channels = 16; + int num_layers = 60; + int64_t attention_head_dim = 128; + int64_t num_attention_heads = 24; + int64_t joint_attention_dim = 3584; + int theta = 10000; + std::vector axes_dim = {16, 56, 56}; + int axes_dim_sum = 128; + bool zero_cond_t = false; + bool use_additional_t_cond = false; + + static QwenImageConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) { + QwenImageConfig config; + config.num_layers = 0; + for (const auto& [name, _] : tensor_storage_map) { + if (!starts_with(name, prefix)) { + continue; + } + if (name.find("__index_timestep_zero__") != std::string::npos) { + config.zero_cond_t = true; + } + size_t pos = name.find("transformer_blocks."); + if (pos == std::string::npos) { + continue; + } + auto items = split_string(name.substr(pos), '.'); + if (items.size() > 1) { + int block_index = atoi(items[1].c_str()); + if (block_index + 1 > config.num_layers) { + config.num_layers = block_index + 1; + } + } + } + LOG_DEBUG("qwen_image: num_layers = %d, zero_cond_t = %s", + config.num_layers, + config.zero_cond_t ? "true" : "false"); + return config; + } + }; + struct TimestepEmbedding : public GGMLBlock { public: TimestepEmbedding(int64_t in_channels, @@ -365,47 +409,32 @@ namespace Qwen { } }; - struct QwenImageParams { - int patch_size = 2; - int64_t in_channels = 64; - int64_t out_channels = 16; - int num_layers = 60; - int64_t attention_head_dim = 128; - int64_t num_attention_heads = 24; - int64_t joint_attention_dim = 3584; - int theta = 10000; - std::vector axes_dim = {16, 56, 56}; - int axes_dim_sum = 128; - bool zero_cond_t = false; - bool use_additional_t_cond = false; - }; - class QwenImageModel : public GGMLBlock { protected: - QwenImageParams params; + QwenImageConfig config; public: QwenImageModel() {} - QwenImageModel(QwenImageParams params) - : params(params) { - int64_t inner_dim = params.num_attention_heads * params.attention_head_dim; - blocks["time_text_embed"] = std::shared_ptr(new QwenTimestepProjEmbeddings(inner_dim, params.use_additional_t_cond)); - blocks["txt_norm"] = std::shared_ptr(new RMSNorm(params.joint_attention_dim, 1e-6f)); - blocks["img_in"] = std::shared_ptr(new Linear(params.in_channels, inner_dim)); - blocks["txt_in"] = std::shared_ptr(new Linear(params.joint_attention_dim, inner_dim)); + QwenImageModel(QwenImageConfig config) + : config(config) { + int64_t inner_dim = config.num_attention_heads * config.attention_head_dim; + blocks["time_text_embed"] = std::shared_ptr(new QwenTimestepProjEmbeddings(inner_dim, config.use_additional_t_cond)); + blocks["txt_norm"] = std::shared_ptr(new RMSNorm(config.joint_attention_dim, 1e-6f)); + blocks["img_in"] = std::shared_ptr(new Linear(config.in_channels, inner_dim)); + blocks["txt_in"] = std::shared_ptr(new Linear(config.joint_attention_dim, inner_dim)); // blocks - for (int i = 0; i < params.num_layers; i++) { + for (int i = 0; i < config.num_layers; i++) { auto block = std::shared_ptr(new QwenImageTransformerBlock(inner_dim, - params.num_attention_heads, - params.attention_head_dim, + config.num_attention_heads, + config.attention_head_dim, 1e-6f, - params.zero_cond_t)); + config.zero_cond_t)); blocks["transformer_blocks." + std::to_string(i)] = block; } blocks["norm_out"] = std::shared_ptr(new AdaLayerNormContinuous(inner_dim, inner_dim, false, 1e-6f)); - blocks["proj_out"] = std::shared_ptr(new Linear(inner_dim, params.patch_size * params.patch_size * params.out_channels)); + blocks["proj_out"] = std::shared_ptr(new Linear(inner_dim, config.patch_size * config.patch_size * config.out_channels)); } ggml_tensor* forward_orig(GGMLRunnerContext* ctx, @@ -423,7 +452,7 @@ namespace Qwen { auto proj_out = std::dynamic_pointer_cast(blocks["proj_out"]); auto t_emb = time_text_embed->forward(ctx, timestep, addition_t_cond); - if (params.zero_cond_t) { + if (config.zero_cond_t) { auto t_emb_0 = time_text_embed->forward(ctx, ggml_ext_zeros_like(ctx->ggml_ctx, timestep), addition_t_cond); t_emb = ggml_concat(ctx->ggml_ctx, t_emb, t_emb_0, 1); } @@ -434,7 +463,7 @@ namespace Qwen { sd::ggml_graph_cut::mark_graph_cut(txt, "qwen_image.prelude", "txt"); // sd::ggml_graph_cut::mark_graph_cut(t_emb, "qwen_image.prelude", "t_emb"); - for (int i = 0; i < params.num_layers; i++) { + for (int i = 0; i < config.num_layers; i++) { auto block = std::dynamic_pointer_cast(blocks["transformer_blocks." + std::to_string(i)]); auto result = block->forward(ctx, img, txt, t_emb, pe, modulate_index); @@ -444,7 +473,7 @@ namespace Qwen { sd::ggml_graph_cut::mark_graph_cut(txt, "qwen_image.transformer_blocks." + std::to_string(i), "txt"); } - if (params.zero_cond_t) { + if (config.zero_cond_t) { t_emb = ggml_ext_chunk(ctx->ggml_ctx, t_emb, 2, 1)[0]; } @@ -480,14 +509,14 @@ namespace Qwen { } auto patchify_input = [&](ggml_tensor* input) -> ggml_tensor* { - input = DiT::pad_to_patch_size(ctx, input, params.patch_size, params.patch_size); + input = DiT::pad_to_patch_size(ctx, input, config.patch_size, config.patch_size); if (!has_time_axis) { - return DiT::patchify(ctx->ggml_ctx, input, params.patch_size, params.patch_size); + return DiT::patchify(ctx->ggml_ctx, input, config.patch_size, config.patch_size); } if (input->ne[3] == 1) { input = ggml_reshape_4d(ctx->ggml_ctx, input, input->ne[0], input->ne[1], 1, input->ne[2]); } - return DiT::patchify(ctx->ggml_ctx, input, 1, params.patch_size, params.patch_size, N); + return DiT::patchify(ctx->ggml_ctx, input, 1, config.patch_size, config.patch_size, N); }; auto img = patchify_input(x); @@ -509,15 +538,15 @@ namespace Qwen { } if (has_time_axis) { - int pad_h = (params.patch_size - H % params.patch_size) % params.patch_size; - int pad_w = (params.patch_size - W % params.patch_size) % params.patch_size; - int h_len = static_cast((H + pad_h) / params.patch_size); - int w_len = static_cast((W + pad_w) / params.patch_size); - out = DiT::unpatchify(ctx->ggml_ctx, out, T, h_len, w_len, 1, params.patch_size, params.patch_size); + int pad_h = (config.patch_size - H % config.patch_size) % config.patch_size; + int pad_w = (config.patch_size - W % config.patch_size) % config.patch_size; + int h_len = static_cast((H + pad_h) / config.patch_size); + int w_len = static_cast((W + pad_w) / config.patch_size); + out = DiT::unpatchify(ctx->ggml_ctx, out, T, h_len, w_len, 1, config.patch_size, config.patch_size); out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, H); // [N*C, T, H, W + pad_w] out = ggml_ext_slice(ctx->ggml_ctx, out, 0, 0, W); // [N*C, T, H, W] } else { - out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, params.patch_size, params.patch_size); // [N, C, H, W] + out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, config.patch_size, config.patch_size); // [N, C, H, W] } return out; @@ -526,7 +555,7 @@ namespace Qwen { struct QwenImageRunner : public DiffusionModelRunner { public: - QwenImageParams qwen_image_params; + QwenImageConfig config; QwenImageModel qwen_image; std::vector pe_vec; std::vector modulate_index_vec; @@ -534,43 +563,19 @@ namespace Qwen { SDVersion version; QwenImageRunner(ggml_backend_t backend, - ggml_backend_t params_backend, - const String2TensorStorage& tensor_storage_map = {}, - const std::string prefix = "", - SDVersion version = VERSION_QWEN_IMAGE, - bool zero_cond_t = false) - : DiffusionModelRunner(backend, params_backend, prefix), + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "", + SDVersion version = VERSION_QWEN_IMAGE, + bool zero_cond_t = false, + std::shared_ptr weight_manager = nullptr) + : DiffusionModelRunner(backend, prefix, weight_manager), + config(QwenImageConfig::detect_from_weights(tensor_storage_map, prefix)), version(version) { - qwen_image_params.num_layers = 0; - qwen_image_params.zero_cond_t = zero_cond_t; - for (auto pair : tensor_storage_map) { - std::string tensor_name = pair.first; - if (tensor_name.find(prefix) == std::string::npos) - continue; - if (tensor_name.find("__index_timestep_zero__") != std::string::npos) { - qwen_image_params.zero_cond_t = true; - } - size_t pos = tensor_name.find("transformer_blocks."); - if (pos != std::string::npos) { - tensor_name = tensor_name.substr(pos); // remove prefix - auto items = split_string(tensor_name, '.'); - if (items.size() > 1) { - int block_index = atoi(items[1].c_str()); - if (block_index + 1 > qwen_image_params.num_layers) { - qwen_image_params.num_layers = block_index + 1; - } - } - continue; - } - } + config.zero_cond_t = config.zero_cond_t || zero_cond_t; if (version == VERSION_QWEN_IMAGE_LAYERED) { - qwen_image_params.use_additional_t_cond = true; + config.use_additional_t_cond = true; } - LOG_INFO("qwen_image_params.num_layers: %ld", qwen_image_params.num_layers); - if (qwen_image_params.zero_cond_t) { - LOG_INFO("use zero_cond_t"); - } - qwen_image = QwenImageModel(qwen_image_params); + qwen_image = QwenImageModel(config); qwen_image.init(params_ctx, tensor_storage_map, prefix); } @@ -617,36 +622,36 @@ namespace Qwen { pe_vec = Rope::gen_qwen_image_pe(time_len, static_cast(x->ne[1]), static_cast(x->ne[0]), - qwen_image_params.patch_size, + config.patch_size, batch_size, static_cast(context->ne[1]), ref_latents, ref_index_mode, - qwen_image_params.theta, + config.theta, circular_y_enabled, circular_x_enabled, - qwen_image_params.axes_dim); - int pos_len = static_cast(pe_vec.size() / qwen_image_params.axes_dim_sum / 2); + config.axes_dim); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); // LOG_DEBUG("pos_len %d", pos_len); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, qwen_image_params.axes_dim_sum / 2, pos_len); + auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); // pe->data = pe_vec.data(); // print_ggml_tensor(pe, true, "pe"); // pe->data = nullptr; set_backend_tensor_data(pe, pe_vec.data()); ggml_tensor* modulate_index = nullptr; - if (qwen_image_params.zero_cond_t) { + if (config.zero_cond_t) { modulate_index_vec.clear(); - int64_t h_len = ((x->ne[1] + (qwen_image_params.patch_size / 2)) / qwen_image_params.patch_size); - int64_t w_len = ((x->ne[0] + (qwen_image_params.patch_size / 2)) / qwen_image_params.patch_size); + int64_t h_len = ((x->ne[1] + (config.patch_size / 2)) / config.patch_size); + int64_t w_len = ((x->ne[0] + (config.patch_size / 2)) / config.patch_size); int64_t num_img_tokens = h_len * w_len; modulate_index_vec.insert(modulate_index_vec.end(), num_img_tokens, 0.f); int64_t num_ref_img_tokens = 0; for (ggml_tensor* ref : ref_latents) { - int64_t h_len = ((ref->ne[1] + (qwen_image_params.patch_size / 2)) / qwen_image_params.patch_size); - int64_t w_len = ((ref->ne[0] + (qwen_image_params.patch_size / 2)) / qwen_image_params.patch_size); + int64_t h_len = ((ref->ne[1] + (config.patch_size / 2)) / config.patch_size); + int64_t w_len = ((ref->ne[0] + (config.patch_size / 2)) / config.patch_size); num_ref_img_tokens += h_len * w_len; } @@ -688,7 +693,7 @@ namespace Qwen { return build_graph(x, timesteps, context, ref_latents, ref_index_mode); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); } sd::Tensor compute(int n_threads, @@ -752,7 +757,8 @@ namespace Qwen { ggml_backend_t backend = sd_backend_cpu_init(); ggml_type model_data_type = GGML_TYPE_Q8_0; - ModelLoader model_loader; + auto model_manager = std::make_shared(); + ModelLoader& model_loader = model_manager->loader(); if (!model_loader.init_from_file_and_convert_name(file_path, "model.diffusion_model.")) { LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str()); return; @@ -766,23 +772,20 @@ namespace Qwen { } std::shared_ptr qwen_image = std::make_shared(backend, - backend, tensor_storage_map, "model.diffusion_model", - VERSION_QWEN_IMAGE); + VERSION_QWEN_IMAGE, + false, + model_manager); - if (!qwen_image->alloc_params_buffer()) { - LOG_ERROR("qwen_image buffer allocation failed"); - return; - } - - std::map tensors; - qwen_image->get_param_tensors(tensors, "model.diffusion_model"); - - bool success = model_loader.load_tensors(tensors); - - if (!success) { - LOG_ERROR("load tensors from model loader failed"); + if (!model_manager->register_runner_params("Qwen image test", + *qwen_image, + "model.diffusion_model", + ModelManager::ResidencyMode::ParamBackend, + backend, + backend) || + !model_manager->validate_registered_tensors()) { + LOG_ERROR("register qwen_image tensors with model manager failed"); return; } @@ -793,4 +796,4 @@ namespace Qwen { } // namespace name -#endif // __QWEN_IMAGE_HPP__ +#endif // __SD_MODEL_DIFFUSION_QWEN_IMAGE_HPP__ diff --git a/src/model/diffusion/sefi_image.hpp b/src/model/diffusion/sefi_image.hpp new file mode 100644 index 00000000..27191988 --- /dev/null +++ b/src/model/diffusion/sefi_image.hpp @@ -0,0 +1,91 @@ +#ifndef __SD_MODEL_DIFFUSION_SEFI_IMAGE_HPP__ +#define __SD_MODEL_DIFFUSION_SEFI_IMAGE_HPP__ + +#include + +#include "model/common/block.hpp" + +namespace SefiImage { + struct SefiImageConfig { + int64_t semantic_channels = 16; + int64_t texture_latent_channels = 32; + int64_t timestep_guidance_in_dim = 256; + int64_t hidden_size = 3072; + float timestep_shift_alpha = 0.3f; + float delta_t = 0.1f; + + int64_t packed_texture_channels(int patch_size) const { + return texture_latent_channels * patch_size * patch_size; + } + + int64_t packed_input_channels(int patch_size) const { + return semantic_channels + packed_texture_channels(patch_size); + } + + static SefiImageConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, + const std::string& prefix) { + SefiImageConfig config; + for (const auto& [name, tensor_storage] : tensor_storage_map) { + if (!starts_with(name, prefix)) { + continue; + } + if (ends_with(name, "dual_time_embed.semantic_embedder.linear_1.weight") && tensor_storage.n_dims == 2) { + config.timestep_guidance_in_dim = tensor_storage.ne[0]; + config.hidden_size = tensor_storage.ne[1] * 2; + } + } + LOG_DEBUG("sefi_image: semantic_channels = %" PRId64 ", texture_latent_channels = %" PRId64 ", hidden_size = %" PRId64, + config.semantic_channels, + config.texture_latent_channels, + config.hidden_size); + return config; + } + }; + + struct SefiTimestepEmbedding : public GGMLBlock { + public: + SefiTimestepEmbedding(int64_t in_channels, int64_t time_embed_dim) { + blocks["linear_1"] = std::shared_ptr(new Linear(in_channels, time_embed_dim, false)); + blocks["linear_2"] = std::shared_ptr(new Linear(time_embed_dim, time_embed_dim, false)); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* sample) { + auto linear_1 = std::dynamic_pointer_cast(blocks["linear_1"]); + auto linear_2 = std::dynamic_pointer_cast(blocks["linear_2"]); + + sample = linear_1->forward(ctx, sample); + sample = ggml_silu_inplace(ctx->ggml_ctx, sample); + sample = linear_2->forward(ctx, sample); + return sample; + } + }; + + struct SefiDualTimestepEmbeddings : public GGMLBlock { + public: + SefiDualTimestepEmbeddings(int64_t in_channels, int64_t embedding_dim) { + GGML_ASSERT(embedding_dim % 2 == 0); + int64_t half_dim = embedding_dim / 2; + blocks["semantic_embedder"] = std::make_shared(in_channels, half_dim); + blocks["texture_embedder"] = std::make_shared(in_channels, half_dim); + timestep_guidance_in_dim = in_channels; + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* timestep_sem, + ggml_tensor* timestep_tex) { + auto semantic_embedder = std::dynamic_pointer_cast(blocks["semantic_embedder"]); + auto texture_embedder = std::dynamic_pointer_cast(blocks["texture_embedder"]); + + auto sem_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_sem, (int)timestep_guidance_in_dim, 10000, 1.f); + auto tex_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_tex, (int)timestep_guidance_in_dim, 10000, 1.f); + auto sem_emb = semantic_embedder->forward(ctx, sem_proj); + auto tex_emb = texture_embedder->forward(ctx, tex_proj); + return ggml_concat(ctx->ggml_ctx, sem_emb, tex_emb, 0); + } + + private: + int64_t timestep_guidance_in_dim = 256; + }; +} // namespace SefiImage + +#endif // __SD_MODEL_DIFFUSION_SEFI_IMAGE_HPP__ diff --git a/src/unet.hpp b/src/model/diffusion/unet.hpp similarity index 84% rename from src/unet.hpp rename to src/model/diffusion/unet.hpp index ef468741..253b3b4b 100644 --- a/src/unet.hpp +++ b/src/model/diffusion/unet.hpp @@ -1,14 +1,136 @@ -#ifndef __UNET_HPP__ -#define __UNET_HPP__ +#ifndef __SD_MODEL_DIFFUSION_UNET_HPP__ +#define __SD_MODEL_DIFFUSION_UNET_HPP__ + +#include +#include -#include "common_block.hpp" -#include "diffusion_model.hpp" #include "model.h" +#include "model/common/block.hpp" +#include "model/diffusion/model.hpp" /*==================================================== UnetModel =====================================================*/ #define UNET_GRAPH_SIZE 102400 +struct UNetConfig { + SDVersion version = VERSION_SD1; + // network hparams + int in_channels = 4; + int out_channels = 4; + int num_res_blocks = 2; + std::vector attention_resolutions = {4, 2, 1}; + std::vector channel_mult = {1, 2, 4, 4}; + std::vector transformer_depth = {1, 1, 1, 1}; + int time_embed_dim = 1280; // model_channels*4 + int num_heads = 8; + int num_head_channels = -1; // channels // num_heads + int context_dim = 768; // 1024 for VERSION_SD2, 2048 for VERSION_SDXL + bool use_linear_projection = false; + bool tiny_unet = false; + int model_channels = 320; + int adm_in_channels = 2816; // only for VERSION_SDXL/SVD + + static UNetConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, + const std::string& prefix, + SDVersion version = VERSION_SD1) { + UNetConfig config; + config.version = version; + + if (sd_version_is_sd2(version)) { + config.context_dim = 1024; + config.num_head_channels = 64; + config.num_heads = -1; + config.use_linear_projection = true; + } else if (sd_version_is_sdxl(version)) { + config.context_dim = 2048; + config.attention_resolutions = {4, 2}; + config.channel_mult = {1, 2, 4}; + config.transformer_depth = {1, 2, 10}; + config.num_head_channels = 64; + config.num_heads = -1; + config.use_linear_projection = true; + if (version == VERSION_SDXL_VEGA) { + config.transformer_depth = {1, 1, 2}; + } + } else if (version == VERSION_SVD) { + config.in_channels = 8; + config.out_channels = 4; + config.context_dim = 1024; + config.adm_in_channels = 768; + config.num_head_channels = 64; + config.num_heads = -1; + config.use_linear_projection = true; + } + if (sd_version_is_inpaint(version)) { + config.in_channels = 9; + } else if (sd_version_is_unet_edit(version)) { + config.in_channels = 8; + } + if (version == VERSION_SD1_TINY_UNET || version == VERSION_SD2_TINY_UNET || version == VERSION_SDXS_512_DS || version == VERSION_SDXS_09) { + config.num_res_blocks = 1; + config.channel_mult = {1, 2, 4}; + config.tiny_unet = true; + if (version == VERSION_SDXS_512_DS) { + config.attention_resolutions = {4, 2}; // here just like SDXL + } + } + + auto find_weight = [&](const std::string& suffix) -> const TensorStorage* { + std::string name = prefix.empty() ? suffix : prefix + "." + suffix; + auto it = tensor_storage_map.find(name); + if (it == tensor_storage_map.end()) { + return nullptr; + } + return &it->second; + }; + + if (const TensorStorage* input = find_weight("input_blocks.0.0.weight")) { + if (input->n_dims == 4) { + config.in_channels = static_cast(input->ne[2]); + config.model_channels = static_cast(input->ne[3]); + config.time_embed_dim = config.model_channels * 4; + } + } + if (const TensorStorage* time_embed = find_weight("time_embed.0.weight")) { + if (time_embed->n_dims == 2) { + config.model_channels = static_cast(time_embed->ne[0]); + config.time_embed_dim = static_cast(time_embed->ne[1]); + } + } + if (const TensorStorage* label_emb = find_weight("label_emb.0.0.weight")) { + if (label_emb->n_dims == 2) { + config.adm_in_channels = static_cast(label_emb->ne[0]); + config.time_embed_dim = static_cast(label_emb->ne[1]); + } + } + if (const TensorStorage* out = find_weight("out.2.weight")) { + if (out->n_dims == 4) { + config.out_channels = static_cast(out->ne[3]); + } + } + for (const auto& [name, tensor_storage] : tensor_storage_map) { + if (!starts_with(name, prefix)) { + continue; + } + if (name.find("attn2.to_k.weight") != std::string::npos && tensor_storage.n_dims == 2) { + config.context_dim = static_cast(tensor_storage.ne[0]); + break; + } + } + + LOG_DEBUG("unet: in_channels = %d, out_channels = %d, model_channels = %d, time_embed_dim = %d, context_dim = %d, adm_in_channels = %d, num_res_blocks = %d, tiny_unet = %s", + config.in_channels, + config.out_channels, + config.model_channels, + config.time_embed_dim, + config.context_dim, + config.adm_in_channels, + config.num_res_blocks, + config.tiny_unet ? "true" : "false"); + return config; + } +}; + class SpatialVideoTransformer : public SpatialTransformer { protected: int64_t time_depth; @@ -166,66 +288,26 @@ public: // ldm.modules.diffusionmodules.openaimodel.UNetModel class UnetModelBlock : public GGMLBlock { -protected: - SDVersion version = VERSION_SD1; - // network hparams - int in_channels = 4; - int out_channels = 4; - int num_res_blocks = 2; - std::vector attention_resolutions = {4, 2, 1}; - std::vector channel_mult = {1, 2, 4, 4}; - std::vector transformer_depth = {1, 1, 1, 1}; - int time_embed_dim = 1280; // model_channels*4 - int num_heads = 8; - int num_head_channels = -1; // channels // num_heads - int context_dim = 768; // 1024 for VERSION_SD2, 2048 for VERSION_SDXL - bool use_linear_projection = false; - bool tiny_unet = false; - public: - int model_channels = 320; - int adm_in_channels = 2816; // only for VERSION_SDXL/SVD + UNetConfig config; - UnetModelBlock(SDVersion version = VERSION_SD1, const String2TensorStorage& tensor_storage_map = {}) - : version(version) { - if (sd_version_is_sd2(version)) { - context_dim = 1024; - num_head_channels = 64; - num_heads = -1; - use_linear_projection = true; - } else if (sd_version_is_sdxl(version)) { - context_dim = 2048; - attention_resolutions = {4, 2}; - channel_mult = {1, 2, 4}; - transformer_depth = {1, 2, 10}; - num_head_channels = 64; - num_heads = -1; - use_linear_projection = true; - if (version == VERSION_SDXL_VEGA) { - transformer_depth = {1, 1, 2}; - } - } else if (version == VERSION_SVD) { - in_channels = 8; - out_channels = 4; - context_dim = 1024; - adm_in_channels = 768; - num_head_channels = 64; - num_heads = -1; - use_linear_projection = true; - } - if (sd_version_is_inpaint(version)) { - in_channels = 9; - } else if (sd_version_is_unet_edit(version)) { - in_channels = 8; - } - if (version == VERSION_SD1_TINY_UNET || version == VERSION_SD2_TINY_UNET || version == VERSION_SDXS_512_DS || version == VERSION_SDXS_09) { - num_res_blocks = 1; - channel_mult = {1, 2, 4}; - tiny_unet = true; - if (version == VERSION_SDXS_512_DS) { - attention_resolutions = {4, 2}; // here just like SDXL - } - } + explicit UnetModelBlock(UNetConfig config = {}) + : config(config) { + const SDVersion version = this->config.version; + const int in_channels = this->config.in_channels; + const int out_channels = this->config.out_channels; + const int num_res_blocks = this->config.num_res_blocks; + const auto& attention_resolutions = this->config.attention_resolutions; + const auto& channel_mult = this->config.channel_mult; + const auto& transformer_depth = this->config.transformer_depth; + const int time_embed_dim = this->config.time_embed_dim; + const int num_heads = this->config.num_heads; + const int num_head_channels = this->config.num_head_channels; + const int context_dim = this->config.context_dim; + const bool use_linear_projection = this->config.use_linear_projection; + const bool tiny_unet = this->config.tiny_unet; + const int model_channels = this->config.model_channels; + const int adm_in_channels = this->config.adm_in_channels; // dims is always 2 // use_temporal_attention is always True for SVD @@ -398,7 +480,7 @@ public: ggml_tensor* x, ggml_tensor* emb, int num_video_frames) { - if (version == VERSION_SVD) { + if (config.version == VERSION_SVD) { auto block = std::dynamic_pointer_cast(blocks[name]); return block->forward(ctx, x, emb, num_video_frames); @@ -414,7 +496,7 @@ public: ggml_tensor* x, ggml_tensor* context, int timesteps) { - if (version == VERSION_SVD) { + if (config.version == VERSION_SVD) { auto block = std::dynamic_pointer_cast(blocks[name]); return block->forward(ctx, x, context, timesteps); @@ -440,6 +522,13 @@ public: // c_concat: [N, in_channels, h, w] or [1, in_channels, h, w] // y: [N, adm_in_channels] or [1, adm_in_channels] // return: [N, out_channels, h, w] + const SDVersion version = config.version; + const int model_channels = config.model_channels; + const int num_res_blocks = config.num_res_blocks; + const auto& attention_resolutions = config.attention_resolutions; + const auto& channel_mult = config.channel_mult; + const bool tiny_unet = config.tiny_unet; + if (context != nullptr) { if (context->ne[2] != x->ne[3]) { context = ggml_repeat(ctx->ggml_ctx, context, ggml_new_tensor_3d(ctx->ggml_ctx, GGML_TYPE_F32, context->ne[0], context->ne[1], x->ne[3])); @@ -601,14 +690,17 @@ public: }; struct UNetModelRunner : public DiffusionModelRunner { + UNetConfig config; UnetModelBlock unet; UNetModelRunner(ggml_backend_t backend, - ggml_backend_t params_backend, const String2TensorStorage& tensor_storage_map, const std::string prefix, - SDVersion version = VERSION_SD1) - : DiffusionModelRunner(backend, params_backend, prefix), unet(version, tensor_storage_map) { + SDVersion version = VERSION_SD1, + std::shared_ptr weight_manager = nullptr) + : DiffusionModelRunner(backend, prefix, weight_manager), + config(UNetConfig::detect_from_weights(tensor_storage_map, prefix, version)), + unet(config) { unet.init(params_ctx, tensor_storage_map, prefix); } @@ -680,7 +772,7 @@ struct UNetModelRunner : public DiffusionModelRunner { return build_graph(x, timesteps, context, c_concat, y, num_video_frames, controls, control_strength); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); } sd::Tensor compute(int n_threads, @@ -752,4 +844,4 @@ struct UNetModelRunner : public DiffusionModelRunner { } }; -#endif // __UNET_HPP__ +#endif // __SD_MODEL_DIFFUSION_UNET_HPP__ diff --git a/src/model/diffusion/wan.hpp b/src/model/diffusion/wan.hpp new file mode 100644 index 00000000..9a907dcf --- /dev/null +++ b/src/model/diffusion/wan.hpp @@ -0,0 +1,1059 @@ +#ifndef __SD_MODEL_DIFFUSION_WAN_HPP__ +#define __SD_MODEL_DIFFUSION_WAN_HPP__ + +#include +#include +#include + +#include "model/common/block.hpp" +#include "model/common/rope.hpp" +#include "model/diffusion/flux.hpp" +#include "model/diffusion/model.hpp" +#include "model_loader.h" + +namespace WAN { + + constexpr int WAN_GRAPH_SIZE = 10240; + + struct WanConfig { + std::string model_type = "t2v"; + std::tuple patch_size = {1, 2, 2}; + int64_t text_len = 512; + int64_t in_dim = 16; + int64_t dim = 2048; + int64_t ffn_dim = 8192; + int freq_dim = 256; + int64_t text_dim = 4096; + int64_t out_dim = 16; + int64_t num_heads = 16; + int num_layers = 32; + int vace_layers = 0; + int64_t vace_in_dim = 96; + std::map vace_layers_mapping = {}; + bool qk_norm = true; + bool cross_attn_norm = true; + float eps = 1e-6f; + int64_t flf_pos_embed_token_number = 0; + int theta = 10000; + // wan2.1 1.3B: 1536/12, wan2.1/2.2 14B: 5120/40, wan2.2 5B: 3074/24 + std::vector axes_dim = {44, 42, 42}; + int64_t axes_dim_sum = 128; + + static WanConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) { + WanConfig config; + config.num_layers = 0; + for (const auto& [name, _] : tensor_storage_map) { + if (!starts_with(name, prefix)) { + continue; + } + size_t pos = name.find("vace_blocks."); + if (pos != std::string::npos) { + auto items = split_string(name.substr(pos), '.'); + if (items.size() > 1) { + int block_index = atoi(items[1].c_str()); + if (block_index + 1 > config.vace_layers) { + config.vace_layers = block_index + 1; + } + } + continue; + } + pos = name.find("blocks."); + if (pos != std::string::npos) { + auto items = split_string(name.substr(pos), '.'); + if (items.size() > 1) { + int block_index = atoi(items[1].c_str()); + if (block_index + 1 > config.num_layers) { + config.num_layers = block_index + 1; + } + } + continue; + } + if (name.find("img_emb") != std::string::npos) { + config.model_type = "i2v"; + } + if (name.find("img_emb.emb_pos") != std::string::npos) { + config.flf_pos_embed_token_number = 514; + } + } + LOG_DEBUG("wan: model_type = %s, num_layers = %d, vace_layers = %d, dim = %" PRId64 ", ffn_dim = %" PRId64 ", num_heads = %" PRId64, + config.model_type.c_str(), + config.num_layers, + config.vace_layers, + config.dim, + config.ffn_dim, + config.num_heads); + return config; + } + }; + + class WanSelfAttention : public GGMLBlock { + public: + int64_t num_heads; + int64_t head_dim; + + public: + WanSelfAttention(int64_t dim, + int64_t num_heads, + bool qk_norm = true, + float eps = 1e-6) + : num_heads(num_heads) { + head_dim = dim / num_heads; + blocks["q"] = std::shared_ptr(new Linear(dim, dim)); + blocks["k"] = std::shared_ptr(new Linear(dim, dim)); + blocks["v"] = std::shared_ptr(new Linear(dim, dim)); + blocks["o"] = std::shared_ptr(new Linear(dim, dim)); + + if (qk_norm) { + blocks["norm_q"] = std::shared_ptr(new RMSNorm(dim, eps)); + blocks["norm_k"] = std::shared_ptr(new RMSNorm(dim, eps)); + } else { + blocks["norm_q"] = std::shared_ptr(new Identity()); + blocks["norm_k"] = std::shared_ptr(new Identity()); + } + } + + virtual ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* pe, + ggml_tensor* mask = nullptr) { + // x: [N, n_token, dim] + // pe: [n_token, d_head/2, 2, 2] + // return [N, n_token, dim] + int64_t N = x->ne[2]; + int64_t n_token = x->ne[1]; + + auto q_proj = std::dynamic_pointer_cast(blocks["q"]); + auto k_proj = std::dynamic_pointer_cast(blocks["k"]); + auto v_proj = std::dynamic_pointer_cast(blocks["v"]); + auto o_proj = std::dynamic_pointer_cast(blocks["o"]); + auto norm_q = std::dynamic_pointer_cast(blocks["norm_q"]); + auto norm_k = std::dynamic_pointer_cast(blocks["norm_k"]); + + auto q = q_proj->forward(ctx, x); + q = norm_q->forward(ctx, q); + auto k = k_proj->forward(ctx, x); + k = norm_k->forward(ctx, k); + auto v = v_proj->forward(ctx, x); // [N, n_token, n_head*d_head] + + q = ggml_reshape_4d(ctx->ggml_ctx, q, head_dim, num_heads, n_token, N); // [N, n_token, n_head, d_head] + k = ggml_reshape_4d(ctx->ggml_ctx, k, head_dim, num_heads, n_token, N); // [N, n_token, n_head, d_head] + v = ggml_reshape_4d(ctx->ggml_ctx, v, head_dim, num_heads, n_token, N); // [N, n_token, n_head, d_head] + + x = Rope::attention(ctx, q, k, v, pe, mask); // [N, n_token, dim] + + x = o_proj->forward(ctx, x); // [N, n_token, dim] + return x; + } + }; + + class WanCrossAttention : public WanSelfAttention { + public: + WanCrossAttention(int64_t dim, + int64_t num_heads, + bool qk_norm = true, + float eps = 1e-6) + : WanSelfAttention(dim, num_heads, qk_norm, eps) {} + virtual ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* context, + int64_t context_img_len) = 0; + }; + + class WanT2VCrossAttention : public WanCrossAttention { + public: + WanT2VCrossAttention(int64_t dim, + int64_t num_heads, + bool qk_norm = true, + float eps = 1e-6) + : WanCrossAttention(dim, num_heads, qk_norm, eps) {} + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* context, + int64_t context_img_len) override { + // x: [N, n_token, dim] + // context: [N, n_context, dim] + // context_img_len: unused + // return [N, n_token, dim] + int64_t N = x->ne[2]; + int64_t n_token = x->ne[1]; + + auto q_proj = std::dynamic_pointer_cast(blocks["q"]); + auto k_proj = std::dynamic_pointer_cast(blocks["k"]); + auto v_proj = std::dynamic_pointer_cast(blocks["v"]); + auto o_proj = std::dynamic_pointer_cast(blocks["o"]); + auto norm_q = std::dynamic_pointer_cast(blocks["norm_q"]); + auto norm_k = std::dynamic_pointer_cast(blocks["norm_k"]); + + auto q = q_proj->forward(ctx, x); + q = norm_q->forward(ctx, q); + auto k = k_proj->forward(ctx, context); // [N, n_context, dim] + k = norm_k->forward(ctx, k); + auto v = v_proj->forward(ctx, context); // [N, n_context, dim] + + x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim] + + x = o_proj->forward(ctx, x); // [N, n_token, dim] + return x; + } + }; + + class WanI2VCrossAttention : public WanCrossAttention { + public: + WanI2VCrossAttention(int64_t dim, + int64_t num_heads, + bool qk_norm = true, + float eps = 1e-6) + : WanCrossAttention(dim, num_heads, qk_norm, eps) { + blocks["k_img"] = std::shared_ptr(new Linear(dim, dim)); + blocks["v_img"] = std::shared_ptr(new Linear(dim, dim)); + + if (qk_norm) { + blocks["norm_k_img"] = std::shared_ptr(new RMSNorm(dim, eps)); + } else { + blocks["norm_k_img"] = std::shared_ptr(new Identity()); + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* context, + int64_t context_img_len) override { + // x: [N, n_token, dim] + // context: [N, context_img_len + context_txt_len, dim] + // return [N, n_token, dim] + + auto q_proj = std::dynamic_pointer_cast(blocks["q"]); + auto k_proj = std::dynamic_pointer_cast(blocks["k"]); + auto v_proj = std::dynamic_pointer_cast(blocks["v"]); + auto o_proj = std::dynamic_pointer_cast(blocks["o"]); + + auto k_img_proj = std::dynamic_pointer_cast(blocks["k_img"]); + auto v_img_proj = std::dynamic_pointer_cast(blocks["v_img"]); + + auto norm_q = std::dynamic_pointer_cast(blocks["norm_q"]); + auto norm_k = std::dynamic_pointer_cast(blocks["norm_k"]); + auto norm_k_img = std::dynamic_pointer_cast(blocks["norm_k_img"]); + + int64_t N = x->ne[2]; + int64_t n_token = x->ne[1]; + int64_t dim = x->ne[0]; + int64_t context_txt_len = context->ne[1] - context_img_len; + + auto context_img = ggml_view_3d(ctx->ggml_ctx, context, dim, context_img_len, N, context->nb[1], context->nb[2], 0); // [N, context_img_len, dim] + auto context_txt = ggml_view_3d(ctx->ggml_ctx, context, dim, context_txt_len, N, context->nb[1], context->nb[2], context_img_len * context->nb[1]); // [N, context_txt_len, dim] + + auto q = q_proj->forward(ctx, x); + q = norm_q->forward(ctx, q); + auto k = k_proj->forward(ctx, context_txt); // [N, context_txt_len, dim] + k = norm_k->forward(ctx, k); + auto v = v_proj->forward(ctx, context_txt); // [N, context_txt_len, dim] + + auto k_img = k_img_proj->forward(ctx, context_img); // [N, context_img_len, dim] + k_img = norm_k_img->forward(ctx, k_img); + auto v_img = v_img_proj->forward(ctx, context_img); // [N, context_img_len, dim] + + auto img_x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k_img, v_img, num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim] + x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim] + + x = ggml_add(ctx->ggml_ctx, x, img_x); + + x = o_proj->forward(ctx, x); // [N, n_token, dim] + return x; + } + }; + + static ggml_tensor* modulate_add(ggml_context* ctx, ggml_tensor* x, ggml_tensor* e) { + // x: [N, n_token, dim] + // e: [N, 1, dim] or [N, T, 1, dim] + if (ggml_n_dims(e) == 3) { + int64_t T = e->ne[2]; + x = ggml_reshape_4d(ctx, x, x->ne[0], x->ne[1] / T, T, x->ne[2]); // [N, T, n_token/T, dim] + x = ggml_add(ctx, x, e); + x = ggml_reshape_3d(ctx, x, x->ne[0], x->ne[1] * x->ne[2], x->ne[3]); // [N, n_token, dim] + } else { + x = ggml_add(ctx, x, e); + } + return x; + } + + static ggml_tensor* modulate_mul(ggml_context* ctx, ggml_tensor* x, ggml_tensor* e) { + // x: [N, n_token, dim] + // e: [N, 1, dim] or [N, T, 1, dim] + if (ggml_n_dims(e) == 3) { + int64_t T = e->ne[2]; + x = ggml_reshape_4d(ctx, x, x->ne[0], x->ne[1] / T, T, x->ne[2]); // [N, T, n_token/T, dim] + x = ggml_mul(ctx, x, e); + x = ggml_reshape_3d(ctx, x, x->ne[0], x->ne[1] * x->ne[2], x->ne[3]); // [N, n_token, dim] + } else { + x = ggml_mul(ctx, x, e); + } + return x; + } + + class WanAttentionBlock : public GGMLBlock { + protected: + int64_t dim; + + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { + enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32); + params["modulation"] = ggml_new_tensor_3d(ctx, wtype, dim, 6, 1); + } + + public: + WanAttentionBlock(bool t2v_cross_attn, + int64_t dim, + int64_t ffn_dim, + int64_t num_heads, + bool qk_norm = true, + bool cross_attn_norm = false, + float eps = 1e-6) + : dim(dim) { + blocks["norm1"] = std::shared_ptr(new LayerNorm(dim, eps, false)); + blocks["self_attn"] = std::shared_ptr(new WanSelfAttention(dim, num_heads, qk_norm, eps)); + if (cross_attn_norm) { + blocks["norm3"] = std::shared_ptr(new LayerNorm(dim, eps, true)); + } else { + blocks["norm3"] = std::shared_ptr(new Identity()); + } + if (t2v_cross_attn) { + blocks["cross_attn"] = std::shared_ptr(new WanT2VCrossAttention(dim, num_heads, qk_norm, eps)); + } else { + blocks["cross_attn"] = std::shared_ptr(new WanI2VCrossAttention(dim, num_heads, qk_norm, eps)); + } + + blocks["norm2"] = std::shared_ptr(new LayerNorm(dim, eps, false)); + + blocks["ffn.0"] = std::shared_ptr(new Linear(dim, ffn_dim)); + // ffn.1 is nn.GELU(approximate='tanh') + blocks["ffn.2"] = std::shared_ptr(new Linear(ffn_dim, dim)); + } + + virtual ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* e, + ggml_tensor* pe, + ggml_tensor* context, + int64_t context_img_len = 257) { + // x: [N, n_token, dim] + // e: [N, 6, dim] or [N, T, 6, dim] + // context: [N, context_img_len + context_txt_len, dim] + // return [N, n_token, dim] + + auto modulation = params["modulation"]; + e = ggml_add(ctx->ggml_ctx, e, modulation); // [N, 6, dim] or [N, T, 6, dim] + auto es = ggml_ext_chunk(ctx->ggml_ctx, e, 6, 1); // ([N, 1, dim], ...) or [N, T, 1, dim] + + auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); + auto self_attn = std::dynamic_pointer_cast(blocks["self_attn"]); + auto norm3 = std::dynamic_pointer_cast(blocks["norm3"]); + auto cross_attn = std::dynamic_pointer_cast(blocks["cross_attn"]); + auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); + auto ffn_0 = std::dynamic_pointer_cast(blocks["ffn.0"]); + auto ffn_2 = std::dynamic_pointer_cast(blocks["ffn.2"]); + + // self-attention + auto y = norm1->forward(ctx, x); + y = ggml_add(ctx->ggml_ctx, y, modulate_mul(ctx->ggml_ctx, y, es[1])); + y = modulate_add(ctx->ggml_ctx, y, es[0]); + y = self_attn->forward(ctx, y, pe); + + x = ggml_add(ctx->ggml_ctx, x, modulate_mul(ctx->ggml_ctx, y, es[2])); + + // cross-attention + x = ggml_add(ctx->ggml_ctx, + x, + cross_attn->forward(ctx, norm3->forward(ctx, x), context, context_img_len)); + + // ffn + y = norm2->forward(ctx, x); + y = ggml_add(ctx->ggml_ctx, y, modulate_mul(ctx->ggml_ctx, y, es[4])); + y = modulate_add(ctx->ggml_ctx, y, es[3]); + + y = ffn_0->forward(ctx, y); + y = ggml_ext_gelu(ctx->ggml_ctx, y, true); + y = ffn_2->forward(ctx, y); + + x = ggml_add(ctx->ggml_ctx, x, modulate_mul(ctx->ggml_ctx, y, es[5])); + + return x; + } + }; + + class VaceWanAttentionBlock : public WanAttentionBlock { + protected: + int block_id; + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { + enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32); + params["modulation"] = ggml_new_tensor_3d(ctx, wtype, dim, 6, 1); + } + + public: + VaceWanAttentionBlock(bool t2v_cross_attn, + int64_t dim, + int64_t ffn_dim, + int64_t num_heads, + bool qk_norm = true, + bool cross_attn_norm = false, + float eps = 1e-6, + int block_id = 0) + : WanAttentionBlock(t2v_cross_attn, dim, ffn_dim, num_heads, qk_norm, cross_attn_norm, eps), block_id(block_id) { + if (block_id == 0) { + blocks["before_proj"] = std::shared_ptr(new Linear(dim, dim)); + } + blocks["after_proj"] = std::shared_ptr(new Linear(dim, dim)); + } + + std::pair forward(GGMLRunnerContext* ctx, + ggml_tensor* c, + ggml_tensor* x, + ggml_tensor* e, + ggml_tensor* pe, + ggml_tensor* context, + int64_t context_img_len = 257) { + // x: [N, n_token, dim] + // e: [N, 6, dim] or [N, T, 6, dim] + // context: [N, context_img_len + context_txt_len, dim] + // return [N, n_token, dim] + if (block_id == 0) { + auto before_proj = std::dynamic_pointer_cast(blocks["before_proj"]); + + c = before_proj->forward(ctx, c); + c = ggml_add(ctx->ggml_ctx, c, x); + } + + auto after_proj = std::dynamic_pointer_cast(blocks["after_proj"]); + + c = WanAttentionBlock::forward(ctx, c, e, pe, context, context_img_len); + auto c_skip = after_proj->forward(ctx, c); + + return {c_skip, c}; + } + }; + + class Head : public GGMLBlock { + protected: + int64_t dim; + + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { + enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32); + params["modulation"] = ggml_new_tensor_3d(ctx, wtype, dim, 2, 1); + } + + public: + Head(int64_t dim, + int64_t out_dim, + std::tuple patch_size, + float eps = 1e-6) + : dim(dim) { + out_dim = out_dim * std::get<0>(patch_size) * std::get<1>(patch_size) * std::get<2>(patch_size); + + blocks["norm"] = std::shared_ptr(new LayerNorm(dim, eps, false)); + blocks["head"] = std::shared_ptr(new Linear(dim, out_dim)); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* e) { + // x: [N, n_token, dim] + // e: [N, dim] or [N, T, dim] + // return [N, n_token, out_dim] + + auto modulation = params["modulation"]; + e = ggml_reshape_4d(ctx->ggml_ctx, e, e->ne[0], 1, e->ne[1], e->ne[2]); // [N, 1, dim] or [N, T, 1, dim] + e = ggml_repeat_4d(ctx->ggml_ctx, e, e->ne[0], 2, e->ne[2], e->ne[3]); // [N, 2, dim] or [N, T, 2, dim] + + e = ggml_add(ctx->ggml_ctx, e, modulation); // [N, 2, dim] or [N, T, 2, dim] + auto es = ggml_ext_chunk(ctx->ggml_ctx, e, 2, 1); // ([N, 1, dim], ...) or ([N, T, 1, dim], ...) + + auto norm = std::dynamic_pointer_cast(blocks["norm"]); + auto head = std::dynamic_pointer_cast(blocks["head"]); + + x = norm->forward(ctx, x); + x = ggml_add(ctx->ggml_ctx, x, modulate_mul(ctx->ggml_ctx, x, es[1])); + x = modulate_add(ctx->ggml_ctx, x, es[0]); + x = head->forward(ctx, x); + return x; + } + }; + + class MLPProj : public GGMLBlock { + protected: + int64_t in_dim; + int64_t flf_pos_embed_token_number; + + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { + if (flf_pos_embed_token_number > 0) { + params["emb_pos"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, in_dim, flf_pos_embed_token_number, 1); + } + } + + public: + MLPProj(int64_t in_dim, + int64_t out_dim, + int64_t flf_pos_embed_token_number = 0) + : in_dim(in_dim), flf_pos_embed_token_number(flf_pos_embed_token_number) { + blocks["proj.0"] = std::shared_ptr(new LayerNorm(in_dim)); + blocks["proj.1"] = std::shared_ptr(new Linear(in_dim, in_dim)); + // proj.2 is nn.GELU() + blocks["proj.3"] = std::shared_ptr(new Linear(in_dim, out_dim)); + blocks["proj.4"] = std::shared_ptr(new LayerNorm(out_dim)); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* image_embeds) { + if (flf_pos_embed_token_number > 0) { + auto emb_pos = params["emb_pos"]; + + auto a = ggml_ext_slice(ctx->ggml_ctx, image_embeds, 1, 0, emb_pos->ne[1]); + auto b = ggml_ext_slice(ctx->ggml_ctx, emb_pos, 1, 0, image_embeds->ne[1]); + + image_embeds = ggml_add(ctx->ggml_ctx, a, b); + } + + auto proj_0 = std::dynamic_pointer_cast(blocks["proj.0"]); + auto proj_1 = std::dynamic_pointer_cast(blocks["proj.1"]); + auto proj_3 = std::dynamic_pointer_cast(blocks["proj.3"]); + auto proj_4 = std::dynamic_pointer_cast(blocks["proj.4"]); + + auto x = proj_0->forward(ctx, image_embeds); + x = proj_1->forward(ctx, x); + x = ggml_ext_gelu(ctx->ggml_ctx, x, true); + x = proj_3->forward(ctx, x); + x = proj_4->forward(ctx, x); + + return x; // clip_extra_context_tokens + } + }; + + class Wan : public GGMLBlock { + protected: + WanConfig config; + + public: + Wan() {} + Wan(WanConfig config) + : config(config) { + // patch_embedding + blocks["patch_embedding"] = std::shared_ptr(new Conv3d(config.in_dim, config.dim, config.patch_size, config.patch_size)); + + // text_embedding + blocks["text_embedding.0"] = std::shared_ptr(new Linear(config.text_dim, config.dim)); + // text_embedding.1 is nn.GELU() + blocks["text_embedding.2"] = std::shared_ptr(new Linear(config.dim, config.dim)); + + // time_embedding + blocks["time_embedding.0"] = std::shared_ptr(new Linear(config.freq_dim, config.dim)); + // time_embedding.1 is nn.SiLU() + blocks["time_embedding.2"] = std::shared_ptr(new Linear(config.dim, config.dim)); + + // time_projection.0 is nn.SiLU() + blocks["time_projection.1"] = std::shared_ptr(new Linear(config.dim, config.dim * 6)); + + // blocks + for (int i = 0; i < config.num_layers; i++) { + auto block = std::shared_ptr(new WanAttentionBlock(config.model_type == "t2v", + config.dim, + config.ffn_dim, + config.num_heads, + config.qk_norm, + config.cross_attn_norm, + config.eps)); + blocks["blocks." + std::to_string(i)] = block; + } + + // head + blocks["head"] = std::shared_ptr(new Head(config.dim, config.out_dim, config.patch_size, config.eps)); + + // img_emb + if (config.model_type == "i2v") { + blocks["img_emb"] = std::shared_ptr(new MLPProj(1280, config.dim, config.flf_pos_embed_token_number)); + } + + // vace + if (config.vace_layers > 0) { + for (int i = 0; i < config.vace_layers; i++) { + auto block = std::shared_ptr(new VaceWanAttentionBlock(config.model_type == "t2v", + config.dim, + config.ffn_dim, + config.num_heads, + config.qk_norm, + config.cross_attn_norm, + config.eps, + i)); + blocks["vace_blocks." + std::to_string(i)] = block; + } + + int step = config.num_layers / config.vace_layers; + int n = 0; + for (int i = 0; i < config.num_layers; i += step) { + this->config.vace_layers_mapping[i] = n; + n++; + } + + blocks["vace_patch_embedding"] = std::shared_ptr(new Conv3d(config.vace_in_dim, config.dim, config.patch_size, config.patch_size)); + } + } + + ggml_tensor* pad_to_patch_size(GGMLRunnerContext* ctx, + ggml_tensor* x) { + int64_t W = x->ne[0]; + int64_t H = x->ne[1]; + int64_t T = x->ne[2]; + + int pad_t = (std::get<0>(config.patch_size) - T % std::get<0>(config.patch_size)) % std::get<0>(config.patch_size); + int pad_h = (std::get<1>(config.patch_size) - H % std::get<1>(config.patch_size)) % std::get<1>(config.patch_size); + int pad_w = (std::get<2>(config.patch_size) - W % std::get<2>(config.patch_size)) % std::get<2>(config.patch_size); + ggml_ext_pad(ctx->ggml_ctx, x, pad_w, pad_h, pad_t, 0, ctx->circular_x_enabled, ctx->circular_y_enabled); + return x; + } + + ggml_tensor* unpatchify(ggml_context* ctx, + ggml_tensor* x, + int64_t t_len, + int64_t h_len, + int64_t w_len) { + // x: [N, t_len*h_len*w_len, pt*ph*pw*C] + // return: [N*C, t_len*pt, h_len*ph, w_len*pw] + int64_t N = x->ne[3]; + int64_t pt = std::get<0>(config.patch_size); + int64_t ph = std::get<1>(config.patch_size); + int64_t pw = std::get<2>(config.patch_size); + int64_t C = x->ne[0] / pt / ph / pw; + + GGML_ASSERT(C * pt * ph * pw == x->ne[0]); + + x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N); // [N, t_len*h_len*w_len, pt*ph*pw, C] + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3)); // [N, C, t_len*h_len*w_len, pt*ph*pw] + x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N); // [N*C*t_len*h_len, w_len, pt*ph, pw] + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, pt*ph, w_len, pw] + x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N); // [N*C*t_len*h_len, pt, ph, w_len*pw] + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, ph, pt, w_len*pw] + x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph * h_len, t_len * C * N); // [N*C*t_len, h_len*ph, pt, w_len*pw] + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len, pt, h_len*ph, w_len*pw] + x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt * t_len, C * N); // [N*C, t_len*pt, h_len*ph, w_len*pw] + return x; + } + + ggml_tensor* forward_orig(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* timestep, + ggml_tensor* context, + ggml_tensor* pe, + ggml_tensor* clip_fea = nullptr, + ggml_tensor* vace_context = nullptr, + float vace_strength = 1.f, + int64_t N = 1) { + // x: [N*C, T, H, W], C => in_dim + // vace_context: [N*vace_in_dim, T, H, W] + // timestep: [N,] or [T] + // context: [N, L, text_dim] + // return: [N, t_len*h_len*w_len, out_dim*pt*ph*pw] + + GGML_ASSERT(N == 1); + + auto patch_embedding = std::dynamic_pointer_cast(blocks["patch_embedding"]); + + auto text_embedding_0 = std::dynamic_pointer_cast(blocks["text_embedding.0"]); + auto text_embedding_2 = std::dynamic_pointer_cast(blocks["text_embedding.2"]); + + auto time_embedding_0 = std::dynamic_pointer_cast(blocks["time_embedding.0"]); + auto time_embedding_2 = std::dynamic_pointer_cast(blocks["time_embedding.2"]); + auto time_projection_1 = std::dynamic_pointer_cast(blocks["time_projection.1"]); + + auto head = std::dynamic_pointer_cast(blocks["head"]); + + // patch_embedding + x = patch_embedding->forward(ctx, x); // [N*dim, t_len, h_len, w_len] + x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1] * x->ne[2], x->ne[3] / N, N); // [N, dim, t_len*h_len*w_len] + x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); // [N, t_len*h_len*w_len, dim] + + // time_embedding + auto e = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, config.freq_dim); + e = time_embedding_0->forward(ctx, e); + e = ggml_silu_inplace(ctx->ggml_ctx, e); + e = time_embedding_2->forward(ctx, e); // [N, dim] or [N, T, dim] + + // time_projection + auto e0 = ggml_silu(ctx->ggml_ctx, e); + e0 = time_projection_1->forward(ctx, e0); + e0 = ggml_reshape_4d(ctx->ggml_ctx, e0, e0->ne[0] / 6, 6, e0->ne[1], e0->ne[2]); // [N, 6, dim] or [N, T, 6, dim] + + context = text_embedding_0->forward(ctx, context); + context = ggml_ext_gelu(ctx->ggml_ctx, context); + context = text_embedding_2->forward(ctx, context); // [N, context_txt_len, dim] + + int64_t context_img_len = 0; + if (clip_fea != nullptr) { + if (config.model_type == "i2v") { + auto img_emb = std::dynamic_pointer_cast(blocks["img_emb"]); + auto context_img = img_emb->forward(ctx, clip_fea); // [N, context_img_len, dim] + context = ggml_concat(ctx->ggml_ctx, context_img, context, 1); // [N, context_img_len + context_txt_len, dim] + } + context_img_len = clip_fea->ne[1]; // 257 + } + + // vace_patch_embedding + ggml_tensor* c = nullptr; + if (config.vace_layers > 0) { + auto vace_patch_embedding = std::dynamic_pointer_cast(blocks["vace_patch_embedding"]); + + c = vace_patch_embedding->forward(ctx, vace_context); // [N*dim, t_len, h_len, w_len] + c = ggml_reshape_3d(ctx->ggml_ctx, c, c->ne[0] * c->ne[1] * c->ne[2], c->ne[3] / N, N); // [N, dim, t_len*h_len*w_len] + c = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, c, 1, 0, 2, 3)); // [N, t_len*h_len*w_len, dim] + } + sd::ggml_graph_cut::mark_graph_cut(x, "wan.prelude", "x"); + // sd::ggml_graph_cut::mark_graph_cut(e, "wan.prelude", "e"); + // sd::ggml_graph_cut::mark_graph_cut(e0, "wan.prelude", "e0"); + // sd::ggml_graph_cut::mark_graph_cut(context, "wan.prelude", "context"); + if (c != nullptr) { + sd::ggml_graph_cut::mark_graph_cut(c, "wan.prelude", "c"); + } + + auto x_orig = x; + + for (int i = 0; i < config.num_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["blocks." + std::to_string(i)]); + + x = block->forward(ctx, x, e0, pe, context, context_img_len); + + auto iter = config.vace_layers_mapping.find(i); + if (iter != config.vace_layers_mapping.end()) { + int n = iter->second; + + auto vace_block = std::dynamic_pointer_cast(blocks["vace_blocks." + std::to_string(n)]); + + auto result = vace_block->forward(ctx, c, x_orig, e0, pe, context, context_img_len); + auto c_skip = result.first; + c = result.second; + c_skip = ggml_ext_scale(ctx->ggml_ctx, c_skip, vace_strength); + x = ggml_add(ctx->ggml_ctx, x, c_skip); + } + sd::ggml_graph_cut::mark_graph_cut(x, "wan.blocks." + std::to_string(i), "x"); + if (c != nullptr) { + sd::ggml_graph_cut::mark_graph_cut(c, "wan.blocks." + std::to_string(i), "c"); + } + } + + x = head->forward(ctx, x, e); // [N, t_len*h_len*w_len, pt*ph*pw*out_dim] + + return x; + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* timestep, + ggml_tensor* context, + ggml_tensor* pe, + ggml_tensor* clip_fea = nullptr, + ggml_tensor* time_dim_concat = nullptr, + ggml_tensor* vace_context = nullptr, + float vace_strength = 1.f, + int64_t N = 1) { + // Forward pass of DiT. + // x: [N*C, T, H, W] + // timestep: [N,] + // context: [N, L, D] + // pe: [L, d_head/2, 2, 2] + // time_dim_concat: [N*C, T2, H, W] + // return: [N*C, T, H, W] + + GGML_ASSERT(N == 1); + + int64_t W = x->ne[0]; + int64_t H = x->ne[1]; + int64_t T = x->ne[2]; + int64_t C = x->ne[3]; + + x = pad_to_patch_size(ctx, x); + + int64_t t_len = ((T + (std::get<0>(config.patch_size) / 2)) / std::get<0>(config.patch_size)); + int64_t h_len = ((H + (std::get<1>(config.patch_size) / 2)) / std::get<1>(config.patch_size)); + int64_t w_len = ((W + (std::get<2>(config.patch_size) / 2)) / std::get<2>(config.patch_size)); + + if (time_dim_concat != nullptr) { + time_dim_concat = pad_to_patch_size(ctx, time_dim_concat); + x = ggml_concat(ctx->ggml_ctx, x, time_dim_concat, 2); // [N*C, (T+pad_t) + (T2+pad_t2), H + pad_h, W + pad_w] + t_len = ((x->ne[2] + (std::get<0>(config.patch_size) / 2)) / std::get<0>(config.patch_size)); + } + + auto out = forward_orig(ctx, x, timestep, context, pe, clip_fea, vace_context, vace_strength, N); // [N, t_len*h_len*w_len, pt*ph*pw*C] + + out = unpatchify(ctx->ggml_ctx, out, t_len, h_len, w_len); // [N*C, (T+pad_t) + (T2+pad_t2), H + pad_h, W + pad_w] + + // slice + + out = ggml_ext_slice(ctx->ggml_ctx, out, 2, 0, T); // [N*C, T, H + pad_h, W + pad_w] + out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, H); // [N*C, T, H, W + pad_w] + out = ggml_ext_slice(ctx->ggml_ctx, out, 0, 0, W); // [N*C, T, H, W] + + return out; + } + }; + + struct WanRunner : public DiffusionModelRunner { + public: + std::string desc = "wan"; + WanConfig config; + Wan wan; + std::vector pe_vec; + SDVersion version; + + WanRunner(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "", + SDVersion version = VERSION_WAN2, + std::shared_ptr weight_manager = nullptr) + : DiffusionModelRunner(backend, prefix, weight_manager), + config(WanConfig::detect_from_weights(tensor_storage_map, prefix)) { + if (config.num_layers == 30) { + if (version == VERSION_WAN2_2_TI2V) { + desc = "Wan2.2-TI2V-5B"; + config.dim = 3072; + config.eps = 1e-06f; + config.ffn_dim = 14336; + config.freq_dim = 256; + config.in_dim = 48; + config.num_heads = 24; + config.out_dim = 48; + config.text_len = 512; + } else { + if (config.vace_layers > 0) { + desc = "Wan2.1-VACE-1.3B"; + config.in_dim = 16; + } else if (config.model_type == "i2v") { + desc = "Wan2.1-I2V-1.3B"; + config.in_dim = 36; + } else { + desc = "Wan2.1-T2V-1.3B"; + config.in_dim = 16; + } + config.dim = 1536; + config.eps = 1e-06f; + config.ffn_dim = 8960; + config.freq_dim = 256; + config.num_heads = 12; + config.out_dim = 16; + config.text_len = 512; + } + } else if (config.num_layers == 40) { + if (config.model_type == "t2v") { + if (version == VERSION_WAN2_2_I2V) { + desc = "Wan2.2-I2V-14B"; + config.in_dim = 36; + } else { + if (config.vace_layers > 0) { + desc = "Wan2.x-VACE-14B"; + } else { + desc = "Wan2.x-T2V-14B"; + } + config.in_dim = 16; + } + } else { + config.in_dim = 36; + if (config.flf_pos_embed_token_number > 0) { + desc = "Wan2.1-FLF2V-14B"; + } else { + desc = "Wan2.1-I2V-14B"; + } + } + config.dim = 5120; + config.eps = 1e-06f; + config.ffn_dim = 13824; + config.freq_dim = 256; + config.num_heads = 40; + config.out_dim = 16; + config.text_len = 512; + } else { + GGML_ABORT("invalid num_layers(%d) of wan", config.num_layers); + } + + LOG_INFO("%s", desc.c_str()); + + wan = Wan(config); + wan.init(params_ctx, tensor_storage_map, prefix); + } + + std::string get_desc() override { + return desc; + } + + void get_param_tensors(std::map& tensors, const std::string& prefix) override { + wan.get_param_tensors(tensors, prefix); + } + + ggml_cgraph* build_graph(const sd::Tensor& x_tensor, + const sd::Tensor& timesteps_tensor, + const sd::Tensor& context_tensor = {}, + const sd::Tensor& clip_fea_tensor = {}, + const sd::Tensor& c_concat_tensor = {}, + const sd::Tensor& time_dim_concat_tensor = {}, + const sd::Tensor& vace_context_tensor = {}, + float vace_strength = 1.f) { + ggml_cgraph* gf = new_graph_custom(WAN_GRAPH_SIZE); + + ggml_tensor* x = make_input(x_tensor); + ggml_tensor* timesteps = make_input(timesteps_tensor); + ggml_tensor* context = make_optional_input(context_tensor); + ggml_tensor* clip_fea = make_optional_input(clip_fea_tensor); + ggml_tensor* c_concat = make_optional_input(c_concat_tensor); + ggml_tensor* time_dim_concat = make_optional_input(time_dim_concat_tensor); + ggml_tensor* vace_context = make_optional_input(vace_context_tensor); + + pe_vec = Rope::gen_wan_pe(static_cast(x->ne[2]), + static_cast(x->ne[1]), + static_cast(x->ne[0]), + std::get<0>(config.patch_size), + std::get<1>(config.patch_size), + std::get<2>(config.patch_size), + 1, + config.theta, + config.axes_dim); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + // LOG_DEBUG("pos_len %d", pos_len); + auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); + // pe->data = pe_vec.data(); + // print_ggml_tensor(pe); + // pe->data = nullptr; + set_backend_tensor_data(pe, pe_vec.data()); + + if (c_concat != nullptr) { + x = ggml_concat(compute_ctx, x, c_concat, 3); + } + + auto runner_ctx = get_context(); + + ggml_tensor* out = wan.forward(&runner_ctx, + x, + timesteps, + context, + pe, + clip_fea, + time_dim_concat, + vace_context, + vace_strength); + + ggml_build_forward_expand(gf, out); + + return gf; + } + + sd::Tensor compute(int n_threads, + const sd::Tensor& x, + const sd::Tensor& timesteps, + const sd::Tensor& context = {}, + const sd::Tensor& clip_fea = {}, + const sd::Tensor& c_concat = {}, + const sd::Tensor& time_dim_concat = {}, + const sd::Tensor& vace_context = {}, + float vace_strength = 1.f) { + auto get_graph = [&]() -> ggml_cgraph* { + return build_graph(x, timesteps, context, clip_fea, c_concat, time_dim_concat, vace_context, vace_strength); + }; + + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + } + + sd::Tensor compute(int n_threads, + const DiffusionParams& diffusion_params) override { + GGML_ASSERT(diffusion_params.x != nullptr); + GGML_ASSERT(diffusion_params.timesteps != nullptr); + const auto* extra = diffusion_extra_as(diffusion_params); + return compute(n_threads, + *diffusion_params.x, + *diffusion_params.timesteps, + tensor_or_empty(diffusion_params.context), + tensor_or_empty(diffusion_params.y), + tensor_or_empty(diffusion_params.c_concat), + sd::Tensor(), + tensor_or_empty(extra->vace_context), + extra->vace_strength); + } + + void test() { + ggml_init_params params; + params.mem_size = static_cast(200 * 1024 * 1024); // 200 MB + params.mem_buffer = nullptr; + params.no_alloc = false; + + ggml_context* ctx = ggml_init(params); + GGML_ASSERT(ctx != nullptr); + + { + // cpu f16: pass + // cuda f16: pass + // cpu q8_0: pass + // auto x = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 104, 60, 1, 16); + // ggml_set_f32(x, 0.01f); + auto x = sd::load_tensor_from_file_as_tensor("wan_dit_x.bin"); + print_sd_tensor(x); + + std::vector timesteps_vec(3, 1000.f); + timesteps_vec[0] = 0.f; + auto timesteps = sd::Tensor::from_vector(timesteps_vec); + + // auto context = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 4096, 512, 1); + // ggml_set_f32(context, 0.01f); + auto context = sd::load_tensor_from_file_as_tensor("wan_dit_context.bin"); + print_sd_tensor(context); + // auto clip_fea = load_tensor_from_file(ctx, "wan_dit_clip_fea.bin"); + // print_ggml_tensor(clip_fea); + + sd::Tensor out; + + int64_t t0 = ggml_time_ms(); + auto out_opt = compute(8, x, timesteps, context, {}, {}, {}, {}, 1.f); + int64_t t1 = ggml_time_ms(); + + GGML_ASSERT(!out_opt.empty()); + out = std::move(out_opt); + print_sd_tensor(out); + LOG_DEBUG("wan test done in %lldms", t1 - t0); + } + } + + static void load_from_file_and_test(const std::string& file_path) { + // ggml_backend_t backend = ggml_backend_cuda_init(0); + ggml_backend_t backend = sd_backend_cpu_init(); + ggml_type model_data_type = GGML_TYPE_F16; + LOG_INFO("loading from '%s'", file_path.c_str()); + + auto model_manager = std::make_shared(); + ModelLoader& model_loader = model_manager->loader(); + if (!model_loader.init_from_file_and_convert_name(file_path, "model.diffusion_model.")) { + LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str()); + return; + } + + auto& tensor_storage_map = model_loader.get_tensor_storage_map(); + for (auto& [name, tensor_storage] : tensor_storage_map) { + if (ends_with(name, "weight")) { + tensor_storage.expected_type = model_data_type; + } + } + + std::shared_ptr wan = std::make_shared(backend, + tensor_storage_map, + "model.diffusion_model", + VERSION_WAN2_2_TI2V, + model_manager); + + if (!model_manager->register_runner_params("Wan test", + *wan, + "model.diffusion_model", + ModelManager::ResidencyMode::ParamBackend, + backend, + backend) || + !model_manager->validate_registered_tensors()) { + LOG_ERROR("register wan tensors with model manager failed"); + return; + } + + LOG_INFO("wan model loaded"); + + wan->test(); + } + }; + +} // namespace WAN + +#endif // __SD_MODEL_DIFFUSION_WAN_HPP__ diff --git a/src/z_image.hpp b/src/model/diffusion/z_image.hpp similarity index 79% rename from src/z_image.hpp rename to src/model/diffusion/z_image.hpp index f39e5621..362192f1 100644 --- a/src/z_image.hpp +++ b/src/model/diffusion/z_image.hpp @@ -1,14 +1,15 @@ -#ifndef __Z_IMAGE_HPP__ -#define __Z_IMAGE_HPP__ +#ifndef __SD_MODEL_DIFFUSION_Z_IMAGE_HPP__ +#define __SD_MODEL_DIFFUSION_Z_IMAGE_HPP__ #include -#include "diffusion_model.hpp" -#include "flux.hpp" -#include "ggml_extend.hpp" -#include "mmdit.hpp" +#include "core/ggml_extend.hpp" +#include "model/diffusion/flux.hpp" +#include "model/diffusion/mmdit.hpp" +#include "model/diffusion/model.hpp" +#include "model_loader.h" -// Ref: https://github.com/Alpha-VLLM/Lumina-Image-2.0/blob/main/models/model.py +// Ref: https://github.com/Alpha-VLLM/Lumina-Image-2.0/blob/main/model/model.py // Ref: https://github.com/huggingface/diffusers/pull/12703 #ifndef MIN @@ -20,6 +21,104 @@ namespace ZImage { constexpr int ADALN_EMBED_DIM = 256; constexpr int SEQ_MULTI_OF = 32; + struct ZImageConfig { + int patch_size = 2; + int64_t hidden_size = 3840; + int64_t in_channels = 16; + int64_t out_channels = 16; + int64_t num_layers = 30; + int64_t num_refiner_layers = 2; + int64_t head_dim = 128; + int64_t num_heads = 30; + int64_t num_kv_heads = 30; + int64_t multiple_of = 256; + float ffn_dim_multiplier = 8.0f / 3.0f; + float norm_eps = 1e-5f; + bool qk_norm = true; + int64_t cap_feat_dim = 2560; + int theta = 256; + std::vector axes_dim = {32, 48, 48}; + int64_t axes_dim_sum = 128; + + static ZImageConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) { + ZImageConfig config; + int64_t detected_layers = 0; + int64_t detected_refiner_layers = 0; + int64_t detected_context_refiner = 0; + int64_t detected_head_dim = 0; + int64_t detected_qkv_dim = 0; + + for (const auto& [name, tensor_storage] : tensor_storage_map) { + if (!starts_with(name, prefix)) { + continue; + } + if (ends_with(name, "x_embedder.weight") && tensor_storage.n_dims == 2) { + int64_t patch_area = config.patch_size * config.patch_size; + config.in_channels = tensor_storage.ne[0] / patch_area; + config.hidden_size = tensor_storage.ne[1]; + } else if (ends_with(name, "cap_embedder.1.weight") && tensor_storage.n_dims == 2) { + config.cap_feat_dim = tensor_storage.ne[0]; + config.hidden_size = tensor_storage.ne[1]; + } else if (ends_with(name, "layers.0.attention.q_norm.weight") && tensor_storage.n_dims == 1) { + detected_head_dim = tensor_storage.ne[0]; + } else if (ends_with(name, "layers.0.attention.qkv.weight") && tensor_storage.n_dims == 2) { + detected_qkv_dim = tensor_storage.ne[1]; + } else if (ends_with(name, "final_layer.linear.weight") && tensor_storage.n_dims == 2) { + int64_t patch_area = config.patch_size * config.patch_size; + config.out_channels = tensor_storage.ne[1] / patch_area; + } + + size_t pos = name.find("layers."); + if (pos != std::string::npos) { + auto items = split_string(name.substr(pos), '.'); + if (items.size() > 1) { + int block_index = atoi(items[1].c_str()); + detected_layers = std::max(detected_layers, block_index + 1); + } + } + pos = name.find("noise_refiner."); + if (pos != std::string::npos) { + auto items = split_string(name.substr(pos), '.'); + if (items.size() > 1) { + int block_index = atoi(items[1].c_str()); + detected_refiner_layers = std::max(detected_refiner_layers, block_index + 1); + } + } + pos = name.find("context_refiner."); + if (pos != std::string::npos) { + auto items = split_string(name.substr(pos), '.'); + if (items.size() > 1) { + int block_index = atoi(items[1].c_str()); + detected_context_refiner = std::max(detected_context_refiner, block_index + 1); + } + } + } + if (detected_layers > 0) { + config.num_layers = detected_layers; + } + if (detected_refiner_layers > 0 || detected_context_refiner > 0) { + config.num_refiner_layers = std::max(detected_refiner_layers, detected_context_refiner); + } + if (detected_head_dim > 0) { + config.head_dim = detected_head_dim; + config.num_heads = config.hidden_size / config.head_dim; + if (detected_qkv_dim > 0) { + int64_t qkv_heads = detected_qkv_dim / config.head_dim; + config.num_kv_heads = std::max(1, (qkv_heads - config.num_heads) / 2); + } + } + LOG_DEBUG("z_image: num_layers = %" PRId64 ", num_refiner_layers = %" PRId64 ", hidden_size = %" PRId64 ", num_heads = %" PRId64 ", num_kv_heads = %" PRId64 ", in_channels = %" PRId64 ", out_channels = %" PRId64, + config.num_layers, + config.num_refiner_layers, + config.hidden_size, + config.num_heads, + config.num_kv_heads, + config.in_channels, + config.out_channels); + return config; + } + }; + struct JointAttention : public GGMLBlock { protected: int64_t head_dim; @@ -263,90 +362,70 @@ namespace ZImage { } }; - struct ZImageParams { - int patch_size = 2; - int64_t hidden_size = 3840; - int64_t in_channels = 16; - int64_t out_channels = 16; - int64_t num_layers = 30; - int64_t num_refiner_layers = 2; - int64_t head_dim = 128; - int64_t num_heads = 30; - int64_t num_kv_heads = 30; - int64_t multiple_of = 256; - float ffn_dim_multiplier = 8.0f / 3.0f; - float norm_eps = 1e-5f; - bool qk_norm = true; - int64_t cap_feat_dim = 2560; - int theta = 256; - std::vector axes_dim = {32, 48, 48}; - int64_t axes_dim_sum = 128; - }; - class ZImageModel : public GGMLBlock { protected: - ZImageParams z_image_params; + ZImageConfig config; void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { - params["cap_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, z_image_params.hidden_size); - params["x_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, z_image_params.hidden_size); + params["cap_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size); + params["x_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size); } public: ZImageModel() = default; - ZImageModel(ZImageParams z_image_params) - : z_image_params(z_image_params) { - blocks["x_embedder"] = std::make_shared(z_image_params.patch_size * z_image_params.patch_size * z_image_params.in_channels, z_image_params.hidden_size); - blocks["t_embedder"] = std::make_shared(MIN(z_image_params.hidden_size, 1024), 256, 256); - blocks["cap_embedder.0"] = std::make_shared(z_image_params.cap_feat_dim, z_image_params.norm_eps); - blocks["cap_embedder.1"] = std::make_shared(z_image_params.cap_feat_dim, z_image_params.hidden_size); + ZImageModel(ZImageConfig config) + : config(config) { + blocks["x_embedder"] = std::make_shared(config.patch_size * config.patch_size * config.in_channels, config.hidden_size); + blocks["t_embedder"] = std::make_shared(MIN(config.hidden_size, 1024), 256, 256); + blocks["cap_embedder.0"] = std::make_shared(config.cap_feat_dim, config.norm_eps); + blocks["cap_embedder.1"] = std::make_shared(config.cap_feat_dim, config.hidden_size); - for (int i = 0; i < z_image_params.num_refiner_layers; i++) { + for (int i = 0; i < config.num_refiner_layers; i++) { auto block = std::make_shared(i, - z_image_params.hidden_size, - z_image_params.head_dim, - z_image_params.num_heads, - z_image_params.num_kv_heads, - z_image_params.multiple_of, - z_image_params.ffn_dim_multiplier, - z_image_params.norm_eps, - z_image_params.qk_norm, + config.hidden_size, + config.head_dim, + config.num_heads, + config.num_kv_heads, + config.multiple_of, + config.ffn_dim_multiplier, + config.norm_eps, + config.qk_norm, true); blocks["noise_refiner." + std::to_string(i)] = block; } - for (int i = 0; i < z_image_params.num_refiner_layers; i++) { + for (int i = 0; i < config.num_refiner_layers; i++) { auto block = std::make_shared(i, - z_image_params.hidden_size, - z_image_params.head_dim, - z_image_params.num_heads, - z_image_params.num_kv_heads, - z_image_params.multiple_of, - z_image_params.ffn_dim_multiplier, - z_image_params.norm_eps, - z_image_params.qk_norm, + config.hidden_size, + config.head_dim, + config.num_heads, + config.num_kv_heads, + config.multiple_of, + config.ffn_dim_multiplier, + config.norm_eps, + config.qk_norm, false); blocks["context_refiner." + std::to_string(i)] = block; } - for (int i = 0; i < z_image_params.num_layers; i++) { + for (int i = 0; i < config.num_layers; i++) { auto block = std::make_shared(i, - z_image_params.hidden_size, - z_image_params.head_dim, - z_image_params.num_heads, - z_image_params.num_kv_heads, - z_image_params.multiple_of, - z_image_params.ffn_dim_multiplier, - z_image_params.norm_eps, - z_image_params.qk_norm, + config.hidden_size, + config.head_dim, + config.num_heads, + config.num_kv_heads, + config.multiple_of, + config.ffn_dim_multiplier, + config.norm_eps, + config.qk_norm, true); blocks["layers." + std::to_string(i)] = block; } - blocks["final_layer"] = std::make_shared(z_image_params.hidden_size, z_image_params.patch_size, z_image_params.out_channels); + blocks["final_layer"] = std::make_shared(config.hidden_size, config.patch_size, config.out_channels); } ggml_tensor* forward_core(GGMLRunnerContext* ctx, @@ -393,14 +472,14 @@ namespace ZImage { auto txt_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, 0, txt->ne[1]); auto img_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt->ne[1], pe->ne[3]); - for (int i = 0; i < z_image_params.num_refiner_layers; i++) { + for (int i = 0; i < config.num_refiner_layers; i++) { auto block = std::dynamic_pointer_cast(blocks["context_refiner." + std::to_string(i)]); txt = block->forward(ctx, txt, txt_pe, nullptr, nullptr); sd::ggml_graph_cut::mark_graph_cut(txt, "z_image.context_refiner." + std::to_string(i), "txt"); } - for (int i = 0; i < z_image_params.num_refiner_layers; i++) { + for (int i = 0; i < config.num_refiner_layers; i++) { auto block = std::dynamic_pointer_cast(blocks["noise_refiner." + std::to_string(i)]); img = block->forward(ctx, img, img_pe, nullptr, t_emb); @@ -410,7 +489,7 @@ namespace ZImage { auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1); // [N, n_txt_token + n_txt_pad_token + n_img_token + n_img_pad_token, hidden_size] sd::ggml_graph_cut::mark_graph_cut(txt_img, "z_image.prelude", "txt_img"); - for (int i = 0; i < z_image_params.num_layers; i++) { + for (int i = 0; i < config.num_layers; i++) { auto block = std::dynamic_pointer_cast(blocks["layers." + std::to_string(i)]); txt_img = block->forward(ctx, txt_img, pe, nullptr, t_emb); @@ -442,7 +521,7 @@ namespace ZImage { int64_t C = x->ne[2]; int64_t N = x->ne[3]; - int patch_size = z_image_params.patch_size; + int patch_size = config.patch_size; auto img = DiT::pad_and_patchify(ctx, x, patch_size, patch_size, false); uint64_t n_img_token = img->ne[1]; @@ -467,19 +546,20 @@ namespace ZImage { struct ZImageRunner : public DiffusionModelRunner { public: - ZImageParams z_image_params; + ZImageConfig config; ZImageModel z_image; std::vector pe_vec; std::vector timestep_vec; SDVersion version; ZImageRunner(ggml_backend_t backend, - ggml_backend_t params_backend, - const String2TensorStorage& tensor_storage_map = {}, - const std::string prefix = "", - SDVersion version = VERSION_Z_IMAGE) - : DiffusionModelRunner(backend, params_backend, prefix) { - z_image = ZImageModel(z_image_params); + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "", + SDVersion version = VERSION_Z_IMAGE, + std::shared_ptr weight_manager = nullptr) + : DiffusionModelRunner(backend, prefix, weight_manager), + config(ZImageConfig::detect_from_weights(tensor_storage_map, prefix)) { + z_image = ZImageModel(config); z_image.init(params_ctx, tensor_storage_map, prefix); } @@ -510,19 +590,19 @@ namespace ZImage { pe_vec = Rope::gen_z_image_pe(static_cast(x->ne[1]), static_cast(x->ne[0]), - z_image_params.patch_size, + config.patch_size, static_cast(x->ne[3]), static_cast(context->ne[1]), SEQ_MULTI_OF, ref_latents, ref_index_mode, - z_image_params.theta, + config.theta, circular_y_enabled, circular_x_enabled, - z_image_params.axes_dim); - int pos_len = static_cast(pe_vec.size() / z_image_params.axes_dim_sum / 2); + config.axes_dim); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); // LOG_DEBUG("pos_len %d", pos_len); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, z_image_params.axes_dim_sum / 2, pos_len); + auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); // pe->data = pe_vec.data(); // print_ggml_tensor(pe, true, "pe"); // pe->data = nullptr; @@ -554,7 +634,7 @@ namespace ZImage { return build_graph(x, timesteps, context, ref_latents, ref_index_mode); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); } sd::Tensor compute(int n_threads, @@ -618,7 +698,8 @@ namespace ZImage { ggml_backend_t backend = sd_backend_cpu_init(); ggml_type model_data_type = GGML_TYPE_Q8_0; - ModelLoader model_loader; + auto model_manager = std::make_shared(); + ModelLoader& model_loader = model_manager->loader(); if (!model_loader.init_from_file_and_convert_name(file_path, "model.diffusion_model.")) { LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str()); return; @@ -634,22 +715,19 @@ namespace ZImage { } std::shared_ptr z_image = std::make_shared(backend, - backend, tensor_storage_map, "model.diffusion_model", - VERSION_QWEN_IMAGE); + VERSION_QWEN_IMAGE, + model_manager); - if (!z_image->alloc_params_buffer()) { - LOG_ERROR("z_image buffer allocation failed"); - return; - } - std::map tensors; - z_image->get_param_tensors(tensors, "model.diffusion_model"); - - bool success = model_loader.load_tensors(tensors); - - if (!success) { - LOG_ERROR("load tensors from model loader failed"); + if (!model_manager->register_runner_params("ZImage test", + *z_image, + "model.diffusion_model", + ModelManager::ResidencyMode::ParamBackend, + backend, + backend) || + !model_manager->validate_registered_tensors()) { + LOG_ERROR("register z_image tensors with model manager failed"); return; } @@ -660,4 +738,4 @@ namespace ZImage { } // namespace ZImage -#endif // __Z_IMAGE_HPP__ +#endif // __SD_MODEL_DIFFUSION_Z_IMAGE_HPP__ diff --git a/src/clip.hpp b/src/model/te/clip.hpp similarity index 96% rename from src/clip.hpp rename to src/model/te/clip.hpp index a3567324..2fde3de7 100644 --- a/src/clip.hpp +++ b/src/model/te/clip.hpp @@ -1,13 +1,13 @@ -#ifndef __CLIP_HPP__ -#define __CLIP_HPP__ +#ifndef __SD_MODEL_TE_CLIP_HPP__ +#define __SD_MODEL_TE_CLIP_HPP__ -#include "ggml_extend.hpp" +#include "core/ggml_extend.hpp" #include "model.h" #include "tokenizers/clip_tokenizer.h" /*================================================ FrozenCLIPEmbedder ================================================*/ -// Ref: https://github.com/huggingface/transformers/blob/main/src/transformers/models/clip/modeling_clip.py +// Ref: https://github.com/huggingface/transformers/blob/main/src/transformers/model/clip/modeling_clip.py struct CLIPMLP : public GGMLBlock { protected: @@ -469,13 +469,13 @@ struct CLIPTextModelRunner : public GGMLRunner { std::vector attention_mask_vec; CLIPTextModelRunner(ggml_backend_t backend, - ggml_backend_t params_backend, const String2TensorStorage& tensor_storage_map, const std::string prefix, - CLIPVersion version = OPENAI_CLIP_VIT_L_14, - bool with_final_ln = true, - bool force_clip_f32 = false) - : GGMLRunner(backend, params_backend) { + CLIPVersion version = OPENAI_CLIP_VIT_L_14, + bool with_final_ln = true, + bool force_clip_f32 = false, + std::shared_ptr weight_manager = nullptr) + : GGMLRunner(backend, weight_manager) { bool proj_in = false; for (const auto& [name, tensor_storage] : tensor_storage_map) { if (!starts_with(name, prefix)) { @@ -567,11 +567,14 @@ struct CLIPTextModelRunner : public GGMLRunner { void* custom_embeddings_data, size_t max_token_idx, bool return_pooled, - int clip_skip) { + int clip_skip, + bool auto_free = true, + bool free_compute_buffer = true, + bool free_compute_params = true) { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(input_ids, num_custom_embeddings, custom_embeddings_data, max_token_idx, return_pooled, clip_skip); }; - auto result = GGMLRunner::compute(get_graph, n_threads, true); + auto result = GGMLRunner::compute(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params); if (return_pooled) { return take_or_empty(std::move(result)); } @@ -579,4 +582,4 @@ struct CLIPTextModelRunner : public GGMLRunner { } }; -#endif // __CLIP_HPP__ +#endif // __SD_MODEL_TE_CLIP_HPP__ diff --git a/src/llm.hpp b/src/model/te/llm.hpp similarity index 86% rename from src/llm.hpp rename to src/model/te/llm.hpp index 0dbd37d9..3c5e9bb5 100644 --- a/src/llm.hpp +++ b/src/model/te/llm.hpp @@ -1,5 +1,5 @@ -#ifndef __LLM_HPP__ -#define __LLM_HPP__ +#ifndef __SD_MODEL_TE_LLM_HPP__ +#define __SD_MODEL_TE_LLM_HPP__ #include #include @@ -18,9 +18,11 @@ #include #include -#include "ggml_extend.hpp" +#include "core/ggml_extend.hpp" #include "json.hpp" -#include "rope.hpp" +#include "model/common/rope.hpp" +#include "model_loader.h" +#include "model_manager.h" #include "tokenizers/bpe_tokenizer.h" #include "tokenizers/gemma_tokenizer.h" #include "tokenizers/gpt_oss_tokenizer.h" @@ -63,7 +65,7 @@ namespace LLM { QWEN3_VL, }; - struct LLMVisionParams { + struct LLMVisionConfig { LLMVisionArch arch = LLMVisionArch::QWEN2_5_VL; int num_layers = 32; int64_t hidden_size = 1280; @@ -77,9 +79,10 @@ namespace LLM { int window_size = 112; int num_position_embeddings = 0; std::set fullatt_block_indexes = {7, 15, 23, 31}; + bool split_patch_embed = false; }; - struct LLMParams { + struct LLMConfig { LLMArch arch = LLMArch::QWEN2_5_VL; int64_t num_layers = 28; int64_t hidden_size = 3584; @@ -101,7 +104,165 @@ namespace LLM { std::vector sliding_attention; int64_t num_experts = 0; int64_t num_experts_per_tok = 0; - LLMVisionParams vision; + LLMVisionConfig vision; + bool have_vision_weight = false; + bool llama_cpp_style = false; + + static LLMConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, + const std::string& prefix, + LLMArch arch) { + LLMConfig config; + config.arch = arch; + if (arch == LLMArch::MISTRAL_SMALL_3_2 || arch == LLMArch::MINISTRAL_3_3B) { + config.head_dim = 128; + config.num_heads = 32; + config.num_kv_heads = 8; + config.qkv_bias = false; + config.rms_norm_eps = 1e-5f; + } else if (arch == LLMArch::QWEN3 || arch == LLMArch::QWEN3_VL) { + config.head_dim = 128; + config.num_heads = 32; + config.num_kv_heads = 8; + config.qkv_bias = false; + config.qk_norm = true; + config.rms_norm_eps = 1e-6f; + if (arch == LLMArch::QWEN3_VL) { + config.max_position_embeddings = 262144; + config.rope_thetas = {5000000.f}; + config.vision.arch = LLMVisionArch::QWEN3_VL; + } + } else if (arch == LLMArch::GEMMA3_12B) { + config.head_dim = 256; + config.num_heads = 16; + config.num_kv_heads = 8; + config.qkv_bias = false; + config.qk_norm = true; + config.rms_norm_eps = 1e-6f; + config.rms_norm_add = false; + config.normalize_input = true; + config.max_position_embeddings = 131072; + config.mlp_activation = MLPActivation::GELU_TANH; + config.rope_thetas = {1000000.f, 10000.f}; + config.rope_scales = {8.f, 1.f}; + config.sliding_attention = {1024, 1024, 1024, 1024, 1024, 0}; + } else if (arch == LLMArch::GEMMA2_2B) { + config.head_dim = 256; + config.num_heads = 8; + config.num_kv_heads = 4; + config.qkv_bias = false; + config.qk_norm = false; + config.rms_norm_eps = 1e-6f; + config.rms_norm_add = true; + config.normalize_input = true; + config.max_position_embeddings = 8192; + config.mlp_activation = MLPActivation::GELU_TANH; + config.hidden_size = 2304; + config.intermediate_size = 9216; + config.num_layers = 26; + config.vocab_size = 256000; + } else if (arch == LLMArch::GPT_OSS_20B) { + config.head_dim = 64; + config.num_heads = 64; + config.num_kv_heads = 8; + config.qkv_bias = true; + config.attention_out_bias = true; + config.qk_norm = false; + config.rms_norm_eps = 1e-5f; + config.hidden_size = 2880; + config.intermediate_size = 2880; + config.num_layers = 24; + config.vocab_size = 201088; + config.max_position_embeddings = 131072; + config.rope_thetas = {150000.f}; + config.rope_scales = {32.f}; + config.sliding_attention = {128, 0}; + config.num_experts = 32; + config.num_experts_per_tok = 4; + } + + config.num_layers = 0; + int detected_vision_layers = 0; + for (const auto& [name, tensor_storage] : tensor_storage_map) { + if (!starts_with(name, prefix)) { + continue; + } + size_t pos = name.find("visual."); + if (pos != std::string::npos) { + config.have_vision_weight = true; + if (contains(name, "attn.q_proj")) { + config.llama_cpp_style = true; + } + if (contains(name, "visual.patch_embed.proj.1.weight")) { + config.vision.split_patch_embed = true; + } + if (contains(name, "visual.patch_embed.proj.0.weight")) { + config.vision.patch_size = static_cast(tensor_storage.ne[0]); + config.vision.in_channels = tensor_storage.ne[2]; + config.vision.hidden_size = tensor_storage.ne[3]; + } + if (contains(name, "visual.patch_embed.bias")) { + config.vision.hidden_size = tensor_storage.ne[0]; + } + if (contains(name, "visual.pos_embed.weight")) { + config.vision.hidden_size = tensor_storage.ne[0]; + config.vision.num_position_embeddings = static_cast(tensor_storage.ne[1]); + } + if (contains(name, "visual.blocks.")) { + auto items = split_string(name.substr(pos), '.'); + if (items.size() > 2) { + int block_index = atoi(items[2].c_str()); + if (block_index + 1 > detected_vision_layers) { + detected_vision_layers = block_index + 1; + } + } + } + if (contains(name, "visual.blocks.0.mlp.linear_fc1.weight") || + contains(name, "visual.blocks.0.mlp.gate_proj.weight")) { + config.vision.intermediate_size = tensor_storage.ne[1]; + } + if (contains(name, "visual.merger.linear_fc2.weight") || + contains(name, "visual.merger.mlp.2.weight")) { + config.vision.out_hidden_size = tensor_storage.ne[1]; + } + continue; + } + pos = name.find("layers."); + if (pos != std::string::npos) { + auto items = split_string(name.substr(pos), '.'); + if (items.size() > 1) { + int block_index = atoi(items[1].c_str()); + if (block_index + 1 > config.num_layers) { + config.num_layers = block_index + 1; + } + } + } + if (contains(name, "embed_tokens.weight")) { + config.hidden_size = tensor_storage.ne[0]; + config.vocab_size = tensor_storage.ne[1]; + } + if (contains(name, "layers.0.mlp.gate_proj.weight")) { + config.intermediate_size = tensor_storage.ne[1]; + } + if (contains(name, "layers.0.mlp.experts.gate_up_proj.weight")) { + config.intermediate_size = tensor_storage.ne[1] / 2; + } + if (contains(name, "layers.0.mlp.experts.gate_proj.weight")) { + config.intermediate_size = tensor_storage.ne[1]; + } + } + if ((arch == LLMArch::QWEN3 || arch == LLMArch::QWEN3_VL) && config.num_layers == 28) { + config.num_heads = 16; + } + if (detected_vision_layers > 0) { + config.vision.num_layers = detected_vision_layers; + } + LOG_DEBUG("llm: num_layers = %" PRId64 ", vocab_size = %" PRId64 ", hidden_size = %" PRId64 ", intermediate_size = %" PRId64, + config.num_layers, + config.vocab_size, + config.hidden_size, + config.intermediate_size); + return config; + } }; struct LLMRMSNorm : public UnaryBlock { @@ -232,11 +393,11 @@ namespace LLM { } 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) {} + GPTOSSMLP(const LLMConfig& config) + : hidden_size(config.hidden_size), + intermediate_size(config.intermediate_size), + num_experts(config.num_experts), + num_experts_per_tok(config.num_experts_per_tok) {} ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { // x: [N, n_token, hidden_size] @@ -415,40 +576,51 @@ namespace LLM { struct VisionPatchEmbed : public GGMLBlock { protected: - bool llama_cpp_style; + bool split_patch_embed; + bool bias; int patch_size; int temporal_patch_size; int64_t in_channels; int64_t embed_dim; + void init_params(ggml_context* ctx, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "") override { + GGML_UNUSED(tensor_storage_map); + GGML_UNUSED(prefix); + if (split_patch_embed && bias) { + params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, embed_dim); + } + } + public: - VisionPatchEmbed(bool llama_cpp_style, + VisionPatchEmbed(bool split_patch_embed, LLMVisionArch arch, int patch_size = 14, int temporal_patch_size = 2, int64_t in_channels = 3, int64_t embed_dim = 1152) - : llama_cpp_style(llama_cpp_style), + : split_patch_embed(split_patch_embed), + bias(arch == LLMVisionArch::QWEN3_VL), patch_size(patch_size), temporal_patch_size(temporal_patch_size), in_channels(in_channels), embed_dim(embed_dim) { - bool bias = arch == LLMVisionArch::QWEN3_VL; - if (llama_cpp_style) { + if (split_patch_embed) { blocks["proj.0"] = std::shared_ptr(new Conv2d(in_channels, embed_dim, {patch_size, patch_size}, {patch_size, patch_size}, {0, 0}, {1, 1}, - bias)); + false)); blocks["proj.1"] = std::shared_ptr(new Conv2d(in_channels, embed_dim, {patch_size, patch_size}, {patch_size, patch_size}, {0, 0}, {1, 1}, - bias)); + false)); } else { std::tuple kernel_size = {(int)temporal_patch_size, (int)patch_size, (int)patch_size}; blocks["proj"] = std::shared_ptr(new Conv3d(in_channels, @@ -469,7 +641,7 @@ namespace LLM { temporal_patch_size, ggml_nelements(x) / (temporal_patch_size * patch_size * patch_size)); - if (llama_cpp_style) { + if (split_patch_embed) { auto proj_0 = std::dynamic_pointer_cast(blocks["proj.0"]); auto proj_1 = std::dynamic_pointer_cast(blocks["proj.1"]); @@ -482,6 +654,10 @@ namespace LLM { x1 = proj_1->forward(ctx, x1); x = ggml_add(ctx->ggml_ctx, x0, x1); + if (bias) { + auto b = ggml_reshape_4d(ctx->ggml_ctx, params["bias"], 1, 1, embed_dim, 1); + x = ggml_add_inplace(ctx->ggml_ctx, x, b); + } } else { auto proj = std::dynamic_pointer_cast(blocks["proj"]); @@ -667,14 +843,14 @@ namespace LLM { public: VisionModel(bool llama_cpp_style, - const LLMVisionParams& vision_params, + const LLMVisionConfig& vision_params, float eps = 1e-6f) : arch_(vision_params.arch), num_layers(vision_params.num_layers), spatial_merge_size(vision_params.spatial_merge_size), num_grid_per_side(vision_params.num_position_embeddings > 0 ? static_cast(std::sqrt(vision_params.num_position_embeddings)) : 0), fullatt_block_indexes(vision_params.fullatt_block_indexes) { - blocks["patch_embed"] = std::shared_ptr(new VisionPatchEmbed(llama_cpp_style, + blocks["patch_embed"] = std::shared_ptr(new VisionPatchEmbed(vision_params.split_patch_embed, arch_, vision_params.patch_size, vision_params.temporal_patch_size, @@ -784,23 +960,23 @@ namespace LLM { } public: - Attention(const LLMParams& params) - : arch(params.arch), - num_heads(params.num_heads), - num_kv_heads(params.num_kv_heads), - head_dim(params.head_dim), - qk_norm(params.qk_norm), - max_position_embeddings(params.max_position_embeddings), - rope_thetas(params.rope_thetas), - rope_scales(params.rope_scales), - has_attention_sinks(params.arch == LLMArch::GPT_OSS_20B) { - blocks["q_proj"] = std::make_shared(params.hidden_size, num_heads * head_dim, params.qkv_bias); - blocks["k_proj"] = std::make_shared(params.hidden_size, num_kv_heads * head_dim, params.qkv_bias); - blocks["v_proj"] = std::make_shared(params.hidden_size, num_kv_heads * head_dim, params.qkv_bias); - blocks["o_proj"] = std::make_shared(num_heads * head_dim, params.hidden_size, params.attention_out_bias); - if (params.qk_norm) { - blocks["q_norm"] = std::make_shared(head_dim, params.rms_norm_eps, params.rms_norm_add); - blocks["k_norm"] = std::make_shared(head_dim, params.rms_norm_eps, params.rms_norm_add); + Attention(const LLMConfig& config) + : arch(config.arch), + num_heads(config.num_heads), + num_kv_heads(config.num_kv_heads), + head_dim(config.head_dim), + qk_norm(config.qk_norm), + max_position_embeddings(config.max_position_embeddings), + rope_thetas(config.rope_thetas), + rope_scales(config.rope_scales), + has_attention_sinks(config.arch == LLMArch::GPT_OSS_20B) { + blocks["q_proj"] = std::make_shared(config.hidden_size, num_heads * head_dim, config.qkv_bias); + blocks["k_proj"] = std::make_shared(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias); + blocks["v_proj"] = std::make_shared(config.hidden_size, num_kv_heads * head_dim, config.qkv_bias); + blocks["o_proj"] = std::make_shared(num_heads * head_dim, config.hidden_size, config.attention_out_bias); + if (config.qk_norm) { + blocks["q_norm"] = std::make_shared(head_dim, config.rms_norm_eps, config.rms_norm_add); + blocks["k_norm"] = std::make_shared(head_dim, config.rms_norm_eps, config.rms_norm_add); } } @@ -982,42 +1158,42 @@ namespace LLM { std::string post_ffw_norm_name; public: - TransformerBlock(const LLMParams& params, int layer_index) - : arch(params.arch), + TransformerBlock(const LLMConfig& config, int layer_index) + : arch(config.arch), sliding_attention(0) { - if (params.arch == LLMArch::GEMMA3_12B) { + if (config.arch == LLMArch::GEMMA3_12B) { post_attention_norm_name = "post_attention_norm"; // attn_post_norm pre_ffw_norm_name = "post_attention_layernorm"; // ffn_norm post_ffw_norm_name = "post_ffw_norm"; // ffn_post_norm - } else if (params.arch == LLMArch::GEMMA2_2B) { + } else if (config.arch == LLMArch::GEMMA2_2B) { post_attention_norm_name = "post_attention_layernorm"; // ffn_norm pre_ffw_norm_name = "pre_feedforward_layernorm"; post_ffw_norm_name = "post_feedforward_layernorm"; - } else if (params.arch == LLMArch::GPT_OSS_20B) { + } else if (config.arch == LLMArch::GPT_OSS_20B) { pre_ffw_norm_name = "post_attention_norm"; // attn_post_norm } else { pre_ffw_norm_name = "post_attention_layernorm"; // ffn_norm } - blocks["self_attn"] = std::make_shared(params); - if (params.arch == LLMArch::GPT_OSS_20B) { - blocks["mlp"] = std::make_shared(params); + blocks["self_attn"] = std::make_shared(config); + if (config.arch == LLMArch::GPT_OSS_20B) { + blocks["mlp"] = std::make_shared(config); } else { - blocks["mlp"] = std::make_shared(params.hidden_size, - params.intermediate_size, + blocks["mlp"] = std::make_shared(config.hidden_size, + config.intermediate_size, false, - params.mlp_activation); + config.mlp_activation); } - blocks["input_layernorm"] = std::make_shared(params.hidden_size, params.rms_norm_eps, params.rms_norm_add); - blocks[pre_ffw_norm_name] = std::make_shared(params.hidden_size, params.rms_norm_eps, params.rms_norm_add); + blocks["input_layernorm"] = std::make_shared(config.hidden_size, config.rms_norm_eps, config.rms_norm_add); + blocks[pre_ffw_norm_name] = std::make_shared(config.hidden_size, config.rms_norm_eps, config.rms_norm_add); if (!post_attention_norm_name.empty()) { - blocks[post_attention_norm_name] = std::make_shared(params.hidden_size, params.rms_norm_eps, params.rms_norm_add); + blocks[post_attention_norm_name] = std::make_shared(config.hidden_size, config.rms_norm_eps, config.rms_norm_add); } if (!post_ffw_norm_name.empty()) { - blocks[post_ffw_norm_name] = std::make_shared(params.hidden_size, params.rms_norm_eps, params.rms_norm_add); + blocks[post_ffw_norm_name] = std::make_shared(config.hidden_size, config.rms_norm_eps, config.rms_norm_add); } - if (!params.sliding_attention.empty()) { - sliding_attention = params.sliding_attention[layer_index % params.sliding_attention.size()]; + if (!config.sliding_attention.empty()) { + sliding_attention = config.sliding_attention[layer_index % config.sliding_attention.size()]; } } @@ -1074,16 +1250,16 @@ namespace LLM { struct TextModel : public GGMLBlock { protected: int64_t num_layers; - LLMParams params; + LLMConfig config; public: - TextModel(const LLMParams& params) - : num_layers(params.num_layers), params(params) { - blocks["embed_tokens"] = std::shared_ptr(new Embedding(params.vocab_size, params.hidden_size)); + TextModel(const LLMConfig& config) + : num_layers(config.num_layers), config(config) { + blocks["embed_tokens"] = std::shared_ptr(new Embedding(config.vocab_size, config.hidden_size)); for (int i = 0; i < num_layers; i++) { - blocks["layers." + std::to_string(i)] = std::shared_ptr(new TransformerBlock(params, i)); + blocks["layers." + std::to_string(i)] = std::shared_ptr(new TransformerBlock(config, i)); } - blocks["norm"] = std::shared_ptr(new LLMRMSNorm(params.hidden_size, params.rms_norm_eps, params.rms_norm_add)); + blocks["norm"] = std::shared_ptr(new LLMRMSNorm(config.hidden_size, config.rms_norm_eps, config.rms_norm_add)); } ggml_tensor* embed(GGMLRunnerContext* ctx, @@ -1103,8 +1279,8 @@ namespace LLM { auto norm = std::dynamic_pointer_cast(blocks["norm"]); std::vector intermediate_outputs; - if (params.normalize_input) { - x = ggml_ext_scale(ctx->ggml_ctx, x, std::sqrt(static_cast(params.hidden_size)), true); + if (config.normalize_input) { + x = ggml_ext_scale(ctx->ggml_ctx, x, std::sqrt(static_cast(config.hidden_size)), true); } if (return_all_hidden_states) { intermediate_outputs.push_back(x); @@ -1174,15 +1350,15 @@ namespace LLM { struct LLM : public GGMLBlock { bool enable_vision; - LLMParams params; + LLMConfig config; public: LLM() = default; - LLM(LLMParams params, bool enable_vision = false, bool llama_cpp_style = false) - : enable_vision(enable_vision), params(params) { - blocks["model"] = std::shared_ptr(new TextModel(params)); + LLM(LLMConfig config, bool enable_vision = false, bool llama_cpp_style = false) + : enable_vision(enable_vision), config(config) { + blocks["model"] = std::shared_ptr(new TextModel(config)); if (enable_vision) { - blocks["visual"] = std::shared_ptr(new VisionModel(llama_cpp_style, params.vision)); + blocks["visual"] = std::shared_ptr(new VisionModel(llama_cpp_style, config.vision)); } } @@ -1226,7 +1402,7 @@ namespace LLM { }; struct LLMRunner : public GGMLRunner { - LLMParams params; + LLMConfig config; bool enable_vision; LLM model; @@ -1242,7 +1418,7 @@ namespace LLM { static ggml_tensor* process_image_common(ggml_context* ctx, ggml_tensor* image, - const LLMVisionParams& vision_params) { + const LLMVisionConfig& vision_params) { // image: [C, H, W] // return: [grid_t*(H/mh/ph)*(W/mw/pw)*mh*mw, C*pt*ph*pw], grid_t == 1 int64_t C = image->ne[2]; @@ -1337,7 +1513,7 @@ namespace LLM { ggml_context* compute_ctx, GGMLRunnerContext* runner_ctx, ggml_tensor* image, - const LLMVisionParams& vision_params, + const LLMVisionConfig& vision_params, std::shared_ptr vision_model, std::vector& window_index_vec, std::vector& window_inverse_index_vec, @@ -1448,140 +1624,29 @@ namespace LLM { public: LLMRunner(LLMArch arch, ggml_backend_t backend, - ggml_backend_t params_backend, const String2TensorStorage& tensor_storage_map, const std::string prefix, - bool enable_vision_ = false) - : GGMLRunner(backend, params_backend), enable_vision(enable_vision_) { - params.arch = arch; - if (arch == LLMArch::MISTRAL_SMALL_3_2 || arch == LLMArch::MINISTRAL_3_3B) { - params.head_dim = 128; - params.num_heads = 32; - params.num_kv_heads = 8; - params.qkv_bias = false; - params.rms_norm_eps = 1e-5f; - } else if (arch == LLMArch::QWEN3) { - params.head_dim = 128; - params.num_heads = 32; - params.num_kv_heads = 8; - params.qkv_bias = false; - params.qk_norm = true; - params.rms_norm_eps = 1e-6f; - } else if (arch == LLMArch::GEMMA3_12B) { - params.head_dim = 256; - params.num_heads = 16; - params.num_kv_heads = 8; - params.qkv_bias = false; - params.qk_norm = true; - params.rms_norm_eps = 1e-6f; - // llama.cpp adds +1 to Gemma3 norm.weight when exporting GGUF, so GGUF loading - // must keep rms_norm_add disabled here or the offset gets applied twice. - // Convenient for the converter, less convenient for whoever gets to debug it later. - params.rms_norm_add = false; - params.normalize_input = true; - params.max_position_embeddings = 131072; - params.mlp_activation = MLPActivation::GELU_TANH; - 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::GEMMA2_2B) { - params.head_dim = 256; - params.num_heads = 8; - params.num_kv_heads = 4; - params.qkv_bias = false; - params.qk_norm = false; - params.rms_norm_eps = 1e-6f; - params.rms_norm_add = true; - params.normalize_input = true; - params.max_position_embeddings = 8192; - params.mlp_activation = MLPActivation::GELU_TANH; - params.hidden_size = 2304; - params.intermediate_size = 9216; - params.num_layers = 26; - params.vocab_size = 256000; - } 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; - params.num_layers = 0; - for (auto pair : tensor_storage_map) { - std::string tensor_name = pair.first; - if (tensor_name.find(prefix) == std::string::npos) - continue; - size_t pos = tensor_name.find("visual."); - if (pos != std::string::npos) { - have_vision_weight = true; - if (contains(tensor_name, "attn.q_proj")) { - llama_cpp_style = true; - } - continue; - } - pos = tensor_name.find("layers."); - if (pos != std::string::npos) { - tensor_name = tensor_name.substr(pos); // remove prefix - auto items = split_string(tensor_name, '.'); - if (items.size() > 1) { - int block_index = atoi(items[1].c_str()); - if (block_index + 1 > params.num_layers) { - params.num_layers = block_index + 1; - } - } - } - if (contains(tensor_name, "embed_tokens.weight")) { - params.hidden_size = pair.second.ne[0]; - params.vocab_size = pair.second.ne[1]; - } - 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; - } - LOG_DEBUG("llm: num_layers = %" PRId64 ", vocab_size = %" PRId64 ", hidden_size = %" PRId64 ", intermediate_size = %" PRId64, - params.num_layers, - params.vocab_size, - params.hidden_size, - params.intermediate_size); - if (enable_vision && !have_vision_weight) { + bool enable_vision_ = false, + std::shared_ptr weight_manager = nullptr) + : GGMLRunner(backend, weight_manager), + config(LLMConfig::detect_from_weights(tensor_storage_map, prefix, arch)), + enable_vision(enable_vision_) { + if (enable_vision && !config.have_vision_weight) { LOG_WARN("no vision weights detected, vision disabled"); enable_vision = false; } if (enable_vision) { LOG_DEBUG("enable llm vision"); - if (llama_cpp_style) { + if (config.llama_cpp_style) { LOG_DEBUG("llama.cpp style vision weight"); } } - model = LLM(params, enable_vision, llama_cpp_style); + model = LLM(config, enable_vision, config.llama_cpp_style); model.init(params_ctx, tensor_storage_map, prefix); } std::string get_desc() override { - return llm_arch_to_str[static_cast(params.arch)]; + return llm_arch_to_str[static_cast(config.arch)]; } void get_param_tensors(std::map& tensors, const std::string prefix) { @@ -1633,12 +1698,12 @@ namespace LLM { } int64_t n_tokens = input_ids->ne[0]; - 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::GEMMA2_2B || - params.arch == LLMArch::GPT_OSS_20B) { + if (config.arch == LLMArch::MISTRAL_SMALL_3_2 || + config.arch == LLMArch::MINISTRAL_3_3B || + config.arch == LLMArch::QWEN3 || + config.arch == LLMArch::GEMMA3_12B || + config.arch == LLMArch::GEMMA2_2B || + config.arch == LLMArch::GPT_OSS_20B) { input_pos_vec.resize(n_tokens); for (int i = 0; i < n_tokens; ++i) { input_pos_vec[i] = i; @@ -1677,9 +1742,9 @@ namespace LLM { set_backend_tensor_data(attention_mask, attention_mask_vec.data()); } - if (params.arch == LLMArch::GEMMA3_12B || params.arch == LLMArch::GPT_OSS_20B) { + if (config.arch == LLMArch::GEMMA3_12B || config.arch == LLMArch::GPT_OSS_20B) { int sliding_window = 0; - for (int window : params.sliding_attention) { + for (int window : config.sliding_attention) { sliding_window = std::max(sliding_window, window); } sliding_attention_mask_vec.resize(n_tokens * n_tokens); @@ -1721,7 +1786,10 @@ namespace LLM { const sd::Tensor& attention_mask, const std::vector>>& image_embeds, std::set out_layers, - bool return_all_hidden_states = false) { + bool return_all_hidden_states = false, + bool auto_free = true, + bool free_compute_buffer = true, + bool free_compute_params = true) { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(input_ids, attention_mask, @@ -1729,21 +1797,21 @@ namespace LLM { out_layers, return_all_hidden_states); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true), + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params), input_ids.dim() + 1); } int64_t get_num_image_tokens(int64_t t, int64_t h, int64_t w) { int64_t grid_t = 1; - int64_t grid_h = h / params.vision.patch_size; - int64_t grid_w = w / params.vision.patch_size; - int64_t llm_grid_h = grid_h / params.vision.spatial_merge_size; - int64_t llm_grid_w = grid_w / params.vision.spatial_merge_size; + int64_t grid_h = h / config.vision.patch_size; + int64_t grid_w = w / config.vision.patch_size; + int64_t llm_grid_h = grid_h / config.vision.spatial_merge_size; + int64_t llm_grid_w = grid_w / config.vision.spatial_merge_size; return grid_t * grid_h * grid_w; } ggml_tensor* process_image(ggml_context* ctx, ggml_tensor* image) { - return process_image_common(ctx, image, params.vision); + return process_image_common(ctx, image, config.vision); } ggml_tensor* build_patch_pos_embeds(GGMLRunnerContext* runner_ctx, @@ -1765,7 +1833,7 @@ namespace LLM { compute_ctx, runner_ctx, image, - params.vision, + config.vision, model.vision_model(), window_index_vec, window_inverse_index_vec, @@ -1779,8 +1847,8 @@ namespace LLM { ggml_cgraph* gf = new_graph_custom(LLM_GRAPH_SIZE); ggml_tensor* image = make_input(image_tensor); - GGML_ASSERT(image->ne[1] % (params.vision.patch_size * params.vision.spatial_merge_size) == 0); - GGML_ASSERT(image->ne[0] % (params.vision.patch_size * params.vision.spatial_merge_size) == 0); + GGML_ASSERT(image->ne[1] % (config.vision.patch_size * config.vision.spatial_merge_size) == 0); + GGML_ASSERT(image->ne[0] % (config.vision.patch_size * config.vision.spatial_merge_size) == 0); auto runnter_ctx = get_context(); ggml_tensor* hidden_states = encode_image(&runnter_ctx, image); @@ -1790,11 +1858,14 @@ namespace LLM { } sd::Tensor encode_image(const int n_threads, - const sd::Tensor& image) { + const sd::Tensor& image, + bool auto_free = false, + bool free_compute_buffer = false, + bool free_compute_params = false) { auto get_graph = [&]() -> ggml_cgraph* { return build_encode_image_graph(image); }; - return take_or_empty(GGMLRunner::compute(get_graph, n_threads, false)); + return take_or_empty(GGMLRunner::compute(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params)); } }; @@ -1804,11 +1875,11 @@ namespace LLM { LLMEmbedder(LLMArch arch, ggml_backend_t backend, - ggml_backend_t params_backend, - const String2TensorStorage& tensor_storage_map = {}, - const std::string prefix = "", - bool enable_vision = false) - : model(arch, backend, params_backend, tensor_storage_map, prefix, enable_vision) { + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "", + bool enable_vision = false, + std::shared_ptr weight_manager = nullptr) + : model(arch, backend, tensor_storage_map, prefix, enable_vision, weight_manager) { if (arch == LLMArch::MISTRAL_SMALL_3_2 || arch == LLMArch::MINISTRAL_3_3B) { tokenizer = std::make_shared(); } else if (arch == LLMArch::GPT_OSS_20B) { @@ -1822,13 +1893,6 @@ namespace LLM { model.get_param_tensors(tensors, prefix); } - bool alloc_params_buffer() { - if (!model.alloc_params_buffer()) { - return false; - } - return true; - } - std::tuple, std::vector> tokenize(std::string text, std::pair attn_range, size_t max_length = 0, @@ -2044,7 +2108,8 @@ namespace LLM { ggml_backend_t backend = sd_backend_cpu_init(); ggml_type model_data_type = GGML_TYPE_COUNT; - ModelLoader model_loader; + auto model_manager = std::make_shared(); + ModelLoader& model_loader = model_manager->loader(); if (!model_loader.init_from_file_and_convert_name(file_path, "text_encoders.llm.")) { LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str()); return; @@ -2062,24 +2127,20 @@ namespace LLM { LLMArch arch = LLMArch::QWEN3; std::shared_ptr llm = std::make_shared(arch, - backend, backend, tensor_storage_map, "text_encoders.llm", - true); + true, + model_manager); - if (!llm->alloc_params_buffer()) { - LOG_ERROR("llm model allocation failed"); - return; - } - - std::map tensors; - llm->get_param_tensors(tensors, "text_encoders.llm"); - - bool success = model_loader.load_tensors(tensors); - - if (!success) { - LOG_ERROR("load tensors from model loader failed"); + if (!model_manager->register_runner_params("LLM test", + *llm, + "text_encoders.llm", + ModelManager::ResidencyMode::ParamBackend, + backend, + backend) || + !model_manager->validate_registered_tensors()) { + LOG_ERROR("register llm tensors with model manager failed"); return; } @@ -2089,4 +2150,4 @@ namespace LLM { }; }; // LLM -#endif // __LLM_HPP__ +#endif // __SD_MODEL_TE_LLM_HPP__ diff --git a/src/t5.hpp b/src/model/te/t5.hpp similarity index 81% rename from src/t5.hpp rename to src/model/te/t5.hpp index 9b2bdaef..6d2326f9 100644 --- a/src/t5.hpp +++ b/src/model/te/t5.hpp @@ -1,5 +1,5 @@ -#ifndef __T5_HPP__ -#define __T5_HPP__ +#ifndef __SD_MODEL_TE_T5_HPP__ +#define __SD_MODEL_TE_T5_HPP__ #include #include @@ -10,10 +10,86 @@ #include #include -#include "ggml_extend.hpp" -#include "model.h" +#include "core/ggml_extend.hpp" +#include "model_loader.h" +#include "model_manager.h" #include "tokenizers/t5_unigram_tokenizer.h" +struct T5Config { + int64_t num_layers = 24; + int64_t model_dim = 4096; + int64_t ff_dim = 10240; + int64_t num_heads = 64; + int64_t vocab_size = 32128; + bool relative_attention = true; + + static T5Config detect_from_weights(const String2TensorStorage& tensor_storage_map, + const std::string& prefix, + bool is_umt5 = false) { + T5Config config; + if (is_umt5) { + config.vocab_size = 256384; + config.relative_attention = false; + } + auto find_tensor = [&](const std::string& suffix) -> const TensorStorage* { + auto it = tensor_storage_map.find(prefix + "." + suffix); + if (it != tensor_storage_map.end()) { + return &it->second; + } + it = tensor_storage_map.find(prefix + suffix); + if (it != tensor_storage_map.end()) { + return &it->second; + } + return nullptr; + }; + + if (const TensorStorage* shared = find_tensor("shared.weight")) { + if (shared->n_dims == 2) { + config.vocab_size = shared->ne[1]; + config.model_dim = shared->ne[0]; + } + } + if (const TensorStorage* q = find_tensor("encoder.block.0.layer.0.SelfAttention.q.weight")) { + if (q->n_dims == 2) { + config.model_dim = q->ne[0]; + int64_t inner_dim = q->ne[1]; + // Flan-T5/T5 uses d_kv=64 for common sizes. + if (inner_dim % 64 == 0) { + config.num_heads = inner_dim / 64; + } + } + } + if (const TensorStorage* wi = find_tensor("encoder.block.0.layer.1.DenseReluDense.wi_0.weight")) { + if (wi->n_dims == 2) { + config.model_dim = wi->ne[0]; + config.ff_dim = wi->ne[1]; + } + } + int64_t detected_layers = 0; + for (const auto& [name, _] : tensor_storage_map) { + std::string base = prefix; + if (!base.empty() && base.back() != '.') { + base += "."; + } + std::string layer_prefix = base + "encoder.block."; + if (!starts_with(name, layer_prefix)) { + continue; + } + size_t pos = layer_prefix.size(); + size_t dot = name.find('.', pos); + if (dot == std::string::npos) { + continue; + } + int64_t layer = atoi(name.substr(pos, dot - pos).c_str()); + detected_layers = std::max(detected_layers, layer + 1); + } + if (detected_layers > 0) { + config.num_layers = detected_layers; + } + return config; + } +}; + class T5LayerNorm : public UnaryBlock { protected: int64_t hidden_size; @@ -272,30 +348,21 @@ public: } }; -struct T5Params { - int64_t num_layers = 24; - int64_t model_dim = 4096; - int64_t ff_dim = 10240; - int64_t num_heads = 64; - int64_t vocab_size = 32128; - bool relative_attention = true; -}; - struct T5 : public GGMLBlock { - T5Params params; + T5Config config; public: T5() {} - T5(T5Params params) - : params(params) { - blocks["encoder"] = std::shared_ptr(new T5Stack(params.num_layers, - params.model_dim, - params.model_dim, - params.ff_dim, - params.num_heads, - params.relative_attention)); - blocks["shared"] = std::shared_ptr(new Embedding(params.vocab_size, - params.model_dim)); + T5(T5Config config) + : config(config) { + blocks["encoder"] = std::shared_ptr(new T5Stack(config.num_layers, + config.model_dim, + config.model_dim, + config.ff_dim, + config.num_heads, + config.relative_attention)); + blocks["shared"] = std::shared_ptr(new Embedding(config.vocab_size, + config.model_dim)); } ggml_tensor* forward(GGMLRunnerContext* ctx, @@ -316,21 +383,18 @@ public: }; struct T5Runner : public GGMLRunner { - T5Params params; + T5Config config; T5 model; std::vector relative_position_bucket_vec; T5Runner(ggml_backend_t backend, - ggml_backend_t params_backend, const String2TensorStorage& tensor_storage_map, const std::string prefix, - bool is_umt5 = false) - : GGMLRunner(backend, params_backend) { - if (is_umt5) { - params.vocab_size = 256384; - params.relative_attention = false; - } - model = T5(params); + bool is_umt5 = false, + std::shared_ptr weight_manager = nullptr) + : GGMLRunner(backend, weight_manager), + config(T5Config::detect_from_weights(tensor_storage_map, prefix, is_umt5)) { + model = T5(config); model.init(params_ctx, tensor_storage_map, prefix); } @@ -384,11 +448,14 @@ struct T5Runner : public GGMLRunner { sd::Tensor compute(const int n_threads, const sd::Tensor& input_ids, - const sd::Tensor& attention_mask) { + const sd::Tensor& attention_mask, + bool auto_free = true, + bool free_compute_buffer = true, + bool free_compute_params = true) { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(input_ids, attention_mask); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true), 3); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, auto_free, free_compute_buffer, free_compute_params), 3); } static std::vector _relative_position_bucket(const std::vector& relative_position, @@ -464,24 +531,17 @@ struct T5Embedder { T5Runner model; T5Embedder(ggml_backend_t backend, - ggml_backend_t params_backend, - const String2TensorStorage& tensor_storage_map = {}, - const std::string prefix = "", - bool is_umt5 = false) - : model(backend, params_backend, tensor_storage_map, prefix, is_umt5), tokenizer(is_umt5) { + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "", + bool is_umt5 = false, + std::shared_ptr weight_manager = nullptr) + : model(backend, tensor_storage_map, prefix, is_umt5, weight_manager), tokenizer(is_umt5) { } void get_param_tensors(std::map& tensors, const std::string prefix) { model.get_param_tensors(tensors, prefix); } - bool alloc_params_buffer() { - if (!model.alloc_params_buffer()) { - return false; - } - return true; - } - std::tuple, std::vector, std::vector> tokenize(std::string text, size_t max_length = 0, bool padding = false) { @@ -566,7 +626,8 @@ struct T5Embedder { ggml_backend_t backend = sd_backend_cpu_init(); ggml_type model_data_type = GGML_TYPE_F16; - ModelLoader model_loader; + auto model_manager = std::make_shared(); + ModelLoader& model_loader = model_manager->loader(); if (!model_loader.init_from_file_and_convert_name(file_path)) { LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str()); return; @@ -579,19 +640,16 @@ struct T5Embedder { } } - std::shared_ptr t5 = std::make_shared(backend, backend, tensor_storage_map, "", true); + std::shared_ptr t5 = std::make_shared(backend, tensor_storage_map, "", true, model_manager); - if (!t5->alloc_params_buffer()) { - LOG_ERROR("t5 params buffer allocation failed"); - return; - } - std::map tensors; - t5->get_param_tensors(tensors, ""); - - bool success = model_loader.load_tensors(tensors); - - if (!success) { - LOG_ERROR("load tensors from model loader failed"); + if (!model_manager->register_runner_params("T5 test", + *t5, + "", + ModelManager::ResidencyMode::ParamBackend, + backend, + backend) || + !model_manager->validate_registered_tensors()) { + LOG_ERROR("register t5 tensors with model manager failed"); return; } @@ -600,4 +658,4 @@ struct T5Embedder { } }; -#endif // __T5_HPP__ +#endif // __SD_MODEL_TE_T5_HPP__ diff --git a/src/model/upscaler/esrgan.hpp b/src/model/upscaler/esrgan.hpp new file mode 100644 index 00000000..21c97712 --- /dev/null +++ b/src/model/upscaler/esrgan.hpp @@ -0,0 +1,273 @@ +#ifndef __SD_MODEL_UPSCALER_ESRGAN_HPP__ +#define __SD_MODEL_UPSCALER_ESRGAN_HPP__ + +#include +#include +#include +#include +#include + +#include "core/ggml_extend.hpp" +#include "core/util.h" + +/* + =================================== ESRGAN =================================== + References: + https://github.com/xinntao/Real-ESRGAN/blob/master/inference_realesrgan.py + https://github.com/XPixelGroup/BasicSR/blob/v1.4.2/basicsr/archs/rrdbnet_arch.py + +*/ + +struct ESRGANConfig { + int scale = 4; + int num_block = 23; + int num_in_ch = 3; + int num_out_ch = 3; + int num_feat = 64; + int num_grow_ch = 32; + + static ESRGANConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, + const std::string& prefix = "") { + ESRGANConfig config; + auto find_weight = [&](const std::string& suffix) -> const TensorStorage* { + std::string name = prefix.empty() ? suffix : prefix + "." + suffix; + auto iter = tensor_storage_map.find(name); + if (iter == tensor_storage_map.end()) { + return nullptr; + } + return &iter->second; + }; + + int detected_num_block = 0; + const std::string body_prefix = prefix.empty() ? "body." : prefix + ".body."; + for (const auto& [name, _] : tensor_storage_map) { + if (!starts_with(name, body_prefix)) { + continue; + } + size_t pos = name.find('.', body_prefix.size()); + if (pos == std::string::npos) { + continue; + } + try { + int idx = std::stoi(name.substr(body_prefix.size(), pos - body_prefix.size())); + detected_num_block = std::max(detected_num_block, idx + 1); + } catch (...) { + } + } + if (detected_num_block > 0) { + config.num_block = detected_num_block; + } + + bool has_conv_up2 = find_weight("conv_up2.weight") != nullptr; + bool has_conv_up1 = find_weight("conv_up1.weight") != nullptr; + bool has_model_tensor = + detected_num_block > 0 || + find_weight("conv_first.weight") != nullptr || + find_weight("conv_hr.weight") != nullptr || + find_weight("conv_last.weight") != nullptr; + if (has_conv_up2) { + config.scale = 4; + } else if (has_conv_up1) { + config.scale = 2; + } else if (has_model_tensor) { + config.scale = 1; + } + + if (has_model_tensor || has_conv_up1 || has_conv_up2) { + LOG_DEBUG("esrgan: scale = %d, num_block = %d, num_in_ch = %d, num_out_ch = %d, num_feat = %d, num_grow_ch = %d", + config.scale, + config.num_block, + config.num_in_ch, + config.num_out_ch, + config.num_feat, + config.num_grow_ch); + } + return config; + } +}; + +class ResidualDenseBlock : public GGMLBlock { +protected: + int num_feat; + int num_grow_ch; + +public: + ResidualDenseBlock(int num_feat = 64, int num_grow_ch = 32) + : num_feat(num_feat), num_grow_ch(num_grow_ch) { + blocks["conv1"] = std::shared_ptr(new Conv2d(num_feat, num_grow_ch, {3, 3}, {1, 1}, {1, 1})); + blocks["conv2"] = std::shared_ptr(new Conv2d(num_feat + num_grow_ch, num_grow_ch, {3, 3}, {1, 1}, {1, 1})); + blocks["conv3"] = std::shared_ptr(new Conv2d(num_feat + 2 * num_grow_ch, num_grow_ch, {3, 3}, {1, 1}, {1, 1})); + blocks["conv4"] = std::shared_ptr(new Conv2d(num_feat + 3 * num_grow_ch, num_grow_ch, {3, 3}, {1, 1}, {1, 1})); + blocks["conv5"] = std::shared_ptr(new Conv2d(num_feat + 4 * num_grow_ch, num_feat, {3, 3}, {1, 1}, {1, 1})); + } + + ggml_tensor* lrelu(GGMLRunnerContext* ctx, ggml_tensor* x) { + return ggml_leaky_relu(ctx->ggml_ctx, x, 0.2f, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + // x: [n, num_feat, h, w] + // return: [n, num_feat, h, w] + + auto conv1 = std::dynamic_pointer_cast(blocks["conv1"]); + auto conv2 = std::dynamic_pointer_cast(blocks["conv2"]); + auto conv3 = std::dynamic_pointer_cast(blocks["conv3"]); + auto conv4 = std::dynamic_pointer_cast(blocks["conv4"]); + auto conv5 = std::dynamic_pointer_cast(blocks["conv5"]); + + auto x1 = lrelu(ctx, conv1->forward(ctx, x)); + auto x_cat = ggml_concat(ctx->ggml_ctx, x, x1, 2); + auto x2 = lrelu(ctx, conv2->forward(ctx, x_cat)); + x_cat = ggml_concat(ctx->ggml_ctx, x_cat, x2, 2); + auto x3 = lrelu(ctx, conv3->forward(ctx, x_cat)); + x_cat = ggml_concat(ctx->ggml_ctx, x_cat, x3, 2); + auto x4 = lrelu(ctx, conv4->forward(ctx, x_cat)); + x_cat = ggml_concat(ctx->ggml_ctx, x_cat, x4, 2); + auto x5 = conv5->forward(ctx, x_cat); + + x5 = ggml_add(ctx->ggml_ctx, ggml_ext_scale(ctx->ggml_ctx, x5, 0.2f), x); + return x5; + } +}; + +class RRDB : public GGMLBlock { +public: + RRDB(int num_feat, int num_grow_ch = 32) { + blocks["rdb1"] = std::shared_ptr(new ResidualDenseBlock(num_feat, num_grow_ch)); + blocks["rdb2"] = std::shared_ptr(new ResidualDenseBlock(num_feat, num_grow_ch)); + blocks["rdb3"] = std::shared_ptr(new ResidualDenseBlock(num_feat, num_grow_ch)); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + // x: [n, num_feat, h, w] + // return: [n, num_feat, h, w] + + auto rdb1 = std::dynamic_pointer_cast(blocks["rdb1"]); + auto rdb2 = std::dynamic_pointer_cast(blocks["rdb2"]); + auto rdb3 = std::dynamic_pointer_cast(blocks["rdb3"]); + + auto out = rdb1->forward(ctx, x); + out = rdb2->forward(ctx, out); + out = rdb3->forward(ctx, out); + + out = ggml_add(ctx->ggml_ctx, ggml_ext_scale(ctx->ggml_ctx, out, 0.2f), x); + return out; + } +}; + +class RRDBNet : public GGMLBlock { +protected: + ESRGANConfig config; + +public: + explicit RRDBNet(ESRGANConfig config) + : config(std::move(config)) { + blocks["conv_first"] = std::shared_ptr(new Conv2d(this->config.num_in_ch, this->config.num_feat, {3, 3}, {1, 1}, {1, 1})); + for (int i = 0; i < this->config.num_block; i++) { + std::string name = "body." + std::to_string(i); + blocks[name] = std::shared_ptr(new RRDB(this->config.num_feat, this->config.num_grow_ch)); + } + blocks["conv_body"] = std::shared_ptr(new Conv2d(this->config.num_feat, this->config.num_feat, {3, 3}, {1, 1}, {1, 1})); + if (this->config.scale >= 2) { + blocks["conv_up1"] = std::shared_ptr(new Conv2d(this->config.num_feat, this->config.num_feat, {3, 3}, {1, 1}, {1, 1})); + } + if (this->config.scale == 4) { + blocks["conv_up2"] = std::shared_ptr(new Conv2d(this->config.num_feat, this->config.num_feat, {3, 3}, {1, 1}, {1, 1})); + } + blocks["conv_hr"] = std::shared_ptr(new Conv2d(this->config.num_feat, this->config.num_feat, {3, 3}, {1, 1}, {1, 1})); + blocks["conv_last"] = std::shared_ptr(new Conv2d(this->config.num_feat, this->config.num_out_ch, {3, 3}, {1, 1}, {1, 1})); + } + + int get_scale() { return config.scale; } + int get_num_block() { return config.num_block; } + + ggml_tensor* lrelu(GGMLRunnerContext* ctx, ggml_tensor* x) { + return ggml_leaky_relu(ctx->ggml_ctx, x, 0.2f, true); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + // x: [n, num_in_ch, h, w] + // return: [n, num_out_ch, h*scale, w*scale] + auto conv_first = std::dynamic_pointer_cast(blocks["conv_first"]); + auto conv_body = std::dynamic_pointer_cast(blocks["conv_body"]); + auto conv_hr = std::dynamic_pointer_cast(blocks["conv_hr"]); + auto conv_last = std::dynamic_pointer_cast(blocks["conv_last"]); + + auto feat = conv_first->forward(ctx, x); + sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.prelude", "feat"); + auto body_feat = feat; + for (int i = 0; i < config.num_block; i++) { + std::string name = "body." + std::to_string(i); + auto block = std::dynamic_pointer_cast(blocks[name]); + + body_feat = block->forward(ctx, body_feat); + sd::ggml_graph_cut::mark_graph_cut(body_feat, "esrgan.body." + std::to_string(i), "feat"); + } + body_feat = conv_body->forward(ctx, body_feat); + feat = ggml_add(ctx->ggml_ctx, feat, body_feat); + sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.body.out", "feat"); + // upsample + if (config.scale >= 2) { + auto conv_up1 = std::dynamic_pointer_cast(blocks["conv_up1"]); + feat = lrelu(ctx, conv_up1->forward(ctx, ggml_upscale(ctx->ggml_ctx, feat, 2, GGML_SCALE_MODE_NEAREST))); + sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.up1", "feat"); + if (config.scale == 4) { + auto conv_up2 = std::dynamic_pointer_cast(blocks["conv_up2"]); + feat = lrelu(ctx, conv_up2->forward(ctx, ggml_upscale(ctx->ggml_ctx, feat, 2, GGML_SCALE_MODE_NEAREST))); + sd::ggml_graph_cut::mark_graph_cut(feat, "esrgan.up2", "feat"); + } + } + // for all scales + auto out = conv_last->forward(ctx, lrelu(ctx, conv_hr->forward(ctx, feat))); + sd::ggml_graph_cut::mark_graph_cut(out, "esrgan.final", "out"); + return out; + } +}; + +struct ESRGAN : public GGMLRunner { + ESRGANConfig config; + std::unique_ptr rrdb_net; + + ESRGAN(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map = {}, + std::shared_ptr weight_manager = nullptr) + : GGMLRunner(backend, weight_manager), + config(ESRGANConfig::detect_from_weights(tensor_storage_map)), + rrdb_net(std::make_unique(config)) { + rrdb_net->init(params_ctx, tensor_storage_map, ""); + } + + std::string get_desc() override { + return "esrgan"; + } + + void get_param_tensors(std::map& tensors) { + if (!rrdb_net) { + return; + } + + rrdb_net->get_param_tensors(tensors); + } + + ggml_cgraph* build_graph(const sd::Tensor& x_tensor) { + if (!rrdb_net) + return nullptr; + constexpr int kGraphNodes = 1 << 16; // 65k + ggml_cgraph* gf = new_graph_custom(kGraphNodes); + ggml_tensor* x = make_input(x_tensor); + + auto runner_ctx = get_context(); + ggml_tensor* out = rrdb_net->forward(&runner_ctx, x); + ggml_build_forward_expand(gf, out); + return gf; + } + + sd::Tensor compute(const int n_threads, + const sd::Tensor& x) { + auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x); }; + auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), x.dim()); + return result; + } +}; + +#endif // __SD_MODEL_UPSCALER_ESRGAN_HPP__ diff --git a/src/ltx_latent_upscaler.hpp b/src/model/upscaler/ltx_latent_upscaler.hpp similarity index 71% rename from src/ltx_latent_upscaler.hpp rename to src/model/upscaler/ltx_latent_upscaler.hpp index ea4a830c..b70e1613 100644 --- a/src/ltx_latent_upscaler.hpp +++ b/src/model/upscaler/ltx_latent_upscaler.hpp @@ -1,9 +1,9 @@ -#ifndef __SD_LTX_LATENT_UPSCALER_HPP__ -#define __SD_LTX_LATENT_UPSCALER_HPP__ +#ifndef __SD_MODEL_UPSCALER_LTX_LATENT_UPSCALER_HPP__ +#define __SD_MODEL_UPSCALER_LTX_LATENT_UPSCALER_HPP__ +#include #include #include -#include #include #include #include @@ -11,11 +11,11 @@ #include #include -#include "common_dit.hpp" -#include "ggml_extend.hpp" -#include "ggml_graph_cut.h" -#include "model.h" -#include "util.h" +#include "core/ggml_extend.hpp" +#include "core/ggml_graph_cut.h" +#include "core/util.h" +#include "model/diffusion/dit.hpp" +#include "model_loader.h" namespace LTXVUpsampler { constexpr int LTX_UPSAMPLER_GRAPH_SIZE = 10240; @@ -32,90 +32,100 @@ namespace LTXVUpsampler { 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 LatentUpsamplerConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, + const std::string& prefix = "") { + LatentUpsamplerConfig config; + auto find_weight = [&](const std::string& suffix) -> const TensorStorage* { + std::string name = prefix.empty() ? suffix : prefix + "." + suffix; + auto iter = tensor_storage_map.find(name); + if (iter == tensor_storage_map.end()) { + return nullptr; + } + return &iter->second; + }; - 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]; - } + bool inferred = false; - 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; + const TensorStorage* initial_norm = find_weight("initial_norm.weight"); + if (initial_norm != nullptr) { + config.mid_channels = initial_norm->ne[0]; + inferred = true; } - 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(std::round(std::sqrt(static_cast(ratio)))); - if (num > 0 && static_cast(num) * num == ratio) { - config.spatial_up_num = num; + const TensorStorage* final_conv = find_weight("final_conv.bias"); + if (final_conv != nullptr) { + config.in_channels = final_conv->ne[0]; + inferred = true; + } + + int detected_blocks = 0; + const std::string res_blocks_prefix = prefix.empty() ? "res_blocks." : prefix + ".res_blocks."; + for (const auto& [name, _] : tensor_storage_map) { + if (!starts_with(name, res_blocks_prefix)) { + continue; + } + size_t begin = res_blocks_prefix.size(); + size_t end = name.find('.', begin); + if (end == std::string::npos) { + continue; + } + try { + int idx = std::stoi(name.substr(begin, end - begin)); + detected_blocks = std::max(detected_blocks, idx + 1); + } catch (...) { } } - 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(config.spatial_up_num); + if (detected_blocks > 0) { + config.num_blocks_per_stage = detected_blocks; + inferred = true; } + + const TensorStorage* rational_upsampler_weight = find_weight("upsampler.conv.weight"); + const TensorStorage* upsampler_bias = find_weight("upsampler.0.bias"); + config.rational_resampler = rational_upsampler_weight != nullptr; + int64_t upsampler_out_channels = upsampler_bias == nullptr ? 0 : upsampler_bias->ne[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.rational_resampler || upsampler_out_channels > 0) { + inferred = true; + } + if (config.temporal_upsample) { + config.temporal_up_factor = 2; + } + if (rational_upsampler_weight != nullptr) { + int64_t out_channels = rational_upsampler_weight->ne[3]; + if (config.mid_channels > 0 && out_channels % config.mid_channels == 0) { + int64_t ratio = out_channels / config.mid_channels; + int num = static_cast(std::round(std::sqrt(static_cast(ratio)))); + if (num > 0 && static_cast(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(config.spatial_up_num); + } + } + + if (inferred) { + LOG_DEBUG("ltx latent upsampler: in_channels = %" PRId64 ", mid_channels = %" PRId64 ", num_blocks_per_stage = %d, spatial_scale = %.3f, temporal_up_factor = %d, rational_resampler = %d", + config.in_channels, + config.mid_channels, + config.num_blocks_per_stage, + config.spatial_scale, + config.temporal_up_factor, + config.rational_resampler); + } + return config; } - return config; - } + }; class VideoGroupNorm : public GGMLBlock { protected: @@ -240,20 +250,25 @@ namespace LTXVUpsampler { protected: int64_t channels; int stride; - ggml_tensor* kernel = nullptr; std::vector kernel_data; + std::string kernel_name; void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { + SD_UNUSED(ctx); 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()); + kernel_name = prefix + "kernel"; + } + public: + BlurDownsample(int64_t channels, int stride) + : channels(channels), + stride(stride) { + GGML_ASSERT(stride >= 1); static const float binomial[5] = {1.f, 4.f, 6.f, 4.f, 1.f}; kernel_data.resize(static_cast(5 * 5 * channels)); for (int64_t c = 0; c < channels; ++c) { @@ -266,26 +281,16 @@ namespace LTXVUpsampler { } } - 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(ctx != nullptr); + GGML_ASSERT(!kernel_data.empty()); GGML_ASSERT(x->ne[2] == channels); + ggml_tensor* kernel = ggml_new_tensor_4d(ctx->ggml_ctx, GGML_TYPE_F32, 5, 5, 1, channels); + ggml_set_name(kernel, kernel_name.empty() ? "blur_down.kernel" : kernel_name.c_str()); + ctx->bind_backend_tensor_data(kernel, kernel_data.data()); if (ctx->conv2d_direct_enabled) { return ggml_conv_2d_dw_direct(ctx->ggml_ctx, kernel, x, stride, stride, 2, 2, 1, 1); } @@ -311,11 +316,6 @@ namespace LTXVUpsampler { blocks["blur_down"] = std::shared_ptr(new BlurDownsample(mid_channels, den)); } - void load_fixed_tensors() { - auto blur_down = std::dynamic_pointer_cast(blocks["blur_down"]); - blur_down->load_fixed_tensors(); - } - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { auto conv = std::dynamic_pointer_cast(blocks["conv"]); auto pixel_shuffle = std::dynamic_pointer_cast(blocks["pixel_shuffle"]); @@ -426,45 +426,17 @@ namespace LTXVUpsampler { 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(blocks["upsampler"]); - upsampler->load_fixed_tensors(); - } }; struct LatentUpsamplerRunner : public GGMLRunner { + LatentUpsamplerConfig config; std::unique_ptr 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); + const String2TensorStorage& tensor_storage_map, + std::shared_ptr weight_manager = nullptr) + : GGMLRunner(backend, weight_manager), + config(LatentUpsamplerConfig::detect_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", @@ -474,36 +446,21 @@ namespace LTXVUpsampler { config.rational_resampler, config.spatial_scale, config.temporal_up_factor); - return false; + return; } model = std::make_unique(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 tensors; - model->get_param_tensors(tensors); - std::set 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(); + std::string get_desc() override { + return "ltx_latent_upsampler"; + } - 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; + void get_param_tensors(std::map& tensors) { + if (model) { + model->get_param_tensors(tensors); + } } ggml_cgraph* build_graph(const sd::Tensor& x_tensor) { @@ -534,18 +491,18 @@ namespace LTXVUpsampler { (long long)x.shape()[4]); return {}; } - if (x.shape()[3] != model->config.in_channels) { + if (x.shape()[3] != config.in_channels) { LOG_ERROR("LTX latent upsampler expected %" PRId64 " channels, got %lld", - model->config.in_channels, + config.in_channels, (long long)x.shape()[3]); return {}; } size_t expected_dim = static_cast(x.dim()); auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), expected_dim); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), expected_dim); } }; } // namespace LTXVUpsampler -#endif // __SD_LTX_LATENT_UPSCALER_HPP__ +#endif // __SD_MODEL_UPSCALER_LTX_LATENT_UPSCALER_HPP__ diff --git a/src/auto_encoder_kl.hpp b/src/model/vae/auto_encoder_kl.hpp similarity index 96% rename from src/auto_encoder_kl.hpp rename to src/model/vae/auto_encoder_kl.hpp index 13396e73..604347d7 100644 --- a/src/auto_encoder_kl.hpp +++ b/src/model/vae/auto_encoder_kl.hpp @@ -1,7 +1,7 @@ -#ifndef __AUTO_ENCODER_KL_HPP__ -#define __AUTO_ENCODER_KL_HPP__ +#ifndef __SD_MODEL_VAE_AUTO_ENCODER_KL_HPP__ +#define __SD_MODEL_VAE_AUTO_ENCODER_KL_HPP__ -#include "vae.hpp" +#include "model/vae/vae.hpp" /*================================================== AutoEncoderKL ===================================================*/ @@ -213,9 +213,9 @@ protected: params["mix_factor"] = ggml_new_tensor_1d(ctx, wtype, 1); } - float get_alpha() { - float alpha = ggml_ext_backend_tensor_get_f32(params["mix_factor"]); - return sigmoid(alpha); + ggml_tensor* get_alpha(GGMLRunnerContext* ctx) { + auto mix_factor = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, params["mix_factor"]); + return ggml_sigmoid(ctx->ggml_ctx, mix_factor); } public: @@ -250,10 +250,12 @@ public: x = time_stack->forward(ctx, x); // b t c (h w) - float alpha = get_alpha(); - x = ggml_add(ctx->ggml_ctx, - ggml_ext_scale(ctx->ggml_ctx, x, alpha), - ggml_ext_scale(ctx->ggml_ctx, x_mix, 1.0f - alpha)); + auto alpha = get_alpha(ctx); + x = ggml_add(ctx->ggml_ctx, + x_mix, + ggml_mul(ctx->ggml_ctx, + ggml_sub(ctx->ggml_ctx, x, x_mix), + alpha)); x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 2, 1, 3)); // b c t (h w) -> b t c (h w) x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H, C, T * B); // b t c (h w) -> (b t) c h w @@ -664,13 +666,13 @@ struct AutoEncoderKL : public VAE { AutoEncoderKLModel ae; AutoEncoderKL(ggml_backend_t backend, - ggml_backend_t params_backend, const String2TensorStorage& tensor_storage_map, const std::string prefix, - bool decode_only = false, - bool use_video_decoder = false, - SDVersion version = VERSION_SD1) - : decode_only(decode_only), VAE(version, backend, params_backend) { + bool decode_only = false, + bool use_video_decoder = false, + SDVersion version = VERSION_SD1, + std::shared_ptr weight_manager = nullptr) + : VAE(version, backend, prefix, weight_manager), decode_only(decode_only) { if (sd_version_is_sd1(version) || sd_version_is_sd2(version)) { scale_factor = 0.18215f; shift_factor = 0.f; @@ -680,7 +682,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) || sd_version_is_longcat(version)) { + } else if (sd_version_uses_flux_vae(version)) { scale_factor = 0.3611f; shift_factor = 0.1159f; } else if (sd_version_uses_flux2_vae(version)) { @@ -718,8 +720,8 @@ struct AutoEncoderKL : public VAE { return "vae"; } - void get_param_tensors(std::map& tensors, const std::string prefix) override { - ae.get_param_tensors(tensors, prefix); + void get_param_tensors(std::map& tensors) override { + ae.get_param_tensors(tensors, weight_prefix); } ggml_cgraph* build_graph(const sd::Tensor& z_tensor, bool decode_graph) { @@ -742,7 +744,7 @@ struct AutoEncoderKL : public VAE { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(z, decode_graph); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), z.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), z.dim()); } sd::Tensor gaussian_latent_sample(const sd::Tensor& moments, std::shared_ptr rng) { @@ -814,12 +816,13 @@ struct AutoEncoderKL : public VAE { } sd::Tensor diffusion_to_vae_latents(const sd::Tensor& latents) override { + auto latents_ = sd_version_is_sefi_image(version) ? sd::ops::slice(latents, 2, 16, 144) : latents; if (sd_version_uses_flux2_vae(version)) { int channel_dim = 2; - auto [mean_tensor, std_tensor] = get_latents_mean_std(latents, channel_dim); - return (latents * std_tensor) / scale_factor + mean_tensor; + auto [mean_tensor, std_tensor] = get_latents_mean_std(latents_, channel_dim); + return (latents_ * std_tensor) / scale_factor + mean_tensor; } - return (latents / scale_factor) + shift_factor; + return (latents_ / scale_factor) + shift_factor; } sd::Tensor vae_to_diffusion_latents(const sd::Tensor& latents) override { @@ -886,4 +889,4 @@ struct AutoEncoderKL : public VAE { }; }; -#endif // __AUTO_ENCODER_KL_HPP__ +#endif // __SD_MODEL_VAE_AUTO_ENCODER_KL_HPP__ diff --git a/src/ltx_audio_vae.h b/src/model/vae/ltx_audio_vae.hpp similarity index 96% rename from src/ltx_audio_vae.h rename to src/model/vae/ltx_audio_vae.hpp index 88c37631..2f8c03d1 100644 --- a/src/ltx_audio_vae.h +++ b/src/model/vae/ltx_audio_vae.hpp @@ -1,5 +1,5 @@ -#ifndef __SD_LTX_AUDIO_VAE_H__ -#define __SD_LTX_AUDIO_VAE_H__ +#ifndef __SD_MODEL_VAE_LTX_AUDIO_VAE_HPP__ +#define __SD_MODEL_VAE_LTX_AUDIO_VAE_HPP__ #include #include @@ -7,7 +7,9 @@ #include #include -#include "ggml_extend.hpp" +#include "core/ggml_extend.hpp" +#include "model_loader.h" +#include "model_manager.h" namespace LTXV { @@ -58,11 +60,12 @@ namespace LTXV { return base_output_sample_rate(); } - static LTXAudioVAEConfig detect_from_weights(const String2TensorStorage& tensor_storage_map) { + static LTXAudioVAEConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix = "") { LTXAudioVAEConfig config; auto require = [&](const std::string& name) -> const TensorStorage* { - auto iter = tensor_storage_map.find(name); + std::string tensor_name = prefix.empty() ? name : prefix + "." + name; + auto iter = tensor_storage_map.find(tensor_name); if (iter == tensor_storage_map.end()) { return nullptr; } @@ -168,6 +171,12 @@ namespace LTXV { if (config.audio_channels != 2 || config.latent_channels != 8 || config.mel_bins != 64) { return config; } + LOG_DEBUG("ltx_audio_vae: sample_rate = %d, mel_bins = %d, latent_channels = %d, latent_frequency_bins = %d, has_bwe = %s", + config.sample_rate, + config.mel_bins, + config.latent_channels, + config.latent_frequency_bins, + config.has_bwe ? "true" : "false"); return config; } }; @@ -989,13 +998,15 @@ namespace LTXV { struct LTXAudioVAERunner : public GGMLRunner { LTXAudioVAEConfig config; LTXAudioVAE model; + std::string weight_prefix; sd::Tensor bwe_skip_filter_tensor; LTXAudioVAERunner(ggml_backend_t backend, - ggml_backend_t params_backend, const String2TensorStorage& tensor_storage_map, - const std::string& prefix = "") - : GGMLRunner(backend, params_backend), + const std::string& prefix = "", + std::shared_ptr weight_manager = nullptr) + : GGMLRunner(backend, weight_manager), + weight_prefix(prefix), config(LTXAudioVAEConfig::detect_from_weights(tensor_storage_map)), model(config) { model.init(params_ctx, tensor_storage_map, prefix); @@ -1005,11 +1016,11 @@ namespace LTXV { } } - void get_param_tensors(std::map& tensors, const std::string prefix) { - model.get_param_tensors(tensors, prefix); + void get_param_tensors(std::map& tensors) { + model.get_param_tensors(tensors, weight_prefix); } - size_t get_params_buffer_size() { + size_t get_params_mem_size() { return model.get_params_mem_size(); } @@ -1029,7 +1040,7 @@ namespace LTXV { ggml_build_forward_expand(gf, waveform); return gf; }; - auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), 4); + auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), 4); int64_t t1 = ggml_time_ms(); LOG_INFO("ltx audio vae decode completed, taking %.2fs", (t1 - t0) * 1.0f / 1000); return result; @@ -1056,7 +1067,8 @@ namespace LTXV { // ggml_backend_t backend = ggml_backend_cuda_init(0); LOG_INFO("loading ltx audio vae from '%s'", model_path.c_str()); - ModelLoader model_loader; + auto model_manager = std::make_shared(); + ModelLoader& model_loader = model_manager->loader(); if (!model_loader.init_from_file(model_path)) { LOG_ERROR("init model loader from file failed: '%s'", model_path.c_str()); return; @@ -1064,20 +1076,17 @@ namespace LTXV { auto& tensor_storage_map = model_loader.get_tensor_storage_map(); auto ltx_audio_vae = std::make_shared(backend, - backend, tensor_storage_map, - prefix); + prefix, + model_manager); - if (!ltx_audio_vae->alloc_params_buffer()) { - LOG_ERROR("ltx audio vae buffer allocation failed"); - return; - } - - std::map tensors; - ltx_audio_vae->get_param_tensors(tensors, ""); - - if (!model_loader.load_tensors(tensors)) { - LOG_ERROR("load tensors from model loader failed"); + if (!model_manager->register_runner_params("LTX audio VAE test", + *ltx_audio_vae, + ModelManager::ResidencyMode::ParamBackend, + backend, + backend) || + !model_manager->validate_registered_tensors()) { + LOG_ERROR("register ltx audio vae tensors with model manager failed"); return; } @@ -1088,4 +1097,4 @@ namespace LTXV { } // namespace LTXV -#endif // __SD_LTX_AUDIO_VAE_H__ +#endif // __SD_MODEL_VAE_LTX_AUDIO_VAE_HPP__ diff --git a/src/ltx_vae.hpp b/src/model/vae/ltx_vae.hpp similarity index 97% rename from src/ltx_vae.hpp rename to src/model/vae/ltx_vae.hpp index 13ec0e39..a19ce820 100644 --- a/src/ltx_vae.hpp +++ b/src/model/vae/ltx_vae.hpp @@ -1,5 +1,5 @@ -#ifndef __SD_LTX_VAE_HPP__ -#define __SD_LTX_VAE_HPP__ +#ifndef __SD_MODEL_VAE_LTX_VAE_HPP__ +#define __SD_MODEL_VAE_LTX_VAE_HPP__ #include #include @@ -9,9 +9,10 @@ #include #include -#include "ltxv.hpp" -#include "vae.hpp" -#include "wan.hpp" +#include "model/diffusion/ltxv.hpp" +#include "model/vae/vae.hpp" +#include "model/vae/wan_vae.hpp" +#include "model_loader.h" namespace LTXVAE { @@ -956,8 +957,8 @@ namespace LTXVAE { ggml_tensor* scaled_timestep = timestep; if (timestep_conditioning) { - auto multiplier = ggml_ext_backend_tensor_get_f32(params["timestep_scale_multiplier"]); - scaled_timestep = ggml_ext_scale(ctx->ggml_ctx, timestep, multiplier); + auto multiplier = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, params["timestep_scale_multiplier"]); + scaled_timestep = ggml_mul(ctx->ggml_ctx, timestep, multiplier); } x = conv_in->forward(ctx, x, causal_decoder); @@ -1007,8 +1008,8 @@ namespace LTXVAE { ggml_tensor* scaled_timestep = timestep; if (timestep_conditioning && timestep != nullptr) { - auto multiplier = ggml_ext_backend_tensor_get_f32(params["timestep_scale_multiplier"]); - scaled_timestep = ggml_ext_scale(ctx->ggml_ctx, timestep, multiplier); + auto multiplier = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, params["timestep_scale_multiplier"]); + scaled_timestep = ggml_mul(ctx->ggml_ctx, timestep, multiplier); } // conv_in with feat_map for left temporal context @@ -1222,11 +1223,11 @@ struct LTXVideoVAE : public VAE { LTXVAE::VideoVAE vae; LTXVideoVAE(ggml_backend_t backend, - ggml_backend_t params_backend, const String2TensorStorage& tensor_storage_map, const std::string& prefix, - bool decode_only = true, - SDVersion version = VERSION_LTXAV) + bool decode_only = true, + SDVersion version = VERSION_LTXAV, + std::shared_ptr weight_manager = nullptr) : decode_only(decode_only), ltx_vae_version(LTXVAE::detect_ltx_vae_version(tensor_storage_map, prefix)), timestep_conditioning(LTXVAE::detect_ltx_vae_timestep_conditioning(tensor_storage_map, prefix)), @@ -1238,7 +1239,7 @@ struct LTXVideoVAE : public VAE { patch_size, tensor_storage_map, prefix), - VAE(version, backend, params_backend) { + VAE(version, backend, prefix, weight_manager) { vae.init(params_ctx, tensor_storage_map, prefix); decode_timestep_tensor.values()[0] = vae.decode_timestep; } @@ -1270,8 +1271,8 @@ struct LTXVideoVAE : public VAE { } } - void get_param_tensors(std::map& tensors, const std::string prefix) override { - vae.get_param_tensors(tensors, prefix); + void get_param_tensors(std::map& tensors) override { + vae.get_param_tensors(tensors, weight_prefix); } struct TemporalTilePlan { @@ -1395,7 +1396,7 @@ struct LTXVideoVAE : public VAE { static_cast(start), chunk_overlap); }; - auto chunk = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true), + auto chunk = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true, true, true), expected_dim); if (chunk.empty()) { free_cache_ctx_and_buffer(); @@ -1425,7 +1426,7 @@ struct LTXVideoVAE : public VAE { const sd::Tensor& z, bool decode_graph) override { if (!decode_graph && decode_only) { - LOG_ERROR("LTX video VAE encode requires encoder weights; create the context with vae_decode_only=false"); + LOG_ERROR("LTX video VAE encode requires encoder weights"); return {}; } sd::Tensor input = z; @@ -1451,7 +1452,7 @@ struct LTXVideoVAE : public VAE { auto get_graph = [&]() -> ggml_cgraph* { return build_graph(input, decode_graph); }; - auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), expected_dim); + auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), expected_dim); if (result.empty()) { return {}; } @@ -1464,7 +1465,7 @@ struct LTXVideoVAE : public VAE { auto get_graph = [&]() -> ggml_cgraph* { return build_latent_statistics_graph(z, normalize); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), static_cast(z.dim())); } @@ -1520,7 +1521,8 @@ struct LTXVideoVAE : public VAE { ggml_backend_t backend = sd_backend_cpu_init(); LOG_INFO("loading ltx vae from '%s'", model_path.c_str()); - ModelLoader model_loader; + auto model_manager = std::make_shared(); + ModelLoader& model_loader = model_manager->loader(); if (!model_loader.init_from_file_and_convert_name(model_path, "vae.")) { LOG_ERROR("init model loader from file failed: '%s'", model_path.c_str()); return; @@ -1528,22 +1530,19 @@ struct LTXVideoVAE : public VAE { auto& tensor_storage_map = model_loader.get_tensor_storage_map(); std::shared_ptr vae = std::make_shared(backend, - backend, tensor_storage_map, "first_stage_model", true, - VERSION_LTXAV); + VERSION_LTXAV, + model_manager); - if (!vae->alloc_params_buffer()) { - LOG_ERROR("vae buffer allocation failed"); - return; - } - - std::map tensors; - vae->get_param_tensors(tensors, "first_stage_model"); - - if (!model_loader.load_tensors(tensors)) { - LOG_ERROR("load tensors from model loader failed"); + if (!model_manager->register_runner_params("LTX VAE test", + *vae, + ModelManager::ResidencyMode::ParamBackend, + backend, + backend) || + !model_manager->validate_registered_tensors()) { + LOG_ERROR("register ltx vae tensors with model manager failed"); return; } @@ -1552,4 +1551,4 @@ struct LTXVideoVAE : public VAE { } }; -#endif // __SD_LTX_VAE_HPP__ +#endif // __SD_MODEL_VAE_LTX_VAE_HPP__ diff --git a/src/tae.hpp b/src/model/vae/tae.hpp similarity index 94% rename from src/tae.hpp rename to src/model/vae/tae.hpp index f43ef48f..a78e5e96 100644 --- a/src/tae.hpp +++ b/src/model/vae/tae.hpp @@ -1,13 +1,13 @@ -#ifndef __TAE_HPP__ -#define __TAE_HPP__ +#ifndef __SD_MODEL_VAE_TAE_HPP__ +#define __SD_MODEL_VAE_TAE_HPP__ -#include "ggml_extend.hpp" +#include "core/ggml_extend.hpp" #include "model.h" /* =================================== TinyAutoEncoder =================================== References: - https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/autoencoders/vae.py + https://github.com/huggingface/diffusers/blob/main/src/diffusers/model/autoencoders/vae.py https://github.com/madebyollin/taesd/blob/main/taesd.py */ @@ -548,7 +548,7 @@ public: } auto result = decoder->forward(ctx, z); if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) { - // (W, H, C, T) -> (W, H, T, C) + // (W, H, T, C) -> (W, H, C, T) result = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, result, 0, 1, 3, 2)); } return result; @@ -556,8 +556,10 @@ public: ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* x) { auto encoder = std::dynamic_pointer_cast(blocks["encoder"]); - // (W, H, T, C) -> (W, H, C, T) - x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2)); + if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) { + // (W, H, T, C) -> (W, H, C, T) + x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2)); + } int64_t num_frames = x->ne[3]; if (num_frames % encoder->t_downscale) { // pad to multiple of encoder->t_downscale at the end @@ -567,7 +569,10 @@ public: } } x = encoder->forward(ctx, x); - x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2)); + if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) { + // (W, H, C, T) -> (W, H, T, C) + x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2)); + } return x; } }; @@ -623,14 +628,14 @@ struct TinyImageAutoEncoder : public VAE { bool decode_only = false; TinyImageAutoEncoder(ggml_backend_t backend, - ggml_backend_t params_backend, const String2TensorStorage& tensor_storage_map, const std::string prefix, - bool decoder_only = true, - SDVersion version = VERSION_SD1) - : decode_only(decoder_only), - taesd(decoder_only, version), - VAE(version, backend, params_backend) { + bool decoder_only = true, + SDVersion version = VERSION_SD1, + std::shared_ptr weight_manager = nullptr) + : VAE(version, backend, "tae", weight_manager), + decode_only(decoder_only), + taesd(decoder_only, version) { scale_input = false; taesd.init(params_ctx, tensor_storage_map, prefix); } @@ -639,8 +644,8 @@ struct TinyImageAutoEncoder : public VAE { return "taesd"; } - void get_param_tensors(std::map& tensors, const std::string prefix) { - taesd.get_param_tensors(tensors, prefix); + void get_param_tensors(std::map& tensors) override { + taesd.get_param_tensors(tensors, weight_prefix); } sd::Tensor vae_output_to_latents(const sd::Tensor& vae_output, std::shared_ptr rng) override { @@ -676,7 +681,7 @@ struct TinyImageAutoEncoder : public VAE { return build_graph(z_tensor, decode_graph); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), z_tensor.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), z_tensor.dim()); } }; @@ -686,13 +691,13 @@ struct TinyVideoAutoEncoder : public VAE { bool is_wide = false; TinyVideoAutoEncoder(ggml_backend_t backend, - ggml_backend_t params_backend, const String2TensorStorage& tensor_storage_map, const std::string prefix, - bool decoder_only = true, - SDVersion version = VERSION_WAN2) - : decode_only(decoder_only), - VAE(version, backend, params_backend) { + bool decoder_only = true, + SDVersion version = VERSION_WAN2, + std::shared_ptr weight_manager = nullptr) + : VAE(version, backend, "tae", weight_manager), + decode_only(decoder_only) { for (auto tensor_storage : tensor_storage_map) { if (tensor_storage.first.find(prefix + ".3.conv.6.weight") != std::string::npos) { is_wide = true; @@ -708,8 +713,8 @@ struct TinyVideoAutoEncoder : public VAE { return "taehv"; } - void get_param_tensors(std::map& tensors, const std::string prefix) { - taehv.get_param_tensors(tensors, prefix); + void get_param_tensors(std::map& tensors) override { + taehv.get_param_tensors(tensors, weight_prefix); } sd::Tensor vae_output_to_latents(const sd::Tensor& vae_output, std::shared_ptr rng) override { @@ -746,8 +751,8 @@ struct TinyVideoAutoEncoder : public VAE { return build_graph(z_tensor, decode_graph); }; - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), z_tensor.dim()); + return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false, false, false), z_tensor.dim()); } }; -#endif // __TAE_HPP__ +#endif // __SD_MODEL_VAE_TAE_HPP__ diff --git a/src/vae.hpp b/src/model/vae/vae.hpp similarity index 92% rename from src/vae.hpp rename to src/model/vae/vae.hpp index cc4cd967..8b8c46de 100644 --- a/src/vae.hpp +++ b/src/model/vae/vae.hpp @@ -1,12 +1,14 @@ -#ifndef __VAE_HPP__ -#define __VAE_HPP__ +#ifndef __SD_MODEL_VAE_VAE_HPP__ +#define __SD_MODEL_VAE_VAE_HPP__ -#include "common_block.hpp" -#include "tensor_ggml.hpp" +#include "core/tensor_ggml.hpp" +#include "model/common/block.hpp" +#include "model_manager.h" struct VAE : public GGMLRunner { protected: SDVersion version; + std::string weight_prefix; bool scale_input = true; virtual sd::Tensor _compute(const int n_threads, const sd::Tensor& z, @@ -62,8 +64,11 @@ protected: } public: - VAE(SDVersion version, ggml_backend_t backend, ggml_backend_t params_backend) - : version(version), GGMLRunner(backend, params_backend) {} + VAE(SDVersion version, + ggml_backend_t backend, + const std::string& weight_prefix = "", + std::shared_ptr weight_manager = nullptr) + : version(version), weight_prefix(weight_prefix), GGMLRunner(backend, weight_manager) {} int get_scale_factor() { int scale_factor = 8; @@ -73,7 +78,7 @@ public: scale_factor = 16; } else if (sd_version_uses_flux2_vae(version)) { scale_factor = 16; - } else if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1) { + } else if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1 || sd_version_is_minit2i(version)) { scale_factor = 1; } return scale_factor; @@ -214,7 +219,7 @@ public: virtual sd::Tensor vae_output_to_latents(const sd::Tensor& vae_output, std::shared_ptr rng) = 0; virtual sd::Tensor diffusion_to_vae_latents(const sd::Tensor& latents) = 0; virtual sd::Tensor vae_to_diffusion_latents(const sd::Tensor& latents) = 0; - virtual void get_param_tensors(std::map& tensors, const std::string prefix) = 0; + virtual void get_param_tensors(std::map& tensors) = 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) { @@ -223,8 +228,10 @@ public: }; struct FakeVAE : public VAE { - FakeVAE(SDVersion version, ggml_backend_t backend, ggml_backend_t params_backend) - : VAE(version, backend, params_backend) {} + FakeVAE(SDVersion version, + ggml_backend_t backend, + std::shared_ptr weight_manager = nullptr) + : VAE(version, backend, "", weight_manager) {} int get_encoder_output_channels(int input_channels) { return input_channels; @@ -251,11 +258,11 @@ struct FakeVAE : public VAE { return latents; } - void get_param_tensors(std::map& tensors, const std::string prefix) override {} + void get_param_tensors(std::map& tensors) override {} std::string get_desc() override { return "fake_vae"; } }; -#endif // __VAE_HPP__ +#endif // __SD_MODEL_VAE_VAE_HPP__ diff --git a/src/wan.hpp b/src/model/vae/wan_vae.hpp similarity index 52% rename from src/wan.hpp rename to src/model/vae/wan_vae.hpp index f7ac472e..a5fd0936 100644 --- a/src/wan.hpp +++ b/src/model/vae/wan_vae.hpp @@ -1,20 +1,17 @@ -#ifndef __WAN_HPP__ -#define __WAN_HPP__ +#ifndef __SD_MODEL_VAE_WAN_VAE_HPP__ +#define __SD_MODEL_VAE_WAN_VAE_HPP__ #include #include #include -#include "common_block.hpp" -#include "diffusion_model.hpp" -#include "flux.hpp" -#include "rope.hpp" -#include "vae.hpp" +#include "model/common/block.hpp" +#include "model/vae/vae.hpp" +#include "model_loader.h" namespace WAN { - constexpr int CACHE_T = 2; - constexpr int WAN_GRAPH_SIZE = 10240; + constexpr int CACHE_T = 2; class CausalConv3d : public GGMLBlock { protected: @@ -116,6 +113,24 @@ namespace WAN { } }; + class Conv2dBut3d : public Conv2d { + public: + using Conv2d::Conv2d; + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + ggml_tensor* x_swapped = ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2); + x_swapped = ggml_cont(ctx->ggml_ctx, x_swapped); + + ggml_tensor* out = Conv2d::forward(ctx, x_swapped); + + ggml_tensor* out_swapped = ggml_permute(ctx->ggml_ctx, out, 0, 1, 3, 2); + + out_swapped = ggml_cont(ctx->ggml_ctx, out_swapped); + + return out_swapped; + } + }; + class Resample : public GGMLBlock { protected: int64_t dim; @@ -341,19 +356,32 @@ namespace WAN { protected: int64_t in_dim; int64_t out_dim; + bool is_2D; public: - ResidualBlock(int64_t in_dim, int64_t out_dim) - : in_dim(in_dim), out_dim(out_dim) { + ResidualBlock(int64_t in_dim, int64_t out_dim, bool is_2D = false) + : in_dim(in_dim), out_dim(out_dim), is_2D(is_2D) { blocks["residual.0"] = std::shared_ptr(new RMS_norm(in_dim)); // residual.1 is nn.SiLU() - blocks["residual.2"] = std::shared_ptr(new CausalConv3d(in_dim, out_dim, {3, 3, 3}, {1, 1, 1}, {1, 1, 1})); + if (is_2D) { + blocks["residual.2"] = std::shared_ptr(new Conv2dBut3d(in_dim, out_dim, {3, 3}, {1, 1}, {1, 1})); + } else { + blocks["residual.2"] = std::shared_ptr(new CausalConv3d(in_dim, out_dim, {3, 3, 3}, {1, 1, 1}, {1, 1, 1})); + } blocks["residual.3"] = std::shared_ptr(new RMS_norm(out_dim)); // residual.4 is nn.SiLU() // residual.5 is nn.Dropout() - blocks["residual.6"] = std::shared_ptr(new CausalConv3d(out_dim, out_dim, {3, 3, 3}, {1, 1, 1}, {1, 1, 1})); + if (is_2D) { + blocks["residual.6"] = std::shared_ptr(new Conv2dBut3d(out_dim, out_dim, {3, 3}, {1, 1}, {1, 1})); + } else { + blocks["residual.6"] = std::shared_ptr(new CausalConv3d(out_dim, out_dim, {3, 3, 3}, {1, 1, 1}, {1, 1, 1})); + } if (in_dim != out_dim) { - blocks["shortcut"] = std::shared_ptr(new CausalConv3d(in_dim, out_dim, {1, 1, 1})); + if (is_2D) { + blocks["shortcut"] = std::shared_ptr(new Conv2dBut3d(in_dim, out_dim, {1, 1})); + } else { + blocks["shortcut"] = std::shared_ptr(new CausalConv3d(in_dim, out_dim, {1, 1, 1})); + } } } @@ -366,9 +394,15 @@ namespace WAN { GGML_ASSERT(b == 1); ggml_tensor* h = x; if (in_dim != out_dim) { - auto shortcut = std::dynamic_pointer_cast(blocks["shortcut"]); + if (is_2D) { + auto shortcut = std::dynamic_pointer_cast(blocks["shortcut"]); - h = shortcut->forward(ctx, x); + h = shortcut->forward(ctx, x); + } else { + auto shortcut = std::dynamic_pointer_cast(blocks["shortcut"]); + + h = shortcut->forward(ctx, x); + } } for (int i = 0; i < 7; i++) { @@ -388,8 +422,13 @@ namespace WAN { cache_x, 2); } + if (is_2D) { + auto layer = std::dynamic_pointer_cast(blocks["residual." + std::to_string(i)]); - x = layer->forward(ctx, x, feat_cache[idx]); + x = layer->forward(ctx, x); + } else { + x = layer->forward(ctx, x, feat_cache[idx]); + } feat_cache[idx] = cache_x; feat_idx += 1; } @@ -415,13 +454,14 @@ namespace WAN { int64_t out_dim, int mult, bool temperal_downsample = false, - bool down_flag = false) + bool down_flag = false, + bool is_2D = false) : mult(mult), down_flag(down_flag) { blocks["avg_shortcut"] = std::shared_ptr(new AvgDown3D(in_dim, out_dim, temperal_downsample ? 2 : 1, down_flag ? 2 : 1)); int i = 0; for (; i < mult; i++) { - blocks["downsamples." + std::to_string(i)] = std::shared_ptr(new ResidualBlock(in_dim, out_dim)); + blocks["downsamples." + std::to_string(i)] = std::shared_ptr(new ResidualBlock(in_dim, out_dim, is_2D)); in_dim = out_dim; } if (down_flag) { @@ -475,7 +515,8 @@ namespace WAN { int64_t out_dim, int mult, bool temperal_upsample = false, - bool up_flag = false) + bool up_flag = false, + bool is_2D = false) : mult(mult), up_flag(up_flag) { if (up_flag) { blocks["avg_shortcut"] = std::shared_ptr(new DupUp3D(in_dim, out_dim, temperal_upsample ? 2 : 1, up_flag ? 2 : 1)); @@ -483,7 +524,7 @@ namespace WAN { int i = 0; for (; i < mult; i++) { - blocks["upsamples." + std::to_string(i)] = std::shared_ptr(new ResidualBlock(in_dim, out_dim)); + blocks["upsamples." + std::to_string(i)] = std::shared_ptr(new ResidualBlock(in_dim, out_dim, is_2D)); in_dim = out_dim; } if (up_flag) { @@ -596,6 +637,7 @@ namespace WAN { std::vector dim_mult; int num_res_blocks; std::vector temperal_downsample; + bool is_2D = false; public: Encoder3d(int64_t dim = 128, @@ -604,21 +646,27 @@ namespace WAN { std::vector dim_mult = {1, 2, 4, 4}, int num_res_blocks = 2, std::vector temperal_downsample = {false, true, true}, - bool wan2_2 = false) + bool wan2_2 = false, + bool is_2D = false) : in_channels(in_channels), dim(dim), z_dim(z_dim), dim_mult(dim_mult), num_res_blocks(num_res_blocks), temperal_downsample(temperal_downsample), - wan2_2(wan2_2) { + wan2_2(wan2_2), + is_2D(is_2D) { // attn_scales is always [] std::vector dims = {dim}; for (int u : dim_mult) { dims.push_back(dim * u); } - blocks["conv1"] = std::shared_ptr(new CausalConv3d(in_channels, dims[0], {3, 3, 3}, {1, 1, 1}, {1, 1, 1})); + if (is_2D) { + blocks["conv1"] = std::shared_ptr(new Conv2dBut3d(in_channels, dims[0], {3, 3}, {1, 1}, {1, 1})); + } else { + blocks["conv1"] = std::shared_ptr(new CausalConv3d(in_channels, dims[0], {3, 3, 3}, {1, 1, 1}, {1, 1, 1})); + } int index = 0; int64_t in_dim; @@ -632,12 +680,13 @@ namespace WAN { out_dim, num_res_blocks, t_down_flag, - i != dim_mult.size() - 1)); + i != dim_mult.size() - 1, + is_2D)); blocks["downsamples." + std::to_string(index++)] = block; } else { for (int j = 0; j < num_res_blocks; j++) { - auto block = std::shared_ptr(new ResidualBlock(in_dim, out_dim)); + auto block = std::shared_ptr(new ResidualBlock(in_dim, out_dim, is_2D)); blocks["downsamples." + std::to_string(index++)] = block; in_dim = out_dim; } @@ -650,13 +699,17 @@ namespace WAN { } } - blocks["middle.0"] = std::shared_ptr(new ResidualBlock(out_dim, out_dim)); + blocks["middle.0"] = std::shared_ptr(new ResidualBlock(out_dim, out_dim, is_2D)); blocks["middle.1"] = std::shared_ptr(new AttentionBlock(out_dim)); - blocks["middle.2"] = std::shared_ptr(new ResidualBlock(out_dim, out_dim)); + blocks["middle.2"] = std::shared_ptr(new ResidualBlock(out_dim, out_dim, is_2D)); blocks["head.0"] = std::shared_ptr(new RMS_norm(out_dim)); // head.1 is nn.SiLU() - blocks["head.2"] = std::shared_ptr(new CausalConv3d(out_dim, z_dim, {3, 3, 3}, {1, 1, 1}, {1, 1, 1})); + if (is_2D) { + blocks["head.2"] = std::shared_ptr(new Conv2dBut3d(out_dim, z_dim, {3, 3}, {1, 1}, {1, 1})); + } else { + blocks["head.2"] = std::shared_ptr(new CausalConv3d(out_dim, z_dim, {3, 3, 3}, {1, 1, 1}, {1, 1, 1})); + } } ggml_tensor* forward(GGMLRunnerContext* ctx, @@ -675,7 +728,10 @@ namespace WAN { auto head_2 = std::dynamic_pointer_cast(blocks["head.2"]); // conv1 - if (feat_cache.size() > 0) { + if (is_2D) { + auto conv1 = std::dynamic_pointer_cast(blocks["conv1"]); + x = conv1->forward(ctx, x); + } else if (feat_cache.size() > 0) { int idx = feat_idx; auto cache_x = ggml_ext_slice(ctx->ggml_ctx, x, 2, -CACHE_T, x->ne[2]); if (cache_x->ne[2] < 2 && feat_cache[idx] != nullptr) { @@ -730,7 +786,10 @@ namespace WAN { // head x = head_0->forward(ctx, x); x = ggml_silu(ctx->ggml_ctx, x); - if (feat_cache.size() > 0) { + if (is_2D) { + auto head_2 = std::dynamic_pointer_cast(blocks["head.2"]); + x = head_2->forward(ctx, x); + } else if (feat_cache.size() > 0) { int idx = feat_idx; auto cache_x = ggml_ext_slice(ctx->ggml_ctx, x, 2, -CACHE_T, x->ne[2]); if (cache_x->ne[2] < 2 && feat_cache[idx] != nullptr) { @@ -761,6 +820,7 @@ namespace WAN { std::vector dim_mult; int num_res_blocks; std::vector temperal_upsample; + bool is_2D = false; public: Decoder3d(int64_t dim = 128, @@ -769,14 +829,16 @@ namespace WAN { std::vector dim_mult = {1, 2, 4, 4}, int num_res_blocks = 2, std::vector temperal_upsample = {true, true, false}, - bool wan2_2 = false) + bool wan2_2 = false, + bool is_2D = false) : out_channels(out_channels), dim(dim), z_dim(z_dim), dim_mult(dim_mult), num_res_blocks(num_res_blocks), temperal_upsample(temperal_upsample), - wan2_2(wan2_2) { + wan2_2(wan2_2), + is_2D(is_2D) { // attn_scales is always [] std::vector dims = {dim_mult[dim_mult.size() - 1] * dim}; for (int i = static_cast(dim_mult.size()) - 1; i >= 0; i--) { @@ -784,12 +846,16 @@ namespace WAN { } // init block - blocks["conv1"] = std::shared_ptr(new CausalConv3d(z_dim, dims[0], {3, 3, 3}, {1, 1, 1}, {1, 1, 1})); + if (is_2D) { + blocks["conv1"] = std::shared_ptr(new Conv2dBut3d(z_dim, dims[0], {3, 3}, {1, 1}, {1, 1})); + } else { + blocks["conv1"] = std::shared_ptr(new CausalConv3d(z_dim, dims[0], {3, 3, 3}, {1, 1, 1}, {1, 1, 1})); + } // middle blocks - blocks["middle.0"] = std::shared_ptr(new ResidualBlock(dims[0], dims[0])); + blocks["middle.0"] = std::shared_ptr(new ResidualBlock(dims[0], dims[0], is_2D)); blocks["middle.1"] = std::shared_ptr(new AttentionBlock(dims[0])); - blocks["middle.2"] = std::shared_ptr(new ResidualBlock(dims[0], dims[0])); + blocks["middle.2"] = std::shared_ptr(new ResidualBlock(dims[0], dims[0], is_2D)); // upsample blocks int index = 0; @@ -804,7 +870,8 @@ namespace WAN { out_dim, num_res_blocks + 1, t_up_flag, - i != dim_mult.size() - 1)); + i != dim_mult.size() - 1, + is_2D)); blocks["upsamples." + std::to_string(index++)] = block; } else { @@ -812,7 +879,7 @@ namespace WAN { in_dim = in_dim / 2; } for (int j = 0; j < num_res_blocks + 1; j++) { - auto block = std::shared_ptr(new ResidualBlock(in_dim, out_dim)); + auto block = std::shared_ptr(new ResidualBlock(in_dim, out_dim, is_2D)); blocks["upsamples." + std::to_string(index++)] = block; in_dim = out_dim; } @@ -826,9 +893,15 @@ namespace WAN { } // output blocks - blocks["head.0"] = std::shared_ptr(new RMS_norm(out_dim)); + blocks["head.0"] = std::shared_ptr(new RMS_norm(out_dim)); + int64_t final_dim = out_channels; // head.1 is nn.SiLU() - blocks["head.2"] = std::shared_ptr(new CausalConv3d(out_dim, out_channels, {3, 3, 3}, {1, 1, 1}, {1, 1, 1})); + if (is_2D) { + blocks["head.2"] = std::shared_ptr(new Conv2dBut3d(out_dim, final_dim, {3, 3}, {1, 1}, {1, 1})); + + } else { + blocks["head.2"] = std::shared_ptr(new CausalConv3d(out_dim, final_dim, {3, 3, 3}, {1, 1, 1}, {1, 1, 1})); + } } ggml_tensor* forward(GGMLRunnerContext* ctx, @@ -847,7 +920,10 @@ namespace WAN { auto head_2 = std::dynamic_pointer_cast(blocks["head.2"]); // conv1 - if (feat_cache.size() > 0) { + if (is_2D) { + auto conv1 = std::dynamic_pointer_cast(blocks["conv1"]); + x = conv1->forward(ctx, x); + } else if (feat_cache.size() > 0) { int idx = feat_idx; auto cache_x = ggml_ext_slice(ctx->ggml_ctx, x, 2, -CACHE_T, x->ne[2]); if (cache_x->ne[2] < 2 && feat_cache[idx] != nullptr) { @@ -902,7 +978,10 @@ namespace WAN { // head x = head_0->forward(ctx, x); x = ggml_silu(ctx->ggml_ctx, x); - if (feat_cache.size() > 0) { + if (is_2D) { + auto head_2 = std::dynamic_pointer_cast(blocks["head.2"]); + x = head_2->forward(ctx, x); + } else if (feat_cache.size() > 0) { int idx = feat_idx; auto cache_x = ggml_ext_slice(ctx->ggml_ctx, x, 2, -CACHE_T, x->ne[2]); if (cache_x->ne[2] < 2 && feat_cache[idx] != nullptr) { @@ -937,6 +1016,7 @@ namespace WAN { int num_res_blocks = 2; std::vector temperal_upsample = {true, true, false}; std::vector temperal_downsample = {false, true, true}; + bool is_2D = false; int _conv_num = 33; int _conv_idx = 0; @@ -953,9 +1033,10 @@ namespace WAN { } public: - WanVAE(bool decode_only = true, SDVersion version = VERSION_WAN2) + WanVAE(bool decode_only = true, SDVersion version = VERSION_WAN2, bool is_2D = false) : decode_only(decode_only), - wan2_2(version == VERSION_WAN2_2_TI2V) { + wan2_2(version == VERSION_WAN2_2_TI2V), + is_2D(is_2D) { // attn_scales is always [] if (wan2_2) { dim = 160; @@ -969,12 +1050,26 @@ namespace WAN { } else if (version == VERSION_QWEN_IMAGE_LAYERED) { input_channels = 4; } - if (!decode_only) { - blocks["encoder"] = std::shared_ptr(new Encoder3d(dim, z_dim * 2, input_channels, dim_mult, num_res_blocks, temperal_downsample, wan2_2)); - blocks["conv1"] = std::shared_ptr(new CausalConv3d(z_dim * 2, z_dim * 2, {1, 1, 1})); + + if (is_2D) { + temperal_upsample = {false, false, false}; + temperal_downsample = {false, false, false}; + } + + if (!decode_only) { + blocks["encoder"] = std::shared_ptr(new Encoder3d(dim, z_dim * 2, input_channels, dim_mult, num_res_blocks, temperal_downsample, wan2_2, is_2D)); + if (is_2D) { + blocks["conv1"] = std::shared_ptr(new Conv2dBut3d(z_dim * 2, z_dim * 2, {1, 1})); + } else { + blocks["conv1"] = std::shared_ptr(new CausalConv3d(z_dim * 2, z_dim * 2, {1, 1, 1})); + } + } + blocks["decoder"] = std::shared_ptr(new Decoder3d(dec_dim, z_dim, input_channels, dim_mult, num_res_blocks, temperal_upsample, wan2_2, is_2D)); + if (is_2D) { + blocks["conv2"] = std::shared_ptr(new Conv2dBut3d(z_dim, z_dim, {1, 1})); + } else { + blocks["conv2"] = std::shared_ptr(new CausalConv3d(z_dim, z_dim, {1, 1, 1})); } - blocks["decoder"] = std::shared_ptr(new Decoder3d(dec_dim, z_dim, input_channels, dim_mult, num_res_blocks, temperal_upsample, wan2_2)); - blocks["conv2"] = std::shared_ptr(new CausalConv3d(z_dim, z_dim, {1, 1, 1})); } static ggml_tensor* patchify(ggml_context* ctx, @@ -1037,6 +1132,10 @@ namespace WAN { GGML_ASSERT(b == 1); GGML_ASSERT(decode_only == false); + if (x->ne[2] > 1 && is_2D) { + LOG_WARN("Using 2D VAE to encode video, expect poor results"); + } + clear_cache(); x = patchify(ctx->ggml_ctx, x, patch_size, b); @@ -1054,12 +1153,18 @@ namespace WAN { auto in = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1); // [b*c, 1, h, w] out = encoder->forward(ctx, in, b, _enc_feat_map, _enc_conv_idx, i); } else { - auto in = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1 + 4 * (i - 1), 1 + 4 * i); // [b*c, 4, h, w] + // if is_2D, drop 3 out of 4 frames + auto in = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1 + 4 * (i - 1), (is_2D ? 1 - 3 : 1) + 4 * i); // [b*c, 4, h, w] auto out_ = encoder->forward(ctx, in, b, _enc_feat_map, _enc_conv_idx, i); out = ggml_concat(ctx->ggml_ctx, out, out_, 2); } } - out = conv1->forward(ctx, out); + if (is_2D) { + auto conv1 = std::dynamic_pointer_cast(blocks["conv1"]); + out = conv1->forward(ctx, out); + } else { + out = conv1->forward(ctx, out); + } auto mu = ggml_ext_chunk(ctx->ggml_ctx, out, 2, 3)[0]; // sd::ggml_graph_cut::mark_graph_cut(mu, "wan_vae.encode.final", "mu"); clear_cache(); @@ -1072,13 +1177,23 @@ namespace WAN { // z: [b*c, t, h, w] GGML_ASSERT(b == 1); + if (z->ne[2] > 1 && is_2D) { + LOG_WARN("Using 2D VAE to decode video, expect poor results"); + } + clear_cache(); auto decoder = std::dynamic_pointer_cast(blocks["decoder"]); auto conv2 = std::dynamic_pointer_cast(blocks["conv2"]); int64_t iter_ = z->ne[2]; - auto x = conv2->forward(ctx, z); + auto x = z; + if (is_2D) { + auto conv2 = std::dynamic_pointer_cast(blocks["conv2"]); + x = conv2->forward(ctx, z); + } else { + x = conv2->forward(ctx, z); + } // sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.decode.prelude", "x"); ggml_tensor* out; for (int i = 0; i < iter_; i++) { @@ -1090,6 +1205,12 @@ namespace WAN { auto in = ggml_ext_slice(ctx->ggml_ctx, x, 2, i, i + 1); // [b*c, 1, h, w] auto out_ = decoder->forward(ctx, in, b, _feat_map, _conv_idx, i); out = ggml_concat(ctx->ggml_ctx, out, out_, 2); + if (is_2D) { + // repeat frames to avoid mismatch + for (int j = 0; j < 4 - 1; j++) { + out = ggml_concat(ctx->ggml_ctx, out, out_, 2); + } + } } } out = unpatchify(ctx->ggml_ctx, out, patch_size, b); @@ -1125,12 +1246,25 @@ namespace WAN { WanVAE ae; WanVAERunner(ggml_backend_t backend, - ggml_backend_t params_backend, - const String2TensorStorage& tensor_storage_map = {}, - const std::string prefix = "", - bool decode_only = false, - SDVersion version = VERSION_WAN2) - : decode_only(decode_only), ae(decode_only, version), VAE(version, backend, params_backend) { + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "", + bool decode_only = false, + SDVersion version = VERSION_WAN2, + std::shared_ptr weight_manager = nullptr) + : VAE(version, backend, prefix, weight_manager), decode_only(decode_only) { + bool is_2D = false; + for (const auto& [name, tensor_storage] : tensor_storage_map) { + if (ends_with(name, "decoder.conv1.weight")) { + if (tensor_storage.ne[2] > 3) { + is_2D = true; + } + break; + } + } + if (is_2D) { + LOG_DEBUG("USING 2D VAE"); + } + ae = WanVAE(decode_only, version, is_2D); ae.init(params_ctx, tensor_storage_map, prefix); } @@ -1138,8 +1272,8 @@ namespace WAN { return "wan_vae"; } - void get_param_tensors(std::map& tensors, const std::string prefix) override { - ae.get_param_tensors(tensors, prefix); + void get_param_tensors(std::map& tensors) override { + ae.get_param_tensors(tensors, weight_prefix); } sd::Tensor vae_output_to_latents(const sd::Tensor& vae_output, std::shared_ptr rng) override { @@ -1256,7 +1390,7 @@ namespace WAN { return build_graph(input, decode_graph); } }; - auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true), + auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, true, true, true), input.empty() ? z.dim() : input.dim()); if (!result.empty() && z.dim() == 4) { result.squeeze_(2); @@ -1269,7 +1403,7 @@ namespace WAN { auto get_graph = [&]() -> ggml_cgraph* { return build_graph_partial(z, decode_graph, i); }; - auto out_opt = GGMLRunner::compute(get_graph, n_threads, true); + auto out_opt = GGMLRunner::compute(get_graph, n_threads, true, true, true); if (!out_opt.has_value()) { return {}; } @@ -1282,7 +1416,7 @@ namespace WAN { sd::Tensor output = std::move(out); for (i = 1; i < t; i++) { - auto chunk_opt = GGMLRunner::compute(get_graph, n_threads, true); + auto chunk_opt = GGMLRunner::compute(get_graph, n_threads, true, true, true); if (!chunk_opt.has_value()) { return {}; } @@ -1328,27 +1462,24 @@ namespace WAN { // ggml_backend_t backend = ggml_backend_cuda_init(0); ggml_backend_t backend = sd_backend_cpu_init(); ggml_type model_data_type = GGML_TYPE_F16; - std::shared_ptr vae = std::make_shared(backend, backend, String2TensorStorage{}, "", false, VERSION_WAN2_2_TI2V); + auto model_manager = std::make_shared(); + std::shared_ptr vae = std::make_shared(backend, String2TensorStorage{}, "first_stage_model", false, VERSION_WAN2_2_TI2V, model_manager); { LOG_INFO("loading from '%s'", file_path.c_str()); - if (!vae->alloc_params_buffer()) { - LOG_ERROR("vae buffer allocation failed"); - return; - } - std::map tensors; - vae->get_param_tensors(tensors, "first_stage_model"); - - ModelLoader model_loader; + ModelLoader& model_loader = model_manager->loader(); if (!model_loader.init_from_file_and_convert_name(file_path, "vae.")) { LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str()); return; } - bool success = model_loader.load_tensors(tensors); - - if (!success) { - LOG_ERROR("load tensors from model loader failed"); + if (!model_manager->register_runner_params("Wan VAE test", + *vae, + ModelManager::ResidencyMode::ParamBackend, + backend, + backend) || + !model_manager->validate_registered_tensors()) { + LOG_ERROR("register wan vae tensors with model manager failed"); return; } @@ -1358,1038 +1489,6 @@ namespace WAN { } }; - class WanSelfAttention : public GGMLBlock { - public: - int64_t num_heads; - int64_t head_dim; - - public: - WanSelfAttention(int64_t dim, - int64_t num_heads, - bool qk_norm = true, - float eps = 1e-6) - : num_heads(num_heads) { - head_dim = dim / num_heads; - blocks["q"] = std::shared_ptr(new Linear(dim, dim)); - blocks["k"] = std::shared_ptr(new Linear(dim, dim)); - blocks["v"] = std::shared_ptr(new Linear(dim, dim)); - blocks["o"] = std::shared_ptr(new Linear(dim, dim)); - - if (qk_norm) { - blocks["norm_q"] = std::shared_ptr(new RMSNorm(dim, eps)); - blocks["norm_k"] = std::shared_ptr(new RMSNorm(dim, eps)); - } else { - blocks["norm_q"] = std::shared_ptr(new Identity()); - blocks["norm_k"] = std::shared_ptr(new Identity()); - } - } - - virtual ggml_tensor* forward(GGMLRunnerContext* ctx, - ggml_tensor* x, - ggml_tensor* pe, - ggml_tensor* mask = nullptr) { - // x: [N, n_token, dim] - // pe: [n_token, d_head/2, 2, 2] - // return [N, n_token, dim] - int64_t N = x->ne[2]; - int64_t n_token = x->ne[1]; - - auto q_proj = std::dynamic_pointer_cast(blocks["q"]); - auto k_proj = std::dynamic_pointer_cast(blocks["k"]); - auto v_proj = std::dynamic_pointer_cast(blocks["v"]); - auto o_proj = std::dynamic_pointer_cast(blocks["o"]); - auto norm_q = std::dynamic_pointer_cast(blocks["norm_q"]); - auto norm_k = std::dynamic_pointer_cast(blocks["norm_k"]); - - auto q = q_proj->forward(ctx, x); - q = norm_q->forward(ctx, q); - auto k = k_proj->forward(ctx, x); - k = norm_k->forward(ctx, k); - auto v = v_proj->forward(ctx, x); // [N, n_token, n_head*d_head] - - q = ggml_reshape_4d(ctx->ggml_ctx, q, head_dim, num_heads, n_token, N); // [N, n_token, n_head, d_head] - k = ggml_reshape_4d(ctx->ggml_ctx, k, head_dim, num_heads, n_token, N); // [N, n_token, n_head, d_head] - v = ggml_reshape_4d(ctx->ggml_ctx, v, head_dim, num_heads, n_token, N); // [N, n_token, n_head, d_head] - - x = Rope::attention(ctx, q, k, v, pe, mask); // [N, n_token, dim] - - x = o_proj->forward(ctx, x); // [N, n_token, dim] - return x; - } - }; - - class WanCrossAttention : public WanSelfAttention { - public: - WanCrossAttention(int64_t dim, - int64_t num_heads, - bool qk_norm = true, - float eps = 1e-6) - : WanSelfAttention(dim, num_heads, qk_norm, eps) {} - virtual ggml_tensor* forward(GGMLRunnerContext* ctx, - ggml_tensor* x, - ggml_tensor* context, - int64_t context_img_len) = 0; - }; - - class WanT2VCrossAttention : public WanCrossAttention { - public: - WanT2VCrossAttention(int64_t dim, - int64_t num_heads, - bool qk_norm = true, - float eps = 1e-6) - : WanCrossAttention(dim, num_heads, qk_norm, eps) {} - ggml_tensor* forward(GGMLRunnerContext* ctx, - ggml_tensor* x, - ggml_tensor* context, - int64_t context_img_len) override { - // x: [N, n_token, dim] - // context: [N, n_context, dim] - // context_img_len: unused - // return [N, n_token, dim] - int64_t N = x->ne[2]; - int64_t n_token = x->ne[1]; - - auto q_proj = std::dynamic_pointer_cast(blocks["q"]); - auto k_proj = std::dynamic_pointer_cast(blocks["k"]); - auto v_proj = std::dynamic_pointer_cast(blocks["v"]); - auto o_proj = std::dynamic_pointer_cast(blocks["o"]); - auto norm_q = std::dynamic_pointer_cast(blocks["norm_q"]); - auto norm_k = std::dynamic_pointer_cast(blocks["norm_k"]); - - auto q = q_proj->forward(ctx, x); - q = norm_q->forward(ctx, q); - auto k = k_proj->forward(ctx, context); // [N, n_context, dim] - k = norm_k->forward(ctx, k); - auto v = v_proj->forward(ctx, context); // [N, n_context, dim] - - x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim] - - x = o_proj->forward(ctx, x); // [N, n_token, dim] - return x; - } - }; - - class WanI2VCrossAttention : public WanCrossAttention { - public: - WanI2VCrossAttention(int64_t dim, - int64_t num_heads, - bool qk_norm = true, - float eps = 1e-6) - : WanCrossAttention(dim, num_heads, qk_norm, eps) { - blocks["k_img"] = std::shared_ptr(new Linear(dim, dim)); - blocks["v_img"] = std::shared_ptr(new Linear(dim, dim)); - - if (qk_norm) { - blocks["norm_k_img"] = std::shared_ptr(new RMSNorm(dim, eps)); - } else { - blocks["norm_k_img"] = std::shared_ptr(new Identity()); - } - } - - ggml_tensor* forward(GGMLRunnerContext* ctx, - ggml_tensor* x, - ggml_tensor* context, - int64_t context_img_len) override { - // x: [N, n_token, dim] - // context: [N, context_img_len + context_txt_len, dim] - // return [N, n_token, dim] - - auto q_proj = std::dynamic_pointer_cast(blocks["q"]); - auto k_proj = std::dynamic_pointer_cast(blocks["k"]); - auto v_proj = std::dynamic_pointer_cast(blocks["v"]); - auto o_proj = std::dynamic_pointer_cast(blocks["o"]); - - auto k_img_proj = std::dynamic_pointer_cast(blocks["k_img"]); - auto v_img_proj = std::dynamic_pointer_cast(blocks["v_img"]); - - auto norm_q = std::dynamic_pointer_cast(blocks["norm_q"]); - auto norm_k = std::dynamic_pointer_cast(blocks["norm_k"]); - auto norm_k_img = std::dynamic_pointer_cast(blocks["norm_k_img"]); - - int64_t N = x->ne[2]; - int64_t n_token = x->ne[1]; - int64_t dim = x->ne[0]; - int64_t context_txt_len = context->ne[1] - context_img_len; - - auto context_img = ggml_view_3d(ctx->ggml_ctx, context, dim, context_img_len, N, context->nb[1], context->nb[2], 0); // [N, context_img_len, dim] - auto context_txt = ggml_view_3d(ctx->ggml_ctx, context, dim, context_txt_len, N, context->nb[1], context->nb[2], context_img_len * context->nb[1]); // [N, context_txt_len, dim] - - auto q = q_proj->forward(ctx, x); - q = norm_q->forward(ctx, q); - auto k = k_proj->forward(ctx, context_txt); // [N, context_txt_len, dim] - k = norm_k->forward(ctx, k); - auto v = v_proj->forward(ctx, context_txt); // [N, context_txt_len, dim] - - auto k_img = k_img_proj->forward(ctx, context_img); // [N, context_img_len, dim] - k_img = norm_k_img->forward(ctx, k_img); - auto v_img = v_img_proj->forward(ctx, context_img); // [N, context_img_len, dim] - - auto img_x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k_img, v_img, num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim] - x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, nullptr, false, ctx->flash_attn_enabled); // [N, n_token, dim] - - x = ggml_add(ctx->ggml_ctx, x, img_x); - - x = o_proj->forward(ctx, x); // [N, n_token, dim] - return x; - } - }; - - static ggml_tensor* modulate_add(ggml_context* ctx, ggml_tensor* x, ggml_tensor* e) { - // x: [N, n_token, dim] - // e: [N, 1, dim] or [N, T, 1, dim] - if (ggml_n_dims(e) == 3) { - int64_t T = e->ne[2]; - x = ggml_reshape_4d(ctx, x, x->ne[0], x->ne[1] / T, T, x->ne[2]); // [N, T, n_token/T, dim] - x = ggml_add(ctx, x, e); - x = ggml_reshape_3d(ctx, x, x->ne[0], x->ne[1] * x->ne[2], x->ne[3]); // [N, n_token, dim] - } else { - x = ggml_add(ctx, x, e); - } - return x; - } - - static ggml_tensor* modulate_mul(ggml_context* ctx, ggml_tensor* x, ggml_tensor* e) { - // x: [N, n_token, dim] - // e: [N, 1, dim] or [N, T, 1, dim] - if (ggml_n_dims(e) == 3) { - int64_t T = e->ne[2]; - x = ggml_reshape_4d(ctx, x, x->ne[0], x->ne[1] / T, T, x->ne[2]); // [N, T, n_token/T, dim] - x = ggml_mul(ctx, x, e); - x = ggml_reshape_3d(ctx, x, x->ne[0], x->ne[1] * x->ne[2], x->ne[3]); // [N, n_token, dim] - } else { - x = ggml_mul(ctx, x, e); - } - return x; - } - - class WanAttentionBlock : public GGMLBlock { - protected: - int64_t dim; - - void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { - enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32); - params["modulation"] = ggml_new_tensor_3d(ctx, wtype, dim, 6, 1); - } - - public: - WanAttentionBlock(bool t2v_cross_attn, - int64_t dim, - int64_t ffn_dim, - int64_t num_heads, - bool qk_norm = true, - bool cross_attn_norm = false, - float eps = 1e-6) - : dim(dim) { - blocks["norm1"] = std::shared_ptr(new LayerNorm(dim, eps, false)); - blocks["self_attn"] = std::shared_ptr(new WanSelfAttention(dim, num_heads, qk_norm, eps)); - if (cross_attn_norm) { - blocks["norm3"] = std::shared_ptr(new LayerNorm(dim, eps, true)); - } else { - blocks["norm3"] = std::shared_ptr(new Identity()); - } - if (t2v_cross_attn) { - blocks["cross_attn"] = std::shared_ptr(new WanT2VCrossAttention(dim, num_heads, qk_norm, eps)); - } else { - blocks["cross_attn"] = std::shared_ptr(new WanI2VCrossAttention(dim, num_heads, qk_norm, eps)); - } - - blocks["norm2"] = std::shared_ptr(new LayerNorm(dim, eps, false)); - - blocks["ffn.0"] = std::shared_ptr(new Linear(dim, ffn_dim)); - // ffn.1 is nn.GELU(approximate='tanh') - blocks["ffn.2"] = std::shared_ptr(new Linear(ffn_dim, dim)); - } - - virtual ggml_tensor* forward(GGMLRunnerContext* ctx, - ggml_tensor* x, - ggml_tensor* e, - ggml_tensor* pe, - ggml_tensor* context, - int64_t context_img_len = 257) { - // x: [N, n_token, dim] - // e: [N, 6, dim] or [N, T, 6, dim] - // context: [N, context_img_len + context_txt_len, dim] - // return [N, n_token, dim] - - auto modulation = params["modulation"]; - e = ggml_add(ctx->ggml_ctx, e, modulation); // [N, 6, dim] or [N, T, 6, dim] - auto es = ggml_ext_chunk(ctx->ggml_ctx, e, 6, 1); // ([N, 1, dim], ...) or [N, T, 1, dim] - - auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); - auto self_attn = std::dynamic_pointer_cast(blocks["self_attn"]); - auto norm3 = std::dynamic_pointer_cast(blocks["norm3"]); - auto cross_attn = std::dynamic_pointer_cast(blocks["cross_attn"]); - auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); - auto ffn_0 = std::dynamic_pointer_cast(blocks["ffn.0"]); - auto ffn_2 = std::dynamic_pointer_cast(blocks["ffn.2"]); - - // self-attention - auto y = norm1->forward(ctx, x); - y = ggml_add(ctx->ggml_ctx, y, modulate_mul(ctx->ggml_ctx, y, es[1])); - y = modulate_add(ctx->ggml_ctx, y, es[0]); - y = self_attn->forward(ctx, y, pe); - - x = ggml_add(ctx->ggml_ctx, x, modulate_mul(ctx->ggml_ctx, y, es[2])); - - // cross-attention - x = ggml_add(ctx->ggml_ctx, - x, - cross_attn->forward(ctx, norm3->forward(ctx, x), context, context_img_len)); - - // ffn - y = norm2->forward(ctx, x); - y = ggml_add(ctx->ggml_ctx, y, modulate_mul(ctx->ggml_ctx, y, es[4])); - y = modulate_add(ctx->ggml_ctx, y, es[3]); - - y = ffn_0->forward(ctx, y); - y = ggml_ext_gelu(ctx->ggml_ctx, y, true); - y = ffn_2->forward(ctx, y); - - x = ggml_add(ctx->ggml_ctx, x, modulate_mul(ctx->ggml_ctx, y, es[5])); - - return x; - } - }; - - class VaceWanAttentionBlock : public WanAttentionBlock { - protected: - int block_id; - void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { - enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32); - params["modulation"] = ggml_new_tensor_3d(ctx, wtype, dim, 6, 1); - } - - public: - VaceWanAttentionBlock(bool t2v_cross_attn, - int64_t dim, - int64_t ffn_dim, - int64_t num_heads, - bool qk_norm = true, - bool cross_attn_norm = false, - float eps = 1e-6, - int block_id = 0) - : WanAttentionBlock(t2v_cross_attn, dim, ffn_dim, num_heads, qk_norm, cross_attn_norm, eps), block_id(block_id) { - if (block_id == 0) { - blocks["before_proj"] = std::shared_ptr(new Linear(dim, dim)); - } - blocks["after_proj"] = std::shared_ptr(new Linear(dim, dim)); - } - - std::pair forward(GGMLRunnerContext* ctx, - ggml_tensor* c, - ggml_tensor* x, - ggml_tensor* e, - ggml_tensor* pe, - ggml_tensor* context, - int64_t context_img_len = 257) { - // x: [N, n_token, dim] - // e: [N, 6, dim] or [N, T, 6, dim] - // context: [N, context_img_len + context_txt_len, dim] - // return [N, n_token, dim] - if (block_id == 0) { - auto before_proj = std::dynamic_pointer_cast(blocks["before_proj"]); - - c = before_proj->forward(ctx, c); - c = ggml_add(ctx->ggml_ctx, c, x); - } - - auto after_proj = std::dynamic_pointer_cast(blocks["after_proj"]); - - c = WanAttentionBlock::forward(ctx, c, e, pe, context, context_img_len); - auto c_skip = after_proj->forward(ctx, c); - - return {c_skip, c}; - } - }; - - class Head : public GGMLBlock { - protected: - int64_t dim; - - void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { - enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32); - params["modulation"] = ggml_new_tensor_3d(ctx, wtype, dim, 2, 1); - } - - public: - Head(int64_t dim, - int64_t out_dim, - std::tuple patch_size, - float eps = 1e-6) - : dim(dim) { - out_dim = out_dim * std::get<0>(patch_size) * std::get<1>(patch_size) * std::get<2>(patch_size); - - blocks["norm"] = std::shared_ptr(new LayerNorm(dim, eps, false)); - blocks["head"] = std::shared_ptr(new Linear(dim, out_dim)); - } - - ggml_tensor* forward(GGMLRunnerContext* ctx, - ggml_tensor* x, - ggml_tensor* e) { - // x: [N, n_token, dim] - // e: [N, dim] or [N, T, dim] - // return [N, n_token, out_dim] - - auto modulation = params["modulation"]; - e = ggml_reshape_4d(ctx->ggml_ctx, e, e->ne[0], 1, e->ne[1], e->ne[2]); // [N, 1, dim] or [N, T, 1, dim] - e = ggml_repeat_4d(ctx->ggml_ctx, e, e->ne[0], 2, e->ne[2], e->ne[3]); // [N, 2, dim] or [N, T, 2, dim] - - e = ggml_add(ctx->ggml_ctx, e, modulation); // [N, 2, dim] or [N, T, 2, dim] - auto es = ggml_ext_chunk(ctx->ggml_ctx, e, 2, 1); // ([N, 1, dim], ...) or ([N, T, 1, dim], ...) - - auto norm = std::dynamic_pointer_cast(blocks["norm"]); - auto head = std::dynamic_pointer_cast(blocks["head"]); - - x = norm->forward(ctx, x); - x = ggml_add(ctx->ggml_ctx, x, modulate_mul(ctx->ggml_ctx, x, es[1])); - x = modulate_add(ctx->ggml_ctx, x, es[0]); - x = head->forward(ctx, x); - return x; - } - }; - - class MLPProj : public GGMLBlock { - protected: - int64_t in_dim; - int64_t flf_pos_embed_token_number; - - void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { - if (flf_pos_embed_token_number > 0) { - params["emb_pos"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, in_dim, flf_pos_embed_token_number, 1); - } - } - - public: - MLPProj(int64_t in_dim, - int64_t out_dim, - int64_t flf_pos_embed_token_number = 0) - : in_dim(in_dim), flf_pos_embed_token_number(flf_pos_embed_token_number) { - blocks["proj.0"] = std::shared_ptr(new LayerNorm(in_dim)); - blocks["proj.1"] = std::shared_ptr(new Linear(in_dim, in_dim)); - // proj.2 is nn.GELU() - blocks["proj.3"] = std::shared_ptr(new Linear(in_dim, out_dim)); - blocks["proj.4"] = std::shared_ptr(new LayerNorm(out_dim)); - } - - ggml_tensor* forward(GGMLRunnerContext* ctx, - ggml_tensor* image_embeds) { - if (flf_pos_embed_token_number > 0) { - auto emb_pos = params["emb_pos"]; - - auto a = ggml_ext_slice(ctx->ggml_ctx, image_embeds, 1, 0, emb_pos->ne[1]); - auto b = ggml_ext_slice(ctx->ggml_ctx, emb_pos, 1, 0, image_embeds->ne[1]); - - image_embeds = ggml_add(ctx->ggml_ctx, a, b); - } - - auto proj_0 = std::dynamic_pointer_cast(blocks["proj.0"]); - auto proj_1 = std::dynamic_pointer_cast(blocks["proj.1"]); - auto proj_3 = std::dynamic_pointer_cast(blocks["proj.3"]); - auto proj_4 = std::dynamic_pointer_cast(blocks["proj.4"]); - - auto x = proj_0->forward(ctx, image_embeds); - x = proj_1->forward(ctx, x); - x = ggml_ext_gelu(ctx->ggml_ctx, x, true); - x = proj_3->forward(ctx, x); - x = proj_4->forward(ctx, x); - - return x; // clip_extra_context_tokens - } - }; - - struct WanParams { - std::string model_type = "t2v"; - std::tuple patch_size = {1, 2, 2}; - int64_t text_len = 512; - int64_t in_dim = 16; - int64_t dim = 2048; - int64_t ffn_dim = 8192; - int freq_dim = 256; - int64_t text_dim = 4096; - int64_t out_dim = 16; - int64_t num_heads = 16; - int num_layers = 32; - int vace_layers = 0; - int64_t vace_in_dim = 96; - std::map vace_layers_mapping = {}; - bool qk_norm = true; - bool cross_attn_norm = true; - float eps = 1e-6f; - int64_t flf_pos_embed_token_number = 0; - int theta = 10000; - // wan2.1 1.3B: 1536/12, wan2.1/2.2 14B: 5120/40, wan2.2 5B: 3074/24 - std::vector axes_dim = {44, 42, 42}; - int64_t axes_dim_sum = 128; - }; - - class Wan : public GGMLBlock { - protected: - WanParams params; - - public: - Wan() {} - Wan(WanParams params) - : params(params) { - // patch_embedding - blocks["patch_embedding"] = std::shared_ptr(new Conv3d(params.in_dim, params.dim, params.patch_size, params.patch_size)); - - // text_embedding - blocks["text_embedding.0"] = std::shared_ptr(new Linear(params.text_dim, params.dim)); - // text_embedding.1 is nn.GELU() - blocks["text_embedding.2"] = std::shared_ptr(new Linear(params.dim, params.dim)); - - // time_embedding - blocks["time_embedding.0"] = std::shared_ptr(new Linear(params.freq_dim, params.dim)); - // time_embedding.1 is nn.SiLU() - blocks["time_embedding.2"] = std::shared_ptr(new Linear(params.dim, params.dim)); - - // time_projection.0 is nn.SiLU() - blocks["time_projection.1"] = std::shared_ptr(new Linear(params.dim, params.dim * 6)); - - // blocks - for (int i = 0; i < params.num_layers; i++) { - auto block = std::shared_ptr(new WanAttentionBlock(params.model_type == "t2v", - params.dim, - params.ffn_dim, - params.num_heads, - params.qk_norm, - params.cross_attn_norm, - params.eps)); - blocks["blocks." + std::to_string(i)] = block; - } - - // head - blocks["head"] = std::shared_ptr(new Head(params.dim, params.out_dim, params.patch_size, params.eps)); - - // img_emb - if (params.model_type == "i2v") { - blocks["img_emb"] = std::shared_ptr(new MLPProj(1280, params.dim, params.flf_pos_embed_token_number)); - } - - // vace - if (params.vace_layers > 0) { - for (int i = 0; i < params.vace_layers; i++) { - auto block = std::shared_ptr(new VaceWanAttentionBlock(params.model_type == "t2v", - params.dim, - params.ffn_dim, - params.num_heads, - params.qk_norm, - params.cross_attn_norm, - params.eps, - i)); - blocks["vace_blocks." + std::to_string(i)] = block; - } - - int step = params.num_layers / params.vace_layers; - int n = 0; - for (int i = 0; i < params.num_layers; i += step) { - this->params.vace_layers_mapping[i] = n; - n++; - } - - blocks["vace_patch_embedding"] = std::shared_ptr(new Conv3d(params.vace_in_dim, params.dim, params.patch_size, params.patch_size)); - } - } - - ggml_tensor* pad_to_patch_size(GGMLRunnerContext* ctx, - ggml_tensor* x) { - int64_t W = x->ne[0]; - int64_t H = x->ne[1]; - int64_t T = x->ne[2]; - - int pad_t = (std::get<0>(params.patch_size) - T % std::get<0>(params.patch_size)) % std::get<0>(params.patch_size); - int pad_h = (std::get<1>(params.patch_size) - H % std::get<1>(params.patch_size)) % std::get<1>(params.patch_size); - int pad_w = (std::get<2>(params.patch_size) - W % std::get<2>(params.patch_size)) % std::get<2>(params.patch_size); - ggml_ext_pad(ctx->ggml_ctx, x, pad_w, pad_h, pad_t, 0, ctx->circular_x_enabled, ctx->circular_y_enabled); - return x; - } - - ggml_tensor* unpatchify(ggml_context* ctx, - ggml_tensor* x, - int64_t t_len, - int64_t h_len, - int64_t w_len) { - // x: [N, t_len*h_len*w_len, pt*ph*pw*C] - // return: [N*C, t_len*pt, h_len*ph, w_len*pw] - int64_t N = x->ne[3]; - int64_t pt = std::get<0>(params.patch_size); - int64_t ph = std::get<1>(params.patch_size); - int64_t pw = std::get<2>(params.patch_size); - int64_t C = x->ne[0] / pt / ph / pw; - - GGML_ASSERT(C * pt * ph * pw == x->ne[0]); - - x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N); // [N, t_len*h_len*w_len, pt*ph*pw, C] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3)); // [N, C, t_len*h_len*w_len, pt*ph*pw] - x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N); // [N*C*t_len*h_len, w_len, pt*ph, pw] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, pt*ph, w_len, pw] - x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N); // [N*C*t_len*h_len, pt, ph, w_len*pw] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, ph, pt, w_len*pw] - x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph * h_len, t_len * C * N); // [N*C*t_len, h_len*ph, pt, w_len*pw] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len, pt, h_len*ph, w_len*pw] - x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt * t_len, C * N); // [N*C, t_len*pt, h_len*ph, w_len*pw] - return x; - } - - ggml_tensor* forward_orig(GGMLRunnerContext* ctx, - ggml_tensor* x, - ggml_tensor* timestep, - ggml_tensor* context, - ggml_tensor* pe, - ggml_tensor* clip_fea = nullptr, - ggml_tensor* vace_context = nullptr, - float vace_strength = 1.f, - int64_t N = 1) { - // x: [N*C, T, H, W], C => in_dim - // vace_context: [N*vace_in_dim, T, H, W] - // timestep: [N,] or [T] - // context: [N, L, text_dim] - // return: [N, t_len*h_len*w_len, out_dim*pt*ph*pw] - - GGML_ASSERT(N == 1); - - auto patch_embedding = std::dynamic_pointer_cast(blocks["patch_embedding"]); - - auto text_embedding_0 = std::dynamic_pointer_cast(blocks["text_embedding.0"]); - auto text_embedding_2 = std::dynamic_pointer_cast(blocks["text_embedding.2"]); - - auto time_embedding_0 = std::dynamic_pointer_cast(blocks["time_embedding.0"]); - auto time_embedding_2 = std::dynamic_pointer_cast(blocks["time_embedding.2"]); - auto time_projection_1 = std::dynamic_pointer_cast(blocks["time_projection.1"]); - - auto head = std::dynamic_pointer_cast(blocks["head"]); - - // patch_embedding - x = patch_embedding->forward(ctx, x); // [N*dim, t_len, h_len, w_len] - x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1] * x->ne[2], x->ne[3] / N, N); // [N, dim, t_len*h_len*w_len] - x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); // [N, t_len*h_len*w_len, dim] - - // time_embedding - auto e = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, params.freq_dim); - e = time_embedding_0->forward(ctx, e); - e = ggml_silu_inplace(ctx->ggml_ctx, e); - e = time_embedding_2->forward(ctx, e); // [N, dim] or [N, T, dim] - - // time_projection - auto e0 = ggml_silu(ctx->ggml_ctx, e); - e0 = time_projection_1->forward(ctx, e0); - e0 = ggml_reshape_4d(ctx->ggml_ctx, e0, e0->ne[0] / 6, 6, e0->ne[1], e0->ne[2]); // [N, 6, dim] or [N, T, 6, dim] - - context = text_embedding_0->forward(ctx, context); - context = ggml_ext_gelu(ctx->ggml_ctx, context); - context = text_embedding_2->forward(ctx, context); // [N, context_txt_len, dim] - - int64_t context_img_len = 0; - if (clip_fea != nullptr) { - if (params.model_type == "i2v") { - auto img_emb = std::dynamic_pointer_cast(blocks["img_emb"]); - auto context_img = img_emb->forward(ctx, clip_fea); // [N, context_img_len, dim] - context = ggml_concat(ctx->ggml_ctx, context_img, context, 1); // [N, context_img_len + context_txt_len, dim] - } - context_img_len = clip_fea->ne[1]; // 257 - } - - // vace_patch_embedding - ggml_tensor* c = nullptr; - if (params.vace_layers > 0) { - auto vace_patch_embedding = std::dynamic_pointer_cast(blocks["vace_patch_embedding"]); - - c = vace_patch_embedding->forward(ctx, vace_context); // [N*dim, t_len, h_len, w_len] - c = ggml_reshape_3d(ctx->ggml_ctx, c, c->ne[0] * c->ne[1] * c->ne[2], c->ne[3] / N, N); // [N, dim, t_len*h_len*w_len] - c = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, c, 1, 0, 2, 3)); // [N, t_len*h_len*w_len, dim] - } - sd::ggml_graph_cut::mark_graph_cut(x, "wan.prelude", "x"); - // sd::ggml_graph_cut::mark_graph_cut(e, "wan.prelude", "e"); - // sd::ggml_graph_cut::mark_graph_cut(e0, "wan.prelude", "e0"); - // sd::ggml_graph_cut::mark_graph_cut(context, "wan.prelude", "context"); - if (c != nullptr) { - sd::ggml_graph_cut::mark_graph_cut(c, "wan.prelude", "c"); - } - - auto x_orig = x; - - for (int i = 0; i < params.num_layers; i++) { - auto block = std::dynamic_pointer_cast(blocks["blocks." + std::to_string(i)]); - - x = block->forward(ctx, x, e0, pe, context, context_img_len); - - auto iter = params.vace_layers_mapping.find(i); - if (iter != params.vace_layers_mapping.end()) { - int n = iter->second; - - auto vace_block = std::dynamic_pointer_cast(blocks["vace_blocks." + std::to_string(n)]); - - auto result = vace_block->forward(ctx, c, x_orig, e0, pe, context, context_img_len); - auto c_skip = result.first; - c = result.second; - c_skip = ggml_ext_scale(ctx->ggml_ctx, c_skip, vace_strength); - x = ggml_add(ctx->ggml_ctx, x, c_skip); - } - sd::ggml_graph_cut::mark_graph_cut(x, "wan.blocks." + std::to_string(i), "x"); - if (c != nullptr) { - sd::ggml_graph_cut::mark_graph_cut(c, "wan.blocks." + std::to_string(i), "c"); - } - } - - x = head->forward(ctx, x, e); // [N, t_len*h_len*w_len, pt*ph*pw*out_dim] - - return x; - } - - ggml_tensor* forward(GGMLRunnerContext* ctx, - ggml_tensor* x, - ggml_tensor* timestep, - ggml_tensor* context, - ggml_tensor* pe, - ggml_tensor* clip_fea = nullptr, - ggml_tensor* time_dim_concat = nullptr, - ggml_tensor* vace_context = nullptr, - float vace_strength = 1.f, - int64_t N = 1) { - // Forward pass of DiT. - // x: [N*C, T, H, W] - // timestep: [N,] - // context: [N, L, D] - // pe: [L, d_head/2, 2, 2] - // time_dim_concat: [N*C, T2, H, W] - // return: [N*C, T, H, W] - - GGML_ASSERT(N == 1); - - int64_t W = x->ne[0]; - int64_t H = x->ne[1]; - int64_t T = x->ne[2]; - int64_t C = x->ne[3]; - - x = pad_to_patch_size(ctx, x); - - int64_t t_len = ((T + (std::get<0>(params.patch_size) / 2)) / std::get<0>(params.patch_size)); - int64_t h_len = ((H + (std::get<1>(params.patch_size) / 2)) / std::get<1>(params.patch_size)); - int64_t w_len = ((W + (std::get<2>(params.patch_size) / 2)) / std::get<2>(params.patch_size)); - - if (time_dim_concat != nullptr) { - time_dim_concat = pad_to_patch_size(ctx, time_dim_concat); - x = ggml_concat(ctx->ggml_ctx, x, time_dim_concat, 2); // [N*C, (T+pad_t) + (T2+pad_t2), H + pad_h, W + pad_w] - t_len = ((x->ne[2] + (std::get<0>(params.patch_size) / 2)) / std::get<0>(params.patch_size)); - } - - auto out = forward_orig(ctx, x, timestep, context, pe, clip_fea, vace_context, vace_strength, N); // [N, t_len*h_len*w_len, pt*ph*pw*C] - - out = unpatchify(ctx->ggml_ctx, out, t_len, h_len, w_len); // [N*C, (T+pad_t) + (T2+pad_t2), H + pad_h, W + pad_w] - - // slice - - out = ggml_ext_slice(ctx->ggml_ctx, out, 2, 0, T); // [N*C, T, H + pad_h, W + pad_w] - out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, H); // [N*C, T, H, W + pad_w] - out = ggml_ext_slice(ctx->ggml_ctx, out, 0, 0, W); // [N*C, T, H, W] - - return out; - } - }; - - struct WanRunner : public DiffusionModelRunner { - public: - std::string desc = "wan"; - WanParams wan_params; - Wan wan; - std::vector pe_vec; - SDVersion version; - - WanRunner(ggml_backend_t backend, - ggml_backend_t params_backend, - const String2TensorStorage& tensor_storage_map = {}, - const std::string prefix = "", - SDVersion version = VERSION_WAN2) - : DiffusionModelRunner(backend, params_backend, prefix) { - wan_params.num_layers = 0; - for (auto pair : tensor_storage_map) { - std::string tensor_name = pair.first; - if (tensor_name.find(prefix) == std::string::npos) - continue; - size_t pos = tensor_name.find("vace_blocks."); - if (pos != std::string::npos) { - tensor_name = tensor_name.substr(pos); // remove prefix - auto items = split_string(tensor_name, '.'); - if (items.size() > 1) { - int block_index = atoi(items[1].c_str()); - if (block_index + 1 > wan_params.vace_layers) { - wan_params.vace_layers = block_index + 1; - } - } - continue; - } - pos = tensor_name.find("blocks."); - if (pos != std::string::npos) { - tensor_name = tensor_name.substr(pos); // remove prefix - auto items = split_string(tensor_name, '.'); - if (items.size() > 1) { - int block_index = atoi(items[1].c_str()); - if (block_index + 1 > wan_params.num_layers) { - wan_params.num_layers = block_index + 1; - } - } - continue; - } - if (tensor_name.find("img_emb") != std::string::npos) { - wan_params.model_type = "i2v"; - } - if (tensor_name.find("img_emb.emb_pos") != std::string::npos) { - wan_params.flf_pos_embed_token_number = 514; - } - } - - if (wan_params.num_layers == 30) { - if (version == VERSION_WAN2_2_TI2V) { - desc = "Wan2.2-TI2V-5B"; - wan_params.dim = 3072; - wan_params.eps = 1e-06f; - wan_params.ffn_dim = 14336; - wan_params.freq_dim = 256; - wan_params.in_dim = 48; - wan_params.num_heads = 24; - wan_params.out_dim = 48; - wan_params.text_len = 512; - } else { - if (wan_params.vace_layers > 0) { - desc = "Wan2.1-VACE-1.3B"; - wan_params.in_dim = 16; - } else if (wan_params.model_type == "i2v") { - desc = "Wan2.1-I2V-1.3B"; - wan_params.in_dim = 36; - } else { - desc = "Wan2.1-T2V-1.3B"; - wan_params.in_dim = 16; - } - wan_params.dim = 1536; - wan_params.eps = 1e-06f; - wan_params.ffn_dim = 8960; - wan_params.freq_dim = 256; - wan_params.num_heads = 12; - wan_params.out_dim = 16; - wan_params.text_len = 512; - } - } else if (wan_params.num_layers == 40) { - if (wan_params.model_type == "t2v") { - if (version == VERSION_WAN2_2_I2V) { - desc = "Wan2.2-I2V-14B"; - wan_params.in_dim = 36; - } else { - if (wan_params.vace_layers > 0) { - desc = "Wan2.x-VACE-14B"; - } else { - desc = "Wan2.x-T2V-14B"; - } - wan_params.in_dim = 16; - } - } else { - wan_params.in_dim = 36; - if (wan_params.flf_pos_embed_token_number > 0) { - desc = "Wan2.1-FLF2V-14B"; - } else { - desc = "Wan2.1-I2V-14B"; - } - } - wan_params.dim = 5120; - wan_params.eps = 1e-06f; - wan_params.ffn_dim = 13824; - wan_params.freq_dim = 256; - wan_params.num_heads = 40; - wan_params.out_dim = 16; - wan_params.text_len = 512; - } else { - GGML_ABORT("invalid num_layers(%d) of wan", wan_params.num_layers); - } - - LOG_INFO("%s", desc.c_str()); - - wan = Wan(wan_params); - wan.init(params_ctx, tensor_storage_map, prefix); - } - - std::string get_desc() override { - return desc; - } - - void get_param_tensors(std::map& tensors, const std::string& prefix) override { - wan.get_param_tensors(tensors, prefix); - } - - ggml_cgraph* build_graph(const sd::Tensor& x_tensor, - const sd::Tensor& timesteps_tensor, - const sd::Tensor& context_tensor = {}, - const sd::Tensor& clip_fea_tensor = {}, - const sd::Tensor& c_concat_tensor = {}, - const sd::Tensor& time_dim_concat_tensor = {}, - const sd::Tensor& vace_context_tensor = {}, - float vace_strength = 1.f) { - ggml_cgraph* gf = new_graph_custom(WAN_GRAPH_SIZE); - - ggml_tensor* x = make_input(x_tensor); - ggml_tensor* timesteps = make_input(timesteps_tensor); - ggml_tensor* context = make_optional_input(context_tensor); - ggml_tensor* clip_fea = make_optional_input(clip_fea_tensor); - ggml_tensor* c_concat = make_optional_input(c_concat_tensor); - ggml_tensor* time_dim_concat = make_optional_input(time_dim_concat_tensor); - ggml_tensor* vace_context = make_optional_input(vace_context_tensor); - - pe_vec = Rope::gen_wan_pe(static_cast(x->ne[2]), - static_cast(x->ne[1]), - static_cast(x->ne[0]), - std::get<0>(wan_params.patch_size), - std::get<1>(wan_params.patch_size), - std::get<2>(wan_params.patch_size), - 1, - wan_params.theta, - wan_params.axes_dim); - int pos_len = static_cast(pe_vec.size() / wan_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, wan_params.axes_dim_sum / 2, pos_len); - // pe->data = pe_vec.data(); - // print_ggml_tensor(pe); - // pe->data = nullptr; - set_backend_tensor_data(pe, pe_vec.data()); - - if (c_concat != nullptr) { - x = ggml_concat(compute_ctx, x, c_concat, 3); - } - - auto runner_ctx = get_context(); - - ggml_tensor* out = wan.forward(&runner_ctx, - x, - timesteps, - context, - pe, - clip_fea, - time_dim_concat, - vace_context, - vace_strength); - - ggml_build_forward_expand(gf, out); - - return gf; - } - - sd::Tensor compute(int n_threads, - const sd::Tensor& x, - const sd::Tensor& timesteps, - const sd::Tensor& context = {}, - const sd::Tensor& clip_fea = {}, - const sd::Tensor& c_concat = {}, - const sd::Tensor& time_dim_concat = {}, - const sd::Tensor& vace_context = {}, - float vace_strength = 1.f) { - auto get_graph = [&]() -> ggml_cgraph* { - return build_graph(x, timesteps, context, clip_fea, c_concat, time_dim_concat, vace_context, vace_strength); - }; - - return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); - } - - sd::Tensor compute(int n_threads, - const DiffusionParams& diffusion_params) override { - GGML_ASSERT(diffusion_params.x != nullptr); - GGML_ASSERT(diffusion_params.timesteps != nullptr); - const auto* extra = diffusion_extra_as(diffusion_params); - return compute(n_threads, - *diffusion_params.x, - *diffusion_params.timesteps, - tensor_or_empty(diffusion_params.context), - tensor_or_empty(diffusion_params.y), - tensor_or_empty(diffusion_params.c_concat), - sd::Tensor(), - tensor_or_empty(extra->vace_context), - extra->vace_strength); - } - - void test() { - ggml_init_params params; - params.mem_size = static_cast(200 * 1024 * 1024); // 200 MB - params.mem_buffer = nullptr; - params.no_alloc = false; - - ggml_context* ctx = ggml_init(params); - GGML_ASSERT(ctx != nullptr); - - { - // cpu f16: pass - // cuda f16: pass - // cpu q8_0: pass - // auto x = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 104, 60, 1, 16); - // ggml_set_f32(x, 0.01f); - auto x = sd::load_tensor_from_file_as_tensor("wan_dit_x.bin"); - print_sd_tensor(x); - - std::vector timesteps_vec(3, 1000.f); - timesteps_vec[0] = 0.f; - auto timesteps = sd::Tensor::from_vector(timesteps_vec); - - // auto context = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 4096, 512, 1); - // ggml_set_f32(context, 0.01f); - auto context = sd::load_tensor_from_file_as_tensor("wan_dit_context.bin"); - print_sd_tensor(context); - // auto clip_fea = load_tensor_from_file(ctx, "wan_dit_clip_fea.bin"); - // print_ggml_tensor(clip_fea); - - sd::Tensor out; - - int64_t t0 = ggml_time_ms(); - auto out_opt = compute(8, x, timesteps, context, {}, {}, {}, {}, 1.f); - int64_t t1 = ggml_time_ms(); - - GGML_ASSERT(!out_opt.empty()); - out = std::move(out_opt); - print_sd_tensor(out); - LOG_DEBUG("wan test done in %lldms", t1 - t0); - } - } - - static void load_from_file_and_test(const std::string& file_path) { - // ggml_backend_t backend = ggml_backend_cuda_init(0); - ggml_backend_t backend = sd_backend_cpu_init(); - ggml_type model_data_type = GGML_TYPE_F16; - LOG_INFO("loading from '%s'", file_path.c_str()); - - ModelLoader model_loader; - if (!model_loader.init_from_file_and_convert_name(file_path, "model.diffusion_model.")) { - LOG_ERROR("init model loader from file failed: '%s'", file_path.c_str()); - return; - } - - auto& tensor_storage_map = model_loader.get_tensor_storage_map(); - for (auto& [name, tensor_storage] : tensor_storage_map) { - if (ends_with(name, "weight")) { - tensor_storage.expected_type = model_data_type; - } - } - - std::shared_ptr wan = std::make_shared(backend, - backend, - tensor_storage_map, - "model.diffusion_model", - VERSION_WAN2_2_TI2V); - - if (!wan->alloc_params_buffer()) { - LOG_ERROR("wan buffer allocation failed"); - return; - } - - std::map tensors; - wan->get_param_tensors(tensors, "model.diffusion_model"); - - bool success = model_loader.load_tensors(tensors); - - if (!success) { - LOG_ERROR("load tensors from model loader failed"); - return; - } - - LOG_INFO("wan model loaded"); - - wan->test(); - } - }; - } // namespace WAN -#endif // __WAN_HPP__ +#endif // __SD_MODEL_VAE_WAN_VAE_HPP__ diff --git a/src/model_io/gguf_io.cpp b/src/model_io/gguf_io.cpp index 378694d8..c701d01f 100644 --- a/src/model_io/gguf_io.cpp +++ b/src/model_io/gguf_io.cpp @@ -5,9 +5,9 @@ #include #include +#include "core/util.h" #include "gguf.h" #include "gguf_reader_ext.h" -#include "util.h" static void set_error(std::string* error, const std::string& message) { if (error != nullptr) { diff --git a/src/model_io/gguf_reader_ext.h b/src/model_io/gguf_reader_ext.h index 95f0027f..7da20d0a 100644 --- a/src/model_io/gguf_reader_ext.h +++ b/src/model_io/gguf_reader_ext.h @@ -6,8 +6,8 @@ #include #include +#include "core/util.h" #include "ggml.h" -#include "util.h" struct GGUFTensorInfo { std::string name; diff --git a/src/model_io/pickle_io.cpp b/src/model_io/pickle_io.cpp index 3a978178..1ea7a902 100644 --- a/src/model_io/pickle_io.cpp +++ b/src/model_io/pickle_io.cpp @@ -8,7 +8,7 @@ #include #include "binary_io.h" -#include "util.h" +#include "core/util.h" // $ python -m pickletools sd-v1-4/archive/data.pkl | head -n 100 // 0: \x80 PROTO 2 diff --git a/src/model_io/safetensors_io.cpp b/src/model_io/safetensors_io.cpp index 88935221..39131dbd 100644 --- a/src/model_io/safetensors_io.cpp +++ b/src/model_io/safetensors_io.cpp @@ -7,8 +7,8 @@ #include #include "binary_io.h" +#include "core/util.h" #include "json.hpp" -#include "util.h" static constexpr size_t ST_HEADER_SIZE_LEN = 8; diff --git a/src/model_io/torch_legacy_io.cpp b/src/model_io/torch_legacy_io.cpp index 81654725..0b0c78f8 100644 --- a/src/model_io/torch_legacy_io.cpp +++ b/src/model_io/torch_legacy_io.cpp @@ -7,8 +7,8 @@ #include #include +#include "core/util.h" #include "pickle_io.h" -#include "util.h" // torch.save format background: // diff --git a/src/model.cpp b/src/model_loader.cpp similarity index 84% rename from src/model.cpp rename to src/model_loader.cpp index 00370e6f..ba8e090d 100644 --- a/src/model.cpp +++ b/src/model_loader.cpp @@ -1,6 +1,7 @@ #include #include #include +#include #include #include #include @@ -13,18 +14,19 @@ #include #include -#include "model.h" +#include "core/util.h" #include "model_io/gguf_io.h" #include "model_io/safetensors_io.h" #include "model_io/torch_legacy_io.h" #include "model_io/torch_zip_io.h" +#include "model_loader.h" +#include "runtime/imatrix.h" #include "stable-diffusion.h" -#include "util.h" +#include "core/ggml_extend_backend.h" #include "ggml-alloc.h" #include "ggml-backend.h" #include "ggml.h" -#include "ggml_extend_backend.h" #include "zip.h" #include "name_conversion.h" @@ -65,7 +67,6 @@ const char* unused_tensors[] = { // "v_pred", // Used to detect SDXL vpred models "text_encoders.llm.output.weight", "text_encoders.llm.lm_head.", - "first_stage_model.bn.", }; bool is_unused_tensor(const std::string& name) { @@ -156,7 +157,8 @@ void convert_tensor(void* src, void* dst, ggml_type dst_type, int nrows, - int n_per_row) { + int n_per_row, + std::vector imatrix = {}) { int n = nrows * n_per_row; if (src_type == dst_type) { size_t nbytes = n * ggml_type_size(src_type) / ggml_blck_size(src_type); @@ -165,7 +167,7 @@ void convert_tensor(void* src, if (dst_type == GGML_TYPE_F16) { ggml_fp32_to_fp16_row((float*)src, (ggml_fp16_t*)dst, n); } else { - std::vector imatrix(n_per_row, 1.0f); // dummy importance matrix + imatrix.resize(n_per_row, 1.0f); const float* im = imatrix.data(); ggml_quantize_chunk(dst_type, (float*)src, dst, 0, nrows, n_per_row, im); } @@ -195,7 +197,7 @@ void convert_tensor(void* src, if (dst_type == GGML_TYPE_F16) { ggml_fp32_to_fp16_row((float*)src_data_f32, (ggml_fp16_t*)dst, n); } else { - std::vector imatrix(n_per_row, 1.0f); // dummy importance matrix + imatrix.resize(n_per_row, 1.0f); const float* im = imatrix.data(); ggml_quantize_chunk(dst_type, (float*)src_data_f32, dst, 0, nrows, n_per_row, im); } @@ -204,10 +206,28 @@ void convert_tensor(void* src, /*================================================= ModelLoader ==================================================*/ +ModelLoader::ModelLoader() + : n_threads_(sd_get_num_physical_cores()) { +} + +size_t ModelLoader::add_file_path(const std::string& file_path) { + if (model_files_processed) { + file_data.clear(); + model_files_processed = false; + } + file_paths_.push_back(file_path); + return file_paths_.size() - 1; +} + void ModelLoader::add_tensor_storage(const TensorStorage& tensor_storage) { tensor_storage_map[tensor_storage.name] = tensor_storage; } +void ModelLoader::set_n_threads(int n_threads) { + n_threads_ = n_threads > 0 ? n_threads : sd_get_num_physical_cores(); + LOG_DEBUG("using %d threads for model loading", n_threads_); +} + bool ModelLoader::init_from_file(const std::string& file_path, const std::string& prefix) { if (is_directory(file_path)) { LOG_INFO("load %s using diffusers format", file_path.c_str()); @@ -271,8 +291,7 @@ bool ModelLoader::init_from_gguf_file(const std::string& file_path, const std::s return false; } - file_paths_.push_back(file_path); - size_t file_index = file_paths_.size() - 1; + size_t file_index = add_file_path(file_path); for (auto& tensor_storage : tensor_storages) { // LOG_DEBUG("%s", tensor_storage.name.c_str()); @@ -300,8 +319,7 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const return false; } - file_paths_.push_back(file_path); - size_t file_index = file_paths_.size() - 1; + size_t file_index = add_file_path(file_path); for (auto& tensor_storage : tensor_storages) { if (is_unused_tensor(tensor_storage.name)) { @@ -335,8 +353,7 @@ bool ModelLoader::init_from_torch_legacy_file(const std::string& file_path, cons return false; } - file_paths_.push_back(file_path); - size_t file_index = file_paths_.size() - 1; + size_t file_index = add_file_path(file_path); for (auto& tensor_storage : tensor_storages) { if (is_unused_tensor(tensor_storage.name)) { @@ -366,8 +383,7 @@ bool ModelLoader::init_from_torch_zip_file(const std::string& file_path, const s return false; } - file_paths_.push_back(file_path); - size_t file_index = file_paths_.size() - 1; + size_t file_index = add_file_path(file_path); for (auto& tensor_storage : tensor_storages) { if (!starts_with(tensor_storage.name, prefix)) { @@ -435,6 +451,13 @@ SDVersion ModelLoader::get_sd_version() { if (tensor_storage.name.find("model.diffusion_model.net.lq_proj.latent_proj.0.weight") != std::string::npos) { return VERSION_PID; } + if (tensor_storage.name.find("embed_image_indicator.weight") != std::string::npos) { + return VERSION_IDEOGRAM4; + } + if (tensor_storage.name.find("model.diffusion_model.txtfusion.projector.weight") != std::string::npos || + tensor_storage.name.find("model.diffusion_model.text_fusion.projector.weight") != std::string::npos) { + return VERSION_KREA2; + } if (tensor_storage.name.find("model.diffusion_model.nerf_final_layer_conv.") != std::string::npos) { return VERSION_CHROMA_RADIANCE; } @@ -449,6 +472,9 @@ SDVersion ModelLoader::get_sd_version() { 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("net.img_embedder.proj1.weight") != std::string::npos) { + return VERSION_MINIT2I; + } if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.img_mod.1.weight") != std::string::npos) { if (tensor_storage_map.find("model.diffusion_model.time_text_embed.addition_t_embedding.weight") != tensor_storage_map.end()) { return VERSION_QWEN_IMAGE_LAYERED; @@ -461,6 +487,9 @@ SDVersion ModelLoader::get_sd_version() { if (tensor_storage.name.find("model.diffusion_model.double_stream_modulation_img.lin.weight") != std::string::npos) { is_flux2 = true; } + if (tensor_storage.name.find("dual_time_embed.semantic_embedder.linear_1.weight") != std::string::npos) { + return VERSION_SEFI_IMAGE; + } if (tensor_storage.name.find("single_blocks.47.linear1.weight") != std::string::npos) { has_single_block_47 = true; } @@ -470,6 +499,9 @@ SDVersion ModelLoader::get_sd_version() { if (tensor_storage.name.find("model.diffusion_model.cap_embedder.0.weight") != std::string::npos) { return VERSION_Z_IMAGE; } + if (tensor_storage.name.find("double_stream_layers.0.img_instruct_attn.processor.img_to_q.weight") != std::string::npos) { + return VERSION_BOOGU_IMAGE; + } if (tensor_storage.name.find("model.diffusion_model.layers.0.adaLN_sa_ln.weight") != std::string::npos) { return VERSION_ERNIE_IMAGE; } @@ -760,8 +792,6 @@ void ModelLoader::process_model_files(bool enable_mmap, bool writable_mmap) { return; } - int64_t start_time = ggml_time_ms(); - std::vector processed_tensor_storages; for (const auto& [name, tensor_storage] : tensor_storage_map) { if (is_unused_tensor(tensor_storage.name)) { @@ -812,20 +842,12 @@ void ModelLoader::process_model_files(bool enable_mmap, bool writable_mmap) { } else { LOG_WARN("failed to memory-map '%s' (falling back to read())", file_path.c_str()); } - } else if (!is_zip) { - LOG_INFO("NOT using mmap for '%s' (mmap disabled by caller)", - file_path.c_str()); } file_data.push_back(std::move(fdata)); } model_files_processed = true; - - int64_t end_time = ggml_time_ms(); - int64_t process_time_ms = end_time - start_time; - - LOG_INFO("model files processing completed in %.2fs", process_time_ms / 1000.f); } std::vector ModelLoader::mmap_tensors(std::map& tensors, @@ -919,7 +941,9 @@ std::vector ModelLoader::mmap_tensors(std::map* target_tensor_names) { process_model_files(enable_mmap, false); std::atomic read_time_ms(0); @@ -928,32 +952,66 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread std::atomic convert_time_ms(0); std::atomic bytes_processed(0); - int num_threads_to_use = n_threads_p > 0 ? n_threads_p : sd_get_num_physical_cores(); - LOG_DEBUG("using %d threads for model loading", num_threads_to_use); + int num_threads_to_use = n_threads_; int64_t start_time = ggml_time_ms(); size_t total_tensors_to_process = 0; + std::vector file_tensors_to_process_counts; + file_tensors_to_process_counts.reserve(file_data.size()); for (const auto& fdata : file_data) { - total_tensors_to_process += fdata.tensors.size(); + size_t file_tensors_to_process = 0; + if (target_tensor_names == nullptr) { + file_tensors_to_process = fdata.tensors.size(); + } else { + for (const TensorStorage& tensor_storage : fdata.tensors) { + if (target_tensor_names->find(tensor_storage.name) != target_tensor_names->end()) { + file_tensors_to_process++; + } + } + } + file_tensors_to_process_counts.push_back(file_tensors_to_process); + total_tensors_to_process += file_tensors_to_process; } bool success = true; size_t total_tensors_processed = 0; const int64_t t_start = start_time; int last_n_threads = 1; + SDVersion imatrix_version = (version_ == VERSION_COUNT) ? get_sd_version() : version_; - for (auto& fdata : file_data) { + for (size_t file_index = 0; file_index < file_data.size(); ++file_index) { + auto& fdata = file_data[file_index]; const std::string& file_path = fdata.path; - LOG_DEBUG("loading tensors from %s", file_path.c_str()); const std::vector& file_tensors = fdata.tensors; + std::vector tensors_to_process; + size_t file_tensors_to_process = file_tensors_to_process_counts[file_index]; + tensors_to_process.reserve(file_tensors_to_process); + if (target_tensor_names == nullptr) { + for (const TensorStorage& tensor_storage : file_tensors) { + tensors_to_process.push_back(&tensor_storage); + } + } else { + for (const TensorStorage& tensor_storage : file_tensors) { + if (target_tensor_names->find(tensor_storage.name) != target_tensor_names->end()) { + tensors_to_process.push_back(&tensor_storage); + } + } + } + if (tensors_to_process.empty()) { + continue; + } + LOG_DEBUG("loading %zu/%zu tensors from %s", + tensors_to_process.size(), + file_tensors.size(), + file_path.c_str()); bool is_zip = fdata.is_zip; std::shared_ptr mmapped = fdata.mmapped; - int n_threads = is_zip ? 1 : std::min(num_threads_to_use, (int)file_tensors.size()); + int n_threads = is_zip ? 1 : std::min(num_threads_to_use, (int)tensors_to_process.size()); if (n_threads < 1) { n_threads = 1; } @@ -962,6 +1020,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread std::atomic tensor_idx(0); std::atomic failed(false); std::vector workers; + std::mutex rpc_backend_mutex; for (int i = 0; i < n_threads; ++i) { workers.emplace_back([&, file_path, is_zip]() { @@ -989,11 +1048,11 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread while (true) { int64_t t0, t1; size_t idx = tensor_idx.fetch_add(1); - if (idx >= file_tensors.size() || failed) { + if (idx >= tensors_to_process.size() || failed) { break; } - const TensorStorage& tensor_storage = file_tensors[idx]; + const TensorStorage& tensor_storage = *tensors_to_process[idx]; ggml_tensor* dst_tensor = nullptr; t0 = ggml_time_ms(); @@ -1104,12 +1163,15 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread failed = true; return; } + std::string processed_name = convert_tensor_name(tensor_storage.name, imatrix_version); + std::vector imatrix = get_imatrix_collector().get_values(processed_name); convert_tensor((void*)target_buf, tensor_storage.type, convert_buf, dst_tensor->type, (int)tensor_storage.nelements() / (int)tensor_storage.ne[0], - (int)tensor_storage.ne[0]); + (int)tensor_storage.ne[0], + std::move(imatrix)); } else { convert_buf = read_buf; } @@ -1118,7 +1180,19 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread if (dst_tensor->buffer != nullptr && !ggml_backend_buffer_is_host(dst_tensor->buffer)) { t0 = ggml_time_ms(); - ggml_backend_tensor_set(dst_tensor, convert_buf, 0, ggml_nbytes(dst_tensor)); + + // RPC backends require serialized access to prevent concurrency issues + const char* buffer_type_name = ggml_backend_buft_name(ggml_backend_buffer_get_type(dst_tensor->buffer)); + bool is_rpc_buffer = buffer_type_name != nullptr && + std::string(buffer_type_name).find("RPC") != std::string::npos; + + if (is_rpc_buffer) { + std::lock_guard lock(rpc_backend_mutex); + ggml_backend_tensor_set(dst_tensor, convert_buf, 0, ggml_nbytes(dst_tensor)); + } else { + ggml_backend_tensor_set(dst_tensor, convert_buf, 0, ggml_nbytes(dst_tensor)); + } + t1 = ggml_time_ms(); copy_to_backend_time_ms.fetch_add(t1 - t0); } @@ -1133,16 +1207,18 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread while (true) { size_t current_idx = tensor_idx.load(); - if (current_idx >= file_tensors.size() || failed) { + if (current_idx >= tensors_to_process.size() || failed) { break; } size_t curr_num = total_tensors_processed + current_idx; float elapsed_seconds = (ggml_time_ms() - t_start) / 1000.0f; - pretty_bytes_progress(static_cast(curr_num), - static_cast(total_tensors_to_process), - bytes_processed.load(), - elapsed_seconds); - std::this_thread::sleep_for(std::chrono::milliseconds(200)); + if (total_tensors_to_process > 0) { + pretty_bytes_progress(static_cast(curr_num), + static_cast(total_tensors_to_process), + bytes_processed.load(), + elapsed_seconds); + } + std::this_thread::sleep_for(std::chrono::milliseconds(total_tensors_to_process <= 4 ? 10 : 200)); } for (auto& w : workers) { @@ -1153,12 +1229,14 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread success = false; break; } - total_tensors_processed += file_tensors.size(); - pretty_bytes_progress(static_cast(total_tensors_processed), - static_cast(total_tensors_to_process), - bytes_processed.load(), - (ggml_time_ms() - t_start) / 1000.0f); - if (total_tensors_processed < total_tensors_to_process) { + total_tensors_processed += tensors_to_process.size(); + if (total_tensors_to_process > 0) { + pretty_bytes_progress(static_cast(total_tensors_processed), + static_cast(total_tensors_to_process), + bytes_processed.load(), + (ggml_time_ms() - t_start) / 1000.0f); + } + if (total_tensors_processed < total_tensors_to_process && total_tensors_to_process > 0) { printf("\n"); } } @@ -1173,9 +1251,77 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, int n_thread return success; } +bool ModelLoader::load_float_tensor(const std::string& name, + std::vector& data, + int n_threads, + bool use_mmap) { + data.clear(); + + auto tensor_storage_it = tensor_storage_map.find(name); + if (tensor_storage_it == tensor_storage_map.end()) { + return false; + } + + const TensorStorage& tensor_storage = tensor_storage_it->second; + int64_t n_elements = tensor_storage.nelements(); + if (n_elements <= 0) { + LOG_ERROR("tensor '%s' has invalid element count: %" PRId64, name.c_str(), n_elements); + return false; + } + if (tensor_storage.n_dims <= 0 || tensor_storage.n_dims > GGML_MAX_DIMS) { + LOG_ERROR("tensor '%s' has unsupported dims: %d", name.c_str(), tensor_storage.n_dims); + return false; + } + + std::vector loaded_data(static_cast(n_elements)); + ggml_init_params params; + params.mem_size = ggml_tensor_overhead(); + params.mem_buffer = nullptr; + params.no_alloc = true; + + ggml_context* ctx = ggml_init(params); + if (ctx == nullptr) { + LOG_ERROR("failed to create context for tensor '%s'", name.c_str()); + return false; + } + + ggml_tensor* tensor = ggml_new_tensor(ctx, GGML_TYPE_F32, tensor_storage.n_dims, tensor_storage.ne); + ggml_set_name(tensor, name.c_str()); + tensor->data = loaded_data.data(); + + bool loaded = false; + auto on_new_tensor_cb = [&](const TensorStorage& current_tensor_storage, ggml_tensor** dst_tensor) -> bool { + *dst_tensor = nullptr; + if (current_tensor_storage.name != name) { + return true; + } + if (current_tensor_storage.nelements() != n_elements) { + LOG_ERROR("tensor '%s' element count changed during load", name.c_str()); + return false; + } + *dst_tensor = tensor; + loaded = true; + return true; + }; + + std::set target_tensor_names{name}; + if (n_threads > 0) { + set_n_threads(n_threads); + } + bool success = load_tensors(on_new_tensor_cb, use_mmap, &target_tensor_names); + ggml_free(ctx); + + if (!success || !loaded) { + data.clear(); + return false; + } + + data = std::move(loaded_data); + return true; +} + bool ModelLoader::load_tensors(std::map& tensors, std::set ignore_tensors, - int n_threads, bool enable_mmap) { std::set tensor_names_in_file; std::mutex tensor_names_mutex; @@ -1219,7 +1365,7 @@ bool ModelLoader::load_tensors(std::map& tensors, return true; }; - bool success = load_tensors(on_new_tensor_cb, n_threads, enable_mmap); + bool success = load_tensors(on_new_tensor_cb, enable_mmap); if (!success) { LOG_ERROR("load tensors from file failed"); return false; @@ -1257,6 +1403,8 @@ bool ModelLoader::tensor_should_be_converted(const TensorStorage& tensor_storage // Pass, do not convert } else if (ends_with(name, ".scale")) { // Pass, do not convert + } else if (ends_with(name, ".weight_scale")) { + // Pass, do not convert } else if (contains(name, "img_in.") || contains(name, "txt_in.") || contains(name, "time_in.") || diff --git a/src/model_loader.h b/src/model_loader.h new file mode 100644 index 00000000..4dc700f2 --- /dev/null +++ b/src/model_loader.h @@ -0,0 +1,93 @@ +#ifndef __MODEL_LOADER_H__ +#define __MODEL_LOADER_H__ + +#include +#include +#include +#include +#include +#include + +#include "model.h" + +TensorTypeRules parse_tensor_type_rules(const std::string& tensor_type_rules); + +class MmapWrapper; + +struct ModelFileData { + std::string path; + std::vector tensors; + std::shared_ptr mmapped; + std::shared_ptr mmbuffer; + bool is_zip; +}; + +struct MmapTensorStore { + std::shared_ptr mmapped; + std::shared_ptr mmbuffer; +}; + +class ModelLoader { +protected: + SDVersion version_ = VERSION_COUNT; + std::vector file_paths_; + std::vector file_data; + bool model_files_processed = false; + String2TensorStorage tensor_storage_map; + int n_threads_; + + size_t add_file_path(const std::string& file_path); + void add_tensor_storage(const TensorStorage& tensor_storage); + + bool init_from_gguf_file(const std::string& file_path, const std::string& prefix = ""); + bool init_from_safetensors_file(const std::string& file_path, const std::string& prefix = ""); + bool init_from_torch_zip_file(const std::string& file_path, const std::string& prefix = ""); + bool init_from_torch_legacy_file(const std::string& file_path, const std::string& prefix = ""); + bool init_from_diffusers_file(const std::string& file_path, const std::string& prefix = ""); + +public: + ModelLoader(); + + bool init_from_file(const std::string& file_path, const std::string& prefix = ""); + void convert_tensors_name(); + bool init_from_file_and_convert_name(const std::string& file_path, + const std::string& prefix = "", + SDVersion version = VERSION_COUNT); + SDVersion get_sd_version(); + std::map get_wtype_stat(); + std::map get_conditioner_wtype_stat(); + std::map get_diffusion_model_wtype_stat(); + std::map get_vae_wtype_stat(); + String2TensorStorage& get_tensor_storage_map() { return tensor_storage_map; } + const String2TensorStorage& get_tensor_storage_map() const { return tensor_storage_map; } + void set_n_threads(int n_threads); + void set_wtype_override(ggml_type wtype, std::string tensor_type_rules = ""); + void process_model_files(bool enable_mmap = false, bool writable_mmap = true); + std::vector mmap_tensors(std::map& tensors, + std::set ignore_tensors = {}, + bool writable = true); + bool load_tensors(on_new_tensor_cb_t on_new_tensor_cb, + bool use_mmap = false, + const std::set* target_tensor_names = nullptr); + bool load_tensors(std::map& tensors, + std::set ignore_tensors = {}, + bool use_mmap = false); + bool load_float_tensor(const std::string& name, + std::vector& data, + int n_threads = 0, + bool use_mmap = false); + + std::vector get_tensor_names() const { + std::vector names; + for (const auto& [name, tensor_storage] : tensor_storage_map) { + names.push_back(name); + } + return names; + } + + bool tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type); + int64_t get_params_mem_size(ggml_backend_t backend, ggml_type type = GGML_TYPE_COUNT); + ~ModelLoader() = default; +}; + +#endif // __MODEL_LOADER_H__ diff --git a/src/model_manager.cpp b/src/model_manager.cpp new file mode 100644 index 00000000..7095ec6a --- /dev/null +++ b/src/model_manager.cpp @@ -0,0 +1,950 @@ +#include "model_manager.h" + +#include +#include +#include +#include +#include + +#include "core/ggml_extend_backend.h" +#include "core/util.h" +#include "model/adapter/lora.hpp" + +static size_t aligned_offset(const void* buffer, size_t offset, size_t alignment) { + GGML_ASSERT(alignment != 0 && (alignment & (alignment - 1)) == 0); + size_t align = (alignment - ((reinterpret_cast(buffer) + offset) % alignment)) % alignment; + return offset + align; +} + +static bool lora_specs_equal(const std::vector& lhs, + const std::vector& rhs) { + if (lhs.size() != rhs.size()) { + return false; + } + for (size_t i = 0; i < lhs.size(); ++i) { + if (lhs[i].path != rhs[i].path || + lhs[i].multiplier != rhs[i].multiplier || + lhs[i].is_high_noise != rhs[i].is_high_noise || + lhs[i].tensor_name_prefix_filter != rhs[i].tensor_name_prefix_filter || + lhs[i].required != rhs[i].required) { + return false; + } + } + return true; +} + +static std::string lora_id(const ModelManager::LoraSpec& lora) { + return lora.is_high_noise ? "|high_noise|" + lora.path : lora.path; +} + +static bool backend_supports_host_buffer(ggml_backend_t backend) { + if (backend == nullptr) { + return false; + } + if (sd_backend_is_cpu(backend)) { + return true; + } + ggml_backend_dev_t dev = ggml_backend_get_device(backend); + if (dev == nullptr) { + return false; + } + ggml_backend_dev_props props; + ggml_backend_dev_get_props(dev, &props); + return props.caps.buffer_from_host_ptr; +} + +ModelManager::~ModelManager() { + release_all(); +} + +void ModelManager::set_common_ignore_tensors(std::set ignore_tensors) { + common_ignore_tensors_ = std::move(ignore_tensors); +} + +void ModelManager::set_loras(std::vector loras, SDVersion version) { + if (loras.empty() && loras_.empty()) { + lora_version_ = version; + return; + } + if (lora_version_ == version && lora_specs_equal(loras_, loras)) { + return; + } + + loras_ = std::move(loras); + lora_version_ = version; + current_lora_epoch_++; + reset_lora_applied_params(); +} + +std::set ModelManager::tensor_names() const { + std::set names; + for (const auto& state : tensor_states_) { + if (state != nullptr) { + names.insert(state->name); + } + } + return names; +} + +size_t estimate_tensors_size(const std::map& tensors) { + size_t size = 0; + std::unordered_set seen; + for (const auto& pair : tensors) { + ggml_tensor* tensor = pair.second; + if (tensor == nullptr || seen.find(tensor) != seen.end()) { + continue; + } + seen.insert(tensor); + size += ggml_nbytes(tensor); + } + return size; +} + +bool ModelManager::register_param_tensors(const std::string& desc, + std::map tensors, + ResidencyMode residency_mode, + ggml_backend_t compute_backend, + ggml_backend_t params_backend, + size_t* registered_tensor_size) { + if (desc.empty()) { + LOG_ERROR("model manager tensor desc is empty"); + return false; + } + if (registered_tensor_size != nullptr) { + *registered_tensor_size += estimate_tensors_size(tensors); + } + + std::vector> new_states; + new_states.reserve(tensors.size()); + + for (const auto& pair : tensors) { + const std::string& name = pair.first; + ggml_tensor* tensor = pair.second; + if (tensor == nullptr) { + continue; + } + if (tensor_states_by_name_.find(name) != tensor_states_by_name_.end()) { + LOG_ERROR("model manager tensor name '%s' is already registered", name.c_str()); + return false; + } + ggml_set_name(tensor, name.c_str()); + + auto state = std::make_unique(); + state->name = name; + state->tensor = tensor; + state->desc = desc; + state->residency_mode = residency_mode; + state->compute_backend = compute_backend; + state->params_backend = params_backend; + new_states.push_back(std::move(state)); + } + + for (auto& state : new_states) { + TensorState* registered_state = state.get(); + tensor_states_by_name_[registered_state->name] = registered_state; + tensor_states_.push_back(std::move(state)); + } + return true; +} + +bool ModelManager::load_all_params_eagerly() { + std::vector all_states; + all_states.reserve(tensor_states_.size()); + for (const auto& s : tensor_states_) { + if (s != nullptr) { + all_states.push_back(s.get()); + } + } + return load_tensors_to_params_backend(all_states); +} + +bool ModelManager::validate_registered_tensors() { + bool ok = true; + for (const auto& state : tensor_states_) { + if (state == nullptr) { + ok = false; + continue; + } + bool state_ok = validate_tensor(*state); + if (state_ok) { + state->metadata_validated = true; + } + ok = state_ok && ok; + } + return ok; +} + +bool ModelManager::load_tensors_to_params_backend(const std::vector& states) { + std::vector need_load; + need_load.reserve(states.size()); + for (TensorState* state : states) { + if (state == nullptr || should_ignore(*state) || is_optional_missing_tensor(state->name)) { + continue; + } + if (!state->metadata_validated) { + if (!validate_tensor(*state)) { + return false; + } + state->metadata_validated = true; + } + if (!state->loaded_to_params_backend) { + need_load.push_back(state); + } + } + if (need_load.empty()) { + return true; + } + + std::vector created_storage_blocks; + if (!mmap_params(need_load, created_storage_blocks)) { + for (ParamsStorageBlock* block : created_storage_blocks) { + if (block != nullptr) { + free_params_storage_block(*block); + erase_params_storage_block(block); + } + } + return false; + } + + std::vector need_alloc; + need_alloc.reserve(need_load.size()); + for (TensorState* state : need_load) { + if (state->tensor != nullptr && state->tensor->data == nullptr && state->tensor->view_src == nullptr) { + need_alloc.push_back(state); + } + } + + if (!alloc_params_buffers(need_alloc, created_storage_blocks) || + !load_tensors(need_load)) { + for (ParamsStorageBlock* block : created_storage_blocks) { + if (block != nullptr) { + free_params_storage_block(*block); + erase_params_storage_block(block); + } + } + return false; + } + for (ParamsStorageBlock* block : created_storage_blocks) { + if (block != nullptr && block->buffer != nullptr) { + LOG_DEBUG("model manager prepared params backend buffer (%6.2f MB, %zu tensors, %s)", + ggml_backend_buffer_get_size(block->buffer) / (1024.f * 1024.f), + block->states.size(), + ggml_backend_buffer_is_host(block->buffer) ? "RAM" : "VRAM"); + } + } + + return true; +} + +bool ModelManager::stage_tensors_to_compute_backend(const std::vector& states) { + std::map> states_by_compute_backend; + for (TensorState* state : states) { + if (state == nullptr || should_ignore(*state) || is_optional_missing_tensor(state->name)) { + continue; + } + if (state->compute_backend == nullptr) { + LOG_ERROR("model manager compute backend is null for tensor '%s'", state->name.c_str()); + return false; + } + if (state->params_backend == nullptr) { + LOG_ERROR("model manager params backend is null for tensor '%s'", state->name.c_str()); + return false; + } + if (state->compute_backend == state->params_backend || state->staged_to_compute_backend) { + continue; + } + if (!state->loaded_to_params_backend || state->tensor == nullptr || state->tensor->data == nullptr) { + LOG_ERROR("model manager tensor '%s' is not loaded to params backend", state->name.c_str()); + return false; + } + states_by_compute_backend[state->compute_backend].push_back(state); + } + + for (const auto& pair : states_by_compute_backend) { + ggml_backend_t compute_backend = pair.first; + const std::vector& states = pair.second; + if (states.empty()) { + continue; + } + + int64_t t0 = ggml_time_ms(); + + ggml_init_params init_params; + init_params.mem_size = std::max(1, states.size()) * ggml_tensor_overhead(); + init_params.mem_buffer = nullptr; + init_params.no_alloc = true; + + ggml_context* staging_ctx = ggml_init(init_params); + GGML_ASSERT(staging_ctx != nullptr); + + std::vector> staged_tensors; + staged_tensors.reserve(states.size()); + for (TensorState* state : states) { + ggml_tensor* staging_tensor = ggml_dup_tensor(staging_ctx, state->tensor); + ggml_set_name(staging_tensor, state->tensor->name); + staged_tensors.push_back({state, staging_tensor}); + } + + ggml_backend_buffer_t compute_buffer = ggml_backend_alloc_ctx_tensors(staging_ctx, compute_backend); + if (compute_buffer == nullptr) { + LOG_ERROR("model manager alloc compute params backend buffer failed, num_tensors = %zu", + staged_tensors.size()); + ggml_free(staging_ctx); + return false; + } + ggml_backend_buffer_set_usage(compute_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + + for (auto& staged_tensor : staged_tensors) { + TensorState* state = staged_tensor.first; + ggml_tensor* managed_tensor = state->tensor; + ggml_tensor* staging_tensor = staged_tensor.second; + ggml_backend_tensor_copy(managed_tensor, staging_tensor); + std::swap(managed_tensor->buffer, staging_tensor->buffer); + std::swap(managed_tensor->data, staging_tensor->data); + std::swap(managed_tensor->extra, staging_tensor->extra); + } + ggml_backend_synchronize(compute_backend); + + auto block = std::make_unique(); + block->compute_backend = compute_backend; + block->buffer = compute_buffer; + block->staging_ctx = staging_ctx; + block->staged_tensors = std::move(staged_tensors); + for (auto& staged_tensor : block->staged_tensors) { + TensorState* state = staged_tensor.first; + state->staged_to_compute_backend = true; + } + compute_staging_blocks_.push_back(std::move(block)); + + int64_t t1 = ggml_time_ms(); + LOG_DEBUG("model manager staged compute params (%6.2f MB, %zu tensors) to %s, taking %.2fs", + ggml_backend_buffer_get_size(compute_buffer) / (1024.f * 1024.f), + states.size(), + ggml_backend_name(compute_backend), + (t1 - t0) * 1.0f / 1000); + } + + return true; +} + +bool ModelManager::apply_loras_to_params(const std::vector& states) { + if (loras_.empty()) { + return true; + } + + struct LoraApplyGroup { + std::map model_tensors; + std::vector states; + }; + + std::map groups; + for (TensorState* state : states) { + if (state == nullptr || state->tensor == nullptr || + should_ignore(*state) || is_optional_missing_tensor(state->name)) { + continue; + } + if (state->applied_lora_epoch == current_lora_epoch_) { + continue; + } + if (state->compute_backend == nullptr) { + LOG_ERROR("model manager compute backend is null for lora target tensor '%s'", state->name.c_str()); + return false; + } + if (state->tensor->data == nullptr) { + LOG_ERROR("model manager lora target tensor '%s' is not prepared", state->name.c_str()); + return false; + } + LoraApplyGroup& group = groups[state->compute_backend]; + group.model_tensors[state->name] = state->tensor; + group.states.push_back(state); + } + + if (groups.empty()) { + return true; + } + + std::set all_tensor_names = tensor_names(); + for (auto& group_pair : groups) { + ggml_backend_t compute_backend = group_pair.first; + LoraApplyGroup& group = group_pair.second; + for (const LoraSpec& lora_spec : loras_) { + if (group.model_tensors.empty()) { + continue; + } + + std::string id = lora_id(lora_spec); + auto lora = std::make_shared(id, + compute_backend, + compute_backend, + lora_spec.path, + lora_spec.is_high_noise ? "model.high_noise_" : "", + lora_version_); + + LoraModel::filter_t lora_tensor_filter = nullptr; + if (!lora_spec.tensor_name_prefix_filter.empty()) { + lora_tensor_filter = [&](const std::string& tensor_name) { + return starts_with(tensor_name, lora_spec.tensor_name_prefix_filter); + }; + } + if (!lora->load_from_file(n_threads_, lora_tensor_filter)) { + LOG_WARN("load lora tensors from %s failed", lora_spec.path.c_str()); + if (lora_spec.required) { + return false; + } + continue; + } + if (lora->lora_tensors.empty()) { + if (lora_spec.required) { + LOG_ERROR("required lora has no tensors: %s", lora_spec.path.c_str()); + return false; + } + continue; + } + lora->multiplier = lora_spec.multiplier; + lora->apply(group.model_tensors, all_tensor_names, lora_version_, n_threads_, false); + lora->release_loaded_tensors(); + } + + for (TensorState* state : group.states) { + if (state != nullptr) { + state->applied_lora_epoch = current_lora_epoch_; + } + } + } + return true; +} + +void ModelManager::reset_lora_applied_params() { + release_compute_staging_blocks(true); + release_params_storage_blocks(true); + for (auto& state : tensor_states_) { + state->applied_lora_epoch = UINT64_MAX; + } +} + +bool ModelManager::should_ignore(const TensorState& state) const { + for (const auto& ignore_prefix : common_ignore_tensors_) { + if (starts_with(state.name, ignore_prefix)) { + return true; + } + } + return false; +} + +bool ModelManager::is_optional_missing_tensor(const std::string& name) const { + return name.find("cond_stage_model.transformer.text_model.encoder.layers.23") != std::string::npos || + name.find("alphas_cumprod") != std::string::npos; +} + +bool ModelManager::validate_tensor(const TensorState& state) const { + if (state.tensor == nullptr || should_ignore(state) || is_optional_missing_tensor(state.name)) { + return true; + } + + const auto& tensor_storage_map = model_loader_.get_tensor_storage_map(); + auto ts_it = tensor_storage_map.find(state.name); + if (ts_it == tensor_storage_map.end()) { + LOG_ERROR("%s tensor '%s' not in model metadata", state.desc.c_str(), state.name.c_str()); + return false; + } + + const TensorStorage& tensor_storage = ts_it->second; + if (state.tensor->ne[0] != tensor_storage.ne[0] || + state.tensor->ne[1] != tensor_storage.ne[1] || + state.tensor->ne[2] != tensor_storage.ne[2] || + state.tensor->ne[3] != tensor_storage.ne[3]) { + LOG_ERROR( + "%s tensor '%s' has wrong shape in model metadata: got [%d, %d, %d, %d], expected [%d, %d, %d, %d]", + state.desc.c_str(), + state.name.c_str(), + (int)tensor_storage.ne[0], (int)tensor_storage.ne[1], (int)tensor_storage.ne[2], (int)tensor_storage.ne[3], + (int)state.tensor->ne[0], (int)state.tensor->ne[1], (int)state.tensor->ne[2], (int)state.tensor->ne[3]); + return false; + } + return true; +} + +bool ModelManager::mmap_params(const std::vector& states, + std::vector& created_storage_blocks) { + std::map mmap_candidates; + std::map mmap_states; + for (TensorState* state : states) { + if (state == nullptr || !can_mmap_storage(*state) || state->tensor == nullptr || + state->tensor->data != nullptr || state->tensor->view_src != nullptr) { + continue; + } + mmap_candidates[state->name] = state->tensor; + mmap_states[state->name] = state; + } + if (mmap_candidates.empty()) { + return true; + } + + auto mmap_store = model_loader_.mmap_tensors(mmap_candidates, {}, writable_mmap_); + if (mmap_store.empty()) { + return true; + } + + auto block = std::make_unique(); + block->mmap_tensor_stores = std::move(mmap_store); + ParamsStorageBlock* raw = block.get(); + for (const auto& pair : mmap_states) { + TensorState* state = pair.second; + if (state != nullptr && state->tensor != nullptr && state->tensor->data != nullptr) { + block->states.push_back(state); + } + } + + if (!block->states.empty()) { + params_storage_blocks_.push_back(std::move(block)); + created_storage_blocks.push_back(raw); + } + return true; +} + +bool ModelManager::can_mmap_storage(const TensorState& state) const { + if (!enable_mmap_ || state.residency_mode != ResidencyMode::ParamBackend) { + return false; + } + if (state.compute_backend == nullptr || state.params_backend == nullptr) { + return false; + } + return sd_backend_is_cpu(state.compute_backend) || + sd_backend_is_cpu(state.params_backend) || + backend_supports_host_buffer(state.compute_backend); +} + +bool ModelManager::alloc_params_buffers(const std::vector& states, + std::vector& created_storage_blocks) { + std::map, std::vector> states_by_buffer_type; + for (TensorState* state : states) { + if (state == nullptr || state->tensor == nullptr) { + continue; + } + ggml_backend_buffer_type_t params_buft = params_buffer_type_for(*state); + if (params_buft == nullptr) { + return false; + } + states_by_buffer_type[{params_buft, static_cast(state->residency_mode)}].push_back(state); + } + + for (const auto& pair : states_by_buffer_type) { + ggml_backend_buffer_type_t params_buft = pair.first.first; + const std::vector& states = pair.second; + size_t alignment = ggml_backend_buft_get_alignment(params_buft); + size_t max_size = ggml_backend_buft_get_max_size(params_buft); + + auto alloc_chunk = [&](const std::vector& chunk, size_t chunk_size) -> bool { + if (chunk.empty() || chunk_size == 0) { + return true; + } + + ggml_backend_buffer_t buffer = ggml_backend_buft_alloc_buffer(params_buft, chunk_size); + if (buffer == nullptr) { + LOG_ERROR("model manager alloc params backend buffer failed, size = %.2fMB", + chunk_size / (1024.0 * 1024.0)); + return false; + } + ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + + std::vector initialized_tensors; + void* base = ggml_backend_buffer_get_base(buffer); + size_t offset = aligned_offset(base, 0, ggml_backend_buffer_get_alignment(buffer)); + for (TensorState* state : chunk) { + ggml_tensor* tensor = state->tensor; + size_t tensor_size = GGML_PAD(ggml_backend_buffer_get_alloc_size(buffer, tensor), + ggml_backend_buffer_get_alignment(buffer)); + enum ggml_status status = ggml_backend_tensor_alloc(buffer, tensor, static_cast(base) + offset); + if (status != GGML_STATUS_SUCCESS) { + LOG_ERROR("model manager failed to initialize params tensor '%s'", ggml_get_name(tensor)); + for (ggml_tensor* initialized : initialized_tensors) { + initialized->buffer = nullptr; + initialized->data = nullptr; + initialized->extra = nullptr; + } + LOG_DEBUG("model manager releasing params backend buffer (%6.2f MB, %zu tensors, %s)", + ggml_backend_buffer_get_size(buffer) / (1024.f * 1024.f), + initialized_tensors.size(), + ggml_backend_buffer_is_host(buffer) ? "RAM" : "VRAM"); + ggml_backend_buffer_free(buffer); + return false; + } + initialized_tensors.push_back(tensor); + offset += tensor_size; + } + + auto block = std::make_unique(); + block->buffer = buffer; + block->states = chunk; + ParamsStorageBlock* raw = block.get(); + params_storage_blocks_.push_back(std::move(block)); + created_storage_blocks.push_back(raw); + + return true; + }; + + std::vector chunk; + size_t chunk_size = 0; + for (TensorState* state : states) { + ggml_tensor* tensor = state->tensor; + size_t tensor_size = GGML_PAD(ggml_backend_buft_get_alloc_size(params_buft, tensor), alignment); + // Some backends, e.g. Vulkan, report a preferred chunk size here rather than a + // hard per-tensor allocation limit. Oversized tensors are allocated alone. + if (!chunk.empty() && max_size > 0 && chunk_size + tensor_size > max_size) { + if (!alloc_chunk(chunk, chunk_size)) { + return false; + } + chunk.clear(); + chunk_size = 0; + } + chunk.push_back(state); + chunk_size += tensor_size; + } + + if (!alloc_chunk(chunk, chunk_size)) { + return false; + } + } + + return true; +} + +bool ModelManager::load_tensors(const std::vector& states) { + std::map states_by_name; + std::set target_tensor_names; + for (TensorState* state : states) { + if (state == nullptr) { + continue; + } + states_by_name[state->name] = state; + target_tensor_names.insert(state->name); + } + if (states_by_name.empty()) { + return true; + } + + std::set loaded_names; + std::mutex loaded_names_mutex; + auto on_new_tensor_cb = [&](const TensorStorage& tensor_storage, ggml_tensor** dst_tensor) -> bool { + const std::string& name = tensor_storage.name; + *dst_tensor = nullptr; + + auto state_it = states_by_name.find(name); + if (state_it == states_by_name.end()) { + return true; + } + + TensorState* state = state_it->second; + if (state == nullptr || state->tensor == nullptr) { + LOG_ERROR("model manager tensor '%s' is null", name.c_str()); + return false; + } + + if (state->tensor->ne[0] != tensor_storage.ne[0] || + state->tensor->ne[1] != tensor_storage.ne[1] || + state->tensor->ne[2] != tensor_storage.ne[2] || + state->tensor->ne[3] != tensor_storage.ne[3]) { + LOG_ERROR( + "model manager tensor '%s' has wrong shape in model file: got [%d, %d, %d, %d], expected [%d, %d, %d, %d]", + name.c_str(), + (int)tensor_storage.ne[0], (int)tensor_storage.ne[1], (int)tensor_storage.ne[2], (int)tensor_storage.ne[3], + (int)state->tensor->ne[0], (int)state->tensor->ne[1], (int)state->tensor->ne[2], (int)state->tensor->ne[3]); + return false; + } + + { + std::lock_guard lock(loaded_names_mutex); + loaded_names.insert(name); + } + *dst_tensor = state->tensor; + return true; + }; + + if (!model_loader_.load_tensors(on_new_tensor_cb, enable_mmap_, &target_tensor_names)) { + LOG_ERROR("model manager load tensors failed"); + return false; + } + + bool missing = false; + for (const auto& pair : states_by_name) { + const std::string& name = pair.first; + if (loaded_names.find(name) == loaded_names.end()) { + LOG_ERROR("model manager tensor '%s' was not loaded", name.c_str()); + missing = true; + } + } + if (missing) { + return false; + } + + for (const auto& pair : states_by_name) { + pair.second->loaded_to_params_backend = true; + } + return true; +} + +ggml_backend_buffer_type_t ModelManager::params_buffer_type_for(const TensorState& state) const { + if (state.params_backend == nullptr) { + LOG_ERROR("model manager params backend is null for tensor '%s'", state.name.c_str()); + return nullptr; + } + ggml_backend_buffer_type_t params_buft = nullptr; + if (state.compute_backend != nullptr && state.params_backend != state.compute_backend) { + ggml_backend_dev_t compute_dev = ggml_backend_get_device(state.compute_backend); + if (compute_dev != nullptr) { + params_buft = ggml_backend_dev_host_buffer_type(compute_dev); + } + } + if (params_buft == nullptr) { + params_buft = ggml_backend_get_default_buffer_type(state.params_backend); + } + return params_buft; +} + +void ModelManager::free_compute_staging_block(ComputeStagingBlock& block) { + for (auto& staged_tensor : block.staged_tensors) { + TensorState* state = staged_tensor.first; + ggml_tensor* staging_tensor = staged_tensor.second; + if (state == nullptr || state->tensor == nullptr || staging_tensor == nullptr) { + continue; + } + ggml_tensor* managed_tensor = state->tensor; + managed_tensor->buffer = staging_tensor->buffer; + managed_tensor->data = staging_tensor->data; + managed_tensor->extra = staging_tensor->extra; + staging_tensor->buffer = nullptr; + staging_tensor->data = nullptr; + staging_tensor->extra = nullptr; + + state->staged_to_compute_backend = false; + state->applied_lora_epoch = UINT64_MAX; + } + + if (block.buffer != nullptr) { + LOG_DEBUG("model manager releasing compute params (%6.2f MB, %zu tensors) from %s", + ggml_backend_buffer_get_size(block.buffer) / (1024.f * 1024.f), + block.staged_tensors.size(), + block.compute_backend != nullptr ? ggml_backend_name(block.compute_backend) : "unknown"); + ggml_backend_buffer_free(block.buffer); + block.buffer = nullptr; + } + if (block.staging_ctx != nullptr) { + ggml_free(block.staging_ctx); + block.staging_ctx = nullptr; + } + block.staged_tensors.clear(); +} + +void ModelManager::release_compute_staging_blocks(bool force, + const std::unordered_set* target_states) { + for (auto it = compute_staging_blocks_.begin(); it != compute_staging_blocks_.end();) { + ComputeStagingBlock* block = it->get(); + bool can_release = force; + if (!can_release) { + can_release = std::all_of(block->staged_tensors.begin(), + block->staged_tensors.end(), + [target_states](const std::pair& pair) { + TensorState* state = pair.first; + if (state == nullptr) { + return true; + } + if (target_states != nullptr && + target_states->find(state) == target_states->end()) { + return false; + } + return state->active_prepare_count == 0; + }); + } + + if (can_release) { + free_compute_staging_block(*block); + it = compute_staging_blocks_.erase(it); + } else { + ++it; + } + } +} + +void ModelManager::free_params_storage_block(ParamsStorageBlock& block) { + if (block.buffer != nullptr) { + LOG_DEBUG("model manager releasing params backend buffer (%6.2f MB, %zu tensors, %s)", + ggml_backend_buffer_get_size(block.buffer) / (1024.f * 1024.f), + block.states.size(), + ggml_backend_buffer_is_host(block.buffer) ? "RAM" : "VRAM"); + ggml_backend_buffer_free(block.buffer); + block.buffer = nullptr; + } + block.mmap_tensor_stores.clear(); + + for (TensorState* state : block.states) { + if (state == nullptr || state->tensor == nullptr) { + continue; + } + state->tensor->buffer = nullptr; + state->tensor->data = nullptr; + state->tensor->extra = nullptr; + + state->loaded_to_params_backend = false; + state->applied_lora_epoch = UINT64_MAX; + } + block.states.clear(); +} + +void ModelManager::release_params_storage_blocks(bool force, + const std::unordered_set* target_states) { + for (auto it = params_storage_blocks_.begin(); it != params_storage_blocks_.end();) { + ParamsStorageBlock* block = it->get(); + bool can_release = force; + if (!can_release) { + can_release = std::all_of(block->states.begin(), + block->states.end(), + [target_states](TensorState* state) { + if (state == nullptr) { + return true; + } + if (target_states != nullptr && + target_states->find(state) == target_states->end()) { + return false; + } + return state->active_prepare_count == 0 && + !state->staged_to_compute_backend && + state->residency_mode == ResidencyMode::Disk; + }); + } + + if (can_release) { + free_params_storage_block(*block); + it = params_storage_blocks_.erase(it); + } else { + ++it; + } + } +} + +void ModelManager::erase_params_storage_block(ParamsStorageBlock* block) { + auto it = std::find_if(params_storage_blocks_.begin(), + params_storage_blocks_.end(), + [block](const std::unique_ptr& item) { + return item.get() == block; + }); + if (it != params_storage_blocks_.end()) { + params_storage_blocks_.erase(it); + } +} + +void ModelManager::release_all() { + for (auto& state : tensor_states_) { + state->active_prepare_count = 0; + state->applied_lora_epoch = UINT64_MAX; + } + release_compute_staging_blocks(true); + release_params_storage_blocks(true); +} + +bool ModelManager::resolve_required_tensor_states(const std::vector& tensors, + std::vector& required_states) const { + required_states.clear(); + std::unordered_set seen; + for (ggml_tensor* tensor : tensors) { + if (tensor == nullptr) { + continue; + } + const char* raw_name = ggml_get_name(tensor); + if (raw_name == nullptr || raw_name[0] == '\0') { + LOG_ERROR("model manager unnamed tensor is not registered"); + return false; + } + auto state_it = tensor_states_by_name_.find(raw_name); + if (state_it == tensor_states_by_name_.end()) { + LOG_ERROR("model manager tensor '%s' is not registered", raw_name); + return false; + } + TensorState* state = state_it->second; + if (state == nullptr) { + LOG_ERROR("model manager tensor '%s' has no tensor state", raw_name); + return false; + } + if (seen.insert(state).second) { + required_states.push_back(state); + } + } + return true; +} + +bool ModelManager::prepare_params(const std::vector& tensors) { + if (tensors.empty()) { + return true; + } + + std::vector required_states; + if (!resolve_required_tensor_states(tensors, required_states)) { + return false; + } + + if (!load_tensors_to_params_backend(required_states)) { + return false; + } + + if (!stage_tensors_to_compute_backend(required_states)) { + release_compute_staging_blocks(false); + release_params_storage_blocks(false); + return false; + } + + if (!apply_loras_to_params(required_states)) { + release_compute_staging_blocks(false); + release_params_storage_blocks(false); + return false; + } + + for (TensorState* state : required_states) { + if (state == nullptr) { + continue; + } + state->active_prepare_count++; + } + return true; +} + +void ModelManager::finish_compute_backend_usage(const std::vector& states) { + if (states.empty()) { + return; + } + + std::unordered_set target_states; + for (TensorState* state : states) { + if (state == nullptr || !target_states.insert(state).second) { + continue; + } + if (state->active_prepare_count > 0) { + state->active_prepare_count--; + } + } + release_compute_staging_blocks(false, &target_states); +} + +void ModelManager::release_compute_backend_params(const std::vector& tensors) { + if (tensors.empty()) { + return; + } + std::vector required_states; + if (!resolve_required_tensor_states(tensors, required_states)) { + return; + } + finish_compute_backend_usage(required_states); +} + +void ModelManager::release_params_backend_params(const std::vector& tensors) { + if (tensors.empty()) { + return; + } + std::vector required_states; + if (!resolve_required_tensor_states(tensors, required_states)) { + return; + } + if (required_states.empty()) { + return; + } + std::unordered_set target_states(required_states.begin(), required_states.end()); + release_params_storage_blocks(false, &target_states); +} diff --git a/src/model_manager.h b/src/model_manager.h new file mode 100644 index 00000000..9225e3ea --- /dev/null +++ b/src/model_manager.h @@ -0,0 +1,170 @@ +#ifndef __MODEL_MANAGER_H__ +#define __MODEL_MANAGER_H__ + +#include +#include +#include +#include +#include +#include +#include + +#include "model_loader.h" +#include "weight_manager.h" + +class ModelManager : public RunnerWeightManager { +public: + enum class ResidencyMode { + Disk, + ParamBackend, + }; + + struct LoraSpec { + std::string path; + float multiplier = 1.0f; + bool is_high_noise = false; + std::string tensor_name_prefix_filter; + bool required = false; + }; + +private: + struct TensorState { + std::string name; + ggml_tensor* tensor = nullptr; + std::string desc; + + ResidencyMode residency_mode = ResidencyMode::ParamBackend; + ggml_backend_t compute_backend = nullptr; + ggml_backend_t params_backend = nullptr; + bool metadata_validated = false; + + int active_prepare_count = 0; + + bool loaded_to_params_backend = false; + bool staged_to_compute_backend = false; + uint64_t applied_lora_epoch = UINT64_MAX; + }; + + struct ParamsStorageBlock { + ggml_backend_buffer_t buffer = nullptr; + std::vector mmap_tensor_stores; + std::vector states; + }; + + struct ComputeStagingBlock { + ggml_backend_t compute_backend = nullptr; + ggml_backend_buffer_t buffer = nullptr; + ggml_context* staging_ctx = nullptr; + std::vector> staged_tensors; + }; + + ModelLoader model_loader_; + std::vector> tensor_states_; + std::map tensor_states_by_name_; + std::vector> params_storage_blocks_; + std::vector> compute_staging_blocks_; + std::set common_ignore_tensors_; + std::vector loras_; + SDVersion lora_version_ = VERSION_COUNT; + uint64_t current_lora_epoch_ = 0; + int n_threads_ = 0; + bool enable_mmap_ = false; + bool writable_mmap_ = false; + + void finish_compute_backend_usage(const std::vector& states); + void release_all(); + + bool resolve_required_tensor_states(const std::vector& tensors, + std::vector& required_states) const; + bool should_ignore(const TensorState& state) const; + bool is_optional_missing_tensor(const std::string& name) const; + bool validate_tensor(const TensorState& state) const; + + bool load_tensors_to_params_backend(const std::vector& states); + bool apply_loras_to_params(const std::vector& states); + bool mmap_params(const std::vector& states, + std::vector& created_storage_blocks); + bool can_mmap_storage(const TensorState& state) const; + bool alloc_params_buffers(const std::vector& states, + std::vector& created_storage_blocks); + bool load_tensors(const std::vector& states); + bool stage_tensors_to_compute_backend(const std::vector& states); + + ggml_backend_buffer_type_t params_buffer_type_for(const TensorState& state) const; + void release_compute_staging_blocks(bool force = false, + const std::unordered_set* target_states = nullptr); + void release_params_storage_blocks(bool force = false, + const std::unordered_set* target_states = nullptr); + void free_compute_staging_block(ComputeStagingBlock& block); + void free_params_storage_block(ParamsStorageBlock& block); + void erase_params_storage_block(ParamsStorageBlock* block); + void reset_lora_applied_params(); + +public: + ~ModelManager() override; + + ModelLoader& loader() { return model_loader_; } + const ModelLoader& loader() const { return model_loader_; } + + void set_n_threads(int n_threads) { + n_threads_ = n_threads; + model_loader_.set_n_threads(n_threads); + } + void set_enable_mmap(bool enable_mmap) { enable_mmap_ = enable_mmap; } + void set_writable_mmap(bool writable_mmap) { writable_mmap_ = writable_mmap; } + void set_common_ignore_tensors(std::set ignore_tensors); + void set_loras(std::vector loras, SDVersion version); + + std::set tensor_names() const; + + bool register_param_tensors(const std::string& desc, + std::map tensors, + ResidencyMode residency_mode, + ggml_backend_t compute_backend, + ggml_backend_t params_backend, + size_t* registered_tensor_size = nullptr); + + template + bool register_runner_params(const std::string& desc, + Runner& runner, + ResidencyMode residency_mode, + ggml_backend_t compute_backend, + ggml_backend_t params_backend, + size_t* registered_tensor_size = nullptr) { + std::map tensors; + runner.get_param_tensors(tensors); + return register_param_tensors(desc, + std::move(tensors), + residency_mode, + compute_backend, + params_backend, + registered_tensor_size); + } + + template + bool register_runner_params(const std::string& desc, + Runner& runner, + const std::string& prefix, + ResidencyMode residency_mode, + ggml_backend_t compute_backend, + ggml_backend_t params_backend, + size_t* registered_tensor_size = nullptr) { + std::map tensors; + runner.get_param_tensors(tensors, prefix); + return register_param_tensors(desc, + std::move(tensors), + residency_mode, + compute_backend, + params_backend, + registered_tensor_size); + } + + bool validate_registered_tensors(); + bool load_all_params_eagerly(); + + bool prepare_params(const std::vector& tensors) override; + void release_compute_backend_params(const std::vector& tensors) override; + void release_params_backend_params(const std::vector& tensors) override; +}; + +#endif // __MODEL_MANAGER_H__ diff --git a/src/name_conversion.cpp b/src/name_conversion.cpp index a9cae0a8..64b7c681 100644 --- a/src/name_conversion.cpp +++ b/src/name_conversion.cpp @@ -1,8 +1,9 @@ +#include #include #include +#include "core/util.h" #include "name_conversion.h" -#include "util.h" void replace_with_name_map(std::string& name, const std::vector>& name_map) { for (auto kv : name_map) { @@ -183,6 +184,27 @@ std::string convert_cond_stage_model_name(std::string name, std::string prefix) return name; } +std::string convert_qwen3_vl_vision_name(std::string name) { + static const std::vector> qwen3_vl_vision_name_map{ + {"mm.0.", "merger.linear_fc1."}, + {"mm.2.", "merger.linear_fc2."}, + {"v.post_ln.", "merger.norm."}, + {"v.position_embd.weight", "pos_embed.weight"}, + {"v.patch_embd.weight.1", "patch_embed.proj.1.weight"}, + {"v.patch_embd.weight", "patch_embed.proj.0.weight"}, + {"v.patch_embd.bias", "patch_embed.bias"}, + {"v.blk.", "blocks."}, + {"attn_qkv.", "attn.qkv."}, + {"attn_out.", "attn.proj."}, + {"ffn_up.", "mlp.linear_fc1."}, + {"ffn_down.", "mlp.linear_fc2."}, + {"ln1.", "norm1."}, + {"ln2.", "norm2."}, + }; + replace_with_name_map(name, qwen3_vl_vision_name_map); + return name; +} + // ref: https://github.com/huggingface/diffusers/blob/main/scripts/convert_diffusers_to_original_stable_diffusion.py std::string convert_diffusers_unet_to_original_sd1(std::string name) { // (stable-diffusion, HF Diffusers) @@ -682,6 +704,38 @@ std::string convert_other_dit_to_original_anima(std::string name) { return name; } +std::string convert_diffusers_dit_to_original_krea2(std::string name) { + static const std::vector> prefix_map = { + {"img_in.", "first."}, + {"time_embed.linear_1.", "tmlp.0."}, + {"time_embed.linear_2.", "tmlp.2."}, + {"time_mod_proj.", "tproj.1."}, + {"txt_in.linear_1.", "txtmlp.1."}, + {"txt_in.linear_2.", "txtmlp.3."}, + {"text_fusion.", "txtfusion."}, + {"transformer_blocks.", "blocks."}, + {"final_layer.", "last."}, + }; + static const std::vector> name_map = { + {"attn.to_out.0.", "attn.wo."}, + {"attn.to_out.", "attn.wo."}, + {"attn.to_gate.", "attn.gate."}, + {"attn.to_q.", "attn.wq."}, + {"attn.to_k.", "attn.wk."}, + {"attn.to_v.", "attn.wv."}, + {"ff.gate.", "mlp.gate."}, + {"ff.up.", "mlp.up."}, + {"ff.down.", "mlp.down."}, + {"txt_in.norm.", "txtmlp.0."}, + {"last.norm.weight", "last.norm.scale"}, + {"last.modulation.weight", "last.modulation.lin"}, + }; + + replace_with_prefix_map(name, prefix_map); + replace_with_name_map(name, name_map); + return name; +} + std::string convert_diffusion_model_name(std::string name, std::string prefix, SDVersion version) { if (sd_version_is_sd1(version) || sd_version_is_sd2(version)) { name = convert_diffusers_unet_to_original_sd1(name); @@ -689,12 +743,14 @@ 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) || sd_version_is_longcat(version)) { + } else if (sd_version_is_flux(version) || sd_version_is_flux2(version) || sd_version_is_longcat(version) || sd_version_is_sefi_image(version)) { name = convert_diffusers_dit_to_original_flux(name); } else if (sd_version_is_z_image(version)) { name = convert_diffusers_dit_to_original_lumina2(name); } else if (sd_version_is_anima(version)) { name = convert_other_dit_to_original_anima(name); + } else if (sd_version_is_krea2(version)) { + name = convert_diffusers_dit_to_original_krea2(name); } return name; } @@ -794,7 +850,77 @@ std::string convert_diffusers_vae_to_original_sd1(std::string name) { return result; } -std::string convert_first_stage_model_name(std::string name, std::string prefix) { +std::string convert_diffusers_to_original_wan_vae(std::string name) { + static const std::vector> prefix_map = { + {"quant_conv.", "conv1."}, + {"post_quant_conv.", "conv2."}, + + {"decoder.up_blocks.0.resnets.0.", "decoder.upsamples.0.residual."}, + {"decoder.up_blocks.0.resnets.1.", "decoder.upsamples.1.residual."}, + {"decoder.up_blocks.0.resnets.2.", "decoder.upsamples.2.residual."}, + {"decoder.up_blocks.0.upsamplers.0.", "decoder.upsamples.3."}, + + {"decoder.up_blocks.1.resnets.0.conv_shortcut.", "decoder.upsamples.4.shortcut."}, + {"decoder.up_blocks.1.resnets.0.", "decoder.upsamples.4.residual."}, + {"decoder.up_blocks.1.resnets.1.", "decoder.upsamples.5.residual."}, + {"decoder.up_blocks.1.resnets.2.", "decoder.upsamples.6.residual."}, + {"decoder.up_blocks.1.upsamplers.0.", "decoder.upsamples.7."}, + {"decoder.up_blocks.2.resnets.0.", "decoder.upsamples.8.residual."}, + {"decoder.up_blocks.2.resnets.1.", "decoder.upsamples.9.residual."}, + {"decoder.up_blocks.2.resnets.2.", "decoder.upsamples.10.residual."}, + {"decoder.up_blocks.2.upsamplers.0.", "decoder.upsamples.11."}, + {"decoder.up_blocks.3.resnets.0.", "decoder.upsamples.12.residual."}, + {"decoder.up_blocks.3.resnets.1.", "decoder.upsamples.13.residual."}, + {"decoder.up_blocks.3.resnets.2.", "decoder.upsamples.14.residual."}, + + {"encoder.down_blocks.0.", "encoder.downsamples.0.residual."}, + {"encoder.down_blocks.1.", "encoder.downsamples.1.residual."}, + {"encoder.down_blocks.2.", "encoder.downsamples.2."}, + {"encoder.down_blocks.3.conv_shortcut.", "encoder.downsamples.3.shortcut."}, + {"encoder.down_blocks.3.", "encoder.downsamples.3.residual."}, + {"encoder.down_blocks.4.", "encoder.downsamples.4.residual."}, + {"encoder.down_blocks.5.", "encoder.downsamples.5."}, + {"encoder.down_blocks.6.conv_shortcut.", "encoder.downsamples.6.shortcut."}, + {"encoder.down_blocks.6.", "encoder.downsamples.6.residual."}, + {"encoder.down_blocks.7.", "encoder.downsamples.7.residual."}, + {"encoder.down_blocks.8.", "encoder.downsamples.8."}, + {"encoder.down_blocks.9.", "encoder.downsamples.9.residual."}, + {"encoder.down_blocks.10.", "encoder.downsamples.10.residual."}, + }; + + static const std::vector> shared_name_map = { + {".conv_in.", ".conv1."}, + {".norm_out.", ".head.0."}, + {".conv_out.", ".head.2."}, + + {".mid_block.attentions.0.", ".middle.1."}, + {".mid_block.resnets.0.", ".middle.0.residual."}, + {".mid_block.resnets.1.", ".middle.2.residual."}, + }; + + static const std::vector> resnet_name_map = { + {".norm1.", ".0."}, + {".conv1.", ".2."}, + {".norm2.", ".3."}, + {".conv2.", ".6."}, + }; + + replace_with_name_map(name, shared_name_map); + replace_with_prefix_map(name, prefix_map); + + // Only apply the ResNet-specific renaming if the tensor belongs to a ResNet block. + // This prevents generic ".conv1." or ".conv2." matching on top-level encoder/decoder convolutions. + if (name.find(".residual.") != std::string::npos) { + replace_with_name_map(name, resnet_name_map); + } + + return name; +} + +std::string convert_first_stage_model_name(std::string name, std::string prefix, SDVersion version) { + if (sd_version_uses_wan_vae(version)) { + return convert_diffusers_to_original_wan_vae(name); + } static std::unordered_map vae_name_map = { {"decoder.post_quant_conv.", "post_quant_conv."}, {"encoder.quant_conv.", "quant_conv."}, @@ -989,7 +1115,46 @@ bool is_first_stage_model_name(const std::string& name) { return false; } +static std::string convert_esrgan_tensor_name(std::string name) { + static std::unordered_map esrgan_name_map; + + if (esrgan_name_map.empty()) { + esrgan_name_map["model.0."] = "conv_first."; + + constexpr int max_num_blocks = 64; + for (int i = 0; i < max_num_blocks; i++) { + std::string block_prefix = "model.1.sub." + std::to_string(i) + "."; + for (int rdb = 1; rdb <= 3; rdb++) { + for (int conv = 1; conv <= 5; conv++) { + esrgan_name_map[block_prefix + "RDB" + std::to_string(rdb) + ".conv" + std::to_string(conv) + ".0."] = + "body." + std::to_string(i) + ".rdb" + std::to_string(rdb) + ".conv" + std::to_string(conv) + "."; + } + } + esrgan_name_map[block_prefix + "weight"] = "conv_body.weight"; + esrgan_name_map[block_prefix + "bias"] = "conv_body.bias"; + } + + // RealESRGAN stores only the learned layers in a Sequential. These indices + // cover the common x1, x2 and x4 layouts. + esrgan_name_map["model.2."] = "conv_hr."; + esrgan_name_map["model.3."] = "conv_up1."; + esrgan_name_map["model.4."] = "conv_last."; + esrgan_name_map["model.5."] = "conv_hr."; + esrgan_name_map["model.6."] = "conv_up2."; + esrgan_name_map["model.7."] = "conv_last."; + esrgan_name_map["model.8."] = "conv_hr."; + esrgan_name_map["model.10."] = "conv_last."; + } + + replace_with_prefix_map(name, esrgan_name_map); + return name; +} + std::string convert_tensor_name(std::string name, SDVersion version) { + if (version == VERSION_ESRGAN) { + return convert_esrgan_tensor_name(std::move(name)); + } + bool is_lora = false; bool is_lycoris_underline = false; bool is_underline = false; @@ -1114,6 +1279,10 @@ std::string convert_tensor_name(std::string name, SDVersion version) { replace_with_prefix_map(name, prefix_map); + if ((sd_version_is_boogu_image(version) || sd_version_is_krea2(version)) && starts_with(name, "text_encoders.llm.visual.")) { + name = convert_qwen3_vl_vision_name(std::move(name)); + } + // diffusion model { for (const auto& prefix : diffuison_model_prefix_vec) { @@ -1140,7 +1309,7 @@ std::string convert_tensor_name(std::string name, SDVersion version) { { for (const auto& prefix : first_stage_model_prefix_vec) { if (starts_with(name, prefix)) { - name = convert_first_stage_model_name(name.substr(prefix.size()), prefix); + name = convert_first_stage_model_name(name.substr(prefix.size()), prefix, version); if (version == VERSION_SDXS_512_DS || version == VERSION_SDXS_09) { name = "tae." + name; } else { diff --git a/src/cache_dit.hpp b/src/runtime/cache_dit.hpp similarity index 99% rename from src/cache_dit.hpp rename to src/runtime/cache_dit.hpp index dad67d45..bec6e811 100644 --- a/src/cache_dit.hpp +++ b/src/runtime/cache_dit.hpp @@ -1,5 +1,5 @@ -#ifndef __CACHE_DIT_HPP__ -#define __CACHE_DIT_HPP__ +#ifndef __SD_RUNTIME_CACHE_DIT_HPP__ +#define __SD_RUNTIME_CACHE_DIT_HPP__ #include #include @@ -8,9 +8,9 @@ #include #include -#include "condition_cache_utils.hpp" -#include "ggml_extend.hpp" -#include "tensor.hpp" +#include "core/ggml_extend.hpp" +#include "core/tensor.hpp" +#include "runtime/condition_cache_utils.hpp" struct DBCacheConfig { bool enabled = false; @@ -893,4 +893,4 @@ struct CacheDitConditionState { } }; -#endif +#endif // __SD_RUNTIME_CACHE_DIT_HPP__ diff --git a/src/condition_cache_utils.hpp b/src/runtime/condition_cache_utils.hpp similarity index 91% rename from src/condition_cache_utils.hpp rename to src/runtime/condition_cache_utils.hpp index 903d64e3..8003df12 100644 --- a/src/condition_cache_utils.hpp +++ b/src/runtime/condition_cache_utils.hpp @@ -1,9 +1,9 @@ -#ifndef __CONDITION_CACHE_UTILS_HPP__ -#define __CONDITION_CACHE_UTILS_HPP__ +#ifndef __SD_RUNTIME_CONDITION_CACHE_UTILS_HPP__ +#define __SD_RUNTIME_CONDITION_CACHE_UTILS_HPP__ #include -#include "tensor.hpp" +#include "core/tensor.hpp" namespace sd { @@ -61,4 +61,4 @@ namespace sd { } // namespace sd -#endif // __CONDITION_CACHE_UTILS_HPP__ +#endif // __SD_RUNTIME_CONDITION_CACHE_UTILS_HPP__ diff --git a/src/denoiser.hpp b/src/runtime/denoiser.hpp similarity index 75% rename from src/denoiser.hpp rename to src/runtime/denoiser.hpp index 365b1dbd..812eebe6 100644 --- a/src/denoiser.hpp +++ b/src/runtime/denoiser.hpp @@ -1,5 +1,5 @@ -#ifndef __DENOISER_HPP__ -#define __DENOISER_HPP__ +#ifndef __SD_RUNTIME_DENOISER_HPP__ +#define __SD_RUNTIME_DENOISER_HPP__ #include #include @@ -8,10 +8,10 @@ #include #include -#include "ggml_extend.hpp" -#include "gits_noise.inl" -#include "guidance.h" -#include "tensor.hpp" +#include "core/ggml_extend.hpp" +#include "core/tensor.hpp" +#include "runtime/gits_noise.h" +#include "runtime/guidance.h" /*================================================= CompVisDenoiser ==================================================*/ @@ -302,6 +302,137 @@ struct KarrasScheduler : SigmaScheduler { } }; +struct BetaScheduler : SigmaScheduler { + static constexpr double alpha = 0.6; + static constexpr double beta = 0.6; + + static double log_beta(double a, double b) { + return std::lgamma(a) + std::lgamma(b) - std::lgamma(a + b); + } + + static double incbeta(double x, double a, double b) { + if (x <= 0.0) { + return 0.0; + } + if (x >= 1.0) { + return 1.0; + } + + // Continued fraction approximation using Lentz's method. + const int max_iter = 200; + const double epsilon = 3.0e-7; + const double tiny = 1e-30; + + const double qab = a + b; + const double qap = a + 1.0; + const double qam = a - 1.0; + + double c = 1.0; + double d = 1.0 - qab * x / qap; + if (std::abs(d) < tiny) { + d = tiny; + } + d = 1.0 / d; + double h = d; + + for (int m = 1; m <= max_iter; m++) { + const int m2 = 2 * m; + + double aa = m * (b - m) * x / ((qam + m2) * (a + m2)); + d = 1.0 + aa * d; + if (std::abs(d) < tiny) { + d = tiny; + } + c = 1.0 + aa / c; + if (std::abs(c) < tiny) { + c = tiny; + } + d = 1.0 / d; + h *= d * c; + + aa = -(a + m) * (qab + m) * x / ((a + m2) * (qap + m2)); + d = 1.0 + aa * d; + if (std::abs(d) < tiny) { + d = tiny; + } + c = 1.0 + aa / c; + if (std::abs(c) < tiny) { + c = tiny; + } + d = 1.0 / d; + const double del = d * c; + h *= del; + + if (std::abs(del - 1.0) < epsilon) { + break; + } + } + + return std::exp(a * std::log(x) + b * std::log(1.0 - x) - log_beta(a, b)) / a * h; + } + + static double beta_cdf(double x, double a, double b) { + if (x == 0.0) { + return 0.0; + } + if (x == 1.0) { + return 1.0; + } + if (x < (a + 1.0) / (a + b + 2.0)) { + return incbeta(x, a, b); + } + return 1.0 - incbeta(1.0 - x, b, a); + } + + static double beta_ppf(double u, double a, double b, int max_iter = 30) { + double x = 0.5; + for (int i = 0; i < max_iter; i++) { + const double f = beta_cdf(x, a, b) - u; + if (std::abs(f) < 1e-10) { + break; + } + const double df = std::exp((a - 1.0) * std::log(x) + (b - 1.0) * std::log(1.0 - x) - log_beta(a, b)); + x -= f / df; + if (x <= 0.0) { + x = 1e-10; + } + if (x >= 1.0) { + x = 1.0 - 1e-10; + } + } + return x; + } + + std::vector get_sigmas(uint32_t n, float /*sigma_min*/, float /*sigma_max*/, t_to_sigma_t t_to_sigma) override { + std::vector result; + result.reserve(n + 1); + + const int t_max = TIMESTEPS - 1; + if (n == 0) { + return result; + } else if (n == 1) { + result.push_back(t_to_sigma(static_cast(t_max))); + result.push_back(0.f); + return result; + } + + int last_t = -1; + for (uint32_t i = 0; i < n; i++) { + const double u = 1.0 - static_cast(i) / static_cast(n); + const double t_cont = beta_ppf(u, alpha, beta) * t_max; + const int t = static_cast(std::lround(t_cont)); + + if (t != last_t) { + result.push_back(t_to_sigma(static_cast(t))); + last_t = t; + } + } + + result.push_back(0.f); + return result; + } +}; + struct SimpleScheduler : SigmaScheduler { std::vector get_sigmas(uint32_t n, float sigma_min, float sigma_max, t_to_sigma_t t_to_sigma) override { std::vector result_sigmas; @@ -559,6 +690,318 @@ struct LTX2Scheduler : SigmaScheduler { } }; +inline float flux_time_shift(float mu, float sigma, float t) { + return ::expf(mu) / (::expf(mu) + ::powf((1.0f / t - 1.0f), sigma)); +} + +// https://github.com/black-forest-labs/flux/blob/main/src/flux/sampling.py#L289 +struct FluxScheduler : SigmaScheduler { + int image_seq_len = 0; + float base_shift = 0.5f; + float max_shift = 1.15f; + + explicit FluxScheduler(int image_seq_len, const char* extra_sample_args = nullptr) + : image_seq_len(image_seq_len) { + parse_extra_sample_args(extra_sample_args); + } + + void parse_extra_sample_args(const char* extra_sample_args) { + for (const auto& [key, value] : parse_key_value_args(extra_sample_args, "flux scheduler arg")) { + if (key == "base_shift") { + if (!parse_strict_float(value, base_shift)) { + LOG_WARN("ignoring invalid flux scheduler arg '%s=%s'", key.c_str(), value.c_str()); + } + } else if (key == "max_shift") { + if (!parse_strict_float(value, max_shift)) { + LOG_WARN("ignoring invalid flux scheduler arg '%s=%s'", key.c_str(), value.c_str()); + } + } + } + } + + float compute_mu() const { + constexpr float base_shift_anchor = 256.0f; + constexpr float max_shift_anchor = 4096.0f; + float m = (max_shift - base_shift) / (max_shift_anchor - base_shift_anchor); + float b = base_shift - m * base_shift_anchor; + return static_cast(image_seq_len) * m + b; + } + + std::vector get_sigmas(uint32_t n, float /*sigma_min*/, float /*sigma_max*/, t_to_sigma_t /*t_to_sigma*/) override { + std::vector sigmas; + sigmas.reserve(n + 1); + + float mu = compute_mu(); + LOG_DEBUG("Flux scheduler: image_seq_len=%d, steps=%u, mu=%.3f", image_seq_len, n, mu); + + if (n == 0) { + sigmas.push_back(1.0f); + return sigmas; + } + + for (uint32_t i = 0; i <= n; ++i) { + float t = 1.0f - static_cast(i) / static_cast(n); + if (t <= 0.0f) { + sigmas.push_back(0.0f); + } else { + sigmas.push_back(flux_time_shift(mu, 1.0f, t)); + } + } + + sigmas[n] = 0.0f; + return sigmas; + } +}; + +// https://github.com/black-forest-labs/flux2/blob/main/src/flux2/sampling.py#L244 +struct Flux2Scheduler : SigmaScheduler { + int image_seq_len = 0; + + explicit Flux2Scheduler(int image_seq_len) + : image_seq_len(image_seq_len) {} + + static float compute_empirical_mu(int image_seq_len, uint32_t num_steps) { + const float a1 = 8.73809524e-05f; + const float b1 = 1.89833333f; + const float a2 = 0.00016927f; + const float b2 = 0.45666666f; + + if (image_seq_len > 4300) { + return a2 * image_seq_len + b2; + } + + float m_200 = a2 * image_seq_len + b2; + float m_10 = a1 * image_seq_len + b1; + + float a = (m_200 - m_10) / 190.0f; + float b = m_200 - 200.0f * a; + return a * num_steps + b; + } + + std::vector get_sigmas(uint32_t n, float /*sigma_min*/, float /*sigma_max*/, t_to_sigma_t /*t_to_sigma*/) override { + std::vector sigmas; + sigmas.reserve(n + 1); + + float mu = compute_empirical_mu(image_seq_len, n); + LOG_DEBUG("Flux2 scheduler: image_seq_len=%d, steps=%u, mu=%.3f", image_seq_len, n, mu); + + if (n == 0) { + sigmas.push_back(1.0f); + return sigmas; + } + + for (uint32_t i = 0; i <= n; ++i) { + float t = 1.0f - static_cast(i) / static_cast(n); + if (t <= 0.0f) { + sigmas.push_back(0.0f); + } else if (t >= 1.0f) { + sigmas.push_back(1.0f); + } else { + sigmas.push_back(flux_time_shift(mu, 1.0f, t)); + } + } + + sigmas[n] = 0.0f; + return sigmas; + } +}; + +/* + * Logit-Normal Scheduler + * Based on: https://github.com/ideogram-oss/ideogram4/blob/main/src/ideogram4/scheduler.py + */ +struct LogitNormalScheduler : SigmaScheduler { + float mean = 0.0f; + float std = 1.75f; + float logsnr_min = -15.0f; + float logsnr_max = 18.0f; + + bool resolution_aware = true; + + float one_minus_t_min, one_minus_t_max; + + void parse_extra_sample_args(int image_seq_len = 0, const char* extra_sample_args = nullptr) { + const int known_seq_len = (512 * 512) / (16 * 16); + if (extra_sample_args) { + for (const auto& [key, value] : parse_key_value_args(extra_sample_args, "logit-normal scheduler arg")) { + if (key == "mu") { + if (!parse_strict_float(value, mean)) { + LOG_WARN("ignoring invalid logit-normal scheduler arg '%s=%s'", key.c_str(), value.c_str()); + } + } else if (key == "std") { + if (!parse_strict_float(value, std)) { + LOG_WARN("ignoring invalid logit-normal scheduler arg '%s=%s'", key.c_str(), value.c_str()); + } + } + if (key == "logsnr_min") { + if (!parse_strict_float(value, logsnr_min)) { + LOG_WARN("ignoring invalid logit-normal scheduler arg '%s=%s'", key.c_str(), value.c_str()); + } + } else if (key == "logsnr_max") { + if (!parse_strict_float(value, logsnr_max)) { + LOG_WARN("ignoring invalid logit-normal scheduler arg '%s=%s'", key.c_str(), value.c_str()); + } + } else if (key == "resolution_aware") { + if (!parse_strict_bool(value, resolution_aware)) { + LOG_WARN("ignoring invalid logit-normal scheduler arg '%s=%s'", key.c_str(), value.c_str()); + } + } + } + } + if (image_seq_len > 0 && resolution_aware) { + mean += 0.5f * std::log(static_cast(image_seq_len) / static_cast(known_seq_len)); + } + } + + float sigmoid(float x) { + return 1.0f / (1.0f + std::exp(-x)); + } + + LogitNormalScheduler(float mean = 0.0f, float std = 1.75f, float logsnr_min = -18.0f, float logsnr_max = 15.0f) + : mean(mean), std(std), logsnr_min(logsnr_min), logsnr_max(logsnr_max) { + // t_min = 1.0f / (1.0f + std::exp(0.5f * logsnr_max)); + one_minus_t_min = sigmoid(0.5f * logsnr_max); + // t_max = 1.0f / (1.0f + std::exp(0.5f * logsnr_min)); + one_minus_t_max = sigmoid(0.5f * logsnr_min); + } + + LogitNormalScheduler(int image_seq_len = 0, const char* extra_sample_args = nullptr) { + mean = 0.0f; + std = 1.75f; + logsnr_min = -15.0f; + logsnr_max = 18.0f; + + parse_extra_sample_args(image_seq_len, extra_sample_args); + // t_min = 1.0f / (1.0f + std::exp(0.5f * logsnr_max)); + one_minus_t_min = sigmoid(0.5f * logsnr_max); + // t_max = 1.0f / (1.0f + std::exp(0.5f * logsnr_min)); + one_minus_t_max = sigmoid(0.5f * logsnr_min); + } + + // https://stackedboxes.org/2017/05/01/acklams-normal-quantile-function/ + double ndtri(double p) { + if (p <= 0.0) { + return -std::numeric_limits::infinity(); + } else if (p >= 1.0) { + return std::numeric_limits::infinity(); + } + + static const double p_low = 0.02425; + static const double p_high = 1.0 - p_low; + + static const double c[6] = {-7.784894002430293e-03, + -3.223964580411365e-01, + -2.400758277161838e+00, + -2.549732539343734e+00, + 4.374664141464968e+00, + 2.938163982698783e+00}; + + static const double d[5] = {7.784695709041462e-03, + 3.224671290700398e-01, + 2.445134137142996e+00, + 3.754408661907416e+00, + 1.0}; + + // Coefficients for the central region + static const double a[6] = {-3.969683028665376e+01, + 2.209460984245205e+02, + -2.759285104469687e+02, + 1.383577518672690e+02, + -3.066479806614716e+01, + 2.506628277459239e+00}; + + static const double b[6] = {-5.447609879822406e+01, + 1.615858368580409e+02, + -1.556989798598866e+02, + 6.680131188771972e+01, + -1.328068155288572e+01, + 1.0}; + + double x = 0.0; + + if (p < p_low) { + // Lower region + double q = std::sqrt(-2.0 * std::log(p)); + + // Numerator: c[0]*q^5 + c[1]*q^4 + ... + c[5] + double numerator = c[0]; + for (int i = 1; i < 6; ++i) { + numerator = numerator * q + c[i]; + } + + // Denominator: d[0]*q^4 + d[1]*q^3 + ... + d[3]*q + 1 + double denominator = d[0]; + for (int i = 1; i < 5; ++i) { + denominator = denominator * q + d[i]; + } + + x = numerator / denominator; + } else if (p > p_high) { + // Upper region + double q = std::sqrt(-2.0 * std::log(1.0 - p)); + + double numerator = c[0]; + for (int i = 1; i < 6; ++i) { + numerator = numerator * q + c[i]; + } + + double denominator = d[0]; + for (int i = 1; i < 5; ++i) { + denominator = denominator * q + d[i]; + } + + x = -(numerator / denominator); + } else { + // Central region + double q = p - 0.5; + double r = q * q; + + // Numerator: (a[0]*r^5 + a[1]*r^4 + ... + a[5])*q + double numerator = a[0]; + for (int i = 1; i < 6; ++i) { + numerator = numerator * r + a[i]; + } + numerator *= q; + + // Denominator: b[0]*r^4 + b[1]*r^3 + ... + b[4]*r + 1 + double denominator = b[0]; + for (int i = 1; i < 6; ++i) { + denominator = denominator * r + b[i]; + } + + x = numerator / denominator; + } + return x; + } + + std::vector get_sigmas(uint32_t n, float /*sigma_min*/, float /*sigma_max*/, t_to_sigma_t /*t_to_sigma*/) override { + std::vector sigmas; + LOG_INFO("LOGIT_NORMAL_SCHEDULER using mean=%.4f, std=%.4f, logsnr_min=%.4f, logsnr_max=%.4f", mean, std, logsnr_min, logsnr_max); + sigmas.reserve(n + 1); + for (uint32_t i = 0; i <= n; ++i) { + float t = static_cast(i) / static_cast(n); + + // ndtri(1-t) == -ndtri(t) + float z = static_cast(-ndtri(t)); + + float y = mean + std * z; + + float timestep = sigmoid(y); + + if (timestep > one_minus_t_min) + timestep = one_minus_t_min; + if (timestep < one_minus_t_max) + timestep = one_minus_t_max; + + float sigma = timestep; + + sigmas.push_back(sigma); + } + sigmas[n] = 0.0f; + return sigmas; + } +}; + struct Denoiser { virtual float sigma_min() = 0; virtual float sigma_max() = 0; @@ -583,6 +1026,10 @@ struct Denoiser { LOG_INFO("get_sigmas with Karras scheduler"); scheduler = std::make_shared(); break; + case BETA_SCHEDULER: + LOG_INFO("get_sigmas with Beta scheduler"); + scheduler = std::make_shared(); + break; case EXPONENTIAL_SCHEDULER: LOG_INFO("get_sigmas exponential scheduler"); scheduler = std::make_shared(); @@ -623,6 +1070,21 @@ struct Denoiser { LOG_INFO("get_sigmas with LTX2 scheduler"); scheduler = std::make_shared(image_seq_len, extra_sample_args); break; + case LOGIT_NORMAL_SCHEDULER: { + LOG_INFO("get_sigmas with Logit-Normal scheduler"); + scheduler = std::make_shared(image_seq_len, extra_sample_args); + break; + } + case FLUX2_SCHEDULER: { + LOG_INFO("get_sigmas with Flux2 scheduler"); + scheduler = std::make_shared(image_seq_len); + break; + } + case FLUX_SCHEDULER: { + LOG_INFO("get_sigmas with Flux scheduler"); + scheduler = std::make_shared(image_seq_len, extra_sample_args); + break; + } default: LOG_INFO("get_sigmas with discrete scheduler (default)"); scheduler = std::make_shared(); @@ -787,10 +1249,6 @@ struct DiscreteFlowDenoiser : public Denoiser { } }; -inline float flux_time_shift(float mu, float sigma, float t) { - return ::expf(mu) / (::expf(mu) + ::powf((1.0f / t - 1.0f), sigma)); -} - struct FluxFlowDenoiser : public DiscreteFlowDenoiser { FluxFlowDenoiser() = default; @@ -804,35 +1262,141 @@ struct FluxFlowDenoiser : public DiscreteFlowDenoiser { } }; -struct Flux2FlowDenoiser : public FluxFlowDenoiser { - Flux2FlowDenoiser() = default; +struct SefiFlowDenoiser; - float compute_empirical_mu(uint32_t n, int image_seq_len) { - const float a1 = 8.73809524e-05f; - const float b1 = 1.89833333f; - const float a2 = 0.00016927f; - const float b2 = 0.45666666f; +struct SefiFlowDenoiser : public FluxFlowDenoiser { + static constexpr int kNumTrainTimesteps = 1000; + static constexpr int kSemChannels = 16; + static constexpr int kTotalChannels = 144; - if (image_seq_len > 4300) { - float mu = a2 * image_seq_len + b2; - return mu; + float delta_t = 0.1f; + float timestep_shift_alpha = 1.0f; + + std::vector sem_sigmas; + std::vector tex_sigmas; + std::vector sem_timesteps; + std::vector tex_timesteps; + + SefiFlowDenoiser() = default; + + static float apply_alpha_shift(float u_unit, float alpha) { + if (alpha == 1.0f) { + return u_unit; + } + float denom = 1.0f + (alpha - 1.0f) * u_unit; + return (alpha * u_unit) / denom; + } + + std::vector get_sigmas(uint32_t n, + int image_seq_len, + scheduler_t scheduler_type, + SDVersion version, + const char* extra_sample_args = nullptr) override { + sem_sigmas.clear(); + tex_sigmas.clear(); + sem_timesteps.clear(); + tex_timesteps.clear(); + + for (const auto& [key, value] : parse_key_value_args(extra_sample_args, "sefi scheduler arg")) { + if (key == "sefi_alpha") { + if (!parse_strict_float(value, timestep_shift_alpha)) { + LOG_WARN("ignoring invalid sefi scheduler arg '%s=%s'", key.c_str(), value.c_str()); + } + } else if (key == "sefi_delta_t") { + if (!parse_strict_float(value, delta_t)) { + LOG_WARN("ignoring invalid sefi scheduler arg '%s=%s'", key.c_str(), value.c_str()); + } + } } - float m_200 = a2 * image_seq_len + b2; - float m_10 = a1 * image_seq_len + b1; + for (uint32_t i = 0; i <= n; ++i) { + float u_base = static_cast(i) / static_cast(n); + float u_shifted = apply_alpha_shift(u_base, timestep_shift_alpha); + float u_sem_raw = u_shifted * (1.0f + delta_t); - float a = (m_200 - m_10) / 190.0f; - float b = m_200 - 200.0f * a; - float mu = a * n + b; + float u_sem = std::min(u_sem_raw, 1.0f); + float u_tex = std::max(0.0f, std::min(u_sem_raw - delta_t, 1.0f)); - return mu; + int idx_sem = std::min(kNumTrainTimesteps - 1, + std::max(0, static_cast(u_sem * (kNumTrainTimesteps - 1)))); + int idx_tex = std::min(kNumTrainTimesteps - 1, + std::max(0, static_cast(u_tex * (kNumTrainTimesteps - 1)))); + + float t_sem = static_cast(kNumTrainTimesteps - idx_sem); + float t_tex = static_cast(kNumTrainTimesteps - idx_tex); + float sigma_sem = t_sem / static_cast(kNumTrainTimesteps); + float sigma_tex = t_tex / static_cast(kNumTrainTimesteps); + + sem_timesteps.push_back(t_sem); + tex_timesteps.push_back(t_tex); + sem_sigmas.push_back(sigma_sem); + tex_sigmas.push_back(sigma_tex); + } + LOG_DEBUG("SefiFlowDenoiser: built %u-step dual schedule (alpha=%.2f delta_t=%.2f)", + n, timestep_shift_alpha, delta_t); + return tex_sigmas; + } +}; + +// MiniT2I predicts x0 directly and integrates a linear flow ODE: +// x_{t+dt} = x_t + (x0 - x_t)/(1 - t) * dt, t in [0, 1), x0 = start = noise * 2. +// Mapping sigma = 1 - t makes the generic Euler update +// x += (x - denoised)/sigma * (sigma_next - sigma) +// exactly reproduce that step when denoised == x0. To make the generic +// `denoised = pred * c_out + x * c_skip` yield x0 from the model's raw x0 +// prediction we use c_skip = 0, c_out = 1, c_in = 1. Sigmas run linearly 1 -> 0. +struct MiniT2IFlowDenoiser : public Denoiser { + float sigma_min() override { + return 0.0f; + } + + float sigma_max() override { + return 1.0f; + } + + float sigma_to_t(float sigma) override { + return 1.0f - sigma; + } + + float t_to_sigma(float t) override { + return 1.0f - t; + } + + std::vector get_scalings(float sigma) override { + SD_UNUSED(sigma); + float c_skip = 0.0f; + float c_out = 1.0f; + float c_in = 1.0f; + return {c_skip, c_out, c_in}; + } + + sd::Tensor noise_scaling(float sigma, + const sd::Tensor& noise, + const sd::Tensor& latent) override { + SD_UNUSED(sigma); + SD_UNUSED(latent); + // Sampling starts from x0_init = noise * 2 (see MiniT2I reference). + return noise * 2.0f; + } + + sd::Tensor inverse_noise_scaling(float sigma, const sd::Tensor& latent) override { + SD_UNUSED(sigma); + return latent; } std::vector get_sigmas(uint32_t n, int image_seq_len, scheduler_t scheduler_type, SDVersion version, const char* extra_sample_args = nullptr) override { - float mu = compute_empirical_mu(n, image_seq_len); - LOG_DEBUG("Flux2FlowDenoiser: set shift to %.3f", mu); - set_shift(mu); - return Denoiser::get_sigmas(n, image_seq_len, scheduler_type, version, extra_sample_args); + SD_UNUSED(image_seq_len); + SD_UNUSED(scheduler_type); + SD_UNUSED(version); + SD_UNUSED(extra_sample_args); + // Uniform t schedule 0 -> 1 => sigma 1 -> 0, matching the reference loop. + std::vector sigmas; + sigmas.reserve(n + 1); + for (uint32_t i = 0; i < n; ++i) { + sigmas.push_back(1.0f - static_cast(i) / static_cast(n)); + } + sigmas.push_back(0.0f); + return sigmas; } }; @@ -939,6 +1503,40 @@ static sd::Tensor sample_euler_ancestral(denoise_cb_t model, return x; } +static sd::Tensor sample_sefi_euler(SefiFlowDenoiser* sefi, + denoise_cb_t model, + sd::Tensor x) { + const std::vector& sigma_tex_vec = sefi->tex_sigmas; + const std::vector& sigma_sem_vec = sefi->sem_sigmas; + int steps = static_cast(sigma_tex_vec.size()) - 1; + for (int i = 0; i < steps; i++) { + float sigma_tex_cur = sigma_tex_vec[i]; + float sigma_tex_next = sigma_tex_vec[i + 1]; + float sigma_sem_cur = sigma_sem_vec[i]; + float sigma_sem_next = sigma_sem_vec[i + 1]; + if (sigma_tex_cur <= 1e-9f) { + continue; + } + auto denoised_opt = model(x, sigma_tex_cur, i + 1); + if (denoised_opt.pred.empty()) { + return {}; + } + sd::Tensor denoised = std::move(denoised_opt.pred); + sd::Tensor velocity = (x - denoised) / sigma_tex_cur; + + auto x_sem = sd::ops::slice(x, 2, 0, SefiFlowDenoiser::kSemChannels); + auto x_tex = sd::ops::slice(x, 2, SefiFlowDenoiser::kSemChannels, SefiFlowDenoiser::kTotalChannels); + auto vel_sem = sd::ops::slice(velocity, 2, 0, SefiFlowDenoiser::kSemChannels); + auto vel_tex = sd::ops::slice(velocity, 2, SefiFlowDenoiser::kSemChannels, SefiFlowDenoiser::kTotalChannels); + auto x_sem_next = x_sem + vel_sem * (sigma_sem_next - sigma_sem_cur); + auto x_tex_next = x_tex + vel_tex * (sigma_tex_next - sigma_tex_cur); + + sd::ops::slice_assign(&x, 2, 0, SefiFlowDenoiser::kSemChannels, x_sem_next); + sd::ops::slice_assign(&x, 2, SefiFlowDenoiser::kSemChannels, SefiFlowDenoiser::kTotalChannels, x_tex_next); + } + return x; +} + static sd::Tensor sample_euler(denoise_cb_t model, sd::Tensor x, const std::vector& sigmas) { @@ -1854,7 +2452,13 @@ static sd::Tensor sample_k_diffusion(sample_method_t method, std::shared_ptr rng, float eta, bool is_flow_denoiser, - const char* extra_sample_args) { + const char* extra_sample_args, + std::shared_ptr denoiser_for_dispatch = nullptr) { + if (denoiser_for_dispatch) { + if (auto sefi = std::dynamic_pointer_cast(denoiser_for_dispatch)) { + return sample_sefi_euler(sefi.get(), model, std::move(x)); + } + } SamplerExtraArgs extra_args = parse_key_value_args(extra_sample_args, "extra sample arg"); switch (method) { case EULER_A_SAMPLE_METHOD: @@ -1902,4 +2506,4 @@ static sd::Tensor sample_k_diffusion(sample_method_t method, } } -#endif // __DENOISER_HPP__ +#endif // __SD_RUNTIME_DENOISER_HPP__ diff --git a/src/easycache.hpp b/src/runtime/easycache.hpp similarity index 97% rename from src/easycache.hpp rename to src/runtime/easycache.hpp index 409a464e..75ae3dda 100644 --- a/src/easycache.hpp +++ b/src/runtime/easycache.hpp @@ -1,15 +1,15 @@ -#ifndef __EASYCACHE_HPP__ -#define __EASYCACHE_HPP__ +#ifndef __SD_RUNTIME_EASYCACHE_HPP__ +#define __SD_RUNTIME_EASYCACHE_HPP__ #include #include #include #include -#include "condition_cache_utils.hpp" -#include "denoiser.hpp" -#include "ggml_extend.hpp" -#include "tensor.hpp" +#include "core/ggml_extend.hpp" +#include "core/tensor.hpp" +#include "runtime/condition_cache_utils.hpp" +#include "runtime/denoiser.hpp" struct EasyCacheConfig { bool enabled = false; @@ -258,4 +258,4 @@ struct EasyCacheState { } }; -#endif +#endif // __SD_RUNTIME_EASYCACHE_HPP__ diff --git a/src/runtime/gits_noise.h b/src/runtime/gits_noise.h new file mode 100644 index 00000000..10b18df3 --- /dev/null +++ b/src/runtime/gits_noise.h @@ -0,0 +1,333 @@ +#ifndef __SD_RUNTIME_GITS_NOISE_H__ +#define __SD_RUNTIME_GITS_NOISE_H__ + +const std::vector> GITS_NOISE_0_80 = { + {14.61464119f, 7.49001646f, 0.02916753f}, + {14.61464119f, 11.54541874f, 6.77309084f, 0.02916753f}, + {14.61464119f, 11.54541874f, 7.49001646f, 3.07277966f, 0.02916753f}, + {14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 2.05039096f, 0.02916753f}, + {14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.85520077f, 2.05039096f, 0.02916753f}, + {14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.85520077f, 3.07277966f, 1.56271636f, 0.02916753f}, + {14.61464119f, 12.96784878f, 11.54541874f, 8.75849152f, 7.49001646f, 5.85520077f, 3.07277966f, 1.56271636f, 0.02916753f}, + {14.61464119f, 13.76078796f, 12.23089790f, 10.90732002f, 8.75849152f, 7.49001646f, 5.85520077f, 3.07277966f, 1.56271636f, 0.02916753f}, + {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 10.90732002f, 8.75849152f, 7.49001646f, 5.85520077f, 3.07277966f, 1.56271636f, 0.02916753f}, + {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 5.85520077f, 3.07277966f, 1.56271636f, 0.02916753f}, + {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.07277966f, 1.56271636f, 0.02916753f}, + {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.07277966f, 1.56271636f, 0.02916753f}, + {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.07277966f, 1.56271636f, 0.02916753f}, + {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.07277966f, 1.56271636f, 0.02916753f}, + {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.75859547f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.19567990f, 1.98035145f, 0.86115354f, 0.02916753f}, + {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.75859547f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.19567990f, 1.98035145f, 0.86115354f, 0.02916753f}, + {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.75859547f, 9.24142551f, 8.75849152f, 8.30717278f, 7.88507891f, 7.49001646f, 6.77309084f, 5.85520077f, 4.65472794f, 3.07277966f, 1.84880662f, 0.83188516f, 0.02916753f}}; + +const std::vector> GITS_NOISE_0_85 = { + {14.61464119f, 7.49001646f, 0.02916753f}, + {14.61464119f, 7.49001646f, 1.84880662f, 0.02916753f}, + {14.61464119f, 11.54541874f, 6.77309084f, 1.56271636f, 0.02916753f}, + {14.61464119f, 11.54541874f, 7.11996698f, 3.07277966f, 1.24153244f, 0.02916753f}, + {14.61464119f, 11.54541874f, 7.49001646f, 5.09240818f, 2.84484982f, 0.95350921f, 0.02916753f}, + {14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.09240818f, 2.84484982f, 0.95350921f, 0.02916753f}, + {14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.58536053f, 3.19567990f, 1.84880662f, 0.803307f, 0.02916753f}, + {14.61464119f, 12.96784878f, 11.54541874f, 8.75849152f, 7.49001646f, 5.58536053f, 3.19567990f, 1.84880662f, 0.803307f, 0.02916753f}, + {14.61464119f, 12.96784878f, 11.54541874f, 8.75849152f, 7.49001646f, 6.14220476f, 4.65472794f, 3.07277966f, 1.84880662f, 0.803307f, 0.02916753f}, + {14.61464119f, 13.76078796f, 12.23089790f, 10.90732002f, 8.75849152f, 7.49001646f, 6.14220476f, 4.65472794f, 3.07277966f, 1.84880662f, 0.803307f, 0.02916753f}, + {14.61464119f, 13.76078796f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.65472794f, 3.07277966f, 1.84880662f, 0.803307f, 0.02916753f}, + {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.65472794f, 3.07277966f, 1.84880662f, 0.803307f, 0.02916753f}, + {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.65472794f, 3.07277966f, 1.84880662f, 0.803307f, 0.02916753f}, + {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.60512662f, 2.63833880f, 1.56271636f, 0.72133851f, 0.02916753f}, + {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 4.65472794f, 3.46139455f, 2.45070267f, 1.56271636f, 0.72133851f, 0.02916753f}, + {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 4.65472794f, 3.46139455f, 2.45070267f, 1.56271636f, 0.72133851f, 0.02916753f}, + {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 4.65472794f, 3.46139455f, 2.45070267f, 1.56271636f, 0.72133851f, 0.02916753f}, + {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.75859547f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 4.65472794f, 3.46139455f, 2.45070267f, 1.56271636f, 0.72133851f, 0.02916753f}, + {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.90732002f, 10.31284904f, 9.75859547f, 9.24142551f, 8.75849152f, 8.30717278f, 7.88507891f, 7.49001646f, 6.77309084f, 5.85520077f, 4.65472794f, 3.46139455f, 2.45070267f, 1.56271636f, 0.72133851f, 0.02916753f}}; + +const std::vector> GITS_NOISE_0_90 = { + {14.61464119f, 6.77309084f, 0.02916753f}, + {14.61464119f, 7.49001646f, 1.56271636f, 0.02916753f}, + {14.61464119f, 7.49001646f, 3.07277966f, 0.95350921f, 0.02916753f}, + {14.61464119f, 7.49001646f, 4.86714602f, 2.54230714f, 0.89115214f, 0.02916753f}, + {14.61464119f, 11.54541874f, 7.49001646f, 4.86714602f, 2.54230714f, 0.89115214f, 0.02916753f}, + {14.61464119f, 11.54541874f, 7.49001646f, 5.09240818f, 3.07277966f, 1.61558151f, 0.69515091f, 0.02916753f}, + {14.61464119f, 12.23089790f, 8.75849152f, 7.11996698f, 4.86714602f, 3.07277966f, 1.61558151f, 0.69515091f, 0.02916753f}, + {14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.85520077f, 4.45427561f, 2.95596409f, 1.61558151f, 0.69515091f, 0.02916753f}, + {14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.19988537f, 1.24153244f, 0.57119018f, 0.02916753f}, + {14.61464119f, 12.96784878f, 10.90732002f, 8.75849152f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.19988537f, 1.24153244f, 0.57119018f, 0.02916753f}, + {14.61464119f, 12.96784878f, 11.54541874f, 9.24142551f, 8.30717278f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.19988537f, 1.24153244f, 0.57119018f, 0.02916753f}, + {14.61464119f, 12.96784878f, 11.54541874f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.75677586f, 2.84484982f, 1.84880662f, 1.08895338f, 0.52423614f, 0.02916753f}, + {14.61464119f, 13.76078796f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 4.86714602f, 3.75677586f, 2.84484982f, 1.84880662f, 1.08895338f, 0.52423614f, 0.02916753f}, + {14.61464119f, 13.76078796f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.44769001f, 5.58536053f, 4.45427561f, 3.32507086f, 2.45070267f, 1.61558151f, 0.95350921f, 0.45573691f, 0.02916753f}, + {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.44769001f, 5.58536053f, 4.45427561f, 3.32507086f, 2.45070267f, 1.61558151f, 0.95350921f, 0.45573691f, 0.02916753f}, + {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 4.86714602f, 3.91689563f, 3.07277966f, 2.27973175f, 1.56271636f, 0.95350921f, 0.45573691f, 0.02916753f}, + {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 4.86714602f, 3.91689563f, 3.07277966f, 2.27973175f, 1.56271636f, 0.95350921f, 0.45573691f, 0.02916753f}, + {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 4.86714602f, 3.91689563f, 3.07277966f, 2.27973175f, 1.56271636f, 0.95350921f, 0.45573691f, 0.02916753f}, + {14.61464119f, 13.76078796f, 12.96784878f, 12.23089790f, 11.54541874f, 10.31284904f, 9.24142551f, 8.75849152f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 5.09240818f, 4.45427561f, 3.60512662f, 2.95596409f, 2.19988537f, 1.51179266f, 0.89115214f, 0.43325692f, 0.02916753f}}; + +const std::vector> GITS_NOISE_0_95 = { + {14.61464119f, 6.77309084f, 0.02916753f}, + {14.61464119f, 6.77309084f, 1.56271636f, 0.02916753f}, + {14.61464119f, 7.49001646f, 2.84484982f, 0.89115214f, 0.02916753f}, + {14.61464119f, 7.49001646f, 4.86714602f, 2.36326075f, 0.803307f, 0.02916753f}, + {14.61464119f, 7.49001646f, 4.86714602f, 2.95596409f, 1.56271636f, 0.64427125f, 0.02916753f}, + {14.61464119f, 11.54541874f, 7.49001646f, 4.86714602f, 2.95596409f, 1.56271636f, 0.64427125f, 0.02916753f}, + {14.61464119f, 11.54541874f, 7.49001646f, 4.86714602f, 3.07277966f, 1.91321158f, 1.08895338f, 0.50118381f, 0.02916753f}, + {14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.07277966f, 1.91321158f, 1.08895338f, 0.50118381f, 0.02916753f}, + {14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.85520077f, 4.45427561f, 3.07277966f, 1.91321158f, 1.08895338f, 0.50118381f, 0.02916753f}, + {14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.19988537f, 1.41535246f, 0.803307f, 0.38853383f, 0.02916753f}, + {14.61464119f, 12.23089790f, 8.75849152f, 7.49001646f, 5.85520077f, 4.65472794f, 3.46139455f, 2.63833880f, 1.84880662f, 1.24153244f, 0.72133851f, 0.34370604f, 0.02916753f}, + {14.61464119f, 12.96784878f, 10.90732002f, 8.75849152f, 7.49001646f, 5.85520077f, 4.65472794f, 3.46139455f, 2.63833880f, 1.84880662f, 1.24153244f, 0.72133851f, 0.34370604f, 0.02916753f}, + {14.61464119f, 12.96784878f, 10.90732002f, 8.75849152f, 7.49001646f, 6.14220476f, 4.86714602f, 3.75677586f, 2.95596409f, 2.19988537f, 1.56271636f, 1.05362725f, 0.64427125f, 0.32104823f, 0.02916753f}, + {14.61464119f, 12.96784878f, 10.90732002f, 8.75849152f, 7.49001646f, 6.44769001f, 5.58536053f, 4.65472794f, 3.60512662f, 2.95596409f, 2.19988537f, 1.56271636f, 1.05362725f, 0.64427125f, 0.32104823f, 0.02916753f}, + {14.61464119f, 12.96784878f, 11.54541874f, 9.24142551f, 8.30717278f, 7.49001646f, 6.44769001f, 5.58536053f, 4.65472794f, 3.60512662f, 2.95596409f, 2.19988537f, 1.56271636f, 1.05362725f, 0.64427125f, 0.32104823f, 0.02916753f}, + {14.61464119f, 12.96784878f, 11.54541874f, 9.24142551f, 8.30717278f, 7.49001646f, 6.44769001f, 5.58536053f, 4.65472794f, 3.75677586f, 3.07277966f, 2.45070267f, 1.78698075f, 1.24153244f, 0.83188516f, 0.50118381f, 0.22545385f, 0.02916753f}, + {14.61464119f, 12.96784878f, 11.54541874f, 9.24142551f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 5.09240818f, 4.45427561f, 3.60512662f, 2.95596409f, 2.36326075f, 1.72759056f, 1.24153244f, 0.83188516f, 0.50118381f, 0.22545385f, 0.02916753f}, + {14.61464119f, 13.76078796f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 5.09240818f, 4.45427561f, 3.60512662f, 2.95596409f, 2.36326075f, 1.72759056f, 1.24153244f, 0.83188516f, 0.50118381f, 0.22545385f, 0.02916753f}, + {14.61464119f, 13.76078796f, 12.23089790f, 10.90732002f, 9.24142551f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 5.09240818f, 4.45427561f, 3.75677586f, 3.07277966f, 2.45070267f, 1.91321158f, 1.46270394f, 1.05362725f, 0.72133851f, 0.43325692f, 0.19894916f, 0.02916753f}}; + +const std::vector> GITS_NOISE_1_00 = { + {14.61464119f, 1.56271636f, 0.02916753f}, + {14.61464119f, 6.77309084f, 0.95350921f, 0.02916753f}, + {14.61464119f, 6.77309084f, 2.36326075f, 0.803307f, 0.02916753f}, + {14.61464119f, 7.11996698f, 3.07277966f, 1.56271636f, 0.59516323f, 0.02916753f}, + {14.61464119f, 7.49001646f, 4.86714602f, 2.84484982f, 1.41535246f, 0.57119018f, 0.02916753f}, + {14.61464119f, 7.49001646f, 4.86714602f, 2.84484982f, 1.61558151f, 0.86115354f, 0.38853383f, 0.02916753f}, + {14.61464119f, 11.54541874f, 7.49001646f, 4.86714602f, 2.84484982f, 1.61558151f, 0.86115354f, 0.38853383f, 0.02916753f}, + {14.61464119f, 11.54541874f, 7.49001646f, 4.86714602f, 3.07277966f, 1.98035145f, 1.24153244f, 0.72133851f, 0.34370604f, 0.02916753f}, + {14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.07277966f, 1.98035145f, 1.24153244f, 0.72133851f, 0.34370604f, 0.02916753f}, + {14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.27973175f, 1.51179266f, 0.95350921f, 0.54755926f, 0.25053367f, 0.02916753f}, + {14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.36326075f, 1.61558151f, 1.08895338f, 0.72133851f, 0.41087446f, 0.17026083f, 0.02916753f}, + {14.61464119f, 11.54541874f, 8.75849152f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.36326075f, 1.61558151f, 1.08895338f, 0.72133851f, 0.41087446f, 0.17026083f, 0.02916753f}, + {14.61464119f, 11.54541874f, 8.75849152f, 7.49001646f, 5.85520077f, 4.65472794f, 3.60512662f, 2.84484982f, 2.12350607f, 1.56271636f, 1.08895338f, 0.72133851f, 0.41087446f, 0.17026083f, 0.02916753f}, + {14.61464119f, 11.54541874f, 8.75849152f, 7.49001646f, 5.85520077f, 4.65472794f, 3.60512662f, 2.84484982f, 2.19988537f, 1.61558151f, 1.162866f, 0.803307f, 0.50118381f, 0.27464288f, 0.09824532f, 0.02916753f}, + {14.61464119f, 11.54541874f, 8.75849152f, 7.49001646f, 5.85520077f, 4.65472794f, 3.75677586f, 3.07277966f, 2.45070267f, 1.84880662f, 1.36964464f, 1.01931262f, 0.72133851f, 0.45573691f, 0.25053367f, 0.09824532f, 0.02916753f}, + {14.61464119f, 11.54541874f, 8.75849152f, 7.49001646f, 6.14220476f, 5.09240818f, 4.26497746f, 3.46139455f, 2.84484982f, 2.19988537f, 1.67050016f, 1.24153244f, 0.92192322f, 0.64427125f, 0.43325692f, 0.25053367f, 0.09824532f, 0.02916753f}, + {14.61464119f, 11.54541874f, 8.75849152f, 7.49001646f, 6.14220476f, 5.09240818f, 4.26497746f, 3.60512662f, 2.95596409f, 2.45070267f, 1.91321158f, 1.51179266f, 1.12534678f, 0.83188516f, 0.59516323f, 0.38853383f, 0.22545385f, 0.09824532f, 0.02916753f}, + {14.61464119f, 12.23089790f, 9.24142551f, 8.30717278f, 7.49001646f, 6.14220476f, 5.09240818f, 4.26497746f, 3.60512662f, 2.95596409f, 2.45070267f, 1.91321158f, 1.51179266f, 1.12534678f, 0.83188516f, 0.59516323f, 0.38853383f, 0.22545385f, 0.09824532f, 0.02916753f}, + {14.61464119f, 12.23089790f, 9.24142551f, 8.30717278f, 7.49001646f, 6.77309084f, 5.85520077f, 5.09240818f, 4.26497746f, 3.60512662f, 2.95596409f, 2.45070267f, 1.91321158f, 1.51179266f, 1.12534678f, 0.83188516f, 0.59516323f, 0.38853383f, 0.22545385f, 0.09824532f, 0.02916753f}}; + +const std::vector> GITS_NOISE_1_05 = { + {14.61464119f, 0.95350921f, 0.02916753f}, + {14.61464119f, 6.77309084f, 0.89115214f, 0.02916753f}, + {14.61464119f, 6.77309084f, 2.05039096f, 0.72133851f, 0.02916753f}, + {14.61464119f, 6.77309084f, 2.84484982f, 1.28281462f, 0.52423614f, 0.02916753f}, + {14.61464119f, 6.77309084f, 3.07277966f, 1.61558151f, 0.803307f, 0.34370604f, 0.02916753f}, + {14.61464119f, 7.49001646f, 4.86714602f, 2.84484982f, 1.56271636f, 0.803307f, 0.34370604f, 0.02916753f}, + {14.61464119f, 7.49001646f, 4.86714602f, 2.84484982f, 1.61558151f, 0.95350921f, 0.52423614f, 0.22545385f, 0.02916753f}, + {14.61464119f, 7.49001646f, 4.86714602f, 3.07277966f, 1.98035145f, 1.24153244f, 0.74807048f, 0.41087446f, 0.17026083f, 0.02916753f}, + {14.61464119f, 7.49001646f, 4.86714602f, 3.19567990f, 2.27973175f, 1.51179266f, 0.95350921f, 0.59516323f, 0.34370604f, 0.13792117f, 0.02916753f}, + {14.61464119f, 7.49001646f, 5.09240818f, 3.46139455f, 2.45070267f, 1.61558151f, 1.08895338f, 0.72133851f, 0.45573691f, 0.25053367f, 0.09824532f, 0.02916753f}, + {14.61464119f, 11.54541874f, 7.49001646f, 5.09240818f, 3.46139455f, 2.45070267f, 1.61558151f, 1.08895338f, 0.72133851f, 0.45573691f, 0.25053367f, 0.09824532f, 0.02916753f}, + {14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.36326075f, 1.61558151f, 1.08895338f, 0.72133851f, 0.45573691f, 0.25053367f, 0.09824532f, 0.02916753f}, + {14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.45070267f, 1.72759056f, 1.24153244f, 0.86115354f, 0.59516323f, 0.38853383f, 0.22545385f, 0.09824532f, 0.02916753f}, + {14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.65472794f, 3.60512662f, 2.84484982f, 2.19988537f, 1.61558151f, 1.162866f, 0.83188516f, 0.59516323f, 0.38853383f, 0.22545385f, 0.09824532f, 0.02916753f}, + {14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.65472794f, 3.60512662f, 2.84484982f, 2.19988537f, 1.67050016f, 1.28281462f, 0.95350921f, 0.72133851f, 0.52423614f, 0.34370604f, 0.19894916f, 0.09824532f, 0.02916753f}, + {14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.65472794f, 3.60512662f, 2.95596409f, 2.36326075f, 1.84880662f, 1.41535246f, 1.08895338f, 0.83188516f, 0.61951244f, 0.45573691f, 0.32104823f, 0.19894916f, 0.09824532f, 0.02916753f}, + {14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.65472794f, 3.60512662f, 2.95596409f, 2.45070267f, 1.91321158f, 1.51179266f, 1.20157266f, 0.95350921f, 0.74807048f, 0.57119018f, 0.43325692f, 0.29807833f, 0.19894916f, 0.09824532f, 0.02916753f}, + {14.61464119f, 11.54541874f, 8.30717278f, 7.11996698f, 5.85520077f, 4.65472794f, 3.60512662f, 2.95596409f, 2.45070267f, 1.91321158f, 1.51179266f, 1.20157266f, 0.95350921f, 0.74807048f, 0.57119018f, 0.43325692f, 0.29807833f, 0.19894916f, 0.09824532f, 0.02916753f}, + {14.61464119f, 11.54541874f, 8.30717278f, 7.11996698f, 5.85520077f, 4.65472794f, 3.60512662f, 2.95596409f, 2.45070267f, 1.98035145f, 1.61558151f, 1.32549286f, 1.08895338f, 0.86115354f, 0.69515091f, 0.54755926f, 0.41087446f, 0.29807833f, 0.19894916f, 0.09824532f, 0.02916753f}}; + +const std::vector> GITS_NOISE_1_10 = { + {14.61464119f, 0.89115214f, 0.02916753f}, + {14.61464119f, 2.36326075f, 0.72133851f, 0.02916753f}, + {14.61464119f, 5.85520077f, 1.61558151f, 0.57119018f, 0.02916753f}, + {14.61464119f, 6.77309084f, 2.45070267f, 1.08895338f, 0.45573691f, 0.02916753f}, + {14.61464119f, 6.77309084f, 2.95596409f, 1.56271636f, 0.803307f, 0.34370604f, 0.02916753f}, + {14.61464119f, 6.77309084f, 3.07277966f, 1.61558151f, 0.89115214f, 0.4783645f, 0.19894916f, 0.02916753f}, + {14.61464119f, 6.77309084f, 3.07277966f, 1.84880662f, 1.08895338f, 0.64427125f, 0.34370604f, 0.13792117f, 0.02916753f}, + {14.61464119f, 7.49001646f, 4.86714602f, 2.84484982f, 1.61558151f, 0.95350921f, 0.54755926f, 0.27464288f, 0.09824532f, 0.02916753f}, + {14.61464119f, 7.49001646f, 4.86714602f, 2.95596409f, 1.91321158f, 1.24153244f, 0.803307f, 0.4783645f, 0.25053367f, 0.09824532f, 0.02916753f}, + {14.61464119f, 7.49001646f, 4.86714602f, 3.07277966f, 2.05039096f, 1.41535246f, 0.95350921f, 0.64427125f, 0.41087446f, 0.22545385f, 0.09824532f, 0.02916753f}, + {14.61464119f, 7.49001646f, 4.86714602f, 3.19567990f, 2.27973175f, 1.61558151f, 1.12534678f, 0.803307f, 0.54755926f, 0.36617002f, 0.22545385f, 0.09824532f, 0.02916753f}, + {14.61464119f, 7.49001646f, 4.86714602f, 3.32507086f, 2.45070267f, 1.72759056f, 1.24153244f, 0.89115214f, 0.64427125f, 0.45573691f, 0.32104823f, 0.19894916f, 0.09824532f, 0.02916753f}, + {14.61464119f, 7.49001646f, 5.09240818f, 3.60512662f, 2.84484982f, 2.05039096f, 1.51179266f, 1.08895338f, 0.803307f, 0.59516323f, 0.43325692f, 0.29807833f, 0.19894916f, 0.09824532f, 0.02916753f}, + {14.61464119f, 7.49001646f, 5.09240818f, 3.60512662f, 2.84484982f, 2.12350607f, 1.61558151f, 1.24153244f, 0.95350921f, 0.72133851f, 0.54755926f, 0.41087446f, 0.29807833f, 0.19894916f, 0.09824532f, 0.02916753f}, + {14.61464119f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.45070267f, 1.84880662f, 1.41535246f, 1.08895338f, 0.83188516f, 0.64427125f, 0.50118381f, 0.36617002f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f}, + {14.61464119f, 7.49001646f, 5.85520077f, 4.45427561f, 3.19567990f, 2.45070267f, 1.91321158f, 1.51179266f, 1.20157266f, 0.95350921f, 0.74807048f, 0.59516323f, 0.45573691f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f}, + {14.61464119f, 7.49001646f, 5.85520077f, 4.45427561f, 3.46139455f, 2.84484982f, 2.19988537f, 1.72759056f, 1.36964464f, 1.08895338f, 0.86115354f, 0.69515091f, 0.54755926f, 0.43325692f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f}, + {14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.46139455f, 2.84484982f, 2.19988537f, 1.72759056f, 1.36964464f, 1.08895338f, 0.86115354f, 0.69515091f, 0.54755926f, 0.43325692f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f}, + {14.61464119f, 11.54541874f, 7.49001646f, 5.85520077f, 4.45427561f, 3.46139455f, 2.84484982f, 2.19988537f, 1.72759056f, 1.36964464f, 1.08895338f, 0.89115214f, 0.72133851f, 0.59516323f, 0.4783645f, 0.38853383f, 0.29807833f, 0.22545385f, 0.17026083f, 0.09824532f, 0.02916753f}}; + +const std::vector> GITS_NOISE_1_15 = { + {14.61464119f, 0.83188516f, 0.02916753f}, + {14.61464119f, 1.84880662f, 0.59516323f, 0.02916753f}, + {14.61464119f, 5.85520077f, 1.56271636f, 0.52423614f, 0.02916753f}, + {14.61464119f, 5.85520077f, 1.91321158f, 0.83188516f, 0.34370604f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.45070267f, 1.24153244f, 0.59516323f, 0.25053367f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.84484982f, 1.51179266f, 0.803307f, 0.41087446f, 0.17026083f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.84484982f, 1.56271636f, 0.89115214f, 0.50118381f, 0.25053367f, 0.09824532f, 0.02916753f}, + {14.61464119f, 6.77309084f, 3.07277966f, 1.84880662f, 1.12534678f, 0.72133851f, 0.43325692f, 0.22545385f, 0.09824532f, 0.02916753f}, + {14.61464119f, 6.77309084f, 3.07277966f, 1.91321158f, 1.24153244f, 0.803307f, 0.52423614f, 0.34370604f, 0.19894916f, 0.09824532f, 0.02916753f}, + {14.61464119f, 7.49001646f, 4.86714602f, 2.95596409f, 1.91321158f, 1.24153244f, 0.803307f, 0.52423614f, 0.34370604f, 0.19894916f, 0.09824532f, 0.02916753f}, + {14.61464119f, 7.49001646f, 4.86714602f, 3.07277966f, 2.05039096f, 1.36964464f, 0.95350921f, 0.69515091f, 0.4783645f, 0.32104823f, 0.19894916f, 0.09824532f, 0.02916753f}, + {14.61464119f, 7.49001646f, 4.86714602f, 3.07277966f, 2.12350607f, 1.51179266f, 1.08895338f, 0.803307f, 0.59516323f, 0.43325692f, 0.29807833f, 0.19894916f, 0.09824532f, 0.02916753f}, + {14.61464119f, 7.49001646f, 4.86714602f, 3.07277966f, 2.12350607f, 1.51179266f, 1.08895338f, 0.803307f, 0.59516323f, 0.45573691f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f}, + {14.61464119f, 7.49001646f, 4.86714602f, 3.07277966f, 2.19988537f, 1.61558151f, 1.24153244f, 0.95350921f, 0.74807048f, 0.59516323f, 0.45573691f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f}, + {14.61464119f, 7.49001646f, 4.86714602f, 3.19567990f, 2.45070267f, 1.78698075f, 1.32549286f, 1.01931262f, 0.803307f, 0.64427125f, 0.50118381f, 0.38853383f, 0.29807833f, 0.22545385f, 0.17026083f, 0.09824532f, 0.02916753f}, + {14.61464119f, 7.49001646f, 4.86714602f, 3.19567990f, 2.45070267f, 1.78698075f, 1.32549286f, 1.01931262f, 0.803307f, 0.64427125f, 0.52423614f, 0.41087446f, 0.32104823f, 0.25053367f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 7.49001646f, 4.86714602f, 3.19567990f, 2.45070267f, 1.84880662f, 1.41535246f, 1.12534678f, 0.89115214f, 0.72133851f, 0.59516323f, 0.4783645f, 0.38853383f, 0.32104823f, 0.25053367f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 7.49001646f, 4.86714602f, 3.19567990f, 2.45070267f, 1.84880662f, 1.41535246f, 1.12534678f, 0.89115214f, 0.72133851f, 0.59516323f, 0.50118381f, 0.41087446f, 0.34370604f, 0.29807833f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}}; + +const std::vector> GITS_NOISE_1_20 = { + {14.61464119f, 0.803307f, 0.02916753f}, + {14.61464119f, 1.56271636f, 0.52423614f, 0.02916753f}, + {14.61464119f, 2.36326075f, 0.92192322f, 0.36617002f, 0.02916753f}, + {14.61464119f, 2.84484982f, 1.24153244f, 0.59516323f, 0.25053367f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.05039096f, 0.95350921f, 0.45573691f, 0.17026083f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.45070267f, 1.24153244f, 0.64427125f, 0.29807833f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.45070267f, 1.36964464f, 0.803307f, 0.45573691f, 0.25053367f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.84484982f, 1.61558151f, 0.95350921f, 0.59516323f, 0.36617002f, 0.19894916f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.84484982f, 1.67050016f, 1.08895338f, 0.74807048f, 0.50118381f, 0.32104823f, 0.19894916f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.95596409f, 1.84880662f, 1.24153244f, 0.83188516f, 0.59516323f, 0.41087446f, 0.27464288f, 0.17026083f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 3.07277966f, 1.98035145f, 1.36964464f, 0.95350921f, 0.69515091f, 0.50118381f, 0.36617002f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f}, + {14.61464119f, 6.77309084f, 3.46139455f, 2.36326075f, 1.56271636f, 1.08895338f, 0.803307f, 0.59516323f, 0.45573691f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f}, + {14.61464119f, 6.77309084f, 3.46139455f, 2.45070267f, 1.61558151f, 1.162866f, 0.86115354f, 0.64427125f, 0.50118381f, 0.38853383f, 0.29807833f, 0.22545385f, 0.17026083f, 0.09824532f, 0.02916753f}, + {14.61464119f, 7.49001646f, 4.65472794f, 3.07277966f, 2.12350607f, 1.51179266f, 1.08895338f, 0.83188516f, 0.64427125f, 0.50118381f, 0.38853383f, 0.29807833f, 0.22545385f, 0.17026083f, 0.09824532f, 0.02916753f}, + {14.61464119f, 7.49001646f, 4.65472794f, 3.07277966f, 2.12350607f, 1.51179266f, 1.08895338f, 0.83188516f, 0.64427125f, 0.50118381f, 0.41087446f, 0.32104823f, 0.25053367f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 7.49001646f, 4.65472794f, 3.07277966f, 2.12350607f, 1.51179266f, 1.08895338f, 0.83188516f, 0.64427125f, 0.50118381f, 0.41087446f, 0.34370604f, 0.27464288f, 0.22545385f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 7.49001646f, 4.65472794f, 3.07277966f, 2.19988537f, 1.61558151f, 1.20157266f, 0.92192322f, 0.72133851f, 0.57119018f, 0.45573691f, 0.36617002f, 0.29807833f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 7.49001646f, 4.65472794f, 3.07277966f, 2.19988537f, 1.61558151f, 1.24153244f, 0.95350921f, 0.74807048f, 0.59516323f, 0.4783645f, 0.38853383f, 0.32104823f, 0.27464288f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 7.49001646f, 4.65472794f, 3.07277966f, 2.19988537f, 1.61558151f, 1.24153244f, 0.95350921f, 0.74807048f, 0.59516323f, 0.50118381f, 0.41087446f, 0.34370604f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}}; + +const std::vector> GITS_NOISE_1_25 = { + {14.61464119f, 0.72133851f, 0.02916753f}, + {14.61464119f, 1.56271636f, 0.50118381f, 0.02916753f}, + {14.61464119f, 2.05039096f, 0.803307f, 0.32104823f, 0.02916753f}, + {14.61464119f, 2.36326075f, 0.95350921f, 0.43325692f, 0.17026083f, 0.02916753f}, + {14.61464119f, 2.84484982f, 1.24153244f, 0.59516323f, 0.27464288f, 0.09824532f, 0.02916753f}, + {14.61464119f, 3.07277966f, 1.51179266f, 0.803307f, 0.43325692f, 0.22545385f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.36326075f, 1.24153244f, 0.72133851f, 0.41087446f, 0.22545385f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.45070267f, 1.36964464f, 0.83188516f, 0.52423614f, 0.34370604f, 0.19894916f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.84484982f, 1.61558151f, 0.98595673f, 0.64427125f, 0.43325692f, 0.27464288f, 0.17026083f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.84484982f, 1.67050016f, 1.08895338f, 0.74807048f, 0.52423614f, 0.36617002f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.84484982f, 1.72759056f, 1.162866f, 0.803307f, 0.59516323f, 0.45573691f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.95596409f, 1.84880662f, 1.24153244f, 0.86115354f, 0.64427125f, 0.4783645f, 0.36617002f, 0.27464288f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.95596409f, 1.84880662f, 1.28281462f, 0.92192322f, 0.69515091f, 0.52423614f, 0.41087446f, 0.32104823f, 0.25053367f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.95596409f, 1.91321158f, 1.32549286f, 0.95350921f, 0.72133851f, 0.54755926f, 0.43325692f, 0.34370604f, 0.27464288f, 0.22545385f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.95596409f, 1.91321158f, 1.32549286f, 0.95350921f, 0.72133851f, 0.57119018f, 0.45573691f, 0.36617002f, 0.29807833f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.95596409f, 1.91321158f, 1.32549286f, 0.95350921f, 0.74807048f, 0.59516323f, 0.4783645f, 0.38853383f, 0.32104823f, 0.27464288f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 3.07277966f, 2.05039096f, 1.41535246f, 1.05362725f, 0.803307f, 0.61951244f, 0.50118381f, 0.41087446f, 0.34370604f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 3.07277966f, 2.05039096f, 1.41535246f, 1.05362725f, 0.803307f, 0.64427125f, 0.52423614f, 0.43325692f, 0.36617002f, 0.32104823f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 3.07277966f, 2.05039096f, 1.46270394f, 1.08895338f, 0.83188516f, 0.66947293f, 0.54755926f, 0.45573691f, 0.38853383f, 0.34370604f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}}; + +const std::vector> GITS_NOISE_1_30 = { + {14.61464119f, 0.72133851f, 0.02916753f}, + {14.61464119f, 1.24153244f, 0.43325692f, 0.02916753f}, + {14.61464119f, 1.56271636f, 0.59516323f, 0.22545385f, 0.02916753f}, + {14.61464119f, 1.84880662f, 0.803307f, 0.36617002f, 0.13792117f, 0.02916753f}, + {14.61464119f, 2.36326075f, 1.01931262f, 0.52423614f, 0.25053367f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.84484982f, 1.36964464f, 0.74807048f, 0.41087446f, 0.22545385f, 0.09824532f, 0.02916753f}, + {14.61464119f, 3.07277966f, 1.56271636f, 0.89115214f, 0.54755926f, 0.34370604f, 0.19894916f, 0.09824532f, 0.02916753f}, + {14.61464119f, 3.07277966f, 1.61558151f, 0.95350921f, 0.61951244f, 0.41087446f, 0.27464288f, 0.17026083f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.45070267f, 1.36964464f, 0.83188516f, 0.54755926f, 0.36617002f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.45070267f, 1.41535246f, 0.92192322f, 0.64427125f, 0.45573691f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.6383388f, 1.56271636f, 1.01931262f, 0.72133851f, 0.50118381f, 0.36617002f, 0.27464288f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.84484982f, 1.61558151f, 1.05362725f, 0.74807048f, 0.54755926f, 0.41087446f, 0.32104823f, 0.25053367f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.84484982f, 1.61558151f, 1.08895338f, 0.77538133f, 0.57119018f, 0.43325692f, 0.34370604f, 0.27464288f, 0.22545385f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.84484982f, 1.61558151f, 1.08895338f, 0.803307f, 0.59516323f, 0.45573691f, 0.36617002f, 0.29807833f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.84484982f, 1.61558151f, 1.08895338f, 0.803307f, 0.59516323f, 0.4783645f, 0.38853383f, 0.32104823f, 0.27464288f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.84484982f, 1.72759056f, 1.162866f, 0.83188516f, 0.64427125f, 0.50118381f, 0.41087446f, 0.34370604f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.84484982f, 1.72759056f, 1.162866f, 0.83188516f, 0.64427125f, 0.52423614f, 0.43325692f, 0.36617002f, 0.32104823f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.84484982f, 1.78698075f, 1.24153244f, 0.92192322f, 0.72133851f, 0.57119018f, 0.45573691f, 0.38853383f, 0.34370604f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.84484982f, 1.78698075f, 1.24153244f, 0.92192322f, 0.72133851f, 0.57119018f, 0.4783645f, 0.41087446f, 0.36617002f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}}; + +const std::vector> GITS_NOISE_1_35 = { + {14.61464119f, 0.69515091f, 0.02916753f}, + {14.61464119f, 0.95350921f, 0.34370604f, 0.02916753f}, + {14.61464119f, 1.56271636f, 0.57119018f, 0.19894916f, 0.02916753f}, + {14.61464119f, 1.61558151f, 0.69515091f, 0.29807833f, 0.09824532f, 0.02916753f}, + {14.61464119f, 1.84880662f, 0.83188516f, 0.43325692f, 0.22545385f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.45070267f, 1.162866f, 0.64427125f, 0.36617002f, 0.19894916f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.84484982f, 1.36964464f, 0.803307f, 0.50118381f, 0.32104823f, 0.19894916f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.84484982f, 1.41535246f, 0.83188516f, 0.54755926f, 0.36617002f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.84484982f, 1.56271636f, 0.95350921f, 0.64427125f, 0.45573691f, 0.32104823f, 0.22545385f, 0.17026083f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.84484982f, 1.56271636f, 0.95350921f, 0.64427125f, 0.45573691f, 0.34370604f, 0.25053367f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 3.07277966f, 1.61558151f, 1.01931262f, 0.72133851f, 0.52423614f, 0.38853383f, 0.29807833f, 0.22545385f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 3.07277966f, 1.61558151f, 1.01931262f, 0.72133851f, 0.52423614f, 0.41087446f, 0.32104823f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 3.07277966f, 1.61558151f, 1.05362725f, 0.74807048f, 0.54755926f, 0.43325692f, 0.34370604f, 0.27464288f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 3.07277966f, 1.72759056f, 1.12534678f, 0.803307f, 0.59516323f, 0.45573691f, 0.36617002f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 3.07277966f, 1.72759056f, 1.12534678f, 0.803307f, 0.59516323f, 0.4783645f, 0.38853383f, 0.32104823f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.45070267f, 1.51179266f, 1.01931262f, 0.74807048f, 0.57119018f, 0.45573691f, 0.36617002f, 0.32104823f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.6383388f, 1.61558151f, 1.08895338f, 0.803307f, 0.61951244f, 0.50118381f, 0.41087446f, 0.34370604f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.6383388f, 1.61558151f, 1.08895338f, 0.803307f, 0.64427125f, 0.52423614f, 0.43325692f, 0.36617002f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 5.85520077f, 2.6383388f, 1.61558151f, 1.08895338f, 0.803307f, 0.64427125f, 0.52423614f, 0.45573691f, 0.38853383f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}}; + +const std::vector> GITS_NOISE_1_40 = { + {14.61464119f, 0.59516323f, 0.02916753f}, + {14.61464119f, 0.95350921f, 0.34370604f, 0.02916753f}, + {14.61464119f, 1.08895338f, 0.43325692f, 0.13792117f, 0.02916753f}, + {14.61464119f, 1.56271636f, 0.64427125f, 0.27464288f, 0.09824532f, 0.02916753f}, + {14.61464119f, 1.61558151f, 0.803307f, 0.43325692f, 0.22545385f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.05039096f, 0.95350921f, 0.54755926f, 0.34370604f, 0.19894916f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.45070267f, 1.24153244f, 0.72133851f, 0.43325692f, 0.27464288f, 0.17026083f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.45070267f, 1.24153244f, 0.74807048f, 0.50118381f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.45070267f, 1.28281462f, 0.803307f, 0.52423614f, 0.36617002f, 0.27464288f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.45070267f, 1.28281462f, 0.803307f, 0.54755926f, 0.38853383f, 0.29807833f, 0.22545385f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.84484982f, 1.41535246f, 0.86115354f, 0.59516323f, 0.43325692f, 0.32104823f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.84484982f, 1.51179266f, 0.95350921f, 0.64427125f, 0.45573691f, 0.34370604f, 0.27464288f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.84484982f, 1.51179266f, 0.95350921f, 0.64427125f, 0.4783645f, 0.36617002f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.84484982f, 1.56271636f, 0.98595673f, 0.69515091f, 0.52423614f, 0.41087446f, 0.34370604f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.84484982f, 1.56271636f, 1.01931262f, 0.72133851f, 0.54755926f, 0.43325692f, 0.36617002f, 0.32104823f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.84484982f, 1.61558151f, 1.05362725f, 0.74807048f, 0.57119018f, 0.45573691f, 0.38853383f, 0.34370604f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.84484982f, 1.61558151f, 1.08895338f, 0.803307f, 0.61951244f, 0.50118381f, 0.41087446f, 0.36617002f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.84484982f, 1.61558151f, 1.08895338f, 0.803307f, 0.61951244f, 0.50118381f, 0.43325692f, 0.38853383f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.84484982f, 1.61558151f, 1.08895338f, 0.803307f, 0.64427125f, 0.52423614f, 0.45573691f, 0.41087446f, 0.36617002f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}}; + +const std::vector> GITS_NOISE_1_45 = { + {14.61464119f, 0.59516323f, 0.02916753f}, + {14.61464119f, 0.803307f, 0.25053367f, 0.02916753f}, + {14.61464119f, 0.95350921f, 0.34370604f, 0.09824532f, 0.02916753f}, + {14.61464119f, 1.24153244f, 0.54755926f, 0.25053367f, 0.09824532f, 0.02916753f}, + {14.61464119f, 1.56271636f, 0.72133851f, 0.36617002f, 0.19894916f, 0.09824532f, 0.02916753f}, + {14.61464119f, 1.61558151f, 0.803307f, 0.45573691f, 0.27464288f, 0.17026083f, 0.09824532f, 0.02916753f}, + {14.61464119f, 1.91321158f, 0.95350921f, 0.57119018f, 0.36617002f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.19988537f, 1.08895338f, 0.64427125f, 0.41087446f, 0.27464288f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.45070267f, 1.24153244f, 0.74807048f, 0.50118381f, 0.34370604f, 0.25053367f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.45070267f, 1.24153244f, 0.74807048f, 0.50118381f, 0.36617002f, 0.27464288f, 0.22545385f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.45070267f, 1.28281462f, 0.803307f, 0.54755926f, 0.41087446f, 0.32104823f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.45070267f, 1.28281462f, 0.803307f, 0.57119018f, 0.43325692f, 0.34370604f, 0.27464288f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.45070267f, 1.28281462f, 0.83188516f, 0.59516323f, 0.45573691f, 0.36617002f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.45070267f, 1.28281462f, 0.83188516f, 0.59516323f, 0.45573691f, 0.36617002f, 0.32104823f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.84484982f, 1.51179266f, 0.95350921f, 0.69515091f, 0.52423614f, 0.41087446f, 0.34370604f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.84484982f, 1.51179266f, 0.95350921f, 0.69515091f, 0.52423614f, 0.43325692f, 0.36617002f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.84484982f, 1.56271636f, 0.98595673f, 0.72133851f, 0.54755926f, 0.45573691f, 0.38853383f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.84484982f, 1.56271636f, 1.01931262f, 0.74807048f, 0.57119018f, 0.4783645f, 0.41087446f, 0.36617002f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.84484982f, 1.56271636f, 1.01931262f, 0.74807048f, 0.59516323f, 0.50118381f, 0.43325692f, 0.38853383f, 0.36617002f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}}; + +const std::vector> GITS_NOISE_1_50 = { + {14.61464119f, 0.54755926f, 0.02916753f}, + {14.61464119f, 0.803307f, 0.25053367f, 0.02916753f}, + {14.61464119f, 0.86115354f, 0.32104823f, 0.09824532f, 0.02916753f}, + {14.61464119f, 1.24153244f, 0.54755926f, 0.25053367f, 0.09824532f, 0.02916753f}, + {14.61464119f, 1.56271636f, 0.72133851f, 0.36617002f, 0.19894916f, 0.09824532f, 0.02916753f}, + {14.61464119f, 1.61558151f, 0.803307f, 0.45573691f, 0.27464288f, 0.17026083f, 0.09824532f, 0.02916753f}, + {14.61464119f, 1.61558151f, 0.83188516f, 0.52423614f, 0.34370604f, 0.25053367f, 0.17026083f, 0.09824532f, 0.02916753f}, + {14.61464119f, 1.84880662f, 0.95350921f, 0.59516323f, 0.38853383f, 0.27464288f, 0.19894916f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 1.84880662f, 0.95350921f, 0.59516323f, 0.41087446f, 0.29807833f, 0.22545385f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 1.84880662f, 0.95350921f, 0.61951244f, 0.43325692f, 0.32104823f, 0.25053367f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.19988537f, 1.12534678f, 0.72133851f, 0.50118381f, 0.36617002f, 0.27464288f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.19988537f, 1.12534678f, 0.72133851f, 0.50118381f, 0.36617002f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.36326075f, 1.24153244f, 0.803307f, 0.57119018f, 0.43325692f, 0.34370604f, 0.29807833f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.36326075f, 1.24153244f, 0.803307f, 0.57119018f, 0.43325692f, 0.34370604f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.36326075f, 1.24153244f, 0.803307f, 0.59516323f, 0.45573691f, 0.36617002f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.36326075f, 1.24153244f, 0.803307f, 0.59516323f, 0.45573691f, 0.38853383f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.45070267f, 1.32549286f, 0.86115354f, 0.64427125f, 0.50118381f, 0.41087446f, 0.36617002f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.45070267f, 1.36964464f, 0.92192322f, 0.69515091f, 0.54755926f, 0.45573691f, 0.41087446f, 0.36617002f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}, + {14.61464119f, 2.45070267f, 1.41535246f, 0.95350921f, 0.72133851f, 0.57119018f, 0.4783645f, 0.43325692f, 0.38853383f, 0.36617002f, 0.34370604f, 0.32104823f, 0.29807833f, 0.27464288f, 0.25053367f, 0.22545385f, 0.19894916f, 0.17026083f, 0.13792117f, 0.09824532f, 0.02916753f}}; + +const std::vector>*> GITS_NOISE = { + &GITS_NOISE_0_80, + &GITS_NOISE_0_85, + &GITS_NOISE_0_90, + &GITS_NOISE_0_95, + &GITS_NOISE_1_00, + &GITS_NOISE_1_05, + &GITS_NOISE_1_10, + &GITS_NOISE_1_15, + &GITS_NOISE_1_20, + &GITS_NOISE_1_25, + &GITS_NOISE_1_30, + &GITS_NOISE_1_35, + &GITS_NOISE_1_40, + &GITS_NOISE_1_45, + &GITS_NOISE_1_50}; + +#endif // __SD_RUNTIME_GITS_NOISE_H__ diff --git a/src/guidance.cpp b/src/runtime/guidance.cpp similarity index 72% rename from src/guidance.cpp rename to src/runtime/guidance.cpp index 2cb305a5..bfb773b0 100644 --- a/src/guidance.cpp +++ b/src/runtime/guidance.cpp @@ -1,12 +1,13 @@ -#include "guidance.h" +#include "runtime/guidance.h" #include #include #include +#include #include #include -#include "util.h" +#include "core/util.h" namespace sd::guidance { @@ -63,6 +64,82 @@ namespace sd::guidance { return uncond; } + std::vector parse_guidance_schedule_from_spec(std::string spec) { + std::vector schedule; + + while (!spec.empty()) { + auto sep = spec.find('+'); + auto segment = spec.substr(0, sep); + + auto x = segment.find('x'); + if (x == std::string::npos) { + LOG_ERROR("Invalid guidance schedule segment: '%s' (expected x)", segment.c_str()); + return {}; + } + + float guidance; + int count; + + auto guidance_str = segment.substr(0, x); + auto count_str = segment.substr(x + 1); + + try { + size_t idx = 0; + guidance = std::stof(guidance_str, &idx); + if (idx != guidance_str.size()) { + LOG_ERROR("Invalid guidance value in guidance schedule: '%s'", guidance_str.c_str()); + return {}; + } + } catch (const std::exception&) { + LOG_ERROR("Invalid guidance value in guidance schedule: '%s'", guidance_str.c_str()); + return {}; + } + + try { + size_t idx = 0; + count = std::stoi(count_str, &idx); + if (idx != count_str.size()) { + LOG_ERROR("Invalid count in guidance schedule: '%s'", count_str.c_str()); + return {}; + } + } catch (const std::exception&) { + LOG_ERROR("Invalid count in guidance schedule: '%s'", count_str.c_str()); + return {}; + } + + if (count <= 0) { + LOG_ERROR("Guidance schedule count must be positive"); + return {}; + } + + schedule.insert(schedule.end(), count, guidance); + + if (sep == std::string::npos) { + break; + } + + spec = spec.substr(sep + 1); + } + + return schedule; + } + + std::vector parse_guidance_schedule(const char* extra_sample_args) { + std::vector guidance_schedule; + std::string guidance_schedule_str = ""; + for (const auto& [key, value] : parse_key_value_args(extra_sample_args, "extra sample arg")) { + float parsed = 0.0f; + if (key == "guidance_schedule") { + guidance_schedule_str = value; + } + } + + if (!guidance_schedule_str.empty()) { + guidance_schedule = parse_guidance_schedule_from_spec(guidance_schedule_str); + } + return guidance_schedule; + } + ClassifierFreeGuidance::ClassifierFreeGuidance(float guidance_scale, float image_guidance_scale) : guidance_scale_(guidance_scale), @@ -70,8 +147,10 @@ namespace sd::guidance { } GuiderOutput ClassifierFreeGuidance::forward(const GuidanceInput& input, - GuiderOutput previous) const { + GuiderOutput previous, + std::optional scale_override) const { (void)previous; + float guidance_scale = scale_override.value_or(guidance_scale_); GuiderOutput output; if (!has_tensor(input.pred_cond)) { @@ -86,14 +165,14 @@ namespace sd::guidance { const sd::Tensor& pred_img_uncond = *input.pred_img_uncond; output.pred = pred_img_uncond + image_guidance_scale_ * (pred_uncond - pred_img_uncond) + - guidance_scale_ * (pred_cond - pred_uncond); + guidance_scale * (pred_cond - pred_uncond); } else { - output.pred = pred_uncond + guidance_scale_ * (pred_cond - pred_uncond); + output.pred = pred_uncond + guidance_scale * (pred_cond - pred_uncond); } } else if (has_tensor(input.pred_img_uncond)) { const sd::Tensor& pred_img_uncond = *input.pred_img_uncond; - output.pred = pred_img_uncond + guidance_scale_ * (pred_cond - pred_img_uncond); + output.pred = pred_img_uncond + guidance_scale * (pred_cond - pred_img_uncond); } return output; @@ -128,8 +207,10 @@ namespace sd::guidance { } GuiderOutput AdaptiveProjectedGuidance::forward(const GuidanceInput& input, - GuiderOutput previous) const { + GuiderOutput previous, + std::optional scale_override) const { (void)previous; + float guidance_scale = scale_override.value_or(guidance_scale_); GuiderOutput output; if (!has_tensor(input.pred_cond)) { @@ -144,13 +225,13 @@ namespace sd::guidance { const sd::Tensor& pred_img_uncond = *input.pred_img_uncond; output.pred = pred_img_uncond + image_guidance_scale_ * (pred_uncond - pred_img_uncond) + - guidance_scale_ * (pred_cond - pred_uncond); + guidance_scale * (pred_cond - pred_uncond); } else { - output.pred = pred_uncond + guidance_scale_ * (pred_cond - pred_uncond); + output.pred = pred_uncond + guidance_scale * (pred_cond - pred_uncond); } } else if (has_tensor(input.pred_img_uncond)) { const sd::Tensor& pred_img_uncond = *input.pred_img_uncond; - output.pred = pred_img_uncond + guidance_scale_ * (pred_cond - pred_img_uncond); + output.pred = pred_img_uncond + guidance_scale * (pred_cond - pred_img_uncond); } if (!has_tensor(input.pred_uncond) && !has_tensor(input.pred_img_uncond)) { return output; @@ -162,7 +243,7 @@ namespace sd::guidance { sd::Tensor deltas = calculate_guidance_delta(pred_cond, pred_uncond, pred_img_uncond, - guidance_scale_, + guidance_scale, image_guidance_scale_); if (params_.momentum != 0.0f) { if (momentum_buffer_.shape() != deltas.shape()) { @@ -172,9 +253,10 @@ namespace sd::guidance { momentum_buffer_ = deltas; } - float diff_norm = 0.0f; + float diff_norm = 0.0f; + const int standard_res = 2 * 1024 / 8; // Use SDXL as the standard resolution (1024x1024, 8x8 patches, 4=2x2 channels) if (params_.norm_threshold > 0.0f) { - diff_norm = std::sqrt((deltas * deltas).sum()); + diff_norm = std::sqrt((deltas * deltas).sum()) * standard_res / std::sqrt(static_cast(deltas.numel())); } float apg_scale_factor = 1.0f; @@ -238,7 +320,8 @@ namespace sd::guidance { } GuiderOutput SkipLayerGuidance::forward(const GuidanceInput& input, - GuiderOutput output) const { + GuiderOutput output, + std::optional /*scale_override*/) const { if (scale_ == 0.0f || !is_enabled_for_step(input) || !input.predict_skip_layer) { return output; } diff --git a/src/guidance.h b/src/runtime/guidance.h similarity index 76% rename from src/guidance.h rename to src/runtime/guidance.h index 26b79de9..3de33704 100644 --- a/src/guidance.h +++ b/src/runtime/guidance.h @@ -1,11 +1,12 @@ -#ifndef __SD_GUIDANCE_H__ -#define __SD_GUIDANCE_H__ +#ifndef __SD_RUNTIME_GUIDANCE_H__ +#define __SD_RUNTIME_GUIDANCE_H__ #include #include +#include #include -#include "tensor.hpp" +#include "core/tensor.hpp" namespace sd::guidance { @@ -27,6 +28,7 @@ namespace sd::guidance { AdaptiveProjectedGuidanceParams parse_adaptive_projected_guidance_args(const char* extra_sample_args); bool is_adaptive_projected_guidance_enabled(const AdaptiveProjectedGuidanceParams& params); bool parse_skip_layer_guidance_uncond_arg(const char* extra_sample_args); + std::vector parse_guidance_schedule(const char* extra_sample_args); struct GuidanceInput { int step = 0; @@ -40,9 +42,10 @@ namespace sd::guidance { class BaseGuidance { public: - virtual ~BaseGuidance() = default; + virtual ~BaseGuidance() = default; virtual GuiderOutput forward(const GuidanceInput& input, - GuiderOutput previous) const = 0; + GuiderOutput previous, + std::optional scale_override = std::nullopt) const = 0; }; class ClassifierFreeGuidance : public BaseGuidance { @@ -54,7 +57,8 @@ namespace sd::guidance { float image_guidance_scale); GuiderOutput forward(const GuidanceInput& input, - GuiderOutput previous) const override; + GuiderOutput previous, + std::optional scale_override = std::nullopt) const override; }; class AdaptiveProjectedGuidance : public BaseGuidance { @@ -69,7 +73,8 @@ namespace sd::guidance { AdaptiveProjectedGuidanceParams params); GuiderOutput forward(const GuidanceInput& input, - GuiderOutput previous) const override; + GuiderOutput previous, + std::optional scale_override = std::nullopt) const override; }; class SkipLayerGuidance : public BaseGuidance { @@ -88,9 +93,10 @@ namespace sd::guidance { const std::vector& layers() const; GuiderOutput forward(const GuidanceInput& input, - GuiderOutput previous) const override; + GuiderOutput previous, + std::optional scale_override = std::nullopt) const override; }; } // namespace sd::guidance -#endif // __SD_GUIDANCE_H__ +#endif // __SD_RUNTIME_GUIDANCE_H__ diff --git a/src/runtime/imatrix.cpp b/src/runtime/imatrix.cpp new file mode 100644 index 00000000..313eadc6 --- /dev/null +++ b/src/runtime/imatrix.cpp @@ -0,0 +1,308 @@ +#include "runtime/imatrix.h" + +/* Adapted from llama.cpp (credits: Kawrakow). */ + +#include "core/util.h" +#include "ggml-backend.h" +#include "ggml.h" +#include "stable-diffusion.h" + +#include +#include +#include + +static IMatrixCollector imatrix_collector; + +IMatrixCollector& get_imatrix_collector() { + return imatrix_collector; +} + +// remove any prefix and suffixes from the name +// CUDA0#blk.0.attn_k.weight#0 => blk.0.attn_k.weight +static std::string filter_tensor_name(const char* name) { + std::string wname; + const char* p = strchr(name, '#'); + if (p != NULL) { + p = p + 1; + const char* q = strchr(p, '#'); + if (q != NULL) { + wname = std::string(p, q - p); + } else { + wname = p; + } + } else { + wname = name; + } + return wname; +} + +bool IMatrixCollector::collect_imatrix(struct ggml_tensor* t, bool ask, void* user_data) { + GGML_UNUSED(user_data); + if (t == nullptr) { + return false; + } + if (t->op != GGML_OP_MUL_MAT && t->op != GGML_OP_MUL_MAT_ID) { + return false; + } + + const struct ggml_tensor* src0 = t->src[0]; + const struct ggml_tensor* src1 = t->src[1]; + if (src0 == nullptr || src1 == nullptr) { + return false; + } + std::string wname = filter_tensor_name(src0->name); + + // when ask is true, the scheduler wants to know if we are interested in data from this tensor + // if we return true, a follow-up call will be made with ask=false in which we can do the actual collection + if (ask) { + if (t->op == GGML_OP_MUL_MAT_ID) { + return true; // collect all indirect matrix multiplications + } + // why are small batches ignored (<16 tokens)? + // if (src1->ne[1] < 16 || src1->type != GGML_TYPE_F32) return false; + if (!(wname.substr(0, 6) == "model." || wname.substr(0, 17) == "cond_stage_model." || wname.substr(0, 14) == "text_encoders.")) { + return false; + } + return true; + } + std::lock_guard lock(mutex_); + + // copy the data from the GPU memory if needed + const bool is_host = src1->buffer == NULL || ggml_backend_buffer_is_host(src1->buffer); + + if (!is_host) { + src1_data_.resize(ggml_nelements(src1)); + ggml_backend_tensor_get(src1, src1_data_.data(), 0, ggml_nbytes(src1)); + } + + const float* data = is_host ? (const float*)src1->data : src1_data_.data(); + + // this has been adapted to the new format of storing merged experts in a single 3d tensor + // ref: https://github.com/ggml-org/llama.cpp/pull/6387 + if (t->op == GGML_OP_MUL_MAT_ID) { + // ids -> [n_experts_used, n_tokens] + // src1 -> [cols, n_expert_used, n_tokens] + const ggml_tensor* ids = t->src[2]; + const int n_as = static_cast(src0->ne[2]); + const int n_ids = static_cast(ids->ne[0]); + + // the top-k selected expert ids are stored in the ids tensor + // for simplicity, always copy ids to host, because it is small + // take into account that ids is not contiguous! + + GGML_ASSERT(ids->ne[1] == src1->ne[2]); + + ids_.resize(ggml_nbytes(ids)); + ggml_backend_tensor_get(ids, ids_.data(), 0, ggml_nbytes(ids)); + + auto& e = stats_[wname]; + + ++e.ncall; + + if (e.values.empty()) { + e.values.resize(src1->ne[0] * n_as, 0); + e.counts.resize(src1->ne[0] * n_as, 0); + } else if (e.values.size() != (size_t)src1->ne[0] * n_as) { + LOG_ERROR("inconsistent size for %s (%d vs %d)\n", wname.c_str(), (int)e.values.size(), (int)src1->ne[0] * n_as); + exit(1); // GGML_ABORT("fatal error"); + } + // loop over all possible experts, regardless if they are used or not in the batch + for (int ex = 0; ex < n_as; ++ex) { + size_t e_start = ex * src1->ne[0]; + + for (int idx = 0; idx < n_ids; ++idx) { + for (int row = 0; row < (int)src1->ne[2]; ++row) { + const int excur = *(const int32_t*)(ids_.data() + row * ids->nb[1] + idx * ids->nb[0]); + + GGML_ASSERT(excur >= 0 && excur < n_as); // sanity check + + if (excur != ex) + continue; + + const int64_t i11 = idx % src1->ne[1]; + const int64_t i12 = row; + const float* x = (const float*)((const char*)data + i11 * src1->nb[1] + i12 * src1->nb[2]); + + for (int j = 0; j < (int)src1->ne[0]; ++j) { + e.values[e_start + j] += x[j] * x[j]; + e.counts[e_start + j]++; + if (!std::isfinite(e.values[e_start + j])) { + LOG_ERROR("%f detected in %s\n", e.values[e_start + j], wname.c_str()); + exit(1); + } + } + } + } + } + } else { + auto& e = stats_[wname]; + if (e.values.empty()) { + e.values.resize(src1->ne[0], 0); + e.counts.resize(src1->ne[0], 0); + } else if (e.values.size() != (size_t)src1->ne[0]) { + LOG_WARN("inconsistent size for %s (%d vs %d)\n", wname.c_str(), (int)e.values.size(), (int)src1->ne[0]); + exit(1); // GGML_ABORT("fatal error"); + } + + ++e.ncall; + for (int row = 0; row < (int)src1->ne[1]; ++row) { + const float* x = data + row * src1->ne[0]; + for (int j = 0; j < (int)src1->ne[0]; ++j) { + if (std::isfinite(x[j])) { + e.values[j] += x[j] * x[j]; + e.counts[j]++; + if (!std::isfinite(e.values[j])) { + LOG_WARN("%f detected in %s\n", e.values[j], wname.c_str()); + exit(1); + } + } else { + // Likely something from an attention mask? + } + } + } + } + return true; +} + +bool load_imatrix(const char* imatrix_path) { + return imatrix_collector.load_imatrix(imatrix_path); +} + +void save_imatrix(const char* imatrix_path) { + imatrix_collector.save_imatrix(imatrix_path); +} + +static bool collect_imatrix(struct ggml_tensor* t, bool ask, void* user_data) { + return imatrix_collector.collect_imatrix(t, ask, user_data); +} + +void enable_imatrix_collection() { + sd_set_backend_eval_callback(collect_imatrix, nullptr); +} + +void disable_imatrix_collection() { + sd_set_backend_eval_callback(nullptr, nullptr); +} + +void IMatrixCollector::save_imatrix(std::string fname, int ncall) const { + if (ncall > 0) { + fname += ".at_"; + fname += std::to_string(ncall); + } + // avoid writing imatrix entries that do not have full data + // this can happen with MoE models where some of the experts end up not being exercised by the provided training data + + int n_entries = 0; + std::vector to_store; + + for (const auto& kv : stats_) { + const int n_all = static_cast(kv.second.counts.size()); + + if (n_all == 0) { + continue; + } + + int n_zeros = 0; + for (const int c : kv.second.counts) { + if (c == 0) { + n_zeros++; + } + } + + if (n_zeros == n_all) { + LOG_WARN("entry '%40s' has no data - skipping\n", kv.first.c_str()); + continue; + } + + if (n_zeros > 0) { + LOG_WARN("entry '%40s' has partial data (%.2f%%) - skipping\n", kv.first.c_str(), 100.0f * (n_all - n_zeros) / n_all); + continue; + } + + n_entries++; + to_store.push_back(kv.first); + } + + if (to_store.size() < stats_.size()) { + LOG_WARN("storing only %zu out of %zu entries\n", to_store.size(), stats_.size()); + } + + std::ofstream out(fname, std::ios::binary); + out.write((const char*)&n_entries, sizeof(n_entries)); + for (const auto& name : to_store) { + const auto& stat = stats_.at(name); + int len = static_cast(name.size()); + out.write((const char*)&len, sizeof(len)); + out.write(name.c_str(), len); + out.write((const char*)&stat.ncall, sizeof(stat.ncall)); + int nval = static_cast(stat.values.size()); + out.write((const char*)&nval, sizeof(nval)); + if (nval > 0) { + std::vector tmp(nval); + for (int i = 0; i < nval; i++) { + tmp[i] = (stat.values[i] / static_cast(stat.counts[i])) * static_cast(stat.ncall); + } + out.write((const char*)tmp.data(), nval * sizeof(float)); + } + } + + // Write the number of call the matrix was computed with + out.write((const char*)&last_call_, sizeof(last_call_)); +} + +bool IMatrixCollector::load_imatrix(const char* fname) { + std::ifstream in(fname, std::ios::binary); + if (!in) { + LOG_ERROR("failed to open %s\n", fname); + return false; + } + int n_entries; + in.read((char*)&n_entries, sizeof(n_entries)); + if (in.fail() || n_entries < 1) { + LOG_ERROR("no data in file %s\n", fname); + return false; + } + for (int i = 0; i < n_entries; ++i) { + int len; + in.read((char*)&len, sizeof(len)); + std::vector name_as_vec(len + 1); + in.read((char*)name_as_vec.data(), len); + if (in.fail()) { + LOG_ERROR("failed reading name for entry %d from %s\n", i + 1, fname); + return false; + } + name_as_vec[len] = 0; + std::string name{name_as_vec.data()}; + auto& e = stats_[std::move(name)]; + int ncall; + in.read((char*)&ncall, sizeof(ncall)); + int nval; + in.read((char*)&nval, sizeof(nval)); + if (in.fail() || nval < 1) { + LOG_ERROR("failed reading number of values for entry %d\n", i); + stats_ = {}; + return false; + } + + if (e.values.empty()) { + e.values.resize(nval, 0); + e.counts.resize(nval, 0); + } + + std::vector tmp(nval); + in.read((char*)tmp.data(), nval * sizeof(float)); + if (in.fail()) { + LOG_ERROR("failed reading data for entry %d\n", i); + stats_ = {}; + return false; + } + + // Recreate the state as expected by save_imatrix(), and correct for weighted sum. + for (int i = 0; i < nval; i++) { + e.values[i] += tmp[i]; + e.counts[i] += ncall; + } + e.ncall += ncall; + } + return true; +} diff --git a/src/runtime/imatrix.h b/src/runtime/imatrix.h new file mode 100644 index 00000000..341c18c8 --- /dev/null +++ b/src/runtime/imatrix.h @@ -0,0 +1,45 @@ +#ifndef __SD_RUNTIME_IMATRIX_H__ +#define __SD_RUNTIME_IMATRIX_H__ + +#include +#include +#include +#include +#include + +/* Adapted from llama.cpp (credits: Kawrakow). */ + +struct ggml_tensor; + +struct IMatrixStats { + std::vector values{}; + std::vector counts{}; + int ncall = 0; +}; + +class IMatrixCollector { +private: + std::unordered_map stats_ = {}; + std::mutex mutex_; + int last_call_ = 0; + std::vector src1_data_; + std::vector ids_; // the expert ids from ggml_mul_mat_id + +public: + IMatrixCollector() = default; + bool collect_imatrix(struct ggml_tensor* t, bool ask, void* user_data); + void save_imatrix(std::string fname, int ncall = -1) const; + bool load_imatrix(const char* fname); + std::vector get_values(const std::string& key) const { + auto it = stats_.find(key); + if (it != stats_.end()) { + return it->second.values; + } else { + return {}; + } + } +}; + +IMatrixCollector& get_imatrix_collector(); + +#endif // __SD_RUNTIME_IMATRIX_H__ diff --git a/src/latent-preview.h b/src/runtime/latent-preview.h similarity index 99% rename from src/latent-preview.h rename to src/runtime/latent-preview.h index 7b90b370..15239908 100644 --- a/src/latent-preview.h +++ b/src/runtime/latent-preview.h @@ -1,8 +1,8 @@ #include #include #include +#include "core/tensor.hpp" #include "ggml.h" -#include "tensor.hpp" const float ltxav_latent_rgb_proj[128][3] = { {-0.0293802f, -0.0362516f, -0.0291386f}, diff --git a/src/preprocessing.hpp b/src/runtime/preprocessing.hpp similarity index 98% rename from src/preprocessing.hpp rename to src/runtime/preprocessing.hpp index 57ab0cec..b39a9a9d 100644 --- a/src/preprocessing.hpp +++ b/src/runtime/preprocessing.hpp @@ -1,10 +1,10 @@ -#ifndef __PREPROCESSING_HPP__ -#define __PREPROCESSING_HPP__ +#ifndef __SD_RUNTIME_PREPROCESSING_HPP__ +#define __SD_RUNTIME_PREPROCESSING_HPP__ #include #include -#include "ggml_extend.hpp" +#include "core/ggml_extend.hpp" #define M_PI_ 3.14159265358979323846f @@ -331,4 +331,4 @@ bool preprocess_canny(sd_image_t img, float high_threshold, float low_threshold, return true; } -#endif // __PREPROCESSING_HPP__ +#endif // __SD_RUNTIME_PREPROCESSING_HPP__ diff --git a/src/sample-cache.cpp b/src/runtime/sample-cache.cpp similarity index 99% rename from src/sample-cache.cpp rename to src/runtime/sample-cache.cpp index 5739178d..bfa3d448 100644 --- a/src/sample-cache.cpp +++ b/src/runtime/sample-cache.cpp @@ -1,4 +1,4 @@ -#include "sample-cache.h" +#include "runtime/sample-cache.h" namespace sd_sample { diff --git a/src/sample-cache.h b/src/runtime/sample-cache.h similarity index 82% rename from src/sample-cache.h rename to src/runtime/sample-cache.h index 398ad065..cb0fe7bc 100644 --- a/src/sample-cache.h +++ b/src/runtime/sample-cache.h @@ -1,16 +1,16 @@ -#ifndef __SAMPLE_CACHE_H__ -#define __SAMPLE_CACHE_H__ +#ifndef __SD_RUNTIME_SAMPLE_CACHE_H__ +#define __SD_RUNTIME_SAMPLE_CACHE_H__ #include -#include "cache_dit.hpp" -#include "denoiser.hpp" -#include "easycache.hpp" +#include "core/tensor.hpp" +#include "core/util.h" #include "model.h" -#include "spectrum.hpp" -#include "tensor.hpp" -#include "ucache.hpp" -#include "util.h" +#include "runtime/cache_dit.hpp" +#include "runtime/denoiser.hpp" +#include "runtime/easycache.hpp" +#include "runtime/spectrum.hpp" +#include "runtime/ucache.hpp" namespace sd_sample { @@ -58,4 +58,4 @@ namespace sd_sample { } // namespace sd_sample -#endif // __SAMPLE_CACHE_H__ +#endif // __SD_RUNTIME_SAMPLE_CACHE_H__ diff --git a/src/spectrum.hpp b/src/runtime/spectrum.hpp similarity index 97% rename from src/spectrum.hpp rename to src/runtime/spectrum.hpp index add1796f..7246d6f0 100644 --- a/src/spectrum.hpp +++ b/src/runtime/spectrum.hpp @@ -1,12 +1,12 @@ -#ifndef __SPECTRUM_HPP__ -#define __SPECTRUM_HPP__ +#ifndef __SD_RUNTIME_SPECTRUM_HPP__ +#define __SD_RUNTIME_SPECTRUM_HPP__ #include #include #include -#include "ggml_extend.hpp" -#include "tensor.hpp" +#include "core/ggml_extend.hpp" +#include "core/tensor.hpp" struct SpectrumConfig { float w = 0.40f; @@ -184,4 +184,4 @@ private: } }; -#endif // __SPECTRUM_HPP__ +#endif // __SD_RUNTIME_SPECTRUM_HPP__ diff --git a/src/ucache.hpp b/src/runtime/ucache.hpp similarity index 98% rename from src/ucache.hpp rename to src/runtime/ucache.hpp index 3d785c5e..187e1e78 100644 --- a/src/ucache.hpp +++ b/src/runtime/ucache.hpp @@ -1,15 +1,15 @@ -#ifndef __UCACHE_HPP__ -#define __UCACHE_HPP__ +#ifndef __SD_RUNTIME_UCACHE_HPP__ +#define __SD_RUNTIME_UCACHE_HPP__ #include #include #include #include -#include "condition_cache_utils.hpp" -#include "denoiser.hpp" -#include "ggml_extend.hpp" -#include "tensor.hpp" +#include "core/ggml_extend.hpp" +#include "core/tensor.hpp" +#include "runtime/condition_cache_utils.hpp" +#include "runtime/denoiser.hpp" struct UCacheConfig { bool enabled = false; @@ -420,4 +420,4 @@ struct UCacheState { } }; -#endif // __UCACHE_HPP__ +#endif // __SD_RUNTIME_UCACHE_HPP__ diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index 610be601..7b23a2d3 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -1,52 +1,63 @@ #include #include #include +#include +#include +#include -#include "ggml_extend.hpp" -#include "ggml_graph_cut.h" +#include "core/ggml_extend.hpp" +#include "core/ggml_graph_cut.h" -#include "model.h" -#include "rng.hpp" -#include "rng_mt19937.hpp" -#include "rng_philox.hpp" +#include "core/rng.hpp" +#include "core/rng_mt19937.hpp" +#include "core/rng_philox.hpp" +#include "core/util.h" +#include "model_loader.h" +#include "model_manager.h" #include "stable-diffusion.h" -#include "util.h" -#include "anima.hpp" -#include "auto_encoder_kl.hpp" -#include "conditioner.hpp" -#include "control.hpp" -#include "denoiser.hpp" -#include "diffusion_model.hpp" -#include "ernie_image.hpp" -#include "esrgan.hpp" -#include "flux.hpp" -#include "guidance.h" -#include "hidream_o1.hpp" -#include "lens.hpp" -#include "lora.hpp" -#include "ltx_audio_vae.h" -#include "ltx_latent_upscaler.hpp" -#include "ltx_vae.hpp" -#include "ltxv.hpp" -#include "mmdit.hpp" -#include "pid.hpp" -#include "pmid.hpp" -#include "qwen_image.hpp" -#include "sample-cache.h" -#include "tae.hpp" -#include "unet.hpp" +#include "conditioning/conditioner.hpp" +#include "extensions/generation_extension.h" +#include "model/adapter/lora.hpp" +#include "model/diffusion/anima.hpp" +#include "model/diffusion/boogu.hpp" +#include "model/diffusion/control.hpp" +#include "model/diffusion/ernie_image.hpp" +#include "model/diffusion/flux.hpp" +#include "model/diffusion/hidream_o1.hpp" +#include "model/diffusion/ideogram4.hpp" +#include "model/diffusion/krea2.hpp" +#include "model/diffusion/lens.hpp" +#include "model/diffusion/ltxv.hpp" +#include "model/diffusion/minit2i.hpp" +#include "model/diffusion/mmdit.hpp" +#include "model/diffusion/model.hpp" +#include "model/diffusion/pid.hpp" +#include "model/diffusion/qwen_image.hpp" +#include "model/diffusion/unet.hpp" +#include "model/diffusion/wan.hpp" +#include "model/diffusion/z_image.hpp" +#include "model/upscaler/esrgan.hpp" +#include "model/upscaler/ltx_latent_upscaler.hpp" +#include "model/vae/auto_encoder_kl.hpp" +#include "model/vae/ltx_audio_vae.hpp" +#include "model/vae/ltx_vae.hpp" +#include "model/vae/tae.hpp" +#include "model/vae/vae.hpp" +#include "model/vae/wan_vae.hpp" +#include "runtime/denoiser.hpp" +#include "runtime/guidance.h" +#include "runtime/sample-cache.h" #include "upscaler.h" -#include "vae.hpp" -#include "wan.hpp" -#include "z_image.hpp" -#include "latent-preview.h" #include "name_conversion.h" +#include "runtime/latent-preview.h" const char* sd_vae_format_name(enum sd_vae_format_t format); static SDVersion sd_vae_format_to_version(enum sd_vae_format_t format, SDVersion fallback); +#include + const char* model_version_to_str[] = { "SD 1.x", "SD 1.x Inpaint", @@ -80,11 +91,17 @@ const char* model_version_to_str[] = { "LTXAV", "HiDream O1", "Z-Image", + "Boogu Image", "Ovis Image", "Ernie Image", "Lens", + "MiniT2I", "Longcat-Image", "PiD", + "Ideogram 4", + "SeFi-Image", + "Krea2", + "ESRGAN", }; const char* sampling_methods_str[] = { @@ -115,7 +132,8 @@ static bool sd_version_supports_ref_latent_img_cfg(SDVersion version) { sd_version_is_flux2(version) || sd_version_is_qwen_image(version) || sd_version_is_longcat(version) || - sd_version_is_z_image(version); + sd_version_is_z_image(version) || + sd_version_is_boogu_image(version); } static bool sd_version_supports_img_cfg(SDVersion version, bool has_ref_images) { @@ -152,15 +170,15 @@ static float get_cache_reuse_threshold(const sd_cache_params_t& params) { /*=============================================== StableDiffusionGGML ================================================*/ +static_assert(std::atomic::is_always_lock_free, + "sd_cancel_mode_t must be lock-free"); + class StableDiffusionGGML { public: - std::vector mmap_tensor_store; SDBackendManager backend_manager; SDVersion version; - bool vae_decode_only = false; bool external_vae_is_invalid = false; - bool free_params_immediately = false; bool circular_x = false; bool circular_y = false; @@ -178,32 +196,26 @@ public: std::shared_ptr preview_vae; std::shared_ptr audio_vae_model; std::shared_ptr control_net; - std::shared_ptr pmid_model; - std::shared_ptr pmid_lora; - std::shared_ptr pmid_id_embeds; - std::vector> cond_stage_lora_models; - std::vector> diffusion_lora_models; - std::vector> first_stage_lora_models; + std::vector> generation_extensions; + std::vector> runtime_lora_models; bool apply_lora_immediately = false; std::string taesd_path; sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0, 0, nullptr}; - bool offload_params_to_cpu = false; - float max_vram = 0.f; - bool stream_layers = false; - bool use_pmid = false; + bool enable_mmap = false; + sd::ggml_graph_cut::MaxVramAssignment max_vram_assignment; + bool stream_layers = false; + bool eager_load = false; std::string backend_spec; std::string params_backend_spec; bool is_using_v_parameterization = false; bool is_using_edm_v_parameterization = false; - std::map tensors; - - // lora_name => multiplier - std::unordered_map curr_lora_state; + std::shared_ptr model_manager; std::shared_ptr denoiser = std::make_shared(); + std::vector file_alphas_cumprod; StableDiffusionGGML() = default; @@ -225,6 +237,24 @@ public: return module_backend; } + std::atomic cancellation_flag = SD_CANCEL_RESET; + + void set_cancel_flag(enum sd_cancel_mode_t flag) { + cancellation_flag.store(flag, std::memory_order_release); + } + + void reset_cancel_flag() { + set_cancel_flag(SD_CANCEL_RESET); + } + + enum sd_cancel_mode_t get_cancel_flag() { + return cancellation_flag.load(std::memory_order_acquire); + } + + size_t max_graph_vram_bytes_for_module(SDBackendModule module) { + return max_vram_assignment.bytes_for_backend(backend_for(module)); + } + bool ensure_backend_pair(SDBackendModule module) { if (backend_for(module) == nullptr) { return false; @@ -232,14 +262,31 @@ public: return params_backend_for(module) != nullptr; } - bool init_backend(const sd_ctx_params_t* sd_ctx_params) { + template + bool register_runner_params(const std::string& desc, + const std::shared_ptr& model, + SDBackendModule module, + size_t* params_mem_size = nullptr) { + if (model == nullptr) { + return true; + } + std::map group_tensors; + model->get_param_tensors(group_tensors); + if (model_manager == nullptr) { + return true; + } + return model_manager->register_param_tensors(desc, + std::move(group_tensors), + backend_manager.params_backend_is_disk(module) ? ModelManager::ResidencyMode::Disk : ModelManager::ResidencyMode::ParamBackend, + backend_for(module), + params_backend_for(module), + params_mem_size); + } + + bool init_backend() { std::string error; - if (!backend_manager.init(sd_ctx_params->backend, - sd_ctx_params->params_backend, - offload_params_to_cpu, - sd_ctx_params->keep_clip_on_cpu, - sd_ctx_params->keep_vae_on_cpu, - sd_ctx_params->keep_control_net_on_cpu, + if (!backend_manager.init(backend_spec.c_str(), + params_backend_spec.c_str(), &error)) { LOG_ERROR("backend config failed: %s", error.c_str()); return false; @@ -257,27 +304,69 @@ public: } } - bool init(const sd_ctx_params_t* sd_ctx_params) { - n_threads = sd_ctx_params->n_threads; - vae_decode_only = sd_ctx_params->vae_decode_only; - free_params_immediately = sd_ctx_params->free_params_immediately; - offload_params_to_cpu = sd_ctx_params->offload_params_to_cpu; - max_vram = sd_ctx_params->max_vram; - stream_layers = sd_ctx_params->stream_layers; - backend_spec = SAFE_STR(sd_ctx_params->backend); - params_backend_spec = SAFE_STR(sd_ctx_params->params_backend); - if (stream_layers && max_vram == 0.f) { - LOG_WARN("--stream-layers has no effect without --max-vram set; ignoring"); - stream_layers = false; + void refresh_compvis_denoiser_sigmas() { + auto comp_vis_denoiser = std::dynamic_pointer_cast(denoiser); + if (!comp_vis_denoiser) { + return; } - if (stream_layers && !offload_params_to_cpu && params_backend_spec.empty()) { - // Streaming needs CPU-resident params. - LOG_WARN("--stream-layers has no effect without --offload-to-cpu (or --params-backend); ignoring"); - stream_layers = false; + std::vector alphas_cumprod(TIMESTEPS); + if (file_alphas_cumprod.size() == TIMESTEPS) { + alphas_cumprod = file_alphas_cumprod; + } else { + calculate_alphas_cumprod(alphas_cumprod.data()); + } + for (int i = 0; i < TIMESTEPS; i++) { + comp_vis_denoiser->sigmas[i] = std::sqrt((1 - alphas_cumprod[i]) / alphas_cumprod[i]); + comp_vis_denoiser->log_sigmas[i] = std::log(comp_vis_denoiser->sigmas[i]); + } + } + + void load_alphas_cumprod(ModelLoader& model_loader) { + file_alphas_cumprod.clear(); + + std::vector loaded_alphas; + if (!model_loader.load_float_tensor("alphas_cumprod", loaded_alphas, n_threads, enable_mmap)) { + return; + } + if (loaded_alphas.size() != TIMESTEPS) { + LOG_WARN("ignore alphas_cumprod from model file: expected %d values, got %zu", + TIMESTEPS, + loaded_alphas.size()); + return; + } + for (float alpha : loaded_alphas) { + if (!std::isfinite(alpha) || alpha <= 0.0f || alpha > 1.0f) { + LOG_WARN("ignore invalid alphas_cumprod from model file"); + return; + } } - bool use_tae = false; - bool use_audio_vae = false; + file_alphas_cumprod = std::move(loaded_alphas); + LOG_DEBUG("loaded alphas_cumprod from model file"); + } + + bool init(const sd_ctx_params_t* sd_ctx_params) { + n_threads = sd_ctx_params->n_threads; + enable_mmap = sd_ctx_params->enable_mmap; + stream_layers = sd_ctx_params->stream_layers; + eager_load = sd_ctx_params->eager_load; + backend_spec = SAFE_STR(sd_ctx_params->backend); + params_backend_spec = SAFE_STR(sd_ctx_params->params_backend); + max_vram_assignment.reset(0.f); + { + std::string error; + if (!max_vram_assignment.parse(SAFE_STR(sd_ctx_params->max_vram), &error)) { + LOG_ERROR("%s", error.c_str()); + return false; + } + } + + std::string rpc_servers_spec = SAFE_STR(sd_ctx_params->rpc_servers); + add_rpc_devices(rpc_servers_spec); + + bool use_tae = false; + bool use_audio_vae = false; + bool use_control_net = false; rng = get_rng(sd_ctx_params->rng_type); if (sd_ctx_params->sampler_rng_type != RNG_TYPE_COUNT && sd_ctx_params->sampler_rng_type != sd_ctx_params->rng_type) { @@ -288,12 +377,25 @@ public: ggml_log_set(ggml_log_callback_default, nullptr); - if (!init_backend(sd_ctx_params)) { + if (!init_backend()) { return false; } - max_vram = sd::ggml_graph_cut::resolve_max_vram_gib(max_vram, backend_for(SDBackendModule::DIFFUSION)); + { + std::string error; + if (!max_vram_assignment.canonicalize_backend_keys(&error)) { + LOG_ERROR("%s", error.c_str()); + return false; + } + } + if (stream_layers && !backend_manager.params_backend_is_cpu(SDBackendModule::DIFFUSION)) { + LOG_WARN("--stream-layers has no effect unless diffusion params backend is cpu; ignoring"); + stream_layers = false; + } - ModelLoader model_loader; + model_manager = std::make_shared(); + model_manager->set_n_threads(n_threads); + model_manager->set_enable_mmap(enable_mmap); + ModelLoader& model_loader = model_manager->loader(); if (strlen(SAFE_STR(sd_ctx_params->model_path)) > 0) { LOG_INFO("loading model from '%s'", sd_ctx_params->model_path); @@ -316,6 +418,13 @@ public: } } + if (strlen(SAFE_STR(sd_ctx_params->uncond_diffusion_model_path)) > 0) { + LOG_INFO("loading unconditional diffusion model from '%s'", sd_ctx_params->uncond_diffusion_model_path); + if (!model_loader.init_from_file(sd_ctx_params->uncond_diffusion_model_path, "model.diffusion_model.uncond.")) { + LOG_WARN("loading unconditional diffusion model from '%s' failed", sd_ctx_params->uncond_diffusion_model_path); + } + } + bool is_unet = sd_version_is_unet(model_loader.get_sd_version()); if (strlen(SAFE_STR(sd_ctx_params->clip_l_path)) > 0) { @@ -349,6 +458,14 @@ public: } } + if (strlen(SAFE_STR(sd_ctx_params->pulid_weights_path)) > 0) { + LOG_INFO("loading PuLID weights from '%s'", sd_ctx_params->pulid_weights_path); + if (!model_loader.init_from_file(sd_ctx_params->pulid_weights_path, + "model.diffusion_model.")) { + LOG_WARN("loading PuLID weights from '%s' failed", sd_ctx_params->pulid_weights_path); + } + } + if (strlen(SAFE_STR(sd_ctx_params->llm_path)) > 0) { LOG_INFO("loading llm from '%s'", sd_ctx_params->llm_path); if (!model_loader.init_from_file(sd_ctx_params->llm_path, "text_encoders.llm.")) { @@ -396,6 +513,15 @@ public: } } + if (strlen(SAFE_STR(sd_ctx_params->control_net_path)) > 0) { + if (!model_loader.init_from_file(sd_ctx_params->control_net_path)) { + LOG_ERROR("init control net model loader from file failed: '%s'", sd_ctx_params->control_net_path); + return false; + } else { + use_control_net = true; + } + } + model_loader.convert_tensors_name(); version = model_loader.get_sd_version(); @@ -407,9 +533,7 @@ public: auto& tensor_storage_map = model_loader.get_tensor_storage_map(); LOG_INFO("Version: %s ", model_version_to_str[version]); - ggml_type wtype = (int)sd_ctx_params->wtype < std::min(SD_TYPE_COUNT, GGML_TYPE_COUNT) - ? (ggml_type)sd_ctx_params->wtype - : GGML_TYPE_COUNT; + ggml_type wtype = sd_type_to_ggml_type(sd_ctx_params->wtype); std::string tensor_type_rules = SAFE_STR(sd_ctx_params->tensor_type_rules); if (wtype != GGML_TYPE_COUNT || tensor_type_rules.size() > 0) { model_loader.set_wtype_override(wtype, tensor_type_rules); @@ -454,8 +578,8 @@ public: } } // Avoid full-model LoRA merge buffers on constrained setups. - const bool streaming_constrained = stream_layers || - sd_ctx_params->offload_params_to_cpu; + const bool params_offloaded = params_backend_for(SDBackendModule::DIFFUSION) != backend_for(SDBackendModule::DIFFUSION); + const bool streaming_constrained = stream_layers || params_offloaded; if (have_quantized_weight || streaming_constrained) { apply_lora_immediately = false; } else { @@ -467,51 +591,20 @@ public: apply_lora_immediately = false; } - std::map mmap_able_tensors; - bool enable_mmap_tensors = false; - bool needs_writable_mmap = false; - if (sd_ctx_params->enable_mmap) { - if (apply_lora_immediately) { - needs_writable_mmap = true; - LOG_WARN("in mode 'immediately', LoRAs will cause extra memory usage with mmap"); - } - enable_mmap_tensors = true; + bool needs_writable_mmap = enable_mmap && apply_lora_immediately; + model_manager->set_writable_mmap(needs_writable_mmap); + if (enable_mmap && apply_lora_immediately) { + LOG_WARN("in mode 'immediately', LoRAs will cause extra memory usage with mmap"); } + model_loader.process_model_files(enable_mmap, needs_writable_mmap); + load_alphas_cumprod(model_loader); - // split definition to avoid msvc choking on the extra parameter handling - auto module_can_mmap = [&](SDBackendModule module) { - return enable_mmap_tensors && - (backend_manager.runtime_backend_is_cpu(module) || - backend_manager.params_backend_is_cpu(module) || - backend_manager.runtime_backend_supports_host_buffer(module)); - }; + size_t text_encoder_params_mem_size = 0; + size_t unet_params_mem_size = 0; + size_t vae_params_mem_size = 0; + size_t control_net_params_mem_size = 0; + size_t extension_params_mem_size = 0; - auto get_param_tensors_p = [&](auto&& model, bool do_mmap, const char* prefix) { - std::map temp; - model->get_param_tensors(temp, prefix); - for (const auto& [key, tensor] : temp) { - tensors[key] = tensor; - if (do_mmap) { - mmap_able_tensors[key] = tensor; - } - } - }; - - auto get_param_tensors = [&](auto&& model, bool do_mmap) { - std::map temp; - model->get_param_tensors(temp); - for (const auto& [key, tensor] : temp) { - tensors[key] = tensor; - if (do_mmap) { - mmap_able_tensors[key] = tensor; - } - } - }; - - if (sd_version_is_control(version)) { - // Might need vae encode for control cond - vae_decode_only = false; - } bool tae_preview_only = sd_ctx_params->tae_preview_only; if (version == VERSION_SDXS_512_DS || version == VERSION_SDXS_09) { tae_preview_only = false; @@ -522,8 +615,6 @@ public: LOG_INFO("Using circular padding for convolutions"); } - const size_t max_graph_vram_bytes = sd::ggml_graph_cut::max_vram_gib_to_bytes(max_vram); - { if (!ensure_backend_pair(SDBackendModule::TE) || !ensure_backend_pair(SDBackendModule::DIFFUSION)) { @@ -532,22 +623,45 @@ public: if (sd_version_is_sd3(version)) { cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - params_backend_for(SDBackendModule::TE), - tensor_storage_map); + tensor_storage_map, + model_manager); diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - params_backend_for(SDBackendModule::DIFFUSION), tensor_storage_map, - "model.diffusion_model"); + "model.diffusion_model", + model_manager); } else if (sd_version_is_pid(version)) { - vae_decode_only = false; cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - params_backend_for(SDBackendModule::TE), tensor_storage_map, - version); + version, + "", + false, + model_manager); diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - params_backend_for(SDBackendModule::DIFFUSION), tensor_storage_map, - "model.diffusion_model.net"); + "model.diffusion_model.net", + model_manager); + } else if (sd_version_is_ideogram4(version)) { + cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + model_manager); + diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + model_manager); + } else if (sd_version_is_krea2(version)) { + cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + model_manager); + diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + model_manager); } else if (sd_version_is_flux(version)) { bool is_chroma = false; for (auto pair : tensor_storage_map) { @@ -557,76 +671,72 @@ public: } } if (is_chroma) { - if ((sd_ctx_params->flash_attn || sd_ctx_params->diffusion_flash_attn) && sd_ctx_params->chroma_use_dit_mask) { - LOG_WARN( - "!!!It looks like you are using Chroma with flash attention. " - "This is currently unsupported. " - "If you find that the generated images are broken, " - "try either disabling flash attention or specifying " - "--chroma-disable-dit-mask as a workaround."); - } - cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - params_backend_for(SDBackendModule::TE), tensor_storage_map, sd_ctx_params->chroma_use_t5_mask, - sd_ctx_params->chroma_t5_mask_pad); + sd_ctx_params->chroma_t5_mask_pad, + false, + model_manager); } else if (version == VERSION_OVIS_IMAGE) { cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - params_backend_for(SDBackendModule::TE), tensor_storage_map, version, "", - false); + false, + model_manager); } else { cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - params_backend_for(SDBackendModule::TE), - tensor_storage_map); + tensor_storage_map, + model_manager); } diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - params_backend_for(SDBackendModule::DIFFUSION), tensor_storage_map, "model.diffusion_model", version, - sd_ctx_params->chroma_use_dit_mask); - } else if (sd_version_is_flux2(version)) { + sd_ctx_params->chroma_use_dit_mask, + model_manager); + } else if (sd_version_is_flux2(version) || sd_version_is_sefi_image(version)) { bool is_chroma = false; cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - params_backend_for(SDBackendModule::TE), tensor_storage_map, - version); + version, + "", + false, + model_manager); diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - params_backend_for(SDBackendModule::DIFFUSION), tensor_storage_map, "model.diffusion_model", version, - sd_ctx_params->chroma_use_dit_mask); + sd_ctx_params->chroma_use_dit_mask, + model_manager); } else if (sd_version_is_ltxav(version)) { cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - params_backend_for(SDBackendModule::TE), - tensor_storage_map); + tensor_storage_map, + "text_encoders.llm", + "text_embedding_projection", + model_manager); diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - params_backend_for(SDBackendModule::DIFFUSION), tensor_storage_map, - "model.diffusion_model"); + "model.diffusion_model", + model_manager); } else if (sd_version_is_wan(version)) { cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - params_backend_for(SDBackendModule::TE), tensor_storage_map, true, 0, - true); + true, + model_manager); diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - params_backend_for(SDBackendModule::DIFFUSION), tensor_storage_map, "model.diffusion_model", - version); + version, + model_manager); if (strlen(SAFE_STR(sd_ctx_params->high_noise_diffusion_model_path)) > 0) { high_noise_diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - params_backend_for(SDBackendModule::DIFFUSION), tensor_storage_map, "model.high_noise_diffusion_model", - version); + version, + model_manager); } if (diffusion_model->get_desc() == "Wan2.1-I2V-14B" || diffusion_model->get_desc() == "Wan2.1-FLF2V-14B" || @@ -635,159 +745,181 @@ public: return false; } clip_vision = std::make_shared(backend_for(SDBackendModule::CLIP_VISION), - params_backend_for(SDBackendModule::CLIP_VISION), - tensor_storage_map); - clip_vision->set_max_graph_vram_bytes(max_graph_vram_bytes); - get_param_tensors(clip_vision, module_can_mmap(SDBackendModule::CLIP_VISION)); + tensor_storage_map, + model_manager); + clip_vision->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::CLIP_VISION)); + if (!register_runner_params("CLIP vision", + clip_vision, + SDBackendModule::CLIP_VISION)) { + return false; + } } } else if (sd_version_is_qwen_image(version)) { - bool enable_vision = false; - if (!vae_decode_only && version != VERSION_QWEN_IMAGE_LAYERED) { - enable_vision = true; - } + bool enable_vision = version != VERSION_QWEN_IMAGE_LAYERED; cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - params_backend_for(SDBackendModule::TE), tensor_storage_map, version, "", - enable_vision); + enable_vision, + model_manager); diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - params_backend_for(SDBackendModule::DIFFUSION), tensor_storage_map, "model.diffusion_model", version, - sd_ctx_params->qwen_image_zero_cond_t); + sd_ctx_params->qwen_image_zero_cond_t, + model_manager); } else if (sd_version_is_longcat(version)) { - bool enable_vision = false; - if (!vae_decode_only) { - enable_vision = true; - } cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - params_backend_for(SDBackendModule::TE), tensor_storage_map, version, "", - enable_vision); + true, + model_manager); diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - params_backend_for(SDBackendModule::DIFFUSION), tensor_storage_map, "model.diffusion_model", version, - sd_ctx_params->chroma_use_dit_mask); + sd_ctx_params->chroma_use_dit_mask, + model_manager); } else if (version == VERSION_HIDREAM_O1) { cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - params_backend_for(SDBackendModule::TE), - tensor_storage_map); + tensor_storage_map, + model_manager); diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - params_backend_for(SDBackendModule::DIFFUSION), tensor_storage_map, - "model"); + "model", + model_manager); + } else if (sd_version_is_minit2i(version)) { + cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), + tensor_storage_map, + model_manager); + diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model.model.net", + model_manager); } else if (sd_version_is_anima(version)) { cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - params_backend_for(SDBackendModule::TE), - tensor_storage_map); + tensor_storage_map, + model_manager); diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - params_backend_for(SDBackendModule::DIFFUSION), tensor_storage_map, - "model.diffusion_model"); + "model.diffusion_model", + model_manager); } else if (sd_version_is_z_image(version)) { cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - params_backend_for(SDBackendModule::TE), tensor_storage_map, - version); + version, + "", + false, + model_manager); diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - params_backend_for(SDBackendModule::DIFFUSION), tensor_storage_map, "model.diffusion_model", - version); + version, + model_manager); + } else if (sd_version_is_boogu_image(version)) { + cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), + tensor_storage_map, + version, + "", + true, + model_manager); + diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + version, + model_manager); } else if (sd_version_is_ernie_image(version)) { cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - params_backend_for(SDBackendModule::TE), tensor_storage_map, - version); + version, + "", + false, + model_manager); diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - params_backend_for(SDBackendModule::DIFFUSION), tensor_storage_map, - "model.diffusion_model"); + "model.diffusion_model", + model_manager); } else if (sd_version_is_lens(version)) { cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - params_backend_for(SDBackendModule::TE), tensor_storage_map, - version); + version, + "", + false, + model_manager); diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - params_backend_for(SDBackendModule::DIFFUSION), tensor_storage_map, - "model.diffusion_model"); + "model.diffusion_model", + model_manager); } else { // SD1.x SD2.x SDXL std::map embbeding_map; for (uint32_t i = 0; i < sd_ctx_params->embedding_count; i++) { embbeding_map.emplace(SAFE_STR(sd_ctx_params->embeddings[i].name), SAFE_STR(sd_ctx_params->embeddings[i].path)); } - if (strstr(SAFE_STR(sd_ctx_params->photo_maker_path), "v2")) { - cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - params_backend_for(SDBackendModule::TE), - tensor_storage_map, - embbeding_map, - version, - PM_VERSION_2); - } else { - cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), - params_backend_for(SDBackendModule::TE), - tensor_storage_map, - embbeding_map, - version); - } - diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), - params_backend_for(SDBackendModule::DIFFUSION), + cond_stage_model = std::make_shared(backend_for(SDBackendModule::TE), + tensor_storage_map, + embbeding_map, + version, + model_manager); + diffusion_model = std::make_shared(backend_for(SDBackendModule::DIFFUSION), tensor_storage_map, "model.diffusion_model", - version); + version, + model_manager); if (sd_ctx_params->diffusion_conv_direct) { LOG_INFO("Using Conv2d direct in the diffusion model"); diffusion_model->set_conv2d_direct_enabled(true); } } - cond_stage_model->set_max_graph_vram_bytes(max_graph_vram_bytes); - get_param_tensors(cond_stage_model, module_can_mmap(SDBackendModule::TE)); + cond_stage_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::TE)); + if (!register_runner_params("Conditioner model", + cond_stage_model, + SDBackendModule::TE, + &text_encoder_params_mem_size)) { + return false; + } - diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes); + diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::DIFFUSION)); diffusion_model->set_stream_layers_enabled(stream_layers); - get_param_tensors(diffusion_model, module_can_mmap(SDBackendModule::DIFFUSION)); - - if (sd_version_is_unet_edit(version)) { - vae_decode_only = false; + if (!register_runner_params("Diffusion model", + diffusion_model, + SDBackendModule::DIFFUSION, + &unet_params_mem_size)) { + return false; } if (high_noise_diffusion_model) { - high_noise_diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes); + high_noise_diffusion_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::DIFFUSION)); high_noise_diffusion_model->set_stream_layers_enabled(stream_layers); - get_param_tensors(high_noise_diffusion_model, module_can_mmap(SDBackendModule::DIFFUSION)); + if (!register_runner_params("High noise diffusion model", + high_noise_diffusion_model, + SDBackendModule::DIFFUSION, + &unet_params_mem_size)) { + return false; + } } if (!ensure_backend_pair(SDBackendModule::VAE)) { return false; } - auto create_tae = [&]() -> std::shared_ptr { - if (sd_version_is_wan(version) || - sd_version_is_qwen_image(version) || - sd_version_is_anima(version) || - sd_version_is_ltxav(version)) { + auto create_tae = [&](bool decode_only) -> std::shared_ptr { + if (sd_version_uses_wan_vae(version) || sd_version_is_ltxav(version)) { return std::make_shared(backend_for(SDBackendModule::VAE), - params_backend_for(SDBackendModule::VAE), tensor_storage_map, "decoder", - vae_decode_only, - version); + decode_only, + version, + model_manager); } else { auto model = std::make_shared(backend_for(SDBackendModule::VAE), - params_backend_for(SDBackendModule::VAE), tensor_storage_map, "decoder.layers", - vae_decode_only, - version); + decode_only, + version, + model_manager); return model; } }; @@ -805,28 +937,26 @@ public: auto create_vae = [&]() -> std::shared_ptr { if (sd_version_is_ltxav(version)) { return std::make_shared(backend_for(SDBackendModule::VAE), - params_backend_for(SDBackendModule::VAE), tensor_storage_map, "first_stage_model", - vae_decode_only, - version); - } else if (sd_version_is_wan(version) || - sd_version_is_qwen_image(version) || - sd_version_is_anima(version)) { + false, + version, + model_manager); + } else if (sd_version_uses_wan_vae(version)) { return std::make_shared(backend_for(SDBackendModule::VAE), - params_backend_for(SDBackendModule::VAE), tensor_storage_map, "first_stage_model", - vae_decode_only, - version); + false, + version, + model_manager); } else { auto model = std::make_shared(backend_for(SDBackendModule::VAE), - params_backend_for(SDBackendModule::VAE), tensor_storage_map, "first_stage_model", - vae_decode_only, false, - vae_version); + false, + vae_version, + model_manager); if (sd_version_is_sdxl(version) && (strlen(SAFE_STR(sd_ctx_params->vae_path)) == 0 || sd_ctx_params->force_sdxl_vae_conv_scale || external_vae_is_invalid)) { float vae_conv_2d_scale = 1.f / 32.f; @@ -840,36 +970,61 @@ public: } }; - bool vae_mmap = module_can_mmap(SDBackendModule::VAE); - - if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1) { + if (version == VERSION_CHROMA_RADIANCE || version == VERSION_HIDREAM_O1 || sd_version_is_minit2i(version)) { LOG_INFO("using FakeVAE"); first_stage_model = std::make_shared(version, backend_for(SDBackendModule::VAE), - params_backend_for(SDBackendModule::VAE)); + model_manager); + if (!register_runner_params("VAE", + first_stage_model, + SDBackendModule::VAE, + &vae_params_mem_size)) { + return false; + } } else if (use_tae && !tae_preview_only) { LOG_INFO("using TAE for encoding / decoding"); - first_stage_model = create_tae(); - first_stage_model->set_max_graph_vram_bytes(max_graph_vram_bytes); - get_param_tensors_p(first_stage_model, vae_mmap, "tae"); + first_stage_model = create_tae(false); + first_stage_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::VAE)); + if (!register_runner_params("VAE", + first_stage_model, + SDBackendModule::VAE, + &vae_params_mem_size)) { + return false; + } } else { LOG_INFO("using VAE for encoding / decoding"); first_stage_model = create_vae(); - first_stage_model->set_max_graph_vram_bytes(max_graph_vram_bytes); - get_param_tensors_p(first_stage_model, vae_mmap, "first_stage_model"); + first_stage_model->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::VAE)); + if (!register_runner_params("VAE", + first_stage_model, + SDBackendModule::VAE, + &vae_params_mem_size)) { + return false; + } if (use_tae && tae_preview_only) { LOG_INFO("using TAE for preview"); - preview_vae = create_tae(); - preview_vae->set_max_graph_vram_bytes(max_graph_vram_bytes); - get_param_tensors_p(preview_vae, vae_mmap, "tae"); + preview_vae = create_tae(true); + preview_vae->set_max_graph_vram_bytes(max_graph_vram_bytes_for_module(SDBackendModule::VAE)); + if (!register_runner_params("preview VAE", + preview_vae, + SDBackendModule::VAE, + &vae_params_mem_size)) { + return false; + } } } if (use_audio_vae) { audio_vae_model = std::make_shared(backend_for(SDBackendModule::VAE), - params_backend_for(SDBackendModule::VAE), - tensor_storage_map); - get_param_tensors_p(audio_vae_model, vae_mmap, ""); + tensor_storage_map, + "", + model_manager); + if (!register_runner_params("LTX audio VAE", + audio_vae_model, + SDBackendModule::VAE, + &vae_params_mem_size)) { + return false; + } } if (sd_ctx_params->vae_conv_direct) { @@ -880,64 +1035,64 @@ public: } } - if (strlen(SAFE_STR(sd_ctx_params->control_net_path)) > 0) { + if (use_control_net) { if (!ensure_backend_pair(SDBackendModule::CONTROL_NET)) { return false; } control_net = std::make_shared(backend_for(SDBackendModule::CONTROL_NET), params_backend_for(SDBackendModule::CONTROL_NET), - tensor_storage_map, - version); + model_loader.get_tensor_storage_map(), + version, + "", + model_manager); if (sd_ctx_params->diffusion_conv_direct) { LOG_INFO("Using Conv2d direct in the control net"); control_net->set_conv2d_direct_enabled(true); } + if (!register_runner_params("ControlNet", + control_net, + SDBackendModule::CONTROL_NET, + &control_net_params_mem_size)) { + return false; + } } - if (strlen(SAFE_STR(sd_ctx_params->photo_maker_path)) > 0) { - if (!ensure_backend_pair(SDBackendModule::PHOTOMAKER)) { - return false; - } - if (strstr(SAFE_STR(sd_ctx_params->photo_maker_path), "v2")) { - pmid_model = std::make_shared(backend_for(SDBackendModule::PHOTOMAKER), - params_backend_for(SDBackendModule::PHOTOMAKER), - tensor_storage_map, - "pmid", - version, - PM_VERSION_2); - LOG_INFO("using PhotoMaker Version 2"); - } else { - pmid_model = std::make_shared(backend_for(SDBackendModule::PHOTOMAKER), - params_backend_for(SDBackendModule::PHOTOMAKER), - tensor_storage_map, - "pmid", - version); - } - pmid_lora = std::make_shared("pmid", - backend_for(SDBackendModule::PHOTOMAKER), - params_backend_for(SDBackendModule::PHOTOMAKER), - sd_ctx_params->photo_maker_path, - "", - version); - auto lora_tensor_filter = [&](const std::string& tensor_name) { - if (starts_with(tensor_name, "lora.model")) { - return true; - } - return false; + { + generation_extensions.clear(); + auto photomaker_extension = create_photomaker_extension(); + GenerationExtensionInitContext extension_ctx{ + sd_ctx_params, + version, + tensor_storage_map, + model_loader, + model_manager, + n_threads, + [this](SDBackendModule module) { return ensure_backend_pair(module); }, + [this](SDBackendModule module) { return backend_for(module); }, + [this](SDBackendModule module) { return params_backend_for(module); }, }; - if (!pmid_lora->load_from_file(n_threads, lora_tensor_filter)) { - LOG_WARN("load photomaker lora tensors from %s failed", sd_ctx_params->photo_maker_path); + if (!photomaker_extension->init(extension_ctx)) { return false; } - LOG_INFO("loading stacked ID embedding (PHOTOMAKER) model file from '%s'", sd_ctx_params->photo_maker_path); - if (!model_loader.init_from_file_and_convert_name(sd_ctx_params->photo_maker_path, "pmid.")) { - LOG_WARN("loading stacked ID embedding from '%s' failed", sd_ctx_params->photo_maker_path); - } else { - use_pmid = true; + if (photomaker_extension->is_enabled()) { + generation_extensions.push_back(photomaker_extension); + } + + auto pulid_extension = create_pulid_extension(); + if (!pulid_extension->init(extension_ctx)) { + return false; + } + if (pulid_extension->is_enabled()) { + generation_extensions.push_back(pulid_extension); } } - if (use_pmid) { - get_param_tensors_p(pmid_model, module_can_mmap(SDBackendModule::PHOTOMAKER), "pmid"); + for (auto& extension : generation_extensions) { + if (!register_runner_params(extension->name(), + extension, + SDBackendModule::PHOTOMAKER, + &extension_params_mem_size)) { + return false; + } } if (sd_ctx_params->flash_attn) { @@ -973,38 +1128,19 @@ public: circular_y = sd_ctx_params->circular_y; } - ggml_init_params params; - params.mem_size = static_cast(10 * 1024) * 1024; // 10M - params.mem_buffer = nullptr; - params.no_alloc = false; - // LOG_DEBUG("mem_size %u ", params.mem_size); - ggml_context* ctx = ggml_init(params); // for alphas_cumprod and is_using_v_parameterization check - GGML_ASSERT(ctx != nullptr); - ggml_tensor* alphas_cumprod_tensor = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, TIMESTEPS); - calculate_alphas_cumprod((float*)alphas_cumprod_tensor->data); - - // load weights - LOG_DEBUG("loading weights"); + LOG_DEBUG("validating model metadata"); std::set ignore_tensors; - tensors["alphas_cumprod"] = alphas_cumprod_tensor; if (use_tae && !tae_preview_only) { ignore_tensors.insert("first_stage_model."); } - if (use_pmid) { - ignore_tensors.insert("pmid.unet."); + for (auto& extension : generation_extensions) { + extension->add_ignore_tensors(ignore_tensors); } ignore_tensors.insert("model.diffusion_model.__x0__"); ignore_tensors.insert("model.diffusion_model.__32x32__"); ignore_tensors.insert("model.diffusion_model.__index_timestep_zero__"); - if (vae_decode_only) { - ignore_tensors.insert("first_stage_model.encoder"); - ignore_tensors.insert("first_stage_model.conv1"); - ignore_tensors.insert("first_stage_model.quant"); - ignore_tensors.insert("tae.encoder"); - ignore_tensors.insert("text_encoders.llm.visual."); - } if (audio_vae_model) { ignore_tensors.insert("audio_vae.encoder"); } @@ -1025,96 +1161,34 @@ public: ignore_tensors.insert("text_encoders.llm.model.layers.0.mlp.experts.gate_up_proj.weight_scale_2"); ignore_tensors.insert("text_encoders.llm.model.layers.0.mlp.experts.down_proj.weight_scale_2"); } + if (sd_version_is_ideogram4(version)) { + ignore_tensors.insert("text_encoders.llm.lm_head."); + ignore_tensors.insert("text_encoders.llm.visual."); + ignore_tensors.insert("text_encoders.llm.vision_model."); + ignore_tensors.insert("text_encoders.llm.tokenizer_json"); + } if (version == VERSION_HIDREAM_O1) { ignore_tensors.insert("lm_head."); ignore_tensors.insert("model.visual.deepstack_merger_list."); } - if (enable_mmap_tensors) { - if (mmap_able_tensors.empty()) { - LOG_DEBUG("no tensors could be memory-mapped"); - } else { - mmap_tensor_store = model_loader.mmap_tensors(mmap_able_tensors, ignore_tensors, needs_writable_mmap); + model_manager->set_common_ignore_tensors(ignore_tensors); + if (!model_manager->validate_registered_tensors()) { + LOG_ERROR("model metadata validation failed"); + return false; + } + + if (eager_load) { + if (!model_manager->load_all_params_eagerly()) { + LOG_ERROR("model params eager load failed"); + return false; } + LOG_DEBUG("model metadata validated; weights pre-loaded to params backend"); + } else { + LOG_DEBUG("model metadata validated; weights will be prepared lazily"); } - if (clip_vision && !clip_vision->alloc_params_buffer()) { - LOG_ERROR("CLIP vision params buffer allocation failed"); - ggml_free(ctx); - return false; - } - if (cond_stage_model && !cond_stage_model->alloc_params_buffer()) { - LOG_ERROR("Conditioner model params buffer allocation failed"); - ggml_free(ctx); - return false; - } - if (diffusion_model && !diffusion_model->alloc_params_buffer()) { - LOG_ERROR("Diffusion model params buffer allocation failed"); - ggml_free(ctx); - return false; - } - if (high_noise_diffusion_model && !high_noise_diffusion_model->alloc_params_buffer()) { - LOG_ERROR("High noise diffusion model params buffer allocation failed"); - ggml_free(ctx); - return false; - } - if (first_stage_model && !first_stage_model->alloc_params_buffer()) { - LOG_ERROR("VAE params buffer allocation failed"); - ggml_free(ctx); - return false; - } - if (preview_vae && !preview_vae->alloc_params_buffer()) { - LOG_ERROR("preview VAE params buffer allocation failed"); - ggml_free(ctx); - return false; - } - if (audio_vae_model && !audio_vae_model->alloc_params_buffer()) { - LOG_ERROR("LTX audio VAE params buffer allocation failed"); - ggml_free(ctx); - return false; - } - if (use_pmid && pmid_model && !pmid_model->alloc_params_buffer()) { - LOG_ERROR("PhotoMaker params buffer allocation failed"); - ggml_free(ctx); - return false; - } - - bool success = model_loader.load_tensors(tensors, ignore_tensors, n_threads, sd_ctx_params->enable_mmap); - if (!success) { - LOG_ERROR("load tensors from model loader failed"); - ggml_free(ctx); - return false; - } - - LOG_DEBUG("finished loaded file"); - { - size_t clip_params_mem_size = cond_stage_model->get_params_buffer_size(); - size_t unet_params_mem_size = diffusion_model->get_params_buffer_size(); - if (high_noise_diffusion_model) { - unet_params_mem_size += high_noise_diffusion_model->get_params_buffer_size(); - } - size_t vae_params_mem_size = 0; - vae_params_mem_size = first_stage_model->get_params_buffer_size(); - if (preview_vae) { - vae_params_mem_size += preview_vae->get_params_buffer_size(); - } - if (audio_vae_model) { - vae_params_mem_size += audio_vae_model->get_params_buffer_size(); - } - size_t control_net_params_mem_size = 0; - if (control_net) { - if (!control_net->load_from_file(SAFE_STR(sd_ctx_params->control_net_path), n_threads)) { - ggml_free(ctx); - return false; - } - control_net_params_mem_size = control_net->get_params_buffer_size(); - } - size_t pmid_params_mem_size = 0; - if (use_pmid) { - pmid_params_mem_size = pmid_model->get_params_buffer_size(); - } - size_t total_params_ram_size = 0; size_t total_params_vram_size = 0; auto add_params_memory = [&](size_t size, SDBackendModule module) { @@ -1143,32 +1217,31 @@ public: return sd_backend_is_cpu(module_backend) ? "RAM" : "VRAM"; }; - if (!add_params_memory(clip_params_mem_size, SDBackendModule::TE) || - !add_params_memory(pmid_params_mem_size, SDBackendModule::PHOTOMAKER) || + if (!add_params_memory(text_encoder_params_mem_size, SDBackendModule::TE) || + !add_params_memory(extension_params_mem_size, SDBackendModule::PHOTOMAKER) || !add_params_memory(unet_params_mem_size, SDBackendModule::DIFFUSION) || !add_params_memory(vae_params_mem_size, SDBackendModule::VAE) || !add_params_memory(control_net_params_mem_size, SDBackendModule::CONTROL_NET)) { - ggml_free(ctx); return false; } size_t total_params_size = total_params_ram_size + total_params_vram_size; LOG_INFO( "total params memory size = %.2fMB (VRAM %.2fMB, RAM %.2fMB): " - "text_encoders %.2fMB(%s), diffusion_model %.2fMB(%s), vae %.2fMB(%s), controlnet %.2fMB(%s), pmid %.2fMB(%s)", + "text_encoders %.2fMB(%s), diffusion_model %.2fMB(%s), vae %.2fMB(%s), controlnet %.2fMB(%s), extensions %.2fMB(%s)", total_params_size / 1024.0 / 1024.0, total_params_vram_size / 1024.0 / 1024.0, total_params_ram_size / 1024.0 / 1024.0, - clip_params_mem_size / 1024.0 / 1024.0, - params_memory_location(clip_params_mem_size, SDBackendModule::TE), + text_encoder_params_mem_size / 1024.0 / 1024.0, + params_memory_location(text_encoder_params_mem_size, SDBackendModule::TE), unet_params_mem_size / 1024.0 / 1024.0, params_memory_location(unet_params_mem_size, SDBackendModule::DIFFUSION), vae_params_mem_size / 1024.0 / 1024.0, params_memory_location(vae_params_mem_size, SDBackendModule::VAE), control_net_params_mem_size / 1024.0 / 1024.0, params_memory_location(control_net_params_mem_size, SDBackendModule::CONTROL_NET), - pmid_params_mem_size / 1024.0 / 1024.0, - params_memory_location(pmid_params_mem_size, SDBackendModule::PHOTOMAKER)); + extension_params_mem_size / 1024.0 / 1024.0, + params_memory_location(extension_params_mem_size, SDBackendModule::PHOTOMAKER)); } // init denoiser @@ -1177,12 +1250,7 @@ public: if (pred_type == PREDICTION_COUNT) { if (sd_version_is_sd2(version)) { - // check is_using_v_parameterization_for_sd2 - if (is_using_v_parameterization_for_sd2(sd_version_is_inpaint(version))) { - pred_type = V_PRED; - } else { - pred_type = EPS_PRED; - } + pred_type = is_using_v_parameterization_for_sd2(sd_version_is_inpaint(version)) ? V_PRED : EPS_PRED; } else if (sd_version_is_sdxl(version)) { if (tensor_storage_map.find("edm_vpred.sigma_max") != tensor_storage_map.end()) { // CosXL models @@ -1200,7 +1268,9 @@ public: sd_version_is_anima(version) || sd_version_is_ernie_image(version) || sd_version_is_z_image(version) || - sd_version_is_pid(version)) { + sd_version_is_boogu_image(version) || + sd_version_is_pid(version) || + sd_version_is_ideogram4(version)) { pred_type = FLOW_PRED; if (sd_version_is_wan(version)) { default_flow_shift = 5.f; @@ -1208,13 +1278,19 @@ public: default_flow_shift = 4.f; } else if (sd_version_is_pid(version)) { default_flow_shift = 1.5f; + } else if (sd_version_is_ideogram4(version)) { + default_flow_shift = 1.0f; + } else if (sd_version_is_boogu_image(version)) { + default_flow_shift = 3.16f; } else { default_flow_shift = 3.f; } } else if (sd_version_is_flux(version) || + sd_version_is_flux2(version) || sd_version_is_longcat(version) || sd_version_is_lens(version) || - sd_version_is_ltxav(version)) { + sd_version_is_ltxav(version) || + sd_version_is_krea2(version)) { pred_type = FLUX_FLOW_PRED; default_flow_shift = 1.0f; // TODO: validate @@ -1230,9 +1306,13 @@ public: default_flow_shift = 1.83f; } else if (sd_version_is_ltxav(version)) { default_flow_shift = 2.37f; + } else if (sd_version_is_krea2(version)) { + default_flow_shift = 1.15f; } - } else if (sd_version_is_flux2(version)) { - pred_type = FLUX2_FLOW_PRED; + } else if (sd_version_is_sefi_image(version)) { + pred_type = SEFI_FLOW_PRED; + } else if (sd_version_is_minit2i(version)) { + pred_type = MINIT2I_FLOW_PRED; } else { pred_type = EPS_PRED; } @@ -1265,32 +1345,39 @@ public: denoiser = std::make_shared(); break; } - case FLUX2_FLOW_PRED: { - LOG_INFO("running in Flux2 FLOW mode"); - denoiser = std::make_shared(); + case SEFI_FLOW_PRED: { + LOG_INFO("running in SeFi-Image dual-time FLOW mode"); + denoiser = std::make_shared(); + break; + } + case MINIT2I_FLOW_PRED: { + LOG_INFO("running in MiniT2I FLOW mode"); + denoiser = std::make_shared(); break; } default: { LOG_ERROR("Unknown predition type %i", pred_type); - ggml_free(ctx); return false; } } - auto comp_vis_denoiser = std::dynamic_pointer_cast(denoiser); - if (comp_vis_denoiser) { - for (int i = 0; i < TIMESTEPS; i++) { - comp_vis_denoiser->sigmas[i] = std::sqrt((1 - ((float*)alphas_cumprod_tensor->data)[i]) / ((float*)alphas_cumprod_tensor->data)[i]); - comp_vis_denoiser->log_sigmas[i] = std::log(comp_vis_denoiser->sigmas[i]); - } - } + refresh_compvis_denoiser_sigmas(); } - ggml_free(ctx); return true; } bool is_using_v_parameterization_for_sd2(bool is_inpaint = false) { + struct RunnerDoneOnExit { + GGMLRunner* runner = nullptr; + ~RunnerDoneOnExit() { + if (runner != nullptr) { + runner->runner_done(); + } + } + }; + RunnerDoneOnExit diffusion_runner_done{diffusion_model.get()}; + sd::Tensor x_t = sd::full({8, 8, 4, 1}, 0.5f); sd::Tensor c = sd::full({1024, 2, 1, 1}, 0.5f); sd::Tensor steps = sd::full({1}, 999.0f); @@ -1312,7 +1399,6 @@ public: auto out_opt = diffusion_model->compute(n_threads, diffusion_params); GGML_ASSERT(!out_opt.empty()); out = std::move(out_opt); - diffusion_model->free_compute_buffer(); double result = static_cast((out - x_t).mean()); int64_t t1 = ggml_time_ms(); @@ -1320,84 +1406,41 @@ public: return result < -1; } - std::shared_ptr load_lora_model_from_file(const std::string& lora_id, - float multiplier, - SDBackendModule module, - LoraModel::filter_t lora_tensor_filter = nullptr) { - std::string lora_path = lora_id; - static std::string high_noise_tag = "|high_noise|"; - bool is_high_noise = false; - if (starts_with(lora_path, high_noise_tag)) { - lora_path = lora_path.substr(high_noise_tag.size()); - is_high_noise = true; - LOG_DEBUG("high noise lora: %s", lora_path.c_str()); - } + static std::string lora_log_id(const ModelManager::LoraSpec& lora) { + return lora.is_high_noise ? "|high_noise|" + lora.path : lora.path; + } + + std::shared_ptr load_lora_model(const ModelManager::LoraSpec& lora_spec, + SDBackendModule module, + LoraModel::filter_t module_filter = nullptr) { if (!ensure_backend_pair(module)) { return nullptr; } - auto lora = std::make_shared(lora_id, + if (lora_spec.is_high_noise) { + LOG_DEBUG("high noise lora: %s", lora_spec.path.c_str()); + } + auto lora = std::make_shared(lora_log_id(lora_spec), backend_for(module), backend_for(module), - lora_path, - is_high_noise ? "model.high_noise_" : "", + lora_spec.path, + lora_spec.is_high_noise ? "model.high_noise_" : "", version); + LoraModel::filter_t lora_tensor_filter = module_filter; + if (!lora_spec.tensor_name_prefix_filter.empty()) { + lora_tensor_filter = [module_filter, prefix = lora_spec.tensor_name_prefix_filter](const std::string& tensor_name) { + return starts_with(tensor_name, prefix) && (!module_filter || module_filter(tensor_name)); + }; + } if (!lora->load_from_file(n_threads, lora_tensor_filter)) { - LOG_WARN("load lora tensors from %s failed", lora_path.c_str()); + LOG_WARN("load lora tensors from %s failed", lora_spec.path.c_str()); return nullptr; } - lora->multiplier = multiplier; + lora->multiplier = lora_spec.multiplier; return lora; } - void apply_loras_immediately(const std::unordered_map& lora_state) { - std::unordered_map lora_state_diff; - for (auto& kv : lora_state) { - const std::string& lora_name = kv.first; - float multiplier = kv.second; - lora_state_diff[lora_name] += multiplier; - } - for (auto& kv : curr_lora_state) { - const std::string& lora_name = kv.first; - float curr_multiplier = kv.second; - lora_state_diff[lora_name] -= curr_multiplier; - } - - if (lora_state_diff.empty()) { - return; - } - - LOG_INFO("apply lora immediately"); - - size_t rm = lora_state_diff.size() - lora_state.size(); - if (rm != 0) { - LOG_INFO("attempting to apply %lu LoRAs (removing %lu applied LoRAs)", lora_state.size(), rm); - } else { - LOG_INFO("attempting to apply %lu LoRAs", lora_state.size()); - } - - for (auto& kv : lora_state_diff) { - int64_t t0 = ggml_time_ms(); - - auto lora = load_lora_model_from_file(kv.first, kv.second, SDBackendModule::DIFFUSION); - if (!lora || lora->lora_tensors.empty()) { - continue; - } - lora->apply(tensors, version, n_threads); - lora->free_params_buffer(); - - int64_t t1 = ggml_time_ms(); - - LOG_INFO("lora '%s' applied, taking %.2fs", kv.first.c_str(), (t1 - t0) * 1.0f / 1000); - } - - curr_lora_state = lora_state; - } - - void apply_loras_at_runtime(const std::unordered_map& lora_state) { - cond_stage_lora_models.clear(); - diffusion_lora_models.clear(); - first_stage_lora_models.clear(); + void clear_lora_adapters() { if (cond_stage_model) { cond_stage_model->set_weight_adapter(nullptr); } @@ -1410,39 +1453,74 @@ public: if (first_stage_model) { first_stage_model->set_weight_adapter(nullptr); } - if (lora_state.empty()) { + } + + std::vector> load_runtime_loras_for_module(const std::vector& loras, + const std::set& model_tensor_names, + SDBackendModule module, + LoraModel::filter_t module_filter = nullptr) { + std::vector> module_lora_models; + for (const auto& lora_spec : loras) { + auto lora = load_lora_model(lora_spec, module, module_filter); + if (lora == nullptr) { + if (lora_spec.required) { + LOG_ERROR("required lora load failed: %s", lora_spec.path.c_str()); + } + continue; + } + if (lora->lora_tensors.empty()) { + continue; + } + + lora->preprocess_lora_tensors(model_tensor_names); + runtime_lora_models.push_back(lora); + module_lora_models.push_back(std::move(lora)); + } + return module_lora_models; + } + + void apply_loras_immediately(const std::vector& loras) { + if (model_manager == nullptr) { + if (!loras.empty()) { + LOG_WARN("model manager is not available for immediate lora"); + } return; } + + clear_lora_adapters(); + runtime_lora_models.clear(); + + model_manager->set_loras(loras, version); + } + + void apply_loras_at_runtime(const std::vector& loras) { + if (model_manager != nullptr) { + model_manager->set_loras({}, version); + } + runtime_lora_models.clear(); + clear_lora_adapters(); + if (loras.empty()) { + return; + } + + std::set model_tensor_names; + if (model_manager != nullptr) { + model_tensor_names = model_manager->tensor_names(); + } + LOG_INFO("apply lora at runtime"); if (cond_stage_model) { - std::vector> lora_models; - auto lora_state_diff = lora_state; - for (auto& lora_model : cond_stage_lora_models) { - auto iter = lora_state_diff.find(lora_model->lora_id); - - if (iter != lora_state_diff.end()) { - lora_model->multiplier = iter->second; - lora_models.push_back(lora_model); - lora_state_diff.erase(iter); - } - } - cond_stage_lora_models = lora_models; auto lora_tensor_filter = [&](const std::string& tensor_name) { if (is_cond_stage_model_name(tensor_name)) { return true; } return false; }; - for (auto& kv : lora_state_diff) { - const std::string& lora_id = kv.first; - float multiplier = kv.second; - - auto lora = load_lora_model_from_file(lora_id, multiplier, SDBackendModule::TE, lora_tensor_filter); - if (lora && !lora->lora_tensors.empty()) { - lora->preprocess_lora_tensors(tensors); - cond_stage_lora_models.push_back(lora); - } - } + auto cond_stage_lora_models = + load_runtime_loras_for_module(loras, + model_tensor_names, + SDBackendModule::TE, + lora_tensor_filter); // Only attach the adapter when there are LoRAs targeting the cond_stage model. // An empty MultiLoraAdapter still routes every linear/conv through // forward_with_lora() instead of the direct kernel path — slower for no benefit. @@ -1452,34 +1530,17 @@ public: } } if (diffusion_model) { - std::vector> lora_models; - auto lora_state_diff = lora_state; - for (auto& lora_model : diffusion_lora_models) { - auto iter = lora_state_diff.find(lora_model->lora_id); - - if (iter != lora_state_diff.end()) { - lora_model->multiplier = iter->second; - lora_models.push_back(lora_model); - lora_state_diff.erase(iter); - } - } - diffusion_lora_models = lora_models; auto lora_tensor_filter = [&](const std::string& tensor_name) { if (is_diffusion_model_name(tensor_name)) { return true; } return false; }; - for (auto& kv : lora_state_diff) { - const std::string& lora_name = kv.first; - float multiplier = kv.second; - - auto lora = load_lora_model_from_file(lora_name, multiplier, SDBackendModule::DIFFUSION, lora_tensor_filter); - if (lora && !lora->lora_tensors.empty()) { - lora->preprocess_lora_tensors(tensors); - diffusion_lora_models.push_back(lora); - } - } + auto diffusion_lora_models = + load_runtime_loras_for_module(loras, + model_tensor_names, + SDBackendModule::DIFFUSION, + lora_tensor_filter); if (!diffusion_lora_models.empty()) { auto multi_lora_adapter = std::make_shared(diffusion_lora_models); diffusion_model->set_weight_adapter(multi_lora_adapter); @@ -1490,34 +1551,17 @@ public: } if (first_stage_model) { - std::vector> lora_models; - auto lora_state_diff = lora_state; - for (auto& lora_model : first_stage_lora_models) { - auto iter = lora_state_diff.find(lora_model->lora_id); - - if (iter != lora_state_diff.end()) { - lora_model->multiplier = iter->second; - lora_models.push_back(lora_model); - lora_state_diff.erase(iter); - } - } - first_stage_lora_models = lora_models; auto lora_tensor_filter = [&](const std::string& tensor_name) { if (is_first_stage_model_name(tensor_name)) { return true; } return false; }; - for (auto& kv : lora_state_diff) { - const std::string& lora_name = kv.first; - float multiplier = kv.second; - - auto lora = load_lora_model_from_file(lora_name, multiplier, SDBackendModule::VAE, lora_tensor_filter); - if (lora && !lora->lora_tensors.empty()) { - lora->preprocess_lora_tensors(tensors); - first_stage_lora_models.push_back(lora); - } - } + auto first_stage_lora_models = + load_runtime_loras_for_module(loras, + model_tensor_names, + SDBackendModule::VAE, + lora_tensor_filter); if (!first_stage_lora_models.empty()) { auto multi_lora_adapter = std::make_shared(first_stage_lora_models); first_stage_model->set_weight_adapter(multi_lora_adapter); @@ -1526,132 +1570,69 @@ public: } void lora_stat() { - if (!cond_stage_lora_models.empty()) { - LOG_INFO("cond_stage_lora_models:"); - for (auto& lora_model : cond_stage_lora_models) { - lora_model->stat(); - } - } - - if (!diffusion_lora_models.empty()) { - LOG_INFO("diffusion_lora_models:"); - for (auto& lora_model : diffusion_lora_models) { - lora_model->stat(); - } - } - - if (!first_stage_lora_models.empty()) { - LOG_INFO("first_stage_lora_models:"); - for (auto& lora_model : first_stage_lora_models) { + if (!runtime_lora_models.empty()) { + LOG_INFO("runtime_lora_models:"); + for (auto& lora_model : runtime_lora_models) { lora_model->stat(); } } } void apply_loras(const sd_lora_t* loras, uint32_t lora_count) { - std::unordered_map lora_f2m; + std::vector all_loras; + all_loras.reserve(lora_count); for (uint32_t i = 0; i < lora_count; i++) { std::string lora_id = SAFE_STR(loras[i].path); + ModelManager::LoraSpec lora_spec; + lora_spec.path = lora_id; + lora_spec.multiplier = loras[i].multiplier; + lora_spec.is_high_noise = loras[i].is_high_noise; + all_loras.push_back(std::move(lora_spec)); if (loras[i].is_high_noise) { lora_id = "|high_noise|" + lora_id; } - lora_f2m[lora_id] = loras[i].multiplier; LOG_DEBUG("lora %s:%.2f", lora_id.c_str(), loras[i].multiplier); } + + for (auto& extension : generation_extensions) { + extension->collect_loras(all_loras); + } + int64_t t0 = ggml_time_ms(); if (apply_lora_immediately) { - apply_loras_immediately(lora_f2m); + apply_loras_immediately(all_loras); } else { - apply_loras_at_runtime(lora_f2m); + apply_loras_at_runtime(all_loras); } int64_t t1 = ggml_time_ms(); - if (!lora_f2m.empty()) { + if (!all_loras.empty()) { LOG_INFO("apply_loras completed, taking %.2fs", (t1 - t0) * 1.0f / 1000); } } - SDCondition get_pmid_conditon(sd_pm_params_t pm_params, - ConditionerParams& condition_params) { - SDCondition id_cond; - if (use_pmid) { - if (!pmid_lora->applied) { - int64_t t0 = ggml_time_ms(); - pmid_lora->apply(tensors, version, n_threads); - int64_t t1 = ggml_time_ms(); - pmid_lora->applied = true; - LOG_INFO("pmid_lora apply completed, taking %.2fs", (t1 - t0) * 1.0f / 1000); - if (free_params_immediately) { - pmid_lora->free_params_buffer(); - } - } - // preprocess input id images - bool pmv2 = pmid_model->get_version() == PM_VERSION_2; - if (pm_params.id_images_count > 0) { - int clip_image_size = 224; - pmid_model->style_strength = pm_params.style_strength; - sd::Tensor id_image_tensor; - for (int i = 0; i < pm_params.id_images_count; i++) { - auto id_image = sd_image_to_tensor(pm_params.id_images[i]); - auto processed_id_image = clip_preprocess(id_image, clip_image_size, clip_image_size); - if (id_image_tensor.empty()) { - id_image_tensor = processed_id_image; - } else { - id_image_tensor = sd::ops::concat(id_image_tensor, processed_id_image, 3); - } - } - - int64_t t0 = ggml_time_ms(); - condition_params.num_input_imgs = pm_params.id_images_count; - auto cond_tup = cond_stage_model->get_learned_condition_with_trigger(n_threads, - condition_params); - id_cond = std::get<0>(cond_tup); - auto class_tokens_mask = std::get<1>(cond_tup); - sd::Tensor id_embeds; - if (pmv2 && pm_params.id_embed_path != nullptr) { - try { - id_embeds = sd::load_tensor_from_file_as_tensor(pm_params.id_embed_path); - } catch (const std::exception&) { - id_embeds = {}; - } - } - if (pmv2 && id_embeds.empty()) { - LOG_WARN("Provided PhotoMaker images, but NO valid ID embeds file for PM v2"); - LOG_WARN("Turn off PhotoMaker"); - use_pmid = false; - } else { - if (pmv2 && pm_params.id_images_count != id_embeds.shape()[1]) { - LOG_WARN("PhotoMaker image count (%d) does NOT match ID embeds (%d). You should run face_detect.py again.", pm_params.id_images_count, static_cast(id_embeds.shape()[1])); - LOG_WARN("Turn off PhotoMaker"); - use_pmid = false; - } else { - auto res = pmid_model->compute(n_threads, - id_image_tensor, - id_cond.c_crossattn, - id_embeds, - class_tokens_mask); - if (res.empty()) { - LOG_ERROR("Photomaker ID Stacking failed"); - LOG_WARN("Turn off PhotoMaker"); - use_pmid = false; - } else { - id_cond.c_crossattn = std::move(res); - int64_t t1 = ggml_time_ms(); - LOG_INFO("Photomaker ID Stacking, taking %" PRId64 " ms", t1 - t0); - // Encode input prompt without the trigger word for delayed conditioning - condition_params.text = cond_stage_model->remove_trigger_from_prompt(condition_params.text); - } - if (free_params_immediately) { - pmid_model->free_params_buffer(); - } - } - } - } else { - LOG_WARN("Provided PhotoMaker model file, but NO input ID images"); - LOG_WARN("Turn off PhotoMaker"); - use_pmid = false; - } + void reset_generation_extensions() { + for (auto& extension : generation_extensions) { + extension->reset_runtime_condition(); + } + } + + void prepare_generation_extensions(const sd_pm_params_t& pm_params, + const sd_pulid_params_t& pulid_params, + ConditionerParams& condition_params, + int total_steps) { + reset_generation_extensions(); + GenerationExtensionConditionContext ctx{ + cond_stage_model.get(), + condition_params, + pm_params, + pulid_params, + n_threads, + total_steps, + }; + + for (auto& extension : generation_extensions) { + extension->prepare_condition(ctx); } - return id_cond; } sd::Tensor get_clip_vision_output(const sd::Tensor& image, @@ -1679,7 +1660,16 @@ public: std::vector process_timesteps(const std::vector& timesteps, const sd::Tensor& init_latent, - const sd::Tensor& denoise_mask) { + const sd::Tensor& denoise_mask, + int step) { + if (auto sefi_denoiser = std::dynamic_pointer_cast(denoiser)) { + int sched_idx = step > 0 ? step - 1 : 0; + if (sched_idx >= static_cast(sefi_denoiser->tex_timesteps.size())) { + sched_idx = static_cast(sefi_denoiser->tex_timesteps.size()) - 1; + } + return {sefi_denoiser->sem_timesteps[sched_idx], + sefi_denoiser->tex_timesteps[sched_idx]}; + } if (diffusion_model->get_desc() == "Wan2.2-TI2V-5B") { int64_t frame_count = init_latent.shape()[2]; auto new_timesteps = std::vector(static_cast(frame_count), timesteps[0]); @@ -1772,10 +1762,10 @@ public: if (sd_version_is_sd3(version)) { latent_rgb_proj = sd3_latent_rgb_proj; latent_rgb_bias = sd3_latent_rgb_bias; - } else if (sd_version_is_flux(version) || sd_version_is_z_image(version) || sd_version_is_longcat(version)) { + } else if (sd_version_uses_flux_vae(version)) { latent_rgb_proj = flux_latent_rgb_proj; latent_rgb_bias = flux_latent_rgb_bias; - } else if (sd_version_is_wan(version) || sd_version_is_qwen_image(version) || sd_version_is_anima(version)) { + } else if (sd_version_uses_wan_vae(version)) { latent_rgb_proj = wan_21_latent_rgb_proj; latent_rgb_bias = wan_21_latent_rgb_bias; } else { @@ -1867,10 +1857,13 @@ public: if (sd_version_is_anima(version)) { return std::vector{t / static_cast(TIMESTEPS)}; } + if (sd_version_is_boogu_image(version)) { + return std::vector{t / static_cast(TIMESTEPS)}; + } if (version == VERSION_HIDREAM_O1) { return std::vector{1.0f - (t / static_cast(TIMESTEPS))}; } - if (sd_version_is_z_image(version)) { + if (sd_version_is_z_image(version) || sd_version_is_ideogram4(version)) { return std::vector{1000.f - t}; } return std::vector{t}; @@ -1910,11 +1903,17 @@ public: sd_get_preview_mode()}; } - void report_sample_progress(int step, size_t total_steps, int64_t t0) { - int64_t t1 = ggml_time_us(); + void report_sample_progress(int step, size_t total_steps, int64_t* last_progress_us) { if (step > 0 || step == -(int)total_steps) { - int showstep = std::abs(step); - pretty_progress(showstep, (int)total_steps, (t1 - t0) / 1000000.f / showstep); + int64_t now = ggml_time_us(); + int showstep = std::abs(step); + float step_seconds = last_progress_us != nullptr && *last_progress_us > 0 + ? (now - *last_progress_us) / 1000000.f + : 0.f; + pretty_progress(showstep, (int)total_steps, step_seconds); + if (last_progress_us != nullptr) { + *last_progress_us = now; + } } } @@ -1950,7 +1949,6 @@ public: const SDCondition& cond, const SDCondition& uncond, const SDCondition& img_uncond, - const SDCondition& id_cond, const sd::Tensor& control_image, float control_strength, const sd_guidance_params_t& guidance, @@ -1960,7 +1958,6 @@ public: bool is_flow_denoiser, const char* extra_sample_args, const std::vector& sigmas, - int start_merge_step, const std::vector>& ref_latents, bool increase_ref_index, const sd::Tensor& denoise_mask, @@ -1970,12 +1967,50 @@ public: float frame_rate, const sd_cache_params_t* cache_params, const sd::Tensor& video_positions = {}) { + struct RunnerDoneOnExit { + GGMLRunner* runner = nullptr; + ~RunnerDoneOnExit() { + if (runner != nullptr) { + runner->runner_done(); + } + } + }; + RunnerDoneOnExit sample_diffusion_runner_done{work_diffusion_model.get()}; + + RunnerDoneOnExit sample_control_runner_done{!control_image.empty() && control_net != nullptr ? control_net.get() : nullptr}; + std::vector skip_layers(guidance.slg.layers, guidance.slg.layers + guidance.slg.layer_count); float cfg_scale = guidance.txt_cfg; float img_cfg_scale = guidance.img_cfg; float slg_scale = guidance.slg.scale; bool slg_uncond = sd::guidance::parse_skip_layer_guidance_uncond_arg(extra_sample_args); + std::vector guidance_schedule = sd::guidance::parse_guidance_schedule(extra_sample_args); + if (!guidance_schedule.empty() && guidance_schedule.size() != sigmas.size() - 1) { + if (guidance_schedule.size() > sigmas.size()) { + LOG_WARN("guidance_schedule length (%zu) is greater than number of steps (%zu)", guidance_schedule.size(), sigmas.size() - 1); + LOG_WARN("truncating guidance_schedule to match step count"); + guidance_schedule.resize(sigmas.size() - 1); + } else { + LOG_INFO("padding guidance_schedule with cfg_scale"); + while (guidance_schedule.size() < sigmas.size() - 1) { + guidance_schedule.push_back(cfg_scale); + } + } + } + + if (!guidance_schedule.empty()) { + std::string schedule_str = "["; + for (size_t i = 0; i < guidance_schedule.size(); ++i) { + schedule_str += std::to_string(guidance_schedule[i]); + if (i < guidance_schedule.size() - 1) { + schedule_str += ", "; + } + } + schedule_str += "]"; + LOG_DEBUG("using guidance schedule: %s", schedule_str.c_str()); + } + sd_sample::SampleCacheRuntime cache_runtime = sd_sample::init_sample_cache_runtime(version, cache_params, denoiser.get(), @@ -2015,16 +2050,23 @@ public: noise *= eta; } - int64_t t0 = ggml_time_us(); - sd::Tensor x_t = !noise.empty() - ? denoiser->noise_scaling(sigmas[0], noise, init_latent) - : init_latent; - sd::Tensor denoised = x_t; + int64_t last_progress_us = ggml_time_us(); SamplePreviewContext preview = prepare_sample_preview_context(); + sd::Tensor x_t = !noise.empty() + ? denoiser->noise_scaling(sigmas[0], noise, init_latent) + : init_latent; + sd::Tensor denoised = x_t; + auto denoise = [&](const sd::Tensor& x, float sigma, int step) -> sd::guidance::GuiderOutput { + if (get_cancel_flag() == SD_CANCEL_ALL) { + LOG_DEBUG("cancelling generation"); + return {}; + } + if (step == 1 || step == -1) { pretty_progress(0, (int)steps, 0); + last_progress_us = ggml_time_us(); } std::vector scaling = denoiser->get_scalings(sigma); @@ -2040,7 +2082,7 @@ public: timesteps_vec = process_ltxav_video_timesteps(base_timesteps_vec, init_latent, denoise_mask); audio_timesteps_tensor = sd::Tensor({static_cast(base_timesteps_vec.size())}, base_timesteps_vec); } else { - timesteps_vec = process_timesteps(timesteps_vec, init_latent, denoise_mask); + timesteps_vec = process_timesteps(timesteps_vec, init_latent, denoise_mask, step); } const std::vector& scaling_timesteps_vec = (sd_version_is_ltxav(version) && !denoise_mask.empty()) ? base_timesteps_vec @@ -2062,7 +2104,7 @@ public: if (sd_should_preview_denoised() && preview.callback != nullptr) { preview_image(step, denoised, version, preview.mode, preview.callback, preview.data, false); } - report_sample_progress(step, steps, t0); + report_sample_progress(step, steps, &last_progress_us); sd::guidance::GuiderOutput output; output.pred = denoised; return output; @@ -2110,7 +2152,7 @@ public: diffusion_params.extra = UNetDiffusionExtra{-1, &controls, control_strength}; } else if (sd_version_is_sd3(version)) { diffusion_params.extra = SkipLayerDiffusionExtra{local_skip_layers}; - } else if (sd_version_is_flux(version) || sd_version_is_flux2(version) || sd_version_is_longcat(version)) { + } else if (sd_version_is_flux(version) || sd_version_is_flux2(version) || sd_version_is_longcat(version) || sd_version_is_sefi_image(version)) { diffusion_params.extra = FluxDiffusionExtra{&guidance_tensor, local_skip_layers}; } else if (sd_version_is_anima(version)) { @@ -2133,6 +2175,9 @@ public: audio_length, frame_rate, video_positions.empty() ? nullptr : &video_positions}; + } else if (sd_version_is_minit2i(version)) { + diffusion_params.extra = MiniT2IDiffusionExtra{ + condition.c_vector.empty() ? nullptr : &condition.c_vector}; } else { diffusion_params.extra = std::monostate{}; } @@ -2142,6 +2187,10 @@ public: return std::move(cached_output); } + for (const auto& extension : generation_extensions) { + extension->before_diffusion(diffusion_params, step); + } + auto output_opt = work_diffusion_model->compute(n_threads, diffusion_params); if (output_opt.empty()) { LOG_ERROR("diffusion model compute failed"); @@ -2152,20 +2201,24 @@ public: return output_opt; }; - if (start_merge_step == -1 || step <= start_merge_step) { - cond_out = run_condition(cond); - if (cond_out.empty()) { - return {}; - } - } else { - GGML_ASSERT(!id_cond.empty()); - cond_out = run_condition(id_cond, - cond.c_concat.empty() ? nullptr : &cond.c_concat); - if (cond_out.empty()) { - return {}; + const SDCondition* positive_condition = &cond; + const sd::Tensor* c_concat_override = nullptr; + for (const auto& extension : generation_extensions) { + const SDCondition& next_condition = extension->before_condition(step, *positive_condition); + if (&next_condition != positive_condition) { + positive_condition = &next_condition; + if (positive_condition != &cond) { + c_concat_override = cond.c_concat.empty() ? nullptr : &cond.c_concat; + } + break; } } + cond_out = run_condition(*positive_condition, c_concat_override); + if (cond_out.empty()) { + return {}; + } + if (!uncond.empty()) { if (!step_cache.is_step_skipped()) { compute_sample_controls(control_image, @@ -2202,7 +2255,7 @@ public: guidance_input.pred_uncond = uncond_out.empty() ? nullptr : &uncond_out; guidance_input.pred_img_uncond = img_uncond_out.empty() ? nullptr : &img_uncond_out; - sd::guidance::GuiderOutput guided = primary_guidance.forward(guidance_input, {}); + sd::guidance::GuiderOutput guided = guidance_schedule.empty() ? primary_guidance.forward(guidance_input, {}) : primary_guidance.forward(guidance_input, {}, guidance_schedule[guidance_schedule.size() - 1 - step]); if (guided.pred.empty()) { return {}; } @@ -2241,12 +2294,12 @@ public: if (sd_should_preview_denoised() && preview.callback != nullptr) { preview_image(step, denoised, version, preview.mode, preview.callback, preview.data, false); } - report_sample_progress(step, steps, t0); + report_sample_progress(step, steps, &last_progress_us); output.pred = denoised; return output; }; - auto x0_opt = sample_k_diffusion(method, denoise, x_t, sigmas, sampler_rng, eta, is_flow_denoiser, extra_sample_args); + auto x0_opt = sample_k_diffusion(method, denoise, x_t, sigmas, sampler_rng, eta, is_flow_denoiser, extra_sample_args, denoiser); if (x0_opt.empty()) { LOG_ERROR("Diffusion model sampling failed"); if (control_net) { @@ -2305,8 +2358,12 @@ public: latent_channel = 3; } else if (version == VERSION_CHROMA_RADIANCE) { latent_channel = 3; + } else if (sd_version_is_minit2i(version)) { + latent_channel = 3; } else if (sd_version_is_pid(version)) { latent_channel = 3; + } else if (sd_version_is_sefi_image(version)) { + latent_channel = 144; } else if (sd_version_uses_flux2_vae(version)) { latent_channel = 128; } else { @@ -2388,18 +2445,9 @@ public: } sd::Tensor decode_first_stage(const sd::Tensor& x, bool decode_video = false) { - if (sd_version_is_pid(version)) { + if (sd_version_is_pid(version) || sd_version_is_minit2i(version)) { return sd::ops::clamp((x + 1.f) * 0.5f, 0.0f, 1.0f); } - // Free resident diffusion params before VAE allocates its compute buffer. - if (stream_layers) { - if (diffusion_model) { - diffusion_model->release_streaming_residency(); - } - if (high_noise_diffusion_model) { - high_noise_diffusion_model->release_streaming_residency(); - } - } auto latents = first_stage_model->diffusion_to_vae_latents(x); first_stage_model->set_temporal_tiling_enabled(vae_tiling_params.temporal_tiling); return first_stage_model->decode(n_threads, latents, vae_tiling_params, decode_video, circular_x, circular_y); @@ -2428,9 +2476,6 @@ public: return {}; } auto waveform = audio_vae_model->decode(n_threads, audio_latent); - if (free_params_immediately) { - audio_vae_model->free_params_buffer(); - } return waveform; } @@ -2543,6 +2588,10 @@ const char* scheduler_to_str[] = { "lcm", "bong_tangent", "ltx2", + "logit_normal", + "flux2", + "flux", + "beta", }; const char* sd_scheduler_name(enum scheduler_t scheduler) { @@ -2553,6 +2602,9 @@ const char* sd_scheduler_name(enum scheduler_t scheduler) { } enum scheduler_t str_to_scheduler(const char* str) { + if (!strcmp(str, "normal")) { + return DISCRETE_SCHEDULER; + } for (int i = 0; i < SCHEDULER_COUNT; i++) { if (!strcmp(str, scheduler_to_str[i])) { return (enum scheduler_t)i; @@ -2567,7 +2619,8 @@ const char* prediction_to_str[] = { "edm_v", "sd3_flow", "flux_flow", - "flux2_flow", + "sefi_flow", + "minit2i_flow", }; const char* sd_prediction_name(enum prediction_t prediction) { @@ -2733,31 +2786,28 @@ void sd_hires_params_init(sd_hires_params_t* hires_params) { } void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) { - *sd_ctx_params = {}; - sd_ctx_params->vae_decode_only = true; - sd_ctx_params->free_params_immediately = true; - sd_ctx_params->n_threads = sd_get_num_physical_cores(); - sd_ctx_params->wtype = SD_TYPE_COUNT; - sd_ctx_params->rng_type = CUDA_RNG; - sd_ctx_params->sampler_rng_type = RNG_TYPE_COUNT; - sd_ctx_params->prediction = PREDICTION_COUNT; - sd_ctx_params->lora_apply_mode = LORA_APPLY_AUTO; - sd_ctx_params->offload_params_to_cpu = false; - sd_ctx_params->max_vram = 0.f; - sd_ctx_params->stream_layers = false; - sd_ctx_params->enable_mmap = false; - sd_ctx_params->keep_clip_on_cpu = false; - sd_ctx_params->keep_control_net_on_cpu = false; - sd_ctx_params->keep_vae_on_cpu = false; - sd_ctx_params->diffusion_flash_attn = false; - sd_ctx_params->circular_x = false; - sd_ctx_params->circular_y = false; - sd_ctx_params->chroma_use_dit_mask = true; - sd_ctx_params->chroma_use_t5_mask = false; - sd_ctx_params->chroma_t5_mask_pad = 1; - sd_ctx_params->vae_format = SD_VAE_FORMAT_AUTO; - sd_ctx_params->backend = nullptr; - sd_ctx_params->params_backend = nullptr; + *sd_ctx_params = {}; + sd_ctx_params->n_threads = sd_get_num_physical_cores(); + sd_ctx_params->wtype = SD_TYPE_COUNT; + sd_ctx_params->rng_type = CUDA_RNG; + sd_ctx_params->sampler_rng_type = RNG_TYPE_COUNT; + sd_ctx_params->prediction = PREDICTION_COUNT; + sd_ctx_params->lora_apply_mode = LORA_APPLY_AUTO; + sd_ctx_params->max_vram = nullptr; + sd_ctx_params->stream_layers = false; + sd_ctx_params->eager_load = false; + sd_ctx_params->enable_mmap = false; + sd_ctx_params->diffusion_flash_attn = false; + sd_ctx_params->circular_x = false; + sd_ctx_params->circular_y = false; + sd_ctx_params->chroma_use_dit_mask = true; + sd_ctx_params->chroma_use_t5_mask = false; + sd_ctx_params->chroma_t5_mask_pad = 1; + sd_ctx_params->vae_format = SD_VAE_FORMAT_AUTO; + sd_ctx_params->backend = nullptr; + sd_ctx_params->params_backend = nullptr; + sd_ctx_params->rpc_servers = nullptr; + sd_ctx_params->pulid_weights_path = nullptr; } char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { @@ -2776,28 +2826,25 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { "llm_vision_path: %s\n" "diffusion_model_path: %s\n" "high_noise_diffusion_model_path: %s\n" + "uncond_diffusion_model_path: %s\n" "embeddings_connectors_path: %s\n" "vae_path: %s\n" "audio_vae_path: %s\n" "taesd_path: %s\n" "control_net_path: %s\n" "photo_maker_path: %s\n" + "pulid_weights_path: %s\n" "tensor_type_rules: %s\n" - "vae_decode_only: %s\n" - "free_params_immediately: %s\n" "n_threads: %d\n" "wtype: %s\n" "rng_type: %s\n" "sampler_rng_type: %s\n" "prediction: %s\n" - "offload_params_to_cpu: %s\n" - "max_vram: %.3f\n" + "max_vram: %s\n" "stream_layers: %s\n" + "eager_load: %s\n" "backend: %s\n" "params_backend: %s\n" - "keep_clip_on_cpu: %s\n" - "keep_control_net_on_cpu: %s\n" - "keep_vae_on_cpu: %s\n" "flash_attn: %s\n" "diffusion_flash_attn: %s\n" "circular_x: %s\n" @@ -2815,28 +2862,25 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) { SAFE_STR(sd_ctx_params->llm_vision_path), SAFE_STR(sd_ctx_params->diffusion_model_path), SAFE_STR(sd_ctx_params->high_noise_diffusion_model_path), + SAFE_STR(sd_ctx_params->uncond_diffusion_model_path), SAFE_STR(sd_ctx_params->embeddings_connectors_path), SAFE_STR(sd_ctx_params->vae_path), SAFE_STR(sd_ctx_params->audio_vae_path), SAFE_STR(sd_ctx_params->taesd_path), SAFE_STR(sd_ctx_params->control_net_path), SAFE_STR(sd_ctx_params->photo_maker_path), + SAFE_STR(sd_ctx_params->pulid_weights_path), SAFE_STR(sd_ctx_params->tensor_type_rules), - BOOL_STR(sd_ctx_params->vae_decode_only), - BOOL_STR(sd_ctx_params->free_params_immediately), sd_ctx_params->n_threads, sd_type_name(sd_ctx_params->wtype), sd_rng_type_name(sd_ctx_params->rng_type), sd_rng_type_name(sd_ctx_params->sampler_rng_type), sd_prediction_name(sd_ctx_params->prediction), - BOOL_STR(sd_ctx_params->offload_params_to_cpu), - sd_ctx_params->max_vram, + SAFE_STR(sd_ctx_params->max_vram), BOOL_STR(sd_ctx_params->stream_layers), + BOOL_STR(sd_ctx_params->eager_load), SAFE_STR(sd_ctx_params->backend), SAFE_STR(sd_ctx_params->params_backend), - BOOL_STR(sd_ctx_params->keep_clip_on_cpu), - BOOL_STR(sd_ctx_params->keep_control_net_on_cpu), - BOOL_STR(sd_ctx_params->keep_vae_on_cpu), BOOL_STR(sd_ctx_params->flash_attn), BOOL_STR(sd_ctx_params->diffusion_flash_attn), BOOL_STR(sd_ctx_params->circular_x), @@ -2921,6 +2965,7 @@ void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params) { sd_img_gen_params->batch_count = 1; sd_img_gen_params->control_strength = 0.9f; sd_img_gen_params->pm_params = {nullptr, 0, nullptr, 20.f}; + sd_img_gen_params->pulid_params = {nullptr, 1.0f}; sd_img_gen_params->vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr}; sd_cache_params_init(&sd_img_gen_params->cache); sd_hires_params_init(&sd_img_gen_params->hires); @@ -3063,6 +3108,15 @@ void free_sd_ctx(sd_ctx_t* sd_ctx) { free(sd_ctx); } +SD_API void sd_cancel_generation(sd_ctx_t* sd_ctx, enum sd_cancel_mode_t mode) { + if (sd_ctx && sd_ctx->sd) { + if (mode < SD_CANCEL_ALL || mode > SD_CANCEL_RESET) { + mode = SD_CANCEL_ALL; + } + sd_ctx->sd->set_cancel_flag(mode); + } +} + static sd_audio_t* waveform_to_sd_audio(const StableDiffusionGGML* sd, const sd::Tensor& waveform) { if (sd == nullptr || waveform.empty()) { @@ -3142,8 +3196,14 @@ enum scheduler_t sd_get_default_scheduler(const sd_ctx_t* sd_ctx, enum sample_me return LCM_SCHEDULER; } else if (sample_method == DDIM_TRAILING_SAMPLE_METHOD) { return SIMPLE_SCHEDULER; + } else if (sd_ctx != nullptr && sd_ctx->sd != nullptr && sd_version_is_flux(sd_ctx->sd->version)) { + return FLUX_SCHEDULER; + } else if (sd_ctx != nullptr && sd_ctx->sd != nullptr && sd_version_is_flux2(sd_ctx->sd->version)) { + return FLUX2_SCHEDULER; } else if (sd_ctx != nullptr && sd_ctx->sd != nullptr && sd_version_is_ltxav(sd_ctx->sd->version)) { return LTX2_SCHEDULER; + } else if (sd_ctx != nullptr && sd_ctx->sd != nullptr && sd_version_is_ideogram4(sd_ctx->sd->version)) { + return LOGIT_NORMAL_SCHEDULER; } return DISCRETE_SCHEDULER; } @@ -3222,6 +3282,7 @@ struct GenerationRequest { sd_guidance_params_t guidance = {}; sd_guidance_params_t high_noise_guidance = {}; sd_pm_params_t pm_params = {}; + sd_pulid_params_t pulid_params = {}; sd_hires_params_t hires = {}; int frames = -1; int requested_frames = -1; @@ -3247,6 +3308,7 @@ struct GenerationRequest { has_ref_images = sd_img_gen_params->ref_images_count > 0; guidance = sd_img_gen_params->sample_params.guidance; pm_params = sd_img_gen_params->pm_params; + pulid_params = sd_img_gen_params->pulid_params; hires = sd_img_gen_params->hires; cache_params = &sd_img_gen_params->cache; resolve(sd_ctx); @@ -3443,7 +3505,6 @@ struct SamplePlan { int high_noise_sample_steps = 0; int total_steps = 0; float moe_boundary = 0.f; - int start_merge_step = -1; std::vector sigmas; SamplePlan(sd_ctx_t* sd_ctx, @@ -3528,11 +3589,6 @@ struct SamplePlan { high_noise_eta = resolve_eta(sd_ctx, high_noise_eta, high_noise_sample_method); LOG_INFO("sampling(high noise) using %s method", sampling_methods_str[high_noise_sample_method]); } - - if (sd_ctx->sd->use_pmid) { - start_merge_step = int(sd_ctx->sd->pmid_model->style_strength / 100.f * total_steps); - LOG_INFO("PHOTOMAKER: start_merge_step: %d", start_merge_step); - } } }; @@ -3863,7 +3919,15 @@ struct ImageGenerationEmbeds { SDCondition cond; SDCondition uncond; SDCondition img_uncond; - SDCondition id_cond; +}; + +struct ConditionerRunnerDoneOnExit { + Conditioner* conditioner = nullptr; + ~ConditionerRunnerDoneOnExit() { + if (conditioner != nullptr) { + conditioner->runner_done(); + } + } }; struct CircularAxesState { @@ -4038,7 +4102,7 @@ static std::optional prepare_image_generation_latents(sd } } - if (!control_image_tensor.empty() && !sd_ctx->sd->vae_decode_only) { + if (!control_image_tensor.empty()) { control_latent = sd_ctx->sd->encode_first_stage(control_image_tensor); if (control_latent.empty()) { LOG_ERROR("failed to encode control image"); @@ -4199,6 +4263,8 @@ static std::optional prepare_image_generation_embeds(sd_c GenerationRequest* request, SamplePlan* plan, ImageGenerationLatents* latents) { + ConditionerRunnerDoneOnExit conditioner_runner_done{sd_ctx->sd->cond_stage_model.get()}; + ConditionerParams condition_params; condition_params.text = request->prompt; condition_params.clip_skip = request->clip_skip; @@ -4206,7 +4272,10 @@ static std::optional prepare_image_generation_embeds(sd_c condition_params.height = request->height; condition_params.ref_images = &latents->ref_images; - auto id_cond = sd_ctx->sd->get_pmid_conditon(request->pm_params, condition_params); + sd_ctx->sd->prepare_generation_extensions(request->pm_params, + request->pulid_params, + condition_params, + plan->total_steps); int64_t prepare_start_ms = ggml_time_ms(); condition_params.zero_out_masked = false; auto cond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads, @@ -4221,16 +4290,25 @@ static std::optional prepare_image_generation_embeds(sd_c SDCondition uncond; if (request->use_uncond || request->use_high_noise_uncond) { - bool zero_out_masked = false; - if (sd_version_is_sdxl(sd_ctx->sd->version) && - request->negative_prompt.empty() && - !sd_ctx->sd->is_using_edm_v_parameterization) { - zero_out_masked = true; + if (sd_version_is_ideogram4(sd_ctx->sd->version)) { + uncond.c_vector = sd::Tensor::from_vector({1.0f}); + } else if (sd_version_is_minit2i(sd_ctx->sd->version)) { + // MiniT2I derives the unconditional signal from the same T5 hidden + // states with a zeroed prompt mask, so no extra text encode is needed. + uncond.c_crossattn = cond.c_crossattn; + uncond.c_vector = sd::Tensor::zeros_like(cond.c_vector); + } else { + bool zero_out_masked = false; + if (sd_version_is_sdxl(sd_ctx->sd->version) && + request->negative_prompt.empty() && + !sd_ctx->sd->is_using_edm_v_parameterization) { + zero_out_masked = true; + } + condition_params.text = request->negative_prompt; + condition_params.zero_out_masked = zero_out_masked; + uncond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads, + condition_params); } - condition_params.text = request->negative_prompt; - condition_params.zero_out_masked = zero_out_masked; - uncond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads, - condition_params); if (uncond.c_concat.empty()) { uncond.c_concat = latents->concat_latent; // TODO: optimize } @@ -4264,31 +4342,41 @@ static std::optional prepare_image_generation_embeds(sd_c int64_t t1 = ggml_time_ms(); LOG_INFO("get_learned_condition completed, taking %.2fs", (t1 - prepare_start_ms) * 1.0f / 1000); - if (sd_ctx->sd->free_params_immediately) { - sd_ctx->sd->cond_stage_model->free_params_buffer(); - } - ImageGenerationEmbeds embeds; embeds.img_uncond = std::move(img_uncond); embeds.cond = std::move(cond); embeds.uncond = std::move(uncond); - embeds.id_cond = std::move(id_cond); return embeds; } static sd_image_t* decode_image_outputs(sd_ctx_t* sd_ctx, const GenerationRequest& request, - const std::vector>& final_latents) { - if (final_latents.size() != static_cast(request.batch_count)) { - LOG_ERROR("expected %d latents, got %zu", request.batch_count, final_latents.size()); + const std::vector>& final_latents, + int* num_images_out) { + if (final_latents.empty()) { + LOG_ERROR("no latent images to decode"); return nullptr; } - LOG_INFO("decoding %zu latents", final_latents.size()); + if (final_latents.size() > static_cast(request.batch_count)) { + LOG_ERROR("expected at most %d latents, got %zu", request.batch_count, final_latents.size()); + return nullptr; + } + if (final_latents.size() < static_cast(request.batch_count)) { + LOG_INFO("decoding %zu/%d latents", final_latents.size(), request.batch_count); + } else { + LOG_INFO("decoding %zu latents", final_latents.size()); + } std::vector> decoded_images; - int64_t t0 = ggml_time_ms(); + int64_t t0 = ggml_time_ms(); + bool cancelled = false; for (size_t i = 0; i < final_latents.size(); i++) { + if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling latent decodings"); + cancelled = true; + break; + } int64_t t1 = ggml_time_ms(); if (sd_ctx->sd->version == VERSION_QWEN_IMAGE_LAYERED) { constexpr int kLayerCount = 4; @@ -4296,31 +4384,30 @@ static sd_image_t* decode_image_outputs(sd_ctx_t* sd_ctx, LOG_ERROR("qwen image layered expected at least %d latent layers, got shape dim=%d", kLayerCount + 1, final_latents[i].dim()); - if (sd_ctx->sd->free_params_immediately) { - sd_ctx->sd->first_stage_model->free_params_buffer(); - } return nullptr; } for (int layer_index = 0; layer_index < kLayerCount; layer_index++) { + if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling latent decodings"); + cancelled = true; + break; + } sd::Tensor layer_latent = sd::ops::slice(final_latents[i], 2, layer_index + 1, layer_index + 2); layer_latent.squeeze_(2); sd::Tensor image = sd_ctx->sd->decode_first_stage(layer_latent); if (image.empty()) { LOG_ERROR("decode_first_stage failed for latent %zu layer %d", i + 1, layer_index + 1); - if (sd_ctx->sd->free_params_immediately) { - sd_ctx->sd->first_stage_model->free_params_buffer(); - } return nullptr; } decoded_images.push_back(std::move(image)); } + if (cancelled) { + break; + } } else { sd::Tensor image = sd_ctx->sd->decode_first_stage(final_latents[i]); if (image.empty()) { LOG_ERROR("decode_first_stage failed for latent %" PRId64, i + 1); - if (sd_ctx->sd->free_params_immediately) { - sd_ctx->sd->first_stage_model->free_params_buffer(); - } return nullptr; } decoded_images.push_back(std::move(image)); @@ -4331,15 +4418,19 @@ static sd_image_t* decode_image_outputs(sd_ctx_t* sd_ctx, int64_t t4 = ggml_time_ms(); LOG_INFO("decode_first_stage completed, taking %.2fs", (t4 - t0) * 1.0f / 1000); - if (sd_ctx->sd->free_params_immediately) { - sd_ctx->sd->first_stage_model->free_params_buffer(); + if (decoded_images.empty()) { + LOG_ERROR(cancelled ? "cancelled before any latent images were decoded" : "no decoded images"); + return nullptr; } - sd_image_t* result_images = (sd_image_t*)calloc(decoded_images.size(), sizeof(sd_image_t)); + int image_count = static_cast(decoded_images.size()); + sd_image_t* result_images = (sd_image_t*)calloc(image_count, sizeof(sd_image_t)); if (result_images == nullptr) { return nullptr; } - memset(result_images, 0, decoded_images.size() * sizeof(sd_image_t)); + if (num_images_out != nullptr) { + *num_images_out = image_count; + } for (size_t i = 0; i < decoded_images.size(); i++) { result_images[i] = tensor_to_sd_image(decoded_images[i]); @@ -4352,6 +4443,11 @@ static sd::Tensor upscale_hires_latent(sd_ctx_t* sd_ctx, const sd::Tensor& latent, const GenerationRequest& request, UpscalerGGML* upscaler) { + if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling hires latent upscale"); + return {}; + } + auto get_hires_latent_target_shape = [&]() { std::vector target_shape = latent.shape(); if (target_shape.size() < 2) { @@ -4413,11 +4509,6 @@ static sd::Tensor upscale_hires_latent(sd_ctx_t* sd_ctx, } else if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL || request.hires.upscaler == SD_HIRES_UPSCALER_LANCZOS || request.hires.upscaler == SD_HIRES_UPSCALER_NEAREST) { - if (sd_ctx->sd->vae_decode_only) { - LOG_ERROR("hires %s upscaler requires VAE encoder weights; create the context with vae_decode_only=false", - sd_hires_upscaler_name(request.hires.upscaler)); - return {}; - } if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL && upscaler == nullptr) { LOG_ERROR("hires model upscaler context is null"); return {}; @@ -4429,6 +4520,10 @@ static sd::Tensor upscale_hires_latent(sd_ctx_t* sd_ctx, sd_hires_upscaler_name(request.hires.upscaler)); return {}; } + if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling hires image upscale"); + return {}; + } sd::Tensor upscaled_tensor; if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL) { @@ -4465,6 +4560,10 @@ static sd::Tensor upscale_hires_latent(sd_ctx_t* sd_ctx, upscaled_tensor = sd::ops::clamp(upscaled_tensor, 0.0f, 1.0f); } + if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling hires latent encode"); + return {}; + } sd::Tensor upscaled_latent = sd_ctx->sd->encode_first_stage(upscaled_tensor); if (upscaled_latent.empty()) { LOG_ERROR("encode_first_stage failed after hires %s upscale", @@ -4524,22 +4623,22 @@ static std::vector make_hires_sigma_schedule(sd_ctx_t* sd_ctx, sigmas.end()); } -SD_API int32_t sd_get_image_result_count(const sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_gen_params) { - if (sd_ctx == nullptr || sd_ctx->sd == nullptr || sd_img_gen_params == nullptr) { - return 0; +SD_API bool generate_image(sd_ctx_t* sd_ctx, + const sd_img_gen_params_t* sd_img_gen_params, + sd_image_t** images_out, + int* num_images_out) { + if (images_out != nullptr) { + *images_out = nullptr; } - int32_t batch_count = sd_img_gen_params->batch_count > 0 ? sd_img_gen_params->batch_count : 1; - if (sd_ctx->sd->version == VERSION_QWEN_IMAGE_LAYERED) { - return batch_count * 4; + if (num_images_out != nullptr) { + *num_images_out = 0; } - return batch_count; -} - -SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_gen_params) { if (sd_ctx == nullptr || sd_img_gen_params == nullptr) { - return nullptr; + return false; } + sd_ctx->sd->reset_cancel_flag(); + int64_t t0 = ggml_time_ms(); sd_ctx->sd->vae_tiling_params = sd_img_gen_params->vae_tiling_params; GenerationRequest request(sd_ctx, sd_img_gen_params); @@ -4558,7 +4657,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s &request, &plan); if (!latents_opt.has_value()) { - return nullptr; + return false; } ImageGenerationLatents latents = std::move(*latents_opt); @@ -4568,13 +4667,25 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s &plan, &latents); if (!embeds_opt.has_value()) { - return nullptr; + return false; } ImageGenerationEmbeds embeds = std::move(*embeds_opt); std::vector> final_latents; int64_t denoise_start = ggml_time_ms(); for (int b = 0; b < request.batch_count; b++) { + sd_cancel_mode_t cancel = sd_ctx->sd->get_cancel_flag(); + if (cancel == SD_CANCEL_ALL) { + LOG_ERROR("cancelling generation"); + return false; + } + if (cancel == SD_CANCEL_NEW_LATENTS) { + LOG_INFO("cancelling new latent generation, returning %zu/%d completed latents", + final_latents.size(), + request.batch_count); + break; + } + int64_t sampling_start = ggml_time_ms(); int64_t cur_seed = request.seed + b; LOG_INFO("generating image: %i/%i - seed %" PRId64, b + 1, request.batch_count, cur_seed); @@ -4590,7 +4701,6 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s embeds.cond, embeds.uncond, embeds.img_uncond, - embeds.id_cond, latents.control_image, request.control_strength, request.guidance, @@ -4600,7 +4710,6 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s sd_ctx->sd->is_flow_denoiser(), plan.extra_sample_args, plan.sigmas, - plan.start_merge_step, latents.ref_latents, request.increase_ref_index, latents.denoise_mask, @@ -4620,40 +4729,42 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s b + 1, request.batch_count, (sampling_end - sampling_start) * 1.0f / 1000); - if (sd_ctx->sd->free_params_immediately) { - sd_ctx->sd->diffusion_model->free_params_buffer(); - } - return nullptr; - } - if (sd_ctx->sd->free_params_immediately && !request.hires.enabled) { - sd_ctx->sd->diffusion_model->free_params_buffer(); + return false; } int64_t denoise_end = ggml_time_ms(); LOG_INFO("generating %zu latent images completed, taking %.2fs", final_latents.size(), (denoise_end - denoise_start) * 1.0f / 1000); + if (final_latents.empty()) { + LOG_ERROR("no latent images generated"); + return false; + } if (request.hires.enabled && request.hires.target_width > 0) { + if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling generation before hires fix"); + return false; + } LOG_INFO("hires fix: upscaling to %dx%d", request.hires.target_width, request.hires.target_height); std::unique_ptr hires_upscaler; if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL) { + if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling generation before hires model load"); + return false; + } LOG_INFO("hires fix: loading model upscaler from '%s'", request.hires.model_path); hires_upscaler = std::make_unique(sd_ctx->sd->n_threads, false, request.hires.upscale_tile_size, sd_ctx->sd->backend_spec, sd_ctx->sd->params_backend_spec); - const size_t max_graph_vram_bytes = sd::ggml_graph_cut::max_vram_gib_to_bytes(sd_ctx->sd->max_vram); + const size_t max_graph_vram_bytes = sd_ctx->sd->max_graph_vram_bytes_for_module(SDBackendModule::UPSCALER); hires_upscaler->set_max_graph_vram_bytes(max_graph_vram_bytes); if (!hires_upscaler->load_from_file(request.hires.model_path, - sd_ctx->sd->offload_params_to_cpu, sd_ctx->sd->n_threads)) { LOG_ERROR("load hires model upscaler failed"); - if (sd_ctx->sd->free_params_immediately) { - sd_ctx->sd->diffusion_model->free_params_buffer(); - } - return nullptr; + return false; } } @@ -4675,6 +4786,10 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s std::vector> hires_final_latents; int64_t hires_denoise_start = ggml_time_ms(); for (int b = 0; b < (int)final_latents.size(); b++) { + if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling generation during hires fix"); + return false; + } int64_t cur_seed = request.seed + b; sd_ctx->sd->rng->manual_seed(cur_seed); sd_ctx->sd->sampler_rng->manual_seed(cur_seed); @@ -4684,10 +4799,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s request, hires_upscaler.get()); if (upscaled.empty()) { - if (sd_ctx->sd->free_params_immediately) { - sd_ctx->sd->diffusion_model->free_params_buffer(); - } - return nullptr; + return false; } sd::Tensor noise = sd::randn_like(upscaled, sd_ctx->sd->rng); @@ -4710,7 +4822,6 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s embeds.cond, embeds.uncond, embeds.img_uncond, - embeds.id_cond, latents.control_image, request.control_strength, request.guidance, @@ -4720,7 +4831,6 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s sd_ctx->sd->is_flow_denoiser(), plan.extra_sample_args, hires_sigma_sched, - plan.start_merge_step, latents.ref_latents, request.increase_ref_index, hires_denoise_mask, @@ -4743,13 +4853,7 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s b + 1, (int)final_latents.size(), (hires_sample_end - hires_sample_start) * 1.0f / 1000); - if (sd_ctx->sd->free_params_immediately) { - sd_ctx->sd->diffusion_model->free_params_buffer(); - } - return nullptr; - } - if (sd_ctx->sd->free_params_immediately) { - sd_ctx->sd->diffusion_model->free_params_buffer(); + return false; } int64_t hires_denoise_end = ggml_time_ms(); LOG_INFO("hires fix completed, taking %.2fs", (hires_denoise_end - hires_denoise_start) * 1.0f / 1000); @@ -4757,16 +4861,25 @@ SD_API sd_image_t* generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* s final_latents = std::move(hires_final_latents); } - auto result = decode_image_outputs(sd_ctx, request, final_latents); + int num_images = 0; + auto result = decode_image_outputs(sd_ctx, request, final_latents, &num_images); if (result == nullptr) { - return nullptr; + return false; } sd_ctx->sd->lora_stat(); int64_t t1 = ggml_time_ms(); LOG_INFO("generate_image completed in %.2fs", (t1 - t0) * 1.0f / 1000); - return result; + if (num_images_out != nullptr) { + *num_images_out = num_images; + } + if (images_out != nullptr) { + *images_out = result; + } else { + free_sd_images(result, num_images); + } + return true; } static std::optional prepare_video_generation_latents(sd_ctx_t* sd_ctx, @@ -4798,11 +4911,6 @@ static std::optional prepare_video_generation_latents(sd } if (!start_image.empty() || !end_image.empty()) { - if (sd_ctx->sd->vae_decode_only) { - LOG_ERROR("LTXAV image conditioning requires VAE encoder weights; create the context with vae_decode_only=false"); - return std::nullopt; - } - if (!start_image.empty() && !end_image.empty()) { LOG_INFO("FLF2V"); } else if (!start_image.empty()) { @@ -5086,6 +5194,8 @@ static ImageGenerationEmbeds prepare_video_generation_embeds(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* sd_vid_gen_params, const GenerationRequest& request, const ImageGenerationLatents& latents) { + ConditionerRunnerDoneOnExit conditioner_runner_done{sd_ctx->sd->cond_stage_model.get()}; + ImageGenerationEmbeds embeds; ConditionerParams condition_params; condition_params.clip_skip = request.clip_skip; @@ -5108,9 +5218,6 @@ static ImageGenerationEmbeds prepare_video_generation_embeds(sd_ctx_t* sd_ctx, int64_t t1 = ggml_time_ms(); LOG_INFO("get_learned_condition completed, taking %.2fs", (t1 - prepare_start_ms) * 1.0f / 1000); - if (sd_ctx->sd->free_params_immediately) { - sd_ctx->sd->cond_stage_model->free_params_buffer(); - } return embeds; } @@ -5122,6 +5229,10 @@ static sd_image_t* decode_video_outputs(sd_ctx_t* sd_ctx, LOG_ERROR("no latent video to decode"); return nullptr; } + if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling video decode"); + return nullptr; + } sd::Tensor video_latent = final_latent; if (sd_version_is_ltxav(sd_ctx->sd->version) && video_latent.shape()[3] > sd_ctx->sd->get_latent_channel()) { @@ -5137,9 +5248,6 @@ static sd_image_t* decode_video_outputs(sd_ctx_t* sd_ctx, sd::Tensor vid = sd_ctx->sd->decode_first_stage(video_latent, true); int64_t t5 = ggml_time_ms(); LOG_INFO("decode_first_stage completed, taking %.2fs", (t5 - t4) * 1.0f / 1000); - if (sd_ctx->sd->free_params_immediately) { - sd_ctx->sd->first_stage_model->free_params_buffer(); - } if (vid.empty()) { LOG_ERROR("decode_first_stage failed for video"); return nullptr; @@ -5204,17 +5312,40 @@ static sd::Tensor upscale_ltx_spatial_video_latent(sd_ctx_t* sd_ctx, return {}; } + auto upsampler_manager = std::make_shared(); + upsampler_manager->set_n_threads(sd_ctx->sd->n_threads); + upsampler_manager->set_enable_mmap(sd_ctx->sd->enable_mmap); + ModelLoader& model_loader = upsampler_manager->loader(); + if (!model_loader.init_from_file(model_path)) { + LOG_ERROR("init LTX latent upsampler model loader from file failed: '%s'", model_path); + return {}; + } + std::unique_ptr upsampler = std::make_unique(sd_ctx->sd->backend_for(SDBackendModule::UPSCALER), - sd_ctx->sd->params_backend_for(SDBackendModule::UPSCALER)); - const size_t max_graph_vram_bytes = sd::ggml_graph_cut::max_vram_gib_to_bytes(sd_ctx->sd->max_vram); + model_loader.get_tensor_storage_map(), + upsampler_manager); + const size_t max_graph_vram_bytes = sd_ctx->sd->max_graph_vram_bytes_for_module(SDBackendModule::UPSCALER); upsampler->set_max_graph_vram_bytes(max_graph_vram_bytes); - if (!upsampler->load_from_file(model_path, sd_ctx->sd->n_threads)) { - LOG_ERROR("load LTX latent upsampler failed"); + if (upsampler->model == nullptr) { + LOG_ERROR("init LTX latent upsampler from metadata failed"); + return {}; + } + + std::map tensors; + upsampler->get_param_tensors(tensors); + if (!upsampler_manager->register_param_tensors("LTX latent upsampler", + std::move(tensors), + ModelManager::ResidencyMode::ParamBackend, + sd_ctx->sd->backend_for(SDBackendModule::UPSCALER), + sd_ctx->sd->params_backend_for(SDBackendModule::UPSCALER)) || + !upsampler_manager->validate_registered_tensors()) { + LOG_ERROR("register LTX latent upsampler tensors with model manager failed"); return {}; } sd::Tensor upscaled = upsampler->compute(sd_ctx->sd->n_threads, unnormalized); + upsampler_manager.reset(); upsampler.reset(); if (upscaled.empty()) { LOG_ERROR("LTX latent spatial upscale failed"); @@ -5248,11 +5379,6 @@ static bool apply_ltxv_refine_image_conditioning(sd_ctx_t* sd_ctx, sd_vid_gen_params->end_image.data == nullptr) { return true; } - if (sd_ctx->sd->vae_decode_only) { - LOG_ERROR("LTXV refine image conditioning requires VAE encoder weights; create the context with vae_decode_only=false"); - return false; - } - constexpr float conditioning_strength = 1.f; int latent_channels = sd_ctx->sd->get_latent_channel(); sd::Tensor video_latent = *latent; @@ -5352,6 +5478,9 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, if (audio_out != nullptr) { *audio_out = nullptr; } + + sd_ctx->sd->reset_cancel_flag(); + if (num_frames_out != nullptr) { *num_frames_out = 0; } @@ -5379,6 +5508,7 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, sd_ctx->sd->sampler_rng->manual_seed(request.seed); sd_ctx->sd->set_flow_shift(sd_vid_gen_params->sample_params.flow_shift); sd_ctx->sd->apply_loras(sd_vid_gen_params->loras, sd_vid_gen_params->lora_count); + sd_ctx->sd->reset_generation_extensions(); SamplePlan plan(sd_ctx, sd_vid_gen_params, request); auto latent_inputs_opt = prepare_video_generation_latents(sd_ctx, sd_vid_gen_params, &request); @@ -5412,6 +5542,10 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, sd::Tensor noise = sd::Tensor::randn_like(x_t, sd_ctx->sd->rng); if (plan.high_noise_sample_steps > 0) { + if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling generation before high-noise sampling"); + return false; + } LOG_DEBUG("sample(high noise) %dx%dx%d", W, H, T); int64_t sampling_start = ggml_time_ms(); @@ -5425,7 +5559,6 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, embeds.cond, request.use_high_noise_uncond ? embeds.uncond : SDCondition(), embeds.img_uncond, - embeds.id_cond, sd::Tensor(), 0.f, request.high_noise_guidance, @@ -5435,7 +5568,6 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, sd_ctx->sd->is_flow_denoiser(), plan.high_noise_extra_sample_args, high_noise_sigmas, - -1, std::vector>{}, false, latents.denoise_mask, @@ -5448,20 +5580,18 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, int64_t sampling_end = ggml_time_ms(); if (x_t_sampled.empty()) { LOG_ERROR("sampling(high noise) failed after %.2fs", (sampling_end - sampling_start) * 1.0f / 1000); - if (sd_ctx->sd->free_params_immediately) { - sd_ctx->sd->high_noise_diffusion_model->free_params_buffer(); - } return false; } x_t = std::move(x_t_sampled); noise = {}; LOG_INFO("sampling(high noise) completed, taking %.2fs", (sampling_end - sampling_start) * 1.0f / 1000); - if (sd_ctx->sd->free_params_immediately) { - sd_ctx->sd->high_noise_diffusion_model->free_params_buffer(); - } } + if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling generation before sampling"); + return false; + } LOG_DEBUG("sample %dx%dx%d", W, H, T); int64_t sampling_start = ggml_time_ms(); sd::Tensor final_latent = sd_ctx->sd->sample(sd_ctx->sd->diffusion_model, @@ -5471,7 +5601,6 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, embeds.cond, request.use_uncond ? embeds.uncond : SDCondition(), embeds.img_uncond, - embeds.id_cond, sd::Tensor(), 0.f, sd_vid_gen_params->sample_params.guidance, @@ -5481,7 +5610,6 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, sd_ctx->sd->is_flow_denoiser(), plan.extra_sample_args, plan.sigmas, - -1, std::vector>{}, false, latents.denoise_mask, @@ -5494,15 +5622,16 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, int64_t sampling_end = ggml_time_ms(); if (final_latent.empty()) { - if (sd_ctx->sd->free_params_immediately) { - sd_ctx->sd->diffusion_model->free_params_buffer(); - } LOG_ERROR("sampling failed after %.2fs", (sampling_end - sampling_start) * 1.0f / 1000); return false; } LOG_INFO("sampling completed, taking %.2fs", (sampling_end - sampling_start) * 1.0f / 1000); if (latent_upscale_enabled) { + if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling generation before latent upscale"); + return false; + } int64_t upscale_start = ggml_time_ms(); sd::Tensor upscaled_latent = upscale_ltx_spatial_video_latent(sd_ctx, request.hires.model_path, @@ -5510,9 +5639,6 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, latents.audio_length); int64_t upscale_end = ggml_time_ms(); if (upscaled_latent.empty()) { - if (sd_ctx->sd->free_params_immediately) { - sd_ctx->sd->diffusion_model->free_params_buffer(); - } return false; } LOG_INFO("LTX latent spatial upscale completed, taking %.2fs", @@ -5545,9 +5671,6 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, LOG_ERROR("failed to resize LTX audio latent for latent upscale: %d -> %d", latents.audio_length, target_audio_length); - if (sd_ctx->sd->free_params_immediately) { - sd_ctx->sd->diffusion_model->free_params_buffer(); - } return false; } x_t = pack_ltxav_audio_and_video_latents(video_latent, audio_latent); @@ -5568,6 +5691,10 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, } sd::Tensor hires_denoise_mask; sd::Tensor hires_video_positions; + if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling generation before latent upscale refine"); + return false; + } if (!apply_ltxv_refine_image_conditioning(sd_ctx, sd_vid_gen_params, hires_request, @@ -5575,9 +5702,6 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, &x_t, &hires_denoise_mask, &hires_video_positions)) { - if (sd_ctx->sd->free_params_immediately) { - sd_ctx->sd->diffusion_model->free_params_buffer(); - } return false; } noise = sd::Tensor::randn_like(x_t, sd_ctx->sd->rng); @@ -5615,7 +5739,6 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, embeds.cond, hires_request.use_uncond ? embeds.uncond : SDCondition(), embeds.img_uncond, - embeds.id_cond, sd::Tensor(), 0.f, sd_vid_gen_params->sample_params.guidance, @@ -5625,7 +5748,6 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, sd_ctx->sd->is_flow_denoiser(), plan.extra_sample_args, hires_sigma_sched, - -1, std::vector>{}, false, hires_denoise_mask, @@ -5636,9 +5758,6 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, hires_request.cache_params, hires_video_positions); sampling_end = ggml_time_ms(); - if (sd_ctx->sd->free_params_immediately) { - sd_ctx->sd->diffusion_model->free_params_buffer(); - } if (final_latent.empty()) { LOG_ERROR("sampling(latent upscale) failed after %.2fs", (sampling_end - sampling_start) * 1.0f / 1000); @@ -5646,8 +5765,6 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, } LOG_INFO("sampling(latent upscale) completed, taking %.2fs", (sampling_end - sampling_start) * 1.0f / 1000); - } else if (sd_ctx->sd->free_params_immediately) { - sd_ctx->sd->diffusion_model->free_params_buffer(); } int64_t latent_end = ggml_time_ms(); @@ -5657,6 +5774,10 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, if (sd_version_is_ltxav(sd_ctx->sd->version) && latents.audio_length > 0 && sd_ctx->sd->audio_vae_model != nullptr) { + if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling generation before audio decode"); + return false; + } int64_t audio_latent_decode_start = ggml_time_ms(); auto audio_latent = unpack_ltxav_audio_latent(final_latent, @@ -5689,6 +5810,11 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, final_latent = sd::ops::slice(final_latent, 2, latents.ref_image_num, final_latent.shape()[2]); } + if (sd_ctx->sd->get_cancel_flag() == SD_CANCEL_ALL) { + LOG_ERROR("cancelling generation before video decode"); + free_sd_audio(generated_audio); + return false; + } auto result = decode_video_outputs(sd_ctx, latent_upscale_enabled ? hires_request : request, final_latent, num_frames_out); if (result == nullptr) { free_sd_audio(generated_audio); @@ -5709,3 +5835,18 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, } return true; } + +SD_API void free_sd_images(sd_image_t* result_images, int num_images) { + if (result_images == nullptr) { + return; + } + + for (int i = 0; i < num_images; ++i) { + if (result_images[i].data != nullptr) { + free(result_images[i].data); + result_images[i].data = nullptr; + } + } + + free(result_images); +} diff --git a/src/tokenizers/bpe_tokenizer.cpp b/src/tokenizers/bpe_tokenizer.cpp index ed4e1878..7733f00d 100644 --- a/src/tokenizers/bpe_tokenizer.cpp +++ b/src/tokenizers/bpe_tokenizer.cpp @@ -3,8 +3,8 @@ #include #include +#include "core/util.h" #include "tokenize_util.h" -#include "util.h" std::vector> BPETokenizer::bytes_to_unicode() { std::vector> byte_unicode_pairs; @@ -134,7 +134,8 @@ std::vector BPETokenizer::encode(const std::string& text, on_new_token_cb_t std::vector bpe_tokens; std::vector token_strs; - auto splited_texts = split_with_special_tokens(text, special_tokens); + std::string normalized_text = normalize_before_split ? normalize(text) : text; + auto splited_texts = split_with_special_tokens(normalized_text, special_tokens); for (auto& splited_text : splited_texts) { if (is_special_token(splited_text)) { @@ -159,7 +160,7 @@ std::vector BPETokenizer::encode(const std::string& text, on_new_token_cb_t } } - std::string token_str = normalize(token); + std::string token_str = normalize_before_split ? token : normalize(token); std::u32string utf32_token; if (byte_level_bpe) { for (int i = 0; i < token_str.length(); i++) { diff --git a/src/tokenizers/clip_tokenizer.cpp b/src/tokenizers/clip_tokenizer.cpp index 70d63772..d51eadec 100644 --- a/src/tokenizers/clip_tokenizer.cpp +++ b/src/tokenizers/clip_tokenizer.cpp @@ -6,9 +6,9 @@ #include #include +#include "core/util.h" #include "ggml.h" #include "tokenize_util.h" -#include "util.h" #include "vocab/vocab.h" CLIPTokenizer::CLIPTokenizer(int pad_token_id, const std::string& merges_utf8_str) { @@ -22,9 +22,10 @@ CLIPTokenizer::CLIPTokenizer(int pad_token_id, const std::string& merges_utf8_st EOS_TOKEN_ID = 49407; PAD_TOKEN_ID = pad_token_id; - end_of_word_suffix = ""; - add_bos_token = true; - add_eos_token = true; + end_of_word_suffix = ""; + add_bos_token = true; + add_eos_token = true; + normalize_before_split = true; if (merges_utf8_str.size() > 0) { load_from_merges(merges_utf8_str); diff --git a/src/tokenizers/gemma_tokenizer.cpp b/src/tokenizers/gemma_tokenizer.cpp index acac8ccd..a7b67ef1 100644 --- a/src/tokenizers/gemma_tokenizer.cpp +++ b/src/tokenizers/gemma_tokenizer.cpp @@ -1,8 +1,8 @@ #include "gemma_tokenizer.h" +#include "core/util.h" #include "ggml.h" #include "json.hpp" -#include "util.h" #include "vocab/vocab.h" std::string GemmaTokenizer::normalize(const std::string& text) const { diff --git a/src/tokenizers/gpt_oss_tokenizer.cpp b/src/tokenizers/gpt_oss_tokenizer.cpp index 8ec109a6..9779734c 100644 --- a/src/tokenizers/gpt_oss_tokenizer.cpp +++ b/src/tokenizers/gpt_oss_tokenizer.cpp @@ -1,7 +1,7 @@ #include "gpt_oss_tokenizer.h" +#include "core/util.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) { diff --git a/src/tokenizers/mistral_tokenizer.cpp b/src/tokenizers/mistral_tokenizer.cpp index 9b0624e3..cc418710 100644 --- a/src/tokenizers/mistral_tokenizer.cpp +++ b/src/tokenizers/mistral_tokenizer.cpp @@ -1,8 +1,8 @@ #include "mistral_tokenizer.h" +#include "core/util.h" #include "ggml.h" #include "json.hpp" -#include "util.h" #include "vocab/vocab.h" void MistralTokenizer::load_from_merges(const std::string& merges_utf8_str, const std::string& vocab_utf8_str) { diff --git a/src/tokenizers/qwen2_tokenizer.cpp b/src/tokenizers/qwen2_tokenizer.cpp index 46ee3117..79e683e7 100644 --- a/src/tokenizers/qwen2_tokenizer.cpp +++ b/src/tokenizers/qwen2_tokenizer.cpp @@ -1,6 +1,6 @@ #include "qwen2_tokenizer.h" -#include "util.h" +#include "core/util.h" #include "vocab/vocab.h" void Qwen2Tokenizer::load_from_merges(const std::string& merges_utf8_str) { diff --git a/src/tokenizers/t5_unigram_tokenizer.cpp b/src/tokenizers/t5_unigram_tokenizer.cpp index 8ed4df53..64e9e071 100644 --- a/src/tokenizers/t5_unigram_tokenizer.cpp +++ b/src/tokenizers/t5_unigram_tokenizer.cpp @@ -6,9 +6,9 @@ #include #include +#include "core/util.h" #include "json.hpp" #include "tokenize_util.h" -#include "util.h" #include "vocab/vocab.h" // Port from: https://github.com/google/sentencepiece/blob/master/src/unigram_model.h diff --git a/src/tokenizers/tokenizer.cpp b/src/tokenizers/tokenizer.cpp index 556cadd8..cb40e9b4 100644 --- a/src/tokenizers/tokenizer.cpp +++ b/src/tokenizers/tokenizer.cpp @@ -4,7 +4,7 @@ #include #include -#include "util.h" +#include "core/util.h" void Tokenizer::add_special_token(const std::string& token) { special_tokens.push_back(token); diff --git a/src/tokenizers/tokenizer.h b/src/tokenizers/tokenizer.h index e044285b..893759e0 100644 --- a/src/tokenizers/tokenizer.h +++ b/src/tokenizers/tokenizer.h @@ -12,9 +12,10 @@ using on_new_token_cb_t = std::function& class Tokenizer { protected: std::vector special_tokens; - bool add_bos_token = false; - bool add_eos_token = false; - bool pad_left = false; + bool add_bos_token = false; + bool add_eos_token = false; + bool pad_left = false; + bool normalize_before_split = false; std::string end_of_word_suffix; virtual std::string decode_token(int token_id) const = 0; diff --git a/src/upscaler.cpp b/src/upscaler.cpp index cef26e36..88a8a633 100644 --- a/src/upscaler.cpp +++ b/src/upscaler.cpp @@ -1,9 +1,10 @@ #include "upscaler.h" -#include "ggml_extend.hpp" -#include "model.h" +#include "core/ggml_extend.hpp" +#include "core/util.h" +#include "model_loader.h" #include "stable-diffusion.h" -#include "util.h" +#include #include UpscalerGGML::UpscalerGGML(int n_threads, @@ -18,6 +19,12 @@ UpscalerGGML::UpscalerGGML(int n_threads, params_backend_spec(std::move(params_backend_spec)) { } +UpscalerGGML::~UpscalerGGML() { + // ModelManager holds raw ggml tensor pointers owned by the runner context. + model_manager.reset(); + esrgan_upscaler.reset(); +} + void UpscalerGGML::set_max_graph_vram_bytes(size_t max_vram_bytes) { max_graph_vram_bytes = max_vram_bytes; if (esrgan_upscaler) { @@ -33,17 +40,12 @@ void UpscalerGGML::set_stream_layers_enabled(bool enabled) { } bool UpscalerGGML::load_from_file(const std::string& esrgan_path, - bool offload_params_to_cpu, int n_threads) { ggml_log_set(ggml_log_callback_default, nullptr); std::string error; if (!backend_manager.init(backend_spec.c_str(), params_backend_spec.c_str(), - offload_params_to_cpu, - false, - false, - false, &error)) { LOG_ERROR("upscaler backend config failed: %s", error.c_str()); return false; @@ -72,22 +74,39 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path, return false; } - ModelLoader model_loader; - if (!model_loader.init_from_file_and_convert_name(esrgan_path)) { + model_manager = std::make_shared(); + model_manager->set_n_threads(n_threads); + model_manager->set_enable_mmap(false); + + ModelLoader& model_loader = model_manager->loader(); + if (!model_loader.init_from_file_and_convert_name(esrgan_path, "", VERSION_ESRGAN)) { LOG_ERROR("init model loader from file failed: '%s'", esrgan_path.c_str()); + return false; } model_loader.set_wtype_override(model_data_type); LOG_INFO("Upscaler weight type: %s", ggml_type_name(model_data_type)); esrgan_upscaler = std::make_shared(backend_for(SDBackendModule::UPSCALER), - params_backend_for(SDBackendModule::UPSCALER), - tile_size, - model_loader.get_tensor_storage_map()); + model_loader.get_tensor_storage_map(), + model_manager); + if (esrgan_upscaler == nullptr || esrgan_upscaler->rrdb_net == nullptr) { + LOG_ERROR("init esrgan model from metadata failed: '%s'", esrgan_path.c_str()); + return false; + } esrgan_upscaler->set_max_graph_vram_bytes(max_graph_vram_bytes); esrgan_upscaler->set_stream_layers_enabled(stream_layers_enabled); if (direct) { esrgan_upscaler->set_conv2d_direct_enabled(true); } - if (!esrgan_upscaler->load_from_file(esrgan_path, n_threads)) { + + std::map tensors; + esrgan_upscaler->get_param_tensors(tensors); + if (!model_manager->register_param_tensors("ESRGAN", + std::move(tensors), + backend_manager.params_backend_is_disk(SDBackendModule::UPSCALER) ? ModelManager::ResidencyMode::Disk : ModelManager::ResidencyMode::ParamBackend, + backend_for(SDBackendModule::UPSCALER), + params_backend_for(SDBackendModule::UPSCALER)) || + !model_manager->validate_registered_tensors()) { + LOG_ERROR("register esrgan tensors with model manager failed"); return false; } return true; @@ -95,6 +114,7 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path, sd::Tensor UpscalerGGML::upscale_tensor(const sd::Tensor& input_tensor) { sd::Tensor upscaled; + const int scale = esrgan_upscaler->config.scale; if (tile_size <= 0 || (input_tensor.shape()[0] <= tile_size && input_tensor.shape()[1] <= tile_size)) { upscaled = esrgan_upscaler->compute(n_threads, input_tensor); } else { @@ -108,9 +128,9 @@ sd::Tensor UpscalerGGML::upscale_tensor(const sd::Tensor& input_te }; upscaled = process_tiles_2d(input_tensor, - static_cast(input_tensor.shape()[0] * esrgan_upscaler->scale), - static_cast(input_tensor.shape()[1] * esrgan_upscaler->scale), - esrgan_upscaler->scale, + static_cast(input_tensor.shape()[0] * scale), + static_cast(input_tensor.shape()[1] * scale), + scale, tile_size, tile_size, 0.25f, @@ -129,8 +149,9 @@ sd::Tensor UpscalerGGML::upscale_tensor(const sd::Tensor& input_te sd_image_t UpscalerGGML::upscale(sd_image_t input_image, uint32_t upscale_factor) { // upscale_factor, unused for RealESRGAN_x4plus_anime_6B.pth sd_image_t upscaled_image = {0, 0, 0, nullptr}; - int output_width = (int)input_image.width * esrgan_upscaler->scale; - int output_height = (int)input_image.height * esrgan_upscaler->scale; + const int scale = esrgan_upscaler->config.scale; + int output_width = (int)input_image.width * scale; + int output_height = (int)input_image.height * scale; LOG_INFO("upscaling from (%i x %i) to (%i x %i)", input_image.width, input_image.height, output_width, output_height); @@ -153,7 +174,6 @@ struct upscaler_ctx_t { }; upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path_c_str, - bool offload_params_to_cpu, bool direct, int n_threads, int tile_size, @@ -170,7 +190,7 @@ upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path_c_str, return nullptr; } - if (!upscaler_ctx->upscaler->load_from_file(esrgan_path, offload_params_to_cpu, n_threads)) { + if (!upscaler_ctx->upscaler->load_from_file(esrgan_path, n_threads)) { delete upscaler_ctx->upscaler; upscaler_ctx->upscaler = nullptr; free(upscaler_ctx); @@ -179,15 +199,48 @@ upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path_c_str, return upscaler_ctx; } -sd_image_t upscale(upscaler_ctx_t* upscaler_ctx, sd_image_t input_image, uint32_t upscale_factor) { - return upscaler_ctx->upscaler->upscale(input_image, upscale_factor); +bool upscale(upscaler_ctx_t* upscaler_ctx, + sd_image_t input_image, + uint32_t upscale_factor, + sd_image_t** images_out, + int* num_images_out) { + if (images_out != nullptr) { + *images_out = nullptr; + } + if (num_images_out != nullptr) { + *num_images_out = 0; + } + if (upscaler_ctx == nullptr || upscaler_ctx->upscaler == nullptr) { + return false; + } + + sd_image_t* result_images = (sd_image_t*)calloc(1, sizeof(sd_image_t)); + if (result_images == nullptr) { + return false; + } + + result_images[0] = upscaler_ctx->upscaler->upscale(input_image, upscale_factor); + if (result_images[0].data == nullptr) { + free(result_images); + return false; + } + + if (num_images_out != nullptr) { + *num_images_out = 1; + } + if (images_out != nullptr) { + *images_out = result_images; + } else { + free_sd_images(result_images, 1); + } + return true; } int get_upscale_factor(upscaler_ctx_t* upscaler_ctx) { if (upscaler_ctx == nullptr || upscaler_ctx->upscaler == nullptr || upscaler_ctx->upscaler->esrgan_upscaler == nullptr) { return 1; } - return upscaler_ctx->upscaler->esrgan_upscaler->scale; + return upscaler_ctx->upscaler->esrgan_upscaler->config.scale; } void free_upscaler_ctx(upscaler_ctx_t* upscaler_ctx) { diff --git a/src/upscaler.h b/src/upscaler.h index e2197694..38150f59 100644 --- a/src/upscaler.h +++ b/src/upscaler.h @@ -1,16 +1,18 @@ #ifndef __SD_UPSCALER_H__ #define __SD_UPSCALER_H__ -#include "esrgan.hpp" -#include "ggml_extend_backend.h" +#include "core/ggml_extend_backend.h" +#include "core/tensor.hpp" +#include "model/upscaler/esrgan.hpp" +#include "model_manager.h" #include "stable-diffusion.h" -#include "tensor.hpp" #include #include struct UpscalerGGML { SDBackendManager backend_manager; + std::shared_ptr model_manager; ggml_type model_data_type = GGML_TYPE_F16; std::shared_ptr esrgan_upscaler; std::string esrgan_path; @@ -27,9 +29,9 @@ struct UpscalerGGML { int tile_size = 128, std::string backend_spec = "", std::string params_backend_spec = ""); + ~UpscalerGGML(); bool load_from_file(const std::string& esrgan_path, - bool offload_params_to_cpu, int n_threads); void set_max_graph_vram_bytes(size_t max_vram_bytes); void set_stream_layers_enabled(bool enabled); diff --git a/src/weight_manager.h b/src/weight_manager.h new file mode 100644 index 00000000..28d6cf5c --- /dev/null +++ b/src/weight_manager.h @@ -0,0 +1,15 @@ +#ifndef __WEIGHT_MANAGER_H__ +#define __WEIGHT_MANAGER_H__ + +#include + +struct ggml_tensor; + +struct RunnerWeightManager { + virtual ~RunnerWeightManager() = default; + virtual bool prepare_params(const std::vector& tensors) = 0; + virtual void release_compute_backend_params(const std::vector& tensors) = 0; + virtual void release_params_backend_params(const std::vector& tensors) = 0; +}; + +#endif // __WEIGHT_MANAGER_H__