diff --git a/.ci/windows_amd_base_files/README_VERY_IMPORTANT.txt b/.ci/windows_amd_base_files/README_VERY_IMPORTANT.txt
new file mode 100755
index 000000000..96a500be2
--- /dev/null
+++ b/.ci/windows_amd_base_files/README_VERY_IMPORTANT.txt
@@ -0,0 +1,27 @@
+As of the time of writing this you need this preview driver for best results:
+https://www.amd.com/en/resources/support-articles/release-notes/RN-AMDGPU-WINDOWS-PYTORCH-PREVIEW.html
+
+HOW TO RUN:
+
+If you have a AMD gpu:
+
+run_amd_gpu.bat
+
+If you have memory issues you can try disabling the smart memory management by running comfyui with:
+
+run_amd_gpu_disable_smart_memory.bat
+
+IF YOU GET A RED ERROR IN THE UI MAKE SURE YOU HAVE A MODEL/CHECKPOINT IN: ComfyUI\models\checkpoints
+
+You can download the stable diffusion XL one from: https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/sd_xl_base_1.0_0.9vae.safetensors
+
+
+RECOMMENDED WAY TO UPDATE:
+To update the ComfyUI code: update\update_comfyui.bat
+
+
+TO SHARE MODELS BETWEEN COMFYUI AND ANOTHER UI:
+In the ComfyUI directory you will find a file: extra_model_paths.yaml.example
+Rename this file to: extra_model_paths.yaml and edit it with your favorite text editor.
+
+
diff --git a/.ci/windows_base_files/run_nvidia_gpu.bat b/.ci/windows_amd_base_files/run_amd_gpu.bat
similarity index 100%
rename from .ci/windows_base_files/run_nvidia_gpu.bat
rename to .ci/windows_amd_base_files/run_amd_gpu.bat
diff --git a/.ci/windows_base_files/run_nvidia_gpu_fast_fp16_accumulation.bat b/.ci/windows_amd_base_files/run_amd_gpu_disable_smart_memory.bat
old mode 100644
new mode 100755
similarity index 65%
rename from .ci/windows_base_files/run_nvidia_gpu_fast_fp16_accumulation.bat
rename to .ci/windows_amd_base_files/run_amd_gpu_disable_smart_memory.bat
index 38f06ecb2..cece0aeb2
--- a/.ci/windows_base_files/run_nvidia_gpu_fast_fp16_accumulation.bat
+++ b/.ci/windows_amd_base_files/run_amd_gpu_disable_smart_memory.bat
@@ -1,2 +1,2 @@
-.\python_embeded\python.exe -s ComfyUI\main.py --windows-standalone-build --fast fp16_accumulation
+.\python_embeded\python.exe -s ComfyUI\main.py --windows-standalone-build --disable-smart-memory
pause
diff --git a/.ci/windows_base_files/README_VERY_IMPORTANT.txt b/.ci/windows_nvidia_base_files/README_VERY_IMPORTANT.txt
similarity index 82%
rename from .ci/windows_base_files/README_VERY_IMPORTANT.txt
rename to .ci/windows_nvidia_base_files/README_VERY_IMPORTANT.txt
index d46acbcbf..8ab70c890 100755
--- a/.ci/windows_base_files/README_VERY_IMPORTANT.txt
+++ b/.ci/windows_nvidia_base_files/README_VERY_IMPORTANT.txt
@@ -4,6 +4,9 @@ if you have a NVIDIA gpu:
run_nvidia_gpu.bat
+if you want to enable the fast fp16 accumulation (faster for fp16 models with slightly less quality):
+
+run_nvidia_gpu_fast_fp16_accumulation.bat
To run it in slow CPU mode:
diff --git a/.ci/windows_nvidia_base_files/advanced/run_nvidia_gpu_disable_api_nodes.bat b/.ci/windows_nvidia_base_files/advanced/run_nvidia_gpu_disable_api_nodes.bat
new file mode 100644
index 000000000..ed00583b6
--- /dev/null
+++ b/.ci/windows_nvidia_base_files/advanced/run_nvidia_gpu_disable_api_nodes.bat
@@ -0,0 +1,3 @@
+..\python_embeded\python.exe -s ..\ComfyUI\main.py --windows-standalone-build --disable-api-nodes
+echo If you see this and ComfyUI did not start try updating your Nvidia Drivers to the latest.
+pause
diff --git a/.ci/windows_base_files/run_cpu.bat b/.ci/windows_nvidia_base_files/run_cpu.bat
similarity index 100%
rename from .ci/windows_base_files/run_cpu.bat
rename to .ci/windows_nvidia_base_files/run_cpu.bat
diff --git a/.ci/windows_nvidia_base_files/run_nvidia_gpu.bat b/.ci/windows_nvidia_base_files/run_nvidia_gpu.bat
new file mode 100755
index 000000000..4898a424f
--- /dev/null
+++ b/.ci/windows_nvidia_base_files/run_nvidia_gpu.bat
@@ -0,0 +1,3 @@
+.\python_embeded\python.exe -s ComfyUI\main.py --windows-standalone-build
+echo If you see this and ComfyUI did not start try updating your Nvidia Drivers to the latest.
+pause
diff --git a/.ci/windows_nvidia_base_files/run_nvidia_gpu_fast_fp16_accumulation.bat b/.ci/windows_nvidia_base_files/run_nvidia_gpu_fast_fp16_accumulation.bat
new file mode 100644
index 000000000..32611e4af
--- /dev/null
+++ b/.ci/windows_nvidia_base_files/run_nvidia_gpu_fast_fp16_accumulation.bat
@@ -0,0 +1,3 @@
+.\python_embeded\python.exe -s ComfyUI\main.py --windows-standalone-build --fast fp16_accumulation
+echo If you see this and ComfyUI did not start try updating your Nvidia Drivers to the latest.
+pause
diff --git a/.gitattributes b/.gitattributes
index 4391de678..5b3c15bb4 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -1,2 +1,3 @@
/web/assets/** linguist-generated
/web/** linguist-vendored
+comfy_api_nodes/apis/__init__.py linguist-generated
diff --git a/.github/ISSUE_TEMPLATE/bug-report.yml b/.github/ISSUE_TEMPLATE/bug-report.yml
index 69ce998eb..3cf2717b7 100644
--- a/.github/ISSUE_TEMPLATE/bug-report.yml
+++ b/.github/ISSUE_TEMPLATE/bug-report.yml
@@ -22,7 +22,7 @@ body:
description: Please confirm you have tried to reproduce the issue with all custom nodes disabled.
options:
- label: I have tried disabling custom nodes and the issue persists (see [how to disable custom nodes](https://docs.comfy.org/troubleshooting/custom-node-issues#step-1%3A-test-with-all-custom-nodes-disabled) if you need help)
- required: true
+ required: false
- type: textarea
attributes:
label: Expected Behavior
diff --git a/.github/ISSUE_TEMPLATE/user-support.yml b/.github/ISSUE_TEMPLATE/user-support.yml
index 50657d493..281661f92 100644
--- a/.github/ISSUE_TEMPLATE/user-support.yml
+++ b/.github/ISSUE_TEMPLATE/user-support.yml
@@ -18,7 +18,7 @@ body:
description: Please confirm you have tried to reproduce the issue with all custom nodes disabled.
options:
- label: I have tried disabling custom nodes and the issue persists (see [how to disable custom nodes](https://docs.comfy.org/troubleshooting/custom-node-issues#step-1%3A-test-with-all-custom-nodes-disabled) if you need help)
- required: true
+ required: false
- type: textarea
attributes:
label: Your question
diff --git a/.github/workflows/check-line-endings.yml b/.github/workflows/check-line-endings.yml
new file mode 100644
index 000000000..eeb594d6c
--- /dev/null
+++ b/.github/workflows/check-line-endings.yml
@@ -0,0 +1,40 @@
+name: Check for Windows Line Endings
+
+on:
+ pull_request:
+ branches: ['*'] # Trigger on all pull requests to any branch
+
+jobs:
+ check-line-endings:
+ runs-on: ubuntu-latest
+
+ steps:
+ - name: Checkout code
+ uses: actions/checkout@v4
+ with:
+ fetch-depth: 0 # Fetch all history to compare changes
+
+ - name: Check for Windows line endings (CRLF)
+ run: |
+ # Get the list of changed files in the PR
+ CHANGED_FILES=$(git diff --name-only ${{ github.event.pull_request.base.sha }}..${{ github.event.pull_request.head.sha }})
+
+ # Flag to track if CRLF is found
+ CRLF_FOUND=false
+
+ # Loop through each changed file
+ for FILE in $CHANGED_FILES; do
+ # Check if the file exists and is a text file
+ if [ -f "$FILE" ] && file "$FILE" | grep -q "text"; then
+ # Check for CRLF line endings
+ if grep -UP '\r$' "$FILE"; then
+ echo "Error: Windows line endings (CRLF) detected in $FILE"
+ CRLF_FOUND=true
+ fi
+ fi
+ done
+
+ # Exit with error if CRLF was found
+ if [ "$CRLF_FOUND" = true ]; then
+ exit 1
+ fi
diff --git a/.github/workflows/release-stable-all.yml b/.github/workflows/release-stable-all.yml
new file mode 100644
index 000000000..7dca7277b
--- /dev/null
+++ b/.github/workflows/release-stable-all.yml
@@ -0,0 +1,61 @@
+name: "Release Stable All Portable Versions"
+
+on:
+ workflow_dispatch:
+ inputs:
+ git_tag:
+ description: 'Git tag'
+ required: true
+ type: string
+
+jobs:
+ release_nvidia_default:
+ permissions:
+ contents: "write"
+ packages: "write"
+ pull-requests: "read"
+ name: "Release NVIDIA Default (cu129)"
+ uses: ./.github/workflows/stable-release.yml
+ with:
+ git_tag: ${{ inputs.git_tag }}
+ cache_tag: "cu130"
+ python_minor: "13"
+ python_patch: "9"
+ rel_name: "nvidia"
+ rel_extra_name: ""
+ test_release: true
+ secrets: inherit
+
+ release_nvidia_cu128:
+ permissions:
+ contents: "write"
+ packages: "write"
+ pull-requests: "read"
+ name: "Release NVIDIA cu128"
+ uses: ./.github/workflows/stable-release.yml
+ with:
+ git_tag: ${{ inputs.git_tag }}
+ cache_tag: "cu128"
+ python_minor: "12"
+ python_patch: "10"
+ rel_name: "nvidia"
+ rel_extra_name: "_cu128"
+ test_release: true
+ secrets: inherit
+
+ release_amd_rocm:
+ permissions:
+ contents: "write"
+ packages: "write"
+ pull-requests: "read"
+ name: "Release AMD ROCm 6.4.4"
+ uses: ./.github/workflows/stable-release.yml
+ with:
+ git_tag: ${{ inputs.git_tag }}
+ cache_tag: "rocm644"
+ python_minor: "12"
+ python_patch: "10"
+ rel_name: "amd"
+ rel_extra_name: ""
+ test_release: false
+ secrets: inherit
diff --git a/.github/workflows/ruff.yml b/.github/workflows/ruff.yml
index 4c1a02594..b24d86a6b 100644
--- a/.github/workflows/ruff.yml
+++ b/.github/workflows/ruff.yml
@@ -21,3 +21,28 @@ jobs:
- name: Run Ruff
run: ruff check .
+
+ pylint:
+ name: Run Pylint
+ runs-on: ubuntu-latest
+
+ steps:
+ - name: Checkout repository
+ uses: actions/checkout@v4
+
+ - name: Set up Python
+ uses: actions/setup-python@v4
+ with:
+ python-version: '3.12'
+
+ - name: Install requirements
+ run: |
+ python -m pip install --upgrade pip
+ pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
+ pip install -r requirements.txt
+
+ - name: Install Pylint
+ run: pip install pylint
+
+ - name: Run Pylint
+ run: pylint comfy_api_nodes
diff --git a/.github/workflows/stable-release.yml b/.github/workflows/stable-release.yml
index 61105abe4..28484a9d1 100644
--- a/.github/workflows/stable-release.yml
+++ b/.github/workflows/stable-release.yml
@@ -2,28 +2,78 @@
name: "Release Stable Version"
on:
+ workflow_call:
+ inputs:
+ git_tag:
+ description: 'Git tag'
+ required: true
+ type: string
+ cache_tag:
+ description: 'Cached dependencies tag'
+ required: true
+ type: string
+ default: "cu129"
+ python_minor:
+ description: 'Python minor version'
+ required: true
+ type: string
+ default: "13"
+ python_patch:
+ description: 'Python patch version'
+ required: true
+ type: string
+ default: "6"
+ rel_name:
+ description: 'Release name'
+ required: true
+ type: string
+ default: "nvidia"
+ rel_extra_name:
+ description: 'Release extra name'
+ required: false
+ type: string
+ default: ""
+ test_release:
+ description: 'Test Release'
+ required: true
+ type: boolean
+ default: true
workflow_dispatch:
inputs:
git_tag:
description: 'Git tag'
required: true
type: string
- cu:
- description: 'CUDA version'
+ cache_tag:
+ description: 'Cached dependencies tag'
required: true
type: string
- default: "128"
+ default: "cu129"
python_minor:
description: 'Python minor version'
required: true
type: string
- default: "12"
+ default: "13"
python_patch:
description: 'Python patch version'
required: true
type: string
- default: "10"
-
+ default: "6"
+ rel_name:
+ description: 'Release name'
+ required: true
+ type: string
+ default: "nvidia"
+ rel_extra_name:
+ description: 'Release extra name'
+ required: false
+ type: string
+ default: ""
+ test_release:
+ description: 'Test Release'
+ required: true
+ type: boolean
+ default: true
jobs:
package_comfy_windows:
@@ -42,15 +92,15 @@ jobs:
id: cache
with:
path: |
- cu${{ inputs.cu }}_python_deps.tar
+ ${{ inputs.cache_tag }}_python_deps.tar
update_comfyui_and_python_dependencies.bat
- key: ${{ runner.os }}-build-cu${{ inputs.cu }}-${{ inputs.python_minor }}
+ key: ${{ runner.os }}-build-${{ inputs.cache_tag }}-${{ inputs.python_minor }}
- shell: bash
run: |
- mv cu${{ inputs.cu }}_python_deps.tar ../
+ mv ${{ inputs.cache_tag }}_python_deps.tar ../
mv update_comfyui_and_python_dependencies.bat ../
cd ..
- tar xf cu${{ inputs.cu }}_python_deps.tar
+ tar xf ${{ inputs.cache_tag }}_python_deps.tar
pwd
ls
@@ -65,9 +115,21 @@ jobs:
echo 'import site' >> ./python3${{ inputs.python_minor }}._pth
curl https://bootstrap.pypa.io/get-pip.py -o get-pip.py
./python.exe get-pip.py
- ./python.exe -s -m pip install ../cu${{ inputs.cu }}_python_deps/*
- sed -i '1i../ComfyUI' ./python3${{ inputs.python_minor }}._pth
- cd ..
+ ./python.exe -s -m pip install ../${{ inputs.cache_tag }}_python_deps/*
+
+ grep comfyui ../ComfyUI/requirements.txt > ./requirements_comfyui.txt
+ ./python.exe -s -m pip install -r requirements_comfyui.txt
+ rm requirements_comfyui.txt
+
+ sed -i '1i../ComfyUI' ./python3${{ inputs.python_minor }}._pth
+
+ if test -f ./Lib/site-packages/torch/lib/dnnl.lib; then
+ rm ./Lib/site-packages/torch/lib/dnnl.lib #I don't think this is actually used and I need the space
+ rm ./Lib/site-packages/torch/lib/libprotoc.lib
+ rm ./Lib/site-packages/torch/lib/libprotobuf.lib
+ fi
+
+ cd ..
git clone --depth 1 https://github.com/comfyanonymous/taesd
cp taesd/*.safetensors ./ComfyUI_copy/models/vae_approx/
@@ -80,14 +142,18 @@ jobs:
mkdir update
cp -r ComfyUI/.ci/update_windows/* ./update/
- cp -r ComfyUI/.ci/windows_base_files/* ./
+ cp -r ComfyUI/.ci/windows_${{ inputs.rel_name }}_base_files/* ./
cp ../update_comfyui_and_python_dependencies.bat ./update/
cd ..
- "C:\Program Files\7-Zip\7z.exe" a -t7z -m0=lzma2 -mx=9 -mfb=128 -md=512m -ms=on -mf=BCJ2 ComfyUI_windows_portable.7z ComfyUI_windows_portable
- mv ComfyUI_windows_portable.7z ComfyUI/ComfyUI_windows_portable_nvidia.7z
+ "C:\Program Files\7-Zip\7z.exe" a -t7z -m0=lzma2 -mx=9 -mfb=128 -md=768m -ms=on -mf=BCJ2 ComfyUI_windows_portable.7z ComfyUI_windows_portable
+ mv ComfyUI_windows_portable.7z ComfyUI/ComfyUI_windows_portable_${{ inputs.rel_name }}${{ inputs.rel_extra_name }}.7z
+ - shell: bash
+ if: ${{ inputs.test_release }}
+ run: |
+ cd ..
cd ComfyUI_windows_portable
python_embeded/python.exe -s ComfyUI/main.py --quick-test-for-ci --cpu
@@ -96,10 +162,9 @@ jobs:
ls
- name: Upload binaries to release
- uses: svenstaro/upload-release-action@v2
+ uses: softprops/action-gh-release@v2
with:
- repo_token: ${{ secrets.GITHUB_TOKEN }}
- file: ComfyUI_windows_portable_nvidia.7z
- tag: ${{ inputs.git_tag }}
- overwrite: true
+ files: ComfyUI_windows_portable_${{ inputs.rel_name }}${{ inputs.rel_extra_name }}.7z
+ tag_name: ${{ inputs.git_tag }}
draft: true
+ overwrite_files: true
diff --git a/.github/workflows/test-execution.yml b/.github/workflows/test-execution.yml
new file mode 100644
index 000000000..00ef07ebf
--- /dev/null
+++ b/.github/workflows/test-execution.yml
@@ -0,0 +1,30 @@
+name: Execution Tests
+
+on:
+ push:
+ branches: [ main, master ]
+ pull_request:
+ branches: [ main, master ]
+
+jobs:
+ test:
+ strategy:
+ matrix:
+ os: [ubuntu-latest, windows-latest, macos-latest]
+ runs-on: ${{ matrix.os }}
+ continue-on-error: true
+ steps:
+ - uses: actions/checkout@v4
+ - name: Set up Python
+ uses: actions/setup-python@v4
+ with:
+ python-version: '3.12'
+ - name: Install requirements
+ run: |
+ python -m pip install --upgrade pip
+ pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
+ pip install -r requirements.txt
+ pip install -r tests-unit/requirements.txt
+ - name: Run Execution Tests
+ run: |
+ python -m pytest tests/execution -v --skip-timing-checks
diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml
index 78c918031..00caf5b8a 100644
--- a/.github/workflows/test-unit.yml
+++ b/.github/workflows/test-unit.yml
@@ -10,7 +10,7 @@ jobs:
test:
strategy:
matrix:
- os: [ubuntu-latest, windows-latest, macos-latest]
+ os: [ubuntu-latest, windows-2022, macos-latest]
runs-on: ${{ matrix.os }}
continue-on-error: true
steps:
diff --git a/.github/workflows/windows_release_dependencies.yml b/.github/workflows/windows_release_dependencies.yml
index dfdb96d50..f61ee21a2 100644
--- a/.github/workflows/windows_release_dependencies.yml
+++ b/.github/workflows/windows_release_dependencies.yml
@@ -17,19 +17,19 @@ on:
description: 'cuda version'
required: true
type: string
- default: "128"
+ default: "130"
python_minor:
description: 'python minor version'
required: true
type: string
- default: "12"
+ default: "13"
python_patch:
description: 'python patch version'
required: true
type: string
- default: "10"
+ default: "9"
# push:
# branches:
# - master
@@ -56,7 +56,8 @@ jobs:
..\python_embeded\python.exe -s -m pip install --upgrade torch torchvision torchaudio ${{ inputs.xformers }} --extra-index-url https://download.pytorch.org/whl/cu${{ inputs.cu }} -r ../ComfyUI/requirements.txt pygit2
pause" > update_comfyui_and_python_dependencies.bat
- python -m pip wheel --no-cache-dir torch torchvision torchaudio ${{ inputs.xformers }} ${{ inputs.extra_dependencies }} --extra-index-url https://download.pytorch.org/whl/cu${{ inputs.cu }} -r requirements.txt pygit2 -w ./temp_wheel_dir
+ grep -v comfyui requirements.txt > requirements_nocomfyui.txt
+ python -m pip wheel --no-cache-dir torch torchvision torchaudio ${{ inputs.xformers }} ${{ inputs.extra_dependencies }} --extra-index-url https://download.pytorch.org/whl/cu${{ inputs.cu }} -r requirements_nocomfyui.txt pygit2 -w ./temp_wheel_dir
python -m pip install --no-cache-dir ./temp_wheel_dir/*
echo installed basic
ls -lah temp_wheel_dir
diff --git a/.github/workflows/windows_release_dependencies_manual.yml b/.github/workflows/windows_release_dependencies_manual.yml
new file mode 100644
index 000000000..0799feef1
--- /dev/null
+++ b/.github/workflows/windows_release_dependencies_manual.yml
@@ -0,0 +1,64 @@
+name: "Windows Release dependencies Manual"
+
+on:
+ workflow_dispatch:
+ inputs:
+ torch_dependencies:
+ description: 'torch dependencies'
+ required: false
+ type: string
+ default: "torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu128"
+ cache_tag:
+ description: 'Cached dependencies tag'
+ required: true
+ type: string
+ default: "cu128"
+
+ python_minor:
+ description: 'python minor version'
+ required: true
+ type: string
+ default: "12"
+
+ python_patch:
+ description: 'python patch version'
+ required: true
+ type: string
+ default: "10"
+
+jobs:
+ build_dependencies:
+ runs-on: windows-latest
+ steps:
+ - uses: actions/checkout@v4
+ - uses: actions/setup-python@v5
+ with:
+ python-version: 3.${{ inputs.python_minor }}.${{ inputs.python_patch }}
+
+ - shell: bash
+ run: |
+ echo "@echo off
+ call update_comfyui.bat nopause
+ echo -
+ echo This will try to update pytorch and all python dependencies.
+ echo -
+ echo If you just want to update normally, close this and run update_comfyui.bat instead.
+ echo -
+ pause
+ ..\python_embeded\python.exe -s -m pip install --upgrade ${{ inputs.torch_dependencies }} -r ../ComfyUI/requirements.txt pygit2
+ pause" > update_comfyui_and_python_dependencies.bat
+
+ grep -v comfyui requirements.txt > requirements_nocomfyui.txt
+ python -m pip wheel --no-cache-dir ${{ inputs.torch_dependencies }} -r requirements_nocomfyui.txt pygit2 -w ./temp_wheel_dir
+ python -m pip install --no-cache-dir ./temp_wheel_dir/*
+ echo installed basic
+ ls -lah temp_wheel_dir
+ mv temp_wheel_dir ${{ inputs.cache_tag }}_python_deps
+ tar cf ${{ inputs.cache_tag }}_python_deps.tar ${{ inputs.cache_tag }}_python_deps
+
+ - uses: actions/cache/save@v4
+ with:
+ path: |
+ ${{ inputs.cache_tag }}_python_deps.tar
+ update_comfyui_and_python_dependencies.bat
+ key: ${{ runner.os }}-build-${{ inputs.cache_tag }}-${{ inputs.python_minor }}
diff --git a/.github/workflows/windows_release_nightly_pytorch.yml b/.github/workflows/windows_release_nightly_pytorch.yml
index 5bdc940de..ca1ef71ae 100644
--- a/.github/workflows/windows_release_nightly_pytorch.yml
+++ b/.github/workflows/windows_release_nightly_pytorch.yml
@@ -68,7 +68,7 @@ jobs:
mkdir update
cp -r ComfyUI/.ci/update_windows/* ./update/
- cp -r ComfyUI/.ci/windows_base_files/* ./
+ cp -r ComfyUI/.ci/windows_nvidia_base_files/* ./
cp -r ComfyUI/.ci/windows_nightly_base_files/* ./
echo "call update_comfyui.bat nopause
diff --git a/.github/workflows/windows_release_package.yml b/.github/workflows/windows_release_package.yml
index 3926a65f3..7955325fc 100644
--- a/.github/workflows/windows_release_package.yml
+++ b/.github/workflows/windows_release_package.yml
@@ -7,19 +7,19 @@ on:
description: 'cuda version'
required: true
type: string
- default: "128"
+ default: "129"
python_minor:
description: 'python minor version'
required: true
type: string
- default: "12"
+ default: "13"
python_patch:
description: 'python patch version'
required: true
type: string
- default: "10"
+ default: "6"
# push:
# branches:
# - master
@@ -64,6 +64,10 @@ jobs:
./python.exe get-pip.py
./python.exe -s -m pip install ../cu${{ inputs.cu }}_python_deps/*
sed -i '1i../ComfyUI' ./python3${{ inputs.python_minor }}._pth
+
+ rm ./Lib/site-packages/torch/lib/dnnl.lib #I don't think this is actually used and I need the space
+ rm ./Lib/site-packages/torch/lib/libprotoc.lib
+ rm ./Lib/site-packages/torch/lib/libprotobuf.lib
cd ..
git clone --depth 1 https://github.com/comfyanonymous/taesd
@@ -77,12 +81,12 @@ jobs:
mkdir update
cp -r ComfyUI/.ci/update_windows/* ./update/
- cp -r ComfyUI/.ci/windows_base_files/* ./
+ cp -r ComfyUI/.ci/windows_nvidia_base_files/* ./
cp ../update_comfyui_and_python_dependencies.bat ./update/
cd ..
- "C:\Program Files\7-Zip\7z.exe" a -t7z -m0=lzma2 -mx=9 -mfb=128 -md=512m -ms=on -mf=BCJ2 ComfyUI_windows_portable.7z ComfyUI_windows_portable
+ "C:\Program Files\7-Zip\7z.exe" a -t7z -m0=lzma2 -mx=9 -mfb=128 -md=768m -ms=on -mf=BCJ2 ComfyUI_windows_portable.7z ComfyUI_windows_portable
mv ComfyUI_windows_portable.7z ComfyUI/new_ComfyUI_windows_portable_nvidia_cu${{ inputs.cu }}_or_cpu.7z
cd ComfyUI_windows_portable
diff --git a/CODEOWNERS b/CODEOWNERS
index c4acbf06e..b7aca9b26 100644
--- a/CODEOWNERS
+++ b/CODEOWNERS
@@ -1,24 +1,3 @@
# Admins
* @comfyanonymous
-
-# Note: Github teams syntax cannot be used here as the repo is not owned by Comfy-Org.
-# Inlined the team members for now.
-
-# Maintainers
-*.md @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @Kosinkadink @christian-byrne
-/tests/ @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @Kosinkadink @christian-byrne
-/tests-unit/ @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @Kosinkadink @christian-byrne
-/notebooks/ @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @Kosinkadink @christian-byrne
-/script_examples/ @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @Kosinkadink @christian-byrne
-/.github/ @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @Kosinkadink @christian-byrne
-/requirements.txt @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @Kosinkadink @christian-byrne
-/pyproject.toml @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @Kosinkadink @christian-byrne
-
-# Python web server
-/api_server/ @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @christian-byrne
-/app/ @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @christian-byrne
-/utils/ @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @christian-byrne
-
-# Node developers
-/comfy_extras/ @yoland68 @robinjhuang @pythongosssss @ltdrdata @Kosinkadink @webfiltered @christian-byrne
-/comfy/comfy_types/ @yoland68 @robinjhuang @pythongosssss @ltdrdata @Kosinkadink @webfiltered @christian-byrne
+* @kosinkadink
diff --git a/README.md b/README.md
index ba8892b17..4204777e9 100644
--- a/README.md
+++ b/README.md
@@ -39,7 +39,7 @@ ComfyUI lets you design and execute advanced stable diffusion pipelines using a
## Get Started
#### [Desktop Application](https://www.comfy.org/download)
-- The easiest way to get started.
+- The easiest way to get started.
- Available on Windows & macOS.
#### [Windows Portable Package](#installing)
@@ -55,7 +55,7 @@ See what ComfyUI can do with the [example workflows](https://comfyanonymous.gith
## Features
- Nodes/graph/flowchart interface to experiment and create complex Stable Diffusion workflows without needing to code anything.
- Image Models
- - SD1.x, SD2.x,
+ - SD1.x, SD2.x ([unCLIP](https://comfyanonymous.github.io/ComfyUI_examples/unclip/))
- [SDXL](https://comfyanonymous.github.io/ComfyUI_examples/sdxl/), [SDXL Turbo](https://comfyanonymous.github.io/ComfyUI_examples/sdturbo/)
- [Stable Cascade](https://comfyanonymous.github.io/ComfyUI_examples/stable_cascade/)
- [SD3 and SD3.5](https://comfyanonymous.github.io/ComfyUI_examples/sd3/)
@@ -65,17 +65,20 @@ See what ComfyUI can do with the [example workflows](https://comfyanonymous.gith
- [Flux](https://comfyanonymous.github.io/ComfyUI_examples/flux/)
- [Lumina Image 2.0](https://comfyanonymous.github.io/ComfyUI_examples/lumina2/)
- [HiDream](https://comfyanonymous.github.io/ComfyUI_examples/hidream/)
- - [Cosmos Predict2](https://comfyanonymous.github.io/ComfyUI_examples/cosmos_predict2/)
+ - [Qwen Image](https://comfyanonymous.github.io/ComfyUI_examples/qwen_image/)
+ - [Hunyuan Image 2.1](https://comfyanonymous.github.io/ComfyUI_examples/hunyuan_image/)
- Image Editing Models
- [Omnigen 2](https://comfyanonymous.github.io/ComfyUI_examples/omnigen/)
- [Flux Kontext](https://comfyanonymous.github.io/ComfyUI_examples/flux/#flux-kontext-image-editing-model)
+ - [HiDream E1.1](https://comfyanonymous.github.io/ComfyUI_examples/hidream/#hidream-e11)
+ - [Qwen Image Edit](https://comfyanonymous.github.io/ComfyUI_examples/qwen_image/#edit-model)
- Video Models
- [Stable Video Diffusion](https://comfyanonymous.github.io/ComfyUI_examples/video/)
- [Mochi](https://comfyanonymous.github.io/ComfyUI_examples/mochi/)
- [LTX-Video](https://comfyanonymous.github.io/ComfyUI_examples/ltxv/)
- [Hunyuan Video](https://comfyanonymous.github.io/ComfyUI_examples/hunyuan_video/)
- - [Nvidia Cosmos](https://comfyanonymous.github.io/ComfyUI_examples/cosmos/) and [Cosmos Predict2](https://comfyanonymous.github.io/ComfyUI_examples/cosmos_predict2/)
- [Wan 2.1](https://comfyanonymous.github.io/ComfyUI_examples/wan/)
+ - [Wan 2.2](https://comfyanonymous.github.io/ComfyUI_examples/wan22/)
- Audio Models
- [Stable Audio](https://comfyanonymous.github.io/ComfyUI_examples/audio/)
- [ACE Step](https://comfyanonymous.github.io/ComfyUI_examples/audio/)
@@ -83,9 +86,9 @@ See what ComfyUI can do with the [example workflows](https://comfyanonymous.gith
- [Hunyuan3D 2.0](https://docs.comfy.org/tutorials/3d/hunyuan3D-2)
- Asynchronous Queue system
- Many optimizations: Only re-executes the parts of the workflow that changes between executions.
-- Smart memory management: can automatically run models on GPUs with as low as 1GB vram.
+- Smart memory management: can automatically run large models on GPUs with as low as 1GB vram with smart offloading.
- Works even if you don't have a GPU with: ```--cpu``` (slow)
-- Can load ckpt, safetensors and diffusers models/checkpoints. Standalone VAEs and CLIP models.
+- Can load ckpt and safetensors: All in one checkpoints or standalone diffusion models, VAEs and CLIP models.
- Safe loading of ckpt, pt, pth, etc.. files.
- Embeddings/Textual inversion
- [Loras (regular, locon and loha)](https://comfyanonymous.github.io/ComfyUI_examples/lora/)
@@ -97,7 +100,6 @@ See what ComfyUI can do with the [example workflows](https://comfyanonymous.gith
- [Inpainting](https://comfyanonymous.github.io/ComfyUI_examples/inpaint/) with both regular and inpainting models.
- [ControlNet and T2I-Adapter](https://comfyanonymous.github.io/ComfyUI_examples/controlnet/)
- [Upscale Models (ESRGAN, ESRGAN variants, SwinIR, Swin2SR, etc...)](https://comfyanonymous.github.io/ComfyUI_examples/upscale_models/)
-- [unCLIP Models](https://comfyanonymous.github.io/ComfyUI_examples/unclip/)
- [GLIGEN](https://comfyanonymous.github.io/ComfyUI_examples/gligen/)
- [Model Merging](https://comfyanonymous.github.io/ComfyUI_examples/model_merging/)
- [LCM models and Loras](https://comfyanonymous.github.io/ComfyUI_examples/lcm/)
@@ -110,7 +112,7 @@ Workflow examples can be found on the [Examples page](https://comfyanonymous.git
## Release Process
-ComfyUI follows a weekly release cycle every Friday, with three interconnected repositories:
+ComfyUI follows a weekly release cycle targeting Friday but this regularly changes because of model releases or large changes to the codebase. There are three interconnected repositories:
1. **[ComfyUI Core](https://github.com/comfyanonymous/ComfyUI)**
- Releases a new stable version (e.g., v0.7.0)
@@ -174,14 +176,18 @@ Simply download, extract with [7-Zip](https://7-zip.org) and run. Make sure you
If you have trouble extracting it, right click the file -> properties -> unblock
+Update your Nvidia drivers if it doesn't start.
+
+#### Alternative Downloads:
+
+[Experimental portable for AMD GPUs](https://github.com/comfyanonymous/ComfyUI/releases/latest/download/ComfyUI_windows_portable_amd.7z)
+
+[Portable with pytorch cuda 12.8 and python 3.12](https://github.com/comfyanonymous/ComfyUI/releases/latest/download/ComfyUI_windows_portable_nvidia_cu128.7z) (Supports Nvidia 10 series and older GPUs).
+
#### How do I share models between another UI and ComfyUI?
See the [Config file](extra_model_paths.yaml.example) to set the search paths for models. In the standalone windows build you can find this file in the ComfyUI directory. Rename this file to extra_model_paths.yaml and edit it with your favorite text editor.
-## Jupyter Notebook
-
-To run it on services like paperspace, kaggle or colab you can use my [Jupyter Notebook](notebooks/comfyui_colab.ipynb)
-
## [comfy-cli](https://docs.comfy.org/comfy-cli/getting-started)
@@ -193,7 +199,11 @@ comfy install
## Manual Install (Windows, Linux)
-python 3.13 is supported but using 3.12 is recommended because some custom nodes and their dependencies might not support it yet.
+Python 3.14 will work if you comment out the `kornia` dependency in the requirements.txt file (breaks the canny node) but it is not recommended.
+
+Python 3.13 is very well supported. If you have trouble with some custom node dependencies on 3.13 you can try 3.12
+
+### Instructions:
Git clone this repo.
@@ -202,48 +212,58 @@ Put your SD checkpoints (the huge ckpt/safetensors files) in: models/checkpoints
Put your VAE in: models/vae
-### AMD GPUs (Linux only)
+### AMD GPUs (Linux)
+
AMD users can install rocm and pytorch with pip if you don't have it already installed, this is the command to install the stable version:
-```pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.3```
+```pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.4```
-This is the command to install the nightly with ROCm 6.4 which might have some performance improvements:
+This is the command to install the nightly with ROCm 7.0 which might have some performance improvements:
-```pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm6.4```
+```pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm7.0```
+
+
+### AMD GPUs (Experimental: Windows and Linux), RDNA 3, 3.5 and 4 only.
+
+These have less hardware support than the builds above but they work on windows. You also need to install the pytorch version specific to your hardware.
+
+RDNA 3 (RX 7000 series):
+
+```pip install --pre torch torchvision torchaudio --index-url https://rocm.nightlies.amd.com/v2/gfx110X-dgpu/```
+
+RDNA 3.5 (Strix halo/Ryzen AI Max+ 365):
+
+```pip install --pre torch torchvision torchaudio --index-url https://rocm.nightlies.amd.com/v2/gfx1151/```
+
+RDNA 4 (RX 9000 series):
+
+```pip install --pre torch torchvision torchaudio --index-url https://rocm.nightlies.amd.com/v2/gfx120X-all/```
### Intel GPUs (Windows and Linux)
-(Option 1) Intel Arc GPU users can install native PyTorch with torch.xpu support using pip (currently available in PyTorch nightly builds). More information can be found [here](https://pytorch.org/docs/main/notes/get_start_xpu.html)
-
-1. To install PyTorch nightly, use the following command:
+(Option 1) Intel Arc GPU users can install native PyTorch with torch.xpu support using pip. More information can be found [here](https://pytorch.org/docs/main/notes/get_start_xpu.html)
+
+1. To install PyTorch xpu, use the following command:
+
+```pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/xpu```
+
+This is the command to install the Pytorch xpu nightly which might have some performance improvements:
```pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/xpu```
-2. Launch ComfyUI by running `python main.py`
-
-
(Option 2) Alternatively, Intel GPUs supported by Intel Extension for PyTorch (IPEX) can leverage IPEX for improved performance.
-1. For Intel® Arc™ A-Series Graphics utilizing IPEX, create a conda environment and use the commands below:
-
-```
-conda install libuv
-pip install torch==2.3.1.post0+cxx11.abi torchvision==0.18.1.post0+cxx11.abi torchaudio==2.3.1.post0+cxx11.abi intel-extension-for-pytorch==2.3.110.post0+xpu --extra-index-url https://pytorch-extension.intel.com/release-whl/stable/xpu/us/ --extra-index-url https://pytorch-extension.intel.com/release-whl/stable/xpu/cn/
-```
-
-For other supported Intel GPUs with IPEX, visit [Installation](https://intel.github.io/intel-extension-for-pytorch/index.html#installation?platform=gpu) for more information.
-
-Additional discussion and help can be found [here](https://github.com/comfyanonymous/ComfyUI/discussions/476).
+1. visit [Installation](https://intel.github.io/intel-extension-for-pytorch/index.html#installation?platform=gpu) for more information.
### NVIDIA
Nvidia users should install stable pytorch using this command:
-```pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu128```
+```pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu130```
This is the command to install pytorch nightly instead which might have performance improvements.
-```pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu129```
+```pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu130```
#### Troubleshooting
@@ -274,12 +294,6 @@ You can install ComfyUI in Apple Mac silicon (M1 or M2) with any recent macOS ve
> **Note**: Remember to add your models, VAE, LoRAs etc. to the corresponding Comfy folders, as discussed in [ComfyUI manual installation](#manual-install-windows-linux).
-#### DirectML (AMD Cards on Windows)
-
-This is very badly supported and is not recommended. There are some unofficial builds of pytorch ROCm on windows that exist that will give you a much better experience than this. This readme will be updated once official pytorch ROCm builds for windows come out.
-
-```pip install torch-directml``` Then you can launch ComfyUI with: ```python main.py --directml```
-
#### Ascend NPUs
For models compatible with Ascend Extension for PyTorch (torch_npu). To get started, ensure your environment meets the prerequisites outlined on the [installation](https://ascend.github.io/docs/sources/ascend/quick_install.html) page. Here's a step-by-step guide tailored to your platform and installation method:
@@ -297,6 +311,13 @@ For models compatible with Cambricon Extension for PyTorch (torch_mlu). Here's a
2. Next, install the PyTorch(torch_mlu) following the instructions on the [Installation](https://www.cambricon.com/docs/sdk_1.15.0/cambricon_pytorch_1.17.0/user_guide_1.9/index.html)
3. Launch ComfyUI by running `python main.py`
+#### Iluvatar Corex
+
+For models compatible with Iluvatar Extension for PyTorch. Here's a step-by-step guide tailored to your platform and installation method:
+
+1. Install the Iluvatar Corex Toolkit by adhering to the platform-specific instructions on the [Installation](https://support.iluvatar.com/#/DocumentCentre?id=1&nameCenter=2&productId=520117912052801536)
+2. Launch ComfyUI by running `python main.py`
+
# Running
```python main.py```
@@ -347,7 +368,7 @@ Generate a self-signed certificate (not appropriate for shared/production use) a
Use `--tls-keyfile key.pem --tls-certfile cert.pem` to enable TLS/SSL, the app will now be accessible with `https://...` instead of `http://...`.
-> Note: Windows users can use [alexisrolland/docker-openssl](https://github.com/alexisrolland/docker-openssl) or one of the [3rd party binary distributions](https://wiki.openssl.org/index.php/Binaries) to run the command example above.
+> Note: Windows users can use [alexisrolland/docker-openssl](https://github.com/alexisrolland/docker-openssl) or one of the [3rd party binary distributions](https://wiki.openssl.org/index.php/Binaries) to run the command example above.
If you use a container, note that the volume mount `-v` can be a relative path so `... -v ".\:/openssl-certs" ...` would create the key & cert files in the current directory of your command prompt or powershell terminal.
## Support and dev channel
diff --git a/app/frontend_management.py b/app/frontend_management.py
index 001ebbecb..cce0c117d 100644
--- a/app/frontend_management.py
+++ b/app/frontend_management.py
@@ -29,18 +29,50 @@ def frontend_install_warning_message():
This error is happening because the ComfyUI frontend is no longer shipped as part of the main repo but as a pip package instead.
""".strip()
+def parse_version(version: str) -> tuple[int, int, int]:
+ return tuple(map(int, version.split(".")))
+
+def is_valid_version(version: str) -> bool:
+ """Validate if a string is a valid semantic version (X.Y.Z format)."""
+ pattern = r"^(\d+)\.(\d+)\.(\d+)$"
+ return bool(re.match(pattern, version))
+
+def get_installed_frontend_version():
+ """Get the currently installed frontend package version."""
+ frontend_version_str = version("comfyui-frontend-package")
+ return frontend_version_str
+
+
+def get_required_frontend_version():
+ """Get the required frontend version from requirements.txt."""
+ try:
+ with open(requirements_path, "r", encoding="utf-8") as f:
+ for line in f:
+ line = line.strip()
+ if line.startswith("comfyui-frontend-package=="):
+ version_str = line.split("==")[-1]
+ if not is_valid_version(version_str):
+ logging.error(f"Invalid version format in requirements.txt: {version_str}")
+ return None
+ return version_str
+ logging.error("comfyui-frontend-package not found in requirements.txt")
+ return None
+ except FileNotFoundError:
+ logging.error("requirements.txt not found. Cannot determine required frontend version.")
+ return None
+ except Exception as e:
+ logging.error(f"Error reading requirements.txt: {e}")
+ return None
+
def check_frontend_version():
"""Check if the frontend version is up to date."""
- def parse_version(version: str) -> tuple[int, int, int]:
- return tuple(map(int, version.split(".")))
-
try:
- frontend_version_str = version("comfyui-frontend-package")
+ frontend_version_str = get_installed_frontend_version()
frontend_version = parse_version(frontend_version_str)
- with open(requirements_path, "r", encoding="utf-8") as f:
- required_frontend = parse_version(f.readline().split("=")[-1])
+ required_frontend_str = get_required_frontend_version()
+ required_frontend = parse_version(required_frontend_str)
if frontend_version < required_frontend:
app.logger.log_startup_warning(
f"""
@@ -168,6 +200,42 @@ def download_release_asset_zip(release: Release, destination_path: str) -> None:
class FrontendManager:
CUSTOM_FRONTENDS_ROOT = str(Path(__file__).parents[1] / "web_custom_versions")
+ @classmethod
+ def get_required_frontend_version(cls) -> str:
+ """Get the required frontend package version."""
+ return get_required_frontend_version()
+
+ @classmethod
+ def get_installed_templates_version(cls) -> str:
+ """Get the currently installed workflow templates package version."""
+ try:
+ templates_version_str = version("comfyui-workflow-templates")
+ return templates_version_str
+ except Exception:
+ return None
+
+ @classmethod
+ def get_required_templates_version(cls) -> str:
+ """Get the required workflow templates version from requirements.txt."""
+ try:
+ with open(requirements_path, "r", encoding="utf-8") as f:
+ for line in f:
+ line = line.strip()
+ if line.startswith("comfyui-workflow-templates=="):
+ version_str = line.split("==")[-1]
+ if not is_valid_version(version_str):
+ logging.error(f"Invalid templates version format in requirements.txt: {version_str}")
+ return None
+ return version_str
+ logging.error("comfyui-workflow-templates not found in requirements.txt")
+ return None
+ except FileNotFoundError:
+ logging.error("requirements.txt not found. Cannot determine required templates version.")
+ return None
+ except Exception as e:
+ logging.error(f"Error reading requirements.txt: {e}")
+ return None
+
@classmethod
def default_frontend_path(cls) -> str:
try:
diff --git a/app/model_manager.py b/app/model_manager.py
index 74d942fb8..ab36bca74 100644
--- a/app/model_manager.py
+++ b/app/model_manager.py
@@ -130,10 +130,21 @@ class ModelFileManager:
for file_name in filenames:
try:
- relative_path = os.path.relpath(os.path.join(dirpath, file_name), directory)
- result.append(relative_path)
- except:
- logging.warning(f"Warning: Unable to access {file_name}. Skipping this file.")
+ full_path = os.path.join(dirpath, file_name)
+ relative_path = os.path.relpath(full_path, directory)
+
+ # Get file metadata
+ file_info = {
+ "name": relative_path,
+ "pathIndex": pathIndex,
+ "modified": os.path.getmtime(full_path), # Add modification time
+ "created": os.path.getctime(full_path), # Add creation time
+ "size": os.path.getsize(full_path) # Add file size
+ }
+ result.append(file_info)
+
+ except Exception as e:
+ logging.warning(f"Warning: Unable to access {file_name}. Error: {e}. Skipping this file.")
continue
for d in subdirs:
@@ -144,7 +155,7 @@ class ModelFileManager:
logging.warning(f"Warning: Unable to access {path}. Skipping this path.")
continue
- return [{"name": f, "pathIndex": pathIndex} for f in result], dirs, time.perf_counter()
+ return result, dirs, time.perf_counter()
def get_model_previews(self, filepath: str) -> list[str | BytesIO]:
dirname = os.path.dirname(filepath)
diff --git a/app/subgraph_manager.py b/app/subgraph_manager.py
new file mode 100644
index 000000000..dbe404541
--- /dev/null
+++ b/app/subgraph_manager.py
@@ -0,0 +1,112 @@
+from __future__ import annotations
+
+from typing import TypedDict
+import os
+import folder_paths
+import glob
+from aiohttp import web
+import hashlib
+
+
+class Source:
+ custom_node = "custom_node"
+
+class SubgraphEntry(TypedDict):
+ source: str
+ """
+ Source of subgraph - custom_nodes vs templates.
+ """
+ path: str
+ """
+ Relative path of the subgraph file.
+ For custom nodes, will be the relative directory like /subgraphs/.json
+ """
+ name: str
+ """
+ Name of subgraph file.
+ """
+ info: CustomNodeSubgraphEntryInfo
+ """
+ Additional info about subgraph; in the case of custom_nodes, will contain nodepack name
+ """
+ data: str
+
+class CustomNodeSubgraphEntryInfo(TypedDict):
+ node_pack: str
+ """Node pack name."""
+
+class SubgraphManager:
+ def __init__(self):
+ self.cached_custom_node_subgraphs: dict[SubgraphEntry] | None = None
+
+ async def load_entry_data(self, entry: SubgraphEntry):
+ with open(entry['path'], 'r') as f:
+ entry['data'] = f.read()
+ return entry
+
+ async def sanitize_entry(self, entry: SubgraphEntry | None, remove_data=False) -> SubgraphEntry | None:
+ if entry is None:
+ return None
+ entry = entry.copy()
+ entry.pop('path', None)
+ if remove_data:
+ entry.pop('data', None)
+ return entry
+
+ async def sanitize_entries(self, entries: dict[str, SubgraphEntry], remove_data=False) -> dict[str, SubgraphEntry]:
+ entries = entries.copy()
+ for key in list(entries.keys()):
+ entries[key] = await self.sanitize_entry(entries[key], remove_data)
+ return entries
+
+ async def get_custom_node_subgraphs(self, loadedModules, force_reload=False):
+ # if not forced to reload and cached, return cache
+ if not force_reload and self.cached_custom_node_subgraphs is not None:
+ return self.cached_custom_node_subgraphs
+ # Load subgraphs from custom nodes
+ subfolder = "subgraphs"
+ subgraphs_dict: dict[SubgraphEntry] = {}
+
+ for folder in folder_paths.get_folder_paths("custom_nodes"):
+ pattern = os.path.join(folder, f"*/{subfolder}/*.json")
+ matched_files = glob.glob(pattern)
+ for file in matched_files:
+ # replace backslashes with forward slashes
+ file = file.replace('\\', '/')
+ info: CustomNodeSubgraphEntryInfo = {
+ "node_pack": "custom_nodes." + file.split('/')[-3]
+ }
+ source = Source.custom_node
+ # hash source + path to make sure id will be as unique as possible, but
+ # reproducible across backend reloads
+ id = hashlib.sha256(f"{source}{file}".encode()).hexdigest()
+ entry: SubgraphEntry = {
+ "source": Source.custom_node,
+ "name": os.path.splitext(os.path.basename(file))[0],
+ "path": file,
+ "info": info,
+ }
+ subgraphs_dict[id] = entry
+ self.cached_custom_node_subgraphs = subgraphs_dict
+ return subgraphs_dict
+
+ async def get_custom_node_subgraph(self, id: str, loadedModules):
+ subgraphs = await self.get_custom_node_subgraphs(loadedModules)
+ entry: SubgraphEntry = subgraphs.get(id, None)
+ if entry is not None and entry.get('data', None) is None:
+ await self.load_entry_data(entry)
+ return entry
+
+ def add_routes(self, routes, loadedModules):
+ @routes.get("/global_subgraphs")
+ async def get_global_subgraphs(request):
+ subgraphs_dict = await self.get_custom_node_subgraphs(loadedModules)
+ # NOTE: we may want to include other sources of global subgraphs such as templates in the future;
+ # that's the reasoning for the current implementation
+ return web.json_response(await self.sanitize_entries(subgraphs_dict, remove_data=True))
+
+ @routes.get("/global_subgraphs/{id}")
+ async def get_global_subgraph(request):
+ id = request.match_info.get("id", None)
+ subgraph = await self.get_custom_node_subgraph(id, loadedModules)
+ return web.json_response(await self.sanitize_entry(subgraph))
diff --git a/app/user_manager.py b/app/user_manager.py
index d31da5b9b..a2d376c0c 100644
--- a/app/user_manager.py
+++ b/app/user_manager.py
@@ -20,13 +20,15 @@ class FileInfo(TypedDict):
path: str
size: int
modified: int
+ created: int
def get_file_info(path: str, relative_to: str) -> FileInfo:
return {
"path": os.path.relpath(path, relative_to).replace(os.sep, '/'),
"size": os.path.getsize(path),
- "modified": os.path.getmtime(path)
+ "modified": os.path.getmtime(path),
+ "created": os.path.getctime(path)
}
@@ -361,10 +363,17 @@ class UserManager():
if not overwrite and os.path.exists(path):
return web.Response(status=409, text="File already exists")
- body = await request.read()
+ try:
+ body = await request.read()
- with open(path, "wb") as f:
- f.write(body)
+ with open(path, "wb") as f:
+ f.write(body)
+ except OSError as e:
+ logging.warning(f"Error saving file '{path}': {e}")
+ return web.Response(
+ status=400,
+ reason="Invalid filename. Please avoid special characters like :\\/*?\"<>|"
+ )
user_path = self.get_request_user_filepath(request, None)
if full_info:
diff --git a/comfy/audio_encoders/audio_encoders.py b/comfy/audio_encoders/audio_encoders.py
new file mode 100644
index 000000000..46ef21c95
--- /dev/null
+++ b/comfy/audio_encoders/audio_encoders.py
@@ -0,0 +1,91 @@
+from .wav2vec2 import Wav2Vec2Model
+from .whisper import WhisperLargeV3
+import comfy.model_management
+import comfy.ops
+import comfy.utils
+import logging
+import torchaudio
+
+
+class AudioEncoderModel():
+ def __init__(self, config):
+ self.load_device = comfy.model_management.text_encoder_device()
+ offload_device = comfy.model_management.text_encoder_offload_device()
+ self.dtype = comfy.model_management.text_encoder_dtype(self.load_device)
+ model_type = config.pop("model_type")
+ model_config = dict(config)
+ model_config.update({
+ "dtype": self.dtype,
+ "device": offload_device,
+ "operations": comfy.ops.manual_cast
+ })
+
+ if model_type == "wav2vec2":
+ self.model = Wav2Vec2Model(**model_config)
+ elif model_type == "whisper3":
+ self.model = WhisperLargeV3(**model_config)
+ self.model.eval()
+ self.patcher = comfy.model_patcher.ModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device)
+ self.model_sample_rate = 16000
+
+ def load_sd(self, sd):
+ return self.model.load_state_dict(sd, strict=False)
+
+ def get_sd(self):
+ return self.model.state_dict()
+
+ def encode_audio(self, audio, sample_rate):
+ comfy.model_management.load_model_gpu(self.patcher)
+ audio = torchaudio.functional.resample(audio, sample_rate, self.model_sample_rate)
+ out, all_layers = self.model(audio.to(self.load_device))
+ outputs = {}
+ outputs["encoded_audio"] = out
+ outputs["encoded_audio_all_layers"] = all_layers
+ outputs["audio_samples"] = audio.shape[2]
+ return outputs
+
+
+def load_audio_encoder_from_sd(sd, prefix=""):
+ sd = comfy.utils.state_dict_prefix_replace(sd, {"wav2vec2.": ""})
+ if "encoder.layer_norm.bias" in sd: #wav2vec2
+ embed_dim = sd["encoder.layer_norm.bias"].shape[0]
+ if embed_dim == 1024:# large
+ config = {
+ "model_type": "wav2vec2",
+ "embed_dim": 1024,
+ "num_heads": 16,
+ "num_layers": 24,
+ "conv_norm": True,
+ "conv_bias": True,
+ "do_normalize": True,
+ "do_stable_layer_norm": True
+ }
+ elif embed_dim == 768: # base
+ config = {
+ "model_type": "wav2vec2",
+ "embed_dim": 768,
+ "num_heads": 12,
+ "num_layers": 12,
+ "conv_norm": False,
+ "conv_bias": False,
+ "do_normalize": False, # chinese-wav2vec2-base has this False
+ "do_stable_layer_norm": False
+ }
+ else:
+ raise RuntimeError("ERROR: audio encoder file is invalid or unsupported embed_dim: {}".format(embed_dim))
+ elif "model.encoder.embed_positions.weight" in sd:
+ sd = comfy.utils.state_dict_prefix_replace(sd, {"model.": ""})
+ config = {
+ "model_type": "whisper3",
+ }
+ else:
+ raise RuntimeError("ERROR: audio encoder not supported.")
+
+ audio_encoder = AudioEncoderModel(config)
+ m, u = audio_encoder.load_sd(sd)
+ if len(m) > 0:
+ logging.warning("missing audio encoder: {}".format(m))
+ if len(u) > 0:
+ logging.warning("unexpected audio encoder: {}".format(u))
+
+ return audio_encoder
diff --git a/comfy/audio_encoders/wav2vec2.py b/comfy/audio_encoders/wav2vec2.py
new file mode 100644
index 000000000..4e34a40a7
--- /dev/null
+++ b/comfy/audio_encoders/wav2vec2.py
@@ -0,0 +1,252 @@
+import torch
+import torch.nn as nn
+from comfy.ldm.modules.attention import optimized_attention_masked
+
+
+class LayerNormConv(nn.Module):
+ def __init__(self, in_channels, out_channels, kernel_size, stride, bias=False, dtype=None, device=None, operations=None):
+ super().__init__()
+ self.conv = operations.Conv1d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, bias=bias, device=device, dtype=dtype)
+ self.layer_norm = operations.LayerNorm(out_channels, elementwise_affine=True, device=device, dtype=dtype)
+
+ def forward(self, x):
+ x = self.conv(x)
+ return torch.nn.functional.gelu(self.layer_norm(x.transpose(-2, -1)).transpose(-2, -1))
+
+class LayerGroupNormConv(nn.Module):
+ def __init__(self, in_channels, out_channels, kernel_size, stride, bias=False, dtype=None, device=None, operations=None):
+ super().__init__()
+ self.conv = operations.Conv1d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, bias=bias, device=device, dtype=dtype)
+ self.layer_norm = operations.GroupNorm(num_groups=out_channels, num_channels=out_channels, affine=True, device=device, dtype=dtype)
+
+ def forward(self, x):
+ x = self.conv(x)
+ return torch.nn.functional.gelu(self.layer_norm(x))
+
+class ConvNoNorm(nn.Module):
+ def __init__(self, in_channels, out_channels, kernel_size, stride, bias=False, dtype=None, device=None, operations=None):
+ super().__init__()
+ self.conv = operations.Conv1d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, bias=bias, device=device, dtype=dtype)
+
+ def forward(self, x):
+ x = self.conv(x)
+ return torch.nn.functional.gelu(x)
+
+
+class ConvFeatureEncoder(nn.Module):
+ def __init__(self, conv_dim, conv_bias=False, conv_norm=True, dtype=None, device=None, operations=None):
+ super().__init__()
+ if conv_norm:
+ self.conv_layers = nn.ModuleList([
+ LayerNormConv(1, conv_dim, kernel_size=10, stride=5, bias=True, device=device, dtype=dtype, operations=operations),
+ LayerNormConv(conv_dim, conv_dim, kernel_size=3, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
+ LayerNormConv(conv_dim, conv_dim, kernel_size=3, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
+ LayerNormConv(conv_dim, conv_dim, kernel_size=3, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
+ LayerNormConv(conv_dim, conv_dim, kernel_size=3, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
+ LayerNormConv(conv_dim, conv_dim, kernel_size=2, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
+ LayerNormConv(conv_dim, conv_dim, kernel_size=2, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
+ ])
+ else:
+ self.conv_layers = nn.ModuleList([
+ LayerGroupNormConv(1, conv_dim, kernel_size=10, stride=5, bias=conv_bias, device=device, dtype=dtype, operations=operations),
+ ConvNoNorm(conv_dim, conv_dim, kernel_size=3, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
+ ConvNoNorm(conv_dim, conv_dim, kernel_size=3, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
+ ConvNoNorm(conv_dim, conv_dim, kernel_size=3, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
+ ConvNoNorm(conv_dim, conv_dim, kernel_size=3, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
+ ConvNoNorm(conv_dim, conv_dim, kernel_size=2, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
+ ConvNoNorm(conv_dim, conv_dim, kernel_size=2, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
+ ])
+
+ def forward(self, x):
+ x = x.unsqueeze(1)
+
+ for conv in self.conv_layers:
+ x = conv(x)
+
+ return x.transpose(1, 2)
+
+
+class FeatureProjection(nn.Module):
+ def __init__(self, conv_dim, embed_dim, dtype=None, device=None, operations=None):
+ super().__init__()
+ self.layer_norm = operations.LayerNorm(conv_dim, eps=1e-05, device=device, dtype=dtype)
+ self.projection = operations.Linear(conv_dim, embed_dim, device=device, dtype=dtype)
+
+ def forward(self, x):
+ x = self.layer_norm(x)
+ x = self.projection(x)
+ return x
+
+
+class PositionalConvEmbedding(nn.Module):
+ def __init__(self, embed_dim=768, kernel_size=128, groups=16):
+ super().__init__()
+ self.conv = nn.Conv1d(
+ embed_dim,
+ embed_dim,
+ kernel_size=kernel_size,
+ padding=kernel_size // 2,
+ groups=groups,
+ )
+ self.conv = torch.nn.utils.parametrizations.weight_norm(self.conv, name="weight", dim=2)
+ self.activation = nn.GELU()
+
+ def forward(self, x):
+ x = x.transpose(1, 2)
+ x = self.conv(x)[:, :, :-1]
+ x = self.activation(x)
+ x = x.transpose(1, 2)
+ return x
+
+
+class TransformerEncoder(nn.Module):
+ def __init__(
+ self,
+ embed_dim=768,
+ num_heads=12,
+ num_layers=12,
+ mlp_ratio=4.0,
+ do_stable_layer_norm=True,
+ dtype=None, device=None, operations=None
+ ):
+ super().__init__()
+
+ self.pos_conv_embed = PositionalConvEmbedding(embed_dim=embed_dim)
+ self.layers = nn.ModuleList([
+ TransformerEncoderLayer(
+ embed_dim=embed_dim,
+ num_heads=num_heads,
+ mlp_ratio=mlp_ratio,
+ do_stable_layer_norm=do_stable_layer_norm,
+ device=device, dtype=dtype, operations=operations
+ )
+ for _ in range(num_layers)
+ ])
+
+ self.layer_norm = operations.LayerNorm(embed_dim, eps=1e-05, device=device, dtype=dtype)
+ self.do_stable_layer_norm = do_stable_layer_norm
+
+ def forward(self, x, mask=None):
+ x = x + self.pos_conv_embed(x)
+ all_x = ()
+ if not self.do_stable_layer_norm:
+ x = self.layer_norm(x)
+ for layer in self.layers:
+ all_x += (x,)
+ x = layer(x, mask)
+ if self.do_stable_layer_norm:
+ x = self.layer_norm(x)
+ all_x += (x,)
+ return x, all_x
+
+
+class Attention(nn.Module):
+ def __init__(self, embed_dim, num_heads, bias=True, dtype=None, device=None, operations=None):
+ super().__init__()
+ self.embed_dim = embed_dim
+ self.num_heads = num_heads
+ self.head_dim = embed_dim // num_heads
+
+ self.k_proj = operations.Linear(embed_dim, embed_dim, bias=bias, device=device, dtype=dtype)
+ self.v_proj = operations.Linear(embed_dim, embed_dim, bias=bias, device=device, dtype=dtype)
+ self.q_proj = operations.Linear(embed_dim, embed_dim, bias=bias, device=device, dtype=dtype)
+ self.out_proj = operations.Linear(embed_dim, embed_dim, bias=bias, device=device, dtype=dtype)
+
+ def forward(self, x, mask=None):
+ assert (mask is None) # TODO?
+ q = self.q_proj(x)
+ k = self.k_proj(x)
+ v = self.v_proj(x)
+
+ out = optimized_attention_masked(q, k, v, self.num_heads)
+ return self.out_proj(out)
+
+
+class FeedForward(nn.Module):
+ def __init__(self, embed_dim, mlp_ratio, dtype=None, device=None, operations=None):
+ super().__init__()
+ self.intermediate_dense = operations.Linear(embed_dim, int(embed_dim * mlp_ratio), device=device, dtype=dtype)
+ self.output_dense = operations.Linear(int(embed_dim * mlp_ratio), embed_dim, device=device, dtype=dtype)
+
+ def forward(self, x):
+ x = self.intermediate_dense(x)
+ x = torch.nn.functional.gelu(x)
+ x = self.output_dense(x)
+ return x
+
+
+class TransformerEncoderLayer(nn.Module):
+ def __init__(
+ self,
+ embed_dim=768,
+ num_heads=12,
+ mlp_ratio=4.0,
+ do_stable_layer_norm=True,
+ dtype=None, device=None, operations=None
+ ):
+ super().__init__()
+
+ self.attention = Attention(embed_dim, num_heads, device=device, dtype=dtype, operations=operations)
+
+ self.layer_norm = operations.LayerNorm(embed_dim, device=device, dtype=dtype)
+ self.feed_forward = FeedForward(embed_dim, mlp_ratio, device=device, dtype=dtype, operations=operations)
+ self.final_layer_norm = operations.LayerNorm(embed_dim, device=device, dtype=dtype)
+ self.do_stable_layer_norm = do_stable_layer_norm
+
+ def forward(self, x, mask=None):
+ residual = x
+ if self.do_stable_layer_norm:
+ x = self.layer_norm(x)
+ x = self.attention(x, mask=mask)
+ x = residual + x
+ if not self.do_stable_layer_norm:
+ x = self.layer_norm(x)
+ return self.final_layer_norm(x + self.feed_forward(x))
+ else:
+ return x + self.feed_forward(self.final_layer_norm(x))
+
+
+class Wav2Vec2Model(nn.Module):
+ """Complete Wav2Vec 2.0 model."""
+
+ def __init__(
+ self,
+ embed_dim=1024,
+ final_dim=256,
+ num_heads=16,
+ num_layers=24,
+ conv_norm=True,
+ conv_bias=True,
+ do_normalize=True,
+ do_stable_layer_norm=True,
+ dtype=None, device=None, operations=None
+ ):
+ super().__init__()
+
+ conv_dim = 512
+ self.feature_extractor = ConvFeatureEncoder(conv_dim, conv_norm=conv_norm, conv_bias=conv_bias, device=device, dtype=dtype, operations=operations)
+ self.feature_projection = FeatureProjection(conv_dim, embed_dim, device=device, dtype=dtype, operations=operations)
+
+ self.masked_spec_embed = nn.Parameter(torch.empty(embed_dim, device=device, dtype=dtype))
+ self.do_normalize = do_normalize
+
+ self.encoder = TransformerEncoder(
+ embed_dim=embed_dim,
+ num_heads=num_heads,
+ num_layers=num_layers,
+ do_stable_layer_norm=do_stable_layer_norm,
+ device=device, dtype=dtype, operations=operations
+ )
+
+ def forward(self, x, mask_time_indices=None, return_dict=False):
+ x = torch.mean(x, dim=1)
+
+ if self.do_normalize:
+ x = (x - x.mean()) / torch.sqrt(x.var() + 1e-7)
+
+ features = self.feature_extractor(x)
+ features = self.feature_projection(features)
+ batch_size, seq_len, _ = features.shape
+
+ x, all_x = self.encoder(features)
+ return x, all_x
diff --git a/comfy/audio_encoders/whisper.py b/comfy/audio_encoders/whisper.py
new file mode 100755
index 000000000..93d3782f1
--- /dev/null
+++ b/comfy/audio_encoders/whisper.py
@@ -0,0 +1,186 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchaudio
+from typing import Optional
+from comfy.ldm.modules.attention import optimized_attention_masked
+import comfy.ops
+
+class WhisperFeatureExtractor(nn.Module):
+ def __init__(self, n_mels=128, device=None):
+ super().__init__()
+ self.sample_rate = 16000
+ self.n_fft = 400
+ self.hop_length = 160
+ self.n_mels = n_mels
+ self.chunk_length = 30
+ self.n_samples = 480000
+
+ self.mel_spectrogram = torchaudio.transforms.MelSpectrogram(
+ sample_rate=self.sample_rate,
+ n_fft=self.n_fft,
+ hop_length=self.hop_length,
+ n_mels=self.n_mels,
+ f_min=0,
+ f_max=8000,
+ norm="slaney",
+ mel_scale="slaney",
+ ).to(device)
+
+ def __call__(self, audio):
+ audio = torch.mean(audio, dim=1)
+ batch_size = audio.shape[0]
+ processed_audio = []
+
+ for i in range(batch_size):
+ aud = audio[i]
+ if aud.shape[0] > self.n_samples:
+ aud = aud[:self.n_samples]
+ elif aud.shape[0] < self.n_samples:
+ aud = F.pad(aud, (0, self.n_samples - aud.shape[0]))
+ processed_audio.append(aud)
+
+ audio = torch.stack(processed_audio)
+
+ mel_spec = self.mel_spectrogram(audio.to(self.mel_spectrogram.spectrogram.window.device))[:, :, :-1].to(audio.device)
+
+ log_mel_spec = torch.clamp(mel_spec, min=1e-10).log10()
+ log_mel_spec = torch.maximum(log_mel_spec, log_mel_spec.max() - 8.0)
+ log_mel_spec = (log_mel_spec + 4.0) / 4.0
+
+ return log_mel_spec
+
+
+class MultiHeadAttention(nn.Module):
+ def __init__(self, d_model: int, n_heads: int, dtype=None, device=None, operations=None):
+ super().__init__()
+ assert d_model % n_heads == 0
+
+ self.d_model = d_model
+ self.n_heads = n_heads
+ self.d_k = d_model // n_heads
+
+ self.q_proj = operations.Linear(d_model, d_model, dtype=dtype, device=device)
+ self.k_proj = operations.Linear(d_model, d_model, bias=False, dtype=dtype, device=device)
+ self.v_proj = operations.Linear(d_model, d_model, dtype=dtype, device=device)
+ self.out_proj = operations.Linear(d_model, d_model, dtype=dtype, device=device)
+
+ def forward(
+ self,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ mask: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ batch_size, seq_len, _ = query.shape
+
+ q = self.q_proj(query)
+ k = self.k_proj(key)
+ v = self.v_proj(value)
+
+ attn_output = optimized_attention_masked(q, k, v, self.n_heads, mask)
+ attn_output = self.out_proj(attn_output)
+
+ return attn_output
+
+
+class EncoderLayer(nn.Module):
+ def __init__(self, d_model: int, n_heads: int, d_ff: int, dtype=None, device=None, operations=None):
+ super().__init__()
+
+ self.self_attn = MultiHeadAttention(d_model, n_heads, dtype=dtype, device=device, operations=operations)
+ self.self_attn_layer_norm = operations.LayerNorm(d_model, dtype=dtype, device=device)
+
+ self.fc1 = operations.Linear(d_model, d_ff, dtype=dtype, device=device)
+ self.fc2 = operations.Linear(d_ff, d_model, dtype=dtype, device=device)
+ self.final_layer_norm = operations.LayerNorm(d_model, dtype=dtype, device=device)
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None
+ ) -> torch.Tensor:
+ residual = x
+ x = self.self_attn_layer_norm(x)
+ x = self.self_attn(x, x, x, attention_mask)
+ x = residual + x
+
+ residual = x
+ x = self.final_layer_norm(x)
+ x = self.fc1(x)
+ x = F.gelu(x)
+ x = self.fc2(x)
+ x = residual + x
+
+ return x
+
+
+class AudioEncoder(nn.Module):
+ def __init__(
+ self,
+ n_mels: int = 128,
+ n_ctx: int = 1500,
+ n_state: int = 1280,
+ n_head: int = 20,
+ n_layer: int = 32,
+ dtype=None,
+ device=None,
+ operations=None
+ ):
+ super().__init__()
+
+ self.conv1 = operations.Conv1d(n_mels, n_state, kernel_size=3, padding=1, dtype=dtype, device=device)
+ self.conv2 = operations.Conv1d(n_state, n_state, kernel_size=3, stride=2, padding=1, dtype=dtype, device=device)
+
+ self.embed_positions = operations.Embedding(n_ctx, n_state, dtype=dtype, device=device)
+
+ self.layers = nn.ModuleList([
+ EncoderLayer(n_state, n_head, n_state * 4, dtype=dtype, device=device, operations=operations)
+ for _ in range(n_layer)
+ ])
+
+ self.layer_norm = operations.LayerNorm(n_state, dtype=dtype, device=device)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x = F.gelu(self.conv1(x))
+ x = F.gelu(self.conv2(x))
+
+ x = x.transpose(1, 2)
+
+ x = x + comfy.ops.cast_to_input(self.embed_positions.weight[:, :x.shape[1]], x)
+
+ all_x = ()
+ for layer in self.layers:
+ all_x += (x,)
+ x = layer(x)
+
+ x = self.layer_norm(x)
+ all_x += (x,)
+ return x, all_x
+
+
+class WhisperLargeV3(nn.Module):
+ def __init__(
+ self,
+ n_mels: int = 128,
+ n_audio_ctx: int = 1500,
+ n_audio_state: int = 1280,
+ n_audio_head: int = 20,
+ n_audio_layer: int = 32,
+ dtype=None,
+ device=None,
+ operations=None
+ ):
+ super().__init__()
+
+ self.feature_extractor = WhisperFeatureExtractor(n_mels=n_mels, device=device)
+
+ self.encoder = AudioEncoder(
+ n_mels, n_audio_ctx, n_audio_state, n_audio_head, n_audio_layer,
+ dtype=dtype, device=device, operations=operations
+ )
+
+ def forward(self, audio):
+ mel = self.feature_extractor(audio)
+ x, all_x = self.encoder(mel)
+ return x, all_x
diff --git a/comfy/cli_args.py b/comfy/cli_args.py
index 7234a7ba0..3d5bc7c90 100644
--- a/comfy/cli_args.py
+++ b/comfy/cli_args.py
@@ -49,7 +49,8 @@ parser.add_argument("--temp-directory", type=str, default=None, help="Set the Co
parser.add_argument("--input-directory", type=str, default=None, help="Set the ComfyUI input directory. Overrides --base-directory.")
parser.add_argument("--auto-launch", action="store_true", help="Automatically launch ComfyUI in the default browser.")
parser.add_argument("--disable-auto-launch", action="store_true", help="Disable auto launching the browser.")
-parser.add_argument("--cuda-device", type=int, default=None, metavar="DEVICE_ID", help="Set the id of the cuda device this instance will use.")
+parser.add_argument("--cuda-device", type=int, default=None, metavar="DEVICE_ID", help="Set the id of the cuda device this instance will use. All other devices will not be visible.")
+parser.add_argument("--default-device", type=int, default=None, metavar="DEFAULT_DEVICE_ID", help="Set the id of the default device, all other devices will stay visible.")
cm_group = parser.add_mutually_exclusive_group()
cm_group.add_argument("--cuda-malloc", action="store_true", help="Enable cudaMallocAsync (enabled by default for torch 2.0 and up).")
cm_group.add_argument("--disable-cuda-malloc", action="store_true", help="Disable cudaMallocAsync.")
@@ -104,6 +105,7 @@ cache_group = parser.add_mutually_exclusive_group()
cache_group.add_argument("--cache-classic", action="store_true", help="Use the old style (aggressive) caching.")
cache_group.add_argument("--cache-lru", type=int, default=0, help="Use LRU caching with a maximum of N node results cached. May use more RAM/VRAM.")
cache_group.add_argument("--cache-none", action="store_true", help="Reduced RAM/VRAM usage at the expense of executing every node for each run.")
+cache_group.add_argument("--cache-ram", nargs='?', const=4.0, type=float, default=0, help="Use RAM pressure caching with the specified headroom threshold. If available RAM drops below the threhold the cache remove large items to free RAM. Default 4GB")
attn_group = parser.add_mutually_exclusive_group()
attn_group.add_argument("--use-split-cross-attention", action="store_true", help="Use the split cross attention optimization. Ignored when xformers is used.")
@@ -131,6 +133,8 @@ parser.add_argument("--reserve-vram", type=float, default=None, help="Set the am
parser.add_argument("--async-offload", action="store_true", help="Use async weight offloading.")
+parser.add_argument("--force-non-blocking", action="store_true", help="Force ComfyUI to use non-blocking operations for all applicable tensors. This may improve performance on some non-Nvidia systems but can cause issues with some workflows.")
+
parser.add_argument("--default-hashing-function", type=str, choices=['md5', 'sha1', 'sha256', 'sha512'], default='sha256', help="Allows you to choose the hash function to use for duplicate filename / contents comparison. Default is sha256.")
parser.add_argument("--disable-smart-memory", action="store_true", help="Force ComfyUI to agressively offload to regular ram instead of keeping models in vram when it can.")
@@ -140,10 +144,13 @@ class PerformanceFeature(enum.Enum):
Fp16Accumulation = "fp16_accumulation"
Fp8MatrixMultiplication = "fp8_matrix_mult"
CublasOps = "cublas_ops"
+ AutoTune = "autotune"
+ PinnedMem = "pinned_memory"
-parser.add_argument("--fast", nargs="*", type=PerformanceFeature, help="Enable some untested and potentially quality deteriorating optimizations. --fast with no arguments enables everything. You can pass a list specific optimizations if you only want to enable specific ones. Current valid optimizations: fp16_accumulation fp8_matrix_mult cublas_ops")
+parser.add_argument("--fast", nargs="*", type=PerformanceFeature, help="Enable some untested and potentially quality deteriorating optimizations. --fast with no arguments enables everything. You can pass a list specific optimizations if you only want to enable specific ones. Current valid optimizations: {}".format(" ".join(map(lambda c: c.value, PerformanceFeature))))
parser.add_argument("--mmap-torch-files", action="store_true", help="Use mmap when loading ckpt/pt files.")
+parser.add_argument("--disable-mmap", action="store_true", help="Don't use mmap when loading safetensors.")
parser.add_argument("--dont-print-server", action="store_true", help="Don't print server output.")
parser.add_argument("--quick-test-for-ci", action="store_true", help="Quick test for CI.")
diff --git a/comfy/clip_model.py b/comfy/clip_model.py
index c8294d483..7c0cadab5 100644
--- a/comfy/clip_model.py
+++ b/comfy/clip_model.py
@@ -61,8 +61,12 @@ class CLIPEncoder(torch.nn.Module):
def forward(self, x, mask=None, intermediate_output=None):
optimized_attention = optimized_attention_for_device(x.device, mask=mask is not None, small_input=True)
+ all_intermediate = None
if intermediate_output is not None:
- if intermediate_output < 0:
+ if intermediate_output == "all":
+ all_intermediate = []
+ intermediate_output = None
+ elif intermediate_output < 0:
intermediate_output = len(self.layers) + intermediate_output
intermediate = None
@@ -70,6 +74,12 @@ class CLIPEncoder(torch.nn.Module):
x = l(x, mask, optimized_attention)
if i == intermediate_output:
intermediate = x.clone()
+ if all_intermediate is not None:
+ all_intermediate.append(x.unsqueeze(1).clone())
+
+ if all_intermediate is not None:
+ intermediate = torch.cat(all_intermediate, dim=1)
+
return x, intermediate
class CLIPEmbeddings(torch.nn.Module):
@@ -97,7 +107,7 @@ class CLIPTextModel_(torch.nn.Module):
self.encoder = CLIPEncoder(num_layers, embed_dim, heads, intermediate_size, intermediate_activation, dtype, device, operations)
self.final_layer_norm = operations.LayerNorm(embed_dim, dtype=dtype, device=device)
- def forward(self, input_tokens=None, attention_mask=None, embeds=None, num_tokens=None, intermediate_output=None, final_layer_norm_intermediate=True, dtype=torch.float32):
+ def forward(self, input_tokens=None, attention_mask=None, embeds=None, num_tokens=None, intermediate_output=None, final_layer_norm_intermediate=True, dtype=torch.float32, embeds_info=[]):
if embeds is not None:
x = embeds + comfy.ops.cast_to(self.embeddings.position_embedding.weight, dtype=dtype, device=embeds.device)
else:
diff --git a/comfy/clip_vision.py b/comfy/clip_vision.py
index 00aab9164..447b1ce4a 100644
--- a/comfy/clip_vision.py
+++ b/comfy/clip_vision.py
@@ -50,7 +50,13 @@ class ClipVisionModel():
self.image_size = config.get("image_size", 224)
self.image_mean = config.get("image_mean", [0.48145466, 0.4578275, 0.40821073])
self.image_std = config.get("image_std", [0.26862954, 0.26130258, 0.27577711])
- model_class = IMAGE_ENCODERS.get(config.get("model_type", "clip_vision_model"))
+ model_type = config.get("model_type", "clip_vision_model")
+ model_class = IMAGE_ENCODERS.get(model_type)
+ if model_type == "siglip_vision_model":
+ self.return_all_hidden_states = True
+ else:
+ self.return_all_hidden_states = False
+
self.load_device = comfy.model_management.text_encoder_device()
offload_device = comfy.model_management.text_encoder_offload_device()
self.dtype = comfy.model_management.text_encoder_dtype(self.load_device)
@@ -68,12 +74,18 @@ class ClipVisionModel():
def encode_image(self, image, crop=True):
comfy.model_management.load_model_gpu(self.patcher)
pixel_values = clip_preprocess(image.to(self.load_device), size=self.image_size, mean=self.image_mean, std=self.image_std, crop=crop).float()
- out = self.model(pixel_values=pixel_values, intermediate_output=-2)
+ out = self.model(pixel_values=pixel_values, intermediate_output='all' if self.return_all_hidden_states else -2)
outputs = Output()
outputs["last_hidden_state"] = out[0].to(comfy.model_management.intermediate_device())
outputs["image_embeds"] = out[2].to(comfy.model_management.intermediate_device())
- outputs["penultimate_hidden_states"] = out[1].to(comfy.model_management.intermediate_device())
+ if self.return_all_hidden_states:
+ all_hs = out[1].to(comfy.model_management.intermediate_device())
+ outputs["penultimate_hidden_states"] = all_hs[:, -2]
+ outputs["all_hidden_states"] = all_hs
+ else:
+ outputs["penultimate_hidden_states"] = out[1].to(comfy.model_management.intermediate_device())
+
outputs["mm_projected"] = out[3]
return outputs
@@ -124,8 +136,12 @@ def load_clipvision_from_sd(sd, prefix="", convert_keys=False):
json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "clip_vision_config_vitl_336.json")
else:
json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "clip_vision_config_vitl.json")
- elif "embeddings.patch_embeddings.projection.weight" in sd:
+
+ # Dinov2
+ elif 'encoder.layer.39.layer_scale2.lambda1' in sd:
json_config = os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "image_encoders"), "dino2_giant.json")
+ elif 'encoder.layer.23.layer_scale2.lambda1' in sd:
+ json_config = os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "image_encoders"), "dino2_large.json")
else:
return None
diff --git a/comfy/conds.py b/comfy/conds.py
index 2af2a43a3..5af3e93ea 100644
--- a/comfy/conds.py
+++ b/comfy/conds.py
@@ -1,6 +1,7 @@
import torch
import math
import comfy.utils
+import logging
class CONDRegular:
@@ -10,12 +11,15 @@ class CONDRegular:
def _copy_with(self, cond):
return self.__class__(cond)
- def process_cond(self, batch_size, device, **kwargs):
- return self._copy_with(comfy.utils.repeat_to_batch_size(self.cond, batch_size).to(device))
+ def process_cond(self, batch_size, **kwargs):
+ return self._copy_with(comfy.utils.repeat_to_batch_size(self.cond, batch_size))
def can_concat(self, other):
if self.cond.shape != other.cond.shape:
return False
+ if self.cond.device != other.cond.device:
+ logging.warning("WARNING: conds not on same device, skipping concat.")
+ return False
return True
def concat(self, others):
@@ -29,14 +33,14 @@ class CONDRegular:
class CONDNoiseShape(CONDRegular):
- def process_cond(self, batch_size, device, area, **kwargs):
+ def process_cond(self, batch_size, area, **kwargs):
data = self.cond
if area is not None:
dims = len(area) // 2
for i in range(dims):
data = data.narrow(i + 2, area[i + dims], area[i])
- return self._copy_with(comfy.utils.repeat_to_batch_size(data, batch_size).to(device))
+ return self._copy_with(comfy.utils.repeat_to_batch_size(data, batch_size))
class CONDCrossAttn(CONDRegular):
@@ -51,6 +55,9 @@ class CONDCrossAttn(CONDRegular):
diff = mult_min // min(s1[1], s2[1])
if diff > 4: #arbitrary limit on the padding because it's probably going to impact performance negatively if it's too much
return False
+ if self.cond.device != other.cond.device:
+ logging.warning("WARNING: conds not on same device: skipping concat.")
+ return False
return True
def concat(self, others):
@@ -73,7 +80,7 @@ class CONDConstant(CONDRegular):
def __init__(self, cond):
self.cond = cond
- def process_cond(self, batch_size, device, **kwargs):
+ def process_cond(self, batch_size, **kwargs):
return self._copy_with(self.cond)
def can_concat(self, other):
@@ -92,10 +99,10 @@ class CONDList(CONDRegular):
def __init__(self, cond):
self.cond = cond
- def process_cond(self, batch_size, device, **kwargs):
+ def process_cond(self, batch_size, **kwargs):
out = []
for c in self.cond:
- out.append(comfy.utils.repeat_to_batch_size(c, batch_size).to(device))
+ out.append(comfy.utils.repeat_to_batch_size(c, batch_size))
return self._copy_with(out)
diff --git a/comfy/context_windows.py b/comfy/context_windows.py
new file mode 100644
index 000000000..041f380f9
--- /dev/null
+++ b/comfy/context_windows.py
@@ -0,0 +1,540 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING, Callable
+import torch
+import numpy as np
+import collections
+from dataclasses import dataclass
+from abc import ABC, abstractmethod
+import logging
+import comfy.model_management
+import comfy.patcher_extension
+if TYPE_CHECKING:
+ from comfy.model_base import BaseModel
+ from comfy.model_patcher import ModelPatcher
+ from comfy.controlnet import ControlBase
+
+
+class ContextWindowABC(ABC):
+ def __init__(self):
+ ...
+
+ @abstractmethod
+ def get_tensor(self, full: torch.Tensor) -> torch.Tensor:
+ """
+ Get torch.Tensor applicable to current window.
+ """
+ raise NotImplementedError("Not implemented.")
+
+ @abstractmethod
+ def add_window(self, full: torch.Tensor, to_add: torch.Tensor) -> torch.Tensor:
+ """
+ Apply torch.Tensor of window to the full tensor, in place. Returns reference to updated full tensor, not a copy.
+ """
+ raise NotImplementedError("Not implemented.")
+
+class ContextHandlerABC(ABC):
+ def __init__(self):
+ ...
+
+ @abstractmethod
+ def should_use_context(self, model: BaseModel, conds: list[list[dict]], x_in: torch.Tensor, timestep: torch.Tensor, model_options: dict[str]) -> bool:
+ raise NotImplementedError("Not implemented.")
+
+ @abstractmethod
+ def get_resized_cond(self, cond_in: list[dict], x_in: torch.Tensor, window: ContextWindowABC, device=None) -> list:
+ raise NotImplementedError("Not implemented.")
+
+ @abstractmethod
+ def execute(self, calc_cond_batch: Callable, model: BaseModel, conds: list[list[dict]], x_in: torch.Tensor, timestep: torch.Tensor, model_options: dict[str]):
+ raise NotImplementedError("Not implemented.")
+
+
+
+class IndexListContextWindow(ContextWindowABC):
+ def __init__(self, index_list: list[int], dim: int=0):
+ self.index_list = index_list
+ self.context_length = len(index_list)
+ self.dim = dim
+
+ def get_tensor(self, full: torch.Tensor, device=None, dim=None) -> torch.Tensor:
+ if dim is None:
+ dim = self.dim
+ if dim == 0 and full.shape[dim] == 1:
+ return full
+ idx = [slice(None)] * dim + [self.index_list]
+ return full[idx].to(device)
+
+ def add_window(self, full: torch.Tensor, to_add: torch.Tensor, dim=None) -> torch.Tensor:
+ if dim is None:
+ dim = self.dim
+ idx = [slice(None)] * dim + [self.index_list]
+ full[idx] += to_add
+ return full
+
+
+class IndexListCallbacks:
+ EVALUATE_CONTEXT_WINDOWS = "evaluate_context_windows"
+ COMBINE_CONTEXT_WINDOW_RESULTS = "combine_context_window_results"
+ EXECUTE_START = "execute_start"
+ EXECUTE_CLEANUP = "execute_cleanup"
+
+ def init_callbacks(self):
+ return {}
+
+
+@dataclass
+class ContextSchedule:
+ name: str
+ func: Callable
+
+@dataclass
+class ContextFuseMethod:
+ name: str
+ func: Callable
+
+ContextResults = collections.namedtuple("ContextResults", ['window_idx', 'sub_conds_out', 'sub_conds', 'window'])
+class IndexListContextHandler(ContextHandlerABC):
+ def __init__(self, context_schedule: ContextSchedule, fuse_method: ContextFuseMethod, context_length: int=1, context_overlap: int=0, context_stride: int=1, closed_loop=False, dim=0):
+ self.context_schedule = context_schedule
+ self.fuse_method = fuse_method
+ self.context_length = context_length
+ self.context_overlap = context_overlap
+ self.context_stride = context_stride
+ self.closed_loop = closed_loop
+ self.dim = dim
+ self._step = 0
+
+ self.callbacks = {}
+
+ def should_use_context(self, model: BaseModel, conds: list[list[dict]], x_in: torch.Tensor, timestep: torch.Tensor, model_options: dict[str]) -> bool:
+ # for now, assume first dim is batch - should have stored on BaseModel in actual implementation
+ if x_in.size(self.dim) > self.context_length:
+ logging.info(f"Using context windows {self.context_length} for {x_in.size(self.dim)} frames.")
+ return True
+ return False
+
+ def prepare_control_objects(self, control: ControlBase, device=None) -> ControlBase:
+ if control.previous_controlnet is not None:
+ self.prepare_control_objects(control.previous_controlnet, device)
+ return control
+
+ def get_resized_cond(self, cond_in: list[dict], x_in: torch.Tensor, window: IndexListContextWindow, device=None) -> list:
+ if cond_in is None:
+ return None
+ # reuse or resize cond items to match context requirements
+ resized_cond = []
+ # cond object is a list containing a dict - outer list is irrelevant, so just loop through it
+ for actual_cond in cond_in:
+ resized_actual_cond = actual_cond.copy()
+ # now we are in the inner dict - "pooled_output" is a tensor, "control" is a ControlBase object, "model_conds" is dictionary
+ for key in actual_cond:
+ try:
+ cond_item = actual_cond[key]
+ if isinstance(cond_item, torch.Tensor):
+ # check that tensor is the expected length - x.size(0)
+ if self.dim < cond_item.ndim and cond_item.size(self.dim) == x_in.size(self.dim):
+ # if so, it's subsetting time - tell controls the expected indeces so they can handle them
+ actual_cond_item = window.get_tensor(cond_item)
+ resized_actual_cond[key] = actual_cond_item.to(device)
+ else:
+ resized_actual_cond[key] = cond_item.to(device)
+ # look for control
+ elif key == "control":
+ resized_actual_cond[key] = self.prepare_control_objects(cond_item, device)
+ elif isinstance(cond_item, dict):
+ new_cond_item = cond_item.copy()
+ # when in dictionary, look for tensors and CONDCrossAttn [comfy/conds.py] (has cond attr that is a tensor)
+ for cond_key, cond_value in new_cond_item.items():
+ if isinstance(cond_value, torch.Tensor):
+ if cond_value.ndim < self.dim and cond_value.size(0) == x_in.size(self.dim):
+ new_cond_item[cond_key] = window.get_tensor(cond_value, device)
+ # if has cond that is a Tensor, check if needs to be subset
+ elif hasattr(cond_value, "cond") and isinstance(cond_value.cond, torch.Tensor):
+ if cond_value.cond.ndim < self.dim and cond_value.cond.size(0) == x_in.size(self.dim):
+ new_cond_item[cond_key] = cond_value._copy_with(window.get_tensor(cond_value.cond, device))
+ elif cond_key == "num_video_frames": # for SVD
+ new_cond_item[cond_key] = cond_value._copy_with(cond_value.cond)
+ new_cond_item[cond_key].cond = window.context_length
+ resized_actual_cond[key] = new_cond_item
+ else:
+ resized_actual_cond[key] = cond_item
+ finally:
+ del cond_item # just in case to prevent VRAM issues
+ resized_cond.append(resized_actual_cond)
+ return resized_cond
+
+ def set_step(self, timestep: torch.Tensor, model_options: dict[str]):
+ mask = torch.isclose(model_options["transformer_options"]["sample_sigmas"], timestep, rtol=0.0001)
+ matches = torch.nonzero(mask)
+ if torch.numel(matches) == 0:
+ raise Exception("No sample_sigmas matched current timestep; something went wrong.")
+ self._step = int(matches[0].item())
+
+ def get_context_windows(self, model: BaseModel, x_in: torch.Tensor, model_options: dict[str]) -> list[IndexListContextWindow]:
+ full_length = x_in.size(self.dim) # TODO: choose dim based on model
+ context_windows = self.context_schedule.func(full_length, self, model_options)
+ context_windows = [IndexListContextWindow(window, dim=self.dim) for window in context_windows]
+ return context_windows
+
+ def execute(self, calc_cond_batch: Callable, model: BaseModel, conds: list[list[dict]], x_in: torch.Tensor, timestep: torch.Tensor, model_options: dict[str]):
+ self.set_step(timestep, model_options)
+ context_windows = self.get_context_windows(model, x_in, model_options)
+ enumerated_context_windows = list(enumerate(context_windows))
+
+ conds_final = [torch.zeros_like(x_in) for _ in conds]
+ if self.fuse_method.name == ContextFuseMethods.RELATIVE:
+ counts_final = [torch.ones(get_shape_for_dim(x_in, self.dim), device=x_in.device) for _ in conds]
+ else:
+ counts_final = [torch.zeros(get_shape_for_dim(x_in, self.dim), device=x_in.device) for _ in conds]
+ biases_final = [([0.0] * x_in.shape[self.dim]) for _ in conds]
+
+ for callback in comfy.patcher_extension.get_all_callbacks(IndexListCallbacks.EXECUTE_START, self.callbacks):
+ callback(self, model, x_in, conds, timestep, model_options)
+
+ for enum_window in enumerated_context_windows:
+ results = self.evaluate_context_windows(calc_cond_batch, model, x_in, conds, timestep, [enum_window], model_options)
+ for result in results:
+ self.combine_context_window_results(x_in, result.sub_conds_out, result.sub_conds, result.window, result.window_idx, len(enumerated_context_windows), timestep,
+ conds_final, counts_final, biases_final)
+ try:
+ # finalize conds
+ if self.fuse_method.name == ContextFuseMethods.RELATIVE:
+ # relative is already normalized, so return as is
+ del counts_final
+ return conds_final
+ else:
+ # normalize conds via division by context usage counts
+ for i in range(len(conds_final)):
+ conds_final[i] /= counts_final[i]
+ del counts_final
+ return conds_final
+ finally:
+ for callback in comfy.patcher_extension.get_all_callbacks(IndexListCallbacks.EXECUTE_CLEANUP, self.callbacks):
+ callback(self, model, x_in, conds, timestep, model_options)
+
+ def evaluate_context_windows(self, calc_cond_batch: Callable, model: BaseModel, x_in: torch.Tensor, conds, timestep: torch.Tensor, enumerated_context_windows: list[tuple[int, IndexListContextWindow]],
+ model_options, device=None, first_device=None):
+ results: list[ContextResults] = []
+ for window_idx, window in enumerated_context_windows:
+ # allow processing to end between context window executions for faster Cancel
+ comfy.model_management.throw_exception_if_processing_interrupted()
+
+ for callback in comfy.patcher_extension.get_all_callbacks(IndexListCallbacks.EVALUATE_CONTEXT_WINDOWS, self.callbacks):
+ callback(self, model, x_in, conds, timestep, model_options, window_idx, window, model_options, device, first_device)
+
+ # update exposed params
+ model_options["transformer_options"]["context_window"] = window
+ # get subsections of x, timestep, conds
+ sub_x = window.get_tensor(x_in, device)
+ sub_timestep = window.get_tensor(timestep, device, dim=0)
+ sub_conds = [self.get_resized_cond(cond, x_in, window, device) for cond in conds]
+
+ sub_conds_out = calc_cond_batch(model, sub_conds, sub_x, sub_timestep, model_options)
+ if device is not None:
+ for i in range(len(sub_conds_out)):
+ sub_conds_out[i] = sub_conds_out[i].to(x_in.device)
+ results.append(ContextResults(window_idx, sub_conds_out, sub_conds, window))
+ return results
+
+
+ def combine_context_window_results(self, x_in: torch.Tensor, sub_conds_out, sub_conds, window: IndexListContextWindow, window_idx: int, total_windows: int, timestep: torch.Tensor,
+ conds_final: list[torch.Tensor], counts_final: list[torch.Tensor], biases_final: list[torch.Tensor]):
+ if self.fuse_method.name == ContextFuseMethods.RELATIVE:
+ for pos, idx in enumerate(window.index_list):
+ # bias is the influence of a specific index in relation to the whole context window
+ bias = 1 - abs(idx - (window.index_list[0] + window.index_list[-1]) / 2) / ((window.index_list[-1] - window.index_list[0] + 1e-2) / 2)
+ bias = max(1e-2, bias)
+ # take weighted average relative to total bias of current idx
+ for i in range(len(sub_conds_out)):
+ bias_total = biases_final[i][idx]
+ prev_weight = (bias_total / (bias_total + bias))
+ new_weight = (bias / (bias_total + bias))
+ # account for dims of tensors
+ idx_window = [slice(None)] * self.dim + [idx]
+ pos_window = [slice(None)] * self.dim + [pos]
+ # apply new values
+ conds_final[i][idx_window] = conds_final[i][idx_window] * prev_weight + sub_conds_out[i][pos_window] * new_weight
+ biases_final[i][idx] = bias_total + bias
+ else:
+ # add conds and counts based on weights of fuse method
+ weights = get_context_weights(window.context_length, x_in.shape[self.dim], window.index_list, self, sigma=timestep)
+ weights_tensor = match_weights_to_dim(weights, x_in, self.dim, device=x_in.device)
+ for i in range(len(sub_conds_out)):
+ window.add_window(conds_final[i], sub_conds_out[i] * weights_tensor)
+ window.add_window(counts_final[i], weights_tensor)
+
+ for callback in comfy.patcher_extension.get_all_callbacks(IndexListCallbacks.COMBINE_CONTEXT_WINDOW_RESULTS, self.callbacks):
+ callback(self, x_in, sub_conds_out, sub_conds, window, window_idx, total_windows, timestep, conds_final, counts_final, biases_final)
+
+
+def _prepare_sampling_wrapper(executor, model, noise_shape: torch.Tensor, *args, **kwargs):
+ # limit noise_shape length to context_length for more accurate vram use estimation
+ model_options = kwargs.get("model_options", None)
+ if model_options is None:
+ raise Exception("model_options not found in prepare_sampling_wrapper; this should never happen, something went wrong.")
+ handler: IndexListContextHandler = model_options.get("context_handler", None)
+ if handler is not None:
+ noise_shape = list(noise_shape)
+ noise_shape[handler.dim] = min(noise_shape[handler.dim], handler.context_length)
+ return executor(model, noise_shape, *args, **kwargs)
+
+
+def create_prepare_sampling_wrapper(model: ModelPatcher):
+ model.add_wrapper_with_key(
+ comfy.patcher_extension.WrappersMP.PREPARE_SAMPLING,
+ "ContextWindows_prepare_sampling",
+ _prepare_sampling_wrapper
+ )
+
+
+def match_weights_to_dim(weights: list[float], x_in: torch.Tensor, dim: int, device=None) -> torch.Tensor:
+ total_dims = len(x_in.shape)
+ weights_tensor = torch.Tensor(weights).to(device=device)
+ for _ in range(dim):
+ weights_tensor = weights_tensor.unsqueeze(0)
+ for _ in range(total_dims - dim - 1):
+ weights_tensor = weights_tensor.unsqueeze(-1)
+ return weights_tensor
+
+def get_shape_for_dim(x_in: torch.Tensor, dim: int) -> list[int]:
+ total_dims = len(x_in.shape)
+ shape = []
+ for _ in range(dim):
+ shape.append(1)
+ shape.append(x_in.shape[dim])
+ for _ in range(total_dims - dim - 1):
+ shape.append(1)
+ return shape
+
+class ContextSchedules:
+ UNIFORM_LOOPED = "looped_uniform"
+ UNIFORM_STANDARD = "standard_uniform"
+ STATIC_STANDARD = "standard_static"
+ BATCHED = "batched"
+
+
+# from https://github.com/neggles/animatediff-cli/blob/main/src/animatediff/pipelines/context.py
+def create_windows_uniform_looped(num_frames: int, handler: IndexListContextHandler, model_options: dict[str]):
+ windows = []
+ if num_frames < handler.context_length:
+ windows.append(list(range(num_frames)))
+ return windows
+
+ context_stride = min(handler.context_stride, int(np.ceil(np.log2(num_frames / handler.context_length))) + 1)
+ # obtain uniform windows as normal, looping and all
+ for context_step in 1 << np.arange(context_stride):
+ pad = int(round(num_frames * ordered_halving(handler._step)))
+ for j in range(
+ int(ordered_halving(handler._step) * context_step) + pad,
+ num_frames + pad + (0 if handler.closed_loop else -handler.context_overlap),
+ (handler.context_length * context_step - handler.context_overlap),
+ ):
+ windows.append([e % num_frames for e in range(j, j + handler.context_length * context_step, context_step)])
+
+ return windows
+
+def create_windows_uniform_standard(num_frames: int, handler: IndexListContextHandler, model_options: dict[str]):
+ # unlike looped, uniform_straight does NOT allow windows that loop back to the beginning;
+ # instead, they get shifted to the corresponding end of the frames.
+ # in the case that a window (shifted or not) is identical to the previous one, it gets skipped.
+ windows = []
+ if num_frames <= handler.context_length:
+ windows.append(list(range(num_frames)))
+ return windows
+
+ context_stride = min(handler.context_stride, int(np.ceil(np.log2(num_frames / handler.context_length))) + 1)
+ # first, obtain uniform windows as normal, looping and all
+ for context_step in 1 << np.arange(context_stride):
+ pad = int(round(num_frames * ordered_halving(handler._step)))
+ for j in range(
+ int(ordered_halving(handler._step) * context_step) + pad,
+ num_frames + pad + (-handler.context_overlap),
+ (handler.context_length * context_step - handler.context_overlap),
+ ):
+ windows.append([e % num_frames for e in range(j, j + handler.context_length * context_step, context_step)])
+
+ # now that windows are created, shift any windows that loop, and delete duplicate windows
+ delete_idxs = []
+ win_i = 0
+ while win_i < len(windows):
+ # if window is rolls over itself, need to shift it
+ is_roll, roll_idx = does_window_roll_over(windows[win_i], num_frames)
+ if is_roll:
+ roll_val = windows[win_i][roll_idx] # roll_val might not be 0 for windows of higher strides
+ shift_window_to_end(windows[win_i], num_frames=num_frames)
+ # check if next window (cyclical) is missing roll_val
+ if roll_val not in windows[(win_i+1) % len(windows)]:
+ # need to insert new window here - just insert window starting at roll_val
+ windows.insert(win_i+1, list(range(roll_val, roll_val + handler.context_length)))
+ # delete window if it's not unique
+ for pre_i in range(0, win_i):
+ if windows[win_i] == windows[pre_i]:
+ delete_idxs.append(win_i)
+ break
+ win_i += 1
+
+ # reverse delete_idxs so that they will be deleted in an order that doesn't break idx correlation
+ delete_idxs.reverse()
+ for i in delete_idxs:
+ windows.pop(i)
+
+ return windows
+
+
+def create_windows_static_standard(num_frames: int, handler: IndexListContextHandler, model_options: dict[str]):
+ windows = []
+ if num_frames <= handler.context_length:
+ windows.append(list(range(num_frames)))
+ return windows
+ # always return the same set of windows
+ delta = handler.context_length - handler.context_overlap
+ for start_idx in range(0, num_frames, delta):
+ # if past the end of frames, move start_idx back to allow same context_length
+ ending = start_idx + handler.context_length
+ if ending >= num_frames:
+ final_delta = ending - num_frames
+ final_start_idx = start_idx - final_delta
+ windows.append(list(range(final_start_idx, final_start_idx + handler.context_length)))
+ break
+ windows.append(list(range(start_idx, start_idx + handler.context_length)))
+ return windows
+
+
+def create_windows_batched(num_frames: int, handler: IndexListContextHandler, model_options: dict[str]):
+ windows = []
+ if num_frames <= handler.context_length:
+ windows.append(list(range(num_frames)))
+ return windows
+ # always return the same set of windows;
+ # no overlap, just cut up based on context_length;
+ # last window size will be different if num_frames % opts.context_length != 0
+ for start_idx in range(0, num_frames, handler.context_length):
+ windows.append(list(range(start_idx, min(start_idx + handler.context_length, num_frames))))
+ return windows
+
+
+def create_windows_default(num_frames: int, handler: IndexListContextHandler):
+ return [list(range(num_frames))]
+
+
+CONTEXT_MAPPING = {
+ ContextSchedules.UNIFORM_LOOPED: create_windows_uniform_looped,
+ ContextSchedules.UNIFORM_STANDARD: create_windows_uniform_standard,
+ ContextSchedules.STATIC_STANDARD: create_windows_static_standard,
+ ContextSchedules.BATCHED: create_windows_batched,
+}
+
+
+def get_matching_context_schedule(context_schedule: str) -> ContextSchedule:
+ func = CONTEXT_MAPPING.get(context_schedule, None)
+ if func is None:
+ raise ValueError(f"Unknown context_schedule '{context_schedule}'.")
+ return ContextSchedule(context_schedule, func)
+
+
+def get_context_weights(length: int, full_length: int, idxs: list[int], handler: IndexListContextHandler, sigma: torch.Tensor=None):
+ return handler.fuse_method.func(length, sigma=sigma, handler=handler, full_length=full_length, idxs=idxs)
+
+
+def create_weights_flat(length: int, **kwargs) -> list[float]:
+ # weight is the same for all
+ return [1.0] * length
+
+def create_weights_pyramid(length: int, **kwargs) -> list[float]:
+ # weight is based on the distance away from the edge of the context window;
+ # based on weighted average concept in FreeNoise paper
+ if length % 2 == 0:
+ max_weight = length // 2
+ weight_sequence = list(range(1, max_weight + 1, 1)) + list(range(max_weight, 0, -1))
+ else:
+ max_weight = (length + 1) // 2
+ weight_sequence = list(range(1, max_weight, 1)) + [max_weight] + list(range(max_weight - 1, 0, -1))
+ return weight_sequence
+
+def create_weights_overlap_linear(length: int, full_length: int, idxs: list[int], handler: IndexListContextHandler, **kwargs):
+ # based on code in Kijai's WanVideoWrapper: https://github.com/kijai/ComfyUI-WanVideoWrapper/blob/dbb2523b37e4ccdf45127e5ae33e31362f755c8e/nodes.py#L1302
+ # only expected overlap is given different weights
+ weights_torch = torch.ones((length))
+ # blend left-side on all except first window
+ if min(idxs) > 0:
+ ramp_up = torch.linspace(1e-37, 1, handler.context_overlap)
+ weights_torch[:handler.context_overlap] = ramp_up
+ # blend right-side on all except last window
+ if max(idxs) < full_length-1:
+ ramp_down = torch.linspace(1, 1e-37, handler.context_overlap)
+ weights_torch[-handler.context_overlap:] = ramp_down
+ return weights_torch
+
+class ContextFuseMethods:
+ FLAT = "flat"
+ PYRAMID = "pyramid"
+ RELATIVE = "relative"
+ OVERLAP_LINEAR = "overlap-linear"
+
+ LIST = [PYRAMID, FLAT, OVERLAP_LINEAR]
+ LIST_STATIC = [PYRAMID, RELATIVE, FLAT, OVERLAP_LINEAR]
+
+
+FUSE_MAPPING = {
+ ContextFuseMethods.FLAT: create_weights_flat,
+ ContextFuseMethods.PYRAMID: create_weights_pyramid,
+ ContextFuseMethods.RELATIVE: create_weights_pyramid,
+ ContextFuseMethods.OVERLAP_LINEAR: create_weights_overlap_linear,
+}
+
+def get_matching_fuse_method(fuse_method: str) -> ContextFuseMethod:
+ func = FUSE_MAPPING.get(fuse_method, None)
+ if func is None:
+ raise ValueError(f"Unknown fuse_method '{fuse_method}'.")
+ return ContextFuseMethod(fuse_method, func)
+
+# Returns fraction that has denominator that is a power of 2
+def ordered_halving(val):
+ # get binary value, padded with 0s for 64 bits
+ bin_str = f"{val:064b}"
+ # flip binary value, padding included
+ bin_flip = bin_str[::-1]
+ # convert binary to int
+ as_int = int(bin_flip, 2)
+ # divide by 1 << 64, equivalent to 2**64, or 18446744073709551616,
+ # or b10000000000000000000000000000000000000000000000000000000000000000 (1 with 64 zero's)
+ return as_int / (1 << 64)
+
+
+def get_missing_indexes(windows: list[list[int]], num_frames: int) -> list[int]:
+ all_indexes = list(range(num_frames))
+ for w in windows:
+ for val in w:
+ try:
+ all_indexes.remove(val)
+ except ValueError:
+ pass
+ return all_indexes
+
+
+def does_window_roll_over(window: list[int], num_frames: int) -> tuple[bool, int]:
+ prev_val = -1
+ for i, val in enumerate(window):
+ val = val % num_frames
+ if val < prev_val:
+ return True, i
+ prev_val = val
+ return False, -1
+
+
+def shift_window_to_start(window: list[int], num_frames: int):
+ start_val = window[0]
+ for i in range(len(window)):
+ # 1) subtract each element by start_val to move vals relative to the start of all frames
+ # 2) add num_frames and take modulus to get adjusted vals
+ window[i] = ((window[i] - start_val) + num_frames) % num_frames
+
+
+def shift_window_to_end(window: list[int], num_frames: int):
+ # 1) shift window to start
+ shift_window_to_start(window, num_frames)
+ end_val = window[-1]
+ end_delta = num_frames - end_val - 1
+ for i in range(len(window)):
+ # 2) add end_delta to each val to slide windows to end
+ window[i] = window[i] + end_delta
diff --git a/comfy/controlnet.py b/comfy/controlnet.py
index 9a47b86f2..0b5e30f52 100644
--- a/comfy/controlnet.py
+++ b/comfy/controlnet.py
@@ -28,6 +28,7 @@ import comfy.model_detection
import comfy.model_patcher
import comfy.ops
import comfy.latent_formats
+import comfy.model_base
import comfy.cldm.cldm
import comfy.t2i_adapter.adapter
@@ -35,6 +36,7 @@ import comfy.ldm.cascade.controlnet
import comfy.cldm.mmdit
import comfy.ldm.hydit.controlnet
import comfy.ldm.flux.controlnet
+import comfy.ldm.qwen_image.controlnet
import comfy.cldm.dit_embedder
from typing import TYPE_CHECKING
if TYPE_CHECKING:
@@ -43,7 +45,6 @@ if TYPE_CHECKING:
def broadcast_image_to(tensor, target_batch_size, batched_number):
current_batch_size = tensor.shape[0]
- #print(current_batch_size, target_batch_size)
if current_batch_size == 1:
return tensor
@@ -236,11 +237,11 @@ class ControlNet(ControlBase):
self.cond_hint = None
compression_ratio = self.compression_ratio
if self.vae is not None:
- compression_ratio *= self.vae.downscale_ratio
+ compression_ratio *= self.vae.spacial_compression_encode()
else:
if self.latent_format is not None:
raise ValueError("This Controlnet needs a VAE but none was provided, please use a ControlNetApply node with a VAE input and connect it.")
- self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, x_noisy.shape[3] * compression_ratio, x_noisy.shape[2] * compression_ratio, self.upscale_algorithm, "center")
+ self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, x_noisy.shape[-1] * compression_ratio, x_noisy.shape[-2] * compression_ratio, self.upscale_algorithm, "center")
self.cond_hint = self.preprocess_image(self.cond_hint)
if self.vae is not None:
loaded_models = comfy.model_management.loaded_models(only_currently_used=True)
@@ -252,7 +253,10 @@ class ControlNet(ControlBase):
to_concat = []
for c in self.extra_concat_orig:
c = c.to(self.cond_hint.device)
- c = comfy.utils.common_upscale(c, self.cond_hint.shape[3], self.cond_hint.shape[2], self.upscale_algorithm, "center")
+ c = comfy.utils.common_upscale(c, self.cond_hint.shape[-1], self.cond_hint.shape[-2], self.upscale_algorithm, "center")
+ if c.ndim < self.cond_hint.ndim:
+ c = c.unsqueeze(2)
+ c = comfy.utils.repeat_to_batch_size(c, self.cond_hint.shape[2], dim=2)
to_concat.append(comfy.utils.repeat_to_batch_size(c, self.cond_hint.shape[0]))
self.cond_hint = torch.cat([self.cond_hint] + to_concat, dim=1)
@@ -265,12 +269,12 @@ class ControlNet(ControlBase):
for c in self.extra_conds:
temp = cond.get(c, None)
if temp is not None:
- extra[c] = temp.to(dtype)
+ extra[c] = comfy.model_base.convert_tensor(temp, dtype, x_noisy.device)
timestep = self.model_sampling_current.timestep(t)
x_noisy = self.model_sampling_current.calculate_input(t, x_noisy)
- control = self.control_model(x=x_noisy.to(dtype), hint=self.cond_hint, timesteps=timestep.to(dtype), context=context.to(dtype), **extra)
+ control = self.control_model(x=x_noisy.to(dtype), hint=self.cond_hint, timesteps=timestep.to(dtype), context=comfy.model_management.cast_to_device(context, x_noisy.device, dtype), **extra)
return self.control_merge(control, control_prev, output_dtype=None)
def copy(self):
@@ -306,11 +310,13 @@ class ControlLoraOps:
self.bias = None
def forward(self, input):
- weight, bias = comfy.ops.cast_bias_weight(self, input)
+ weight, bias, offload_stream = comfy.ops.cast_bias_weight(self, input, offloadable=True)
if self.up is not None:
- return torch.nn.functional.linear(input, weight + (torch.mm(self.up.flatten(start_dim=1), self.down.flatten(start_dim=1))).reshape(self.weight.shape).type(input.dtype), bias)
+ x = torch.nn.functional.linear(input, weight + (torch.mm(self.up.flatten(start_dim=1), self.down.flatten(start_dim=1))).reshape(self.weight.shape).type(input.dtype), bias)
else:
- return torch.nn.functional.linear(input, weight, bias)
+ x = torch.nn.functional.linear(input, weight, bias)
+ comfy.ops.uncast_bias_weight(self, weight, bias, offload_stream)
+ return x
class Conv2d(torch.nn.Module, comfy.ops.CastWeightBiasOp):
def __init__(
@@ -346,12 +352,13 @@ class ControlLoraOps:
def forward(self, input):
- weight, bias = comfy.ops.cast_bias_weight(self, input)
+ weight, bias, offload_stream = comfy.ops.cast_bias_weight(self, input, offloadable=True)
if self.up is not None:
- return torch.nn.functional.conv2d(input, weight + (torch.mm(self.up.flatten(start_dim=1), self.down.flatten(start_dim=1))).reshape(self.weight.shape).type(input.dtype), bias, self.stride, self.padding, self.dilation, self.groups)
+ x = torch.nn.functional.conv2d(input, weight + (torch.mm(self.up.flatten(start_dim=1), self.down.flatten(start_dim=1))).reshape(self.weight.shape).type(input.dtype), bias, self.stride, self.padding, self.dilation, self.groups)
else:
- return torch.nn.functional.conv2d(input, weight, bias, self.stride, self.padding, self.dilation, self.groups)
-
+ x = torch.nn.functional.conv2d(input, weight, bias, self.stride, self.padding, self.dilation, self.groups)
+ comfy.ops.uncast_bias_weight(self, weight, bias, offload_stream)
+ return x
class ControlLora(ControlNet):
def __init__(self, control_weights, global_average_pooling=False, model_options={}): #TODO? model_options
@@ -582,6 +589,22 @@ def load_controlnet_flux_instantx(sd, model_options={}):
control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, concat_mask=concat_mask, load_device=load_device, manual_cast_dtype=manual_cast_dtype, extra_conds=extra_conds)
return control
+def load_controlnet_qwen_instantx(sd, model_options={}):
+ model_config, operations, load_device, unet_dtype, manual_cast_dtype, offload_device = controlnet_config(sd, model_options=model_options)
+ control_latent_channels = sd.get("controlnet_x_embedder.weight").shape[1]
+
+ extra_condition_channels = 0
+ concat_mask = False
+ if control_latent_channels == 68: #inpaint controlnet
+ extra_condition_channels = control_latent_channels - 64
+ concat_mask = True
+ control_model = comfy.ldm.qwen_image.controlnet.QwenImageControlNetModel(extra_condition_channels=extra_condition_channels, operations=operations, device=offload_device, dtype=unet_dtype, **model_config.unet_config)
+ control_model = controlnet_load_state_dict(control_model, sd)
+ latent_format = comfy.latent_formats.Wan21()
+ extra_conds = []
+ control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, concat_mask=concat_mask, load_device=load_device, manual_cast_dtype=manual_cast_dtype, extra_conds=extra_conds)
+ return control
+
def convert_mistoline(sd):
return comfy.utils.state_dict_prefix_replace(sd, {"single_controlnet_blocks.": "controlnet_single_blocks."})
@@ -655,8 +678,11 @@ def load_controlnet_state_dict(state_dict, model=None, model_options={}):
return load_controlnet_sd35(controlnet_data, model_options=model_options) #Stability sd3.5 format
else:
return load_controlnet_mmdit(controlnet_data, model_options=model_options) #SD3 diffusers controlnet
+ elif "transformer_blocks.0.img_mlp.net.0.proj.weight" in controlnet_data:
+ return load_controlnet_qwen_instantx(controlnet_data, model_options=model_options)
elif "controlnet_x_embedder.weight" in controlnet_data:
return load_controlnet_flux_instantx(controlnet_data, model_options=model_options)
+
elif "controlnet_blocks.0.linear.weight" in controlnet_data: #mistoline flux
return load_controlnet_flux_xlabs_mistoline(convert_mistoline(controlnet_data), mistoline=True, model_options=model_options)
diff --git a/comfy/gligen.py b/comfy/gligen.py
index 161d8a5e5..1d7b6c2f4 100644
--- a/comfy/gligen.py
+++ b/comfy/gligen.py
@@ -1,55 +1,10 @@
import math
import torch
from torch import nn
-from .ldm.modules.attention import CrossAttention
-from inspect import isfunction
+from .ldm.modules.attention import CrossAttention, FeedForward
import comfy.ops
ops = comfy.ops.manual_cast
-def exists(val):
- return val is not None
-
-
-def uniq(arr):
- return{el: True for el in arr}.keys()
-
-
-def default(val, d):
- if exists(val):
- return val
- return d() if isfunction(d) else d
-
-
-# feedforward
-class GEGLU(nn.Module):
- def __init__(self, dim_in, dim_out):
- super().__init__()
- self.proj = ops.Linear(dim_in, dim_out * 2)
-
- def forward(self, x):
- x, gate = self.proj(x).chunk(2, dim=-1)
- return x * torch.nn.functional.gelu(gate)
-
-
-class FeedForward(nn.Module):
- def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.):
- super().__init__()
- inner_dim = int(dim * mult)
- dim_out = default(dim_out, dim)
- project_in = nn.Sequential(
- ops.Linear(dim, inner_dim),
- nn.GELU()
- ) if not glu else GEGLU(dim, inner_dim)
-
- self.net = nn.Sequential(
- project_in,
- nn.Dropout(dropout),
- ops.Linear(inner_dim, dim_out)
- )
-
- def forward(self, x):
- return self.net(x)
-
class GatedCrossAttentionDense(nn.Module):
def __init__(self, query_dim, context_dim, n_heads, d_head):
diff --git a/comfy/image_encoders/dino2.py b/comfy/image_encoders/dino2.py
index 976f98c65..9b6dace9d 100644
--- a/comfy/image_encoders/dino2.py
+++ b/comfy/image_encoders/dino2.py
@@ -31,6 +31,20 @@ class LayerScale(torch.nn.Module):
def forward(self, x):
return x * comfy.model_management.cast_to_device(self.lambda1, x.device, x.dtype)
+class Dinov2MLP(torch.nn.Module):
+ def __init__(self, hidden_size: int, dtype, device, operations):
+ super().__init__()
+
+ mlp_ratio = 4
+ hidden_features = int(hidden_size * mlp_ratio)
+ self.fc1 = operations.Linear(hidden_size, hidden_features, bias = True, device=device, dtype=dtype)
+ self.fc2 = operations.Linear(hidden_features, hidden_size, bias = True, device=device, dtype=dtype)
+
+ def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
+ hidden_state = self.fc1(hidden_state)
+ hidden_state = torch.nn.functional.gelu(hidden_state)
+ hidden_state = self.fc2(hidden_state)
+ return hidden_state
class SwiGLUFFN(torch.nn.Module):
def __init__(self, dim, dtype, device, operations):
@@ -50,12 +64,15 @@ class SwiGLUFFN(torch.nn.Module):
class Dino2Block(torch.nn.Module):
- def __init__(self, dim, num_heads, layer_norm_eps, dtype, device, operations):
+ def __init__(self, dim, num_heads, layer_norm_eps, dtype, device, operations, use_swiglu_ffn):
super().__init__()
self.attention = Dino2AttentionBlock(dim, num_heads, layer_norm_eps, dtype, device, operations)
self.layer_scale1 = LayerScale(dim, dtype, device, operations)
self.layer_scale2 = LayerScale(dim, dtype, device, operations)
- self.mlp = SwiGLUFFN(dim, dtype, device, operations)
+ if use_swiglu_ffn:
+ self.mlp = SwiGLUFFN(dim, dtype, device, operations)
+ else:
+ self.mlp = Dinov2MLP(dim, dtype, device, operations)
self.norm1 = operations.LayerNorm(dim, eps=layer_norm_eps, dtype=dtype, device=device)
self.norm2 = operations.LayerNorm(dim, eps=layer_norm_eps, dtype=dtype, device=device)
@@ -66,9 +83,10 @@ class Dino2Block(torch.nn.Module):
class Dino2Encoder(torch.nn.Module):
- def __init__(self, dim, num_heads, layer_norm_eps, num_layers, dtype, device, operations):
+ def __init__(self, dim, num_heads, layer_norm_eps, num_layers, dtype, device, operations, use_swiglu_ffn):
super().__init__()
- self.layer = torch.nn.ModuleList([Dino2Block(dim, num_heads, layer_norm_eps, dtype, device, operations) for _ in range(num_layers)])
+ self.layer = torch.nn.ModuleList([Dino2Block(dim, num_heads, layer_norm_eps, dtype, device, operations, use_swiglu_ffn = use_swiglu_ffn)
+ for _ in range(num_layers)])
def forward(self, x, intermediate_output=None):
optimized_attention = optimized_attention_for_device(x.device, False, small_input=True)
@@ -78,8 +96,8 @@ class Dino2Encoder(torch.nn.Module):
intermediate_output = len(self.layer) + intermediate_output
intermediate = None
- for i, l in enumerate(self.layer):
- x = l(x, optimized_attention)
+ for i, layer in enumerate(self.layer):
+ x = layer(x, optimized_attention)
if i == intermediate_output:
intermediate = x.clone()
return x, intermediate
@@ -128,9 +146,10 @@ class Dinov2Model(torch.nn.Module):
dim = config_dict["hidden_size"]
heads = config_dict["num_attention_heads"]
layer_norm_eps = config_dict["layer_norm_eps"]
+ use_swiglu_ffn = config_dict["use_swiglu_ffn"]
self.embeddings = Dino2Embeddings(dim, dtype, device, operations)
- self.encoder = Dino2Encoder(dim, heads, layer_norm_eps, num_layers, dtype, device, operations)
+ self.encoder = Dino2Encoder(dim, heads, layer_norm_eps, num_layers, dtype, device, operations, use_swiglu_ffn = use_swiglu_ffn)
self.layernorm = operations.LayerNorm(dim, eps=layer_norm_eps, dtype=dtype, device=device)
def forward(self, pixel_values, attention_mask=None, intermediate_output=None):
diff --git a/comfy/image_encoders/dino2_large.json b/comfy/image_encoders/dino2_large.json
new file mode 100644
index 000000000..43fbb58ff
--- /dev/null
+++ b/comfy/image_encoders/dino2_large.json
@@ -0,0 +1,22 @@
+{
+ "hidden_size": 1024,
+ "use_mask_token": true,
+ "patch_size": 14,
+ "image_size": 518,
+ "num_channels": 3,
+ "num_attention_heads": 16,
+ "initializer_range": 0.02,
+ "attention_probs_dropout_prob": 0.0,
+ "hidden_dropout_prob": 0.0,
+ "hidden_act": "gelu",
+ "mlp_ratio": 4,
+ "model_type": "dinov2",
+ "num_hidden_layers": 24,
+ "layer_norm_eps": 1e-6,
+ "qkv_bias": true,
+ "use_swiglu_ffn": false,
+ "layerscale_value": 1.0,
+ "drop_path_rate": 0.0,
+ "image_mean": [0.485, 0.456, 0.406],
+ "image_std": [0.229, 0.224, 0.225]
+}
diff --git a/comfy/k_diffusion/sa_solver.py b/comfy/k_diffusion/sa_solver.py
new file mode 100644
index 000000000..0c6821b60
--- /dev/null
+++ b/comfy/k_diffusion/sa_solver.py
@@ -0,0 +1,121 @@
+# SA-Solver: Stochastic Adams Solver (NeurIPS 2023, arXiv:2309.05019)
+# Conference: https://proceedings.neurips.cc/paper_files/paper/2023/file/f4a6806490d31216a3ba667eb240c897-Paper-Conference.pdf
+# Codebase ref: https://github.com/scxue/SA-Solver
+
+import math
+from typing import Union, Callable
+import torch
+
+
+def compute_exponential_coeffs(s: torch.Tensor, t: torch.Tensor, solver_order: int, tau_t: float) -> torch.Tensor:
+ """Compute (1 + tau^2) * integral of exp((1 + tau^2) * x) * x^p dx from s to t with exp((1 + tau^2) * t) factored out, using integration by parts.
+
+ Integral of exp((1 + tau^2) * x) * x^p dx
+ = product_terms[p] - (p / (1 + tau^2)) * integral of exp((1 + tau^2) * x) * x^(p-1) dx,
+ with base case p=0 where integral equals product_terms[0].
+
+ where
+ product_terms[p] = x^p * exp((1 + tau^2) * x) / (1 + tau^2).
+
+ Construct a recursive coefficient matrix following the above recursive relation to compute all integral terms up to p = (solver_order - 1).
+ Return coefficients used by the SA-Solver in data prediction mode.
+
+ Args:
+ s: Start time s.
+ t: End time t.
+ solver_order: Current order of the solver.
+ tau_t: Stochastic strength parameter in the SDE.
+
+ Returns:
+ Exponential coefficients used in data prediction, with exp((1 + tau^2) * t) factored out, ordered from p=0 to p=solver_order−1, shape (solver_order,).
+ """
+ tau_mul = 1 + tau_t ** 2
+ h = t - s
+ p = torch.arange(solver_order, dtype=s.dtype, device=s.device)
+
+ # product_terms after factoring out exp((1 + tau^2) * t)
+ # Includes (1 + tau^2) factor from outside the integral
+ product_terms_factored = (t ** p - s ** p * (-tau_mul * h).exp())
+
+ # Lower triangular recursive coefficient matrix
+ # Accumulates recursive coefficients based on p / (1 + tau^2)
+ recursive_depth_mat = p.unsqueeze(1) - p.unsqueeze(0)
+ log_factorial = (p + 1).lgamma()
+ recursive_coeff_mat = log_factorial.unsqueeze(1) - log_factorial.unsqueeze(0)
+ if tau_t > 0:
+ recursive_coeff_mat = recursive_coeff_mat - (recursive_depth_mat * math.log(tau_mul))
+ signs = torch.where(recursive_depth_mat % 2 == 0, 1.0, -1.0)
+ recursive_coeff_mat = (recursive_coeff_mat.exp() * signs).tril()
+
+ return recursive_coeff_mat @ product_terms_factored
+
+
+def compute_simple_stochastic_adams_b_coeffs(sigma_next: torch.Tensor, curr_lambdas: torch.Tensor, lambda_s: torch.Tensor, lambda_t: torch.Tensor, tau_t: float, is_corrector_step: bool = False) -> torch.Tensor:
+ """Compute simple order-2 b coefficients from SA-Solver paper (Appendix D. Implementation Details)."""
+ tau_mul = 1 + tau_t ** 2
+ h = lambda_t - lambda_s
+ alpha_t = sigma_next * lambda_t.exp()
+ if is_corrector_step:
+ # Simplified 1-step (order-2) corrector
+ b_1 = alpha_t * (0.5 * tau_mul * h)
+ b_2 = alpha_t * (-h * tau_mul).expm1().neg() - b_1
+ else:
+ # Simplified 2-step predictor
+ b_2 = alpha_t * (0.5 * tau_mul * h ** 2) / (curr_lambdas[-2] - lambda_s)
+ b_1 = alpha_t * (-h * tau_mul).expm1().neg() - b_2
+ return torch.stack([b_2, b_1])
+
+
+def compute_stochastic_adams_b_coeffs(sigma_next: torch.Tensor, curr_lambdas: torch.Tensor, lambda_s: torch.Tensor, lambda_t: torch.Tensor, tau_t: float, simple_order_2: bool = False, is_corrector_step: bool = False) -> torch.Tensor:
+ """Compute b_i coefficients for the SA-Solver (see eqs. 15 and 18).
+
+ The solver order corresponds to the number of input lambdas (half-logSNR points).
+
+ Args:
+ sigma_next: Sigma at end time t.
+ curr_lambdas: Lambda time points used to construct the Lagrange basis, shape (N,).
+ lambda_s: Lambda at start time s.
+ lambda_t: Lambda at end time t.
+ tau_t: Stochastic strength parameter in the SDE.
+ simple_order_2: Whether to enable the simple order-2 scheme.
+ is_corrector_step: Flag for corrector step in simple order-2 mode.
+
+ Returns:
+ b_i coefficients for the SA-Solver, shape (N,), where N is the solver order.
+ """
+ num_timesteps = curr_lambdas.shape[0]
+
+ if simple_order_2 and num_timesteps == 2:
+ return compute_simple_stochastic_adams_b_coeffs(sigma_next, curr_lambdas, lambda_s, lambda_t, tau_t, is_corrector_step)
+
+ # Compute coefficients by solving a linear system from Lagrange basis interpolation
+ exp_integral_coeffs = compute_exponential_coeffs(lambda_s, lambda_t, num_timesteps, tau_t)
+ vandermonde_matrix_T = torch.vander(curr_lambdas, num_timesteps, increasing=True).T
+ lagrange_integrals = torch.linalg.solve(vandermonde_matrix_T, exp_integral_coeffs)
+
+ # (sigma_t * exp(-tau^2 * lambda_t)) * exp((1 + tau^2) * lambda_t)
+ # = sigma_t * exp(lambda_t) = alpha_t
+ # exp((1 + tau^2) * lambda_t) is extracted from the integral
+ alpha_t = sigma_next * lambda_t.exp()
+ return alpha_t * lagrange_integrals
+
+
+def get_tau_interval_func(start_sigma: float, end_sigma: float, eta: float = 1.0) -> Callable[[Union[torch.Tensor, float]], float]:
+ """Return a function that controls the stochasticity of SA-Solver.
+
+ When eta = 0, SA-Solver runs as ODE. The official approach uses
+ time t to determine the SDE interval, while here we use sigma instead.
+
+ See:
+ https://github.com/scxue/SA-Solver/blob/main/README.md
+ """
+
+ def tau_func(sigma: Union[torch.Tensor, float]) -> float:
+ if eta <= 0:
+ return 0.0 # ODE
+
+ if isinstance(sigma, torch.Tensor):
+ sigma = sigma.item()
+ return eta if start_sigma >= sigma >= end_sigma else 0.0
+
+ return tau_func
diff --git a/comfy/k_diffusion/sampling.py b/comfy/k_diffusion/sampling.py
index 34218337a..0e2cda291 100644
--- a/comfy/k_diffusion/sampling.py
+++ b/comfy/k_diffusion/sampling.py
@@ -9,6 +9,7 @@ from tqdm.auto import trange, tqdm
from . import utils
from . import deis
+from . import sa_solver
import comfy.model_patcher
import comfy.model_sampling
@@ -85,24 +86,24 @@ class BatchedBrownianTree:
"""A wrapper around torchsde.BrownianTree that enables batches of entropy."""
def __init__(self, x, t0, t1, seed=None, **kwargs):
- self.cpu_tree = True
- if "cpu" in kwargs:
- self.cpu_tree = kwargs.pop("cpu")
+ self.cpu_tree = kwargs.pop("cpu", True)
t0, t1, self.sign = self.sort(t0, t1)
- w0 = kwargs.get('w0', torch.zeros_like(x))
+ w0 = kwargs.pop('w0', None)
+ if w0 is None:
+ w0 = torch.zeros_like(x)
+ self.batched = False
if seed is None:
- seed = torch.randint(0, 2 ** 63 - 1, []).item()
- self.batched = True
- try:
- assert len(seed) == x.shape[0]
+ seed = (torch.randint(0, 2 ** 63 - 1, ()).item(),)
+ elif isinstance(seed, (tuple, list)):
+ if len(seed) != x.shape[0]:
+ raise ValueError("Passing a list or tuple of seeds to BatchedBrownianTree requires a length matching the batch size.")
+ self.batched = True
w0 = w0[0]
- except TypeError:
- seed = [seed]
- self.batched = False
- if self.cpu_tree:
- self.trees = [torchsde.BrownianTree(t0.cpu(), w0.cpu(), t1.cpu(), entropy=s, **kwargs) for s in seed]
else:
- self.trees = [torchsde.BrownianTree(t0, w0, t1, entropy=s, **kwargs) for s in seed]
+ seed = (seed,)
+ if self.cpu_tree:
+ t0, w0, t1 = t0.detach().cpu(), w0.detach().cpu(), t1.detach().cpu()
+ self.trees = tuple(torchsde.BrownianTree(t0, w0, t1, entropy=s, **kwargs) for s in seed)
@staticmethod
def sort(a, b):
@@ -110,11 +111,10 @@ class BatchedBrownianTree:
def __call__(self, t0, t1):
t0, t1, sign = self.sort(t0, t1)
+ device, dtype = t0.device, t0.dtype
if self.cpu_tree:
- w = torch.stack([tree(t0.cpu().float(), t1.cpu().float()).to(t0.dtype).to(t0.device) for tree in self.trees]) * (self.sign * sign)
- else:
- w = torch.stack([tree(t0, t1) for tree in self.trees]) * (self.sign * sign)
-
+ t0, t1 = t0.detach().cpu().float(), t1.detach().cpu().float()
+ w = torch.stack([tree(t0, t1) for tree in self.trees]).to(device=device, dtype=dtype) * (self.sign * sign)
return w if self.batched else w[0]
@@ -170,6 +170,16 @@ def offset_first_sigma_for_snr(sigmas, model_sampling, percent_offset=1e-4):
return sigmas
+def ei_h_phi_1(h: torch.Tensor) -> torch.Tensor:
+ """Compute the result of h*phi_1(h) in exponential integrator methods."""
+ return torch.expm1(h)
+
+
+def ei_h_phi_2(h: torch.Tensor) -> torch.Tensor:
+ """Compute the result of h*phi_2(h) in exponential integrator methods."""
+ return (torch.expm1(h) - h) / h
+
+
@torch.no_grad()
def sample_euler(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
"""Implements Algorithm 2 (Euler steps) from Karras et al. (2022)."""
@@ -852,6 +862,11 @@ def sample_dpmpp_2m_sde(model, x, sigmas, extra_args=None, callback=None, disabl
return x
+@torch.no_grad()
+def sample_dpmpp_2m_sde_heun(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, solver_type='heun'):
+ return sample_dpmpp_2m_sde(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler, solver_type=solver_type)
+
+
@torch.no_grad()
def sample_dpmpp_3m_sde(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
"""DPM-Solver++(3M) SDE."""
@@ -924,6 +939,16 @@ def sample_dpmpp_3m_sde_gpu(model, x, sigmas, extra_args=None, callback=None, di
return sample_dpmpp_3m_sde(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler)
+@torch.no_grad()
+def sample_dpmpp_2m_sde_heun_gpu(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, solver_type='heun'):
+ if len(sigmas) <= 1:
+ return x
+ extra_args = {} if extra_args is None else extra_args
+ sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
+ noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=extra_args.get("seed", None), cpu=False) if noise_sampler is None else noise_sampler
+ return sample_dpmpp_2m_sde_heun(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler, solver_type=solver_type)
+
+
@torch.no_grad()
def sample_dpmpp_2m_sde_gpu(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, solver_type='midpoint'):
if len(sigmas) <= 1:
@@ -1209,39 +1234,21 @@ def sample_deis(model, x, sigmas, extra_args=None, callback=None, disable=None,
return x_next
-@torch.no_grad()
-def sample_euler_cfg_pp(model, x, sigmas, extra_args=None, callback=None, disable=None):
- extra_args = {} if extra_args is None else extra_args
-
- temp = [0]
- def post_cfg_function(args):
- temp[0] = args["uncond_denoised"]
- return args["denoised"]
-
- model_options = extra_args.get("model_options", {}).copy()
- extra_args["model_options"] = comfy.model_patcher.set_model_options_post_cfg_function(model_options, post_cfg_function, disable_cfg1_optimization=True)
-
- s_in = x.new_ones([x.shape[0]])
- for i in trange(len(sigmas) - 1, disable=disable):
- sigma_hat = sigmas[i]
- denoised = model(x, sigma_hat * s_in, **extra_args)
- d = to_d(x, sigma_hat, temp[0])
- if callback is not None:
- callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
- # Euler method
- x = denoised + d * sigmas[i + 1]
- return x
-
@torch.no_grad()
def sample_euler_ancestral_cfg_pp(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
- """Ancestral sampling with Euler method steps."""
+ """Ancestral sampling with Euler method steps (CFG++)."""
extra_args = {} if extra_args is None else extra_args
seed = extra_args.get("seed", None)
noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
- temp = [0]
+ model_sampling = model.inner_model.model_patcher.get_model_object("model_sampling")
+ lambda_fn = partial(sigma_to_half_log_snr, model_sampling=model_sampling)
+
+ uncond_denoised = None
+
def post_cfg_function(args):
- temp[0] = args["uncond_denoised"]
+ nonlocal uncond_denoised
+ uncond_denoised = args["uncond_denoised"]
return args["denoised"]
model_options = extra_args.get("model_options", {}).copy()
@@ -1250,15 +1257,33 @@ def sample_euler_ancestral_cfg_pp(model, x, sigmas, extra_args=None, callback=No
s_in = x.new_ones([x.shape[0]])
for i in trange(len(sigmas) - 1, disable=disable):
denoised = model(x, sigmas[i] * s_in, **extra_args)
- sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
if callback is not None:
callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
- d = to_d(x, sigmas[i], temp[0])
- # Euler method
- x = denoised + d * sigma_down
- if sigmas[i + 1] > 0:
- x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
+ if sigmas[i + 1] == 0:
+ # Denoising step
+ x = denoised
+ else:
+ alpha_s = sigmas[i] * lambda_fn(sigmas[i]).exp()
+ alpha_t = sigmas[i + 1] * lambda_fn(sigmas[i + 1]).exp()
+ d = to_d(x, sigmas[i], alpha_s * uncond_denoised) # to noise
+
+ # DDIM stochastic sampling
+ sigma_down, sigma_up = get_ancestral_step(sigmas[i] / alpha_s, sigmas[i + 1] / alpha_t, eta=eta)
+ sigma_down = alpha_t * sigma_down
+
+ # Euler method
+ x = alpha_t * denoised + sigma_down * d
+ if eta > 0 and s_noise > 0:
+ x = x + alpha_t * noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
return x
+
+
+@torch.no_grad()
+def sample_euler_cfg_pp(model, x, sigmas, extra_args=None, callback=None, disable=None):
+ """Euler method steps (CFG++)."""
+ return sample_euler_ancestral_cfg_pp(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=0.0, s_noise=0.0, noise_sampler=None)
+
+
@torch.no_grad()
def sample_dpmpp_2s_ancestral_cfg_pp(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
"""Ancestral sampling with DPM-Solver++(2S) second-order steps."""
@@ -1534,13 +1559,12 @@ def sample_er_sde(model, x, sigmas, extra_args=None, callback=None, disable=None
@torch.no_grad()
def sample_seeds_2(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, r=0.5):
"""SEEDS-2 - Stochastic Explicit Exponential Derivative-free Solvers (VP Data Prediction) stage 2.
- arXiv: https://arxiv.org/abs/2305.14267
+ arXiv: https://arxiv.org/abs/2305.14267 (NeurIPS 2023)
"""
extra_args = {} if extra_args is None else extra_args
seed = extra_args.get("seed", None)
noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
s_in = x.new_ones([x.shape[0]])
-
inject_noise = eta > 0 and s_noise > 0
model_sampling = model.inner_model.model_patcher.get_model_object('model_sampling')
@@ -1548,55 +1572,53 @@ def sample_seeds_2(model, x, sigmas, extra_args=None, callback=None, disable=Non
lambda_fn = partial(sigma_to_half_log_snr, model_sampling=model_sampling)
sigmas = offset_first_sigma_for_snr(sigmas, model_sampling)
+ fac = 1 / (2 * r)
+
for i in trange(len(sigmas) - 1, disable=disable):
denoised = model(x, sigmas[i] * s_in, **extra_args)
if callback is not None:
callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+
if sigmas[i + 1] == 0:
x = denoised
- else:
- lambda_s, lambda_t = lambda_fn(sigmas[i]), lambda_fn(sigmas[i + 1])
- h = lambda_t - lambda_s
- h_eta = h * (eta + 1)
- lambda_s_1 = lambda_s + r * h
- fac = 1 / (2 * r)
- sigma_s_1 = sigma_fn(lambda_s_1)
+ continue
- # alpha_t = sigma_t * exp(log(alpha_t / sigma_t)) = sigma_t * exp(lambda_t)
- alpha_s_1 = sigma_s_1 * lambda_s_1.exp()
- alpha_t = sigmas[i + 1] * lambda_t.exp()
+ lambda_s, lambda_t = lambda_fn(sigmas[i]), lambda_fn(sigmas[i + 1])
+ h = lambda_t - lambda_s
+ h_eta = h * (eta + 1)
+ lambda_s_1 = torch.lerp(lambda_s, lambda_t, r)
+ sigma_s_1 = sigma_fn(lambda_s_1)
- coeff_1, coeff_2 = (-r * h_eta).expm1(), (-h_eta).expm1()
- if inject_noise:
- # 0 < r < 1
- noise_coeff_1 = (-2 * r * h * eta).expm1().neg().sqrt()
- noise_coeff_2 = (-r * h * eta).exp() * (-2 * (1 - r) * h * eta).expm1().neg().sqrt()
- noise_1, noise_2 = noise_sampler(sigmas[i], sigma_s_1), noise_sampler(sigma_s_1, sigmas[i + 1])
+ alpha_s_1 = sigma_s_1 * lambda_s_1.exp()
+ alpha_t = sigmas[i + 1] * lambda_t.exp()
- # Step 1
- x_2 = sigma_s_1 / sigmas[i] * (-r * h * eta).exp() * x - alpha_s_1 * coeff_1 * denoised
- if inject_noise:
- x_2 = x_2 + sigma_s_1 * (noise_coeff_1 * noise_1) * s_noise
- denoised_2 = model(x_2, sigma_s_1 * s_in, **extra_args)
+ # Step 1
+ x_2 = sigma_s_1 / sigmas[i] * (-r * h * eta).exp() * x - alpha_s_1 * ei_h_phi_1(-r * h_eta) * denoised
+ if inject_noise:
+ sde_noise = (-2 * r * h * eta).expm1().neg().sqrt() * noise_sampler(sigmas[i], sigma_s_1)
+ x_2 = x_2 + sde_noise * sigma_s_1 * s_noise
+ denoised_2 = model(x_2, sigma_s_1 * s_in, **extra_args)
- # Step 2
- denoised_d = (1 - fac) * denoised + fac * denoised_2
- x = sigmas[i + 1] / sigmas[i] * (-h * eta).exp() * x - alpha_t * coeff_2 * denoised_d
- if inject_noise:
- x = x + sigmas[i + 1] * (noise_coeff_2 * noise_1 + noise_coeff_1 * noise_2) * s_noise
+ # Step 2
+ denoised_d = torch.lerp(denoised, denoised_2, fac)
+ x = sigmas[i + 1] / sigmas[i] * (-h * eta).exp() * x - alpha_t * ei_h_phi_1(-h_eta) * denoised_d
+ if inject_noise:
+ segment_factor = (r - 1) * h * eta
+ sde_noise = sde_noise * segment_factor.exp()
+ sde_noise = sde_noise + segment_factor.mul(2).expm1().neg().sqrt() * noise_sampler(sigma_s_1, sigmas[i + 1])
+ x = x + sde_noise * sigmas[i + 1] * s_noise
return x
@torch.no_grad()
def sample_seeds_3(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, r_1=1./3, r_2=2./3):
"""SEEDS-3 - Stochastic Explicit Exponential Derivative-free Solvers (VP Data Prediction) stage 3.
- arXiv: https://arxiv.org/abs/2305.14267
+ arXiv: https://arxiv.org/abs/2305.14267 (NeurIPS 2023)
"""
extra_args = {} if extra_args is None else extra_args
seed = extra_args.get("seed", None)
noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
s_in = x.new_ones([x.shape[0]])
-
inject_noise = eta > 0 and s_noise > 0
model_sampling = model.inner_model.model_patcher.get_model_object('model_sampling')
@@ -1608,43 +1630,157 @@ def sample_seeds_3(model, x, sigmas, extra_args=None, callback=None, disable=Non
denoised = model(x, sigmas[i] * s_in, **extra_args)
if callback is not None:
callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+
if sigmas[i + 1] == 0:
x = denoised
- else:
- lambda_s, lambda_t = lambda_fn(sigmas[i]), lambda_fn(sigmas[i + 1])
- h = lambda_t - lambda_s
- h_eta = h * (eta + 1)
- lambda_s_1 = lambda_s + r_1 * h
- lambda_s_2 = lambda_s + r_2 * h
- sigma_s_1, sigma_s_2 = sigma_fn(lambda_s_1), sigma_fn(lambda_s_2)
+ continue
- # alpha_t = sigma_t * exp(log(alpha_t / sigma_t)) = sigma_t * exp(lambda_t)
- alpha_s_1 = sigma_s_1 * lambda_s_1.exp()
- alpha_s_2 = sigma_s_2 * lambda_s_2.exp()
- alpha_t = sigmas[i + 1] * lambda_t.exp()
+ lambda_s, lambda_t = lambda_fn(sigmas[i]), lambda_fn(sigmas[i + 1])
+ h = lambda_t - lambda_s
+ h_eta = h * (eta + 1)
+ lambda_s_1 = torch.lerp(lambda_s, lambda_t, r_1)
+ lambda_s_2 = torch.lerp(lambda_s, lambda_t, r_2)
+ sigma_s_1, sigma_s_2 = sigma_fn(lambda_s_1), sigma_fn(lambda_s_2)
- coeff_1, coeff_2, coeff_3 = (-r_1 * h_eta).expm1(), (-r_2 * h_eta).expm1(), (-h_eta).expm1()
- if inject_noise:
- # 0 < r_1 < r_2 < 1
- noise_coeff_1 = (-2 * r_1 * h * eta).expm1().neg().sqrt()
- noise_coeff_2 = (-r_1 * h * eta).exp() * (-2 * (r_2 - r_1) * h * eta).expm1().neg().sqrt()
- noise_coeff_3 = (-r_2 * h * eta).exp() * (-2 * (1 - r_2) * h * eta).expm1().neg().sqrt()
- noise_1, noise_2, noise_3 = noise_sampler(sigmas[i], sigma_s_1), noise_sampler(sigma_s_1, sigma_s_2), noise_sampler(sigma_s_2, sigmas[i + 1])
+ alpha_s_1 = sigma_s_1 * lambda_s_1.exp()
+ alpha_s_2 = sigma_s_2 * lambda_s_2.exp()
+ alpha_t = sigmas[i + 1] * lambda_t.exp()
- # Step 1
- x_2 = sigma_s_1 / sigmas[i] * (-r_1 * h * eta).exp() * x - alpha_s_1 * coeff_1 * denoised
- if inject_noise:
- x_2 = x_2 + sigma_s_1 * (noise_coeff_1 * noise_1) * s_noise
- denoised_2 = model(x_2, sigma_s_1 * s_in, **extra_args)
+ # Step 1
+ x_2 = sigma_s_1 / sigmas[i] * (-r_1 * h * eta).exp() * x - alpha_s_1 * ei_h_phi_1(-r_1 * h_eta) * denoised
+ if inject_noise:
+ sde_noise = (-2 * r_1 * h * eta).expm1().neg().sqrt() * noise_sampler(sigmas[i], sigma_s_1)
+ x_2 = x_2 + sde_noise * sigma_s_1 * s_noise
+ denoised_2 = model(x_2, sigma_s_1 * s_in, **extra_args)
- # Step 2
- x_3 = sigma_s_2 / sigmas[i] * (-r_2 * h * eta).exp() * x - alpha_s_2 * coeff_2 * denoised + (r_2 / r_1) * alpha_s_2 * (coeff_2 / (r_2 * h_eta) + 1) * (denoised_2 - denoised)
- if inject_noise:
- x_3 = x_3 + sigma_s_2 * (noise_coeff_2 * noise_1 + noise_coeff_1 * noise_2) * s_noise
- denoised_3 = model(x_3, sigma_s_2 * s_in, **extra_args)
+ # Step 2
+ a3_2 = r_2 / r_1 * ei_h_phi_2(-r_2 * h_eta)
+ a3_1 = ei_h_phi_1(-r_2 * h_eta) - a3_2
+ x_3 = sigma_s_2 / sigmas[i] * (-r_2 * h * eta).exp() * x - alpha_s_2 * (a3_1 * denoised + a3_2 * denoised_2)
+ if inject_noise:
+ segment_factor = (r_1 - r_2) * h * eta
+ sde_noise = sde_noise * segment_factor.exp()
+ sde_noise = sde_noise + segment_factor.mul(2).expm1().neg().sqrt() * noise_sampler(sigma_s_1, sigma_s_2)
+ x_3 = x_3 + sde_noise * sigma_s_2 * s_noise
+ denoised_3 = model(x_3, sigma_s_2 * s_in, **extra_args)
- # Step 3
- x = sigmas[i + 1] / sigmas[i] * (-h * eta).exp() * x - alpha_t * coeff_3 * denoised + (1. / r_2) * alpha_t * (coeff_3 / h_eta + 1) * (denoised_3 - denoised)
- if inject_noise:
- x = x + sigmas[i + 1] * (noise_coeff_3 * noise_1 + noise_coeff_2 * noise_2 + noise_coeff_1 * noise_3) * s_noise
+ # Step 3
+ b3 = ei_h_phi_2(-h_eta) / r_2
+ b1 = ei_h_phi_1(-h_eta) - b3
+ x = sigmas[i + 1] / sigmas[i] * (-h * eta).exp() * x - alpha_t * (b1 * denoised + b3 * denoised_3)
+ if inject_noise:
+ segment_factor = (r_2 - 1) * h * eta
+ sde_noise = sde_noise * segment_factor.exp()
+ sde_noise = sde_noise + segment_factor.mul(2).expm1().neg().sqrt() * noise_sampler(sigma_s_2, sigmas[i + 1])
+ x = x + sde_noise * sigmas[i + 1] * s_noise
return x
+
+
+@torch.no_grad()
+def sample_sa_solver(model, x, sigmas, extra_args=None, callback=None, disable=False, tau_func=None, s_noise=1.0, noise_sampler=None, predictor_order=3, corrector_order=4, use_pece=False, simple_order_2=False):
+ """Stochastic Adams Solver with predictor-corrector method (NeurIPS 2023)."""
+ if len(sigmas) <= 1:
+ return x
+ extra_args = {} if extra_args is None else extra_args
+ seed = extra_args.get("seed", None)
+ noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
+ s_in = x.new_ones([x.shape[0]])
+
+ model_sampling = model.inner_model.model_patcher.get_model_object("model_sampling")
+ sigmas = offset_first_sigma_for_snr(sigmas, model_sampling)
+ lambdas = sigma_to_half_log_snr(sigmas, model_sampling=model_sampling)
+
+ if tau_func is None:
+ # Use default interval for stochastic sampling
+ start_sigma = model_sampling.percent_to_sigma(0.2)
+ end_sigma = model_sampling.percent_to_sigma(0.8)
+ tau_func = sa_solver.get_tau_interval_func(start_sigma, end_sigma, eta=1.0)
+
+ max_used_order = max(predictor_order, corrector_order)
+ x_pred = x # x: current state, x_pred: predicted next state
+
+ h = 0.0
+ tau_t = 0.0
+ noise = 0.0
+ pred_list = []
+
+ # Lower order near the end to improve stability
+ lower_order_to_end = sigmas[-1].item() == 0
+
+ for i in trange(len(sigmas) - 1, disable=disable):
+ # Evaluation
+ denoised = model(x_pred, sigmas[i] * s_in, **extra_args)
+ if callback is not None:
+ callback({"x": x_pred, "i": i, "sigma": sigmas[i], "sigma_hat": sigmas[i], "denoised": denoised})
+ pred_list.append(denoised)
+ pred_list = pred_list[-max_used_order:]
+
+ predictor_order_used = min(predictor_order, len(pred_list))
+ if i == 0 or (sigmas[i + 1] == 0 and not use_pece):
+ corrector_order_used = 0
+ else:
+ corrector_order_used = min(corrector_order, len(pred_list))
+
+ if lower_order_to_end:
+ predictor_order_used = min(predictor_order_used, len(sigmas) - 2 - i)
+ corrector_order_used = min(corrector_order_used, len(sigmas) - 1 - i)
+
+ # Corrector
+ if corrector_order_used == 0:
+ # Update by the predicted state
+ x = x_pred
+ else:
+ curr_lambdas = lambdas[i - corrector_order_used + 1:i + 1]
+ b_coeffs = sa_solver.compute_stochastic_adams_b_coeffs(
+ sigmas[i],
+ curr_lambdas,
+ lambdas[i - 1],
+ lambdas[i],
+ tau_t,
+ simple_order_2,
+ is_corrector_step=True,
+ )
+ pred_mat = torch.stack(pred_list[-corrector_order_used:], dim=1) # (B, K, ...)
+ corr_res = torch.tensordot(pred_mat, b_coeffs, dims=([1], [0])) # (B, ...)
+ x = sigmas[i] / sigmas[i - 1] * (-(tau_t ** 2) * h).exp() * x + corr_res
+
+ if tau_t > 0 and s_noise > 0:
+ # The noise from the previous predictor step
+ x = x + noise
+
+ if use_pece:
+ # Evaluate the corrected state
+ denoised = model(x, sigmas[i] * s_in, **extra_args)
+ pred_list[-1] = denoised
+
+ # Predictor
+ if sigmas[i + 1] == 0:
+ # Denoising step
+ x = denoised
+ else:
+ tau_t = tau_func(sigmas[i + 1])
+ curr_lambdas = lambdas[i - predictor_order_used + 1:i + 1]
+ b_coeffs = sa_solver.compute_stochastic_adams_b_coeffs(
+ sigmas[i + 1],
+ curr_lambdas,
+ lambdas[i],
+ lambdas[i + 1],
+ tau_t,
+ simple_order_2,
+ is_corrector_step=False,
+ )
+ pred_mat = torch.stack(pred_list[-predictor_order_used:], dim=1) # (B, K, ...)
+ pred_res = torch.tensordot(pred_mat, b_coeffs, dims=([1], [0])) # (B, ...)
+ h = lambdas[i + 1] - lambdas[i]
+ x_pred = sigmas[i + 1] / sigmas[i] * (-(tau_t ** 2) * h).exp() * x + pred_res
+
+ if tau_t > 0 and s_noise > 0:
+ noise = noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * tau_t ** 2 * h).expm1().neg().sqrt() * s_noise
+ x_pred = x_pred + noise
+ return x
+
+
+@torch.no_grad()
+def sample_sa_solver_pece(model, x, sigmas, extra_args=None, callback=None, disable=False, tau_func=None, s_noise=1.0, noise_sampler=None, predictor_order=3, corrector_order=4, simple_order_2=False):
+ """Stochastic Adams Solver with PECE (Predict–Evaluate–Correct–Evaluate) mode (NeurIPS 2023)."""
+ return sample_sa_solver(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, tau_func=tau_func, s_noise=s_noise, noise_sampler=noise_sampler, predictor_order=predictor_order, corrector_order=corrector_order, use_pece=True, simple_order_2=simple_order_2)
diff --git a/comfy/latent_formats.py b/comfy/latent_formats.py
index 82d9f9bb8..77e642a94 100644
--- a/comfy/latent_formats.py
+++ b/comfy/latent_formats.py
@@ -457,11 +457,170 @@ class Wan21(LatentFormat):
latents_std = self.latents_std.to(latent.device, latent.dtype)
return latent * latents_std / self.scale_factor + latents_mean
+class Wan22(Wan21):
+ latent_channels = 48
+ latent_dimensions = 3
+
+ latent_rgb_factors = [
+ [ 0.0119, 0.0103, 0.0046],
+ [-0.1062, -0.0504, 0.0165],
+ [ 0.0140, 0.0409, 0.0491],
+ [-0.0813, -0.0677, 0.0607],
+ [ 0.0656, 0.0851, 0.0808],
+ [ 0.0264, 0.0463, 0.0912],
+ [ 0.0295, 0.0326, 0.0590],
+ [-0.0244, -0.0270, 0.0025],
+ [ 0.0443, -0.0102, 0.0288],
+ [-0.0465, -0.0090, -0.0205],
+ [ 0.0359, 0.0236, 0.0082],
+ [-0.0776, 0.0854, 0.1048],
+ [ 0.0564, 0.0264, 0.0561],
+ [ 0.0006, 0.0594, 0.0418],
+ [-0.0319, -0.0542, -0.0637],
+ [-0.0268, 0.0024, 0.0260],
+ [ 0.0539, 0.0265, 0.0358],
+ [-0.0359, -0.0312, -0.0287],
+ [-0.0285, -0.1032, -0.1237],
+ [ 0.1041, 0.0537, 0.0622],
+ [-0.0086, -0.0374, -0.0051],
+ [ 0.0390, 0.0670, 0.2863],
+ [ 0.0069, 0.0144, 0.0082],
+ [ 0.0006, -0.0167, 0.0079],
+ [ 0.0313, -0.0574, -0.0232],
+ [-0.1454, -0.0902, -0.0481],
+ [ 0.0714, 0.0827, 0.0447],
+ [-0.0304, -0.0574, -0.0196],
+ [ 0.0401, 0.0384, 0.0204],
+ [-0.0758, -0.0297, -0.0014],
+ [ 0.0568, 0.1307, 0.1372],
+ [-0.0055, -0.0310, -0.0380],
+ [ 0.0239, -0.0305, 0.0325],
+ [-0.0663, -0.0673, -0.0140],
+ [-0.0416, -0.0047, -0.0023],
+ [ 0.0166, 0.0112, -0.0093],
+ [-0.0211, 0.0011, 0.0331],
+ [ 0.1833, 0.1466, 0.2250],
+ [-0.0368, 0.0370, 0.0295],
+ [-0.3441, -0.3543, -0.2008],
+ [-0.0479, -0.0489, -0.0420],
+ [-0.0660, -0.0153, 0.0800],
+ [-0.0101, 0.0068, 0.0156],
+ [-0.0690, -0.0452, -0.0927],
+ [-0.0145, 0.0041, 0.0015],
+ [ 0.0421, 0.0451, 0.0373],
+ [ 0.0504, -0.0483, -0.0356],
+ [-0.0837, 0.0168, 0.0055]
+ ]
+
+ latent_rgb_factors_bias = [0.0317, -0.0878, -0.1388]
+
+ def __init__(self):
+ self.scale_factor = 1.0
+ self.latents_mean = torch.tensor([
+ -0.2289, -0.0052, -0.1323, -0.2339, -0.2799, 0.0174, 0.1838, 0.1557,
+ -0.1382, 0.0542, 0.2813, 0.0891, 0.1570, -0.0098, 0.0375, -0.1825,
+ -0.2246, -0.1207, -0.0698, 0.5109, 0.2665, -0.2108, -0.2158, 0.2502,
+ -0.2055, -0.0322, 0.1109, 0.1567, -0.0729, 0.0899, -0.2799, -0.1230,
+ -0.0313, -0.1649, 0.0117, 0.0723, -0.2839, -0.2083, -0.0520, 0.3748,
+ 0.0152, 0.1957, 0.1433, -0.2944, 0.3573, -0.0548, -0.1681, -0.0667,
+ ]).view(1, self.latent_channels, 1, 1, 1)
+ self.latents_std = torch.tensor([
+ 0.4765, 1.0364, 0.4514, 1.1677, 0.5313, 0.4990, 0.4818, 0.5013,
+ 0.8158, 1.0344, 0.5894, 1.0901, 0.6885, 0.6165, 0.8454, 0.4978,
+ 0.5759, 0.3523, 0.7135, 0.6804, 0.5833, 1.4146, 0.8986, 0.5659,
+ 0.7069, 0.5338, 0.4889, 0.4917, 0.4069, 0.4999, 0.6866, 0.4093,
+ 0.5709, 0.6065, 0.6415, 0.4944, 0.5726, 1.2042, 0.5458, 1.6887,
+ 0.3971, 1.0600, 0.3943, 0.5537, 0.5444, 0.4089, 0.7468, 0.7744
+ ]).view(1, self.latent_channels, 1, 1, 1)
+
+class HunyuanImage21(LatentFormat):
+ latent_channels = 64
+ latent_dimensions = 2
+ scale_factor = 0.75289
+
+ latent_rgb_factors = [
+ [-0.0154, -0.0397, -0.0521],
+ [ 0.0005, 0.0093, 0.0006],
+ [-0.0805, -0.0773, -0.0586],
+ [-0.0494, -0.0487, -0.0498],
+ [-0.0212, -0.0076, -0.0261],
+ [-0.0179, -0.0417, -0.0505],
+ [ 0.0158, 0.0310, 0.0239],
+ [ 0.0409, 0.0516, 0.0201],
+ [ 0.0350, 0.0553, 0.0036],
+ [-0.0447, -0.0327, -0.0479],
+ [-0.0038, -0.0221, -0.0365],
+ [-0.0423, -0.0718, -0.0654],
+ [ 0.0039, 0.0368, 0.0104],
+ [ 0.0655, 0.0217, 0.0122],
+ [ 0.0490, 0.1638, 0.2053],
+ [ 0.0932, 0.0829, 0.0650],
+ [-0.0186, -0.0209, -0.0135],
+ [-0.0080, -0.0076, -0.0148],
+ [-0.0284, -0.0201, 0.0011],
+ [-0.0642, -0.0294, -0.0777],
+ [-0.0035, 0.0076, -0.0140],
+ [ 0.0519, 0.0731, 0.0887],
+ [-0.0102, 0.0095, 0.0704],
+ [ 0.0068, 0.0218, -0.0023],
+ [-0.0726, -0.0486, -0.0519],
+ [ 0.0260, 0.0295, 0.0263],
+ [ 0.0250, 0.0333, 0.0341],
+ [ 0.0168, -0.0120, -0.0174],
+ [ 0.0226, 0.1037, 0.0114],
+ [ 0.2577, 0.1906, 0.1604],
+ [-0.0646, -0.0137, -0.0018],
+ [-0.0112, 0.0309, 0.0358],
+ [-0.0347, 0.0146, -0.0481],
+ [ 0.0234, 0.0179, 0.0201],
+ [ 0.0157, 0.0313, 0.0225],
+ [ 0.0423, 0.0675, 0.0524],
+ [-0.0031, 0.0027, -0.0255],
+ [ 0.0447, 0.0555, 0.0330],
+ [-0.0152, 0.0103, 0.0299],
+ [-0.0755, -0.0489, -0.0635],
+ [ 0.0853, 0.0788, 0.1017],
+ [-0.0272, -0.0294, -0.0471],
+ [ 0.0440, 0.0400, -0.0137],
+ [ 0.0335, 0.0317, -0.0036],
+ [-0.0344, -0.0621, -0.0984],
+ [-0.0127, -0.0630, -0.0620],
+ [-0.0648, 0.0360, 0.0924],
+ [-0.0781, -0.0801, -0.0409],
+ [ 0.0363, 0.0613, 0.0499],
+ [ 0.0238, 0.0034, 0.0041],
+ [-0.0135, 0.0258, 0.0310],
+ [ 0.0614, 0.1086, 0.0589],
+ [ 0.0428, 0.0350, 0.0205],
+ [ 0.0153, 0.0173, -0.0018],
+ [-0.0288, -0.0455, -0.0091],
+ [ 0.0344, 0.0109, -0.0157],
+ [-0.0205, -0.0247, -0.0187],
+ [ 0.0487, 0.0126, 0.0064],
+ [-0.0220, -0.0013, 0.0074],
+ [-0.0203, -0.0094, -0.0048],
+ [-0.0719, 0.0429, -0.0442],
+ [ 0.1042, 0.0497, 0.0356],
+ [-0.0659, -0.0578, -0.0280],
+ [-0.0060, -0.0322, -0.0234]]
+
+ latent_rgb_factors_bias = [0.0007, -0.0256, -0.0206]
+
+class HunyuanImage21Refiner(LatentFormat):
+ latent_channels = 64
+ latent_dimensions = 3
+ scale_factor = 1.03682
+
class Hunyuan3Dv2(LatentFormat):
latent_channels = 64
latent_dimensions = 1
scale_factor = 0.9990943042622529
+class Hunyuan3Dv2_1(LatentFormat):
+ scale_factor = 1.0039506158752403
+ latent_channels = 64
+ latent_dimensions = 1
+
class Hunyuan3Dv2mini(LatentFormat):
latent_channels = 64
latent_dimensions = 1
@@ -470,3 +629,20 @@ class Hunyuan3Dv2mini(LatentFormat):
class ACEAudio(LatentFormat):
latent_channels = 8
latent_dimensions = 2
+
+class ChromaRadiance(LatentFormat):
+ latent_channels = 3
+
+ def __init__(self):
+ self.latent_rgb_factors = [
+ # R G B
+ [ 1.0, 0.0, 0.0 ],
+ [ 0.0, 1.0, 0.0 ],
+ [ 0.0, 0.0, 1.0 ]
+ ]
+
+ def process_in(self, latent):
+ return latent
+
+ def process_out(self, latent):
+ return latent
diff --git a/comfy/ldm/ace/attention.py b/comfy/ldm/ace/attention.py
index f20a01669..670eb9783 100644
--- a/comfy/ldm/ace/attention.py
+++ b/comfy/ldm/ace/attention.py
@@ -133,6 +133,7 @@ class Attention(nn.Module):
hidden_states: torch.Tensor,
encoder_hidden_states: Optional[torch.Tensor] = None,
attention_mask: Optional[torch.Tensor] = None,
+ transformer_options={},
**cross_attention_kwargs,
) -> torch.Tensor:
return self.processor(
@@ -140,6 +141,7 @@ class Attention(nn.Module):
hidden_states,
encoder_hidden_states=encoder_hidden_states,
attention_mask=attention_mask,
+ transformer_options=transformer_options,
**cross_attention_kwargs,
)
@@ -366,6 +368,7 @@ class CustomerAttnProcessor2_0:
encoder_attention_mask: Optional[torch.FloatTensor] = None,
rotary_freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]] = None,
rotary_freqs_cis_cross: Union[torch.Tensor, Tuple[torch.Tensor]] = None,
+ transformer_options={},
*args,
**kwargs,
) -> torch.Tensor:
@@ -433,7 +436,7 @@ class CustomerAttnProcessor2_0:
# the output of sdp = (batch, num_heads, seq_len, head_dim)
hidden_states = optimized_attention(
- query, key, value, heads=query.shape[1], mask=attention_mask, skip_reshape=True,
+ query, key, value, heads=query.shape[1], mask=attention_mask, skip_reshape=True, transformer_options=transformer_options,
).to(query.dtype)
# linear proj
@@ -697,6 +700,7 @@ class LinearTransformerBlock(nn.Module):
rotary_freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]] = None,
rotary_freqs_cis_cross: Union[torch.Tensor, Tuple[torch.Tensor]] = None,
temb: torch.FloatTensor = None,
+ transformer_options={},
):
N = hidden_states.shape[0]
@@ -720,6 +724,7 @@ class LinearTransformerBlock(nn.Module):
encoder_attention_mask=encoder_attention_mask,
rotary_freqs_cis=rotary_freqs_cis,
rotary_freqs_cis_cross=rotary_freqs_cis_cross,
+ transformer_options=transformer_options,
)
else:
attn_output, _ = self.attn(
@@ -729,6 +734,7 @@ class LinearTransformerBlock(nn.Module):
encoder_attention_mask=None,
rotary_freqs_cis=rotary_freqs_cis,
rotary_freqs_cis_cross=None,
+ transformer_options=transformer_options,
)
if self.use_adaln_single:
@@ -743,6 +749,7 @@ class LinearTransformerBlock(nn.Module):
encoder_attention_mask=encoder_attention_mask,
rotary_freqs_cis=rotary_freqs_cis,
rotary_freqs_cis_cross=rotary_freqs_cis_cross,
+ transformer_options=transformer_options,
)
hidden_states = attn_output + hidden_states
diff --git a/comfy/ldm/ace/model.py b/comfy/ldm/ace/model.py
index 12c524701..399329853 100644
--- a/comfy/ldm/ace/model.py
+++ b/comfy/ldm/ace/model.py
@@ -19,6 +19,7 @@ import torch
from torch import nn
import comfy.model_management
+import comfy.patcher_extension
from comfy.ldm.lightricks.model import TimestepEmbedding, Timesteps
from .attention import LinearTransformerBlock, t2i_modulate
@@ -313,6 +314,7 @@ class ACEStepTransformer2DModel(nn.Module):
output_length: int = 0,
block_controlnet_hidden_states: Optional[Union[List[torch.Tensor], torch.Tensor]] = None,
controlnet_scale: Union[float, torch.Tensor] = 1.0,
+ transformer_options={},
):
embedded_timestep = self.timestep_embedder(self.time_proj(timestep).to(dtype=hidden_states.dtype))
temb = self.t_block(embedded_timestep)
@@ -338,12 +340,34 @@ class ACEStepTransformer2DModel(nn.Module):
rotary_freqs_cis=rotary_freqs_cis,
rotary_freqs_cis_cross=encoder_rotary_freqs_cis,
temb=temb,
+ transformer_options=transformer_options,
)
output = self.final_layer(hidden_states, embedded_timestep, output_length)
return output
- def forward(
+ def forward(self,
+ x,
+ timestep,
+ attention_mask=None,
+ context: Optional[torch.Tensor] = None,
+ text_attention_mask: Optional[torch.LongTensor] = None,
+ speaker_embeds: Optional[torch.FloatTensor] = None,
+ lyric_token_idx: Optional[torch.LongTensor] = None,
+ lyric_mask: Optional[torch.LongTensor] = None,
+ block_controlnet_hidden_states: Optional[Union[List[torch.Tensor], torch.Tensor]] = None,
+ controlnet_scale: Union[float, torch.Tensor] = 1.0,
+ lyrics_strength=1.0,
+ **kwargs
+ ):
+ return comfy.patcher_extension.WrapperExecutor.new_class_executor(
+ self._forward,
+ self,
+ comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, kwargs.get("transformer_options", {}))
+ ).execute(x, timestep, attention_mask, context, text_attention_mask, speaker_embeds, lyric_token_idx, lyric_mask, block_controlnet_hidden_states,
+ controlnet_scale, lyrics_strength, **kwargs)
+
+ def _forward(
self,
x,
timestep,
@@ -371,6 +395,7 @@ class ACEStepTransformer2DModel(nn.Module):
output_length = hidden_states.shape[-1]
+ transformer_options = kwargs.get("transformer_options", {})
output = self.decode(
hidden_states=hidden_states,
attention_mask=attention_mask,
@@ -380,6 +405,7 @@ class ACEStepTransformer2DModel(nn.Module):
output_length=output_length,
block_controlnet_hidden_states=block_controlnet_hidden_states,
controlnet_scale=controlnet_scale,
+ transformer_options=transformer_options,
)
return output
diff --git a/comfy/ldm/ace/vae/music_dcae_pipeline.py b/comfy/ldm/ace/vae/music_dcae_pipeline.py
index af81280eb..3c8830c17 100644
--- a/comfy/ldm/ace/vae/music_dcae_pipeline.py
+++ b/comfy/ldm/ace/vae/music_dcae_pipeline.py
@@ -23,8 +23,6 @@ class MusicDCAE(torch.nn.Module):
else:
self.source_sample_rate = source_sample_rate
- # self.resampler = torchaudio.transforms.Resample(source_sample_rate, 44100)
-
self.transform = transforms.Compose([
transforms.Normalize(0.5, 0.5),
])
@@ -37,10 +35,6 @@ class MusicDCAE(torch.nn.Module):
self.scale_factor = 0.1786
self.shift_factor = -1.9091
- def load_audio(self, audio_path):
- audio, sr = torchaudio.load(audio_path)
- return audio, sr
-
def forward_mel(self, audios):
mels = []
for i in range(len(audios)):
@@ -73,10 +67,8 @@ class MusicDCAE(torch.nn.Module):
latent = self.dcae.encoder(mel.unsqueeze(0))
latents.append(latent)
latents = torch.cat(latents, dim=0)
- # latent_lengths = (audio_lengths / sr * 44100 / 512 / self.time_dimention_multiple).long()
latents = (latents - self.shift_factor) * self.scale_factor
return latents
- # return latents, latent_lengths
@torch.no_grad()
def decode(self, latents, audio_lengths=None, sr=None):
@@ -91,9 +83,7 @@ class MusicDCAE(torch.nn.Module):
wav = self.vocoder.decode(mels[0]).squeeze(1)
if sr is not None:
- # resampler = torchaudio.transforms.Resample(44100, sr).to(latents.device).to(latents.dtype)
wav = torchaudio.functional.resample(wav, 44100, sr)
- # wav = resampler(wav)
else:
sr = 44100
pred_wavs.append(wav)
@@ -101,7 +91,6 @@ class MusicDCAE(torch.nn.Module):
if audio_lengths is not None:
pred_wavs = [wav[:, :length].cpu() for wav, length in zip(pred_wavs, audio_lengths)]
return torch.stack(pred_wavs)
- # return sr, pred_wavs
def forward(self, audios, audio_lengths=None, sr=None):
latents, latent_lengths = self.encode(audios=audios, audio_lengths=audio_lengths, sr=sr)
diff --git a/comfy/ldm/audio/dit.py b/comfy/ldm/audio/dit.py
index 179c5b67e..ca865189e 100644
--- a/comfy/ldm/audio/dit.py
+++ b/comfy/ldm/audio/dit.py
@@ -298,7 +298,8 @@ class Attention(nn.Module):
mask = None,
context_mask = None,
rotary_pos_emb = None,
- causal = None
+ causal = None,
+ transformer_options={},
):
h, kv_h, has_context = self.num_heads, self.kv_heads, context is not None
@@ -363,7 +364,7 @@ class Attention(nn.Module):
heads_per_kv_head = h // kv_h
k, v = map(lambda t: t.repeat_interleave(heads_per_kv_head, dim = 1), (k, v))
- out = optimized_attention(q, k, v, h, skip_reshape=True)
+ out = optimized_attention(q, k, v, h, skip_reshape=True, transformer_options=transformer_options)
out = self.to_out(out)
if mask is not None:
@@ -488,7 +489,8 @@ class TransformerBlock(nn.Module):
global_cond=None,
mask = None,
context_mask = None,
- rotary_pos_emb = None
+ rotary_pos_emb = None,
+ transformer_options={}
):
if self.global_cond_dim is not None and self.global_cond_dim > 0 and global_cond is not None:
@@ -498,12 +500,12 @@ class TransformerBlock(nn.Module):
residual = x
x = self.pre_norm(x)
x = x * (1 + scale_self) + shift_self
- x = self.self_attn(x, mask = mask, rotary_pos_emb = rotary_pos_emb)
+ x = self.self_attn(x, mask = mask, rotary_pos_emb = rotary_pos_emb, transformer_options=transformer_options)
x = x * torch.sigmoid(1 - gate_self)
x = x + residual
if context is not None:
- x = x + self.cross_attn(self.cross_attend_norm(x), context = context, context_mask = context_mask)
+ x = x + self.cross_attn(self.cross_attend_norm(x), context = context, context_mask = context_mask, transformer_options=transformer_options)
if self.conformer is not None:
x = x + self.conformer(x)
@@ -517,10 +519,10 @@ class TransformerBlock(nn.Module):
x = x + residual
else:
- x = x + self.self_attn(self.pre_norm(x), mask = mask, rotary_pos_emb = rotary_pos_emb)
+ x = x + self.self_attn(self.pre_norm(x), mask = mask, rotary_pos_emb = rotary_pos_emb, transformer_options=transformer_options)
if context is not None:
- x = x + self.cross_attn(self.cross_attend_norm(x), context = context, context_mask = context_mask)
+ x = x + self.cross_attn(self.cross_attend_norm(x), context = context, context_mask = context_mask, transformer_options=transformer_options)
if self.conformer is not None:
x = x + self.conformer(x)
@@ -606,7 +608,8 @@ class ContinuousTransformer(nn.Module):
return_info = False,
**kwargs
):
- patches_replace = kwargs.get("transformer_options", {}).get("patches_replace", {})
+ transformer_options = kwargs.get("transformer_options", {})
+ patches_replace = transformer_options.get("patches_replace", {})
batch, seq, device = *x.shape[:2], x.device
context = kwargs["context"]
@@ -632,7 +635,7 @@ class ContinuousTransformer(nn.Module):
# Attention layers
if self.rotary_pos_emb is not None:
- rotary_pos_emb = self.rotary_pos_emb.forward_from_seq_len(x.shape[1], dtype=x.dtype, device=x.device)
+ rotary_pos_emb = self.rotary_pos_emb.forward_from_seq_len(x.shape[1], dtype=torch.float, device=x.device)
else:
rotary_pos_emb = None
@@ -645,13 +648,13 @@ class ContinuousTransformer(nn.Module):
if ("double_block", i) in blocks_replace:
def block_wrap(args):
out = {}
- out["img"] = layer(args["img"], rotary_pos_emb=args["pe"], global_cond=args["vec"], context=args["txt"])
+ out["img"] = layer(args["img"], rotary_pos_emb=args["pe"], global_cond=args["vec"], context=args["txt"], transformer_options=args["transformer_options"])
return out
- out = blocks_replace[("double_block", i)]({"img": x, "txt": context, "vec": global_cond, "pe": rotary_pos_emb}, {"original_block": block_wrap})
+ out = blocks_replace[("double_block", i)]({"img": x, "txt": context, "vec": global_cond, "pe": rotary_pos_emb, "transformer_options": transformer_options}, {"original_block": block_wrap})
x = out["img"]
else:
- x = layer(x, rotary_pos_emb = rotary_pos_emb, global_cond=global_cond, context=context)
+ x = layer(x, rotary_pos_emb = rotary_pos_emb, global_cond=global_cond, context=context, transformer_options=transformer_options)
# x = checkpoint(layer, x, rotary_pos_emb = rotary_pos_emb, global_cond=global_cond, **kwargs)
if return_info:
diff --git a/comfy/ldm/aura/mmdit.py b/comfy/ldm/aura/mmdit.py
index 1258ae11f..66d9613b6 100644
--- a/comfy/ldm/aura/mmdit.py
+++ b/comfy/ldm/aura/mmdit.py
@@ -9,6 +9,7 @@ import torch.nn.functional as F
from comfy.ldm.modules.attention import optimized_attention
import comfy.ops
+import comfy.patcher_extension
import comfy.ldm.common_dit
def modulate(x, shift, scale):
@@ -84,7 +85,7 @@ class SingleAttention(nn.Module):
)
#@torch.compile()
- def forward(self, c):
+ def forward(self, c, transformer_options={}):
bsz, seqlen1, _ = c.shape
@@ -94,7 +95,7 @@ class SingleAttention(nn.Module):
v = v.view(bsz, seqlen1, self.n_heads, self.head_dim)
q, k = self.q_norm1(q), self.k_norm1(k)
- output = optimized_attention(q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3), self.n_heads, skip_reshape=True)
+ output = optimized_attention(q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3), self.n_heads, skip_reshape=True, transformer_options=transformer_options)
c = self.w1o(output)
return c
@@ -143,7 +144,7 @@ class DoubleAttention(nn.Module):
#@torch.compile()
- def forward(self, c, x):
+ def forward(self, c, x, transformer_options={}):
bsz, seqlen1, _ = c.shape
bsz, seqlen2, _ = x.shape
@@ -167,7 +168,7 @@ class DoubleAttention(nn.Module):
torch.cat([cv, xv], dim=1),
)
- output = optimized_attention(q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3), self.n_heads, skip_reshape=True)
+ output = optimized_attention(q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3), self.n_heads, skip_reshape=True, transformer_options=transformer_options)
c, x = output.split([seqlen1, seqlen2], dim=1)
c = self.w1o(c)
@@ -206,7 +207,7 @@ class MMDiTBlock(nn.Module):
self.is_last = is_last
#@torch.compile()
- def forward(self, c, x, global_cond, **kwargs):
+ def forward(self, c, x, global_cond, transformer_options={}, **kwargs):
cres, xres = c, x
@@ -224,7 +225,7 @@ class MMDiTBlock(nn.Module):
x = modulate(self.normX1(x), xshift_msa, xscale_msa)
# attention
- c, x = self.attn(c, x)
+ c, x = self.attn(c, x, transformer_options=transformer_options)
c = self.normC2(cres + cgate_msa.unsqueeze(1) * c)
@@ -254,13 +255,13 @@ class DiTBlock(nn.Module):
self.mlp = MLP(dim, hidden_dim=dim * 4, dtype=dtype, device=device, operations=operations)
#@torch.compile()
- def forward(self, cx, global_cond, **kwargs):
+ def forward(self, cx, global_cond, transformer_options={}, **kwargs):
cxres = cx
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.modCX(
global_cond
).chunk(6, dim=1)
cx = modulate(self.norm1(cx), shift_msa, scale_msa)
- cx = self.attn(cx)
+ cx = self.attn(cx, transformer_options=transformer_options)
cx = self.norm2(cxres + gate_msa.unsqueeze(1) * cx)
mlpout = self.mlp(modulate(cx, shift_mlp, scale_mlp))
cx = gate_mlp.unsqueeze(1) * mlpout
@@ -436,6 +437,13 @@ class MMDiT(nn.Module):
return x + pos_encoding.reshape(1, -1, self.positional_encoding.shape[-1])
def forward(self, x, timestep, context, transformer_options={}, **kwargs):
+ return comfy.patcher_extension.WrapperExecutor.new_class_executor(
+ self._forward,
+ self,
+ comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options)
+ ).execute(x, timestep, context, transformer_options, **kwargs)
+
+ def _forward(self, x, timestep, context, transformer_options={}, **kwargs):
patches_replace = transformer_options.get("patches_replace", {})
# patchify x, add PE
b, c, h, w = x.shape
@@ -465,13 +473,14 @@ class MMDiT(nn.Module):
out = {}
out["txt"], out["img"] = layer(args["txt"],
args["img"],
- args["vec"])
+ args["vec"],
+ transformer_options=args["transformer_options"])
return out
- out = blocks_replace[("double_block", i)]({"img": x, "txt": c, "vec": global_cond}, {"original_block": block_wrap})
+ out = blocks_replace[("double_block", i)]({"img": x, "txt": c, "vec": global_cond, "transformer_options": transformer_options}, {"original_block": block_wrap})
c = out["txt"]
x = out["img"]
else:
- c, x = layer(c, x, global_cond, **kwargs)
+ c, x = layer(c, x, global_cond, transformer_options=transformer_options, **kwargs)
if len(self.single_layers) > 0:
c_len = c.size(1)
@@ -480,13 +489,13 @@ class MMDiT(nn.Module):
if ("single_block", i) in blocks_replace:
def block_wrap(args):
out = {}
- out["img"] = layer(args["img"], args["vec"])
+ out["img"] = layer(args["img"], args["vec"], transformer_options=args["transformer_options"])
return out
- out = blocks_replace[("single_block", i)]({"img": cx, "vec": global_cond}, {"original_block": block_wrap})
+ out = blocks_replace[("single_block", i)]({"img": cx, "vec": global_cond, "transformer_options": transformer_options}, {"original_block": block_wrap})
cx = out["img"]
else:
- cx = layer(cx, global_cond, **kwargs)
+ cx = layer(cx, global_cond, transformer_options=transformer_options, **kwargs)
x = cx[:, c_len:]
diff --git a/comfy/ldm/cascade/common.py b/comfy/ldm/cascade/common.py
index 3eaa0c821..42ef98c7a 100644
--- a/comfy/ldm/cascade/common.py
+++ b/comfy/ldm/cascade/common.py
@@ -32,12 +32,12 @@ class OptimizedAttention(nn.Module):
self.out_proj = operations.Linear(c, c, bias=True, dtype=dtype, device=device)
- def forward(self, q, k, v):
+ def forward(self, q, k, v, transformer_options={}):
q = self.to_q(q)
k = self.to_k(k)
v = self.to_v(v)
- out = optimized_attention(q, k, v, self.heads)
+ out = optimized_attention(q, k, v, self.heads, transformer_options=transformer_options)
return self.out_proj(out)
@@ -47,13 +47,13 @@ class Attention2D(nn.Module):
self.attn = OptimizedAttention(c, nhead, dtype=dtype, device=device, operations=operations)
# self.attn = nn.MultiheadAttention(c, nhead, dropout=dropout, bias=True, batch_first=True, dtype=dtype, device=device)
- def forward(self, x, kv, self_attn=False):
+ def forward(self, x, kv, self_attn=False, transformer_options={}):
orig_shape = x.shape
x = x.view(x.size(0), x.size(1), -1).permute(0, 2, 1) # Bx4xHxW -> Bx(HxW)x4
if self_attn:
kv = torch.cat([x, kv], dim=1)
# x = self.attn(x, kv, kv, need_weights=False)[0]
- x = self.attn(x, kv, kv)
+ x = self.attn(x, kv, kv, transformer_options=transformer_options)
x = x.permute(0, 2, 1).view(*orig_shape)
return x
@@ -114,9 +114,9 @@ class AttnBlock(nn.Module):
operations.Linear(c_cond, c, dtype=dtype, device=device)
)
- def forward(self, x, kv):
+ def forward(self, x, kv, transformer_options={}):
kv = self.kv_mapper(kv)
- x = x + self.attention(self.norm(x), kv, self_attn=self.self_attn)
+ x = x + self.attention(self.norm(x), kv, self_attn=self.self_attn, transformer_options=transformer_options)
return x
diff --git a/comfy/ldm/cascade/stage_b.py b/comfy/ldm/cascade/stage_b.py
index 773830956..428c67fdf 100644
--- a/comfy/ldm/cascade/stage_b.py
+++ b/comfy/ldm/cascade/stage_b.py
@@ -173,7 +173,7 @@ class StageB(nn.Module):
clip = self.clip_norm(clip)
return clip
- def _down_encode(self, x, r_embed, clip):
+ def _down_encode(self, x, r_embed, clip, transformer_options={}):
level_outputs = []
block_group = zip(self.down_blocks, self.down_downscalers, self.down_repeat_mappers)
for down_block, downscaler, repmap in block_group:
@@ -187,7 +187,7 @@ class StageB(nn.Module):
elif isinstance(block, AttnBlock) or (
hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
AttnBlock)):
- x = block(x, clip)
+ x = block(x, clip, transformer_options=transformer_options)
elif isinstance(block, TimestepBlock) or (
hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
TimestepBlock)):
@@ -199,7 +199,7 @@ class StageB(nn.Module):
level_outputs.insert(0, x)
return level_outputs
- def _up_decode(self, level_outputs, r_embed, clip):
+ def _up_decode(self, level_outputs, r_embed, clip, transformer_options={}):
x = level_outputs[0]
block_group = zip(self.up_blocks, self.up_upscalers, self.up_repeat_mappers)
for i, (up_block, upscaler, repmap) in enumerate(block_group):
@@ -216,7 +216,7 @@ class StageB(nn.Module):
elif isinstance(block, AttnBlock) or (
hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
AttnBlock)):
- x = block(x, clip)
+ x = block(x, clip, transformer_options=transformer_options)
elif isinstance(block, TimestepBlock) or (
hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
TimestepBlock)):
@@ -228,7 +228,7 @@ class StageB(nn.Module):
x = upscaler(x)
return x
- def forward(self, x, r, effnet, clip, pixels=None, **kwargs):
+ def forward(self, x, r, effnet, clip, pixels=None, transformer_options={}, **kwargs):
if pixels is None:
pixels = x.new_zeros(x.size(0), 3, 8, 8)
@@ -245,8 +245,8 @@ class StageB(nn.Module):
nn.functional.interpolate(effnet, size=x.shape[-2:], mode='bilinear', align_corners=True))
x = x + nn.functional.interpolate(self.pixels_mapper(pixels), size=x.shape[-2:], mode='bilinear',
align_corners=True)
- level_outputs = self._down_encode(x, r_embed, clip)
- x = self._up_decode(level_outputs, r_embed, clip)
+ level_outputs = self._down_encode(x, r_embed, clip, transformer_options=transformer_options)
+ x = self._up_decode(level_outputs, r_embed, clip, transformer_options=transformer_options)
return self.clf(x)
def update_weights_ema(self, src_model, beta=0.999):
diff --git a/comfy/ldm/cascade/stage_c.py b/comfy/ldm/cascade/stage_c.py
index b952d0349..ebc4434e2 100644
--- a/comfy/ldm/cascade/stage_c.py
+++ b/comfy/ldm/cascade/stage_c.py
@@ -182,7 +182,7 @@ class StageC(nn.Module):
clip = self.clip_norm(clip)
return clip
- def _down_encode(self, x, r_embed, clip, cnet=None):
+ def _down_encode(self, x, r_embed, clip, cnet=None, transformer_options={}):
level_outputs = []
block_group = zip(self.down_blocks, self.down_downscalers, self.down_repeat_mappers)
for down_block, downscaler, repmap in block_group:
@@ -201,7 +201,7 @@ class StageC(nn.Module):
elif isinstance(block, AttnBlock) or (
hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
AttnBlock)):
- x = block(x, clip)
+ x = block(x, clip, transformer_options=transformer_options)
elif isinstance(block, TimestepBlock) or (
hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
TimestepBlock)):
@@ -213,7 +213,7 @@ class StageC(nn.Module):
level_outputs.insert(0, x)
return level_outputs
- def _up_decode(self, level_outputs, r_embed, clip, cnet=None):
+ def _up_decode(self, level_outputs, r_embed, clip, cnet=None, transformer_options={}):
x = level_outputs[0]
block_group = zip(self.up_blocks, self.up_upscalers, self.up_repeat_mappers)
for i, (up_block, upscaler, repmap) in enumerate(block_group):
@@ -235,7 +235,7 @@ class StageC(nn.Module):
elif isinstance(block, AttnBlock) or (
hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
AttnBlock)):
- x = block(x, clip)
+ x = block(x, clip, transformer_options=transformer_options)
elif isinstance(block, TimestepBlock) or (
hasattr(block, '_fsdp_wrapped_module') and isinstance(block._fsdp_wrapped_module,
TimestepBlock)):
@@ -247,7 +247,7 @@ class StageC(nn.Module):
x = upscaler(x)
return x
- def forward(self, x, r, clip_text, clip_text_pooled, clip_img, control=None, **kwargs):
+ def forward(self, x, r, clip_text, clip_text_pooled, clip_img, control=None, transformer_options={}, **kwargs):
# Process the conditioning embeddings
r_embed = self.gen_r_embedding(r).to(dtype=x.dtype)
for c in self.t_conds:
@@ -262,8 +262,8 @@ class StageC(nn.Module):
# Model Blocks
x = self.embedding(x)
- level_outputs = self._down_encode(x, r_embed, clip, cnet)
- x = self._up_decode(level_outputs, r_embed, clip, cnet)
+ level_outputs = self._down_encode(x, r_embed, clip, cnet, transformer_options=transformer_options)
+ x = self._up_decode(level_outputs, r_embed, clip, cnet, transformer_options=transformer_options)
return self.clf(x)
def update_weights_ema(self, src_model, beta=0.999):
diff --git a/comfy/ldm/chroma/layers.py b/comfy/ldm/chroma/layers.py
index 2a0dec606..fc7110cce 100644
--- a/comfy/ldm/chroma/layers.py
+++ b/comfy/ldm/chroma/layers.py
@@ -76,7 +76,7 @@ class DoubleStreamBlock(nn.Module):
)
self.flipped_img_txt = flipped_img_txt
- def forward(self, img: Tensor, txt: Tensor, pe: Tensor, vec: Tensor, attn_mask=None):
+ def forward(self, img: Tensor, txt: Tensor, pe: Tensor, vec: Tensor, attn_mask=None, transformer_options={}):
(img_mod1, img_mod2), (txt_mod1, txt_mod2) = vec
# prepare image for attention
@@ -95,7 +95,7 @@ class DoubleStreamBlock(nn.Module):
attn = attention(torch.cat((txt_q, img_q), dim=2),
torch.cat((txt_k, img_k), dim=2),
torch.cat((txt_v, img_v), dim=2),
- pe=pe, mask=attn_mask)
+ pe=pe, mask=attn_mask, transformer_options=transformer_options)
txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :]
@@ -148,7 +148,7 @@ class SingleStreamBlock(nn.Module):
self.mlp_act = nn.GELU(approximate="tanh")
- def forward(self, x: Tensor, pe: Tensor, vec: Tensor, attn_mask=None) -> Tensor:
+ def forward(self, x: Tensor, pe: Tensor, vec: Tensor, attn_mask=None, transformer_options={}) -> Tensor:
mod = vec
x_mod = torch.addcmul(mod.shift, 1 + mod.scale, self.pre_norm(x))
qkv, mlp = torch.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1)
@@ -157,7 +157,7 @@ class SingleStreamBlock(nn.Module):
q, k = self.norm(q, k, v)
# compute attention
- attn = attention(q, k, v, pe=pe, mask=attn_mask)
+ attn = attention(q, k, v, pe=pe, mask=attn_mask, transformer_options=transformer_options)
# compute activation in mlp stream, cat again and run second linear layer
output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
x.addcmul_(mod.gate, output)
diff --git a/comfy/ldm/chroma/model.py b/comfy/ldm/chroma/model.py
index c75023a31..ad1c523fe 100644
--- a/comfy/ldm/chroma/model.py
+++ b/comfy/ldm/chroma/model.py
@@ -5,6 +5,7 @@ from dataclasses import dataclass
import torch
from torch import Tensor, nn
from einops import rearrange, repeat
+import comfy.patcher_extension
import comfy.ldm.common_dit
from comfy.ldm.flux.layers import (
@@ -150,8 +151,6 @@ class Chroma(nn.Module):
attn_mask: Tensor = None,
) -> Tensor:
patches_replace = transformer_options.get("patches_replace", {})
- if img.ndim != 3 or txt.ndim != 3:
- raise ValueError("Input img and txt tensors must have 3 dimensions.")
# running on sequences img
img = self.img_in(img)
@@ -192,14 +191,16 @@ class Chroma(nn.Module):
txt=args["txt"],
vec=args["vec"],
pe=args["pe"],
- attn_mask=args.get("attn_mask"))
+ attn_mask=args.get("attn_mask"),
+ transformer_options=args.get("transformer_options"))
return out
out = blocks_replace[("double_block", i)]({"img": img,
"txt": txt,
"vec": double_mod,
"pe": pe,
- "attn_mask": attn_mask},
+ "attn_mask": attn_mask,
+ "transformer_options": transformer_options},
{"original_block": block_wrap})
txt = out["txt"]
img = out["img"]
@@ -208,7 +209,8 @@ class Chroma(nn.Module):
txt=txt,
vec=double_mod,
pe=pe,
- attn_mask=attn_mask)
+ attn_mask=attn_mask,
+ transformer_options=transformer_options)
if control is not None: # Controlnet
control_i = control.get("input")
@@ -228,17 +230,19 @@ class Chroma(nn.Module):
out["img"] = block(args["img"],
vec=args["vec"],
pe=args["pe"],
- attn_mask=args.get("attn_mask"))
+ attn_mask=args.get("attn_mask"),
+ transformer_options=args.get("transformer_options"))
return out
out = blocks_replace[("single_block", i)]({"img": img,
"vec": single_mod,
"pe": pe,
- "attn_mask": attn_mask},
+ "attn_mask": attn_mask,
+ "transformer_options": transformer_options},
{"original_block": block_wrap})
img = out["img"]
else:
- img = block(img, vec=single_mod, pe=pe, attn_mask=attn_mask)
+ img = block(img, vec=single_mod, pe=pe, attn_mask=attn_mask, transformer_options=transformer_options)
if control is not None: # Controlnet
control_o = control.get("output")
@@ -248,19 +252,29 @@ class Chroma(nn.Module):
img[:, txt.shape[1] :, ...] += add
img = img[:, txt.shape[1] :, ...]
- final_mod = self.get_modulations(mod_vectors, "final")
- img = self.final_layer(img, vec=final_mod) # (N, T, patch_size ** 2 * out_channels)
+ if hasattr(self, "final_layer"):
+ final_mod = self.get_modulations(mod_vectors, "final")
+ img = self.final_layer(img, vec=final_mod) # (N, T, patch_size ** 2 * out_channels)
return img
def forward(self, x, timestep, context, guidance, control=None, transformer_options={}, **kwargs):
+ return comfy.patcher_extension.WrapperExecutor.new_class_executor(
+ self._forward,
+ self,
+ comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options)
+ ).execute(x, timestep, context, guidance, control, transformer_options, **kwargs)
+
+ def _forward(self, x, timestep, context, guidance, control=None, transformer_options={}, **kwargs):
bs, c, h, w = x.shape
- patch_size = 2
- x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch_size, patch_size))
+ x = comfy.ldm.common_dit.pad_to_patch_size(x, (self.patch_size, self.patch_size))
- img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch_size, pw=patch_size)
+ img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=self.patch_size, pw=self.patch_size)
- h_len = ((h + (patch_size // 2)) // patch_size)
- w_len = ((w + (patch_size // 2)) // patch_size)
+ if img.ndim != 3 or context.ndim != 3:
+ raise ValueError("Input img and txt tensors must have 3 dimensions.")
+
+ h_len = ((h + (self.patch_size // 2)) // self.patch_size)
+ w_len = ((w + (self.patch_size // 2)) // self.patch_size)
img_ids = torch.zeros((h_len, w_len, 3), device=x.device, dtype=x.dtype)
img_ids[:, :, 1] = img_ids[:, :, 1] + torch.linspace(0, h_len - 1, steps=h_len, device=x.device, dtype=x.dtype).unsqueeze(1)
img_ids[:, :, 2] = img_ids[:, :, 2] + torch.linspace(0, w_len - 1, steps=w_len, device=x.device, dtype=x.dtype).unsqueeze(0)
@@ -268,4 +282,4 @@ class Chroma(nn.Module):
txt_ids = torch.zeros((bs, context.shape[1], 3), device=x.device, dtype=x.dtype)
out = self.forward_orig(img, img_ids, context, txt_ids, timestep, guidance, control, transformer_options, attn_mask=kwargs.get("attention_mask", None))
- return rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=h_len, w=w_len, ph=2, pw=2)[:,:,:h,:w]
+ return rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=h_len, w=w_len, ph=self.patch_size, pw=self.patch_size)[:,:,:h,:w]
diff --git a/comfy/ldm/chroma_radiance/layers.py b/comfy/ldm/chroma_radiance/layers.py
new file mode 100644
index 000000000..3c7bc9b6b
--- /dev/null
+++ b/comfy/ldm/chroma_radiance/layers.py
@@ -0,0 +1,206 @@
+# Adapted from https://github.com/lodestone-rock/flow
+from functools import lru_cache
+
+import torch
+from torch import nn
+
+from comfy.ldm.flux.layers import RMSNorm
+
+
+class NerfEmbedder(nn.Module):
+ """
+ An embedder module that combines input features with a 2D positional
+ encoding that mimics the Discrete Cosine Transform (DCT).
+
+ This module takes an input tensor of shape (B, P^2, C), where P is the
+ patch size, and enriches it with positional information before projecting
+ it to a new hidden size.
+ """
+ def __init__(
+ self,
+ in_channels: int,
+ hidden_size_input: int,
+ max_freqs: int,
+ dtype=None,
+ device=None,
+ operations=None,
+ ):
+ """
+ Initializes the NerfEmbedder.
+
+ Args:
+ in_channels (int): The number of channels in the input tensor.
+ hidden_size_input (int): The desired dimension of the output embedding.
+ max_freqs (int): The number of frequency components to use for both
+ the x and y dimensions of the positional encoding.
+ The total number of positional features will be max_freqs^2.
+ """
+ super().__init__()
+ self.dtype = dtype
+ self.max_freqs = max_freqs
+ self.hidden_size_input = hidden_size_input
+
+ # A linear layer to project the concatenated input features and
+ # positional encodings to the final output dimension.
+ self.embedder = nn.Sequential(
+ operations.Linear(in_channels + max_freqs**2, hidden_size_input, dtype=dtype, device=device)
+ )
+
+ @lru_cache(maxsize=4)
+ def fetch_pos(self, patch_size: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+ """
+ Generates and caches 2D DCT-like positional embeddings for a given patch size.
+
+ The LRU cache is a performance optimization that avoids recomputing the
+ same positional grid on every forward pass.
+
+ Args:
+ patch_size (int): The side length of the square input patch.
+ device: The torch device to create the tensors on.
+ dtype: The torch dtype for the tensors.
+
+ Returns:
+ A tensor of shape (1, patch_size^2, max_freqs^2) containing the
+ positional embeddings.
+ """
+ # Create normalized 1D coordinate grids from 0 to 1.
+ pos_x = torch.linspace(0, 1, patch_size, device=device, dtype=dtype)
+ pos_y = torch.linspace(0, 1, patch_size, device=device, dtype=dtype)
+
+ # Create a 2D meshgrid of coordinates.
+ pos_y, pos_x = torch.meshgrid(pos_y, pos_x, indexing="ij")
+
+ # Reshape positions to be broadcastable with frequencies.
+ # Shape becomes (patch_size^2, 1, 1).
+ pos_x = pos_x.reshape(-1, 1, 1)
+ pos_y = pos_y.reshape(-1, 1, 1)
+
+ # Create a 1D tensor of frequency values from 0 to max_freqs-1.
+ freqs = torch.linspace(0, self.max_freqs - 1, self.max_freqs, dtype=dtype, device=device)
+
+ # Reshape frequencies to be broadcastable for creating 2D basis functions.
+ # freqs_x shape: (1, max_freqs, 1)
+ # freqs_y shape: (1, 1, max_freqs)
+ freqs_x = freqs[None, :, None]
+ freqs_y = freqs[None, None, :]
+
+ # A custom weighting coefficient, not part of standard DCT.
+ # This seems to down-weight the contribution of higher-frequency interactions.
+ coeffs = (1 + freqs_x * freqs_y) ** -1
+
+ # Calculate the 1D cosine basis functions for x and y coordinates.
+ # This is the core of the DCT formulation.
+ dct_x = torch.cos(pos_x * freqs_x * torch.pi)
+ dct_y = torch.cos(pos_y * freqs_y * torch.pi)
+
+ # Combine the 1D basis functions to create 2D basis functions by element-wise
+ # multiplication, and apply the custom coefficients. Broadcasting handles the
+ # combination of all (pos_x, freqs_x) with all (pos_y, freqs_y).
+ # The result is flattened into a feature vector for each position.
+ dct = (dct_x * dct_y * coeffs).view(1, -1, self.max_freqs ** 2)
+
+ return dct
+
+ def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+ """
+ Forward pass for the embedder.
+
+ Args:
+ inputs (Tensor): The input tensor of shape (B, P^2, C).
+
+ Returns:
+ Tensor: The output tensor of shape (B, P^2, hidden_size_input).
+ """
+ # Get the batch size, number of pixels, and number of channels.
+ B, P2, C = inputs.shape
+
+ # Infer the patch side length from the number of pixels (P^2).
+ patch_size = int(P2 ** 0.5)
+
+ input_dtype = inputs.dtype
+ inputs = inputs.to(dtype=self.dtype)
+
+ # Fetch the pre-computed or cached positional embeddings.
+ dct = self.fetch_pos(patch_size, inputs.device, self.dtype)
+
+ # Repeat the positional embeddings for each item in the batch.
+ dct = dct.repeat(B, 1, 1)
+
+ # Concatenate the original input features with the positional embeddings
+ # along the feature dimension.
+ inputs = torch.cat((inputs, dct), dim=-1)
+
+ # Project the combined tensor to the target hidden size.
+ return self.embedder(inputs).to(dtype=input_dtype)
+
+
+class NerfGLUBlock(nn.Module):
+ """
+ A NerfBlock using a Gated Linear Unit (GLU) like MLP.
+ """
+ def __init__(self, hidden_size_s: int, hidden_size_x: int, mlp_ratio, dtype=None, device=None, operations=None):
+ super().__init__()
+ # The total number of parameters for the MLP is increased to accommodate
+ # the gate, value, and output projection matrices.
+ # We now need to generate parameters for 3 matrices.
+ total_params = 3 * hidden_size_x**2 * mlp_ratio
+ self.param_generator = operations.Linear(hidden_size_s, total_params, dtype=dtype, device=device)
+ self.norm = RMSNorm(hidden_size_x, dtype=dtype, device=device, operations=operations)
+ self.mlp_ratio = mlp_ratio
+
+
+ def forward(self, x: torch.Tensor, s: torch.Tensor) -> torch.Tensor:
+ batch_size, num_x, hidden_size_x = x.shape
+ mlp_params = self.param_generator(s)
+
+ # Split the generated parameters into three parts for the gate, value, and output projection.
+ fc1_gate_params, fc1_value_params, fc2_params = mlp_params.chunk(3, dim=-1)
+
+ # Reshape the parameters into matrices for batch matrix multiplication.
+ fc1_gate = fc1_gate_params.view(batch_size, hidden_size_x, hidden_size_x * self.mlp_ratio)
+ fc1_value = fc1_value_params.view(batch_size, hidden_size_x, hidden_size_x * self.mlp_ratio)
+ fc2 = fc2_params.view(batch_size, hidden_size_x * self.mlp_ratio, hidden_size_x)
+
+ # Normalize the generated weight matrices as in the original implementation.
+ fc1_gate = torch.nn.functional.normalize(fc1_gate, dim=-2)
+ fc1_value = torch.nn.functional.normalize(fc1_value, dim=-2)
+ fc2 = torch.nn.functional.normalize(fc2, dim=-2)
+
+ res_x = x
+ x = self.norm(x)
+
+ # Apply the final output projection.
+ x = torch.bmm(torch.nn.functional.silu(torch.bmm(x, fc1_gate)) * torch.bmm(x, fc1_value), fc2)
+
+ return x + res_x
+
+
+class NerfFinalLayer(nn.Module):
+ def __init__(self, hidden_size, out_channels, dtype=None, device=None, operations=None):
+ super().__init__()
+ self.norm = RMSNorm(hidden_size, dtype=dtype, device=device, operations=operations)
+ self.linear = operations.Linear(hidden_size, out_channels, dtype=dtype, device=device)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ # RMSNorm normalizes over the last dimension, but our channel dim (C) is at dim=1.
+ # So we temporarily move the channel dimension to the end for the norm operation.
+ return self.linear(self.norm(x.movedim(1, -1))).movedim(-1, 1)
+
+
+class NerfFinalLayerConv(nn.Module):
+ def __init__(self, hidden_size: int, out_channels: int, dtype=None, device=None, operations=None):
+ super().__init__()
+ self.norm = RMSNorm(hidden_size, dtype=dtype, device=device, operations=operations)
+ self.conv = operations.Conv2d(
+ in_channels=hidden_size,
+ out_channels=out_channels,
+ kernel_size=3,
+ padding=1,
+ dtype=dtype,
+ device=device,
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ # RMSNorm normalizes over the last dimension, but our channel dim (C) is at dim=1.
+ # So we temporarily move the channel dimension to the end for the norm operation.
+ return self.conv(self.norm(x.movedim(1, -1)).movedim(-1, 1))
diff --git a/comfy/ldm/chroma_radiance/model.py b/comfy/ldm/chroma_radiance/model.py
new file mode 100644
index 000000000..7d7be80f5
--- /dev/null
+++ b/comfy/ldm/chroma_radiance/model.py
@@ -0,0 +1,320 @@
+# Credits:
+# Original Flux code can be found on: https://github.com/black-forest-labs/flux
+# Chroma Radiance adaption referenced from https://github.com/lodestone-rock/flow
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+from torch import Tensor, nn
+from einops import repeat
+import comfy.ldm.common_dit
+
+from comfy.ldm.flux.layers import EmbedND
+
+from comfy.ldm.chroma.model import Chroma, ChromaParams
+from comfy.ldm.chroma.layers import (
+ DoubleStreamBlock,
+ SingleStreamBlock,
+ Approximator,
+)
+from .layers import (
+ NerfEmbedder,
+ NerfGLUBlock,
+ NerfFinalLayer,
+ NerfFinalLayerConv,
+)
+
+
+@dataclass
+class ChromaRadianceParams(ChromaParams):
+ patch_size: int
+ nerf_hidden_size: int
+ nerf_mlp_ratio: int
+ nerf_depth: int
+ nerf_max_freqs: int
+ # Setting nerf_tile_size to 0 disables tiling.
+ nerf_tile_size: int
+ # Currently one of linear (legacy) or conv.
+ nerf_final_head_type: str
+ # None means use the same dtype as the model.
+ nerf_embedder_dtype: Optional[torch.dtype]
+
+
+class ChromaRadiance(Chroma):
+ """
+ Transformer model for flow matching on sequences.
+ """
+
+ def __init__(self, image_model=None, final_layer=True, dtype=None, device=None, operations=None, **kwargs):
+ if operations is None:
+ raise RuntimeError("Attempt to create ChromaRadiance object without setting operations")
+ nn.Module.__init__(self)
+ self.dtype = dtype
+ params = ChromaRadianceParams(**kwargs)
+ self.params = params
+ self.patch_size = params.patch_size
+ self.in_channels = params.in_channels
+ self.out_channels = params.out_channels
+ if params.hidden_size % params.num_heads != 0:
+ raise ValueError(
+ f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}"
+ )
+ pe_dim = params.hidden_size // params.num_heads
+ if sum(params.axes_dim) != pe_dim:
+ raise ValueError(f"Got {params.axes_dim} but expected positional dim {pe_dim}")
+ self.hidden_size = params.hidden_size
+ self.num_heads = params.num_heads
+ self.in_dim = params.in_dim
+ self.out_dim = params.out_dim
+ self.hidden_dim = params.hidden_dim
+ self.n_layers = params.n_layers
+ self.pe_embedder = EmbedND(dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim)
+ self.img_in_patch = operations.Conv2d(
+ params.in_channels,
+ params.hidden_size,
+ kernel_size=params.patch_size,
+ stride=params.patch_size,
+ bias=True,
+ dtype=dtype,
+ device=device,
+ )
+ self.txt_in = operations.Linear(params.context_in_dim, self.hidden_size, dtype=dtype, device=device)
+ # set as nn identity for now, will overwrite it later.
+ self.distilled_guidance_layer = Approximator(
+ in_dim=self.in_dim,
+ hidden_dim=self.hidden_dim,
+ out_dim=self.out_dim,
+ n_layers=self.n_layers,
+ dtype=dtype, device=device, operations=operations
+ )
+
+
+ self.double_blocks = nn.ModuleList(
+ [
+ DoubleStreamBlock(
+ self.hidden_size,
+ self.num_heads,
+ mlp_ratio=params.mlp_ratio,
+ qkv_bias=params.qkv_bias,
+ dtype=dtype, device=device, operations=operations
+ )
+ for _ in range(params.depth)
+ ]
+ )
+
+ self.single_blocks = nn.ModuleList(
+ [
+ SingleStreamBlock(
+ self.hidden_size,
+ self.num_heads,
+ mlp_ratio=params.mlp_ratio,
+ dtype=dtype, device=device, operations=operations,
+ )
+ for _ in range(params.depth_single_blocks)
+ ]
+ )
+
+ # pixel channel concat with DCT
+ self.nerf_image_embedder = NerfEmbedder(
+ in_channels=params.in_channels,
+ hidden_size_input=params.nerf_hidden_size,
+ max_freqs=params.nerf_max_freqs,
+ dtype=params.nerf_embedder_dtype or dtype,
+ device=device,
+ operations=operations,
+ )
+
+ self.nerf_blocks = nn.ModuleList([
+ NerfGLUBlock(
+ hidden_size_s=params.hidden_size,
+ hidden_size_x=params.nerf_hidden_size,
+ mlp_ratio=params.nerf_mlp_ratio,
+ dtype=dtype,
+ device=device,
+ operations=operations,
+ ) for _ in range(params.nerf_depth)
+ ])
+
+ if params.nerf_final_head_type == "linear":
+ self.nerf_final_layer = NerfFinalLayer(
+ params.nerf_hidden_size,
+ out_channels=params.in_channels,
+ dtype=dtype,
+ device=device,
+ operations=operations,
+ )
+ elif params.nerf_final_head_type == "conv":
+ self.nerf_final_layer_conv = NerfFinalLayerConv(
+ params.nerf_hidden_size,
+ out_channels=params.in_channels,
+ dtype=dtype,
+ device=device,
+ operations=operations,
+ )
+ else:
+ errstr = f"Unsupported nerf_final_head_type {params.nerf_final_head_type}"
+ raise ValueError(errstr)
+
+ self.skip_mmdit = []
+ self.skip_dit = []
+ self.lite = False
+
+ @property
+ def _nerf_final_layer(self) -> nn.Module:
+ if self.params.nerf_final_head_type == "linear":
+ return self.nerf_final_layer
+ if self.params.nerf_final_head_type == "conv":
+ return self.nerf_final_layer_conv
+ # Impossible to get here as we raise an error on unexpected types on initialization.
+ raise NotImplementedError
+
+ def img_in(self, img: Tensor) -> Tensor:
+ img = self.img_in_patch(img) # -> [B, Hidden, H/P, W/P]
+ # flatten into a sequence for the transformer.
+ return img.flatten(2).transpose(1, 2) # -> [B, NumPatches, Hidden]
+
+ def forward_nerf(
+ self,
+ img_orig: Tensor,
+ img_out: Tensor,
+ params: ChromaRadianceParams,
+ ) -> Tensor:
+ B, C, H, W = img_orig.shape
+ num_patches = img_out.shape[1]
+ patch_size = params.patch_size
+
+ # Store the raw pixel values of each patch for the NeRF head later.
+ # unfold creates patches: [B, C * P * P, NumPatches]
+ nerf_pixels = nn.functional.unfold(img_orig, kernel_size=patch_size, stride=patch_size)
+ nerf_pixels = nerf_pixels.transpose(1, 2) # -> [B, NumPatches, C * P * P]
+
+ # Reshape for per-patch processing
+ nerf_hidden = img_out.reshape(B * num_patches, params.hidden_size)
+ nerf_pixels = nerf_pixels.reshape(B * num_patches, C, patch_size**2).transpose(1, 2)
+
+ if params.nerf_tile_size > 0 and num_patches > params.nerf_tile_size:
+ # Enable tiling if nerf_tile_size isn't 0 and we actually have more patches than
+ # the tile size.
+ img_dct = self.forward_tiled_nerf(nerf_hidden, nerf_pixels, B, C, num_patches, patch_size, params)
+ else:
+ # Get DCT-encoded pixel embeddings [pixel-dct]
+ img_dct = self.nerf_image_embedder(nerf_pixels)
+
+ # Pass through the dynamic MLP blocks (the NeRF)
+ for block in self.nerf_blocks:
+ img_dct = block(img_dct, nerf_hidden)
+
+ # Reassemble the patches into the final image.
+ img_dct = img_dct.transpose(1, 2) # -> [B*NumPatches, C, P*P]
+ # Reshape to combine with batch dimension for fold
+ img_dct = img_dct.reshape(B, num_patches, -1) # -> [B, NumPatches, C*P*P]
+ img_dct = img_dct.transpose(1, 2) # -> [B, C*P*P, NumPatches]
+ img_dct = nn.functional.fold(
+ img_dct,
+ output_size=(H, W),
+ kernel_size=patch_size,
+ stride=patch_size,
+ )
+ return self._nerf_final_layer(img_dct)
+
+ def forward_tiled_nerf(
+ self,
+ nerf_hidden: Tensor,
+ nerf_pixels: Tensor,
+ batch: int,
+ channels: int,
+ num_patches: int,
+ patch_size: int,
+ params: ChromaRadianceParams,
+ ) -> Tensor:
+ """
+ Processes the NeRF head in tiles to save memory.
+ nerf_hidden has shape [B, L, D]
+ nerf_pixels has shape [B, L, C * P * P]
+ """
+ tile_size = params.nerf_tile_size
+ output_tiles = []
+ # Iterate over the patches in tiles. The dimension L (num_patches) is at index 1.
+ for i in range(0, num_patches, tile_size):
+ end = min(i + tile_size, num_patches)
+
+ # Slice the current tile from the input tensors
+ nerf_hidden_tile = nerf_hidden[i * batch:end * batch]
+ nerf_pixels_tile = nerf_pixels[i * batch:end * batch]
+
+ # get DCT-encoded pixel embeddings [pixel-dct]
+ img_dct_tile = self.nerf_image_embedder(nerf_pixels_tile)
+
+ # pass through the dynamic MLP blocks (the NeRF)
+ for block in self.nerf_blocks:
+ img_dct_tile = block(img_dct_tile, nerf_hidden_tile)
+
+ output_tiles.append(img_dct_tile)
+
+ # Concatenate the processed tiles along the patch dimension
+ return torch.cat(output_tiles, dim=0)
+
+ def radiance_get_override_params(self, overrides: dict) -> ChromaRadianceParams:
+ params = self.params
+ if not overrides:
+ return params
+ params_dict = {k: getattr(params, k) for k in params.__dataclass_fields__}
+ nullable_keys = frozenset(("nerf_embedder_dtype",))
+ bad_keys = tuple(k for k in overrides if k not in params_dict)
+ if bad_keys:
+ e = f"Unknown key(s) in transformer_options chroma_radiance_options: {', '.join(bad_keys)}"
+ raise ValueError(e)
+ bad_keys = tuple(
+ k
+ for k, v in overrides.items()
+ if type(v) != type(getattr(params, k)) and (v is not None or k not in nullable_keys)
+ )
+ if bad_keys:
+ e = f"Invalid value(s) in transformer_options chroma_radiance_options: {', '.join(bad_keys)}"
+ raise ValueError(e)
+ # At this point it's all valid keys and values so we can merge with the existing params.
+ params_dict |= overrides
+ return params.__class__(**params_dict)
+
+ def _forward(
+ self,
+ x: Tensor,
+ timestep: Tensor,
+ context: Tensor,
+ guidance: Optional[Tensor],
+ control: Optional[dict]=None,
+ transformer_options: dict={},
+ **kwargs: dict,
+ ) -> Tensor:
+ bs, c, h, w = x.shape
+ img = comfy.ldm.common_dit.pad_to_patch_size(x, (self.patch_size, self.patch_size))
+
+ if img.ndim != 4:
+ raise ValueError("Input img tensor must be in [B, C, H, W] format.")
+ if context.ndim != 3:
+ raise ValueError("Input txt tensors must have 3 dimensions.")
+
+ params = self.radiance_get_override_params(transformer_options.get("chroma_radiance_options", {}))
+
+ h_len = (img.shape[-2] // self.patch_size)
+ w_len = (img.shape[-1] // self.patch_size)
+
+ img_ids = torch.zeros((h_len, w_len, 3), device=x.device, dtype=x.dtype)
+ img_ids[:, :, 1] = img_ids[:, :, 1] + torch.linspace(0, h_len - 1, steps=h_len, device=x.device, dtype=x.dtype).unsqueeze(1)
+ img_ids[:, :, 2] = img_ids[:, :, 2] + torch.linspace(0, w_len - 1, steps=w_len, device=x.device, dtype=x.dtype).unsqueeze(0)
+ img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
+ txt_ids = torch.zeros((bs, context.shape[1], 3), device=x.device, dtype=x.dtype)
+
+ img_out = self.forward_orig(
+ img,
+ img_ids,
+ context,
+ txt_ids,
+ timestep,
+ guidance,
+ control,
+ transformer_options,
+ attn_mask=kwargs.get("attention_mask", None),
+ )
+ return self.forward_nerf(img, img_out, params)[:, :, :h, :w]
diff --git a/comfy/ldm/cosmos/blocks.py b/comfy/ldm/cosmos/blocks.py
index 5c4356a3f..afb43d469 100644
--- a/comfy/ldm/cosmos/blocks.py
+++ b/comfy/ldm/cosmos/blocks.py
@@ -176,6 +176,7 @@ class Attention(nn.Module):
context=None,
mask=None,
rope_emb=None,
+ transformer_options={},
**kwargs,
):
"""
@@ -184,7 +185,7 @@ class Attention(nn.Module):
context (Optional[Tensor]): The key tensor of shape [B, Mk, K] or use x as context [self attention] if None
"""
q, k, v = self.cal_qkv(x, context, mask, rope_emb=rope_emb, **kwargs)
- out = optimized_attention(q, k, v, self.heads, skip_reshape=True, mask=mask, skip_output_reshape=True)
+ out = optimized_attention(q, k, v, self.heads, skip_reshape=True, mask=mask, skip_output_reshape=True, transformer_options=transformer_options)
del q, k, v
out = rearrange(out, " b n s c -> s b (n c)")
return self.to_out(out)
@@ -546,6 +547,7 @@ class VideoAttn(nn.Module):
context: Optional[torch.Tensor] = None,
crossattn_mask: Optional[torch.Tensor] = None,
rope_emb_L_1_1_D: Optional[torch.Tensor] = None,
+ transformer_options: Optional[dict] = {},
) -> torch.Tensor:
"""
Forward pass for video attention.
@@ -571,6 +573,7 @@ class VideoAttn(nn.Module):
context_M_B_D,
crossattn_mask,
rope_emb=rope_emb_L_1_1_D,
+ transformer_options=transformer_options,
)
x_T_H_W_B_D = rearrange(x_THW_B_D, "(t h w) b d -> t h w b d", h=H, w=W)
return x_T_H_W_B_D
@@ -665,6 +668,7 @@ class DITBuildingBlock(nn.Module):
crossattn_mask: Optional[torch.Tensor] = None,
rope_emb_L_1_1_D: Optional[torch.Tensor] = None,
adaln_lora_B_3D: Optional[torch.Tensor] = None,
+ transformer_options: Optional[dict] = {},
) -> torch.Tensor:
"""
Forward pass for dynamically configured blocks with adaptive normalization.
@@ -702,6 +706,7 @@ class DITBuildingBlock(nn.Module):
adaln_norm_state(self.norm_state, x, scale_1_1_1_B_D, shift_1_1_1_B_D),
context=None,
rope_emb_L_1_1_D=rope_emb_L_1_1_D,
+ transformer_options=transformer_options,
)
elif self.block_type in ["cross_attn", "ca"]:
x = x + gate_1_1_1_B_D * self.block(
@@ -709,6 +714,7 @@ class DITBuildingBlock(nn.Module):
context=crossattn_emb,
crossattn_mask=crossattn_mask,
rope_emb_L_1_1_D=rope_emb_L_1_1_D,
+ transformer_options=transformer_options,
)
else:
raise ValueError(f"Unknown block type: {self.block_type}")
@@ -784,6 +790,7 @@ class GeneralDITTransformerBlock(nn.Module):
crossattn_mask: Optional[torch.Tensor] = None,
rope_emb_L_1_1_D: Optional[torch.Tensor] = None,
adaln_lora_B_3D: Optional[torch.Tensor] = None,
+ transformer_options: Optional[dict] = {},
) -> torch.Tensor:
for block in self.blocks:
x = block(
@@ -793,5 +800,6 @@ class GeneralDITTransformerBlock(nn.Module):
crossattn_mask,
rope_emb_L_1_1_D=rope_emb_L_1_1_D,
adaln_lora_B_3D=adaln_lora_B_3D,
+ transformer_options=transformer_options,
)
return x
diff --git a/comfy/ldm/cosmos/cosmos_tokenizer/utils.py b/comfy/ldm/cosmos/cosmos_tokenizer/utils.py
index 3af8d0d05..ca993006f 100644
--- a/comfy/ldm/cosmos/cosmos_tokenizer/utils.py
+++ b/comfy/ldm/cosmos/cosmos_tokenizer/utils.py
@@ -58,7 +58,8 @@ def is_odd(n: int) -> bool:
def nonlinearity(x):
- return x * torch.sigmoid(x)
+ # x * sigmoid(x)
+ return torch.nn.functional.silu(x)
def Normalize(in_channels, num_groups=32):
diff --git a/comfy/ldm/cosmos/model.py b/comfy/ldm/cosmos/model.py
index 4836e0b69..52ef7ef43 100644
--- a/comfy/ldm/cosmos/model.py
+++ b/comfy/ldm/cosmos/model.py
@@ -27,6 +27,8 @@ from torchvision import transforms
from enum import Enum
import logging
+import comfy.patcher_extension
+
from .blocks import (
FinalLayer,
GeneralDITTransformerBlock,
@@ -435,6 +437,42 @@ class GeneralDIT(nn.Module):
latent_condition_sigma: Optional[torch.Tensor] = None,
condition_video_augment_sigma: Optional[torch.Tensor] = None,
**kwargs,
+ ):
+ return comfy.patcher_extension.WrapperExecutor.new_class_executor(
+ self._forward,
+ self,
+ comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, kwargs.get("transformer_options", {}))
+ ).execute(x,
+ timesteps,
+ context,
+ attention_mask,
+ fps,
+ image_size,
+ padding_mask,
+ scalar_feature,
+ data_type,
+ latent_condition,
+ latent_condition_sigma,
+ condition_video_augment_sigma,
+ **kwargs)
+
+ def _forward(
+ self,
+ x: torch.Tensor,
+ timesteps: torch.Tensor,
+ context: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ # crossattn_emb: torch.Tensor,
+ # crossattn_mask: Optional[torch.Tensor] = None,
+ fps: Optional[torch.Tensor] = None,
+ image_size: Optional[torch.Tensor] = None,
+ padding_mask: Optional[torch.Tensor] = None,
+ scalar_feature: Optional[torch.Tensor] = None,
+ data_type: Optional[DataType] = DataType.VIDEO,
+ latent_condition: Optional[torch.Tensor] = None,
+ latent_condition_sigma: Optional[torch.Tensor] = None,
+ condition_video_augment_sigma: Optional[torch.Tensor] = None,
+ **kwargs,
):
"""
Args:
@@ -482,6 +520,7 @@ class GeneralDIT(nn.Module):
x.shape == extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D.shape
), f"{x.shape} != {extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D.shape} {original_shape}"
+ transformer_options = kwargs.get("transformer_options", {})
for _, block in self.blocks.items():
assert (
self.blocks["block0"].x_format == block.x_format
@@ -496,6 +535,7 @@ class GeneralDIT(nn.Module):
crossattn_mask,
rope_emb_L_1_1_D=rope_emb_L_1_1_D,
adaln_lora_B_3D=adaln_lora_B_3D,
+ transformer_options=transformer_options,
)
x_B_T_H_W_D = rearrange(x, "T H W B D -> B T H W D")
diff --git a/comfy/ldm/cosmos/predict2.py b/comfy/ldm/cosmos/predict2.py
index 316117f77..07a4fc79f 100644
--- a/comfy/ldm/cosmos/predict2.py
+++ b/comfy/ldm/cosmos/predict2.py
@@ -11,6 +11,7 @@ import math
from .position_embedding import VideoRopePosition3DEmb, LearnablePosEmbAxis
from torchvision import transforms
+import comfy.patcher_extension
from comfy.ldm.modules.attention import optimized_attention
def apply_rotary_pos_emb(
@@ -43,7 +44,7 @@ class GPT2FeedForward(nn.Module):
return x
-def torch_attention_op(q_B_S_H_D: torch.Tensor, k_B_S_H_D: torch.Tensor, v_B_S_H_D: torch.Tensor) -> torch.Tensor:
+def torch_attention_op(q_B_S_H_D: torch.Tensor, k_B_S_H_D: torch.Tensor, v_B_S_H_D: torch.Tensor, transformer_options: Optional[dict] = {}) -> torch.Tensor:
"""Computes multi-head attention using PyTorch's native implementation.
This function provides a PyTorch backend alternative to Transformer Engine's attention operation.
@@ -70,7 +71,7 @@ def torch_attention_op(q_B_S_H_D: torch.Tensor, k_B_S_H_D: torch.Tensor, v_B_S_H
q_B_H_S_D = rearrange(q_B_S_H_D, "b ... h k -> b h ... k").view(in_q_shape[0], in_q_shape[-2], -1, in_q_shape[-1])
k_B_H_S_D = rearrange(k_B_S_H_D, "b ... h v -> b h ... v").view(in_k_shape[0], in_k_shape[-2], -1, in_k_shape[-1])
v_B_H_S_D = rearrange(v_B_S_H_D, "b ... h v -> b h ... v").view(in_k_shape[0], in_k_shape[-2], -1, in_k_shape[-1])
- return optimized_attention(q_B_H_S_D, k_B_H_S_D, v_B_H_S_D, in_q_shape[-2], skip_reshape=True)
+ return optimized_attention(q_B_H_S_D, k_B_H_S_D, v_B_H_S_D, in_q_shape[-2], skip_reshape=True, transformer_options=transformer_options)
class Attention(nn.Module):
@@ -179,8 +180,8 @@ class Attention(nn.Module):
return q, k, v
- def compute_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
- result = self.attn_op(q, k, v) # [B, S, H, D]
+ def compute_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, transformer_options: Optional[dict] = {}) -> torch.Tensor:
+ result = self.attn_op(q, k, v, transformer_options=transformer_options) # [B, S, H, D]
return self.output_dropout(self.output_proj(result))
def forward(
@@ -188,6 +189,7 @@ class Attention(nn.Module):
x: torch.Tensor,
context: Optional[torch.Tensor] = None,
rope_emb: Optional[torch.Tensor] = None,
+ transformer_options: Optional[dict] = {},
) -> torch.Tensor:
"""
Args:
@@ -195,7 +197,7 @@ class Attention(nn.Module):
context (Optional[Tensor]): The key tensor of shape [B, Mk, K] or use x as context [self attention] if None
"""
q, k, v = self.compute_qkv(x, context, rope_emb=rope_emb)
- return self.compute_attention(q, k, v)
+ return self.compute_attention(q, k, v, transformer_options=transformer_options)
class Timesteps(nn.Module):
@@ -458,6 +460,7 @@ class Block(nn.Module):
rope_emb_L_1_1_D: Optional[torch.Tensor] = None,
adaln_lora_B_T_3D: Optional[torch.Tensor] = None,
extra_per_block_pos_emb: Optional[torch.Tensor] = None,
+ transformer_options: Optional[dict] = {},
) -> torch.Tensor:
if extra_per_block_pos_emb is not None:
x_B_T_H_W_D = x_B_T_H_W_D + extra_per_block_pos_emb
@@ -511,6 +514,7 @@ class Block(nn.Module):
rearrange(normalized_x_B_T_H_W_D, "b t h w d -> b (t h w) d"),
None,
rope_emb=rope_emb_L_1_1_D,
+ transformer_options=transformer_options,
),
"b (t h w) d -> b t h w d",
t=T,
@@ -524,6 +528,7 @@ class Block(nn.Module):
layer_norm_cross_attn: Callable,
_scale_cross_attn_B_T_1_1_D: torch.Tensor,
_shift_cross_attn_B_T_1_1_D: torch.Tensor,
+ transformer_options: Optional[dict] = {},
) -> torch.Tensor:
_normalized_x_B_T_H_W_D = _fn(
_x_B_T_H_W_D, layer_norm_cross_attn, _scale_cross_attn_B_T_1_1_D, _shift_cross_attn_B_T_1_1_D
@@ -533,6 +538,7 @@ class Block(nn.Module):
rearrange(_normalized_x_B_T_H_W_D, "b t h w d -> b (t h w) d"),
crossattn_emb,
rope_emb=rope_emb_L_1_1_D,
+ transformer_options=transformer_options,
),
"b (t h w) d -> b t h w d",
t=T,
@@ -546,6 +552,7 @@ class Block(nn.Module):
self.layer_norm_cross_attn,
scale_cross_attn_B_T_1_1_D,
shift_cross_attn_B_T_1_1_D,
+ transformer_options=transformer_options,
)
x_B_T_H_W_D = result_B_T_H_W_D * gate_cross_attn_B_T_1_1_D + x_B_T_H_W_D
@@ -805,7 +812,21 @@ class MiniTrainDIT(nn.Module):
)
return x_B_C_Tt_Hp_Wp
- def forward(
+ def forward(self,
+ x: torch.Tensor,
+ timesteps: torch.Tensor,
+ context: torch.Tensor,
+ fps: Optional[torch.Tensor] = None,
+ padding_mask: Optional[torch.Tensor] = None,
+ **kwargs,
+ ):
+ return comfy.patcher_extension.WrapperExecutor.new_class_executor(
+ self._forward,
+ self,
+ comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, kwargs.get("transformer_options", {}))
+ ).execute(x, timesteps, context, fps, padding_mask, **kwargs)
+
+ def _forward(
self,
x: torch.Tensor,
timesteps: torch.Tensor,
@@ -850,6 +871,7 @@ class MiniTrainDIT(nn.Module):
"rope_emb_L_1_1_D": rope_emb_L_1_1_D.unsqueeze(1).unsqueeze(0),
"adaln_lora_B_T_3D": adaln_lora_B_T_3D,
"extra_per_block_pos_emb": extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D,
+ "transformer_options": kwargs.get("transformer_options", {}),
}
for block in self.blocks:
x_B_T_H_W_D = block(
diff --git a/comfy/ldm/flux/layers.py b/comfy/ldm/flux/layers.py
index 113eb2096..ef21b416b 100644
--- a/comfy/ldm/flux/layers.py
+++ b/comfy/ldm/flux/layers.py
@@ -159,7 +159,7 @@ class DoubleStreamBlock(nn.Module):
)
self.flipped_img_txt = flipped_img_txt
- def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor, attn_mask=None, modulation_dims_img=None, modulation_dims_txt=None):
+ def forward(self, img: Tensor, txt: Tensor, vec: Tensor, pe: Tensor, attn_mask=None, modulation_dims_img=None, modulation_dims_txt=None, transformer_options={}):
img_mod1, img_mod2 = self.img_mod(vec)
txt_mod1, txt_mod2 = self.txt_mod(vec)
@@ -182,7 +182,7 @@ class DoubleStreamBlock(nn.Module):
attn = attention(torch.cat((img_q, txt_q), dim=2),
torch.cat((img_k, txt_k), dim=2),
torch.cat((img_v, txt_v), dim=2),
- pe=pe, mask=attn_mask)
+ pe=pe, mask=attn_mask, transformer_options=transformer_options)
img_attn, txt_attn = attn[:, : img.shape[1]], attn[:, img.shape[1]:]
else:
@@ -190,7 +190,7 @@ class DoubleStreamBlock(nn.Module):
attn = attention(torch.cat((txt_q, img_q), dim=2),
torch.cat((txt_k, img_k), dim=2),
torch.cat((txt_v, img_v), dim=2),
- pe=pe, mask=attn_mask)
+ pe=pe, mask=attn_mask, transformer_options=transformer_options)
txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1]:]
@@ -244,7 +244,7 @@ class SingleStreamBlock(nn.Module):
self.mlp_act = nn.GELU(approximate="tanh")
self.modulation = Modulation(hidden_size, double=False, dtype=dtype, device=device, operations=operations)
- def forward(self, x: Tensor, vec: Tensor, pe: Tensor, attn_mask=None, modulation_dims=None) -> Tensor:
+ def forward(self, x: Tensor, vec: Tensor, pe: Tensor, attn_mask=None, modulation_dims=None, transformer_options={}) -> Tensor:
mod, _ = self.modulation(vec)
qkv, mlp = torch.split(self.linear1(apply_mod(self.pre_norm(x), (1 + mod.scale), mod.shift, modulation_dims)), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1)
@@ -252,7 +252,7 @@ class SingleStreamBlock(nn.Module):
q, k = self.norm(q, k, v)
# compute attention
- attn = attention(q, k, v, pe=pe, mask=attn_mask)
+ attn = attention(q, k, v, pe=pe, mask=attn_mask, transformer_options=transformer_options)
# compute activation in mlp stream, cat again and run second linear layer
output = self.linear2(torch.cat((attn, self.mlp_act(mlp)), 2))
x += apply_mod(output, mod.gate, None, modulation_dims)
diff --git a/comfy/ldm/flux/math.py b/comfy/ldm/flux/math.py
index 3e0978176..8deda0d4a 100644
--- a/comfy/ldm/flux/math.py
+++ b/comfy/ldm/flux/math.py
@@ -6,7 +6,7 @@ from comfy.ldm.modules.attention import optimized_attention
import comfy.model_management
-def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask=None) -> Tensor:
+def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask=None, transformer_options={}) -> Tensor:
q_shape = q.shape
k_shape = k.shape
@@ -17,7 +17,7 @@ def attention(q: Tensor, k: Tensor, v: Tensor, pe: Tensor, mask=None) -> Tensor:
k = (pe[..., 0] * k[..., 0] + pe[..., 1] * k[..., 1]).reshape(*k_shape).type_as(v)
heads = q.shape[1]
- x = optimized_attention(q, k, v, heads, skip_reshape=True, mask=mask)
+ x = optimized_attention(q, k, v, heads, skip_reshape=True, mask=mask, transformer_options=transformer_options)
return x
@@ -35,11 +35,13 @@ def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
return out.to(dtype=torch.float32, device=pos.device)
+def apply_rope1(x: Tensor, freqs_cis: Tensor):
+ x_ = x.to(dtype=freqs_cis.dtype).reshape(*x.shape[:-1], -1, 1, 2)
+
+ x_out = freqs_cis[..., 0] * x_[..., 0]
+ x_out.addcmul_(freqs_cis[..., 1], x_[..., 1])
+
+ return x_out.reshape(*x.shape).type_as(x)
def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor):
- xq_ = xq.to(dtype=freqs_cis.dtype).reshape(*xq.shape[:-1], -1, 1, 2)
- xk_ = xk.to(dtype=freqs_cis.dtype).reshape(*xk.shape[:-1], -1, 1, 2)
- xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
- xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
- return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
-
+ return apply_rope1(xq, freqs_cis), apply_rope1(xk, freqs_cis)
diff --git a/comfy/ldm/flux/model.py b/comfy/ldm/flux/model.py
index 8f4d99f54..14f90cea5 100644
--- a/comfy/ldm/flux/model.py
+++ b/comfy/ldm/flux/model.py
@@ -6,6 +6,7 @@ import torch
from torch import Tensor, nn
from einops import rearrange, repeat
import comfy.ldm.common_dit
+import comfy.patcher_extension
from .layers import (
DoubleStreamBlock,
@@ -105,6 +106,7 @@ class Flux(nn.Module):
if y is None:
y = torch.zeros((img.shape[0], self.params.vec_in_dim), device=img.device, dtype=img.dtype)
+ patches = transformer_options.get("patches", {})
patches_replace = transformer_options.get("patches_replace", {})
if img.ndim != 3 or txt.ndim != 3:
raise ValueError("Input img and txt tensors must have 3 dimensions.")
@@ -116,9 +118,17 @@ class Flux(nn.Module):
if guidance is not None:
vec = vec + self.guidance_in(timestep_embedding(guidance, 256).to(img.dtype))
- vec = vec + self.vector_in(y[:,:self.params.vec_in_dim])
+ vec = vec + self.vector_in(y[:, :self.params.vec_in_dim])
txt = self.txt_in(txt)
+ if "post_input" in patches:
+ for p in patches["post_input"]:
+ out = p({"img": img, "txt": txt, "img_ids": img_ids, "txt_ids": txt_ids})
+ img = out["img"]
+ txt = out["txt"]
+ img_ids = out["img_ids"]
+ txt_ids = out["txt_ids"]
+
if img_ids is not None:
ids = torch.cat((txt_ids, img_ids), dim=1)
pe = self.pe_embedder(ids)
@@ -134,14 +144,16 @@ class Flux(nn.Module):
txt=args["txt"],
vec=args["vec"],
pe=args["pe"],
- attn_mask=args.get("attn_mask"))
+ attn_mask=args.get("attn_mask"),
+ transformer_options=args.get("transformer_options"))
return out
out = blocks_replace[("double_block", i)]({"img": img,
"txt": txt,
"vec": vec,
"pe": pe,
- "attn_mask": attn_mask},
+ "attn_mask": attn_mask,
+ "transformer_options": transformer_options},
{"original_block": block_wrap})
txt = out["txt"]
img = out["img"]
@@ -150,14 +162,15 @@ class Flux(nn.Module):
txt=txt,
vec=vec,
pe=pe,
- attn_mask=attn_mask)
+ attn_mask=attn_mask,
+ transformer_options=transformer_options)
if control is not None: # Controlnet
control_i = control.get("input")
if i < len(control_i):
add = control_i[i]
if add is not None:
- img += add
+ img[:, :add.shape[1]] += add
if img.dtype == torch.float16:
img = torch.nan_to_num(img, nan=0.0, posinf=65504, neginf=-65504)
@@ -171,24 +184,26 @@ class Flux(nn.Module):
out["img"] = block(args["img"],
vec=args["vec"],
pe=args["pe"],
- attn_mask=args.get("attn_mask"))
+ attn_mask=args.get("attn_mask"),
+ transformer_options=args.get("transformer_options"))
return out
out = blocks_replace[("single_block", i)]({"img": img,
"vec": vec,
"pe": pe,
- "attn_mask": attn_mask},
+ "attn_mask": attn_mask,
+ "transformer_options": transformer_options},
{"original_block": block_wrap})
img = out["img"]
else:
- img = block(img, vec=vec, pe=pe, attn_mask=attn_mask)
+ img = block(img, vec=vec, pe=pe, attn_mask=attn_mask, transformer_options=transformer_options)
if control is not None: # Controlnet
control_o = control.get("output")
if i < len(control_o):
add = control_o[i]
if add is not None:
- img[:, txt.shape[1] :, ...] += add
+ img[:, txt.shape[1] : txt.shape[1] + add.shape[1], ...] += add
img = img[:, txt.shape[1] :, ...]
@@ -214,6 +229,13 @@ class Flux(nn.Module):
return img, repeat(img_ids, "h w c -> b (h w) c", b=bs)
def forward(self, x, timestep, context, y=None, guidance=None, ref_latents=None, control=None, transformer_options={}, **kwargs):
+ return comfy.patcher_extension.WrapperExecutor.new_class_executor(
+ self._forward,
+ self,
+ comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options)
+ ).execute(x, timestep, context, y, guidance, ref_latents, control, transformer_options, **kwargs)
+
+ def _forward(self, x, timestep, context, y=None, guidance=None, ref_latents=None, control=None, transformer_options={}, **kwargs):
bs, c, h_orig, w_orig = x.shape
patch_size = self.patch_size
@@ -224,19 +246,33 @@ class Flux(nn.Module):
if ref_latents is not None:
h = 0
w = 0
+ index = 0
+ ref_latents_method = kwargs.get("ref_latents_method", "offset")
for ref in ref_latents:
- h_offset = 0
- w_offset = 0
- if ref.shape[-2] + h > ref.shape[-1] + w:
- w_offset = w
+ if ref_latents_method == "index":
+ index += 1
+ h_offset = 0
+ w_offset = 0
+ elif ref_latents_method == "uxo":
+ index = 0
+ h_offset = h_len * patch_size + h
+ w_offset = w_len * patch_size + w
+ h += ref.shape[-2]
+ w += ref.shape[-1]
else:
- h_offset = h
+ index = 1
+ h_offset = 0
+ w_offset = 0
+ if ref.shape[-2] + h > ref.shape[-1] + w:
+ w_offset = w
+ else:
+ h_offset = h
+ h = max(h, ref.shape[-2] + h_offset)
+ w = max(w, ref.shape[-1] + w_offset)
- kontext, kontext_ids = self.process_img(ref, index=1, h_offset=h_offset, w_offset=w_offset)
+ kontext, kontext_ids = self.process_img(ref, index=index, h_offset=h_offset, w_offset=w_offset)
img = torch.cat([img, kontext], dim=1)
img_ids = torch.cat([img_ids, kontext_ids], dim=1)
- h = max(h, ref.shape[-2] + h_offset)
- w = max(w, ref.shape[-1] + w_offset)
txt_ids = torch.zeros((bs, context.shape[1], 3), device=x.device, dtype=x.dtype)
out = self.forward_orig(img, img_ids, context, txt_ids, timestep, y, guidance, control, transformer_options, attn_mask=kwargs.get("attention_mask", None))
diff --git a/comfy/ldm/genmo/joint_model/asymm_models_joint.py b/comfy/ldm/genmo/joint_model/asymm_models_joint.py
index 366a8b713..5c1bb4d42 100644
--- a/comfy/ldm/genmo/joint_model/asymm_models_joint.py
+++ b/comfy/ldm/genmo/joint_model/asymm_models_joint.py
@@ -109,6 +109,7 @@ class AsymmetricAttention(nn.Module):
scale_x: torch.Tensor, # (B, dim_x), modulation for pre-RMSNorm.
scale_y: torch.Tensor, # (B, dim_y), modulation for pre-RMSNorm.
crop_y,
+ transformer_options={},
**rope_rotation,
) -> Tuple[torch.Tensor, torch.Tensor]:
rope_cos = rope_rotation.get("rope_cos")
@@ -143,7 +144,7 @@ class AsymmetricAttention(nn.Module):
xy = optimized_attention(q,
k,
- v, self.num_heads, skip_reshape=True)
+ v, self.num_heads, skip_reshape=True, transformer_options=transformer_options)
x, y = torch.tensor_split(xy, (q_x.shape[1],), dim=1)
x = self.proj_x(x)
@@ -224,6 +225,7 @@ class AsymmetricJointBlock(nn.Module):
x: torch.Tensor,
c: torch.Tensor,
y: torch.Tensor,
+ transformer_options={},
**attn_kwargs,
):
"""Forward pass of a block.
@@ -256,6 +258,7 @@ class AsymmetricJointBlock(nn.Module):
y,
scale_x=scale_msa_x,
scale_y=scale_msa_y,
+ transformer_options=transformer_options,
**attn_kwargs,
)
@@ -524,10 +527,11 @@ class AsymmDiTJoint(nn.Module):
args["txt"],
rope_cos=args["rope_cos"],
rope_sin=args["rope_sin"],
- crop_y=args["num_tokens"]
+ crop_y=args["num_tokens"],
+ transformer_options=args["transformer_options"]
)
return out
- out = blocks_replace[("double_block", i)]({"img": x, "txt": y_feat, "vec": c, "rope_cos": rope_cos, "rope_sin": rope_sin, "num_tokens": num_tokens}, {"original_block": block_wrap})
+ out = blocks_replace[("double_block", i)]({"img": x, "txt": y_feat, "vec": c, "rope_cos": rope_cos, "rope_sin": rope_sin, "num_tokens": num_tokens, "transformer_options": transformer_options}, {"original_block": block_wrap})
y_feat = out["txt"]
x = out["img"]
else:
@@ -538,6 +542,7 @@ class AsymmDiTJoint(nn.Module):
rope_cos=rope_cos,
rope_sin=rope_sin,
crop_y=num_tokens,
+ transformer_options=transformer_options,
) # (B, M, D), (B, L, D)
del y_feat # Final layers don't use dense text features.
diff --git a/comfy/ldm/hidream/model.py b/comfy/ldm/hidream/model.py
index 0305747bf..28d81c79e 100644
--- a/comfy/ldm/hidream/model.py
+++ b/comfy/ldm/hidream/model.py
@@ -13,6 +13,7 @@ from comfy.ldm.flux.layers import LastLayer
from comfy.ldm.modules.attention import optimized_attention
import comfy.model_management
+import comfy.patcher_extension
import comfy.ldm.common_dit
@@ -71,8 +72,8 @@ class TimestepEmbed(nn.Module):
return t_emb
-def attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor):
- return optimized_attention(query.view(query.shape[0], -1, query.shape[-1] * query.shape[-2]), key.view(key.shape[0], -1, key.shape[-1] * key.shape[-2]), value.view(value.shape[0], -1, value.shape[-1] * value.shape[-2]), query.shape[2])
+def attention(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, transformer_options={}):
+ return optimized_attention(query.view(query.shape[0], -1, query.shape[-1] * query.shape[-2]), key.view(key.shape[0], -1, key.shape[-1] * key.shape[-2]), value.view(value.shape[0], -1, value.shape[-1] * value.shape[-2]), query.shape[2], transformer_options=transformer_options)
class HiDreamAttnProcessor_flashattn:
@@ -85,6 +86,7 @@ class HiDreamAttnProcessor_flashattn:
image_tokens_masks: Optional[torch.FloatTensor] = None,
text_tokens: Optional[torch.FloatTensor] = None,
rope: torch.FloatTensor = None,
+ transformer_options={},
*args,
**kwargs,
) -> torch.FloatTensor:
@@ -132,7 +134,7 @@ class HiDreamAttnProcessor_flashattn:
query = torch.cat([query_1, query_2], dim=-1)
key = torch.cat([key_1, key_2], dim=-1)
- hidden_states = attention(query, key, value)
+ hidden_states = attention(query, key, value, transformer_options=transformer_options)
if not attn.single:
hidden_states_i, hidden_states_t = torch.split(hidden_states, [num_image_tokens, num_text_tokens], dim=1)
@@ -198,6 +200,7 @@ class HiDreamAttention(nn.Module):
image_tokens_masks: torch.FloatTensor = None,
norm_text_tokens: torch.FloatTensor = None,
rope: torch.FloatTensor = None,
+ transformer_options={},
) -> torch.Tensor:
return self.processor(
self,
@@ -205,6 +208,7 @@ class HiDreamAttention(nn.Module):
image_tokens_masks = image_tokens_masks,
text_tokens = norm_text_tokens,
rope = rope,
+ transformer_options=transformer_options,
)
@@ -405,7 +409,7 @@ class HiDreamImageSingleTransformerBlock(nn.Module):
text_tokens: Optional[torch.FloatTensor] = None,
adaln_input: Optional[torch.FloatTensor] = None,
rope: torch.FloatTensor = None,
-
+ transformer_options={},
) -> torch.FloatTensor:
wtype = image_tokens.dtype
shift_msa_i, scale_msa_i, gate_msa_i, shift_mlp_i, scale_mlp_i, gate_mlp_i = \
@@ -418,6 +422,7 @@ class HiDreamImageSingleTransformerBlock(nn.Module):
norm_image_tokens,
image_tokens_masks,
rope = rope,
+ transformer_options=transformer_options,
)
image_tokens = gate_msa_i * attn_output_i + image_tokens
@@ -482,6 +487,7 @@ class HiDreamImageTransformerBlock(nn.Module):
text_tokens: Optional[torch.FloatTensor] = None,
adaln_input: Optional[torch.FloatTensor] = None,
rope: torch.FloatTensor = None,
+ transformer_options={},
) -> torch.FloatTensor:
wtype = image_tokens.dtype
shift_msa_i, scale_msa_i, gate_msa_i, shift_mlp_i, scale_mlp_i, gate_mlp_i, \
@@ -499,6 +505,7 @@ class HiDreamImageTransformerBlock(nn.Module):
image_tokens_masks,
norm_text_tokens,
rope = rope,
+ transformer_options=transformer_options,
)
image_tokens = gate_msa_i * attn_output_i + image_tokens
@@ -549,6 +556,7 @@ class HiDreamImageBlock(nn.Module):
text_tokens: Optional[torch.FloatTensor] = None,
adaln_input: torch.FloatTensor = None,
rope: torch.FloatTensor = None,
+ transformer_options={},
) -> torch.FloatTensor:
return self.block(
image_tokens,
@@ -556,6 +564,7 @@ class HiDreamImageBlock(nn.Module):
text_tokens,
adaln_input,
rope,
+ transformer_options=transformer_options,
)
@@ -692,7 +701,23 @@ class HiDreamImageTransformer2DModel(nn.Module):
raise NotImplementedError
return x, x_masks, img_sizes
- def forward(
+ def forward(self,
+ x: torch.Tensor,
+ t: torch.Tensor,
+ y: Optional[torch.Tensor] = None,
+ context: Optional[torch.Tensor] = None,
+ encoder_hidden_states_llama3=None,
+ image_cond=None,
+ control = None,
+ transformer_options = {},
+ ):
+ return comfy.patcher_extension.WrapperExecutor.new_class_executor(
+ self._forward,
+ self,
+ comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options)
+ ).execute(x, t, y, context, encoder_hidden_states_llama3, image_cond, control, transformer_options)
+
+ def _forward(
self,
x: torch.Tensor,
t: torch.Tensor,
@@ -769,6 +794,7 @@ class HiDreamImageTransformer2DModel(nn.Module):
text_tokens = cur_encoder_hidden_states,
adaln_input = adaln_input,
rope = rope,
+ transformer_options=transformer_options,
)
initial_encoder_hidden_states = initial_encoder_hidden_states[:, :initial_encoder_hidden_states_seq_len]
block_id += 1
@@ -792,6 +818,7 @@ class HiDreamImageTransformer2DModel(nn.Module):
text_tokens=None,
adaln_input=adaln_input,
rope=rope,
+ transformer_options=transformer_options,
)
hidden_states = hidden_states[:, :hidden_states_seq_len]
block_id += 1
diff --git a/comfy/ldm/hunyuan3d/model.py b/comfy/ldm/hunyuan3d/model.py
index 4e18358f0..4991b1645 100644
--- a/comfy/ldm/hunyuan3d/model.py
+++ b/comfy/ldm/hunyuan3d/model.py
@@ -7,6 +7,7 @@ from comfy.ldm.flux.layers import (
SingleStreamBlock,
timestep_embedding,
)
+import comfy.patcher_extension
class Hunyuan3Dv2(nn.Module):
@@ -67,6 +68,13 @@ class Hunyuan3Dv2(nn.Module):
self.final_layer = LastLayer(hidden_size, 1, in_channels, dtype=dtype, device=device, operations=operations)
def forward(self, x, timestep, context, guidance=None, transformer_options={}, **kwargs):
+ return comfy.patcher_extension.WrapperExecutor.new_class_executor(
+ self._forward,
+ self,
+ comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options)
+ ).execute(x, timestep, context, guidance, transformer_options, **kwargs)
+
+ def _forward(self, x, timestep, context, guidance=None, transformer_options={}, **kwargs):
x = x.movedim(-1, -2)
timestep = 1.0 - timestep
txt = context
@@ -91,14 +99,16 @@ class Hunyuan3Dv2(nn.Module):
txt=args["txt"],
vec=args["vec"],
pe=args["pe"],
- attn_mask=args.get("attn_mask"))
+ attn_mask=args.get("attn_mask"),
+ transformer_options=args["transformer_options"])
return out
out = blocks_replace[("double_block", i)]({"img": img,
"txt": txt,
"vec": vec,
"pe": pe,
- "attn_mask": attn_mask},
+ "attn_mask": attn_mask,
+ "transformer_options": transformer_options},
{"original_block": block_wrap})
txt = out["txt"]
img = out["img"]
@@ -107,7 +117,8 @@ class Hunyuan3Dv2(nn.Module):
txt=txt,
vec=vec,
pe=pe,
- attn_mask=attn_mask)
+ attn_mask=attn_mask,
+ transformer_options=transformer_options)
img = torch.cat((txt, img), 1)
@@ -118,17 +129,19 @@ class Hunyuan3Dv2(nn.Module):
out["img"] = block(args["img"],
vec=args["vec"],
pe=args["pe"],
- attn_mask=args.get("attn_mask"))
+ attn_mask=args.get("attn_mask"),
+ transformer_options=args["transformer_options"])
return out
out = blocks_replace[("single_block", i)]({"img": img,
"vec": vec,
"pe": pe,
- "attn_mask": attn_mask},
+ "attn_mask": attn_mask,
+ "transformer_options": transformer_options},
{"original_block": block_wrap})
img = out["img"]
else:
- img = block(img, vec=vec, pe=pe, attn_mask=attn_mask)
+ img = block(img, vec=vec, pe=pe, attn_mask=attn_mask, transformer_options=transformer_options)
img = img[:, txt.shape[1]:, ...]
img = self.final_layer(img, vec)
diff --git a/comfy/ldm/hunyuan3d/vae.py b/comfy/ldm/hunyuan3d/vae.py
index 5eb2c6548..760944827 100644
--- a/comfy/ldm/hunyuan3d/vae.py
+++ b/comfy/ldm/hunyuan3d/vae.py
@@ -4,81 +4,458 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
-
-
-from typing import Union, Tuple, List, Callable, Optional
-
import numpy as np
-from einops import repeat, rearrange
+import math
from tqdm import tqdm
+
+from typing import Optional
+
import logging
import comfy.ops
ops = comfy.ops.disable_weight_init
-def generate_dense_grid_points(
- bbox_min: np.ndarray,
- bbox_max: np.ndarray,
- octree_resolution: int,
- indexing: str = "ij",
-):
- length = bbox_max - bbox_min
- num_cells = octree_resolution
+def fps(src: torch.Tensor, batch: torch.Tensor, sampling_ratio: float, start_random: bool = True):
- x = np.linspace(bbox_min[0], bbox_max[0], int(num_cells) + 1, dtype=np.float32)
- y = np.linspace(bbox_min[1], bbox_max[1], int(num_cells) + 1, dtype=np.float32)
- z = np.linspace(bbox_min[2], bbox_max[2], int(num_cells) + 1, dtype=np.float32)
- [xs, ys, zs] = np.meshgrid(x, y, z, indexing=indexing)
- xyz = np.stack((xs, ys, zs), axis=-1)
- grid_size = [int(num_cells) + 1, int(num_cells) + 1, int(num_cells) + 1]
+ # manually create the pointer vector
+ assert src.size(0) == batch.numel()
- return xyz, grid_size, length
+ batch_size = int(batch.max()) + 1
+ deg = src.new_zeros(batch_size, dtype = torch.long)
+
+ deg.scatter_add_(0, batch, torch.ones_like(batch))
+
+ ptr_vec = deg.new_zeros(batch_size + 1)
+ torch.cumsum(deg, 0, out=ptr_vec[1:])
+
+ #return fps_sampling(src, ptr_vec, ratio)
+ sampled_indicies = []
+
+ for b in range(batch_size):
+ # start and the end of each batch
+ start, end = ptr_vec[b].item(), ptr_vec[b + 1].item()
+ # points from the point cloud
+ points = src[start:end]
+
+ num_points = points.size(0)
+ num_samples = max(1, math.ceil(num_points * sampling_ratio))
+
+ selected = torch.zeros(num_samples, device = src.device, dtype = torch.long)
+ distances = torch.full((num_points,), float("inf"), device = src.device)
+
+ # select a random start point
+ if start_random:
+ farthest = torch.randint(0, num_points, (1,), device = src.device)
+ else:
+ farthest = torch.tensor([0], device = src.device, dtype = torch.long)
+
+ for i in range(num_samples):
+ selected[i] = farthest
+ centroid = points[farthest].squeeze(0)
+ dist = torch.norm(points - centroid, dim = 1) # compute euclidean distance
+ distances = torch.minimum(distances, dist)
+ farthest = torch.argmax(distances)
+
+ sampled_indicies.append(torch.arange(start, end)[selected])
+
+ return torch.cat(sampled_indicies, dim = 0)
+class PointCrossAttention(nn.Module):
+ def __init__(self,
+ num_latents: int,
+ downsample_ratio: float,
+ pc_size: int,
+ pc_sharpedge_size: int,
+ point_feats: int,
+ width: int,
+ heads: int,
+ layers: int,
+ fourier_embedder,
+ normal_pe: bool = False,
+ qkv_bias: bool = False,
+ use_ln_post: bool = True,
+ qk_norm: bool = True):
+
+ super().__init__()
+
+ self.fourier_embedder = fourier_embedder
+
+ self.pc_size = pc_size
+ self.normal_pe = normal_pe
+ self.downsample_ratio = downsample_ratio
+ self.pc_sharpedge_size = pc_sharpedge_size
+ self.num_latents = num_latents
+ self.point_feats = point_feats
+
+ self.input_proj = nn.Linear(self.fourier_embedder.out_dim + point_feats, width)
+
+ self.cross_attn = ResidualCrossAttentionBlock(
+ width = width,
+ heads = heads,
+ qkv_bias = qkv_bias,
+ qk_norm = qk_norm
+ )
+
+ self.self_attn = None
+ if layers > 0:
+ self.self_attn = Transformer(
+ width = width,
+ heads = heads,
+ qkv_bias = qkv_bias,
+ qk_norm = qk_norm,
+ layers = layers
+ )
+
+ if use_ln_post:
+ self.ln_post = nn.LayerNorm(width)
+ else:
+ self.ln_post = None
+
+ def sample_points_and_latents(self, point_cloud: torch.Tensor, features: torch.Tensor):
+
+ """
+ Subsample points randomly from the point cloud (input_pc)
+ Further sample the subsampled points to get query_pc
+ take the fourier embeddings for both input and query pc
+
+ Mental Note: FPS-sampled points (query_pc) act as latent tokens that attend to and learn from the broader context in input_pc.
+ Goal: get a smaller represenation (query_pc) to represent the entire scence structure by learning from a broader subset (input_pc).
+ More computationally efficient.
+
+ Features are additional information for each point in the cloud
+ """
+
+ B, _, D = point_cloud.shape
+
+ num_latents = int(self.num_latents)
+
+ num_random_query = self.pc_size / (self.pc_size + self.pc_sharpedge_size) * num_latents
+ num_sharpedge_query = num_latents - num_random_query
+
+ # Split random and sharpedge surface points
+ random_pc, sharpedge_pc = torch.split(point_cloud, [self.pc_size, self.pc_sharpedge_size], dim=1)
+
+ # assert statements
+ assert random_pc.shape[1] <= self.pc_size, "Random surface points size must be less than or equal to pc_size"
+ assert sharpedge_pc.shape[1] <= self.pc_sharpedge_size, "Sharpedge surface points size must be less than or equal to pc_sharpedge_size"
+
+ input_random_pc_size = int(num_random_query * self.downsample_ratio)
+ random_query_pc, random_input_pc, random_idx_pc, random_idx_query = \
+ self.subsample(pc = random_pc, num_query = num_random_query, input_pc_size = input_random_pc_size)
+
+ input_sharpedge_pc_size = int(num_sharpedge_query * self.downsample_ratio)
+
+ if input_sharpedge_pc_size == 0:
+ sharpedge_input_pc = torch.zeros(B, 0, D, dtype = random_input_pc.dtype).to(point_cloud.device)
+ sharpedge_query_pc = torch.zeros(B, 0, D, dtype= random_query_pc.dtype).to(point_cloud.device)
+
+ else:
+ sharpedge_query_pc, sharpedge_input_pc, sharpedge_idx_pc, sharpedge_idx_query = \
+ self.subsample(pc = sharpedge_pc, num_query = num_sharpedge_query, input_pc_size = input_sharpedge_pc_size)
+
+ # concat the random and sharpedges
+ query_pc = torch.cat([random_query_pc, sharpedge_query_pc], dim = 1)
+ input_pc = torch.cat([random_input_pc, sharpedge_input_pc], dim = 1)
+
+ query = self.fourier_embedder(query_pc)
+ data = self.fourier_embedder(input_pc)
+
+ if self.point_feats > 0:
+ random_surface_features, sharpedge_surface_features = torch.split(features, [self.pc_size, self.pc_sharpedge_size], dim = 1)
+
+ input_random_surface_features, query_random_features = \
+ self.handle_features(features = random_surface_features, idx_pc = random_idx_pc, batch_size = B,
+ input_pc_size = input_random_pc_size, idx_query = random_idx_query)
+
+ if input_sharpedge_pc_size == 0:
+ input_sharpedge_surface_features = torch.zeros(B, 0, self.point_feats,
+ dtype = input_random_surface_features.dtype, device = point_cloud.device)
+
+ query_sharpedge_features = torch.zeros(B, 0, self.point_feats,
+ dtype = query_random_features.dtype, device = point_cloud.device)
+ else:
+
+ input_sharpedge_surface_features, query_sharpedge_features = \
+ self.handle_features(idx_pc = sharpedge_idx_pc, features = sharpedge_surface_features,
+ batch_size = B, idx_query = sharpedge_idx_query, input_pc_size = input_sharpedge_pc_size)
+
+ query_features = torch.cat([query_random_features, query_sharpedge_features], dim = 1)
+ input_features = torch.cat([input_random_surface_features, input_sharpedge_surface_features], dim = 1)
+
+ if self.normal_pe:
+ # apply the fourier embeddings on the first 3 dims (xyz)
+ input_features_pe = self.fourier_embedder(input_features[..., :3])
+ query_features_pe = self.fourier_embedder(query_features[..., :3])
+ # replace the first 3 dims with the new PE ones
+ input_features = torch.cat([input_features_pe, input_features[..., :3]], dim = -1)
+ query_features = torch.cat([query_features_pe, query_features[..., :3]], dim = -1)
+
+ # concat at the channels dim
+ query = torch.cat([query, query_features], dim = -1)
+ data = torch.cat([data, input_features], dim = -1)
+
+ # don't return pc_info to avoid unnecessary memory usuage
+ return query.view(B, -1, query.shape[-1]), data.view(B, -1, data.shape[-1])
+
+ def forward(self, point_cloud: torch.Tensor, features: torch.Tensor):
+
+ query, data = self.sample_points_and_latents(point_cloud = point_cloud, features = features)
+
+ # apply projections
+ query = self.input_proj(query)
+ data = self.input_proj(data)
+
+ # apply cross attention between query and data
+ latents = self.cross_attn(query, data)
+
+ if self.self_attn is not None:
+ latents = self.self_attn(latents)
+
+ if self.ln_post is not None:
+ latents = self.ln_post(latents)
+
+ return latents
-class VanillaVolumeDecoder:
+ def subsample(self, pc, num_query, input_pc_size: int):
+
+ """
+ num_query: number of points to keep after FPS
+ input_pc_size: number of points to select before FPS
+ """
+
+ B, _, D = pc.shape
+ query_ratio = num_query / input_pc_size
+
+ # random subsampling of points inside the point cloud
+ idx_pc = torch.randperm(pc.shape[1], device = pc.device)[:input_pc_size]
+ input_pc = pc[:, idx_pc, :]
+
+ # flatten to allow applying fps across the whole batch
+ flattent_input_pc = input_pc.view(B * input_pc_size, D)
+
+ # construct a batch_down tensor to tell fps
+ # which points belong to which batch
+ N_down = int(flattent_input_pc.shape[0] / B)
+ batch_down = torch.arange(B).to(pc.device)
+ batch_down = torch.repeat_interleave(batch_down, N_down)
+
+ idx_query = fps(flattent_input_pc, batch_down, sampling_ratio = query_ratio)
+ query_pc = flattent_input_pc[idx_query].view(B, -1, D)
+
+ return query_pc, input_pc, idx_pc, idx_query
+
+ def handle_features(self, features, idx_pc, input_pc_size, batch_size: int, idx_query):
+
+ B = batch_size
+
+ input_surface_features = features[:, idx_pc, :]
+ flattent_input_features = input_surface_features.view(B * input_pc_size, -1)
+ query_features = flattent_input_features[idx_query].view(B, -1,
+ flattent_input_features.shape[-1])
+
+ return input_surface_features, query_features
+
+def normalize_mesh(mesh, scale = 0.9999):
+ """Normalize mesh to fit in [-scale, scale]. Translate mesh so its center is [0,0,0]"""
+
+ bbox = mesh.bounds
+ center = (bbox[1] + bbox[0]) / 2
+
+ max_extent = (bbox[1] - bbox[0]).max()
+ mesh.apply_translation(-center)
+ mesh.apply_scale((2 * scale) / max_extent)
+
+ return mesh
+
+def sample_pointcloud(mesh, num = 200000):
+ """ Uniformly sample points from the surface of the mesh """
+
+ points, face_idx = mesh.sample(num, return_index = True)
+ normals = mesh.face_normals[face_idx]
+ return torch.from_numpy(points.astype(np.float32)), torch.from_numpy(normals.astype(np.float32))
+
+def detect_sharp_edges(mesh, threshold=0.985):
+ """Return edge indices (a, b) that lie on sharp boundaries of the mesh."""
+
+ V, F = mesh.vertices, mesh.faces
+ VN, FN = mesh.vertex_normals, mesh.face_normals
+
+ sharp_mask = np.ones(V.shape[0])
+ for i in range(3):
+ indices = F[:, i]
+ alignment = np.einsum('ij,ij->i', VN[indices], FN)
+ dot_stack = np.stack((sharp_mask[indices], alignment), axis=-1)
+ sharp_mask[indices] = np.min(dot_stack, axis=-1)
+
+ edge_a = np.concatenate([F[:, 0], F[:, 1], F[:, 2]])
+ edge_b = np.concatenate([F[:, 1], F[:, 2], F[:, 0]])
+ sharp_edges = (sharp_mask[edge_a] < threshold) & (sharp_mask[edge_b] < threshold)
+
+ return edge_a[sharp_edges], edge_b[sharp_edges]
+
+
+def sharp_sample_pointcloud(mesh, num = 16384):
+ """ Sample points preferentially from sharp edges in the mesh. """
+
+ edge_a, edge_b = detect_sharp_edges(mesh)
+ V, VN = mesh.vertices, mesh.vertex_normals
+
+ va, vb = V[edge_a], V[edge_b]
+ na, nb = VN[edge_a], VN[edge_b]
+
+ edge_lengths = np.linalg.norm(vb - va, axis=-1)
+ weights = edge_lengths / edge_lengths.sum()
+
+ indices = np.searchsorted(np.cumsum(weights), np.random.rand(num))
+ t = np.random.rand(num, 1)
+
+ samples = t * va[indices] + (1 - t) * vb[indices]
+ normals = t * na[indices] + (1 - t) * nb[indices]
+
+ return samples.astype(np.float32), normals.astype(np.float32)
+
+def load_surface_sharpedge(mesh, num_points=4096, num_sharp_points=4096, sharpedge_flag = True, device = "cuda"):
+ """Load a surface with optional sharp-edge annotations from a trimesh mesh."""
+
+ import trimesh
+
+ try:
+ mesh_full = trimesh.util.concatenate(mesh.dump())
+ except Exception:
+ mesh_full = trimesh.util.concatenate(mesh)
+
+ mesh_full = normalize_mesh(mesh_full)
+
+ faces = mesh_full.faces
+ vertices = mesh_full.vertices
+ origin_face_count = faces.shape[0]
+
+ mesh_surface = trimesh.Trimesh(vertices=vertices, faces=faces[:origin_face_count])
+ mesh_fill = trimesh.Trimesh(vertices=vertices, faces=faces[origin_face_count:])
+
+ area_surface = mesh_surface.area
+ area_fill = mesh_fill.area
+ total_area = area_surface + area_fill
+
+ sample_num = 499712 // 2
+ fill_ratio = area_fill / total_area if total_area > 0 else 0
+
+ num_fill = int(sample_num * fill_ratio)
+ num_surface = sample_num - num_fill
+
+ surf_pts, surf_normals = sample_pointcloud(mesh_surface, num_surface)
+ fill_pts, fill_normals = (torch.zeros(0, 3), torch.zeros(0, 3)) if num_fill == 0 else sample_pointcloud(mesh_fill, num_fill)
+
+ sharp_pts, sharp_normals = sharp_sample_pointcloud(mesh_surface, sample_num)
+
+ def assemble_tensor(points, normals, label=None):
+
+ data = torch.cat([points, normals], dim=1).half().to(device)
+
+ if label is not None:
+ label_tensor = torch.full((data.shape[0], 1), float(label), dtype=torch.float16).to(device)
+ data = torch.cat([data, label_tensor], dim=1)
+
+ return data
+
+ surface = assemble_tensor(torch.cat([surf_pts.to(device), fill_pts.to(device)], dim=0),
+ torch.cat([surf_normals.to(device), fill_normals.to(device)], dim=0),
+ label = 0 if sharpedge_flag else None)
+
+ sharp_surface = assemble_tensor(torch.from_numpy(sharp_pts), torch.from_numpy(sharp_normals),
+ label = 1 if sharpedge_flag else None)
+
+ rng = np.random.default_rng()
+
+ surface = surface[rng.choice(surface.shape[0], num_points, replace = False)]
+ sharp_surface = sharp_surface[rng.choice(sharp_surface.shape[0], num_sharp_points, replace = False)]
+
+ full = torch.cat([surface, sharp_surface], dim = 0).unsqueeze(0)
+
+ return full
+
+class SharpEdgeSurfaceLoader:
+ """ Load mesh surface and sharp edge samples. """
+
+ def __init__(self, num_uniform_points = 8192, num_sharp_points = 8192):
+
+ self.num_uniform_points = num_uniform_points
+ self.num_sharp_points = num_sharp_points
+ self.total_points = num_uniform_points + num_sharp_points
+
+ def __call__(self, mesh_input, device = "cuda"):
+ mesh = self._load_mesh(mesh_input)
+ return load_surface_sharpedge(mesh, self.num_uniform_points, self.num_sharp_points, device = device)
+
+ @staticmethod
+ def _load_mesh(mesh_input):
+ import trimesh
+
+ if isinstance(mesh_input, str):
+ mesh = trimesh.load(mesh_input, force="mesh", merge_primitives = True)
+ else:
+ mesh = mesh_input
+
+ if isinstance(mesh, trimesh.Scene):
+ combined = None
+ for obj in mesh.geometry.values():
+ combined = obj if combined is None else combined + obj
+ return combined
+
+ return mesh
+
+class DiagonalGaussianDistribution:
+ def __init__(self, params: torch.Tensor, feature_dim: int = -1):
+
+ # divide quant channels (8) into mean and log variance
+ self.mean, self.logvar = torch.chunk(params, 2, dim = feature_dim)
+
+ self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
+ self.std = torch.exp(0.5 * self.logvar)
+
+ def sample(self):
+
+ eps = torch.randn_like(self.std)
+ z = self.mean + eps * self.std
+
+ return z
+
+################################################
+# Volume Decoder
+################################################
+
+class VanillaVolumeDecoder():
@torch.no_grad()
- def __call__(
- self,
- latents: torch.FloatTensor,
- geo_decoder: Callable,
- bounds: Union[Tuple[float], List[float], float] = 1.01,
- num_chunks: int = 10000,
- octree_resolution: int = None,
- enable_pbar: bool = True,
- **kwargs,
- ):
- device = latents.device
- dtype = latents.dtype
- batch_size = latents.shape[0]
+ def __call__(self, latents: torch.Tensor, geo_decoder: callable, octree_resolution: int, bounds = 1.01,
+ num_chunks: int = 10_000, enable_pbar: bool = True, **kwargs):
- # 1. generate query points
if isinstance(bounds, float):
bounds = [-bounds, -bounds, -bounds, bounds, bounds, bounds]
- bbox_min, bbox_max = np.array(bounds[0:3]), np.array(bounds[3:6])
- xyz_samples, grid_size, length = generate_dense_grid_points(
- bbox_min=bbox_min,
- bbox_max=bbox_max,
- octree_resolution=octree_resolution,
- indexing="ij"
- )
- xyz_samples = torch.from_numpy(xyz_samples).to(device, dtype=dtype).contiguous().reshape(-1, 3)
+ bbox_min, bbox_max = torch.tensor(bounds[:3]), torch.tensor(bounds[3:])
+
+ x = torch.linspace(bbox_min[0], bbox_max[0], int(octree_resolution) + 1, dtype = torch.float32)
+ y = torch.linspace(bbox_min[1], bbox_max[1], int(octree_resolution) + 1, dtype = torch.float32)
+ z = torch.linspace(bbox_min[2], bbox_max[2], int(octree_resolution) + 1, dtype = torch.float32)
+
+ [xs, ys, zs] = torch.meshgrid(x, y, z, indexing = "ij")
+ xyz = torch.stack((xs, ys, zs), axis=-1).to(latents.device, dtype = latents.dtype).contiguous().reshape(-1, 3)
+ grid_size = [int(octree_resolution) + 1, int(octree_resolution) + 1, int(octree_resolution) + 1]
- # 2. latents to 3d volume
batch_logits = []
- for start in tqdm(range(0, xyz_samples.shape[0], num_chunks), desc="Volume Decoding",
+ for start in tqdm(range(0, xyz.shape[0], num_chunks), desc="Volume Decoding",
disable=not enable_pbar):
- chunk_queries = xyz_samples[start: start + num_chunks, :]
- chunk_queries = repeat(chunk_queries, "p c -> b p c", b=batch_size)
- logits = geo_decoder(queries=chunk_queries, latents=latents)
+
+ chunk_queries = xyz[start: start + num_chunks, :]
+ chunk_queries = chunk_queries.unsqueeze(0).repeat(latents.shape[0], 1, 1)
+ logits = geo_decoder(queries = chunk_queries, latents = latents)
batch_logits.append(logits)
- grid_logits = torch.cat(batch_logits, dim=1)
- grid_logits = grid_logits.view((batch_size, *grid_size)).float()
+ grid_logits = torch.cat(batch_logits, dim = 1)
+ grid_logits = grid_logits.view((latents.shape[0], *grid_size)).float()
return grid_logits
-
class FourierEmbedder(nn.Module):
"""The sin/cosine positional embedding. Given an input tensor `x` of shape [n_batch, ..., c_dim], it converts
each feature dimension of `x[..., i]` into:
@@ -175,13 +552,11 @@ class FourierEmbedder(nn.Module):
else:
return x
-
class CrossAttentionProcessor:
def __call__(self, attn, q, k, v):
- out = F.scaled_dot_product_attention(q, k, v)
+ out = comfy.ops.scaled_dot_product_attention(q, k, v)
return out
-
class DropPath(nn.Module):
"""Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
"""
@@ -232,38 +607,41 @@ class MLP(nn.Module):
def forward(self, x):
return self.drop_path(self.c_proj(self.gelu(self.c_fc(x))))
-
class QKVMultiheadCrossAttention(nn.Module):
def __init__(
self,
- *,
heads: int,
+ n_data = None,
width=None,
qk_norm=False,
norm_layer=ops.LayerNorm
):
super().__init__()
self.heads = heads
+ self.n_data = n_data
self.q_norm = norm_layer(width // heads, elementwise_affine=True, eps=1e-6) if qk_norm else nn.Identity()
self.k_norm = norm_layer(width // heads, elementwise_affine=True, eps=1e-6) if qk_norm else nn.Identity()
- self.attn_processor = CrossAttentionProcessor()
-
def forward(self, q, kv):
+
_, n_ctx, _ = q.shape
bs, n_data, width = kv.shape
+
attn_ch = width // self.heads // 2
q = q.view(bs, n_ctx, self.heads, -1)
+
kv = kv.view(bs, n_data, self.heads, -1)
k, v = torch.split(kv, attn_ch, dim=-1)
q = self.q_norm(q)
k = self.k_norm(k)
- q, k, v = map(lambda t: rearrange(t, 'b n h d -> b h n d', h=self.heads), (q, k, v))
- out = self.attn_processor(self, q, k, v)
- out = out.transpose(1, 2).reshape(bs, n_ctx, -1)
- return out
+ q, k, v = [t.permute(0, 2, 1, 3) for t in (q, k, v)]
+ out = F.scaled_dot_product_attention(q, k, v)
+
+ out = out.transpose(1, 2).reshape(bs, n_ctx, -1)
+
+ return out
class MultiheadCrossAttention(nn.Module):
def __init__(
@@ -306,7 +684,6 @@ class MultiheadCrossAttention(nn.Module):
x = self.c_proj(x)
return x
-
class ResidualCrossAttentionBlock(nn.Module):
def __init__(
self,
@@ -366,7 +743,7 @@ class QKVMultiheadAttention(nn.Module):
q = self.q_norm(q)
k = self.k_norm(k)
- q, k, v = map(lambda t: rearrange(t, 'b n h d -> b h n d', h=self.heads), (q, k, v))
+ q, k, v = [t.permute(0, 2, 1, 3) for t in (q, k, v)]
out = F.scaled_dot_product_attention(q, k, v).transpose(1, 2).reshape(bs, n_ctx, -1)
return out
@@ -383,8 +760,7 @@ class MultiheadAttention(nn.Module):
drop_path_rate: float = 0.0
):
super().__init__()
- self.width = width
- self.heads = heads
+
self.c_qkv = ops.Linear(width, width * 3, bias=qkv_bias)
self.c_proj = ops.Linear(width, width)
self.attention = QKVMultiheadAttention(
@@ -491,7 +867,7 @@ class CrossAttentionDecoder(nn.Module):
self.query_proj = ops.Linear(self.fourier_embedder.out_dim, width)
if self.downsample_ratio != 1:
self.latents_proj = ops.Linear(width * downsample_ratio, width)
- if self.enable_ln_post == False:
+ if not self.enable_ln_post:
qk_norm = False
self.cross_attn_decoder = ResidualCrossAttentionBlock(
width=width,
@@ -522,28 +898,44 @@ class CrossAttentionDecoder(nn.Module):
class ShapeVAE(nn.Module):
def __init__(
- self,
- *,
- embed_dim: int,
- width: int,
- heads: int,
- num_decoder_layers: int,
- geo_decoder_downsample_ratio: int = 1,
- geo_decoder_mlp_expand_ratio: int = 4,
- geo_decoder_ln_post: bool = True,
- num_freqs: int = 8,
- include_pi: bool = True,
- qkv_bias: bool = True,
- qk_norm: bool = False,
- label_type: str = "binary",
- drop_path_rate: float = 0.0,
- scale_factor: float = 1.0,
+ self,
+ *,
+ num_latents: int = 4096,
+ embed_dim: int = 64,
+ width: int = 1024,
+ heads: int = 16,
+ num_decoder_layers: int = 16,
+ num_encoder_layers: int = 8,
+ pc_size: int = 81920,
+ pc_sharpedge_size: int = 0,
+ point_feats: int = 4,
+ downsample_ratio: int = 20,
+ geo_decoder_downsample_ratio: int = 1,
+ geo_decoder_mlp_expand_ratio: int = 4,
+ geo_decoder_ln_post: bool = True,
+ num_freqs: int = 8,
+ qkv_bias: bool = False,
+ qk_norm: bool = True,
+ drop_path_rate: float = 0.0,
+ include_pi: bool = False,
+ scale_factor: float = 1.0039506158752403,
+ label_type: str = "binary",
):
super().__init__()
self.geo_decoder_ln_post = geo_decoder_ln_post
self.fourier_embedder = FourierEmbedder(num_freqs=num_freqs, include_pi=include_pi)
+ self.encoder = PointCrossAttention(layers = num_encoder_layers,
+ num_latents = num_latents,
+ downsample_ratio = downsample_ratio,
+ heads = heads,
+ pc_size = pc_size,
+ width = width,
+ point_feats = point_feats,
+ fourier_embedder = self.fourier_embedder,
+ pc_sharpedge_size = pc_sharpedge_size)
+
self.post_kl = ops.Linear(embed_dim, width)
self.transformer = Transformer(
@@ -583,5 +975,14 @@ class ShapeVAE(nn.Module):
grid_logits = self.volume_decoder(latents, self.geo_decoder, bounds=bounds, num_chunks=num_chunks, octree_resolution=octree_resolution, enable_pbar=enable_pbar)
return grid_logits.movedim(-2, -1)
- def encode(self, x):
- return None
+ def encode(self, surface):
+
+ pc, feats = surface[:, :, :3], surface[:, :, 3:]
+ latents = self.encoder(pc, feats)
+
+ moments = self.pre_kl(latents)
+ posterior = DiagonalGaussianDistribution(moments, feature_dim = -1)
+
+ latents = posterior.sample()
+
+ return latents
diff --git a/comfy/ldm/hunyuan3dv2_1/hunyuandit.py b/comfy/ldm/hunyuan3dv2_1/hunyuandit.py
new file mode 100644
index 000000000..d48d9d642
--- /dev/null
+++ b/comfy/ldm/hunyuan3dv2_1/hunyuandit.py
@@ -0,0 +1,659 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from comfy.ldm.modules.attention import optimized_attention
+import comfy.model_management
+
+class GELU(nn.Module):
+
+ def __init__(self, dim_in: int, dim_out: int, operations, device, dtype):
+ super().__init__()
+ self.proj = operations.Linear(dim_in, dim_out, device = device, dtype = dtype)
+
+ def gelu(self, gate: torch.Tensor) -> torch.Tensor:
+
+ if gate.device.type == "mps":
+ return F.gelu(gate.to(dtype = torch.float32)).to(dtype = gate.dtype)
+
+ return F.gelu(gate)
+
+ def forward(self, hidden_states):
+
+ hidden_states = self.proj(hidden_states)
+ hidden_states = self.gelu(hidden_states)
+
+ return hidden_states
+
+class FeedForward(nn.Module):
+
+ def __init__(self, dim: int, dim_out = None, mult: int = 4,
+ dropout: float = 0.0, inner_dim = None, operations = None, device = None, dtype = None):
+
+ super().__init__()
+ if inner_dim is None:
+ inner_dim = int(dim * mult)
+
+ dim_out = dim_out if dim_out is not None else dim
+
+ act_fn = GELU(dim, inner_dim, operations = operations, device = device, dtype = dtype)
+
+ self.net = nn.ModuleList([])
+ self.net.append(act_fn)
+
+ self.net.append(nn.Dropout(dropout))
+ self.net.append(operations.Linear(inner_dim, dim_out, device = device, dtype = dtype))
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ for module in self.net:
+ hidden_states = module(hidden_states)
+ return hidden_states
+
+class AddAuxLoss(torch.autograd.Function):
+
+ @staticmethod
+ def forward(ctx, x, loss):
+ # do nothing in forward (no computation)
+ ctx.requires_aux_loss = loss.requires_grad
+ ctx.dtype = loss.dtype
+
+ return x
+
+ @staticmethod
+ def backward(ctx, grad_output):
+ # add the aux loss gradients
+ grad_loss = None
+ # put the aux grad the same as the main grad loss
+ # aux grad contributes equally
+ if ctx.requires_aux_loss:
+ grad_loss = torch.ones(1, dtype = ctx.dtype, device = grad_output.device)
+
+ return grad_output, grad_loss
+
+class MoEGate(nn.Module):
+
+ def __init__(self, embed_dim, num_experts=16, num_experts_per_tok=2, aux_loss_alpha=0.01, device = None, dtype = None):
+
+ super().__init__()
+ self.top_k = num_experts_per_tok
+ self.n_routed_experts = num_experts
+
+ self.alpha = aux_loss_alpha
+
+ self.gating_dim = embed_dim
+ self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim), device = device, dtype = dtype))
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+
+ # flatten hidden states
+ hidden_states = hidden_states.view(-1, hidden_states.size(-1))
+
+ # get logits and pass it to softmax
+ logits = F.linear(hidden_states, comfy.model_management.cast_to(self.weight, dtype=hidden_states.dtype, device=hidden_states.device), bias = None)
+ scores = logits.softmax(dim = -1)
+
+ topk_weight, topk_idx = torch.topk(scores, k = self.top_k, dim = -1, sorted = False)
+
+ if self.training and self.alpha > 0.0:
+ scores_for_aux = scores
+
+ # used bincount instead of one hot encoding
+ counts = torch.bincount(topk_idx.view(-1), minlength = self.n_routed_experts).float()
+ ce = counts / topk_idx.numel() # normalized expert usage
+
+ # mean expert score
+ Pi = scores_for_aux.mean(0)
+
+ # expert balance loss
+ aux_loss = (Pi * ce * self.n_routed_experts).sum() * self.alpha
+ else:
+ aux_loss = None
+
+ return topk_idx, topk_weight, aux_loss
+
+class MoEBlock(nn.Module):
+ def __init__(self, dim, num_experts: int = 6, moe_top_k: int = 2, dropout: float = 0.0,
+ ff_inner_dim: int = None, operations = None, device = None, dtype = None):
+ super().__init__()
+
+ self.moe_top_k = moe_top_k
+ self.num_experts = num_experts
+
+ self.experts = nn.ModuleList([
+ FeedForward(dim, dropout = dropout, inner_dim = ff_inner_dim, operations = operations, device = device, dtype = dtype)
+ for _ in range(num_experts)
+ ])
+
+ self.gate = MoEGate(dim, num_experts = num_experts, num_experts_per_tok = moe_top_k, device = device, dtype = dtype)
+ self.shared_experts = FeedForward(dim, dropout = dropout, inner_dim = ff_inner_dim, operations = operations, device = device, dtype = dtype)
+
+ def forward(self, hidden_states) -> torch.Tensor:
+
+ identity = hidden_states
+ orig_shape = hidden_states.shape
+ topk_idx, topk_weight, aux_loss = self.gate(hidden_states)
+
+ hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
+ flat_topk_idx = topk_idx.view(-1)
+
+ if self.training:
+
+ hidden_states = hidden_states.repeat_interleave(self.moe_top_k, dim = 0)
+ y = torch.empty_like(hidden_states, dtype = hidden_states.dtype)
+
+ for i, expert in enumerate(self.experts):
+ tmp = expert(hidden_states[flat_topk_idx == i])
+ y[flat_topk_idx == i] = tmp.to(hidden_states.dtype)
+
+ y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim = 1)
+ y = y.view(*orig_shape)
+
+ y = AddAuxLoss.apply(y, aux_loss)
+ else:
+ y = self.moe_infer(hidden_states, flat_expert_indices = flat_topk_idx,flat_expert_weights = topk_weight.view(-1, 1)).view(*orig_shape)
+
+ y = y + self.shared_experts(identity)
+
+ return y
+
+ @torch.no_grad()
+ def moe_infer(self, x, flat_expert_indices, flat_expert_weights):
+
+ expert_cache = torch.zeros_like(x)
+ idxs = flat_expert_indices.argsort()
+
+ # no need for .numpy().cpu() here
+ tokens_per_expert = flat_expert_indices.bincount().cumsum(0)
+ token_idxs = idxs // self.moe_top_k
+
+ for i, end_idx in enumerate(tokens_per_expert):
+
+ start_idx = 0 if i == 0 else tokens_per_expert[i-1]
+
+ if start_idx == end_idx:
+ continue
+
+ expert = self.experts[i]
+ exp_token_idx = token_idxs[start_idx:end_idx]
+
+ expert_tokens = x[exp_token_idx]
+ expert_out = expert(expert_tokens)
+
+ expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]])
+
+ # use index_add_ with a 1-D index tensor directly avoids building a large [N, D] index map and extra memcopy required by scatter_reduce_
+ # + avoid dtype conversion
+ expert_cache.index_add_(0, exp_token_idx, expert_out)
+
+ return expert_cache
+
+class Timesteps(nn.Module):
+ def __init__(self, num_channels: int, downscale_freq_shift: float = 0.0,
+ scale: float = 1.0, max_period: int = 10000):
+ super().__init__()
+
+ self.num_channels = num_channels
+ half_dim = num_channels // 2
+
+ # precompute the “inv_freq” vector once
+ exponent = -math.log(max_period) * torch.arange(
+ half_dim, dtype=torch.float32
+ ) / (half_dim - downscale_freq_shift)
+
+ inv_freq = torch.exp(exponent)
+
+ # pad
+ if num_channels % 2 == 1:
+ # we’ll pad a zero at the end of the cos-half
+ inv_freq = torch.cat([inv_freq, inv_freq.new_zeros(1)])
+
+ # register to buffer so it moves with the device
+ self.register_buffer("inv_freq", inv_freq, persistent = False)
+ self.scale = scale
+
+ def forward(self, timesteps: torch.Tensor):
+
+ x = timesteps.float().unsqueeze(1) * self.inv_freq.to(timesteps.device).unsqueeze(0)
+
+
+ # fused CUDA kernels for sin and cos
+ sin_emb = x.sin()
+ cos_emb = x.cos()
+
+ emb = torch.cat([sin_emb, cos_emb], dim = 1)
+
+ # scale factor
+ if self.scale != 1.0:
+ emb = emb * self.scale
+
+ # If we padded inv_freq for odd, emb is already wide enough; otherwise:
+ if emb.shape[1] > self.num_channels:
+ emb = emb[:, :self.num_channels]
+
+ return emb
+
+class TimestepEmbedder(nn.Module):
+ def __init__(self, hidden_size, frequency_embedding_size = 256, cond_proj_dim = None, operations = None, device = None, dtype = None):
+ super().__init__()
+
+ self.mlp = nn.Sequential(
+ operations.Linear(hidden_size, frequency_embedding_size, bias=True, device = device, dtype = dtype),
+ nn.GELU(),
+ operations.Linear(frequency_embedding_size, hidden_size, bias=True, device = device, dtype = dtype),
+ )
+ self.frequency_embedding_size = frequency_embedding_size
+
+ if cond_proj_dim is not None:
+ self.cond_proj = operations.Linear(cond_proj_dim, frequency_embedding_size, bias=False, device = device, dtype = dtype)
+
+ self.time_embed = Timesteps(hidden_size)
+
+ def forward(self, timesteps, condition):
+
+ timestep_embed = self.time_embed(timesteps).type(self.mlp[0].weight.dtype)
+
+ if condition is not None:
+ cond_embed = self.cond_proj(condition)
+ timestep_embed = timestep_embed + cond_embed
+
+ time_conditioned = self.mlp(timestep_embed)
+
+ # for broadcasting with image tokens
+ return time_conditioned.unsqueeze(1)
+
+class MLP(nn.Module):
+ def __init__(self, *, width: int, operations = None, device = None, dtype = None):
+ super().__init__()
+ self.width = width
+ self.fc1 = operations.Linear(width, width * 4, device = device, dtype = dtype)
+ self.fc2 = operations.Linear(width * 4, width, device = device, dtype = dtype)
+ self.gelu = nn.GELU()
+
+ def forward(self, x):
+ return self.fc2(self.gelu(self.fc1(x)))
+
+class CrossAttention(nn.Module):
+ def __init__(
+ self,
+ qdim,
+ kdim,
+ num_heads,
+ qkv_bias=True,
+ qk_norm=False,
+ norm_layer=nn.LayerNorm,
+ use_fp16: bool = False,
+ operations = None,
+ dtype = None,
+ device = None,
+ **kwargs,
+ ):
+ super().__init__()
+ self.qdim = qdim
+ self.kdim = kdim
+
+ self.num_heads = num_heads
+ self.head_dim = self.qdim // num_heads
+
+ self.scale = self.head_dim ** -0.5
+
+ self.to_q = operations.Linear(qdim, qdim, bias=qkv_bias, device = device, dtype = dtype)
+ self.to_k = operations.Linear(kdim, qdim, bias=qkv_bias, device = device, dtype = dtype)
+ self.to_v = operations.Linear(kdim, qdim, bias=qkv_bias, device = device, dtype = dtype)
+
+ if use_fp16:
+ eps = 1.0 / 65504
+ else:
+ eps = 1e-6
+
+ if norm_layer == nn.LayerNorm:
+ norm_layer = operations.LayerNorm
+ else:
+ norm_layer = operations.RMSNorm
+
+ self.q_norm = norm_layer(self.head_dim, elementwise_affine=True, eps = eps, device = device, dtype = dtype) if qk_norm else nn.Identity()
+ self.k_norm = norm_layer(self.head_dim, elementwise_affine=True, eps = eps, device = device, dtype = dtype) if qk_norm else nn.Identity()
+ self.out_proj = operations.Linear(qdim, qdim, bias=True, device = device, dtype = dtype)
+
+ def forward(self, x, y):
+
+ b, s1, _ = x.shape
+ _, s2, _ = y.shape
+
+ y = y.to(next(self.to_k.parameters()).dtype)
+
+ q = self.to_q(x)
+ k = self.to_k(y)
+ v = self.to_v(y)
+
+ kv = torch.cat((k, v), dim=-1)
+ split_size = kv.shape[-1] // self.num_heads // 2
+
+ kv = kv.view(1, -1, self.num_heads, split_size * 2)
+ k, v = torch.split(kv, split_size, dim=-1)
+
+ q = q.view(b, s1, self.num_heads, self.head_dim)
+ k = k.view(b, s2, self.num_heads, self.head_dim)
+ v = v.reshape(b, s2, self.num_heads * self.head_dim)
+
+ q = self.q_norm(q)
+ k = self.k_norm(k)
+
+ x = optimized_attention(
+ q.reshape(b, s1, self.num_heads * self.head_dim),
+ k.reshape(b, s2, self.num_heads * self.head_dim),
+ v,
+ heads=self.num_heads,
+ )
+
+ out = self.out_proj(x)
+
+ return out
+
+class Attention(nn.Module):
+
+ def __init__(
+ self,
+ dim,
+ num_heads,
+ qkv_bias = True,
+ qk_norm = False,
+ norm_layer = nn.LayerNorm,
+ use_fp16: bool = False,
+ operations = None,
+ device = None,
+ dtype = None
+ ):
+ super().__init__()
+ self.dim = dim
+ self.num_heads = num_heads
+ self.head_dim = self.dim // num_heads
+ self.scale = self.head_dim ** -0.5
+
+ self.to_q = operations.Linear(dim, dim, bias = qkv_bias, device = device, dtype = dtype)
+ self.to_k = operations.Linear(dim, dim, bias = qkv_bias, device = device, dtype = dtype)
+ self.to_v = operations.Linear(dim, dim, bias = qkv_bias, device = device, dtype = dtype)
+
+ if use_fp16:
+ eps = 1.0 / 65504
+ else:
+ eps = 1e-6
+
+ if norm_layer == nn.LayerNorm:
+ norm_layer = operations.LayerNorm
+ else:
+ norm_layer = operations.RMSNorm
+
+ self.q_norm = norm_layer(self.head_dim, elementwise_affine=True, eps = eps, device = device, dtype = dtype) if qk_norm else nn.Identity()
+ self.k_norm = norm_layer(self.head_dim, elementwise_affine=True, eps = eps, device = device, dtype = dtype) if qk_norm else nn.Identity()
+ self.out_proj = operations.Linear(dim, dim, device = device, dtype = dtype)
+
+ def forward(self, x):
+ B, N, _ = x.shape
+
+ query = self.to_q(x)
+ key = self.to_k(x)
+ value = self.to_v(x)
+
+ qkv_combined = torch.cat((query, key, value), dim=-1)
+ split_size = qkv_combined.shape[-1] // self.num_heads // 3
+
+ qkv = qkv_combined.view(1, -1, self.num_heads, split_size * 3)
+ query, key, value = torch.split(qkv, split_size, dim=-1)
+
+ query = query.reshape(B, N, self.num_heads, self.head_dim)
+ key = key.reshape(B, N, self.num_heads, self.head_dim)
+ value = value.reshape(B, N, self.num_heads * self.head_dim)
+
+ query = self.q_norm(query)
+ key = self.k_norm(key)
+
+ x = optimized_attention(
+ query.reshape(B, N, self.num_heads * self.head_dim),
+ key.reshape(B, N, self.num_heads * self.head_dim),
+ value,
+ heads=self.num_heads,
+ )
+
+ x = self.out_proj(x)
+ return x
+
+class HunYuanDiTBlock(nn.Module):
+ def __init__(
+ self,
+ hidden_size,
+ c_emb_size,
+ num_heads,
+ text_states_dim=1024,
+ qk_norm=False,
+ norm_layer=nn.LayerNorm,
+ qk_norm_layer=True,
+ qkv_bias=True,
+ skip_connection=True,
+ timested_modulate=False,
+ use_moe: bool = False,
+ num_experts: int = 8,
+ moe_top_k: int = 2,
+ use_fp16: bool = False,
+ operations = None,
+ device = None, dtype = None
+ ):
+ super().__init__()
+
+ # eps can't be 1e-6 in fp16 mode because of numerical stability issues
+ if use_fp16:
+ eps = 1.0 / 65504
+ else:
+ eps = 1e-6
+
+ self.norm1 = norm_layer(hidden_size, elementwise_affine = True, eps = eps, device = device, dtype = dtype)
+
+ self.attn1 = Attention(hidden_size, num_heads=num_heads, qkv_bias=qkv_bias, qk_norm=qk_norm,
+ norm_layer=qk_norm_layer, use_fp16 = use_fp16, device = device, dtype = dtype, operations = operations)
+
+ self.norm2 = norm_layer(hidden_size, elementwise_affine = True, eps = eps, device = device, dtype = dtype)
+
+ self.timested_modulate = timested_modulate
+ if self.timested_modulate:
+ self.default_modulation = nn.Sequential(
+ nn.SiLU(),
+ operations.Linear(c_emb_size, hidden_size, bias=True, device = device, dtype = dtype)
+ )
+
+ self.attn2 = CrossAttention(hidden_size, text_states_dim, num_heads=num_heads, qkv_bias=qkv_bias,
+ qk_norm=qk_norm, norm_layer=qk_norm_layer, use_fp16 = use_fp16,
+ device = device, dtype = dtype, operations = operations)
+
+ self.norm3 = norm_layer(hidden_size, elementwise_affine = True, eps = eps, device = device, dtype = dtype)
+
+ if skip_connection:
+ self.skip_norm = norm_layer(hidden_size, elementwise_affine = True, eps = eps, device = device, dtype = dtype)
+ self.skip_linear = operations.Linear(2 * hidden_size, hidden_size, device = device, dtype = dtype)
+ else:
+ self.skip_linear = None
+
+ self.use_moe = use_moe
+
+ if self.use_moe:
+ self.moe = MoEBlock(
+ hidden_size,
+ num_experts = num_experts,
+ moe_top_k = moe_top_k,
+ dropout = 0.0,
+ ff_inner_dim = int(hidden_size * 4.0),
+ device = device, dtype = dtype,
+ operations = operations
+ )
+ else:
+ self.mlp = MLP(width=hidden_size, operations=operations, device = device, dtype = dtype)
+
+ def forward(self, hidden_states, conditioning=None, text_states=None, skip_tensor=None):
+
+ if self.skip_linear is not None:
+ combined = torch.cat([skip_tensor, hidden_states], dim=-1)
+ hidden_states = self.skip_linear(combined)
+ hidden_states = self.skip_norm(hidden_states)
+
+ # self attention
+ if self.timested_modulate:
+ modulation_shift = self.default_modulation(conditioning).unsqueeze(dim=1)
+ hidden_states = hidden_states + modulation_shift
+
+ self_attn_out = self.attn1(self.norm1(hidden_states))
+ hidden_states = hidden_states + self_attn_out
+
+ # cross attention
+ hidden_states = hidden_states + self.attn2(self.norm2(hidden_states), text_states)
+
+ # MLP Layer
+ mlp_input = self.norm3(hidden_states)
+
+ if self.use_moe:
+ hidden_states = hidden_states + self.moe(mlp_input)
+ else:
+ hidden_states = hidden_states + self.mlp(mlp_input)
+
+ return hidden_states
+
+class FinalLayer(nn.Module):
+
+ def __init__(self, final_hidden_size, out_channels, operations, use_fp16: bool = False, device = None, dtype = None):
+ super().__init__()
+
+ if use_fp16:
+ eps = 1.0 / 65504
+ else:
+ eps = 1e-6
+
+ self.norm_final = operations.LayerNorm(final_hidden_size, elementwise_affine = True, eps = eps, device = device, dtype = dtype)
+ self.linear = operations.Linear(final_hidden_size, out_channels, bias = True, device = device, dtype = dtype)
+
+ def forward(self, x):
+ x = self.norm_final(x)
+ x = x[:, 1:]
+ x = self.linear(x)
+ return x
+
+class HunYuanDiTPlain(nn.Module):
+
+ # init with the defaults values from https://huggingface.co/tencent/Hunyuan3D-2.1/blob/main/hunyuan3d-dit-v2-1/config.yaml
+ def __init__(
+ self,
+ in_channels: int = 64,
+ hidden_size: int = 2048,
+ context_dim: int = 1024,
+ depth: int = 21,
+ num_heads: int = 16,
+ qk_norm: bool = True,
+ qkv_bias: bool = False,
+ num_moe_layers: int = 6,
+ guidance_cond_proj_dim = 2048,
+ norm_type = 'layer',
+ num_experts: int = 8,
+ moe_top_k: int = 2,
+ use_fp16: bool = False,
+ dtype = None,
+ device = None,
+ operations = None,
+ **kwargs
+ ):
+
+ self.dtype = dtype
+
+ super().__init__()
+
+ self.depth = depth
+
+ self.in_channels = in_channels
+ self.out_channels = in_channels
+
+ self.num_heads = num_heads
+ self.hidden_size = hidden_size
+
+ norm = operations.LayerNorm if norm_type == 'layer' else operations.RMSNorm
+ qk_norm = operations.RMSNorm
+
+ self.context_dim = context_dim
+ self.guidance_cond_proj_dim = guidance_cond_proj_dim
+
+ self.x_embedder = operations.Linear(in_channels, hidden_size, bias = True, device = device, dtype = dtype)
+ self.t_embedder = TimestepEmbedder(hidden_size, hidden_size * 4, cond_proj_dim = guidance_cond_proj_dim, device = device, dtype = dtype, operations = operations)
+
+
+ # HUnYuanDiT Blocks
+ self.blocks = nn.ModuleList([
+ HunYuanDiTBlock(hidden_size=hidden_size,
+ c_emb_size=hidden_size,
+ num_heads=num_heads,
+ text_states_dim=context_dim,
+ qk_norm=qk_norm,
+ norm_layer = norm,
+ qk_norm_layer = qk_norm,
+ skip_connection=layer > depth // 2,
+ qkv_bias=qkv_bias,
+ use_moe=True if depth - layer <= num_moe_layers else False,
+ num_experts=num_experts,
+ moe_top_k=moe_top_k,
+ use_fp16 = use_fp16,
+ device = device, dtype = dtype, operations = operations)
+ for layer in range(depth)
+ ])
+
+ self.depth = depth
+
+ self.final_layer = FinalLayer(hidden_size, self.out_channels, use_fp16 = use_fp16, operations = operations, device = device, dtype = dtype)
+
+ def forward(self, x, t, context, transformer_options = {}, **kwargs):
+
+ x = x.movedim(-1, -2)
+ uncond_emb, cond_emb = context.chunk(2, dim = 0)
+
+ context = torch.cat([cond_emb, uncond_emb], dim = 0)
+ main_condition = context
+
+ t = 1.0 - t
+
+ time_embedded = self.t_embedder(t, condition = kwargs.get('guidance_cond'))
+
+ x = x.to(dtype = next(self.x_embedder.parameters()).dtype)
+ x_embedded = self.x_embedder(x)
+
+ combined = torch.cat([time_embedded, x_embedded], dim=1)
+
+ def block_wrap(args):
+ return block(
+ args["x"],
+ args["t"],
+ args["cond"],
+ skip_tensor=args.get("skip"),)
+
+ skip_stack = []
+ patches_replace = transformer_options.get("patches_replace", {})
+ blocks_replace = patches_replace.get("dit", {})
+ for idx, block in enumerate(self.blocks):
+ if idx <= self.depth // 2:
+ skip_input = None
+ else:
+ skip_input = skip_stack.pop()
+
+ if ("block", idx) in blocks_replace:
+
+ combined = blocks_replace[("block", idx)](
+ {
+ "x": combined,
+ "t": time_embedded,
+ "cond": main_condition,
+ "skip": skip_input,
+ },
+ {"original_block": block_wrap},
+ )
+ else:
+ combined = block(combined, time_embedded, main_condition, skip_tensor=skip_input)
+
+ if idx < self.depth // 2:
+ skip_stack.append(combined)
+
+ output = self.final_layer(combined)
+ output = output.movedim(-2, -1) * (-1.0)
+
+ cond_emb, uncond_emb = output.chunk(2, dim = 0)
+ return torch.cat([uncond_emb, cond_emb])
diff --git a/comfy/ldm/hunyuan_video/model.py b/comfy/ldm/hunyuan_video/model.py
index fbd8d4196..5132e6c07 100644
--- a/comfy/ldm/hunyuan_video/model.py
+++ b/comfy/ldm/hunyuan_video/model.py
@@ -1,6 +1,7 @@
#Based on Flux code because of weird hunyuan video code license.
import torch
+import comfy.patcher_extension
import comfy.ldm.flux.layers
import comfy.ldm.modules.diffusionmodules.mmdit
from comfy.ldm.modules.attention import optimized_attention
@@ -39,6 +40,8 @@ class HunyuanVideoParams:
patch_size: list
qkv_bias: bool
guidance_embed: bool
+ byt5: bool
+ meanflow: bool
class SelfAttentionRef(nn.Module):
@@ -77,13 +80,13 @@ class TokenRefinerBlock(nn.Module):
operations.Linear(mlp_hidden_dim, hidden_size, bias=True, dtype=dtype, device=device),
)
- def forward(self, x, c, mask):
+ def forward(self, x, c, mask, transformer_options={}):
mod1, mod2 = self.adaLN_modulation(c).chunk(2, dim=1)
norm_x = self.norm1(x)
qkv = self.self_attn.qkv(norm_x)
q, k, v = qkv.reshape(qkv.shape[0], qkv.shape[1], 3, self.heads, -1).permute(2, 0, 3, 1, 4)
- attn = optimized_attention(q, k, v, self.heads, mask=mask, skip_reshape=True)
+ attn = optimized_attention(q, k, v, self.heads, mask=mask, skip_reshape=True, transformer_options=transformer_options)
x = x + self.self_attn.proj(attn) * mod1.unsqueeze(1)
x = x + self.mlp(self.norm2(x)) * mod2.unsqueeze(1)
@@ -114,14 +117,14 @@ class IndividualTokenRefiner(nn.Module):
]
)
- def forward(self, x, c, mask):
+ def forward(self, x, c, mask, transformer_options={}):
m = None
if mask is not None:
m = mask.view(mask.shape[0], 1, 1, mask.shape[1]).repeat(1, 1, mask.shape[1], 1)
m = m + m.transpose(2, 3)
for block in self.blocks:
- x = block(x, c, m)
+ x = block(x, c, m, transformer_options=transformer_options)
return x
@@ -149,6 +152,7 @@ class TokenRefiner(nn.Module):
x,
timesteps,
mask,
+ transformer_options={},
):
t = self.t_embedder(timestep_embedding(timesteps, 256, time_factor=1.0).to(x.dtype))
# m = mask.float().unsqueeze(-1)
@@ -157,9 +161,33 @@ class TokenRefiner(nn.Module):
c = t + self.c_embedder(c.to(x.dtype))
x = self.input_embedder(x)
- x = self.individual_token_refiner(x, c, mask)
+ x = self.individual_token_refiner(x, c, mask, transformer_options=transformer_options)
return x
+
+class ByT5Mapper(nn.Module):
+ def __init__(self, in_dim, out_dim, hidden_dim, out_dim1, use_res=False, dtype=None, device=None, operations=None):
+ super().__init__()
+ self.layernorm = operations.LayerNorm(in_dim, dtype=dtype, device=device)
+ self.fc1 = operations.Linear(in_dim, hidden_dim, dtype=dtype, device=device)
+ self.fc2 = operations.Linear(hidden_dim, out_dim, dtype=dtype, device=device)
+ self.fc3 = operations.Linear(out_dim, out_dim1, dtype=dtype, device=device)
+ self.use_res = use_res
+ self.act_fn = nn.GELU()
+
+ def forward(self, x):
+ if self.use_res:
+ res = x
+ x = self.layernorm(x)
+ x = self.fc1(x)
+ x = self.act_fn(x)
+ x = self.fc2(x)
+ x2 = self.act_fn(x)
+ x2 = self.fc3(x2)
+ if self.use_res:
+ x2 = x2 + res
+ return x2
+
class HunyuanVideo(nn.Module):
"""
Transformer model for flow matching on sequences.
@@ -184,9 +212,13 @@ class HunyuanVideo(nn.Module):
self.num_heads = params.num_heads
self.pe_embedder = EmbedND(dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim)
- self.img_in = comfy.ldm.modules.diffusionmodules.mmdit.PatchEmbed(None, self.patch_size, self.in_channels, self.hidden_size, conv3d=True, dtype=dtype, device=device, operations=operations)
+ self.img_in = comfy.ldm.modules.diffusionmodules.mmdit.PatchEmbed(None, self.patch_size, self.in_channels, self.hidden_size, conv3d=len(self.patch_size) == 3, dtype=dtype, device=device, operations=operations)
self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size, dtype=dtype, device=device, operations=operations)
- self.vector_in = MLPEmbedder(params.vec_in_dim, self.hidden_size, dtype=dtype, device=device, operations=operations)
+ if params.vec_in_dim is not None:
+ self.vector_in = MLPEmbedder(params.vec_in_dim, self.hidden_size, dtype=dtype, device=device, operations=operations)
+ else:
+ self.vector_in = None
+
self.guidance_in = (
MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size, dtype=dtype, device=device, operations=operations) if params.guidance_embed else nn.Identity()
)
@@ -214,6 +246,23 @@ class HunyuanVideo(nn.Module):
]
)
+ if params.byt5:
+ self.byt5_in = ByT5Mapper(
+ in_dim=1472,
+ out_dim=2048,
+ hidden_dim=2048,
+ out_dim1=self.hidden_size,
+ use_res=False,
+ dtype=dtype, device=device, operations=operations
+ )
+ else:
+ self.byt5_in = None
+
+ if params.meanflow:
+ self.time_r_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size, dtype=dtype, device=device, operations=operations)
+ else:
+ self.time_r_in = None
+
if final_layer:
self.final_layer = LastLayer(self.hidden_size, self.patch_size[-1], self.out_channels, dtype=dtype, device=device, operations=operations)
@@ -225,10 +274,12 @@ class HunyuanVideo(nn.Module):
txt_ids: Tensor,
txt_mask: Tensor,
timesteps: Tensor,
- y: Tensor,
+ y: Tensor = None,
+ txt_byt5=None,
guidance: Tensor = None,
guiding_frame_index=None,
ref_latent=None,
+ disable_time_r=False,
control=None,
transformer_options={},
) -> Tensor:
@@ -239,6 +290,14 @@ class HunyuanVideo(nn.Module):
img = self.img_in(img)
vec = self.time_in(timestep_embedding(timesteps, 256, time_factor=1.0).to(img.dtype))
+ if (self.time_r_in is not None) and (not disable_time_r):
+ w = torch.where(transformer_options['sigmas'][0] == transformer_options['sample_sigmas'])[0] # This most likely could be improved
+ if len(w) > 0:
+ timesteps_r = transformer_options['sample_sigmas'][w[0] + 1]
+ timesteps_r = timesteps_r.unsqueeze(0).to(device=timesteps.device, dtype=timesteps.dtype)
+ vec_r = self.time_r_in(timestep_embedding(timesteps_r, 256, time_factor=1000.0).to(img.dtype))
+ vec = (vec + vec_r) / 2
+
if ref_latent is not None:
ref_latent_ids = self.img_ids(ref_latent)
ref_latent = self.img_in(ref_latent)
@@ -249,13 +308,17 @@ class HunyuanVideo(nn.Module):
if guiding_frame_index is not None:
token_replace_vec = self.time_in(timestep_embedding(guiding_frame_index, 256, time_factor=1.0))
- vec_ = self.vector_in(y[:, :self.params.vec_in_dim])
- vec = torch.cat([(vec_ + token_replace_vec).unsqueeze(1), (vec_ + vec).unsqueeze(1)], dim=1)
+ if self.vector_in is not None:
+ vec_ = self.vector_in(y[:, :self.params.vec_in_dim])
+ vec = torch.cat([(vec_ + token_replace_vec).unsqueeze(1), (vec_ + vec).unsqueeze(1)], dim=1)
+ else:
+ vec = torch.cat([(token_replace_vec).unsqueeze(1), (vec).unsqueeze(1)], dim=1)
frame_tokens = (initial_shape[-1] // self.patch_size[-1]) * (initial_shape[-2] // self.patch_size[-2])
modulation_dims = [(0, frame_tokens, 0), (frame_tokens, None, 1)]
modulation_dims_txt = [(0, None, 1)]
else:
- vec = vec + self.vector_in(y[:, :self.params.vec_in_dim])
+ if self.vector_in is not None:
+ vec = vec + self.vector_in(y[:, :self.params.vec_in_dim])
modulation_dims = None
modulation_dims_txt = None
@@ -266,7 +329,13 @@ class HunyuanVideo(nn.Module):
if txt_mask is not None and not torch.is_floating_point(txt_mask):
txt_mask = (txt_mask - 1).to(img.dtype) * torch.finfo(img.dtype).max
- txt = self.txt_in(txt, timesteps, txt_mask)
+ txt = self.txt_in(txt, timesteps, txt_mask, transformer_options=transformer_options)
+
+ if self.byt5_in is not None and txt_byt5 is not None:
+ txt_byt5 = self.byt5_in(txt_byt5)
+ txt_byt5_ids = torch.zeros((txt_ids.shape[0], txt_byt5.shape[1], txt_ids.shape[-1]), device=txt_ids.device, dtype=txt_ids.dtype)
+ txt = torch.cat((txt, txt_byt5), dim=1)
+ txt_ids = torch.cat((txt_ids, txt_byt5_ids), dim=1)
ids = torch.cat((img_ids, txt_ids), dim=1)
pe = self.pe_embedder(ids)
@@ -284,14 +353,14 @@ class HunyuanVideo(nn.Module):
if ("double_block", i) in blocks_replace:
def block_wrap(args):
out = {}
- out["img"], out["txt"] = block(img=args["img"], txt=args["txt"], vec=args["vec"], pe=args["pe"], attn_mask=args["attention_mask"], modulation_dims_img=args["modulation_dims_img"], modulation_dims_txt=args["modulation_dims_txt"])
+ out["img"], out["txt"] = block(img=args["img"], txt=args["txt"], vec=args["vec"], pe=args["pe"], attn_mask=args["attention_mask"], modulation_dims_img=args["modulation_dims_img"], modulation_dims_txt=args["modulation_dims_txt"], transformer_options=args["transformer_options"])
return out
- out = blocks_replace[("double_block", i)]({"img": img, "txt": txt, "vec": vec, "pe": pe, "attention_mask": attn_mask, 'modulation_dims_img': modulation_dims, 'modulation_dims_txt': modulation_dims_txt}, {"original_block": block_wrap})
+ out = blocks_replace[("double_block", i)]({"img": img, "txt": txt, "vec": vec, "pe": pe, "attention_mask": attn_mask, 'modulation_dims_img': modulation_dims, 'modulation_dims_txt': modulation_dims_txt, 'transformer_options': transformer_options}, {"original_block": block_wrap})
txt = out["txt"]
img = out["img"]
else:
- img, txt = block(img=img, txt=txt, vec=vec, pe=pe, attn_mask=attn_mask, modulation_dims_img=modulation_dims, modulation_dims_txt=modulation_dims_txt)
+ img, txt = block(img=img, txt=txt, vec=vec, pe=pe, attn_mask=attn_mask, modulation_dims_img=modulation_dims, modulation_dims_txt=modulation_dims_txt, transformer_options=transformer_options)
if control is not None: # Controlnet
control_i = control.get("input")
@@ -306,13 +375,13 @@ class HunyuanVideo(nn.Module):
if ("single_block", i) in blocks_replace:
def block_wrap(args):
out = {}
- out["img"] = block(args["img"], vec=args["vec"], pe=args["pe"], attn_mask=args["attention_mask"], modulation_dims=args["modulation_dims"])
+ out["img"] = block(args["img"], vec=args["vec"], pe=args["pe"], attn_mask=args["attention_mask"], modulation_dims=args["modulation_dims"], transformer_options=args["transformer_options"])
return out
- out = blocks_replace[("single_block", i)]({"img": img, "vec": vec, "pe": pe, "attention_mask": attn_mask, 'modulation_dims': modulation_dims}, {"original_block": block_wrap})
+ out = blocks_replace[("single_block", i)]({"img": img, "vec": vec, "pe": pe, "attention_mask": attn_mask, 'modulation_dims': modulation_dims, 'transformer_options': transformer_options}, {"original_block": block_wrap})
img = out["img"]
else:
- img = block(img, vec=vec, pe=pe, attn_mask=attn_mask, modulation_dims=modulation_dims)
+ img = block(img, vec=vec, pe=pe, attn_mask=attn_mask, modulation_dims=modulation_dims, transformer_options=transformer_options)
if control is not None: # Controlnet
control_o = control.get("output")
@@ -327,12 +396,16 @@ class HunyuanVideo(nn.Module):
img = self.final_layer(img, vec, modulation_dims=modulation_dims) # (N, T, patch_size ** 2 * out_channels)
- shape = initial_shape[-3:]
+ shape = initial_shape[-len(self.patch_size):]
for i in range(len(shape)):
shape[i] = shape[i] // self.patch_size[i]
img = img.reshape([img.shape[0]] + shape + [self.out_channels] + self.patch_size)
- img = img.permute(0, 4, 1, 5, 2, 6, 3, 7)
- img = img.reshape(initial_shape[0], self.out_channels, initial_shape[2], initial_shape[3], initial_shape[4])
+ if img.ndim == 8:
+ img = img.permute(0, 4, 1, 5, 2, 6, 3, 7)
+ img = img.reshape(initial_shape[0], self.out_channels, initial_shape[2], initial_shape[3], initial_shape[4])
+ else:
+ img = img.permute(0, 3, 1, 4, 2, 5)
+ img = img.reshape(initial_shape[0], self.out_channels, initial_shape[2], initial_shape[3])
return img
def img_ids(self, x):
@@ -347,9 +420,30 @@ class HunyuanVideo(nn.Module):
img_ids[:, :, :, 2] = img_ids[:, :, :, 2] + torch.linspace(0, w_len - 1, steps=w_len, device=x.device, dtype=x.dtype).reshape(1, 1, -1)
return repeat(img_ids, "t h w c -> b (t h w) c", b=bs)
- def forward(self, x, timestep, context, y, guidance=None, attention_mask=None, guiding_frame_index=None, ref_latent=None, control=None, transformer_options={}, **kwargs):
- bs, c, t, h, w = x.shape
- img_ids = self.img_ids(x)
- txt_ids = torch.zeros((bs, context.shape[1], 3), device=x.device, dtype=x.dtype)
- out = self.forward_orig(x, img_ids, context, txt_ids, attention_mask, timestep, y, guidance, guiding_frame_index, ref_latent, control=control, transformer_options=transformer_options)
+ def img_ids_2d(self, x):
+ bs, c, h, w = x.shape
+ patch_size = self.patch_size
+ h_len = ((h + (patch_size[0] // 2)) // patch_size[0])
+ w_len = ((w + (patch_size[1] // 2)) // patch_size[1])
+ img_ids = torch.zeros((h_len, w_len, 2), device=x.device, dtype=x.dtype)
+ img_ids[:, :, 0] = img_ids[:, :, 0] + torch.linspace(0, h_len - 1, steps=h_len, device=x.device, dtype=x.dtype).unsqueeze(1)
+ img_ids[:, :, 1] = img_ids[:, :, 1] + torch.linspace(0, w_len - 1, steps=w_len, device=x.device, dtype=x.dtype).unsqueeze(0)
+ return repeat(img_ids, "h w c -> b (h w) c", b=bs)
+
+ def forward(self, x, timestep, context, y=None, txt_byt5=None, guidance=None, attention_mask=None, guiding_frame_index=None, ref_latent=None, disable_time_r=False, control=None, transformer_options={}, **kwargs):
+ return comfy.patcher_extension.WrapperExecutor.new_class_executor(
+ self._forward,
+ self,
+ comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options)
+ ).execute(x, timestep, context, y, txt_byt5, guidance, attention_mask, guiding_frame_index, ref_latent, disable_time_r, control, transformer_options, **kwargs)
+
+ def _forward(self, x, timestep, context, y=None, txt_byt5=None, guidance=None, attention_mask=None, guiding_frame_index=None, ref_latent=None, disable_time_r=False, control=None, transformer_options={}, **kwargs):
+ bs = x.shape[0]
+ if len(self.patch_size) == 3:
+ img_ids = self.img_ids(x)
+ txt_ids = torch.zeros((bs, context.shape[1], 3), device=x.device, dtype=x.dtype)
+ else:
+ img_ids = self.img_ids_2d(x)
+ txt_ids = torch.zeros((bs, context.shape[1], 2), device=x.device, dtype=x.dtype)
+ out = self.forward_orig(x, img_ids, context, txt_ids, attention_mask, timestep, y, txt_byt5, guidance, guiding_frame_index, ref_latent, disable_time_r=disable_time_r, control=control, transformer_options=transformer_options)
return out
diff --git a/comfy/ldm/hunyuan_video/vae.py b/comfy/ldm/hunyuan_video/vae.py
new file mode 100644
index 000000000..40c12b183
--- /dev/null
+++ b/comfy/ldm/hunyuan_video/vae.py
@@ -0,0 +1,136 @@
+import torch.nn as nn
+import torch.nn.functional as F
+from comfy.ldm.modules.diffusionmodules.model import ResnetBlock, AttnBlock
+import comfy.ops
+ops = comfy.ops.disable_weight_init
+
+
+class PixelShuffle2D(nn.Module):
+ def __init__(self, in_dim, out_dim, op=ops.Conv2d):
+ super().__init__()
+ self.conv = op(in_dim, out_dim >> 2, 3, 1, 1)
+ self.ratio = (in_dim << 2) // out_dim
+
+ def forward(self, x):
+ b, c, h, w = x.shape
+ h2, w2 = h >> 1, w >> 1
+ y = self.conv(x).view(b, -1, h2, 2, w2, 2).permute(0, 3, 5, 1, 2, 4).reshape(b, -1, h2, w2)
+ r = x.view(b, c, h2, 2, w2, 2).permute(0, 3, 5, 1, 2, 4).reshape(b, c << 2, h2, w2)
+ return y + r.view(b, y.shape[1], self.ratio, h2, w2).mean(2)
+
+
+class PixelUnshuffle2D(nn.Module):
+ def __init__(self, in_dim, out_dim, op=ops.Conv2d):
+ super().__init__()
+ self.conv = op(in_dim, out_dim << 2, 3, 1, 1)
+ self.scale = (out_dim << 2) // in_dim
+
+ def forward(self, x):
+ b, c, h, w = x.shape
+ h2, w2 = h << 1, w << 1
+ y = self.conv(x).view(b, 2, 2, -1, h, w).permute(0, 3, 4, 1, 5, 2).reshape(b, -1, h2, w2)
+ r = x.repeat_interleave(self.scale, 1).view(b, 2, 2, -1, h, w).permute(0, 3, 4, 1, 5, 2).reshape(b, -1, h2, w2)
+ return y + r
+
+
+class Encoder(nn.Module):
+ def __init__(self, in_channels, z_channels, block_out_channels, num_res_blocks,
+ ffactor_spatial, downsample_match_channel=True, **_):
+ super().__init__()
+ self.z_channels = z_channels
+ self.block_out_channels = block_out_channels
+ self.num_res_blocks = num_res_blocks
+ self.conv_in = ops.Conv2d(in_channels, block_out_channels[0], 3, 1, 1)
+
+ self.down = nn.ModuleList()
+ ch = block_out_channels[0]
+ depth = (ffactor_spatial >> 1).bit_length()
+
+ for i, tgt in enumerate(block_out_channels):
+ stage = nn.Module()
+ stage.block = nn.ModuleList([ResnetBlock(in_channels=ch if j == 0 else tgt,
+ out_channels=tgt,
+ temb_channels=0,
+ conv_op=ops.Conv2d)
+ for j in range(num_res_blocks)])
+ ch = tgt
+ if i < depth:
+ nxt = block_out_channels[i + 1] if i + 1 < len(block_out_channels) and downsample_match_channel else ch
+ stage.downsample = PixelShuffle2D(ch, nxt, ops.Conv2d)
+ ch = nxt
+ self.down.append(stage)
+
+ self.mid = nn.Module()
+ self.mid.block_1 = ResnetBlock(in_channels=ch, out_channels=ch, temb_channels=0, conv_op=ops.Conv2d)
+ self.mid.attn_1 = AttnBlock(ch, conv_op=ops.Conv2d)
+ self.mid.block_2 = ResnetBlock(in_channels=ch, out_channels=ch, temb_channels=0, conv_op=ops.Conv2d)
+
+ self.norm_out = ops.GroupNorm(32, ch, 1e-6, True)
+ self.conv_out = ops.Conv2d(ch, z_channels << 1, 3, 1, 1)
+
+ def forward(self, x):
+ x = self.conv_in(x)
+
+ for stage in self.down:
+ for blk in stage.block:
+ x = blk(x)
+ if hasattr(stage, 'downsample'):
+ x = stage.downsample(x)
+
+ x = self.mid.block_2(self.mid.attn_1(self.mid.block_1(x)))
+
+ b, c, h, w = x.shape
+ grp = c // (self.z_channels << 1)
+ skip = x.view(b, c // grp, grp, h, w).mean(2)
+
+ return self.conv_out(F.silu(self.norm_out(x))) + skip
+
+
+class Decoder(nn.Module):
+ def __init__(self, z_channels, out_channels, block_out_channels, num_res_blocks,
+ ffactor_spatial, upsample_match_channel=True, **_):
+ super().__init__()
+ block_out_channels = block_out_channels[::-1]
+ self.z_channels = z_channels
+ self.block_out_channels = block_out_channels
+ self.num_res_blocks = num_res_blocks
+
+ ch = block_out_channels[0]
+ self.conv_in = ops.Conv2d(z_channels, ch, 3, 1, 1)
+
+ self.mid = nn.Module()
+ self.mid.block_1 = ResnetBlock(in_channels=ch, out_channels=ch, temb_channels=0, conv_op=ops.Conv2d)
+ self.mid.attn_1 = AttnBlock(ch, conv_op=ops.Conv2d)
+ self.mid.block_2 = ResnetBlock(in_channels=ch, out_channels=ch, temb_channels=0, conv_op=ops.Conv2d)
+
+ self.up = nn.ModuleList()
+ depth = (ffactor_spatial >> 1).bit_length()
+
+ for i, tgt in enumerate(block_out_channels):
+ stage = nn.Module()
+ stage.block = nn.ModuleList([ResnetBlock(in_channels=ch if j == 0 else tgt,
+ out_channels=tgt,
+ temb_channels=0,
+ conv_op=ops.Conv2d)
+ for j in range(num_res_blocks + 1)])
+ ch = tgt
+ if i < depth:
+ nxt = block_out_channels[i + 1] if i + 1 < len(block_out_channels) and upsample_match_channel else ch
+ stage.upsample = PixelUnshuffle2D(ch, nxt, ops.Conv2d)
+ ch = nxt
+ self.up.append(stage)
+
+ self.norm_out = ops.GroupNorm(32, ch, 1e-6, True)
+ self.conv_out = ops.Conv2d(ch, out_channels, 3, 1, 1)
+
+ def forward(self, z):
+ x = self.conv_in(z) + z.repeat_interleave(self.block_out_channels[0] // self.z_channels, 1)
+ x = self.mid.block_2(self.mid.attn_1(self.mid.block_1(x)))
+
+ for stage in self.up:
+ for blk in stage.block:
+ x = blk(x)
+ if hasattr(stage, 'upsample'):
+ x = stage.upsample(x)
+
+ return self.conv_out(F.silu(self.norm_out(x)))
diff --git a/comfy/ldm/hunyuan_video/vae_refiner.py b/comfy/ldm/hunyuan_video/vae_refiner.py
new file mode 100644
index 000000000..c2a0b507d
--- /dev/null
+++ b/comfy/ldm/hunyuan_video/vae_refiner.py
@@ -0,0 +1,301 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from comfy.ldm.modules.diffusionmodules.model import ResnetBlock, AttnBlock, VideoConv3d, Normalize
+import comfy.ops
+import comfy.ldm.models.autoencoder
+ops = comfy.ops.disable_weight_init
+
+class RMS_norm(nn.Module):
+ def __init__(self, dim):
+ super().__init__()
+ shape = (dim, 1, 1, 1)
+ self.scale = dim**0.5
+ self.gamma = nn.Parameter(torch.empty(shape))
+
+ def forward(self, x):
+ return F.normalize(x, dim=1) * self.scale * self.gamma
+
+class DnSmpl(nn.Module):
+ def __init__(self, ic, oc, tds=True, refiner_vae=True, op=VideoConv3d):
+ super().__init__()
+ fct = 2 * 2 * 2 if tds else 1 * 2 * 2
+ assert oc % fct == 0
+ self.conv = op(ic, oc // fct, kernel_size=3, stride=1, padding=1)
+ self.refiner_vae = refiner_vae
+
+ self.tds = tds
+ self.gs = fct * ic // oc
+
+ def forward(self, x):
+ r1 = 2 if self.tds else 1
+ h = self.conv(x)
+
+ if self.tds and self.refiner_vae:
+ hf = h[:, :, :1, :, :]
+ b, c, f, ht, wd = hf.shape
+ hf = hf.reshape(b, c, f, ht // 2, 2, wd // 2, 2)
+ hf = hf.permute(0, 4, 6, 1, 2, 3, 5)
+ hf = hf.reshape(b, 2 * 2 * c, f, ht // 2, wd // 2)
+ hf = torch.cat([hf, hf], dim=1)
+
+ hn = h[:, :, 1:, :, :]
+ b, c, frms, ht, wd = hn.shape
+ nf = frms // r1
+ hn = hn.reshape(b, c, nf, r1, ht // 2, 2, wd // 2, 2)
+ hn = hn.permute(0, 3, 5, 7, 1, 2, 4, 6)
+ hn = hn.reshape(b, r1 * 2 * 2 * c, nf, ht // 2, wd // 2)
+
+ h = torch.cat([hf, hn], dim=2)
+
+ xf = x[:, :, :1, :, :]
+ b, ci, f, ht, wd = xf.shape
+ xf = xf.reshape(b, ci, f, ht // 2, 2, wd // 2, 2)
+ xf = xf.permute(0, 4, 6, 1, 2, 3, 5)
+ xf = xf.reshape(b, 2 * 2 * ci, f, ht // 2, wd // 2)
+ B, C, T, H, W = xf.shape
+ xf = xf.view(B, h.shape[1], self.gs // 2, T, H, W).mean(dim=2)
+
+ xn = x[:, :, 1:, :, :]
+ b, ci, frms, ht, wd = xn.shape
+ nf = frms // r1
+ xn = xn.reshape(b, ci, nf, r1, ht // 2, 2, wd // 2, 2)
+ xn = xn.permute(0, 3, 5, 7, 1, 2, 4, 6)
+ xn = xn.reshape(b, r1 * 2 * 2 * ci, nf, ht // 2, wd // 2)
+ B, C, T, H, W = xn.shape
+ xn = xn.view(B, h.shape[1], self.gs, T, H, W).mean(dim=2)
+ sc = torch.cat([xf, xn], dim=2)
+ else:
+ b, c, frms, ht, wd = h.shape
+
+ nf = frms // r1
+ h = h.reshape(b, c, nf, r1, ht // 2, 2, wd // 2, 2)
+ h = h.permute(0, 3, 5, 7, 1, 2, 4, 6)
+ h = h.reshape(b, r1 * 2 * 2 * c, nf, ht // 2, wd // 2)
+
+ b, ci, frms, ht, wd = x.shape
+ nf = frms // r1
+ sc = x.reshape(b, ci, nf, r1, ht // 2, 2, wd // 2, 2)
+ sc = sc.permute(0, 3, 5, 7, 1, 2, 4, 6)
+ sc = sc.reshape(b, r1 * 2 * 2 * ci, nf, ht // 2, wd // 2)
+ B, C, T, H, W = sc.shape
+ sc = sc.view(B, h.shape[1], self.gs, T, H, W).mean(dim=2)
+
+ return h + sc
+
+
+class UpSmpl(nn.Module):
+ def __init__(self, ic, oc, tus=True, refiner_vae=True, op=VideoConv3d):
+ super().__init__()
+ fct = 2 * 2 * 2 if tus else 1 * 2 * 2
+ self.conv = op(ic, oc * fct, kernel_size=3, stride=1, padding=1)
+ self.refiner_vae = refiner_vae
+
+ self.tus = tus
+ self.rp = fct * oc // ic
+
+ def forward(self, x):
+ r1 = 2 if self.tus else 1
+ h = self.conv(x)
+
+ if self.tus and self.refiner_vae:
+ hf = h[:, :, :1, :, :]
+ b, c, f, ht, wd = hf.shape
+ nc = c // (2 * 2)
+ hf = hf.reshape(b, 2, 2, nc, f, ht, wd)
+ hf = hf.permute(0, 3, 4, 5, 1, 6, 2)
+ hf = hf.reshape(b, nc, f, ht * 2, wd * 2)
+ hf = hf[:, : hf.shape[1] // 2]
+
+ hn = h[:, :, 1:, :, :]
+ b, c, frms, ht, wd = hn.shape
+ nc = c // (r1 * 2 * 2)
+ hn = hn.reshape(b, r1, 2, 2, nc, frms, ht, wd)
+ hn = hn.permute(0, 4, 5, 1, 6, 2, 7, 3)
+ hn = hn.reshape(b, nc, frms * r1, ht * 2, wd * 2)
+
+ h = torch.cat([hf, hn], dim=2)
+
+ xf = x[:, :, :1, :, :]
+ b, ci, f, ht, wd = xf.shape
+ xf = xf.repeat_interleave(repeats=self.rp // 2, dim=1)
+ b, c, f, ht, wd = xf.shape
+ nc = c // (2 * 2)
+ xf = xf.reshape(b, 2, 2, nc, f, ht, wd)
+ xf = xf.permute(0, 3, 4, 5, 1, 6, 2)
+ xf = xf.reshape(b, nc, f, ht * 2, wd * 2)
+
+ xn = x[:, :, 1:, :, :]
+ xn = xn.repeat_interleave(repeats=self.rp, dim=1)
+ b, c, frms, ht, wd = xn.shape
+ nc = c // (r1 * 2 * 2)
+ xn = xn.reshape(b, r1, 2, 2, nc, frms, ht, wd)
+ xn = xn.permute(0, 4, 5, 1, 6, 2, 7, 3)
+ xn = xn.reshape(b, nc, frms * r1, ht * 2, wd * 2)
+ sc = torch.cat([xf, xn], dim=2)
+ else:
+ b, c, frms, ht, wd = h.shape
+ nc = c // (r1 * 2 * 2)
+ h = h.reshape(b, r1, 2, 2, nc, frms, ht, wd)
+ h = h.permute(0, 4, 5, 1, 6, 2, 7, 3)
+ h = h.reshape(b, nc, frms * r1, ht * 2, wd * 2)
+
+ sc = x.repeat_interleave(repeats=self.rp, dim=1)
+ b, c, frms, ht, wd = sc.shape
+ nc = c // (r1 * 2 * 2)
+ sc = sc.reshape(b, r1, 2, 2, nc, frms, ht, wd)
+ sc = sc.permute(0, 4, 5, 1, 6, 2, 7, 3)
+ sc = sc.reshape(b, nc, frms * r1, ht * 2, wd * 2)
+
+ return h + sc
+
+class Encoder(nn.Module):
+ def __init__(self, in_channels, z_channels, block_out_channels, num_res_blocks,
+ ffactor_spatial, ffactor_temporal, downsample_match_channel=True, refiner_vae=True, **_):
+ super().__init__()
+ self.z_channels = z_channels
+ self.block_out_channels = block_out_channels
+ self.num_res_blocks = num_res_blocks
+ self.ffactor_temporal = ffactor_temporal
+
+ self.refiner_vae = refiner_vae
+ if self.refiner_vae:
+ conv_op = VideoConv3d
+ norm_op = RMS_norm
+ else:
+ conv_op = ops.Conv3d
+ norm_op = Normalize
+
+ self.conv_in = conv_op(in_channels, block_out_channels[0], 3, 1, 1)
+
+ self.down = nn.ModuleList()
+ ch = block_out_channels[0]
+ depth = (ffactor_spatial >> 1).bit_length()
+ depth_temporal = ((ffactor_spatial // self.ffactor_temporal) >> 1).bit_length()
+
+ for i, tgt in enumerate(block_out_channels):
+ stage = nn.Module()
+ stage.block = nn.ModuleList([ResnetBlock(in_channels=ch if j == 0 else tgt,
+ out_channels=tgt,
+ temb_channels=0,
+ conv_op=conv_op, norm_op=norm_op)
+ for j in range(num_res_blocks)])
+ ch = tgt
+ if i < depth:
+ nxt = block_out_channels[i + 1] if i + 1 < len(block_out_channels) and downsample_match_channel else ch
+ stage.downsample = DnSmpl(ch, nxt, tds=i >= depth_temporal, refiner_vae=self.refiner_vae, op=conv_op)
+ ch = nxt
+ self.down.append(stage)
+
+ self.mid = nn.Module()
+ self.mid.block_1 = ResnetBlock(in_channels=ch, out_channels=ch, temb_channels=0, conv_op=conv_op, norm_op=norm_op)
+ self.mid.attn_1 = AttnBlock(ch, conv_op=ops.Conv3d, norm_op=norm_op)
+ self.mid.block_2 = ResnetBlock(in_channels=ch, out_channels=ch, temb_channels=0, conv_op=conv_op, norm_op=norm_op)
+
+ self.norm_out = norm_op(ch)
+ self.conv_out = conv_op(ch, z_channels << 1, 3, 1, 1)
+
+ self.regul = comfy.ldm.models.autoencoder.DiagonalGaussianRegularizer()
+
+ def forward(self, x):
+ if not self.refiner_vae and x.shape[2] == 1:
+ x = x.expand(-1, -1, self.ffactor_temporal, -1, -1)
+
+ x = self.conv_in(x)
+
+ for stage in self.down:
+ for blk in stage.block:
+ x = blk(x)
+ if hasattr(stage, 'downsample'):
+ x = stage.downsample(x)
+
+ x = self.mid.block_2(self.mid.attn_1(self.mid.block_1(x)))
+
+ b, c, t, h, w = x.shape
+ grp = c // (self.z_channels << 1)
+ skip = x.view(b, c // grp, grp, t, h, w).mean(2)
+
+ out = self.conv_out(F.silu(self.norm_out(x))) + skip
+
+ if self.refiner_vae:
+ out = self.regul(out)[0]
+
+ out = torch.cat((out[:, :, :1], out), dim=2)
+ out = out.permute(0, 2, 1, 3, 4)
+ b, f_times_2, c, h, w = out.shape
+ out = out.reshape(b, f_times_2 // 2, 2 * c, h, w)
+ out = out.permute(0, 2, 1, 3, 4).contiguous()
+
+ return out
+
+class Decoder(nn.Module):
+ def __init__(self, z_channels, out_channels, block_out_channels, num_res_blocks,
+ ffactor_spatial, ffactor_temporal, upsample_match_channel=True, refiner_vae=True, **_):
+ super().__init__()
+ block_out_channels = block_out_channels[::-1]
+ self.z_channels = z_channels
+ self.block_out_channels = block_out_channels
+ self.num_res_blocks = num_res_blocks
+
+ self.refiner_vae = refiner_vae
+ if self.refiner_vae:
+ conv_op = VideoConv3d
+ norm_op = RMS_norm
+ else:
+ conv_op = ops.Conv3d
+ norm_op = Normalize
+
+ ch = block_out_channels[0]
+ self.conv_in = conv_op(z_channels, ch, kernel_size=3, stride=1, padding=1)
+
+ self.mid = nn.Module()
+ self.mid.block_1 = ResnetBlock(in_channels=ch, out_channels=ch, temb_channels=0, conv_op=conv_op, norm_op=norm_op)
+ self.mid.attn_1 = AttnBlock(ch, conv_op=ops.Conv3d, norm_op=norm_op)
+ self.mid.block_2 = ResnetBlock(in_channels=ch, out_channels=ch, temb_channels=0, conv_op=conv_op, norm_op=norm_op)
+
+ self.up = nn.ModuleList()
+ depth = (ffactor_spatial >> 1).bit_length()
+ depth_temporal = (ffactor_temporal >> 1).bit_length()
+
+ for i, tgt in enumerate(block_out_channels):
+ stage = nn.Module()
+ stage.block = nn.ModuleList([ResnetBlock(in_channels=ch if j == 0 else tgt,
+ out_channels=tgt,
+ temb_channels=0,
+ conv_op=conv_op, norm_op=norm_op)
+ for j in range(num_res_blocks + 1)])
+ ch = tgt
+ if i < depth:
+ nxt = block_out_channels[i + 1] if i + 1 < len(block_out_channels) and upsample_match_channel else ch
+ stage.upsample = UpSmpl(ch, nxt, tus=i < depth_temporal, refiner_vae=self.refiner_vae, op=conv_op)
+ ch = nxt
+ self.up.append(stage)
+
+ self.norm_out = norm_op(ch)
+ self.conv_out = conv_op(ch, out_channels, 3, stride=1, padding=1)
+
+ def forward(self, z):
+ if self.refiner_vae:
+ z = z.permute(0, 2, 1, 3, 4)
+ b, f, c, h, w = z.shape
+ z = z.reshape(b, f, 2, c // 2, h, w)
+ z = z.permute(0, 1, 2, 3, 4, 5).reshape(b, f * 2, c // 2, h, w)
+ z = z.permute(0, 2, 1, 3, 4)
+ z = z[:, :, 1:]
+
+ x = self.conv_in(z) + z.repeat_interleave(self.block_out_channels[0] // self.z_channels, 1)
+ x = self.mid.block_2(self.mid.attn_1(self.mid.block_1(x)))
+
+ for stage in self.up:
+ for blk in stage.block:
+ x = blk(x)
+ if hasattr(stage, 'upsample'):
+ x = stage.upsample(x)
+
+ out = self.conv_out(F.silu(self.norm_out(x)))
+
+ if not self.refiner_vae:
+ if z.shape[-3] == 1:
+ out = out[:, :, -1:]
+
+ return out
diff --git a/comfy/ldm/lightricks/model.py b/comfy/ldm/lightricks/model.py
index ad9a7daea..def365ba7 100644
--- a/comfy/ldm/lightricks/model.py
+++ b/comfy/ldm/lightricks/model.py
@@ -1,5 +1,6 @@
import torch
from torch import nn
+import comfy.patcher_extension
import comfy.ldm.modules.attention
import comfy.ldm.common_dit
from einops import rearrange
@@ -270,7 +271,7 @@ class CrossAttention(nn.Module):
self.to_out = nn.Sequential(operations.Linear(inner_dim, query_dim, dtype=dtype, device=device), nn.Dropout(dropout))
- def forward(self, x, context=None, mask=None, pe=None):
+ def forward(self, x, context=None, mask=None, pe=None, transformer_options={}):
q = self.to_q(x)
context = x if context is None else context
k = self.to_k(context)
@@ -284,9 +285,9 @@ class CrossAttention(nn.Module):
k = apply_rotary_emb(k, pe)
if mask is None:
- out = comfy.ldm.modules.attention.optimized_attention(q, k, v, self.heads, attn_precision=self.attn_precision)
+ out = comfy.ldm.modules.attention.optimized_attention(q, k, v, self.heads, attn_precision=self.attn_precision, transformer_options=transformer_options)
else:
- out = comfy.ldm.modules.attention.optimized_attention_masked(q, k, v, self.heads, mask, attn_precision=self.attn_precision)
+ out = comfy.ldm.modules.attention.optimized_attention_masked(q, k, v, self.heads, mask, attn_precision=self.attn_precision, transformer_options=transformer_options)
return self.to_out(out)
@@ -302,12 +303,12 @@ class BasicTransformerBlock(nn.Module):
self.scale_shift_table = nn.Parameter(torch.empty(6, dim, device=device, dtype=dtype))
- def forward(self, x, context=None, attention_mask=None, timestep=None, pe=None):
+ def forward(self, x, context=None, attention_mask=None, timestep=None, pe=None, transformer_options={}):
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (self.scale_shift_table[None, None].to(device=x.device, dtype=x.dtype) + timestep.reshape(x.shape[0], timestep.shape[1], self.scale_shift_table.shape[0], -1)).unbind(dim=2)
- x += self.attn1(comfy.ldm.common_dit.rms_norm(x) * (1 + scale_msa) + shift_msa, pe=pe) * gate_msa
+ x += self.attn1(comfy.ldm.common_dit.rms_norm(x) * (1 + scale_msa) + shift_msa, pe=pe, transformer_options=transformer_options) * gate_msa
- x += self.attn2(x, context=context, mask=attention_mask)
+ x += self.attn2(x, context=context, mask=attention_mask, transformer_options=transformer_options)
y = comfy.ldm.common_dit.rms_norm(x) * (1 + scale_mlp) + shift_mlp
x += self.ff(y) * gate_mlp
@@ -420,6 +421,13 @@ class LTXVModel(torch.nn.Module):
self.patchifier = SymmetricPatchifier(1)
def forward(self, x, timestep, context, attention_mask, frame_rate=25, transformer_options={}, keyframe_idxs=None, **kwargs):
+ return comfy.patcher_extension.WrapperExecutor.new_class_executor(
+ self._forward,
+ self,
+ comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options)
+ ).execute(x, timestep, context, attention_mask, frame_rate, transformer_options, keyframe_idxs, **kwargs)
+
+ def _forward(self, x, timestep, context, attention_mask, frame_rate=25, transformer_options={}, keyframe_idxs=None, **kwargs):
patches_replace = transformer_options.get("patches_replace", {})
orig_shape = list(x.shape)
@@ -471,10 +479,10 @@ class LTXVModel(torch.nn.Module):
if ("double_block", i) in blocks_replace:
def block_wrap(args):
out = {}
- out["img"] = block(args["img"], context=args["txt"], attention_mask=args["attention_mask"], timestep=args["vec"], pe=args["pe"])
+ out["img"] = block(args["img"], context=args["txt"], attention_mask=args["attention_mask"], timestep=args["vec"], pe=args["pe"], transformer_options=args["transformer_options"])
return out
- out = blocks_replace[("double_block", i)]({"img": x, "txt": context, "attention_mask": attention_mask, "vec": timestep, "pe": pe}, {"original_block": block_wrap})
+ out = blocks_replace[("double_block", i)]({"img": x, "txt": context, "attention_mask": attention_mask, "vec": timestep, "pe": pe, "transformer_options": transformer_options}, {"original_block": block_wrap})
x = out["img"]
else:
x = block(
@@ -482,7 +490,8 @@ class LTXVModel(torch.nn.Module):
context=context,
attention_mask=attention_mask,
timestep=timestep,
- pe=pe
+ pe=pe,
+ transformer_options=transformer_options,
)
# 3. Output
diff --git a/comfy/ldm/lightricks/vae/causal_video_autoencoder.py b/comfy/ldm/lightricks/vae/causal_video_autoencoder.py
index f91870d71..75ed069ad 100644
--- a/comfy/ldm/lightricks/vae/causal_video_autoencoder.py
+++ b/comfy/ldm/lightricks/vae/causal_video_autoencoder.py
@@ -973,7 +973,7 @@ class VideoVAE(nn.Module):
norm_layer=config.get("norm_layer", "group_norm"),
causal=config.get("causal_decoder", False),
timestep_conditioning=self.timestep_conditioning,
- spatial_padding_mode=config.get("spatial_padding_mode", "zeros"),
+ spatial_padding_mode=config.get("spatial_padding_mode", "reflect"),
)
self.per_channel_statistics = processor()
diff --git a/comfy/ldm/lumina/model.py b/comfy/ldm/lumina/model.py
index f8dc4d7db..b4494a51d 100644
--- a/comfy/ldm/lumina/model.py
+++ b/comfy/ldm/lumina/model.py
@@ -11,6 +11,7 @@ import comfy.ldm.common_dit
from comfy.ldm.modules.diffusionmodules.mmdit import TimestepEmbedder
from comfy.ldm.modules.attention import optimized_attention_masked
from comfy.ldm.flux.layers import EmbedND
+import comfy.patcher_extension
def modulate(x, scale):
@@ -103,6 +104,7 @@ class JointAttention(nn.Module):
x: torch.Tensor,
x_mask: torch.Tensor,
freqs_cis: torch.Tensor,
+ transformer_options={},
) -> torch.Tensor:
"""
@@ -139,7 +141,7 @@ class JointAttention(nn.Module):
if n_rep >= 1:
xk = xk.unsqueeze(3).repeat(1, 1, 1, n_rep, 1).flatten(2, 3)
xv = xv.unsqueeze(3).repeat(1, 1, 1, n_rep, 1).flatten(2, 3)
- output = optimized_attention_masked(xq.movedim(1, 2), xk.movedim(1, 2), xv.movedim(1, 2), self.n_local_heads, x_mask, skip_reshape=True)
+ output = optimized_attention_masked(xq.movedim(1, 2), xk.movedim(1, 2), xv.movedim(1, 2), self.n_local_heads, x_mask, skip_reshape=True, transformer_options=transformer_options)
return self.out(output)
@@ -267,6 +269,7 @@ class JointTransformerBlock(nn.Module):
x_mask: torch.Tensor,
freqs_cis: torch.Tensor,
adaln_input: Optional[torch.Tensor]=None,
+ transformer_options={},
):
"""
Perform a forward pass through the TransformerBlock.
@@ -289,6 +292,7 @@ class JointTransformerBlock(nn.Module):
modulate(self.attention_norm1(x), scale_msa),
x_mask,
freqs_cis,
+ transformer_options=transformer_options,
)
)
x = x + gate_mlp.unsqueeze(1).tanh() * self.ffn_norm2(
@@ -303,6 +307,7 @@ class JointTransformerBlock(nn.Module):
self.attention_norm1(x),
x_mask,
freqs_cis,
+ transformer_options=transformer_options,
)
)
x = x + self.ffn_norm2(
@@ -493,7 +498,7 @@ class NextDiT(nn.Module):
return imgs
def patchify_and_embed(
- self, x: List[torch.Tensor] | torch.Tensor, cap_feats: torch.Tensor, cap_mask: torch.Tensor, t: torch.Tensor, num_tokens
+ self, x: List[torch.Tensor] | torch.Tensor, cap_feats: torch.Tensor, cap_mask: torch.Tensor, t: torch.Tensor, num_tokens, transformer_options={}
) -> Tuple[torch.Tensor, torch.Tensor, List[Tuple[int, int]], List[int], torch.Tensor]:
bsz = len(x)
pH = pW = self.patch_size
@@ -517,7 +522,7 @@ class NextDiT(nn.Module):
max_cap_len = max(l_effective_cap_len)
max_img_len = max(l_effective_img_len)
- position_ids = torch.zeros(bsz, max_seq_len, 3, dtype=torch.int32, device=device)
+ position_ids = torch.zeros(bsz, max_seq_len, 3, dtype=torch.float32, device=device)
for i in range(bsz):
cap_len = l_effective_cap_len[i]
@@ -526,10 +531,22 @@ class NextDiT(nn.Module):
H_tokens, W_tokens = H // pH, W // pW
assert H_tokens * W_tokens == img_len
- position_ids[i, :cap_len, 0] = torch.arange(cap_len, dtype=torch.int32, device=device)
+ rope_options = transformer_options.get("rope_options", None)
+ h_scale = 1.0
+ w_scale = 1.0
+ h_start = 0
+ w_start = 0
+ if rope_options is not None:
+ h_scale = rope_options.get("scale_y", 1.0)
+ w_scale = rope_options.get("scale_x", 1.0)
+
+ h_start = rope_options.get("shift_y", 0.0)
+ w_start = rope_options.get("shift_x", 0.0)
+
+ position_ids[i, :cap_len, 0] = torch.arange(cap_len, dtype=torch.float32, device=device)
position_ids[i, cap_len:cap_len+img_len, 0] = cap_len
- row_ids = torch.arange(H_tokens, dtype=torch.int32, device=device).view(-1, 1).repeat(1, W_tokens).flatten()
- col_ids = torch.arange(W_tokens, dtype=torch.int32, device=device).view(1, -1).repeat(H_tokens, 1).flatten()
+ row_ids = (torch.arange(H_tokens, dtype=torch.float32, device=device) * h_scale + h_start).view(-1, 1).repeat(1, W_tokens).flatten()
+ col_ids = (torch.arange(W_tokens, dtype=torch.float32, device=device) * w_scale + w_start).view(1, -1).repeat(H_tokens, 1).flatten()
position_ids[i, cap_len:cap_len+img_len, 1] = row_ids
position_ids[i, cap_len:cap_len+img_len, 2] = col_ids
@@ -553,7 +570,7 @@ class NextDiT(nn.Module):
# refine context
for layer in self.context_refiner:
- cap_feats = layer(cap_feats, cap_mask, cap_freqs_cis)
+ cap_feats = layer(cap_feats, cap_mask, cap_freqs_cis, transformer_options=transformer_options)
# refine image
flat_x = []
@@ -572,7 +589,7 @@ class NextDiT(nn.Module):
padded_img_embed = self.x_embedder(padded_img_embed)
padded_img_mask = padded_img_mask.unsqueeze(1)
for layer in self.noise_refiner:
- padded_img_embed = layer(padded_img_embed, padded_img_mask, img_freqs_cis, t)
+ padded_img_embed = layer(padded_img_embed, padded_img_mask, img_freqs_cis, t, transformer_options=transformer_options)
if cap_mask is not None:
mask = torch.zeros(bsz, max_seq_len, dtype=dtype, device=device)
@@ -590,8 +607,15 @@ class NextDiT(nn.Module):
return padded_full_embed, mask, img_sizes, l_effective_cap_len, freqs_cis
- # def forward(self, x, t, cap_feats, cap_mask):
def forward(self, x, timesteps, context, num_tokens, attention_mask=None, **kwargs):
+ return comfy.patcher_extension.WrapperExecutor.new_class_executor(
+ self._forward,
+ self,
+ comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, kwargs.get("transformer_options", {}))
+ ).execute(x, timesteps, context, num_tokens, attention_mask, **kwargs)
+
+ # def forward(self, x, t, cap_feats, cap_mask):
+ def _forward(self, x, timesteps, context, num_tokens, attention_mask=None, **kwargs):
t = 1.0 - timesteps
cap_feats = context
cap_mask = attention_mask
@@ -608,12 +632,13 @@ class NextDiT(nn.Module):
cap_feats = self.cap_embedder(cap_feats) # (N, L, D) # todo check if able to batchify w.o. redundant compute
+ transformer_options = kwargs.get("transformer_options", {})
x_is_tensor = isinstance(x, torch.Tensor)
- x, mask, img_size, cap_size, freqs_cis = self.patchify_and_embed(x, cap_feats, cap_mask, t, num_tokens)
+ x, mask, img_size, cap_size, freqs_cis = self.patchify_and_embed(x, cap_feats, cap_mask, t, num_tokens, transformer_options=transformer_options)
freqs_cis = freqs_cis.to(x.device)
for layer in self.layers:
- x = layer(x, mask, freqs_cis, adaln_input)
+ x = layer(x, mask, freqs_cis, adaln_input, transformer_options=transformer_options)
x = self.final_layer(x, adaln_input)
x = self.unpatchify(x, img_size, cap_size, return_tensor=x_is_tensor)[:,:,:h,:w]
diff --git a/comfy/ldm/mmaudio/vae/__init__.py b/comfy/ldm/mmaudio/vae/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/comfy/ldm/mmaudio/vae/activations.py b/comfy/ldm/mmaudio/vae/activations.py
new file mode 100644
index 000000000..db9192e3e
--- /dev/null
+++ b/comfy/ldm/mmaudio/vae/activations.py
@@ -0,0 +1,120 @@
+# Implementation adapted from https://github.com/EdwardDixon/snake under the MIT license.
+# LICENSE is in incl_licenses directory.
+
+import torch
+from torch import nn, sin, pow
+from torch.nn import Parameter
+import comfy.model_management
+
+class Snake(nn.Module):
+ '''
+ Implementation of a sine-based periodic activation function
+ Shape:
+ - Input: (B, C, T)
+ - Output: (B, C, T), same shape as the input
+ Parameters:
+ - alpha - trainable parameter
+ References:
+ - This activation function is from this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
+ https://arxiv.org/abs/2006.08195
+ Examples:
+ >>> a1 = snake(256)
+ >>> x = torch.randn(256)
+ >>> x = a1(x)
+ '''
+ def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False):
+ '''
+ Initialization.
+ INPUT:
+ - in_features: shape of the input
+ - alpha: trainable parameter
+ alpha is initialized to 1 by default, higher values = higher-frequency.
+ alpha will be trained along with the rest of your model.
+ '''
+ super(Snake, self).__init__()
+ self.in_features = in_features
+
+ # initialize alpha
+ self.alpha_logscale = alpha_logscale
+ if self.alpha_logscale:
+ self.alpha = Parameter(torch.empty(in_features))
+ else:
+ self.alpha = Parameter(torch.empty(in_features))
+
+ self.alpha.requires_grad = alpha_trainable
+
+ self.no_div_by_zero = 0.000000001
+
+ def forward(self, x):
+ '''
+ Forward pass of the function.
+ Applies the function to the input elementwise.
+ Snake ∶= x + 1/a * sin^2 (xa)
+ '''
+ alpha = comfy.model_management.cast_to(self.alpha, dtype=x.dtype, device=x.device).unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T]
+ if self.alpha_logscale:
+ alpha = torch.exp(alpha)
+ x = x + (1.0 / (alpha + self.no_div_by_zero)) * pow(sin(x * alpha), 2)
+
+ return x
+
+
+class SnakeBeta(nn.Module):
+ '''
+ A modified Snake function which uses separate parameters for the magnitude of the periodic components
+ Shape:
+ - Input: (B, C, T)
+ - Output: (B, C, T), same shape as the input
+ Parameters:
+ - alpha - trainable parameter that controls frequency
+ - beta - trainable parameter that controls magnitude
+ References:
+ - This activation function is a modified version based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda:
+ https://arxiv.org/abs/2006.08195
+ Examples:
+ >>> a1 = snakebeta(256)
+ >>> x = torch.randn(256)
+ >>> x = a1(x)
+ '''
+ def __init__(self, in_features, alpha=1.0, alpha_trainable=True, alpha_logscale=False):
+ '''
+ Initialization.
+ INPUT:
+ - in_features: shape of the input
+ - alpha - trainable parameter that controls frequency
+ - beta - trainable parameter that controls magnitude
+ alpha is initialized to 1 by default, higher values = higher-frequency.
+ beta is initialized to 1 by default, higher values = higher-magnitude.
+ alpha will be trained along with the rest of your model.
+ '''
+ super(SnakeBeta, self).__init__()
+ self.in_features = in_features
+
+ # initialize alpha
+ self.alpha_logscale = alpha_logscale
+ if self.alpha_logscale:
+ self.alpha = Parameter(torch.empty(in_features))
+ self.beta = Parameter(torch.empty(in_features))
+ else:
+ self.alpha = Parameter(torch.empty(in_features))
+ self.beta = Parameter(torch.empty(in_features))
+
+ self.alpha.requires_grad = alpha_trainable
+ self.beta.requires_grad = alpha_trainable
+
+ self.no_div_by_zero = 0.000000001
+
+ def forward(self, x):
+ '''
+ Forward pass of the function.
+ Applies the function to the input elementwise.
+ SnakeBeta ∶= x + 1/b * sin^2 (xa)
+ '''
+ alpha = comfy.model_management.cast_to(self.alpha, dtype=x.dtype, device=x.device).unsqueeze(0).unsqueeze(-1) # line up with x to [B, C, T]
+ beta = comfy.model_management.cast_to(self.beta, dtype=x.dtype, device=x.device).unsqueeze(0).unsqueeze(-1)
+ if self.alpha_logscale:
+ alpha = torch.exp(alpha)
+ beta = torch.exp(beta)
+ x = x + (1.0 / (beta + self.no_div_by_zero)) * pow(sin(x * alpha), 2)
+
+ return x
diff --git a/comfy/ldm/mmaudio/vae/alias_free_torch.py b/comfy/ldm/mmaudio/vae/alias_free_torch.py
new file mode 100644
index 000000000..35c70b897
--- /dev/null
+++ b/comfy/ldm/mmaudio/vae/alias_free_torch.py
@@ -0,0 +1,157 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+import comfy.model_management
+
+if 'sinc' in dir(torch):
+ sinc = torch.sinc
+else:
+ # This code is adopted from adefossez's julius.core.sinc under the MIT License
+ # https://adefossez.github.io/julius/julius/core.html
+ # LICENSE is in incl_licenses directory.
+ def sinc(x: torch.Tensor):
+ """
+ Implementation of sinc, i.e. sin(pi * x) / (pi * x)
+ __Warning__: Different to julius.sinc, the input is multiplied by `pi`!
+ """
+ return torch.where(x == 0,
+ torch.tensor(1., device=x.device, dtype=x.dtype),
+ torch.sin(math.pi * x) / math.pi / x)
+
+
+# This code is adopted from adefossez's julius.lowpass.LowPassFilters under the MIT License
+# https://adefossez.github.io/julius/julius/lowpass.html
+# LICENSE is in incl_licenses directory.
+def kaiser_sinc_filter1d(cutoff, half_width, kernel_size): # return filter [1,1,kernel_size]
+ even = (kernel_size % 2 == 0)
+ half_size = kernel_size // 2
+
+ #For kaiser window
+ delta_f = 4 * half_width
+ A = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
+ if A > 50.:
+ beta = 0.1102 * (A - 8.7)
+ elif A >= 21.:
+ beta = 0.5842 * (A - 21)**0.4 + 0.07886 * (A - 21.)
+ else:
+ beta = 0.
+ window = torch.kaiser_window(kernel_size, beta=beta, periodic=False)
+
+ # ratio = 0.5/cutoff -> 2 * cutoff = 1 / ratio
+ if even:
+ time = (torch.arange(-half_size, half_size) + 0.5)
+ else:
+ time = torch.arange(kernel_size) - half_size
+ if cutoff == 0:
+ filter_ = torch.zeros_like(time)
+ else:
+ filter_ = 2 * cutoff * window * sinc(2 * cutoff * time)
+ # Normalize filter to have sum = 1, otherwise we will have a small leakage
+ # of the constant component in the input signal.
+ filter_ /= filter_.sum()
+ filter = filter_.view(1, 1, kernel_size)
+
+ return filter
+
+
+class LowPassFilter1d(nn.Module):
+ def __init__(self,
+ cutoff=0.5,
+ half_width=0.6,
+ stride: int = 1,
+ padding: bool = True,
+ padding_mode: str = 'replicate',
+ kernel_size: int = 12):
+ # kernel_size should be even number for stylegan3 setup,
+ # in this implementation, odd number is also possible.
+ super().__init__()
+ if cutoff < -0.:
+ raise ValueError("Minimum cutoff must be larger than zero.")
+ if cutoff > 0.5:
+ raise ValueError("A cutoff above 0.5 does not make sense.")
+ self.kernel_size = kernel_size
+ self.even = (kernel_size % 2 == 0)
+ self.pad_left = kernel_size // 2 - int(self.even)
+ self.pad_right = kernel_size // 2
+ self.stride = stride
+ self.padding = padding
+ self.padding_mode = padding_mode
+ filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size)
+ self.register_buffer("filter", filter)
+
+ #input [B, C, T]
+ def forward(self, x):
+ _, C, _ = x.shape
+
+ if self.padding:
+ x = F.pad(x, (self.pad_left, self.pad_right),
+ mode=self.padding_mode)
+ out = F.conv1d(x, comfy.model_management.cast_to(self.filter.expand(C, -1, -1), dtype=x.dtype, device=x.device),
+ stride=self.stride, groups=C)
+
+ return out
+
+
+class UpSample1d(nn.Module):
+ def __init__(self, ratio=2, kernel_size=None):
+ super().__init__()
+ self.ratio = ratio
+ self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+ self.stride = ratio
+ self.pad = self.kernel_size // ratio - 1
+ self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
+ self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
+ filter = kaiser_sinc_filter1d(cutoff=0.5 / ratio,
+ half_width=0.6 / ratio,
+ kernel_size=self.kernel_size)
+ self.register_buffer("filter", filter)
+
+ # x: [B, C, T]
+ def forward(self, x):
+ _, C, _ = x.shape
+
+ x = F.pad(x, (self.pad, self.pad), mode='replicate')
+ x = self.ratio * F.conv_transpose1d(
+ x, comfy.model_management.cast_to(self.filter.expand(C, -1, -1), dtype=x.dtype, device=x.device), stride=self.stride, groups=C)
+ x = x[..., self.pad_left:-self.pad_right]
+
+ return x
+
+
+class DownSample1d(nn.Module):
+ def __init__(self, ratio=2, kernel_size=None):
+ super().__init__()
+ self.ratio = ratio
+ self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+ self.lowpass = LowPassFilter1d(cutoff=0.5 / ratio,
+ half_width=0.6 / ratio,
+ stride=ratio,
+ kernel_size=self.kernel_size)
+
+ def forward(self, x):
+ xx = self.lowpass(x)
+
+ return xx
+
+class Activation1d(nn.Module):
+ def __init__(self,
+ activation,
+ up_ratio: int = 2,
+ down_ratio: int = 2,
+ up_kernel_size: int = 12,
+ down_kernel_size: int = 12):
+ super().__init__()
+ self.up_ratio = up_ratio
+ self.down_ratio = down_ratio
+ self.act = activation
+ self.upsample = UpSample1d(up_ratio, up_kernel_size)
+ self.downsample = DownSample1d(down_ratio, down_kernel_size)
+
+ # x: [B,C,T]
+ def forward(self, x):
+ x = self.upsample(x)
+ x = self.act(x)
+ x = self.downsample(x)
+
+ return x
diff --git a/comfy/ldm/mmaudio/vae/autoencoder.py b/comfy/ldm/mmaudio/vae/autoencoder.py
new file mode 100644
index 000000000..cbb9de302
--- /dev/null
+++ b/comfy/ldm/mmaudio/vae/autoencoder.py
@@ -0,0 +1,156 @@
+from typing import Literal
+
+import torch
+import torch.nn as nn
+
+from .distributions import DiagonalGaussianDistribution
+from .vae import VAE_16k
+from .bigvgan import BigVGANVocoder
+import logging
+
+try:
+ import torchaudio
+except:
+ logging.warning("torchaudio missing, MMAudio VAE model will be broken")
+
+def dynamic_range_compression_torch(x, C=1, clip_val=1e-5, *, norm_fn):
+ return norm_fn(torch.clamp(x, min=clip_val) * C)
+
+
+def spectral_normalize_torch(magnitudes, norm_fn):
+ output = dynamic_range_compression_torch(magnitudes, norm_fn=norm_fn)
+ return output
+
+class MelConverter(nn.Module):
+
+ def __init__(
+ self,
+ *,
+ sampling_rate: float,
+ n_fft: int,
+ num_mels: int,
+ hop_size: int,
+ win_size: int,
+ fmin: float,
+ fmax: float,
+ norm_fn,
+ ):
+ super().__init__()
+ self.sampling_rate = sampling_rate
+ self.n_fft = n_fft
+ self.num_mels = num_mels
+ self.hop_size = hop_size
+ self.win_size = win_size
+ self.fmin = fmin
+ self.fmax = fmax
+ self.norm_fn = norm_fn
+
+ # mel = librosa_mel_fn(sr=self.sampling_rate,
+ # n_fft=self.n_fft,
+ # n_mels=self.num_mels,
+ # fmin=self.fmin,
+ # fmax=self.fmax)
+ # mel_basis = torch.from_numpy(mel).float()
+ mel_basis = torch.empty((num_mels, 1 + n_fft // 2))
+ hann_window = torch.hann_window(self.win_size)
+
+ self.register_buffer('mel_basis', mel_basis)
+ self.register_buffer('hann_window', hann_window)
+
+ @property
+ def device(self):
+ return self.mel_basis.device
+
+ def forward(self, waveform: torch.Tensor, center: bool = False) -> torch.Tensor:
+ waveform = waveform.clamp(min=-1., max=1.).to(self.device)
+
+ waveform = torch.nn.functional.pad(
+ waveform.unsqueeze(1),
+ [int((self.n_fft - self.hop_size) / 2),
+ int((self.n_fft - self.hop_size) / 2)],
+ mode='reflect')
+ waveform = waveform.squeeze(1)
+
+ spec = torch.stft(waveform,
+ self.n_fft,
+ hop_length=self.hop_size,
+ win_length=self.win_size,
+ window=self.hann_window,
+ center=center,
+ pad_mode='reflect',
+ normalized=False,
+ onesided=True,
+ return_complex=True)
+
+ spec = torch.view_as_real(spec)
+ spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9))
+ spec = torch.matmul(self.mel_basis, spec)
+ spec = spectral_normalize_torch(spec, self.norm_fn)
+
+ return spec
+
+class AudioAutoencoder(nn.Module):
+
+ def __init__(
+ self,
+ *,
+ # ckpt_path: str,
+ mode=Literal['16k', '44k'],
+ need_vae_encoder: bool = True,
+ ):
+ super().__init__()
+
+ assert mode == "16k", "Only 16k mode is supported currently."
+ self.mel_converter = MelConverter(sampling_rate=16_000,
+ n_fft=1024,
+ num_mels=80,
+ hop_size=256,
+ win_size=1024,
+ fmin=0,
+ fmax=8_000,
+ norm_fn=torch.log10)
+
+ self.vae = VAE_16k().eval()
+
+ bigvgan_config = {
+ "resblock": "1",
+ "num_mels": 80,
+ "upsample_rates": [4, 4, 2, 2, 2, 2],
+ "upsample_kernel_sizes": [8, 8, 4, 4, 4, 4],
+ "upsample_initial_channel": 1536,
+ "resblock_kernel_sizes": [3, 7, 11],
+ "resblock_dilation_sizes": [
+ [1, 3, 5],
+ [1, 3, 5],
+ [1, 3, 5],
+ ],
+ "activation": "snakebeta",
+ "snake_logscale": True,
+ }
+
+ self.vocoder = BigVGANVocoder(
+ bigvgan_config
+ ).eval()
+
+ @torch.inference_mode()
+ def encode_audio(self, x) -> DiagonalGaussianDistribution:
+ # x: (B * L)
+ mel = self.mel_converter(x)
+ dist = self.vae.encode(mel)
+
+ return dist
+
+ @torch.no_grad()
+ def decode(self, z):
+ mel_decoded = self.vae.decode(z)
+ audio = self.vocoder(mel_decoded)
+
+ audio = torchaudio.functional.resample(audio, 16000, 44100)
+ return audio
+
+ @torch.no_grad()
+ def encode(self, audio):
+ audio = audio.mean(dim=1)
+ audio = torchaudio.functional.resample(audio, 44100, 16000)
+ dist = self.encode_audio(audio)
+ return dist.mean
diff --git a/comfy/ldm/mmaudio/vae/bigvgan.py b/comfy/ldm/mmaudio/vae/bigvgan.py
new file mode 100644
index 000000000..3a24337f6
--- /dev/null
+++ b/comfy/ldm/mmaudio/vae/bigvgan.py
@@ -0,0 +1,219 @@
+# Copyright (c) 2022 NVIDIA CORPORATION.
+# Licensed under the MIT license.
+
+# Adapted from https://github.com/jik876/hifi-gan under the MIT license.
+# LICENSE is in incl_licenses directory.
+
+import torch
+import torch.nn as nn
+from types import SimpleNamespace
+from . import activations
+from .alias_free_torch import Activation1d
+import comfy.ops
+ops = comfy.ops.disable_weight_init
+
+def get_padding(kernel_size, dilation=1):
+ return int((kernel_size * dilation - dilation) / 2)
+
+class AMPBlock1(torch.nn.Module):
+
+ def __init__(self, h, channels, kernel_size=3, dilation=(1, 3, 5), activation=None):
+ super(AMPBlock1, self).__init__()
+ self.h = h
+
+ self.convs1 = nn.ModuleList([
+ ops.Conv1d(channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=dilation[0],
+ padding=get_padding(kernel_size, dilation[0])),
+ ops.Conv1d(channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=dilation[1],
+ padding=get_padding(kernel_size, dilation[1])),
+ ops.Conv1d(channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=dilation[2],
+ padding=get_padding(kernel_size, dilation[2]))
+ ])
+
+ self.convs2 = nn.ModuleList([
+ ops.Conv1d(channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=1,
+ padding=get_padding(kernel_size, 1)),
+ ops.Conv1d(channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=1,
+ padding=get_padding(kernel_size, 1)),
+ ops.Conv1d(channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=1,
+ padding=get_padding(kernel_size, 1))
+ ])
+
+ self.num_layers = len(self.convs1) + len(self.convs2) # total number of conv layers
+
+ if activation == 'snake': # periodic nonlinearity with snake function and anti-aliasing
+ self.activations = nn.ModuleList([
+ Activation1d(
+ activation=activations.Snake(channels, alpha_logscale=h.snake_logscale))
+ for _ in range(self.num_layers)
+ ])
+ elif activation == 'snakebeta': # periodic nonlinearity with snakebeta function and anti-aliasing
+ self.activations = nn.ModuleList([
+ Activation1d(
+ activation=activations.SnakeBeta(channels, alpha_logscale=h.snake_logscale))
+ for _ in range(self.num_layers)
+ ])
+ else:
+ raise NotImplementedError(
+ "activation incorrectly specified. check the config file and look for 'activation'."
+ )
+
+ def forward(self, x):
+ acts1, acts2 = self.activations[::2], self.activations[1::2]
+ for c1, c2, a1, a2 in zip(self.convs1, self.convs2, acts1, acts2):
+ xt = a1(x)
+ xt = c1(xt)
+ xt = a2(xt)
+ xt = c2(xt)
+ x = xt + x
+
+ return x
+
+
+class AMPBlock2(torch.nn.Module):
+
+ def __init__(self, h, channels, kernel_size=3, dilation=(1, 3), activation=None):
+ super(AMPBlock2, self).__init__()
+ self.h = h
+
+ self.convs = nn.ModuleList([
+ ops.Conv1d(channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=dilation[0],
+ padding=get_padding(kernel_size, dilation[0])),
+ ops.Conv1d(channels,
+ channels,
+ kernel_size,
+ 1,
+ dilation=dilation[1],
+ padding=get_padding(kernel_size, dilation[1]))
+ ])
+
+ self.num_layers = len(self.convs) # total number of conv layers
+
+ if activation == 'snake': # periodic nonlinearity with snake function and anti-aliasing
+ self.activations = nn.ModuleList([
+ Activation1d(
+ activation=activations.Snake(channels, alpha_logscale=h.snake_logscale))
+ for _ in range(self.num_layers)
+ ])
+ elif activation == 'snakebeta': # periodic nonlinearity with snakebeta function and anti-aliasing
+ self.activations = nn.ModuleList([
+ Activation1d(
+ activation=activations.SnakeBeta(channels, alpha_logscale=h.snake_logscale))
+ for _ in range(self.num_layers)
+ ])
+ else:
+ raise NotImplementedError(
+ "activation incorrectly specified. check the config file and look for 'activation'."
+ )
+
+ def forward(self, x):
+ for c, a in zip(self.convs, self.activations):
+ xt = a(x)
+ xt = c(xt)
+ x = xt + x
+
+ return x
+
+
+class BigVGANVocoder(torch.nn.Module):
+ # this is our main BigVGAN model. Applies anti-aliased periodic activation for resblocks.
+ def __init__(self, h):
+ super().__init__()
+ if isinstance(h, dict):
+ h = SimpleNamespace(**h)
+ self.h = h
+
+ self.num_kernels = len(h.resblock_kernel_sizes)
+ self.num_upsamples = len(h.upsample_rates)
+
+ # pre conv
+ self.conv_pre = ops.Conv1d(h.num_mels, h.upsample_initial_channel, 7, 1, padding=3)
+
+ # define which AMPBlock to use. BigVGAN uses AMPBlock1 as default
+ resblock = AMPBlock1 if h.resblock == '1' else AMPBlock2
+
+ # transposed conv-based upsamplers. does not apply anti-aliasing
+ self.ups = nn.ModuleList()
+ for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)):
+ self.ups.append(
+ nn.ModuleList([
+ ops.ConvTranspose1d(h.upsample_initial_channel // (2**i),
+ h.upsample_initial_channel // (2**(i + 1)),
+ k,
+ u,
+ padding=(k - u) // 2)
+ ]))
+
+ # residual blocks using anti-aliased multi-periodicity composition modules (AMP)
+ self.resblocks = nn.ModuleList()
+ for i in range(len(self.ups)):
+ ch = h.upsample_initial_channel // (2**(i + 1))
+ for j, (k, d) in enumerate(zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)):
+ self.resblocks.append(resblock(h, ch, k, d, activation=h.activation))
+
+ # post conv
+ if h.activation == "snake": # periodic nonlinearity with snake function and anti-aliasing
+ activation_post = activations.Snake(ch, alpha_logscale=h.snake_logscale)
+ self.activation_post = Activation1d(activation=activation_post)
+ elif h.activation == "snakebeta": # periodic nonlinearity with snakebeta function and anti-aliasing
+ activation_post = activations.SnakeBeta(ch, alpha_logscale=h.snake_logscale)
+ self.activation_post = Activation1d(activation=activation_post)
+ else:
+ raise NotImplementedError(
+ "activation incorrectly specified. check the config file and look for 'activation'."
+ )
+
+ self.conv_post = ops.Conv1d(ch, 1, 7, 1, padding=3)
+
+
+ def forward(self, x):
+ # pre conv
+ x = self.conv_pre(x)
+
+ for i in range(self.num_upsamples):
+ # upsampling
+ for i_up in range(len(self.ups[i])):
+ x = self.ups[i][i_up](x)
+ # AMP blocks
+ xs = None
+ for j in range(self.num_kernels):
+ if xs is None:
+ xs = self.resblocks[i * self.num_kernels + j](x)
+ else:
+ xs += self.resblocks[i * self.num_kernels + j](x)
+ x = xs / self.num_kernels
+
+ # post conv
+ x = self.activation_post(x)
+ x = self.conv_post(x)
+ x = torch.tanh(x)
+
+ return x
diff --git a/comfy/ldm/mmaudio/vae/distributions.py b/comfy/ldm/mmaudio/vae/distributions.py
new file mode 100644
index 000000000..df987c5ec
--- /dev/null
+++ b/comfy/ldm/mmaudio/vae/distributions.py
@@ -0,0 +1,92 @@
+import torch
+import numpy as np
+
+
+class AbstractDistribution:
+ def sample(self):
+ raise NotImplementedError()
+
+ def mode(self):
+ raise NotImplementedError()
+
+
+class DiracDistribution(AbstractDistribution):
+ def __init__(self, value):
+ self.value = value
+
+ def sample(self):
+ return self.value
+
+ def mode(self):
+ return self.value
+
+
+class DiagonalGaussianDistribution(object):
+ def __init__(self, parameters, deterministic=False):
+ self.parameters = parameters
+ self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
+ self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
+ self.deterministic = deterministic
+ self.std = torch.exp(0.5 * self.logvar)
+ self.var = torch.exp(self.logvar)
+ if self.deterministic:
+ self.var = self.std = torch.zeros_like(self.mean, device=self.parameters.device)
+
+ def sample(self):
+ x = self.mean + self.std * torch.randn(self.mean.shape, device=self.parameters.device)
+ return x
+
+ def kl(self, other=None):
+ if self.deterministic:
+ return torch.Tensor([0.])
+ else:
+ if other is None:
+ return 0.5 * torch.sum(torch.pow(self.mean, 2)
+ + self.var - 1.0 - self.logvar,
+ dim=[1, 2, 3])
+ else:
+ return 0.5 * torch.sum(
+ torch.pow(self.mean - other.mean, 2) / other.var
+ + self.var / other.var - 1.0 - self.logvar + other.logvar,
+ dim=[1, 2, 3])
+
+ def nll(self, sample, dims=[1,2,3]):
+ if self.deterministic:
+ return torch.Tensor([0.])
+ logtwopi = np.log(2.0 * np.pi)
+ return 0.5 * torch.sum(
+ logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var,
+ dim=dims)
+
+ def mode(self):
+ return self.mean
+
+
+def normal_kl(mean1, logvar1, mean2, logvar2):
+ """
+ source: https://github.com/openai/guided-diffusion/blob/27c20a8fab9cb472df5d6bdd6c8d11c8f430b924/guided_diffusion/losses.py#L12
+ Compute the KL divergence between two gaussians.
+ Shapes are automatically broadcasted, so batches can be compared to
+ scalars, among other use cases.
+ """
+ tensor = None
+ for obj in (mean1, logvar1, mean2, logvar2):
+ if isinstance(obj, torch.Tensor):
+ tensor = obj
+ break
+ assert tensor is not None, "at least one argument must be a Tensor"
+
+ # Force variances to be Tensors. Broadcasting helps convert scalars to
+ # Tensors, but it does not work for torch.exp().
+ logvar1, logvar2 = [
+ x if isinstance(x, torch.Tensor) else torch.tensor(x).to(tensor)
+ for x in (logvar1, logvar2)
+ ]
+
+ return 0.5 * (
+ -1.0
+ + logvar2
+ - logvar1
+ + torch.exp(logvar1 - logvar2)
+ + ((mean1 - mean2) ** 2) * torch.exp(-logvar2)
+ )
diff --git a/comfy/ldm/mmaudio/vae/vae.py b/comfy/ldm/mmaudio/vae/vae.py
new file mode 100644
index 000000000..62f24606c
--- /dev/null
+++ b/comfy/ldm/mmaudio/vae/vae.py
@@ -0,0 +1,358 @@
+import logging
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from .vae_modules import (AttnBlock1D, Downsample1D, ResnetBlock1D,
+ Upsample1D, nonlinearity)
+from .distributions import DiagonalGaussianDistribution
+
+import comfy.ops
+ops = comfy.ops.disable_weight_init
+
+log = logging.getLogger()
+
+DATA_MEAN_80D = [
+ -1.6058, -1.3676, -1.2520, -1.2453, -1.2078, -1.2224, -1.2419, -1.2439, -1.2922, -1.2927,
+ -1.3170, -1.3543, -1.3401, -1.3836, -1.3907, -1.3912, -1.4313, -1.4152, -1.4527, -1.4728,
+ -1.4568, -1.5101, -1.5051, -1.5172, -1.5623, -1.5373, -1.5746, -1.5687, -1.6032, -1.6131,
+ -1.6081, -1.6331, -1.6489, -1.6489, -1.6700, -1.6738, -1.6953, -1.6969, -1.7048, -1.7280,
+ -1.7361, -1.7495, -1.7658, -1.7814, -1.7889, -1.8064, -1.8221, -1.8377, -1.8417, -1.8643,
+ -1.8857, -1.8929, -1.9173, -1.9379, -1.9531, -1.9673, -1.9824, -2.0042, -2.0215, -2.0436,
+ -2.0766, -2.1064, -2.1418, -2.1855, -2.2319, -2.2767, -2.3161, -2.3572, -2.3954, -2.4282,
+ -2.4659, -2.5072, -2.5552, -2.6074, -2.6584, -2.7107, -2.7634, -2.8266, -2.8981, -2.9673
+]
+
+DATA_STD_80D = [
+ 1.0291, 1.0411, 1.0043, 0.9820, 0.9677, 0.9543, 0.9450, 0.9392, 0.9343, 0.9297, 0.9276, 0.9263,
+ 0.9242, 0.9254, 0.9232, 0.9281, 0.9263, 0.9315, 0.9274, 0.9247, 0.9277, 0.9199, 0.9188, 0.9194,
+ 0.9160, 0.9161, 0.9146, 0.9161, 0.9100, 0.9095, 0.9145, 0.9076, 0.9066, 0.9095, 0.9032, 0.9043,
+ 0.9038, 0.9011, 0.9019, 0.9010, 0.8984, 0.8983, 0.8986, 0.8961, 0.8962, 0.8978, 0.8962, 0.8973,
+ 0.8993, 0.8976, 0.8995, 0.9016, 0.8982, 0.8972, 0.8974, 0.8949, 0.8940, 0.8947, 0.8936, 0.8939,
+ 0.8951, 0.8956, 0.9017, 0.9167, 0.9436, 0.9690, 1.0003, 1.0225, 1.0381, 1.0491, 1.0545, 1.0604,
+ 1.0761, 1.0929, 1.1089, 1.1196, 1.1176, 1.1156, 1.1117, 1.1070
+]
+
+DATA_MEAN_128D = [
+ -3.3462, -2.6723, -2.4893, -2.3143, -2.2664, -2.3317, -2.1802, -2.4006, -2.2357, -2.4597,
+ -2.3717, -2.4690, -2.5142, -2.4919, -2.6610, -2.5047, -2.7483, -2.5926, -2.7462, -2.7033,
+ -2.7386, -2.8112, -2.7502, -2.9594, -2.7473, -3.0035, -2.8891, -2.9922, -2.9856, -3.0157,
+ -3.1191, -2.9893, -3.1718, -3.0745, -3.1879, -3.2310, -3.1424, -3.2296, -3.2791, -3.2782,
+ -3.2756, -3.3134, -3.3509, -3.3750, -3.3951, -3.3698, -3.4505, -3.4509, -3.5089, -3.4647,
+ -3.5536, -3.5788, -3.5867, -3.6036, -3.6400, -3.6747, -3.7072, -3.7279, -3.7283, -3.7795,
+ -3.8259, -3.8447, -3.8663, -3.9182, -3.9605, -3.9861, -4.0105, -4.0373, -4.0762, -4.1121,
+ -4.1488, -4.1874, -4.2461, -4.3170, -4.3639, -4.4452, -4.5282, -4.6297, -4.7019, -4.7960,
+ -4.8700, -4.9507, -5.0303, -5.0866, -5.1634, -5.2342, -5.3242, -5.4053, -5.4927, -5.5712,
+ -5.6464, -5.7052, -5.7619, -5.8410, -5.9188, -6.0103, -6.0955, -6.1673, -6.2362, -6.3120,
+ -6.3926, -6.4797, -6.5565, -6.6511, -6.8130, -6.9961, -7.1275, -7.2457, -7.3576, -7.4663,
+ -7.6136, -7.7469, -7.8815, -8.0132, -8.1515, -8.3071, -8.4722, -8.7418, -9.3975, -9.6628,
+ -9.7671, -9.8863, -9.9992, -10.0860, -10.1709, -10.5418, -11.2795, -11.3861
+]
+
+DATA_STD_128D = [
+ 2.3804, 2.4368, 2.3772, 2.3145, 2.2803, 2.2510, 2.2316, 2.2083, 2.1996, 2.1835, 2.1769, 2.1659,
+ 2.1631, 2.1618, 2.1540, 2.1606, 2.1571, 2.1567, 2.1612, 2.1579, 2.1679, 2.1683, 2.1634, 2.1557,
+ 2.1668, 2.1518, 2.1415, 2.1449, 2.1406, 2.1350, 2.1313, 2.1415, 2.1281, 2.1352, 2.1219, 2.1182,
+ 2.1327, 2.1195, 2.1137, 2.1080, 2.1179, 2.1036, 2.1087, 2.1036, 2.1015, 2.1068, 2.0975, 2.0991,
+ 2.0902, 2.1015, 2.0857, 2.0920, 2.0893, 2.0897, 2.0910, 2.0881, 2.0925, 2.0873, 2.0960, 2.0900,
+ 2.0957, 2.0958, 2.0978, 2.0936, 2.0886, 2.0905, 2.0845, 2.0855, 2.0796, 2.0840, 2.0813, 2.0817,
+ 2.0838, 2.0840, 2.0917, 2.1061, 2.1431, 2.1976, 2.2482, 2.3055, 2.3700, 2.4088, 2.4372, 2.4609,
+ 2.4731, 2.4847, 2.5072, 2.5451, 2.5772, 2.6147, 2.6529, 2.6596, 2.6645, 2.6726, 2.6803, 2.6812,
+ 2.6899, 2.6916, 2.6931, 2.6998, 2.7062, 2.7262, 2.7222, 2.7158, 2.7041, 2.7485, 2.7491, 2.7451,
+ 2.7485, 2.7233, 2.7297, 2.7233, 2.7145, 2.6958, 2.6788, 2.6439, 2.6007, 2.4786, 2.2469, 2.1877,
+ 2.1392, 2.0717, 2.0107, 1.9676, 1.9140, 1.7102, 0.9101, 0.7164
+]
+
+
+class VAE(nn.Module):
+
+ def __init__(
+ self,
+ *,
+ data_dim: int,
+ embed_dim: int,
+ hidden_dim: int,
+ ):
+ super().__init__()
+
+ if data_dim == 80:
+ self.data_mean = nn.Buffer(torch.tensor(DATA_MEAN_80D, dtype=torch.float32))
+ self.data_std = nn.Buffer(torch.tensor(DATA_STD_80D, dtype=torch.float32))
+ elif data_dim == 128:
+ self.data_mean = nn.Buffer(torch.tensor(DATA_MEAN_128D, dtype=torch.float32))
+ self.data_std = nn.Buffer(torch.tensor(DATA_STD_128D, dtype=torch.float32))
+
+ self.data_mean = self.data_mean.view(1, -1, 1)
+ self.data_std = self.data_std.view(1, -1, 1)
+
+ self.encoder = Encoder1D(
+ dim=hidden_dim,
+ ch_mult=(1, 2, 4),
+ num_res_blocks=2,
+ attn_layers=[3],
+ down_layers=[0],
+ in_dim=data_dim,
+ embed_dim=embed_dim,
+ )
+ self.decoder = Decoder1D(
+ dim=hidden_dim,
+ ch_mult=(1, 2, 4),
+ num_res_blocks=2,
+ attn_layers=[3],
+ down_layers=[0],
+ in_dim=data_dim,
+ out_dim=data_dim,
+ embed_dim=embed_dim,
+ )
+
+ self.embed_dim = embed_dim
+ # self.quant_conv = nn.Conv1d(2 * embed_dim, 2 * embed_dim, 1)
+ # self.post_quant_conv = nn.Conv1d(embed_dim, embed_dim, 1)
+
+ self.initialize_weights()
+
+ def initialize_weights(self):
+ pass
+
+ def encode(self, x: torch.Tensor, normalize: bool = True) -> DiagonalGaussianDistribution:
+ if normalize:
+ x = self.normalize(x)
+ moments = self.encoder(x)
+ posterior = DiagonalGaussianDistribution(moments)
+ return posterior
+
+ def decode(self, z: torch.Tensor, unnormalize: bool = True) -> torch.Tensor:
+ dec = self.decoder(z)
+ if unnormalize:
+ dec = self.unnormalize(dec)
+ return dec
+
+ def normalize(self, x: torch.Tensor) -> torch.Tensor:
+ return (x - comfy.model_management.cast_to(self.data_mean, dtype=x.dtype, device=x.device)) / comfy.model_management.cast_to(self.data_std, dtype=x.dtype, device=x.device)
+
+ def unnormalize(self, x: torch.Tensor) -> torch.Tensor:
+ return x * comfy.model_management.cast_to(self.data_std, dtype=x.dtype, device=x.device) + comfy.model_management.cast_to(self.data_mean, dtype=x.dtype, device=x.device)
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ sample_posterior: bool = True,
+ rng: Optional[torch.Generator] = None,
+ normalize: bool = True,
+ unnormalize: bool = True,
+ ) -> tuple[torch.Tensor, DiagonalGaussianDistribution]:
+
+ posterior = self.encode(x, normalize=normalize)
+ if sample_posterior:
+ z = posterior.sample(rng)
+ else:
+ z = posterior.mode()
+ dec = self.decode(z, unnormalize=unnormalize)
+ return dec, posterior
+
+ def load_weights(self, src_dict) -> None:
+ self.load_state_dict(src_dict, strict=True)
+
+ @property
+ def device(self) -> torch.device:
+ return next(self.parameters()).device
+
+ def get_last_layer(self):
+ return self.decoder.conv_out.weight
+
+ def remove_weight_norm(self):
+ return self
+
+
+class Encoder1D(nn.Module):
+
+ def __init__(self,
+ *,
+ dim: int,
+ ch_mult: tuple[int] = (1, 2, 4, 8),
+ num_res_blocks: int,
+ attn_layers: list[int] = [],
+ down_layers: list[int] = [],
+ resamp_with_conv: bool = True,
+ in_dim: int,
+ embed_dim: int,
+ double_z: bool = True,
+ kernel_size: int = 3,
+ clip_act: float = 256.0):
+ super().__init__()
+ self.dim = dim
+ self.num_layers = len(ch_mult)
+ self.num_res_blocks = num_res_blocks
+ self.in_channels = in_dim
+ self.clip_act = clip_act
+ self.down_layers = down_layers
+ self.attn_layers = attn_layers
+ self.conv_in = ops.Conv1d(in_dim, self.dim, kernel_size=kernel_size, padding=kernel_size // 2, bias=False)
+
+ in_ch_mult = (1, ) + tuple(ch_mult)
+ self.in_ch_mult = in_ch_mult
+ # downsampling
+ self.down = nn.ModuleList()
+ for i_level in range(self.num_layers):
+ block = nn.ModuleList()
+ attn = nn.ModuleList()
+ block_in = dim * in_ch_mult[i_level]
+ block_out = dim * ch_mult[i_level]
+ for i_block in range(self.num_res_blocks):
+ block.append(
+ ResnetBlock1D(in_dim=block_in,
+ out_dim=block_out,
+ kernel_size=kernel_size,
+ use_norm=True))
+ block_in = block_out
+ if i_level in attn_layers:
+ attn.append(AttnBlock1D(block_in))
+ down = nn.Module()
+ down.block = block
+ down.attn = attn
+ if i_level in down_layers:
+ down.downsample = Downsample1D(block_in, resamp_with_conv)
+ self.down.append(down)
+
+ # middle
+ self.mid = nn.Module()
+ self.mid.block_1 = ResnetBlock1D(in_dim=block_in,
+ out_dim=block_in,
+ kernel_size=kernel_size,
+ use_norm=True)
+ self.mid.attn_1 = AttnBlock1D(block_in)
+ self.mid.block_2 = ResnetBlock1D(in_dim=block_in,
+ out_dim=block_in,
+ kernel_size=kernel_size,
+ use_norm=True)
+
+ # end
+ self.conv_out = ops.Conv1d(block_in,
+ 2 * embed_dim if double_z else embed_dim,
+ kernel_size=kernel_size, padding=kernel_size // 2, bias=False)
+
+ self.learnable_gain = nn.Parameter(torch.zeros([]))
+
+ def forward(self, x):
+
+ # downsampling
+ h = self.conv_in(x)
+ for i_level in range(self.num_layers):
+ for i_block in range(self.num_res_blocks):
+ h = self.down[i_level].block[i_block](h)
+ if len(self.down[i_level].attn) > 0:
+ h = self.down[i_level].attn[i_block](h)
+ h = h.clamp(-self.clip_act, self.clip_act)
+ if i_level in self.down_layers:
+ h = self.down[i_level].downsample(h)
+
+ # middle
+ h = self.mid.block_1(h)
+ h = self.mid.attn_1(h)
+ h = self.mid.block_2(h)
+ h = h.clamp(-self.clip_act, self.clip_act)
+
+ # end
+ h = nonlinearity(h)
+ h = self.conv_out(h) * (self.learnable_gain + 1)
+ return h
+
+
+class Decoder1D(nn.Module):
+
+ def __init__(self,
+ *,
+ dim: int,
+ out_dim: int,
+ ch_mult: tuple[int] = (1, 2, 4, 8),
+ num_res_blocks: int,
+ attn_layers: list[int] = [],
+ down_layers: list[int] = [],
+ kernel_size: int = 3,
+ resamp_with_conv: bool = True,
+ in_dim: int,
+ embed_dim: int,
+ clip_act: float = 256.0):
+ super().__init__()
+ self.ch = dim
+ self.num_layers = len(ch_mult)
+ self.num_res_blocks = num_res_blocks
+ self.in_channels = in_dim
+ self.clip_act = clip_act
+ self.down_layers = [i + 1 for i in down_layers] # each downlayer add one
+
+ # compute in_ch_mult, block_in and curr_res at lowest res
+ block_in = dim * ch_mult[self.num_layers - 1]
+
+ # z to block_in
+ self.conv_in = ops.Conv1d(embed_dim, block_in, kernel_size=kernel_size, padding=kernel_size // 2, bias=False)
+
+ # middle
+ self.mid = nn.Module()
+ self.mid.block_1 = ResnetBlock1D(in_dim=block_in, out_dim=block_in, use_norm=True)
+ self.mid.attn_1 = AttnBlock1D(block_in)
+ self.mid.block_2 = ResnetBlock1D(in_dim=block_in, out_dim=block_in, use_norm=True)
+
+ # upsampling
+ self.up = nn.ModuleList()
+ for i_level in reversed(range(self.num_layers)):
+ block = nn.ModuleList()
+ attn = nn.ModuleList()
+ block_out = dim * ch_mult[i_level]
+ for i_block in range(self.num_res_blocks + 1):
+ block.append(ResnetBlock1D(in_dim=block_in, out_dim=block_out, use_norm=True))
+ block_in = block_out
+ if i_level in attn_layers:
+ attn.append(AttnBlock1D(block_in))
+ up = nn.Module()
+ up.block = block
+ up.attn = attn
+ if i_level in self.down_layers:
+ up.upsample = Upsample1D(block_in, resamp_with_conv)
+ self.up.insert(0, up) # prepend to get consistent order
+
+ # end
+ self.conv_out = ops.Conv1d(block_in, out_dim, kernel_size=kernel_size, padding=kernel_size // 2, bias=False)
+ self.learnable_gain = nn.Parameter(torch.zeros([]))
+
+ def forward(self, z):
+ # z to block_in
+ h = self.conv_in(z)
+
+ # middle
+ h = self.mid.block_1(h)
+ h = self.mid.attn_1(h)
+ h = self.mid.block_2(h)
+ h = h.clamp(-self.clip_act, self.clip_act)
+
+ # upsampling
+ for i_level in reversed(range(self.num_layers)):
+ for i_block in range(self.num_res_blocks + 1):
+ h = self.up[i_level].block[i_block](h)
+ if len(self.up[i_level].attn) > 0:
+ h = self.up[i_level].attn[i_block](h)
+ h = h.clamp(-self.clip_act, self.clip_act)
+ if i_level in self.down_layers:
+ h = self.up[i_level].upsample(h)
+
+ h = nonlinearity(h)
+ h = self.conv_out(h) * (self.learnable_gain + 1)
+ return h
+
+
+def VAE_16k(**kwargs) -> VAE:
+ return VAE(data_dim=80, embed_dim=20, hidden_dim=384, **kwargs)
+
+
+def VAE_44k(**kwargs) -> VAE:
+ return VAE(data_dim=128, embed_dim=40, hidden_dim=512, **kwargs)
+
+
+def get_my_vae(name: str, **kwargs) -> VAE:
+ if name == '16k':
+ return VAE_16k(**kwargs)
+ if name == '44k':
+ return VAE_44k(**kwargs)
+ raise ValueError(f'Unknown model: {name}')
+
diff --git a/comfy/ldm/mmaudio/vae/vae_modules.py b/comfy/ldm/mmaudio/vae/vae_modules.py
new file mode 100644
index 000000000..3ad05134b
--- /dev/null
+++ b/comfy/ldm/mmaudio/vae/vae_modules.py
@@ -0,0 +1,121 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from comfy.ldm.modules.diffusionmodules.model import vae_attention
+import math
+import comfy.ops
+ops = comfy.ops.disable_weight_init
+
+def nonlinearity(x):
+ # swish
+ return torch.nn.functional.silu(x) / 0.596
+
+def mp_sum(a, b, t=0.5):
+ return a.lerp(b, t) / math.sqrt((1 - t)**2 + t**2)
+
+def normalize(x, dim=None, eps=1e-4):
+ if dim is None:
+ dim = list(range(1, x.ndim))
+ norm = torch.linalg.vector_norm(x, dim=dim, keepdim=True, dtype=torch.float32)
+ norm = torch.add(eps, norm, alpha=math.sqrt(norm.numel() / x.numel()))
+ return x / norm.to(x.dtype)
+
+class ResnetBlock1D(nn.Module):
+
+ def __init__(self, *, in_dim, out_dim=None, conv_shortcut=False, kernel_size=3, use_norm=True):
+ super().__init__()
+ self.in_dim = in_dim
+ out_dim = in_dim if out_dim is None else out_dim
+ self.out_dim = out_dim
+ self.use_conv_shortcut = conv_shortcut
+ self.use_norm = use_norm
+
+ self.conv1 = ops.Conv1d(in_dim, out_dim, kernel_size=kernel_size, padding=kernel_size // 2, bias=False)
+ self.conv2 = ops.Conv1d(out_dim, out_dim, kernel_size=kernel_size, padding=kernel_size // 2, bias=False)
+ if self.in_dim != self.out_dim:
+ if self.use_conv_shortcut:
+ self.conv_shortcut = ops.Conv1d(in_dim, out_dim, kernel_size=kernel_size, padding=kernel_size // 2, bias=False)
+ else:
+ self.nin_shortcut = ops.Conv1d(in_dim, out_dim, kernel_size=1, padding=0, bias=False)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+
+ # pixel norm
+ if self.use_norm:
+ x = normalize(x, dim=1)
+
+ h = x
+ h = nonlinearity(h)
+ h = self.conv1(h)
+
+ h = nonlinearity(h)
+ h = self.conv2(h)
+
+ if self.in_dim != self.out_dim:
+ if self.use_conv_shortcut:
+ x = self.conv_shortcut(x)
+ else:
+ x = self.nin_shortcut(x)
+
+ return mp_sum(x, h, t=0.3)
+
+
+class AttnBlock1D(nn.Module):
+
+ def __init__(self, in_channels, num_heads=1):
+ super().__init__()
+ self.in_channels = in_channels
+
+ self.num_heads = num_heads
+ self.qkv = ops.Conv1d(in_channels, in_channels * 3, kernel_size=1, padding=0, bias=False)
+ self.proj_out = ops.Conv1d(in_channels, in_channels, kernel_size=1, padding=0, bias=False)
+ self.optimized_attention = vae_attention()
+
+ def forward(self, x):
+ h = x
+ y = self.qkv(h)
+ y = y.reshape(y.shape[0], -1, 3, y.shape[-1])
+ q, k, v = normalize(y, dim=1).unbind(2)
+
+ h = self.optimized_attention(q, k, v)
+ h = self.proj_out(h)
+
+ return mp_sum(x, h, t=0.3)
+
+
+class Upsample1D(nn.Module):
+
+ def __init__(self, in_channels, with_conv):
+ super().__init__()
+ self.with_conv = with_conv
+ if self.with_conv:
+ self.conv = ops.Conv1d(in_channels, in_channels, kernel_size=3, padding=1, bias=False)
+
+ def forward(self, x):
+ x = F.interpolate(x, scale_factor=2.0, mode='nearest-exact') # support 3D tensor(B,C,T)
+ if self.with_conv:
+ x = self.conv(x)
+ return x
+
+
+class Downsample1D(nn.Module):
+
+ def __init__(self, in_channels, with_conv):
+ super().__init__()
+ self.with_conv = with_conv
+ if self.with_conv:
+ # no asymmetric padding in torch conv, must do it ourselves
+ self.conv1 = ops.Conv1d(in_channels, in_channels, kernel_size=1, padding=0, bias=False)
+ self.conv2 = ops.Conv1d(in_channels, in_channels, kernel_size=1, padding=0, bias=False)
+
+ def forward(self, x):
+
+ if self.with_conv:
+ x = self.conv1(x)
+
+ x = F.avg_pool1d(x, kernel_size=2, stride=2)
+
+ if self.with_conv:
+ x = self.conv2(x)
+
+ return x
diff --git a/comfy/ldm/models/autoencoder.py b/comfy/ldm/models/autoencoder.py
index 13bd6e16b..611d36a1b 100644
--- a/comfy/ldm/models/autoencoder.py
+++ b/comfy/ldm/models/autoencoder.py
@@ -26,6 +26,12 @@ class DiagonalGaussianRegularizer(torch.nn.Module):
z = posterior.mode()
return z, None
+class EmptyRegularizer(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+
+ def forward(self, z: torch.Tensor) -> Tuple[torch.Tensor, dict]:
+ return z, None
class AbstractAutoencoder(torch.nn.Module):
"""
diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py
index 35d2270ee..7437e0567 100644
--- a/comfy/ldm/modules/attention.py
+++ b/comfy/ldm/modules/attention.py
@@ -5,8 +5,9 @@ import torch
import torch.nn.functional as F
from torch import nn, einsum
from einops import rearrange, repeat
-from typing import Optional
+from typing import Optional, Any, Callable, Union
import logging
+import functools
from .diffusionmodules.util import AlphaBlender, timestep_embedding
from .sub_quadratic_attention import efficient_dot_product_attention
@@ -17,23 +18,45 @@ if model_management.xformers_enabled():
import xformers
import xformers.ops
-if model_management.sage_attention_enabled():
- try:
- from sageattention import sageattn
- except ModuleNotFoundError as e:
+SAGE_ATTENTION_IS_AVAILABLE = False
+try:
+ from sageattention import sageattn
+ SAGE_ATTENTION_IS_AVAILABLE = True
+except ImportError as e:
+ if model_management.sage_attention_enabled():
if e.name == "sageattention":
logging.error(f"\n\nTo use the `--use-sage-attention` feature, the `sageattention` package must be installed first.\ncommand:\n\t{sys.executable} -m pip install sageattention")
else:
raise e
exit(-1)
-if model_management.flash_attention_enabled():
- try:
- from flash_attn import flash_attn_func
- except ModuleNotFoundError:
+FLASH_ATTENTION_IS_AVAILABLE = False
+try:
+ from flash_attn import flash_attn_func
+ FLASH_ATTENTION_IS_AVAILABLE = True
+except ImportError:
+ if model_management.flash_attention_enabled():
logging.error(f"\n\nTo use the `--use-flash-attention` feature, the `flash-attn` package must be installed first.\ncommand:\n\t{sys.executable} -m pip install flash-attn")
exit(-1)
+REGISTERED_ATTENTION_FUNCTIONS = {}
+def register_attention_function(name: str, func: Callable):
+ # avoid replacing existing functions
+ if name not in REGISTERED_ATTENTION_FUNCTIONS:
+ REGISTERED_ATTENTION_FUNCTIONS[name] = func
+ else:
+ logging.warning(f"Attention function {name} already registered, skipping registration.")
+
+def get_attention_function(name: str, default: Any=...) -> Union[Callable, None]:
+ if name == "optimized":
+ return optimized_attention
+ elif name not in REGISTERED_ATTENTION_FUNCTIONS:
+ if default is ...:
+ raise KeyError(f"Attention function {name} not found.")
+ else:
+ return default
+ return REGISTERED_ATTENTION_FUNCTIONS[name]
+
from comfy.cli_args import args
import comfy.ops
ops = comfy.ops.disable_weight_init
@@ -91,7 +114,27 @@ class FeedForward(nn.Module):
def Normalize(in_channels, dtype=None, device=None):
return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True, dtype=dtype, device=device)
-def attention_basic(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False):
+
+def wrap_attn(func):
+ @functools.wraps(func)
+ def wrapper(*args, **kwargs):
+ remove_attn_wrapper_key = False
+ try:
+ if "_inside_attn_wrapper" not in kwargs:
+ transformer_options = kwargs.get("transformer_options", None)
+ remove_attn_wrapper_key = True
+ kwargs["_inside_attn_wrapper"] = True
+ if transformer_options is not None:
+ if "optimized_attention_override" in transformer_options:
+ return transformer_options["optimized_attention_override"](func, *args, **kwargs)
+ return func(*args, **kwargs)
+ finally:
+ if remove_attn_wrapper_key:
+ del kwargs["_inside_attn_wrapper"]
+ return wrapper
+
+@wrap_attn
+def attention_basic(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs):
attn_precision = get_attn_precision(attn_precision, q.dtype)
if skip_reshape:
@@ -159,8 +202,8 @@ def attention_basic(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
)
return out
-
-def attention_sub_quad(query, key, value, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False):
+@wrap_attn
+def attention_sub_quad(query, key, value, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs):
attn_precision = get_attn_precision(attn_precision, query.dtype)
if skip_reshape:
@@ -230,7 +273,8 @@ def attention_sub_quad(query, key, value, heads, mask=None, attn_precision=None,
hidden_states = hidden_states.unflatten(0, (-1, heads)).transpose(1,2).flatten(start_dim=2)
return hidden_states
-def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False):
+@wrap_attn
+def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs):
attn_precision = get_attn_precision(attn_precision, q.dtype)
if skip_reshape:
@@ -359,7 +403,8 @@ try:
except:
pass
-def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False):
+@wrap_attn
+def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs):
b = q.shape[0]
dim_head = q.shape[-1]
# check to make sure xformers isn't broken
@@ -374,7 +419,7 @@ def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_resh
disabled_xformers = True
if disabled_xformers:
- return attention_pytorch(q, k, v, heads, mask, skip_reshape=skip_reshape)
+ return attention_pytorch(q, k, v, heads, mask, skip_reshape=skip_reshape, **kwargs)
if skip_reshape:
# b h k d -> b k h d
@@ -427,8 +472,8 @@ else:
#TODO: other GPUs ?
SDP_BATCH_LIMIT = 2**31
-
-def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False):
+@wrap_attn
+def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs):
if skip_reshape:
b, _, _, dim_head = q.shape
else:
@@ -448,7 +493,7 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha
mask = mask.unsqueeze(1)
if SDP_BATCH_LIMIT >= b:
- out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)
+ out = comfy.ops.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)
if not skip_output_reshape:
out = (
out.transpose(1, 2).reshape(b, -1, heads * dim_head)
@@ -461,7 +506,7 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha
if mask.shape[0] > 1:
m = mask[i : i + SDP_BATCH_LIMIT]
- out[i : i + SDP_BATCH_LIMIT] = torch.nn.functional.scaled_dot_product_attention(
+ out[i : i + SDP_BATCH_LIMIT] = comfy.ops.scaled_dot_product_attention(
q[i : i + SDP_BATCH_LIMIT],
k[i : i + SDP_BATCH_LIMIT],
v[i : i + SDP_BATCH_LIMIT],
@@ -470,8 +515,8 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha
).transpose(1, 2).reshape(-1, q.shape[2], heads * dim_head)
return out
-
-def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False):
+@wrap_attn
+def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs):
if skip_reshape:
b, _, _, dim_head = q.shape
tensor_layout = "HND"
@@ -501,7 +546,7 @@ def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=
lambda t: t.transpose(1, 2),
(q, k, v),
)
- return attention_pytorch(q, k, v, heads, mask=mask, skip_reshape=True, skip_output_reshape=skip_output_reshape)
+ return attention_pytorch(q, k, v, heads, mask=mask, skip_reshape=True, skip_output_reshape=skip_output_reshape, **kwargs)
if tensor_layout == "HND":
if not skip_output_reshape:
@@ -534,8 +579,8 @@ except AttributeError as error:
dropout_p: float = 0.0, causal: bool = False) -> torch.Tensor:
assert False, f"Could not define flash_attn_wrapper: {FLASH_ATTN_ERROR}"
-
-def attention_flash(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False):
+@wrap_attn
+def attention_flash(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs):
if skip_reshape:
b, _, _, dim_head = q.shape
else:
@@ -555,7 +600,8 @@ def attention_flash(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
mask = mask.unsqueeze(1)
try:
- assert mask is None
+ if mask is not None:
+ raise RuntimeError("Mask must not be set for Flash attention")
out = flash_attn_wrapper(
q.transpose(1, 2),
k.transpose(1, 2),
@@ -597,6 +643,19 @@ else:
optimized_attention_masked = optimized_attention
+
+# register core-supported attention functions
+if SAGE_ATTENTION_IS_AVAILABLE:
+ register_attention_function("sage", attention_sage)
+if FLASH_ATTENTION_IS_AVAILABLE:
+ register_attention_function("flash", attention_flash)
+if model_management.xformers_enabled():
+ register_attention_function("xformers", attention_xformers)
+register_attention_function("pytorch", attention_pytorch)
+register_attention_function("sub_quad", attention_sub_quad)
+register_attention_function("split", attention_split)
+
+
def optimized_attention_for_device(device, mask=False, small_input=False):
if small_input:
if model_management.pytorch_attention_enabled():
@@ -629,7 +688,7 @@ class CrossAttention(nn.Module):
self.to_out = nn.Sequential(operations.Linear(inner_dim, query_dim, dtype=dtype, device=device), nn.Dropout(dropout))
- def forward(self, x, context=None, value=None, mask=None):
+ def forward(self, x, context=None, value=None, mask=None, transformer_options={}):
q = self.to_q(x)
context = default(context, x)
k = self.to_k(context)
@@ -640,9 +699,9 @@ class CrossAttention(nn.Module):
v = self.to_v(context)
if mask is None:
- out = optimized_attention(q, k, v, self.heads, attn_precision=self.attn_precision)
+ out = optimized_attention(q, k, v, self.heads, attn_precision=self.attn_precision, transformer_options=transformer_options)
else:
- out = optimized_attention_masked(q, k, v, self.heads, mask, attn_precision=self.attn_precision)
+ out = optimized_attention_masked(q, k, v, self.heads, mask, attn_precision=self.attn_precision, transformer_options=transformer_options)
return self.to_out(out)
@@ -746,7 +805,7 @@ class BasicTransformerBlock(nn.Module):
n = attn1_replace_patch[block_attn1](n, context_attn1, value_attn1, extra_options)
n = self.attn1.to_out(n)
else:
- n = self.attn1(n, context=context_attn1, value=value_attn1)
+ n = self.attn1(n, context=context_attn1, value=value_attn1, transformer_options=transformer_options)
if "attn1_output_patch" in transformer_patches:
patch = transformer_patches["attn1_output_patch"]
@@ -786,7 +845,7 @@ class BasicTransformerBlock(nn.Module):
n = attn2_replace_patch[block_attn2](n, context_attn2, value_attn2, extra_options)
n = self.attn2.to_out(n)
else:
- n = self.attn2(n, context=context_attn2, value=value_attn2)
+ n = self.attn2(n, context=context_attn2, value=value_attn2, transformer_options=transformer_options)
if "attn2_output_patch" in transformer_patches:
patch = transformer_patches["attn2_output_patch"]
@@ -1017,7 +1076,7 @@ class SpatialVideoTransformer(SpatialTransformer):
B, S, C = x_mix.shape
x_mix = rearrange(x_mix, "(b t) s c -> (b s) t c", t=timesteps)
- x_mix = mix_block(x_mix, context=time_context) #TODO: transformer_options
+ x_mix = mix_block(x_mix, context=time_context, transformer_options=transformer_options)
x_mix = rearrange(
x_mix, "(b s) t c -> (b t) s c", s=S, b=B // timesteps, c=C, t=timesteps
)
diff --git a/comfy/ldm/modules/diffusionmodules/mmdit.py b/comfy/ldm/modules/diffusionmodules/mmdit.py
index eaf3e73a4..42f406f1a 100644
--- a/comfy/ldm/modules/diffusionmodules/mmdit.py
+++ b/comfy/ldm/modules/diffusionmodules/mmdit.py
@@ -109,7 +109,7 @@ class PatchEmbed(nn.Module):
def modulate(x, shift, scale):
if shift is None:
shift = torch.zeros_like(scale)
- return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+ return torch.addcmul(shift.unsqueeze(1), x, 1+ scale.unsqueeze(1))
#################################################################################
@@ -564,10 +564,7 @@ class DismantledBlock(nn.Module):
assert not self.pre_only
attn1 = self.attn.post_attention(attn)
attn2 = self.attn2.post_attention(attn2)
- out1 = gate_msa.unsqueeze(1) * attn1
- out2 = gate_msa2.unsqueeze(1) * attn2
- x = x + out1
- x = x + out2
+ x = gate_cat(x, gate_msa, gate_msa2, attn1, attn2)
x = x + gate_mlp.unsqueeze(1) * self.mlp(
modulate(self.norm2(x), shift_mlp, scale_mlp)
)
@@ -594,6 +591,11 @@ class DismantledBlock(nn.Module):
)
return self.post_attention(attn, *intermediates)
+def gate_cat(x, gate_msa, gate_msa2, attn1, attn2):
+ out1 = gate_msa.unsqueeze(1) * attn1
+ out2 = gate_msa2.unsqueeze(1) * attn2
+ x = torch.stack([x, out1, out2], dim=0).sum(dim=0)
+ return x
def block_mixing(*args, use_checkpoint=True, **kwargs):
if use_checkpoint:
@@ -604,7 +606,7 @@ def block_mixing(*args, use_checkpoint=True, **kwargs):
return _block_mixing(*args, **kwargs)
-def _block_mixing(context, x, context_block, x_block, c):
+def _block_mixing(context, x, context_block, x_block, c, transformer_options={}):
context_qkv, context_intermediates = context_block.pre_attention(context, c)
if x_block.x_block_self_attn:
@@ -620,6 +622,7 @@ def _block_mixing(context, x, context_block, x_block, c):
attn = optimized_attention(
qkv[0], qkv[1], qkv[2],
heads=x_block.attn.num_heads,
+ transformer_options=transformer_options,
)
context_attn, x_attn = (
attn[:, : context_qkv[0].shape[1]],
@@ -635,6 +638,7 @@ def _block_mixing(context, x, context_block, x_block, c):
attn2 = optimized_attention(
x_qkv2[0], x_qkv2[1], x_qkv2[2],
heads=x_block.attn2.num_heads,
+ transformer_options=transformer_options,
)
x = x_block.post_attention_x(x_attn, attn2, *x_intermediates)
else:
@@ -956,10 +960,10 @@ class MMDiT(nn.Module):
if ("double_block", i) in blocks_replace:
def block_wrap(args):
out = {}
- out["txt"], out["img"] = self.joint_blocks[i](args["txt"], args["img"], c=args["vec"])
+ out["txt"], out["img"] = self.joint_blocks[i](args["txt"], args["img"], c=args["vec"], transformer_options=args["transformer_options"])
return out
- out = blocks_replace[("double_block", i)]({"img": x, "txt": context, "vec": c_mod}, {"original_block": block_wrap})
+ out = blocks_replace[("double_block", i)]({"img": x, "txt": context, "vec": c_mod, "transformer_options": transformer_options}, {"original_block": block_wrap})
context = out["txt"]
x = out["img"]
else:
@@ -968,6 +972,7 @@ class MMDiT(nn.Module):
x,
c=c_mod,
use_checkpoint=self.use_checkpoint,
+ transformer_options=transformer_options,
)
if control is not None:
control_o = control.get("output")
diff --git a/comfy/ldm/modules/diffusionmodules/model.py b/comfy/ldm/modules/diffusionmodules/model.py
index 8162742cf..4245eedca 100644
--- a/comfy/ldm/modules/diffusionmodules/model.py
+++ b/comfy/ldm/modules/diffusionmodules/model.py
@@ -36,7 +36,7 @@ def get_timestep_embedding(timesteps, embedding_dim):
def nonlinearity(x):
# swish
- return x*torch.sigmoid(x)
+ return torch.nn.functional.silu(x)
def Normalize(in_channels, num_groups=32):
@@ -145,7 +145,7 @@ class Downsample(nn.Module):
class ResnetBlock(nn.Module):
def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False,
- dropout, temb_channels=512, conv_op=ops.Conv2d):
+ dropout=0.0, temb_channels=512, conv_op=ops.Conv2d, norm_op=Normalize):
super().__init__()
self.in_channels = in_channels
out_channels = in_channels if out_channels is None else out_channels
@@ -153,7 +153,7 @@ class ResnetBlock(nn.Module):
self.use_conv_shortcut = conv_shortcut
self.swish = torch.nn.SiLU(inplace=True)
- self.norm1 = Normalize(in_channels)
+ self.norm1 = norm_op(in_channels)
self.conv1 = conv_op(in_channels,
out_channels,
kernel_size=3,
@@ -162,7 +162,7 @@ class ResnetBlock(nn.Module):
if temb_channels > 0:
self.temb_proj = ops.Linear(temb_channels,
out_channels)
- self.norm2 = Normalize(out_channels)
+ self.norm2 = norm_op(out_channels)
self.dropout = torch.nn.Dropout(dropout, inplace=True)
self.conv2 = conv_op(out_channels,
out_channels,
@@ -183,7 +183,7 @@ class ResnetBlock(nn.Module):
stride=1,
padding=0)
- def forward(self, x, temb):
+ def forward(self, x, temb=None):
h = x
h = self.norm1(h)
h = self.swish(h)
@@ -285,7 +285,7 @@ def pytorch_attention(q, k, v):
)
try:
- out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False)
+ out = comfy.ops.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False)
out = out.transpose(2, 3).reshape(orig_shape)
except model_management.OOM_EXCEPTION:
logging.warning("scaled_dot_product_attention OOMed: switched to slice attention")
@@ -305,11 +305,11 @@ def vae_attention():
return normal_attention
class AttnBlock(nn.Module):
- def __init__(self, in_channels, conv_op=ops.Conv2d):
+ def __init__(self, in_channels, conv_op=ops.Conv2d, norm_op=Normalize):
super().__init__()
self.in_channels = in_channels
- self.norm = Normalize(in_channels)
+ self.norm = norm_op(in_channels)
self.q = conv_op(in_channels,
in_channels,
kernel_size=1,
diff --git a/comfy/ldm/omnigen/omnigen2.py b/comfy/ldm/omnigen/omnigen2.py
index 4884449f8..82edc92da 100644
--- a/comfy/ldm/omnigen/omnigen2.py
+++ b/comfy/ldm/omnigen/omnigen2.py
@@ -120,7 +120,7 @@ class Attention(nn.Module):
nn.Dropout(0.0)
)
- def forward(self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, image_rotary_emb: Optional[torch.Tensor] = None) -> torch.Tensor:
+ def forward(self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, image_rotary_emb: Optional[torch.Tensor] = None, transformer_options={}) -> torch.Tensor:
batch_size, sequence_length, _ = hidden_states.shape
query = self.to_q(hidden_states)
@@ -146,7 +146,7 @@ class Attention(nn.Module):
key = key.repeat_interleave(self.heads // self.kv_heads, dim=1)
value = value.repeat_interleave(self.heads // self.kv_heads, dim=1)
- hidden_states = optimized_attention_masked(query, key, value, self.heads, attention_mask, skip_reshape=True)
+ hidden_states = optimized_attention_masked(query, key, value, self.heads, attention_mask, skip_reshape=True, transformer_options=transformer_options)
hidden_states = self.to_out[0](hidden_states)
return hidden_states
@@ -182,16 +182,16 @@ class OmniGen2TransformerBlock(nn.Module):
self.norm2 = operations.RMSNorm(dim, eps=norm_eps, dtype=dtype, device=device)
self.ffn_norm2 = operations.RMSNorm(dim, eps=norm_eps, dtype=dtype, device=device)
- def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, image_rotary_emb: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
+ def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, image_rotary_emb: torch.Tensor, temb: Optional[torch.Tensor] = None, transformer_options={}) -> torch.Tensor:
if self.modulation:
norm_hidden_states, gate_msa, scale_mlp, gate_mlp = self.norm1(hidden_states, temb)
- attn_output = self.attn(norm_hidden_states, norm_hidden_states, attention_mask, image_rotary_emb)
+ attn_output = self.attn(norm_hidden_states, norm_hidden_states, attention_mask, image_rotary_emb, transformer_options=transformer_options)
hidden_states = hidden_states + gate_msa.unsqueeze(1).tanh() * self.norm2(attn_output)
mlp_output = self.feed_forward(self.ffn_norm1(hidden_states) * (1 + scale_mlp.unsqueeze(1)))
hidden_states = hidden_states + gate_mlp.unsqueeze(1).tanh() * self.ffn_norm2(mlp_output)
else:
norm_hidden_states = self.norm1(hidden_states)
- attn_output = self.attn(norm_hidden_states, norm_hidden_states, attention_mask, image_rotary_emb)
+ attn_output = self.attn(norm_hidden_states, norm_hidden_states, attention_mask, image_rotary_emb, transformer_options=transformer_options)
hidden_states = hidden_states + self.norm2(attn_output)
mlp_output = self.feed_forward(self.ffn_norm1(hidden_states))
hidden_states = hidden_states + self.ffn_norm2(mlp_output)
@@ -390,7 +390,7 @@ class OmniGen2Transformer2DModel(nn.Module):
ref_img_sizes, img_sizes,
)
- def img_patch_embed_and_refine(self, hidden_states, ref_image_hidden_states, padded_img_mask, padded_ref_img_mask, noise_rotary_emb, ref_img_rotary_emb, l_effective_ref_img_len, l_effective_img_len, temb):
+ def img_patch_embed_and_refine(self, hidden_states, ref_image_hidden_states, padded_img_mask, padded_ref_img_mask, noise_rotary_emb, ref_img_rotary_emb, l_effective_ref_img_len, l_effective_img_len, temb, transformer_options={}):
batch_size = len(hidden_states)
hidden_states = self.x_embedder(hidden_states)
@@ -405,17 +405,17 @@ class OmniGen2Transformer2DModel(nn.Module):
shift += ref_img_len
for layer in self.noise_refiner:
- hidden_states = layer(hidden_states, padded_img_mask, noise_rotary_emb, temb)
+ hidden_states = layer(hidden_states, padded_img_mask, noise_rotary_emb, temb, transformer_options=transformer_options)
if ref_image_hidden_states is not None:
for layer in self.ref_image_refiner:
- ref_image_hidden_states = layer(ref_image_hidden_states, padded_ref_img_mask, ref_img_rotary_emb, temb)
+ ref_image_hidden_states = layer(ref_image_hidden_states, padded_ref_img_mask, ref_img_rotary_emb, temb, transformer_options=transformer_options)
hidden_states = torch.cat([ref_image_hidden_states, hidden_states], dim=1)
return hidden_states
- def forward(self, x, timesteps, context, num_tokens, ref_latents=None, attention_mask=None, **kwargs):
+ def forward(self, x, timesteps, context, num_tokens, ref_latents=None, attention_mask=None, transformer_options={}, **kwargs):
B, C, H, W = x.shape
hidden_states = comfy.ldm.common_dit.pad_to_patch_size(x, (self.patch_size, self.patch_size))
_, _, H_padded, W_padded = hidden_states.shape
@@ -444,7 +444,7 @@ class OmniGen2Transformer2DModel(nn.Module):
)
for layer in self.context_refiner:
- text_hidden_states = layer(text_hidden_states, text_attention_mask, context_rotary_emb)
+ text_hidden_states = layer(text_hidden_states, text_attention_mask, context_rotary_emb, transformer_options=transformer_options)
img_len = hidden_states.shape[1]
combined_img_hidden_states = self.img_patch_embed_and_refine(
@@ -453,13 +453,14 @@ class OmniGen2Transformer2DModel(nn.Module):
noise_rotary_emb, ref_img_rotary_emb,
l_effective_ref_img_len, l_effective_img_len,
temb,
+ transformer_options=transformer_options,
)
hidden_states = torch.cat([text_hidden_states, combined_img_hidden_states], dim=1)
attention_mask = None
for layer in self.layers:
- hidden_states = layer(hidden_states, attention_mask, rotary_emb, temb)
+ hidden_states = layer(hidden_states, attention_mask, rotary_emb, temb, transformer_options=transformer_options)
hidden_states = self.norm_out(hidden_states, temb)
diff --git a/comfy/ldm/pixart/pixartms.py b/comfy/ldm/pixart/pixartms.py
index 7d4eebdce..d1ac49d84 100644
--- a/comfy/ldm/pixart/pixartms.py
+++ b/comfy/ldm/pixart/pixartms.py
@@ -1,256 +1,256 @@
-# Based on:
-# https://github.com/PixArt-alpha/PixArt-alpha [Apache 2.0 license]
-# https://github.com/PixArt-alpha/PixArt-sigma [Apache 2.0 license]
-import torch
-import torch.nn as nn
-
-from .blocks import (
- t2i_modulate,
- CaptionEmbedder,
- AttentionKVCompress,
- MultiHeadCrossAttention,
- T2IFinalLayer,
- SizeEmbedder,
-)
-from comfy.ldm.modules.diffusionmodules.mmdit import TimestepEmbedder, PatchEmbed, Mlp, get_1d_sincos_pos_embed_from_grid_torch
-
-
-def get_2d_sincos_pos_embed_torch(embed_dim, w, h, pe_interpolation=1.0, base_size=16, device=None, dtype=torch.float32):
- grid_h, grid_w = torch.meshgrid(
- torch.arange(h, device=device, dtype=dtype) / (h/base_size) / pe_interpolation,
- torch.arange(w, device=device, dtype=dtype) / (w/base_size) / pe_interpolation,
- indexing='ij'
- )
- emb_h = get_1d_sincos_pos_embed_from_grid_torch(embed_dim // 2, grid_h, device=device, dtype=dtype)
- emb_w = get_1d_sincos_pos_embed_from_grid_torch(embed_dim // 2, grid_w, device=device, dtype=dtype)
- emb = torch.cat([emb_w, emb_h], dim=1) # (H*W, D)
- return emb
-
-class PixArtMSBlock(nn.Module):
- """
- A PixArt block with adaptive layer norm zero (adaLN-Zero) conditioning.
- """
- def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, drop_path=0., input_size=None,
- sampling=None, sr_ratio=1, qk_norm=False, dtype=None, device=None, operations=None, **block_kwargs):
- super().__init__()
- self.hidden_size = hidden_size
- self.norm1 = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
- self.attn = AttentionKVCompress(
- hidden_size, num_heads=num_heads, qkv_bias=True, sampling=sampling, sr_ratio=sr_ratio,
- qk_norm=qk_norm, dtype=dtype, device=device, operations=operations, **block_kwargs
- )
- self.cross_attn = MultiHeadCrossAttention(
- hidden_size, num_heads, dtype=dtype, device=device, operations=operations, **block_kwargs
- )
- self.norm2 = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
- # to be compatible with lower version pytorch
- approx_gelu = lambda: nn.GELU(approximate="tanh")
- self.mlp = Mlp(
- in_features=hidden_size, hidden_features=int(hidden_size * mlp_ratio), act_layer=approx_gelu,
- dtype=dtype, device=device, operations=operations
- )
- self.scale_shift_table = nn.Parameter(torch.randn(6, hidden_size) / hidden_size ** 0.5)
-
- def forward(self, x, y, t, mask=None, HW=None, **kwargs):
- B, N, C = x.shape
-
- shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (self.scale_shift_table[None].to(dtype=x.dtype, device=x.device) + t.reshape(B, 6, -1)).chunk(6, dim=1)
- x = x + (gate_msa * self.attn(t2i_modulate(self.norm1(x), shift_msa, scale_msa), HW=HW))
- x = x + self.cross_attn(x, y, mask)
- x = x + (gate_mlp * self.mlp(t2i_modulate(self.norm2(x), shift_mlp, scale_mlp)))
-
- return x
-
-
-### Core PixArt Model ###
-class PixArtMS(nn.Module):
- """
- Diffusion model with a Transformer backbone.
- """
- def __init__(
- self,
- input_size=32,
- patch_size=2,
- in_channels=4,
- hidden_size=1152,
- depth=28,
- num_heads=16,
- mlp_ratio=4.0,
- class_dropout_prob=0.1,
- learn_sigma=True,
- pred_sigma=True,
- drop_path: float = 0.,
- caption_channels=4096,
- pe_interpolation=None,
- pe_precision=None,
- config=None,
- model_max_length=120,
- micro_condition=True,
- qk_norm=False,
- kv_compress_config=None,
- dtype=None,
- device=None,
- operations=None,
- **kwargs,
- ):
- nn.Module.__init__(self)
- self.dtype = dtype
- self.pred_sigma = pred_sigma
- self.in_channels = in_channels
- self.out_channels = in_channels * 2 if pred_sigma else in_channels
- self.patch_size = patch_size
- self.num_heads = num_heads
- self.pe_interpolation = pe_interpolation
- self.pe_precision = pe_precision
- self.hidden_size = hidden_size
- self.depth = depth
-
- approx_gelu = lambda: nn.GELU(approximate="tanh")
- self.t_block = nn.Sequential(
- nn.SiLU(),
- operations.Linear(hidden_size, 6 * hidden_size, bias=True, dtype=dtype, device=device)
- )
- self.x_embedder = PatchEmbed(
- patch_size=patch_size,
- in_chans=in_channels,
- embed_dim=hidden_size,
- bias=True,
- dtype=dtype,
- device=device,
- operations=operations
- )
- self.t_embedder = TimestepEmbedder(
- hidden_size, dtype=dtype, device=device, operations=operations,
- )
- self.y_embedder = CaptionEmbedder(
- in_channels=caption_channels, hidden_size=hidden_size, uncond_prob=class_dropout_prob,
- act_layer=approx_gelu, token_num=model_max_length,
- dtype=dtype, device=device, operations=operations,
- )
-
- self.micro_conditioning = micro_condition
- if self.micro_conditioning:
- self.csize_embedder = SizeEmbedder(hidden_size//3, dtype=dtype, device=device, operations=operations)
- self.ar_embedder = SizeEmbedder(hidden_size//3, dtype=dtype, device=device, operations=operations)
-
- # For fixed sin-cos embedding:
- # num_patches = (input_size // patch_size) * (input_size // patch_size)
- # self.base_size = input_size // self.patch_size
- # self.register_buffer("pos_embed", torch.zeros(1, num_patches, hidden_size))
-
- drop_path = [x.item() for x in torch.linspace(0, drop_path, depth)] # stochastic depth decay rule
- if kv_compress_config is None:
- kv_compress_config = {
- 'sampling': None,
- 'scale_factor': 1,
- 'kv_compress_layer': [],
- }
- self.blocks = nn.ModuleList([
- PixArtMSBlock(
- hidden_size, num_heads, mlp_ratio=mlp_ratio, drop_path=drop_path[i],
- sampling=kv_compress_config['sampling'],
- sr_ratio=int(kv_compress_config['scale_factor']) if i in kv_compress_config['kv_compress_layer'] else 1,
- qk_norm=qk_norm,
- dtype=dtype,
- device=device,
- operations=operations,
- )
- for i in range(depth)
- ])
- self.final_layer = T2IFinalLayer(
- hidden_size, patch_size, self.out_channels, dtype=dtype, device=device, operations=operations
- )
-
- def forward_orig(self, x, timestep, y, mask=None, c_size=None, c_ar=None, **kwargs):
- """
- Original forward pass of PixArt.
- x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images)
- t: (N,) tensor of diffusion timesteps
- y: (N, 1, 120, C) conditioning
- ar: (N, 1): aspect ratio
- cs: (N ,2) size conditioning for height/width
- """
- B, C, H, W = x.shape
- c_res = (H + W) // 2
- pe_interpolation = self.pe_interpolation
- if pe_interpolation is None or self.pe_precision is not None:
- # calculate pe_interpolation on-the-fly
- pe_interpolation = round(c_res / (512/8.0), self.pe_precision or 0)
-
- pos_embed = get_2d_sincos_pos_embed_torch(
- self.hidden_size,
- h=(H // self.patch_size),
- w=(W // self.patch_size),
- pe_interpolation=pe_interpolation,
- base_size=((round(c_res / 64) * 64) // self.patch_size),
- device=x.device,
- dtype=x.dtype,
- ).unsqueeze(0)
-
- x = self.x_embedder(x) + pos_embed # (N, T, D), where T = H * W / patch_size ** 2
- t = self.t_embedder(timestep, x.dtype) # (N, D)
-
- if self.micro_conditioning and (c_size is not None and c_ar is not None):
- bs = x.shape[0]
- c_size = self.csize_embedder(c_size, bs) # (N, D)
- c_ar = self.ar_embedder(c_ar, bs) # (N, D)
- t = t + torch.cat([c_size, c_ar], dim=1)
-
- t0 = self.t_block(t)
- y = self.y_embedder(y, self.training) # (N, D)
-
- if mask is not None:
- if mask.shape[0] != y.shape[0]:
- mask = mask.repeat(y.shape[0] // mask.shape[0], 1)
- mask = mask.squeeze(1).squeeze(1)
- y = y.squeeze(1).masked_select(mask.unsqueeze(-1) != 0).view(1, -1, x.shape[-1])
- y_lens = mask.sum(dim=1).tolist()
- else:
- y_lens = None
- y = y.squeeze(1).view(1, -1, x.shape[-1])
- for block in self.blocks:
- x = block(x, y, t0, y_lens, (H, W), **kwargs) # (N, T, D)
-
- x = self.final_layer(x, t) # (N, T, patch_size ** 2 * out_channels)
- x = self.unpatchify(x, H, W) # (N, out_channels, H, W)
-
- return x
-
- def forward(self, x, timesteps, context, c_size=None, c_ar=None, **kwargs):
- B, C, H, W = x.shape
-
- # Fallback for missing microconds
- if self.micro_conditioning:
- if c_size is None:
- c_size = torch.tensor([H*8, W*8], dtype=x.dtype, device=x.device).repeat(B, 1)
-
- if c_ar is None:
- c_ar = torch.tensor([H/W], dtype=x.dtype, device=x.device).repeat(B, 1)
-
- ## Still accepts the input w/o that dim but returns garbage
- if len(context.shape) == 3:
- context = context.unsqueeze(1)
-
- ## run original forward pass
- out = self.forward_orig(x, timesteps, context, c_size=c_size, c_ar=c_ar)
-
- ## only return EPS
- if self.pred_sigma:
- return out[:, :self.in_channels]
- return out
-
- def unpatchify(self, x, h, w):
- """
- x: (N, T, patch_size**2 * C)
- imgs: (N, H, W, C)
- """
- c = self.out_channels
- p = self.x_embedder.patch_size[0]
- h = h // self.patch_size
- w = w // self.patch_size
- assert h * w == x.shape[1]
-
- x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
- x = torch.einsum('nhwpqc->nchpwq', x)
- imgs = x.reshape(shape=(x.shape[0], c, h * p, w * p))
- return imgs
+# Based on:
+# https://github.com/PixArt-alpha/PixArt-alpha [Apache 2.0 license]
+# https://github.com/PixArt-alpha/PixArt-sigma [Apache 2.0 license]
+import torch
+import torch.nn as nn
+
+from .blocks import (
+ t2i_modulate,
+ CaptionEmbedder,
+ AttentionKVCompress,
+ MultiHeadCrossAttention,
+ T2IFinalLayer,
+ SizeEmbedder,
+)
+from comfy.ldm.modules.diffusionmodules.mmdit import TimestepEmbedder, PatchEmbed, Mlp, get_1d_sincos_pos_embed_from_grid_torch
+
+
+def get_2d_sincos_pos_embed_torch(embed_dim, w, h, pe_interpolation=1.0, base_size=16, device=None, dtype=torch.float32):
+ grid_h, grid_w = torch.meshgrid(
+ torch.arange(h, device=device, dtype=dtype) / (h/base_size) / pe_interpolation,
+ torch.arange(w, device=device, dtype=dtype) / (w/base_size) / pe_interpolation,
+ indexing='ij'
+ )
+ emb_h = get_1d_sincos_pos_embed_from_grid_torch(embed_dim // 2, grid_h, device=device, dtype=dtype)
+ emb_w = get_1d_sincos_pos_embed_from_grid_torch(embed_dim // 2, grid_w, device=device, dtype=dtype)
+ emb = torch.cat([emb_w, emb_h], dim=1) # (H*W, D)
+ return emb
+
+class PixArtMSBlock(nn.Module):
+ """
+ A PixArt block with adaptive layer norm zero (adaLN-Zero) conditioning.
+ """
+ def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, drop_path=0., input_size=None,
+ sampling=None, sr_ratio=1, qk_norm=False, dtype=None, device=None, operations=None, **block_kwargs):
+ super().__init__()
+ self.hidden_size = hidden_size
+ self.norm1 = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
+ self.attn = AttentionKVCompress(
+ hidden_size, num_heads=num_heads, qkv_bias=True, sampling=sampling, sr_ratio=sr_ratio,
+ qk_norm=qk_norm, dtype=dtype, device=device, operations=operations, **block_kwargs
+ )
+ self.cross_attn = MultiHeadCrossAttention(
+ hidden_size, num_heads, dtype=dtype, device=device, operations=operations, **block_kwargs
+ )
+ self.norm2 = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
+ # to be compatible with lower version pytorch
+ approx_gelu = lambda: nn.GELU(approximate="tanh")
+ self.mlp = Mlp(
+ in_features=hidden_size, hidden_features=int(hidden_size * mlp_ratio), act_layer=approx_gelu,
+ dtype=dtype, device=device, operations=operations
+ )
+ self.scale_shift_table = nn.Parameter(torch.randn(6, hidden_size) / hidden_size ** 0.5)
+
+ def forward(self, x, y, t, mask=None, HW=None, **kwargs):
+ B, N, C = x.shape
+
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (self.scale_shift_table[None].to(dtype=x.dtype, device=x.device) + t.reshape(B, 6, -1)).chunk(6, dim=1)
+ x = x + (gate_msa * self.attn(t2i_modulate(self.norm1(x), shift_msa, scale_msa), HW=HW))
+ x = x + self.cross_attn(x, y, mask)
+ x = x + (gate_mlp * self.mlp(t2i_modulate(self.norm2(x), shift_mlp, scale_mlp)))
+
+ return x
+
+
+### Core PixArt Model ###
+class PixArtMS(nn.Module):
+ """
+ Diffusion model with a Transformer backbone.
+ """
+ def __init__(
+ self,
+ input_size=32,
+ patch_size=2,
+ in_channels=4,
+ hidden_size=1152,
+ depth=28,
+ num_heads=16,
+ mlp_ratio=4.0,
+ class_dropout_prob=0.1,
+ learn_sigma=True,
+ pred_sigma=True,
+ drop_path: float = 0.,
+ caption_channels=4096,
+ pe_interpolation=None,
+ pe_precision=None,
+ config=None,
+ model_max_length=120,
+ micro_condition=True,
+ qk_norm=False,
+ kv_compress_config=None,
+ dtype=None,
+ device=None,
+ operations=None,
+ **kwargs,
+ ):
+ nn.Module.__init__(self)
+ self.dtype = dtype
+ self.pred_sigma = pred_sigma
+ self.in_channels = in_channels
+ self.out_channels = in_channels * 2 if pred_sigma else in_channels
+ self.patch_size = patch_size
+ self.num_heads = num_heads
+ self.pe_interpolation = pe_interpolation
+ self.pe_precision = pe_precision
+ self.hidden_size = hidden_size
+ self.depth = depth
+
+ approx_gelu = lambda: nn.GELU(approximate="tanh")
+ self.t_block = nn.Sequential(
+ nn.SiLU(),
+ operations.Linear(hidden_size, 6 * hidden_size, bias=True, dtype=dtype, device=device)
+ )
+ self.x_embedder = PatchEmbed(
+ patch_size=patch_size,
+ in_chans=in_channels,
+ embed_dim=hidden_size,
+ bias=True,
+ dtype=dtype,
+ device=device,
+ operations=operations
+ )
+ self.t_embedder = TimestepEmbedder(
+ hidden_size, dtype=dtype, device=device, operations=operations,
+ )
+ self.y_embedder = CaptionEmbedder(
+ in_channels=caption_channels, hidden_size=hidden_size, uncond_prob=class_dropout_prob,
+ act_layer=approx_gelu, token_num=model_max_length,
+ dtype=dtype, device=device, operations=operations,
+ )
+
+ self.micro_conditioning = micro_condition
+ if self.micro_conditioning:
+ self.csize_embedder = SizeEmbedder(hidden_size//3, dtype=dtype, device=device, operations=operations)
+ self.ar_embedder = SizeEmbedder(hidden_size//3, dtype=dtype, device=device, operations=operations)
+
+ # For fixed sin-cos embedding:
+ # num_patches = (input_size // patch_size) * (input_size // patch_size)
+ # self.base_size = input_size // self.patch_size
+ # self.register_buffer("pos_embed", torch.zeros(1, num_patches, hidden_size))
+
+ drop_path = [x.item() for x in torch.linspace(0, drop_path, depth)] # stochastic depth decay rule
+ if kv_compress_config is None:
+ kv_compress_config = {
+ 'sampling': None,
+ 'scale_factor': 1,
+ 'kv_compress_layer': [],
+ }
+ self.blocks = nn.ModuleList([
+ PixArtMSBlock(
+ hidden_size, num_heads, mlp_ratio=mlp_ratio, drop_path=drop_path[i],
+ sampling=kv_compress_config['sampling'],
+ sr_ratio=int(kv_compress_config['scale_factor']) if i in kv_compress_config['kv_compress_layer'] else 1,
+ qk_norm=qk_norm,
+ dtype=dtype,
+ device=device,
+ operations=operations,
+ )
+ for i in range(depth)
+ ])
+ self.final_layer = T2IFinalLayer(
+ hidden_size, patch_size, self.out_channels, dtype=dtype, device=device, operations=operations
+ )
+
+ def forward_orig(self, x, timestep, y, mask=None, c_size=None, c_ar=None, **kwargs):
+ """
+ Original forward pass of PixArt.
+ x: (N, C, H, W) tensor of spatial inputs (images or latent representations of images)
+ t: (N,) tensor of diffusion timesteps
+ y: (N, 1, 120, C) conditioning
+ ar: (N, 1): aspect ratio
+ cs: (N ,2) size conditioning for height/width
+ """
+ B, C, H, W = x.shape
+ c_res = (H + W) // 2
+ pe_interpolation = self.pe_interpolation
+ if pe_interpolation is None or self.pe_precision is not None:
+ # calculate pe_interpolation on-the-fly
+ pe_interpolation = round(c_res / (512/8.0), self.pe_precision or 0)
+
+ pos_embed = get_2d_sincos_pos_embed_torch(
+ self.hidden_size,
+ h=(H // self.patch_size),
+ w=(W // self.patch_size),
+ pe_interpolation=pe_interpolation,
+ base_size=((round(c_res / 64) * 64) // self.patch_size),
+ device=x.device,
+ dtype=x.dtype,
+ ).unsqueeze(0)
+
+ x = self.x_embedder(x) + pos_embed # (N, T, D), where T = H * W / patch_size ** 2
+ t = self.t_embedder(timestep, x.dtype) # (N, D)
+
+ if self.micro_conditioning and (c_size is not None and c_ar is not None):
+ bs = x.shape[0]
+ c_size = self.csize_embedder(c_size, bs) # (N, D)
+ c_ar = self.ar_embedder(c_ar, bs) # (N, D)
+ t = t + torch.cat([c_size, c_ar], dim=1)
+
+ t0 = self.t_block(t)
+ y = self.y_embedder(y, self.training) # (N, D)
+
+ if mask is not None:
+ if mask.shape[0] != y.shape[0]:
+ mask = mask.repeat(y.shape[0] // mask.shape[0], 1)
+ mask = mask.squeeze(1).squeeze(1)
+ y = y.squeeze(1).masked_select(mask.unsqueeze(-1) != 0).view(1, -1, x.shape[-1])
+ y_lens = mask.sum(dim=1).tolist()
+ else:
+ y_lens = None
+ y = y.squeeze(1).view(1, -1, x.shape[-1])
+ for block in self.blocks:
+ x = block(x, y, t0, y_lens, (H, W), **kwargs) # (N, T, D)
+
+ x = self.final_layer(x, t) # (N, T, patch_size ** 2 * out_channels)
+ x = self.unpatchify(x, H, W) # (N, out_channels, H, W)
+
+ return x
+
+ def forward(self, x, timesteps, context, c_size=None, c_ar=None, **kwargs):
+ B, C, H, W = x.shape
+
+ # Fallback for missing microconds
+ if self.micro_conditioning:
+ if c_size is None:
+ c_size = torch.tensor([H*8, W*8], dtype=x.dtype, device=x.device).repeat(B, 1)
+
+ if c_ar is None:
+ c_ar = torch.tensor([H/W], dtype=x.dtype, device=x.device).repeat(B, 1)
+
+ ## Still accepts the input w/o that dim but returns garbage
+ if len(context.shape) == 3:
+ context = context.unsqueeze(1)
+
+ ## run original forward pass
+ out = self.forward_orig(x, timesteps, context, c_size=c_size, c_ar=c_ar)
+
+ ## only return EPS
+ if self.pred_sigma:
+ return out[:, :self.in_channels]
+ return out
+
+ def unpatchify(self, x, h, w):
+ """
+ x: (N, T, patch_size**2 * C)
+ imgs: (N, H, W, C)
+ """
+ c = self.out_channels
+ p = self.x_embedder.patch_size[0]
+ h = h // self.patch_size
+ w = w // self.patch_size
+ assert h * w == x.shape[1]
+
+ x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
+ x = torch.einsum('nhwpqc->nchpwq', x)
+ imgs = x.reshape(shape=(x.shape[0], c, h * p, w * p))
+ return imgs
diff --git a/comfy/ldm/qwen_image/controlnet.py b/comfy/ldm/qwen_image/controlnet.py
new file mode 100644
index 000000000..92ac3cf0a
--- /dev/null
+++ b/comfy/ldm/qwen_image/controlnet.py
@@ -0,0 +1,77 @@
+import torch
+import math
+
+from .model import QwenImageTransformer2DModel
+
+
+class QwenImageControlNetModel(QwenImageTransformer2DModel):
+ def __init__(
+ self,
+ extra_condition_channels=0,
+ dtype=None,
+ device=None,
+ operations=None,
+ **kwargs
+ ):
+ super().__init__(final_layer=False, dtype=dtype, device=device, operations=operations, **kwargs)
+ self.main_model_double = 60
+
+ # controlnet_blocks
+ self.controlnet_blocks = torch.nn.ModuleList([])
+ for _ in range(len(self.transformer_blocks)):
+ self.controlnet_blocks.append(operations.Linear(self.inner_dim, self.inner_dim, device=device, dtype=dtype))
+ self.controlnet_x_embedder = operations.Linear(self.in_channels + extra_condition_channels, self.inner_dim, device=device, dtype=dtype)
+
+ def forward(
+ self,
+ x,
+ timesteps,
+ context,
+ attention_mask=None,
+ guidance: torch.Tensor = None,
+ ref_latents=None,
+ hint=None,
+ transformer_options={},
+ **kwargs
+ ):
+ timestep = timesteps
+ encoder_hidden_states = context
+ encoder_hidden_states_mask = attention_mask
+
+ hidden_states, img_ids, orig_shape = self.process_img(x)
+ hint, _, _ = self.process_img(hint)
+
+ txt_start = round(max(((x.shape[-1] + (self.patch_size // 2)) // self.patch_size) // 2, ((x.shape[-2] + (self.patch_size // 2)) // self.patch_size) // 2))
+ txt_ids = torch.arange(txt_start, txt_start + context.shape[1], device=x.device).reshape(1, -1, 1).repeat(x.shape[0], 1, 3)
+ ids = torch.cat((txt_ids, img_ids), dim=1)
+ image_rotary_emb = self.pe_embedder(ids).squeeze(1).unsqueeze(2).to(x.dtype)
+ del ids, txt_ids, img_ids
+
+ hidden_states = self.img_in(hidden_states) + self.controlnet_x_embedder(hint)
+ encoder_hidden_states = self.txt_norm(encoder_hidden_states)
+ encoder_hidden_states = self.txt_in(encoder_hidden_states)
+
+ if guidance is not None:
+ guidance = guidance * 1000
+
+ temb = (
+ self.time_text_embed(timestep, hidden_states)
+ if guidance is None
+ else self.time_text_embed(timestep, guidance, hidden_states)
+ )
+
+ repeat = math.ceil(self.main_model_double / len(self.controlnet_blocks))
+
+ controlnet_block_samples = ()
+ for i, block in enumerate(self.transformer_blocks):
+ encoder_hidden_states, hidden_states = block(
+ hidden_states=hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_hidden_states_mask=encoder_hidden_states_mask,
+ temb=temb,
+ image_rotary_emb=image_rotary_emb,
+ )
+
+ controlnet_block_samples = controlnet_block_samples + (self.controlnet_blocks[i](hidden_states),) * repeat
+
+ return {"input": controlnet_block_samples[:self.main_model_double]}
diff --git a/comfy/ldm/qwen_image/model.py b/comfy/ldm/qwen_image/model.py
new file mode 100644
index 000000000..b9f60c2b7
--- /dev/null
+++ b/comfy/ldm/qwen_image/model.py
@@ -0,0 +1,473 @@
+# https://github.com/QwenLM/Qwen-Image (Apache 2.0)
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple
+from einops import repeat
+
+from comfy.ldm.lightricks.model import TimestepEmbedding, Timesteps
+from comfy.ldm.modules.attention import optimized_attention_masked
+from comfy.ldm.flux.layers import EmbedND
+import comfy.ldm.common_dit
+import comfy.patcher_extension
+
+class GELU(nn.Module):
+ def __init__(self, dim_in: int, dim_out: int, approximate: str = "none", bias: bool = True, dtype=None, device=None, operations=None):
+ super().__init__()
+ self.proj = operations.Linear(dim_in, dim_out, bias=bias, dtype=dtype, device=device)
+ self.approximate = approximate
+
+ def forward(self, hidden_states):
+ hidden_states = self.proj(hidden_states)
+ hidden_states = F.gelu(hidden_states, approximate=self.approximate)
+ return hidden_states
+
+
+class FeedForward(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ dim_out: Optional[int] = None,
+ mult: int = 4,
+ dropout: float = 0.0,
+ inner_dim=None,
+ bias: bool = True,
+ dtype=None, device=None, operations=None
+ ):
+ super().__init__()
+ if inner_dim is None:
+ inner_dim = int(dim * mult)
+ dim_out = dim_out if dim_out is not None else dim
+
+ self.net = nn.ModuleList([])
+ self.net.append(GELU(dim, inner_dim, approximate="tanh", bias=bias, dtype=dtype, device=device, operations=operations))
+ self.net.append(nn.Dropout(dropout))
+ self.net.append(operations.Linear(inner_dim, dim_out, bias=bias, dtype=dtype, device=device))
+
+ def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+ for module in self.net:
+ hidden_states = module(hidden_states)
+ return hidden_states
+
+
+def apply_rotary_emb(x, freqs_cis):
+ if x.shape[1] == 0:
+ return x
+
+ t_ = x.reshape(*x.shape[:-1], -1, 1, 2)
+ t_out = freqs_cis[..., 0] * t_[..., 0] + freqs_cis[..., 1] * t_[..., 1]
+ return t_out.reshape(*x.shape)
+
+
+class QwenTimestepProjEmbeddings(nn.Module):
+ def __init__(self, embedding_dim, pooled_projection_dim, dtype=None, device=None, operations=None):
+ super().__init__()
+ self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0, scale=1000)
+ self.timestep_embedder = TimestepEmbedding(
+ in_channels=256,
+ time_embed_dim=embedding_dim,
+ dtype=dtype,
+ device=device,
+ operations=operations
+ )
+
+ def forward(self, timestep, hidden_states):
+ timesteps_proj = self.time_proj(timestep)
+ timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_states.dtype))
+ return timesteps_emb
+
+
+class Attention(nn.Module):
+ def __init__(
+ self,
+ query_dim: int,
+ dim_head: int = 64,
+ heads: int = 8,
+ dropout: float = 0.0,
+ bias: bool = False,
+ eps: float = 1e-5,
+ out_bias: bool = True,
+ out_dim: int = None,
+ out_context_dim: int = None,
+ dtype=None,
+ device=None,
+ operations=None
+ ):
+ super().__init__()
+ self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+ self.inner_kv_dim = self.inner_dim
+ self.heads = heads
+ self.dim_head = dim_head
+ self.out_dim = out_dim if out_dim is not None else query_dim
+ self.out_context_dim = out_context_dim if out_context_dim is not None else query_dim
+ self.dropout = dropout
+
+ # Q/K normalization
+ self.norm_q = operations.RMSNorm(dim_head, eps=eps, elementwise_affine=True, dtype=dtype, device=device)
+ self.norm_k = operations.RMSNorm(dim_head, eps=eps, elementwise_affine=True, dtype=dtype, device=device)
+ self.norm_added_q = operations.RMSNorm(dim_head, eps=eps, dtype=dtype, device=device)
+ self.norm_added_k = operations.RMSNorm(dim_head, eps=eps, dtype=dtype, device=device)
+
+ # Image stream projections
+ self.to_q = operations.Linear(query_dim, self.inner_dim, bias=bias, dtype=dtype, device=device)
+ self.to_k = operations.Linear(query_dim, self.inner_kv_dim, bias=bias, dtype=dtype, device=device)
+ self.to_v = operations.Linear(query_dim, self.inner_kv_dim, bias=bias, dtype=dtype, device=device)
+
+ # Text stream projections
+ self.add_q_proj = operations.Linear(query_dim, self.inner_dim, bias=bias, dtype=dtype, device=device)
+ self.add_k_proj = operations.Linear(query_dim, self.inner_kv_dim, bias=bias, dtype=dtype, device=device)
+ self.add_v_proj = operations.Linear(query_dim, self.inner_kv_dim, bias=bias, dtype=dtype, device=device)
+
+ # Output projections
+ self.to_out = nn.ModuleList([
+ operations.Linear(self.inner_dim, self.out_dim, bias=out_bias, dtype=dtype, device=device),
+ nn.Dropout(dropout)
+ ])
+ self.to_add_out = operations.Linear(self.inner_dim, self.out_context_dim, bias=out_bias, dtype=dtype, device=device)
+
+ def forward(
+ self,
+ hidden_states: torch.FloatTensor, # Image stream
+ encoder_hidden_states: torch.FloatTensor = None, # Text stream
+ encoder_hidden_states_mask: torch.FloatTensor = None,
+ attention_mask: Optional[torch.FloatTensor] = None,
+ image_rotary_emb: Optional[torch.Tensor] = None,
+ transformer_options={},
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ seq_txt = encoder_hidden_states.shape[1]
+
+ img_query = self.to_q(hidden_states).unflatten(-1, (self.heads, -1))
+ img_key = self.to_k(hidden_states).unflatten(-1, (self.heads, -1))
+ img_value = self.to_v(hidden_states).unflatten(-1, (self.heads, -1))
+
+ txt_query = self.add_q_proj(encoder_hidden_states).unflatten(-1, (self.heads, -1))
+ txt_key = self.add_k_proj(encoder_hidden_states).unflatten(-1, (self.heads, -1))
+ txt_value = self.add_v_proj(encoder_hidden_states).unflatten(-1, (self.heads, -1))
+
+ img_query = self.norm_q(img_query)
+ img_key = self.norm_k(img_key)
+ txt_query = self.norm_added_q(txt_query)
+ txt_key = self.norm_added_k(txt_key)
+
+ joint_query = torch.cat([txt_query, img_query], dim=1)
+ joint_key = torch.cat([txt_key, img_key], dim=1)
+ joint_value = torch.cat([txt_value, img_value], dim=1)
+
+ joint_query = apply_rotary_emb(joint_query, image_rotary_emb)
+ joint_key = apply_rotary_emb(joint_key, image_rotary_emb)
+
+ joint_query = joint_query.flatten(start_dim=2)
+ joint_key = joint_key.flatten(start_dim=2)
+ joint_value = joint_value.flatten(start_dim=2)
+
+ joint_hidden_states = optimized_attention_masked(joint_query, joint_key, joint_value, self.heads, attention_mask, transformer_options=transformer_options)
+
+ txt_attn_output = joint_hidden_states[:, :seq_txt, :]
+ img_attn_output = joint_hidden_states[:, seq_txt:, :]
+
+ img_attn_output = self.to_out[0](img_attn_output)
+ img_attn_output = self.to_out[1](img_attn_output)
+ txt_attn_output = self.to_add_out(txt_attn_output)
+
+ return img_attn_output, txt_attn_output
+
+
+class QwenImageTransformerBlock(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ num_attention_heads: int,
+ attention_head_dim: int,
+ eps: float = 1e-6,
+ dtype=None,
+ device=None,
+ operations=None
+ ):
+ super().__init__()
+ self.dim = dim
+ self.num_attention_heads = num_attention_heads
+ self.attention_head_dim = attention_head_dim
+
+ self.img_mod = nn.Sequential(
+ nn.SiLU(),
+ operations.Linear(dim, 6 * dim, bias=True, dtype=dtype, device=device),
+ )
+ self.img_norm1 = operations.LayerNorm(dim, elementwise_affine=False, eps=eps, dtype=dtype, device=device)
+ self.img_norm2 = operations.LayerNorm(dim, elementwise_affine=False, eps=eps, dtype=dtype, device=device)
+ self.img_mlp = FeedForward(dim=dim, dim_out=dim, dtype=dtype, device=device, operations=operations)
+
+ self.txt_mod = nn.Sequential(
+ nn.SiLU(),
+ operations.Linear(dim, 6 * dim, bias=True, dtype=dtype, device=device),
+ )
+ self.txt_norm1 = operations.LayerNorm(dim, elementwise_affine=False, eps=eps, dtype=dtype, device=device)
+ self.txt_norm2 = operations.LayerNorm(dim, elementwise_affine=False, eps=eps, dtype=dtype, device=device)
+ self.txt_mlp = FeedForward(dim=dim, dim_out=dim, dtype=dtype, device=device, operations=operations)
+
+ self.attn = Attention(
+ query_dim=dim,
+ dim_head=attention_head_dim,
+ heads=num_attention_heads,
+ out_dim=dim,
+ bias=True,
+ eps=eps,
+ dtype=dtype,
+ device=device,
+ operations=operations,
+ )
+
+ def _modulate(self, x: torch.Tensor, mod_params: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+ shift, scale, gate = torch.chunk(mod_params, 3, dim=-1)
+ return torch.addcmul(shift.unsqueeze(1), x, 1 + scale.unsqueeze(1)), gate.unsqueeze(1)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor,
+ encoder_hidden_states_mask: torch.Tensor,
+ temb: torch.Tensor,
+ image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+ transformer_options={},
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ img_mod_params = self.img_mod(temb)
+ txt_mod_params = self.txt_mod(temb)
+ img_mod1, img_mod2 = img_mod_params.chunk(2, dim=-1)
+ txt_mod1, txt_mod2 = txt_mod_params.chunk(2, dim=-1)
+
+ img_normed = self.img_norm1(hidden_states)
+ img_modulated, img_gate1 = self._modulate(img_normed, img_mod1)
+ txt_normed = self.txt_norm1(encoder_hidden_states)
+ txt_modulated, txt_gate1 = self._modulate(txt_normed, txt_mod1)
+
+ img_attn_output, txt_attn_output = self.attn(
+ hidden_states=img_modulated,
+ encoder_hidden_states=txt_modulated,
+ encoder_hidden_states_mask=encoder_hidden_states_mask,
+ image_rotary_emb=image_rotary_emb,
+ transformer_options=transformer_options,
+ )
+
+ hidden_states = hidden_states + img_gate1 * img_attn_output
+ encoder_hidden_states = encoder_hidden_states + txt_gate1 * txt_attn_output
+
+ img_normed2 = self.img_norm2(hidden_states)
+ img_modulated2, img_gate2 = self._modulate(img_normed2, img_mod2)
+ hidden_states = torch.addcmul(hidden_states, img_gate2, self.img_mlp(img_modulated2))
+
+ txt_normed2 = self.txt_norm2(encoder_hidden_states)
+ txt_modulated2, txt_gate2 = self._modulate(txt_normed2, txt_mod2)
+ encoder_hidden_states = torch.addcmul(encoder_hidden_states, txt_gate2, self.txt_mlp(txt_modulated2))
+
+ return encoder_hidden_states, hidden_states
+
+
+class LastLayer(nn.Module):
+ def __init__(
+ self,
+ embedding_dim: int,
+ conditioning_embedding_dim: int,
+ elementwise_affine=False,
+ eps=1e-6,
+ bias=True,
+ dtype=None, device=None, operations=None
+ ):
+ super().__init__()
+ self.silu = nn.SiLU()
+ self.linear = operations.Linear(conditioning_embedding_dim, embedding_dim * 2, bias=bias, dtype=dtype, device=device)
+ self.norm = operations.LayerNorm(embedding_dim, eps, elementwise_affine=False, bias=bias, dtype=dtype, device=device)
+
+ def forward(self, x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torch.Tensor:
+ emb = self.linear(self.silu(conditioning_embedding))
+ scale, shift = torch.chunk(emb, 2, dim=1)
+ x = torch.addcmul(shift[:, None, :], self.norm(x), (1 + scale)[:, None, :])
+ return x
+
+
+class QwenImageTransformer2DModel(nn.Module):
+ def __init__(
+ self,
+ patch_size: int = 2,
+ in_channels: int = 64,
+ out_channels: Optional[int] = 16,
+ num_layers: int = 60,
+ attention_head_dim: int = 128,
+ num_attention_heads: int = 24,
+ joint_attention_dim: int = 3584,
+ pooled_projection_dim: int = 768,
+ guidance_embeds: bool = False,
+ axes_dims_rope: Tuple[int, int, int] = (16, 56, 56),
+ image_model=None,
+ final_layer=True,
+ dtype=None,
+ device=None,
+ operations=None,
+ ):
+ super().__init__()
+ self.dtype = dtype
+ self.patch_size = patch_size
+ self.in_channels = in_channels
+ self.out_channels = out_channels or in_channels
+ self.inner_dim = num_attention_heads * attention_head_dim
+
+ self.pe_embedder = EmbedND(dim=attention_head_dim, theta=10000, axes_dim=list(axes_dims_rope))
+
+ self.time_text_embed = QwenTimestepProjEmbeddings(
+ embedding_dim=self.inner_dim,
+ pooled_projection_dim=pooled_projection_dim,
+ dtype=dtype,
+ device=device,
+ operations=operations
+ )
+
+ self.txt_norm = operations.RMSNorm(joint_attention_dim, eps=1e-6, dtype=dtype, device=device)
+ self.img_in = operations.Linear(in_channels, self.inner_dim, dtype=dtype, device=device)
+ self.txt_in = operations.Linear(joint_attention_dim, self.inner_dim, dtype=dtype, device=device)
+
+ self.transformer_blocks = nn.ModuleList([
+ QwenImageTransformerBlock(
+ dim=self.inner_dim,
+ num_attention_heads=num_attention_heads,
+ attention_head_dim=attention_head_dim,
+ dtype=dtype,
+ device=device,
+ operations=operations
+ )
+ for _ in range(num_layers)
+ ])
+
+ if final_layer:
+ self.norm_out = LastLayer(self.inner_dim, self.inner_dim, dtype=dtype, device=device, operations=operations)
+ self.proj_out = operations.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True, dtype=dtype, device=device)
+
+ def process_img(self, x, index=0, h_offset=0, w_offset=0):
+ bs, c, t, h, w = x.shape
+ patch_size = self.patch_size
+ hidden_states = comfy.ldm.common_dit.pad_to_patch_size(x, (1, self.patch_size, self.patch_size))
+ orig_shape = hidden_states.shape
+ hidden_states = hidden_states.view(orig_shape[0], orig_shape[1], orig_shape[-2] // 2, 2, orig_shape[-1] // 2, 2)
+ hidden_states = hidden_states.permute(0, 2, 4, 1, 3, 5)
+ hidden_states = hidden_states.reshape(orig_shape[0], (orig_shape[-2] // 2) * (orig_shape[-1] // 2), orig_shape[1] * 4)
+ h_len = ((h + (patch_size // 2)) // patch_size)
+ w_len = ((w + (patch_size // 2)) // patch_size)
+
+ h_offset = ((h_offset + (patch_size // 2)) // patch_size)
+ w_offset = ((w_offset + (patch_size // 2)) // patch_size)
+
+ img_ids = torch.zeros((h_len, w_len, 3), device=x.device)
+ img_ids[:, :, 0] = img_ids[:, :, 1] + index
+ img_ids[:, :, 1] = img_ids[:, :, 1] + torch.linspace(h_offset, h_len - 1 + h_offset, steps=h_len, device=x.device, dtype=x.dtype).unsqueeze(1) - (h_len // 2)
+ img_ids[:, :, 2] = img_ids[:, :, 2] + torch.linspace(w_offset, w_len - 1 + w_offset, steps=w_len, device=x.device, dtype=x.dtype).unsqueeze(0) - (w_len // 2)
+ return hidden_states, repeat(img_ids, "h w c -> b (h w) c", b=bs), orig_shape
+
+ def forward(self, x, timestep, context, attention_mask=None, guidance=None, ref_latents=None, transformer_options={}, **kwargs):
+ return comfy.patcher_extension.WrapperExecutor.new_class_executor(
+ self._forward,
+ self,
+ comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options)
+ ).execute(x, timestep, context, attention_mask, guidance, ref_latents, transformer_options, **kwargs)
+
+ def _forward(
+ self,
+ x,
+ timesteps,
+ context,
+ attention_mask=None,
+ guidance: torch.Tensor = None,
+ ref_latents=None,
+ transformer_options={},
+ control=None,
+ **kwargs
+ ):
+ timestep = timesteps
+ encoder_hidden_states = context
+ encoder_hidden_states_mask = attention_mask
+
+ hidden_states, img_ids, orig_shape = self.process_img(x)
+ num_embeds = hidden_states.shape[1]
+
+ if ref_latents is not None:
+ h = 0
+ w = 0
+ index = 0
+ index_ref_method = kwargs.get("ref_latents_method", "index") == "index"
+ for ref in ref_latents:
+ if index_ref_method:
+ index += 1
+ h_offset = 0
+ w_offset = 0
+ else:
+ index = 1
+ h_offset = 0
+ w_offset = 0
+ if ref.shape[-2] + h > ref.shape[-1] + w:
+ w_offset = w
+ else:
+ h_offset = h
+ h = max(h, ref.shape[-2] + h_offset)
+ w = max(w, ref.shape[-1] + w_offset)
+
+ kontext, kontext_ids, _ = self.process_img(ref, index=index, h_offset=h_offset, w_offset=w_offset)
+ hidden_states = torch.cat([hidden_states, kontext], dim=1)
+ img_ids = torch.cat([img_ids, kontext_ids], dim=1)
+
+ txt_start = round(max(((x.shape[-1] + (self.patch_size // 2)) // self.patch_size) // 2, ((x.shape[-2] + (self.patch_size // 2)) // self.patch_size) // 2))
+ txt_ids = torch.arange(txt_start, txt_start + context.shape[1], device=x.device).reshape(1, -1, 1).repeat(x.shape[0], 1, 3)
+ ids = torch.cat((txt_ids, img_ids), dim=1)
+ image_rotary_emb = self.pe_embedder(ids).squeeze(1).unsqueeze(2).to(x.dtype)
+ del ids, txt_ids, img_ids
+
+ hidden_states = self.img_in(hidden_states)
+ encoder_hidden_states = self.txt_norm(encoder_hidden_states)
+ encoder_hidden_states = self.txt_in(encoder_hidden_states)
+
+ if guidance is not None:
+ guidance = guidance * 1000
+
+ temb = (
+ self.time_text_embed(timestep, hidden_states)
+ if guidance is None
+ else self.time_text_embed(timestep, guidance, hidden_states)
+ )
+
+ patches_replace = transformer_options.get("patches_replace", {})
+ patches = transformer_options.get("patches", {})
+ blocks_replace = patches_replace.get("dit", {})
+
+ for i, block in enumerate(self.transformer_blocks):
+ if ("double_block", i) in blocks_replace:
+ def block_wrap(args):
+ out = {}
+ out["txt"], out["img"] = block(hidden_states=args["img"], encoder_hidden_states=args["txt"], encoder_hidden_states_mask=encoder_hidden_states_mask, temb=args["vec"], image_rotary_emb=args["pe"], transformer_options=args["transformer_options"])
+ return out
+ out = blocks_replace[("double_block", i)]({"img": hidden_states, "txt": encoder_hidden_states, "vec": temb, "pe": image_rotary_emb, "transformer_options": transformer_options}, {"original_block": block_wrap})
+ hidden_states = out["img"]
+ encoder_hidden_states = out["txt"]
+ else:
+ encoder_hidden_states, hidden_states = block(
+ hidden_states=hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_hidden_states_mask=encoder_hidden_states_mask,
+ temb=temb,
+ image_rotary_emb=image_rotary_emb,
+ transformer_options=transformer_options,
+ )
+
+ if "double_block" in patches:
+ for p in patches["double_block"]:
+ out = p({"img": hidden_states, "txt": encoder_hidden_states, "x": x, "block_index": i, "transformer_options": transformer_options})
+ hidden_states = out["img"]
+ encoder_hidden_states = out["txt"]
+
+ if control is not None: # Controlnet
+ control_i = control.get("input")
+ if i < len(control_i):
+ add = control_i[i]
+ if add is not None:
+ hidden_states[:, :add.shape[1]] += add
+
+ hidden_states = self.norm_out(hidden_states, temb)
+ hidden_states = self.proj_out(hidden_states)
+
+ hidden_states = hidden_states[:, :num_embeds].view(orig_shape[0], orig_shape[-2] // 2, orig_shape[-1] // 2, orig_shape[1], 2, 2)
+ hidden_states = hidden_states.permute(0, 3, 1, 4, 2, 5)
+ return hidden_states.reshape(orig_shape)[:, :, :, :x.shape[-2], :x.shape[-1]]
diff --git a/comfy/ldm/wan/model.py b/comfy/ldm/wan/model.py
index 1d6edb354..5ec1511ce 100644
--- a/comfy/ldm/wan/model.py
+++ b/comfy/ldm/wan/model.py
@@ -4,13 +4,14 @@ import math
import torch
import torch.nn as nn
-from einops import repeat
+from einops import rearrange
from comfy.ldm.modules.attention import optimized_attention
from comfy.ldm.flux.layers import EmbedND
-from comfy.ldm.flux.math import apply_rope
+from comfy.ldm.flux.math import apply_rope1
import comfy.ldm.common_dit
import comfy.model_management
+import comfy.patcher_extension
def sinusoidal_embedding_1d(dim, position):
@@ -33,7 +34,9 @@ class WanSelfAttention(nn.Module):
num_heads,
window_size=(-1, -1),
qk_norm=True,
- eps=1e-6, operation_settings={}):
+ eps=1e-6,
+ kv_dim=None,
+ operation_settings={}):
assert dim % num_heads == 0
super().__init__()
self.dim = dim
@@ -42,16 +45,18 @@ class WanSelfAttention(nn.Module):
self.window_size = window_size
self.qk_norm = qk_norm
self.eps = eps
+ if kv_dim is None:
+ kv_dim = dim
# layers
self.q = operation_settings.get("operations").Linear(dim, dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
- self.k = operation_settings.get("operations").Linear(dim, dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
- self.v = operation_settings.get("operations").Linear(dim, dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
+ self.k = operation_settings.get("operations").Linear(kv_dim, dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
+ self.v = operation_settings.get("operations").Linear(kv_dim, dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
self.o = operation_settings.get("operations").Linear(dim, dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
self.norm_q = operation_settings.get("operations").RMSNorm(dim, eps=eps, elementwise_affine=True, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")) if qk_norm else nn.Identity()
self.norm_k = operation_settings.get("operations").RMSNorm(dim, eps=eps, elementwise_affine=True, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")) if qk_norm else nn.Identity()
- def forward(self, x, freqs):
+ def forward(self, x, freqs, transformer_options={}):
r"""
Args:
x(Tensor): Shape [B, L, num_heads, C / num_heads]
@@ -59,21 +64,26 @@ class WanSelfAttention(nn.Module):
"""
b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
- # query, key, value function
- def qkv_fn(x):
+ def qkv_fn_q(x):
q = self.norm_q(self.q(x)).view(b, s, n, d)
- k = self.norm_k(self.k(x)).view(b, s, n, d)
- v = self.v(x).view(b, s, n * d)
- return q, k, v
+ return apply_rope1(q, freqs)
- q, k, v = qkv_fn(x)
- q, k = apply_rope(q, k, freqs)
+ def qkv_fn_k(x):
+ k = self.norm_k(self.k(x)).view(b, s, n, d)
+ return apply_rope1(k, freqs)
+
+ #These two are VRAM hogs, so we want to do all of q computation and
+ #have pytorch garbage collect the intermediates on the sub function
+ #return before we touch k
+ q = qkv_fn_q(x)
+ k = qkv_fn_k(x)
x = optimized_attention(
q.view(b, s, n * d),
k.view(b, s, n * d),
- v,
+ self.v(x).view(b, s, n * d),
heads=self.num_heads,
+ transformer_options=transformer_options,
)
x = self.o(x)
@@ -82,7 +92,7 @@ class WanSelfAttention(nn.Module):
class WanT2VCrossAttention(WanSelfAttention):
- def forward(self, x, context, **kwargs):
+ def forward(self, x, context, transformer_options={}, **kwargs):
r"""
Args:
x(Tensor): Shape [B, L1, C]
@@ -94,7 +104,7 @@ class WanT2VCrossAttention(WanSelfAttention):
v = self.v(context)
# compute attention
- x = optimized_attention(q, k, v, heads=self.num_heads)
+ x = optimized_attention(q, k, v, heads=self.num_heads, transformer_options=transformer_options)
x = self.o(x)
return x
@@ -115,7 +125,7 @@ class WanI2VCrossAttention(WanSelfAttention):
# self.alpha = nn.Parameter(torch.zeros((1, )))
self.norm_k_img = operation_settings.get("operations").RMSNorm(dim, eps=eps, elementwise_affine=True, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")) if qk_norm else nn.Identity()
- def forward(self, x, context, context_img_len):
+ def forward(self, x, context, context_img_len, transformer_options={}):
r"""
Args:
x(Tensor): Shape [B, L1, C]
@@ -130,9 +140,9 @@ class WanI2VCrossAttention(WanSelfAttention):
v = self.v(context)
k_img = self.norm_k_img(self.k_img(context_img))
v_img = self.v_img(context_img)
- img_x = optimized_attention(q, k_img, v_img, heads=self.num_heads)
+ img_x = optimized_attention(q, k_img, v_img, heads=self.num_heads, transformer_options=transformer_options)
# compute attention
- x = optimized_attention(q, k, v, heads=self.num_heads)
+ x = optimized_attention(q, k, v, heads=self.num_heads, transformer_options=transformer_options)
# output
x = x + img_x
@@ -146,6 +156,18 @@ WAN_CROSSATTENTION_CLASSES = {
}
+def repeat_e(e, x):
+ repeats = 1
+ if e.size(1) > 1:
+ repeats = x.size(1) // e.size(1)
+ if repeats == 1:
+ return e
+ if repeats * e.size(1) == x.size(1):
+ return torch.repeat_interleave(e, repeats, dim=1)
+ else:
+ return torch.repeat_interleave(e, repeats + 1, dim=1)[:, :x.size(1)]
+
+
class WanAttentionBlock(nn.Module):
def __init__(self,
@@ -193,6 +215,7 @@ class WanAttentionBlock(nn.Module):
freqs,
context,
context_img_len=257,
+ transformer_options={},
):
r"""
Args:
@@ -202,20 +225,24 @@ class WanAttentionBlock(nn.Module):
"""
# assert e.dtype == torch.float32
- e = (comfy.model_management.cast_to(self.modulation, dtype=x.dtype, device=x.device) + e).chunk(6, dim=1)
+ if e.ndim < 4:
+ e = (comfy.model_management.cast_to(self.modulation, dtype=x.dtype, device=x.device) + e).chunk(6, dim=1)
+ else:
+ e = (comfy.model_management.cast_to(self.modulation, dtype=x.dtype, device=x.device).unsqueeze(0) + e).unbind(2)
# assert e[0].dtype == torch.float32
# self-attention
y = self.self_attn(
- self.norm1(x) * (1 + e[1]) + e[0],
- freqs)
+ torch.addcmul(repeat_e(e[0], x), self.norm1(x), 1 + repeat_e(e[1], x)),
+ freqs, transformer_options=transformer_options)
- x = x + y * e[2]
+ x = torch.addcmul(x, y, repeat_e(e[2], x))
+ del y
# cross-attention & ffn
- x = x + self.cross_attn(self.norm3(x), context, context_img_len=context_img_len)
- y = self.ffn(self.norm2(x) * (1 + e[4]) + e[3])
- x = x + y * e[5]
+ x = x + self.cross_attn(self.norm3(x), context, context_img_len=context_img_len, transformer_options=transformer_options)
+ y = self.ffn(torch.addcmul(repeat_e(e[3], x), self.norm2(x), 1 + repeat_e(e[4], x)))
+ x = torch.addcmul(x, y, repeat_e(e[5], x))
return x
@@ -325,8 +352,12 @@ class Head(nn.Module):
e(Tensor): Shape [B, C]
"""
# assert e.dtype == torch.float32
- e = (comfy.model_management.cast_to(self.modulation, dtype=x.dtype, device=x.device) + e.unsqueeze(1)).chunk(2, dim=1)
- x = (self.head(self.norm(x) * (1 + e[1]) + e[0]))
+ if e.ndim < 3:
+ e = (comfy.model_management.cast_to(self.modulation, dtype=x.dtype, device=x.device) + e.unsqueeze(1)).chunk(2, dim=1)
+ else:
+ e = (comfy.model_management.cast_to(self.modulation, dtype=x.dtype, device=x.device).unsqueeze(0) + e.unsqueeze(2)).unbind(2)
+
+ x = (self.head(torch.addcmul(repeat_e(e[0], x), self.norm(x), 1 + repeat_e(e[1], x))))
return x
@@ -375,6 +406,8 @@ class WanModel(torch.nn.Module):
cross_attn_norm=True,
eps=1e-6,
flf_pos_embed_token_number=None,
+ in_dim_ref_conv=None,
+ wan_attn_block_class=WanAttentionBlock,
image_model=None,
device=None,
dtype=None,
@@ -452,8 +485,8 @@ class WanModel(torch.nn.Module):
# blocks
cross_attn_type = 't2v_cross_attn' if model_type == 't2v' else 'i2v_cross_attn'
self.blocks = nn.ModuleList([
- WanAttentionBlock(cross_attn_type, dim, ffn_dim, num_heads,
- window_size, qk_norm, cross_attn_norm, eps, operation_settings=operation_settings)
+ wan_attn_block_class(cross_attn_type, dim, ffn_dim, num_heads,
+ window_size, qk_norm, cross_attn_norm, eps, operation_settings=operation_settings)
for _ in range(num_layers)
])
@@ -468,6 +501,11 @@ class WanModel(torch.nn.Module):
else:
self.img_emb = None
+ if in_dim_ref_conv is not None:
+ self.ref_conv = operations.Conv2d(in_dim_ref_conv, dim, kernel_size=patch_size[1:], stride=patch_size[1:], device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
+ else:
+ self.ref_conv = None
+
def forward_orig(
self,
x,
@@ -506,8 +544,16 @@ class WanModel(torch.nn.Module):
# time embeddings
e = self.time_embedding(
- sinusoidal_embedding_1d(self.freq_dim, t).to(dtype=x[0].dtype))
- e0 = self.time_projection(e).unflatten(1, (6, self.dim))
+ sinusoidal_embedding_1d(self.freq_dim, t.flatten()).to(dtype=x[0].dtype))
+ e = e.reshape(t.shape[0], -1, e.shape[-1])
+ e0 = self.time_projection(e).unflatten(2, (6, self.dim))
+
+ full_ref = None
+ if self.ref_conv is not None:
+ full_ref = kwargs.get("reference_latent", None)
+ if full_ref is not None:
+ full_ref = self.ref_conv(full_ref).flatten(2).transpose(1, 2)
+ x = torch.concat((full_ref, x), dim=1)
# context
context = self.text_embedding(context)
@@ -525,41 +571,78 @@ class WanModel(torch.nn.Module):
if ("double_block", i) in blocks_replace:
def block_wrap(args):
out = {}
- out["img"] = block(args["img"], context=args["txt"], e=args["vec"], freqs=args["pe"], context_img_len=context_img_len)
+ out["img"] = block(args["img"], context=args["txt"], e=args["vec"], freqs=args["pe"], context_img_len=context_img_len, transformer_options=args["transformer_options"])
return out
- out = blocks_replace[("double_block", i)]({"img": x, "txt": context, "vec": e0, "pe": freqs}, {"original_block": block_wrap})
+ out = blocks_replace[("double_block", i)]({"img": x, "txt": context, "vec": e0, "pe": freqs, "transformer_options": transformer_options}, {"original_block": block_wrap})
x = out["img"]
else:
- x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len)
+ x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, transformer_options=transformer_options)
# head
x = self.head(x, e)
+ if full_ref is not None:
+ x = x[:, full_ref.shape[1]:]
+
# unpatchify
x = self.unpatchify(x, grid_sizes)
return x
- def forward(self, x, timestep, context, clip_fea=None, time_dim_concat=None, transformer_options={}, **kwargs):
- bs, c, t, h, w = x.shape
- x = comfy.ldm.common_dit.pad_to_patch_size(x, self.patch_size)
-
+ def rope_encode(self, t, h, w, t_start=0, steps_t=None, steps_h=None, steps_w=None, device=None, dtype=None, transformer_options={}):
patch_size = self.patch_size
t_len = ((t + (patch_size[0] // 2)) // patch_size[0])
h_len = ((h + (patch_size[1] // 2)) // patch_size[1])
w_len = ((w + (patch_size[2] // 2)) // patch_size[2])
+ if steps_t is None:
+ steps_t = t_len
+ if steps_h is None:
+ steps_h = h_len
+ if steps_w is None:
+ steps_w = w_len
+
+ h_start = 0
+ w_start = 0
+ rope_options = transformer_options.get("rope_options", None)
+ if rope_options is not None:
+ t_len = (t_len - 1.0) * rope_options.get("scale_t", 1.0) + 1.0
+ h_len = (h_len - 1.0) * rope_options.get("scale_y", 1.0) + 1.0
+ w_len = (w_len - 1.0) * rope_options.get("scale_x", 1.0) + 1.0
+
+ t_start += rope_options.get("shift_t", 0.0)
+ h_start += rope_options.get("shift_y", 0.0)
+ w_start += rope_options.get("shift_x", 0.0)
+
+ img_ids = torch.zeros((steps_t, steps_h, steps_w, 3), device=device, dtype=dtype)
+ img_ids[:, :, :, 0] = img_ids[:, :, :, 0] + torch.linspace(t_start, t_start + (t_len - 1), steps=steps_t, device=device, dtype=dtype).reshape(-1, 1, 1)
+ img_ids[:, :, :, 1] = img_ids[:, :, :, 1] + torch.linspace(h_start, h_start + (h_len - 1), steps=steps_h, device=device, dtype=dtype).reshape(1, -1, 1)
+ img_ids[:, :, :, 2] = img_ids[:, :, :, 2] + torch.linspace(w_start, w_start + (w_len - 1), steps=steps_w, device=device, dtype=dtype).reshape(1, 1, -1)
+ img_ids = img_ids.reshape(1, -1, img_ids.shape[-1])
+
+ freqs = self.rope_embedder(img_ids).movedim(1, 2)
+ return freqs
+
+ def forward(self, x, timestep, context, clip_fea=None, time_dim_concat=None, transformer_options={}, **kwargs):
+ return comfy.patcher_extension.WrapperExecutor.new_class_executor(
+ self._forward,
+ self,
+ comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options)
+ ).execute(x, timestep, context, clip_fea, time_dim_concat, transformer_options, **kwargs)
+
+ def _forward(self, x, timestep, context, clip_fea=None, time_dim_concat=None, transformer_options={}, **kwargs):
+ bs, c, t, h, w = x.shape
+ x = comfy.ldm.common_dit.pad_to_patch_size(x, self.patch_size)
+
+ t_len = t
if time_dim_concat is not None:
time_dim_concat = comfy.ldm.common_dit.pad_to_patch_size(time_dim_concat, self.patch_size)
x = torch.cat([x, time_dim_concat], dim=2)
- t_len = ((x.shape[2] + (patch_size[0] // 2)) // patch_size[0])
+ t_len = x.shape[2]
- img_ids = torch.zeros((t_len, h_len, w_len, 3), device=x.device, dtype=x.dtype)
- img_ids[:, :, :, 0] = img_ids[:, :, :, 0] + torch.linspace(0, t_len - 1, steps=t_len, device=x.device, dtype=x.dtype).reshape(-1, 1, 1)
- img_ids[:, :, :, 1] = img_ids[:, :, :, 1] + torch.linspace(0, h_len - 1, steps=h_len, device=x.device, dtype=x.dtype).reshape(1, -1, 1)
- img_ids[:, :, :, 2] = img_ids[:, :, :, 2] + torch.linspace(0, w_len - 1, steps=w_len, device=x.device, dtype=x.dtype).reshape(1, 1, -1)
- img_ids = repeat(img_ids, "t h w c -> b (t h w) c", b=bs)
+ if self.ref_conv is not None and "reference_latent" in kwargs:
+ t_len += 1
- freqs = self.rope_embedder(img_ids).movedim(1, 2)
+ freqs = self.rope_encode(t_len, h, w, device=x.device, dtype=x.dtype, transformer_options=transformer_options)
return self.forward_orig(x, timestep, context, clip_fea=clip_fea, freqs=freqs, transformer_options=transformer_options, **kwargs)[:, :, :t, :h, :w]
def unpatchify(self, x, grid_sizes):
@@ -683,17 +766,17 @@ class VaceWanModel(WanModel):
if ("double_block", i) in blocks_replace:
def block_wrap(args):
out = {}
- out["img"] = block(args["img"], context=args["txt"], e=args["vec"], freqs=args["pe"], context_img_len=context_img_len)
+ out["img"] = block(args["img"], context=args["txt"], e=args["vec"], freqs=args["pe"], context_img_len=context_img_len, transformer_options=args["transformer_options"])
return out
- out = blocks_replace[("double_block", i)]({"img": x, "txt": context, "vec": e0, "pe": freqs}, {"original_block": block_wrap})
+ out = blocks_replace[("double_block", i)]({"img": x, "txt": context, "vec": e0, "pe": freqs, "transformer_options": transformer_options}, {"original_block": block_wrap})
x = out["img"]
else:
- x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len)
+ x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, transformer_options=transformer_options)
ii = self.vace_layers_mapping.get(i, None)
if ii is not None:
for iii in range(len(c)):
- c_skip, c[iii] = self.vace_blocks[ii](c[iii], x=x_orig, e=e0, freqs=freqs, context=context, context_img_len=context_img_len)
+ c_skip, c[iii] = self.vace_blocks[ii](c[iii], x=x_orig, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, transformer_options=transformer_options)
x += c_skip * vace_strength[iii]
del c_skip
# head
@@ -732,7 +815,12 @@ class CameraWanModel(WanModel):
operations=None,
):
- super().__init__(model_type='i2v', patch_size=patch_size, text_len=text_len, in_dim=in_dim, dim=dim, ffn_dim=ffn_dim, freq_dim=freq_dim, text_dim=text_dim, out_dim=out_dim, num_heads=num_heads, num_layers=num_layers, window_size=window_size, qk_norm=qk_norm, cross_attn_norm=cross_attn_norm, eps=eps, flf_pos_embed_token_number=flf_pos_embed_token_number, image_model=image_model, device=device, dtype=dtype, operations=operations)
+ if model_type == 'camera':
+ model_type = 'i2v'
+ else:
+ model_type = 't2v'
+
+ super().__init__(model_type=model_type, patch_size=patch_size, text_len=text_len, in_dim=in_dim, dim=dim, ffn_dim=ffn_dim, freq_dim=freq_dim, text_dim=text_dim, out_dim=out_dim, num_heads=num_heads, num_layers=num_layers, window_size=window_size, qk_norm=qk_norm, cross_attn_norm=cross_attn_norm, eps=eps, flf_pos_embed_token_number=flf_pos_embed_token_number, image_model=image_model, device=device, dtype=dtype, operations=operations)
operation_settings = {"operations": operations, "device": device, "dtype": dtype}
self.control_adapter = WanCamAdapter(in_dim_control_adapter, dim, kernel_size=patch_size[1:], stride=patch_size[1:], operation_settings=operation_settings)
@@ -752,8 +840,7 @@ class CameraWanModel(WanModel):
# embeddings
x = self.patch_embedding(x.float()).to(x.dtype)
if self.control_adapter is not None and camera_conditions is not None:
- x_camera = self.control_adapter(camera_conditions).to(x.dtype)
- x = x + x_camera
+ x = x + self.control_adapter(camera_conditions).to(x.dtype)
grid_sizes = x.shape[2:]
x = x.flatten(2).transpose(1, 2)
@@ -778,12 +865,724 @@ class CameraWanModel(WanModel):
if ("double_block", i) in blocks_replace:
def block_wrap(args):
out = {}
- out["img"] = block(args["img"], context=args["txt"], e=args["vec"], freqs=args["pe"], context_img_len=context_img_len)
+ out["img"] = block(args["img"], context=args["txt"], e=args["vec"], freqs=args["pe"], context_img_len=context_img_len, transformer_options=args["transformer_options"])
return out
- out = blocks_replace[("double_block", i)]({"img": x, "txt": context, "vec": e0, "pe": freqs}, {"original_block": block_wrap})
+ out = blocks_replace[("double_block", i)]({"img": x, "txt": context, "vec": e0, "pe": freqs, "transformer_options": transformer_options}, {"original_block": block_wrap})
x = out["img"]
else:
- x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len)
+ x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, transformer_options=transformer_options)
+
+ # head
+ x = self.head(x, e)
+
+ # unpatchify
+ x = self.unpatchify(x, grid_sizes)
+ return x
+
+
+class CausalConv1d(nn.Module):
+
+ def __init__(self,
+ chan_in,
+ chan_out,
+ kernel_size=3,
+ stride=1,
+ dilation=1,
+ pad_mode='replicate',
+ operations=None,
+ **kwargs):
+ super().__init__()
+
+ self.pad_mode = pad_mode
+ padding = (kernel_size - 1, 0) # T
+ self.time_causal_padding = padding
+
+ self.conv = operations.Conv1d(
+ chan_in,
+ chan_out,
+ kernel_size,
+ stride=stride,
+ dilation=dilation,
+ **kwargs)
+
+ def forward(self, x):
+ x = torch.nn.functional.pad(x, self.time_causal_padding, mode=self.pad_mode)
+ return self.conv(x)
+
+
+class MotionEncoder_tc(nn.Module):
+
+ def __init__(self,
+ in_dim: int,
+ hidden_dim: int,
+ num_heads: int,
+ need_global=True,
+ dtype=None,
+ device=None,
+ operations=None,):
+ factory_kwargs = {"dtype": dtype, "device": device}
+ super().__init__()
+
+ self.num_heads = num_heads
+ self.need_global = need_global
+ self.conv1_local = CausalConv1d(in_dim, hidden_dim // 4 * num_heads, 3, stride=1, operations=operations, **factory_kwargs)
+ if need_global:
+ self.conv1_global = CausalConv1d(
+ in_dim, hidden_dim // 4, 3, stride=1, operations=operations, **factory_kwargs)
+ self.norm1 = operations.LayerNorm(
+ hidden_dim // 4,
+ elementwise_affine=False,
+ eps=1e-6,
+ **factory_kwargs)
+ self.act = nn.SiLU()
+ self.conv2 = CausalConv1d(hidden_dim // 4, hidden_dim // 2, 3, stride=2, operations=operations, **factory_kwargs)
+ self.conv3 = CausalConv1d(hidden_dim // 2, hidden_dim, 3, stride=2, operations=operations, **factory_kwargs)
+
+ if need_global:
+ self.final_linear = operations.Linear(hidden_dim, hidden_dim, **factory_kwargs)
+
+ self.norm1 = operations.LayerNorm(
+ hidden_dim // 4,
+ elementwise_affine=False,
+ eps=1e-6,
+ **factory_kwargs)
+
+ self.norm2 = operations.LayerNorm(
+ hidden_dim // 2,
+ elementwise_affine=False,
+ eps=1e-6,
+ **factory_kwargs)
+
+ self.norm3 = operations.LayerNorm(
+ hidden_dim, elementwise_affine=False, eps=1e-6, **factory_kwargs)
+
+ self.padding_tokens = nn.Parameter(torch.empty(1, 1, 1, hidden_dim, **factory_kwargs))
+
+ def forward(self, x):
+ x = rearrange(x, 'b t c -> b c t')
+ x_ori = x.clone()
+ b, c, t = x.shape
+ x = self.conv1_local(x)
+ x = rearrange(x, 'b (n c) t -> (b n) t c', n=self.num_heads)
+ x = self.norm1(x)
+ x = self.act(x)
+ x = rearrange(x, 'b t c -> b c t')
+ x = self.conv2(x)
+ x = rearrange(x, 'b c t -> b t c')
+ x = self.norm2(x)
+ x = self.act(x)
+ x = rearrange(x, 'b t c -> b c t')
+ x = self.conv3(x)
+ x = rearrange(x, 'b c t -> b t c')
+ x = self.norm3(x)
+ x = self.act(x)
+ x = rearrange(x, '(b n) t c -> b t n c', b=b)
+ padding = comfy.model_management.cast_to(self.padding_tokens, dtype=x.dtype, device=x.device).repeat(b, x.shape[1], 1, 1)
+ x = torch.cat([x, padding], dim=-2)
+ x_local = x.clone()
+
+ if not self.need_global:
+ return x_local
+
+ x = self.conv1_global(x_ori)
+ x = rearrange(x, 'b c t -> b t c')
+ x = self.norm1(x)
+ x = self.act(x)
+ x = rearrange(x, 'b t c -> b c t')
+ x = self.conv2(x)
+ x = rearrange(x, 'b c t -> b t c')
+ x = self.norm2(x)
+ x = self.act(x)
+ x = rearrange(x, 'b t c -> b c t')
+ x = self.conv3(x)
+ x = rearrange(x, 'b c t -> b t c')
+ x = self.norm3(x)
+ x = self.act(x)
+ x = self.final_linear(x)
+ x = rearrange(x, '(b n) t c -> b t n c', b=b)
+
+ return x, x_local
+
+
+class CausalAudioEncoder(nn.Module):
+
+ def __init__(self,
+ dim=5120,
+ num_layers=25,
+ out_dim=2048,
+ video_rate=8,
+ num_token=4,
+ need_global=False,
+ dtype=None,
+ device=None,
+ operations=None):
+ super().__init__()
+ self.encoder = MotionEncoder_tc(
+ in_dim=dim,
+ hidden_dim=out_dim,
+ num_heads=num_token,
+ need_global=need_global, dtype=dtype, device=device, operations=operations)
+ weight = torch.empty((1, num_layers, 1, 1), dtype=dtype, device=device)
+
+ self.weights = torch.nn.Parameter(weight)
+ self.act = torch.nn.SiLU()
+
+ def forward(self, features):
+ # features B * num_layers * dim * video_length
+ weights = self.act(comfy.model_management.cast_to(self.weights, dtype=features.dtype, device=features.device))
+ weights_sum = weights.sum(dim=1, keepdims=True)
+ weighted_feat = ((features * weights) / weights_sum).sum(
+ dim=1) # b dim f
+ weighted_feat = weighted_feat.permute(0, 2, 1) # b f dim
+ res = self.encoder(weighted_feat) # b f n dim
+ return res # b f n dim
+
+
+class AdaLayerNorm(nn.Module):
+ def __init__(self, embedding_dim, output_dim=None, norm_elementwise_affine=False, norm_eps=1e-5, dtype=None, device=None, operations=None):
+ super().__init__()
+
+ output_dim = output_dim or embedding_dim * 2
+
+ self.silu = nn.SiLU()
+ self.linear = operations.Linear(embedding_dim, output_dim, dtype=dtype, device=device)
+ self.norm = operations.LayerNorm(output_dim // 2, norm_eps, norm_elementwise_affine, dtype=dtype, device=device)
+
+ def forward(self, x, temb):
+ temb = self.linear(self.silu(temb))
+ shift, scale = temb.chunk(2, dim=1)
+ shift = shift[:, None, :]
+ scale = scale[:, None, :]
+ x = self.norm(x) * (1 + scale) + shift
+ return x
+
+
+class AudioInjector_WAN(nn.Module):
+
+ def __init__(self,
+ dim=2048,
+ num_heads=32,
+ inject_layer=[0, 27],
+ root_net=None,
+ enable_adain=False,
+ adain_dim=2048,
+ adain_mode=None,
+ dtype=None,
+ device=None,
+ operations=None):
+ super().__init__()
+ self.enable_adain = enable_adain
+ self.adain_mode = adain_mode
+ self.injected_block_id = {}
+ audio_injector_id = 0
+ for inject_id in inject_layer:
+ self.injected_block_id[inject_id] = audio_injector_id
+ audio_injector_id += 1
+
+ self.injector = nn.ModuleList([
+ WanT2VCrossAttention(
+ dim=dim,
+ num_heads=num_heads,
+ qk_norm=True, operation_settings={"operations": operations, "device": device, "dtype": dtype}
+ ) for _ in range(audio_injector_id)
+ ])
+ self.injector_pre_norm_feat = nn.ModuleList([
+ operations.LayerNorm(
+ dim,
+ elementwise_affine=False,
+ eps=1e-6, dtype=dtype, device=device
+ ) for _ in range(audio_injector_id)
+ ])
+ self.injector_pre_norm_vec = nn.ModuleList([
+ operations.LayerNorm(
+ dim,
+ elementwise_affine=False,
+ eps=1e-6, dtype=dtype, device=device
+ ) for _ in range(audio_injector_id)
+ ])
+ if enable_adain:
+ self.injector_adain_layers = nn.ModuleList([
+ AdaLayerNorm(
+ output_dim=dim * 2, embedding_dim=adain_dim, dtype=dtype, device=device, operations=operations)
+ for _ in range(audio_injector_id)
+ ])
+ if adain_mode != "attn_norm":
+ self.injector_adain_output_layers = nn.ModuleList(
+ [operations.Linear(dim, dim, dtype=dtype, device=device) for _ in range(audio_injector_id)])
+
+ def forward(self, x, block_id, audio_emb, audio_emb_global, seq_len):
+ audio_attn_id = self.injected_block_id.get(block_id, None)
+ if audio_attn_id is None:
+ return x
+
+ num_frames = audio_emb.shape[1]
+ input_hidden_states = rearrange(x[:, :seq_len], "b (t n) c -> (b t) n c", t=num_frames)
+ if self.enable_adain and self.adain_mode == "attn_norm":
+ audio_emb_global = rearrange(audio_emb_global, "b t n c -> (b t) n c")
+ adain_hidden_states = self.injector_adain_layers[audio_attn_id](input_hidden_states, temb=audio_emb_global[:, 0])
+ attn_hidden_states = adain_hidden_states
+ else:
+ attn_hidden_states = self.injector_pre_norm_feat[audio_attn_id](input_hidden_states)
+ audio_emb = rearrange(audio_emb, "b t n c -> (b t) n c", t=num_frames)
+ attn_audio_emb = audio_emb
+ residual_out = self.injector[audio_attn_id](x=attn_hidden_states, context=attn_audio_emb)
+ residual_out = rearrange(
+ residual_out, "(b t) n c -> b (t n) c", t=num_frames)
+ x[:, :seq_len] = x[:, :seq_len] + residual_out
+ return x
+
+
+class FramePackMotioner(nn.Module):
+ def __init__(
+ self,
+ inner_dim=1024,
+ num_heads=16, # Used to indicate the number of heads in the backbone network; unrelated to this module's design
+ zip_frame_buckets=[
+ 1, 2, 16
+ ], # Three numbers representing the number of frames sampled for patch operations from the nearest to the farthest frames
+ drop_mode="drop", # If not "drop", it will use "padd", meaning padding instead of deletion
+ dtype=None,
+ device=None,
+ operations=None):
+ super().__init__()
+ self.proj = operations.Conv3d(16, inner_dim, kernel_size=(1, 2, 2), stride=(1, 2, 2), dtype=dtype, device=device)
+ self.proj_2x = operations.Conv3d(16, inner_dim, kernel_size=(2, 4, 4), stride=(2, 4, 4), dtype=dtype, device=device)
+ self.proj_4x = operations.Conv3d(16, inner_dim, kernel_size=(4, 8, 8), stride=(4, 8, 8), dtype=dtype, device=device)
+ self.zip_frame_buckets = zip_frame_buckets
+
+ self.inner_dim = inner_dim
+ self.num_heads = num_heads
+
+ self.drop_mode = drop_mode
+
+ def forward(self, motion_latents, rope_embedder, add_last_motion=2):
+ lat_height, lat_width = motion_latents.shape[3], motion_latents.shape[4]
+ padd_lat = torch.zeros(motion_latents.shape[0], 16, sum(self.zip_frame_buckets), lat_height, lat_width).to(device=motion_latents.device, dtype=motion_latents.dtype)
+ overlap_frame = min(padd_lat.shape[2], motion_latents.shape[2])
+ if overlap_frame > 0:
+ padd_lat[:, :, -overlap_frame:] = motion_latents[:, :, -overlap_frame:]
+
+ if add_last_motion < 2 and self.drop_mode != "drop":
+ zero_end_frame = sum(self.zip_frame_buckets[:len(self.zip_frame_buckets) - add_last_motion - 1])
+ padd_lat[:, :, -zero_end_frame:] = 0
+
+ clean_latents_4x, clean_latents_2x, clean_latents_post = padd_lat[:, :, -sum(self.zip_frame_buckets):, :, :].split(self.zip_frame_buckets[::-1], dim=2) # 16, 2 ,1
+
+ # patchfy
+ clean_latents_post = self.proj(clean_latents_post).flatten(2).transpose(1, 2)
+ clean_latents_2x = self.proj_2x(clean_latents_2x)
+ l_2x_shape = clean_latents_2x.shape
+ clean_latents_2x = clean_latents_2x.flatten(2).transpose(1, 2)
+ clean_latents_4x = self.proj_4x(clean_latents_4x)
+ l_4x_shape = clean_latents_4x.shape
+ clean_latents_4x = clean_latents_4x.flatten(2).transpose(1, 2)
+
+ if add_last_motion < 2 and self.drop_mode == "drop":
+ clean_latents_post = clean_latents_post[:, :
+ 0] if add_last_motion < 2 else clean_latents_post
+ clean_latents_2x = clean_latents_2x[:, :
+ 0] if add_last_motion < 1 else clean_latents_2x
+
+ motion_lat = torch.cat([clean_latents_post, clean_latents_2x, clean_latents_4x], dim=1)
+
+ rope_post = rope_embedder.rope_encode(1, lat_height, lat_width, t_start=-1, device=motion_latents.device, dtype=motion_latents.dtype)
+ rope_2x = rope_embedder.rope_encode(1, lat_height, lat_width, t_start=-3, steps_h=l_2x_shape[-2], steps_w=l_2x_shape[-1], device=motion_latents.device, dtype=motion_latents.dtype)
+ rope_4x = rope_embedder.rope_encode(4, lat_height, lat_width, t_start=-19, steps_h=l_4x_shape[-2], steps_w=l_4x_shape[-1], device=motion_latents.device, dtype=motion_latents.dtype)
+
+ rope = torch.cat([rope_post, rope_2x, rope_4x], dim=1)
+ return motion_lat, rope
+
+
+class WanModel_S2V(WanModel):
+ def __init__(self,
+ model_type='s2v',
+ patch_size=(1, 2, 2),
+ text_len=512,
+ in_dim=16,
+ dim=2048,
+ ffn_dim=8192,
+ freq_dim=256,
+ text_dim=4096,
+ out_dim=16,
+ num_heads=16,
+ num_layers=32,
+ window_size=(-1, -1),
+ qk_norm=True,
+ cross_attn_norm=True,
+ eps=1e-6,
+ audio_dim=1024,
+ num_audio_token=4,
+ enable_adain=True,
+ cond_dim=16,
+ audio_inject_layers=[0, 4, 8, 12, 16, 20, 24, 27, 30, 33, 36, 39],
+ adain_mode="attn_norm",
+ framepack_drop_mode="padd",
+ image_model=None,
+ device=None,
+ dtype=None,
+ operations=None,
+ ):
+
+ super().__init__(model_type='t2v', patch_size=patch_size, text_len=text_len, in_dim=in_dim, dim=dim, ffn_dim=ffn_dim, freq_dim=freq_dim, text_dim=text_dim, out_dim=out_dim, num_heads=num_heads, num_layers=num_layers, window_size=window_size, qk_norm=qk_norm, cross_attn_norm=cross_attn_norm, eps=eps, image_model=image_model, device=device, dtype=dtype, operations=operations)
+
+ self.trainable_cond_mask = operations.Embedding(3, self.dim, device=device, dtype=dtype)
+
+ self.casual_audio_encoder = CausalAudioEncoder(
+ dim=audio_dim,
+ out_dim=self.dim,
+ num_token=num_audio_token,
+ need_global=enable_adain, dtype=dtype, device=device, operations=operations)
+
+ if cond_dim > 0:
+ self.cond_encoder = operations.Conv3d(
+ cond_dim,
+ self.dim,
+ kernel_size=self.patch_size,
+ stride=self.patch_size, device=device, dtype=dtype)
+
+ self.audio_injector = AudioInjector_WAN(
+ dim=self.dim,
+ num_heads=self.num_heads,
+ inject_layer=audio_inject_layers,
+ root_net=self,
+ enable_adain=enable_adain,
+ adain_dim=self.dim,
+ adain_mode=adain_mode,
+ dtype=dtype, device=device, operations=operations
+ )
+
+ self.frame_packer = FramePackMotioner(
+ inner_dim=self.dim,
+ num_heads=self.num_heads,
+ zip_frame_buckets=[1, 2, 16],
+ drop_mode=framepack_drop_mode,
+ dtype=dtype, device=device, operations=operations)
+
+ def forward_orig(
+ self,
+ x,
+ t,
+ context,
+ audio_embed=None,
+ reference_latent=None,
+ control_video=None,
+ reference_motion=None,
+ clip_fea=None,
+ freqs=None,
+ transformer_options={},
+ **kwargs,
+ ):
+ if audio_embed is not None:
+ num_embeds = x.shape[-3] * 4
+ audio_emb_global, audio_emb = self.casual_audio_encoder(audio_embed[:, :, :, :num_embeds])
+ else:
+ audio_emb = None
+
+ # embeddings
+ bs, _, time, height, width = x.shape
+ x = self.patch_embedding(x.float()).to(x.dtype)
+ if control_video is not None:
+ x = x + self.cond_encoder(control_video)
+
+ if t.ndim == 1:
+ t = t.unsqueeze(1).repeat(1, x.shape[2])
+
+ grid_sizes = x.shape[2:]
+ x = x.flatten(2).transpose(1, 2)
+ seq_len = x.size(1)
+
+ cond_mask_weight = comfy.model_management.cast_to(self.trainable_cond_mask.weight, dtype=x.dtype, device=x.device).unsqueeze(1).unsqueeze(1)
+ x = x + cond_mask_weight[0]
+
+ if reference_latent is not None:
+ ref = self.patch_embedding(reference_latent.float()).to(x.dtype)
+ ref = ref.flatten(2).transpose(1, 2)
+ freqs_ref = self.rope_encode(reference_latent.shape[-3], reference_latent.shape[-2], reference_latent.shape[-1], t_start=max(30, time + 9), device=x.device, dtype=x.dtype)
+ ref = ref + cond_mask_weight[1]
+ x = torch.cat([x, ref], dim=1)
+ freqs = torch.cat([freqs, freqs_ref], dim=1)
+ t = torch.cat([t, torch.zeros((t.shape[0], reference_latent.shape[-3]), device=t.device, dtype=t.dtype)], dim=1)
+ del ref, freqs_ref
+
+ if reference_motion is not None:
+ motion_encoded, freqs_motion = self.frame_packer(reference_motion, self)
+ motion_encoded = motion_encoded + cond_mask_weight[2]
+ x = torch.cat([x, motion_encoded], dim=1)
+ freqs = torch.cat([freqs, freqs_motion], dim=1)
+
+ t = torch.repeat_interleave(t, 2, dim=1)
+ t = torch.cat([t, torch.zeros((t.shape[0], 3), device=t.device, dtype=t.dtype)], dim=1)
+ del motion_encoded, freqs_motion
+
+ # time embeddings
+ e = self.time_embedding(
+ sinusoidal_embedding_1d(self.freq_dim, t.flatten()).to(dtype=x[0].dtype))
+ e = e.reshape(t.shape[0], -1, e.shape[-1])
+ e0 = self.time_projection(e).unflatten(2, (6, self.dim))
+
+ # context
+ context = self.text_embedding(context)
+
+ patches_replace = transformer_options.get("patches_replace", {})
+ blocks_replace = patches_replace.get("dit", {})
+ for i, block in enumerate(self.blocks):
+ if ("double_block", i) in blocks_replace:
+ def block_wrap(args):
+ out = {}
+ out["img"] = block(args["img"], context=args["txt"], e=args["vec"], freqs=args["pe"])
+ return out
+ out = blocks_replace[("double_block", i)]({"img": x, "txt": context, "vec": e0, "pe": freqs}, {"original_block": block_wrap})
+ x = out["img"]
+ else:
+ x = block(x, e=e0, freqs=freqs, context=context)
+ if audio_emb is not None:
+ x = self.audio_injector(x, i, audio_emb, audio_emb_global, seq_len)
+ # head
+ x = self.head(x, e)
+
+ # unpatchify
+ x = self.unpatchify(x, grid_sizes)
+ return x
+
+
+class WanT2VCrossAttentionGather(WanSelfAttention):
+
+ def forward(self, x, context, transformer_options={}, **kwargs):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L1, C] - video tokens
+ context(Tensor): Shape [B, L2, C] - audio tokens with shape [B, frames*16, 1536]
+ """
+ b, n, d = x.size(0), self.num_heads, self.head_dim
+
+ q = self.norm_q(self.q(x))
+ k = self.norm_k(self.k(context))
+ v = self.v(context)
+
+ # Handle audio temporal structure (16 tokens per frame)
+ k = k.reshape(-1, 16, n, d).transpose(1, 2)
+ v = v.reshape(-1, 16, n, d).transpose(1, 2)
+
+ # Handle video spatial structure
+ q = q.reshape(k.shape[0], -1, n, d).transpose(1, 2)
+
+ x = optimized_attention(q, k, v, heads=self.num_heads, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options)
+
+ x = x.transpose(1, 2).reshape(b, -1, n * d)
+ x = self.o(x)
+ return x
+
+
+class AudioCrossAttentionWrapper(nn.Module):
+ def __init__(self, dim, kv_dim, num_heads, qk_norm=True, eps=1e-6, operation_settings={}):
+ super().__init__()
+
+ self.audio_cross_attn = WanT2VCrossAttentionGather(dim, num_heads, qk_norm=qk_norm, kv_dim=kv_dim, eps=eps, operation_settings=operation_settings)
+ self.norm1_audio = operation_settings.get("operations").LayerNorm(dim, eps, elementwise_affine=True, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
+
+ def forward(self, x, audio, transformer_options={}):
+ x = x + self.audio_cross_attn(self.norm1_audio(x), audio, transformer_options=transformer_options)
+ return x
+
+
+class WanAttentionBlockAudio(WanAttentionBlock):
+
+ def __init__(self,
+ cross_attn_type,
+ dim,
+ ffn_dim,
+ num_heads,
+ window_size=(-1, -1),
+ qk_norm=True,
+ cross_attn_norm=False,
+ eps=1e-6, operation_settings={}):
+ super().__init__(cross_attn_type, dim, ffn_dim, num_heads, window_size, qk_norm, cross_attn_norm, eps, operation_settings)
+ self.audio_cross_attn_wrapper = AudioCrossAttentionWrapper(dim, 1536, num_heads, qk_norm, eps, operation_settings=operation_settings)
+
+ def forward(
+ self,
+ x,
+ e,
+ freqs,
+ context,
+ context_img_len=257,
+ audio=None,
+ transformer_options={},
+ ):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L, C]
+ e(Tensor): Shape [B, 6, C]
+ freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+ """
+ # assert e.dtype == torch.float32
+
+ if e.ndim < 4:
+ e = (comfy.model_management.cast_to(self.modulation, dtype=x.dtype, device=x.device) + e).chunk(6, dim=1)
+ else:
+ e = (comfy.model_management.cast_to(self.modulation, dtype=x.dtype, device=x.device).unsqueeze(0) + e).unbind(2)
+ # assert e[0].dtype == torch.float32
+
+ # self-attention
+ y = self.self_attn(
+ torch.addcmul(repeat_e(e[0], x), self.norm1(x), 1 + repeat_e(e[1], x)),
+ freqs, transformer_options=transformer_options)
+
+ x = torch.addcmul(x, y, repeat_e(e[2], x))
+
+ # cross-attention & ffn
+ x = x + self.cross_attn(self.norm3(x), context, context_img_len=context_img_len, transformer_options=transformer_options)
+ if audio is not None:
+ x = self.audio_cross_attn_wrapper(x, audio, transformer_options=transformer_options)
+ y = self.ffn(torch.addcmul(repeat_e(e[3], x), self.norm2(x), 1 + repeat_e(e[4], x)))
+ x = torch.addcmul(x, y, repeat_e(e[5], x))
+ return x
+
+class DummyAdapterLayer(nn.Module):
+ def __init__(self, layer):
+ super().__init__()
+ self.layer = layer
+
+ def forward(self, *args, **kwargs):
+ return self.layer(*args, **kwargs)
+
+
+class AudioProjModel(nn.Module):
+ def __init__(
+ self,
+ seq_len=5,
+ blocks=13, # add a new parameter blocks
+ channels=768, # add a new parameter channels
+ intermediate_dim=512,
+ output_dim=1536,
+ context_tokens=16,
+ device=None,
+ dtype=None,
+ operations=None,
+ ):
+ super().__init__()
+
+ self.seq_len = seq_len
+ self.blocks = blocks
+ self.channels = channels
+ self.input_dim = seq_len * blocks * channels # update input_dim to be the product of blocks and channels.
+ self.intermediate_dim = intermediate_dim
+ self.context_tokens = context_tokens
+ self.output_dim = output_dim
+
+ # define multiple linear layers
+ self.audio_proj_glob_1 = DummyAdapterLayer(operations.Linear(self.input_dim, intermediate_dim, dtype=dtype, device=device))
+ self.audio_proj_glob_2 = DummyAdapterLayer(operations.Linear(intermediate_dim, intermediate_dim, dtype=dtype, device=device))
+ self.audio_proj_glob_3 = DummyAdapterLayer(operations.Linear(intermediate_dim, context_tokens * output_dim, dtype=dtype, device=device))
+
+ self.audio_proj_glob_norm = DummyAdapterLayer(operations.LayerNorm(output_dim, dtype=dtype, device=device))
+
+ def forward(self, audio_embeds):
+ video_length = audio_embeds.shape[1]
+ audio_embeds = rearrange(audio_embeds, "bz f w b c -> (bz f) w b c")
+ batch_size, window_size, blocks, channels = audio_embeds.shape
+ audio_embeds = audio_embeds.view(batch_size, window_size * blocks * channels)
+
+ audio_embeds = torch.relu(self.audio_proj_glob_1(audio_embeds))
+ audio_embeds = torch.relu(self.audio_proj_glob_2(audio_embeds))
+
+ context_tokens = self.audio_proj_glob_3(audio_embeds).reshape(batch_size, self.context_tokens, self.output_dim)
+
+ context_tokens = self.audio_proj_glob_norm(context_tokens)
+ context_tokens = rearrange(context_tokens, "(bz f) m c -> bz f m c", f=video_length)
+
+ return context_tokens
+
+
+class HumoWanModel(WanModel):
+ r"""
+ Wan diffusion backbone supporting both text-to-video and image-to-video.
+ """
+
+ def __init__(self,
+ model_type='humo',
+ patch_size=(1, 2, 2),
+ text_len=512,
+ in_dim=16,
+ dim=2048,
+ ffn_dim=8192,
+ freq_dim=256,
+ text_dim=4096,
+ out_dim=16,
+ num_heads=16,
+ num_layers=32,
+ window_size=(-1, -1),
+ qk_norm=True,
+ cross_attn_norm=True,
+ eps=1e-6,
+ flf_pos_embed_token_number=None,
+ image_model=None,
+ audio_token_num=16,
+ device=None,
+ dtype=None,
+ operations=None,
+ ):
+
+ super().__init__(model_type='t2v', patch_size=patch_size, text_len=text_len, in_dim=in_dim, dim=dim, ffn_dim=ffn_dim, freq_dim=freq_dim, text_dim=text_dim, out_dim=out_dim, num_heads=num_heads, num_layers=num_layers, window_size=window_size, qk_norm=qk_norm, cross_attn_norm=cross_attn_norm, eps=eps, flf_pos_embed_token_number=flf_pos_embed_token_number, wan_attn_block_class=WanAttentionBlockAudio, image_model=image_model, device=device, dtype=dtype, operations=operations)
+
+ self.audio_proj = AudioProjModel(seq_len=8, blocks=5, channels=1280, intermediate_dim=512, output_dim=1536, context_tokens=audio_token_num, dtype=dtype, device=device, operations=operations)
+
+ def forward_orig(
+ self,
+ x,
+ t,
+ context,
+ freqs=None,
+ audio_embed=None,
+ reference_latent=None,
+ transformer_options={},
+ **kwargs,
+ ):
+ bs, _, time, height, width = x.shape
+
+ # embeddings
+ x = self.patch_embedding(x.float()).to(x.dtype)
+ grid_sizes = x.shape[2:]
+ x = x.flatten(2).transpose(1, 2)
+
+ # time embeddings
+ e = self.time_embedding(
+ sinusoidal_embedding_1d(self.freq_dim, t.flatten()).to(dtype=x[0].dtype))
+ e = e.reshape(t.shape[0], -1, e.shape[-1])
+ e0 = self.time_projection(e).unflatten(2, (6, self.dim))
+
+ if reference_latent is not None:
+ ref = self.patch_embedding(reference_latent.float()).to(x.dtype)
+ ref = ref.flatten(2).transpose(1, 2)
+ freqs_ref = self.rope_encode(reference_latent.shape[-3], reference_latent.shape[-2], reference_latent.shape[-1], t_start=time, device=x.device, dtype=x.dtype)
+ x = torch.cat([x, ref], dim=1)
+ freqs = torch.cat([freqs, freqs_ref], dim=1)
+ del ref, freqs_ref
+
+ # context
+ context = self.text_embedding(context)
+ context_img_len = None
+
+ if audio_embed is not None:
+ if reference_latent is not None:
+ zero_audio_pad = torch.zeros(audio_embed.shape[0], reference_latent.shape[-3], *audio_embed.shape[2:], device=audio_embed.device, dtype=audio_embed.dtype)
+ audio_embed = torch.cat([audio_embed, zero_audio_pad], dim=1)
+ audio = self.audio_proj(audio_embed).permute(0, 3, 1, 2).flatten(2).transpose(1, 2)
+ else:
+ audio = None
+
+ patches_replace = transformer_options.get("patches_replace", {})
+ blocks_replace = patches_replace.get("dit", {})
+ for i, block in enumerate(self.blocks):
+ if ("double_block", i) in blocks_replace:
+ def block_wrap(args):
+ out = {}
+ out["img"] = block(args["img"], context=args["txt"], e=args["vec"], freqs=args["pe"], context_img_len=context_img_len, audio=audio, transformer_options=args["transformer_options"])
+ return out
+ out = blocks_replace[("double_block", i)]({"img": x, "txt": context, "vec": e0, "pe": freqs, "transformer_options": transformer_options}, {"original_block": block_wrap})
+ x = out["img"]
+ else:
+ x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, audio=audio, transformer_options=transformer_options)
# head
x = self.head(x, e)
diff --git a/comfy/ldm/wan/model_animate.py b/comfy/ldm/wan/model_animate.py
new file mode 100644
index 000000000..7c87835d4
--- /dev/null
+++ b/comfy/ldm/wan/model_animate.py
@@ -0,0 +1,548 @@
+from torch import nn
+import torch
+from typing import Tuple, Optional
+from einops import rearrange
+import torch.nn.functional as F
+import math
+from .model import WanModel, sinusoidal_embedding_1d
+from comfy.ldm.modules.attention import optimized_attention
+import comfy.model_management
+
+class CausalConv1d(nn.Module):
+
+ def __init__(self, chan_in, chan_out, kernel_size=3, stride=1, dilation=1, pad_mode="replicate", operations=None, **kwargs):
+ super().__init__()
+
+ self.pad_mode = pad_mode
+ padding = (kernel_size - 1, 0) # T
+ self.time_causal_padding = padding
+
+ self.conv = operations.Conv1d(chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs)
+
+ def forward(self, x):
+ x = F.pad(x, self.time_causal_padding, mode=self.pad_mode)
+ return self.conv(x)
+
+
+class FaceEncoder(nn.Module):
+ def __init__(self, in_dim: int, hidden_dim: int, num_heads=int, dtype=None, device=None, operations=None):
+ factory_kwargs = {"dtype": dtype, "device": device}
+ super().__init__()
+
+ self.num_heads = num_heads
+ self.conv1_local = CausalConv1d(in_dim, 1024 * num_heads, 3, stride=1, operations=operations, **factory_kwargs)
+ self.norm1 = operations.LayerNorm(hidden_dim // 8, elementwise_affine=False, eps=1e-6, **factory_kwargs)
+ self.act = nn.SiLU()
+ self.conv2 = CausalConv1d(1024, 1024, 3, stride=2, operations=operations, **factory_kwargs)
+ self.conv3 = CausalConv1d(1024, 1024, 3, stride=2, operations=operations, **factory_kwargs)
+
+ self.out_proj = operations.Linear(1024, hidden_dim, **factory_kwargs)
+ self.norm1 = operations.LayerNorm(1024, elementwise_affine=False, eps=1e-6, **factory_kwargs)
+
+ self.norm2 = operations.LayerNorm(1024, elementwise_affine=False, eps=1e-6, **factory_kwargs)
+
+ self.norm3 = operations.LayerNorm(1024, elementwise_affine=False, eps=1e-6, **factory_kwargs)
+
+ self.padding_tokens = nn.Parameter(torch.empty(1, 1, 1, hidden_dim, **factory_kwargs))
+
+ def forward(self, x):
+
+ x = rearrange(x, "b t c -> b c t")
+ b, c, t = x.shape
+
+ x = self.conv1_local(x)
+ x = rearrange(x, "b (n c) t -> (b n) t c", n=self.num_heads)
+
+ x = self.norm1(x)
+ x = self.act(x)
+ x = rearrange(x, "b t c -> b c t")
+ x = self.conv2(x)
+ x = rearrange(x, "b c t -> b t c")
+ x = self.norm2(x)
+ x = self.act(x)
+ x = rearrange(x, "b t c -> b c t")
+ x = self.conv3(x)
+ x = rearrange(x, "b c t -> b t c")
+ x = self.norm3(x)
+ x = self.act(x)
+ x = self.out_proj(x)
+ x = rearrange(x, "(b n) t c -> b t n c", b=b)
+ padding = comfy.model_management.cast_to(self.padding_tokens, dtype=x.dtype, device=x.device).repeat(b, x.shape[1], 1, 1)
+ x = torch.cat([x, padding], dim=-2)
+ x_local = x.clone()
+
+ return x_local
+
+
+def get_norm_layer(norm_layer, operations=None):
+ """
+ Get the normalization layer.
+
+ Args:
+ norm_layer (str): The type of normalization layer.
+
+ Returns:
+ norm_layer (nn.Module): The normalization layer.
+ """
+ if norm_layer == "layer":
+ return operations.LayerNorm
+ elif norm_layer == "rms":
+ return operations.RMSNorm
+ else:
+ raise NotImplementedError(f"Norm layer {norm_layer} is not implemented")
+
+
+class FaceAdapter(nn.Module):
+ def __init__(
+ self,
+ hidden_dim: int,
+ heads_num: int,
+ qk_norm: bool = True,
+ qk_norm_type: str = "rms",
+ num_adapter_layers: int = 1,
+ dtype=None, device=None, operations=None
+ ):
+
+ factory_kwargs = {"dtype": dtype, "device": device}
+ super().__init__()
+ self.hidden_size = hidden_dim
+ self.heads_num = heads_num
+ self.fuser_blocks = nn.ModuleList(
+ [
+ FaceBlock(
+ self.hidden_size,
+ self.heads_num,
+ qk_norm=qk_norm,
+ qk_norm_type=qk_norm_type,
+ operations=operations,
+ **factory_kwargs,
+ )
+ for _ in range(num_adapter_layers)
+ ]
+ )
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ motion_embed: torch.Tensor,
+ idx: int,
+ freqs_cis_q: Tuple[torch.Tensor, torch.Tensor] = None,
+ freqs_cis_k: Tuple[torch.Tensor, torch.Tensor] = None,
+ ) -> torch.Tensor:
+
+ return self.fuser_blocks[idx](x, motion_embed, freqs_cis_q, freqs_cis_k)
+
+
+
+class FaceBlock(nn.Module):
+ def __init__(
+ self,
+ hidden_size: int,
+ heads_num: int,
+ qk_norm: bool = True,
+ qk_norm_type: str = "rms",
+ qk_scale: float = None,
+ dtype: Optional[torch.dtype] = None,
+ device: Optional[torch.device] = None,
+ operations=None
+ ):
+ factory_kwargs = {"device": device, "dtype": dtype}
+ super().__init__()
+
+ self.deterministic = False
+ self.hidden_size = hidden_size
+ self.heads_num = heads_num
+ head_dim = hidden_size // heads_num
+ self.scale = qk_scale or head_dim**-0.5
+
+ self.linear1_kv = operations.Linear(hidden_size, hidden_size * 2, **factory_kwargs)
+ self.linear1_q = operations.Linear(hidden_size, hidden_size, **factory_kwargs)
+
+ self.linear2 = operations.Linear(hidden_size, hidden_size, **factory_kwargs)
+
+ qk_norm_layer = get_norm_layer(qk_norm_type, operations=operations)
+ self.q_norm = (
+ qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) if qk_norm else nn.Identity()
+ )
+ self.k_norm = (
+ qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) if qk_norm else nn.Identity()
+ )
+
+ self.pre_norm_feat = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs)
+
+ self.pre_norm_motion = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs)
+
+ def forward(
+ self,
+ x: torch.Tensor,
+ motion_vec: torch.Tensor,
+ motion_mask: Optional[torch.Tensor] = None,
+ # use_context_parallel=False,
+ ) -> torch.Tensor:
+
+ B, T, N, C = motion_vec.shape
+ T_comp = T
+
+ x_motion = self.pre_norm_motion(motion_vec)
+ x_feat = self.pre_norm_feat(x)
+
+ kv = self.linear1_kv(x_motion)
+ q = self.linear1_q(x_feat)
+
+ k, v = rearrange(kv, "B L N (K H D) -> K B L N H D", K=2, H=self.heads_num)
+ q = rearrange(q, "B S (H D) -> B S H D", H=self.heads_num)
+
+ # Apply QK-Norm if needed.
+ q = self.q_norm(q).to(v)
+ k = self.k_norm(k).to(v)
+
+ k = rearrange(k, "B L N H D -> (B L) N H D")
+ v = rearrange(v, "B L N H D -> (B L) N H D")
+
+ q = rearrange(q, "B (L S) H D -> (B L) S (H D)", L=T_comp)
+
+ attn = optimized_attention(q, k, v, heads=self.heads_num)
+
+ attn = rearrange(attn, "(B L) S C -> B (L S) C", L=T_comp)
+
+ output = self.linear2(attn)
+
+ if motion_mask is not None:
+ output = output * rearrange(motion_mask, "B T H W -> B (T H W)").unsqueeze(-1)
+
+ return output
+
+# https://github.com/XPixelGroup/BasicSR/blob/8d56e3a045f9fb3e1d8872f92ee4a4f07f886b0a/basicsr/ops/upfirdn2d/upfirdn2d.py#L162
+def upfirdn2d_native(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1):
+ _, minor, in_h, in_w = input.shape
+ kernel_h, kernel_w = kernel.shape
+
+ out = input.view(-1, minor, in_h, 1, in_w, 1)
+ out = F.pad(out, [0, up_x - 1, 0, 0, 0, up_y - 1, 0, 0])
+ out = out.view(-1, minor, in_h * up_y, in_w * up_x)
+
+ out = F.pad(out, [max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)])
+ out = out[:, :, max(-pad_y0, 0): out.shape[2] - max(-pad_y1, 0), max(-pad_x0, 0): out.shape[3] - max(-pad_x1, 0)]
+
+ out = out.reshape([-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1])
+ w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w)
+ out = F.conv2d(out, w)
+ out = out.reshape(-1, minor, in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1, in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1)
+ return out[:, :, ::down_y, ::down_x]
+
+def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)):
+ return upfirdn2d_native(input, kernel, up, up, down, down, pad[0], pad[1], pad[0], pad[1])
+
+# https://github.com/XPixelGroup/BasicSR/blob/8d56e3a045f9fb3e1d8872f92ee4a4f07f886b0a/basicsr/ops/fused_act/fused_act.py#L81
+class FusedLeakyReLU(torch.nn.Module):
+ def __init__(self, channel, negative_slope=0.2, scale=2 ** 0.5, dtype=None, device=None):
+ super().__init__()
+ self.bias = torch.nn.Parameter(torch.empty(1, channel, 1, 1, dtype=dtype, device=device))
+ self.negative_slope = negative_slope
+ self.scale = scale
+
+ def forward(self, input):
+ return fused_leaky_relu(input, comfy.model_management.cast_to(self.bias, device=input.device, dtype=input.dtype), self.negative_slope, self.scale)
+
+def fused_leaky_relu(input, bias, negative_slope=0.2, scale=2 ** 0.5):
+ return F.leaky_relu(input + bias, negative_slope) * scale
+
+class Blur(torch.nn.Module):
+ def __init__(self, kernel, pad, dtype=None, device=None):
+ super().__init__()
+ kernel = torch.tensor(kernel, dtype=dtype, device=device)
+ kernel = kernel[None, :] * kernel[:, None]
+ kernel = kernel / kernel.sum()
+ self.register_buffer('kernel', kernel)
+ self.pad = pad
+
+ def forward(self, input):
+ return upfirdn2d(input, comfy.model_management.cast_to(self.kernel, dtype=input.dtype, device=input.device), pad=self.pad)
+
+#https://github.com/XPixelGroup/BasicSR/blob/8d56e3a045f9fb3e1d8872f92ee4a4f07f886b0a/basicsr/archs/stylegan2_arch.py#L590
+class ScaledLeakyReLU(torch.nn.Module):
+ def __init__(self, negative_slope=0.2):
+ super().__init__()
+ self.negative_slope = negative_slope
+
+ def forward(self, input):
+ return F.leaky_relu(input, negative_slope=self.negative_slope)
+
+# https://github.com/XPixelGroup/BasicSR/blob/8d56e3a045f9fb3e1d8872f92ee4a4f07f886b0a/basicsr/archs/stylegan2_arch.py#L605
+class EqualConv2d(torch.nn.Module):
+ def __init__(self, in_channel, out_channel, kernel_size, stride=1, padding=0, bias=True, dtype=None, device=None, operations=None):
+ super().__init__()
+ self.weight = torch.nn.Parameter(torch.empty(out_channel, in_channel, kernel_size, kernel_size, device=device, dtype=dtype))
+ self.scale = 1 / math.sqrt(in_channel * kernel_size ** 2)
+ self.stride = stride
+ self.padding = padding
+ self.bias = torch.nn.Parameter(torch.empty(out_channel, device=device, dtype=dtype)) if bias else None
+
+ def forward(self, input):
+ if self.bias is None:
+ bias = None
+ else:
+ bias = comfy.model_management.cast_to(self.bias, device=input.device, dtype=input.dtype)
+
+ return F.conv2d(input, comfy.model_management.cast_to(self.weight, device=input.device, dtype=input.dtype) * self.scale, bias=bias, stride=self.stride, padding=self.padding)
+
+# https://github.com/XPixelGroup/BasicSR/blob/8d56e3a045f9fb3e1d8872f92ee4a4f07f886b0a/basicsr/archs/stylegan2_arch.py#L134
+class EqualLinear(torch.nn.Module):
+ def __init__(self, in_dim, out_dim, bias=True, bias_init=0, lr_mul=1, activation=None, dtype=None, device=None, operations=None):
+ super().__init__()
+ self.weight = torch.nn.Parameter(torch.empty(out_dim, in_dim, device=device, dtype=dtype))
+ self.bias = torch.nn.Parameter(torch.empty(out_dim, device=device, dtype=dtype)) if bias else None
+ self.activation = activation
+ self.scale = (1 / math.sqrt(in_dim)) * lr_mul
+ self.lr_mul = lr_mul
+
+ def forward(self, input):
+ if self.bias is None:
+ bias = None
+ else:
+ bias = comfy.model_management.cast_to(self.bias, device=input.device, dtype=input.dtype) * self.lr_mul
+
+ if self.activation:
+ out = F.linear(input, comfy.model_management.cast_to(self.weight, device=input.device, dtype=input.dtype) * self.scale)
+ return fused_leaky_relu(out, bias)
+ return F.linear(input, comfy.model_management.cast_to(self.weight, device=input.device, dtype=input.dtype) * self.scale, bias=bias)
+
+# https://github.com/XPixelGroup/BasicSR/blob/8d56e3a045f9fb3e1d8872f92ee4a4f07f886b0a/basicsr/archs/stylegan2_arch.py#L654
+class ConvLayer(torch.nn.Sequential):
+ def __init__(self, in_channel, out_channel, kernel_size, downsample=False, blur_kernel=[1, 3, 3, 1], bias=True, activate=True, dtype=None, device=None, operations=None):
+ layers = []
+
+ if downsample:
+ factor = 2
+ p = (len(blur_kernel) - factor) + (kernel_size - 1)
+ layers.append(Blur(blur_kernel, pad=((p + 1) // 2, p // 2)))
+ stride, padding = 2, 0
+ else:
+ stride, padding = 1, kernel_size // 2
+
+ layers.append(EqualConv2d(in_channel, out_channel, kernel_size, padding=padding, stride=stride, bias=bias and not activate, dtype=dtype, device=device, operations=operations))
+
+ if activate:
+ layers.append(FusedLeakyReLU(out_channel) if bias else ScaledLeakyReLU(0.2))
+
+ super().__init__(*layers)
+
+# https://github.com/XPixelGroup/BasicSR/blob/8d56e3a045f9fb3e1d8872f92ee4a4f07f886b0a/basicsr/archs/stylegan2_arch.py#L704
+class ResBlock(torch.nn.Module):
+ def __init__(self, in_channel, out_channel, dtype=None, device=None, operations=None):
+ super().__init__()
+ self.conv1 = ConvLayer(in_channel, in_channel, 3, dtype=dtype, device=device, operations=operations)
+ self.conv2 = ConvLayer(in_channel, out_channel, 3, downsample=True, dtype=dtype, device=device, operations=operations)
+ self.skip = ConvLayer(in_channel, out_channel, 1, downsample=True, activate=False, bias=False, dtype=dtype, device=device, operations=operations)
+
+ def forward(self, input):
+ out = self.conv2(self.conv1(input))
+ skip = self.skip(input)
+ return (out + skip) / math.sqrt(2)
+
+
+class EncoderApp(torch.nn.Module):
+ def __init__(self, w_dim=512, dtype=None, device=None, operations=None):
+ super().__init__()
+ kwargs = {"device": device, "dtype": dtype, "operations": operations}
+
+ self.convs = torch.nn.ModuleList([
+ ConvLayer(3, 32, 1, **kwargs), ResBlock(32, 64, **kwargs),
+ ResBlock(64, 128, **kwargs), ResBlock(128, 256, **kwargs),
+ ResBlock(256, 512, **kwargs), ResBlock(512, 512, **kwargs),
+ ResBlock(512, 512, **kwargs), ResBlock(512, 512, **kwargs),
+ EqualConv2d(512, w_dim, 4, padding=0, bias=False, **kwargs)
+ ])
+
+ def forward(self, x):
+ h = x
+ for conv in self.convs:
+ h = conv(h)
+ return h.squeeze(-1).squeeze(-1)
+
+class Encoder(torch.nn.Module):
+ def __init__(self, dim=512, motion_dim=20, dtype=None, device=None, operations=None):
+ super().__init__()
+ self.net_app = EncoderApp(dim, dtype=dtype, device=device, operations=operations)
+ self.fc = torch.nn.Sequential(*[EqualLinear(dim, dim, dtype=dtype, device=device, operations=operations) for _ in range(4)] + [EqualLinear(dim, motion_dim, dtype=dtype, device=device, operations=operations)])
+
+ def encode_motion(self, x):
+ return self.fc(self.net_app(x))
+
+class Direction(torch.nn.Module):
+ def __init__(self, motion_dim, dtype=None, device=None, operations=None):
+ super().__init__()
+ self.weight = torch.nn.Parameter(torch.empty(512, motion_dim, device=device, dtype=dtype))
+ self.motion_dim = motion_dim
+
+ def forward(self, input):
+ stabilized_weight = comfy.model_management.cast_to(self.weight, device=input.device, dtype=input.dtype) + 1e-8 * torch.eye(512, self.motion_dim, device=input.device, dtype=input.dtype)
+ Q, _ = torch.linalg.qr(stabilized_weight.float())
+ if input is None:
+ return Q
+ return torch.sum(input.unsqueeze(-1) * Q.T.to(input.dtype), dim=1)
+
+class Synthesis(torch.nn.Module):
+ def __init__(self, motion_dim, dtype=None, device=None, operations=None):
+ super().__init__()
+ self.direction = Direction(motion_dim, dtype=dtype, device=device, operations=operations)
+
+class Generator(torch.nn.Module):
+ def __init__(self, style_dim=512, motion_dim=20, dtype=None, device=None, operations=None):
+ super().__init__()
+ self.enc = Encoder(style_dim, motion_dim, dtype=dtype, device=device, operations=operations)
+ self.dec = Synthesis(motion_dim, dtype=dtype, device=device, operations=operations)
+
+ def get_motion(self, img):
+ motion_feat = self.enc.encode_motion(img)
+ return self.dec.direction(motion_feat)
+
+class AnimateWanModel(WanModel):
+ r"""
+ Wan diffusion backbone supporting both text-to-video and image-to-video.
+ """
+
+ def __init__(self,
+ model_type='animate',
+ patch_size=(1, 2, 2),
+ text_len=512,
+ in_dim=16,
+ dim=2048,
+ ffn_dim=8192,
+ freq_dim=256,
+ text_dim=4096,
+ out_dim=16,
+ num_heads=16,
+ num_layers=32,
+ window_size=(-1, -1),
+ qk_norm=True,
+ cross_attn_norm=True,
+ eps=1e-6,
+ flf_pos_embed_token_number=None,
+ motion_encoder_dim=512,
+ image_model=None,
+ device=None,
+ dtype=None,
+ operations=None,
+ ):
+
+ super().__init__(model_type='i2v', patch_size=patch_size, text_len=text_len, in_dim=in_dim, dim=dim, ffn_dim=ffn_dim, freq_dim=freq_dim, text_dim=text_dim, out_dim=out_dim, num_heads=num_heads, num_layers=num_layers, window_size=window_size, qk_norm=qk_norm, cross_attn_norm=cross_attn_norm, eps=eps, flf_pos_embed_token_number=flf_pos_embed_token_number, image_model=image_model, device=device, dtype=dtype, operations=operations)
+
+ self.pose_patch_embedding = operations.Conv3d(
+ 16, dim, kernel_size=patch_size, stride=patch_size, device=device, dtype=dtype
+ )
+
+ self.motion_encoder = Generator(style_dim=512, motion_dim=20, device=device, dtype=dtype, operations=operations)
+
+ self.face_adapter = FaceAdapter(
+ heads_num=self.num_heads,
+ hidden_dim=self.dim,
+ num_adapter_layers=self.num_layers // 5,
+ device=device, dtype=dtype, operations=operations
+ )
+
+ self.face_encoder = FaceEncoder(
+ in_dim=motion_encoder_dim,
+ hidden_dim=self.dim,
+ num_heads=4,
+ device=device, dtype=dtype, operations=operations
+ )
+
+ def after_patch_embedding(self, x, pose_latents, face_pixel_values):
+ if pose_latents is not None:
+ pose_latents = self.pose_patch_embedding(pose_latents)
+ x[:, :, 1:pose_latents.shape[2] + 1] += pose_latents[:, :, :x.shape[2] - 1]
+
+ if face_pixel_values is None:
+ return x, None
+
+ b, c, T, h, w = face_pixel_values.shape
+ face_pixel_values = rearrange(face_pixel_values, "b c t h w -> (b t) c h w")
+ encode_bs = 8
+ face_pixel_values_tmp = []
+ for i in range(math.ceil(face_pixel_values.shape[0] / encode_bs)):
+ face_pixel_values_tmp.append(self.motion_encoder.get_motion(face_pixel_values[i * encode_bs: (i + 1) * encode_bs]))
+
+ motion_vec = torch.cat(face_pixel_values_tmp)
+
+ motion_vec = rearrange(motion_vec, "(b t) c -> b t c", t=T)
+ motion_vec = self.face_encoder(motion_vec)
+
+ B, L, H, C = motion_vec.shape
+ pad_face = torch.zeros(B, 1, H, C).type_as(motion_vec)
+ motion_vec = torch.cat([pad_face, motion_vec], dim=1)
+
+ if motion_vec.shape[1] < x.shape[2]:
+ B, L, H, C = motion_vec.shape
+ pad = torch.zeros(B, x.shape[2] - motion_vec.shape[1], H, C).type_as(motion_vec)
+ motion_vec = torch.cat([motion_vec, pad], dim=1)
+ else:
+ motion_vec = motion_vec[:, :x.shape[2]]
+ return x, motion_vec
+
+ def forward_orig(
+ self,
+ x,
+ t,
+ context,
+ clip_fea=None,
+ pose_latents=None,
+ face_pixel_values=None,
+ freqs=None,
+ transformer_options={},
+ **kwargs,
+ ):
+ # embeddings
+ x = self.patch_embedding(x.float()).to(x.dtype)
+ x, motion_vec = self.after_patch_embedding(x, pose_latents, face_pixel_values)
+ grid_sizes = x.shape[2:]
+ x = x.flatten(2).transpose(1, 2)
+
+ # time embeddings
+ e = self.time_embedding(
+ sinusoidal_embedding_1d(self.freq_dim, t.flatten()).to(dtype=x[0].dtype))
+ e = e.reshape(t.shape[0], -1, e.shape[-1])
+ e0 = self.time_projection(e).unflatten(2, (6, self.dim))
+
+ full_ref = None
+ if self.ref_conv is not None:
+ full_ref = kwargs.get("reference_latent", None)
+ if full_ref is not None:
+ full_ref = self.ref_conv(full_ref).flatten(2).transpose(1, 2)
+ x = torch.concat((full_ref, x), dim=1)
+
+ # context
+ context = self.text_embedding(context)
+
+ context_img_len = None
+ if clip_fea is not None:
+ if self.img_emb is not None:
+ context_clip = self.img_emb(clip_fea) # bs x 257 x dim
+ context = torch.concat([context_clip, context], dim=1)
+ context_img_len = clip_fea.shape[-2]
+
+ patches_replace = transformer_options.get("patches_replace", {})
+ blocks_replace = patches_replace.get("dit", {})
+ for i, block in enumerate(self.blocks):
+ if ("double_block", i) in blocks_replace:
+ def block_wrap(args):
+ out = {}
+ out["img"] = block(args["img"], context=args["txt"], e=args["vec"], freqs=args["pe"], context_img_len=context_img_len, transformer_options=args["transformer_options"])
+ return out
+ out = blocks_replace[("double_block", i)]({"img": x, "txt": context, "vec": e0, "pe": freqs, "transformer_options": transformer_options}, {"original_block": block_wrap})
+ x = out["img"]
+ else:
+ x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, transformer_options=transformer_options)
+
+ if i % 5 == 0 and motion_vec is not None:
+ x = x + self.face_adapter.fuser_blocks[i // 5](x, motion_vec)
+
+ # head
+ x = self.head(x, e)
+
+ if full_ref is not None:
+ x = x[:, full_ref.shape[1]:]
+
+ # unpatchify
+ x = self.unpatchify(x, grid_sizes)
+ return x
diff --git a/comfy/ldm/wan/vae.py b/comfy/ldm/wan/vae.py
index a8ebc5ec6..ccbb25822 100644
--- a/comfy/ldm/wan/vae.py
+++ b/comfy/ldm/wan/vae.py
@@ -24,12 +24,17 @@ class CausalConv3d(ops.Conv3d):
self.padding[1], 2 * self.padding[0], 0)
self.padding = (0, 0, 0)
- def forward(self, x, cache_x=None):
+ def forward(self, x, cache_x=None, cache_list=None, cache_idx=None):
+ if cache_list is not None:
+ cache_x = cache_list[cache_idx]
+ cache_list[cache_idx] = None
+
padding = list(self._padding)
if cache_x is not None and self._padding[4] > 0:
cache_x = cache_x.to(x.device)
x = torch.cat([cache_x, x], dim=2)
padding[4] -= cache_x.shape[2]
+ del cache_x
x = F.pad(x, padding)
return super().forward(x)
@@ -52,15 +57,6 @@ class RMS_norm(nn.Module):
x, dim=(1 if self.channel_first else -1)) * self.scale * self.gamma.to(x) + (self.bias.to(x) if self.bias is not None else 0)
-class Upsample(nn.Upsample):
-
- def forward(self, x):
- """
- Fix bfloat16 support for nearest neighbor interpolation.
- """
- return super().forward(x.float()).type_as(x)
-
-
class Resample(nn.Module):
def __init__(self, dim, mode):
@@ -73,11 +69,11 @@ class Resample(nn.Module):
# layers
if mode == 'upsample2d':
self.resample = nn.Sequential(
- Upsample(scale_factor=(2., 2.), mode='nearest-exact'),
+ nn.Upsample(scale_factor=(2., 2.), mode='nearest-exact'),
ops.Conv2d(dim, dim // 2, 3, padding=1))
elif mode == 'upsample3d':
self.resample = nn.Sequential(
- Upsample(scale_factor=(2., 2.), mode='nearest-exact'),
+ nn.Upsample(scale_factor=(2., 2.), mode='nearest-exact'),
ops.Conv2d(dim, dim // 2, 3, padding=1))
self.time_conv = CausalConv3d(
dim, dim * 2, (3, 1, 1), padding=(1, 0, 0))
@@ -157,29 +153,6 @@ class Resample(nn.Module):
feat_idx[0] += 1
return x
- def init_weight(self, conv):
- conv_weight = conv.weight
- nn.init.zeros_(conv_weight)
- c1, c2, t, h, w = conv_weight.size()
- one_matrix = torch.eye(c1, c2)
- init_matrix = one_matrix
- nn.init.zeros_(conv_weight)
- #conv_weight.data[:,:,-1,1,1] = init_matrix * 0.5
- conv_weight.data[:, :, 1, 0, 0] = init_matrix #* 0.5
- conv.weight.data.copy_(conv_weight)
- nn.init.zeros_(conv.bias.data)
-
- def init_weight2(self, conv):
- conv_weight = conv.weight.data
- nn.init.zeros_(conv_weight)
- c1, c2, t, h, w = conv_weight.size()
- init_matrix = torch.eye(c1 // 2, c2)
- #init_matrix = repeat(init_matrix, 'o ... -> (o 2) ...').permute(1,0,2).contiguous().reshape(c1,c2)
- conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix
- conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix
- conv.weight.data.copy_(conv_weight)
- nn.init.zeros_(conv.bias.data)
-
class ResidualBlock(nn.Module):
@@ -198,7 +171,7 @@ class ResidualBlock(nn.Module):
if in_dim != out_dim else nn.Identity()
def forward(self, x, feat_cache=None, feat_idx=[0]):
- h = self.shortcut(x)
+ old_x = x
for layer in self.residual:
if isinstance(layer, CausalConv3d) and feat_cache is not None:
idx = feat_idx[0]
@@ -210,12 +183,12 @@ class ResidualBlock(nn.Module):
cache_x.device), cache_x
],
dim=2)
- x = layer(x, feat_cache[idx])
+ x = layer(x, cache_list=feat_cache, cache_idx=idx)
feat_cache[idx] = cache_x
feat_idx[0] += 1
else:
x = layer(x)
- return x + h
+ return x + self.shortcut(old_x)
class AttentionBlock(nn.Module):
@@ -494,74 +467,47 @@ class WanVAE(nn.Module):
self.decoder = Decoder3d(dim, z_dim, dim_mult, num_res_blocks,
attn_scales, self.temperal_upsample, dropout)
- def forward(self, x):
- mu, log_var = self.encode(x)
- z = self.reparameterize(mu, log_var)
- x_recon = self.decode(z)
- return x_recon, mu, log_var
-
def encode(self, x):
- self.clear_cache()
+ conv_idx = [0]
+ feat_map = [None] * count_conv3d(self.decoder)
## cache
t = x.shape[2]
iter_ = 1 + (t - 1) // 4
## 对encode输入的x，按时间拆分为1、4、4、4....
for i in range(iter_):
- self._enc_conv_idx = [0]
+ conv_idx = [0]
if i == 0:
out = self.encoder(
x[:, :, :1, :, :],
- feat_cache=self._enc_feat_map,
- feat_idx=self._enc_conv_idx)
+ feat_cache=feat_map,
+ feat_idx=conv_idx)
else:
out_ = self.encoder(
x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :],
- feat_cache=self._enc_feat_map,
- feat_idx=self._enc_conv_idx)
+ feat_cache=feat_map,
+ feat_idx=conv_idx)
out = torch.cat([out, out_], 2)
mu, log_var = self.conv1(out).chunk(2, dim=1)
- self.clear_cache()
return mu
def decode(self, z):
- self.clear_cache()
+ conv_idx = [0]
+ feat_map = [None] * count_conv3d(self.decoder)
# z: [b,c,t,h,w]
iter_ = z.shape[2]
x = self.conv2(z)
for i in range(iter_):
- self._conv_idx = [0]
+ conv_idx = [0]
if i == 0:
out = self.decoder(
x[:, :, i:i + 1, :, :],
- feat_cache=self._feat_map,
- feat_idx=self._conv_idx)
+ feat_cache=feat_map,
+ feat_idx=conv_idx)
else:
out_ = self.decoder(
x[:, :, i:i + 1, :, :],
- feat_cache=self._feat_map,
- feat_idx=self._conv_idx)
+ feat_cache=feat_map,
+ feat_idx=conv_idx)
out = torch.cat([out, out_], 2)
- self.clear_cache()
return out
-
- def reparameterize(self, mu, log_var):
- std = torch.exp(0.5 * log_var)
- eps = torch.randn_like(std)
- return eps * std + mu
-
- def sample(self, imgs, deterministic=False):
- mu, log_var = self.encode(imgs)
- if deterministic:
- return mu
- std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0))
- return mu + std * torch.randn_like(std)
-
- def clear_cache(self):
- self._conv_num = count_conv3d(self.decoder)
- self._conv_idx = [0]
- self._feat_map = [None] * self._conv_num
- #cache encode
- self._enc_conv_num = count_conv3d(self.encoder)
- self._enc_conv_idx = [0]
- self._enc_feat_map = [None] * self._enc_conv_num
diff --git a/comfy/ldm/wan/vae2_2.py b/comfy/ldm/wan/vae2_2.py
new file mode 100644
index 000000000..8e1593a54
--- /dev/null
+++ b/comfy/ldm/wan/vae2_2.py
@@ -0,0 +1,717 @@
+# original version: https://github.com/Wan-Video/Wan2.2/blob/main/wan/modules/vae2_2.py
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from .vae import AttentionBlock, CausalConv3d, RMS_norm
+
+import comfy.ops
+ops = comfy.ops.disable_weight_init
+
+CACHE_T = 2
+
+
+class Resample(nn.Module):
+
+ def __init__(self, dim, mode):
+ assert mode in (
+ "none",
+ "upsample2d",
+ "upsample3d",
+ "downsample2d",
+ "downsample3d",
+ )
+ super().__init__()
+ self.dim = dim
+ self.mode = mode
+
+ # layers
+ if mode == "upsample2d":
+ self.resample = nn.Sequential(
+ nn.Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
+ ops.Conv2d(dim, dim, 3, padding=1),
+ )
+ elif mode == "upsample3d":
+ self.resample = nn.Sequential(
+ nn.Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"),
+ ops.Conv2d(dim, dim, 3, padding=1),
+ # ops.Conv2d(dim, dim//2, 3, padding=1)
+ )
+ self.time_conv = CausalConv3d(
+ dim, dim * 2, (3, 1, 1), padding=(1, 0, 0))
+ elif mode == "downsample2d":
+ self.resample = nn.Sequential(
+ nn.ZeroPad2d((0, 1, 0, 1)),
+ ops.Conv2d(dim, dim, 3, stride=(2, 2)))
+ elif mode == "downsample3d":
+ self.resample = nn.Sequential(
+ nn.ZeroPad2d((0, 1, 0, 1)),
+ ops.Conv2d(dim, dim, 3, stride=(2, 2)))
+ self.time_conv = CausalConv3d(
+ dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))
+ else:
+ self.resample = nn.Identity()
+
+ def forward(self, x, feat_cache=None, feat_idx=[0]):
+ b, c, t, h, w = x.size()
+ if self.mode == "upsample3d":
+ if feat_cache is not None:
+ idx = feat_idx[0]
+ if feat_cache[idx] is None:
+ feat_cache[idx] = "Rep"
+ feat_idx[0] += 1
+ else:
+ cache_x = x[:, :, -CACHE_T:, :, :].clone()
+ if (cache_x.shape[2] < 2 and feat_cache[idx] is not None and
+ feat_cache[idx] != "Rep"):
+ # cache last frame of last two chunk
+ cache_x = torch.cat(
+ [
+ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+ cache_x.device),
+ cache_x,
+ ],
+ dim=2,
+ )
+ if (cache_x.shape[2] < 2 and feat_cache[idx] is not None and
+ feat_cache[idx] == "Rep"):
+ cache_x = torch.cat(
+ [
+ torch.zeros_like(cache_x).to(cache_x.device),
+ cache_x
+ ],
+ dim=2,
+ )
+ if feat_cache[idx] == "Rep":
+ x = self.time_conv(x)
+ else:
+ x = self.time_conv(x, feat_cache[idx])
+ feat_cache[idx] = cache_x
+ feat_idx[0] += 1
+ x = x.reshape(b, 2, c, t, h, w)
+ x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]),
+ 3)
+ x = x.reshape(b, c, t * 2, h, w)
+ t = x.shape[2]
+ x = rearrange(x, "b c t h w -> (b t) c h w")
+ x = self.resample(x)
+ x = rearrange(x, "(b t) c h w -> b c t h w", t=t)
+
+ if self.mode == "downsample3d":
+ if feat_cache is not None:
+ idx = feat_idx[0]
+ if feat_cache[idx] is None:
+ feat_cache[idx] = x.clone()
+ feat_idx[0] += 1
+ else:
+ cache_x = x[:, :, -1:, :, :].clone()
+ x = self.time_conv(
+ torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))
+ feat_cache[idx] = cache_x
+ feat_idx[0] += 1
+ return x
+
+
+class ResidualBlock(nn.Module):
+
+ def __init__(self, in_dim, out_dim, dropout=0.0):
+ super().__init__()
+ self.in_dim = in_dim
+ self.out_dim = out_dim
+
+ # layers
+ self.residual = nn.Sequential(
+ RMS_norm(in_dim, images=False),
+ nn.SiLU(),
+ CausalConv3d(in_dim, out_dim, 3, padding=1),
+ RMS_norm(out_dim, images=False),
+ nn.SiLU(),
+ nn.Dropout(dropout),
+ CausalConv3d(out_dim, out_dim, 3, padding=1),
+ )
+ self.shortcut = (
+ CausalConv3d(in_dim, out_dim, 1)
+ if in_dim != out_dim else nn.Identity())
+
+ def forward(self, x, feat_cache=None, feat_idx=[0]):
+ old_x = x
+ for layer in self.residual:
+ if isinstance(layer, CausalConv3d) and feat_cache is not None:
+ idx = feat_idx[0]
+ cache_x = x[:, :, -CACHE_T:, :, :].clone()
+ if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+ # cache last frame of last two chunk
+ cache_x = torch.cat(
+ [
+ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+ cache_x.device),
+ cache_x,
+ ],
+ dim=2,
+ )
+ x = layer(x, cache_list=feat_cache, cache_idx=idx)
+ feat_cache[idx] = cache_x
+ feat_idx[0] += 1
+ else:
+ x = layer(x)
+ return x + self.shortcut(old_x)
+
+
+def patchify(x, patch_size):
+ if patch_size == 1:
+ return x
+ if x.dim() == 4:
+ x = rearrange(
+ x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size)
+ elif x.dim() == 5:
+ x = rearrange(
+ x,
+ "b c f (h q) (w r) -> b (c r q) f h w",
+ q=patch_size,
+ r=patch_size,
+ )
+ else:
+ raise ValueError(f"Invalid input shape: {x.shape}")
+
+ return x
+
+
+def unpatchify(x, patch_size):
+ if patch_size == 1:
+ return x
+
+ if x.dim() == 4:
+ x = rearrange(
+ x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size)
+ elif x.dim() == 5:
+ x = rearrange(
+ x,
+ "b (c r q) f h w -> b c f (h q) (w r)",
+ q=patch_size,
+ r=patch_size,
+ )
+ return x
+
+
+class AvgDown3D(nn.Module):
+
+ def __init__(
+ self,
+ in_channels,
+ out_channels,
+ factor_t,
+ factor_s=1,
+ ):
+ super().__init__()
+ self.in_channels = in_channels
+ self.out_channels = out_channels
+ self.factor_t = factor_t
+ self.factor_s = factor_s
+ self.factor = self.factor_t * self.factor_s * self.factor_s
+
+ assert in_channels * self.factor % out_channels == 0
+ self.group_size = in_channels * self.factor // out_channels
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t
+ pad = (0, 0, 0, 0, pad_t, 0)
+ x = F.pad(x, pad)
+ B, C, T, H, W = x.shape
+ x = x.view(
+ B,
+ C,
+ T // self.factor_t,
+ self.factor_t,
+ H // self.factor_s,
+ self.factor_s,
+ W // self.factor_s,
+ self.factor_s,
+ )
+ x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous()
+ x = x.view(
+ B,
+ C * self.factor,
+ T // self.factor_t,
+ H // self.factor_s,
+ W // self.factor_s,
+ )
+ x = x.view(
+ B,
+ self.out_channels,
+ self.group_size,
+ T // self.factor_t,
+ H // self.factor_s,
+ W // self.factor_s,
+ )
+ x = x.mean(dim=2)
+ return x
+
+
+class DupUp3D(nn.Module):
+
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ factor_t,
+ factor_s=1,
+ ):
+ super().__init__()
+ self.in_channels = in_channels
+ self.out_channels = out_channels
+
+ self.factor_t = factor_t
+ self.factor_s = factor_s
+ self.factor = self.factor_t * self.factor_s * self.factor_s
+
+ assert out_channels * self.factor % in_channels == 0
+ self.repeats = out_channels * self.factor // in_channels
+
+ def forward(self, x: torch.Tensor, first_chunk=False) -> torch.Tensor:
+ x = x.repeat_interleave(self.repeats, dim=1)
+ x = x.view(
+ x.size(0),
+ self.out_channels,
+ self.factor_t,
+ self.factor_s,
+ self.factor_s,
+ x.size(2),
+ x.size(3),
+ x.size(4),
+ )
+ x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()
+ x = x.view(
+ x.size(0),
+ self.out_channels,
+ x.size(2) * self.factor_t,
+ x.size(4) * self.factor_s,
+ x.size(6) * self.factor_s,
+ )
+ if first_chunk:
+ x = x[:, :, self.factor_t - 1:, :, :]
+ return x
+
+
+class Down_ResidualBlock(nn.Module):
+
+ def __init__(self,
+ in_dim,
+ out_dim,
+ dropout,
+ mult,
+ temperal_downsample=False,
+ down_flag=False):
+ super().__init__()
+
+ # Shortcut path with downsample
+ self.avg_shortcut = AvgDown3D(
+ in_dim,
+ out_dim,
+ factor_t=2 if temperal_downsample else 1,
+ factor_s=2 if down_flag else 1,
+ )
+
+ # Main path with residual blocks and downsample
+ downsamples = []
+ for _ in range(mult):
+ downsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+ in_dim = out_dim
+
+ # Add the final downsample block
+ if down_flag:
+ mode = "downsample3d" if temperal_downsample else "downsample2d"
+ downsamples.append(Resample(out_dim, mode=mode))
+
+ self.downsamples = nn.Sequential(*downsamples)
+
+ def forward(self, x, feat_cache=None, feat_idx=[0]):
+ x_copy = x
+ for module in self.downsamples:
+ x = module(x, feat_cache, feat_idx)
+
+ return x + self.avg_shortcut(x_copy)
+
+
+class Up_ResidualBlock(nn.Module):
+
+ def __init__(self,
+ in_dim,
+ out_dim,
+ dropout,
+ mult,
+ temperal_upsample=False,
+ up_flag=False):
+ super().__init__()
+ # Shortcut path with upsample
+ if up_flag:
+ self.avg_shortcut = DupUp3D(
+ in_dim,
+ out_dim,
+ factor_t=2 if temperal_upsample else 1,
+ factor_s=2 if up_flag else 1,
+ )
+ else:
+ self.avg_shortcut = None
+
+ # Main path with residual blocks and upsample
+ upsamples = []
+ for _ in range(mult):
+ upsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+ in_dim = out_dim
+
+ # Add the final upsample block
+ if up_flag:
+ mode = "upsample3d" if temperal_upsample else "upsample2d"
+ upsamples.append(Resample(out_dim, mode=mode))
+
+ self.upsamples = nn.Sequential(*upsamples)
+
+ def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False):
+ x_main = x
+ for module in self.upsamples:
+ x_main = module(x_main, feat_cache, feat_idx)
+ if self.avg_shortcut is not None:
+ x_shortcut = self.avg_shortcut(x, first_chunk)
+ return x_main + x_shortcut
+ else:
+ return x_main
+
+
+class Encoder3d(nn.Module):
+
+ def __init__(
+ self,
+ dim=128,
+ z_dim=4,
+ dim_mult=[1, 2, 4, 4],
+ num_res_blocks=2,
+ attn_scales=[],
+ temperal_downsample=[True, True, False],
+ dropout=0.0,
+ ):
+ super().__init__()
+ self.dim = dim
+ self.z_dim = z_dim
+ self.dim_mult = dim_mult
+ self.num_res_blocks = num_res_blocks
+ self.attn_scales = attn_scales
+ self.temperal_downsample = temperal_downsample
+
+ # dimensions
+ dims = [dim * u for u in [1] + dim_mult]
+ scale = 1.0
+
+ # init block
+ self.conv1 = CausalConv3d(12, dims[0], 3, padding=1)
+
+ # downsample blocks
+ downsamples = []
+ for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+ t_down_flag = (
+ temperal_downsample[i]
+ if i < len(temperal_downsample) else False)
+ downsamples.append(
+ Down_ResidualBlock(
+ in_dim=in_dim,
+ out_dim=out_dim,
+ dropout=dropout,
+ mult=num_res_blocks,
+ temperal_downsample=t_down_flag,
+ down_flag=i != len(dim_mult) - 1,
+ ))
+ scale /= 2.0
+ self.downsamples = nn.Sequential(*downsamples)
+
+ # middle blocks
+ self.middle = nn.Sequential(
+ ResidualBlock(out_dim, out_dim, dropout),
+ AttentionBlock(out_dim),
+ ResidualBlock(out_dim, out_dim, dropout),
+ )
+
+ # # output blocks
+ self.head = nn.Sequential(
+ RMS_norm(out_dim, images=False),
+ nn.SiLU(),
+ CausalConv3d(out_dim, z_dim, 3, padding=1),
+ )
+
+ def forward(self, x, feat_cache=None, feat_idx=[0]):
+
+ if feat_cache is not None:
+ idx = feat_idx[0]
+ cache_x = x[:, :, -CACHE_T:, :, :].clone()
+ if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+ cache_x = torch.cat(
+ [
+ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+ cache_x.device),
+ cache_x,
+ ],
+ dim=2,
+ )
+ x = self.conv1(x, feat_cache[idx])
+ feat_cache[idx] = cache_x
+ feat_idx[0] += 1
+ else:
+ x = self.conv1(x)
+
+ ## downsamples
+ for layer in self.downsamples:
+ if feat_cache is not None:
+ x = layer(x, feat_cache, feat_idx)
+ else:
+ x = layer(x)
+
+ ## middle
+ for layer in self.middle:
+ if isinstance(layer, ResidualBlock) and feat_cache is not None:
+ x = layer(x, feat_cache, feat_idx)
+ else:
+ x = layer(x)
+
+ ## head
+ for layer in self.head:
+ if isinstance(layer, CausalConv3d) and feat_cache is not None:
+ idx = feat_idx[0]
+ cache_x = x[:, :, -CACHE_T:, :, :].clone()
+ if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+ cache_x = torch.cat(
+ [
+ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+ cache_x.device),
+ cache_x,
+ ],
+ dim=2,
+ )
+ x = layer(x, feat_cache[idx])
+ feat_cache[idx] = cache_x
+ feat_idx[0] += 1
+ else:
+ x = layer(x)
+
+ return x
+
+
+class Decoder3d(nn.Module):
+
+ def __init__(
+ self,
+ dim=128,
+ z_dim=4,
+ dim_mult=[1, 2, 4, 4],
+ num_res_blocks=2,
+ attn_scales=[],
+ temperal_upsample=[False, True, True],
+ dropout=0.0,
+ ):
+ super().__init__()
+ self.dim = dim
+ self.z_dim = z_dim
+ self.dim_mult = dim_mult
+ self.num_res_blocks = num_res_blocks
+ self.attn_scales = attn_scales
+ self.temperal_upsample = temperal_upsample
+
+ # dimensions
+ dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]
+ # init block
+ self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1)
+
+ # middle blocks
+ self.middle = nn.Sequential(
+ ResidualBlock(dims[0], dims[0], dropout),
+ AttentionBlock(dims[0]),
+ ResidualBlock(dims[0], dims[0], dropout),
+ )
+
+ # upsample blocks
+ upsamples = []
+ for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+ t_up_flag = temperal_upsample[i] if i < len(
+ temperal_upsample) else False
+ upsamples.append(
+ Up_ResidualBlock(
+ in_dim=in_dim,
+ out_dim=out_dim,
+ dropout=dropout,
+ mult=num_res_blocks + 1,
+ temperal_upsample=t_up_flag,
+ up_flag=i != len(dim_mult) - 1,
+ ))
+ self.upsamples = nn.Sequential(*upsamples)
+
+ # output blocks
+ self.head = nn.Sequential(
+ RMS_norm(out_dim, images=False),
+ nn.SiLU(),
+ CausalConv3d(out_dim, 12, 3, padding=1),
+ )
+
+ def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False):
+ if feat_cache is not None:
+ idx = feat_idx[0]
+ cache_x = x[:, :, -CACHE_T:, :, :].clone()
+ if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+ cache_x = torch.cat(
+ [
+ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+ cache_x.device),
+ cache_x,
+ ],
+ dim=2,
+ )
+ x = self.conv1(x, feat_cache[idx])
+ feat_cache[idx] = cache_x
+ feat_idx[0] += 1
+ else:
+ x = self.conv1(x)
+
+ for layer in self.middle:
+ if isinstance(layer, ResidualBlock) and feat_cache is not None:
+ x = layer(x, feat_cache, feat_idx)
+ else:
+ x = layer(x)
+
+ ## upsamples
+ for layer in self.upsamples:
+ if feat_cache is not None:
+ x = layer(x, feat_cache, feat_idx, first_chunk)
+ else:
+ x = layer(x)
+
+ ## head
+ for layer in self.head:
+ if isinstance(layer, CausalConv3d) and feat_cache is not None:
+ idx = feat_idx[0]
+ cache_x = x[:, :, -CACHE_T:, :, :].clone()
+ if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+ cache_x = torch.cat(
+ [
+ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+ cache_x.device),
+ cache_x,
+ ],
+ dim=2,
+ )
+ x = layer(x, feat_cache[idx])
+ feat_cache[idx] = cache_x
+ feat_idx[0] += 1
+ else:
+ x = layer(x)
+ return x
+
+
+def count_conv3d(model):
+ count = 0
+ for m in model.modules():
+ if isinstance(m, CausalConv3d):
+ count += 1
+ return count
+
+
+class WanVAE(nn.Module):
+
+ def __init__(
+ self,
+ dim=160,
+ dec_dim=256,
+ z_dim=16,
+ dim_mult=[1, 2, 4, 4],
+ num_res_blocks=2,
+ attn_scales=[],
+ temperal_downsample=[True, True, False],
+ dropout=0.0,
+ ):
+ super().__init__()
+ self.dim = dim
+ self.z_dim = z_dim
+ self.dim_mult = dim_mult
+ self.num_res_blocks = num_res_blocks
+ self.attn_scales = attn_scales
+ self.temperal_downsample = temperal_downsample
+ self.temperal_upsample = temperal_downsample[::-1]
+
+ # modules
+ self.encoder = Encoder3d(
+ dim,
+ z_dim * 2,
+ dim_mult,
+ num_res_blocks,
+ attn_scales,
+ self.temperal_downsample,
+ dropout,
+ )
+ self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1)
+ self.conv2 = CausalConv3d(z_dim, z_dim, 1)
+ self.decoder = Decoder3d(
+ dec_dim,
+ z_dim,
+ dim_mult,
+ num_res_blocks,
+ attn_scales,
+ self.temperal_upsample,
+ dropout,
+ )
+
+ def encode(self, x):
+ conv_idx = [0]
+ feat_map = [None] * count_conv3d(self.encoder)
+ x = patchify(x, patch_size=2)
+ t = x.shape[2]
+ iter_ = 1 + (t - 1) // 4
+ for i in range(iter_):
+ conv_idx = [0]
+ if i == 0:
+ out = self.encoder(
+ x[:, :, :1, :, :],
+ feat_cache=feat_map,
+ feat_idx=conv_idx,
+ )
+ else:
+ out_ = self.encoder(
+ x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :],
+ feat_cache=feat_map,
+ feat_idx=conv_idx,
+ )
+ out = torch.cat([out, out_], 2)
+ mu, log_var = self.conv1(out).chunk(2, dim=1)
+ return mu
+
+ def decode(self, z):
+ conv_idx = [0]
+ feat_map = [None] * count_conv3d(self.decoder)
+ iter_ = z.shape[2]
+ x = self.conv2(z)
+ for i in range(iter_):
+ conv_idx = [0]
+ if i == 0:
+ out = self.decoder(
+ x[:, :, i:i + 1, :, :],
+ feat_cache=feat_map,
+ feat_idx=conv_idx,
+ first_chunk=True,
+ )
+ else:
+ out_ = self.decoder(
+ x[:, :, i:i + 1, :, :],
+ feat_cache=feat_map,
+ feat_idx=conv_idx,
+ )
+ out = torch.cat([out, out_], 2)
+ out = unpatchify(out, patch_size=2)
+ return out
+
+ def reparameterize(self, mu, log_var):
+ std = torch.exp(0.5 * log_var)
+ eps = torch.randn_like(std)
+ return eps * std + mu
+
+ def sample(self, imgs, deterministic=False):
+ mu, log_var = self.encode(imgs)
+ if deterministic:
+ return mu
+ std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0))
+ return mu + std * torch.randn_like(std)
diff --git a/comfy/lora.py b/comfy/lora.py
index 387d5c52a..36d26293a 100644
--- a/comfy/lora.py
+++ b/comfy/lora.py
@@ -260,6 +260,10 @@ def model_lora_keys_unet(model, key_map={}):
key_map["transformer.{}".format(k[:-len(".weight")])] = to #simpletrainer and probably regular diffusers flux lora format
key_map["lycoris_{}".format(k[:-len(".weight")].replace(".", "_"))] = to #simpletrainer lycoris
key_map["lora_transformer_{}".format(k[:-len(".weight")].replace(".", "_"))] = to #onetrainer
+ for k in sdk:
+ hidden_size = model.model_config.unet_config.get("hidden_size", 0)
+ if k.endswith(".weight") and ".linear1." in k:
+ key_map["{}".format(k.replace(".linear1.weight", ".linear1_qkv"))] = (k, (0, 0, hidden_size * 3))
if isinstance(model, comfy.model_base.GenmoMochi):
for k in sdk:
@@ -293,6 +297,22 @@ def model_lora_keys_unet(model, key_map={}):
key_lora = k[len("diffusion_model."):-len(".weight")]
key_map["{}".format(key_lora)] = k
+ if isinstance(model, comfy.model_base.Omnigen2):
+ for k in sdk:
+ if k.startswith("diffusion_model.") and k.endswith(".weight"):
+ key_lora = k[len("diffusion_model."):-len(".weight")]
+ key_map["{}".format(key_lora)] = k
+
+ if isinstance(model, comfy.model_base.QwenImage):
+ for k in sdk:
+ if k.startswith("diffusion_model.") and k.endswith(".weight"): #QwenImage lora format
+ key_lora = k[len("diffusion_model."):-len(".weight")]
+ # Direct mapping for transformer_blocks format (QwenImage LoRA format)
+ key_map["{}".format(key_lora)] = k
+ # Support transformer prefix format
+ key_map["transformer.{}".format(key_lora)] = k
+ key_map["lycoris_{}".format(key_lora.replace(".", "_"))] = k #SimpleTuner lycoris format
+
return key_map
diff --git a/comfy/lora_convert.py b/comfy/lora_convert.py
index 3e00b63db..9d8d21efe 100644
--- a/comfy/lora_convert.py
+++ b/comfy/lora_convert.py
@@ -15,10 +15,29 @@ def convert_lora_bfl_control(sd): #BFL loras for Flux
def convert_lora_wan_fun(sd): #Wan Fun loras
return comfy.utils.state_dict_prefix_replace(sd, {"lora_unet__": "lora_unet_"})
+def convert_uso_lora(sd):
+ sd_out = {}
+ for k in sd:
+ tensor = sd[k]
+ k_to = "diffusion_model.{}".format(k.replace(".down.weight", ".lora_down.weight")
+ .replace(".up.weight", ".lora_up.weight")
+ .replace(".qkv_lora2.", ".txt_attn.qkv.")
+ .replace(".qkv_lora1.", ".img_attn.qkv.")
+ .replace(".proj_lora1.", ".img_attn.proj.")
+ .replace(".proj_lora2.", ".txt_attn.proj.")
+ .replace(".qkv_lora.", ".linear1_qkv.")
+ .replace(".proj_lora.", ".linear2.")
+ .replace(".processor.", ".")
+ )
+ sd_out[k_to] = tensor
+ return sd_out
+
def convert_lora(sd):
if "img_in.lora_A.weight" in sd and "single_blocks.0.norm.key_norm.scale" in sd:
return convert_lora_bfl_control(sd)
if "lora_unet__blocks_0_cross_attn_k.lora_down.weight" in sd:
return convert_lora_wan_fun(sd)
+ if "single_blocks.37.processor.qkv_lora.up.weight" in sd and "double_blocks.18.processor.qkv_lora2.up.weight" in sd:
+ return convert_uso_lora(sd)
return sd
diff --git a/comfy/model_base.py b/comfy/model_base.py
index c9fe4e4ef..56391257f 100644
--- a/comfy/model_base.py
+++ b/comfy/model_base.py
@@ -16,6 +16,8 @@
along with this program. If not, see .
"""
+import comfy.ldm.hunyuan3dv2_1
+import comfy.ldm.hunyuan3dv2_1.hunyuandit
import torch
import logging
from comfy.ldm.modules.diffusionmodules.openaimodel import UNetModel, Timestep
@@ -37,12 +39,15 @@ import comfy.ldm.cosmos.model
import comfy.ldm.cosmos.predict2
import comfy.ldm.lumina.model
import comfy.ldm.wan.model
+import comfy.ldm.wan.model_animate
import comfy.ldm.hunyuan3d.model
import comfy.ldm.hidream.model
import comfy.ldm.chroma.model
+import comfy.ldm.chroma_radiance.model
import comfy.ldm.ace.model
import comfy.ldm.omnigen.omnigen2
import comfy.ldm.hunyuan_image_3.model
+import comfy.ldm.qwen_image.model
import comfy.model_management
import comfy.patcher_extension
@@ -107,10 +112,12 @@ def model_sampling(model_config, model_type):
return ModelSampling(model_config)
-def convert_tensor(extra, dtype):
+def convert_tensor(extra, dtype, device):
if hasattr(extra, "dtype"):
if extra.dtype != torch.int and extra.dtype != torch.long:
- extra = extra.to(dtype)
+ extra = comfy.model_management.cast_to_device(extra, device, dtype)
+ else:
+ extra = comfy.model_management.cast_to_device(extra, device, None)
return extra
@@ -128,10 +135,11 @@ class BaseModel(torch.nn.Module):
if not unet_config.get("disable_unet_model_creation", False):
if model_config.custom_operations is None:
fp8 = model_config.optimizations.get("fp8", False)
- operations = comfy.ops.pick_operations(unet_config.get("dtype", None), self.manual_cast_dtype, fp8_optimizations=fp8, scaled_fp8=model_config.scaled_fp8)
+ operations = comfy.ops.pick_operations(unet_config.get("dtype", None), self.manual_cast_dtype, fp8_optimizations=fp8, scaled_fp8=model_config.scaled_fp8, model_config=model_config)
else:
operations = model_config.custom_operations
self.diffusion_model = unet_model(**unet_config, device=device, operations=operations)
+ self.diffusion_model.eval()
if comfy.model_management.force_channels_last():
self.diffusion_model.to(memory_format=torch.channels_last)
logging.debug("using channels last mode for diffusion model")
@@ -148,6 +156,7 @@ class BaseModel(torch.nn.Module):
logging.debug("adm {}".format(self.adm_channels))
self.memory_usage_factor = model_config.memory_usage_factor
self.memory_usage_factor_conds = ()
+ self.memory_usage_shape_process = {}
def apply_model(self, x, t, c_concat=None, c_crossattn=None, control=None, transformer_options={}, **kwargs):
return comfy.patcher_extension.WrapperExecutor.new_class_executor(
@@ -161,7 +170,7 @@ class BaseModel(torch.nn.Module):
xc = self.model_sampling.calculate_input(sigma, x)
if c_concat is not None:
- xc = torch.cat([xc] + [c_concat], dim=1)
+ xc = torch.cat([xc] + [comfy.model_management.cast_to_device(c_concat, xc.device, xc.dtype)], dim=1)
context = c_crossattn
dtype = self.get_dtype()
@@ -170,26 +179,33 @@ class BaseModel(torch.nn.Module):
dtype = self.manual_cast_dtype
xc = xc.to(dtype)
+ device = xc.device
t = self.model_sampling.timestep(t).float()
if context is not None:
- context = context.to(dtype)
+ context = comfy.model_management.cast_to_device(context, device, dtype)
extra_conds = {}
for o in kwargs:
extra = kwargs[o]
if hasattr(extra, "dtype"):
- extra = convert_tensor(extra, dtype)
+ extra = convert_tensor(extra, dtype, device)
elif isinstance(extra, list):
ex = []
for ext in extra:
- ex.append(convert_tensor(ext, dtype))
+ ex.append(convert_tensor(ext, dtype, device))
extra = ex
extra_conds[o] = extra
t = self.process_timestep(t, x=x, **extra_conds)
- model_output = self.diffusion_model(xc, t, context=context, control=control, transformer_options=transformer_options, **extra_conds).float()
- return self.model_sampling.calculate_denoised(sigma, model_output, x)
+ if "latent_shapes" in extra_conds:
+ xc = utils.unpack_latents(xc, extra_conds.pop("latent_shapes"))
+
+ model_output = self.diffusion_model(xc, t, context=context, control=control, transformer_options=transformer_options, **extra_conds)
+ if len(model_output) > 1 and not torch.is_tensor(model_output):
+ model_output, _ = utils.pack_latents(model_output)
+
+ return self.model_sampling.calculate_denoised(sigma, model_output.float(), x)
def process_timestep(self, timestep, **kwargs):
return timestep
@@ -318,6 +334,14 @@ class BaseModel(torch.nn.Module):
if self.model_config.scaled_fp8 is not None:
unet_state_dict["scaled_fp8"] = torch.tensor([], dtype=self.model_config.scaled_fp8)
+ # Save mixed precision metadata
+ if hasattr(self.model_config, 'layer_quant_config') and self.model_config.layer_quant_config:
+ metadata = {
+ "format_version": "1.0",
+ "layers": self.model_config.layer_quant_config
+ }
+ unet_state_dict["_quantization_metadata"] = metadata
+
unet_state_dict = self.model_config.process_unet_state_dict_for_saving(unet_state_dict)
if self.model_type == ModelType.V_PREDICTION:
@@ -347,8 +371,15 @@ class BaseModel(torch.nn.Module):
input_shapes = [input_shape]
for c in self.memory_usage_factor_conds:
shape = cond_shapes.get(c, None)
- if shape is not None and len(shape) > 0:
- input_shapes += shape
+ if shape is not None:
+ if c in self.memory_usage_shape_process:
+ out = []
+ for s in shape:
+ out.append(self.memory_usage_shape_process[c](s))
+ shape = out
+
+ if len(shape) > 0:
+ input_shapes += shape
if comfy.model_management.xformers_enabled() or comfy.model_management.pytorch_attention_flash_attention():
dtype = self.get_dtype()
@@ -399,7 +430,7 @@ class SD21UNCLIP(BaseModel):
unclip_conditioning = kwargs.get("unclip_conditioning", None)
device = kwargs["device"]
if unclip_conditioning is None:
- return torch.zeros((1, self.adm_channels))
+ return torch.zeros((1, self.adm_channels), device=device)
else:
return unclip_adm(unclip_conditioning, device, self.noise_augmentor, kwargs.get("unclip_noise_augment_merge", 0.05), kwargs.get("seed", 0) - 10)
@@ -613,9 +644,11 @@ class IP2P:
if image is None:
image = torch.zeros_like(noise)
+ else:
+ image = image.to(device=device)
if image.shape[1:] != noise.shape[1:]:
- image = utils.common_upscale(image.to(device), noise.shape[-1], noise.shape[-2], "bilinear", "center")
+ image = utils.common_upscale(image, noise.shape[-1], noise.shape[-2], "bilinear", "center")
image = utils.resize_to_batch_size(image, noise.shape[0])
return self.process_ip2p_image_in(image)
@@ -652,7 +685,6 @@ class Lotus(BaseModel):
class StableCascade_C(BaseModel):
def __init__(self, model_config, model_type=ModelType.STABLE_CASCADE, device=None):
super().__init__(model_config, model_type, device=device, unet_model=StageC)
- self.diffusion_model.eval().requires_grad_(False)
def extra_conds(self, **kwargs):
out = {}
@@ -681,7 +713,6 @@ class StableCascade_C(BaseModel):
class StableCascade_B(BaseModel):
def __init__(self, model_config, model_type=ModelType.STABLE_CASCADE, device=None):
super().__init__(model_config, model_type, device=device, unet_model=StageB)
- self.diffusion_model.eval().requires_grad_(False)
def extra_conds(self, **kwargs):
out = {}
@@ -694,7 +725,7 @@ class StableCascade_B(BaseModel):
#size of prior doesn't really matter if zeros because it gets resized but I still want it to get batched
prior = kwargs.get("stable_cascade_prior", torch.zeros((1, 16, (noise.shape[2] * 4) // 42, (noise.shape[3] * 4) // 42), dtype=noise.dtype, layout=noise.layout, device=noise.device))
- out["effnet"] = comfy.conds.CONDRegular(prior)
+ out["effnet"] = comfy.conds.CONDRegular(prior.to(device=noise.device))
out["sca"] = comfy.conds.CONDRegular(torch.zeros((1,)))
return out
@@ -885,6 +916,10 @@ class Flux(BaseModel):
for lat in ref_latents:
latents.append(self.process_latent_in(lat))
out['ref_latents'] = comfy.conds.CONDList(latents)
+
+ ref_latents_method = kwargs.get("reference_latents_method", None)
+ if ref_latents_method is not None:
+ out['ref_latents_method'] = comfy.conds.CONDConstant(ref_latents_method)
return out
def extra_conds_shapes(self, **kwargs):
@@ -1093,13 +1128,15 @@ class WAN21(BaseModel):
shape_image[1] = extra_channels
image = torch.zeros(shape_image, dtype=noise.dtype, layout=noise.layout, device=noise.device)
else:
+ latent_dim = self.latent_format.latent_channels
image = utils.common_upscale(image.to(device), noise.shape[-1], noise.shape[-2], "bilinear", "center")
- for i in range(0, image.shape[1], 16):
- image[:, i: i + 16] = self.process_latent_in(image[:, i: i + 16])
+ for i in range(0, image.shape[1], latent_dim):
+ image[:, i: i + latent_dim] = self.process_latent_in(image[:, i: i + latent_dim])
image = utils.resize_to_batch_size(image, noise.shape[0])
- if not self.image_to_video or extra_channels == image.shape[1]:
- return image
+ if extra_channels != image.shape[1] + 4:
+ if not self.image_to_video or extra_channels == image.shape[1]:
+ return image
if image.shape[1] > (extra_channels - 4):
image = image[:, :(extra_channels - 4)]
@@ -1118,7 +1155,11 @@ class WAN21(BaseModel):
mask = mask.repeat(1, 4, 1, 1, 1)
mask = utils.resize_to_batch_size(mask, noise.shape[0])
- return torch.cat((mask, image), dim=1)
+ concat_mask_index = kwargs.get("concat_mask_index", 0)
+ if concat_mask_index != 0:
+ return torch.cat((image[:, :concat_mask_index], mask, image[:, concat_mask_index:]), dim=1)
+ else:
+ return torch.cat((mask, image), dim=1)
def extra_conds(self, **kwargs):
out = super().extra_conds(**kwargs)
@@ -1134,6 +1175,10 @@ class WAN21(BaseModel):
if time_dim_concat is not None:
out['time_dim_concat'] = comfy.conds.CONDRegular(self.process_latent_in(time_dim_concat))
+ reference_latents = kwargs.get("reference_latents", None)
+ if reference_latents is not None:
+ out['reference_latent'] = comfy.conds.CONDRegular(self.process_latent_in(reference_latents[-1])[:, :, 0])
+
return out
@@ -1158,10 +1203,10 @@ class WAN21_Vace(WAN21):
vace_frames_out = []
for j in range(len(vace_frames)):
- vf = vace_frames[j].clone()
+ vf = vace_frames[j].to(device=noise.device, dtype=noise.dtype, copy=True)
for i in range(0, vf.shape[1], 16):
vf[:, i:i + 16] = self.process_latent_in(vf[:, i:i + 16])
- vf = torch.cat([vf, mask[j]], dim=1)
+ vf = torch.cat([vf, mask[j].to(device=noise.device, dtype=noise.dtype)], dim=1)
vace_frames_out.append(vf)
vace_frames = torch.stack(vace_frames_out, dim=1)
@@ -1183,6 +1228,120 @@ class WAN21_Camera(WAN21):
out['camera_conditions'] = comfy.conds.CONDRegular(camera_conditions)
return out
+class WAN21_HuMo(WAN21):
+ def __init__(self, model_config, model_type=ModelType.FLOW, image_to_video=False, device=None):
+ super(WAN21, self).__init__(model_config, model_type, device=device, unet_model=comfy.ldm.wan.model.HumoWanModel)
+ self.image_to_video = image_to_video
+
+ def extra_conds(self, **kwargs):
+ out = super().extra_conds(**kwargs)
+ noise = kwargs.get("noise", None)
+
+ audio_embed = kwargs.get("audio_embed", None)
+ if audio_embed is not None:
+ out['audio_embed'] = comfy.conds.CONDRegular(audio_embed)
+
+ if "c_concat" not in out: # 1.7B model
+ reference_latents = kwargs.get("reference_latents", None)
+ if reference_latents is not None:
+ out['reference_latent'] = comfy.conds.CONDRegular(self.process_latent_in(reference_latents[-1]))
+ else:
+ noise_shape = list(noise.shape)
+ noise_shape[1] += 4
+ concat_latent = torch.zeros(noise_shape, device=noise.device, dtype=noise.dtype)
+ zero_vae_values_first = torch.tensor([0.8660, -0.4326, -0.0017, -0.4884, -0.5283, 0.9207, -0.9896, 0.4433, -0.5543, -0.0113, 0.5753, -0.6000, -0.8346, -0.3497, -0.1926, -0.6938]).view(1, 16, 1, 1, 1)
+ zero_vae_values_second = torch.tensor([1.0869, -1.2370, 0.0206, -0.4357, -0.6411, 2.0307, -1.5972, 1.2659, -0.8595, -0.4654, 0.9638, -1.6330, -1.4310, -0.1098, -0.3856, -1.4583]).view(1, 16, 1, 1, 1)
+ zero_vae_values = torch.tensor([0.8642, -1.8583, 0.1577, 0.1350, -0.3641, 2.5863, -1.9670, 1.6065, -1.0475, -0.8678, 1.1734, -1.8138, -1.5933, -0.7721, -0.3289, -1.3745]).view(1, 16, 1, 1, 1)
+ concat_latent[:, 4:] = zero_vae_values
+ concat_latent[:, 4:, :1] = zero_vae_values_first
+ concat_latent[:, 4:, 1:2] = zero_vae_values_second
+ out['c_concat'] = comfy.conds.CONDNoiseShape(concat_latent)
+ reference_latents = kwargs.get("reference_latents", None)
+ if reference_latents is not None:
+ ref_latent = self.process_latent_in(reference_latents[-1])
+ ref_latent_shape = list(ref_latent.shape)
+ ref_latent_shape[1] += 4 + ref_latent_shape[1]
+ ref_latent_full = torch.zeros(ref_latent_shape, device=ref_latent.device, dtype=ref_latent.dtype)
+ ref_latent_full[:, 20:] = ref_latent
+ ref_latent_full[:, 16:20] = 1.0
+ out['reference_latent'] = comfy.conds.CONDRegular(ref_latent_full)
+
+ return out
+
+class WAN22_Animate(WAN21):
+ def __init__(self, model_config, model_type=ModelType.FLOW, image_to_video=False, device=None):
+ super(WAN21, self).__init__(model_config, model_type, device=device, unet_model=comfy.ldm.wan.model_animate.AnimateWanModel)
+ self.image_to_video = image_to_video
+
+ def extra_conds(self, **kwargs):
+ out = super().extra_conds(**kwargs)
+
+ face_video_pixels = kwargs.get("face_video_pixels", None)
+ if face_video_pixels is not None:
+ out['face_pixel_values'] = comfy.conds.CONDRegular(face_video_pixels)
+
+ pose_latents = kwargs.get("pose_video_latent", None)
+ if pose_latents is not None:
+ out['pose_latents'] = comfy.conds.CONDRegular(self.process_latent_in(pose_latents))
+ return out
+
+class WAN22_S2V(WAN21):
+ def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
+ super(WAN21, self).__init__(model_config, model_type, device=device, unet_model=comfy.ldm.wan.model.WanModel_S2V)
+ self.memory_usage_factor_conds = ("reference_latent", "reference_motion")
+ self.memory_usage_shape_process = {"reference_motion": lambda shape: [shape[0], shape[1], 1.5, shape[-2], shape[-1]]}
+
+ def extra_conds(self, **kwargs):
+ out = super().extra_conds(**kwargs)
+ audio_embed = kwargs.get("audio_embed", None)
+ if audio_embed is not None:
+ out['audio_embed'] = comfy.conds.CONDRegular(audio_embed)
+
+ reference_latents = kwargs.get("reference_latents", None)
+ if reference_latents is not None:
+ out['reference_latent'] = comfy.conds.CONDRegular(self.process_latent_in(reference_latents[-1]))
+
+ reference_motion = kwargs.get("reference_motion", None)
+ if reference_motion is not None:
+ out['reference_motion'] = comfy.conds.CONDRegular(self.process_latent_in(reference_motion))
+
+ control_video = kwargs.get("control_video", None)
+ if control_video is not None:
+ out['control_video'] = comfy.conds.CONDRegular(self.process_latent_in(control_video))
+ return out
+
+ def extra_conds_shapes(self, **kwargs):
+ out = {}
+ ref_latents = kwargs.get("reference_latents", None)
+ if ref_latents is not None:
+ out['reference_latent'] = list([1, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16])
+
+ reference_motion = kwargs.get("reference_motion", None)
+ if reference_motion is not None:
+ out['reference_motion'] = reference_motion.shape
+ return out
+
+class WAN22(WAN21):
+ def __init__(self, model_config, model_type=ModelType.FLOW, image_to_video=False, device=None):
+ super(WAN21, self).__init__(model_config, model_type, device=device, unet_model=comfy.ldm.wan.model.WanModel)
+ self.image_to_video = image_to_video
+
+ def extra_conds(self, **kwargs):
+ out = super().extra_conds(**kwargs)
+ denoise_mask = kwargs.get("denoise_mask", None)
+ if denoise_mask is not None:
+ out["denoise_mask"] = comfy.conds.CONDRegular(denoise_mask)
+ return out
+
+ def process_timestep(self, timestep, x, denoise_mask=None, **kwargs):
+ if denoise_mask is None:
+ return timestep
+ temp_ts = (torch.mean(denoise_mask[:, :, :, :, :], dim=(1, 3, 4), keepdim=True) * timestep.view([timestep.shape[0]] + [1] * (denoise_mask.ndim - 1))).reshape(timestep.shape[0], -1)
+ return temp_ts
+
+ def scale_latent_inpaint(self, sigma, noise, latent_image, **kwargs):
+ return latent_image
+
class Hunyuan3Dv2(BaseModel):
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.hunyuan3d.model.Hunyuan3Dv2)
@@ -1202,6 +1361,21 @@ class HunyuanImage3(BaseModel):
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
super().__init__(model_config, model_type, device, unet_model = comfy.ldm.hunyuan_image_3.model.HunyuanImage3ForCausalMM)
+class Hunyuan3Dv2_1(BaseModel):
+ def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
+ super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.hunyuan3dv2_1.hunyuandit.HunYuanDiTPlain)
+
+ def extra_conds(self, **kwargs):
+ out = super().extra_conds(**kwargs)
+ cross_attn = kwargs.get("cross_attn", None)
+ if cross_attn is not None:
+ out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
+
+ guidance = kwargs.get("guidance", 5.0)
+ if guidance is not None:
+ out['guidance'] = comfy.conds.CONDRegular(torch.FloatTensor([guidance]))
+ return out
+
class HiDream(BaseModel):
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.hidream.model.HiDreamImageTransformer2DModel)
@@ -1223,8 +1397,8 @@ class HiDream(BaseModel):
return out
class Chroma(Flux):
- def __init__(self, model_config, model_type=ModelType.FLUX, device=None):
- super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.chroma.model.Chroma)
+ def __init__(self, model_config, model_type=ModelType.FLUX, device=None, unet_model=comfy.ldm.chroma.model.Chroma):
+ super().__init__(model_config, model_type, device=device, unet_model=unet_model)
def extra_conds(self, **kwargs):
out = super().extra_conds(**kwargs)
@@ -1234,6 +1408,10 @@ class Chroma(Flux):
out['guidance'] = comfy.conds.CONDRegular(torch.FloatTensor([guidance]))
return out
+class ChromaRadiance(Chroma):
+ def __init__(self, model_config, model_type=ModelType.FLUX, device=None):
+ super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.chroma_radiance.model.ChromaRadiance)
+
class ACEStep(BaseModel):
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.ace.model.ACEStepTransformer2DModel)
@@ -1282,3 +1460,84 @@ class Omnigen2(BaseModel):
if ref_latents is not None:
out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16])
return out
+
+class QwenImage(BaseModel):
+ def __init__(self, model_config, model_type=ModelType.FLUX, device=None):
+ super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.qwen_image.model.QwenImageTransformer2DModel)
+ self.memory_usage_factor_conds = ("ref_latents",)
+
+ def extra_conds(self, **kwargs):
+ out = super().extra_conds(**kwargs)
+ cross_attn = kwargs.get("cross_attn", None)
+ if cross_attn is not None:
+ out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
+ ref_latents = kwargs.get("reference_latents", None)
+ if ref_latents is not None:
+ latents = []
+ for lat in ref_latents:
+ latents.append(self.process_latent_in(lat))
+ out['ref_latents'] = comfy.conds.CONDList(latents)
+
+ ref_latents_method = kwargs.get("reference_latents_method", None)
+ if ref_latents_method is not None:
+ out['ref_latents_method'] = comfy.conds.CONDConstant(ref_latents_method)
+ return out
+
+ def extra_conds_shapes(self, **kwargs):
+ out = {}
+ ref_latents = kwargs.get("reference_latents", None)
+ if ref_latents is not None:
+ out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16])
+ return out
+
+class HunyuanImage21(BaseModel):
+ def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
+ super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.hunyuan_video.model.HunyuanVideo)
+
+ def extra_conds(self, **kwargs):
+ out = super().extra_conds(**kwargs)
+ attention_mask = kwargs.get("attention_mask", None)
+ if attention_mask is not None:
+ if torch.numel(attention_mask) != attention_mask.sum():
+ out['attention_mask'] = comfy.conds.CONDRegular(attention_mask)
+ cross_attn = kwargs.get("cross_attn", None)
+ if cross_attn is not None:
+ out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
+
+ conditioning_byt5small = kwargs.get("conditioning_byt5small", None)
+ if conditioning_byt5small is not None:
+ out['txt_byt5'] = comfy.conds.CONDRegular(conditioning_byt5small)
+
+ guidance = kwargs.get("guidance", 6.0)
+ if guidance is not None:
+ out['guidance'] = comfy.conds.CONDRegular(torch.FloatTensor([guidance]))
+
+ return out
+
+class HunyuanImage21Refiner(HunyuanImage21):
+ def concat_cond(self, **kwargs):
+ noise = kwargs.get("noise", None)
+ image = kwargs.get("concat_latent_image", None)
+ noise_augmentation = kwargs.get("noise_augmentation", 0.0)
+ device = kwargs["device"]
+
+ if image is None:
+ shape_image = list(noise.shape)
+ image = torch.zeros(shape_image, dtype=noise.dtype, layout=noise.layout, device=noise.device)
+ else:
+ image = utils.common_upscale(image.to(device), noise.shape[-1], noise.shape[-2], "bilinear", "center")
+ image = self.process_latent_in(image)
+ image = utils.resize_to_batch_size(image, noise.shape[0])
+ if noise_augmentation > 0:
+ generator = torch.Generator(device="cpu")
+ generator.manual_seed(kwargs.get("seed", 0) - 10)
+ noise = torch.randn(image.shape, generator=generator, dtype=image.dtype, device="cpu").to(image.device)
+ image = noise_augmentation * noise + min(1.0 - noise_augmentation, 0.75) * image
+ else:
+ image = 0.75 * image
+ return image
+
+ def extra_conds(self, **kwargs):
+ out = super().extra_conds(**kwargs)
+ out['disable_time_r'] = comfy.conds.CONDConstant(True)
+ return out
diff --git a/comfy/model_detection.py b/comfy/model_detection.py
index 850a3c0ac..4669eb14b 100644
--- a/comfy/model_detection.py
+++ b/comfy/model_detection.py
@@ -6,6 +6,20 @@ import math
import logging
import torch
+
+def detect_layer_quantization(metadata):
+ quant_key = "_quantization_metadata"
+ if metadata is not None and quant_key in metadata:
+ quant_metadata = metadata.pop(quant_key)
+ quant_metadata = json.loads(quant_metadata)
+ if isinstance(quant_metadata, dict) and "layers" in quant_metadata:
+ logging.info(f"Found quantization metadata (version {quant_metadata.get('format_version', 'unknown')})")
+ return quant_metadata["layers"]
+ else:
+ raise ValueError("Invalid quantization metadata format")
+ return None
+
+
def count_blocks(state_dict_keys, prefix_string):
count = 0
while True:
@@ -136,25 +150,45 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
if '{}txt_in.individual_token_refiner.blocks.0.norm1.weight'.format(key_prefix) in state_dict_keys: #Hunyuan Video
dit_config = {}
+ in_w = state_dict['{}img_in.proj.weight'.format(key_prefix)]
+ out_w = state_dict['{}final_layer.linear.weight'.format(key_prefix)]
dit_config["image_model"] = "hunyuan_video"
- dit_config["in_channels"] = state_dict['{}img_in.proj.weight'.format(key_prefix)].shape[1] #SkyReels img2video has 32 input channels
- dit_config["patch_size"] = [1, 2, 2]
- dit_config["out_channels"] = 16
- dit_config["vec_in_dim"] = 768
- dit_config["context_in_dim"] = 4096
- dit_config["hidden_size"] = 3072
+ dit_config["in_channels"] = in_w.shape[1] #SkyReels img2video has 32 input channels
+ dit_config["patch_size"] = list(in_w.shape[2:])
+ dit_config["out_channels"] = out_w.shape[0] // math.prod(dit_config["patch_size"])
+ if any(s.startswith('{}vector_in.'.format(key_prefix)) for s in state_dict_keys):
+ dit_config["vec_in_dim"] = 768
+ else:
+ dit_config["vec_in_dim"] = None
+
+ if len(dit_config["patch_size"]) == 2:
+ dit_config["axes_dim"] = [64, 64]
+ else:
+ dit_config["axes_dim"] = [16, 56, 56]
+
+ if any(s.startswith('{}time_r_in.'.format(key_prefix)) for s in state_dict_keys):
+ dit_config["meanflow"] = True
+ else:
+ dit_config["meanflow"] = False
+
+ dit_config["context_in_dim"] = state_dict['{}txt_in.input_embedder.weight'.format(key_prefix)].shape[1]
+ dit_config["hidden_size"] = in_w.shape[0]
dit_config["mlp_ratio"] = 4.0
- dit_config["num_heads"] = 24
+ dit_config["num_heads"] = in_w.shape[0] // 128
dit_config["depth"] = count_blocks(state_dict_keys, '{}double_blocks.'.format(key_prefix) + '{}.')
dit_config["depth_single_blocks"] = count_blocks(state_dict_keys, '{}single_blocks.'.format(key_prefix) + '{}.')
- dit_config["axes_dim"] = [16, 56, 56]
dit_config["theta"] = 256
dit_config["qkv_bias"] = True
+ if '{}byt5_in.fc1.weight'.format(key_prefix) in state_dict:
+ dit_config["byt5"] = True
+ else:
+ dit_config["byt5"] = False
+
guidance_keys = list(filter(lambda a: a.startswith("{}guidance_in.".format(key_prefix)), state_dict_keys))
dit_config["guidance_embed"] = len(guidance_keys) > 0
return dit_config
- if '{}double_blocks.0.img_attn.norm.key_norm.scale'.format(key_prefix) in state_dict_keys and '{}img_in.weight'.format(key_prefix) in state_dict_keys: #Flux
+ if '{}double_blocks.0.img_attn.norm.key_norm.scale'.format(key_prefix) in state_dict_keys and ('{}img_in.weight'.format(key_prefix) in state_dict_keys or f"{key_prefix}distilled_guidance_layer.norms.0.scale" in state_dict_keys): #Flux, Chroma or Chroma Radiance (has no img_in.weight)
dit_config = {}
dit_config["image_model"] = "flux"
dit_config["in_channels"] = 16
@@ -184,6 +218,18 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
dit_config["out_dim"] = 3072
dit_config["hidden_dim"] = 5120
dit_config["n_layers"] = 5
+ if f"{key_prefix}nerf_blocks.0.norm.scale" in state_dict_keys: #Chroma Radiance
+ dit_config["image_model"] = "chroma_radiance"
+ dit_config["in_channels"] = 3
+ dit_config["out_channels"] = 3
+ dit_config["patch_size"] = 16
+ dit_config["nerf_hidden_size"] = 64
+ dit_config["nerf_mlp_ratio"] = 4
+ dit_config["nerf_depth"] = 4
+ dit_config["nerf_max_freqs"] = 8
+ dit_config["nerf_tile_size"] = 512
+ dit_config["nerf_final_head_type"] = "conv" if f"{key_prefix}nerf_final_layer_conv.norm.scale" in state_dict_keys else "linear"
+ dit_config["nerf_embedder_dtype"] = torch.float32
else:
dit_config["guidance_embed"] = "{}guidance_in.in_layer.weight".format(key_prefix) in state_dict_keys
return dit_config
@@ -333,8 +379,8 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
dit_config["patch_size"] = 2
dit_config["in_channels"] = 16
dit_config["dim"] = 2304
- dit_config["cap_feat_dim"] = 2304
- dit_config["n_layers"] = 26
+ dit_config["cap_feat_dim"] = state_dict['{}cap_embedder.1.weight'.format(key_prefix)].shape[1]
+ dit_config["n_layers"] = count_blocks(state_dict_keys, '{}layers.'.format(key_prefix) + '{}.')
dit_config["n_heads"] = 24
dit_config["n_kv_heads"] = 8
dit_config["qk_norm"] = True
@@ -346,7 +392,9 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
dit_config = {}
dit_config["image_model"] = "wan2.1"
dim = state_dict['{}head.modulation'.format(key_prefix)].shape[-1]
+ out_dim = state_dict['{}head.head.weight'.format(key_prefix)].shape[0] // 4
dit_config["dim"] = dim
+ dit_config["out_dim"] = out_dim
dit_config["num_heads"] = dim // 128
dit_config["ffn_dim"] = state_dict['{}blocks.0.ffn.0.weight'.format(key_prefix)].shape[0]
dit_config["num_layers"] = count_blocks(state_dict_keys, '{}blocks.'.format(key_prefix) + '{}.')
@@ -362,7 +410,16 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
dit_config["vace_in_dim"] = state_dict['{}vace_patch_embedding.weight'.format(key_prefix)].shape[1]
dit_config["vace_layers"] = count_blocks(state_dict_keys, '{}vace_blocks.'.format(key_prefix) + '{}.')
elif '{}control_adapter.conv.weight'.format(key_prefix) in state_dict_keys:
- dit_config["model_type"] = "camera"
+ if '{}img_emb.proj.0.bias'.format(key_prefix) in state_dict_keys:
+ dit_config["model_type"] = "camera"
+ else:
+ dit_config["model_type"] = "camera_2.2"
+ elif '{}casual_audio_encoder.encoder.final_linear.weight'.format(key_prefix) in state_dict_keys:
+ dit_config["model_type"] = "s2v"
+ elif '{}audio_proj.audio_proj_glob_1.layer.bias'.format(key_prefix) in state_dict_keys:
+ dit_config["model_type"] = "humo"
+ elif '{}face_adapter.fuser_blocks.0.k_norm.weight'.format(key_prefix) in state_dict_keys:
+ dit_config["model_type"] = "animate"
else:
if '{}img_emb.proj.0.bias'.format(key_prefix) in state_dict_keys:
dit_config["model_type"] = "i2v"
@@ -371,6 +428,11 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
flf_weight = state_dict.get('{}img_emb.emb_pos'.format(key_prefix))
if flf_weight is not None:
dit_config["flf_pos_embed_token_number"] = flf_weight.shape[1]
+
+ ref_conv_weight = state_dict.get('{}ref_conv.weight'.format(key_prefix))
+ if ref_conv_weight is not None:
+ dit_config["in_dim_ref_conv"] = ref_conv_weight.shape[1]
+
return dit_config
if '{}latent_in.weight'.format(key_prefix) in state_dict_keys: # Hunyuan 3D
@@ -388,6 +450,20 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
dit_config["guidance_embed"] = "{}guidance_in.in_layer.weight".format(key_prefix) in state_dict_keys
return dit_config
+ if f"{key_prefix}t_embedder.mlp.2.weight" in state_dict_keys: # Hunyuan 3D 2.1
+
+ dit_config = {}
+ dit_config["image_model"] = "hunyuan3d2_1"
+ dit_config["in_channels"] = state_dict[f"{key_prefix}x_embedder.weight"].shape[1]
+ dit_config["context_dim"] = 1024
+ dit_config["hidden_size"] = state_dict[f"{key_prefix}x_embedder.weight"].shape[0]
+ dit_config["mlp_ratio"] = 4.0
+ dit_config["num_heads"] = 16
+ dit_config["depth"] = count_blocks(state_dict_keys, f"{key_prefix}blocks.{{}}")
+ dit_config["qkv_bias"] = False
+ dit_config["guidance_cond_proj_dim"] = None#f"{key_prefix}t_embedder.cond_proj.weight" in state_dict_keys
+ return dit_config
+
if '{}caption_projection.0.linear.weight'.format(key_prefix) in state_dict_keys: # HiDream
dit_config = {}
dit_config["image_model"] = "hidream"
@@ -488,6 +564,13 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
dit_config["timestep_scale"] = 1000.0
return dit_config
+ if '{}txt_norm.weight'.format(key_prefix) in state_dict_keys: # Qwen Image
+ dit_config = {}
+ dit_config["image_model"] = "qwen_image"
+ dit_config["in_channels"] = state_dict['{}img_in.weight'.format(key_prefix)].shape[1]
+ dit_config["num_layers"] = count_blocks(state_dict_keys, '{}transformer_blocks.'.format(key_prefix) + '{}.')
+ return dit_config
+
if '{}input_blocks.0.0.weight'.format(key_prefix) not in state_dict_keys:
return None
@@ -641,6 +724,12 @@ def model_config_from_unet(state_dict, unet_key_prefix, use_base_if_no_match=Fal
else:
model_config.optimizations["fp8"] = True
+ # Detect per-layer quantization (mixed precision)
+ layer_quant_config = detect_layer_quantization(metadata)
+ if layer_quant_config:
+ model_config.layer_quant_config = layer_quant_config
+ logging.info(f"Detected mixed precision quantization: {len(layer_quant_config)} layers quantized")
+
return model_config
def unet_prefix_from_state_dict(state_dict):
@@ -874,7 +963,7 @@ def convert_diffusers_mmdit(state_dict, output_prefix=""):
depth_single_blocks = count_blocks(state_dict, 'single_transformer_blocks.{}.')
hidden_size = state_dict["x_embedder.bias"].shape[0]
sd_map = comfy.utils.flux_to_diffusers({"depth": depth, "depth_single_blocks": depth_single_blocks, "hidden_size": hidden_size}, output_prefix=output_prefix)
- elif 'transformer_blocks.0.attn.add_q_proj.weight' in state_dict: #SD3
+ elif 'transformer_blocks.0.attn.add_q_proj.weight' in state_dict and 'pos_embed.proj.weight' in state_dict: #SD3
num_blocks = count_blocks(state_dict, 'transformer_blocks.{}.')
depth = state_dict["pos_embed.proj.weight"].shape[0] // 64
sd_map = comfy.utils.mmdit_to_diffusers({"depth": depth, "num_blocks": num_blocks}, output_prefix=output_prefix)
diff --git a/comfy/model_management.py b/comfy/model_management.py
index 816caf18f..79c0dfdb4 100644
--- a/comfy/model_management.py
+++ b/comfy/model_management.py
@@ -22,6 +22,7 @@ from enum import Enum
from comfy.cli_args import args, PerformanceFeature
import torch
import sys
+import importlib
import platform
import weakref
import gc
@@ -78,7 +79,6 @@ try:
torch_version = torch.version.__version__
temp = torch_version.split(".")
torch_version_numeric = (int(temp[0]), int(temp[1]))
- xpu_available = (torch_version_numeric[0] < 2 or (torch_version_numeric[0] == 2 and torch_version_numeric[1] <= 4)) and torch.xpu.is_available()
except:
pass
@@ -89,6 +89,7 @@ if args.deterministic:
directml_enabled = False
if args.directml is not None:
+ logging.warning("WARNING: torch-directml barely works, is very slow, has not been updated in over 1 year and might be removed soon, please don't use it, there are better options.")
import torch_directml
directml_enabled = True
device_index = args.directml
@@ -101,11 +102,15 @@ if args.directml is not None:
lowvram_available = False #TODO: need to find a way to get free memory in directml before this can be enabled by default.
try:
- import intel_extension_for_pytorch as ipex
- _ = torch.xpu.device_count()
- xpu_available = xpu_available or torch.xpu.is_available()
+ import intel_extension_for_pytorch as ipex # noqa: F401
except:
- xpu_available = xpu_available or (hasattr(torch, "xpu") and torch.xpu.is_available())
+ pass
+
+try:
+ _ = torch.xpu.device_count()
+ xpu_available = torch.xpu.is_available()
+except:
+ xpu_available = False
try:
if torch.backends.mps.is_available():
@@ -128,6 +133,11 @@ try:
except:
mlu_available = False
+try:
+ ixuca_available = hasattr(torch, "corex")
+except:
+ ixuca_available = False
+
if args.cpu:
cpu_state = CPUState.CPU
@@ -151,6 +161,12 @@ def is_mlu():
return True
return False
+def is_ixuca():
+ global ixuca_available
+ if ixuca_available:
+ return True
+ return False
+
def get_torch_device():
global directml_enabled
global cpu_state
@@ -186,8 +202,9 @@ def get_total_memory(dev=None, torch_total_too=False):
elif is_intel_xpu():
stats = torch.xpu.memory_stats(dev)
mem_reserved = stats['reserved_bytes.all.current']
+ mem_total_xpu = torch.xpu.get_device_properties(dev).total_memory
mem_total_torch = mem_reserved
- mem_total = torch.xpu.get_device_properties(dev).total_memory
+ mem_total = mem_total_xpu
elif is_ascend_npu():
stats = torch.npu.memory_stats(dev)
mem_reserved = stats['reserved_bytes.all.current']
@@ -274,6 +291,24 @@ def is_amd():
return True
return False
+def amd_min_version(device=None, min_rdna_version=0):
+ if not is_amd():
+ return False
+
+ if is_device_cpu(device):
+ return False
+
+ arch = torch.cuda.get_device_properties(device).gcnArchName
+ if arch.startswith('gfx') and len(arch) == 7:
+ try:
+ cmp_rdna_version = int(arch[4]) + 2
+ except:
+ cmp_rdna_version = 0
+ if cmp_rdna_version >= min_rdna_version:
+ return True
+
+ return False
+
MIN_WEIGHT_MEMORY_RATIO = 0.4
if is_nvidia():
MIN_WEIGHT_MEMORY_RATIO = 0.0
@@ -288,7 +323,7 @@ try:
if torch_version_numeric[0] >= 2:
if ENABLE_PYTORCH_ATTENTION == False and args.use_split_cross_attention == False and args.use_quad_cross_attention == False:
ENABLE_PYTORCH_ATTENTION = True
- if is_intel_xpu() or is_ascend_npu() or is_mlu():
+ if is_intel_xpu() or is_ascend_npu() or is_mlu() or is_ixuca():
if args.use_split_cross_attention == False and args.use_quad_cross_attention == False:
ENABLE_PYTORCH_ATTENTION = True
except:
@@ -296,21 +331,33 @@ except:
SUPPORT_FP8_OPS = args.supports_fp8_compute
+
+AMD_RDNA2_AND_OLDER_ARCH = ["gfx1030", "gfx1031", "gfx1010", "gfx1011", "gfx1012", "gfx906", "gfx900", "gfx803"]
+
try:
if is_amd():
+ arch = torch.cuda.get_device_properties(get_torch_device()).gcnArchName
+ if not (any((a in arch) for a in AMD_RDNA2_AND_OLDER_ARCH)):
+ torch.backends.cudnn.enabled = False # Seems to improve things a lot on AMD
+ logging.info("Set: torch.backends.cudnn.enabled = False for better AMD performance.")
+
try:
rocm_version = tuple(map(int, str(torch.version.hip).split(".")[:2]))
except:
rocm_version = (6, -1)
- arch = torch.cuda.get_device_properties(get_torch_device()).gcnArchName
+
logging.info("AMD arch: {}".format(arch))
logging.info("ROCm version: {}".format(rocm_version))
if args.use_split_cross_attention == False and args.use_quad_cross_attention == False:
- if torch_version_numeric >= (2, 7): # works on 2.6 but doesn't actually seem to improve much
- if any((a in arch) for a in ["gfx90a", "gfx942", "gfx1100", "gfx1101", "gfx1151"]): # TODO: more arches, TODO: gfx1201 and gfx950
- ENABLE_PYTORCH_ATTENTION = True
+ if importlib.util.find_spec('triton') is not None: # AMD efficient attention implementation depends on triton. TODO: better way of detecting if it's compiled in or not.
+ if torch_version_numeric >= (2, 7): # works on 2.6 but doesn't actually seem to improve much
+ if any((a in arch) for a in ["gfx90a", "gfx942", "gfx1100", "gfx1101", "gfx1151"]): # TODO: more arches, TODO: gfx950
+ ENABLE_PYTORCH_ATTENTION = True
+ if rocm_version >= (7, 0):
+ if any((a in arch) for a in ["gfx1201"]):
+ ENABLE_PYTORCH_ATTENTION = True
if torch_version_numeric >= (2, 7) and rocm_version >= (6, 4):
- if any((a in arch) for a in ["gfx1201", "gfx942", "gfx950"]): # TODO: more arches
+ if any((a in arch) for a in ["gfx1200", "gfx1201", "gfx950"]): # TODO: more arches, "gfx942" gives error on pytorch nightly 2.10 1013 rocm7.0
SUPPORT_FP8_OPS = True
except:
@@ -325,13 +372,16 @@ if ENABLE_PYTORCH_ATTENTION:
PRIORITIZE_FP16 = False # TODO: remove and replace with something that shows exactly which dtype is faster than the other
try:
- if is_nvidia() and PerformanceFeature.Fp16Accumulation in args.fast:
+ if (is_nvidia() or is_amd()) and PerformanceFeature.Fp16Accumulation in args.fast:
torch.backends.cuda.matmul.allow_fp16_accumulation = True
PRIORITIZE_FP16 = True # TODO: limit to cards where it actually boosts performance
logging.info("Enabled fp16 accumulation.")
except:
pass
+if torch.cuda.is_available() and torch.backends.cudnn.is_available() and PerformanceFeature.AutoTune in args.fast:
+ torch.backends.cudnn.benchmark = True
+
try:
if torch_version_numeric >= (2, 5):
torch.backends.cuda.allow_fp16_bf16_reduction_math_sdp(True)
@@ -377,6 +427,8 @@ def get_torch_device_name(device):
except:
allocator_backend = ""
return "{} {} : {}".format(device, torch.cuda.get_device_name(device), allocator_backend)
+ elif device.type == "xpu":
+ return "{} {}".format(device, torch.xpu.get_device_name(device))
else:
return "{}".format(device.type)
elif is_intel_xpu():
@@ -512,6 +564,8 @@ WINDOWS = any(platform.win32_ver())
EXTRA_RESERVED_VRAM = 400 * 1024 * 1024
if WINDOWS:
EXTRA_RESERVED_VRAM = 600 * 1024 * 1024 #Windows is higher because of the shared vram issue
+ if total_vram > (15 * 1024): # more extra reserved vram on 16GB+ cards
+ EXTRA_RESERVED_VRAM += 100 * 1024 * 1024
if args.reserve_vram is not None:
EXTRA_RESERVED_VRAM = args.reserve_vram * 1024 * 1024 * 1024
@@ -571,7 +625,13 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimu
else:
minimum_memory_required = max(inference_memory, minimum_memory_required + extra_reserved_memory())
- models = set(models)
+ models_temp = set()
+ for m in models:
+ models_temp.add(m)
+ for mm in m.model_patches_models():
+ models_temp.add(mm)
+
+ models = models_temp
models_to_load = []
@@ -597,7 +657,9 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimu
if loaded_model.model.is_clone(current_loaded_models[i].model):
to_unload = [i] + to_unload
for i in to_unload:
- current_loaded_models.pop(i).model.detach(unpatch_all=False)
+ model_to_unload = current_loaded_models.pop(i)
+ model_to_unload.model.detach(unpatch_all=False)
+ model_to_unload.model_finalizer.detach()
total_memory_required = {}
for loaded_model in models_to_load:
@@ -875,8 +937,7 @@ def vae_dtype(device=None, allowed_dtypes=[]):
if d == torch.float16 and should_use_fp16(device):
return d
- # NOTE: bfloat16 seems to work on AMD for the VAE but is extremely slow in some cases compared to fp32
- if d == torch.bfloat16 and (not is_amd()) and should_use_bf16(device):
+ if d == torch.bfloat16 and should_use_bf16(device):
return d
return torch.float32
@@ -926,9 +987,11 @@ def pick_weight_dtype(dtype, fallback_dtype, device=None):
return dtype
def device_supports_non_blocking(device):
+ if args.force_non_blocking:
+ return True
if is_device_mps(device):
return False #pytorch bug? mps doesn't support non blocking
- if is_intel_xpu():
+ if is_intel_xpu(): #xpu does support non blocking but it is slower on iGPUs for some reason so disable by default until situation changes
return False
if args.deterministic: #TODO: figure out why deterministic breaks non blocking from gpu to cpu (previews)
return False
@@ -936,12 +999,6 @@ def device_supports_non_blocking(device):
return False
return True
-def device_should_use_non_blocking(device):
- if not device_supports_non_blocking(device):
- return False
- return False
- # return True #TODO: figure out why this causes memory issues on Nvidia and possibly others
-
def force_channels_last():
if args.force_channels_last:
return True
@@ -956,6 +1013,16 @@ if args.async_offload:
NUM_STREAMS = 2
logging.info("Using async weight offloading with {} streams".format(NUM_STREAMS))
+def current_stream(device):
+ if device is None:
+ return None
+ if is_device_cuda(device):
+ return torch.cuda.current_stream()
+ elif is_device_xpu(device):
+ return torch.xpu.current_stream()
+ else:
+ return None
+
stream_counters = {}
def get_offload_stream(device):
stream_counter = stream_counters.get(device, 0)
@@ -964,28 +1031,33 @@ def get_offload_stream(device):
if device in STREAMS:
ss = STREAMS[device]
- s = ss[stream_counter]
+ #Sync the oldest stream in the queue with the current
+ ss[stream_counter].wait_stream(current_stream(device))
stream_counter = (stream_counter + 1) % len(ss)
- if is_device_cuda(device):
- ss[stream_counter].wait_stream(torch.cuda.current_stream())
stream_counters[device] = stream_counter
- return s
+ return ss[stream_counter]
elif is_device_cuda(device):
ss = []
for k in range(NUM_STREAMS):
ss.append(torch.cuda.Stream(device=device, priority=0))
STREAMS[device] = ss
s = ss[stream_counter]
- stream_counter = (stream_counter + 1) % len(ss)
+ stream_counters[device] = stream_counter
+ return s
+ elif is_device_xpu(device):
+ ss = []
+ for k in range(NUM_STREAMS):
+ ss.append(torch.xpu.Stream(device=device, priority=0))
+ STREAMS[device] = ss
+ s = ss[stream_counter]
stream_counters[device] = stream_counter
return s
return None
def sync_stream(device, stream):
- if stream is None:
+ if stream is None or current_stream(device) is None:
return
- if is_device_cuda(device):
- torch.cuda.current_stream().wait_stream(stream)
+ current_stream(device).wait_stream(stream)
def cast_to(weight, dtype=None, device=None, non_blocking=False, copy=False, stream=None):
if device is None or weight.device == device:
@@ -1010,6 +1082,36 @@ def cast_to_device(tensor, device, dtype, copy=False):
non_blocking = device_supports_non_blocking(device)
return cast_to(tensor, dtype=dtype, device=device, non_blocking=non_blocking, copy=copy)
+def pin_memory(tensor):
+ if PerformanceFeature.PinnedMem not in args.fast:
+ return False
+
+ if not is_nvidia():
+ return False
+
+ if not is_device_cpu(tensor.device):
+ return False
+
+ if torch.cuda.cudart().cudaHostRegister(tensor.data_ptr(), tensor.numel() * tensor.element_size(), 1) == 0:
+ return True
+
+ return False
+
+def unpin_memory(tensor):
+ if PerformanceFeature.PinnedMem not in args.fast:
+ return False
+
+ if not is_nvidia():
+ return False
+
+ if not is_device_cpu(tensor.device):
+ return False
+
+ if torch.cuda.cudart().cudaHostUnregister(tensor.data_ptr()) == 0:
+ return True
+
+ return False
+
def sage_attention_enabled():
return args.use_sage_attention
@@ -1027,6 +1129,8 @@ def xformers_enabled():
return False
if is_mlu():
return False
+ if is_ixuca():
+ return False
if directml_enabled:
return False
return XFORMERS_IS_AVAILABLE
@@ -1062,6 +1166,8 @@ def pytorch_attention_flash_attention():
return True
if is_amd():
return True #if you have pytorch attention enabled on AMD it probably supports at least mem efficient attention
+ if is_ixuca():
+ return True
return False
def force_upcast_attention_dtype():
@@ -1092,8 +1198,8 @@ def get_free_memory(dev=None, torch_free_too=False):
stats = torch.xpu.memory_stats(dev)
mem_active = stats['active_bytes.all.current']
mem_reserved = stats['reserved_bytes.all.current']
- mem_free_torch = mem_reserved - mem_active
mem_free_xpu = torch.xpu.get_device_properties(dev).total_memory - mem_reserved
+ mem_free_torch = mem_reserved - mem_active
mem_free_total = mem_free_xpu + mem_free_torch
elif is_ascend_npu():
stats = torch.npu.memory_stats(dev)
@@ -1142,6 +1248,9 @@ def is_device_cpu(device):
def is_device_mps(device):
return is_device_type(device, 'mps')
+def is_device_xpu(device):
+ return is_device_type(device, 'xpu')
+
def is_device_cuda(device):
return is_device_type(device, 'cuda')
@@ -1173,7 +1282,10 @@ def should_use_fp16(device=None, model_params=0, prioritize_performance=True, ma
return False
if is_intel_xpu():
- return True
+ if torch_version_numeric < (2, 3):
+ return True
+ else:
+ return torch.xpu.get_device_properties(device).has_fp16
if is_ascend_npu():
return True
@@ -1181,6 +1293,9 @@ def should_use_fp16(device=None, model_params=0, prioritize_performance=True, ma
if is_mlu():
return True
+ if is_ixuca():
+ return True
+
if torch.version.hip:
return True
@@ -1236,14 +1351,20 @@ def should_use_bf16(device=None, model_params=0, prioritize_performance=True, ma
return False
if is_intel_xpu():
- return True
+ if torch_version_numeric < (2, 3):
+ return True
+ else:
+ return torch.xpu.is_bf16_supported()
if is_ascend_npu():
return True
+ if is_ixuca():
+ return True
+
if is_amd():
arch = torch.cuda.get_device_properties(device).gcnArchName
- if any((a in arch) for a in ["gfx1030", "gfx1031", "gfx1010", "gfx1011", "gfx1012", "gfx906", "gfx900", "gfx803"]): # RDNA2 and older don't support bf16
+ if any((a in arch) for a in AMD_RDNA2_AND_OLDER_ARCH): # RDNA2 and older don't support bf16
if manual_cast:
return True
return False
diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py
index 52e76b5f3..3e8799983 100644
--- a/comfy/model_patcher.py
+++ b/comfy/model_patcher.py
@@ -123,16 +123,30 @@ def move_weight_functions(m, device):
return memory
class LowVramPatch:
- def __init__(self, key, patches):
+ def __init__(self, key, patches, convert_func=None, set_func=None):
self.key = key
self.patches = patches
+ self.convert_func = convert_func
+ self.set_func = set_func
+
def __call__(self, weight):
intermediate_dtype = weight.dtype
+ if self.convert_func is not None:
+ weight = self.convert_func(weight.to(dtype=torch.float32, copy=True), inplace=True)
+
if intermediate_dtype not in [torch.float32, torch.float16, torch.bfloat16]: #intermediate_dtype has to be one that is supported in math ops
intermediate_dtype = torch.float32
- return comfy.float.stochastic_rounding(comfy.lora.calculate_weight(self.patches[self.key], weight.to(intermediate_dtype), self.key, intermediate_dtype=intermediate_dtype), weight.dtype, seed=string_to_seed(self.key))
+ out = comfy.lora.calculate_weight(self.patches[self.key], weight.to(intermediate_dtype), self.key, intermediate_dtype=intermediate_dtype)
+ if self.set_func is None:
+ return comfy.float.stochastic_rounding(out, weight.dtype, seed=string_to_seed(self.key))
+ else:
+ return self.set_func(out, seed=string_to_seed(self.key), return_weight=True)
- return comfy.lora.calculate_weight(self.patches[self.key], weight, self.key, intermediate_dtype=intermediate_dtype)
+ out = comfy.lora.calculate_weight(self.patches[self.key], weight, self.key, intermediate_dtype=intermediate_dtype)
+ if self.set_func is not None:
+ return self.set_func(out, seed=string_to_seed(self.key), return_weight=True).to(dtype=intermediate_dtype)
+ else:
+ return out
def get_key_weight(model, key):
set_func = None
@@ -224,6 +238,7 @@ class ModelPatcher:
self.force_cast_weights = False
self.patches_uuid = uuid.uuid4()
self.parent = None
+ self.pinned = set()
self.attachments: dict[str] = {}
self.additional_models: dict[str, list[ModelPatcher]] = {}
@@ -261,6 +276,9 @@ class ModelPatcher:
self.size = comfy.model_management.module_size(self.model)
return self.size
+ def get_ram_usage(self):
+ return self.model_size()
+
def loaded_size(self):
return self.model.model_loaded_weight_memory
@@ -430,6 +448,25 @@ class ModelPatcher:
def set_model_forward_timestep_embed_patch(self, patch):
self.set_model_patch(patch, "forward_timestep_embed_patch")
+ def set_model_double_block_patch(self, patch):
+ self.set_model_patch(patch, "double_block")
+
+ def set_model_post_input_patch(self, patch):
+ self.set_model_patch(patch, "post_input")
+
+ def set_model_rope_options(self, scale_x, shift_x, scale_y, shift_y, scale_t, shift_t, **kwargs):
+ rope_options = self.model_options["transformer_options"].get("rope_options", {})
+ rope_options["scale_x"] = scale_x
+ rope_options["scale_y"] = scale_y
+ rope_options["scale_t"] = scale_t
+
+ rope_options["shift_x"] = shift_x
+ rope_options["shift_y"] = shift_y
+ rope_options["shift_t"] = shift_t
+
+ self.model_options["transformer_options"]["rope_options"] = rope_options
+
+
def add_object_patch(self, name, obj):
self.object_patches[name] = obj
@@ -486,6 +523,30 @@ class ModelPatcher:
if hasattr(wrap_func, "to"):
self.model_options["model_function_wrapper"] = wrap_func.to(device)
+ def model_patches_models(self):
+ to = self.model_options["transformer_options"]
+ models = []
+ if "patches" in to:
+ patches = to["patches"]
+ for name in patches:
+ patch_list = patches[name]
+ for i in range(len(patch_list)):
+ if hasattr(patch_list[i], "models"):
+ models += patch_list[i].models()
+ if "patches_replace" in to:
+ patches = to["patches_replace"]
+ for name in patches:
+ patch_list = patches[name]
+ for k in patch_list:
+ if hasattr(patch_list[k], "models"):
+ models += patch_list[k].models()
+ if "model_function_wrapper" in self.model_options:
+ wrap_func = self.model_options["model_function_wrapper"]
+ if hasattr(wrap_func, "models"):
+ models += wrap_func.models()
+
+ return models
+
def model_dtype(self):
if hasattr(self.model, "get_dtype"):
return self.model.get_dtype()
@@ -574,6 +635,21 @@ class ModelPatcher:
else:
set_func(out_weight, inplace_update=inplace_update, seed=string_to_seed(key))
+ def pin_weight_to_device(self, key):
+ weight, set_func, convert_func = get_key_weight(self.model, key)
+ if comfy.model_management.pin_memory(weight):
+ self.pinned.add(key)
+
+ def unpin_weight(self, key):
+ if key in self.pinned:
+ weight, set_func, convert_func = get_key_weight(self.model, key)
+ comfy.model_management.unpin_memory(weight)
+ self.pinned.remove(key)
+
+ def unpin_all_weights(self):
+ for key in list(self.pinned):
+ self.unpin_weight(key)
+
def _load_list(self):
loading = []
for n, m in self.model.named_modules():
@@ -595,9 +671,11 @@ class ModelPatcher:
mem_counter = 0
patch_counter = 0
lowvram_counter = 0
+ lowvram_mem_counter = 0
loading = self._load_list()
load_completely = []
+ offloaded = []
loading.sort(reverse=True)
for x in loading:
n = x[1]
@@ -614,6 +692,7 @@ class ModelPatcher:
if mem_counter + module_mem >= lowvram_model_memory:
lowvram_weight = True
lowvram_counter += 1
+ lowvram_mem_counter += module_mem
if hasattr(m, "prev_comfy_cast_weights"): #Already lowvramed
continue
@@ -627,16 +706,19 @@ class ModelPatcher:
if force_patch_weights:
self.patch_weight_to_device(weight_key)
else:
- m.weight_function = [LowVramPatch(weight_key, self.patches)]
+ _, set_func, convert_func = get_key_weight(self.model, weight_key)
+ m.weight_function = [LowVramPatch(weight_key, self.patches, convert_func, set_func)]
patch_counter += 1
if bias_key in self.patches:
if force_patch_weights:
self.patch_weight_to_device(bias_key)
else:
- m.bias_function = [LowVramPatch(bias_key, self.patches)]
+ _, set_func, convert_func = get_key_weight(self.model, bias_key)
+ m.bias_function = [LowVramPatch(bias_key, self.patches, convert_func, set_func)]
patch_counter += 1
cast_weight = True
+ offloaded.append((module_mem, n, m, params))
else:
if hasattr(m, "comfy_cast_weights"):
wipe_lowvram_weight(m)
@@ -667,7 +749,9 @@ class ModelPatcher:
continue
for param in params:
- self.patch_weight_to_device("{}.{}".format(n, param), device_to=device_to)
+ key = "{}.{}".format(n, param)
+ self.unpin_weight(key)
+ self.patch_weight_to_device(key, device_to=device_to)
logging.debug("lowvram: loaded module regularly {} {}".format(n, m))
m.comfy_patched_weights = True
@@ -675,11 +759,17 @@ class ModelPatcher:
for x in load_completely:
x[2].to(device_to)
+ for x in offloaded:
+ n = x[1]
+ params = x[3]
+ for param in params:
+ self.pin_weight_to_device("{}.{}".format(n, param))
+
if lowvram_counter > 0:
- logging.info("loaded partially {} {} {}".format(lowvram_model_memory / (1024 * 1024), mem_counter / (1024 * 1024), patch_counter))
+ logging.info("loaded partially; {:.2f} MB usable, {:.2f} MB loaded, {:.2f} MB offloaded, lowvram patches: {}".format(lowvram_model_memory / (1024 * 1024), mem_counter / (1024 * 1024), lowvram_mem_counter / (1024 * 1024), patch_counter))
self.model.model_lowvram = True
else:
- logging.info("loaded completely {} {} {}".format(lowvram_model_memory / (1024 * 1024), mem_counter / (1024 * 1024), full_load))
+ logging.info("loaded completely; {:.2f} MB usable, {:.2f} MB loaded, full load: {}".format(lowvram_model_memory / (1024 * 1024), mem_counter / (1024 * 1024), full_load))
self.model.model_lowvram = False
if full_load:
self.model.to(device_to)
@@ -716,6 +806,7 @@ class ModelPatcher:
self.eject_model()
if unpatch_weights:
self.unpatch_hooks()
+ self.unpin_all_weights()
if self.model.model_lowvram:
for m in self.model.modules():
move_weight_functions(m, device_to)
@@ -795,10 +886,12 @@ class ModelPatcher:
module_mem += move_weight_functions(m, device_to)
if lowvram_possible:
if weight_key in self.patches:
- m.weight_function.append(LowVramPatch(weight_key, self.patches))
+ _, set_func, convert_func = get_key_weight(self.model, weight_key)
+ m.weight_function.append(LowVramPatch(weight_key, self.patches, convert_func, set_func))
patch_counter += 1
if bias_key in self.patches:
- m.bias_function.append(LowVramPatch(bias_key, self.patches))
+ _, set_func, convert_func = get_key_weight(self.model, bias_key)
+ m.bias_function.append(LowVramPatch(bias_key, self.patches, convert_func, set_func))
patch_counter += 1
cast_weight = True
@@ -809,6 +902,9 @@ class ModelPatcher:
memory_freed += module_mem
logging.debug("freed {}".format(n))
+ for param in params:
+ self.pin_weight_to_device("{}.{}".format(n, param))
+
self.model.model_lowvram = True
self.model.lowvram_patch_counter += patch_counter
self.model.model_loaded_weight_memory -= memory_freed
@@ -1211,5 +1307,6 @@ class ModelPatcher:
self.clear_cached_hook_weights()
def __del__(self):
+ self.unpin_all_weights()
self.detach(unpatch_all=False)
diff --git a/comfy/model_sampling.py b/comfy/model_sampling.py
index b240b7f29..2a00ed819 100644
--- a/comfy/model_sampling.py
+++ b/comfy/model_sampling.py
@@ -21,17 +21,23 @@ def rescale_zero_terminal_snr_sigmas(sigmas):
alphas_bar[-1] = 4.8973451890853435e-08
return ((1 - alphas_bar) / alphas_bar) ** 0.5
+def reshape_sigma(sigma, noise_dim):
+ if sigma.nelement() == 1:
+ return sigma.view(())
+ else:
+ return sigma.view(sigma.shape[:1] + (1,) * (noise_dim - 1))
+
class EPS:
def calculate_input(self, sigma, noise):
- sigma = sigma.view(sigma.shape[:1] + (1,) * (noise.ndim - 1))
+ sigma = reshape_sigma(sigma, noise.ndim)
return noise / (sigma ** 2 + self.sigma_data ** 2) ** 0.5
def calculate_denoised(self, sigma, model_output, model_input):
- sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1))
+ sigma = reshape_sigma(sigma, model_output.ndim)
return model_input - model_output * sigma
def noise_scaling(self, sigma, noise, latent_image, max_denoise=False):
- sigma = sigma.view(sigma.shape[:1] + (1,) * (noise.ndim - 1))
+ sigma = reshape_sigma(sigma, noise.ndim)
if max_denoise:
noise = noise * torch.sqrt(1.0 + sigma ** 2.0)
else:
@@ -45,12 +51,12 @@ class EPS:
class V_PREDICTION(EPS):
def calculate_denoised(self, sigma, model_output, model_input):
- sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1))
+ sigma = reshape_sigma(sigma, model_output.ndim)
return model_input * self.sigma_data ** 2 / (sigma ** 2 + self.sigma_data ** 2) - model_output * sigma * self.sigma_data / (sigma ** 2 + self.sigma_data ** 2) ** 0.5
class EDM(V_PREDICTION):
def calculate_denoised(self, sigma, model_output, model_input):
- sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1))
+ sigma = reshape_sigma(sigma, model_output.ndim)
return model_input * self.sigma_data ** 2 / (sigma ** 2 + self.sigma_data ** 2) + model_output * sigma * self.sigma_data / (sigma ** 2 + self.sigma_data ** 2) ** 0.5
class CONST:
@@ -58,15 +64,15 @@ class CONST:
return noise
def calculate_denoised(self, sigma, model_output, model_input):
- sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1))
+ sigma = reshape_sigma(sigma, model_output.ndim)
return model_input - model_output * sigma
def noise_scaling(self, sigma, noise, latent_image, max_denoise=False):
- sigma = sigma.view(sigma.shape[:1] + (1,) * (noise.ndim - 1))
+ sigma = reshape_sigma(sigma, noise.ndim)
return sigma * noise + (1.0 - sigma) * latent_image
def inverse_noise_scaling(self, sigma, latent):
- sigma = sigma.view(sigma.shape[:1] + (1,) * (latent.ndim - 1))
+ sigma = reshape_sigma(sigma, latent.ndim)
return latent / (1.0 - sigma)
class X0(EPS):
@@ -80,16 +86,16 @@ class IMG_TO_IMG(X0):
class COSMOS_RFLOW:
def calculate_input(self, sigma, noise):
sigma = (sigma / (sigma + 1))
- sigma = sigma.view(sigma.shape[:1] + (1,) * (noise.ndim - 1))
+ sigma = reshape_sigma(sigma, noise.ndim)
return noise * (1.0 - sigma)
def calculate_denoised(self, sigma, model_output, model_input):
sigma = (sigma / (sigma + 1))
- sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1))
+ sigma = reshape_sigma(sigma, model_output.ndim)
return model_input * (1.0 - sigma) - model_output * sigma
def noise_scaling(self, sigma, noise, latent_image, max_denoise=False):
- sigma = sigma.view(sigma.shape[:1] + (1,) * (noise.ndim - 1))
+ sigma = reshape_sigma(sigma, noise.ndim)
noise = noise * sigma
noise += latent_image
return noise
diff --git a/comfy/nested_tensor.py b/comfy/nested_tensor.py
new file mode 100644
index 000000000..b700816fa
--- /dev/null
+++ b/comfy/nested_tensor.py
@@ -0,0 +1,91 @@
+import torch
+
+class NestedTensor:
+ def __init__(self, tensors):
+ self.tensors = list(tensors)
+ self.is_nested = True
+
+ def _copy(self):
+ return NestedTensor(self.tensors)
+
+ def apply_operation(self, other, operation):
+ o = self._copy()
+ if isinstance(other, NestedTensor):
+ for i, t in enumerate(o.tensors):
+ o.tensors[i] = operation(t, other.tensors[i])
+ else:
+ for i, t in enumerate(o.tensors):
+ o.tensors[i] = operation(t, other)
+ return o
+
+ def __add__(self, b):
+ return self.apply_operation(b, lambda x, y: x + y)
+
+ def __sub__(self, b):
+ return self.apply_operation(b, lambda x, y: x - y)
+
+ def __mul__(self, b):
+ return self.apply_operation(b, lambda x, y: x * y)
+
+ # def __itruediv__(self, b):
+ # return self.apply_operation(b, lambda x, y: x / y)
+
+ def __truediv__(self, b):
+ return self.apply_operation(b, lambda x, y: x / y)
+
+ def __getitem__(self, *args, **kwargs):
+ return self.apply_operation(None, lambda x, y: x.__getitem__(*args, **kwargs))
+
+ def unbind(self):
+ return self.tensors
+
+ def to(self, *args, **kwargs):
+ o = self._copy()
+ for i, t in enumerate(o.tensors):
+ o.tensors[i] = t.to(*args, **kwargs)
+ return o
+
+ def new_ones(self, *args, **kwargs):
+ return self.tensors[0].new_ones(*args, **kwargs)
+
+ def float(self):
+ return self.to(dtype=torch.float)
+
+ def chunk(self, *args, **kwargs):
+ return self.apply_operation(None, lambda x, y: x.chunk(*args, **kwargs))
+
+ def size(self):
+ return self.tensors[0].size()
+
+ @property
+ def shape(self):
+ return self.tensors[0].shape
+
+ @property
+ def ndim(self):
+ dims = 0
+ for t in self.tensors:
+ dims = max(t.ndim, dims)
+ return dims
+
+ @property
+ def device(self):
+ return self.tensors[0].device
+
+ @property
+ def dtype(self):
+ return self.tensors[0].dtype
+
+ @property
+ def layout(self):
+ return self.tensors[0].layout
+
+
+def cat_nested(tensors, *args, **kwargs):
+ cated_tensors = []
+ for i in range(len(tensors[0].tensors)):
+ tens = []
+ for j in range(len(tensors)):
+ tens.append(tensors[j].tensors[i])
+ cated_tensors.append(torch.cat(tens, *args, **kwargs))
+ return NestedTensor(cated_tensors)
diff --git a/comfy/ops.py b/comfy/ops.py
index 2cc9bbc27..18f6b804b 100644
--- a/comfy/ops.py
+++ b/comfy/ops.py
@@ -24,12 +24,58 @@ import comfy.float
import comfy.rmsnorm
import contextlib
+def run_every_op():
+ if torch.compiler.is_compiling():
+ return
+
+ comfy.model_management.throw_exception_if_processing_interrupted()
+
+def scaled_dot_product_attention(q, k, v, *args, **kwargs):
+ return torch.nn.functional.scaled_dot_product_attention(q, k, v, *args, **kwargs)
+
+
+try:
+ if torch.cuda.is_available():
+ from torch.nn.attention import SDPBackend, sdpa_kernel
+ import inspect
+ if "set_priority" in inspect.signature(sdpa_kernel).parameters:
+ SDPA_BACKEND_PRIORITY = [
+ SDPBackend.FLASH_ATTENTION,
+ SDPBackend.EFFICIENT_ATTENTION,
+ SDPBackend.MATH,
+ ]
+
+ SDPA_BACKEND_PRIORITY.insert(0, SDPBackend.CUDNN_ATTENTION)
+
+ def scaled_dot_product_attention(q, k, v, *args, **kwargs):
+ with sdpa_kernel(SDPA_BACKEND_PRIORITY, set_priority=True):
+ return torch.nn.functional.scaled_dot_product_attention(q, k, v, *args, **kwargs)
+ else:
+ logging.warning("Torch version too old to set sdpa backend priority.")
+except (ModuleNotFoundError, TypeError):
+ logging.warning("Could not set sdpa backend priority.")
+
+NVIDIA_MEMORY_CONV_BUG_WORKAROUND = False
+try:
+ if comfy.model_management.is_nvidia():
+ if torch.backends.cudnn.version() >= 91002 and comfy.model_management.torch_version_numeric >= (2, 9) and comfy.model_management.torch_version_numeric <= (2, 10):
+ #TODO: change upper bound version once it's fixed'
+ NVIDIA_MEMORY_CONV_BUG_WORKAROUND = True
+ logging.info("working around nvidia conv3d memory bug.")
+except:
+ pass
+
cast_to = comfy.model_management.cast_to #TODO: remove once no more references
def cast_to_input(weight, input, non_blocking=False, copy=True):
return comfy.model_management.cast_to(weight, input.dtype, input.device, non_blocking=non_blocking, copy=copy)
-def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None):
+
+@torch.compiler.disable()
+def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None, offloadable=False):
+ # NOTE: offloadable=False is a a legacy and if you are a custom node author reading this please pass
+ # offloadable=True and call uncast_bias_weight() after your last usage of the weight/bias. This
+ # will add async-offload support to your cast and improve performance.
if input is not None:
if dtype is None:
dtype = input.dtype
@@ -38,7 +84,11 @@ def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None):
if device is None:
device = input.device
- offload_stream = comfy.model_management.get_offload_stream(device)
+ if offloadable:
+ offload_stream = comfy.model_management.get_offload_stream(device)
+ else:
+ offload_stream = None
+
if offload_stream is not None:
wf_context = offload_stream
else:
@@ -63,7 +113,24 @@ def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None):
weight = f(weight)
comfy.model_management.sync_stream(device, offload_stream)
- return weight, bias
+ if offloadable:
+ return weight, bias, offload_stream
+ else:
+ #Legacy function signature
+ return weight, bias
+
+
+def uncast_bias_weight(s, weight, bias, offload_stream):
+ if offload_stream is None:
+ return
+ if weight is not None:
+ device = weight.device
+ else:
+ if bias is None:
+ return
+ device = bias.device
+ offload_stream.wait_stream(comfy.model_management.current_stream(device))
+
class CastWeightBiasOp:
comfy_cast_weights = False
@@ -76,10 +143,13 @@ class disable_weight_init:
return None
def forward_comfy_cast_weights(self, input):
- weight, bias = cast_bias_weight(self, input)
- return torch.nn.functional.linear(input, weight, bias)
+ weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
+ x = torch.nn.functional.linear(input, weight, bias)
+ uncast_bias_weight(self, weight, bias, offload_stream)
+ return x
def forward(self, *args, **kwargs):
+ run_every_op()
if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:
return self.forward_comfy_cast_weights(*args, **kwargs)
else:
@@ -90,10 +160,13 @@ class disable_weight_init:
return None
def forward_comfy_cast_weights(self, input):
- weight, bias = cast_bias_weight(self, input)
- return self._conv_forward(input, weight, bias)
+ weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
+ x = self._conv_forward(input, weight, bias)
+ uncast_bias_weight(self, weight, bias, offload_stream)
+ return x
def forward(self, *args, **kwargs):
+ run_every_op()
if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:
return self.forward_comfy_cast_weights(*args, **kwargs)
else:
@@ -104,10 +177,13 @@ class disable_weight_init:
return None
def forward_comfy_cast_weights(self, input):
- weight, bias = cast_bias_weight(self, input)
- return self._conv_forward(input, weight, bias)
+ weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
+ x = self._conv_forward(input, weight, bias)
+ uncast_bias_weight(self, weight, bias, offload_stream)
+ return x
def forward(self, *args, **kwargs):
+ run_every_op()
if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:
return self.forward_comfy_cast_weights(*args, **kwargs)
else:
@@ -117,11 +193,23 @@ class disable_weight_init:
def reset_parameters(self):
return None
+ def _conv_forward(self, input, weight, bias, *args, **kwargs):
+ if NVIDIA_MEMORY_CONV_BUG_WORKAROUND and weight.dtype in (torch.float16, torch.bfloat16):
+ out = torch.cudnn_convolution(input, weight, self.padding, self.stride, self.dilation, self.groups, benchmark=False, deterministic=False, allow_tf32=True)
+ if bias is not None:
+ out += bias.reshape((1, -1) + (1,) * (out.ndim - 2))
+ return out
+ else:
+ return super()._conv_forward(input, weight, bias, *args, **kwargs)
+
def forward_comfy_cast_weights(self, input):
- weight, bias = cast_bias_weight(self, input)
- return self._conv_forward(input, weight, bias)
+ weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
+ x = self._conv_forward(input, weight, bias)
+ uncast_bias_weight(self, weight, bias, offload_stream)
+ return x
def forward(self, *args, **kwargs):
+ run_every_op()
if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:
return self.forward_comfy_cast_weights(*args, **kwargs)
else:
@@ -132,10 +220,13 @@ class disable_weight_init:
return None
def forward_comfy_cast_weights(self, input):
- weight, bias = cast_bias_weight(self, input)
- return torch.nn.functional.group_norm(input, self.num_groups, weight, bias, self.eps)
+ weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
+ x = torch.nn.functional.group_norm(input, self.num_groups, weight, bias, self.eps)
+ uncast_bias_weight(self, weight, bias, offload_stream)
+ return x
def forward(self, *args, **kwargs):
+ run_every_op()
if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:
return self.forward_comfy_cast_weights(*args, **kwargs)
else:
@@ -147,13 +238,17 @@ class disable_weight_init:
def forward_comfy_cast_weights(self, input):
if self.weight is not None:
- weight, bias = cast_bias_weight(self, input)
+ weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
else:
weight = None
bias = None
- return torch.nn.functional.layer_norm(input, self.normalized_shape, weight, bias, self.eps)
+ offload_stream = None
+ x = torch.nn.functional.layer_norm(input, self.normalized_shape, weight, bias, self.eps)
+ uncast_bias_weight(self, weight, bias, offload_stream)
+ return x
def forward(self, *args, **kwargs):
+ run_every_op()
if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:
return self.forward_comfy_cast_weights(*args, **kwargs)
else:
@@ -166,13 +261,18 @@ class disable_weight_init:
def forward_comfy_cast_weights(self, input):
if self.weight is not None:
- weight, bias = cast_bias_weight(self, input)
+ weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
else:
weight = None
- return comfy.rmsnorm.rms_norm(input, weight, self.eps) # TODO: switch to commented out line when old torch is deprecated
- # return torch.nn.functional.rms_norm(input, self.normalized_shape, weight, self.eps)
+ bias = None
+ offload_stream = None
+ x = comfy.rmsnorm.rms_norm(input, weight, self.eps) # TODO: switch to commented out line when old torch is deprecated
+ # x = torch.nn.functional.rms_norm(input, self.normalized_shape, weight, self.eps)
+ uncast_bias_weight(self, weight, bias, offload_stream)
+ return x
def forward(self, *args, **kwargs):
+ run_every_op()
if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:
return self.forward_comfy_cast_weights(*args, **kwargs)
else:
@@ -188,12 +288,15 @@ class disable_weight_init:
input, output_size, self.stride, self.padding, self.kernel_size,
num_spatial_dims, self.dilation)
- weight, bias = cast_bias_weight(self, input)
- return torch.nn.functional.conv_transpose2d(
+ weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
+ x = torch.nn.functional.conv_transpose2d(
input, weight, bias, self.stride, self.padding,
output_padding, self.groups, self.dilation)
+ uncast_bias_weight(self, weight, bias, offload_stream)
+ return x
def forward(self, *args, **kwargs):
+ run_every_op()
if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:
return self.forward_comfy_cast_weights(*args, **kwargs)
else:
@@ -209,12 +312,15 @@ class disable_weight_init:
input, output_size, self.stride, self.padding, self.kernel_size,
num_spatial_dims, self.dilation)
- weight, bias = cast_bias_weight(self, input)
- return torch.nn.functional.conv_transpose1d(
+ weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
+ x = torch.nn.functional.conv_transpose1d(
input, weight, bias, self.stride, self.padding,
output_padding, self.groups, self.dilation)
+ uncast_bias_weight(self, weight, bias, offload_stream)
+ return x
def forward(self, *args, **kwargs):
+ run_every_op()
if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:
return self.forward_comfy_cast_weights(*args, **kwargs)
else:
@@ -229,10 +335,14 @@ class disable_weight_init:
output_dtype = out_dtype
if self.weight.dtype == torch.float16 or self.weight.dtype == torch.bfloat16:
out_dtype = None
- weight, bias = cast_bias_weight(self, device=input.device, dtype=out_dtype)
- return torch.nn.functional.embedding(input, weight, self.padding_idx, self.max_norm, self.norm_type, self.scale_grad_by_freq, self.sparse).to(dtype=output_dtype)
+ weight, bias, offload_stream = cast_bias_weight(self, device=input.device, dtype=out_dtype, offloadable=True)
+ x = torch.nn.functional.embedding(input, weight, self.padding_idx, self.max_norm, self.norm_type, self.scale_grad_by_freq, self.sparse).to(dtype=output_dtype)
+ uncast_bias_weight(self, weight, bias, offload_stream)
+ return x
+
def forward(self, *args, **kwargs):
+ run_every_op()
if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:
return self.forward_comfy_cast_weights(*args, **kwargs)
else:
@@ -283,6 +393,10 @@ class manual_cast(disable_weight_init):
def fp8_linear(self, input):
+ """
+ Legacy FP8 linear function for backward compatibility.
+ Uses QuantizedTensor subclass for dispatch.
+ """
dtype = self.weight.dtype
if dtype not in [torch.float8_e4m3fn]:
return None
@@ -294,9 +408,9 @@ def fp8_linear(self, input):
input_shape = input.shape
input_dtype = input.dtype
+
if len(input.shape) == 3:
- w, bias = cast_bias_weight(self, input, dtype=dtype, bias_dtype=input_dtype)
- w = w.t()
+ w, bias, offload_stream = cast_bias_weight(self, input, dtype=dtype, bias_dtype=input_dtype, offloadable=True)
scale_weight = self.scale_weight
scale_input = self.scale_input
@@ -309,21 +423,22 @@ def fp8_linear(self, input):
scale_input = torch.ones((), device=input.device, dtype=torch.float32)
input = torch.clamp(input, min=-448, max=448, out=input)
input = input.reshape(-1, input_shape[2]).to(dtype).contiguous()
+ layout_params_weight = {'scale': scale_input, 'orig_dtype': input_dtype}
+ quantized_input = QuantizedTensor(input.reshape(-1, input_shape[2]).to(dtype).contiguous(), TensorCoreFP8Layout, layout_params_weight)
else:
scale_input = scale_input.to(input.device)
- input = (input * (1.0 / scale_input).to(input_dtype)).reshape(-1, input_shape[2]).to(dtype).contiguous()
+ quantized_input = QuantizedTensor.from_float(input.reshape(-1, input_shape[2]), TensorCoreFP8Layout, scale=scale_input, dtype=dtype)
- if bias is not None:
- o = torch._scaled_mm(input, w, out_dtype=input_dtype, bias=bias, scale_a=scale_input, scale_b=scale_weight)
- else:
- o = torch._scaled_mm(input, w, out_dtype=input_dtype, scale_a=scale_input, scale_b=scale_weight)
+ # Wrap weight in QuantizedTensor - this enables unified dispatch
+ # Call F.linear - __torch_dispatch__ routes to fp8_linear handler in quant_ops.py!
+ layout_params_weight = {'scale': scale_weight, 'orig_dtype': input_dtype}
+ quantized_weight = QuantizedTensor(w, TensorCoreFP8Layout, layout_params_weight)
+ o = torch.nn.functional.linear(quantized_input, quantized_weight, bias)
- if isinstance(o, tuple):
- o = o[0]
+ uncast_bias_weight(self, w, bias, offload_stream)
if tensor_2d:
return o.reshape(input_shape[0], -1)
-
return o.reshape((-1, input_shape[1], self.weight.shape[0]))
return None
@@ -336,15 +451,18 @@ class fp8_ops(manual_cast):
return None
def forward_comfy_cast_weights(self, input):
- try:
- out = fp8_linear(self, input)
- if out is not None:
- return out
- except Exception as e:
- logging.info("Exception during fp8 op: {}".format(e))
+ if not self.training:
+ try:
+ out = fp8_linear(self, input)
+ if out is not None:
+ return out
+ except Exception as e:
+ logging.info("Exception during fp8 op: {}".format(e))
- weight, bias = cast_bias_weight(self, input)
- return torch.nn.functional.linear(input, weight, bias)
+ weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
+ x = torch.nn.functional.linear(input, weight, bias)
+ uncast_bias_weight(self, weight, bias, offload_stream)
+ return x
def scaled_fp8_ops(fp8_matrix_mult=False, scale_input=False, override_dtype=None):
logging.info("Using scaled fp8: fp8 matrix mult: {}, scale input: {}".format(fp8_matrix_mult, scale_input))
@@ -372,12 +490,14 @@ def scaled_fp8_ops(fp8_matrix_mult=False, scale_input=False, override_dtype=None
if out is not None:
return out
- weight, bias = cast_bias_weight(self, input)
+ weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
if weight.numel() < input.numel(): #TODO: optimize
- return torch.nn.functional.linear(input, weight * self.scale_weight.to(device=weight.device, dtype=weight.dtype), bias)
+ x = torch.nn.functional.linear(input, weight * self.scale_weight.to(device=weight.device, dtype=weight.dtype), bias)
else:
- return torch.nn.functional.linear(input * self.scale_weight.to(device=weight.device, dtype=weight.dtype), weight, bias)
+ x = torch.nn.functional.linear(input * self.scale_weight.to(device=weight.device, dtype=weight.dtype), weight, bias)
+ uncast_bias_weight(self, weight, bias, offload_stream)
+ return x
def convert_weight(self, weight, inplace=False, **kwargs):
if inplace:
@@ -386,8 +506,10 @@ def scaled_fp8_ops(fp8_matrix_mult=False, scale_input=False, override_dtype=None
else:
return weight * self.scale_weight.to(device=weight.device, dtype=weight.dtype)
- def set_weight(self, weight, inplace_update=False, seed=None, **kwargs):
+ def set_weight(self, weight, inplace_update=False, seed=None, return_weight=False, **kwargs):
weight = comfy.float.stochastic_rounding(weight / self.scale_weight.to(device=weight.device, dtype=weight.dtype), self.weight.dtype, seed=seed)
+ if return_weight:
+ return weight
if inplace_update:
self.weight.data.copy_(weight)
else:
@@ -414,7 +536,130 @@ if CUBLAS_IS_AVAILABLE:
def forward(self, *args, **kwargs):
return super().forward(*args, **kwargs)
-def pick_operations(weight_dtype, compute_dtype, load_device=None, disable_fast_fp8=False, fp8_optimizations=False, scaled_fp8=None):
+
+# ==============================================================================
+# Mixed Precision Operations
+# ==============================================================================
+from .quant_ops import QuantizedTensor, TensorCoreFP8Layout
+
+QUANT_FORMAT_MIXINS = {
+ "float8_e4m3fn": {
+ "dtype": torch.float8_e4m3fn,
+ "layout_type": TensorCoreFP8Layout,
+ "parameters": {
+ "weight_scale": torch.nn.Parameter(torch.zeros((), dtype=torch.float32), requires_grad=False),
+ "input_scale": torch.nn.Parameter(torch.zeros((), dtype=torch.float32), requires_grad=False),
+ }
+ }
+}
+
+class MixedPrecisionOps(disable_weight_init):
+ _layer_quant_config = {}
+ _compute_dtype = torch.bfloat16
+
+ class Linear(torch.nn.Module, CastWeightBiasOp):
+ def __init__(
+ self,
+ in_features: int,
+ out_features: int,
+ bias: bool = True,
+ device=None,
+ dtype=None,
+ ) -> None:
+ super().__init__()
+
+ self.factory_kwargs = {"device": device, "dtype": MixedPrecisionOps._compute_dtype}
+ # self.factory_kwargs = {"device": device, "dtype": dtype}
+
+ self.in_features = in_features
+ self.out_features = out_features
+ if bias:
+ self.bias = torch.nn.Parameter(torch.empty(out_features, **self.factory_kwargs))
+ else:
+ self.register_parameter("bias", None)
+
+ self.tensor_class = None
+
+ def reset_parameters(self):
+ return None
+
+ def _load_from_state_dict(self, state_dict, prefix, local_metadata,
+ strict, missing_keys, unexpected_keys, error_msgs):
+
+ device = self.factory_kwargs["device"]
+ layer_name = prefix.rstrip('.')
+ weight_key = f"{prefix}weight"
+ weight = state_dict.pop(weight_key, None)
+ if weight is None:
+ raise ValueError(f"Missing weight for layer {layer_name}")
+
+ manually_loaded_keys = [weight_key]
+
+ if layer_name not in MixedPrecisionOps._layer_quant_config:
+ self.weight = torch.nn.Parameter(weight.to(device=device, dtype=MixedPrecisionOps._compute_dtype), requires_grad=False)
+ else:
+ quant_format = MixedPrecisionOps._layer_quant_config[layer_name].get("format", None)
+ if quant_format is None:
+ raise ValueError(f"Unknown quantization format for layer {layer_name}")
+
+ mixin = QUANT_FORMAT_MIXINS[quant_format]
+ self.layout_type = mixin["layout_type"]
+
+ scale_key = f"{prefix}weight_scale"
+ layout_params = {
+ 'scale': state_dict.pop(scale_key, None),
+ 'orig_dtype': MixedPrecisionOps._compute_dtype
+ }
+ if layout_params['scale'] is not None:
+ manually_loaded_keys.append(scale_key)
+
+ self.weight = torch.nn.Parameter(
+ QuantizedTensor(weight.to(device=device, dtype=mixin["dtype"]), self.layout_type, layout_params),
+ requires_grad=False
+ )
+
+ for param_name, param_value in mixin["parameters"].items():
+ param_key = f"{prefix}{param_name}"
+ _v = state_dict.pop(param_key, None)
+ if _v is None:
+ continue
+ setattr(self, param_name, torch.nn.Parameter(_v.to(device=device), requires_grad=False))
+ manually_loaded_keys.append(param_key)
+
+ super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
+
+ for key in manually_loaded_keys:
+ if key in missing_keys:
+ missing_keys.remove(key)
+
+ def _forward(self, input, weight, bias):
+ return torch.nn.functional.linear(input, weight, bias)
+
+ def forward_comfy_cast_weights(self, input):
+ weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
+ x = self._forward(input, weight, bias)
+ uncast_bias_weight(self, weight, bias, offload_stream)
+ return x
+
+ def forward(self, input, *args, **kwargs):
+ run_every_op()
+
+ if self.comfy_cast_weights or len(self.weight_function) > 0 or len(self.bias_function) > 0:
+ return self.forward_comfy_cast_weights(input, *args, **kwargs)
+ if (getattr(self, 'layout_type', None) is not None and
+ getattr(self, 'input_scale', None) is not None and
+ not isinstance(input, QuantizedTensor)):
+ input = QuantizedTensor.from_float(input, self.layout_type, scale=self.input_scale, fp8_dtype=self.weight.dtype)
+ return self._forward(input, self.weight, self.bias)
+
+
+def pick_operations(weight_dtype, compute_dtype, load_device=None, disable_fast_fp8=False, fp8_optimizations=False, scaled_fp8=None, model_config=None):
+ if model_config and hasattr(model_config, 'layer_quant_config') and model_config.layer_quant_config:
+ MixedPrecisionOps._layer_quant_config = model_config.layer_quant_config
+ MixedPrecisionOps._compute_dtype = compute_dtype
+ logging.info(f"Using mixed precision operations: {len(model_config.layer_quant_config)} quantized layers")
+ return MixedPrecisionOps
+
fp8_compute = comfy.model_management.supports_fp8_compute(load_device)
if scaled_fp8 is not None:
return scaled_fp8_ops(fp8_matrix_mult=fp8_compute and fp8_optimizations, scale_input=fp8_optimizations, override_dtype=scaled_fp8)
diff --git a/comfy/patcher_extension.py b/comfy/patcher_extension.py
index 965958f4c..5ee4d5ee5 100644
--- a/comfy/patcher_extension.py
+++ b/comfy/patcher_extension.py
@@ -50,6 +50,7 @@ class WrappersMP:
OUTER_SAMPLE = "outer_sample"
PREPARE_SAMPLING = "prepare_sampling"
SAMPLER_SAMPLE = "sampler_sample"
+ PREDICT_NOISE = "predict_noise"
CALC_COND_BATCH = "calc_cond_batch"
APPLY_MODEL = "apply_model"
DIFFUSION_MODEL = "diffusion_model"
@@ -149,7 +150,7 @@ def merge_nested_dicts(dict1: dict, dict2: dict, copy_dict1=True):
for key, value in dict2.items():
if isinstance(value, dict):
curr_value = merged_dict.setdefault(key, {})
- merged_dict[key] = merge_nested_dicts(value, curr_value)
+ merged_dict[key] = merge_nested_dicts(curr_value, value)
elif isinstance(value, list):
merged_dict.setdefault(key, []).extend(value)
else:
diff --git a/comfy/pixel_space_convert.py b/comfy/pixel_space_convert.py
new file mode 100644
index 000000000..049bbcfb4
--- /dev/null
+++ b/comfy/pixel_space_convert.py
@@ -0,0 +1,16 @@
+import torch
+
+
+# "Fake" VAE that converts from IMAGE B, H, W, C and values on the scale of 0..1
+# to LATENT B, C, H, W and values on the scale of -1..1.
+class PixelspaceConversionVAE(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.pixel_space_vae = torch.nn.Parameter(torch.tensor(1.0))
+
+ def encode(self, pixels: torch.Tensor, *_args, **_kwargs) -> torch.Tensor:
+ return pixels
+
+ def decode(self, samples: torch.Tensor, *_args, **_kwargs) -> torch.Tensor:
+ return samples
+
diff --git a/comfy/quant_ops.py b/comfy/quant_ops.py
new file mode 100644
index 000000000..c822fe53c
--- /dev/null
+++ b/comfy/quant_ops.py
@@ -0,0 +1,438 @@
+import torch
+import logging
+from typing import Tuple, Dict
+
+_LAYOUT_REGISTRY = {}
+_GENERIC_UTILS = {}
+
+
+def register_layout_op(torch_op, layout_type):
+ """
+ Decorator to register a layout-specific operation handler.
+ Args:
+ torch_op: PyTorch operation (e.g., torch.ops.aten.linear.default)
+ layout_type: Layout class (e.g., TensorCoreFP8Layout)
+ Example:
+ @register_layout_op(torch.ops.aten.linear.default, TensorCoreFP8Layout)
+ def fp8_linear(func, args, kwargs):
+ # FP8-specific linear implementation
+ ...
+ """
+ def decorator(handler_func):
+ if torch_op not in _LAYOUT_REGISTRY:
+ _LAYOUT_REGISTRY[torch_op] = {}
+ _LAYOUT_REGISTRY[torch_op][layout_type] = handler_func
+ return handler_func
+ return decorator
+
+
+def register_generic_util(torch_op):
+ """
+ Decorator to register a generic utility that works for all layouts.
+ Args:
+ torch_op: PyTorch operation (e.g., torch.ops.aten.detach.default)
+
+ Example:
+ @register_generic_util(torch.ops.aten.detach.default)
+ def generic_detach(func, args, kwargs):
+ # Works for any layout
+ ...
+ """
+ def decorator(handler_func):
+ _GENERIC_UTILS[torch_op] = handler_func
+ return handler_func
+ return decorator
+
+
+def _get_layout_from_args(args):
+ for arg in args:
+ if isinstance(arg, QuantizedTensor):
+ return arg._layout_type
+ elif isinstance(arg, (list, tuple)):
+ for item in arg:
+ if isinstance(item, QuantizedTensor):
+ return item._layout_type
+ return None
+
+
+def _move_layout_params_to_device(params, device):
+ new_params = {}
+ for k, v in params.items():
+ if isinstance(v, torch.Tensor):
+ new_params[k] = v.to(device=device)
+ else:
+ new_params[k] = v
+ return new_params
+
+
+def _copy_layout_params(params):
+ new_params = {}
+ for k, v in params.items():
+ if isinstance(v, torch.Tensor):
+ new_params[k] = v.clone()
+ else:
+ new_params[k] = v
+ return new_params
+
+
+class QuantizedLayout:
+ """
+ Base class for quantization layouts.
+
+ A layout encapsulates the format-specific logic for quantization/dequantization
+ and provides a uniform interface for extracting raw tensors needed for computation.
+
+ New quantization formats should subclass this and implement the required methods.
+ """
+ @classmethod
+ def quantize(cls, tensor, **kwargs) -> Tuple[torch.Tensor, Dict]:
+ raise NotImplementedError(f"{cls.__name__} must implement quantize()")
+
+ @staticmethod
+ def dequantize(qdata, **layout_params) -> torch.Tensor:
+ raise NotImplementedError("TensorLayout must implement dequantize()")
+
+ @classmethod
+ def get_plain_tensors(cls, qtensor) -> torch.Tensor:
+ raise NotImplementedError(f"{cls.__name__} must implement get_plain_tensors()")
+
+
+class QuantizedTensor(torch.Tensor):
+ """
+ Universal quantized tensor that works with any layout.
+
+ This tensor subclass uses a pluggable layout system to support multiple
+ quantization formats (FP8, INT4, INT8, etc.) without code duplication.
+
+ The layout_type determines format-specific behavior, while common operations
+ (detach, clone, to) are handled generically.
+
+ Attributes:
+ _qdata: The quantized tensor data
+ _layout_type: Layout class (e.g., TensorCoreFP8Layout)
+ _layout_params: Dict with layout-specific params (scale, zero_point, etc.)
+ """
+
+ @staticmethod
+ def __new__(cls, qdata, layout_type, layout_params):
+ """
+ Create a quantized tensor.
+
+ Args:
+ qdata: The quantized data tensor
+ layout_type: Layout class (subclass of QuantizedLayout)
+ layout_params: Dict with layout-specific parameters
+ """
+ return torch.Tensor._make_subclass(cls, qdata, require_grad=False)
+
+ def __init__(self, qdata, layout_type, layout_params):
+ self._qdata = qdata.contiguous()
+ self._layout_type = layout_type
+ self._layout_params = layout_params
+
+ def __repr__(self):
+ layout_name = self._layout_type.__name__
+ param_str = ", ".join(f"{k}={v}" for k, v in list(self._layout_params.items())[:2])
+ return f"QuantizedTensor(shape={self.shape}, layout={layout_name}, {param_str})"
+
+ @property
+ def layout_type(self):
+ return self._layout_type
+
+ def __tensor_flatten__(self):
+ """
+ Tensor flattening protocol for proper device movement.
+ """
+ inner_tensors = ["_qdata"]
+ ctx = {
+ "layout_type": self._layout_type,
+ }
+
+ tensor_params = {}
+ non_tensor_params = {}
+ for k, v in self._layout_params.items():
+ if isinstance(v, torch.Tensor):
+ tensor_params[k] = v
+ else:
+ non_tensor_params[k] = v
+
+ ctx["tensor_param_keys"] = list(tensor_params.keys())
+ ctx["non_tensor_params"] = non_tensor_params
+
+ for k, v in tensor_params.items():
+ attr_name = f"_layout_param_{k}"
+ object.__setattr__(self, attr_name, v)
+ inner_tensors.append(attr_name)
+
+ return inner_tensors, ctx
+
+ @staticmethod
+ def __tensor_unflatten__(inner_tensors, ctx, outer_size, outer_stride):
+ """
+ Tensor unflattening protocol for proper device movement.
+ Reconstructs the QuantizedTensor after device movement.
+ """
+ layout_type = ctx["layout_type"]
+ layout_params = dict(ctx["non_tensor_params"])
+
+ for key in ctx["tensor_param_keys"]:
+ attr_name = f"_layout_param_{key}"
+ layout_params[key] = inner_tensors[attr_name]
+
+ return QuantizedTensor(inner_tensors["_q_data"], layout_type, layout_params)
+
+ @classmethod
+ def from_float(cls, tensor, layout_type, **quantize_kwargs) -> 'QuantizedTensor':
+ qdata, layout_params = layout_type.quantize(tensor, **quantize_kwargs)
+ return cls(qdata, layout_type, layout_params)
+
+ def dequantize(self) -> torch.Tensor:
+ return self._layout_type.dequantize(self._qdata, **self._layout_params)
+
+ @classmethod
+ def __torch_dispatch__(cls, func, types, args=(), kwargs=None):
+ kwargs = kwargs or {}
+
+ # Step 1: Check generic utilities first (detach, clone, to, etc.)
+ if func in _GENERIC_UTILS:
+ return _GENERIC_UTILS[func](func, args, kwargs)
+
+ # Step 2: Check layout-specific handlers (linear, matmul, etc.)
+ layout_type = _get_layout_from_args(args)
+ if layout_type and func in _LAYOUT_REGISTRY:
+ handler = _LAYOUT_REGISTRY[func].get(layout_type)
+ if handler:
+ return handler(func, args, kwargs)
+
+ # Step 3: Fallback to dequantization
+ if isinstance(args[0] if args else None, QuantizedTensor):
+ logging.info(f"QuantizedTensor: Unhandled operation {func}, falling back to dequantization. kwargs={kwargs}")
+ return cls._dequant_and_fallback(func, args, kwargs)
+
+ @classmethod
+ def _dequant_and_fallback(cls, func, args, kwargs):
+ def dequant_arg(arg):
+ if isinstance(arg, QuantizedTensor):
+ return arg.dequantize()
+ elif isinstance(arg, (list, tuple)):
+ return type(arg)(dequant_arg(a) for a in arg)
+ return arg
+
+ new_args = dequant_arg(args)
+ new_kwargs = dequant_arg(kwargs)
+ return func(*new_args, **new_kwargs)
+
+
+# ==============================================================================
+# Generic Utilities (Layout-Agnostic Operations)
+# ==============================================================================
+
+def _create_transformed_qtensor(qt, transform_fn):
+ new_data = transform_fn(qt._qdata)
+ new_params = _copy_layout_params(qt._layout_params)
+ return QuantizedTensor(new_data, qt._layout_type, new_params)
+
+
+def _handle_device_transfer(qt, target_device, target_dtype=None, target_layout=None, op_name="to"):
+ if target_dtype is not None and target_dtype != qt.dtype:
+ logging.warning(
+ f"QuantizedTensor: dtype conversion requested to {target_dtype}, "
+ f"but not supported for quantized tensors. Ignoring dtype."
+ )
+
+ if target_layout is not None and target_layout != torch.strided:
+ logging.warning(
+ f"QuantizedTensor: layout change requested to {target_layout}, "
+ f"but not supported. Ignoring layout."
+ )
+
+ # Handle device transfer
+ current_device = qt._qdata.device
+ if target_device is not None:
+ # Normalize device for comparison
+ if isinstance(target_device, str):
+ target_device = torch.device(target_device)
+ if isinstance(current_device, str):
+ current_device = torch.device(current_device)
+
+ if target_device != current_device:
+ logging.debug(f"QuantizedTensor.{op_name}: Moving from {current_device} to {target_device}")
+ new_q_data = qt._qdata.to(device=target_device)
+ new_params = _move_layout_params_to_device(qt._layout_params, target_device)
+ new_qt = QuantizedTensor(new_q_data, qt._layout_type, new_params)
+ logging.debug(f"QuantizedTensor.{op_name}: Created new tensor on {target_device}")
+ return new_qt
+
+ logging.debug(f"QuantizedTensor.{op_name}: No device change needed, returning original")
+ return qt
+
+
+@register_generic_util(torch.ops.aten.detach.default)
+def generic_detach(func, args, kwargs):
+ """Detach operation - creates a detached copy of the quantized tensor."""
+ qt = args[0]
+ if isinstance(qt, QuantizedTensor):
+ return _create_transformed_qtensor(qt, lambda x: x.detach())
+ return func(*args, **kwargs)
+
+
+@register_generic_util(torch.ops.aten.clone.default)
+def generic_clone(func, args, kwargs):
+ """Clone operation - creates a deep copy of the quantized tensor."""
+ qt = args[0]
+ if isinstance(qt, QuantizedTensor):
+ return _create_transformed_qtensor(qt, lambda x: x.clone())
+ return func(*args, **kwargs)
+
+
+@register_generic_util(torch.ops.aten._to_copy.default)
+def generic_to_copy(func, args, kwargs):
+ """Device/dtype transfer operation - handles .to(device) calls."""
+ qt = args[0]
+ if isinstance(qt, QuantizedTensor):
+ return _handle_device_transfer(
+ qt,
+ target_device=kwargs.get('device', None),
+ target_dtype=kwargs.get('dtype', None),
+ op_name="_to_copy"
+ )
+ return func(*args, **kwargs)
+
+
+@register_generic_util(torch.ops.aten.to.dtype_layout)
+def generic_to_dtype_layout(func, args, kwargs):
+ """Handle .to(device) calls using the dtype_layout variant."""
+ qt = args[0]
+ if isinstance(qt, QuantizedTensor):
+ return _handle_device_transfer(
+ qt,
+ target_device=kwargs.get('device', None),
+ target_dtype=kwargs.get('dtype', None),
+ target_layout=kwargs.get('layout', None),
+ op_name="to"
+ )
+ return func(*args, **kwargs)
+
+
+@register_generic_util(torch.ops.aten.copy_.default)
+def generic_copy_(func, args, kwargs):
+ qt_dest = args[0]
+ src = args[1]
+
+ if isinstance(qt_dest, QuantizedTensor):
+ if isinstance(src, QuantizedTensor):
+ # Copy from another quantized tensor
+ qt_dest._qdata.copy_(src._qdata)
+ qt_dest._layout_type = src._layout_type
+ qt_dest._layout_params = _copy_layout_params(src._layout_params)
+ else:
+ # Copy from regular tensor - just copy raw data
+ qt_dest._qdata.copy_(src)
+ return qt_dest
+ return func(*args, **kwargs)
+
+
+@register_generic_util(torch.ops.aten._has_compatible_shallow_copy_type.default)
+def generic_has_compatible_shallow_copy_type(func, args, kwargs):
+ return True
+
+# ==============================================================================
+# FP8 Layout + Operation Handlers
+# ==============================================================================
+class TensorCoreFP8Layout(QuantizedLayout):
+ """
+ Storage format:
+ - qdata: FP8 tensor (torch.float8_e4m3fn or torch.float8_e5m2)
+ - scale: Scalar tensor (float32) for dequantization
+ - orig_dtype: Original dtype before quantization (for casting back)
+ """
+ @classmethod
+ def quantize(cls, tensor, scale=None, dtype=torch.float8_e4m3fn):
+ orig_dtype = tensor.dtype
+
+ if scale is None:
+ scale = torch.amax(tensor.abs()) / torch.finfo(dtype).max
+
+ if not isinstance(scale, torch.Tensor):
+ scale = torch.tensor(scale)
+ scale = scale.to(device=tensor.device, dtype=torch.float32)
+
+ tensor_scaled = tensor * (1.0 / scale).to(tensor.dtype)
+ # TODO: uncomment this if it's actually needed because the clamp has a small performance penality'
+ # lp_amax = torch.finfo(dtype).max
+ # torch.clamp(tensor_scaled, min=-lp_amax, max=lp_amax, out=tensor_scaled)
+ qdata = tensor_scaled.to(dtype, memory_format=torch.contiguous_format)
+
+ layout_params = {
+ 'scale': scale,
+ 'orig_dtype': orig_dtype
+ }
+ return qdata, layout_params
+
+ @staticmethod
+ def dequantize(qdata, scale, orig_dtype, **kwargs):
+ plain_tensor = torch.ops.aten._to_copy.default(qdata, dtype=orig_dtype)
+ return plain_tensor * scale
+
+ @classmethod
+ def get_plain_tensors(cls, qtensor):
+ return qtensor._qdata, qtensor._layout_params['scale']
+
+
+@register_layout_op(torch.ops.aten.linear.default, TensorCoreFP8Layout)
+def fp8_linear(func, args, kwargs):
+ input_tensor = args[0]
+ weight = args[1]
+ bias = args[2] if len(args) > 2 else None
+
+ if isinstance(input_tensor, QuantizedTensor) and isinstance(weight, QuantizedTensor):
+ plain_input, scale_a = TensorCoreFP8Layout.get_plain_tensors(input_tensor)
+ plain_weight, scale_b = TensorCoreFP8Layout.get_plain_tensors(weight)
+
+ out_dtype = kwargs.get("out_dtype")
+ if out_dtype is None:
+ out_dtype = input_tensor._layout_params['orig_dtype']
+
+ weight_t = plain_weight.t()
+
+ tensor_2d = False
+ if len(plain_input.shape) == 2:
+ tensor_2d = True
+ plain_input = plain_input.unsqueeze(1)
+
+ input_shape = plain_input.shape
+ if len(input_shape) != 3:
+ return None
+
+ try:
+ output = torch._scaled_mm(
+ plain_input.reshape(-1, input_shape[2]),
+ weight_t,
+ bias=bias,
+ scale_a=scale_a,
+ scale_b=scale_b,
+ out_dtype=out_dtype,
+ )
+ if not tensor_2d:
+ output = output.reshape((-1, input_shape[1], weight.shape[0]))
+
+ if output.dtype in [torch.float8_e4m3fn, torch.float8_e5m2]:
+ output_scale = scale_a * scale_b
+ output_params = {
+ 'scale': output_scale,
+ 'orig_dtype': input_tensor._layout_params['orig_dtype']
+ }
+ return QuantizedTensor(output, TensorCoreFP8Layout, output_params)
+ else:
+ return output
+
+ except Exception as e:
+ raise RuntimeError(f"FP8 _scaled_mm failed, falling back to dequantization: {e}")
+
+ # Case 2: DQ Fallback
+ if isinstance(weight, QuantizedTensor):
+ weight = weight.dequantize()
+ if isinstance(input_tensor, QuantizedTensor):
+ input_tensor = input_tensor.dequantize()
+
+ return torch.nn.functional.linear(input_tensor, weight, bias)
diff --git a/comfy/rmsnorm.py b/comfy/rmsnorm.py
index 66ae8321d..555542a46 100644
--- a/comfy/rmsnorm.py
+++ b/comfy/rmsnorm.py
@@ -1,6 +1,7 @@
import torch
import comfy.model_management
import numbers
+import logging
RMSNorm = None
@@ -9,6 +10,7 @@ try:
RMSNorm = torch.nn.RMSNorm
except:
rms_norm_torch = None
+ logging.warning("Please update pytorch to use native RMSNorm")
def rms_norm(x, weight=None, eps=1e-6):
diff --git a/comfy/sample.py b/comfy/sample.py
index be5a7e246..2f8f3a51c 100644
--- a/comfy/sample.py
+++ b/comfy/sample.py
@@ -4,13 +4,9 @@ import comfy.samplers
import comfy.utils
import numpy as np
import logging
+import comfy.nested_tensor
-def prepare_noise(latent_image, seed, noise_inds=None):
- """
- creates random noise given a latent image and a seed.
- optional arg skip can be used to skip and discard x number of noise generations for a given seed
- """
- generator = torch.manual_seed(seed)
+def prepare_noise_inner(latent_image, generator, noise_inds=None):
if noise_inds is None:
return torch.randn(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, generator=generator, device="cpu")
@@ -21,10 +17,29 @@ def prepare_noise(latent_image, seed, noise_inds=None):
if i in unique_inds:
noises.append(noise)
noises = [noises[i] for i in inverse]
- noises = torch.cat(noises, axis=0)
+ return torch.cat(noises, axis=0)
+
+def prepare_noise(latent_image, seed, noise_inds=None):
+ """
+ creates random noise given a latent image and a seed.
+ optional arg skip can be used to skip and discard x number of noise generations for a given seed
+ """
+ generator = torch.manual_seed(seed)
+
+ if latent_image.is_nested:
+ tensors = latent_image.unbind()
+ noises = []
+ for t in tensors:
+ noises.append(prepare_noise_inner(t, generator, noise_inds))
+ noises = comfy.nested_tensor.NestedTensor(noises)
+ else:
+ noises = prepare_noise_inner(latent_image, generator, noise_inds)
+
return noises
def fix_empty_latent_channels(model, latent_image):
+ if latent_image.is_nested:
+ return latent_image
latent_format = model.get_model_object("latent_format") #Resize the empty latent image so it has the right number of channels
if latent_format.latent_channels != latent_image.shape[1] and torch.count_nonzero(latent_image) == 0:
latent_image = comfy.utils.repeat_to_batch_size(latent_image, latent_format.latent_channels, dim=1)
diff --git a/comfy/sampler_helpers.py b/comfy/sampler_helpers.py
index 8dbc41455..e46971afb 100644
--- a/comfy/sampler_helpers.py
+++ b/comfy/sampler_helpers.py
@@ -149,7 +149,7 @@ def cleanup_models(conds, models):
cleanup_additional_models(set(control_cleanup))
-def prepare_model_patcher(model: 'ModelPatcher', conds, model_options: dict):
+def prepare_model_patcher(model: ModelPatcher, conds, model_options: dict):
'''
Registers hooks from conds.
'''
@@ -158,8 +158,8 @@ def prepare_model_patcher(model: 'ModelPatcher', conds, model_options: dict):
for k in conds:
get_hooks_from_cond(conds[k], hooks)
# add wrappers and callbacks from ModelPatcher to transformer_options
- model_options["transformer_options"]["wrappers"] = comfy.patcher_extension.copy_nested_dicts(model.wrappers)
- model_options["transformer_options"]["callbacks"] = comfy.patcher_extension.copy_nested_dicts(model.callbacks)
+ comfy.patcher_extension.merge_nested_dicts(model_options["transformer_options"].setdefault("wrappers", {}), model.wrappers, copy_dict1=False)
+ comfy.patcher_extension.merge_nested_dicts(model_options["transformer_options"].setdefault("callbacks", {}), model.callbacks, copy_dict1=False)
# begin registering hooks
registered = comfy.hooks.HookGroup()
target_dict = comfy.hooks.create_target_dict(comfy.hooks.EnumWeightTarget.Model)
diff --git a/comfy/samplers.py b/comfy/samplers.py
old mode 100644
new mode 100755
index 25ccaf39f..fa4640842
--- a/comfy/samplers.py
+++ b/comfy/samplers.py
@@ -16,6 +16,8 @@ import comfy.sampler_helpers
import comfy.model_patcher
import comfy.patcher_extension
import comfy.hooks
+import comfy.context_windows
+import comfy.utils
import scipy.stats
import numpy
@@ -60,7 +62,7 @@ def get_area_and_mult(conds, x_in, timestep_in):
if "mask_strength" in conds:
mask_strength = conds["mask_strength"]
mask = conds['mask']
- assert (mask.shape[1:] == x_in.shape[2:])
+ # assert (mask.shape[1:] == x_in.shape[2:])
mask = mask[:input_x.shape[0]]
if area is not None:
@@ -68,7 +70,7 @@ def get_area_and_mult(conds, x_in, timestep_in):
mask = mask.narrow(i + 1, area[len(dims) + i], area[i])
mask = mask * mask_strength
- mask = mask.unsqueeze(1).repeat(input_x.shape[0] // mask.shape[0], input_x.shape[1], 1, 1)
+ mask = mask.unsqueeze(1).repeat((input_x.shape[0] // mask.shape[0], input_x.shape[1]) + (1, ) * (mask.ndim - 1))
else:
mask = torch.ones_like(input_x)
mult = mask * strength
@@ -89,7 +91,7 @@ def get_area_and_mult(conds, x_in, timestep_in):
conditioning = {}
model_conds = conds["model_conds"]
for c in model_conds:
- conditioning[c] = model_conds[c].process_cond(batch_size=x_in.shape[0], device=x_in.device, area=area)
+ conditioning[c] = model_conds[c].process_cond(batch_size=x_in.shape[0], area=area)
hooks = conds.get('hooks', None)
control = conds.get('control', None)
@@ -198,14 +200,20 @@ def finalize_default_conds(model: 'BaseModel', hooked_to_run: dict[comfy.hooks.H
hooked_to_run.setdefault(p.hooks, list())
hooked_to_run[p.hooks] += [(p, i)]
-def calc_cond_batch(model: 'BaseModel', conds: list[list[dict]], x_in: torch.Tensor, timestep, model_options):
+def calc_cond_batch(model: BaseModel, conds: list[list[dict]], x_in: torch.Tensor, timestep, model_options: dict[str]):
+ handler: comfy.context_windows.ContextHandlerABC = model_options.get("context_handler", None)
+ if handler is None or not handler.should_use_context(model, conds, x_in, timestep, model_options):
+ return _calc_cond_batch_outer(model, conds, x_in, timestep, model_options)
+ return handler.execute(_calc_cond_batch_outer, model, conds, x_in, timestep, model_options)
+
+def _calc_cond_batch_outer(model: BaseModel, conds: list[list[dict]], x_in: torch.Tensor, timestep, model_options):
executor = comfy.patcher_extension.WrapperExecutor.new_executor(
_calc_cond_batch,
comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.CALC_COND_BATCH, model_options, is_model_options=True)
)
return executor.execute(model, conds, x_in, timestep, model_options)
-def _calc_cond_batch(model: 'BaseModel', conds: list[list[dict]], x_in: torch.Tensor, timestep, model_options):
+def _calc_cond_batch(model: BaseModel, conds: list[list[dict]], x_in: torch.Tensor, timestep, model_options):
out_conds = []
out_counts = []
# separate conds by matching hooks
@@ -298,17 +306,10 @@ def _calc_cond_batch(model: 'BaseModel', conds: list[list[dict]], x_in: torch.Te
copy_dict1=False)
if patches is not None:
- # TODO: replace with merge_nested_dicts function
- if "patches" in transformer_options:
- cur_patches = transformer_options["patches"].copy()
- for p in patches:
- if p in cur_patches:
- cur_patches[p] = cur_patches[p] + patches[p]
- else:
- cur_patches[p] = patches[p]
- transformer_options["patches"] = cur_patches
- else:
- transformer_options["patches"] = patches
+ transformer_options["patches"] = comfy.patcher_extension.merge_nested_dicts(
+ transformer_options.get("patches", {}),
+ patches
+ )
transformer_options["cond_or_uncond"] = cond_or_uncond[:]
transformer_options["uuids"] = uuids[:]
@@ -352,7 +353,7 @@ def calc_cond_uncond_batch(model, cond, uncond, x_in, timestep, model_options):
def cfg_function(model, cond_pred, uncond_pred, cond_scale, x, timestep, model_options={}, cond=None, uncond=None):
if "sampler_cfg_function" in model_options:
args = {"cond": x - cond_pred, "uncond": x - uncond_pred, "cond_scale": cond_scale, "timestep": timestep, "input": x, "sigma": timestep,
- "cond_denoised": cond_pred, "uncond_denoised": uncond_pred, "model": model, "model_options": model_options}
+ "cond_denoised": cond_pred, "uncond_denoised": uncond_pred, "model": model, "model_options": model_options, "input_cond": cond, "input_uncond": uncond}
cfg_result = x - model_options["sampler_cfg_function"](args)
else:
cfg_result = uncond_pred + (cond_pred - uncond_pred) * cond_scale
@@ -382,7 +383,7 @@ def sampling_function(model, x, timestep, uncond, cond, cond_scale, model_option
for fn in model_options.get("sampler_pre_cfg_function", []):
args = {"conds":conds, "conds_out": out, "cond_scale": cond_scale, "timestep": timestep,
"input": x, "sigma": timestep, "model": model, "model_options": model_options}
- out = fn(args)
+ out = fn(args)
return cfg_function(model, out[0], out[1], cond_scale, x, timestep, model_options=model_options, cond=cond, uncond=uncond_)
@@ -546,7 +547,10 @@ def resolve_areas_and_cond_masks_multidim(conditions, dims, device):
if len(mask.shape) == len(dims):
mask = mask.unsqueeze(0)
if mask.shape[1:] != dims:
- mask = torch.nn.functional.interpolate(mask.unsqueeze(1), size=dims, mode='bilinear', align_corners=False).squeeze(1)
+ if mask.ndim < 4:
+ mask = comfy.utils.common_upscale(mask.unsqueeze(1), dims[-1], dims[-2], 'bilinear', 'none').squeeze(1)
+ else:
+ mask = comfy.utils.common_upscale(mask, dims[-1], dims[-2], 'bilinear', 'none')
if modified.get("set_area_to_bounds", False): #TODO: handle dim != 2
bounds = torch.max(torch.abs(mask),dim=0).values.unsqueeze(0)
@@ -718,9 +722,9 @@ class Sampler:
KSAMPLER_NAMES = ["euler", "euler_cfg_pp", "euler_ancestral", "euler_ancestral_cfg_pp", "heun", "heunpp2","dpm_2", "dpm_2_ancestral",
"lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_2s_ancestral_cfg_pp", "dpmpp_sde", "dpmpp_sde_gpu",
- "dpmpp_2m", "dpmpp_2m_cfg_pp", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm",
+ "dpmpp_2m", "dpmpp_2m_cfg_pp", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_2m_sde_heun", "dpmpp_2m_sde_heun_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm",
"ipndm", "ipndm_v", "deis", "res_multistep", "res_multistep_cfg_pp", "res_multistep_ancestral", "res_multistep_ancestral_cfg_pp",
- "gradient_estimation", "gradient_estimation_cfg_pp", "er_sde", "seeds_2", "seeds_3"]
+ "gradient_estimation", "gradient_estimation_cfg_pp", "er_sde", "seeds_2", "seeds_3", "sa_solver", "sa_solver_pece"]
class KSAMPLER(Sampler):
def __init__(self, sampler_function, extra_options={}, inpaint_options={}):
@@ -778,7 +782,7 @@ def ksampler(sampler_name, extra_options={}, inpaint_options={}):
return KSAMPLER(sampler_function, extra_options, inpaint_options)
-def process_conds(model, noise, conds, device, latent_image=None, denoise_mask=None, seed=None):
+def process_conds(model, noise, conds, device, latent_image=None, denoise_mask=None, seed=None, latent_shapes=None):
for k in conds:
conds[k] = conds[k][:]
resolve_areas_and_cond_masks_multidim(conds[k], noise.shape[2:], device)
@@ -788,7 +792,7 @@ def process_conds(model, noise, conds, device, latent_image=None, denoise_mask=N
if hasattr(model, 'extra_conds'):
for k in conds:
- conds[k] = encode_model_conds(model.extra_conds, conds[k], noise, device, k, latent_image=latent_image, denoise_mask=denoise_mask, seed=seed)
+ conds[k] = encode_model_conds(model.extra_conds, conds[k], noise, device, k, latent_image=latent_image, denoise_mask=denoise_mask, seed=seed, latent_shapes=latent_shapes)
#make sure each cond area has an opposite one with the same area
for k in conds:
@@ -946,16 +950,23 @@ class CFGGuider:
self.original_conds[k] = comfy.sampler_helpers.convert_cond(conds[k])
def __call__(self, *args, **kwargs):
- return self.predict_noise(*args, **kwargs)
+ return self.outer_predict_noise(*args, **kwargs)
+
+ def outer_predict_noise(self, x, timestep, model_options={}, seed=None):
+ return comfy.patcher_extension.WrapperExecutor.new_class_executor(
+ self.predict_noise,
+ self,
+ comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.PREDICT_NOISE, self.model_options, is_model_options=True)
+ ).execute(x, timestep, model_options, seed)
def predict_noise(self, x, timestep, model_options={}, seed=None):
return sampling_function(self.inner_model, x, timestep, self.conds.get("negative", None), self.conds.get("positive", None), self.cfg, model_options=model_options, seed=seed)
- def inner_sample(self, noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed):
+ def inner_sample(self, noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed, latent_shapes=None):
if latent_image is not None and torch.count_nonzero(latent_image) > 0: #Don't shift the empty latent image.
latent_image = self.inner_model.process_latent_in(latent_image)
- self.conds = process_conds(self.inner_model, noise, self.conds, device, latent_image, denoise_mask, seed)
+ self.conds = process_conds(self.inner_model, noise, self.conds, device, latent_image, denoise_mask, seed, latent_shapes=latent_shapes)
extra_model_options = comfy.model_patcher.create_model_options_clone(self.model_options)
extra_model_options.setdefault("transformer_options", {})["sample_sigmas"] = sigmas
@@ -969,7 +980,7 @@ class CFGGuider:
samples = executor.execute(self, sigmas, extra_args, callback, noise, latent_image, denoise_mask, disable_pbar)
return self.inner_model.process_latent_out(samples.to(torch.float32))
- def outer_sample(self, noise, latent_image, sampler, sigmas, denoise_mask=None, callback=None, disable_pbar=False, seed=None):
+ def outer_sample(self, noise, latent_image, sampler, sigmas, denoise_mask=None, callback=None, disable_pbar=False, seed=None, latent_shapes=None):
self.inner_model, self.conds, self.loaded_models = comfy.sampler_helpers.prepare_sampling(self.model_patcher, noise.shape, self.conds, self.model_options)
device = self.model_patcher.load_device
@@ -983,7 +994,7 @@ class CFGGuider:
try:
self.model_patcher.pre_run()
- output = self.inner_sample(noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
+ output = self.inner_sample(noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed, latent_shapes=latent_shapes)
finally:
self.model_patcher.cleanup()
@@ -996,6 +1007,12 @@ class CFGGuider:
if sigmas.shape[-1] == 0:
return latent_image
+ if latent_image.is_nested:
+ latent_image, latent_shapes = comfy.utils.pack_latents(latent_image.unbind())
+ noise, _ = comfy.utils.pack_latents(noise.unbind())
+ else:
+ latent_shapes = [latent_image.shape]
+
self.conds = {}
for k in self.original_conds:
self.conds[k] = list(map(lambda a: a.copy(), self.original_conds[k]))
@@ -1015,7 +1032,7 @@ class CFGGuider:
self,
comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.OUTER_SAMPLE, self.model_options, is_model_options=True)
)
- output = executor.execute(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
+ output = executor.execute(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed, latent_shapes=latent_shapes)
finally:
cast_to_load_options(self.model_options, device=self.model_patcher.offload_device)
self.model_options = orig_model_options
@@ -1023,6 +1040,9 @@ class CFGGuider:
self.model_patcher.restore_hook_patches()
del self.conds
+
+ if len(latent_shapes) > 1:
+ output = comfy.nested_tensor.NestedTensor(comfy.utils.unpack_latents(output, latent_shapes))
return output
diff --git a/comfy/sd.py b/comfy/sd.py
index 5b95cf75a..9e5ebbf15 100644
--- a/comfy/sd.py
+++ b/comfy/sd.py
@@ -14,10 +14,15 @@ import comfy.ldm.genmo.vae.model
import comfy.ldm.lightricks.vae.causal_video_autoencoder
import comfy.ldm.cosmos.vae
import comfy.ldm.wan.vae
+import comfy.ldm.wan.vae2_2
import comfy.ldm.hunyuan3d.vae
import comfy.ldm.ace.vae.music_dcae_pipeline
+import comfy.ldm.hunyuan_video.vae
+import comfy.ldm.mmaudio.vae.autoencoder
+import comfy.pixel_space_convert
import yaml
import math
+import os
import comfy.utils
@@ -45,6 +50,8 @@ import comfy.text_encoders.wan
import comfy.text_encoders.hidream
import comfy.text_encoders.ace
import comfy.text_encoders.omnigen2
+import comfy.text_encoders.qwen_image
+import comfy.text_encoders.hunyuan_image
import comfy.model_patcher
import comfy.lora
@@ -136,6 +143,9 @@ class CLIP:
n.apply_hooks_to_conds = self.apply_hooks_to_conds
return n
+ def get_ram_usage(self):
+ return self.patcher.get_ram_usage()
+
def add_patches(self, patches, strength_patch=1.0, strength_model=1.0):
return self.patcher.add_patches(patches, strength_patch, strength_model)
@@ -269,8 +279,13 @@ class VAE:
if 'decoder.up_blocks.0.resnets.0.norm1.weight' in sd.keys(): #diffusers format
sd = diffusers_convert.convert_vae_state_dict(sd)
- self.memory_used_encode = lambda shape, dtype: (1767 * shape[2] * shape[3]) * model_management.dtype_size(dtype) #These are for AutoencoderKL and need tweaking (should be lower)
- self.memory_used_decode = lambda shape, dtype: (2178 * shape[2] * shape[3] * 64) * model_management.dtype_size(dtype)
+ if model_management.is_amd():
+ VAE_KL_MEM_RATIO = 2.73
+ else:
+ VAE_KL_MEM_RATIO = 1.0
+
+ self.memory_used_encode = lambda shape, dtype: (1767 * shape[2] * shape[3]) * model_management.dtype_size(dtype) * VAE_KL_MEM_RATIO #These are for AutoencoderKL and need tweaking (should be lower)
+ self.memory_used_decode = lambda shape, dtype: (2178 * shape[2] * shape[3] * 64) * model_management.dtype_size(dtype) * VAE_KL_MEM_RATIO
self.downscale_ratio = 8
self.upscale_ratio = 8
self.latent_channels = 4
@@ -280,10 +295,13 @@ class VAE:
self.process_output = lambda image: torch.clamp((image + 1.0) / 2.0, min=0.0, max=1.0)
self.working_dtypes = [torch.bfloat16, torch.float32]
self.disable_offload = False
+ self.not_video = False
+ self.size = None
self.downscale_index_formula = None
self.upscale_index_formula = None
self.extra_1d_channel = None
+ self.crop_input = True
if config is None:
if "decoder.mid.block_1.mix_factor" in sd:
@@ -326,21 +344,50 @@ class VAE:
self.downscale_ratio = 32
self.latent_channels = 16
elif "decoder.conv_in.weight" in sd:
- #default SD1.x/SD2.x VAE parameters
- ddconfig = {'double_z': True, 'z_channels': 4, 'resolution': 256, 'in_channels': 3, 'out_ch': 3, 'ch': 128, 'ch_mult': [1, 2, 4, 4], 'num_res_blocks': 2, 'attn_resolutions': [], 'dropout': 0.0}
-
- if 'encoder.down.2.downsample.conv.weight' not in sd and 'decoder.up.3.upsample.conv.weight' not in sd: #Stable diffusion x4 upscaler VAE
- ddconfig['ch_mult'] = [1, 2, 4]
- self.downscale_ratio = 4
- self.upscale_ratio = 4
-
- self.latent_channels = ddconfig['z_channels'] = sd["decoder.conv_in.weight"].shape[1]
- if 'post_quant_conv.weight' in sd:
- self.first_stage_model = AutoencoderKL(ddconfig=ddconfig, embed_dim=sd['post_quant_conv.weight'].shape[1])
- else:
+ if sd['decoder.conv_in.weight'].shape[1] == 64:
+ ddconfig = {"block_out_channels": [128, 256, 512, 512, 1024, 1024], "in_channels": 3, "out_channels": 3, "num_res_blocks": 2, "ffactor_spatial": 32, "downsample_match_channel": True, "upsample_match_channel": True}
+ self.latent_channels = ddconfig['z_channels'] = sd["decoder.conv_in.weight"].shape[1]
+ self.downscale_ratio = 32
+ self.upscale_ratio = 32
+ self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32]
self.first_stage_model = AutoencodingEngine(regularizer_config={'target': "comfy.ldm.models.autoencoder.DiagonalGaussianRegularizer"},
- encoder_config={'target': "comfy.ldm.modules.diffusionmodules.model.Encoder", 'params': ddconfig},
- decoder_config={'target': "comfy.ldm.modules.diffusionmodules.model.Decoder", 'params': ddconfig})
+ encoder_config={'target': "comfy.ldm.hunyuan_video.vae.Encoder", 'params': ddconfig},
+ decoder_config={'target': "comfy.ldm.hunyuan_video.vae.Decoder", 'params': ddconfig})
+
+ self.memory_used_encode = lambda shape, dtype: (700 * shape[2] * shape[3]) * model_management.dtype_size(dtype)
+ self.memory_used_decode = lambda shape, dtype: (700 * shape[2] * shape[3] * 32 * 32) * model_management.dtype_size(dtype)
+ elif sd['decoder.conv_in.weight'].shape[1] == 32:
+ ddconfig = {"block_out_channels": [128, 256, 512, 1024, 1024], "in_channels": 3, "out_channels": 3, "num_res_blocks": 2, "ffactor_spatial": 16, "ffactor_temporal": 4, "downsample_match_channel": True, "upsample_match_channel": True, "refiner_vae": False}
+ self.latent_channels = ddconfig['z_channels'] = sd["decoder.conv_in.weight"].shape[1]
+ self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32]
+ self.upscale_ratio = (lambda a: max(0, a * 4 - 3), 16, 16)
+ self.upscale_index_formula = (4, 16, 16)
+ self.downscale_ratio = (lambda a: max(0, math.floor((a + 3) / 4)), 16, 16)
+ self.downscale_index_formula = (4, 16, 16)
+ self.latent_dim = 3
+ self.not_video = True
+ self.first_stage_model = AutoencodingEngine(regularizer_config={'target': "comfy.ldm.models.autoencoder.DiagonalGaussianRegularizer"},
+ encoder_config={'target': "comfy.ldm.hunyuan_video.vae_refiner.Encoder", 'params': ddconfig},
+ decoder_config={'target': "comfy.ldm.hunyuan_video.vae_refiner.Decoder", 'params': ddconfig})
+
+ self.memory_used_encode = lambda shape, dtype: (2800 * shape[-2] * shape[-1]) * model_management.dtype_size(dtype)
+ self.memory_used_decode = lambda shape, dtype: (2800 * shape[-3] * shape[-2] * shape[-1] * 16 * 16) * model_management.dtype_size(dtype)
+ else:
+ #default SD1.x/SD2.x VAE parameters
+ ddconfig = {'double_z': True, 'z_channels': 4, 'resolution': 256, 'in_channels': 3, 'out_ch': 3, 'ch': 128, 'ch_mult': [1, 2, 4, 4], 'num_res_blocks': 2, 'attn_resolutions': [], 'dropout': 0.0}
+
+ if 'encoder.down.2.downsample.conv.weight' not in sd and 'decoder.up.3.upsample.conv.weight' not in sd: #Stable diffusion x4 upscaler VAE
+ ddconfig['ch_mult'] = [1, 2, 4]
+ self.downscale_ratio = 4
+ self.upscale_ratio = 4
+
+ self.latent_channels = ddconfig['z_channels'] = sd["decoder.conv_in.weight"].shape[1]
+ if 'post_quant_conv.weight' in sd:
+ self.first_stage_model = AutoencoderKL(ddconfig=ddconfig, embed_dim=sd['post_quant_conv.weight'].shape[1])
+ else:
+ self.first_stage_model = AutoencodingEngine(regularizer_config={'target': "comfy.ldm.models.autoencoder.DiagonalGaussianRegularizer"},
+ encoder_config={'target': "comfy.ldm.modules.diffusionmodules.model.Encoder", 'params': ddconfig},
+ decoder_config={'target': "comfy.ldm.modules.diffusionmodules.model.Decoder", 'params': ddconfig})
elif "decoder.layers.1.layers.0.beta" in sd:
self.first_stage_model = AudioOobleckVAE()
self.memory_used_encode = lambda shape, dtype: (1000 * shape[2]) * model_management.dtype_size(dtype)
@@ -391,6 +438,23 @@ class VAE:
self.downscale_ratio = (lambda a: max(0, math.floor((a + 7) / 8)), 32, 32)
self.downscale_index_formula = (8, 32, 32)
self.working_dtypes = [torch.bfloat16, torch.float32]
+ elif "decoder.conv_in.conv.weight" in sd and sd['decoder.conv_in.conv.weight'].shape[1] == 32:
+ ddconfig = {"block_out_channels": [128, 256, 512, 1024, 1024], "in_channels": 3, "out_channels": 3, "num_res_blocks": 2, "ffactor_spatial": 16, "ffactor_temporal": 4, "downsample_match_channel": True, "upsample_match_channel": True}
+ ddconfig['z_channels'] = sd["decoder.conv_in.conv.weight"].shape[1]
+ self.latent_channels = 64
+ self.upscale_ratio = (lambda a: max(0, a * 4 - 3), 16, 16)
+ self.upscale_index_formula = (4, 16, 16)
+ self.downscale_ratio = (lambda a: max(0, math.floor((a + 3) / 4)), 16, 16)
+ self.downscale_index_formula = (4, 16, 16)
+ self.latent_dim = 3
+ self.not_video = True
+ self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32]
+ self.first_stage_model = AutoencodingEngine(regularizer_config={'target': "comfy.ldm.models.autoencoder.EmptyRegularizer"},
+ encoder_config={'target': "comfy.ldm.hunyuan_video.vae_refiner.Encoder", 'params': ddconfig},
+ decoder_config={'target': "comfy.ldm.hunyuan_video.vae_refiner.Decoder", 'params': ddconfig})
+
+ self.memory_used_encode = lambda shape, dtype: (1400 * shape[-2] * shape[-1]) * model_management.dtype_size(dtype)
+ self.memory_used_decode = lambda shape, dtype: (1400 * shape[-3] * shape[-2] * shape[-1] * 16 * 16) * model_management.dtype_size(dtype)
elif "decoder.conv_in.conv.weight" in sd:
ddconfig = {'double_z': True, 'z_channels': 4, 'resolution': 256, 'in_channels': 3, 'out_ch': 3, 'ch': 128, 'ch_mult': [1, 2, 4, 4], 'num_res_blocks': 2, 'attn_resolutions': [], 'dropout': 0.0}
ddconfig["conv3d"] = True
@@ -419,28 +483,53 @@ class VAE:
self.memory_used_encode = lambda shape, dtype: (50 * (round((shape[2] + 7) / 8) * 8) * shape[3] * shape[4]) * model_management.dtype_size(dtype)
self.working_dtypes = [torch.bfloat16, torch.float32]
elif "decoder.middle.0.residual.0.gamma" in sd:
- self.upscale_ratio = (lambda a: max(0, a * 4 - 3), 8, 8)
- self.upscale_index_formula = (4, 8, 8)
- self.downscale_ratio = (lambda a: max(0, math.floor((a + 3) / 4)), 8, 8)
- self.downscale_index_formula = (4, 8, 8)
- self.latent_dim = 3
- self.latent_channels = 16
- ddconfig = {"dim": 96, "z_dim": self.latent_channels, "dim_mult": [1, 2, 4, 4], "num_res_blocks": 2, "attn_scales": [], "temperal_downsample": [False, True, True], "dropout": 0.0}
- self.first_stage_model = comfy.ldm.wan.vae.WanVAE(**ddconfig)
- self.working_dtypes = [torch.bfloat16, torch.float16, torch.float32]
- self.memory_used_encode = lambda shape, dtype: 6000 * shape[3] * shape[4] * model_management.dtype_size(dtype)
- self.memory_used_decode = lambda shape, dtype: 7000 * shape[3] * shape[4] * (8 * 8) * model_management.dtype_size(dtype)
+ if "decoder.upsamples.0.upsamples.0.residual.2.weight" in sd: # Wan 2.2 VAE
+ self.upscale_ratio = (lambda a: max(0, a * 4 - 3), 16, 16)
+ self.upscale_index_formula = (4, 16, 16)
+ self.downscale_ratio = (lambda a: max(0, math.floor((a + 3) / 4)), 16, 16)
+ self.downscale_index_formula = (4, 16, 16)
+ self.latent_dim = 3
+ self.latent_channels = 48
+ ddconfig = {"dim": 160, "z_dim": self.latent_channels, "dim_mult": [1, 2, 4, 4], "num_res_blocks": 2, "attn_scales": [], "temperal_downsample": [False, True, True], "dropout": 0.0}
+ self.first_stage_model = comfy.ldm.wan.vae2_2.WanVAE(**ddconfig)
+ self.working_dtypes = [torch.bfloat16, torch.float16, torch.float32]
+ self.memory_used_encode = lambda shape, dtype: 3300 * shape[3] * shape[4] * model_management.dtype_size(dtype)
+ self.memory_used_decode = lambda shape, dtype: 8000 * shape[3] * shape[4] * (16 * 16) * model_management.dtype_size(dtype)
+ else: # Wan 2.1 VAE
+ self.upscale_ratio = (lambda a: max(0, a * 4 - 3), 8, 8)
+ self.upscale_index_formula = (4, 8, 8)
+ self.downscale_ratio = (lambda a: max(0, math.floor((a + 3) / 4)), 8, 8)
+ self.downscale_index_formula = (4, 8, 8)
+ self.latent_dim = 3
+ self.latent_channels = 16
+ ddconfig = {"dim": 96, "z_dim": self.latent_channels, "dim_mult": [1, 2, 4, 4], "num_res_blocks": 2, "attn_scales": [], "temperal_downsample": [False, True, True], "dropout": 0.0}
+ self.first_stage_model = comfy.ldm.wan.vae.WanVAE(**ddconfig)
+ self.working_dtypes = [torch.bfloat16, torch.float16, torch.float32]
+ self.memory_used_encode = lambda shape, dtype: 6000 * shape[3] * shape[4] * model_management.dtype_size(dtype)
+ self.memory_used_decode = lambda shape, dtype: 7000 * shape[3] * shape[4] * (8 * 8) * model_management.dtype_size(dtype)
+ # Hunyuan 3d v2 2.0 & 2.1
elif "geo_decoder.cross_attn_decoder.ln_1.bias" in sd:
+
self.latent_dim = 1
- ln_post = "geo_decoder.ln_post.weight" in sd
- inner_size = sd["geo_decoder.output_proj.weight"].shape[1]
- downsample_ratio = sd["post_kl.weight"].shape[0] // inner_size
- mlp_expand = sd["geo_decoder.cross_attn_decoder.mlp.c_fc.weight"].shape[0] // inner_size
- self.memory_used_encode = lambda shape, dtype: (1000 * shape[2]) * model_management.dtype_size(dtype) # TODO
- self.memory_used_decode = lambda shape, dtype: (1024 * 1024 * 1024 * 2.0) * model_management.dtype_size(dtype) # TODO
- ddconfig = {"embed_dim": 64, "num_freqs": 8, "include_pi": False, "heads": 16, "width": 1024, "num_decoder_layers": 16, "qkv_bias": False, "qk_norm": True, "geo_decoder_mlp_expand_ratio": mlp_expand, "geo_decoder_downsample_ratio": downsample_ratio, "geo_decoder_ln_post": ln_post}
- self.first_stage_model = comfy.ldm.hunyuan3d.vae.ShapeVAE(**ddconfig)
+
+ def estimate_memory(shape, dtype, num_layers = 16, kv_cache_multiplier = 2):
+ batch, num_tokens, hidden_dim = shape
+ dtype_size = model_management.dtype_size(dtype)
+
+ total_mem = batch * num_tokens * hidden_dim * dtype_size * (1 + kv_cache_multiplier * num_layers)
+ return total_mem
+
+ # better memory estimations
+ self.memory_used_encode = lambda shape, dtype, num_layers = 8, kv_cache_multiplier = 0:\
+ estimate_memory(shape, dtype, num_layers, kv_cache_multiplier)
+
+ self.memory_used_decode = lambda shape, dtype, num_layers = 16, kv_cache_multiplier = 2: \
+ estimate_memory(shape, dtype, num_layers, kv_cache_multiplier)
+
+ self.first_stage_model = comfy.ldm.hunyuan3d.vae.ShapeVAE()
self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32]
+
+
elif "vocoder.backbone.channel_layers.0.0.bias" in sd: #Ace Step Audio
self.first_stage_model = comfy.ldm.ace.vae.music_dcae_pipeline.MusicDCAE(source_sample_rate=44100)
self.memory_used_encode = lambda shape, dtype: (shape[2] * 330) * model_management.dtype_size(dtype)
@@ -455,6 +544,34 @@ class VAE:
self.working_dtypes = [torch.bfloat16, torch.float16, torch.float32]
self.disable_offload = True
self.extra_1d_channel = 16
+ elif "pixel_space_vae" in sd:
+ self.first_stage_model = comfy.pixel_space_convert.PixelspaceConversionVAE()
+ self.memory_used_encode = lambda shape, dtype: (1 * shape[2] * shape[3]) * model_management.dtype_size(dtype)
+ self.memory_used_decode = lambda shape, dtype: (1 * shape[2] * shape[3]) * model_management.dtype_size(dtype)
+ self.downscale_ratio = 1
+ self.upscale_ratio = 1
+ self.latent_channels = 3
+ self.latent_dim = 2
+ self.output_channels = 3
+ elif "vocoder.activation_post.downsample.lowpass.filter" in sd: #MMAudio VAE
+ sample_rate = 16000
+ if sample_rate == 16000:
+ mode = '16k'
+ else:
+ mode = '44k'
+
+ self.first_stage_model = comfy.ldm.mmaudio.vae.autoencoder.AudioAutoencoder(mode=mode)
+ self.memory_used_encode = lambda shape, dtype: (30 * shape[2]) * model_management.dtype_size(dtype)
+ self.memory_used_decode = lambda shape, dtype: (90 * shape[2] * 1411.2) * model_management.dtype_size(dtype)
+ self.latent_channels = 20
+ self.output_channels = 2
+ self.upscale_ratio = 512 * (44100 / sample_rate)
+ self.downscale_ratio = 512 * (44100 / sample_rate)
+ self.latent_dim = 1
+ self.process_output = lambda audio: audio
+ self.process_input = lambda audio: audio
+ self.working_dtypes = [torch.float32]
+ self.crop_input = False
else:
logging.warning("WARNING: No VAE weights detected, VAE not initalized.")
self.first_stage_model = None
@@ -482,12 +599,25 @@ class VAE:
self.patcher = comfy.model_patcher.ModelPatcher(self.first_stage_model, load_device=self.device, offload_device=offload_device)
logging.info("VAE load device: {}, offload device: {}, dtype: {}".format(self.device, offload_device, self.vae_dtype))
+ self.model_size()
+
+ def model_size(self):
+ if self.size is not None:
+ return self.size
+ self.size = comfy.model_management.module_size(self.first_stage_model)
+ return self.size
+
+ def get_ram_usage(self):
+ return self.model_size()
def throw_exception_if_invalid(self):
if self.first_stage_model is None:
raise RuntimeError("ERROR: VAE is invalid: None\n\nIf the VAE is from a checkpoint loader node your checkpoint does not contain a valid VAE.")
def vae_encode_crop_pixels(self, pixels):
+ if not self.crop_input:
+ return pixels
+
downscale_ratio = self.spacial_compression_encode()
dims = pixels.shape[1:-1]
@@ -565,6 +695,7 @@ class VAE:
def decode(self, samples_in, vae_options={}):
self.throw_exception_if_invalid()
pixel_samples = None
+ do_tile = False
try:
memory_used = self.memory_used_decode(samples_in.shape, self.vae_dtype)
model_management.load_models_gpu([self.patcher], memory_required=memory_used, force_full_load=self.disable_offload)
@@ -580,6 +711,13 @@ class VAE:
pixel_samples[x:x+batch_number] = out
except model_management.OOM_EXCEPTION:
logging.warning("Warning: Ran out of memory when regular VAE decoding, retrying with tiled VAE decoding.")
+ #NOTE: We don't know what tensors were allocated to stack variables at the time of the
+ #exception and the exception itself refs them all until we get out of this except block.
+ #So we just set a flag for tiler fallback so that tensor gc can happen once the
+ #exception is fully off the books.
+ do_tile = True
+
+ if do_tile:
dims = samples_in.ndim - 2
if dims == 1 or self.extra_1d_channel is not None:
pixel_samples = self.decode_tiled_1d(samples_in)
@@ -626,8 +764,12 @@ class VAE:
self.throw_exception_if_invalid()
pixel_samples = self.vae_encode_crop_pixels(pixel_samples)
pixel_samples = pixel_samples.movedim(-1, 1)
+ do_tile = False
if self.latent_dim == 3 and pixel_samples.ndim < 5:
- pixel_samples = pixel_samples.movedim(1, 0).unsqueeze(0)
+ if not self.not_video:
+ pixel_samples = pixel_samples.movedim(1, 0).unsqueeze(0)
+ else:
+ pixel_samples = pixel_samples.unsqueeze(2)
try:
memory_used = self.memory_used_encode(pixel_samples.shape, self.vae_dtype)
model_management.load_models_gpu([self.patcher], memory_required=memory_used, force_full_load=self.disable_offload)
@@ -644,6 +786,13 @@ class VAE:
except model_management.OOM_EXCEPTION:
logging.warning("Warning: Ran out of memory when regular VAE encoding, retrying with tiled VAE encoding.")
+ #NOTE: We don't know what tensors were allocated to stack variables at the time of the
+ #exception and the exception itself refs them all until we get out of this except block.
+ #So we just set a flag for tiler fallback so that tensor gc can happen once the
+ #exception is fully off the books.
+ do_tile = True
+
+ if do_tile:
if self.latent_dim == 3:
tile = 256
overlap = tile // 4
@@ -661,7 +810,10 @@ class VAE:
dims = self.latent_dim
pixel_samples = pixel_samples.movedim(-1, 1)
if dims == 3:
- pixel_samples = pixel_samples.movedim(1, 0).unsqueeze(0)
+ if not self.not_video:
+ pixel_samples = pixel_samples.movedim(1, 0).unsqueeze(0)
+ else:
+ pixel_samples = pixel_samples.unsqueeze(2)
memory_used = self.memory_used_encode(pixel_samples.shape, self.vae_dtype) # TODO: calculate mem required for tile
model_management.load_models_gpu([self.patcher], memory_required=memory_used, force_full_load=self.disable_offload)
@@ -718,6 +870,7 @@ class VAE:
except:
return None
+
class StyleModel:
def __init__(self, model, device="cpu"):
self.model = model
@@ -756,6 +909,8 @@ class CLIPType(Enum):
CHROMA = 15
ACE = 16
OMNIGEN2 = 17
+ QWEN_IMAGE = 18
+ HUNYUAN_IMAGE = 19
def load_clip(ckpt_paths, embedding_directory=None, clip_type=CLIPType.STABLE_DIFFUSION, model_options={}):
@@ -776,6 +931,9 @@ class TEModel(Enum):
T5_XXL_OLD = 8
GEMMA_2_2B = 9
QWEN25_3B = 10
+ QWEN25_7B = 11
+ BYT5_SMALL_GLYPH = 12
+ GEMMA_3_4B = 13
def detect_te_model(sd):
if "text_model.encoder.layers.30.mlp.fc1.weight" in sd:
@@ -793,11 +951,20 @@ def detect_te_model(sd):
if 'encoder.block.23.layer.1.DenseReluDense.wi.weight' in sd:
return TEModel.T5_XXL_OLD
if "encoder.block.0.layer.0.SelfAttention.k.weight" in sd:
+ weight = sd['encoder.block.0.layer.0.SelfAttention.k.weight']
+ if weight.shape[0] == 384:
+ return TEModel.BYT5_SMALL_GLYPH
return TEModel.T5_BASE
if 'model.layers.0.post_feedforward_layernorm.weight' in sd:
+ if 'model.layers.0.self_attn.q_norm.weight' in sd:
+ return TEModel.GEMMA_3_4B
return TEModel.GEMMA_2_2B
if 'model.layers.0.self_attn.k_proj.bias' in sd:
- return TEModel.QWEN25_3B
+ weight = sd['model.layers.0.self_attn.k_proj.bias']
+ if weight.shape[0] == 256:
+ return TEModel.QWEN25_3B
+ if weight.shape[0] == 512:
+ return TEModel.QWEN25_7B
if "model.layers.0.post_attention_layernorm.weight" in sd:
return TEModel.LLAMA3_8
return None
@@ -895,6 +1062,10 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
clip_target.clip = comfy.text_encoders.lumina2.te(**llama_detect(clip_data))
clip_target.tokenizer = comfy.text_encoders.lumina2.LuminaTokenizer
tokenizer_data["spiece_model"] = clip_data[0].get("spiece_model", None)
+ elif te_model == TEModel.GEMMA_3_4B:
+ clip_target.clip = comfy.text_encoders.lumina2.te(**llama_detect(clip_data), model_type="gemma3_4b")
+ clip_target.tokenizer = comfy.text_encoders.lumina2.NTokenizer
+ tokenizer_data["spiece_model"] = clip_data[0].get("spiece_model", None)
elif te_model == TEModel.LLAMA3_8:
clip_target.clip = comfy.text_encoders.hidream.hidream_clip(**llama_detect(clip_data),
clip_l=False, clip_g=False, t5=False, llama=True, dtype_t5=None, t5xxl_scaled_fp8=None)
@@ -902,6 +1073,13 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
elif te_model == TEModel.QWEN25_3B:
clip_target.clip = comfy.text_encoders.omnigen2.te(**llama_detect(clip_data))
clip_target.tokenizer = comfy.text_encoders.omnigen2.Omnigen2Tokenizer
+ elif te_model == TEModel.QWEN25_7B:
+ if clip_type == CLIPType.HUNYUAN_IMAGE:
+ clip_target.clip = comfy.text_encoders.hunyuan_image.te(byt5=False, **llama_detect(clip_data))
+ clip_target.tokenizer = comfy.text_encoders.hunyuan_image.HunyuanImageTokenizer
+ else:
+ clip_target.clip = comfy.text_encoders.qwen_image.te(**llama_detect(clip_data))
+ clip_target.tokenizer = comfy.text_encoders.qwen_image.QwenImageTokenizer
else:
# clip_l
if clip_type == CLIPType.SD3:
@@ -945,6 +1123,9 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
clip_target.clip = comfy.text_encoders.hidream.hidream_clip(clip_l=clip_l, clip_g=clip_g, t5=t5, llama=llama, **t5_kwargs, **llama_kwargs)
clip_target.tokenizer = comfy.text_encoders.hidream.HiDreamTokenizer
+ elif clip_type == CLIPType.HUNYUAN_IMAGE:
+ clip_target.clip = comfy.text_encoders.hunyuan_image.te(**llama_detect(clip_data))
+ clip_target.tokenizer = comfy.text_encoders.hunyuan_image.HunyuanImageTokenizer
else:
clip_target.clip = sdxl_clip.SDXLClipModel
clip_target.tokenizer = sdxl_clip.SDXLTokenizer
@@ -977,6 +1158,12 @@ def load_gligen(ckpt_path):
model = model.half()
return comfy.model_patcher.ModelPatcher(model, load_device=model_management.get_torch_device(), offload_device=model_management.unet_offload_device())
+def model_detection_error_hint(path, state_dict):
+ filename = os.path.basename(path)
+ if 'lora' in filename.lower():
+ return "\nHINT: This seems to be a Lora file and Lora files should be put in the lora folder and loaded with a lora loader node.."
+ return ""
+
def load_checkpoint(config_path=None, ckpt_path=None, output_vae=True, output_clip=True, embedding_directory=None, state_dict=None, config=None):
logging.warning("Warning: The load checkpoint with config function is deprecated and will eventually be removed, please use the other one.")
model, clip, vae, _ = load_checkpoint_guess_config(ckpt_path, output_vae=output_vae, output_clip=output_clip, output_clipvision=False, embedding_directory=embedding_directory, output_model=True)
@@ -1005,7 +1192,7 @@ def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, o
sd, metadata = comfy.utils.load_torch_file(ckpt_path, return_metadata=True)
out = load_state_dict_guess_config(sd, output_vae, output_clip, output_clipvision, embedding_directory, output_model, model_options, te_model_options=te_model_options, metadata=metadata)
if out is None:
- raise RuntimeError("ERROR: Could not detect model type of: {}".format(ckpt_path))
+ raise RuntimeError("ERROR: Could not detect model type of: {}\n{}".format(ckpt_path, model_detection_error_hint(ckpt_path, sd)))
return out
def load_state_dict_guess_config(sd, output_vae=True, output_clip=True, output_clipvision=False, embedding_directory=None, output_model=True, model_options={}, te_model_options={}, metadata=None):
@@ -1089,7 +1276,7 @@ def load_state_dict_guess_config(sd, output_vae=True, output_clip=True, output_c
return (model_patcher, clip, vae, clipvision)
-def load_diffusion_model_state_dict(sd, model_options={}):
+def load_diffusion_model_state_dict(sd, model_options={}, metadata=None):
"""
Loads a UNet diffusion model from a state dictionary, supporting both diffusers and regular formats.
@@ -1123,7 +1310,7 @@ def load_diffusion_model_state_dict(sd, model_options={}):
weight_dtype = comfy.utils.weight_dtype(sd)
load_device = model_management.get_torch_device()
- model_config = model_detection.model_config_from_unet(sd, "")
+ model_config = model_detection.model_config_from_unet(sd, "", metadata=metadata)
if model_config is not None:
new_sd = sd
@@ -1157,7 +1344,10 @@ def load_diffusion_model_state_dict(sd, model_options={}):
else:
unet_dtype = dtype
- manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)
+ if model_config.layer_quant_config is not None:
+ manual_cast_dtype = model_management.unet_manual_cast(None, load_device, model_config.supported_inference_dtypes)
+ else:
+ manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)
model_config.set_inference_dtype(unet_dtype, manual_cast_dtype)
model_config.custom_operations = model_options.get("custom_operations", model_config.custom_operations)
if model_options.get("fp8_optimizations", False):
@@ -1173,11 +1363,11 @@ def load_diffusion_model_state_dict(sd, model_options={}):
def load_diffusion_model(unet_path, model_options={}):
- sd = comfy.utils.load_torch_file(unet_path)
- model = load_diffusion_model_state_dict(sd, model_options=model_options)
+ sd, metadata = comfy.utils.load_torch_file(unet_path, return_metadata=True)
+ model = load_diffusion_model_state_dict(sd, model_options=model_options, metadata=metadata)
if model is None:
logging.error("ERROR UNSUPPORTED DIFFUSION MODEL {}".format(unet_path))
- raise RuntimeError("ERROR: Could not detect model type of: {}".format(unet_path))
+ raise RuntimeError("ERROR: Could not detect model type of: {}\n{}".format(unet_path, model_detection_error_hint(unet_path, sd)))
return model
def load_unet(unet_path, dtype=None):
diff --git a/comfy/sd1_clip.py b/comfy/sd1_clip.py
index ade340fd1..f8a7c2a1b 100644
--- a/comfy/sd1_clip.py
+++ b/comfy/sd1_clip.py
@@ -204,17 +204,19 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
tokens_embed = self.transformer.get_input_embeddings()(tokens_embed, out_dtype=torch.float32)
index = 0
pad_extra = 0
+ embeds_info = []
for o in other_embeds:
emb = o[1]
if torch.is_tensor(emb):
emb = {"type": "embedding", "data": emb}
+ extra = None
emb_type = emb.get("type", None)
if emb_type == "embedding":
emb = emb.get("data", None)
else:
if hasattr(self.transformer, "preprocess_embed"):
- emb = self.transformer.preprocess_embed(emb, device=device)
+ emb, extra = self.transformer.preprocess_embed(emb, device=device)
else:
emb = None
@@ -229,6 +231,7 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
tokens_embed = torch.cat([tokens_embed[:, :ind], emb, tokens_embed[:, ind:]], dim=1)
attention_mask = attention_mask[:ind] + [1] * emb_shape + attention_mask[ind:]
index += emb_shape - 1
+ embeds_info.append({"type": emb_type, "index": ind, "size": emb_shape, "extra": extra})
else:
index += -1
pad_extra += emb_shape
@@ -243,11 +246,11 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
attention_masks.append(attention_mask)
num_tokens.append(sum(attention_mask))
- return torch.cat(embeds_out), torch.tensor(attention_masks, device=device, dtype=torch.long), num_tokens
+ return torch.cat(embeds_out), torch.tensor(attention_masks, device=device, dtype=torch.long), num_tokens, embeds_info
def forward(self, tokens):
device = self.transformer.get_input_embeddings().weight.device
- embeds, attention_mask, num_tokens = self.process_tokens(tokens, device)
+ embeds, attention_mask, num_tokens, embeds_info = self.process_tokens(tokens, device)
attention_mask_model = None
if self.enable_attention_masks:
@@ -258,7 +261,7 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
else:
intermediate_output = self.layer_idx
- outputs = self.transformer(None, attention_mask_model, embeds=embeds, num_tokens=num_tokens, intermediate_output=intermediate_output, final_layer_norm_intermediate=self.layer_norm_hidden_state, dtype=torch.float32)
+ outputs = self.transformer(None, attention_mask_model, embeds=embeds, num_tokens=num_tokens, intermediate_output=intermediate_output, final_layer_norm_intermediate=self.layer_norm_hidden_state, dtype=torch.float32, embeds_info=embeds_info)
if self.layer == "last":
z = outputs[0].float()
@@ -531,7 +534,10 @@ class SDTokenizer:
min_padding = tokenizer_options.get("{}_min_padding".format(self.embedding_key), self.min_padding)
text = escape_important(text)
- parsed_weights = token_weights(text, 1.0)
+ if kwargs.get("disable_weights", False):
+ parsed_weights = [(text, 1.0)]
+ else:
+ parsed_weights = token_weights(text, 1.0)
# tokenize words
tokens = []
diff --git a/comfy/sd1_tokenizer/tokenizer_config.json b/comfy/sd1_tokenizer/tokenizer_config.json
index 5ba7bf706..8f7b3151d 100644
--- a/comfy/sd1_tokenizer/tokenizer_config.json
+++ b/comfy/sd1_tokenizer/tokenizer_config.json
@@ -18,7 +18,7 @@
"single_word": false
},
"errors": "replace",
- "model_max_length": 77,
+ "model_max_length": 8192,
"name_or_path": "openai/clip-vit-large-patch14",
"pad_token": "<|endoftext|>",
"special_tokens_map_file": "./special_tokens_map.json",
diff --git a/comfy/supported_models.py b/comfy/supported_models.py
index 2669ca01e..4064bdae1 100644
--- a/comfy/supported_models.py
+++ b/comfy/supported_models.py
@@ -19,6 +19,8 @@ import comfy.text_encoders.lumina2
import comfy.text_encoders.wan
import comfy.text_encoders.ace
import comfy.text_encoders.omnigen2
+import comfy.text_encoders.qwen_image
+import comfy.text_encoders.hunyuan_image
from . import supported_models_base
from . import latent_formats
@@ -699,7 +701,7 @@ class Flux(supported_models_base.BASE):
unet_extra_config = {}
latent_format = latent_formats.Flux
- memory_usage_factor = 2.8
+ memory_usage_factor = 3.1 # TODO: debug why flux mem usage is so weird on windows.
supported_inference_dtypes = [torch.bfloat16, torch.float16, torch.float32]
@@ -993,7 +995,7 @@ class WAN21_T2V(supported_models_base.BASE):
unet_extra_config = {}
latent_format = latent_formats.Wan21
- memory_usage_factor = 1.0
+ memory_usage_factor = 0.9
supported_inference_dtypes = [torch.float16, torch.bfloat16, torch.float32]
@@ -1002,7 +1004,7 @@ class WAN21_T2V(supported_models_base.BASE):
def __init__(self, unet_config):
super().__init__(unet_config)
- self.memory_usage_factor = self.unet_config.get("dim", 2000) / 2000
+ self.memory_usage_factor = self.unet_config.get("dim", 2000) / 2222
def get_model(self, state_dict, prefix="", device=None):
out = model_base.WAN21(self, device=device)
@@ -1045,6 +1047,18 @@ class WAN21_Camera(WAN21_T2V):
def get_model(self, state_dict, prefix="", device=None):
out = model_base.WAN21_Camera(self, image_to_video=False, device=device)
return out
+
+class WAN22_Camera(WAN21_T2V):
+ unet_config = {
+ "image_model": "wan2.1",
+ "model_type": "camera_2.2",
+ "in_dim": 36,
+ }
+
+ def get_model(self, state_dict, prefix="", device=None):
+ out = model_base.WAN21_Camera(self, image_to_video=False, device=device)
+ return out
+
class WAN21_Vace(WAN21_T2V):
unet_config = {
"image_model": "wan2.1",
@@ -1059,6 +1073,55 @@ class WAN21_Vace(WAN21_T2V):
out = model_base.WAN21_Vace(self, image_to_video=False, device=device)
return out
+class WAN21_HuMo(WAN21_T2V):
+ unet_config = {
+ "image_model": "wan2.1",
+ "model_type": "humo",
+ }
+
+ def get_model(self, state_dict, prefix="", device=None):
+ out = model_base.WAN21_HuMo(self, image_to_video=False, device=device)
+ return out
+
+class WAN22_S2V(WAN21_T2V):
+ unet_config = {
+ "image_model": "wan2.1",
+ "model_type": "s2v",
+ }
+
+ def __init__(self, unet_config):
+ super().__init__(unet_config)
+
+ def get_model(self, state_dict, prefix="", device=None):
+ out = model_base.WAN22_S2V(self, device=device)
+ return out
+
+class WAN22_Animate(WAN21_T2V):
+ unet_config = {
+ "image_model": "wan2.1",
+ "model_type": "animate",
+ }
+
+ def __init__(self, unet_config):
+ super().__init__(unet_config)
+
+ def get_model(self, state_dict, prefix="", device=None):
+ out = model_base.WAN22_Animate(self, device=device)
+ return out
+
+class WAN22_T2V(WAN21_T2V):
+ unet_config = {
+ "image_model": "wan2.1",
+ "model_type": "t2v",
+ "out_dim": 48,
+ }
+
+ latent_format = latent_formats.Wan22
+
+ def get_model(self, state_dict, prefix="", device=None):
+ out = model_base.WAN22(self, image_to_video=True, device=device)
+ return out
+
class Hunyuan3Dv2(supported_models_base.BASE):
unet_config = {
"image_model": "hunyuan3d2",
@@ -1089,6 +1152,17 @@ class Hunyuan3Dv2(supported_models_base.BASE):
def clip_target(self, state_dict={}):
return None
+class Hunyuan3Dv2_1(Hunyuan3Dv2):
+ unet_config = {
+ "image_model": "hunyuan3d2_1",
+ }
+
+ latent_format = latent_formats.Hunyuan3Dv2_1
+
+ def get_model(self, state_dict, prefix="", device=None):
+ out = model_base.Hunyuan3Dv2_1(self, device = device)
+ return out
+
class Hunyuan3Dv2mini(Hunyuan3Dv2):
unet_config = {
"image_model": "hunyuan3d2",
@@ -1154,6 +1228,19 @@ class Chroma(supported_models_base.BASE):
t5_detect = comfy.text_encoders.sd3_clip.t5_xxl_detect(state_dict, "{}t5xxl.transformer.".format(pref))
return supported_models_base.ClipTarget(comfy.text_encoders.pixart_t5.PixArtTokenizer, comfy.text_encoders.pixart_t5.pixart_te(**t5_detect))
+class ChromaRadiance(Chroma):
+ unet_config = {
+ "image_model": "chroma_radiance",
+ }
+
+ latent_format = comfy.latent_formats.ChromaRadiance
+
+ # Pixel-space model, no spatial compression for model input.
+ memory_usage_factor = 0.038
+
+ def get_model(self, state_dict, prefix="", device=None):
+ return model_base.ChromaRadiance(self, device=device)
+
class ACEStep(supported_models_base.BASE):
unet_config = {
"audio_model": "ace",
@@ -1214,9 +1301,79 @@ class Omnigen2(supported_models_base.BASE):
def clip_target(self, state_dict={}):
pref = self.text_encoder_key_prefix[0]
hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen25_3b.transformer.".format(pref))
- return supported_models_base.ClipTarget(comfy.text_encoders.omnigen2.LuminaTokenizer, comfy.text_encoders.omnigen2.te(**hunyuan_detect))
+ return supported_models_base.ClipTarget(comfy.text_encoders.omnigen2.Omnigen2Tokenizer, comfy.text_encoders.omnigen2.te(**hunyuan_detect))
+class QwenImage(supported_models_base.BASE):
+ unet_config = {
+ "image_model": "qwen_image",
+ }
-models = [LotusD, Stable_Zero123, SD15_instructpix2pix, SD15, SD20, SD21UnclipL, SD21UnclipH, SDXL_instructpix2pix, SDXLRefiner, SDXL, SSD1B, KOALA_700M, KOALA_1B, Segmind_Vega, SD_X4Upscaler, Stable_Cascade_C, Stable_Cascade_B, SV3D_u, SV3D_p, SD3, StableAudio, AuraFlow, PixArtAlpha, PixArtSigma, HunyuanDiT, HunyuanDiT1, FluxInpaint, Flux, FluxSchnell, GenmoMochi, LTXV, HunyuanVideoSkyreelsI2V, HunyuanVideoI2V, HunyuanVideo, CosmosT2V, CosmosI2V, CosmosT2IPredict2, CosmosI2VPredict2, Lumina2, WAN21_T2V, WAN21_I2V, WAN21_FunControl2V, WAN21_Vace, WAN21_Camera, Hunyuan3Dv2mini, Hunyuan3Dv2, HiDream, Chroma, ACEStep, Omnigen2]
+ sampling_settings = {
+ "multiplier": 1.0,
+ "shift": 1.15,
+ }
+
+ memory_usage_factor = 1.8 #TODO
+
+ unet_extra_config = {}
+ latent_format = latent_formats.Wan21
+
+ supported_inference_dtypes = [torch.bfloat16, torch.float32]
+
+ vae_key_prefix = ["vae."]
+ text_encoder_key_prefix = ["text_encoders."]
+
+ def get_model(self, state_dict, prefix="", device=None):
+ out = model_base.QwenImage(self, device=device)
+ return out
+
+ def clip_target(self, state_dict={}):
+ pref = self.text_encoder_key_prefix[0]
+ hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen25_7b.transformer.".format(pref))
+ return supported_models_base.ClipTarget(comfy.text_encoders.qwen_image.QwenImageTokenizer, comfy.text_encoders.qwen_image.te(**hunyuan_detect))
+
+class HunyuanImage21(HunyuanVideo):
+ unet_config = {
+ "image_model": "hunyuan_video",
+ "vec_in_dim": None,
+ }
+
+ sampling_settings = {
+ "shift": 5.0,
+ }
+
+ latent_format = latent_formats.HunyuanImage21
+
+ memory_usage_factor = 7.7
+
+ supported_inference_dtypes = [torch.bfloat16, torch.float32]
+
+ def get_model(self, state_dict, prefix="", device=None):
+ out = model_base.HunyuanImage21(self, device=device)
+ return out
+
+ def clip_target(self, state_dict={}):
+ pref = self.text_encoder_key_prefix[0]
+ hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen25_7b.transformer.".format(pref))
+ return supported_models_base.ClipTarget(comfy.text_encoders.hunyuan_image.HunyuanImageTokenizer, comfy.text_encoders.hunyuan_image.te(**hunyuan_detect))
+
+class HunyuanImage21Refiner(HunyuanVideo):
+ unet_config = {
+ "image_model": "hunyuan_video",
+ "patch_size": [1, 1, 1],
+ "vec_in_dim": None,
+ }
+
+ sampling_settings = {
+ "shift": 4.0,
+ }
+
+ latent_format = latent_formats.HunyuanImage21Refiner
+
+ def get_model(self, state_dict, prefix="", device=None):
+ out = model_base.HunyuanImage21Refiner(self, device=device)
+ return out
+
+models = [LotusD, Stable_Zero123, SD15_instructpix2pix, SD15, SD20, SD21UnclipL, SD21UnclipH, SDXL_instructpix2pix, SDXLRefiner, SDXL, SSD1B, KOALA_700M, KOALA_1B, Segmind_Vega, SD_X4Upscaler, Stable_Cascade_C, Stable_Cascade_B, SV3D_u, SV3D_p, SD3, StableAudio, AuraFlow, PixArtAlpha, PixArtSigma, HunyuanDiT, HunyuanDiT1, FluxInpaint, Flux, FluxSchnell, GenmoMochi, LTXV, HunyuanImage21Refiner, HunyuanImage21, HunyuanVideoSkyreelsI2V, HunyuanVideoI2V, HunyuanVideo, CosmosT2V, CosmosI2V, CosmosT2IPredict2, CosmosI2VPredict2, Lumina2, WAN22_T2V, WAN21_T2V, WAN21_I2V, WAN21_FunControl2V, WAN21_Vace, WAN21_Camera, WAN22_Camera, WAN22_S2V, WAN21_HuMo, WAN22_Animate, Hunyuan3Dv2mini, Hunyuan3Dv2, Hunyuan3Dv2_1, HiDream, Chroma, ChromaRadiance, ACEStep, Omnigen2, QwenImage]
models += [SVD_img2vid]
diff --git a/comfy/supported_models_base.py b/comfy/supported_models_base.py
index 54573abb1..e4bd74514 100644
--- a/comfy/supported_models_base.py
+++ b/comfy/supported_models_base.py
@@ -50,6 +50,7 @@ class BASE:
manual_cast_dtype = None
custom_operations = None
scaled_fp8 = None
+ layer_quant_config = None # Per-layer quantization configuration for mixed precision
optimizations = {"fp8": False}
@classmethod
diff --git a/comfy/text_encoders/bert.py b/comfy/text_encoders/bert.py
index 551b03162..ed4638a9a 100644
--- a/comfy/text_encoders/bert.py
+++ b/comfy/text_encoders/bert.py
@@ -116,7 +116,7 @@ class BertModel_(torch.nn.Module):
self.embeddings = BertEmbeddings(config_dict["vocab_size"], config_dict["max_position_embeddings"], config_dict["type_vocab_size"], config_dict["pad_token_id"], embed_dim, layer_norm_eps, dtype, device, operations)
self.encoder = BertEncoder(config_dict["num_hidden_layers"], embed_dim, config_dict["intermediate_size"], config_dict["num_attention_heads"], layer_norm_eps, dtype, device, operations)
- def forward(self, input_tokens, attention_mask=None, embeds=None, num_tokens=None, intermediate_output=None, final_layer_norm_intermediate=True, dtype=None):
+ def forward(self, input_tokens, attention_mask=None, embeds=None, num_tokens=None, intermediate_output=None, final_layer_norm_intermediate=True, dtype=None, embeds_info=[]):
x = self.embeddings(input_tokens, embeds=embeds, dtype=dtype)
mask = None
if attention_mask is not None:
diff --git a/comfy/text_encoders/byt5_config_small_glyph.json b/comfy/text_encoders/byt5_config_small_glyph.json
new file mode 100644
index 000000000..0239c7164
--- /dev/null
+++ b/comfy/text_encoders/byt5_config_small_glyph.json
@@ -0,0 +1,22 @@
+{
+ "d_ff": 3584,
+ "d_kv": 64,
+ "d_model": 1472,
+ "decoder_start_token_id": 0,
+ "dropout_rate": 0.1,
+ "eos_token_id": 1,
+ "dense_act_fn": "gelu_pytorch_tanh",
+ "initializer_factor": 1.0,
+ "is_encoder_decoder": true,
+ "is_gated_act": true,
+ "layer_norm_epsilon": 1e-06,
+ "model_type": "t5",
+ "num_decoder_layers": 4,
+ "num_heads": 6,
+ "num_layers": 12,
+ "output_past": true,
+ "pad_token_id": 0,
+ "relative_attention_num_buckets": 32,
+ "tie_word_embeddings": false,
+ "vocab_size": 1510
+}
diff --git a/comfy/text_encoders/byt5_tokenizer/added_tokens.json b/comfy/text_encoders/byt5_tokenizer/added_tokens.json
new file mode 100644
index 000000000..93c190b56
--- /dev/null
+++ b/comfy/text_encoders/byt5_tokenizer/added_tokens.json
@@ -0,0 +1,127 @@
+{
+ "": 259,
+ "": 359,
+ "": 360,
+ "": 361,
+ "": 362,
+ "": 363,
+ "": 364,
+ "": 365,
+ "": 366,
+ "": 367,
+ "": 368,
+ "": 269,
+ "": 369,
+ "": 370,
+ "": 371,
+ "": 372,
+ "": 373,
+ "": 374,
+ "": 375,
+ "": 376,
+ "": 377,
+ "": 378,
+ "": 270,
+ "": 379,
+ "": 380,
+ "": 381,
+ "": 382,
+ "": 383,
+ "": 271,
+ "": 272,
+ "": 273,
+ "": 274,
+ "": 275,
+ "": 276,
+ "": 277,
+ "": 278,
+ "": 260,
+ "": 279,
+ "": 280,
+ "": 281,
+ "": 282,
+ "": 283,
+ "": 284,
+ "": 285,
+ "": 286,
+ "": 287,
+ "": 288,
+ "": 261,
+ "": 289,
+ "": 290,
+ "": 291,
+ "": 292,
+ "": 293,
+ "": 294,
+ "": 295,
+ "": 296,
+ "": 297,
+ "": 298,
+ "": 262,
+ "": 299,
+ "": 300,
+ "": 301,
+ "": 302,
+ "": 303,
+ "": 304,
+ "": 305,
+ "": 306,
+ "": 307,
+ "": 308,
+ "": 263,
+ "": 309,
+ "": 310,
+ "": 311,
+ "": 312,
+ "": 313,
+ "": 314,
+ "": 315,
+ "": 316,
+ "": 317,
+ "": 318,
+ "": 264,
+ "": 319,
+ "": 320,
+ "": 321,
+ "": 322,
+ "": 323,
+ "": 324,
+ "": 325,
+ "": 326,
+ "": 327,
+ "": 328,
+ "": 265,
+ "": 329,
+ "": 330,
+ "": 331,
+ "": 332,
+ "": 333,
+ "": 334,
+ "": 335,
+ "": 336,
+ "": 337,
+ "": 338,
+ "": 266,
+ "": 339,
+ "": 340,
+ "": 341,
+ "": 342,
+ "": 343,
+ "": 344,
+ "": 345,
+ "": 346,
+ "": 347,
+ "": 348,
+ "": 267,
+ "": 349,
+ "": 350,
+ "": 351,
+ "": 352,
+ "": 353,
+ "": 354,
+ "": 355,
+ "": 356,
+ "": 357,
+ "": 358,
+ "": 268
+}
diff --git a/comfy/text_encoders/byt5_tokenizer/special_tokens_map.json b/comfy/text_encoders/byt5_tokenizer/special_tokens_map.json
new file mode 100644
index 000000000..04fd58b5f
--- /dev/null
+++ b/comfy/text_encoders/byt5_tokenizer/special_tokens_map.json
@@ -0,0 +1,150 @@
+{
+ "additional_special_tokens": [
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ "",
+ ""
+ ],
+ "eos_token": {
+ "content": "",
+ "lstrip": false,
+ "normalized": true,
+ "rstrip": false,
+ "single_word": false
+ },
+ "pad_token": {
+ "content": "",
+ "lstrip": false,
+ "normalized": true,
+ "rstrip": false,
+ "single_word": false
+ },
+ "unk_token": {
+ "content": "