fix: prevent torch checkpoint offset overflow (#1832)

.github/workflows/build.yml

@@ -449,8 +449,8 @@ jobs:
     runs-on: windows-2022
 
     env:
-      ROCM_VERSION: "7.13.0"
-      GPU_TARGETS: "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1150;gfx1151;gfx1152;gfx1200;gfx1201"
+      ROCM_VERSION: "7.14.0"
+      GPU_TARGETS: "gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201"
 
     steps:
       - uses: actions/checkout@v3
@@ -472,34 +472,68 @@ jobs:
         uses: actions/cache@v4
         with:
           path: C:\TheRock\build
-          key: rocm-${{ env.ROCM_VERSION }}-gfx1151-${{ runner.os }}
+          key: rocm-wheels-${{ env.ROCM_VERSION }}-${{ runner.os }}
 
       - name: ccache
         uses: ggml-org/ccache-action@v1.2.16
         with:
-          key: windows-latest-rocm-${{ env.ROCM_VERSION }}-x64
+          key: windows-rocm-${{ env.ROCM_VERSION }}-x64
           evict-old-files: 1d
 
-      - name: Install ROCm
+      - name: Install ROCm with Wheels
         if: steps.cache-rocm.outputs.cache-hit != 'true'
         run: |
           $ErrorActionPreference = "Stop"
-          write-host "Downloading AMD ROCm ${{ env.ROCM_VERSION }} tarball"
-          Invoke-WebRequest -Uri "https://repo.amd.com/rocm/tarball/therock-dist-windows-gfx1151-${{ env.ROCM_VERSION }}.tar.gz" -OutFile "${env:RUNNER_TEMP}\rocm.tar.gz"
-          write-host "Extracting ROCm tarball"
-          mkdir C:\TheRock\build -Force
-          tar -xzf "${env:RUNNER_TEMP}\rocm.tar.gz" -C C:\TheRock\build --strip-components=1
-          write-host "Completed ROCm extraction"
+          write-host "Setting up Python virtual environment"
+
+          # Create the venv directly at the cache location to avoid relocation issues
+          New-Item -Path "C:\TheRock\build" -ItemType Directory -Force | Out-Null
+          python -m venv C:\TheRock\build\.venv
+          & C:\TheRock\build\.venv\Scripts\Activate.ps1
+
+          write-host "Upgrading pip"
+          python -m pip install --upgrade pip
+
+          write-host "Installing ROCm wheels for multi-arch support"
+          # Install ROCm wheels for multi-arch support (this may take several minutes)
+          python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ "rocm[libraries,devel]==${{env.ROCM_VERSION}}"
+
+          # Pre-expand the devel tree so it is included in the cache
+          write-host "Initializing ROCm devel tree"
+          rocm-sdk init
+          if ($LASTEXITCODE -ne 0) { throw "rocm-sdk init failed with exit code $LASTEXITCODE" }
+          write-host "Completed ROCm wheel installation to C:\TheRock\build"
 
       - name: Setup ROCm Environment
         run: |
-          $rocmPath = "C:\TheRock\build"
+          $ErrorActionPreference = "Stop"
+
+          # Activate venv from cache or fresh install
+          & C:\TheRock\build\.venv\Scripts\Activate.ps1
+
+          # Expand the devel tree (idempotent; no-op if already done during install)
+          rocm-sdk init
+          if ($LASTEXITCODE -ne 0) { throw "rocm-sdk init failed with exit code $LASTEXITCODE" }
+
+          # Get ROCm installation paths using the rocm-sdk CLI tool
+          $rocmPath = (rocm-sdk path --root)
+          if (-not $rocmPath) { throw "rocm-sdk path --root returned empty - devel package may not be installed" }
+          $rocmPath = $rocmPath.Trim()
+          $cmakePath = (rocm-sdk path --cmake).Trim()
+          $binPath = (rocm-sdk path --bin).Trim()
+          write-host "ROCm root: $rocmPath"
+          write-host "CMake path: $cmakePath"
+          write-host "Bin path: $binPath"
+
           echo "HIP_PATH=$rocmPath" >> $env:GITHUB_ENV
+          echo "CMAKE_PREFIX_PATH=$cmakePath" >> $env:GITHUB_ENV
           echo "HIP_DEVICE_LIB_PATH=$rocmPath\lib\llvm\amdgcn\bitcode" >> $env:GITHUB_ENV
           echo "HIP_PLATFORM=amd" >> $env:GITHUB_ENV
           echo "LLVM_PATH=$rocmPath\lib\llvm" >> $env:GITHUB_ENV
-          echo "$rocmPath\bin" >> $env:GITHUB_PATH
-          echo "$rocmPath\lib\llvm\bin" >> $env:GITHUB_PATH
+          echo "$binPath" >> $env:GITHUB_PATH
+
+          # Keep venv in PATH for subsequent steps
+          echo "C:\TheRock\build\.venv\Scripts" >> $env:GITHUB_PATH
 
       - name: Build
         run: |
@@ -527,139 +561,6 @@ jobs:
       - name: Pack artifacts
         if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
         run: |
-          $ErrorActionPreference = "Stop"
-          $dst = "build\bin"
-          $rocmBin = Join-Path "${env:HIP_PATH}" "bin"
-          $requiredRocmPaths = @(
-            (Join-Path $rocmBin "rocblas.dll"),
-            (Join-Path $rocmBin "rocblas\library")
-          )
-          foreach ($path in $requiredRocmPaths) {
-            if (!(Test-Path $path)) {
-              throw "Missing ROCm runtime dependency: $path"
-            }
-          }
-
-          foreach ($pattern in @("rocblas*.dll", "hipblas*.dll", "libhipblas*.dll")) {
-            Copy-Item -Path (Join-Path $rocmBin $pattern) -Destination $dst -Force -ErrorAction SilentlyContinue
-          }
-
-          foreach ($dir in @("rocblas", "hipblaslt")) {
-            $src = Join-Path $rocmBin $dir
-            if (Test-Path $src) {
-              Copy-Item -Path $src -Destination $dst -Recurse -Force
-            }
-          }
-
-          7z a sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-rocm-${{ env.ROCM_VERSION }}-x64.zip .\build\bin\*
-
-      - name: Upload artifacts
-        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        uses: actions/upload-artifact@v4
-        with:
-          name: sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-rocm-${{ env.ROCM_VERSION }}-x64.zip
-          path: |
-            sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-rocm-${{ env.ROCM_VERSION }}-x64.zip
-
-  windows-latest-cmake-hip:
-    runs-on: windows-2022
-
-    env:
-      HIPSDK_INSTALLER_VERSION: "26.Q1"
-      ROCM_VERSION: "7.1.1"
-      GPU_TARGETS: "gfx1150;gfx1151;gfx1200;gfx1201;gfx1100;gfx1101;gfx1102;gfx1030;gfx1031;gfx1032"
-
-    steps:
-      - uses: actions/checkout@v3
-        with:
-          submodules: recursive
-
-      - name: Setup Node
-        uses: actions/setup-node@v4
-        with:
-          node-version: 20
-
-      - name: Setup pnpm
-        uses: pnpm/action-setup@v4
-        with:
-          version: 10.15.1
-
-      - name: Cache ROCm Installation
-        id: cache-rocm
-        uses: actions/cache@v4
-        with:
-          path: C:\Program Files\AMD\ROCm
-          key: rocm-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ runner.os }}
-
-      - name: ccache
-        uses: ggml-org/ccache-action@v1.2.16
-        with:
-          key: windows-latest-cmake-hip-${{ env.HIPSDK_INSTALLER_VERSION }}-x64
-          evict-old-files: 1d
-
-      - name: Install ROCm
-        if: steps.cache-rocm.outputs.cache-hit != 'true'
-        run: |
-          $ErrorActionPreference = "Stop"
-          write-host "Downloading AMD HIP SDK Installer"
-          Invoke-WebRequest -Uri "https://download.amd.com/developer/eula/rocm-hub/AMD-Software-PRO-Edition-${{ env.HIPSDK_INSTALLER_VERSION }}-Win11-For-HIP.exe" -OutFile "${env:RUNNER_TEMP}\rocm-install.exe"
-          write-host "Installing AMD HIP SDK"
-          $proc = Start-Process "${env:RUNNER_TEMP}\rocm-install.exe" -ArgumentList '-install' -NoNewWindow -PassThru
-          $completed = $proc.WaitForExit(600000)
-          if (-not $completed) {
-              Write-Error "ROCm installation timed out after 10 minutes. Killing the process"
-              $proc.Kill()
-              exit 1
-          }
-          if ($proc.ExitCode -ne 0) {
-              Write-Error "ROCm installation failed with exit code $($proc.ExitCode)"
-              exit 1
-          }
-          write-host "Completed AMD HIP SDK installation"
-
-      - name: Verify ROCm
-        run: |
-          # Find and test ROCm installation
-          $clangPath = Get-ChildItem 'C:\Program Files\AMD\ROCm\*\bin\clang.exe' | Select-Object -First 1
-          if (-not $clangPath) {
-            Write-Error "ROCm installation not found"
-            exit 1
-          }
-          & $clangPath.FullName --version
-          # Set HIP_PATH environment variable for later steps
-          echo "HIP_PATH=$(Resolve-Path 'C:\Program Files\AMD\ROCm\*\bin\clang.exe' | split-path | split-path)" >> $env:GITHUB_ENV
-
-      - name: Build
-        run: |
-          mkdir build
-          cd build
-          $env:CMAKE_PREFIX_PATH="${env:HIP_PATH}"
-          cmake .. `
-            -G "Unix Makefiles" `
-            -DSD_HIPBLAS=ON `
-            -DSD_BUILD_SHARED_LIBS=ON `
-            -DGGML_NATIVE=OFF `
-            -DCMAKE_C_COMPILER=clang `
-            -DCMAKE_CXX_COMPILER=clang++ `
-            -DCMAKE_BUILD_TYPE=Release `
-            -DGPU_TARGETS="${{ env.GPU_TARGETS }}"
-          cmake --build . --config Release --parallel ${env:NUMBER_OF_PROCESSORS}
-
-      - name: Get commit hash
-        id: commit
-        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        uses: prompt/actions-commit-hash@v2
-
-      - name: Pack artifacts
-        if: ${{ ( github.event_name == 'push' && github.ref == 'refs/heads/master' ) || github.event.inputs.create_release == 'true' }}
-        run: |
-          md "build\bin\rocblas\library\"
-          md "build\bin\hipblaslt\library"
-          cp "${env:HIP_PATH}\bin\libhipblas.dll" "build\bin\"
-          cp "${env:HIP_PATH}\bin\libhipblaslt.dll" "build\bin\"
-          cp "${env:HIP_PATH}\bin\rocblas.dll" "build\bin\"
-          cp "${env:HIP_PATH}\bin\rocblas\library\*" "build\bin\rocblas\library\"
-          cp "${env:HIP_PATH}\bin\hipblaslt\library\*" "build\bin\hipblaslt\library\"
           7z a sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-rocm-${{ env.ROCM_VERSION }}-x64.zip .\build\bin\*
 
       - name: Upload artifacts
@@ -679,11 +580,8 @@ jobs:
     strategy:
       matrix:
         include:
-          - ROCM_VERSION: "7.2.1"
-            gpu_targets: "gfx908;gfx90a;gfx942;gfx1030;gfx1031;gfx1032;gfx1100;gfx1101;gfx1102;gfx1151;gfx1150;gfx1200;gfx1201"
-            build: 'x64'
-          - ROCM_VERSION: "7.13.0"
-            gpu_targets: "gfx906;gfx908;gfx90a;gfx942;gfx950;gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1150;gfx1151;gfx1152;gfx1200;gfx1201"
+          - ROCM_VERSION: "7.14.0"
+            gpu_targets: "gfx900;gfx906;gfx908;gfx90a;gfx90c;gfx942;gfx950;gfx1010;gfx1011;gfx1012;gfx1030;gfx1031;gfx1032;gfx1033;gfx1034;gfx1035;gfx1036;gfx1100;gfx1101;gfx1102;gfx1103;gfx1150;gfx1151;gfx1152;gfx1153;gfx1200;gfx1201"
             build: x64
 
     steps:
@@ -702,7 +600,7 @@ jobs:
       - name: Dependencies
         id: depends
         run: |
-          sudo apt install -y build-essential cmake wget zip ninja-build
+          sudo apt install -y build-essential git cmake wget
 
       - name: Free disk space
         run: |
@@ -723,38 +621,36 @@ jobs:
           sudo apt clean
           df -h
 
-      - name: Setup Legacy ROCm
-        if: matrix.ROCM_VERSION == '7.2.1'
-        id: legacy_env
-        run: |
-          sudo mkdir --parents --mode=0755 /etc/apt/keyrings
-          wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | \
-            gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null
-
-          sudo tee /etc/apt/sources.list.d/rocm.list << EOF
-          deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/${{ matrix.ROCM_VERSION }} noble main
-          EOF
-
-          sudo tee /etc/apt/preferences.d/rocm-pin-600 << EOF
-          Package: *
-          Pin: release o=repo.radeon.com
-          Pin-Priority: 600
-          EOF
-
-          sudo apt update
-          sudo apt-get install -y libssl-dev rocm-hip-sdk
-
-      - name: Setup TheRock
-        if: matrix.ROCM_VERSION != '7.2.1'
+      - name: Setup TheRock with Wheels
         id: therock_env
         run: |
-          wget https://repo.amd.com/rocm/tarball/therock-dist-linux-gfx1151-${{ matrix.ROCM_VERSION }}.tar.gz
-          mkdir install
-          tar -xf *.tar.gz -C install
-          export ROCM_PATH=$(pwd)/install
-          echo ROCM_PATH=$ROCM_PATH >> $GITHUB_ENV
-          echo PATH=$PATH:$ROCM_PATH/bin >> $GITHUB_ENV
-          echo LD_LIBRARY_PATH=$ROCM_PATH/lib:$ROCM_PATH/llvm/lib:$ROCM_PATH/lib/rocprofiler-systems >> $GITHUB_ENV
+          # Create Python virtual environment
+          python3 -m venv .venv
+          source .venv/bin/activate
+
+          # Install ROCm wheels for build
+          # libraries = HIP runtime and CMake configs needed for linking
+          # devel = compilers, headers, static libs
+          python -m pip install --upgrade pip
+          python -m pip install --index-url https://repo.amd.com/rocm/whl-multi-arch/ "rocm[libraries,devel]==${{matrix.ROCM_VERSION}}"
+
+          # Get ROCm installation paths using the rocm-sdk CLI tool
+          ROCM_PATH=$(rocm-sdk path --root)
+          CMAKE_PATH=$(rocm-sdk path --cmake)
+          BIN_PATH=$(rocm-sdk path --bin)
+          echo "ROCM_PATH=$ROCM_PATH"
+          echo "CMAKE_PATH=$CMAKE_PATH"
+          echo "BIN_PATH=$BIN_PATH"
+
+          # Set environment variables
+          echo "ROCM_PATH=$ROCM_PATH" >> $GITHUB_ENV
+          echo "CMAKE_PREFIX_PATH=$CMAKE_PATH" >> $GITHUB_ENV
+          echo "HIP_PATH=$ROCM_PATH" >> $GITHUB_ENV
+          echo "PATH=$BIN_PATH:${PATH}" >> $GITHUB_ENV
+          echo "LD_LIBRARY_PATH=$ROCM_PATH/lib:${LD_LIBRARY_PATH:-}" >> $GITHUB_ENV
+
+          # Keep venv activated for subsequent steps
+          echo "$(pwd)/.venv/bin" >> $GITHUB_PATH
 
       # setup-node installs into /opt/hostedtoolcache, which is removed above.
       # Keep Node/pnpm setup after disk cleanup so the server frontend can be embedded.
@@ -839,7 +735,6 @@ jobs:
       - build-and-push-docker-images
       - macOS-latest-cmake
       - windows-latest-cmake
-      - windows-latest-cmake-hip
       - windows-latest-rocm
 
     steps:

README.md

@@ -54,6 +54,7 @@ API and command-line option may change frequently.***
     - [ERNIE-Image](./docs/ernie_image.md)
     - [Boogu Image](./docs/boogu_image.md)
     - [Krea2](./docs/krea2.md)
+    - [Mage-Flow](./docs/mage_flow.md)
     - [SeFi-Image](./docs/sefi_image.md)
     - [HiDream-O1-Image](./docs/hidream_o1_image.md)
     - [Ideogram4](./docs/ideogram4.md)
@@ -62,11 +63,14 @@ API and command-line option may change frequently.***
     - [Qwen Image Edit series](./docs/qwen_image_edit.md)
     - [LongCat Image Edit](./docs/longcat_image.md)
     - [Boogu Image Edit](./docs/boogu_image.md)
+    - [Mage-Flow-Edit](./docs/mage_flow.md#image-editing)
   - Video Models
     - [Wan2.1/Wan2.2](./docs/wan.md)
     - [LTX-2.3](./docs/ltx2.md)
+    - [HunyuanVideo 1.5](./docs/hunyuan_video.md)
     - [LingBot-Video](./docs/lingbot_video.md)
   - [PhotoMaker](./docs/photo_maker.md) support.
+  - [IP-Adapter](./docs/ip_adapter.md) support (SD 1.5 and SDXL)
   - Control Net support with SD 1.5
   - [ADetailer](./docs/adetailer.md)
   - LoRA support, same as [stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#lora)
@@ -122,7 +126,7 @@ API and command-line option may change frequently.***
     - Stable Diffusion v1.5 from https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5 
 
     ```sh
-    curl -L -O https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors
+    curl -L -O https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors
     ```
 
 ### Generate an image with just one command
@@ -164,6 +168,7 @@ These projects wrap `stable-diffusion.cpp` for easier use in other languages/fra
 
 These projects use `stable-diffusion.cpp` as a backend for their image generation.
 
+- [GIMP Plugins](https://github.com/themanyone/gimp-plugins)
 - [Jellybox](https://jellybox.com)
 - [Stable Diffusion GUI](https://github.com/fszontagh/sd.cpp.gui.wx)
 - [Stable Diffusion CLI-GUI](https://github.com/piallai/stable-diffusion.cpp)

docker/Dockerfile.vulkan

@@ -33,7 +33,7 @@ RUN cmake --build ./build --config Release -j$(nproc)
 FROM ubuntu:$UBUNTU_VERSION AS runtime
 
 RUN apt-get update && \
-    apt-get install --yes --no-install-recommends libgomp1 libvulkan1 mesa-vulkan-drivers && \
+    apt-get install --yes --no-install-recommends libgomp1 libvulkan1 mesa-vulkan-drivers libglvnd0 libgl1 libglx0 libegl1 libgles2 && \
     apt-get clean
 
 COPY --from=build /sd.cpp/build/bin /sd.cpp/bin

docs/edit.md

@@ -16,6 +16,7 @@ Depending on the architecture, different models handle reference images differen
 | [**Flux.2 [Dev] / Flux.2 [Klein]**](./flux2.md) | `flux2` |
 | [**Boogu Image Edit**](./boogu_image.md) | `z_image_omni` |
 | **Krea2 (Community Edit LoRAs)** | `krea2_ostris_edit` |
+| [**Mage-Flow-Edit**](./mage_flow.md#image-editing) | `mage_flow` |
 | **Anima (Community Edit LoRAs)** | `cosmos_reference` |
 
 Stable-diffusion.spp also supports basic Unet-based editing models like instruct-pix2pix or CosXL-Edit. This document is not about those.
@@ -48,6 +49,7 @@ The `--ref-image-args` argument accepts a comma-separated list of key-value pair
 | `qwen_layered` | Qwen Image Layered |
 | `z_image_omni` | Boogu, Z-Image Omni |
 | `krea2_ostris_edit` | Most Krea2 Community edit LoRAs (trained with Ostris script) |
+| `mage_flow` | Mage-Flow-Edit |
 | `krea2_edit` | Specifically for [lbouaraba/krea2edit](https://huggingface.co/conradlocke/krea2-identity-edit). (or similar) |
 | `cosmos_reference` | For Anima |
 | `default` | Uses the automatic detection based on model architecture. |
@@ -83,6 +85,7 @@ For a technical overview of how each preset is configured, see the table below.
 | `flux2` | No | `increase` | `none` | |
 | `qwen` | Yes | `increase` | `area` | |
 | `qwen_layered` | Yes | `decrease` | `area` | |
+| `mage_flow` | Yes | `increase` | `longest` | `vlm_max_size = 384`, VAE input resized to target |
 | `z_image_omni` | Yes | `fixed` | `area` | |
 | `krea2_ostris_edit`| Yes | `increase` | `area` | `force_ref_timestep_zero = true` |
 | `krea2_edit` | Yes | `increase` | `longest` | `vlm_size = 768` |

docs/hunyuan_video.md

@@ -0,0 +1,24 @@
+# HunyuanVideo 1.5
+
+HunyuanVideo 1.5 uses a HunyuanVideo diffusion transformer, a causal video VAE, Qwen2.5-VL 7B for the main text conditioning, 
+and ByT5 Small GlyphXL for glyph-aware text conditioning.
+
+## Download weights
+
+- Download HunyuanVideo 1.5
+    - safetensors: https://huggingface.co/Comfy-Org/HunyuanVideo_1.5_repackaged/tree/main/split_files/diffusion_models
+- Download vae
+    - safetensors: https://huggingface.co/Comfy-Org/HunyuanVideo_1.5_repackaged/tree/main/split_files/vae
+- Download qwen_2.5_vl 7b
+    - safetensors: https://huggingface.co/Comfy-Org/Qwen-Image_ComfyUI/tree/main/split_files/text_encoders
+    - gguf: https://huggingface.co/mradermacher/Qwen2.5-VL-7B-Instruct-GGUF/tree/main
+- Download byt5 small glyphxl
+    - safetensros: https://huggingface.co/Comfy-Org/HunyuanVideo_1.5_repackaged/tree/main/split_files/text_encoders
+
+## Text-to-video example
+
+```shell
+.\bin\Release\sd-cli.exe -M vid_gen --diffusion-model  ..\models\diffusion_models\hunyuanvideo1.5_720p_t2v_fp16.safetensors --vae ..\models\vae\hunyuanvideo15_vae_fp16.safetensors --llm ..\models\text_encoders\qwen_2.5_vl_7b.safetensors --t5xxl ..\models\text_encoders\byt5_small_glyphxl_fp16.safetensors  -p "a lovely cat" --cfg-scale 6.0 --sampling-method euler -v -W 1280 -H 720 --offload-to-cpu --diffusion-fa --video-frames 33 --vae-tiling
+```
+
+<video src=../assets/hunyuan_video/hy1.5_t2v.mp4 controls="controls" muted="muted" type="video/mp4"></video>

docs/ip_adapter.md

@@ -0,0 +1,55 @@
+# IP-Adapter
+
+stable-diffusion.cpp supports [IP-Adapter](https://github.com/tencent-ailab/IP-Adapter)
+image-prompt conditioning for SD 1.5 and SDXL. Given a reference image,
+IP-Adapter transfers the subject and appearance of that image into the
+generation, alongside the text prompt.
+
+IP-Adapter encodes the reference image with a CLIP-Vision (ViT-H/14)
+encoder, projects the embedding into a few image tokens, and injects them
+through a decoupled cross-attention added to every attn2 layer of the
+UNet. It composes with Control Net, so a reference image (appearance) and
+an OpenPose hint (pose) can be combined in a single generation.
+
+## Required weights
+
+1. A base SD 1.5 or SDXL model.
+2. A CLIP-Vision (ViT-H/14) image encoder, passed with `--clip_vision`
+   (for example `clip_vision_h.safetensors`).
+3. An IP-Adapter weight file, passed with `--ip-adapter`. The `vit-h`
+   variants reuse the same ViT-H encoder as above. From
+   [h94/IP-Adapter](https://huggingface.co/h94/IP-Adapter):
+   - SD 1.5: `models/ip-adapter_sd15.safetensors`
+   - SDXL: `sdxl_models/ip-adapter_sdxl_vit-h.safetensors`
+
+## Options
+
+- `--ip-adapter <path>` path to the IP-Adapter weight file.
+- `--ip-adapter-image <path>` path to the reference image.
+- `--ip-adapter-strength <float>` strength of the IP-Adapter injection
+  (default 1.0). Lower values let the text prompt dominate; 0.6 to 0.8 is
+  a good starting range.
+
+## Example (SD 1.5)
+
+```
+sd-cli -m ..\models\sd_v1.5.safetensors --clip_vision ..\models\clip_vision_h.safetensors --ip-adapter ..\models\ip-adapter_sd15.safetensors --ip-adapter-image ..\assets\reference.png --ip-adapter-strength 0.8 -p "a woman, best quality" -n "lowres, bad anatomy" --cfg-scale 7 --steps 30 --sampling-method dpm++2m --scheduler karras -W 512 -H 512
+```
+
+## Example (SDXL)
+
+```
+sd-cli -m ..\models\sdxl.safetensors --clip_vision ..\models\clip_vision_h.safetensors --ip-adapter ..\models\ip-adapter_sdxl_vit-h.safetensors --ip-adapter-image ..\assets\reference.png --ip-adapter-strength 0.8 -p "a woman, best quality" -n "lowres, bad anatomy" --cfg-scale 6 --steps 25 --sampling-method dpm++2m --scheduler karras -W 1024 -H 1024 --diffusion-fa --vae-tiling
+```
+
+The SDXL VAE decode at 1024x1024 is memory heavy; add `--vae-tiling` (and
+`--offload-to-cpu`) on GPUs with limited VRAM.
+
+## Combining with Control Net
+
+Add the usual Control Net options to keep the reference appearance while
+controlling the pose:
+
+```
+sd-cli -m ..\models\sdxl.safetensors --clip_vision ..\models\clip_vision_h.safetensors --ip-adapter ..\models\ip-adapter_sdxl_vit-h.safetensors --ip-adapter-image ..\assets\character.png --ip-adapter-strength 0.9 --control-net ..\models\OpenPoseXL2.safetensors --control-image ..\assets\pose.png --control-strength 0.8 -p "a character, side view" --cfg-scale 6 --steps 25 -W 1024 -H 1024 --diffusion-fa --vae-tiling
+```

docs/mage_flow.md

@@ -0,0 +1,45 @@
+# Mage-Flow
+
+[Mage-Flow](https://github.com/microsoft/Mage) uses a 4B native-resolution multimodal diffusion transformer, Qwen3-VL for text and image conditioning, and the 128-channel Mage-VAE. Both text-to-image and instruction-based image editing checkpoints are supported.
+
+## Download weights
+
+- Download Mage-Flow
+    - safetensors: https://huggingface.co/microsoft/Mage-Flow/tree/main/transformer
+- Download Mage-Flow-Base
+    - safetensors: https://huggingface.co/microsoft/Mage-Flow-Base/tree/main/transformer
+- Download Mage-Flow-Turbo
+    - safetensors: https://huggingface.co/microsoft/Mage-Flow-Turbo/tree/main/transformer
+- Download Mage-Flow-Edit
+    - safetensors: https://huggingface.co/microsoft/Mage-Flow-Edit/tree/main/transformer
+- Download Mage-Flow-Edit-Turbo
+    - safetensors: https://huggingface.co/microsoft/Mage-Flow-Edit-Turbo/tree/main/transformer
+- Download Mage-Flow-Edit-Base
+    - safetensors: https://huggingface.co/microsoft/Mage-Flow-Edit-Base/tree/main/transformer
+- Download Mage-Flow vae
+    - safetensors: https://huggingface.co/microsoft/Mage-Flow/tree/main/vae
+- Download Qwen3-VL 4B
+    - safetensors: https://huggingface.co/Comfy-Org/Krea-2/tree/main/text_encoders
+    - gguf: https://huggingface.co/Qwen/Qwen3-VL-4B-Instruct-GGUF/tree/main
+
+## Text-to-image
+
+Use 30 steps for Base models and 4 steps with `--cfg-scale 1` for Turbo models. Image dimensions must be multiples of 16; the official checkpoints are trained for native resolutions from 512 to 2048 pixels.
+
+```bash
+.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\Mage-Flow-Turbo.safetensors --llm ..\models\text_encoders\Qwen3-VL-4B-Instruct-Q4_K_M.gguf --vae ..\models\vae\mage_vae.safetensors -p "a lovely cat holding a sign says 'mage.cpp'" --cfg-scale 1.0 --steps 4 --diffusion-fa -v --offload-to-cpu
+```
+
+<img width="256" alt="Mage-Flow example" src="../assets/mage_flow/example.png" />
+
+## Image editing
+
+Mage-Flow-Edit accepts one or more reference images. The default `mage_flow` reference preset sends each image to both Qwen3-VL and the diffusion transformer, caps the VLM copy's longest edge at 384 pixels, and keeps the VAE copy at the requested output resolution.
+
+For the Turbo edit checkpoint, use 4 steps and `--cfg-scale 1`.
+
+```bash
+.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\Mage-Flow-Edit.safetensors --llm ..\models\text_encoders\Qwen3-VL-4B-Instruct-Q4_K_M.gguf --llm_vision ..\models\text_encoders\Qwen3-VL-4B-Instruct-mmproj-BF16.gguf --vae ..\models\vae\mage_vae.safetensors -r ..\assets\flux\flux1-dev-q8_0.png -p "change 'flux.cpp' to 'mage.cpp'" --cfg-scale 4.0 --sampling-method euler -v --diffusion-fa --offload-to-cpu
+```
+
+<img width="256" alt="Mage-Flow-Edit example" src="../assets/mage_flow/edit_example.png" />

docs/pid.md

@@ -1,7 +1,7 @@
 # How to Use
 
 PiD is NVIDIA's Pixel Diffusion Decoder. It replaces the usual VAE decode or decode-then-upscale path with a pixel-space diffusion decoder conditioned on a
-source latent and text prompt.
+source latent and text prompt. Both the original PiD checkpoints and PiD 1.5 are supported.
 
 In stable-diffusion.cpp, PiD currently runs as an image edit pipeline: provide a reference image with `-r`/`--ref-image`, encode that image with a matching VAE, then let the PiD diffusion model decode/upscale directly to RGB.
 
@@ -16,6 +16,7 @@ In stable-diffusion.cpp, PiD currently runs as an image edit pipeline: provide a
     - Flux / Z-Image PiD: use the Flux VAE and pass `--vae-format flux`
     - SD3 PiD: use the SD3 VAE and pass `--vae-format sd3`
     - Flux.2 PiD: use the Flux.2 VAE and pass `--vae-format flux2`
+    - Qwen-Image PiD: use the Qwen-Image 2D VAE and pass `--vae-format wan`
 
 The official PiD model card should be checked before use. At the time of the initial PiD release, the official weights are under the NSCLv1 non-commercial license.
 

docs/sd.md

@@ -2,8 +2,8 @@
 
 - download original weights(.ckpt or .safetensors). For example
     - Stable Diffusion v1.4 from https://huggingface.co/CompVis/stable-diffusion-v-1-4-original
-    - Stable Diffusion v1.5 from https://huggingface.co/runwayml/stable-diffusion-v1-5
-    - Stable Diffuison v2.1 from https://huggingface.co/stabilityai/stable-diffusion-2-1
+    - Stable Diffusion v1.5 from https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5
+    - Stable Diffuison v2.1 from https://huggingface.co/Manojb/stable-diffusion-2-1-base
     - Stable Diffusion 3 2B from https://huggingface.co/stabilityai/stable-diffusion-3-medium
 
 ### txt2img example
@@ -34,4 +34,4 @@ Using formats of different precisions will yield results of varying quality.
 
 <p align="center">
   <img src="../assets/img2img_output.png" width="256x">
-</p>
+</p>

examples/CMakeLists.txt

@@ -1,4 +1,4 @@
 include_directories(${CMAKE_CURRENT_SOURCE_DIR})
 
 add_subdirectory(cli)
-add_subdirectory(server)
+add_subdirectory(server)

examples/cli/main.cpp

@@ -817,6 +817,16 @@ int main(int argc, const char* argv[]) {
         }
     }
 
+    if (gen_params.ip_adapter_image_path.size() > 0) {
+        if (!load_sd_image_from_file(gen_params.ip_adapter_image.put(),
+                                     gen_params.ip_adapter_image_path.c_str(),
+                                     0,
+                                     0)) {
+            LOG_ERROR("load image from '%s' failed", gen_params.ip_adapter_image_path.c_str());
+            return 1;
+        }
+    }
+
     if (!gen_params.control_video_path.empty()) {
         gen_params.control_frames.clear();
         if (!load_images_from_dir(gen_params.control_video_path,

examples/common/common.cpp

@@ -50,6 +50,9 @@ static sd_vae_format_t str_to_vae_format(const std::string& value) {
     if (value == "flux2") {
         return SD_VAE_FORMAT_FLUX2;
     }
+    if (value == "wan") {
+        return SD_VAE_FORMAT_WAN;
+    }
     return SD_VAE_FORMAT_COUNT;
 }
 
@@ -401,7 +404,7 @@ ArgOptions SDContextParams::get_options() {
          &vae_path},
         {"",
          "--vae-format",
-         "VAE latent format override: auto, flux, sd3, or flux2 (default: auto)",
+         "VAE latent format override: auto, flux, sd3, flux2, or wan (default: auto)",
          0,
          &vae_format},
         {"",
@@ -424,6 +427,11 @@ ArgOptions SDContextParams::get_options() {
          "path to control net model",
          0,
          &control_net_path},
+        {"",
+         "--ip-adapter",
+         "path to IP-Adapter model (requires --clip_vision)",
+         0,
+         &ip_adapter_path},
         {"",
          "--motion-module",
          "path to AnimateDiff motion module (SD 1.5); enables video generation on --video-frames > 1",
@@ -743,7 +751,7 @@ bool SDContextParams::validate(SDMode mode) {
     }
 
     if (str_to_vae_format(vae_format) == SD_VAE_FORMAT_COUNT) {
-        LOG_ERROR("error: vae_format must be 'auto', 'flux', 'sd3', or 'flux2'");
+        LOG_ERROR("error: vae_format must be 'auto', 'flux', 'sd3', 'flux2', or 'wan'");
         return false;
     }
 
@@ -873,6 +881,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) {
     sd_ctx_params.audio_vae_path                  = audio_vae_path.c_str();
     sd_ctx_params.taesd_path                      = taesd_path.c_str();
     sd_ctx_params.control_net_path                = control_net_path.c_str();
+    sd_ctx_params.ip_adapter_path                 = ip_adapter_path.c_str();
     sd_ctx_params.motion_module_path              = motion_module_path.c_str();
     sd_ctx_params.embeddings                      = embedding_vec.data();
     sd_ctx_params.embedding_count                 = static_cast<uint32_t>(embedding_vec.size());
@@ -963,6 +972,11 @@ ArgOptions SDGenerationParams::get_options() {
          "path to control image, control net",
          0,
          &control_image_path},
+        {"",
+         "--ip-adapter-image",
+         "path to the IP-Adapter reference image",
+         0,
+         &ip_adapter_image_path},
         {"",
          "--control-video",
          "path to control video frames, It must be a directory path. The video frames inside should be stored as images in "
@@ -1155,6 +1169,10 @@ ArgOptions SDGenerationParams::get_options() {
          "--control-strength",
          "strength to apply Control Net (default: 0.9). 1.0 corresponds to full destruction of information in init image",
          &control_strength},
+        {"",
+         "--ip-adapter-strength",
+         "strength to apply IP-Adapter (default: 1.0)",
+         &ip_adapter_strength},
         {"",
          "--moe-boundary",
          "timestep boundary for Wan2.2 MoE model. (default: 0.875). Only enabled if `--high-noise-steps` is set to -1",
@@ -1883,6 +1901,7 @@ bool SDGenerationParams::from_json_str(
 
     load_if_exists("strength", strength);
     load_if_exists("control_strength", control_strength);
+    load_if_exists("ip_adapter_strength", ip_adapter_strength);
     load_if_exists("moe_boundary", moe_boundary);
     load_if_exists("vace_strength", vace_strength);
 
@@ -2054,6 +2073,10 @@ bool SDGenerationParams::from_json_str(
         LOG_ERROR("invalid control_image");
         return false;
     }
+    if (!parse_image_json_field(j, "ip_adapter_image", 3, width, height, ip_adapter_image)) {
+        LOG_ERROR("invalid ip_adapter_image");
+        return false;
+    }
 
     return true;
 }
@@ -2476,29 +2499,31 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
         LOG_WARN("Notice: --increase-ref-index is deprecated. Use --ref-image-args \"ref_index_mode=increase\" instead.");
     }
 
-    params.loras             = lora_vec.empty() ? nullptr : lora_vec.data();
-    params.lora_count        = static_cast<uint32_t>(lora_vec.size());
-    params.prompt            = prompt.c_str();
-    params.negative_prompt   = negative_prompt.c_str();
-    params.clip_skip         = clip_skip;
-    params.init_image        = init_image.get();
-    params.ref_images        = ref_image_views.empty() ? nullptr : ref_image_views.data();
-    params.ref_images_count  = static_cast<int>(ref_image_views.size());
-    params.ref_image_args    = ref_image_args.c_str();
-    params.mask_image        = mask_image.get();
-    params.width             = get_resolved_width();
-    params.height            = get_resolved_height();
-    params.sample_params     = sample_params;
-    params.strength          = strength;
-    params.seed              = seed;
-    params.batch_count       = batch_count;
-    params.qwen_image_layers = qwen_image_layers;
-    params.control_image     = control_image.get();
-    params.control_strength  = control_strength;
-    params.pm_params         = pm_params;
-    params.pulid_params      = pulid_params;
-    params.vae_tiling_params = vae_tiling_params;
-    params.cache             = cache_params;
+    params.loras               = lora_vec.empty() ? nullptr : lora_vec.data();
+    params.lora_count          = static_cast<uint32_t>(lora_vec.size());
+    params.prompt              = prompt.c_str();
+    params.negative_prompt     = negative_prompt.c_str();
+    params.clip_skip           = clip_skip;
+    params.init_image          = init_image.get();
+    params.ref_images          = ref_image_views.empty() ? nullptr : ref_image_views.data();
+    params.ref_images_count    = static_cast<int>(ref_image_views.size());
+    params.ref_image_args      = ref_image_args.c_str();
+    params.mask_image          = mask_image.get();
+    params.width               = get_resolved_width();
+    params.height              = get_resolved_height();
+    params.sample_params       = sample_params;
+    params.strength            = strength;
+    params.seed                = seed;
+    params.batch_count         = batch_count;
+    params.qwen_image_layers   = qwen_image_layers;
+    params.control_image       = control_image.get();
+    params.control_strength    = control_strength;
+    params.ip_adapter_image    = ip_adapter_image.get();
+    params.ip_adapter_strength = ip_adapter_strength;
+    params.pm_params           = pm_params;
+    params.pulid_params        = pulid_params;
+    params.vae_tiling_params   = vae_tiling_params;
+    params.cache               = cache_params;
 
     params.hires.enabled             = hires_enabled;
     params.hires.upscaler            = resolved_hires_upscaler;
@@ -2787,6 +2812,7 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params,
     root["clip_skip"]             = gen_params.clip_skip;
     root["strength"]              = gen_params.strength;
     root["control_strength"]      = gen_params.control_strength;
+    root["ip_adapter_strength"]   = gen_params.ip_adapter_strength;
     root["auto_resize_ref_image"] = gen_params.auto_resize_ref_image;
     root["increase_ref_index"]    = gen_params.increase_ref_index;
     if (mode == VID_GEN) {

examples/common/common.h

@@ -133,6 +133,7 @@ struct SDContextParams {
     std::string taesd_path;
     std::string esrgan_path;
     std::string control_net_path;
+    std::string ip_adapter_path;
     std::string motion_module_path;
     std::string embedding_dir;
     std::string photo_maker_path;
@@ -200,6 +201,7 @@ struct SDGenerationParams {
     int64_t seed               = 42;
     float strength             = 0.75f;
     float control_strength     = 0.9f;
+    float ip_adapter_strength  = 1.0f;
     bool auto_resize_ref_image = true;
     bool increase_ref_index    = false;
     bool embed_image_metadata  = true;
@@ -208,6 +210,7 @@ struct SDGenerationParams {
     std::string end_image_path;
     std::string mask_image_path;
     std::string control_image_path;
+    std::string ip_adapter_image_path;
     std::vector<std::string> ref_image_paths;
     std::string control_video_path;
 
@@ -274,6 +277,7 @@ struct SDGenerationParams {
     std::vector<SDImageOwner> ref_images;
     SDImageOwner mask_image;
     SDImageOwner control_image;
+    SDImageOwner ip_adapter_image;
     std::vector<SDImageOwner> pm_id_images;
     std::vector<SDImageOwner> control_frames;
 

examples/server/api.md

@@ -528,6 +528,7 @@ Shared default fields used by both `img_gen` and `vid_gen`:
 | `auto_resize_ref_image` | `boolean` |
 | `increase_ref_index` | `boolean` |
 | `control_strength` | `number` |
+| `ip_adapter_strength` | `number` |
 | `hires` | `object` |
 | `hires.enabled` | `boolean` |
 | `hires.upscaler` | `string` |
@@ -567,6 +568,7 @@ Fields returned in `features_by_mode.img_gen`:
 - `init_image`
 - `mask_image`
 - `control_image`
+- `ip_adapter_image`
 - `ref_images`
 - `lora`
 - `vae_tiling`
@@ -653,12 +655,14 @@ Example:
   "auto_resize_ref_image": true,
   "increase_ref_index": false,
   "control_strength": 0.9,
+  "ip_adapter_strength": 1.0,
   "embed_image_metadata": true,
 
   "init_image": null,
   "ref_images": [],
   "mask_image": null,
   "control_image": null,
+  "ip_adapter_image": null,
 
   "sample_params": {
     "scheduler": "discrete",
@@ -733,6 +737,7 @@ Channel expectations:
 - `init_image`: 3 channels
 - `ref_images[]`: 3 channels
 - `control_image`: 3 channels
+- `ip_adapter_image`: 3 channels
 - `mask_image`: 1 channel
 
 If omitted or null:
@@ -757,6 +762,7 @@ Top-level scalar fields:
 | `auto_resize_ref_image` | `boolean` |
 | `increase_ref_index` | `boolean` |
 | `control_strength` | `number` |
+| `ip_adapter_strength` | `number` |
 | `embed_image_metadata` | `boolean` |
 
 Image fields:
@@ -767,6 +773,7 @@ Image fields:
 | `ref_images` | `array<string>` |
 | `mask_image` | `string \| null` |
 | `control_image` | `string \| null` |
+| `ip_adapter_image` | `string \| null` |
 
 LoRA fields:
 
@@ -958,7 +965,7 @@ Response fields:
 Compared with `img_gen`, the `vid_gen` request body:
 
 - `vid_gen` is a single video sequence job, so `batch_count` is not part of the request schema
-- `ref_images`, `mask_image`, `control_image`, `control_strength`, and `embed_image_metadata` are not part of the request schema
+- `ref_images`, `mask_image`, `control_image`, `control_strength`, `ip_adapter_image`, `ip_adapter_strength`, and `embed_image_metadata` are not part of the request schema
 - `vid_gen` adds `end_image`, `control_frames`, `high_noise_sample_params`, `video_frames`, `fps`, `moe_boundary`, and `vace_strength`
 
 Example:

examples/server/routes_sdcpp.cpp

@@ -130,6 +130,7 @@ static json make_img_gen_defaults_json(const SDGenerationParams& defaults, const
         {"auto_resize_ref_image", defaults.auto_resize_ref_image},
         {"increase_ref_index", defaults.increase_ref_index},
         {"control_strength", defaults.control_strength},
+        {"ip_adapter_strength", defaults.ip_adapter_strength},
         {"sample_params", make_sample_params_json(defaults.sample_params, defaults.skip_layers)},
         {"hires", make_hires_json(defaults)},
         {"vae_tiling_params", make_vae_tiling_json(defaults.vae_tiling_params)},
@@ -173,6 +174,7 @@ static json make_img_gen_features_json() {
         {"init_image", true},
         {"mask_image", true},
         {"control_image", true},
+        {"ip_adapter_image", true},
         {"ref_images", true},
         {"lora", true},
         {"vae_tiling", true},

include/stable-diffusion.h

@@ -180,6 +180,7 @@ enum sd_vae_format_t {
     SD_VAE_FORMAT_FLUX,
     SD_VAE_FORMAT_SD3,
     SD_VAE_FORMAT_FLUX2,
+    SD_VAE_FORMAT_WAN,
     SD_VAE_FORMAT_COUNT,
 };
 
@@ -199,6 +200,7 @@ typedef struct {
     const char* audio_vae_path;
     const char* taesd_path;
     const char* control_net_path;
+    const char* ip_adapter_path;
     const char* motion_module_path;
     const sd_embedding_t* embeddings;
     uint32_t embedding_count;
@@ -374,6 +376,8 @@ typedef struct {
     int batch_count;
     sd_image_t control_image;
     float control_strength;
+    sd_image_t ip_adapter_image;
+    float ip_adapter_strength;
     sd_pm_params_t pm_params;
     sd_pulid_params_t pulid_params;
     sd_tiling_params_t vae_tiling_params;

src/conditioning/conditioner.hpp

@@ -117,6 +117,7 @@ public:
     virtual SDCondition get_learned_condition(int n_threads,
                                               const ConditionerParams& conditioner_params) = 0;
     virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors)           = 0;
+    virtual void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) {}
     virtual void set_max_graph_vram_bytes(size_t max_vram_bytes) {}
     virtual void set_stream_layers_enabled(bool enabled) {}
     virtual void set_runtime_backends(const std::vector<ggml_backend_t>& backends) {}
@@ -1664,6 +1665,10 @@ struct AnimaConditioner : public Conditioner {
         llm->get_param_tensors(tensors, "text_encoders.llm");
     }
 
+    void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) override {
+        llm->get_param_tensor_ops(tensor_ops);
+    }
+
     void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
         llm->set_max_graph_vram_bytes(max_vram_bytes);
     }
@@ -1782,6 +1787,7 @@ struct LLMEmbedder : public Conditioner {
     SDVersion version;
     std::shared_ptr<BPETokenizer> tokenizer;
     std::shared_ptr<LLM::LLMRunner> llm;
+    std::shared_ptr<T5Runner> byt5;
 
     LLMEmbedder(ggml_backend_t backend,
                 const String2TensorStorage& tensor_storage_map      = {},
@@ -1799,7 +1805,12 @@ struct LLMEmbedder : public Conditioner {
             arch = LLM::LLMArch::GPT_OSS_20B;
         } else if (sd_version_is_pid(version)) {
             arch = LLM::LLMArch::GEMMA2_2B;
-        } else if (sd_version_is_lingbot_video(version) || sd_version_is_ideogram4(version) || sd_version_is_boogu_image(version) || sd_version_is_sefi_image(version) || sd_version_is_krea2(version)) {
+        } else if (sd_version_is_lingbot_video(version) ||
+                   sd_version_is_ideogram4(version) ||
+                   sd_version_is_boogu_image(version) ||
+                   sd_version_is_sefi_image(version) ||
+                   sd_version_is_krea2(version) ||
+                   sd_version_is_mage_flow(version)) {
             arch = LLM::LLMArch::QWEN3_VL;
         } else if (sd_version_is_z_image(version) || version == VERSION_OVIS_IMAGE || version == VERSION_FLUX2_KLEIN) {
             arch = LLM::LLMArch::QWEN3;
@@ -1819,54 +1830,101 @@ struct LLMEmbedder : public Conditioner {
                                                "text_encoders.llm",
                                                enable_vision,
                                                weight_manager);
+        if (sd_version_is_hunyuan_video(version)) {
+            const std::string byt5_prefix = "text_encoders.t5xxl.transformer";
+            for (const auto& [name, _] : tensor_storage_map) {
+                if (starts_with(name, byt5_prefix + ".")) {
+                    byt5 = std::make_shared<T5Runner>(backend,
+                                                      tensor_storage_map,
+                                                      byt5_prefix,
+                                                      false,
+                                                      weight_manager);
+                    break;
+                }
+            }
+        }
     }
 
     void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
         llm->get_param_tensors(tensors, "text_encoders.llm");
+        if (byt5) {
+            byt5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
+        }
+    }
+
+    void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) override {
+        llm->get_param_tensor_ops(tensor_ops);
     }
 
     void set_max_graph_vram_bytes(size_t max_vram_bytes) override {
         llm->set_max_graph_vram_bytes(max_vram_bytes);
+        if (byt5) {
+            byt5->set_max_graph_vram_bytes(max_vram_bytes);
+        }
     }
 
     void set_stream_layers_enabled(bool enabled) override {
         llm->set_stream_layers_enabled(enabled);
+        if (byt5) {
+            byt5->set_stream_layers_enabled(enabled);
+        }
     }
 
     void set_runtime_backends(const std::vector<ggml_backend_t>& backends) override {
         llm->set_runtime_backends(backends);
+        if (byt5) {
+            byt5->set_runtime_backends(backends);
+        }
     }
 
     void set_graph_cut_layer_split_enabled(bool enabled) override {
         if (llm) {
             llm->set_graph_cut_layer_split_enabled(enabled);
         }
+        if (byt5) {
+            byt5->set_graph_cut_layer_split_enabled(enabled);
+        }
     }
 
     void set_graph_cut_layer_split_backend_vram_limits(const std::vector<size_t>& limits) override {
         if (llm) {
             llm->set_graph_cut_layer_split_backend_vram_limits(limits);
         }
+        if (byt5) {
+            byt5->set_graph_cut_layer_split_backend_vram_limits(limits);
+        }
     }
 
     void get_layer_split_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
         llm->get_param_tensors(tensors, "text_encoders.llm");
+        if (byt5) {
+            byt5->get_param_tensors(tensors, "text_encoders.t5xxl.transformer");
+        }
     }
 
     void set_flash_attention_enabled(bool enabled) override {
         llm->set_flash_attention_enabled(enabled);
+        if (byt5) {
+            byt5->set_flash_attention_enabled(enabled);
+        }
     }
 
     void set_weight_adapter(const std::shared_ptr<WeightAdapter>& adapter) override {
         if (llm) {
             llm->set_weight_adapter(adapter);
         }
+        if (byt5) {
+            byt5->set_weight_adapter(adapter);
+        }
     }
 
     void runner_done() override {
         if (llm) {
             llm->runner_done();
         }
+        if (byt5) {
+            byt5->runner_done();
+        }
     }
 
     std::tuple<std::vector<int>, std::vector<float>, std::vector<float>> tokenize(std::string text,
@@ -2060,7 +2118,24 @@ struct LLMEmbedder : public Conditioner {
         int64_t t0                     = ggml_time_ms();
         RefImageResizeMode resize_mode = conditioner_params.ref_image_params.vlm_resize_mode;
 
-        if (sd_version_is_lingbot_video(version)) {
+        if (sd_version_is_hunyuan_video(version)) {
+            prompt_template_encode_start_idx = 98;
+            out_layers                       = {26};
+
+            prompt =
+                "<|im_start|>system\nYou are a helpful assistant. Describe the video by detailing the following aspects:\n"
+                "1. The main content and theme of the video.\n"
+                "2. The color, shape, size, texture, quantity, text, and spatial relationships of the objects.\n"
+                "3. Actions, events, behaviors temporal relationships, physical movement changes of the objects.\n"
+                "4. background environment, light, style and atmosphere.\n"
+                "5. camera angles, movements, and transitions used in the video.<|im_end|>\n"
+                "<|im_start|>user\n";
+
+            prompt_attn_range.first = static_cast<int>(prompt.size());
+            prompt += conditioner_params.text;
+            prompt_attn_range.second = static_cast<int>(prompt.size());
+            prompt += "<|im_end|>\n<|im_start|>assistant\n";
+        } else if (sd_version_is_lingbot_video(version)) {
             const int pad_token = 151643;
             const std::string prompt_prefix =
                 "<|im_start|>system\nGiven a user input that may include a text prompt alone, "
@@ -2144,22 +2219,22 @@ struct LLMEmbedder : public Conditioner {
             prompt += conditioner_params.text;
             prompt_attn_range = {0, 0};
             prompt += "<|im_end|>\n<|im_start|>assistant\n";
-        } else if (sd_version_is_qwen_image(version)) {
+        } else if (sd_version_is_qwen_image(version) || sd_version_is_mage_flow(version)) {
             if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) {
-                LOG_INFO("QwenImageEditPlusPipeline");
+                LOG_INFO("%s", sd_version_is_mage_flow(version) ? "MageFlowEditPipeline" : "QwenImageEditPlusPipeline");
                 prompt_template_encode_start_idx = 64;
                 int image_embed_idx              = 64 + 6;
 
                 int min_pixels = conditioner_params.ref_image_params.vlm_min_size;
                 if (min_pixels <= 0) {
-                    min_pixels = 384;
-                    if (resize_mode == RefImageResizeMode::AREA) {
+                    min_pixels = sd_version_is_mage_flow(version) ? -1 : 384;
+                    if (min_pixels > 0 && resize_mode == RefImageResizeMode::AREA) {
                         min_pixels *= min_pixels;
                     }
                 }
                 int max_pixels = conditioner_params.ref_image_params.vlm_max_size;
                 if (max_pixels <= 0) {
-                    max_pixels = 560;
+                    max_pixels = sd_version_is_mage_flow(version) ? 384 : 560;
                     if (resize_mode == RefImageResizeMode::AREA) {
                         max_pixels *= max_pixels;
                     }
@@ -2187,7 +2262,7 @@ struct LLMEmbedder : public Conditioner {
                     image_embeds.emplace_back(image_embed_idx, image_embed);
                     image_embed_idx += 1 + static_cast<int>(image_embed.shape()[1]) + 6;
 
-                    img_prompt += "Picture " + std::to_string(i + 1) + ": <|vision_start|>";  // [24669, 220, index, 25, 220, 151652]
+                    img_prompt += (sd_version_is_mage_flow(version) ? "Image " : "Picture ") + std::to_string(i + 1) + ": <|vision_start|>";
                     int64_t num_image_tokens = image_embed.shape()[1];
                     img_prompt.reserve(num_image_tokens * placeholder.size());
                     for (int j = 0; j < num_image_tokens; j++) {
@@ -2215,6 +2290,9 @@ struct LLMEmbedder : public Conditioner {
 
                 prompt += "<|im_end|>\n<|im_start|>assistant\n";
             }
+            if (sd_version_is_mage_flow(version)) {
+                max_length = 2048 + prompt_template_encode_start_idx;
+            }
         } else if (sd_version_is_boogu_image(version)) {
             prompt_template_encode_start_idx = 0;
 
@@ -2590,6 +2668,46 @@ struct LLMEmbedder : public Conditioner {
                                            spell_quotes,
                                            max_length);
         std::vector<sd::Tensor<float>> extra_hidden_states_vec;
+        if (sd_version_is_hunyuan_video(version) && byt5) {
+            std::vector<std::string> quoted_texts;
+            auto collect_quoted = [&](const std::string& open, const std::string& close) {
+                size_t begin = 0;
+                while ((begin = conditioner_params.text.find(open, begin)) != std::string::npos) {
+                    size_t content_begin = begin + open.size();
+                    size_t end           = conditioner_params.text.find(close, content_begin);
+                    if (end == std::string::npos) {
+                        break;
+                    }
+                    quoted_texts.push_back(conditioner_params.text.substr(content_begin, end - content_begin));
+                    begin = end + close.size();
+                }
+            };
+            collect_quoted("\"", "\"");
+            collect_quoted("\xE2\x80\x98", "\xE2\x80\x99");
+            collect_quoted("\xE2\x80\x9C", "\xE2\x80\x9D");
+
+            if (!quoted_texts.empty()) {
+                std::string byt5_text;
+                for (const auto& text : quoted_texts) {
+                    byt5_text += "Text \"" + text + "\". ";
+                }
+                std::vector<int> tokens;
+                tokens.reserve(byt5_text.size() + 1);
+                for (unsigned char byte : byt5_text) {
+                    tokens.push_back(static_cast<int>(byte) + 3);
+                }
+                tokens.push_back(1);
+                sd::Tensor<int32_t> input_ids({static_cast<int64_t>(tokens.size())}, tokens);
+                auto byt5_hidden_states = byt5->compute(n_threads,
+                                                        input_ids,
+                                                        sd::Tensor<float>(),
+                                                        false,
+                                                        true,
+                                                        true);
+                GGML_ASSERT(!byt5_hidden_states.empty());
+                extra_hidden_states_vec.push_back(std::move(byt5_hidden_states));
+            }
+        }
         for (int i = 0; i < extra_prompts.size(); i++) {
             auto extra_hidden_states = encode_prompt(n_threads,
                                                      extra_prompts[i],
@@ -2719,6 +2837,10 @@ struct LTXAVEmbedder : public Conditioner {
         projector->get_param_tensors(tensors, "text_embedding_projection");
     }
 
+    void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) override {
+        llm->get_param_tensor_ops(tensor_ops);
+    }
+
     void set_flash_attention_enabled(bool enabled) override {
         llm->set_flash_attention_enabled(enabled);
         projector->set_flash_attention_enabled(enabled);

src/core/ggml_extend.hpp

@@ -1689,6 +1689,8 @@ struct GGMLRunnerContext {
     bool conv2d_direct_enabled                                       = false;
     bool circular_x_enabled                                          = false;
     bool circular_y_enabled                                          = false;
+    ggml_tensor* ip_context                                          = nullptr;
+    float ip_scale                                                   = 1.0f;
     std::shared_ptr<WeightAdapter> weight_adapter                    = nullptr;
     std::vector<std::pair<ggml_tensor*, std::string>>* debug_tensors = nullptr;
     std::function<ggml_tensor*(const std::string&)> get_cache_tensor;
@@ -1751,7 +1753,7 @@ protected:
     std::vector<size_t> graph_cut_layer_split_backend_vram_limits_;
 
     std::vector<ggml_backend_t> extra_runtime_backends;  // borrowed (SDBackendManager-owned)
-    ggml_backend_sched_t sched             = nullptr;    // owned, multi-device only
+    ggml_backend_sched_t sched             = nullptr;    // owned
     ggml_backend_t cpu_fallback_backend    = nullptr;    // owned, sched requires a trailing CPU backend
     bool multi_device_eval_callback_warned = false;
 
@@ -2145,8 +2147,22 @@ protected:
         return !extra_runtime_backends.empty();
     }
 
+    bool graph_requires_backend_fallback(ggml_cgraph* gf) const {
+        if (gf == nullptr || sd_backend_is_cpu(runtime_backend)) {
+            return false;
+        }
+        const int n_nodes = ggml_graph_n_nodes(gf);
+        for (int i = 0; i < n_nodes; ++i) {
+            ggml_tensor* node = ggml_graph_node(gf, i);
+            if (node != nullptr && !ggml_backend_supports_op(runtime_backend, node)) {
+                return true;
+            }
+        }
+        return false;
+    }
+
     bool alloc_compute_buffer(ggml_cgraph* gf) {
-        if (is_multi_device()) {
+        if (sched != nullptr || is_multi_device() || graph_requires_backend_fallback(gf)) {
             // The sched replaces the gallocr. Do NOT ggml_backend_sched_reserve
             // the graph here: reserve runs split_graph, which rewires the
             // graph's src pointers to sched-internal copy tensors, and the
@@ -2154,6 +2170,10 @@ protected:
             // rewired graph, silently corrupting every cross-backend input. A
             // graph must be split at most once; the alloc in execute_graph
             // performs the real allocation.
+            if (compute_allocr != nullptr) {
+                ggml_gallocr_free(compute_allocr);
+                compute_allocr = nullptr;
+            }
             return ensure_sched(gf);
         }
         if (compute_allocr != nullptr) {
@@ -2751,7 +2771,7 @@ protected:
         };
         ComputeBufferGuard compute_buffer_guard(this, free_compute_buffer);
 
-        if (is_multi_device()) {
+        if (sched != nullptr) {
             ggml_backend_sched_reset(sched);
             pin_multi_device_nodes(gf);  // reset clears the pins; re-apply before alloc
             if (!ggml_backend_sched_alloc_graph(sched, gf)) {
@@ -2772,9 +2792,9 @@ protected:
         }
 
         ggml_status status;
-        if (is_multi_device()) {
+        if (sched != nullptr) {
             if (sd_get_backend_eval_callback() != nullptr && !multi_device_eval_callback_warned) {
-                LOG_WARN("%s: eval callback is not supported with multiple runtime backends; ignoring",
+                LOG_WARN("%s: eval callback is not supported with the backend scheduler; ignoring",
                          get_desc().c_str());
                 multi_device_eval_callback_warned = true;
             }
@@ -3016,12 +3036,9 @@ public:
 
     // do copy after alloc graph
     void set_backend_tensor_data(ggml_tensor* tensor, const void* data) {
-        if (is_multi_device()) {
-            // The sched only assigns a backend (and thus a buffer) to tensors
-            // that participate in the graph; flag standalone data tensors as
-            // inputs so they get one.
-            ggml_set_input(tensor);
-        }
+        // The scheduler only allocates standalone data tensors when they are
+        // marked as graph inputs. The flag is harmless for single-backend graphs.
+        ggml_set_input(tensor);
         backend_tensor_data_map[tensor] = data;
     }
 
@@ -3238,6 +3255,11 @@ protected:
 
     virtual void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") {}
 
+    virtual enum ggml_op param_usage_op(const std::string& name) const {
+        (void)name;
+        return GGML_OP_NONE;
+    }
+
 public:
     void init(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") {
         if (prefix.size() > 0) {
@@ -3288,6 +3310,18 @@ public:
         }
     }
 
+    void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) {
+        for (auto& pair : blocks) {
+            pair.second->get_param_tensor_ops(tensor_ops);
+        }
+        for (auto& pair : params) {
+            enum ggml_op op = param_usage_op(pair.first);
+            if (op != GGML_OP_NONE) {
+                tensor_ops[pair.second] = op;
+            }
+        }
+    }
+
     virtual std::string get_desc() {
         return "GGMLBlock";
     }
@@ -3309,7 +3343,7 @@ public:
 
 class Identity : public UnaryBlock {
 public:
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
         return x;
     }
 };
@@ -3368,7 +3402,7 @@ public:
         force_prec_f32 = force_prec_f32_;
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
         ggml_tensor* w = params["weight"];
         ggml_tensor* b = nullptr;
         if (bias) {
@@ -3415,6 +3449,10 @@ protected:
         params["weight"] = ggml_new_tensor_2d(ctx, wtype, embedding_dim, num_embeddings);
     }
 
+    enum ggml_op param_usage_op(const std::string& name) const override {
+        return name == "weight" ? GGML_OP_GET_ROWS : GGML_OP_NONE;
+    }
+
 public:
     Embedding(int64_t num_embeddings, int64_t embedding_dim)
         : embedding_dim(embedding_dim),
@@ -3422,7 +3460,7 @@ public:
     }
 
     ggml_tensor* forward(GGMLRunnerContext* ctx,
-                         ggml_tensor* input_ids) {
+                         ggml_tensor* input_ids) override {
         // input_ids: [N, n_token]
         auto weight = params["weight"];
 
@@ -3482,11 +3520,11 @@ public:
         scale = scale_value;
     }
 
-    std::string get_desc() {
+    std::string get_desc() override {
         return "Conv2d";
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
         ggml_tensor* w = params["weight"];
         ggml_tensor* b = nullptr;
         if (bias) {
@@ -3569,11 +3607,11 @@ public:
         scale = scale_value;
     }
 
-    std::string get_desc() {
+    std::string get_desc() override {
         return "Conv2d_grouped";
     }
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
         ggml_tensor* w = params["weight"];
         ggml_tensor* b = nullptr;
         if (bias) {
@@ -3609,18 +3647,19 @@ public:
         if (groups == in_channels && groups == out_channels) {
             ggml_tensor* res;
             if (ctx->conv2d_direct_enabled) {
-                res = ggml_conv_2d_dw_direct(ctx->ggml_ctx, x, w,
+                res = ggml_conv_2d_dw_direct(ctx->ggml_ctx, w, x,
                                              stride.second, stride.first,
                                              padding.second, padding.first,
                                              dilation.second, dilation.first);
             } else {
-                res = ggml_conv_2d_dw(ctx->ggml_ctx, x, w,
+                res = ggml_conv_2d_dw(ctx->ggml_ctx, w, x,
                                       stride.second, stride.first,
                                       padding.second, padding.first,
                                       dilation.second, dilation.first);
             }
             if (b) {
-                res = ggml_add(ctx->ggml_ctx, res, b);
+                b   = ggml_reshape_4d(ctx->ggml_ctx, b, 1, 1, b->ne[0], 1);
+                res = ggml_add_inplace(ctx->ggml_ctx, res, b);
             }
             return res;
         }
@@ -3725,7 +3764,7 @@ public:
           bias(bias),
           force_prec_f32(force_prec_f32) {}
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
         ggml_tensor* w = params["weight"];
         ggml_tensor* b = nullptr;
         if (ctx->weight_adapter) {
@@ -3778,7 +3817,7 @@ public:
           elementwise_affine(elementwise_affine),
           bias(bias) {}
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
         ggml_tensor* w = nullptr;
         ggml_tensor* b = nullptr;
 
@@ -3865,7 +3904,7 @@ public:
         : hidden_size(hidden_size),
           eps(eps) {}
 
-    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
         ggml_tensor* w = params["weight"];
         if (ctx->weight_adapter) {
             w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");

src/model.h

@@ -38,6 +38,7 @@ enum SDVersion {
     VERSION_LINGBOT_VIDEO,
     VERSION_QWEN_IMAGE,
     VERSION_QWEN_IMAGE_LAYERED,
+    VERSION_HUNYUAN_VIDEO,
     VERSION_ANIMA,
     VERSION_FLUX2,
     VERSION_FLUX2_KLEIN,
@@ -54,6 +55,7 @@ enum SDVersion {
     VERSION_IDEOGRAM4,
     VERSION_SEFI_IMAGE,
     VERSION_KREA2,
+    VERSION_MAGE_FLOW,
     VERSION_ESRGAN,
     VERSION_COUNT,
 };
@@ -142,6 +144,13 @@ static inline bool sd_version_is_qwen_image(SDVersion version) {
     return false;
 }
 
+static inline bool sd_version_is_hunyuan_video(SDVersion version) {
+    if (version == VERSION_HUNYUAN_VIDEO) {
+        return true;
+    }
+    return false;
+}
+
 static inline bool sd_version_is_anima(SDVersion version) {
     if (version == VERSION_ANIMA) {
         return true;
@@ -219,6 +228,10 @@ static inline bool sd_version_is_krea2(SDVersion version) {
     return false;
 }
 
+static inline bool sd_version_is_mage_flow(SDVersion version) {
+    return version == VERSION_MAGE_FLOW;
+}
+
 static inline bool sd_version_uses_flux_vae(SDVersion version) {
     if (sd_version_is_flux(version) || sd_version_is_z_image(version) || sd_version_is_boogu_image(version) || sd_version_is_longcat(version)) {
         return true;
@@ -240,6 +253,10 @@ static inline bool sd_version_uses_wan_vae(SDVersion version) {
     return false;
 }
 
+static inline bool sd_version_uses_hunyuan_video_vae(SDVersion version) {
+    return sd_version_is_hunyuan_video(version);
+}
+
 static inline bool sd_version_is_inpaint(SDVersion version) {
     if (version == VERSION_SD1_INPAINT ||
         version == VERSION_SD2_INPAINT ||
@@ -259,6 +276,7 @@ static inline bool sd_version_is_dit(SDVersion version) {
         sd_version_is_wan(version) ||
         sd_version_is_lingbot_video(version) ||
         sd_version_is_qwen_image(version) ||
+        sd_version_is_hunyuan_video(version) ||
         version == VERSION_HIDREAM_O1 ||
         sd_version_is_anima(version) ||
         sd_version_is_z_image(version) ||
@@ -270,7 +288,8 @@ static inline bool sd_version_is_dit(SDVersion version) {
         sd_version_is_pid(version) ||
         sd_version_is_ideogram4(version) ||
         sd_version_is_sefi_image(version) ||
-        sd_version_is_krea2(version)) {
+        sd_version_is_krea2(version) ||
+        sd_version_is_mage_flow(version)) {
         return true;
     }
     return false;

src/model/adapter/ip_adapter.hpp

@@ -0,0 +1,88 @@
+#ifndef __SD_MODEL_ADAPTER_IP_ADAPTER_HPP__
+#define __SD_MODEL_ADAPTER_IP_ADAPTER_HPP__
+
+#include "core/ggml_extend.hpp"
+#include "model/common/block.hpp"
+#include "model_loader.h"
+
+namespace IPAdapter {
+
+    struct ImageProjModel : public GGMLBlock {
+        int64_t num_tokens = 4;
+        int64_t ctx_dim    = 768;
+        int64_t clip_dim   = 1024;
+
+        ImageProjModel() {}
+        ImageProjModel(int64_t num_tokens, int64_t ctx_dim, int64_t clip_dim)
+            : num_tokens(num_tokens), ctx_dim(ctx_dim), clip_dim(clip_dim) {
+            blocks["proj"] = std::shared_ptr<GGMLBlock>(new Linear(clip_dim, num_tokens * ctx_dim, true));
+            blocks["norm"] = std::shared_ptr<GGMLBlock>(new LayerNorm(ctx_dim));
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* image_embeds) {
+            auto proj = std::dynamic_pointer_cast<Linear>(blocks["proj"]);
+            auto norm = std::dynamic_pointer_cast<LayerNorm>(blocks["norm"]);
+
+            int64_t n = image_embeds->ne[1];
+            auto x    = proj->forward(ctx, image_embeds);
+            x         = ggml_reshape_3d(ctx->ggml_ctx, x, ctx_dim, num_tokens, n);
+            x         = norm->forward(ctx, x);
+            return x;
+        }
+    };
+
+    struct IPAdapterRunner : public GGMLRunner {
+        ImageProjModel image_proj;
+        int64_t num_tokens = 4;
+        std::string prefix;
+
+        IPAdapterRunner(ggml_backend_t backend,
+                        const String2TensorStorage& tensor_storage_map,
+                        const std::string prefix,
+                        std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : GGMLRunner(backend, weight_manager), prefix(prefix) {
+            int64_t ctx_dim  = 768;
+            int64_t clip_dim = 1024;
+            int64_t out_dim  = 3072;
+            auto norm_iter   = tensor_storage_map.find(prefix + ".image_proj.norm.weight");
+            if (norm_iter != tensor_storage_map.end()) {
+                ctx_dim = norm_iter->second.ne[0];
+            }
+            auto proj_iter = tensor_storage_map.find(prefix + ".image_proj.proj.weight");
+            if (proj_iter != tensor_storage_map.end()) {
+                clip_dim = proj_iter->second.ne[0];
+                out_dim  = proj_iter->second.ne[1];
+            }
+            num_tokens = out_dim / ctx_dim;
+            image_proj = ImageProjModel(num_tokens, ctx_dim, clip_dim);
+            image_proj.init(params_ctx, tensor_storage_map, prefix + ".image_proj");
+        }
+
+        std::string get_desc() override {
+            return "ip_adapter";
+        }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string = "") {
+            image_proj.get_param_tensors(tensors, prefix + ".image_proj");
+        }
+
+        ggml_cgraph* build_graph(const sd::Tensor<float>& image_embeds_tensor) {
+            ggml_cgraph* gf     = new_graph_custom(1024);
+            ggml_tensor* embeds = make_input(image_embeds_tensor);
+            auto runner_ctx     = get_context();
+            ggml_tensor* out    = image_proj.forward(&runner_ctx, embeds);
+            ggml_build_forward_expand(gf, out);
+            return gf;
+        }
+
+        sd::Tensor<float> compute(int n_threads, const sd::Tensor<float>& image_embeds) {
+            auto get_graph = [&]() -> ggml_cgraph* {
+                return build_graph(image_embeds);
+            };
+            return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true, true, true));
+        }
+    };
+
+}  // namespace IPAdapter
+
+#endif  // __SD_MODEL_ADAPTER_IP_ADAPTER_HPP__

src/model/adapter/lora.hpp

@@ -14,6 +14,8 @@ struct LoraModel : public GGMLRunner {
     std::unordered_map<std::string, ggml_tensor*> lora_tensors;
     std::map<ggml_tensor*, ggml_tensor*> original_tensor_to_final_tensor;
     std::set<std::string> applied_lora_tensors;
+    std::set<std::string> skipped_incompatible_lora_tensors;
+    std::set<std::string> warned_incompatible_model_tensors;
     std::string file_path;
     std::shared_ptr<ModelManager> model_manager;
     ggml_backend_t params_backend = nullptr;
@@ -133,6 +135,8 @@ struct LoraModel : public GGMLRunner {
         lora_tensors.clear();
         original_tensor_to_final_tensor.clear();
         applied_lora_tensors.clear();
+        skipped_incompatible_lora_tensors.clear();
+        warned_incompatible_model_tensors.clear();
         applied             = false;
         tensor_preprocessed = false;
     }
@@ -338,7 +342,9 @@ struct LoraModel : public GGMLRunner {
             iter = lora_tensors.find(hada_1_mid_name);
             if (iter != lora_tensors.end()) {
                 hada_1_mid = ggml_ext_cast_f32(ctx, backend, iter->second);
-                hada_1_up  = ggml_cont(ctx, ggml_transpose(ctx, hada_1_up));
+                if (hada_1_up != nullptr) {
+                    hada_1_up = ggml_cont(ctx, ggml_transpose(ctx, hada_1_up));
+                }
             }
 
             iter = lora_tensors.find(hada_2_down_name);
@@ -354,7 +360,9 @@ struct LoraModel : public GGMLRunner {
             iter = lora_tensors.find(hada_2_mid_name);
             if (iter != lora_tensors.end()) {
                 hada_2_mid = ggml_ext_cast_f32(ctx, backend, iter->second);
-                hada_2_up  = ggml_cont(ctx, ggml_transpose(ctx, hada_2_up));
+                if (hada_2_up != nullptr) {
+                    hada_2_up = ggml_cont(ctx, ggml_transpose(ctx, hada_2_up));
+                }
             }
 
             if (hada_1_up == nullptr || hada_1_down == nullptr || hada_2_up == nullptr || hada_2_down == nullptr) {
@@ -546,7 +554,27 @@ struct LoraModel : public GGMLRunner {
                 }
             }
 
-            GGML_ASSERT(ggml_nelements(diff) == ggml_nelements(model_tensor));
+            if (ggml_nelements(diff) != ggml_nelements(model_tensor)) {
+                const std::string lora_tensor_prefix = "lora." + model_tensor_name + ".";
+                for (const auto& tensor_name : applied_lora_tensors) {
+                    if (starts_with(tensor_name, lora_tensor_prefix)) {
+                        skipped_incompatible_lora_tensors.insert(tensor_name);
+                    }
+                }
+                if (warned_incompatible_model_tensors.insert(model_tensor_name).second) {
+                    LOG_WARN("skip incompatible LoRA tensor |%s|: model shape = [%lld, %lld, %lld, %lld], LoRA shape = [%lld, %lld, %lld, %lld]",
+                             model_tensor_name.c_str(),
+                             static_cast<long long>(model_tensor->ne[0]),
+                             static_cast<long long>(model_tensor->ne[1]),
+                             static_cast<long long>(model_tensor->ne[2]),
+                             static_cast<long long>(model_tensor->ne[3]),
+                             static_cast<long long>(diff->ne[0]),
+                             static_cast<long long>(diff->ne[1]),
+                             static_cast<long long>(diff->ne[2]),
+                             static_cast<long long>(diff->ne[3]));
+                }
+                return nullptr;
+            }
             diff = ggml_reshape(ctx, diff, model_tensor);
         }
         return diff;
@@ -555,6 +583,7 @@ struct LoraModel : public GGMLRunner {
     ggml_tensor* get_out_diff(ggml_context* ctx,
                               ggml_backend_t backend,
                               ggml_tensor* x,
+                              ggml_tensor* model_weight,
                               WeightAdapter::ForwardParams forward_params,
                               const std::string& model_tensor_name) {
         ggml_tensor* out_diff = nullptr;
@@ -707,6 +736,43 @@ struct LoraModel : public GGMLRunner {
                 break;
             }
 
+            if (!is_conv2d) {
+                const int64_t down_in  = lora_down->ne[0];
+                const int64_t down_out = lora_down->ne[1];
+                const int64_t up_in    = lora_up->ne[0];
+                const int64_t up_out   = lora_up->ne[1];
+
+                bool compatible = down_in == model_weight->ne[0] &&
+                                  up_out == model_weight->ne[1];
+                if (lora_mid != nullptr) {
+                    compatible = compatible &&
+                                 lora_mid->ne[0] == down_out &&
+                                 up_in == lora_mid->ne[1];
+                } else {
+                    compatible = compatible && up_in == down_out;
+                }
+
+                if (!compatible) {
+                    skipped_incompatible_lora_tensors.insert(lora_down_name);
+                    skipped_incompatible_lora_tensors.insert(lora_up_name);
+                    skipped_incompatible_lora_tensors.insert(lora_mid_name);
+                    skipped_incompatible_lora_tensors.insert(scale_name);
+                    skipped_incompatible_lora_tensors.insert(alpha_name);
+                    if (warned_incompatible_model_tensors.insert(model_tensor_name).second) {
+                        LOG_WARN("skip incompatible LoRA tensor |%s|: model shape = [%lld, %lld], down shape = [%lld, %lld], up shape = [%lld, %lld]",
+                                 model_tensor_name.c_str(),
+                                 static_cast<long long>(model_weight->ne[0]),
+                                 static_cast<long long>(model_weight->ne[1]),
+                                 static_cast<long long>(down_in),
+                                 static_cast<long long>(down_out),
+                                 static_cast<long long>(up_in),
+                                 static_cast<long long>(up_out));
+                    }
+                    index++;
+                    continue;
+                }
+            }
+
             applied_lora_tensors.insert(lora_up_name);
             applied_lora_tensors.insert(lora_down_name);
 
@@ -869,10 +935,13 @@ struct LoraModel : public GGMLRunner {
     void stat(bool at_runntime = false) {
         size_t total_lora_tensors_count   = 0;
         size_t applied_lora_tensors_count = 0;
+        size_t skipped_lora_tensors_count = 0;
 
         for (auto& kv : lora_tensors) {
             total_lora_tensors_count++;
-            if (applied_lora_tensors.find(kv.first) == applied_lora_tensors.end()) {
+            if (skipped_incompatible_lora_tensors.find(kv.first) != skipped_incompatible_lora_tensors.end()) {
+                skipped_lora_tensors_count++;
+            } else if (applied_lora_tensors.find(kv.first) == applied_lora_tensors.end()) {
                 if (!at_runntime) {
                     LOG_WARN("unused lora tensor |%s|", kv.first.c_str());
                     print_ggml_tensor(kv.second, true);
@@ -884,12 +953,17 @@ struct LoraModel : public GGMLRunner {
         /* Don't worry if this message shows up twice in the logs per LoRA,
          * this function is called once to calculate the required buffer size
          * and then again to actually generate a graph to be used */
-        if (!at_runntime && applied_lora_tensors_count != total_lora_tensors_count) {
+        size_t compatible_lora_tensors_count = total_lora_tensors_count - skipped_lora_tensors_count;
+        if (!at_runntime && applied_lora_tensors_count != compatible_lora_tensors_count) {
             LOG_WARN("Only (%lu / %lu) LoRA tensors have been applied, lora_file_path = %s",
-                     applied_lora_tensors_count, total_lora_tensors_count, file_path.c_str());
+                     applied_lora_tensors_count, compatible_lora_tensors_count, file_path.c_str());
         } else {
             LOG_INFO("(%lu / %lu) LoRA tensors have been applied, lora_file_path = %s",
-                     applied_lora_tensors_count, total_lora_tensors_count, file_path.c_str());
+                     applied_lora_tensors_count, compatible_lora_tensors_count, file_path.c_str());
+        }
+        if (skipped_lora_tensors_count > 0) {
+            LOG_WARN("(%lu / %lu) incompatible LoRA tensors have been skipped, lora_file_path = %s",
+                     skipped_lora_tensors_count, total_lora_tensors_count, file_path.c_str());
         }
     }
 };
@@ -953,7 +1027,7 @@ public:
                                    forward_params.conv2d.scale);
         }
         for (auto& lora_model : lora_models) {
-            ggml_tensor* out_diff = lora_model->get_out_diff(ctx, backend, x, forward_params, prefix + "weight");
+            ggml_tensor* out_diff = lora_model->get_out_diff(ctx, backend, x, w, forward_params, prefix + "weight");
             if (out_diff == nullptr) {
                 continue;
             }

src/model/common/block.hpp

@@ -310,17 +310,33 @@ protected:
     int64_t context_dim;
     int64_t n_head;
     int64_t d_head;
-    bool xtra_dim = false;
+    bool xtra_dim  = false;
+    bool enable_ip = false;
+    bool has_ip    = false;
+
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
+        GGMLBlock::init_params(ctx, tensor_storage_map, prefix);
+        if (enable_ip &&
+            tensor_storage_map.find(prefix + "to_k_ip.weight") != tensor_storage_map.end()) {
+            has_ip            = true;
+            int64_t inner_dim = d_head * n_head;
+            int64_t ip_dim    = tensor_storage_map.at(prefix + "to_k_ip.weight").ne[0];
+            blocks["to_k_ip"] = std::shared_ptr<GGMLBlock>(new Linear(ip_dim, inner_dim, false));
+            blocks["to_v_ip"] = std::shared_ptr<GGMLBlock>(new Linear(ip_dim, inner_dim, false));
+        }
+    }
 
 public:
     CrossAttention(int64_t query_dim,
                    int64_t context_dim,
                    int64_t n_head,
-                   int64_t d_head)
+                   int64_t d_head,
+                   bool enable_ip = false)
         : n_head(n_head),
           d_head(d_head),
           query_dim(query_dim),
-          context_dim(context_dim) {
+          context_dim(context_dim),
+          enable_ip(enable_ip) {
         int64_t inner_dim = d_head * n_head;
         if (context_dim == 320 && d_head == 320) {
             // LOG_DEBUG("CrossAttention: temp set dim to 1024 for sdxs_09");
@@ -363,6 +379,15 @@ public:
         }
         x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, nullptr, false, ctx->flash_attn_enabled);  // [N, n_token, inner_dim]
 
+        if (has_ip && ctx->ip_context != nullptr && ctx->ip_scale != 0.0f) {
+            auto to_k_ip = std::dynamic_pointer_cast<Linear>(blocks["to_k_ip"]);
+            auto to_v_ip = std::dynamic_pointer_cast<Linear>(blocks["to_v_ip"]);
+            auto k_ip    = to_k_ip->forward(ctx, ctx->ip_context);
+            auto v_ip    = to_v_ip->forward(ctx, ctx->ip_context);
+            auto x_ip    = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k_ip, v_ip, n_head, nullptr, false, ctx->flash_attn_enabled);
+            x            = ggml_add(ctx->ggml_ctx, x, ggml_scale(ctx->ggml_ctx, x_ip, ctx->ip_scale));
+        }
+
         x = to_out_0->forward(ctx, x);  // [N, n_token, query_dim]
         return x;
     }
@@ -387,7 +412,7 @@ public:
         // inner_dim is always None or equal to dim
         // gated_ff is always True
         blocks["attn1"] = std::shared_ptr<GGMLBlock>(new CrossAttention(dim, dim, n_head, d_head));
-        blocks["attn2"] = std::shared_ptr<GGMLBlock>(new CrossAttention(dim, context_dim, n_head, d_head));
+        blocks["attn2"] = std::shared_ptr<GGMLBlock>(new CrossAttention(dim, context_dim, n_head, d_head, true));
         blocks["ff"]    = std::shared_ptr<GGMLBlock>(new FeedForward(dim, dim));
         blocks["norm1"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
         blocks["norm2"] = std::shared_ptr<GGMLBlock>(new LayerNorm(dim));
@@ -450,7 +475,7 @@ protected:
     int64_t context_dim = 768;  // hidden_size, 1024 for VERSION_SD2
     bool use_linear     = false;
 
-    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") {
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override {
         auto iter = tensor_storage_map.find(prefix + "proj_out.weight");
         if (iter != tensor_storage_map.end()) {
             int64_t inner_dim = n_head * d_head;

src/model/common/rope.hpp

@@ -535,6 +535,33 @@ namespace Rope {
         return vid_ids_repeated;
     }
 
+    __STATIC_INLINE__ std::vector<std::vector<float>> gen_hunyuan_video_ids(int t,
+                                                                            int h,
+                                                                            int w,
+                                                                            int patch_t,
+                                                                            int patch_h,
+                                                                            int patch_w,
+                                                                            int bs,
+                                                                            int context_len) {
+        std::vector<std::vector<float>> txt_ids(bs * context_len, std::vector<float>(3, 0.0f));
+        auto img_ids = gen_vid_ids(t, h, w, patch_t, patch_h, patch_w, bs);
+        return concat_ids(txt_ids, img_ids, bs);
+    }
+
+    __STATIC_INLINE__ std::vector<float> gen_hunyuan_video_pe(int t,
+                                                              int h,
+                                                              int w,
+                                                              int patch_t,
+                                                              int patch_h,
+                                                              int patch_w,
+                                                              int bs,
+                                                              int context_len,
+                                                              float theta,
+                                                              const std::vector<int>& axes_dim) {
+        auto ids = gen_hunyuan_video_ids(t, h, w, patch_t, patch_h, patch_w, bs, context_len);
+        return embed_nd(ids, bs, theta, axes_dim);
+    }
+
     __STATIC_INLINE__ std::vector<std::vector<float>> gen_qwen_image_ids(int t,
                                                                          int h,
                                                                          int w,
@@ -627,6 +654,43 @@ namespace Rope {
         return embed_nd(ids, bs, static_cast<float>(theta), axes_dim, wrap_dims);
     }
 
+    __STATIC_INLINE__ std::vector<float> gen_mage_flow_pe(int h,
+                                                          int w,
+                                                          int bs,
+                                                          int context_len,
+                                                          const std::vector<ggml_tensor*>& ref_latents,
+                                                          int theta,
+                                                          const std::vector<int>& axes_dim) {
+        const int axes_dim_num = static_cast<int>(axes_dim.size());
+        auto make_image_ids    = [=](int image_h, int image_w, int image_index) {
+            std::vector<std::vector<float>> image_ids(static_cast<size_t>(bs) * image_h * image_w,
+                                                         std::vector<float>(axes_dim_num, 0.f));
+            int h_start = -(image_h - image_h / 2);
+            int w_start = -(image_w - image_w / 2);
+            for (int b = 0; b < bs; ++b) {
+                for (int y = 0; y < image_h; ++y) {
+                    for (int x = 0; x < image_w; ++x) {
+                        auto& id = image_ids[static_cast<size_t>(b) * image_h * image_w + y * image_w + x];
+                        id[0]    = static_cast<float>(image_index);
+                        id[1]    = static_cast<float>(h_start + y);
+                        id[2]    = static_cast<float>(w_start + x);
+                    }
+                }
+            }
+            return image_ids;
+        };
+        auto ids     = gen_flux_txt_ids(bs, context_len, axes_dim_num, {});
+        auto img_ids = make_image_ids(h, w, 0);
+        ids          = concat_ids(ids, img_ids, bs);
+        for (size_t i = 0; i < ref_latents.size(); ++i) {
+            auto ref_ids = make_image_ids(static_cast<int>(ref_latents[i]->ne[1]),
+                                          static_cast<int>(ref_latents[i]->ne[0]),
+                                          static_cast<int>(i + 1));
+            ids          = concat_ids(ids, ref_ids, bs);
+        }
+        return embed_nd(ids, bs, static_cast<float>(theta), axes_dim);
+    }
+
     __STATIC_INLINE__ std::vector<std::vector<float>> gen_lens_ids(int h,
                                                                    int w,
                                                                    int bs,

src/model/diffusion/flux.hpp

@@ -706,11 +706,13 @@ namespace Flux {
         LastLayer(int64_t hidden_size,
                   int64_t patch_size,
                   int64_t out_channels,
-                  bool prune_mod = false,
-                  bool bias      = true)
+                  bool prune_mod       = false,
+                  bool bias            = true,
+                  int64_t patch_volume = 0)
             : prune_mod(prune_mod) {
             blocks["norm_final"] = std::shared_ptr<GGMLBlock>(new LayerNorm(hidden_size, 1e-06f, false));
-            blocks["linear"]     = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, patch_size * patch_size * out_channels, bias));
+            int64_t out_dim      = (patch_volume > 0 ? patch_volume : patch_size * patch_size) * out_channels;
+            blocks["linear"]     = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, out_dim, bias));
             if (!prune_mod) {
                 blocks["adaLN_modulation.1"] = std::shared_ptr<GGMLBlock>(new Linear(hidden_size, 2 * hidden_size, bias));
             }

src/model/diffusion/hunyuan.hpp

@@ -0,0 +1,681 @@
+#ifndef __SD_MODEL_DIFFUSION_HUNYUAN_HPP__
+#define __SD_MODEL_DIFFUSION_HUNYUAN_HPP__
+
+#include <memory>
+
+#include "model/common/block.hpp"
+#include "model/diffusion/flux.hpp"
+#include "model/diffusion/mmdit.hpp"
+#include "model/diffusion/wan.hpp"
+#include "model_manager.h"
+
+namespace Hunyuan {
+    constexpr int HUNYUAN_VIDEO_GRAPH_SIZE = 65536;
+
+    // Ref: https://github.com/huggingface/diffusers/pull/12696
+    struct IndividualTokenRefinerBlock : public GGMLBlock {
+    protected:
+        int64_t num_heads;
+
+    public:
+        IndividualTokenRefinerBlock(int64_t num_heads,
+                                    int64_t head_dim,
+                                    int64_t mlp_ratio = 4,
+                                    bool attn_bias    = true)
+            : num_heads(num_heads) {
+            int64_t hidden_size      = num_heads * head_dim;
+            blocks["self_attn.qkv"]  = std::make_shared<Linear>(hidden_size, hidden_size * 3, attn_bias);
+            blocks["self_attn.proj"] = std::make_shared<Linear>(hidden_size, hidden_size, attn_bias);
+
+            blocks["norm1"] = std::make_shared<LayerNorm>(hidden_size, 1e-6f, true);
+            blocks["norm2"] = std::make_shared<LayerNorm>(hidden_size, 1e-6f, true);
+
+            blocks["mlp.0"] = std::make_shared<Linear>(hidden_size, hidden_size * mlp_ratio);
+            blocks["mlp.2"] = std::make_shared<Linear>(hidden_size * mlp_ratio, hidden_size);
+
+            // adaLN_modulation.0 is nn.SiLU()
+            blocks["adaLN_modulation.1"] = std::make_shared<Linear>(hidden_size, hidden_size * 2);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* txt, ggml_tensor* t_emb, ggml_tensor* mask) {
+            auto norm1              = std::dynamic_pointer_cast<LayerNorm>(blocks["norm1"]);
+            auto norm2              = std::dynamic_pointer_cast<LayerNorm>(blocks["norm2"]);
+            auto self_attn_qkv      = std::dynamic_pointer_cast<Linear>(blocks["self_attn.qkv"]);
+            auto self_attn_proj     = std::dynamic_pointer_cast<Linear>(blocks["self_attn.proj"]);
+            auto mlp_fc1            = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
+            auto mlp_fc2            = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
+            auto adaLN_modulation_1 = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.1"]);
+
+            // self attn
+            auto qkv     = self_attn_qkv->forward(ctx, norm1->forward(ctx, txt));
+            auto qkv_vec = split_qkv(ctx->ggml_ctx, qkv);
+            auto q       = qkv_vec[0];
+            auto k       = qkv_vec[1];
+            auto v       = qkv_vec[2];
+
+            auto attn_out = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, mask, false, ctx->flash_attn_enabled);
+            attn_out      = self_attn_proj->forward(ctx, attn_out);
+
+            // adaLN_modulation
+            auto emb  = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, t_emb));
+            auto mods = ggml_ext_chunk(ctx->ggml_ctx, emb, 2, 0);
+
+            txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, attn_out, mods[0]));
+
+            // mlp
+            auto mlp_out = mlp_fc1->forward(ctx, norm2->forward(ctx, txt));
+            mlp_out      = ggml_silu_inplace(ctx->ggml_ctx, mlp_out);
+            mlp_out      = mlp_fc2->forward(ctx, mlp_out);
+            txt          = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, mlp_out, mods[1]));
+
+            return txt;
+        }
+    };
+
+    struct IndividualTokenRefiner : public GGMLBlock {
+    protected:
+        int num_layers;
+
+    public:
+        IndividualTokenRefiner(int64_t num_heads,
+                               int64_t head_dim,
+                               int num_layers,
+                               int64_t mlp_ratio = 4,
+                               bool attn_bias    = true)
+            : num_layers(num_layers) {
+            for (int i = 0; i < num_layers; i++) {
+                blocks["blocks." + std::to_string(i)] = std::make_shared<IndividualTokenRefinerBlock>(num_heads, head_dim, mlp_ratio, attn_bias);
+            }
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* txt, ggml_tensor* t_emb, ggml_tensor* mask) {
+            for (int i = 0; i < num_layers; i++) {
+                auto block = std::dynamic_pointer_cast<IndividualTokenRefinerBlock>(blocks["blocks." + std::to_string(i)]);
+
+                txt = block->forward(ctx, txt, t_emb, mask);
+            }
+
+            return txt;
+        }
+    };
+
+    struct TokenRefiner : public GGMLBlock {
+    public:
+        TokenRefiner(int64_t in_channels,
+                     int64_t num_heads,
+                     int64_t head_dim,
+                     int num_layers,
+                     int64_t mlp_ratio = 4,
+                     bool attn_bias    = true) {
+            int64_t hidden_size                = num_heads * head_dim;
+            blocks["input_embedder"]           = std::make_shared<Linear>(in_channels, hidden_size);
+            blocks["t_embedder"]               = std::make_shared<Flux::MLPEmbedder>(256, hidden_size);
+            blocks["c_embedder"]               = std::make_shared<Flux::MLPEmbedder>(in_channels, hidden_size);
+            blocks["individual_token_refiner"] = std::make_shared<IndividualTokenRefiner>(num_heads, head_dim, num_layers, mlp_ratio, attn_bias);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* txt, ggml_tensor* timestep, ggml_tensor* mask) {
+            auto input_embedder           = std::dynamic_pointer_cast<Linear>(blocks["input_embedder"]);
+            auto t_embedder               = std::dynamic_pointer_cast<Flux::MLPEmbedder>(blocks["t_embedder"]);
+            auto c_embedder               = std::dynamic_pointer_cast<Flux::MLPEmbedder>(blocks["c_embedder"]);
+            auto individual_token_refiner = std::dynamic_pointer_cast<IndividualTokenRefiner>(blocks["individual_token_refiner"]);
+
+            auto t_emb = t_embedder->forward(ctx, ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, 256, 10000, 1.f));
+
+            auto h                  = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, txt, 1, 0, 2, 3));
+            auto pooled_projections = ggml_scale(ctx->ggml_ctx, ggml_sum_rows(ctx->ggml_ctx, h), 1.f / txt->ne[1]);
+            pooled_projections      = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, pooled_projections, 1, 0, 2, 3));
+            auto c_emb              = c_embedder->forward(ctx, pooled_projections);
+
+            t_emb = ggml_add(ctx->ggml_ctx, t_emb, c_emb);
+            txt   = input_embedder->forward(ctx, txt);
+            txt   = individual_token_refiner->forward(ctx, txt, t_emb, mask);
+            return txt;
+        }
+    };
+
+    struct ByT5Mapper : public UnaryBlock {
+        ByT5Mapper(int64_t in_dim, int64_t hidden_size) {
+            blocks["layernorm"] = std::make_shared<LayerNorm>(in_dim);
+            blocks["fc1"]       = std::make_shared<Linear>(in_dim, 2048);
+            blocks["fc2"]       = std::make_shared<Linear>(2048, 2048);
+            blocks["fc3"]       = std::make_shared<Linear>(2048, hidden_size);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+            auto layernorm = std::dynamic_pointer_cast<LayerNorm>(blocks["layernorm"]);
+            auto fc1       = std::dynamic_pointer_cast<Linear>(blocks["fc1"]);
+            auto fc2       = std::dynamic_pointer_cast<Linear>(blocks["fc2"]);
+            auto fc3       = std::dynamic_pointer_cast<Linear>(blocks["fc3"]);
+
+            x = fc1->forward(ctx, layernorm->forward(ctx, x));
+            x = ggml_ext_gelu(ctx->ggml_ctx, x);
+            x = fc2->forward(ctx, x);
+            x = ggml_ext_gelu(ctx->ggml_ctx, x);
+            return fc3->forward(ctx, x);
+        }
+    };
+
+    struct HunyuanVideoConfig {
+        std::tuple<int, int, int> patch_size = {1, 2, 2};
+        int64_t in_channels                  = 65;
+        int64_t out_channels                 = 32;
+        int64_t hidden_size                  = 2048;
+        int64_t vec_in_dim                   = 0;
+        int64_t context_in_dim               = 3584;
+        int64_t vision_in_dim                = 0;
+        float mlp_ratio                      = 4.0f;
+        int num_heads                        = 16;
+        int depth                            = 54;
+        int depth_single_blocks              = 0;
+        bool qkv_bias                        = true;
+        bool guidance_embed                  = false;
+        bool use_byt5                        = false;
+        bool use_cond_type_embedding         = false;
+        bool use_meanflow                    = false;
+        bool use_meanflow_sum                = false;
+        float theta                          = 256;
+        std::vector<int> axes_dim            = {16, 56, 56};
+        int axes_dim_sum                     = 128;
+
+        int64_t patch_volume() const {
+            return static_cast<int64_t>(std::get<0>(patch_size)) * std::get<1>(patch_size) * std::get<2>(patch_size);
+        }
+
+        static HunyuanVideoConfig detect_from_weights(const String2TensorStorage& tensor_storage_map,
+                                                      const std::string& prefix) {
+            HunyuanVideoConfig config;
+            config.depth               = 0;
+            config.depth_single_blocks = 0;
+            bool inferred              = false;
+
+            int64_t img_embed_dim = 0;
+            for (const auto& [name, storage] : tensor_storage_map) {
+                if (starts_with(name, prefix) && ends_with(name, "img_in.proj.bias")) {
+                    img_embed_dim = storage.ne[0];
+                    break;
+                }
+            }
+
+            for (const auto& entry : tensor_storage_map) {
+                const auto& name    = entry.first;
+                const auto& storage = entry.second;
+                if (!starts_with(name, prefix)) {
+                    continue;
+                }
+
+                auto update_depth = [&](const char* block_prefix, int* depth) {
+                    size_t pos = name.find(block_prefix);
+                    if (pos == std::string::npos) {
+                        return;
+                    }
+                    pos += strlen(block_prefix);
+                    size_t end = name.find('.', pos);
+                    if (end != std::string::npos) {
+                        *depth = std::max(*depth, atoi(name.substr(pos, end - pos).c_str()) + 1);
+                    }
+                };
+                update_depth("double_blocks.", &config.depth);
+                update_depth("single_blocks.", &config.depth_single_blocks);
+
+                if (ends_with(name, "img_in.proj.weight") && storage.n_dims == 5) {
+                    config.patch_size  = {static_cast<int>(storage.ne[2]),
+                                          static_cast<int>(storage.ne[1]),
+                                          static_cast<int>(storage.ne[0])};
+                    config.in_channels = storage.ne[3];
+                    config.hidden_size = storage.ne[4];
+                    inferred           = true;
+                } else if (ends_with(name, "img_in.proj.weight") && storage.n_dims == 4) {
+                    config.patch_size = {static_cast<int>(storage.ne[2]),
+                                         static_cast<int>(storage.ne[1]),
+                                         static_cast<int>(storage.ne[0])};
+                    if (img_embed_dim > 0 && storage.ne[3] % img_embed_dim == 0) {
+                        config.hidden_size = img_embed_dim;
+                        config.in_channels = storage.ne[3] / img_embed_dim;
+                    }
+                    inferred = true;
+                } else if (ends_with(name, "txt_in.input_embedder.weight")) {
+                    config.context_in_dim = storage.ne[0];
+                    inferred              = true;
+                } else if (ends_with(name, "vector_in.in_layer.weight")) {
+                    config.vec_in_dim = storage.ne[0];
+                } else if (ends_with(name, "vision_in.proj.0.weight")) {
+                    config.vision_in_dim = storage.ne[0];
+                } else if (ends_with(name, "double_blocks.0.img_attn.norm.key_norm.scale") ||
+                           ends_with(name, "double_blocks.0.img_attn.norm.key_norm.weight")) {
+                    config.num_heads = static_cast<int>(config.hidden_size / storage.ne[0]);
+                } else if (ends_with(name, "double_blocks.0.img_mlp.0.weight")) {
+                    config.mlp_ratio = static_cast<float>(storage.ne[1]) / static_cast<float>(storage.ne[0]);
+                }
+
+                config.guidance_embed = config.guidance_embed || name.find("guidance_in.") != std::string::npos;
+                config.use_byt5       = config.use_byt5 || name.find("byt5_in.") != std::string::npos;
+                config.use_meanflow   = config.use_meanflow || name.find("time_r_in.") != std::string::npos;
+            }
+
+            config.use_cond_type_embedding = tensor_storage_map.find(prefix + ".cond_type_embedding.weight") != tensor_storage_map.end();
+            config.use_meanflow_sum        = config.vision_in_dim > 0;
+
+            auto final_iter = tensor_storage_map.find(prefix + ".final_layer.linear.weight");
+            if (final_iter != tensor_storage_map.end()) {
+                config.out_channels = final_iter->second.ne[1] / config.patch_volume();
+            }
+            config.qkv_bias = tensor_storage_map.find(prefix + ".double_blocks.0.img_attn.qkv.bias") != tensor_storage_map.end();
+
+            GGML_ASSERT(config.hidden_size % config.num_heads == 0);
+            GGML_ASSERT(config.hidden_size / config.num_heads == config.axes_dim_sum);
+
+            if (inferred) {
+                LOG_DEBUG("hunyuan video: depth = %d, single depth = %d, in_channels = %" PRId64 ", out_channels = %" PRId64 ", hidden_size = %" PRId64 ", context_in_dim = %" PRId64 ", patch_size = %dx%dx%d",
+                          config.depth,
+                          config.depth_single_blocks,
+                          config.in_channels,
+                          config.out_channels,
+                          config.hidden_size,
+                          config.context_in_dim,
+                          std::get<0>(config.patch_size),
+                          std::get<1>(config.patch_size),
+                          std::get<2>(config.patch_size));
+            }
+            return config;
+        }
+    };
+
+    class HunyuanVideoModel : public GGMLBlock {
+    protected:
+        HunyuanVideoConfig config;
+
+        void init_params(struct ggml_context* ctx,
+                         const String2TensorStorage& tensor_storage_map = {},
+                         const std::string prefix                       = "") override {
+            if (config.use_cond_type_embedding) {
+                ggml_type type                                  = get_type(prefix + "cond_type_embedding.weight", tensor_storage_map, GGML_TYPE_F16);
+                GGMLBlock::params["cond_type_embedding.weight"] = ggml_new_tensor_2d(ctx, type, config.hidden_size, 3);
+            }
+        }
+
+    public:
+        HunyuanVideoModel() {}
+        explicit HunyuanVideoModel(HunyuanVideoConfig config)
+            : config(std::move(config)) {
+            int64_t head_dim  = this->config.hidden_size / this->config.num_heads;
+            blocks["txt_in"]  = std::make_shared<TokenRefiner>(this->config.context_in_dim, this->config.num_heads, head_dim, 2);
+            blocks["img_in"]  = std::make_shared<PatchEmbed>(static_cast<int64_t>(224) /*Not used*/,
+                                                            this->config.patch_size,
+                                                            this->config.in_channels,
+                                                            this->config.hidden_size);
+            blocks["time_in"] = std::make_shared<Flux::MLPEmbedder>(256, this->config.hidden_size);
+            if (this->config.vec_in_dim > 0) {
+                blocks["vector_in"] = std::make_shared<Flux::MLPEmbedder>(this->config.vec_in_dim, this->config.hidden_size);
+            }
+            if (this->config.vision_in_dim > 0) {
+                blocks["vision_in"] = std::make_shared<WAN::MLPProj>(this->config.vision_in_dim, this->config.hidden_size);
+            }
+            if (this->config.guidance_embed) {
+                blocks["guidance_in"] = std::make_shared<Flux::MLPEmbedder>(256, this->config.hidden_size);
+            }
+            if (this->config.use_byt5) {
+                blocks["byt5_in"] = std::make_shared<ByT5Mapper>(1472, this->config.hidden_size);
+            }
+            if (this->config.use_meanflow) {
+                blocks["time_r_in"] = std::make_shared<Flux::MLPEmbedder>(256, this->config.hidden_size);
+            }
+
+            for (int i = 0; i < this->config.depth; i++) {
+                blocks["double_blocks." + std::to_string(i)] = std::make_shared<Flux::DoubleStreamBlock>(this->config.hidden_size,
+                                                                                                         this->config.num_heads,
+                                                                                                         this->config.mlp_ratio,
+                                                                                                         i,
+                                                                                                         this->config.qkv_bias);
+            }
+
+            for (int i = 0; i < this->config.depth_single_blocks; i++) {
+                blocks["single_blocks." + std::to_string(i)] = std::make_shared<Flux::SingleStreamBlock>(this->config.hidden_size,
+                                                                                                         this->config.num_heads,
+                                                                                                         this->config.mlp_ratio,
+                                                                                                         i,
+                                                                                                         0.f);
+            }
+
+            blocks["final_layer"] = std::make_shared<Flux::LastLayer>(this->config.hidden_size,
+                                                                      std::get<2>(this->config.patch_size),
+                                                                      this->config.out_channels,
+                                                                      false,
+                                                                      true,
+                                                                      this->config.patch_volume());
+        }
+
+        ggml_tensor* pad_to_patch_size(struct ggml_context* ctx,
+                                       ggml_tensor* x) {
+            int64_t W = x->ne[0];
+            int64_t H = x->ne[1];
+            int64_t T = x->ne[2];
+
+            int pt    = std::get<0>(config.patch_size);
+            int ph    = std::get<1>(config.patch_size);
+            int pw    = std::get<2>(config.patch_size);
+            int pad_t = (pt - static_cast<int>(T % pt)) % pt;
+            int pad_h = (ph - static_cast<int>(H % ph)) % ph;
+            int pad_w = (pw - static_cast<int>(W % pw)) % pw;
+            x         = ggml_pad(ctx, x, pad_w, pad_h, pad_t, 0);  // [N*C, T + pad_t, H + pad_h, W + pad_w]
+
+            return x;
+        }
+
+        ggml_tensor* unpatchify(struct ggml_context* ctx,
+                                ggml_tensor* x,
+                                int64_t t_len,
+                                int64_t h_len,
+                                int64_t w_len) {
+            // x: [N, t_len*h_len*w_len, C*pt*ph*pw]
+            // return: [N*C, t_len*pt, h_len*ph, w_len*pw]
+            int64_t N  = x->ne[3];
+            int64_t pt = std::get<0>(config.patch_size);
+            int64_t ph = std::get<1>(config.patch_size);
+            int64_t pw = std::get<2>(config.patch_size);
+            int64_t C  = x->ne[0] / pt / ph / pw;
+
+            GGML_ASSERT(C * pt * ph * pw == x->ne[0]);
+
+            x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N);  // [N, t_len*h_len*w_len, pt*ph*pw, C]
+            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3));      // [N, C, t_len*h_len*w_len, pt*ph*pw]
+            x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N);  // [N*C*t_len*h_len, w_len, pt*ph, pw]
+            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len*h_len, pt*ph, w_len, pw]
+            x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N);  // [N*C*t_len*h_len, pt, ph, w_len*pw]
+            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len*h_len, ph, pt, w_len*pw]
+            x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph * h_len, t_len * C * N);  // [N*C*t_len, h_len*ph, pt, w_len*pw]
+            x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3));      // [N*C*t_len, pt, h_len*ph, w_len*pw]
+            x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt * t_len, C * N);  // [N*C, t_len*pt, h_len*ph, w_len*pw]
+            return x;
+        }
+
+        ggml_tensor* add_condition_type(GGMLRunnerContext* ctx, ggml_tensor* x, int type) {
+            if (!config.use_cond_type_embedding) {
+                return x;
+            }
+            auto weight = GGMLBlock::params["cond_type_embedding.weight"];
+            auto row    = ggml_view_1d(ctx->ggml_ctx,
+                                       weight,
+                                       weight->ne[0],
+                                       static_cast<size_t>(type) * weight->nb[1]);
+            auto target = ggml_new_tensor_3d(ctx->ggml_ctx, row->type, config.hidden_size, x->ne[1], x->ne[2]);
+            auto embed  = ggml_repeat(ctx->ggml_ctx, row, target);
+            embed       = ggml_cast(ctx->ggml_ctx, embed, x->type);
+            return ggml_add(ctx->ggml_ctx, x, embed);
+        }
+
+        ggml_tensor* forward_orig(GGMLRunnerContext* ctx,
+                                  ggml_tensor* img,
+                                  ggml_tensor* txt,
+                                  ggml_tensor* timestep,
+                                  ggml_tensor* pe,
+                                  ggml_tensor* guidance   = nullptr,
+                                  ggml_tensor* y          = nullptr,
+                                  ggml_tensor* txt_byt5   = nullptr,
+                                  ggml_tensor* clip_fea   = nullptr,
+                                  ggml_tensor* timestep_r = nullptr,
+                                  int64_t N               = 1) {
+            // img: [N*C, T, H, W], C => in_dim
+            // txt: [N, L, text_dim]
+            // timestep: [N,] or [T]
+            // return: [N, t_len*h_len*w_len, out_dim*pt*ph*pw]
+
+            GGML_ASSERT(N == 1);
+
+            auto img_in      = std::dynamic_pointer_cast<PatchEmbed>(blocks["img_in"]);
+            auto txt_in      = std::dynamic_pointer_cast<TokenRefiner>(blocks["txt_in"]);
+            auto time_in     = std::dynamic_pointer_cast<Flux::MLPEmbedder>(blocks["time_in"]);
+            auto final_layer = std::dynamic_pointer_cast<Flux::LastLayer>(blocks["final_layer"]);
+
+            img      = img_in->forward(ctx, img);                     // [N*C, t_len*h_len*w_len, hidden_size]
+            txt      = txt_in->forward(ctx, txt, timestep, nullptr);  // [N, n_txt_token, hidden_size]
+            auto vec = time_in->forward(ctx, ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, 256, 10000, 1.f));
+            if (config.use_meanflow && timestep_r != nullptr) {
+                auto time_r_in = std::dynamic_pointer_cast<Flux::MLPEmbedder>(blocks["time_r_in"]);
+                auto vec_r     = time_r_in->forward(ctx, ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_r, 256, 10000, 1000.f));
+                vec            = ggml_add(ctx->ggml_ctx, vec, vec_r);
+                if (!config.use_meanflow_sum) {
+                    vec = ggml_scale(ctx->ggml_ctx, vec, 0.5f);
+                }
+            }
+            if (config.vec_in_dim > 0 && y != nullptr) {
+                auto vector_in = std::dynamic_pointer_cast<Flux::MLPEmbedder>(blocks["vector_in"]);
+                vec            = ggml_add(ctx->ggml_ctx, vec, vector_in->forward(ctx, y));
+            }
+            if (config.guidance_embed && guidance != nullptr) {
+                auto guidance_in  = std::dynamic_pointer_cast<Flux::MLPEmbedder>(blocks["guidance_in"]);
+                auto guidance_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, guidance, 256, 10000, 1.f);
+                vec               = ggml_add(ctx->ggml_ctx, vec, guidance_in->forward(ctx, guidance_emb));
+            }
+
+            txt = add_condition_type(ctx, txt, 0);
+            if (config.use_byt5 && txt_byt5 != nullptr) {
+                auto byt5_in = std::dynamic_pointer_cast<ByT5Mapper>(blocks["byt5_in"]);
+                txt_byt5     = add_condition_type(ctx, byt5_in->forward(ctx, txt_byt5), 1);
+                txt          = config.use_cond_type_embedding ? ggml_concat(ctx->ggml_ctx, txt_byt5, txt, 1)
+                                                              : ggml_concat(ctx->ggml_ctx, txt, txt_byt5, 1);
+            }
+            if (config.vision_in_dim > 0 && clip_fea != nullptr) {
+                auto vision_in = std::dynamic_pointer_cast<WAN::MLPProj>(blocks["vision_in"]);
+                clip_fea       = add_condition_type(ctx, vision_in->forward(ctx, clip_fea), 2);
+                txt            = ggml_concat(ctx->ggml_ctx, clip_fea, txt, 1);
+            }
+
+            for (int i = 0; i < config.depth; i++) {
+                auto block = std::dynamic_pointer_cast<Flux::DoubleStreamBlock>(blocks["double_blocks." + std::to_string(i)]);
+
+                auto img_txt = block->forward(ctx, img, txt, vec, pe, nullptr);
+                img          = img_txt.first;   // [N, n_img_token, hidden_size]
+                txt          = img_txt.second;  // [N, n_txt_token, hidden_size]
+            }
+
+            if (config.depth_single_blocks > 0) {
+                auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1);  // [N, n_txt_token + n_img_token, hidden_size]
+                for (int i = 0; i < config.depth_single_blocks; i++) {
+                    auto block = std::dynamic_pointer_cast<Flux::SingleStreamBlock>(blocks["single_blocks." + std::to_string(i)]);
+                    txt_img    = block->forward(ctx, txt_img, vec, pe, nullptr);
+                }
+
+                txt_img = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, txt_img, 0, 2, 1, 3));
+                img     = ggml_view_3d(ctx->ggml_ctx,
+                                       txt_img,
+                                       txt_img->ne[0],
+                                       txt_img->ne[1],
+                                       img->ne[1],
+                                       txt_img->nb[1],
+                                       txt_img->nb[2],
+                                       txt_img->nb[2] * txt->ne[1]);
+                img     = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, img, 0, 2, 1, 3));
+            }
+
+            img = final_layer->forward(ctx, img, vec);  // (N, t_len*h_len*w_len, out_channels * patch_size ** 3)
+
+            return img;
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             ggml_tensor* timestep,
+                             ggml_tensor* context,
+                             ggml_tensor* pe,
+                             ggml_tensor* guidance   = nullptr,
+                             ggml_tensor* y          = nullptr,
+                             ggml_tensor* txt_byt5   = nullptr,
+                             ggml_tensor* clip_fea   = nullptr,
+                             ggml_tensor* timestep_r = nullptr,
+                             int64_t N               = 1) {
+            // Forward pass of DiT.
+            // x: [N*C, T, H, W]
+            // timestep: [N,]
+            // context: [N, L, D]
+            // pe: [L, d_head/2, 2, 2]
+            // return: [N*C, T, H, W]
+
+            GGML_ASSERT(N == 1);
+
+            int64_t W = x->ne[0];
+            int64_t H = x->ne[1];
+            int64_t T = x->ne[2];
+            x         = pad_to_patch_size(ctx->ggml_ctx, x);
+
+            int64_t pt    = std::get<0>(config.patch_size);
+            int64_t ph    = std::get<1>(config.patch_size);
+            int64_t pw    = std::get<2>(config.patch_size);
+            int64_t t_len = (T + pt - 1) / pt;
+            int64_t h_len = (H + ph - 1) / ph;
+            int64_t w_len = (W + pw - 1) / pw;
+
+            auto out = forward_orig(ctx, x, context, timestep, pe, guidance, y, txt_byt5, clip_fea, timestep_r, N);
+
+            out = unpatchify(ctx->ggml_ctx, out, t_len, h_len, w_len);  // [N*C, (T+pad_t) + (T2+pad_t2), H + pad_h, W + pad_w]
+
+            // slice
+            out = ggml_ext_slice(ctx->ggml_ctx, out, 2, 0, T);  // [N*C, T, H + pad_h, W + pad_w]
+            out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, H);  // [N*C, T, H, W + pad_w]
+            out = ggml_ext_slice(ctx->ggml_ctx, out, 0, 0, W);  // [N*C, T, H, W]
+
+            return out;
+        }
+    };
+
+    struct HunyuanVideoRunner : public DiffusionModelRunner {
+    public:
+        HunyuanVideoConfig config;
+        HunyuanVideoModel hunyuan_video;
+        std::vector<float> pe_vec;
+        SDVersion version;
+
+        HunyuanVideoRunner(ggml_backend_t backend,
+                           const String2TensorStorage& tensor_storage_map      = {},
+                           const std::string prefix                            = "",
+                           SDVersion version                                   = VERSION_HUNYUAN_VIDEO,
+                           std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : DiffusionModelRunner(backend, prefix, weight_manager),
+              config(HunyuanVideoConfig::detect_from_weights(tensor_storage_map, prefix)),
+              version(version) {
+            LOG_INFO("HunyuanVideo blocks: %d double, %d single", config.depth, config.depth_single_blocks);
+
+            hunyuan_video = HunyuanVideoModel(config);
+            hunyuan_video.init(params_ctx, tensor_storage_map, prefix);
+        }
+
+        std::string get_desc() override {
+            return "hunyuan_video";
+        }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
+            hunyuan_video.get_param_tensors(tensors, prefix);
+        }
+
+        ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
+                                 const sd::Tensor<float>& timesteps_tensor,
+                                 const sd::Tensor<float>& context_tensor,
+                                 const sd::Tensor<float>& c_concat_tensor   = {},
+                                 const sd::Tensor<float>& y_tensor          = {},
+                                 const sd::Tensor<float>& guidance_tensor   = {},
+                                 const sd::Tensor<float>& byt5_tensor       = {},
+                                 const sd::Tensor<float>& vision_tensor     = {},
+                                 const sd::Tensor<float>& timestep_r_tensor = {}) {
+            ggml_cgraph* gf = new_graph_custom(HUNYUAN_VIDEO_GRAPH_SIZE);
+
+            ggml_tensor* x          = make_input(x_tensor);
+            ggml_tensor* timesteps  = make_input(timesteps_tensor);
+            ggml_tensor* context    = make_input(context_tensor);
+            ggml_tensor* c_concat   = make_optional_input(c_concat_tensor);
+            ggml_tensor* y          = make_optional_input(y_tensor);
+            ggml_tensor* guidance   = make_optional_input(guidance_tensor);
+            ggml_tensor* byt5       = make_optional_input(byt5_tensor);
+            ggml_tensor* vision     = make_optional_input(vision_tensor);
+            ggml_tensor* timestep_r = make_optional_input(timestep_r_tensor);
+
+            GGML_ASSERT(x->ne[3] == config.out_channels);
+            if (c_concat != nullptr) {
+                x = ggml_concat(compute_ctx, x, c_concat, 3);
+            }
+            GGML_ASSERT(x->ne[3] <= config.in_channels);
+            if (x->ne[3] < config.in_channels) {
+                x = ggml_pad(compute_ctx, x, 0, 0, 0, static_cast<int>(config.in_channels - x->ne[3]));
+            }
+
+            int text_len = static_cast<int>(context->ne[1]);
+            if (byt5 != nullptr) {
+                text_len += static_cast<int>(byt5->ne[1]);
+            }
+            if (vision != nullptr) {
+                text_len += static_cast<int>(vision->ne[1]);
+            }
+            pe_vec          = Rope::gen_hunyuan_video_pe(static_cast<int>(x->ne[2]),
+                                                         static_cast<int>(x->ne[1]),
+                                                         static_cast<int>(x->ne[0]),
+                                                         std::get<0>(config.patch_size),
+                                                         std::get<1>(config.patch_size),
+                                                         std::get<2>(config.patch_size),
+                                                         1,
+                                                         text_len,
+                                                         config.theta,
+                                                         config.axes_dim);
+            int64_t pos_len = static_cast<int64_t>(pe_vec.size() / config.axes_dim_sum / 2);
+            // LOG_DEBUG("pos_len %d", pos_len);
+            auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
+            // pe->data = pe_vec.data();
+            // print_ggml_tensor(pe, true, "pe");
+            // pe->data = nullptr;
+            set_backend_tensor_data(pe, pe_vec.data());
+
+            auto runner_ctx = get_context();
+
+            ggml_tensor* out = hunyuan_video.forward(&runner_ctx,
+                                                     x,
+                                                     timesteps,
+                                                     context,
+                                                     pe,
+                                                     guidance,
+                                                     y,
+                                                     byt5,
+                                                     vision,
+                                                     timestep_r);
+
+            ggml_build_forward_expand(gf, out);
+
+            return gf;
+        }
+
+        sd::Tensor<float> compute(int n_threads,
+                                  const sd::Tensor<float>& x,
+                                  const sd::Tensor<float>& timesteps,
+                                  const sd::Tensor<float>& context,
+                                  const sd::Tensor<float>& c_concat   = {},
+                                  const sd::Tensor<float>& y          = {},
+                                  const sd::Tensor<float>& guidance   = {},
+                                  const sd::Tensor<float>& byt5       = {},
+                                  const sd::Tensor<float>& vision     = {},
+                                  const sd::Tensor<float>& timestep_r = {}) {
+            auto get_graph = [&]() -> ggml_cgraph* {
+                return build_graph(x, timesteps, context, c_concat, y, guidance, byt5, vision, timestep_r);
+            };
+
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
+        }
+
+        sd::Tensor<float> compute(int n_threads,
+                                  const DiffusionParams& diffusion_params) override {
+            GGML_ASSERT(diffusion_params.x != nullptr);
+            GGML_ASSERT(diffusion_params.timesteps != nullptr);
+            GGML_ASSERT(diffusion_params.context != nullptr);
+            const auto* extra = diffusion_extra_as<HunyuanVideoDiffusionExtra>(diffusion_params);
+            return compute(n_threads,
+                           *diffusion_params.x,
+                           *diffusion_params.timesteps,
+                           *diffusion_params.context,
+                           tensor_or_empty(diffusion_params.c_concat),
+                           tensor_or_empty(diffusion_params.y),
+                           tensor_or_empty(extra->guidance),
+                           tensor_or_empty(extra->byt5),
+                           tensor_or_empty(extra->vision),
+                           tensor_or_empty(extra->timestep_r));
+        }
+    };
+
+}  // namespace Hunyuan
+
+#endif  // __SD_MODEL_DIFFUSION_HUNYUAN_HPP__

src/model/diffusion/ltxv.hpp

@@ -800,7 +800,7 @@ namespace LTXV {
             auto gate_mlp  = mods[5];
 
             auto x_norm = rms_norm(ctx->ggml_ctx, x);
-            x_norm      = modulate(ctx->ggml_ctx, x_norm, shift_msa, scale_msa);
+            x_norm      = LTXV::modulate(ctx->ggml_ctx, x_norm, shift_msa, scale_msa);
             auto msa    = attn1->forward(ctx, x_norm, nullptr, self_attention_mask, pe);
             x           = ggml_add(ctx->ggml_ctx, x, apply_gate(ctx->ggml_ctx, msa, gate_msa));
 
@@ -810,12 +810,12 @@ namespace LTXV {
                 auto gate_q  = mods[8];
 
                 auto q = rms_norm(ctx->ggml_ctx, x);
-                q      = modulate(ctx->ggml_ctx, q, shift_q, scale_q);
+                q      = LTXV::modulate(ctx->ggml_ctx, q, shift_q, scale_q);
 
                 auto context_mod = context;
                 if (prompt_timestep != nullptr) {
                     auto prompt_mods = get_prompt_scale_shift_values(ctx, prompt_timestep);
-                    context_mod      = modulate(ctx->ggml_ctx, context_mod, prompt_mods[0], prompt_mods[1]);
+                    context_mod      = LTXV::modulate(ctx->ggml_ctx, context_mod, prompt_mods[0], prompt_mods[1]);
                 }
 
                 auto mca = attn2->forward(ctx, q, context_mod, attention_mask, nullptr, nullptr);
@@ -826,7 +826,7 @@ namespace LTXV {
             }
 
             auto y       = rms_norm(ctx->ggml_ctx, x);
-            y            = modulate(ctx->ggml_ctx, y, shift_mlp, scale_mlp);
+            y            = LTXV::modulate(ctx->ggml_ctx, y, shift_mlp, scale_mlp);
             auto mlp_out = ff->forward(ctx, y);
             x            = ggml_add(ctx->ggml_ctx, x, apply_gate(ctx->ggml_ctx, mlp_out, gate_mlp));
             return x;
@@ -1177,11 +1177,11 @@ namespace LTXV {
             if (cross_attention_adaln) {
                 auto q_mods      = get_ada_values(ctx, table, timestep, dim, 9, 6, 3);
                 auto q           = rms_norm(ctx->ggml_ctx, x);
-                q                = modulate(ctx->ggml_ctx, q, q_mods[0], q_mods[1]);
+                q                = LTXV::modulate(ctx->ggml_ctx, q, q_mods[0], q_mods[1]);
                 auto context_mod = context;
                 if (prompt_timestep != nullptr && prompt_table != nullptr) {
                     auto p_mods = get_ada_values(ctx, prompt_table, prompt_timestep, dim, 2);
-                    context_mod = modulate(ctx->ggml_ctx, context_mod, p_mods[0], p_mods[1]);
+                    context_mod = LTXV::modulate(ctx->ggml_ctx, context_mod, p_mods[0], p_mods[1]);
                 }
                 auto out = attn->forward(ctx, q, context_mod, attention_mask, nullptr, nullptr);
                 return apply_gate(ctx->ggml_ctx, out, q_mods[2]);
@@ -1228,7 +1228,7 @@ namespace LTXV {
 
             auto v_mods = get_ada_values(ctx, v_table, v_timestep, v_dim, cross_attention_adaln ? 9 : 6);
             auto v_norm = rms_norm(ctx->ggml_ctx, vx);
-            v_norm      = modulate(ctx->ggml_ctx, v_norm, v_mods[0], v_mods[1]);
+            v_norm      = LTXV::modulate(ctx->ggml_ctx, v_norm, v_mods[0], v_mods[1]);
             auto v_sa   = attn1->forward(ctx, v_norm, nullptr, self_attention_mask, v_pe);
             vx          = ggml_add(ctx->ggml_ctx, vx, apply_gate(ctx->ggml_ctx, v_sa, v_mods[2]));
             auto v_txt  = apply_text_cross_attention(ctx,
@@ -1246,7 +1246,7 @@ namespace LTXV {
             if (run_ax) {
                 auto a_mods = get_ada_values(ctx, a_table, a_timestep, a_dim, cross_attention_adaln ? 9 : 6);
                 auto a_norm = rms_norm(ctx->ggml_ctx, ax);
-                a_norm      = modulate(ctx->ggml_ctx, a_norm, a_mods[0], a_mods[1]);
+                a_norm      = LTXV::modulate(ctx->ggml_ctx, a_norm, a_mods[0], a_mods[1]);
                 auto a_sa   = audio_attn1->forward(ctx, a_norm, nullptr, nullptr, a_pe);
                 ax          = ggml_add(ctx->ggml_ctx, ax, apply_gate(ctx->ggml_ctx, a_sa, a_mods[2]));
                 auto a_txt  = apply_text_cross_attention(ctx,
@@ -1269,8 +1269,8 @@ namespace LTXV {
                     auto a2v_video_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 0, 4);
                     auto a2v_audio       = get_ada_values(ctx, a2v_audio_table, a_cross_scale_shift_timestep, a_dim, 4);
                     auto a2v_video       = get_ada_values(ctx, a2v_video_table, v_cross_scale_shift_timestep, v_dim, 4);
-                    auto vx_scaled       = modulate(ctx->ggml_ctx, vx_norm3, a2v_video[1], a2v_video[0]);
-                    auto ax_scaled       = modulate(ctx->ggml_ctx, ax_norm3, a2v_audio[1], a2v_audio[0]);
+                    auto vx_scaled       = LTXV::modulate(ctx->ggml_ctx, vx_norm3, a2v_video[1], a2v_video[0]);
+                    auto ax_scaled       = LTXV::modulate(ctx->ggml_ctx, ax_norm3, a2v_audio[1], a2v_audio[0]);
                     auto a2v_out         = audio_to_video_attn->forward(ctx, vx_scaled, ax_scaled, nullptr, v_cross_pe, a_cross_pe);
                     auto a2v_gate_table  = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 4, 5);
                     auto a2v_gate        = get_ada_values(ctx, a2v_gate_table, v_cross_gate_timestep, v_dim, 1)[0];
@@ -1282,8 +1282,8 @@ namespace LTXV {
                     auto v2a_video_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 0, 4);
                     auto v2a_audio       = get_ada_values(ctx, v2a_audio_table, a_cross_scale_shift_timestep, a_dim, 4);
                     auto v2a_video       = get_ada_values(ctx, v2a_video_table, v_cross_scale_shift_timestep, v_dim, 4);
-                    auto ax_scaled       = modulate(ctx->ggml_ctx, ax_norm3, v2a_audio[3], v2a_audio[2]);
-                    auto vx_scaled       = modulate(ctx->ggml_ctx, vx_norm3, v2a_video[3], v2a_video[2]);
+                    auto ax_scaled       = LTXV::modulate(ctx->ggml_ctx, ax_norm3, v2a_audio[3], v2a_audio[2]);
+                    auto vx_scaled       = LTXV::modulate(ctx->ggml_ctx, vx_norm3, v2a_video[3], v2a_video[2]);
                     auto v2a_out         = video_to_audio_attn->forward(ctx, ax_scaled, vx_scaled, nullptr, a_cross_pe, v_cross_pe);
                     auto v2a_gate_table  = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_audio"], 1, 4, 5);
                     auto v2a_gate        = get_ada_values(ctx, v2a_gate_table, a_cross_gate_timestep, a_dim, 1)[0];
@@ -1291,14 +1291,14 @@ namespace LTXV {
                 }
                 auto a_ff_mods = get_ada_values(ctx, a_table, a_timestep, a_dim, cross_attention_adaln ? 9 : 6, 3, 3);
                 auto ax_scaled = rms_norm(ctx->ggml_ctx, ax);
-                ax_scaled      = modulate(ctx->ggml_ctx, ax_scaled, a_ff_mods[0], a_ff_mods[1]);
+                ax_scaled      = LTXV::modulate(ctx->ggml_ctx, ax_scaled, a_ff_mods[0], a_ff_mods[1]);
                 auto a_ff_out  = audio_ff->forward(ctx, ax_scaled);
                 ax             = ggml_add(ctx->ggml_ctx, ax, apply_gate(ctx->ggml_ctx, a_ff_out, a_ff_mods[2]));
             }
 
             auto v_ff_mods = get_ada_values(ctx, v_table, v_timestep, v_dim, cross_attention_adaln ? 9 : 6, 3, 3);
             auto vx_scaled = rms_norm(ctx->ggml_ctx, vx);
-            vx_scaled      = modulate(ctx->ggml_ctx, vx_scaled, v_ff_mods[0], v_ff_mods[1]);
+            vx_scaled      = LTXV::modulate(ctx->ggml_ctx, vx_scaled, v_ff_mods[0], v_ff_mods[1]);
             auto v_ff_out  = ff->forward(ctx, vx_scaled);
             vx             = ggml_add(ctx->ggml_ctx, vx, apply_gate(ctx->ggml_ctx, v_ff_out, v_ff_mods[2]));
 
@@ -1657,14 +1657,14 @@ namespace LTXV {
 
             auto v_shift_scale = get_output_scale_shift(ctx, params["scale_shift_table"], v_embedded_time, config.hidden_size);
             vx                 = norm_out->forward(ctx, vx);
-            vx                 = modulate(ctx->ggml_ctx, vx, v_shift_scale[0], v_shift_scale[1]);
+            vx                 = LTXV::modulate(ctx->ggml_ctx, vx, v_shift_scale[0], v_shift_scale[1]);
             vx                 = proj_out->forward(ctx, vx);
             vx                 = unpatchify_video(ctx, vx, width, height, frames);
 
             if (ax != nullptr && audio_time > 0) {
                 auto a_shift_scale = get_output_scale_shift(ctx, params["audio_scale_shift_table"], a_embedded_time, config.audio_hidden_size);
                 ax                 = audio_norm_out->forward(ctx, ax);
-                ax                 = modulate(ctx->ggml_ctx, ax, a_shift_scale[0], a_shift_scale[1]);
+                ax                 = LTXV::modulate(ctx->ggml_ctx, ax, a_shift_scale[0], a_shift_scale[1]);
                 ax                 = audio_proj_out->forward(ctx, ax);
                 ax                 = unpatchify_audio(ctx, ax, audio_time);
             }

src/model/diffusion/mage_flow.hpp

@@ -0,0 +1,162 @@
+#ifndef __SD_MODEL_DIFFUSION_MAGE_FLOW_HPP__
+#define __SD_MODEL_DIFFUSION_MAGE_FLOW_HPP__
+
+#include <cmath>
+#include <memory>
+
+#include "model/diffusion/qwen_image.hpp"
+
+namespace MageFlow {
+    constexpr int MAGE_FLOW_GRAPH_SIZE = 20480;
+
+    // Mage-Flow was trained with BF16-rounded timestep frequencies; using Qwen's F32 projection degrades generation quality.
+    struct MageFlowTimestepProjEmbeddings : public Qwen::QwenTimestepProjEmbeddings {
+        static constexpr int TIMESTEP_DIM = 256;
+        static constexpr int HALF_DIM     = TIMESTEP_DIM / 2;
+
+        std::vector<float> frequencies;
+        std::vector<float> timesteps_proj;
+
+        explicit MageFlowTimestepProjEmbeddings(int64_t embedding_dim)
+            : QwenTimestepProjEmbeddings(embedding_dim), frequencies(HALF_DIM) {
+            for (int i = 0; i < HALF_DIM; ++i) {
+                float frequency = std::exp(-std::log(10000.f) * static_cast<float>(i) / HALF_DIM);
+                frequencies[i]  = ggml_bf16_to_fp32(ggml_fp32_to_bf16(frequency));
+            }
+        }
+
+        void prepare(const sd::Tensor<float>& timesteps) {
+            size_t num_timesteps = static_cast<size_t>(timesteps.numel());
+            timesteps_proj.resize(static_cast<size_t>(TIMESTEP_DIM) * num_timesteps);
+            for (size_t b = 0; b < num_timesteps; ++b) {
+                float sigma = ggml_bf16_to_fp32(ggml_fp32_to_bf16(timesteps.values()[b] / 1000.f));
+                for (int i = 0; i < HALF_DIM; ++i) {
+                    float argument = sigma * frequencies[i] * 1000.f;
+                    timesteps_proj[b * TIMESTEP_DIM + i] =
+                        ggml_bf16_to_fp32(ggml_fp32_to_bf16(std::cos(argument)));
+                    timesteps_proj[b * TIMESTEP_DIM + HALF_DIM + i] =
+                        ggml_bf16_to_fp32(ggml_fp32_to_bf16(std::sin(argument)));
+                }
+            }
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* timesteps,
+                             ggml_tensor* addition_t_cond = nullptr) override {
+            GGML_ASSERT(addition_t_cond == nullptr);
+            GGML_ASSERT(timesteps_proj.size() ==
+                        static_cast<size_t>(TIMESTEP_DIM * ggml_nelements(timesteps)));
+            auto projection = ggml_new_tensor_2d(ctx->ggml_ctx,
+                                                 GGML_TYPE_F32,
+                                                 TIMESTEP_DIM,
+                                                 ggml_nelements(timesteps));
+            ctx->bind_backend_tensor_data(projection, timesteps_proj.data());
+            auto timestep_embedder = std::dynamic_pointer_cast<Qwen::TimestepEmbedding>(blocks["timestep_embedder"]);
+            return timestep_embedder->forward(ctx, projection);
+        }
+    };
+
+    struct MageFlowRunner : public DiffusionModelRunner {
+    public:
+        Qwen::QwenImageConfig config;
+        Qwen::QwenImageModel mage_flow;
+        std::shared_ptr<MageFlowTimestepProjEmbeddings> time_text_embed;
+        std::vector<float> pe_vec;
+
+        MageFlowRunner(ggml_backend_t backend,
+                       const String2TensorStorage& tensor_storage_map      = {},
+                       const std::string prefix                            = "",
+                       std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : DiffusionModelRunner(backend, prefix, weight_manager) {
+            config.patch_size          = 1;
+            config.in_channels         = 128;
+            config.out_channels        = 128;
+            config.num_layers          = 12;
+            config.attention_head_dim  = 128;
+            config.num_attention_heads = 24;
+            config.joint_attention_dim = 2560;
+            config.theta               = 10000;
+            config.axes_dim            = {16, 56, 56};
+            config.axes_dim_sum        = 128;
+            time_text_embed            = std::make_shared<MageFlowTimestepProjEmbeddings>(
+                config.num_attention_heads * config.attention_head_dim);
+            mage_flow = Qwen::QwenImageModel(config, time_text_embed);
+            mage_flow.init(params_ctx, tensor_storage_map, prefix);
+        }
+
+        std::string get_desc() override {
+            return "mage_flow";
+        }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string& prefix) override {
+            mage_flow.get_param_tensors(tensors, prefix);
+        }
+
+        ggml_cgraph* build_graph(const sd::Tensor<float>& x_tensor,
+                                 const sd::Tensor<float>& timesteps_tensor,
+                                 const sd::Tensor<float>& context_tensor,
+                                 const std::vector<sd::Tensor<float>>& ref_latents_tensor = {}) {
+            ggml_cgraph* gf        = new_graph_custom(MAGE_FLOW_GRAPH_SIZE);
+            ggml_tensor* x         = make_input(x_tensor);
+            ggml_tensor* timesteps = make_input(timesteps_tensor);
+            GGML_ASSERT(x->ne[3] == 1);
+            GGML_ASSERT(!context_tensor.empty());
+            ggml_tensor* context = make_input(context_tensor);
+
+            std::vector<ggml_tensor*> ref_latents;
+            ref_latents.reserve(ref_latents_tensor.size());
+            for (const auto& ref_latent_tensor : ref_latents_tensor) {
+                ref_latents.push_back(make_input(ref_latent_tensor));
+            }
+
+            int batch_size = static_cast<int>(x->ne[3]);
+            pe_vec         = Rope::gen_mage_flow_pe(static_cast<int>(x->ne[1]),
+                                                    static_cast<int>(x->ne[0]),
+                                                    batch_size,
+                                                    static_cast<int>(context->ne[1]),
+                                                    ref_latents,
+                                                    config.theta,
+                                                    config.axes_dim);
+            int pos_len    = static_cast<int>(pe_vec.size() / config.axes_dim_sum / 2);
+            auto pe        = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len);
+            set_backend_tensor_data(pe, pe_vec.data());
+
+            time_text_embed->prepare(timesteps_tensor);
+            auto runner_ctx = get_context();
+            auto out        = mage_flow.forward(&runner_ctx,
+                                                x,
+                                                timesteps,
+                                                nullptr,
+                                                context,
+                                                pe,
+                                                ref_latents);
+            ggml_build_forward_expand(gf, out);
+            return gf;
+        }
+
+        sd::Tensor<float> compute(int n_threads,
+                                  const sd::Tensor<float>& x,
+                                  const sd::Tensor<float>& timesteps,
+                                  const sd::Tensor<float>& context,
+                                  const std::vector<sd::Tensor<float>>& ref_latents = {}) {
+            auto get_graph = [&]() -> ggml_cgraph* {
+                return build_graph(x, timesteps, context, ref_latents);
+            };
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
+        }
+
+        sd::Tensor<float> compute(int n_threads,
+                                  const DiffusionParams& diffusion_params) override {
+            GGML_ASSERT(diffusion_params.x != nullptr);
+            GGML_ASSERT(diffusion_params.timesteps != nullptr);
+            static const std::vector<sd::Tensor<float>> empty_ref_latents;
+            return compute(n_threads,
+                           *diffusion_params.x,
+                           *diffusion_params.timesteps,
+                           tensor_or_empty(diffusion_params.context),
+                           diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents);
+        }
+    };
+}  // namespace MageFlow
+
+#endif  // __SD_MODEL_DIFFUSION_MAGE_FLOW_HPP__

src/model/diffusion/mmdit.hpp

@@ -136,11 +136,15 @@ struct MMDiTConfig {
 };
 
 struct PatchEmbed : public GGMLBlock {
-    // 2D Image to Patch Embedding
+    // 2D/3D Image to Patch Embedding
 protected:
+    bool is_3d;
     bool flatten;
     bool dynamic_img_pad;
-    int patch_size;
+    int patch_t;
+    int patch_h;
+    int patch_w;
+    int64_t embed_dim;
 
 public:
     PatchEmbed(int64_t img_size     = 224,
@@ -149,42 +153,90 @@ public:
                int64_t embed_dim    = 1536,
                bool bias            = true,
                bool flatten         = true,
-               bool dynamic_img_pad = true)
-        : patch_size(patch_size),
+               bool dynamic_img_pad = true,
+               bool is_3d           = false)
+        : patch_t(is_3d ? patch_size : 1),
+          patch_h(patch_size),
+          patch_w(patch_size),
+          embed_dim(embed_dim),
           flatten(flatten),
-          dynamic_img_pad(dynamic_img_pad) {
+          dynamic_img_pad(dynamic_img_pad),
+          is_3d(is_3d) {
         // img_size is always None
         // patch_size is always 2
         // in_chans is always 16
         // norm_layer is always False
         // strict_img_size is always true, but not used
 
-        blocks["proj"] = std::shared_ptr<GGMLBlock>(new Conv2d(in_chans,
-                                                               embed_dim,
-                                                               {patch_size, patch_size},
-                                                               {patch_size, patch_size},
-                                                               {0, 0},
-                                                               {1, 1},
-                                                               bias));
+        if (is_3d) {
+            blocks["proj"] = std::make_shared<Conv3d>(in_chans,
+                                                      embed_dim,
+                                                      std::tuple{patch_size, patch_size, patch_size},
+                                                      std::tuple{patch_size, patch_size, patch_size},
+                                                      std::tuple{0, 0, 0},
+                                                      std::tuple{1, 1, 1},
+                                                      bias);
+        } else {
+            blocks["proj"] = std::make_shared<Conv2d>(in_chans,
+                                                      embed_dim,
+                                                      std::pair{patch_size, patch_size},
+                                                      std::pair{patch_size, patch_size},
+                                                      std::pair{0, 0},
+                                                      std::pair{1, 1},
+                                                      bias);
+        }
+    }
+
+    PatchEmbed(int64_t img_size,
+               std::tuple<int, int, int> patch_size,
+               int64_t in_chans,
+               int64_t embed_dim,
+               bool bias            = true,
+               bool flatten         = true,
+               bool dynamic_img_pad = true)
+        : patch_t(std::get<0>(patch_size)),
+          patch_h(std::get<1>(patch_size)),
+          patch_w(std::get<2>(patch_size)),
+          embed_dim(embed_dim),
+          flatten(flatten),
+          dynamic_img_pad(dynamic_img_pad),
+          is_3d(true) {
+        SD_UNUSED(img_size);
+        blocks["proj"] = std::make_shared<Conv3d>(in_chans,
+                                                  embed_dim,
+                                                  patch_size,
+                                                  patch_size,
+                                                  std::tuple{0, 0, 0},
+                                                  std::tuple{1, 1, 1},
+                                                  bias);
     }
 
     ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
-        // x: [N, C, H, W]
-        // return: [N, H*W, embed_dim]
-        auto proj = std::dynamic_pointer_cast<Conv2d>(blocks["proj"]);
+        // x: [N, C, H, W] or [N*C, T, H, W]
+        // return: [N, h_len*w_len, embed_dim] or [N, t_len*h_len*w_len, embed_dim]
+        auto proj = std::dynamic_pointer_cast<UnaryBlock>(blocks["proj"]);
 
         if (dynamic_img_pad) {
             int64_t W = x->ne[0];
             int64_t H = x->ne[1];
-            int pad_h = (patch_size - H % patch_size) % patch_size;
-            int pad_w = (patch_size - W % patch_size) % patch_size;
-            x         = ggml_pad(ctx->ggml_ctx, x, pad_w, pad_h, 0, 0);  // TODO: reflect pad mode
+            int pad_t = 0;
+            int pad_h = (patch_h - static_cast<int>(H % patch_h)) % patch_h;
+            int pad_w = (patch_w - static_cast<int>(W % patch_w)) % patch_w;
+            if (is_3d) {
+                int64_t T = x->ne[2];
+                pad_t     = (patch_t - static_cast<int>(T % patch_t)) % patch_t;
+            }
+            x = ggml_pad(ctx->ggml_ctx, x, pad_w, pad_h, pad_t, 0);  // TODO: reflect pad mode
         }
-        x = proj->forward(ctx, x);
+        x = proj->forward(ctx, x);  // [N, C, h_len, w_len] or [N*C, t_len, h_len, w_len]
 
         if (flatten) {
-            x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1], x->ne[2], x->ne[3]);
-            x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
+            if (is_3d) {
+                x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1] * x->ne[2], embed_dim, x->ne[3] / embed_dim);  // [N, C, t_len*h_len*w_len]
+            } else {
+                x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1], x->ne[2], x->ne[3]);  // [N, C, h_len*w_len]
+            }
+            x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));  // [N, h_len*w_len, C]
         }
         return x;
     }

src/model/diffusion/model.hpp

@@ -26,6 +26,7 @@ struct RefImageParams {
     RefImageResizeMode vlm_resize_mode = RefImageResizeMode::AREA;
     int vlm_min_size                   = -1;
     int vlm_max_size                   = -1;
+    bool resize_vae_to_target          = false;
 };
 
 const std::unordered_map<std::string, RefImageParams> REF_IMAGE_PRESETS = {
@@ -34,6 +35,7 @@ const std::unordered_map<std::string, RefImageParams> REF_IMAGE_PRESETS = {
     {"flux2", {false, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::NONE, -1, -1}},
     {"qwen", {true, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
     {"qwen_layered", {true, true, Rope::RefIndexMode::DECREASE, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
+    {"mage_flow", {true, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::LONGEST_SIDE, -1, 384, true}},
     {"z_image_omni", {true, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::AREA, -1, -1}},
     {"krea2_ostris_edit", {true, true, Rope::RefIndexMode::INCREASE, true, true, -1, RefImageResizeMode::AREA, -1, -1}},
     {"krea2_edit", {true, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::LONGEST_SIDE, 768, 768}},
@@ -44,6 +46,8 @@ struct UNetDiffusionExtra {
     int num_video_frames                           = -1;
     const std::vector<sd::Tensor<float>>* controls = nullptr;
     float control_strength                         = 0.f;
+    const sd::Tensor<float>* ip_context            = nullptr;
+    float ip_scale                                 = 1.f;
 };
 
 struct SkipLayerDiffusionExtra {
@@ -87,6 +91,13 @@ struct MiniT2IDiffusionExtra {
     const sd::Tensor<float>* mask = nullptr;
 };
 
+struct HunyuanVideoDiffusionExtra {
+    const sd::Tensor<float>* guidance   = nullptr;
+    const sd::Tensor<float>* byt5       = nullptr;
+    const sd::Tensor<float>* vision     = nullptr;
+    const sd::Tensor<float>* timestep_r = nullptr;
+};
+
 using DiffusionExtraParams = std::variant<std::monostate,
                                           UNetDiffusionExtra,
                                           SkipLayerDiffusionExtra,
@@ -95,7 +106,8 @@ using DiffusionExtraParams = std::variant<std::monostate,
                                           WanDiffusionExtra,
                                           HiDreamO1DiffusionExtra,
                                           LTXAVDiffusionExtra,
-                                          MiniT2IDiffusionExtra>;
+                                          MiniT2IDiffusionExtra,
+                                          HunyuanVideoDiffusionExtra>;
 
 struct DiffusionParams {
     const sd::Tensor<float>* x                        = nullptr;

src/model/diffusion/pid.hpp

@@ -17,30 +17,38 @@ namespace Pid {
     constexpr float PID_PI       = 3.14159265358979323846f;
 
     struct PixelDiTConfig {
-        int64_t in_channels            = 3;
-        int64_t hidden_size            = 1536;
-        int64_t num_groups             = 24;
-        int64_t patch_mlp_hidden_dim   = 4096;
-        int64_t pixel_hidden_size      = 16;
-        int64_t pixel_attn_hidden_size = 1152;
-        int64_t pixel_num_groups       = 16;
-        int64_t patch_depth            = 14;
-        int64_t pixel_depth            = 2;
-        int64_t patch_size             = 16;
-        int64_t txt_embed_dim          = 2304;
-        int64_t txt_max_length         = 300;
-        float text_rope_theta          = 10000.f;
-        int64_t lq_latent_channels     = 16;
-        int64_t lq_hidden_dim          = 512;
-        int64_t lq_num_res_blocks      = 4;
-        int64_t lq_interval            = 2;
-        int64_t lq_sr_scale            = 4;
-        int64_t lq_latent_down_factor  = 8;
-        int64_t rope_ref_grid_h        = 64;
-        int64_t rope_ref_grid_w        = 64;
+        int64_t in_channels                 = 3;
+        int64_t hidden_size                 = 1536;
+        int64_t num_groups                  = 24;
+        int64_t patch_mlp_hidden_dim        = 4096;
+        int64_t pixel_hidden_size           = 16;
+        int64_t pixel_attn_hidden_size      = 1152;
+        int64_t pixel_num_groups            = 16;
+        int64_t patch_depth                 = 14;
+        int64_t pixel_depth                 = 2;
+        int64_t patch_size                  = 16;
+        int64_t txt_embed_dim               = 2304;
+        int64_t txt_max_length              = 300;
+        float text_rope_theta               = 10000.f;
+        int64_t lq_latent_channels          = 16;
+        int64_t lq_hidden_dim               = 512;
+        int64_t lq_num_res_blocks           = 4;
+        int64_t lq_interval                 = 2;
+        int64_t lq_sr_scale                 = 4;
+        int64_t lq_latent_down_factor       = 8;
+        int64_t lq_latent_unpatchify_factor = 1;
+        bool lq_replicate_padding           = false;
+        bool lq_gate_per_token              = false;
+        bool pit_lq_inject                  = false;
+        int64_t rope_ref_grid_h             = 64;
+        int64_t rope_ref_grid_w             = 64;
 
         static PixelDiTConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, const std::string& prefix) {
             PixelDiTConfig config;
+            int64_t latent_proj_in_channels = config.lq_latent_channels;
+            int64_t num_lq_gates            = 0;
+            const std::string lq_prefix     = prefix + ".lq_proj.";
+            config.pit_lq_inject            = tensor_storage_map.find(lq_prefix + "pit_head.weight") != tensor_storage_map.end();
             for (const auto& [name, tensor_storage] : tensor_storage_map) {
                 if (!starts_with(name, prefix)) {
                     continue;
@@ -61,20 +69,56 @@ namespace Pid {
                         config.pixel_depth = std::max<int64_t>(config.pixel_depth, block_index + 1);
                     }
                 }
-                if (name.find("lq_proj.latent_proj.0.weight") != std::string::npos) {
-                    config.lq_latent_channels    = tensor_storage.ne[2];
-                    config.lq_latent_down_factor = config.lq_latent_channels >= 64 ? 16 : 8;
+                if (name == lq_prefix + "latent_proj.0.weight") {
+                    latent_proj_in_channels = tensor_storage.ne[2];
+                    config.lq_hidden_dim    = tensor_storage.ne[3];
+                }
+                if (starts_with(name, lq_prefix + "gate_modules.")) {
+                    auto items = split_string(name.substr(lq_prefix.size()), '.');
+                    if (items.size() > 1) {
+                        int gate_index = atoi(items[1].c_str());
+                        num_lq_gates   = std::max<int64_t>(num_lq_gates, gate_index + 1);
+                    }
                 }
                 if (name.find("patch_blocks.0.mlp_x.w1.weight") != std::string::npos) {
                     config.patch_mlp_hidden_dim = tensor_storage.ne[1];
                 }
             }
-            LOG_DEBUG("pid: patch_depth = %" PRId64 ", pixel_depth = %" PRId64 ", patch_mlp_hidden_dim = %" PRId64 ", lq_latent_channels = %" PRId64 ", lq_latent_down_factor = %" PRId64,
+            if (num_lq_gates > 0) {
+                config.lq_interval = (config.patch_depth + num_lq_gates - 1) / num_lq_gates;
+            }
+            if (config.pit_lq_inject) {
+                if (latent_proj_in_channels == 16) {
+                    config.lq_latent_channels          = 16;
+                    config.lq_latent_down_factor       = 8;
+                    config.lq_latent_unpatchify_factor = 1;
+                } else {
+                    GGML_ASSERT(latent_proj_in_channels == 32);
+                    config.lq_latent_channels          = 128;
+                    config.lq_latent_down_factor       = 16;
+                    config.lq_latent_unpatchify_factor = 2;
+                }
+                auto gate_weight = tensor_storage_map.find(lq_prefix + "gate_modules.0.content_proj.weight");
+                if (gate_weight != tensor_storage_map.end()) {
+                    config.lq_gate_per_token = gate_weight->second.ne[1] == 1;
+                }
+                config.lq_replicate_padding = true;
+                config.rope_ref_grid_h      = 128;
+                config.rope_ref_grid_w      = 128;
+            } else {
+                config.lq_latent_channels    = latent_proj_in_channels;
+                config.lq_latent_down_factor = latent_proj_in_channels >= 64 ? 16 : 8;
+            }
+            LOG_DEBUG("pid: version = %s, patch_depth = %" PRId64 ", pixel_depth = %" PRId64 ", patch_mlp_hidden_dim = %" PRId64 ", lq_latent_channels = %" PRId64 ", lq_hidden_dim = %" PRId64 ", lq_latent_down_factor = %" PRId64 ", lq_latent_unpatchify_factor = %" PRId64 ", lq_interval = %" PRId64,
+                      config.pit_lq_inject ? "1.5" : "1",
                       config.patch_depth,
                       config.pixel_depth,
                       config.patch_mlp_hidden_dim,
                       config.lq_latent_channels,
-                      config.lq_latent_down_factor);
+                      config.lq_hidden_dim,
+                      config.lq_latent_down_factor,
+                      config.lq_latent_unpatchify_factor,
+                      config.lq_interval);
             return config;
         }
     };
@@ -135,6 +179,18 @@ namespace Pid {
         return ggml_add(ctx, ggml_add(ctx, x, ggml_mul(ctx, x, scale)), shift);
     }
 
+    inline ggml_tensor* replicate_pad_2d(ggml_context* ctx, ggml_tensor* x) {
+        auto left  = ggml_ext_slice(ctx, x, 0, 0, 1);
+        auto right = ggml_ext_slice(ctx, x, 0, x->ne[0] - 1, x->ne[0]);
+        x          = ggml_concat(ctx, left, x, 0);
+        x          = ggml_concat(ctx, x, right, 0);
+
+        auto top    = ggml_ext_slice(ctx, x, 1, 0, 1);
+        auto bottom = ggml_ext_slice(ctx, x, 1, x->ne[1] - 1, x->ne[1]);
+        x           = ggml_concat(ctx, top, x, 1);
+        return ggml_concat(ctx, x, bottom, 1);
+    }
+
     struct PatchTokenEmbedder : public GGMLBlock {
         bool use_rms_norm;
 
@@ -457,9 +513,9 @@ namespace Pid {
     struct SigmaAwareGate : public GGMLBlock {
         int64_t dim;
 
-        SigmaAwareGate(int64_t dim)
+        SigmaAwareGate(int64_t dim, bool per_token = false)
             : dim(dim) {
-            blocks["content_proj"] = std::make_shared<Linear>(dim * 2, dim, true);
+            blocks["content_proj"] = std::make_shared<Linear>(dim * 2, per_token ? 1 : dim, true);
         }
 
         void init_params(ggml_context* ctx,
@@ -479,16 +535,20 @@ namespace Pid {
             auto alpha         = ggml_exp(ctx->ggml_ctx, params["log_alpha"]);
             auto offset        = ggml_neg(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, alpha, sigma));
             auto gate          = ggml_sigmoid(ctx->ggml_ctx, ggml_add(ctx->ggml_ctx, content_logit, offset));
-            return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, gate, lq));
+            return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, lq, gate));
         }
     };
 
     struct PiDResBlock : public GGMLBlock {
-        PiDResBlock(int64_t channels) {
-            blocks["block.0"] = std::make_shared<GroupNorm>(4, channels, 1e-5f);
-            blocks["block.2"] = std::make_shared<Conv2d>(channels, channels, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, std::pair<int, int>{1, 1});
-            blocks["block.3"] = std::make_shared<GroupNorm>(4, channels, 1e-5f);
-            blocks["block.5"] = std::make_shared<Conv2d>(channels, channels, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, std::pair<int, int>{1, 1});
+        bool replicate_padding;
+
+        PiDResBlock(int64_t channels, bool replicate_padding = false)
+            : replicate_padding(replicate_padding) {
+            std::pair<int, int> padding = replicate_padding ? std::pair<int, int>{0, 0} : std::pair<int, int>{1, 1};
+            blocks["block.0"]           = std::make_shared<GroupNorm>(4, channels, 1e-5f);
+            blocks["block.2"]           = std::make_shared<Conv2d>(channels, channels, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, padding);
+            blocks["block.3"]           = std::make_shared<GroupNorm>(4, channels, 1e-5f);
+            blocks["block.5"]           = std::make_shared<Conv2d>(channels, channels, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, padding);
         }
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
@@ -497,9 +557,15 @@ namespace Pid {
             auto norm2 = std::dynamic_pointer_cast<GroupNorm>(blocks["block.3"]);
             auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["block.5"]);
             auto h     = ggml_silu_inplace(ctx->ggml_ctx, norm1->forward(ctx, x));
-            h          = conv1->forward(ctx, h);
-            h          = ggml_silu_inplace(ctx->ggml_ctx, norm2->forward(ctx, h));
-            h          = conv2->forward(ctx, h);
+            if (replicate_padding) {
+                h = replicate_pad_2d(ctx->ggml_ctx, h);
+            }
+            h = conv1->forward(ctx, h);
+            h = ggml_silu_inplace(ctx->ggml_ctx, norm2->forward(ctx, h));
+            if (replicate_padding) {
+                h = replicate_pad_2d(ctx->ggml_ctx, h);
+            }
+            h = conv2->forward(ctx, h);
             return ggml_add(ctx->ggml_ctx, x, h);
         }
     };
@@ -509,16 +575,23 @@ namespace Pid {
 
         LQProjection2D(const PixelDiTConfig& config)
             : config(config) {
-            blocks["latent_proj.0"] = std::make_shared<Conv2d>(config.lq_latent_channels, config.lq_hidden_dim, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, std::pair<int, int>{1, 1});
-            blocks["latent_proj.2"] = std::make_shared<Conv2d>(config.lq_hidden_dim, config.lq_hidden_dim, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, std::pair<int, int>{1, 1});
+            int64_t unpatchify_area = config.lq_latent_unpatchify_factor * config.lq_latent_unpatchify_factor;
+            GGML_ASSERT(config.lq_latent_channels % unpatchify_area == 0);
+            int64_t latent_proj_in_channels = config.lq_latent_channels / unpatchify_area;
+            std::pair<int, int> padding     = config.lq_replicate_padding ? std::pair<int, int>{0, 0} : std::pair<int, int>{1, 1};
+            blocks["latent_proj.0"]         = std::make_shared<Conv2d>(latent_proj_in_channels, config.lq_hidden_dim, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, padding);
+            blocks["latent_proj.2"]         = std::make_shared<Conv2d>(config.lq_hidden_dim, config.lq_hidden_dim, std::pair<int, int>{3, 3}, std::pair<int, int>{1, 1}, padding);
             for (int i = 0; i < config.lq_num_res_blocks; ++i) {
-                blocks["latent_proj." + std::to_string(3 + i)] = std::make_shared<PiDResBlock>(config.lq_hidden_dim);
+                blocks["latent_proj." + std::to_string(3 + i)] = std::make_shared<PiDResBlock>(config.lq_hidden_dim, config.lq_replicate_padding);
             }
 
             int num_outputs = static_cast<int>((config.patch_depth + config.lq_interval - 1) / config.lq_interval);
             for (int i = 0; i < num_outputs; ++i) {
                 blocks["output_heads." + std::to_string(i)] = std::make_shared<Linear>(config.lq_hidden_dim, config.hidden_size, true);
-                blocks["gate_modules." + std::to_string(i)] = std::make_shared<SigmaAwareGate>(config.hidden_size);
+                blocks["gate_modules." + std::to_string(i)] = std::make_shared<SigmaAwareGate>(config.hidden_size, config.lq_gate_per_token);
+            }
+            if (config.pit_lq_inject) {
+                blocks["pit_head"] = std::make_shared<Linear>(config.lq_hidden_dim, config.hidden_size, true);
             }
         }
 
@@ -543,9 +616,29 @@ namespace Pid {
                                           ggml_tensor* lq_latent,
                                           int64_t target_pH,
                                           int64_t target_pW) {
-            auto conv0             = std::dynamic_pointer_cast<Conv2d>(blocks["latent_proj.0"]);
-            auto conv2             = std::dynamic_pointer_cast<Conv2d>(blocks["latent_proj.2"]);
-            float z_to_patch_ratio = static_cast<float>(config.lq_sr_scale * config.lq_latent_down_factor) /
+            auto conv0                = std::dynamic_pointer_cast<Conv2d>(blocks["latent_proj.0"]);
+            auto conv2                = std::dynamic_pointer_cast<Conv2d>(blocks["latent_proj.2"]);
+            int64_t unpatchify_factor = config.lq_latent_unpatchify_factor;
+            if (unpatchify_factor > 1) {
+                int64_t latent_h = lq_latent->ne[1];
+                int64_t latent_w = lq_latent->ne[0];
+                lq_latent        = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, lq_latent, 2, 0, 1, 3));
+                lq_latent        = ggml_reshape_3d(ctx->ggml_ctx,
+                                                   lq_latent,
+                                                   lq_latent->ne[0],
+                                                   lq_latent->ne[1] * lq_latent->ne[2],
+                                                   lq_latent->ne[3]);
+                lq_latent        = DiT::unpatchify(ctx->ggml_ctx,
+                                                   lq_latent,
+                                                   latent_h,
+                                                   latent_w,
+                                                   static_cast<int>(unpatchify_factor),
+                                                   static_cast<int>(unpatchify_factor),
+                                                   true);
+            }
+
+            int64_t effective_down_factor = config.lq_latent_down_factor / unpatchify_factor;
+            float z_to_patch_ratio        = static_cast<float>(config.lq_sr_scale * effective_down_factor) /
                                      static_cast<float>(config.patch_size);
             GGML_ASSERT(z_to_patch_ratio >= 1.0f);
             if (lq_latent->ne[0] != target_pW || lq_latent->ne[1] != target_pH) {
@@ -558,9 +651,15 @@ namespace Pid {
                                              GGML_SCALE_MODE_NEAREST);
             }
 
+            if (config.lq_replicate_padding) {
+                lq_latent = replicate_pad_2d(ctx->ggml_ctx, lq_latent);
+            }
             auto feat = conv0->forward(ctx, lq_latent);
             feat      = ggml_silu_inplace(ctx->ggml_ctx, feat);
-            feat      = conv2->forward(ctx, feat);
+            if (config.lq_replicate_padding) {
+                feat = replicate_pad_2d(ctx->ggml_ctx, feat);
+            }
+            feat = conv2->forward(ctx, feat);
             for (int i = 0; i < config.lq_num_res_blocks; ++i) {
                 auto block = std::dynamic_pointer_cast<PiDResBlock>(blocks["latent_proj." + std::to_string(3 + i)]);
                 feat       = block->forward(ctx, feat);
@@ -574,11 +673,15 @@ namespace Pid {
 
             int num_outputs = static_cast<int>((config.patch_depth + config.lq_interval - 1) / config.lq_interval);
             std::vector<ggml_tensor*> outputs;
-            outputs.reserve(num_outputs);
+            outputs.reserve(num_outputs + (config.pit_lq_inject ? 1 : 0));
             for (int i = 0; i < num_outputs; ++i) {
                 auto head = std::dynamic_pointer_cast<Linear>(blocks["output_heads." + std::to_string(i)]);
                 outputs.push_back(head->forward(ctx, tokens));
             }
+            if (config.pit_lq_inject) {
+                auto pit_head = std::dynamic_pointer_cast<Linear>(blocks["pit_head"]);
+                outputs.push_back(pit_head->forward(ctx, tokens));
+            }
             return outputs;
         }
     };
@@ -606,6 +709,9 @@ namespace Pid {
             }
             blocks["final_layer"] = std::make_shared<FinalLayer>(config.pixel_hidden_size, config.in_channels);
             blocks["lq_proj"]     = std::make_shared<LQProjection2D>(config);
+            if (config.pit_lq_inject) {
+                blocks["pit_lq_gate"] = std::make_shared<SigmaAwareGate>(config.hidden_size, config.lq_gate_per_token);
+            }
         }
 
         void init_params(ggml_context* ctx,
@@ -654,6 +760,11 @@ namespace Pid {
             y_emb        = ggml_add(ctx->ggml_ctx, y_emb, y_pos);
 
             std::vector<ggml_tensor*> lq_features = lq_proj->forward(ctx, lq_latent, Hs, Ws);
+            ggml_tensor* pit_lq_feature           = nullptr;
+            if (config.pit_lq_inject) {
+                pit_lq_feature = lq_features.back();
+                lq_features.pop_back();
+            }
 
             auto s = s_embedder->forward(ctx, x_patches);
 
@@ -677,6 +788,10 @@ namespace Pid {
                 sd::ggml_graph_cut::mark_graph_cut(y_emb, "pid.patch_blocks." + std::to_string(i), "y");
             }
             s = ggml_silu(ctx->ggml_ctx, ggml_add(ctx->ggml_ctx, s, t_emb));
+            if (pit_lq_feature != nullptr) {
+                auto pit_lq_gate = std::dynamic_pointer_cast<SigmaAwareGate>(blocks["pit_lq_gate"]);
+                s                = pit_lq_gate->forward(ctx, s, pit_lq_feature, degrade_sigma);
+            }
 
             auto s_cond = ggml_reshape_2d(ctx->ggml_ctx, s, config.hidden_size, L * B);
             auto pixels = pixel_embedder->forward(ctx, x, config.patch_size, pixel_pos_full);

src/model/diffusion/qwen_image.hpp

@@ -103,9 +103,9 @@ namespace Qwen {
             }
         }
 
-        ggml_tensor* forward(GGMLRunnerContext* ctx,
-                             ggml_tensor* timesteps,
-                             ggml_tensor* addition_t_cond = nullptr) {
+        virtual ggml_tensor* forward(GGMLRunnerContext* ctx,
+                                     ggml_tensor* timesteps,
+                                     ggml_tensor* addition_t_cond = nullptr) {
             // timesteps: [N,]
             // return: [N, embedding_dim]
             auto timestep_embedder = std::dynamic_pointer_cast<TimestepEmbedding>(blocks["timestep_embedder"]);
@@ -416,10 +416,14 @@ namespace Qwen {
 
     public:
         QwenImageModel() {}
-        QwenImageModel(QwenImageConfig config)
+        QwenImageModel(QwenImageConfig config,
+                       std::shared_ptr<QwenTimestepProjEmbeddings> time_text_embed = nullptr)
             : config(config) {
-            int64_t inner_dim         = config.num_attention_heads * config.attention_head_dim;
-            blocks["time_text_embed"] = std::shared_ptr<GGMLBlock>(new QwenTimestepProjEmbeddings(inner_dim, config.use_additional_t_cond));
+            int64_t inner_dim = config.num_attention_heads * config.attention_head_dim;
+            if (time_text_embed == nullptr) {
+                time_text_embed = std::make_shared<QwenTimestepProjEmbeddings>(inner_dim, config.use_additional_t_cond);
+            }
+            blocks["time_text_embed"] = std::move(time_text_embed);
             blocks["txt_norm"]        = std::shared_ptr<GGMLBlock>(new RMSNorm(config.joint_attention_dim, 1e-6f));
             blocks["img_in"]          = std::shared_ptr<GGMLBlock>(new Linear(config.in_channels, inner_dim));
             blocks["txt_in"]          = std::shared_ptr<GGMLBlock>(new Linear(config.joint_attention_dim, inner_dim));

src/model/diffusion/unet.hpp

@@ -775,14 +775,17 @@ struct UNetModelRunner : public DiffusionModelRunner {
                              const sd::Tensor<float>& y_tensor                     = {},
                              int num_video_frames                                  = -1,
                              const std::vector<sd::Tensor<float>>& controls_tensor = {},
-                             float control_strength                                = 0.f) {
+                             float control_strength                                = 0.f,
+                             const sd::Tensor<float>& ip_context_tensor            = {},
+                             float ip_scale                                        = 1.f) {
         ggml_cgraph* gf = new_graph_custom(UNET_GRAPH_SIZE);
 
-        ggml_tensor* x         = make_input(x_tensor);
-        ggml_tensor* timesteps = make_input(timesteps_tensor);
-        ggml_tensor* context   = make_optional_input(context_tensor);
-        ggml_tensor* c_concat  = make_optional_input(c_concat_tensor);
-        ggml_tensor* y         = make_optional_input(y_tensor);
+        ggml_tensor* x          = make_input(x_tensor);
+        ggml_tensor* timesteps  = make_input(timesteps_tensor);
+        ggml_tensor* context    = make_optional_input(context_tensor);
+        ggml_tensor* c_concat   = make_optional_input(c_concat_tensor);
+        ggml_tensor* y          = make_optional_input(y_tensor);
+        ggml_tensor* ip_context = make_optional_input(ip_context_tensor);
         std::vector<ggml_tensor*> controls;
         controls.reserve(controls_tensor.size());
         for (const auto& control_tensor : controls_tensor) {
@@ -793,7 +796,9 @@ struct UNetModelRunner : public DiffusionModelRunner {
             num_video_frames = static_cast<int>(x->ne[3]);
         }
 
-        auto runner_ctx = get_context();
+        auto runner_ctx       = get_context();
+        runner_ctx.ip_context = ip_context;
+        runner_ctx.ip_scale   = ip_scale;
 
         ggml_tensor* out = unet.forward(&runner_ctx,
                                         x,
@@ -818,14 +823,16 @@ struct UNetModelRunner : public DiffusionModelRunner {
                               const sd::Tensor<float>& y                     = {},
                               int num_video_frames                           = -1,
                               const std::vector<sd::Tensor<float>>& controls = {},
-                              float control_strength                         = 0.f) {
+                              float control_strength                         = 0.f,
+                              const sd::Tensor<float>& ip_context            = {},
+                              float ip_scale                                 = 1.f) {
         // x: [N, in_channels, h, w]
         // timesteps: [N, ]
         // context: [N, max_position, hidden_size]([N, 77, 768]) or [1, max_position, hidden_size]
         // c_concat: [N, in_channels, h, w] or [1, in_channels, h, w]
         // y: [N, adm_in_channels] or [1, adm_in_channels]
         auto get_graph = [&]() -> ggml_cgraph* {
-            return build_graph(x, timesteps, context, c_concat, y, num_video_frames, controls, control_strength);
+            return build_graph(x, timesteps, context, c_concat, y, num_video_frames, controls, control_strength, ip_context, ip_scale);
         };
 
         return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), x.dim());
@@ -845,7 +852,9 @@ struct UNetModelRunner : public DiffusionModelRunner {
                        tensor_or_empty(diffusion_params.y),
                        extra->num_video_frames,
                        extra->controls ? *extra->controls : empty_controls,
-                       extra->control_strength);
+                       extra->control_strength,
+                       extra->ip_context ? *extra->ip_context : sd::Tensor<float>{},
+                       extra->ip_scale);
     }
 
     void test() {

src/model/te/llm.hpp

@@ -200,7 +200,11 @@ namespace LLM {
                         config.vision.in_channels = tensor_storage.ne[2];
                         config.vision.hidden_size = tensor_storage.ne[3];
                     }
-                    if (contains(name, "visual.patch_embed.bias")) {
+                    // HF-format checkpoints keep the patch embed unsplit under a single name.
+                    if (contains(name, "visual.patch_embed.proj.weight")) {
+                        config.vision.patch_size = static_cast<int>(tensor_storage.ne[0]);
+                    }
+                    if (contains(name, "visual.patch_embed.bias") || contains(name, "visual.patch_embed.proj.bias")) {
                         config.vision.hidden_size = tensor_storage.ne[0];
                     }
                     if (contains(name, "visual.pos_embed.weight")) {
@@ -285,7 +289,7 @@ namespace LLM {
                    bool add_unit_offset = false)
             : hidden_size(hidden_size), eps(eps), add_unit_offset(add_unit_offset) {}
 
-        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
             ggml_tensor* w = params["weight"];
             if (ctx->weight_adapter) {
                 w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight");
@@ -1653,6 +1657,10 @@ namespace LLM {
             model.get_param_tensors(tensors, prefix);
         }
 
+        void get_param_tensor_ops(std::map<ggml_tensor*, enum ggml_op>& tensor_ops) {
+            model.get_param_tensor_ops(tensor_ops);
+        }
+
         ggml_tensor* forward(GGMLRunnerContext* ctx,
                              ggml_tensor* input_ids,
                              ggml_tensor* input_pos,

src/model/te/t5.hpp

@@ -18,6 +18,7 @@
 struct T5Config {
     int64_t num_layers      = 24;
     int64_t model_dim       = 4096;
+    int64_t inner_dim       = 4096;
     int64_t ff_dim          = 10240;
     int64_t num_heads       = 64;
     int64_t vocab_size      = 32128;
@@ -53,6 +54,7 @@ struct T5Config {
             if (q->n_dims == 2) {
                 config.model_dim  = q->ne[0];
                 int64_t inner_dim = q->ne[1];
+                config.inner_dim  = inner_dim;
                 // Flan-T5/T5 uses d_kv=64 for common sizes.
                 if (inner_dim % 64 == 0) {
                     config.num_heads = inner_dim / 64;
@@ -357,7 +359,7 @@ public:
         : config(config) {
         blocks["encoder"] = std::shared_ptr<GGMLBlock>(new T5Stack(config.num_layers,
                                                                    config.model_dim,
-                                                                   config.model_dim,
+                                                                   config.inner_dim,
                                                                    config.ff_dim,
                                                                    config.num_heads,
                                                                    config.relative_attention));

src/model/vae/hunyuan_vae.hpp

@@ -0,0 +1,834 @@
+#ifndef __SD_MODEL_VAE_HUNYUAN_VAE_HPP__
+#define __SD_MODEL_VAE_HUNYUAN_VAE_HPP__
+
+#include <algorithm>
+#include <cmath>
+#include <map>
+#include <memory>
+#include <string>
+#include <tuple>
+#include <utility>
+#include <vector>
+
+#include "model/vae/wan_vae.hpp"
+#include "model_manager.h"
+
+namespace Hunyuan {
+    constexpr int HUNYUAN_VIDEO_VAE_GRAPH_SIZE                  = 65536;
+    constexpr int HUNYUAN_VIDEO_VAE_GRAPH_SIZE_PER_LATENT_FRAME = 8192;
+    constexpr int HUNYUAN_VIDEO_VAE_TEMPORAL_CHUNK_SIZE         = 1;
+
+    struct TemporalConvCarry {
+        const std::vector<ggml_tensor*>* input = nullptr;
+        std::vector<ggml_tensor*>* output      = nullptr;
+        size_t input_index                     = 0;
+
+        bool is_continuation() const {
+            return input != nullptr;
+        }
+
+        ggml_tensor* take() {
+            GGML_ASSERT(input != nullptr && input_index < input->size());
+            return (*input)[input_index++];
+        }
+
+        void push(ggml_tensor* tensor) {
+            if (output != nullptr) {
+                output->push_back(tensor);
+            }
+        }
+
+        void finish() const {
+            GGML_ASSERT(input == nullptr || input_index == input->size());
+        }
+    };
+
+    static ggml_tensor* repeat_interleave_channels(GGMLRunnerContext* ctx,
+                                                   ggml_tensor* x,
+                                                   int64_t repeats,
+                                                   int64_t width,
+                                                   int64_t height,
+                                                   int64_t frames) {
+        GGML_ASSERT(repeats > 0);
+        GGML_ASSERT(width * height * frames == x->ne[0] * x->ne[1] * x->ne[2]);
+        int64_t channels = x->ne[3];
+        if (repeats == 1) {
+            return ggml_reshape_4d(ctx->ggml_ctx, x, width, height, frames, channels);
+        }
+        x           = ggml_reshape_3d(ctx->ggml_ctx, x, width * height * frames, 1, channels);
+        auto target = ggml_new_tensor_3d(ctx->ggml_ctx, x->type, width * height * frames, repeats, channels);
+        x           = ggml_repeat(ctx->ggml_ctx, x, target);
+        return ggml_reshape_4d(ctx->ggml_ctx, x, width, height, frames, channels * repeats);
+    }
+
+    class CausalConv3d : public GGMLBlock {
+    protected:
+        std::tuple<int, int, int> kernel_size;
+
+    public:
+        CausalConv3d(int64_t in_channels,
+                     int64_t out_channels,
+                     std::tuple<int, int, int> kernel_size,
+                     std::tuple<int, int, int> stride   = {1, 1, 1},
+                     std::tuple<int, int, int> padding  = {0, 0, 0},
+                     std::tuple<int, int, int> dilation = {1, 1, 1},
+                     bool bias                          = true)
+            : kernel_size(kernel_size) {
+            blocks["conv"] = std::make_shared<Conv3d>(in_channels, out_channels, kernel_size, stride, padding, dilation, bias);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             TemporalConvCarry* carry = nullptr) {
+            // x: [N*IC, ID, IH, IW]
+            // result: x: [N*OC, OD, OH, OW]
+            // assert N == 1
+
+            auto conv = std::dynamic_pointer_cast<Conv3d>(blocks["conv"]);
+
+            int pad_w = std::get<2>(kernel_size) / 2;
+            int pad_h = std::get<1>(kernel_size) / 2;
+            int pad_t = std::get<0>(kernel_size) - 1;
+            std::vector<ggml_tensor*> temporal_frames;
+            temporal_frames.reserve(x->ne[2] + pad_t);
+            if (pad_t > 0) {
+                if (carry != nullptr && carry->is_continuation()) {
+                    auto previous = carry->take();
+                    GGML_ASSERT(previous->ne[2] <= pad_t);
+                    for (int64_t frame = 0; frame < previous->ne[2]; frame++) {
+                        temporal_frames.push_back(ggml_ext_slice(ctx->ggml_ctx, previous, 2, frame, frame + 1));
+                    }
+                    for (int64_t frame = previous->ne[2]; frame < pad_t; frame++) {
+                        temporal_frames.push_back(ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1));
+                    }
+                } else {
+                    auto first = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1);
+                    for (int frame = 0; frame < pad_t; frame++) {
+                        temporal_frames.push_back(first);
+                    }
+                }
+            }
+            for (int64_t frame = 0; frame < x->ne[2]; frame++) {
+                temporal_frames.push_back(ggml_ext_slice(ctx->ggml_ctx, x, 2, frame, frame + 1));
+            }
+
+            if (pad_t > 0 && carry != nullptr && carry->output != nullptr) {
+                ggml_tensor* next = nullptr;
+                for (int frame = pad_t; frame > 0; frame--) {
+                    auto item = temporal_frames[temporal_frames.size() - frame];
+                    next      = next == nullptr ? item : ggml_concat(ctx->ggml_ctx, next, item, 2);
+                }
+                carry->push(ggml_cont(ctx->ggml_ctx, next));
+            }
+
+            ggml_tensor* padded = nullptr;
+            for (auto frame : temporal_frames) {
+                padded = padded == nullptr ? frame : ggml_concat(ctx->ggml_ctx, padded, frame, 2);
+            }
+            auto replicate_pad = [&](ggml_tensor* input, int dim, int left, int right) {
+                if (left > 0) {
+                    auto first = ggml_ext_slice(ctx->ggml_ctx, input, dim, 0, 1);
+                    for (int i = 0; i < left; i++) {
+                        input = ggml_concat(ctx->ggml_ctx, first, input, dim);
+                    }
+                }
+                if (right > 0) {
+                    auto last = ggml_ext_slice(ctx->ggml_ctx, input, dim, input->ne[dim] - 1, input->ne[dim]);
+                    for (int i = 0; i < right; i++) {
+                        input = ggml_concat(ctx->ggml_ctx, input, last, dim);
+                    }
+                }
+                return input;
+            };
+            padded = replicate_pad(padded, 0, pad_w, pad_w);
+            padded = replicate_pad(padded, 1, pad_h, pad_h);
+            return conv->forward(ctx, padded);
+        }
+    };
+
+    class AttnBlock : public UnaryBlock {
+    protected:
+        int64_t in_channels;
+
+    public:
+        AttnBlock(int64_t in_channels)
+            : in_channels(in_channels) {
+            blocks["norm"]     = std::make_shared<WAN::RMS_norm>(in_channels);
+            blocks["q"]        = std::make_shared<Conv3d>(in_channels, in_channels, std::tuple{1, 1, 1});
+            blocks["k"]        = std::make_shared<Conv3d>(in_channels, in_channels, std::tuple{1, 1, 1});
+            blocks["v"]        = std::make_shared<Conv3d>(in_channels, in_channels, std::tuple{1, 1, 1});
+            blocks["proj_out"] = std::make_shared<Conv3d>(in_channels, in_channels, std::tuple{1, 1, 1});
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x) override {
+            // x: [b*c, t, h, w]
+            auto norm     = std::dynamic_pointer_cast<WAN::RMS_norm>(blocks["norm"]);
+            auto q_proj   = std::dynamic_pointer_cast<UnaryBlock>(blocks["q"]);
+            auto k_proj   = std::dynamic_pointer_cast<UnaryBlock>(blocks["k"]);
+            auto v_proj   = std::dynamic_pointer_cast<UnaryBlock>(blocks["v"]);
+            auto proj_out = std::dynamic_pointer_cast<UnaryBlock>(blocks["proj_out"]);
+
+            const int64_t b = x->ne[3] / in_channels;
+
+            auto identity = x;
+
+            x = norm->forward(ctx, x);
+
+            const int64_t c = x->ne[3] / b;
+            const int64_t t = x->ne[2];
+            const int64_t h = x->ne[1];
+            const int64_t w = x->ne[0];
+
+            auto q = q_proj->forward(ctx, x);  // [b*c, t, h, w]
+            auto k = k_proj->forward(ctx, x);  // [b*c, t, h, w]
+            auto v = v_proj->forward(ctx, x);  // [b*c, t, h, w]
+
+            q = ggml_reshape_3d(ctx->ggml_ctx, q, w * h * t, c, b);                                  // [b, c, t*h*w]
+            q = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, q, 1, 0, 2, 3));  // [b, t*h*w, c]
+
+            k = ggml_reshape_3d(ctx->ggml_ctx, k, w * h * t, c, b);                                  // [b, c, t*h*w]
+            k = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, k, 1, 0, 2, 3));  // [b, t*h*w, c]
+
+            v = ggml_reshape_3d(ctx->ggml_ctx, v, w * h * t, c, b);                                  // [b, c, t*h*w]
+            v = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, v, 1, 0, 2, 3));  // [b, t*h*w, c]
+
+            x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled);  // [b, t*h*w, c]
+
+            x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));  // [b, c, t*h*w]
+            x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, t, c * b);                         // [b*c, t, h, w]
+
+            x = proj_out->forward(ctx, x);
+
+            x = ggml_add(ctx->ggml_ctx, x, identity);
+            return x;
+        }
+    };
+
+    class ResnetBlock : public UnaryBlock {
+    protected:
+        int64_t in_channels;
+        int64_t out_channels;
+
+    public:
+        ResnetBlock(int64_t in_channels,
+                    int64_t out_channels)
+            : in_channels(in_channels),
+              out_channels(out_channels) {
+            blocks["norm1"] = std::make_shared<WAN::RMS_norm>(in_channels);
+            blocks["conv1"] = std::make_shared<CausalConv3d>(in_channels, out_channels, std::tuple{3, 3, 3});
+
+            blocks["norm2"] = std::make_shared<WAN::RMS_norm>(out_channels);
+            blocks["conv2"] = std::make_shared<CausalConv3d>(out_channels, out_channels, std::tuple{3, 3, 3});
+
+            if (out_channels != in_channels) {
+                blocks["nin_shortcut"] = std::make_shared<CausalConv3d>(in_channels, out_channels, std::tuple{1, 1, 1});
+            }
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+            return forward(ctx, x, nullptr);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             TemporalConvCarry* carry) {
+            // x: [B*IC, IT, OH, OW]
+            // return: [B*OC, OT, OH, OW]
+            auto norm1 = std::dynamic_pointer_cast<WAN::RMS_norm>(blocks["norm1"]);
+            auto conv1 = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv1"]);
+            auto norm2 = std::dynamic_pointer_cast<WAN::RMS_norm>(blocks["norm2"]);
+            auto conv2 = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv2"]);
+
+            auto h = x;
+            h      = norm1->forward(ctx, h);
+            h      = ggml_silu_inplace(ctx->ggml_ctx, h);  // swish
+            h      = conv1->forward(ctx, h, carry);
+
+            h = norm2->forward(ctx, h);
+            h = ggml_silu_inplace(ctx->ggml_ctx, h);  // swish
+            // dropout, skip for inference
+            h = conv2->forward(ctx, h, carry);
+
+            // skip connection
+            if (out_channels != in_channels) {
+                auto nin_shortcut = std::dynamic_pointer_cast<CausalConv3d>(blocks["nin_shortcut"]);
+
+                x = nin_shortcut->forward(ctx, x);  // [B*OC, OT, OH, OW]
+            }
+
+            h = ggml_add(ctx->ggml_ctx, h, x);
+            return h;  // [B*OC, OT, OH, OW]
+        }
+    };
+
+    class Upsample : public GGMLBlock {
+    protected:
+        int64_t in_channels;
+        int64_t out_channels;
+        int64_t factor_t;
+        int64_t factor_s;
+        int64_t factor;
+        int64_t repeats;
+
+    public:
+        Upsample(int64_t in_channels, int64_t out_channels, bool add_temporal_upsample)
+            : in_channels(in_channels), out_channels(out_channels) {
+            if (add_temporal_upsample) {
+                factor_t = 2;
+            } else {
+                factor_t = 1;
+            }
+            factor_s = 2;
+            factor   = factor_t * factor_s * factor_s;
+            GGML_ASSERT(out_channels * factor % in_channels == 0);
+            repeats        = out_channels * factor / in_channels;
+            blocks["conv"] = std::make_shared<CausalConv3d>(in_channels, out_channels * factor, std::tuple{3, 3, 3});
+        }
+
+        static ggml_tensor* _pixel_shuffle_3d(GGMLRunnerContext* ctx,
+                                              ggml_tensor* x,
+                                              int64_t factor_t,
+                                              int64_t factor_s,
+                                              int64_t B = 1) {
+            // x: [B*factor*C, T, H, W]
+            // return: [B*C, T*factor_t, H*factor_s, W*factor_s]
+            GGML_ASSERT(B == 1);
+            int64_t factor = factor_t * factor_s * factor_s;
+            int64_t C      = x->ne[3] / factor;
+            int64_t T      = x->ne[2];
+            int64_t H      = x->ne[1];
+            int64_t W      = x->ne[0];
+
+            x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H * T, C, factor);                          // [factor, C, T*H, W]
+            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));  // [C, factor, T*H, W]
+            x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H * T, factor_s, factor_s * factor_t * C);  // [C*factor_t*factor_s, factor_s, T*H, W]
+            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));  // [C*factor_t*factor_s, T*H, W, factor_s]
+            x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W, H * T, factor_s, factor_t * C);  // [C*factor_t, factor_s, T*H, W*factor_s]
+            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));  // [C*factor_t, T*H, factor_s, W*factor_s]
+            x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W * factor_s * H, T, factor_t, C);  // [C, factor_t, T, H*factor_s*W*factor_s]
+            x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));  // [C, T, factor_t, H*factor_s*W*factor_s]
+            x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W, factor_s * H, factor_t * T, C);  // [C, T*factor_t, H*factor_s, W*factor_s]
+
+            return x;
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             TemporalConvCarry* carry = nullptr) {
+            // x: [B*IC, T, H, W]
+            // return: [B*OC, 1 + (T - 1)*factor_t, H*factor_s, W*factor_s]
+            const int64_t B = x->ne[3] / in_channels;
+            GGML_ASSERT(B == 1);
+
+            auto conv = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv"]);
+
+            const bool continuation = carry != nullptr && carry->is_continuation();
+            auto h                  = conv->forward(ctx, x, carry);  // [B*factor*OC, T, H, W]
+
+            ggml_tensor* shortcut = nullptr;
+            if (factor_t == 2 && !continuation) {
+                auto h_first = ggml_ext_slice(ctx->ggml_ctx, h, 2, 0, 1);                   // [B*factor*OC, 1, H, W]
+                h_first      = _pixel_shuffle_3d(ctx, h_first, 1, factor_s, B);             // [B*2*OC, 1, H*factor_s, W*factor_s]
+                h_first      = ggml_ext_slice(ctx->ggml_ctx, h_first, 3, 0, out_channels);  // [B*OC, 1, H*factor_s, W*factor_s]
+
+                auto x_first = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1);
+                x_first      = repeat_interleave_channels(ctx, x_first, repeats / 2, x->ne[0], x->ne[1], 1);
+                x_first      = _pixel_shuffle_3d(ctx, x_first, 1, factor_s, B);
+
+                if (x->ne[2] == 1) {
+                    return ggml_add(ctx->ggml_ctx, h_first, x_first);
+                }
+
+                auto h_next = ggml_ext_slice(ctx->ggml_ctx, h, 2, 1, h->ne[2]);       // [B*factor*OC, T - 1, H, W]
+                h_next      = _pixel_shuffle_3d(ctx, h_next, factor_t, factor_s, B);  // [B*OC, (T - 1)*factor_t, H*factor_s, W*factor_s]
+
+                h = ggml_concat(ctx->ggml_ctx, h_first, h_next, 2);  // [B*OC, 1 + (T - 1)*factor_t, H*factor_s, W*factor_s]
+
+                auto x_next = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1, x->ne[2]);
+                x_next      = repeat_interleave_channels(ctx, x_next, repeats, x->ne[0], x->ne[1], x->ne[2] - 1);
+                x_next      = _pixel_shuffle_3d(ctx, x_next, factor_t, factor_s, B);
+
+                shortcut = ggml_concat(ctx->ggml_ctx, x_first, x_next, 2);  // [B*OC, 1 + (T - 1)*factor_t, H*factor_s, W*factor_s]
+            } else {
+                h        = _pixel_shuffle_3d(ctx, h, factor_t, factor_s, B);
+                shortcut = repeat_interleave_channels(ctx, x, repeats, x->ne[0], x->ne[1], x->ne[2]);
+                shortcut = _pixel_shuffle_3d(ctx, shortcut, factor_t, factor_s, B);  // [B*OC, T*factor_t, H*factor_s, W*factor_s]
+            }
+
+            return ggml_add(ctx->ggml_ctx, h, shortcut);
+        }
+    };
+
+    static ggml_tensor* pixel_unshuffle_3d(GGMLRunnerContext* ctx,
+                                           ggml_tensor* x,
+                                           int64_t factor_t,
+                                           int64_t factor_s) {
+        GGML_ASSERT(x->ne[0] % factor_s == 0);
+        GGML_ASSERT(x->ne[1] % factor_s == 0);
+        GGML_ASSERT(x->ne[2] % factor_t == 0);
+        int64_t W      = x->ne[0] / factor_s;
+        int64_t H      = x->ne[1] / factor_s;
+        int64_t T      = x->ne[2] / factor_t;
+        int64_t C      = x->ne[3];
+        int64_t factor = factor_t * factor_s * factor_s;
+
+        x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W * factor_s * H, factor_t, T, C);
+        x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));
+        x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W, factor_s, H * T, factor_t * C);
+        x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3));
+        x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s, W, H * T, factor_s * factor_t * C);
+        x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 2, 0, 3));
+        x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H * T, factor, C);
+        x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
+        return ggml_reshape_4d(ctx->ggml_ctx, x, W, H, T, C * factor);
+    }
+
+    static ggml_tensor* mean_channel_groups(GGMLRunnerContext* ctx,
+                                            ggml_tensor* x,
+                                            int64_t group_size) {
+        GGML_ASSERT(group_size > 0);
+        GGML_ASSERT(x->ne[3] % group_size == 0);
+        if (group_size == 1) {
+            return x;
+        }
+        int64_t W       = x->ne[0];
+        int64_t H       = x->ne[1];
+        int64_t T       = x->ne[2];
+        int64_t spatial = W * H * T;
+        int64_t groups  = x->ne[3] / group_size;
+        x               = ggml_reshape_3d(ctx->ggml_ctx, x, spatial, group_size, groups);
+        x               = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
+        x               = ggml_sum_rows(ctx->ggml_ctx, x);
+        x               = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
+        x               = ggml_reshape_4d(ctx->ggml_ctx, x, W, H, T, groups);
+        return ggml_scale(ctx->ggml_ctx, x, 1.f / static_cast<float>(group_size));
+    }
+
+    class Downsample : public GGMLBlock {
+    protected:
+        int64_t in_channels;
+        int64_t out_channels;
+        int64_t factor_t;
+        int64_t factor_s = 2;
+        int64_t factor;
+        int64_t group_size;
+
+    public:
+        Downsample(int64_t in_channels, int64_t out_channels, bool add_temporal_downsample)
+            : in_channels(in_channels),
+              out_channels(out_channels),
+              factor_t(add_temporal_downsample ? 2 : 1),
+              factor(factor_t * factor_s * factor_s),
+              group_size(factor * in_channels / out_channels) {
+            GGML_ASSERT(out_channels % factor == 0);
+            GGML_ASSERT(factor * in_channels % out_channels == 0);
+            blocks["conv"] = std::make_shared<CausalConv3d>(in_channels,
+                                                            out_channels / factor,
+                                                            std::tuple{3, 3, 3});
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            auto conv = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv"]);
+            auto h    = conv->forward(ctx, x);
+
+            ggml_tensor* h_first = nullptr;
+            ggml_tensor* x_first = nullptr;
+            if (factor_t == 2) {
+                h_first = ggml_ext_slice(ctx->ggml_ctx, h, 2, 0, 1);
+                h_first = pixel_unshuffle_3d(ctx, h_first, 1, factor_s);
+                h_first = ggml_concat(ctx->ggml_ctx, h_first, h_first, 3);
+
+                x_first = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1);
+                x_first = pixel_unshuffle_3d(ctx, x_first, 1, factor_s);
+                x_first = mean_channel_groups(ctx, x_first, group_size / 2);
+
+                if (x->ne[2] == 1) {
+                    return ggml_add(ctx->ggml_ctx, h_first, x_first);
+                }
+                h = ggml_ext_slice(ctx->ggml_ctx, h, 2, 1, h->ne[2]);
+                x = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1, x->ne[2]);
+            }
+
+            GGML_ASSERT(h->ne[2] % factor_t == 0);
+            h = pixel_unshuffle_3d(ctx, h, factor_t, factor_s);
+            x = pixel_unshuffle_3d(ctx, x, factor_t, factor_s);
+            x = mean_channel_groups(ctx, x, group_size);
+
+            if (factor_t == 2) {
+                h = ggml_concat(ctx->ggml_ctx, h_first, h, 2);
+                x = ggml_concat(ctx->ggml_ctx, x_first, x, 2);
+            }
+            return ggml_add(ctx->ggml_ctx, h, x);
+        }
+    };
+
+    class MidBlock : public UnaryBlock {
+    protected:
+        int64_t in_channels;
+        int num_layers;
+        bool add_attention;
+
+    public:
+        MidBlock(int64_t in_channels,
+                 int num_layers     = 1,
+                 bool add_attention = true)
+            : in_channels(in_channels),
+              num_layers(num_layers),
+              add_attention(add_attention) {
+            blocks["block_1"] = std::make_shared<ResnetBlock>(in_channels, in_channels);
+            for (int i = 0; i < num_layers; i++) {
+                if (add_attention) {
+                    blocks["attn_" + std::to_string(i + 1)] = std::make_shared<AttnBlock>(in_channels);
+                }
+                blocks["block_" + std::to_string(i + 2)] = std::make_shared<ResnetBlock>(in_channels, in_channels);
+            }
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+            // x: [B*C, T, H, W]
+            // return: [B*C, T, H, W]
+            auto block_1 = std::dynamic_pointer_cast<ResnetBlock>(blocks["block_1"]);
+
+            x = block_1->forward(ctx, x);
+            for (int i = 0; i < num_layers; i++) {
+                if (add_attention) {
+                    auto block = std::dynamic_pointer_cast<AttnBlock>(blocks["attn_" + std::to_string(i + 1)]);
+                    x          = block->forward(ctx, x);
+                }
+                auto block = std::dynamic_pointer_cast<ResnetBlock>(blocks["block_" + std::to_string(i + 2)]);
+                x          = block->forward(ctx, x);
+            }
+            return x;
+        }
+    };
+
+    class UpBlock : public UnaryBlock {
+    protected:
+        int num_layers;
+        int64_t upsample_out_channels;
+
+    public:
+        UpBlock(int64_t in_channels,
+                int64_t out_channels,
+                int num_layers                = 1,
+                int64_t upsample_out_channels = 0,
+                bool add_temporal_upsample    = true)
+            : num_layers(num_layers),
+              upsample_out_channels(upsample_out_channels) {
+            for (int i = 0; i < num_layers; i++) {
+                int64_t IC                           = i == 0 ? in_channels : out_channels;
+                blocks["block." + std::to_string(i)] = std::make_shared<ResnetBlock>(IC, out_channels);
+            }
+            if (upsample_out_channels > 0) {
+                blocks["upsample"] = std::make_shared<Upsample>(out_channels, upsample_out_channels, add_temporal_upsample);
+            }
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+            return forward(ctx, x, nullptr);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx,
+                             ggml_tensor* x,
+                             TemporalConvCarry* carry) {
+            // x: [B*IC, T, H, W]
+            // return: [B*OC, T, H, W] or [B*OC, T, H*2, W*2] or [B*OC, T*2, H*2, W*2]
+            for (int i = 0; i < num_layers; i++) {
+                auto block = std::dynamic_pointer_cast<ResnetBlock>(blocks["block." + std::to_string(i)]);
+                x          = block->forward(ctx, x, carry);
+            }
+            if (upsample_out_channels > 0) {
+                auto upsample = std::dynamic_pointer_cast<Upsample>(blocks["upsample"]);
+                x             = upsample->forward(ctx, x, carry);
+            }
+            return x;
+        }
+    };
+
+    class DownBlock : public UnaryBlock {
+    protected:
+        int num_layers;
+        int64_t downsample_out_channels;
+
+    public:
+        DownBlock(int64_t in_channels,
+                  int64_t out_channels,
+                  int num_layers,
+                  int64_t downsample_out_channels = 0,
+                  bool add_temporal_downsample    = false)
+            : num_layers(num_layers),
+              downsample_out_channels(downsample_out_channels) {
+            for (int i = 0; i < num_layers; i++) {
+                int64_t IC                           = i == 0 ? in_channels : out_channels;
+                blocks["block." + std::to_string(i)] = std::make_shared<ResnetBlock>(IC, out_channels);
+            }
+            if (downsample_out_channels > 0) {
+                blocks["downsample"] = std::make_shared<Downsample>(out_channels,
+                                                                    downsample_out_channels,
+                                                                    add_temporal_downsample);
+            }
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+            for (int i = 0; i < num_layers; i++) {
+                auto block = std::dynamic_pointer_cast<ResnetBlock>(blocks["block." + std::to_string(i)]);
+                x          = block->forward(ctx, x);
+            }
+            if (downsample_out_channels > 0) {
+                auto downsample = std::dynamic_pointer_cast<Downsample>(blocks["downsample"]);
+                x               = downsample->forward(ctx, x);
+            }
+            return x;
+        }
+    };
+
+    class Encoder : public GGMLBlock {
+    protected:
+        int64_t z_channels;
+        std::vector<int64_t> block_out_channels;
+
+    public:
+        Encoder(int64_t in_channels                     = 3,
+                int64_t z_channels                      = 32,
+                std::vector<int64_t> block_out_channels = {128, 256, 512, 1024, 1024},
+                int layers_per_block                    = 2,
+                int spatial_compression_ratio           = 16,
+                int temporal_compression_ratio          = 4,
+                bool downsample_match_channel           = true)
+            : z_channels(z_channels),
+              block_out_channels(std::move(block_out_channels)) {
+            blocks["conv_in"] = std::make_shared<CausalConv3d>(in_channels,
+                                                               this->block_out_channels[0],
+                                                               std::tuple{3, 3, 3});
+
+            int spatial_depth  = static_cast<int>(std::log2(static_cast<double>(spatial_compression_ratio)));
+            int temporal_start = static_cast<int>(std::log2(static_cast<double>(spatial_compression_ratio / temporal_compression_ratio)));
+            int64_t channels   = this->block_out_channels[0];
+            for (int i = 0; i < static_cast<int>(this->block_out_channels.size()); i++) {
+                int64_t out_channels = this->block_out_channels[i];
+                if (i < spatial_depth) {
+                    int64_t next_channels               = downsample_match_channel ? this->block_out_channels[i + 1] : out_channels;
+                    blocks["down." + std::to_string(i)] = std::make_shared<DownBlock>(channels,
+                                                                                      out_channels,
+                                                                                      layers_per_block,
+                                                                                      next_channels,
+                                                                                      i >= temporal_start);
+                    channels                            = next_channels;
+                } else {
+                    blocks["down." + std::to_string(i)] = std::make_shared<DownBlock>(channels,
+                                                                                      out_channels,
+                                                                                      layers_per_block);
+                    channels                            = out_channels;
+                }
+            }
+
+            blocks["mid"]      = std::make_shared<MidBlock>(channels);
+            blocks["norm_out"] = std::make_shared<WAN::RMS_norm>(channels);
+            blocks["conv_out"] = std::make_shared<CausalConv3d>(channels,
+                                                                z_channels * 2,
+                                                                std::tuple{3, 3, 3});
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            auto conv_in  = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv_in"]);
+            auto mid      = std::dynamic_pointer_cast<MidBlock>(blocks["mid"]);
+            auto norm_out = std::dynamic_pointer_cast<WAN::RMS_norm>(blocks["norm_out"]);
+            auto conv_out = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv_out"]);
+
+            x = conv_in->forward(ctx, x);
+            for (int i = 0; i < static_cast<int>(block_out_channels.size()); i++) {
+                auto down = std::dynamic_pointer_cast<DownBlock>(blocks["down." + std::to_string(i)]);
+                x         = down->forward(ctx, x);
+            }
+            x = mid->forward(ctx, x);
+
+            auto shortcut = mean_channel_groups(ctx, x, x->ne[3] / (z_channels * 2));
+            x             = norm_out->forward(ctx, x);
+            x             = ggml_silu_inplace(ctx->ggml_ctx, x);
+            x             = conv_out->forward(ctx, x);
+            x             = ggml_add(ctx->ggml_ctx, x, shortcut);
+            return ggml_ext_slice(ctx->ggml_ctx, x, 3, 0, z_channels);
+        }
+    };
+
+    class Decoder : public GGMLBlock {
+    protected:
+        int64_t repeats;
+        std::vector<int64_t> block_out_channels;
+
+    public:
+        Decoder(int64_t in_channels                     = 32,
+                int64_t out_channels                    = 3,
+                std::vector<int64_t> block_out_channels = {1024, 1024, 512, 256, 128},
+                int layers_per_block                    = 2,
+                int spatial_compression_ratio           = 16,
+                int temporal_compression_ratio          = 4,
+                bool upsample_match_channel             = true)
+            : block_out_channels(std::move(block_out_channels)) {
+            repeats           = this->block_out_channels[0] / in_channels;
+            blocks["conv_in"] = std::make_shared<CausalConv3d>(in_channels, this->block_out_channels[0], std::tuple{3, 3, 3});
+            blocks["mid"]     = std::make_shared<MidBlock>(this->block_out_channels[0]);
+
+            int64_t IC = this->block_out_channels[0];
+            for (int i = 0; i < this->block_out_channels.size(); i++) {
+                int64_t OC                 = this->block_out_channels[i];
+                bool add_spatial_upsample  = i < std::log2(static_cast<double>(spatial_compression_ratio));
+                bool add_temporal_upsample = i < std::log2(static_cast<double>(temporal_compression_ratio));
+
+                if (add_spatial_upsample || add_temporal_upsample) {
+                    int64_t upsample_out_channels     = upsample_match_channel ? this->block_out_channels[i + 1] : OC;
+                    blocks["up." + std::to_string(i)] = std::make_shared<UpBlock>(IC, OC, layers_per_block + 1, upsample_out_channels, add_temporal_upsample);
+                    IC                                = upsample_out_channels;
+                } else {
+                    blocks["up." + std::to_string(i)] = std::make_shared<UpBlock>(IC, OC, layers_per_block + 1, 0, false);
+                }
+            }
+
+            blocks["norm_out"] = std::make_shared<WAN::RMS_norm>(this->block_out_channels.back());
+            blocks["conv_out"] = std::make_shared<CausalConv3d>(this->block_out_channels.back(), out_channels, std::tuple{3, 3, 3});
+        }
+
+        struct ggml_tensor* forward(GGMLRunnerContext* ctx, struct ggml_tensor* z) {
+            auto conv_in   = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv_in"]);
+            auto mid_block = std::dynamic_pointer_cast<MidBlock>(blocks["mid"]);
+            auto norm_out  = std::dynamic_pointer_cast<WAN::RMS_norm>(blocks["norm_out"]);
+            auto conv_out  = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv_out"]);
+
+            auto h = conv_in->forward(ctx, z);
+
+            auto shortcut = repeat_interleave_channels(ctx, z, repeats, z->ne[0], z->ne[1], z->ne[2]);
+            h             = ggml_add(ctx->ggml_ctx, h, shortcut);
+
+            h = mid_block->forward(ctx, h);
+
+            ggml_tensor* output = nullptr;
+            std::vector<ggml_tensor*> carry_input;
+            const int64_t frames = h->ne[2];
+            for (int64_t start = 0; start < frames; start += HUNYUAN_VIDEO_VAE_TEMPORAL_CHUNK_SIZE) {
+                const int64_t end = std::min(start + HUNYUAN_VIDEO_VAE_TEMPORAL_CHUNK_SIZE, frames);
+                auto chunk        = ggml_ext_slice(ctx->ggml_ctx, h, 2, start, end);
+
+                std::vector<ggml_tensor*> carry_output;
+                TemporalConvCarry carry{
+                    start == 0 ? nullptr : &carry_input,
+                    end == frames ? nullptr : &carry_output,
+                };
+
+                for (int i = 0; i < block_out_channels.size(); i++) {
+                    auto up_block = std::dynamic_pointer_cast<UpBlock>(blocks["up." + std::to_string(i)]);
+                    chunk         = up_block->forward(ctx, chunk, &carry);
+                }
+
+                chunk = norm_out->forward(ctx, chunk);
+                chunk = ggml_silu_inplace(ctx->ggml_ctx, chunk);  // nonlinearity/swish
+                chunk = conv_out->forward(ctx, chunk, &carry);
+                carry.finish();
+
+                output      = output == nullptr ? chunk : ggml_concat(ctx->ggml_ctx, output, chunk, 2);
+                carry_input = std::move(carry_output);
+            }
+            return output;
+        }
+    };
+
+    class HunyuanVideoVAERunner : public VAE {
+    protected:
+        bool decode_only;
+        Encoder encoder;
+        Decoder decoder;
+
+    public:
+        HunyuanVideoVAERunner(ggml_backend_t backend,
+                              const String2TensorStorage& tensor_storage_map,
+                              const std::string& prefix,
+                              bool decode_only,
+                              SDVersion version,
+                              std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : VAE(version, backend, prefix, weight_manager),
+              decode_only(decode_only ||
+                          tensor_storage_map.find(prefix + ".encoder.conv_in.conv.weight") == tensor_storage_map.end()) {
+            if (!this->decode_only) {
+                encoder.init(params_ctx, tensor_storage_map, prefix + ".encoder");
+            }
+            decoder.init(params_ctx, tensor_storage_map, prefix + ".decoder");
+        }
+
+        std::string get_desc() override {
+            return "hunyuan_video_vae";
+        }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+            if (!decode_only) {
+                encoder.get_param_tensors(tensors, weight_prefix + ".encoder");
+            }
+            decoder.get_param_tensors(tensors, weight_prefix + ".decoder");
+        }
+
+        int get_encoder_output_channels(int input_channels) override {
+            SD_UNUSED(input_channels);
+            return 32;
+        }
+
+        sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output,
+                                                std::shared_ptr<RNG> rng) override {
+            SD_UNUSED(rng);
+            return vae_output;
+        }
+
+        sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) override {
+            return latents / 1.03682f;
+        }
+
+        sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) override {
+            return latents * 1.03682f;
+        }
+
+        ggml_cgraph* build_graph(const sd::Tensor<float>& input_tensor, bool decode_graph) {
+            size_t graph_size = HUNYUAN_VIDEO_VAE_GRAPH_SIZE;
+            if (decode_graph) {
+                graph_size = std::max(graph_size,
+                                      HUNYUAN_VIDEO_VAE_GRAPH_SIZE_PER_LATENT_FRAME *
+                                          static_cast<size_t>(input_tensor.shape()[2]));
+            }
+            ggml_cgraph* gf     = new_graph_custom(graph_size);
+            ggml_tensor* input  = make_input(input_tensor);
+            auto runner_ctx     = get_context();
+            ggml_tensor* output = decode_graph ? decoder.forward(&runner_ctx, input)
+                                               : encoder.forward(&runner_ctx, input);
+            ggml_build_forward_expand(gf, output);
+            return gf;
+        }
+
+        sd::Tensor<float> _compute(const int n_threads,
+                                   const sd::Tensor<float>& input,
+                                   bool decode_graph) override {
+            if (!decode_graph && decode_only) {
+                LOG_ERROR("Hunyuan Video VAE encoder weights are not available");
+                return {};
+            }
+
+            sd::Tensor<float> expanded;
+            if (input.dim() == 4) {
+                expanded = input.unsqueeze(2);
+            }
+            const auto& graph_input = expanded.empty() ? input : expanded;
+            auto get_graph          = [&]() -> ggml_cgraph* {
+                return build_graph(graph_input, decode_graph);
+            };
+            auto output = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph,
+                                                                                     n_threads,
+                                                                                     true,
+                                                                                     true,
+                                                                                     true),
+                                                          graph_input.dim());
+            if (!output.empty() && input.dim() == 4) {
+                output.squeeze_(2);
+            }
+            return output;
+        }
+    };
+
+}  // namespace Hunyuan
+
+#endif  // __SD_MODEL_VAE_HUNYUAN_VAE_HPP__

src/model/vae/mage_vae.hpp

@@ -0,0 +1,521 @@
+#ifndef __SD_MODEL_VAE_MAGE_VAE_HPP__
+#define __SD_MODEL_VAE_MAGE_VAE_HPP__
+
+#include "model/diffusion/dit.hpp"
+#include "model/vae/vae.hpp"
+
+namespace MageVAE {
+    constexpr int MAGE_VAE_GRAPH_SIZE = 327680;
+    constexpr int HIDDEN_SIZE         = 384;
+    constexpr int LATENT_CHANNELS     = 128;
+    constexpr int PATCH_SIZE          = 16;
+
+    struct LayerNorm2d : public UnaryBlock {
+        int64_t channels;
+        bool affine;
+        std::string prefix;
+
+        void init_params(ggml_context* ctx,
+                         const String2TensorStorage& tensor_storage_map = {},
+                         const std::string prefix                       = "") override {
+            this->prefix = prefix;
+            if (affine) {
+                params["weight"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels);
+                params["bias"]   = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels);
+            }
+        }
+
+        LayerNorm2d(int64_t channels, bool affine = true)
+            : channels(channels), affine(affine) {}
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+            ggml_tensor* weight = affine ? params["weight"] : nullptr;
+            ggml_tensor* bias   = affine ? params["bias"] : nullptr;
+            if (affine && ctx->weight_adapter) {
+                weight = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, weight, prefix + "weight");
+                bias   = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, bias, prefix + "bias");
+            }
+            // [N, C, H, W] -> [N, H, W, C] so layer norm reduces over channels.
+            x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 2, 0, 3));
+            x = ggml_ext_layer_norm(ctx->ggml_ctx, x, weight, bias, 1e-6f);
+            return ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));
+        }
+    };
+
+    inline ggml_tensor* modulate_2d(ggml_context* ctx,
+                                    ggml_tensor* x,
+                                    ggml_tensor* shift,
+                                    ggml_tensor* scale) {
+        shift = ggml_reshape_4d(ctx, shift, 1, 1, shift->ne[0], shift->ne[1]);
+        scale = ggml_reshape_4d(ctx, scale, 1, 1, scale->ne[0], scale->ne[1]);
+        return ggml_add(ctx, ggml_mul(ctx, x, ggml_add(ctx, scale, ggml_ext_ones(ctx, 1, 1, 1, 1))), shift);
+    }
+
+    inline ggml_tensor* channel_attention(GGMLRunnerContext* ctx,
+                                          ggml_tensor* x,
+                                          Conv2d* projection) {
+        auto pooled = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1], x->ne[2], x->ne[3]);
+        pooled      = ggml_mean(ctx->ggml_ctx, pooled);
+        pooled      = ggml_reshape_4d(ctx->ggml_ctx, pooled, 1, 1, x->ne[2], x->ne[3]);
+        pooled      = ggml_sigmoid(ctx->ggml_ctx, projection->forward(ctx, pooled));
+        return ggml_mul(ctx->ggml_ctx, x, pooled);
+    }
+
+    struct TimestepEmbedder : public GGMLBlock {
+        TimestepEmbedder() {
+            blocks["mlp.0"] = std::make_shared<Linear>(256, HIDDEN_SIZE);
+            blocks["mlp.2"] = std::make_shared<Linear>(HIDDEN_SIZE, HIDDEN_SIZE);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* timestep) {
+            auto linear_0 = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"]);
+            auto linear_2 = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"]);
+            auto x        = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, 256, 10000, 1.f);
+            x             = linear_0->forward(ctx, x);
+            x             = ggml_silu_inplace(ctx->ggml_ctx, x);
+            return linear_2->forward(ctx, x);
+        }
+    };
+
+    struct EncoderDiCoBlock : public UnaryBlock {
+        explicit EncoderDiCoBlock(int64_t channels) {
+            blocks["conv1"] = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
+            blocks["conv2"] = std::make_shared<Conv2d_grouped>(channels, channels, static_cast<int>(channels), std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1});
+            blocks["conv3"] = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
+            blocks["ca.1"]  = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
+            blocks["conv4"] = std::make_shared<Conv2d>(channels, channels * 4, std::pair{1, 1});
+            blocks["conv5"] = std::make_shared<Conv2d>(channels * 4, channels, std::pair{1, 1});
+            blocks["norm1"] = std::make_shared<LayerNorm2d>(channels);
+            blocks["norm2"] = std::make_shared<LayerNorm2d>(channels);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input) override {
+            auto conv1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv1"]);
+            auto conv2 = std::dynamic_pointer_cast<Conv2d_grouped>(blocks["conv2"]);
+            auto conv3 = std::dynamic_pointer_cast<Conv2d>(blocks["conv3"]);
+            auto ca    = std::dynamic_pointer_cast<Conv2d>(blocks["ca.1"]);
+            auto conv4 = std::dynamic_pointer_cast<Conv2d>(blocks["conv4"]);
+            auto conv5 = std::dynamic_pointer_cast<Conv2d>(blocks["conv5"]);
+            auto norm1 = std::dynamic_pointer_cast<LayerNorm2d>(blocks["norm1"]);
+            auto norm2 = std::dynamic_pointer_cast<LayerNorm2d>(blocks["norm2"]);
+
+            auto x = norm1->forward(ctx, input);
+            x      = conv1->forward(ctx, x);
+            x      = conv2->forward(ctx, x);
+            x      = ggml_gelu(ctx->ggml_ctx, x);
+            x      = channel_attention(ctx, x, ca.get());
+            x      = conv3->forward(ctx, x);
+            x      = ggml_add(ctx->ggml_ctx, input, x);
+            auto h = norm2->forward(ctx, x);
+            h      = conv4->forward(ctx, h);
+            h      = ggml_gelu(ctx->ggml_ctx, h);
+            h      = conv5->forward(ctx, h);
+            return ggml_add(ctx->ggml_ctx, x, h);
+        }
+    };
+
+    struct DiCoBlock : public GGMLBlock {
+        explicit DiCoBlock(int64_t channels) {
+            blocks["conv1"]              = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
+            blocks["conv2"]              = std::make_shared<Conv2d_grouped>(channels, channels, static_cast<int>(channels), std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1});
+            blocks["conv3"]              = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
+            blocks["ca.1"]               = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
+            blocks["conv4"]              = std::make_shared<Conv2d>(channels, channels * 4, std::pair{1, 1});
+            blocks["conv5"]              = std::make_shared<Conv2d>(channels * 4, channels, std::pair{1, 1});
+            blocks["norm1"]              = std::make_shared<LayerNorm2d>(channels, false);
+            blocks["norm2"]              = std::make_shared<LayerNorm2d>(channels, false);
+            blocks["adaLN_modulation.1"] = std::make_shared<Linear>(channels, channels * 6);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input, ggml_tensor* condition) {
+            auto conv1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv1"]);
+            auto conv2 = std::dynamic_pointer_cast<Conv2d_grouped>(blocks["conv2"]);
+            auto conv3 = std::dynamic_pointer_cast<Conv2d>(blocks["conv3"]);
+            auto ca    = std::dynamic_pointer_cast<Conv2d>(blocks["ca.1"]);
+            auto conv4 = std::dynamic_pointer_cast<Conv2d>(blocks["conv4"]);
+            auto conv5 = std::dynamic_pointer_cast<Conv2d>(blocks["conv5"]);
+            auto norm1 = std::dynamic_pointer_cast<LayerNorm2d>(blocks["norm1"]);
+            auto norm2 = std::dynamic_pointer_cast<LayerNorm2d>(blocks["norm2"]);
+            auto ada   = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.1"]);
+
+            auto params = ada->forward(ctx, ggml_silu(ctx->ggml_ctx, condition));
+            auto chunks = ggml_ext_chunk(ctx->ggml_ctx, params, 6, 0);
+            auto x      = norm1->forward(ctx, input);
+            x           = modulate_2d(ctx->ggml_ctx, x, chunks[0], chunks[1]);
+            x           = conv1->forward(ctx, x);
+            x           = conv2->forward(ctx, x);
+            x           = ggml_gelu(ctx->ggml_ctx, x);
+            x           = channel_attention(ctx, x, ca.get());
+            x           = conv3->forward(ctx, x);
+            auto gate_1 = ggml_reshape_4d(ctx->ggml_ctx, chunks[2], 1, 1, chunks[2]->ne[0], chunks[2]->ne[1]);
+            x           = ggml_add(ctx->ggml_ctx, input, ggml_mul(ctx->ggml_ctx, x, gate_1));
+
+            auto h      = norm2->forward(ctx, x);
+            h           = modulate_2d(ctx->ggml_ctx, h, chunks[3], chunks[4]);
+            h           = conv4->forward(ctx, h);
+            h           = ggml_gelu(ctx->ggml_ctx, h);
+            h           = conv5->forward(ctx, h);
+            auto gate_2 = ggml_reshape_4d(ctx->ggml_ctx, chunks[5], 1, 1, chunks[5]->ne[0], chunks[5]->ne[1]);
+            return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, h, gate_2));
+        }
+    };
+
+    struct MageResnetBlock : public UnaryBlock {
+        explicit MageResnetBlock(int64_t channels) {
+            blocks["norm1"] = std::make_shared<GroupNorm32>(channels);
+            blocks["conv1"] = std::make_shared<Conv2d>(channels, channels, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1});
+            blocks["norm2"] = std::make_shared<GroupNorm32>(channels);
+            blocks["conv2"] = std::make_shared<Conv2d>(channels, channels, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1});
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input) override {
+            auto norm1 = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm1"]);
+            auto conv1 = std::dynamic_pointer_cast<Conv2d>(blocks["conv1"]);
+            auto norm2 = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm2"]);
+            auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["conv2"]);
+            auto x     = conv1->forward(ctx, ggml_silu(ctx->ggml_ctx, norm1->forward(ctx, input)));
+            x          = conv2->forward(ctx, ggml_silu(ctx->ggml_ctx, norm2->forward(ctx, x)));
+            return ggml_add(ctx->ggml_ctx, input, x);
+        }
+    };
+
+    inline ggml_tensor* replicate_pad_right_bottom(ggml_context* ctx,
+                                                   ggml_tensor* x,
+                                                   int pad_w,
+                                                   int pad_h) {
+        if (pad_w > 0) {
+            auto edge = ggml_ext_slice(ctx, x, 0, x->ne[0] - 1, x->ne[0]);
+            edge      = ggml_repeat_4d(ctx, edge, pad_w, x->ne[1], x->ne[2], x->ne[3]);
+            x         = ggml_concat(ctx, x, edge, 0);
+        }
+        if (pad_h > 0) {
+            auto edge = ggml_ext_slice(ctx, x, 1, x->ne[1] - 1, x->ne[1]);
+            edge      = ggml_repeat_4d(ctx, edge, x->ne[0], pad_h, x->ne[2], x->ne[3]);
+            x         = ggml_concat(ctx, x, edge, 1);
+        }
+        return x;
+    }
+
+    struct MageAttnBlock : public UnaryBlock {
+        int64_t channels;
+        int patch_size;
+
+        MageAttnBlock(int64_t channels, int patch_size = 32)
+            : channels(channels), patch_size(patch_size) {
+            blocks["norm"]     = std::make_shared<GroupNorm32>(channels);
+            blocks["q"]        = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
+            blocks["k"]        = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
+            blocks["v"]        = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
+            blocks["proj_out"] = std::make_shared<Conv2d>(channels, channels, std::pair{1, 1});
+        }
+
+        ggml_tensor* to_patches(ggml_context* ctx, ggml_tensor* x) {
+            x = DiT::patchify(ctx, x, patch_size, patch_size);
+            x = ggml_reshape_4d(ctx, x, patch_size * patch_size, channels, x->ne[1], x->ne[2]);
+            // [N, np, C, P] -> [N, np, P, C] for attention over P pixels.
+            x = ggml_ext_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3));
+            return ggml_reshape_3d(ctx, x, channels, patch_size * patch_size, x->ne[2] * x->ne[3]);
+        }
+
+        ggml_tensor* from_patches(ggml_context* ctx,
+                                  ggml_tensor* x,
+                                  int64_t patch_count,
+                                  int64_t batch_size,
+                                  int64_t h_patches,
+                                  int64_t w_patches) {
+            x = ggml_reshape_4d(ctx, x, channels, patch_size * patch_size, patch_count, batch_size);
+            // [N, np, P, C] -> [N, np, C, P] before spatial unpatchify.
+            x = ggml_ext_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3));
+            x = ggml_reshape_3d(ctx, x, patch_size * patch_size * channels, patch_count, batch_size);
+            return DiT::unpatchify(ctx, x, h_patches, w_patches, patch_size, patch_size);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input) override {
+            auto norm     = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm"]);
+            auto q_proj   = std::dynamic_pointer_cast<Conv2d>(blocks["q"]);
+            auto k_proj   = std::dynamic_pointer_cast<Conv2d>(blocks["k"]);
+            auto v_proj   = std::dynamic_pointer_cast<Conv2d>(blocks["v"]);
+            auto proj_out = std::dynamic_pointer_cast<Conv2d>(blocks["proj_out"]);
+
+            int64_t width  = input->ne[0];
+            int64_t height = input->ne[1];
+            int64_t batch  = input->ne[3];
+            int pad_w      = (patch_size - static_cast<int>(width % patch_size)) % patch_size;
+            int pad_h      = (patch_size - static_cast<int>(height % patch_size)) % patch_size;
+            int64_t wp     = (width + pad_w) / patch_size;
+            int64_t hp     = (height + pad_h) / patch_size;
+            int64_t np     = wp * hp;
+
+            auto h = norm->forward(ctx, input);
+            auto q = replicate_pad_right_bottom(ctx->ggml_ctx, q_proj->forward(ctx, h), pad_w, pad_h);
+            auto k = replicate_pad_right_bottom(ctx->ggml_ctx, k_proj->forward(ctx, h), pad_w, pad_h);
+            auto v = replicate_pad_right_bottom(ctx->ggml_ctx, v_proj->forward(ctx, h), pad_w, pad_h);
+            q      = to_patches(ctx->ggml_ctx, q);
+            k      = to_patches(ctx->ggml_ctx, k);
+            v      = to_patches(ctx->ggml_ctx, v);
+            h      = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled);
+            h      = from_patches(ctx->ggml_ctx, h, np, batch, hp, wp);
+            if (pad_h > 0) {
+                h = ggml_ext_slice(ctx->ggml_ctx, h, 1, 0, height);
+            }
+            if (pad_w > 0) {
+                h = ggml_ext_slice(ctx->ggml_ctx, h, 0, 0, width);
+            }
+            return ggml_add(ctx->ggml_ctx, input, proj_out->forward(ctx, h));
+        }
+    };
+
+    struct Decoder : public UnaryBlock {
+        Decoder() {
+            blocks["conv_in"]  = std::make_shared<Conv2d>(LATENT_CHANNELS, HIDDEN_SIZE, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1});
+            blocks["block.0"]  = std::make_shared<MageResnetBlock>(HIDDEN_SIZE);
+            blocks["block.1"]  = std::make_shared<MageAttnBlock>(HIDDEN_SIZE);
+            blocks["block.2"]  = std::make_shared<MageResnetBlock>(HIDDEN_SIZE);
+            blocks["block.3"]  = std::make_shared<MageAttnBlock>(HIDDEN_SIZE);
+            blocks["block.4"]  = std::make_shared<MageResnetBlock>(HIDDEN_SIZE);
+            blocks["norm_out"] = std::make_shared<GroupNorm32>(HIDDEN_SIZE);
+            blocks["conv_out"] = std::make_shared<Conv2d>(HIDDEN_SIZE, HIDDEN_SIZE, std::pair{3, 3}, std::pair{1, 1}, std::pair{1, 1});
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
+            x = std::dynamic_pointer_cast<Conv2d>(blocks["conv_in"])->forward(ctx, x);
+            for (int i = 0; i < 5; ++i) {
+                x = std::dynamic_pointer_cast<UnaryBlock>(blocks["block." + std::to_string(i)])->forward(ctx, x);
+            }
+            x = std::dynamic_pointer_cast<GroupNorm32>(blocks["norm_out"])->forward(ctx, x);
+            x = ggml_silu(ctx->ggml_ctx, x);
+            return std::dynamic_pointer_cast<Conv2d>(blocks["conv_out"])->forward(ctx, x);
+        }
+    };
+
+    struct DConvEncoder : public UnaryBlock {
+        DConvEncoder() {
+            blocks["patch_cond_embed"] = std::make_shared<Conv2d>(3, 768, std::pair{PATCH_SIZE, PATCH_SIZE}, std::pair{PATCH_SIZE, PATCH_SIZE});
+            for (int i = 0; i < 2; ++i) {
+                blocks["head_blocks." + std::to_string(i)] = std::make_shared<EncoderDiCoBlock>(768);
+            }
+            blocks["proj_down"]  = std::make_shared<Conv2d>(768, HIDDEN_SIZE, std::pair{1, 1});
+            blocks["z_proj"]     = std::make_shared<Conv2d>(LATENT_CHANNELS, HIDDEN_SIZE, std::pair{1, 1});
+            blocks["fuse_proj"]  = std::make_shared<Conv2d>(HIDDEN_SIZE * 2, HIDDEN_SIZE, std::pair{1, 1});
+            blocks["t_embedder"] = std::make_shared<TimestepEmbedder>();
+            for (int i = 0; i < 21; ++i) {
+                blocks["blocks." + std::to_string(i)] = std::make_shared<DiCoBlock>(HIDDEN_SIZE);
+            }
+            blocks["norm_out"] = std::make_shared<LayerNorm2d>(HIDDEN_SIZE);
+            blocks["proj_out"] = std::make_shared<Conv2d>(HIDDEN_SIZE, LATENT_CHANNELS * 2, std::pair{1, 1});
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* image) override {
+            auto cond = std::dynamic_pointer_cast<Conv2d>(blocks["patch_cond_embed"])->forward(ctx, image);
+            for (int i = 0; i < 2; ++i) {
+                cond = std::dynamic_pointer_cast<EncoderDiCoBlock>(blocks["head_blocks." + std::to_string(i)])->forward(ctx, cond);
+            }
+            cond   = std::dynamic_pointer_cast<Conv2d>(blocks["proj_down"])->forward(ctx, cond);
+            auto z = ggml_ext_zeros(ctx->ggml_ctx, cond->ne[0], cond->ne[1], LATENT_CHANNELS, cond->ne[3]);
+            z      = std::dynamic_pointer_cast<Conv2d>(blocks["z_proj"])->forward(ctx, z);
+            z      = ggml_concat(ctx->ggml_ctx, cond, z, 2);
+            z      = std::dynamic_pointer_cast<Conv2d>(blocks["fuse_proj"])->forward(ctx, z);
+            auto t = ggml_ext_zeros(ctx->ggml_ctx, image->ne[3], 1, 1, 1);
+            t      = ggml_reshape_1d(ctx->ggml_ctx, t, image->ne[3]);
+            auto c = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"])->forward(ctx, t);
+            for (int i = 0; i < 21; ++i) {
+                z = std::dynamic_pointer_cast<DiCoBlock>(blocks["blocks." + std::to_string(i)])->forward(ctx, z, c);
+            }
+            z = std::dynamic_pointer_cast<LayerNorm2d>(blocks["norm_out"])->forward(ctx, z);
+            return std::dynamic_pointer_cast<Conv2d>(blocks["proj_out"])->forward(ctx, z);
+        }
+    };
+
+    struct MLPResBlock : public GGMLBlock {
+        MLPResBlock() {
+            blocks["in_ln"]              = std::make_shared<LayerNorm>(32, 1e-6f);
+            blocks["mlp.0"]              = std::make_shared<Linear>(32, 32);
+            blocks["mlp.2"]              = std::make_shared<Linear>(32, 32);
+            blocks["adaLN_modulation.1"] = std::make_shared<Linear>(32, 96);
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* condition) {
+            auto params = std::dynamic_pointer_cast<Linear>(blocks["adaLN_modulation.1"])->forward(ctx, ggml_silu(ctx->ggml_ctx, condition));
+            auto chunks = ggml_ext_chunk(ctx->ggml_ctx, params, 3, 0);
+            auto h      = std::dynamic_pointer_cast<LayerNorm>(blocks["in_ln"])->forward(ctx, x);
+            h           = ggml_add(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, h, ggml_add(ctx->ggml_ctx, chunks[1], ggml_ext_ones(ctx->ggml_ctx, 1, 1, 1, 1))), chunks[0]);
+            h           = std::dynamic_pointer_cast<Linear>(blocks["mlp.0"])->forward(ctx, h);
+            h           = ggml_silu(ctx->ggml_ctx, h);
+            h           = std::dynamic_pointer_cast<Linear>(blocks["mlp.2"])->forward(ctx, h);
+            return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, chunks[2], h));
+        }
+    };
+
+    struct DConvDenoiser : public GGMLBlock {
+        DConvDenoiser() {
+            blocks["t_embedder"]            = std::make_shared<TimestepEmbedder>();
+            blocks["y_embedder_x"]          = std::make_shared<Conv2d>(HIDDEN_SIZE, 32 * PATCH_SIZE * PATCH_SIZE, std::pair{1, 1});
+            blocks["x_embedder.embedder.0"] = std::make_shared<Linear>(3 + 32 + 64, 32);
+            blocks["s_embedder.proj1"]      = std::make_shared<Conv2d>(3, LATENT_CHANNELS, std::pair{PATCH_SIZE, PATCH_SIZE}, std::pair{PATCH_SIZE, PATCH_SIZE}, std::pair{0, 0}, std::pair{1, 1}, false);
+            blocks["s_embedder.proj2"]      = std::make_shared<Conv2d>(LATENT_CHANNELS + HIDDEN_SIZE, HIDDEN_SIZE, std::pair{1, 1});
+            for (int i = 0; i < 21; ++i) {
+                blocks["blocks." + std::to_string(i)] = std::make_shared<DiCoBlock>(HIDDEN_SIZE);
+            }
+            blocks["dec_net.cond_embed"] = std::make_shared<Linear>(HIDDEN_SIZE, PATCH_SIZE * PATCH_SIZE * 32);
+            blocks["dec_net.input_proj"] = std::make_shared<Linear>(32, 32);
+            for (int i = 0; i < 3; ++i) {
+                blocks["dec_net.res_blocks." + std::to_string(i)] = std::make_shared<MLPResBlock>();
+            }
+            blocks["final_layer.norm"]   = std::make_shared<RMSNorm>(32);
+            blocks["final_layer.linear"] = std::make_shared<Linear>(32, 3);
+            blocks["y_embedder.decoder"] = std::make_shared<Decoder>();
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* latent, ggml_tensor* dct) {
+            auto cond      = std::dynamic_pointer_cast<Decoder>(blocks["y_embedder.decoder"])->forward(ctx, latent);
+            int64_t w      = cond->ne[0];
+            int64_t h      = cond->ne[1];
+            int64_t n      = cond->ne[3];
+            int64_t length = w * h;
+
+            auto image = ggml_ext_zeros(ctx->ggml_ctx, w * PATCH_SIZE, h * PATCH_SIZE, 3, n);
+            auto t     = ggml_ext_zeros(ctx->ggml_ctx, n, 1, 1, 1);
+            t          = ggml_reshape_1d(ctx->ggml_ctx, t, n);
+            auto c     = std::dynamic_pointer_cast<TimestepEmbedder>(blocks["t_embedder"])->forward(ctx, t);
+
+            auto s0 = std::dynamic_pointer_cast<Conv2d>(blocks["s_embedder.proj1"])->forward(ctx, image);
+            s0      = ggml_concat(ctx->ggml_ctx, s0, cond, 2);
+            auto s  = std::dynamic_pointer_cast<Conv2d>(blocks["s_embedder.proj2"])->forward(ctx, s0);
+            for (int i = 0; i < 21; ++i) {
+                s = std::dynamic_pointer_cast<DiCoBlock>(blocks["blocks." + std::to_string(i)])->forward(ctx, s, c);
+            }
+            // [N, C, H, W] -> [N*H*W, C].
+            s = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, s, 1, 2, 0, 3));
+            s = ggml_reshape_2d(ctx->ggml_ctx, s, HIDDEN_SIZE, length * n);
+
+            auto y = std::dynamic_pointer_cast<Conv2d>(blocks["y_embedder_x"])->forward(ctx, cond);
+            // Split 32*P channels as [32, P], then produce [N*L, P, 32].
+            y          = ggml_reshape_4d(ctx->ggml_ctx, y, length, PATCH_SIZE * PATCH_SIZE, 32, n);
+            y          = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, y, 2, 1, 0, 3));
+            y          = ggml_reshape_3d(ctx->ggml_ctx, y, 32, PATCH_SIZE * PATCH_SIZE, length * n);
+            auto zeros = ggml_ext_zeros(ctx->ggml_ctx, 3, PATCH_SIZE * PATCH_SIZE, length * n, 1);
+            dct        = ggml_repeat_4d(ctx->ggml_ctx, dct, 64, PATCH_SIZE * PATCH_SIZE, length * n, 1);
+            auto x     = ggml_concat(ctx->ggml_ctx, zeros, y, 0);
+            x          = ggml_concat(ctx->ggml_ctx, x, dct, 0);
+            x          = std::dynamic_pointer_cast<Linear>(blocks["x_embedder.embedder.0"])->forward(ctx, x);
+            x          = std::dynamic_pointer_cast<Linear>(blocks["dec_net.input_proj"])->forward(ctx, x);
+
+            auto dec_cond = std::dynamic_pointer_cast<Linear>(blocks["dec_net.cond_embed"])->forward(ctx, s);
+            dec_cond      = ggml_reshape_3d(ctx->ggml_ctx, dec_cond, 32, PATCH_SIZE * PATCH_SIZE, length * n);
+            for (int i = 0; i < 3; ++i) {
+                x = std::dynamic_pointer_cast<MLPResBlock>(blocks["dec_net.res_blocks." + std::to_string(i)])->forward(ctx, x, dec_cond);
+            }
+            x = std::dynamic_pointer_cast<RMSNorm>(blocks["final_layer.norm"])->forward(ctx, x);
+            x = std::dynamic_pointer_cast<Linear>(blocks["final_layer.linear"])->forward(ctx, x);
+            // [N*L, P, 3] -> [N, L, 3*P] for fold/unpatchify.
+            x = ggml_reshape_4d(ctx->ggml_ctx, x, 3, PATCH_SIZE * PATCH_SIZE, length, n);
+            x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3));
+            x = ggml_reshape_3d(ctx->ggml_ctx, x, 3 * PATCH_SIZE * PATCH_SIZE, length, n);
+            return DiT::unpatchify(ctx->ggml_ctx, x, h, w, PATCH_SIZE, PATCH_SIZE);
+        }
+    };
+
+    struct MageVAEModel : public GGMLBlock {
+        MageVAEModel() {
+            blocks["student.dconv_encoder"] = std::make_shared<DConvEncoder>();
+            blocks["pipeline"]              = std::make_shared<DConvDenoiser>();
+        }
+
+        ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* image) {
+            return std::dynamic_pointer_cast<DConvEncoder>(blocks["student.dconv_encoder"])->forward(ctx, image);
+        }
+
+        ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* latent, ggml_tensor* dct) {
+            return std::dynamic_pointer_cast<DConvDenoiser>(blocks["pipeline"])->forward(ctx, latent, dct);
+        }
+    };
+
+    struct MageVAERunner : public VAE {
+        MageVAEModel model;
+        std::vector<float> dct_vec;
+
+        MageVAERunner(ggml_backend_t backend,
+                      const String2TensorStorage& tensor_storage_map,
+                      const std::string& prefix,
+                      std::shared_ptr<RunnerWeightManager> weight_manager = nullptr)
+            : VAE(VERSION_MAGE_FLOW, backend, prefix, weight_manager) {
+            model = MageVAEModel();
+            model.init(params_ctx, tensor_storage_map, prefix);
+            dct_vec.resize(64 * PATCH_SIZE * PATCH_SIZE);
+            constexpr float pi = 3.14159265358979323846f;
+            for (int py = 0; py < PATCH_SIZE; ++py) {
+                float y = static_cast<float>(py) / static_cast<float>(PATCH_SIZE - 1);
+                for (int px = 0; px < PATCH_SIZE; ++px) {
+                    float x = static_cast<float>(px) / static_cast<float>(PATCH_SIZE - 1);
+                    int pos = py * PATCH_SIZE + px;
+                    for (int fy = 0; fy < 8; ++fy) {
+                        for (int fx = 0; fx < 8; ++fx) {
+                            int freq                 = fx * 8 + fy;
+                            float freq_x             = static_cast<float>(fx) * 8.f / 7.f;
+                            float freq_y             = static_cast<float>(fy) * 8.f / 7.f;
+                            float coeff              = 1.f / (1.f + freq_x * freq_y);
+                            dct_vec[freq + 64 * pos] = std::cos(x * freq_x * pi) *
+                                                       std::cos(y * freq_y * pi) * coeff;
+                        }
+                    }
+                }
+            }
+        }
+
+        std::string get_desc() override {
+            return "mage_vae";
+        }
+
+        void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) override {
+            model.get_param_tensors(tensors, weight_prefix);
+        }
+
+        ggml_cgraph* build_graph(const sd::Tensor<float>& input_tensor, bool decode_graph) {
+            ggml_cgraph* gf  = new_graph_custom(MAGE_VAE_GRAPH_SIZE);
+            auto input       = make_input(input_tensor);
+            auto runner_ctx  = get_context();
+            ggml_tensor* dct = nullptr;
+            if (decode_graph) {
+                dct = ggml_new_tensor_3d(compute_ctx, GGML_TYPE_F32, 64, PATCH_SIZE * PATCH_SIZE, 1);
+                set_backend_tensor_data(dct, dct_vec.data());
+            }
+            auto out = decode_graph ? model.decode(&runner_ctx, input, dct) : model.encode(&runner_ctx, input);
+            ggml_build_forward_expand(gf, out);
+            return gf;
+        }
+
+        sd::Tensor<float> _compute(const int n_threads,
+                                   const sd::Tensor<float>& input,
+                                   bool decode_graph) override {
+            auto get_graph = [&]() -> ggml_cgraph* {
+                return build_graph(input, decode_graph);
+            };
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false, false, false), input.dim());
+        }
+
+        int get_encoder_output_channels(int input_channels) override {
+            SD_UNUSED(input_channels);
+            return LATENT_CHANNELS * 2;
+        }
+
+        sd::Tensor<float> vae_output_to_latents(const sd::Tensor<float>& vae_output, std::shared_ptr<RNG> rng) override {
+            const auto chunks         = sd::ops::chunk(vae_output, 2, 2);
+            const auto& mean          = chunks[0];
+            const auto& logvar        = chunks[1];
+            sd::Tensor<float> stddev  = sd::ops::exp(0.5f * sd::ops::clamp(logvar, -20.0f, 10.0f));
+            sd::Tensor<float> noise   = sd::Tensor<float>::randn_like(mean, rng);
+            sd::Tensor<float> latents = mean + stddev * noise;
+            return latents;
+        }
+
+        sd::Tensor<float> diffusion_to_vae_latents(const sd::Tensor<float>& latents) override {
+            return latents;
+        }
+
+        sd::Tensor<float> vae_to_diffusion_latents(const sd::Tensor<float>& latents) override {
+            return latents;
+        }
+    };
+}  // namespace MageVAE
+
+#endif  // __SD_MODEL_VAE_MAGE_VAE_HPP__

src/model/vae/tae.hpp

@@ -528,6 +528,9 @@ public:
         if (version == VERSION_WAN2_2_TI2V) {
             z_channels = 48;
             patch      = 2;
+        } else if (sd_version_is_hunyuan_video(version)) {
+            z_channels = 32;
+            patch      = 2;
         } else if (sd_version_is_ltxav(version)) {
             z_channels     = 128;
             patch          = 4;
@@ -542,12 +545,12 @@ public:
 
     ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* z) {
         auto decoder = std::dynamic_pointer_cast<TinyVideoDecoder>(blocks["decoder"]);
-        if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) {
+        if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) {
             // (W, H, C, T) -> (W, H, T, C)
             z = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, z, 0, 1, 3, 2));
         }
         auto result = decoder->forward(ctx, z);
-        if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) {
+        if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) {
             // (W, H, T, C) -> (W, H, C, T)
             result = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, result, 0, 1, 3, 2));
         }
@@ -556,7 +559,7 @@ public:
 
     ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* x) {
         auto encoder = std::dynamic_pointer_cast<TinyVideoEncoder>(blocks["encoder"]);
-        if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) {
+        if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) {
             // (W, H, T, C) -> (W, H, C, T)
             x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
         }
@@ -569,7 +572,7 @@ public:
             }
         }
         x = encoder->forward(ctx, x);
-        if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) {
+        if (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) {
             // (W, H, C, T) -> (W, H, T, C)
             x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
         }

src/model/vae/vae.hpp

@@ -74,7 +74,7 @@ public:
         int scale_factor = 8;
         if (version == VERSION_LTXAV) {
             scale_factor = 32;
-        } else if (version == VERSION_WAN2_2_TI2V) {
+        } else if (version == VERSION_WAN2_2_TI2V || sd_version_is_hunyuan_video(version) || sd_version_is_mage_flow(version)) {
             scale_factor = 16;
         } else if (sd_version_uses_flux2_vae(version)) {
             scale_factor = 16;
@@ -136,7 +136,7 @@ public:
             // Image VAE encode is more sensitive to tile boundary context than decode.
             // Keep the smaller legacy factor for video VAEs, but default image encode
             // tiles to 64 latent pixels so a 512px SD image is encoded as one tile.
-            const float encode_tile_factor = (sd_version_is_wan(version) || sd_version_is_ltxav(version)) ? 1.30539f : 2.0f;
+            const float encode_tile_factor = (sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) ? 1.30539f : 2.0f;
             get_tile_sizes(tile_size_x, tile_size_y, tile_overlap, tiling_params, W, H, encode_tile_factor);
             LOG_DEBUG("VAE Tile size: %dx%d", tile_size_x, tile_size_y);
             output = tiled_compute(input,

src/model_io/gguf_io.h

@@ -21,7 +21,7 @@ bool write_gguf_file(const std::string& file_path,
 class GGUFStreamingWriter : public StreamingModelWriter {
 public:
     GGUFStreamingWriter() = default;
-    ~GGUFStreamingWriter();
+    ~GGUFStreamingWriter() override;
 
     bool write_metadata(const std::string& file_path,
                         const std::vector<TensorWritePlan>& tensors,

src/model_io/pickle_io.cpp

@@ -2,6 +2,7 @@
 
 #include <cstdlib>
 #include <cstring>
+#include <limits>
 #include <string>
 #include <unordered_map>
 #include <utility>
@@ -512,8 +513,51 @@ static bool parse_storage_type(const std::string& global_name, PickleStorageInfo
     return false;
 }
 
-static bool tensor_is_contiguous(const PickleTensorInfo& tensor) {
-    if (tensor.tensor_storage.nelements() == 0) {
+static bool checked_pickle_byte_count(int64_t element_count,
+                                      uint64_t element_nbytes,
+                                      uint64_t* byte_count) {
+    if (element_count < 0 || element_nbytes == 0) {
+        return false;
+    }
+
+    uint64_t count = static_cast<uint64_t>(element_count);
+    if (count > std::numeric_limits<uint64_t>::max() / element_nbytes) {
+        return false;
+    }
+
+    *byte_count = count * element_nbytes;
+    return true;
+}
+
+static bool tensor_layout_is_valid(const PickleTensorInfo& tensor, uint64_t raw_element_nbytes) {
+    if (raw_element_nbytes == 0) {
+        return false;
+    }
+
+    bool has_zero_dimension = false;
+    uint64_t element_count  = 1;
+    for (int i = 0; i < tensor.tensor_storage.n_dims; ++i) {
+        int64_t dimension = tensor.tensor_storage.ne[i];
+        if (dimension < 0) {
+            return false;
+        }
+        if (dimension == 0) {
+            has_zero_dimension = true;
+            continue;
+        }
+
+        uint64_t size = static_cast<uint64_t>(dimension);
+        if (element_count > static_cast<uint64_t>(std::numeric_limits<int64_t>::max()) / size) {
+            return false;
+        }
+        element_count *= size;
+    }
+
+    if (!has_zero_dimension &&
+        element_count > static_cast<uint64_t>(std::numeric_limits<int64_t>::max()) / raw_element_nbytes) {
+        return false;
+    }
+    if (has_zero_dimension) {
         return true;
     }
     if (tensor.stride_n_dims != tensor.tensor_storage.n_dims) {
@@ -932,7 +976,12 @@ bool parse_torch_state_dict_pickle(const uint8_t* buffer,
                 if (storage.key.empty() || !parse_storage_type(pid.items[1].str_value, &storage)) {
                     return false;
                 }
-                storage.nbytes              = (uint64_t)pid.items[4].int_value * storage.raw_element_nbytes;
+                if (!checked_pickle_byte_count(pid.items[4].int_value,
+                                               storage.raw_element_nbytes,
+                                               &storage.nbytes)) {
+                    set_error(error, "invalid storage size in torch pickle");
+                    return false;
+                }
                 storage_nbytes[storage.key] = storage.nbytes;
                 stack.push_back(make_storage_value(storage));
             } break;
@@ -963,7 +1012,12 @@ bool parse_torch_state_dict_pickle(const uint8_t* buffer,
                     tensor.tensor_storage.is_f64      = args.items[0].storage.is_f64;
                     tensor.tensor_storage.is_i64      = args.items[0].storage.is_i64;
                     tensor.tensor_storage.storage_key = args.items[0].storage.key;
-                    tensor.tensor_storage.offset      = (uint64_t)args.items[1].int_value * args.items[0].storage.raw_element_nbytes;
+                    if (!checked_pickle_byte_count(args.items[1].int_value,
+                                                   args.items[0].storage.raw_element_nbytes,
+                                                   &tensor.tensor_storage.offset)) {
+                        set_error(error, "invalid tensor storage offset in torch pickle");
+                        return false;
+                    }
 
                     for (const auto& item : args.items[2].items) {
                         if (item.kind != PickleValue::INT || tensor.tensor_storage.n_dims >= SD_MAX_DIMS) {
@@ -979,7 +1033,8 @@ bool parse_torch_state_dict_pickle(const uint8_t* buffer,
                         tensor.stride[tensor.stride_n_dims++] = item.int_value;
                     }
 
-                    if (!tensor_is_contiguous(tensor)) {
+                    if (!tensor_layout_is_valid(tensor, args.items[0].storage.raw_element_nbytes)) {
+                        set_error(error, "invalid tensor shape or stride in torch pickle");
                         return false;
                     }
                     stack.push_back(make_tensor_value(tensor));

src/model_io/torch_legacy_io.cpp

@@ -139,11 +139,16 @@ bool read_torch_legacy_file(const std::string& file_path,
             if (it == legacy_storage_map.end()) {
                 return false;
             }
-            if (current_offset + LEGACY_STORAGE_HEADER_SIZE + it->second > file_size) {
+            if (current_offset > file_size ||
+                LEGACY_STORAGE_HEADER_SIZE > file_size - current_offset) {
                 return false;
             }
-            storage_offsets[storage_key] = current_offset + LEGACY_STORAGE_HEADER_SIZE;
-            current_offset += LEGACY_STORAGE_HEADER_SIZE + it->second;
+            uint64_t storage_offset = current_offset + LEGACY_STORAGE_HEADER_SIZE;
+            if (it->second > file_size - storage_offset) {
+                return false;
+            }
+            storage_offsets[storage_key] = storage_offset;
+            current_offset               = storage_offset + it->second;
         }
 
         for (auto& tensor_storage : tensor_storages) {
@@ -159,8 +164,10 @@ bool read_torch_legacy_file(const std::string& file_path,
 
             uint64_t base_offset    = it_offset->second;
             uint64_t storage_nbytes = it_size->second;
-            uint64_t tensor_nbytes  = tensor_storage.nbytes_to_read();
-            if (tensor_storage.offset + tensor_nbytes > storage_nbytes) {
+            int64_t tensor_nbytes   = tensor_storage.nbytes_to_read();
+            if (tensor_nbytes < 0 ||
+                tensor_storage.offset > storage_nbytes ||
+                static_cast<uint64_t>(tensor_nbytes) > storage_nbytes - tensor_storage.offset) {
                 return false;
             }
 

src/model_io/torch_zip_io.cpp

@@ -76,8 +76,10 @@ static bool parse_zip_data_pkl(const uint8_t* buffer,
             return false;
         }
 
-        uint64_t tensor_nbytes = tensor_storage.nbytes_to_read();
-        if (tensor_storage.offset + tensor_nbytes > entry_size) {
+        int64_t tensor_nbytes = tensor_storage.nbytes_to_read();
+        if (tensor_nbytes < 0 ||
+            tensor_storage.offset > entry_size ||
+            static_cast<uint64_t>(tensor_nbytes) > entry_size - tensor_storage.offset) {
             set_error(error, "tensor '" + tensor_storage.name + "' exceeds storage entry '" + entry_name + "'");
             return false;
         }

src/model_loader.cpp

@@ -498,11 +498,18 @@ SDVersion ModelLoader::get_sd_version() {
             return VERSION_MINIT2I;
         }
         if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.img_mod.1.weight") != std::string::npos) {
+            auto img_in = tensor_storage_map.find("model.diffusion_model.img_in.weight");
+            if (img_in != tensor_storage_map.end() && img_in->second.ne[0] == 128) {
+                return VERSION_MAGE_FLOW;
+            }
             if (tensor_storage_map.find("model.diffusion_model.time_text_embed.addition_t_embedding.weight") != tensor_storage_map.end()) {
                 return VERSION_QWEN_IMAGE_LAYERED;
             }
             return VERSION_QWEN_IMAGE;
         }
+        if (tensor_storage.name.find("model.diffusion_model.txt_in.individual_token_refiner.blocks.0.adaLN_modulation.1.weight") != std::string::npos) {
+            return VERSION_HUNYUAN_VIDEO;
+        }
         if (tensor_storage.name.find("llm_adapter.blocks.0.cross_attn.q_proj.weight") != std::string::npos) {
             return VERSION_ANIMA;
         }
@@ -1046,7 +1053,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
         std::atomic<size_t> tensor_idx(0);
         std::atomic<bool> failed(false);
         std::vector<std::thread> workers;
-        std::mutex rpc_backend_mutex;
+        std::mutex backend_tensor_set_mutex;
 
         for (int i = 0; i < n_threads; ++i) {
             workers.emplace_back([&, file_path, is_zip]() {
@@ -1070,6 +1077,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
 
                 std::vector<uint8_t> read_buffer;
                 std::vector<uint8_t> convert_buffer;
+                std::vector<uint8_t> zip_entry_buffer;
 
                 while (true) {
                     int64_t t0, t1;
@@ -1108,34 +1116,60 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
 
                     size_t nbytes_to_read = tensor_storage.nbytes_to_read();
 
-                    auto read_data = [&](char* buf, size_t n) {
+                    auto read_data = [&](char* buf, size_t n) -> bool {
                         if (zip != nullptr) {
-                            zip_entry_openbyindex(zip, tensor_storage.index_in_zip);
+                            if (zip_entry_openbyindex(zip, tensor_storage.index_in_zip) != 0) {
+                                LOG_ERROR("failed to open zip entry for tensor '%s'", tensor_storage.name.c_str());
+                                return false;
+                            }
                             size_t entry_size = zip_entry_size(zip);
+                            if (tensor_storage.offset > entry_size) {
+                                LOG_ERROR("tensor '%s' exceeds its zip storage entry", tensor_storage.name.c_str());
+                                zip_entry_close(zip);
+                                return false;
+                            }
+                            size_t tensor_offset = static_cast<size_t>(tensor_storage.offset);
+                            if (n > entry_size - tensor_offset) {
+                                LOG_ERROR("tensor '%s' exceeds its zip storage entry", tensor_storage.name.c_str());
+                                zip_entry_close(zip);
+                                return false;
+                            }
+
                             if (entry_size != n) {
                                 int64_t t_memcpy_start;
-                                read_buffer.resize(entry_size);
-                                zip_entry_noallocread(zip, (void*)read_buffer.data(), entry_size);
+                                zip_entry_buffer.resize(entry_size);
+                                auto bytes_read = zip_entry_noallocread(zip, (void*)zip_entry_buffer.data(), entry_size);
+                                if (bytes_read < 0 || static_cast<size_t>(bytes_read) != entry_size) {
+                                    LOG_ERROR("failed to read zip entry for tensor '%s'", tensor_storage.name.c_str());
+                                    zip_entry_close(zip);
+                                    return false;
+                                }
                                 t_memcpy_start = ggml_time_ms();
-                                memcpy((void*)buf, (void*)(read_buffer.data() + tensor_storage.offset), n);
+                                memcpy((void*)buf, (void*)(zip_entry_buffer.data() + tensor_offset), n);
                                 memcpy_time_ms.fetch_add(ggml_time_ms() - t_memcpy_start);
                             } else {
-                                zip_entry_noallocread(zip, (void*)buf, n);
+                                auto bytes_read = zip_entry_noallocread(zip, (void*)buf, n);
+                                if (bytes_read < 0 || static_cast<size_t>(bytes_read) != n) {
+                                    LOG_ERROR("failed to read zip entry for tensor '%s'", tensor_storage.name.c_str());
+                                    zip_entry_close(zip);
+                                    return false;
+                                }
                             }
                             zip_entry_close(zip);
                         } else if (mmapped) {
                             if (!mmapped->copy_data(buf, n, tensor_storage.offset)) {
                                 LOG_ERROR("read tensor data failed: '%s'", file_path.c_str());
-                                failed = true;
+                                return false;
                             }
                         } else {
                             file.seekg(tensor_storage.offset);
                             file.read(buf, n);
                             if (!file) {
                                 LOG_ERROR("read tensor data failed: '%s'", file_path.c_str());
-                                failed = true;
+                                return false;
                             }
                         }
+                        return true;
                     };
 
                     char* read_buf    = nullptr;
@@ -1169,7 +1203,10 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
                     }
 
                     t0 = ggml_time_ms();
-                    read_data(read_buf, nbytes_to_read);
+                    if (!read_data(read_buf, nbytes_to_read)) {
+                        failed = true;
+                        break;
+                    }
                     t1 = ggml_time_ms();
                     read_time_ms.fetch_add(t1 - t0);
 
@@ -1207,17 +1244,8 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb,
                     if (dst_tensor->buffer != nullptr && !ggml_backend_buffer_is_host(dst_tensor->buffer)) {
                         t0 = ggml_time_ms();
 
-                        // RPC backends require serialized access to prevent concurrency issues
-                        const char* buffer_type_name = ggml_backend_buft_name(ggml_backend_buffer_get_type(dst_tensor->buffer));
-                        bool is_rpc_buffer           = buffer_type_name != nullptr &&
-                                             std::string(buffer_type_name).find("RPC") != std::string::npos;
-
-                        if (is_rpc_buffer) {
-                            std::lock_guard<std::mutex> lock(rpc_backend_mutex);
-                            ggml_backend_tensor_set(dst_tensor, convert_buf, 0, ggml_nbytes(dst_tensor));
-                        } else {
-                            ggml_backend_tensor_set(dst_tensor, convert_buf, 0, ggml_nbytes(dst_tensor));
-                        }
+                        std::lock_guard<std::mutex> lock(backend_tensor_set_mutex);
+                        ggml_backend_tensor_set(dst_tensor, convert_buf, 0, ggml_nbytes(dst_tensor));
 
                         t1 = ggml_time_ms();
                         copy_to_backend_time_ms.fetch_add(t1 - t0);

src/model_manager.cpp

@@ -53,6 +53,48 @@ static bool backend_supports_host_buffer(ggml_backend_t backend) {
     return props.caps.buffer_from_host_ptr;
 }
 
+static bool device_supports_param_op(ggml_backend_dev_t device,
+                                     ggml_tensor* weight,
+                                     enum ggml_op op,
+                                     ggml_backend_buffer_type_t buft) {
+    if (op == GGML_OP_NONE) {
+        return true;
+    }
+    if (device == nullptr || weight == nullptr || buft == nullptr || weight->buffer != nullptr) {
+        return false;
+    }
+
+    ggml_init_params params;
+    params.mem_size   = ggml_tensor_overhead() * 2;
+    params.mem_buffer = nullptr;
+    params.no_alloc   = true;
+    ggml_context* ctx = ggml_init(params);
+    if (ctx == nullptr) {
+        return false;
+    }
+
+    ggml_tensor* op_tensor = nullptr;
+    if (op == GGML_OP_GET_ROWS) {
+        ggml_tensor* indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1);
+        op_tensor            = ggml_get_rows(ctx, weight, indices);
+    }
+    if (op_tensor == nullptr) {
+        ggml_free(ctx);
+        return false;
+    }
+
+    weight->buffer = ggml_backend_buft_alloc_buffer(buft, 0);
+    if (weight->buffer == nullptr) {
+        ggml_free(ctx);
+        return false;
+    }
+    bool supported = ggml_backend_dev_supports_op(device, op_tensor);
+    ggml_backend_buffer_free(weight->buffer);
+    weight->buffer = nullptr;
+    ggml_free(ctx);
+    return supported;
+}
+
 ModelManager::~ModelManager() {
     release_all();
 }
@@ -135,7 +177,8 @@ bool ModelManager::register_param_tensors(const std::string& desc,
                                           ggml_backend_t params_backend,
                                           size_t* registered_tensor_size,
                                           bool allow_split_buffer,
-                                          bool params_follow_compute_backend) {
+                                          bool params_follow_compute_backend,
+                                          const std::map<ggml_tensor*, enum ggml_op>* tensor_ops) {
     if (desc.empty()) {
         LOG_ERROR("model manager tensor desc is empty");
         return false;
@@ -168,6 +211,12 @@ bool ModelManager::register_param_tensors(const std::string& desc,
         state->params_backend                = params_backend;
         state->allow_split_buffer            = allow_split_buffer;
         state->params_follow_compute_backend = params_follow_compute_backend;
+        if (tensor_ops != nullptr) {
+            auto op_it = tensor_ops->find(tensor);
+            if (op_it != tensor_ops->end()) {
+                state->usage_op = op_it->second;
+            }
+        }
         new_states.push_back(std::move(state));
     }
 
@@ -844,6 +893,22 @@ ggml_backend_buffer_type_t ModelManager::params_buffer_type_for(const TensorStat
     if (params_buft == nullptr) {
         params_buft = ggml_backend_get_default_buffer_type(state.params_backend);
     }
+    if (state.usage_op != GGML_OP_NONE &&
+        state.compute_backend != nullptr) {
+        ggml_backend_dev_t compute_dev = ggml_backend_get_device(state.compute_backend);
+        if (device_supports_param_op(compute_dev, state.tensor, state.usage_op, params_buft)) {
+            return params_buft;
+        }
+
+        ggml_backend_dev_t cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
+        params_buft                = cpu_dev != nullptr ? ggml_backend_dev_buffer_type(cpu_dev) : nullptr;
+        if (!device_supports_param_op(cpu_dev, state.tensor, state.usage_op, params_buft)) {
+            LOG_ERROR("model manager has no compatible buffer for tensor '%s' used by %s",
+                      state.name.c_str(),
+                      ggml_op_name(state.usage_op));
+            return nullptr;
+        }
+    }
     return params_buft;
 }
 

src/model_manager.h

@@ -39,6 +39,7 @@ private:
         bool allow_split_buffer            = false;
         bool params_follow_compute_backend = false;
         bool metadata_validated            = false;
+        enum ggml_op usage_op              = GGML_OP_NONE;
 
         int active_prepare_count = 0;
 
@@ -130,9 +131,10 @@ public:
                                 ResidencyMode residency_mode,
                                 ggml_backend_t compute_backend,
                                 ggml_backend_t params_backend,
-                                size_t* registered_tensor_size     = nullptr,
-                                bool allow_split_buffer            = false,
-                                bool params_follow_compute_backend = false);
+                                size_t* registered_tensor_size                         = nullptr,
+                                bool allow_split_buffer                                = false,
+                                bool params_follow_compute_backend                     = false,
+                                const std::map<ggml_tensor*, enum ggml_op>* tensor_ops = nullptr);
 
     bool unregister_param_tensors(const std::string& desc,
                                   size_t* registered_tensor_size = nullptr);

src/name_conversion.cpp

@@ -664,6 +664,72 @@ std::string convert_diffusers_dit_to_original_flux(std::string name) {
     return name;
 }
 
+std::string convert_hunyuan_video_to_original_flux(std::string name) {
+    int num_layers        = 54;
+    int num_single_layers = 0;
+    static std::unordered_map<std::string, std::string> hy_name_map;
+
+    if (hy_name_map.empty()) {
+        // --- double transformer blocks ---
+        for (int i = 0; i < num_layers; ++i) {
+            std::string block_prefix = "double_blocks." + std::to_string(i) + ".";
+            std::string dst_prefix   = "double_blocks." + std::to_string(i) + ".";
+
+            hy_name_map[block_prefix + "img_mod.linear"] = dst_prefix + "img_mod.lin";
+            hy_name_map[block_prefix + "txt_mod.linear"] = dst_prefix + "txt_mod.lin";
+
+            // attn
+            hy_name_map[block_prefix + "img_attn_qkv"] = dst_prefix + "img_attn.qkv";
+            hy_name_map[block_prefix + "txt_attn_qkv"] = dst_prefix + "txt_attn.qkv";
+
+            // norm
+            hy_name_map[block_prefix + "img_attn_q_norm.weight"] = dst_prefix + "img_attn.norm.query_norm.scale";
+            hy_name_map[block_prefix + "img_attn_k_norm.weight"] = dst_prefix + "img_attn.norm.key_norm.scale";
+            hy_name_map[block_prefix + "txt_attn_q_norm.weight"] = dst_prefix + "txt_attn.norm.query_norm.scale";
+            hy_name_map[block_prefix + "txt_attn_k_norm.weight"] = dst_prefix + "txt_attn.norm.key_norm.scale";
+
+            // ff
+            hy_name_map[block_prefix + "img_mlp.fc1"] = dst_prefix + "img_mlp.0";
+            hy_name_map[block_prefix + "img_mlp.fc2"] = dst_prefix + "img_mlp.2";
+
+            hy_name_map[block_prefix + "txt_mlp.fc1"] = dst_prefix + "txt_mlp.0";
+            hy_name_map[block_prefix + "txt_mlp.fc2"] = dst_prefix + "txt_mlp.2";
+
+            // output projections
+            hy_name_map[block_prefix + "img_attn_proj"] = dst_prefix + "img_attn.proj";
+            hy_name_map[block_prefix + "txt_attn_proj"] = dst_prefix + "txt_attn.proj";
+        }
+    }
+
+    hy_name_map["time_in.mlp.0"]     = "time_in.in_layer";
+    hy_name_map["time_in.mlp.2"]     = "time_in.out_layer";
+    hy_name_map["time_r_in.mlp.0"]   = "time_r_in.in_layer";
+    hy_name_map["time_r_in.mlp.2"]   = "time_r_in.out_layer";
+    hy_name_map["vector_in.mlp.0"]   = "vector_in.in_layer";
+    hy_name_map["vector_in.mlp.2"]   = "vector_in.out_layer";
+    hy_name_map["guidance_in.mlp.0"] = "guidance_in.in_layer";
+    hy_name_map["guidance_in.mlp.2"] = "guidance_in.out_layer";
+
+    hy_name_map["txt_in.c_embedder.linear_1"] = "txt_in.c_embedder.in_layer";
+    hy_name_map["txt_in.c_embedder.linear_2"] = "txt_in.c_embedder.out_layer";
+
+    hy_name_map["txt_in.t_embedder.mlp.0"] = "txt_in.t_embedder.in_layer";
+    hy_name_map["txt_in.t_embedder.mlp.2"] = "txt_in.t_embedder.out_layer";
+
+    replace_with_prefix_map(name, hy_name_map);
+
+    static const std::vector<std::pair<std::string, std::string>> generic_name_map = {
+        {"_attn_qkv.", "_attn.qkv."},
+        {"_attn_proj.", "_attn.proj."},
+        {"mlp.fc1.", "mlp.0."},
+        {"mlp.fc2.", "mlp.2."},
+        {".modulation.linear.", ".modulation.lin."},
+    };
+    replace_with_name_map(name, generic_name_map);
+
+    return name;
+}
+
 std::string convert_diffusers_dit_to_original_lumina2(std::string name) {
     int num_layers         = 30;
     int num_refiner_layers = 2;
@@ -807,6 +873,8 @@ std::string convert_diffusion_model_name(std::string name, std::string prefix, S
         name = convert_diffusers_dit_to_original_sd3(name);
     } else if (sd_version_is_flux(version) || sd_version_is_flux2(version) || sd_version_is_longcat(version) || sd_version_is_sefi_image(version)) {
         name = convert_diffusers_dit_to_original_flux(name);
+    } else if (sd_version_is_hunyuan_video(version)) {
+        name = convert_hunyuan_video_to_original_flux(name);
     } else if (sd_version_is_z_image(version)) {
         name = convert_diffusers_dit_to_original_lumina2(name);
     } else if (sd_version_is_anima(version)) {
@@ -980,6 +1048,9 @@ std::string convert_diffusers_to_original_wan_vae(std::string name) {
 }
 
 std::string convert_first_stage_model_name(std::string name, std::string prefix, SDVersion version) {
+    if (sd_version_is_hunyuan_video(version) || sd_version_is_mage_flow(version)) {
+        return name;
+    }
     if (sd_version_uses_wan_vae(version)) {
         return convert_diffusers_to_original_wan_vae(name);
     }
@@ -1214,11 +1285,54 @@ static std::string convert_esrgan_tensor_name(std::string name) {
     return name;
 }
 
+static const std::map<int, std::string>& ip_adapter_index_map(SDVersion version) {
+    static const std::map<int, std::string> sd15_map = {
+        {1, "input_blocks.1.1.transformer_blocks.0"}, {3, "input_blocks.2.1.transformer_blocks.0"}, {5, "input_blocks.4.1.transformer_blocks.0"}, {7, "input_blocks.5.1.transformer_blocks.0"}, {9, "input_blocks.7.1.transformer_blocks.0"}, {11, "input_blocks.8.1.transformer_blocks.0"}, {13, "output_blocks.3.1.transformer_blocks.0"}, {15, "output_blocks.4.1.transformer_blocks.0"}, {17, "output_blocks.5.1.transformer_blocks.0"}, {19, "output_blocks.6.1.transformer_blocks.0"}, {21, "output_blocks.7.1.transformer_blocks.0"}, {23, "output_blocks.8.1.transformer_blocks.0"}, {25, "output_blocks.9.1.transformer_blocks.0"}, {27, "output_blocks.10.1.transformer_blocks.0"}, {29, "output_blocks.11.1.transformer_blocks.0"}, {31, "middle_block.1.transformer_blocks.0"}};
+
+    static std::map<int, std::string> sdxl_map;
+    if (sdxl_map.empty()) {
+        std::vector<std::pair<std::string, int>> order = {
+            {"input_blocks.4.1", 2}, {"input_blocks.5.1", 2}, {"input_blocks.7.1", 10}, {"input_blocks.8.1", 10}, {"output_blocks.0.1", 10}, {"output_blocks.1.1", 10}, {"output_blocks.2.1", 10}, {"output_blocks.3.1", 2}, {"output_blocks.4.1", 2}, {"output_blocks.5.1", 2}, {"middle_block.1", 10}};
+        int idx = 1;
+        for (const auto& [block, depth] : order) {
+            for (int m = 0; m < depth; m++) {
+                sdxl_map[idx] = block + ".transformer_blocks." + std::to_string(m);
+                idx += 2;
+            }
+        }
+    }
+    return sd_version_is_sdxl(version) ? sdxl_map : sd15_map;
+}
+
+static std::string convert_ip_adapter_name(std::string name, SDVersion version) {
+    if (starts_with(name, "image_proj.")) {
+        return "ip_adapter." + name;
+    }
+    if (starts_with(name, "ip_adapter.")) {
+        auto items = split_string(name, '.');
+        if (items.size() < 4) {
+            return name;
+        }
+        int idx        = atoi(items[1].c_str());
+        const auto& mp = ip_adapter_index_map(version);
+        auto blk       = mp.find(idx);
+        if (blk == mp.end()) {
+            return name;
+        }
+        return "model.diffusion_model." + blk->second + ".attn2." + items[2] + "." + items[3];
+    }
+    return name;
+}
+
 std::string convert_tensor_name(std::string name, SDVersion version) {
     if (version == VERSION_ESRGAN) {
         return convert_esrgan_tensor_name(std::move(name));
     }
 
+    if (starts_with(name, "ip_adapter.") || starts_with(name, "image_proj.")) {
+        return convert_ip_adapter_name(std::move(name), version);
+    }
+
     bool is_lora                             = false;
     bool is_lycoris_underline                = false;
     bool is_underline                        = false;
@@ -1343,8 +1457,14 @@ 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));
+    if (sd_version_is_boogu_image(version) || sd_version_is_krea2(version) || sd_version_is_mage_flow(version)) {
+        const std::string hf_vision_prefix = "text_encoders.llm.model.visual.";
+        if (starts_with(name, hf_vision_prefix)) {
+            name = "text_encoders.llm.visual." + name.substr(hf_vision_prefix.size());
+        }
+        if (starts_with(name, "text_encoders.llm.visual.")) {
+            name = convert_qwen3_vl_vision_name(std::move(name));
+        }
     }
 
     // diffusion model

src/runtime/denoiser.hpp

@@ -1153,9 +1153,9 @@ struct CompVisDenoiser : public Denoiser {
         return {c_skip, c_out, c_in};
     }
 
-    virtual sd::Tensor<float> noise_scaling(float sigma,
-                                            const sd::Tensor<float>& noise,
-                                            const sd::Tensor<float>& latent) override {
+    sd::Tensor<float> noise_scaling(float sigma,
+                                    const sd::Tensor<float>& noise,
+                                    const sd::Tensor<float>& latent) override {
         GGML_ASSERT(noise.numel() == latent.numel());
         return latent + noise * sigma;
     }
@@ -1165,7 +1165,7 @@ struct CompVisDenoiser : public Denoiser {
         return latent;
     }
 
-    float noise_level_to_sigma(float noise_level) {
+    float noise_level_to_sigma(float noise_level) override {
         return noise_level / (1.0f - noise_level);
     }
 };
@@ -1256,7 +1256,7 @@ struct DiscreteFlowDenoiser : public Denoiser {
         return latent * (1.0f / (1.0f - sigma));
     }
 
-    float noise_level_to_sigma(float noise_level) {
+    float noise_level_to_sigma(float noise_level) override {
         return noise_level;
     }
 };
@@ -1396,7 +1396,7 @@ struct MiniT2IFlowDenoiser : public Denoiser {
         return latent;
     }
 
-    float noise_level_to_sigma(float noise_level) {
+    float noise_level_to_sigma(float noise_level) override {
         SD_UNUSED(noise_level);
         return 1.0f;
     }

src/stable-diffusion.cpp

@@ -22,6 +22,7 @@
 #include "conditioning/conditioner.hpp"
 #include "core/backend_fit.h"
 #include "extensions/generation_extension.h"
+#include "model/adapter/ip_adapter.hpp"
 #include "model/adapter/lora.hpp"
 #include "model/diffusion/anima.hpp"
 #include "model/diffusion/animatediff.hpp"
@@ -30,11 +31,13 @@
 #include "model/diffusion/ernie_image.hpp"
 #include "model/diffusion/flux.hpp"
 #include "model/diffusion/hidream_o1.hpp"
+#include "model/diffusion/hunyuan.hpp"
 #include "model/diffusion/ideogram4.hpp"
 #include "model/diffusion/krea2.hpp"
 #include "model/diffusion/lens.hpp"
 #include "model/diffusion/lingbot_video.hpp"
 #include "model/diffusion/ltxv.hpp"
+#include "model/diffusion/mage_flow.hpp"
 #include "model/diffusion/minit2i.hpp"
 #include "model/diffusion/mmdit.hpp"
 #include "model/diffusion/model.hpp"
@@ -46,8 +49,10 @@
 #include "model/upscaler/esrgan.hpp"
 #include "model/upscaler/ltx_latent_upscaler.hpp"
 #include "model/vae/auto_encoder_kl.hpp"
+#include "model/vae/hunyuan_vae.hpp"
 #include "model/vae/ltx_audio_vae.hpp"
 #include "model/vae/ltx_vae.hpp"
+#include "model/vae/mage_vae.hpp"
 #include "model/vae/tae.hpp"
 #include "model/vae/vae.hpp"
 #include "model/vae/wan_vae.hpp"
@@ -96,6 +101,7 @@ const char* model_version_to_str[] = {
     "LingBot Video",
     "Qwen Image",
     "Qwen Image Layered",
+    "Hunyuan Video",
     "Anima",
     "Flux.2",
     "Flux.2 klein",
@@ -112,6 +118,7 @@ const char* model_version_to_str[] = {
     "Ideogram 4",
     "SeFi-Image",
     "Krea2",
+    "Mage Flow",
     "ESRGAN",
 };
 
@@ -134,6 +141,8 @@ const char* sampling_methods_str[] = {
     "Euler CFG++",
     "Euler A CFG++",
     "Euler GE",
+    "DPM++ (2M) SDE",
+    "DPM++ (2M) SDE BT",
 };
 
 /*================================================== Helper Functions ================================================*/
@@ -142,6 +151,7 @@ static bool sd_version_supports_ref_latent_img_cfg(SDVersion version) {
     return version == VERSION_FLUX ||
            sd_version_is_flux2(version) ||
            sd_version_is_qwen_image(version) ||
+           sd_version_is_mage_flow(version) ||
            sd_version_is_longcat(version) ||
            sd_version_is_z_image(version) ||
            sd_version_is_boogu_image(version);
@@ -214,6 +224,10 @@ public:
     std::shared_ptr<VAE> preview_vae;
     std::shared_ptr<LTXV::LTXAudioVAERunner> audio_vae_model;
     std::shared_ptr<ControlNet> control_net;
+    std::shared_ptr<IPAdapter::IPAdapterRunner> ip_adapter;
+    sd::Tensor<float> ip_adapter_tokens;
+    sd::Tensor<float> ip_adapter_uncond_tokens;
+    float ip_adapter_strength = 1.0f;
     std::vector<std::shared_ptr<GenerationExtension>> generation_extensions;
     std::vector<std::shared_ptr<LoraModel>> runtime_lora_models;
     bool apply_lora_immediately = false;
@@ -306,7 +320,11 @@ public:
             return true;
         }
         std::map<std::string, ggml_tensor*> group_tensors;
+        std::map<ggml_tensor*, enum ggml_op> tensor_ops;
         model->get_param_tensors(group_tensors);
+        if constexpr (std::is_base_of_v<Conditioner, T>) {
+            model->get_param_tensor_ops(tensor_ops);
+        }
         if (model_manager == nullptr) {
             return true;
         }
@@ -323,6 +341,7 @@ public:
                                                                 module,
                                                                 module_backends,
                                                                 std::move(group_tensors),
+                                                                tensor_ops,
                                                                 residency_mode,
                                                                 params_mem_size);
                     }
@@ -331,6 +350,7 @@ public:
                                                               module,
                                                               module_backends,
                                                               std::move(group_tensors),
+                                                              tensor_ops,
                                                               residency_mode,
                                                               params_mem_size);
                 }
@@ -344,7 +364,10 @@ public:
                                                      residency_mode,
                                                      backend_for(module),
                                                      params_backend_for(module),
-                                                     params_mem_size);
+                                                     params_mem_size,
+                                                     false,
+                                                     false,
+                                                     &tensor_ops);
     }
 
     template <typename T>
@@ -353,6 +376,7 @@ public:
                                           SDBackendModule module,
                                           const std::vector<ggml_backend_t>& module_backends,
                                           std::map<std::string, ggml_tensor*> group_tensors,
+                                          const std::map<ggml_tensor*, enum ggml_op>& tensor_ops,
                                           ModelManager::ResidencyMode residency_mode,
                                           size_t* params_mem_size) {
         ggml_backend_t main_backend = module_backends[0];
@@ -364,6 +388,7 @@ public:
                                                       module,
                                                       module_backends,
                                                       std::move(group_tensors),
+                                                      tensor_ops,
                                                       residency_mode,
                                                       params_mem_size);
         };
@@ -438,7 +463,9 @@ public:
                                                    main_backend,
                                                    params_backend_for(module),
                                                    params_mem_size,
-                                                   /*allow_split_buffer=*/true)) {
+                                                   /*allow_split_buffer=*/true,
+                                                   false,
+                                                   &tensor_ops)) {
             return false;
         }
         return model_manager->register_param_tensors(desc,
@@ -446,7 +473,10 @@ public:
                                                      residency_mode,
                                                      main_backend,
                                                      params_backend_for(module),
-                                                     params_mem_size);
+                                                     params_mem_size,
+                                                     false,
+                                                     false,
+                                                     &tensor_ops);
     }
 
     // Register graph-cut layer-split tensors on the primary backend first.
@@ -458,6 +488,7 @@ public:
                                             SDBackendModule module,
                                             const std::vector<ggml_backend_t>& module_backends,
                                             std::map<std::string, ggml_tensor*> group_tensors,
+                                            const std::map<ggml_tensor*, enum ggml_op>& tensor_ops,
                                             ModelManager::ResidencyMode residency_mode,
                                             size_t* params_mem_size) {
         bool has_cpu_device = false;
@@ -479,7 +510,10 @@ public:
                                                          residency_mode,
                                                          module_backends[0],
                                                          params_backend_for(module),
-                                                         params_mem_size);
+                                                         params_mem_size,
+                                                         false,
+                                                         false,
+                                                         &tensor_ops);
         }
 
         model->set_runtime_backends(module_backends);
@@ -504,7 +538,8 @@ public:
                                                      initial_params_backend,
                                                      params_mem_size,
                                                      false,
-                                                     params_follow_runtime);
+                                                     params_follow_runtime,
+                                                     &tensor_ops);
     }
 
     bool unload_control_net() {
@@ -834,6 +869,13 @@ public:
             }
         }
 
+        if (strlen(SAFE_STR(sd_ctx_params->ip_adapter_path)) > 0) {
+            if (!model_loader.init_from_file(sd_ctx_params->ip_adapter_path)) {
+                LOG_ERROR("init ip-adapter model loader from file failed: '%s'", sd_ctx_params->ip_adapter_path);
+                return false;
+            }
+        }
+
         model_loader.convert_tensors_name();
 
         version = model_loader.get_sd_version();
@@ -1064,6 +1106,18 @@ public:
                                                                       tensor_storage_map,
                                                                       "model.diffusion_model",
                                                                       model_manager);
+            } else if (sd_version_is_hunyuan_video(version)) {
+                cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
+                                                                 tensor_storage_map,
+                                                                 version,
+                                                                 "",
+                                                                 false,
+                                                                 model_manager);
+                diffusion_model  = std::make_shared<Hunyuan::HunyuanVideoRunner>(backend_for(SDBackendModule::DIFFUSION),
+                                                                                tensor_storage_map,
+                                                                                "model.diffusion_model",
+                                                                                version,
+                                                                                model_manager);
             } else if (sd_version_is_wan(version)) {
                 cond_stage_model = std::make_shared<T5CLIPEmbedder>(backend_for(SDBackendModule::TE),
                                                                     tensor_storage_map,
@@ -1132,6 +1186,17 @@ public:
                                                                           version,
                                                                           model_manager,
                                                                           sd_ctx_params->model_args);
+            } else if (sd_version_is_mage_flow(version)) {
+                cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
+                                                                 tensor_storage_map,
+                                                                 version,
+                                                                 "",
+                                                                 true,
+                                                                 model_manager);
+                diffusion_model  = std::make_shared<MageFlow::MageFlowRunner>(backend_for(SDBackendModule::DIFFUSION),
+                                                                             tensor_storage_map,
+                                                                             "model.diffusion_model",
+                                                                             model_manager);
             } else if (sd_version_is_longcat(version)) {
                 cond_stage_model = std::make_shared<LLMEmbedder>(backend_for(SDBackendModule::TE),
                                                                  tensor_storage_map,
@@ -1264,12 +1329,39 @@ public:
                 }
             }
 
+            if (strlen(SAFE_STR(sd_ctx_params->ip_adapter_path)) > 0 && clip_vision == nullptr) {
+                if (!ensure_backend_pair(SDBackendModule::CLIP_VISION)) {
+                    return false;
+                }
+                clip_vision = std::make_shared<FrozenCLIPVisionEmbedder>(backend_for(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;
+                }
+            }
+
+            if (strlen(SAFE_STR(sd_ctx_params->ip_adapter_path)) > 0) {
+                ip_adapter = std::make_shared<IPAdapter::IPAdapterRunner>(backend_for(SDBackendModule::DIFFUSION),
+                                                                          tensor_storage_map,
+                                                                          "ip_adapter",
+                                                                          model_manager);
+                if (!register_runner_params("IP-Adapter",
+                                            ip_adapter,
+                                            SDBackendModule::DIFFUSION)) {
+                    return false;
+                }
+            }
+
             if (!ensure_backend_pair(SDBackendModule::VAE)) {
                 return false;
             }
 
             auto create_tae = [&](bool decode_only) -> std::shared_ptr<VAE> {
-                if (sd_version_uses_wan_vae(version) || sd_version_is_ltxav(version)) {
+                if (sd_version_uses_wan_vae(version) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version)) {
                     return std::make_shared<TinyVideoAutoEncoder>(backend_for(SDBackendModule::VAE),
                                                                   tensor_storage_map,
                                                                   "decoder",
@@ -1306,12 +1398,24 @@ public:
                                                          false,
                                                          version,
                                                          model_manager);
-                } else if (sd_version_uses_wan_vae(version)) {
+                } else if (sd_version_is_mage_flow(vae_version)) {
+                    return std::make_shared<MageVAE::MageVAERunner>(backend_for(SDBackendModule::VAE),
+                                                                    tensor_storage_map,
+                                                                    "first_stage_model",
+                                                                    model_manager);
+                } else if (sd_version_uses_hunyuan_video_vae(vae_version)) {
+                    return std::make_shared<Hunyuan::HunyuanVideoVAERunner>(backend_for(SDBackendModule::VAE),
+                                                                            tensor_storage_map,
+                                                                            "first_stage_model",
+                                                                            false,
+                                                                            vae_version,
+                                                                            model_manager);
+                } else if (sd_version_uses_wan_vae(vae_version)) {
                     return std::make_shared<WAN::WanVAERunner>(backend_for(SDBackendModule::VAE),
                                                                tensor_storage_map,
                                                                "first_stage_model",
                                                                false,
-                                                               version,
+                                                               vae_version,
                                                                model_manager);
                 } else {
                     auto model = std::make_shared<AutoEncoderKL>(backend_for(SDBackendModule::VAE),
@@ -1617,8 +1721,10 @@ public:
                     }
                 } else if (sd_version_is_sd3(version) ||
                            sd_version_is_wan(version) ||
+                           sd_version_is_hunyuan_video(version) ||
                            sd_version_is_lingbot_video(version) ||
                            sd_version_is_qwen_image(version) ||
+                           sd_version_is_mage_flow(version) ||
                            version == VERSION_HIDREAM_O1 ||
                            sd_version_is_anima(version) ||
                            sd_version_is_ernie_image(version) ||
@@ -1629,6 +1735,8 @@ public:
                     pred_type = FLOW_PRED;
                     if (sd_version_is_wan(version)) {
                         default_flow_shift = 5.f;
+                    } else if (sd_version_is_hunyuan_video(version)) {
+                        default_flow_shift = 7.f;
                     } else if (sd_version_is_ernie_image(version)) {
                         default_flow_shift = 4.f;
                     } else if (sd_version_is_pid(version)) {
@@ -1637,6 +1745,8 @@ public:
                         default_flow_shift = 1.0f;
                     } else if (sd_version_is_boogu_image(version)) {
                         default_flow_shift = 3.16f;
+                    } else if (sd_version_is_mage_flow(version)) {
+                        default_flow_shift = 6.f;
                     } else {
                         default_flow_shift = 3.f;
                     }
@@ -2013,6 +2123,34 @@ public:
         return output;
     }
 
+    void compute_ip_adapter_tokens(const sd_image_t& image, float strength) {
+        ip_adapter_tokens        = {};
+        ip_adapter_uncond_tokens = {};
+        ip_adapter_strength      = strength;
+        if (ip_adapter == nullptr || clip_vision == nullptr || image.data == nullptr) {
+            return;
+        }
+        auto image_tensor = sd_image_to_tensor(image);
+        auto embed        = get_clip_vision_output(image_tensor, true, -1);
+        if (embed.empty()) {
+            return;
+        }
+        ip_adapter_tokens = ip_adapter->compute(n_threads, embed);
+        if (ip_adapter_tokens.empty()) {
+            LOG_ERROR("IP-Adapter conditional image projection failed");
+            return;
+        }
+        auto uncond_embed        = sd::Tensor<float>::zeros_like(embed);
+        ip_adapter_uncond_tokens = ip_adapter->compute(n_threads, uncond_embed);
+        if (ip_adapter_uncond_tokens.empty()) {
+            LOG_ERROR("IP-Adapter unconditional image projection failed");
+            ip_adapter_tokens = {};
+            return;
+        }
+        LOG_INFO("IP-Adapter: %lld image tokens, strength %.2f",
+                 (long long)ip_adapter_tokens.shape()[1], strength);
+    }
+
     std::vector<float> process_timesteps(const std::vector<float>& timesteps,
                                          const sd::Tensor<float>& init_latent,
                                          const sd::Tensor<float>& denoise_mask,
@@ -2446,6 +2584,10 @@ public:
 
             sd::Tensor<float> timesteps_tensor({static_cast<int64_t>(timesteps_vec.size())}, timesteps_vec);
             sd::Tensor<float> guidance_tensor({1}, std::vector<float>{guidance.distilled_guidance});
+            sd::Tensor<float> hunyuan_timestep_r_tensor;
+            if (sd_version_is_hunyuan_video(version) && step + 1 < sigmas.size()) {
+                hunyuan_timestep_r_tensor = sd::Tensor<float>::from_vector({sigmas[step + 1]});
+            }
             sd::Tensor<float> noised_input = x * c_in;
             if (!denoise_mask.empty() && (version == VERSION_WAN2_2_TI2V || sd_version_is_ltxav(version) || sd_version_is_lingbot_video(version))) {
                 noised_input = noised_input * denoise_mask + init_latent * (1.0f - denoise_mask);
@@ -2497,7 +2639,8 @@ public:
             auto run_condition = [&](const SDCondition& condition,
                                      const sd::Tensor<float>* c_concat_override                 = nullptr,
                                      const std::vector<int>* local_skip_layers                  = nullptr,
-                                     const std::vector<sd::Tensor<float>>* ref_latents_override = nullptr) -> sd::Tensor<float> {
+                                     const std::vector<sd::Tensor<float>>* ref_latents_override = nullptr,
+                                     bool use_uncond_ip                                         = false) -> sd::Tensor<float> {
                 diffusion_params.context     = condition.c_crossattn.empty() ? nullptr : &condition.c_crossattn;
                 diffusion_params.c_concat    = c_concat_override != nullptr ? c_concat_override : (condition.c_concat.empty() ? nullptr : &condition.c_concat);
                 diffusion_params.y           = condition.c_vector.empty() ? nullptr : &condition.c_vector;
@@ -2508,7 +2651,13 @@ public:
                     if (animatediff_loaded && noised_input.dim() >= 4 && noised_input.shape()[3] > 1) {
                         nvf = static_cast<int>(noised_input.shape()[3]);
                     }
-                    diffusion_params.extra = UNetDiffusionExtra{nvf, &controls, control_strength};
+                    UNetDiffusionExtra unet_extra{nvf, &controls, control_strength};
+                    const auto& ip_tokens = use_uncond_ip ? ip_adapter_uncond_tokens : ip_adapter_tokens;
+                    if (!ip_tokens.empty()) {
+                        unet_extra.ip_context = &ip_tokens;
+                        unet_extra.ip_scale   = ip_adapter_strength;
+                    }
+                    diffusion_params.extra = unet_extra;
                 } 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) || sd_version_is_sefi_image(version)) {
@@ -2520,6 +2669,12 @@ public:
                 } else if (sd_version_is_wan(version)) {
                     diffusion_params.extra = WanDiffusionExtra{vace_context.empty() ? nullptr : &vace_context,
                                                                vace_strength};
+                } else if (sd_version_is_hunyuan_video(version)) {
+                    diffusion_params.extra = HunyuanVideoDiffusionExtra{
+                        &guidance_tensor,
+                        condition.extra_c_crossattns.empty() ? nullptr : &condition.extra_c_crossattns[0],
+                        condition.c_vector.empty() ? nullptr : &condition.c_vector,
+                        hunyuan_timestep_r_tensor.empty() ? nullptr : &hunyuan_timestep_r_tensor};
                 } else if (version == VERSION_HIDREAM_O1) {
                     diffusion_params.extra = HiDreamO1DiffusionExtra{
                         condition.c_input_ids.empty() ? nullptr : &condition.c_input_ids,
@@ -2593,7 +2748,9 @@ public:
                 }
                 uncond_out = run_condition(uncond,
                                            uncond.c_concat.empty() ? nullptr : &uncond.c_concat,
-                                           uncond_skip_layers);
+                                           uncond_skip_layers,
+                                           nullptr,
+                                           true);
                 if (uncond_out.empty()) {
                     return {};
                 }
@@ -2602,7 +2759,8 @@ public:
                 img_uncond_out = run_condition(img_uncond,
                                                img_uncond.c_concat.empty() ? nullptr : &img_uncond.c_concat,
                                                nullptr,
-                                               uncond_without_ref_latents ? &empty_ref_latents : nullptr);
+                                               uncond_without_ref_latents ? &empty_ref_latents : nullptr,
+                                               true);
                 if (img_uncond_out.empty()) {
                     return {};
                 }
@@ -2713,6 +2871,8 @@ public:
                 latent_channel = 128;
             } else if (version == VERSION_WAN2_2_TI2V) {
                 latent_channel = 48;
+            } else if (sd_version_is_hunyuan_video(version)) {
+                latent_channel = 32;
             } else if (version == VERSION_HIDREAM_O1) {
                 latent_channel = 3;
             } else if (version == VERSION_CHROMA_RADIANCE) {
@@ -2725,6 +2885,8 @@ public:
                 latent_channel = 144;
             } else if (sd_version_uses_flux2_vae(version)) {
                 latent_channel = 128;
+            } else if (sd_version_is_mage_flow(version)) {
+                latent_channel = 128;
             } else {
                 latent_channel = 16;
             }
@@ -2760,7 +2922,7 @@ public:
         int latent_frames = frames;
         if (sd_version_is_ltxav(version)) {
             latent_frames = ((frames - 1) / 8) + 1;
-        } else if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version)) {
+        } else if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_hunyuan_video(version)) {
             latent_frames = ((frames - 1) / 4) + 1;
         }
         return latent_frames;
@@ -2773,7 +2935,7 @@ public:
         if (sd_version_is_ltxav(version)) {
             return (latent_frames - 1) * 8 + 1;
         }
-        if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version)) {
+        if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_hunyuan_video(version)) {
             return (latent_frames - 1) * 4 + 1;
         }
         return latent_frames;
@@ -2873,6 +3035,8 @@ public:
             return "qwen_layered";
         } else if (sd_version_is_qwen_image(version)) {
             return "qwen";
+        } else if (sd_version_is_mage_flow(version)) {
+            return "mage_flow";
         } else if (sd_version_is_z_image(version) || sd_version_is_boogu_image(version)) {
             return "z_image_omni";
         } else if (sd_version_is_krea2(version)) {
@@ -3213,6 +3377,8 @@ const char* sd_vae_format_name(enum sd_vae_format_t format) {
             return "sd3";
         case SD_VAE_FORMAT_FLUX2:
             return "flux2";
+        case SD_VAE_FORMAT_WAN:
+            return "wan";
         default:
             return NONE_STR;
     }
@@ -3226,6 +3392,8 @@ static SDVersion sd_vae_format_to_version(enum sd_vae_format_t format, SDVersion
             return VERSION_SD3;
         case SD_VAE_FORMAT_FLUX2:
             return VERSION_FLUX2;
+        case SD_VAE_FORMAT_WAN:
+            return VERSION_WAN2;
         case SD_VAE_FORMAT_AUTO:
         default:
             return fallback;
@@ -3440,21 +3608,22 @@ char* sd_sample_params_to_str(const sd_sample_params_t* sample_params) {
 void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params) {
     *sd_img_gen_params = {};
     sd_sample_params_init(&sd_img_gen_params->sample_params);
-    sd_img_gen_params->clip_skip         = -1;
-    sd_img_gen_params->ref_images_count  = 0;
-    sd_img_gen_params->ref_image_args    = "";
-    sd_img_gen_params->width             = 512;
-    sd_img_gen_params->height            = 512;
-    sd_img_gen_params->strength          = 0.75f;
-    sd_img_gen_params->seed              = -1;
-    sd_img_gen_params->batch_count       = 1;
-    sd_img_gen_params->control_strength  = 0.9f;
-    sd_img_gen_params->qwen_image_layers = 3;
-    sd_img_gen_params->circular_x        = false;
-    sd_img_gen_params->circular_y        = false;
-    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_img_gen_params->clip_skip           = -1;
+    sd_img_gen_params->ref_images_count    = 0;
+    sd_img_gen_params->ref_image_args      = "";
+    sd_img_gen_params->width               = 512;
+    sd_img_gen_params->height              = 512;
+    sd_img_gen_params->strength            = 0.75f;
+    sd_img_gen_params->seed                = -1;
+    sd_img_gen_params->batch_count         = 1;
+    sd_img_gen_params->control_strength    = 0.9f;
+    sd_img_gen_params->ip_adapter_strength = 1.0f;
+    sd_img_gen_params->qwen_image_layers   = 3;
+    sd_img_gen_params->circular_x          = false;
+    sd_img_gen_params->circular_y          = false;
+    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);
 }
@@ -3566,7 +3735,7 @@ struct sd_ctx_t {
 };
 
 static bool sd_version_supports_video_generation(SDVersion version) {
-    return version == VERSION_SVD || sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_ltxav(version);
+    return version == VERSION_SVD || sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_lingbot_video(version) || sd_version_is_ltxav(version);
 }
 
 static bool sd_version_supports_image_generation(SDVersion version) {
@@ -4732,10 +4901,17 @@ static std::optional<ImageGenerationLatents> prepare_image_generation_latents(sd
         sd::Tensor<float> ref_latent;
         if (ref_image_params.resize_before_vae && !sd_version_is_pid(sd_ctx->sd->version)) {
             LOG_DEBUG("auto resize ref images");
-            int target_pixels  = ref_image_params.vae_input_max_pixels > 0 ? ref_image_params.vae_input_max_pixels : 1024 * 1024;
-            int vae_image_size = std::min(target_pixels, request->width * request->height);
-            double vae_width   = sqrt(vae_image_size * ref_images[i].shape()[0] / ref_images[i].shape()[1]);
-            double vae_height  = vae_width * ref_images[i].shape()[1] / ref_images[i].shape()[0];
+            double vae_width;
+            double vae_height;
+            if (ref_image_params.resize_vae_to_target) {
+                vae_width  = request->width;
+                vae_height = request->height;
+            } else {
+                int target_pixels  = ref_image_params.vae_input_max_pixels > 0 ? ref_image_params.vae_input_max_pixels : 1024 * 1024;
+                int vae_image_size = std::min(target_pixels, request->width * request->height);
+                vae_width          = sqrt(vae_image_size * ref_images[i].shape()[0] / ref_images[i].shape()[1]);
+                vae_height         = vae_width * ref_images[i].shape()[1] / ref_images[i].shape()[0];
+            }
 
             int factor = sd_version_is_qwen_image(sd_ctx->sd->version) ? 32 : 16;
             vae_height = round(vae_height / factor) * factor;
@@ -4879,6 +5055,7 @@ static std::optional<ImageGenerationEmbeds> prepare_image_generation_embeds(sd_c
                                               request->pulid_params,
                                               condition_params,
                                               plan->total_steps);
+    sd_ctx->sd->compute_ip_adapter_tokens(sd_img_gen_params->ip_adapter_image, sd_img_gen_params->ip_adapter_strength);
     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,
@@ -4930,8 +5107,8 @@ static std::optional<ImageGenerationEmbeds> prepare_image_generation_embeds(sd_c
             }
             condition_params.text            = request->negative_prompt;
             condition_params.zero_out_masked = zero_out_masked;
+            std::vector<sd::Tensor<float>> empty_ref_images;
             if (use_ref_latent_img_cfg) {
-                std::vector<sd::Tensor<float>> empty_ref_images;
                 condition_params.ref_images = &empty_ref_images;
             }
             img_uncond = sd_ctx->sd->cond_stage_model->get_learned_condition(sd_ctx->sd->n_threads,
@@ -5623,6 +5800,66 @@ static std::optional<ImageGenerationLatents> prepare_video_generation_latents(sd
         }
     }
 
+    if (sd_version_is_hunyuan_video(sd_ctx->sd->version) &&
+        (!start_image.empty() || !end_image.empty())) {
+        LOG_INFO("Hunyuan Video IMG2VID");
+
+        int64_t t1                  = ggml_time_ms();
+        auto concat_latent          = sd_ctx->sd->generate_init_latent(request->width,
+                                                                       request->height,
+                                                                       request->frames,
+                                                                       true);
+        auto encode_condition_frame = [&](const sd::Tensor<float>& image,
+                                          int64_t latent_frame,
+                                          const char* name) -> bool {
+            auto encoded = sd_ctx->sd->encode_first_stage(image.unsqueeze(2));
+            if (encoded.empty()) {
+                LOG_ERROR("failed to encode Hunyuan Video %s conditioning frame", name);
+                return false;
+            }
+            if (encoded.dim() == 4) {
+                encoded.unsqueeze_(2);
+            }
+            if (encoded.dim() != 5 ||
+                encoded.shape()[0] != concat_latent.shape()[0] ||
+                encoded.shape()[1] != concat_latent.shape()[1] ||
+                encoded.shape()[3] != concat_latent.shape()[3]) {
+                LOG_ERROR("invalid Hunyuan Video %s conditioning latent shape", name);
+                return false;
+            }
+            sd::ops::slice_assign(&concat_latent,
+                                  2,
+                                  latent_frame,
+                                  latent_frame + 1,
+                                  sd::ops::slice(encoded, 2, 0, 1));
+            return true;
+        };
+
+        if (!start_image.empty() && !encode_condition_frame(start_image, 0, "start")) {
+            return std::nullopt;
+        }
+        if (!end_image.empty() &&
+            !encode_condition_frame(end_image, concat_latent.shape()[2] - 1, "end")) {
+            return std::nullopt;
+        }
+
+        sd::Tensor<float> concat_mask = sd::zeros<float>({concat_latent.shape()[0],
+                                                          concat_latent.shape()[1],
+                                                          concat_latent.shape()[2],
+                                                          1,
+                                                          1});
+        if (!start_image.empty()) {
+            sd::ops::fill_slice(&concat_mask, 2, 0, 1, 1.0f);
+        }
+        if (!end_image.empty()) {
+            sd::ops::fill_slice(&concat_mask, 2, concat_mask.shape()[2] - 1, concat_mask.shape()[2], 1.0f);
+        }
+        latents.concat_latent = sd::ops::concat(concat_latent, concat_mask, 3);
+
+        int64_t t2 = ggml_time_ms();
+        LOG_INFO("encode_first_stage completed, taking %" PRId64 " ms", t2 - t1);
+    }
+
     if (sd_ctx->sd->diffusion_model->get_desc() == "Wan2.1-I2V-14B" ||
         sd_ctx->sd->diffusion_model->get_desc() == "Wan2.2-I2V-14B" ||
         sd_ctx->sd->diffusion_model->get_desc() == "Wan2.1-I2V-1.3B" ||

src/tokenizers/bpe_tokenizer.cpp

@@ -205,7 +205,7 @@ std::vector<int> BPETokenizer::encode(const std::string& text, on_new_token_cb_t
         ss << "\"" << token << "\", ";
     }
     ss << "]";
-    LOG_DEBUG("split prompt \"%s\" to tokens %s", text.c_str(), ss.str().c_str());
+    LOG_DEBUG("split prompt \"%s\" to %zu tokens %s", text.c_str(), bpe_tokens.size(), ss.str().c_str());
     return bpe_tokens;
 }