Merge branch 'master' into execution_model_inversion

.github/ISSUE_TEMPLATE/bug-report.yml

@@ -0,0 +1,48 @@
+name: Bug Report
+description: "Something is broken inside of ComfyUI. (Do not use this if you're just having issues and need help, or if the issue relates to a custom node)"
+labels: ["Potential Bug"]
+body:
+  - type: markdown
+    attributes:
+      value: |
+        Before submitting a **Bug Report**, please ensure the following:
+
+        - **1:** You are running the latest version of ComfyUI.
+        - **2:** You have looked at the existing bug reports and made sure this isn't already reported.
+        - **3:** You confirmed that the bug is not caused by a custom node. You can disable all custom nodes by passing
+        `--disable-all-custom-nodes` command line argument.
+        - **4:** This is an actual bug in ComfyUI, not just a support question. A bug is when you can specify exact
+        steps to replicate what went wrong and others will be able to repeat your steps and see the same issue happen.
+
+        If unsure, ask on the [ComfyUI Matrix Space](https://app.element.io/#/room/%23comfyui_space%3Amatrix.org) or the [Comfy Org Discord](https://discord.gg/comfyorg) first.
+  - type: textarea
+    attributes:
+      label: Expected Behavior
+      description: "What you expected to happen."
+    validations:
+      required: true
+  - type: textarea
+    attributes:
+      label: Actual Behavior
+      description: "What actually happened. Please include a screenshot of the issue if possible."
+    validations:
+      required: true
+  - type: textarea
+    attributes:
+      label: Steps to Reproduce
+      description: "Describe how to reproduce the issue. Please be sure to attach a workflow JSON or PNG, ideally one that doesn't require custom nodes to test. If the bug open happens when certain custom nodes are used, most likely that custom node is what has the bug rather than ComfyUI, in which case it should be reported to the node's author."
+    validations:
+      required: true
+  - type: textarea
+    attributes:
+      label: Debug Logs
+      description: "Please copy the output from your terminal logs here."
+      render: powershell
+    validations:
+      required: true
+  - type: textarea
+    attributes:
+      label: Other
+      description: "Any other additional information you think might be helpful."
+    validations:
+      required: false

.github/ISSUE_TEMPLATE/config.yml

@@ -0,0 +1,8 @@
+blank_issues_enabled: true
+contact_links:
+  - name: ComfyUI Matrix Space
+    url: https://app.element.io/#/room/%23comfyui_space%3Amatrix.org
+    about: The ComfyUI Matrix Space is available for support and general discussion related to ComfyUI (Matrix is like Discord but open source).
+  - name: Comfy Org Discord
+    url: https://discord.gg/comfyorg
+    about: The Comfy Org Discord is available for support and general discussion related to ComfyUI.

.github/ISSUE_TEMPLATE/feature-request.yml

@@ -0,0 +1,32 @@
+name: Feature Request
+description: "You have an idea for something new you would like to see added to ComfyUI's core."
+labels: [ "Feature" ]
+body:
+    - type: markdown
+      attributes:
+        value: |
+                Before submitting a **Feature Request**, please ensure the following:
+
+                **1:** You are running the latest version of ComfyUI.
+                **2:** You have looked to make sure there is not already a feature that does what you need, and there is not already a Feature Request listed for the same idea.
+                **3:** This is something that makes sense to add to ComfyUI Core, and wouldn't make more sense as a custom node.
+
+                If unsure, ask on the [ComfyUI Matrix Space](https://app.element.io/#/room/%23comfyui_space%3Amatrix.org) or the [Comfy Org Discord](https://discord.gg/comfyorg) first.
+    - type: textarea
+      attributes:
+            label: Feature Idea
+            description: "Describe the feature you want to see."
+      validations:
+            required: true
+    - type: textarea
+      attributes:
+                label: Existing Solutions
+                description: "Please search through available custom nodes / extensions to see if there are existing custom solutions for this. If so, please link the options you found here as a reference."
+      validations:
+                required: false
+    - type: textarea
+      attributes:
+                label: Other
+                description: "Any other additional information you think might be helpful."
+      validations:
+                required: false

.github/ISSUE_TEMPLATE/user-support.yml

@@ -0,0 +1,32 @@
+name: User Support
+description: "Use this if you need help with something, or you're experiencing an issue."
+labels: [ "User Support" ]
+body:
+    - type: markdown
+      attributes:
+        value: |
+            Before submitting a **User Report** issue, please ensure the following:
+
+            **1:** You are running the latest version of ComfyUI.
+            **2:** You have made an effort to find public answers to your question before asking here. In other words, you googled it first, and scrolled through recent help topics.
+
+                If unsure, ask on the [ComfyUI Matrix Space](https://app.element.io/#/room/%23comfyui_space%3Amatrix.org) or the [Comfy Org Discord](https://discord.gg/comfyorg) first.
+    - type: textarea
+      attributes:
+            label: Your question
+            description: "Post your question here. Please be as detailed as possible."
+      validations:
+            required: true
+    - type: textarea
+      attributes:
+                label: Logs
+                description: "If your question relates to an issue you're experiencing, please go to `Server` -> `Logs` -> potentially set `View Type` to `Debug` as well, then copypaste all the text into here."
+                render: powershell
+      validations:
+                required: false
+    - type: textarea
+      attributes:
+                label: Other
+                description: "Any other additional information you think might be helpful."
+      validations:
+                required: false

.github/workflows/pylint.yml

@@ -0,0 +1,23 @@
+name: Python Linting
+
+on: [push, pull_request]
+
+jobs:
+  pylint:
+    name: Run Pylint
+    runs-on: ubuntu-latest
+
+    steps:
+    - name: Checkout repository
+      uses: actions/checkout@v4
+
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: 3.x
+
+    - name: Install Pylint
+      run: pip install pylint
+
+    - name: Run Pylint
+      run: pylint --rcfile=.pylintrc $(find . -type f -name "*.py")

.github/workflows/stable-release.yml

@@ -0,0 +1,109 @@
+
+name: "Release Stable Version"
+
+on:
+  push:
+    tags:
+      - 'v*'
+
+jobs:
+  package_comfy_windows:
+    permissions:
+      contents: "write"
+      packages: "write"
+      pull-requests: "read"
+    runs-on: windows-latest
+    strategy:
+      matrix:
+        python_version: [3.11.8]
+        cuda_version: [121]
+    steps:
+      - name: Calculate Minor Version
+        shell: bash
+        run: |
+          # Extract the minor version from the Python version
+          MINOR_VERSION=$(echo "${{ matrix.python_version }}" | cut -d'.' -f2)
+          echo "MINOR_VERSION=$MINOR_VERSION" >> $GITHUB_ENV
+      - name: Setup Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python_version }}
+        
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+          persist-credentials: false
+      - 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 torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu${{ matrix.cuda_version }} -r ../ComfyUI/requirements.txt pygit2
+          pause" > update_comfyui_and_python_dependencies.bat
+
+          python -m pip wheel --no-cache-dir torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu${{ matrix.cuda_version }} -r requirements.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 cu${{ matrix.cuda_version }}_python_deps
+          mv cu${{ matrix.cuda_version }}_python_deps ../
+          mv update_comfyui_and_python_dependencies.bat ../
+          cd ..
+          pwd
+          ls
+          
+          cp -r ComfyUI ComfyUI_copy
+          curl https://www.python.org/ftp/python/${{ matrix.python_version }}/python-${{ matrix.python_version }}-embed-amd64.zip -o python_embeded.zip
+          unzip python_embeded.zip -d python_embeded
+          cd python_embeded
+          echo ${{ env.MINOR_VERSION }}
+          echo 'import site' >> ./python3${{ env.MINOR_VERSION }}._pth
+          curl https://bootstrap.pypa.io/get-pip.py -o get-pip.py
+          ./python.exe get-pip.py
+          ./python.exe --version
+          echo "Pip version:"
+          ./python.exe -m pip --version
+
+          set PATH=$PWD/Scripts:$PATH
+          echo $PATH
+          ./python.exe -s -m pip install ../cu${{ matrix.cuda_version }}_python_deps/*
+          sed -i '1i../ComfyUI' ./python3${{ env.MINOR_VERSION }}._pth
+          cd ..
+
+          git clone https://github.com/comfyanonymous/taesd
+          cp taesd/*.pth ./ComfyUI_copy/models/vae_approx/
+
+          mkdir ComfyUI_windows_portable
+          mv python_embeded ComfyUI_windows_portable
+          mv ComfyUI_copy ComfyUI_windows_portable/ComfyUI
+
+          cd ComfyUI_windows_portable
+
+          mkdir update
+          cp -r ComfyUI/.ci/update_windows/* ./update/
+          cp -r ComfyUI/.ci/windows_base_files/* ./
+          cp ../update_comfyui_and_python_dependencies.bat ./update/
+
+          cd ..
+
+          "C:\Program Files\7-Zip\7z.exe" a -t7z -m0=lzma2 -mx=8 -mfb=64 -md=32m -ms=on -mf=BCJ2 ComfyUI_windows_portable.7z ComfyUI_windows_portable
+          mv ComfyUI_windows_portable.7z ComfyUI/ComfyUI_windows_portable_nvidia.7z
+
+          cd ComfyUI_windows_portable
+          python_embeded/python.exe -s ComfyUI/main.py --quick-test-for-ci --cpu
+
+          ls
+
+      - name: Upload binaries to release
+        uses: svenstaro/upload-release-action@v2
+        with:
+          repo_token: ${{ secrets.GITHUB_TOKEN }}
+          file: ComfyUI_windows_portable_nvidia.7z
+          tag: ${{ github.ref }}
+          overwrite: true
+        

.github/workflows/test-browser.yml

@@ -0,0 +1,76 @@
+# This is a temporary action during frontend TS migration.
+# This file should be removed after TS migration is completed.
+# The browser test is here to ensure TS repo is working the same way as the
+# current JS code.
+# If you are adding UI feature, please sync your changes to the TS repo:
+# huchenlei/ComfyUI_frontend and update test expectation files accordingly.
+name: Playwright Browser Tests CI
+
+on:
+  push:
+    branches: [ main, master ]
+  pull_request:
+    branches: [ main, master ]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+    - name: Checkout ComfyUI
+      uses: actions/checkout@v4
+      with:
+        repository: "comfyanonymous/ComfyUI"
+        path: "ComfyUI"
+    - name: Checkout ComfyUI_frontend
+      uses: actions/checkout@v4
+      with:
+        repository: "huchenlei/ComfyUI_frontend"
+        path: "ComfyUI_frontend"
+        ref: "fcc54d803e5b6a9b08a462a1d94899318c96dcbb"
+    - uses: actions/setup-node@v3
+      with:
+        node-version: lts/*
+    - uses: actions/setup-python@v4
+      with:
+        python-version: '3.10'
+    - 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 wait-for-it
+      working-directory: ComfyUI
+    - name: Start ComfyUI server
+      run: |
+        python main.py --cpu 2>&1 | tee console_output.log &
+        wait-for-it --service 127.0.0.1:8188 -t 600
+      working-directory: ComfyUI
+    - name: Install ComfyUI_frontend dependencies
+      run: |
+        npm ci
+      working-directory: ComfyUI_frontend
+    - name: Install Playwright Browsers
+      run: npx playwright install --with-deps
+      working-directory: ComfyUI_frontend
+    - name: Run Playwright tests
+      run: npx playwright test
+      working-directory: ComfyUI_frontend
+    - name: Check for unhandled exceptions in server log
+      run: |
+        if grep -qE "Exception|Error" console_output.log; then
+          echo "Unhandled exception/error found in server log."
+          exit 1
+        fi
+      working-directory: ComfyUI
+    - uses: actions/upload-artifact@v4
+      if: always()
+      with:
+        name: playwright-report
+        path: ComfyUI_frontend/playwright-report/
+        retention-days: 30
+    - uses: actions/upload-artifact@v4
+      if: always()
+      with:
+        name: console-output
+        path: ComfyUI/console_output.log
+        retention-days: 30

.github/workflows/test-ui.yaml

@@ -24,3 +24,7 @@ jobs:
         npm run test:generate
         npm test -- --verbose
       working-directory: ./tests-ui
+    - name: Run Unit Tests
+      run: |
+        pip install -r tests-unit/requirements.txt
+        python -m pytest tests-unit

.github/workflows/windows_release_dependencies.yml

@@ -33,8 +33,8 @@ jobs:
   build_dependencies:
     runs-on: windows-latest
     steps:
-        - uses: actions/checkout@v3
+        - uses: actions/checkout@v4
-        - uses: actions/setup-python@v4
+        - uses: actions/setup-python@v5
           with:
             python-version: 3.${{ inputs.python_minor }}.${{ inputs.python_patch }}
 
@@ -58,7 +58,7 @@ jobs:
             mv temp_wheel_dir cu${{ inputs.cu }}_python_deps
             tar cf cu${{ inputs.cu }}_python_deps.tar cu${{ inputs.cu }}_python_deps
 
-        - uses: actions/cache/save@v3
+        - uses: actions/cache/save@v4
           with:
             path: |
               cu${{ inputs.cu }}_python_deps.tar

.github/workflows/windows_release_nightly_pytorch.yml

@@ -32,11 +32,11 @@ jobs:
         pull-requests: "read"
     runs-on: windows-latest
     steps:
-        - uses: actions/checkout@v3
+        - uses: actions/checkout@v4
           with:
             fetch-depth: 0
             persist-credentials: false
-        - uses: actions/setup-python@v4
+        - uses: actions/setup-python@v5
           with:
             python-version: 3.${{ inputs.python_minor }}.${{ inputs.python_patch }}
         - shell: bash
@@ -73,7 +73,7 @@ jobs:
             pause" > ./update/update_comfyui_and_python_dependencies.bat
             cd ..
 
-            "C:\Program Files\7-Zip\7z.exe" a -t7z -m0=lzma -mx=8 -mfb=64 -md=32m -ms=on -mf=BCJ2 ComfyUI_windows_portable_nightly_pytorch.7z ComfyUI_windows_portable_nightly_pytorch
+            "C:\Program Files\7-Zip\7z.exe" a -t7z -m0=lzma2 -mx=8 -mfb=64 -md=32m -ms=on -mf=BCJ2 ComfyUI_windows_portable_nightly_pytorch.7z ComfyUI_windows_portable_nightly_pytorch
             mv ComfyUI_windows_portable_nightly_pytorch.7z ComfyUI/ComfyUI_windows_portable_nvidia_or_cpu_nightly_pytorch.7z
 
             cd ComfyUI_windows_portable_nightly_pytorch

.github/workflows/windows_release_package.yml

@@ -32,7 +32,7 @@ jobs:
         pull-requests: "read"
     runs-on: windows-latest
     steps:
-        - uses: actions/cache/restore@v3
+        - uses: actions/cache/restore@v4
           id: cache
           with:
             path: |
@@ -48,7 +48,7 @@ jobs:
             pwd
             ls
 
-        - uses: actions/checkout@v3
+        - uses: actions/checkout@v4
           with:
             fetch-depth: 0
             persist-credentials: false
@@ -82,7 +82,7 @@ jobs:
 
             cd ..
 
-            "C:\Program Files\7-Zip\7z.exe" a -t7z -m0=lzma -mx=8 -mfb=64 -md=32m -ms=on -mf=BCJ2 ComfyUI_windows_portable.7z ComfyUI_windows_portable
+            "C:\Program Files\7-Zip\7z.exe" a -t7z -m0=lzma2 -mx=8 -mfb=64 -md=32m -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

.gitignore

@@ -9,6 +9,7 @@ __pycache__/
 !custom_nodes/example_node.py.example
 extra_model_paths.yaml
 /.vs
+.vscode/
 .idea/
 venv/
 /web/extensions/*
@@ -16,3 +17,5 @@ venv/
 !/web/extensions/core/
 /tests-ui/data/object_info.json
 /user/
+*.log
+web_custom_versions/

.pylintrc

@@ -0,0 +1,3 @@
+[MESSAGES CONTROL]
+disable=all
+enable=eval-used

CONTRIBUTING.md

@@ -0,0 +1,41 @@
+# Contributing to ComfyUI
+
+Welcome, and thank you for your interest in contributing to ComfyUI!
+
+There are several ways in which you can contribute, beyond writing code. The goal of this document is to provide a high-level overview of how you can get involved.
+
+## Asking Questions
+
+Have a question? Instead of opening an issue, please ask on [Discord](https://comfy.org/discord) or [Matrix](https://app.element.io/#/room/%23comfyui_space%3Amatrix.org) channels. Our team and the community will help you.
+
+## Providing Feedback
+
+Your comments and feedback are welcome, and the development team is available via a handful of different channels.
+
+See the `#bug-report`, `#feature-request` and `#feedback` channels on Discord.
+
+## Reporting Issues
+
+Have you identified a reproducible problem in ComfyUI? Do you have a feature request? We want to hear about it! Here's how you can report your issue as effectively as possible.
+
+
+### Look For an Existing Issue
+
+Before you create a new issue, please do a search in [open issues](https://github.com/comfyanonymous/ComfyUI/issues) to see if the issue or feature request has already been filed.
+
+If you find your issue already exists, make relevant comments and add your [reaction](https://github.com/blog/2119-add-reactions-to-pull-requests-issues-and-comments). Use a reaction in place of a "+1" comment:
+
+* 👍 - upvote
+* 👎 - downvote
+
+If you cannot find an existing issue that describes your bug or feature, create a new issue. We have an issue template in place to organize new issues.
+
+
+### Creating Pull Requests
+
+* Please refer to the article on [creating pull requests](https://github.com/comfyanonymous/ComfyUI/wiki/How-to-Contribute-Code) and contributing to this project.
+
+
+## Thank You
+
+Your contributions to open source, large or small, make great projects like this possible. Thank you for taking the time to contribute.

README.md

@@ -11,16 +11,16 @@ This ui will let you design and execute advanced stable diffusion pipelines usin
 
 ## Features
 - Nodes/graph/flowchart interface to experiment and create complex Stable Diffusion workflows without needing to code anything.
-- Fully supports SD1.x, SD2.x, [SDXL](https://comfyanonymous.github.io/ComfyUI_examples/sdxl/), [Stable Video Diffusion](https://comfyanonymous.github.io/ComfyUI_examples/video/), [Stable Cascade](https://comfyanonymous.github.io/ComfyUI_examples/stable_cascade/) and [SD3](https://comfyanonymous.github.io/ComfyUI_examples/sd3/)
+- Fully supports SD1.x, SD2.x, [SDXL](https://comfyanonymous.github.io/ComfyUI_examples/sdxl/), [Stable Video Diffusion](https://comfyanonymous.github.io/ComfyUI_examples/video/), [Stable Cascade](https://comfyanonymous.github.io/ComfyUI_examples/stable_cascade/), [SD3](https://comfyanonymous.github.io/ComfyUI_examples/sd3/) and [Stable Audio](https://comfyanonymous.github.io/ComfyUI_examples/audio/)
 - Asynchronous Queue system
 - Many optimizations: Only re-executes the parts of the workflow that changes between executions.
-- Command line option: ```--lowvram``` to make it work on GPUs with less than 3GB vram (enabled automatically on GPUs with low vram)
+- Smart memory management: can automatically run models on GPUs with as low as 1GB vram.
 - 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.
 - Embeddings/Textual inversion
 - [Loras (regular, locon and loha)](https://comfyanonymous.github.io/ComfyUI_examples/lora/)
 - [Hypernetworks](https://comfyanonymous.github.io/ComfyUI_examples/hypernetworks/)
-- Loading full workflows (with seeds) from generated PNG files.
+- Loading full workflows (with seeds) from generated PNG, WebP and FLAC files.
 - Saving/Loading workflows as Json files.
 - Nodes interface can be used to create complex workflows like one for [Hires fix](https://comfyanonymous.github.io/ComfyUI_examples/2_pass_txt2img/) or much more advanced ones.
 - [Area Composition](https://comfyanonymous.github.io/ComfyUI_examples/area_composition/)
@@ -32,6 +32,7 @@ This ui will let you design and execute advanced stable diffusion pipelines usin
 - [Model Merging](https://comfyanonymous.github.io/ComfyUI_examples/model_merging/)
 - [LCM models and Loras](https://comfyanonymous.github.io/ComfyUI_examples/lcm/)
 - [SDXL Turbo](https://comfyanonymous.github.io/ComfyUI_examples/sdturbo/)
+- [AuraFlow](https://comfyanonymous.github.io/ComfyUI_examples/aura_flow/)
 - Latent previews with [TAESD](#how-to-show-high-quality-previews)
 - Starts up very fast.
 - Works fully offline: will never download anything.
@@ -225,12 +226,11 @@ Use `--tls-keyfile key.pem --tls-certfile cert.pem` to enable TLS/SSL, the app w
 
 [Matrix space: #comfyui_space:matrix.org](https://app.element.io/#/room/%23comfyui_space%3Amatrix.org) (it's like discord but open source).
 
+See also: [https://www.comfy.org/](https://www.comfy.org/)
+
 # QA
 
-### Why did you make this?
+### Which GPU should I buy for this?
 
-I wanted to learn how Stable Diffusion worked in detail. I also wanted something clean and powerful that would let me experiment with SD without restrictions.
+[See this page for some recommendations](https://github.com/comfyanonymous/ComfyUI/wiki/Which-GPU-should-I-buy-for-ComfyUI)
 
-### Who is this for?
-
-This is for anyone that wants to make complex workflows with SD or that wants to learn more how SD works. The interface follows closely how SD works and the code should be much more simple to understand than other SD UIs.

app/frontend_management.py

@@ -0,0 +1,188 @@
+from __future__ import annotations
+import argparse
+import logging
+import os
+import re
+import tempfile
+import zipfile
+from dataclasses import dataclass
+from functools import cached_property
+from pathlib import Path
+from typing import TypedDict
+
+import requests
+from typing_extensions import NotRequired
+from comfy.cli_args import DEFAULT_VERSION_STRING
+
+
+REQUEST_TIMEOUT = 10  # seconds
+
+
+class Asset(TypedDict):
+    url: str
+
+
+class Release(TypedDict):
+    id: int
+    tag_name: str
+    name: str
+    prerelease: bool
+    created_at: str
+    published_at: str
+    body: str
+    assets: NotRequired[list[Asset]]
+
+
+@dataclass
+class FrontEndProvider:
+    owner: str
+    repo: str
+
+    @property
+    def folder_name(self) -> str:
+        return f"{self.owner}_{self.repo}"
+
+    @property
+    def release_url(self) -> str:
+        return f"https://api.github.com/repos/{self.owner}/{self.repo}/releases"
+
+    @cached_property
+    def all_releases(self) -> list[Release]:
+        releases = []
+        api_url = self.release_url
+        while api_url:
+            response = requests.get(api_url, timeout=REQUEST_TIMEOUT)
+            response.raise_for_status()  # Raises an HTTPError if the response was an error
+            releases.extend(response.json())
+            # GitHub uses the Link header to provide pagination links. Check if it exists and update api_url accordingly.
+            if "next" in response.links:
+                api_url = response.links["next"]["url"]
+            else:
+                api_url = None
+        return releases
+
+    @cached_property
+    def latest_release(self) -> Release:
+        latest_release_url = f"{self.release_url}/latest"
+        response = requests.get(latest_release_url, timeout=REQUEST_TIMEOUT)
+        response.raise_for_status()  # Raises an HTTPError if the response was an error
+        return response.json()
+
+    def get_release(self, version: str) -> Release:
+        if version == "latest":
+            return self.latest_release
+        else:
+            for release in self.all_releases:
+                if release["tag_name"] in [version, f"v{version}"]:
+                    return release
+            raise ValueError(f"Version {version} not found in releases")
+
+
+def download_release_asset_zip(release: Release, destination_path: str) -> None:
+    """Download dist.zip from github release."""
+    asset_url = None
+    for asset in release.get("assets", []):
+        if asset["name"] == "dist.zip":
+            asset_url = asset["url"]
+            break
+
+    if not asset_url:
+        raise ValueError("dist.zip not found in the release assets")
+
+    # Use a temporary file to download the zip content
+    with tempfile.TemporaryFile() as tmp_file:
+        headers = {"Accept": "application/octet-stream"}
+        response = requests.get(
+            asset_url, headers=headers, allow_redirects=True, timeout=REQUEST_TIMEOUT
+        )
+        response.raise_for_status()  # Ensure we got a successful response
+
+        # Write the content to the temporary file
+        tmp_file.write(response.content)
+
+        # Go back to the beginning of the temporary file
+        tmp_file.seek(0)
+
+        # Extract the zip file content to the destination path
+        with zipfile.ZipFile(tmp_file, "r") as zip_ref:
+            zip_ref.extractall(destination_path)
+
+
+class FrontendManager:
+    DEFAULT_FRONTEND_PATH = str(Path(__file__).parents[1] / "web")
+    CUSTOM_FRONTENDS_ROOT = str(Path(__file__).parents[1] / "web_custom_versions")
+
+    @classmethod
+    def parse_version_string(cls, value: str) -> tuple[str, str, str]:
+        """
+        Args:
+            value (str): The version string to parse.
+
+        Returns:
+            tuple[str, str]: A tuple containing provider name and version.
+
+        Raises:
+            argparse.ArgumentTypeError: If the version string is invalid.
+        """
+        VERSION_PATTERN = r"^([a-zA-Z0-9][a-zA-Z0-9-]{0,38})/([a-zA-Z0-9_.-]+)@(v?\d+\.\d+\.\d+|latest)$"
+        match_result = re.match(VERSION_PATTERN, value)
+        if match_result is None:
+            raise argparse.ArgumentTypeError(f"Invalid version string: {value}")
+
+        return match_result.group(1), match_result.group(2), match_result.group(3)
+
+    @classmethod
+    def init_frontend_unsafe(cls, version_string: str) -> str:
+        """
+        Initializes the frontend for the specified version.
+
+        Args:
+            version_string (str): The version string.
+
+        Returns:
+            str: The path to the initialized frontend.
+
+        Raises:
+            Exception: If there is an error during the initialization process.
+            main error source might be request timeout or invalid URL.
+        """
+        if version_string == DEFAULT_VERSION_STRING:
+            return cls.DEFAULT_FRONTEND_PATH
+
+        repo_owner, repo_name, version = cls.parse_version_string(version_string)
+        provider = FrontEndProvider(repo_owner, repo_name)
+        release = provider.get_release(version)
+
+        semantic_version = release["tag_name"].lstrip("v")
+        web_root = str(
+            Path(cls.CUSTOM_FRONTENDS_ROOT) / provider.folder_name / semantic_version
+        )
+        if not os.path.exists(web_root):
+            os.makedirs(web_root, exist_ok=True)
+            logging.info(
+                "Downloading frontend(%s) version(%s) to (%s)",
+                provider.folder_name,
+                semantic_version,
+                web_root,
+            )
+            logging.debug(release)
+            download_release_asset_zip(release, destination_path=web_root)
+        return web_root
+
+    @classmethod
+    def init_frontend(cls, version_string: str) -> str:
+        """
+        Initializes the frontend with the specified version string.
+
+        Args:
+            version_string (str): The version string to initialize the frontend with.
+
+        Returns:
+            str: The path of the initialized frontend.
+        """
+        try:
+            return cls.init_frontend_unsafe(version_string)
+        except Exception as e:
+            logging.error("Failed to initialize frontend: %s", e)
+            logging.info("Falling back to the default frontend.")
+            return cls.DEFAULT_FRONTEND_PATH

app/user_manager.py

@@ -2,6 +2,8 @@ import json
 import os
 import re
 import uuid
+import glob
+import shutil
 from aiohttp import web
 from comfy.cli_args import args
 from folder_paths import user_directory
@@ -56,16 +58,16 @@ class UserManager():
         if os.path.commonpath((root_dir, user_root)) != root_dir:
             return None
 
-        parent = user_root
-
         if file is not None:
             # prevent leaving /{type}/{user}
             path = os.path.abspath(os.path.join(user_root, file))
             if os.path.commonpath((user_root, path)) != user_root:
                 return None
 
+        parent = os.path.split(path)[0]
+
         if create_dir and not os.path.exists(parent):
-            os.mkdir(parent)
+            os.makedirs(parent, exist_ok=True)
 
         return path
 
@@ -108,24 +110,32 @@ class UserManager():
             user_id = self.add_user(username)
             return web.json_response(user_id)
 
-        @routes.get("/userdata/{file}")
+        @routes.get("/userdata")
-        async def getuserdata(request):
+        async def listuserdata(request):
-            file = request.match_info.get("file", None)
+            directory = request.rel_url.query.get('dir', '')
-            if not file:
+            if not directory:
                 return web.Response(status=400)
                 
-            path = self.get_request_user_filepath(request, file)
+            path = self.get_request_user_filepath(request, directory)
             if not path:
                 return web.Response(status=403)
             
             if not os.path.exists(path):
                 return web.Response(status=404)
             
-            return web.FileResponse(path)
+            recurse = request.rel_url.query.get('recurse', '').lower() == "true"
+            results = glob.glob(os.path.join(
+                glob.escape(path), '**/*'), recursive=recurse)
+            results = [os.path.relpath(x, path) for x in results if os.path.isfile(x)]
             
-        @routes.post("/userdata/{file}")
+            split_path = request.rel_url.query.get('split', '').lower() == "true"
-        async def post_userdata(request):
+            if split_path:
-            file = request.match_info.get("file", None)
+                results = [[x] + x.split(os.sep) for x in results]
+
+            return web.json_response(results)
+
+        def get_user_data_path(request, check_exists = False, param = "file"):
+            file = request.match_info.get(param, None)
             if not file:
                 return web.Response(status=400)
                 
@@ -133,8 +143,63 @@ class UserManager():
             if not path:
                 return web.Response(status=403)
             
+            if check_exists and not os.path.exists(path):
+                return web.Response(status=404)
+            
+            return path
+
+        @routes.get("/userdata/{file}")
+        async def getuserdata(request):
+            path = get_user_data_path(request, check_exists=True)
+            if not isinstance(path, str):
+                return path
+            
+            return web.FileResponse(path)
+
+        @routes.post("/userdata/{file}")
+        async def post_userdata(request):
+            path = get_user_data_path(request)
+            if not isinstance(path, str):
+                return path
+            
+            overwrite = request.query["overwrite"] != "false"
+            if not overwrite and os.path.exists(path):
+                return web.Response(status=409)
+
             body = await request.read()
+
             with open(path, "wb") as f:
                 f.write(body)
                 
-            return web.Response(status=200)
+            resp = os.path.relpath(path, self.get_request_user_filepath(request, None))
+            return web.json_response(resp)
+
+        @routes.delete("/userdata/{file}")
+        async def delete_userdata(request):
+            path = get_user_data_path(request, check_exists=True)
+            if not isinstance(path, str):
+                return path
+
+            os.remove(path)
+                
+            return web.Response(status=204)
+
+        @routes.post("/userdata/{file}/move/{dest}")
+        async def move_userdata(request):
+            source = get_user_data_path(request, check_exists=True)
+            if not isinstance(source, str):
+                return source
+            
+            dest = get_user_data_path(request, check_exists=False, param="dest")
+            if not isinstance(source, str):
+                return dest
+            
+            overwrite = request.query["overwrite"] != "false"
+            if not overwrite and os.path.exists(dest):
+                return web.Response(status=409)
+
+            print(f"moving '{source}' -> '{dest}'")
+            shutil.move(source, dest)
+                
+            resp = os.path.relpath(dest, self.get_request_user_filepath(request, None))
+            return web.json_response(resp)

comfy/cldm/cldm.py

@@ -13,7 +13,46 @@ from ..ldm.modules.diffusionmodules.util import (
 from ..ldm.modules.attention import SpatialTransformer
 from ..ldm.modules.diffusionmodules.openaimodel import UNetModel, TimestepEmbedSequential, ResBlock, Downsample
 from ..ldm.util import exists
+from collections import OrderedDict
 import comfy.ops
+from comfy.ldm.modules.attention import optimized_attention
+
+class OptimizedAttention(nn.Module):
+    def __init__(self, c, nhead, dropout=0.0, dtype=None, device=None, operations=None):
+        super().__init__()
+        self.heads = nhead
+        self.c = c
+
+        self.in_proj = operations.Linear(c, c * 3, bias=True, dtype=dtype, device=device)
+        self.out_proj = operations.Linear(c, c, bias=True, dtype=dtype, device=device)
+
+    def forward(self, x):
+        x = self.in_proj(x)
+        q, k, v = x.split(self.c, dim=2)
+        out = optimized_attention(q, k, v, self.heads)
+        return self.out_proj(out)
+
+class QuickGELU(nn.Module):
+    def forward(self, x: torch.Tensor):
+        return x * torch.sigmoid(1.702 * x)
+
+class ResBlockUnionControlnet(nn.Module):
+    def __init__(self, dim, nhead, dtype=None, device=None, operations=None):
+        super().__init__()
+        self.attn = OptimizedAttention(dim, nhead, dtype=dtype, device=device, operations=operations)
+        self.ln_1 = operations.LayerNorm(dim, dtype=dtype, device=device)
+        self.mlp = nn.Sequential(
+            OrderedDict([("c_fc", operations.Linear(dim, dim * 4, dtype=dtype, device=device)), ("gelu", QuickGELU()),
+                         ("c_proj", operations.Linear(dim * 4, dim, dtype=dtype, device=device))]))
+        self.ln_2 = operations.LayerNorm(dim, dtype=dtype, device=device)
+
+    def attention(self, x: torch.Tensor):
+        return self.attn(x)
+
+    def forward(self, x: torch.Tensor):
+        x = x + self.attention(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
 
 class ControlledUnetModel(UNetModel):
     #implemented in the ldm unet
@@ -53,6 +92,7 @@ class ControlNet(nn.Module):
         transformer_depth_middle=None,
         transformer_depth_output=None,
         attn_precision=None,
+        union_controlnet_num_control_type=None,
         device=None,
         operations=comfy.ops.disable_weight_init,
         **kwargs,
@@ -280,6 +320,65 @@ class ControlNet(nn.Module):
         self.middle_block_out = self.make_zero_conv(ch, operations=operations, dtype=self.dtype, device=device)
         self._feature_size += ch
 
+        if union_controlnet_num_control_type is not None:
+            self.num_control_type = union_controlnet_num_control_type
+            num_trans_channel = 320
+            num_trans_head = 8
+            num_trans_layer = 1
+            num_proj_channel = 320
+            # task_scale_factor = num_trans_channel ** 0.5
+            self.task_embedding = nn.Parameter(torch.empty(self.num_control_type, num_trans_channel, dtype=self.dtype, device=device))
+
+            self.transformer_layes = nn.Sequential(*[ResBlockUnionControlnet(num_trans_channel, num_trans_head, dtype=self.dtype, device=device, operations=operations) for _ in range(num_trans_layer)])
+            self.spatial_ch_projs = operations.Linear(num_trans_channel, num_proj_channel, dtype=self.dtype, device=device)
+            #-----------------------------------------------------------------------------------------------------
+
+            control_add_embed_dim = 256
+            class ControlAddEmbedding(nn.Module):
+                def __init__(self, in_dim, out_dim, num_control_type, dtype=None, device=None, operations=None):
+                    super().__init__()
+                    self.num_control_type = num_control_type
+                    self.in_dim = in_dim
+                    self.linear_1 = operations.Linear(in_dim * num_control_type, out_dim, dtype=dtype, device=device)
+                    self.linear_2 = operations.Linear(out_dim, out_dim, dtype=dtype, device=device)
+                def forward(self, control_type, dtype, device):
+                    c_type = torch.zeros((self.num_control_type,), device=device)
+                    c_type[control_type] = 1.0
+                    c_type = timestep_embedding(c_type.flatten(), self.in_dim, repeat_only=False).to(dtype).reshape((-1, self.num_control_type * self.in_dim))
+                    return self.linear_2(torch.nn.functional.silu(self.linear_1(c_type)))
+
+            self.control_add_embedding = ControlAddEmbedding(control_add_embed_dim, time_embed_dim, self.num_control_type, dtype=self.dtype, device=device, operations=operations)
+        else:
+            self.task_embedding = None
+            self.control_add_embedding = None
+
+    def union_controlnet_merge(self, hint, control_type, emb, context):
+        # Equivalent to: https://github.com/xinsir6/ControlNetPlus/tree/main
+        inputs = []
+        condition_list = []
+
+        for idx in range(min(1, len(control_type))):
+            controlnet_cond = self.input_hint_block(hint[idx], emb, context)
+            feat_seq = torch.mean(controlnet_cond, dim=(2, 3))
+            if idx < len(control_type):
+                feat_seq += self.task_embedding[control_type[idx]].to(dtype=feat_seq.dtype, device=feat_seq.device)
+
+            inputs.append(feat_seq.unsqueeze(1))
+            condition_list.append(controlnet_cond)
+
+        x = torch.cat(inputs, dim=1)
+        x = self.transformer_layes(x)
+        controlnet_cond_fuser = None
+        for idx in range(len(control_type)):
+            alpha = self.spatial_ch_projs(x[:, idx])
+            alpha = alpha.unsqueeze(-1).unsqueeze(-1)
+            o = condition_list[idx] + alpha
+            if controlnet_cond_fuser is None:
+                controlnet_cond_fuser = o
+            else:
+                controlnet_cond_fuser += o
+        return controlnet_cond_fuser
+
     def make_zero_conv(self, channels, operations=None, dtype=None, device=None):
         return TimestepEmbedSequential(operations.conv_nd(self.dims, channels, channels, 1, padding=0, dtype=dtype, device=device))
 
@@ -287,9 +386,21 @@ class ControlNet(nn.Module):
         t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype)
         emb = self.time_embed(t_emb)
 
-        guided_hint = self.input_hint_block(hint, emb, context)
+        guided_hint = None
+        if self.control_add_embedding is not None: #Union Controlnet
+            control_type = kwargs.get("control_type", [])
 
-        outs = []
+            emb += self.control_add_embedding(control_type, emb.dtype, emb.device)
+            if len(control_type) > 0:
+                if len(hint.shape) < 5:
+                    hint = hint.unsqueeze(dim=0)
+                guided_hint = self.union_controlnet_merge(hint, control_type, emb, context)
+
+        if guided_hint is None:
+            guided_hint = self.input_hint_block(hint, emb, context)
+
+        out_output = []
+        out_middle = []
 
         hs = []
         if self.num_classes is not None:
@@ -304,10 +415,10 @@ class ControlNet(nn.Module):
                 guided_hint = None
             else:
                 h = module(h, emb, context)
-            outs.append(zero_conv(h, emb, context))
+            out_output.append(zero_conv(h, emb, context))
 
         h = self.middle_block(h, emb, context)
-        outs.append(self.middle_block_out(h, emb, context))
+        out_middle.append(self.middle_block_out(h, emb, context))
 
-        return outs
+        return {"middle": out_middle, "output": out_output}
 

comfy/cldm/mmdit.py

@@ -0,0 +1,77 @@
+import torch
+from typing import Dict, Optional
+import comfy.ldm.modules.diffusionmodules.mmdit
+
+class ControlNet(comfy.ldm.modules.diffusionmodules.mmdit.MMDiT):
+    def __init__(
+        self,
+        num_blocks = None,
+        dtype = None,
+        device = None,
+        operations = None,
+        **kwargs,
+    ):
+        super().__init__(dtype=dtype, device=device, operations=operations, final_layer=False, num_blocks=num_blocks, **kwargs)
+        # controlnet_blocks
+        self.controlnet_blocks = torch.nn.ModuleList([])
+        for _ in range(len(self.joint_blocks)):
+            self.controlnet_blocks.append(operations.Linear(self.hidden_size, self.hidden_size, device=device, dtype=dtype))
+
+        self.pos_embed_input = comfy.ldm.modules.diffusionmodules.mmdit.PatchEmbed(
+            None,
+            self.patch_size,
+            self.in_channels,
+            self.hidden_size,
+            bias=True,
+            strict_img_size=False,
+            dtype=dtype,
+            device=device,
+            operations=operations
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        timesteps: torch.Tensor,
+        y: Optional[torch.Tensor] = None,
+        context: Optional[torch.Tensor] = None,
+        hint = None,
+    ) -> torch.Tensor:
+
+        #weird sd3 controlnet specific stuff
+        y = torch.zeros_like(y)
+
+        if self.context_processor is not None:
+            context = self.context_processor(context)
+
+        hw = x.shape[-2:]
+        x = self.x_embedder(x) + self.cropped_pos_embed(hw, device=x.device).to(dtype=x.dtype, device=x.device)
+        x += self.pos_embed_input(hint)
+
+        c = self.t_embedder(timesteps, dtype=x.dtype)
+        if y is not None and self.y_embedder is not None:
+            y = self.y_embedder(y)
+            c = c + y
+
+        if context is not None:
+            context = self.context_embedder(context)
+
+        output = []
+
+        blocks = len(self.joint_blocks)
+        for i in range(blocks):
+            context, x = self.joint_blocks[i](
+                context,
+                x,
+                c=c,
+                use_checkpoint=self.use_checkpoint,
+            )
+
+            out = self.controlnet_blocks[i](x)
+            count = self.depth // blocks
+            if i == blocks - 1:
+                count -= 1
+            for j in range(count):
+                output.append(out)
+
+        return {"output": output}

comfy/cli_args.py

@@ -1,7 +1,10 @@
 import argparse
 import enum
+import os
+from typing import Optional
 import comfy.options
 
+
 class EnumAction(argparse.Action):
     """
     Argparse action for handling Enums
@@ -113,6 +116,7 @@ vram_group.add_argument("--lowvram", action="store_true", help="Split the unet i
 vram_group.add_argument("--novram", action="store_true", help="When lowvram isn't enough.")
 vram_group.add_argument("--cpu", action="store_true", help="To use the CPU for everything (slow).")
 
+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.")
 parser.add_argument("--deterministic", action="store_true", help="Make pytorch use slower deterministic algorithms when it can. Note that this might not make images deterministic in all cases.")
@@ -122,11 +126,44 @@ parser.add_argument("--quick-test-for-ci", action="store_true", help="Quick test
 parser.add_argument("--windows-standalone-build", action="store_true", help="Windows standalone build: Enable convenient things that most people using the standalone windows build will probably enjoy (like auto opening the page on startup).")
 
 parser.add_argument("--disable-metadata", action="store_true", help="Disable saving prompt metadata in files.")
+parser.add_argument("--disable-all-custom-nodes", action="store_true", help="Disable loading all custom nodes.")
 
 parser.add_argument("--multi-user", action="store_true", help="Enables per-user storage.")
 
 parser.add_argument("--verbose", action="store_true", help="Enables more debug prints.")
 
+# The default built-in provider hosted under web/
+DEFAULT_VERSION_STRING = "comfyanonymous/ComfyUI@latest"
+
+parser.add_argument(
+    "--front-end-version",
+    type=str,
+    default=DEFAULT_VERSION_STRING,
+    help="""
+    Specifies the version of the frontend to be used. This command needs internet connectivity to query and
+    download available frontend implementations from GitHub releases.
+
+    The version string should be in the format of:
+    [repoOwner]/[repoName]@[version]
+    where version is one of: "latest" or a valid version number (e.g. "1.0.0")
+    """,
+)
+
+def is_valid_directory(path: Optional[str]) -> Optional[str]:
+    """Validate if the given path is a directory."""
+    if path is None:
+        return None
+
+    if not os.path.isdir(path):
+        raise argparse.ArgumentTypeError(f"{path} is not a valid directory.")
+    return path
+
+parser.add_argument(
+    "--front-end-root",
+    type=is_valid_directory,
+    default=None,
+    help="The local filesystem path to the directory where the frontend is located. Overrides --front-end-version.",
+)
 
 if comfy.options.args_parsing:
     args = parser.parse_args()

comfy/clip_vision.py

@@ -34,6 +34,7 @@ class ClipVisionModel():
         with open(json_config) as f:
             config = json.load(f)
 
+        self.image_size = config.get("image_size", 224)
         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)
@@ -50,7 +51,7 @@ class ClipVisionModel():
 
     def encode_image(self, image):
         comfy.model_management.load_model_gpu(self.patcher)
-        pixel_values = clip_preprocess(image.to(self.load_device)).float()
+        pixel_values = clip_preprocess(image.to(self.load_device), size=self.image_size).float()
         out = self.model(pixel_values=pixel_values, intermediate_output=-2)
 
         outputs = Output()
@@ -93,7 +94,10 @@ def load_clipvision_from_sd(sd, prefix="", convert_keys=False):
     elif "vision_model.encoder.layers.30.layer_norm1.weight" in sd:
         json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "clip_vision_config_h.json")
     elif "vision_model.encoder.layers.22.layer_norm1.weight" in sd:
-        json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "clip_vision_config_vitl.json")
+        if sd["vision_model.embeddings.position_embedding.weight"].shape[0] == 577:
+            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")
     else:
         return None
 

comfy/clip_vision_config_vitl_336.json

@@ -0,0 +1,18 @@
+{
+  "attention_dropout": 0.0,
+  "dropout": 0.0,
+  "hidden_act": "quick_gelu",
+  "hidden_size": 1024,
+  "image_size": 336,
+  "initializer_factor": 1.0,
+  "initializer_range": 0.02,
+  "intermediate_size": 4096,
+  "layer_norm_eps": 1e-5,
+  "model_type": "clip_vision_model",
+  "num_attention_heads": 16,
+  "num_channels": 3,
+  "num_hidden_layers": 24,
+  "patch_size": 14,
+  "projection_dim": 768,
+  "torch_dtype": "float32"
+}

comfy/controlnet.py

@@ -7,10 +7,12 @@ import comfy.model_management
 import comfy.model_detection
 import comfy.model_patcher
 import comfy.ops
+import comfy.latent_formats
 
 import comfy.cldm.cldm
 import comfy.t2i_adapter.adapter
 import comfy.ldm.cascade.controlnet
+import comfy.cldm.mmdit
 
 
 def broadcast_image_to(tensor, target_batch_size, batched_number):
@@ -37,20 +39,25 @@ class ControlBase:
         self.cond_hint = None
         self.strength = 1.0
         self.timestep_percent_range = (0.0, 1.0)
+        self.latent_format = None
+        self.vae = None
         self.global_average_pooling = False
         self.timestep_range = None
         self.compression_ratio = 8
         self.upscale_algorithm = 'nearest-exact'
+        self.extra_args = {}
 
         if device is None:
             device = comfy.model_management.get_torch_device()
         self.device = device
         self.previous_controlnet = None
 
-    def set_cond_hint(self, cond_hint, strength=1.0, timestep_percent_range=(0.0, 1.0)):
+    def set_cond_hint(self, cond_hint, strength=1.0, timestep_percent_range=(0.0, 1.0), vae=None):
         self.cond_hint_original = cond_hint
         self.strength = strength
         self.timestep_percent_range = timestep_percent_range
+        if self.latent_format is not None:
+            self.vae = vae
         return self
 
     def pre_run(self, model, percent_to_timestep_function):
@@ -83,43 +90,36 @@ class ControlBase:
         c.global_average_pooling = self.global_average_pooling
         c.compression_ratio = self.compression_ratio
         c.upscale_algorithm = self.upscale_algorithm
+        c.latent_format = self.latent_format
+        c.extra_args = self.extra_args.copy()
+        c.vae = self.vae
 
     def inference_memory_requirements(self, dtype):
         if self.previous_controlnet is not None:
             return self.previous_controlnet.inference_memory_requirements(dtype)
         return 0
 
-    def control_merge(self, control_input, control_output, control_prev, output_dtype):
+    def control_merge(self, control, control_prev, output_dtype):
         out = {'input':[], 'middle':[], 'output': []}
 
-        if control_input is not None:
+        for key in control:
-            for i in range(len(control_input)):
+            control_output = control[key]
-                key = 'input'
+            applied_to = set()
-                x = control_input[i]
-                if x is not None:
-                    x *= self.strength
-                    if x.dtype != output_dtype:
-                        x = x.to(output_dtype)
-                out[key].insert(0, x)
-
-        if control_output is not None:
             for i in range(len(control_output)):
-                if i == (len(control_output) - 1):
-                    key = 'middle'
-                    index = 0
-                else:
-                    key = 'output'
-                    index = i
                 x = control_output[i]
                 if x is not None:
                     if self.global_average_pooling:
                         x = torch.mean(x, dim=(2, 3), keepdim=True).repeat(1, 1, x.shape[2], x.shape[3])
 
-                    x *= self.strength
+                    if x not in applied_to: #memory saving strategy, allow shared tensors and only apply strength to shared tensors once
+                        applied_to.add(x)
+                        x *= self.strength
+
                     if x.dtype != output_dtype:
                         x = x.to(output_dtype)
 
                 out[key].append(x)
+
         if control_prev is not None:
             for x in ['input', 'middle', 'output']:
                 o = out[x]
@@ -134,20 +134,26 @@ class ControlBase:
                             if o[i].shape[0] < prev_val.shape[0]:
                                 o[i] = prev_val + o[i]
                             else:
-                                o[i] += prev_val
+                                o[i] = prev_val + o[i] #TODO: change back to inplace add if shared tensors stop being an issue
         return out
 
+    def set_extra_arg(self, argument, value=None):
+        self.extra_args[argument] = value
+
+
 class ControlNet(ControlBase):
-    def __init__(self, control_model=None, global_average_pooling=False, device=None, load_device=None, manual_cast_dtype=None):
+    def __init__(self, control_model=None, global_average_pooling=False, compression_ratio=8, latent_format=None, device=None, load_device=None, manual_cast_dtype=None):
         super().__init__(device)
         self.control_model = control_model
         self.load_device = load_device
         if control_model is not None:
             self.control_model_wrapped = comfy.model_patcher.ModelPatcher(self.control_model, load_device=load_device, offload_device=comfy.model_management.unet_offload_device())
 
+        self.compression_ratio = compression_ratio
         self.global_average_pooling = global_average_pooling
         self.model_sampling_current = None
         self.manual_cast_dtype = manual_cast_dtype
+        self.latent_format = latent_format
 
     def get_control(self, x_noisy, t, cond, batched_number):
         control_prev = None
@@ -170,7 +176,17 @@ class ControlNet(ControlBase):
             if self.cond_hint is not None:
                 del self.cond_hint
             self.cond_hint = None
-            self.cond_hint = comfy.utils.common_upscale(self.cond_hint_original, x_noisy.shape[3] * self.compression_ratio, x_noisy.shape[2] * self.compression_ratio, self.upscale_algorithm, "center").to(dtype).to(self.device)
+            compression_ratio = self.compression_ratio
+            if self.vae is not None:
+                compression_ratio *= self.vae.downscale_ratio
+            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")
+            if self.vae is not None:
+                loaded_models = comfy.model_management.loaded_models(only_currently_used=True)
+                self.cond_hint = self.vae.encode(self.cond_hint.movedim(1, -1))
+                comfy.model_management.load_models_gpu(loaded_models)
+            if self.latent_format is not None:
+                self.cond_hint = self.latent_format.process_in(self.cond_hint)
+            self.cond_hint = self.cond_hint.to(device=self.device, dtype=dtype)
         if x_noisy.shape[0] != self.cond_hint.shape[0]:
             self.cond_hint = broadcast_image_to(self.cond_hint, x_noisy.shape[0], batched_number)
 
@@ -181,8 +197,8 @@ class ControlNet(ControlBase):
         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.float(), context=context.to(dtype), y=y)
+        control = self.control_model(x=x_noisy.to(dtype), hint=self.cond_hint, timesteps=timestep.float(), context=context.to(dtype), y=y, **self.extra_args)
-        return self.control_merge(None, control, control_prev, output_dtype)
+        return self.control_merge(control, control_prev, output_dtype)
 
     def copy(self):
         c = ControlNet(None, global_average_pooling=self.global_average_pooling, load_device=self.load_device, manual_cast_dtype=self.manual_cast_dtype)
@@ -322,6 +338,39 @@ class ControlLora(ControlNet):
     def inference_memory_requirements(self, dtype):
         return comfy.utils.calculate_parameters(self.control_weights) * comfy.model_management.dtype_size(dtype) + ControlBase.inference_memory_requirements(self, dtype)
 
+def load_controlnet_mmdit(sd):
+    new_sd = comfy.model_detection.convert_diffusers_mmdit(sd, "")
+    model_config = comfy.model_detection.model_config_from_unet(new_sd, "", True)
+    num_blocks = comfy.model_detection.count_blocks(new_sd, 'joint_blocks.{}.')
+    for k in sd:
+        new_sd[k] = sd[k]
+
+    supported_inference_dtypes = model_config.supported_inference_dtypes
+
+    controlnet_config = model_config.unet_config
+    unet_dtype = comfy.model_management.unet_dtype(supported_dtypes=supported_inference_dtypes)
+    load_device = comfy.model_management.get_torch_device()
+    manual_cast_dtype = comfy.model_management.unet_manual_cast(unet_dtype, load_device)
+    if manual_cast_dtype is not None:
+        operations = comfy.ops.manual_cast
+    else:
+        operations = comfy.ops.disable_weight_init
+
+    control_model = comfy.cldm.mmdit.ControlNet(num_blocks=num_blocks, operations=operations, device=load_device, dtype=unet_dtype, **controlnet_config)
+    missing, unexpected = control_model.load_state_dict(new_sd, strict=False)
+
+    if len(missing) > 0:
+        logging.warning("missing controlnet keys: {}".format(missing))
+
+    if len(unexpected) > 0:
+        logging.debug("unexpected controlnet keys: {}".format(unexpected))
+
+    latent_format = comfy.latent_formats.SD3()
+    latent_format.shift_factor = 0 #SD3 controlnet weirdness
+    control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, load_device=load_device, manual_cast_dtype=manual_cast_dtype)
+    return control
+
+
 def load_controlnet(ckpt_path, model=None):
     controlnet_data = comfy.utils.load_torch_file(ckpt_path, safe_load=True)
     if "lora_controlnet" in controlnet_data:
@@ -370,10 +419,18 @@ def load_controlnet(ckpt_path, model=None):
             if k in controlnet_data:
                 new_sd[diffusers_keys[k]] = controlnet_data.pop(k)
 
+        if "control_add_embedding.linear_1.bias" in controlnet_data: #Union Controlnet
+            controlnet_config["union_controlnet_num_control_type"] = controlnet_data["task_embedding"].shape[0]
+            for k in list(controlnet_data.keys()):
+                new_k = k.replace('.attn.in_proj_', '.attn.in_proj.')
+                new_sd[new_k] = controlnet_data.pop(k)
+
         leftover_keys = controlnet_data.keys()
         if len(leftover_keys) > 0:
             logging.warning("leftover keys: {}".format(leftover_keys))
         controlnet_data = new_sd
+    elif "controlnet_blocks.0.weight" in controlnet_data: #SD3 diffusers format
+        return load_controlnet_mmdit(controlnet_data)
 
     pth_key = 'control_model.zero_convs.0.0.weight'
     pth = False
@@ -490,12 +547,11 @@ class T2IAdapter(ControlBase):
             self.control_input = self.t2i_model(self.cond_hint.to(x_noisy.dtype))
             self.t2i_model.cpu()
 
-        control_input = list(map(lambda a: None if a is None else a.clone(), self.control_input))
+        control_input = {}
-        mid = None
+        for k in self.control_input:
-        if self.t2i_model.xl == True:
+            control_input[k] = list(map(lambda a: None if a is None else a.clone(), self.control_input[k]))
-            mid = control_input[-1:]
+
-            control_input = control_input[:-1]
+        return self.control_merge(control_input, control_prev, x_noisy.dtype)
-        return self.control_merge(control_input, mid, control_prev, x_noisy.dtype)
 
     def copy(self):
         c = T2IAdapter(self.t2i_model, self.channels_in, self.compression_ratio, self.upscale_algorithm)

comfy/k_diffusion/deis.py

@@ -0,0 +1,121 @@
+#Taken from: https://github.com/zju-pi/diff-sampler/blob/main/gits-main/solver_utils.py
+#under Apache 2 license
+import torch
+import numpy as np
+
+# A pytorch reimplementation of DEIS (https://github.com/qsh-zh/deis).
+#############################
+### Utils for DEIS solver ###
+#############################
+#----------------------------------------------------------------------------
+# Transfer from the input time (sigma) used in EDM to that (t) used in DEIS.
+
+def edm2t(edm_steps, epsilon_s=1e-3, sigma_min=0.002, sigma_max=80):
+    vp_sigma = lambda beta_d, beta_min: lambda t: (np.e ** (0.5 * beta_d * (t ** 2) + beta_min * t) - 1) ** 0.5
+    vp_sigma_inv = lambda beta_d, beta_min: lambda sigma: ((beta_min ** 2 + 2 * beta_d * (sigma ** 2 + 1).log()).sqrt() - beta_min) / beta_d
+    vp_beta_d = 2 * (np.log(torch.tensor(sigma_min).cpu() ** 2 + 1) / epsilon_s - np.log(torch.tensor(sigma_max).cpu() ** 2 + 1)) / (epsilon_s - 1)
+    vp_beta_min = np.log(torch.tensor(sigma_max).cpu() ** 2 + 1) - 0.5 * vp_beta_d
+    t_steps = vp_sigma_inv(vp_beta_d.clone().detach().cpu(), vp_beta_min.clone().detach().cpu())(edm_steps.clone().detach().cpu())
+    return t_steps, vp_beta_min, vp_beta_d + vp_beta_min
+
+#----------------------------------------------------------------------------
+
+def cal_poly(prev_t, j, taus):
+    poly = 1
+    for k in range(prev_t.shape[0]):
+        if k == j:
+            continue
+        poly *= (taus - prev_t[k]) / (prev_t[j] - prev_t[k])
+    return poly
+
+#----------------------------------------------------------------------------
+# Transfer from t to alpha_t.
+
+def t2alpha_fn(beta_0, beta_1, t):
+    return torch.exp(-0.5 * t ** 2 * (beta_1 - beta_0) - t * beta_0)
+
+#----------------------------------------------------------------------------
+
+def cal_intergrand(beta_0, beta_1, taus):
+    with torch.inference_mode(mode=False):
+        taus = taus.clone()
+        beta_0 = beta_0.clone()
+        beta_1 = beta_1.clone()
+        with torch.enable_grad():
+            taus.requires_grad_(True)
+            alpha = t2alpha_fn(beta_0, beta_1, taus)
+            log_alpha = alpha.log()
+            log_alpha.sum().backward()
+            d_log_alpha_dtau = taus.grad
+    integrand = -0.5 * d_log_alpha_dtau / torch.sqrt(alpha * (1 - alpha))
+    return integrand
+
+#----------------------------------------------------------------------------
+
+def get_deis_coeff_list(t_steps, max_order, N=10000, deis_mode='tab'):
+    """
+    Get the coefficient list for DEIS sampling.
+
+    Args:
+        t_steps: A pytorch tensor. The time steps for sampling.
+        max_order: A `int`. Maximum order of the solver. 1 <= max_order <= 4
+        N: A `int`. Use how many points to perform the numerical integration when deis_mode=='tab'.
+        deis_mode: A `str`. Select between 'tab' and 'rhoab'. Type of DEIS.
+    Returns:
+        A pytorch tensor. A batch of generated samples or sampling trajectories if return_inters=True.
+    """
+    if deis_mode == 'tab':
+        t_steps, beta_0, beta_1 = edm2t(t_steps)
+        C = []
+        for i, (t_cur, t_next) in enumerate(zip(t_steps[:-1], t_steps[1:])):
+            order = min(i+1, max_order)
+            if order == 1:
+                C.append([])
+            else:
+                taus = torch.linspace(t_cur, t_next, N)   # split the interval for integral appximation
+                dtau = (t_next - t_cur) / N
+                prev_t = t_steps[[i - k for k in range(order)]]
+                coeff_temp = []
+                integrand = cal_intergrand(beta_0, beta_1, taus)
+                for j in range(order):
+                    poly = cal_poly(prev_t, j, taus)
+                    coeff_temp.append(torch.sum(integrand * poly) * dtau)
+                C.append(coeff_temp)
+
+    elif deis_mode == 'rhoab':
+        # Analytical solution, second order
+        def get_def_intergral_2(a, b, start, end, c):
+            coeff = (end**3 - start**3) / 3 - (end**2 - start**2) * (a + b) / 2 + (end - start) * a * b
+            return coeff / ((c - a) * (c - b))
+
+        # Analytical solution, third order
+        def get_def_intergral_3(a, b, c, start, end, d):
+            coeff = (end**4 - start**4) / 4 - (end**3 - start**3) * (a + b + c) / 3 \
+                    + (end**2 - start**2) * (a*b + a*c + b*c) / 2 - (end - start) * a * b * c
+            return coeff / ((d - a) * (d - b) * (d - c))
+
+        C = []
+        for i, (t_cur, t_next) in enumerate(zip(t_steps[:-1], t_steps[1:])):
+            order = min(i, max_order)
+            if order == 0:
+                C.append([])
+            else:
+                prev_t = t_steps[[i - k for k in range(order+1)]]
+                if order == 1:
+                    coeff_cur = ((t_next - prev_t[1])**2 - (t_cur - prev_t[1])**2) / (2 * (t_cur - prev_t[1]))
+                    coeff_prev1 = (t_next - t_cur)**2 / (2 * (prev_t[1] - t_cur))
+                    coeff_temp = [coeff_cur, coeff_prev1]
+                elif order == 2:
+                    coeff_cur = get_def_intergral_2(prev_t[1], prev_t[2], t_cur, t_next, t_cur)
+                    coeff_prev1 = get_def_intergral_2(t_cur, prev_t[2], t_cur, t_next, prev_t[1])
+                    coeff_prev2 = get_def_intergral_2(t_cur, prev_t[1], t_cur, t_next, prev_t[2])
+                    coeff_temp = [coeff_cur, coeff_prev1, coeff_prev2]
+                elif order == 3:
+                    coeff_cur = get_def_intergral_3(prev_t[1], prev_t[2], prev_t[3], t_cur, t_next, t_cur)
+                    coeff_prev1 = get_def_intergral_3(t_cur, prev_t[2], prev_t[3], t_cur, t_next, prev_t[1])
+                    coeff_prev2 = get_def_intergral_3(t_cur, prev_t[1], prev_t[3], t_cur, t_next, prev_t[2])
+                    coeff_prev3 = get_def_intergral_3(t_cur, prev_t[1], prev_t[2], t_cur, t_next, prev_t[3])
+                    coeff_temp = [coeff_cur, coeff_prev1, coeff_prev2, coeff_prev3]
+                C.append(coeff_temp)
+    return C
+

comfy/k_diffusion/sampling.py

@@ -7,7 +7,8 @@ import torchsde
 from tqdm.auto import trange, tqdm
 
 from . import utils
-
+from . import deis
+import comfy.model_patcher
 
 def append_zero(x):
     return torch.cat([x, x.new_zeros([1])])
@@ -841,3 +842,209 @@ def sample_heunpp2(model, x, sigmas, extra_args=None, callback=None, disable=Non
             d_prime = w1 * d + w2 * d_2 + w3 * d_3
             x = x + d_prime * dt
     return x
+
+
+#From https://github.com/zju-pi/diff-sampler/blob/main/diff-solvers-main/solvers.py
+#under Apache 2 license
+def sample_ipndm(model, x, sigmas, extra_args=None, callback=None, disable=None, max_order=4):
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+
+    x_next = x
+
+    buffer_model = []
+    for i in trange(len(sigmas) - 1, disable=disable):
+        t_cur = sigmas[i]
+        t_next = sigmas[i + 1]
+
+        x_cur = x_next
+
+        denoised = model(x_cur, t_cur * s_in, **extra_args)
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+
+        d_cur = (x_cur - denoised) / t_cur
+
+        order = min(max_order, i+1)
+        if order == 1:      # First Euler step.
+            x_next = x_cur + (t_next - t_cur) * d_cur
+        elif order == 2:    # Use one history point.
+            x_next = x_cur + (t_next - t_cur) * (3 * d_cur - buffer_model[-1]) / 2
+        elif order == 3:    # Use two history points.
+            x_next = x_cur + (t_next - t_cur) * (23 * d_cur - 16 * buffer_model[-1] + 5 * buffer_model[-2]) / 12
+        elif order == 4:    # Use three history points.
+            x_next = x_cur + (t_next - t_cur) * (55 * d_cur - 59 * buffer_model[-1] + 37 * buffer_model[-2] - 9 * buffer_model[-3]) / 24
+
+        if len(buffer_model) == max_order - 1:
+            for k in range(max_order - 2):
+                buffer_model[k] = buffer_model[k+1]
+            buffer_model[-1] = d_cur
+        else:
+            buffer_model.append(d_cur)
+
+    return x_next
+
+#From https://github.com/zju-pi/diff-sampler/blob/main/diff-solvers-main/solvers.py
+#under Apache 2 license
+def sample_ipndm_v(model, x, sigmas, extra_args=None, callback=None, disable=None, max_order=4):
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+
+    x_next = x
+    t_steps = sigmas
+
+    buffer_model = []
+    for i in trange(len(sigmas) - 1, disable=disable):
+        t_cur = sigmas[i]
+        t_next = sigmas[i + 1]
+
+        x_cur = x_next
+
+        denoised = model(x_cur, t_cur * s_in, **extra_args)
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+
+        d_cur = (x_cur - denoised) / t_cur
+
+        order = min(max_order, i+1)
+        if order == 1:      # First Euler step.
+            x_next = x_cur + (t_next - t_cur) * d_cur
+        elif order == 2:    # Use one history point.
+            h_n = (t_next - t_cur)
+            h_n_1 = (t_cur - t_steps[i-1])
+            coeff1 = (2 + (h_n / h_n_1)) / 2
+            coeff2 = -(h_n / h_n_1) / 2
+            x_next = x_cur + (t_next - t_cur) * (coeff1 * d_cur + coeff2 * buffer_model[-1])
+        elif order == 3:    # Use two history points.
+            h_n = (t_next - t_cur)
+            h_n_1 = (t_cur - t_steps[i-1])
+            h_n_2 = (t_steps[i-1] - t_steps[i-2])
+            temp = (1 - h_n / (3 * (h_n + h_n_1)) * (h_n * (h_n + h_n_1)) / (h_n_1 * (h_n_1 + h_n_2))) / 2
+            coeff1 = (2 + (h_n / h_n_1)) / 2 + temp
+            coeff2 = -(h_n / h_n_1) / 2 - (1 + h_n_1 / h_n_2) * temp
+            coeff3 = temp * h_n_1 / h_n_2
+            x_next = x_cur + (t_next - t_cur) * (coeff1 * d_cur + coeff2 * buffer_model[-1] + coeff3 * buffer_model[-2])
+        elif order == 4:    # Use three history points.
+            h_n = (t_next - t_cur)
+            h_n_1 = (t_cur - t_steps[i-1])
+            h_n_2 = (t_steps[i-1] - t_steps[i-2])
+            h_n_3 = (t_steps[i-2] - t_steps[i-3])
+            temp1 = (1 - h_n / (3 * (h_n + h_n_1)) * (h_n * (h_n + h_n_1)) / (h_n_1 * (h_n_1 + h_n_2))) / 2
+            temp2 = ((1 - h_n / (3 * (h_n + h_n_1))) / 2 + (1 - h_n / (2 * (h_n + h_n_1))) * h_n / (6 * (h_n + h_n_1 + h_n_2))) \
+                   * (h_n * (h_n + h_n_1) * (h_n + h_n_1 + h_n_2)) / (h_n_1 * (h_n_1 + h_n_2) * (h_n_1 + h_n_2 + h_n_3))
+            coeff1 = (2 + (h_n / h_n_1)) / 2 + temp1 + temp2
+            coeff2 = -(h_n / h_n_1) / 2 - (1 + h_n_1 / h_n_2) * temp1 - (1 + (h_n_1 / h_n_2) + (h_n_1 * (h_n_1 + h_n_2) / (h_n_2 * (h_n_2 + h_n_3)))) * temp2
+            coeff3 = temp1 * h_n_1 / h_n_2 + ((h_n_1 / h_n_2) + (h_n_1 * (h_n_1 + h_n_2) / (h_n_2 * (h_n_2 + h_n_3))) * (1 + h_n_2 / h_n_3)) * temp2
+            coeff4 = -temp2 * (h_n_1 * (h_n_1 + h_n_2) / (h_n_2 * (h_n_2 + h_n_3))) * h_n_1 / h_n_2
+            x_next = x_cur + (t_next - t_cur) * (coeff1 * d_cur + coeff2 * buffer_model[-1] + coeff3 * buffer_model[-2] + coeff4 * buffer_model[-3])
+
+        if len(buffer_model) == max_order - 1:
+            for k in range(max_order - 2):
+                buffer_model[k] = buffer_model[k+1]
+            buffer_model[-1] = d_cur.detach()
+        else:
+            buffer_model.append(d_cur.detach())
+
+    return x_next
+
+#From https://github.com/zju-pi/diff-sampler/blob/main/diff-solvers-main/solvers.py
+#under Apache 2 license
+@torch.no_grad()
+def sample_deis(model, x, sigmas, extra_args=None, callback=None, disable=None, max_order=3, deis_mode='tab'):
+    extra_args = {} if extra_args is None else extra_args
+    s_in = x.new_ones([x.shape[0]])
+
+    x_next = x
+    t_steps = sigmas
+
+    coeff_list = deis.get_deis_coeff_list(t_steps, max_order, deis_mode=deis_mode)
+
+    buffer_model = []
+    for i in trange(len(sigmas) - 1, disable=disable):
+        t_cur = sigmas[i]
+        t_next = sigmas[i + 1]
+
+        x_cur = x_next
+
+        denoised = model(x_cur, t_cur * s_in, **extra_args)
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
+
+        d_cur = (x_cur - denoised) / t_cur
+
+        order = min(max_order, i+1)
+        if t_next <= 0:
+            order = 1
+
+        if order == 1:          # First Euler step.
+            x_next = x_cur + (t_next - t_cur) * d_cur
+        elif order == 2:        # Use one history point.
+            coeff_cur, coeff_prev1 = coeff_list[i]
+            x_next = x_cur + coeff_cur * d_cur + coeff_prev1 * buffer_model[-1]
+        elif order == 3:        # Use two history points.
+            coeff_cur, coeff_prev1, coeff_prev2 = coeff_list[i]
+            x_next = x_cur + coeff_cur * d_cur + coeff_prev1 * buffer_model[-1] + coeff_prev2 * buffer_model[-2]
+        elif order == 4:        # Use three history points.
+            coeff_cur, coeff_prev1, coeff_prev2, coeff_prev3 = coeff_list[i]
+            x_next = x_cur + coeff_cur * d_cur + coeff_prev1 * buffer_model[-1] + coeff_prev2 * buffer_model[-2] + coeff_prev3 * buffer_model[-3]
+
+        if len(buffer_model) == max_order - 1:
+            for k in range(max_order - 2):
+                buffer_model[k] = buffer_model[k+1]
+            buffer_model[-1] = d_cur.detach()
+        else:
+            buffer_model.append(d_cur.detach())
+
+    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})
+        dt = sigmas[i + 1] - sigma_hat
+        # 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."""
+    extra_args = {} if extra_args is None else extra_args
+    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
+
+    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):
+        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
+        dt = sigma_down - sigmas[i]
+        x = denoised + d * sigma_down
+        if sigmas[i + 1] > 0:
+            x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
+    return x

comfy/ldm/aura/mmdit.py

@@ -0,0 +1,479 @@
+#AuraFlow MMDiT
+#Originally written by the AuraFlow Authors
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from comfy.ldm.modules.attention import optimized_attention
+
+def modulate(x, shift, scale):
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+def find_multiple(n: int, k: int) -> int:
+    if n % k == 0:
+        return n
+    return n + k - (n % k)
+
+
+class MLP(nn.Module):
+    def __init__(self, dim, hidden_dim=None, dtype=None, device=None, operations=None) -> None:
+        super().__init__()
+        if hidden_dim is None:
+            hidden_dim = 4 * dim
+
+        n_hidden = int(2 * hidden_dim / 3)
+        n_hidden = find_multiple(n_hidden, 256)
+
+        self.c_fc1 = operations.Linear(dim, n_hidden, bias=False, dtype=dtype, device=device)
+        self.c_fc2 = operations.Linear(dim, n_hidden, bias=False, dtype=dtype, device=device)
+        self.c_proj = operations.Linear(n_hidden, dim, bias=False, dtype=dtype, device=device)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = F.silu(self.c_fc1(x)) * self.c_fc2(x)
+        x = self.c_proj(x)
+        return x
+
+
+class MultiHeadLayerNorm(nn.Module):
+    def __init__(self, hidden_size=None, eps=1e-5, dtype=None, device=None):
+        # Copy pasta from https://github.com/huggingface/transformers/blob/e5f71ecaae50ea476d1e12351003790273c4b2ed/src/transformers/models/cohere/modeling_cohere.py#L78
+
+        super().__init__()
+        self.weight = nn.Parameter(torch.empty(hidden_size, dtype=dtype, device=device))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        mean = hidden_states.mean(-1, keepdim=True)
+        variance = (hidden_states - mean).pow(2).mean(-1, keepdim=True)
+        hidden_states = (hidden_states - mean) * torch.rsqrt(
+            variance + self.variance_epsilon
+        )
+        hidden_states = self.weight.to(torch.float32) * hidden_states
+        return hidden_states.to(input_dtype)
+
+class SingleAttention(nn.Module):
+    def __init__(self, dim, n_heads, mh_qknorm=False, dtype=None, device=None, operations=None):
+        super().__init__()
+
+        self.n_heads = n_heads
+        self.head_dim = dim // n_heads
+
+        # this is for cond
+        self.w1q = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
+        self.w1k = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
+        self.w1v = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
+        self.w1o = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
+
+        self.q_norm1 = (
+            MultiHeadLayerNorm((self.n_heads, self.head_dim), dtype=dtype, device=device)
+            if mh_qknorm
+            else operations.LayerNorm(self.head_dim, elementwise_affine=False, dtype=dtype, device=device)
+        )
+        self.k_norm1 = (
+            MultiHeadLayerNorm((self.n_heads, self.head_dim), dtype=dtype, device=device)
+            if mh_qknorm
+            else operations.LayerNorm(self.head_dim, elementwise_affine=False, dtype=dtype, device=device)
+        )
+
+    #@torch.compile()
+    def forward(self, c):
+
+        bsz, seqlen1, _ = c.shape
+
+        q, k, v = self.w1q(c), self.w1k(c), self.w1v(c)
+        q = q.view(bsz, seqlen1, self.n_heads, self.head_dim)
+        k = k.view(bsz, seqlen1, self.n_heads, self.head_dim)
+        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)
+        c = self.w1o(output)
+        return c
+
+
+
+class DoubleAttention(nn.Module):
+    def __init__(self, dim, n_heads, mh_qknorm=False, dtype=None, device=None, operations=None):
+        super().__init__()
+
+        self.n_heads = n_heads
+        self.head_dim = dim // n_heads
+
+        # this is for cond
+        self.w1q = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
+        self.w1k = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
+        self.w1v = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
+        self.w1o = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
+
+        # this is for x
+        self.w2q = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
+        self.w2k = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
+        self.w2v = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
+        self.w2o = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
+
+        self.q_norm1 = (
+            MultiHeadLayerNorm((self.n_heads, self.head_dim), dtype=dtype, device=device)
+            if mh_qknorm
+            else operations.LayerNorm(self.head_dim, elementwise_affine=False, dtype=dtype, device=device)
+        )
+        self.k_norm1 = (
+            MultiHeadLayerNorm((self.n_heads, self.head_dim), dtype=dtype, device=device)
+            if mh_qknorm
+            else operations.LayerNorm(self.head_dim, elementwise_affine=False, dtype=dtype, device=device)
+        )
+
+        self.q_norm2 = (
+            MultiHeadLayerNorm((self.n_heads, self.head_dim), dtype=dtype, device=device)
+            if mh_qknorm
+            else operations.LayerNorm(self.head_dim, elementwise_affine=False, dtype=dtype, device=device)
+        )
+        self.k_norm2 = (
+            MultiHeadLayerNorm((self.n_heads, self.head_dim), dtype=dtype, device=device)
+            if mh_qknorm
+            else operations.LayerNorm(self.head_dim, elementwise_affine=False, dtype=dtype, device=device)
+        )
+
+
+    #@torch.compile()
+    def forward(self, c, x):
+
+        bsz, seqlen1, _ = c.shape
+        bsz, seqlen2, _ = x.shape
+        seqlen = seqlen1 + seqlen2
+
+        cq, ck, cv = self.w1q(c), self.w1k(c), self.w1v(c)
+        cq = cq.view(bsz, seqlen1, self.n_heads, self.head_dim)
+        ck = ck.view(bsz, seqlen1, self.n_heads, self.head_dim)
+        cv = cv.view(bsz, seqlen1, self.n_heads, self.head_dim)
+        cq, ck = self.q_norm1(cq), self.k_norm1(ck)
+
+        xq, xk, xv = self.w2q(x), self.w2k(x), self.w2v(x)
+        xq = xq.view(bsz, seqlen2, self.n_heads, self.head_dim)
+        xk = xk.view(bsz, seqlen2, self.n_heads, self.head_dim)
+        xv = xv.view(bsz, seqlen2, self.n_heads, self.head_dim)
+        xq, xk = self.q_norm2(xq), self.k_norm2(xk)
+
+        # concat all
+        q, k, v = (
+            torch.cat([cq, xq], dim=1),
+            torch.cat([ck, xk], dim=1),
+            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)
+
+        c, x = output.split([seqlen1, seqlen2], dim=1)
+        c = self.w1o(c)
+        x = self.w2o(x)
+
+        return c, x
+
+
+class MMDiTBlock(nn.Module):
+    def __init__(self, dim, heads=8, global_conddim=1024, is_last=False, dtype=None, device=None, operations=None):
+        super().__init__()
+
+        self.normC1 = operations.LayerNorm(dim, elementwise_affine=False, dtype=dtype, device=device)
+        self.normC2 = operations.LayerNorm(dim, elementwise_affine=False, dtype=dtype, device=device)
+        if not is_last:
+            self.mlpC = MLP(dim, hidden_dim=dim * 4, dtype=dtype, device=device, operations=operations)
+            self.modC = nn.Sequential(
+                nn.SiLU(),
+                operations.Linear(global_conddim, 6 * dim, bias=False, dtype=dtype, device=device),
+            )
+        else:
+            self.modC = nn.Sequential(
+                nn.SiLU(),
+                operations.Linear(global_conddim, 2 * dim, bias=False, dtype=dtype, device=device),
+            )
+
+        self.normX1 = operations.LayerNorm(dim, elementwise_affine=False, dtype=dtype, device=device)
+        self.normX2 = operations.LayerNorm(dim, elementwise_affine=False, dtype=dtype, device=device)
+        self.mlpX = MLP(dim, hidden_dim=dim * 4, dtype=dtype, device=device, operations=operations)
+        self.modX = nn.Sequential(
+            nn.SiLU(),
+            operations.Linear(global_conddim, 6 * dim, bias=False, dtype=dtype, device=device),
+        )
+
+        self.attn = DoubleAttention(dim, heads, dtype=dtype, device=device, operations=operations)
+        self.is_last = is_last
+
+    #@torch.compile()
+    def forward(self, c, x, global_cond, **kwargs):
+
+        cres, xres = c, x
+
+        cshift_msa, cscale_msa, cgate_msa, cshift_mlp, cscale_mlp, cgate_mlp = (
+            self.modC(global_cond).chunk(6, dim=1)
+        )
+
+        c = modulate(self.normC1(c), cshift_msa, cscale_msa)
+
+        # xpath
+        xshift_msa, xscale_msa, xgate_msa, xshift_mlp, xscale_mlp, xgate_mlp = (
+            self.modX(global_cond).chunk(6, dim=1)
+        )
+
+        x = modulate(self.normX1(x), xshift_msa, xscale_msa)
+
+        # attention
+        c, x = self.attn(c, x)
+
+
+        c = self.normC2(cres + cgate_msa.unsqueeze(1) * c)
+        c = cgate_mlp.unsqueeze(1) * self.mlpC(modulate(c, cshift_mlp, cscale_mlp))
+        c = cres + c
+
+        x = self.normX2(xres + xgate_msa.unsqueeze(1) * x)
+        x = xgate_mlp.unsqueeze(1) * self.mlpX(modulate(x, xshift_mlp, xscale_mlp))
+        x = xres + x
+
+        return c, x
+
+class DiTBlock(nn.Module):
+    # like MMDiTBlock, but it only has X
+    def __init__(self, dim, heads=8, global_conddim=1024, dtype=None, device=None, operations=None):
+        super().__init__()
+
+        self.norm1 = operations.LayerNorm(dim, elementwise_affine=False, dtype=dtype, device=device)
+        self.norm2 = operations.LayerNorm(dim, elementwise_affine=False, dtype=dtype, device=device)
+
+        self.modCX = nn.Sequential(
+            nn.SiLU(),
+            operations.Linear(global_conddim, 6 * dim, bias=False, dtype=dtype, device=device),
+        )
+
+        self.attn = SingleAttention(dim, heads, dtype=dtype, device=device, operations=operations)
+        self.mlp = MLP(dim, hidden_dim=dim * 4, dtype=dtype, device=device, operations=operations)
+
+    #@torch.compile()
+    def forward(self, cx, global_cond, **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.norm2(cxres + gate_msa.unsqueeze(1) * cx)
+        mlpout = self.mlp(modulate(cx, shift_mlp, scale_mlp))
+        cx = gate_mlp.unsqueeze(1) * mlpout
+
+        cx = cxres + cx
+
+        return cx
+
+
+
+class TimestepEmbedder(nn.Module):
+    def __init__(self, hidden_size, frequency_embedding_size=256, dtype=None, device=None, operations=None):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            operations.Linear(frequency_embedding_size, hidden_size, dtype=dtype, device=device),
+            nn.SiLU(),
+            operations.Linear(hidden_size, hidden_size, dtype=dtype, device=device),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        half = dim // 2
+        freqs = 1000 * torch.exp(
+            -math.log(max_period) * torch.arange(start=0, end=half) / half
+        ).to(t.device)
+        args = t[:, None] * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat(
+                [embedding, torch.zeros_like(embedding[:, :1])], dim=-1
+            )
+        return embedding
+
+    #@torch.compile()
+    def forward(self, t, dtype):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size).to(dtype)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+class MMDiT(nn.Module):
+    def __init__(
+        self,
+        in_channels=4,
+        out_channels=4,
+        patch_size=2,
+        dim=3072,
+        n_layers=36,
+        n_double_layers=4,
+        n_heads=12,
+        global_conddim=3072,
+        cond_seq_dim=2048,
+        max_seq=32 * 32,
+        device=None,
+        dtype=None,
+        operations=None,
+    ):
+        super().__init__()
+        self.dtype = dtype
+
+        self.t_embedder = TimestepEmbedder(global_conddim, dtype=dtype, device=device, operations=operations)
+
+        self.cond_seq_linear = operations.Linear(
+            cond_seq_dim, dim, bias=False, dtype=dtype, device=device
+        )  # linear for something like text sequence.
+        self.init_x_linear = operations.Linear(
+            patch_size * patch_size * in_channels, dim, dtype=dtype, device=device
+        )  # init linear for patchified image.
+
+        self.positional_encoding = nn.Parameter(torch.empty(1, max_seq, dim, dtype=dtype, device=device))
+        self.register_tokens = nn.Parameter(torch.empty(1, 8, dim, dtype=dtype, device=device))
+
+        self.double_layers = nn.ModuleList([])
+        self.single_layers = nn.ModuleList([])
+
+
+        for idx in range(n_double_layers):
+            self.double_layers.append(
+                MMDiTBlock(dim, n_heads, global_conddim, is_last=(idx == n_layers - 1), dtype=dtype, device=device, operations=operations)
+            )
+
+        for idx in range(n_double_layers, n_layers):
+            self.single_layers.append(
+                DiTBlock(dim, n_heads, global_conddim, dtype=dtype, device=device, operations=operations)
+            )
+
+
+        self.final_linear = operations.Linear(
+            dim, patch_size * patch_size * out_channels, bias=False, dtype=dtype, device=device
+        )
+
+        self.modF = nn.Sequential(
+            nn.SiLU(),
+            operations.Linear(global_conddim, 2 * dim, bias=False, dtype=dtype, device=device),
+        )
+
+        self.out_channels = out_channels
+        self.patch_size = patch_size
+        self.n_double_layers = n_double_layers
+        self.n_layers = n_layers
+
+        self.h_max = round(max_seq**0.5)
+        self.w_max = round(max_seq**0.5)
+
+    @torch.no_grad()
+    def extend_pe(self, init_dim=(16, 16), target_dim=(64, 64)):
+        # extend pe
+        pe_data = self.positional_encoding.data.squeeze(0)[: init_dim[0] * init_dim[1]]
+
+        pe_as_2d = pe_data.view(init_dim[0], init_dim[1], -1).permute(2, 0, 1)
+
+        # now we need to extend this to target_dim. for this we will use interpolation.
+        # we will use torch.nn.functional.interpolate
+        pe_as_2d = F.interpolate(
+            pe_as_2d.unsqueeze(0), size=target_dim, mode="bilinear"
+        )
+        pe_new = pe_as_2d.squeeze(0).permute(1, 2, 0).flatten(0, 1)
+        self.positional_encoding.data = pe_new.unsqueeze(0).contiguous()
+        self.h_max, self.w_max = target_dim
+        print("PE extended to", target_dim)
+
+    def pe_selection_index_based_on_dim(self, h, w):
+        h_p, w_p = h // self.patch_size, w // self.patch_size
+        original_pe_indexes = torch.arange(self.positional_encoding.shape[1])
+        original_pe_indexes = original_pe_indexes.view(self.h_max, self.w_max)
+        starth =  self.h_max // 2 - h_p // 2
+        endh =starth + h_p
+        startw = self.w_max // 2 - w_p // 2
+        endw = startw + w_p
+        original_pe_indexes = original_pe_indexes[
+            starth:endh, startw:endw
+        ]
+        return original_pe_indexes.flatten()
+
+    def unpatchify(self, x, h, w):
+        c = self.out_channels
+        p = self.patch_size
+
+        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
+
+    def patchify(self, x):
+        B, C, H, W = x.size()
+        pad_h = (self.patch_size - H % self.patch_size) % self.patch_size
+        pad_w = (self.patch_size - W % self.patch_size) % self.patch_size
+
+        x = torch.nn.functional.pad(x, (0, pad_w, 0, pad_h), mode='reflect')
+        x = x.view(
+            B,
+            C,
+            (H + 1) // self.patch_size,
+            self.patch_size,
+            (W + 1) // self.patch_size,
+            self.patch_size,
+        )
+        x = x.permute(0, 2, 4, 1, 3, 5).flatten(-3).flatten(1, 2)
+        return x
+
+    def apply_pos_embeds(self, x, h, w):
+        h = (h + 1) // self.patch_size
+        w = (w + 1) // self.patch_size
+        max_dim = max(h, w)
+
+        cur_dim = self.h_max
+        pos_encoding = self.positional_encoding.reshape(1, cur_dim, cur_dim, -1).to(device=x.device, dtype=x.dtype)
+
+        if max_dim > cur_dim:
+            pos_encoding = F.interpolate(pos_encoding.movedim(-1, 1), (max_dim, max_dim), mode="bilinear").movedim(1, -1)
+            cur_dim = max_dim
+
+        from_h = (cur_dim - h) // 2
+        from_w = (cur_dim - w) // 2
+        pos_encoding = pos_encoding[:,from_h:from_h+h,from_w:from_w+w]
+        return x + pos_encoding.reshape(1, -1, self.positional_encoding.shape[-1])
+
+    def forward(self, x, timestep, context, **kwargs):
+        # patchify x, add PE
+        b, c, h, w = x.shape
+
+        # pe_indexes = self.pe_selection_index_based_on_dim(h, w)
+        # print(pe_indexes, pe_indexes.shape)
+
+        x = self.init_x_linear(self.patchify(x))  # B, T_x, D
+        x = self.apply_pos_embeds(x, h, w)
+        # x = x + self.positional_encoding[:, : x.size(1)].to(device=x.device, dtype=x.dtype)
+        # x = x + self.positional_encoding[:, pe_indexes].to(device=x.device, dtype=x.dtype)
+
+        # process conditions for MMDiT Blocks
+        c_seq = context  # B, T_c, D_c
+        t = timestep
+
+        c = self.cond_seq_linear(c_seq)  # B, T_c, D
+        c = torch.cat([self.register_tokens.to(device=c.device, dtype=c.dtype).repeat(c.size(0), 1, 1), c], dim=1)
+
+        global_cond = self.t_embedder(t, x.dtype)  # B, D
+
+        if len(self.double_layers) > 0:
+            for layer in self.double_layers:
+                c, x = layer(c, x, global_cond, **kwargs)
+
+        if len(self.single_layers) > 0:
+            c_len = c.size(1)
+            cx = torch.cat([c, x], dim=1)
+            for layer in self.single_layers:
+                cx = layer(cx, global_cond, **kwargs)
+
+            x = cx[:, c_len:]
+
+        fshift, fscale = self.modF(global_cond).chunk(2, dim=1)
+
+        x = modulate(x, fshift, fscale)
+        x = self.final_linear(x)
+        x = self.unpatchify(x, (h + 1) // self.patch_size, (w + 1) // self.patch_size)[:,:,:h,:w]
+        return x

comfy/ldm/cascade/controlnet.py

@@ -90,4 +90,4 @@ class ControlNet(nn.Module):
         proj_outputs = [None for _ in range(max(self.proj_blocks) + 1)]
         for i, idx in enumerate(self.proj_blocks):
             proj_outputs[idx] = self.projections[i](x)
-        return proj_outputs
+        return {"input": proj_outputs[::-1]}

comfy/ldm/modules/diffusionmodules/mmdit.py

@@ -745,6 +745,8 @@ class MMDiT(nn.Module):
         qkv_bias: bool = True,
         context_processor_layers = None,
         context_size = 4096,
+        num_blocks = None,
+        final_layer = True,
         dtype = None, #TODO
         device = None,
         operations = None,
@@ -766,7 +768,10 @@ class MMDiT(nn.Module):
         # apply magic --> this defines a head_size of 64
         self.hidden_size = 64 * depth
         num_heads = depth
+        if num_blocks is None:
+            num_blocks = depth
 
+        self.depth = depth
         self.num_heads = num_heads
 
         self.x_embedder = PatchEmbed(
@@ -821,7 +826,7 @@ class MMDiT(nn.Module):
                     mlp_ratio=mlp_ratio,
                     qkv_bias=qkv_bias,
                     attn_mode=attn_mode,
-                    pre_only=i == depth - 1,
+                    pre_only=(i == num_blocks - 1) and final_layer,
                     rmsnorm=rmsnorm,
                     scale_mod_only=scale_mod_only,
                     swiglu=swiglu,
@@ -830,11 +835,12 @@ class MMDiT(nn.Module):
                     device=device,
                     operations=operations
                 )
-                for i in range(depth)
+                for i in range(num_blocks)
             ]
         )
 
-        self.final_layer = FinalLayer(self.hidden_size, patch_size, self.out_channels, dtype=dtype, device=device, operations=operations)
+        if final_layer:
+            self.final_layer = FinalLayer(self.hidden_size, patch_size, self.out_channels, dtype=dtype, device=device, operations=operations)
 
         if compile_core:
             assert False
@@ -893,6 +899,7 @@ class MMDiT(nn.Module):
         x: torch.Tensor,
         c_mod: torch.Tensor,
         context: Optional[torch.Tensor] = None,
+        control = None,
     ) -> torch.Tensor:
         if self.register_length > 0:
             context = torch.cat(
@@ -905,13 +912,20 @@ class MMDiT(nn.Module):
 
         # context is B, L', D
         # x is B, L, D
-        for block in self.joint_blocks:
+        blocks = len(self.joint_blocks)
-            context, x = block(
+        for i in range(blocks):
+            context, x = self.joint_blocks[i](
                 context,
                 x,
                 c=c_mod,
                 use_checkpoint=self.use_checkpoint,
             )
+            if control is not None:
+                control_o = control.get("output")
+                if i < len(control_o):
+                    add = control_o[i]
+                    if add is not None:
+                        x += add
 
         x = self.final_layer(x, c_mod)  # (N, T, patch_size ** 2 * out_channels)
         return x
@@ -922,6 +936,7 @@ class MMDiT(nn.Module):
         t: torch.Tensor,
         y: Optional[torch.Tensor] = None,
         context: Optional[torch.Tensor] = None,
+        control = None,
     ) -> torch.Tensor:
         """
         Forward pass of DiT.
@@ -943,7 +958,7 @@ class MMDiT(nn.Module):
         if context is not None:
             context = self.context_embedder(context)
 
-        x = self.forward_core_with_concat(x, c, context)
+        x = self.forward_core_with_concat(x, c, context, control)
 
         x = self.unpatchify(x, hw=hw)  # (N, out_channels, H, W)
         return x[:,:,:hw[-2],:hw[-1]]
@@ -956,7 +971,8 @@ class OpenAISignatureMMDITWrapper(MMDiT):
         timesteps: torch.Tensor,
         context: Optional[torch.Tensor] = None,
         y: Optional[torch.Tensor] = None,
+        control = None,
         **kwargs,
     ) -> torch.Tensor:
-        return super().forward(x, timesteps, context=context, y=y)
+        return super().forward(x, timesteps, context=context, y=y, control=control)
 

comfy/lora.py

@@ -30,6 +30,7 @@ def load_lora(lora, to_load):
         regular_lora = "{}.lora_up.weight".format(x)
         diffusers_lora = "{}_lora.up.weight".format(x)
         diffusers2_lora = "{}.lora_B.weight".format(x)
+        diffusers3_lora = "{}.lora.up.weight".format(x)
         transformers_lora = "{}.lora_linear_layer.up.weight".format(x)
         A_name = None
 
@@ -45,6 +46,10 @@ def load_lora(lora, to_load):
             A_name = diffusers2_lora
             B_name = "{}.lora_A.weight".format(x)
             mid_name = None
+        elif diffusers3_lora in lora.keys():
+            A_name = diffusers3_lora
+            B_name = "{}.lora.down.weight".format(x)
+            mid_name = None
         elif transformers_lora in lora.keys():
             A_name = transformers_lora
             B_name ="{}.lora_linear_layer.down.weight".format(x)
@@ -213,12 +218,21 @@ def model_lora_keys_clip(model, key_map={}):
                     lora_key = "lora_prior_te_text_model_encoder_layers_{}_{}".format(b, LORA_CLIP_MAP[c]) #cascade lora: TODO put lora key prefix in the model config
                     key_map[lora_key] = k
 
+    for k in sdk: #OneTrainer SD3 lora
+        if k.startswith("t5xxl.transformer.") and k.endswith(".weight"):
+            l_key = k[len("t5xxl.transformer."):-len(".weight")]
+            lora_key = "lora_te3_{}".format(l_key.replace(".", "_"))
+            key_map[lora_key] = k
 
     k = "clip_g.transformer.text_projection.weight"
     if k in sdk:
         key_map["lora_prior_te_text_projection"] = k #cascade lora?
         # key_map["text_encoder.text_projection"] = k #TODO: check if other lora have the text_projection too
-        # key_map["lora_te_text_projection"] = k
+        key_map["lora_te2_text_projection"] = k #OneTrainer SD3 lora
+
+    k = "clip_l.transformer.text_projection.weight"
+    if k in sdk:
+        key_map["lora_te1_text_projection"] = k #OneTrainer SD3 lora, not necessary but omits warning
 
     return key_map
 
@@ -247,15 +261,25 @@ def model_lora_keys_unet(model, key_map={}):
                 key_map[diffusers_lora_key] = unet_key
 
     if isinstance(model, comfy.model_base.SD3): #Diffusers lora SD3
-        for i in range(model.model_config.unet_config.get("depth", 0)):
+        diffusers_keys = comfy.utils.mmdit_to_diffusers(model.model_config.unet_config, output_prefix="diffusion_model.")
-            k = "transformer.transformer_blocks.{}.attn.".format(i)
+        for k in diffusers_keys:
-            qkv = "diffusion_model.joint_blocks.{}.x_block.attn.qkv.weight".format(i)
+            if k.endswith(".weight"):
-            proj = "diffusion_model.joint_blocks.{}.x_block.attn.proj.weight".format(i)
+                to = diffusers_keys[k]
-            if qkv in sd:
+                key_lora = "transformer.{}".format(k[:-len(".weight")]) #regular diffusers sd3 lora format
-                offset = sd[qkv].shape[0] // 3
+                key_map[key_lora] = to
-                key_map["{}to_q".format(k)] = (qkv, (0, 0, offset))
+
-                key_map["{}to_k".format(k)] = (qkv, (0, offset, offset))
+                key_lora = "base_model.model.{}".format(k[:-len(".weight")]) #format for flash-sd3 lora and others?
-                key_map["{}to_v".format(k)] = (qkv, (0, offset * 2, offset))
+                key_map[key_lora] = to
-                key_map["{}to_out.0".format(k)] = proj
+
+                key_lora = "lora_transformer_{}".format(k[:-len(".weight")].replace(".", "_")) #OneTrainer lora
+                key_map[key_lora] = to
+
+    if isinstance(model, comfy.model_base.AuraFlow): #Diffusers lora AuraFlow
+        diffusers_keys = comfy.utils.auraflow_to_diffusers(model.model_config.unet_config, output_prefix="diffusion_model.")
+        for k in diffusers_keys:
+            if k.endswith(".weight"):
+                to = diffusers_keys[k]
+                key_lora = "transformer.{}".format(k[:-len(".weight")]) #simpletrainer and probably regular diffusers lora format
+                key_map[key_lora] = to
 
     return key_map

comfy/model_base.py

@@ -6,6 +6,7 @@ from comfy.ldm.cascade.stage_b import StageB
 from comfy.ldm.modules.encoders.noise_aug_modules import CLIPEmbeddingNoiseAugmentation
 from comfy.ldm.modules.diffusionmodules.upscaling import ImageConcatWithNoiseAugmentation
 from comfy.ldm.modules.diffusionmodules.mmdit import OpenAISignatureMMDITWrapper
+import comfy.ldm.aura.mmdit
 import comfy.ldm.audio.dit
 import comfy.ldm.audio.embedders
 import comfy.model_management
@@ -598,6 +599,17 @@ class SD3(BaseModel):
             area = input_shape[0] * input_shape[2] * input_shape[3]
             return (area * 0.3) * (1024 * 1024)
 
+class AuraFlow(BaseModel):
+    def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
+        super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.aura.mmdit.MMDiT)
+
+    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)
+        return out
+
 
 class StableAudio1(BaseModel):
     def __init__(self, model_config, seconds_start_embedder_weights, seconds_total_embedder_weights, model_type=ModelType.V_PREDICTION_CONTINUOUS, device=None):
@@ -627,3 +639,12 @@ class StableAudio1(BaseModel):
             cross_attn = torch.cat([cross_attn.to(device), seconds_start_embed.repeat((cross_attn.shape[0], 1, 1)), seconds_total_embed.repeat((cross_attn.shape[0], 1, 1))], dim=1)
             out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
         return out
+
+    def state_dict_for_saving(self, clip_state_dict=None, vae_state_dict=None, clip_vision_state_dict=None):
+        sd = super().state_dict_for_saving(clip_state_dict=clip_state_dict, vae_state_dict=vae_state_dict, clip_vision_state_dict=clip_vision_state_dict)
+        d = {"conditioner.conditioners.seconds_start.": self.seconds_start_embedder.state_dict(), "conditioner.conditioners.seconds_total.": self.seconds_total_embedder.state_dict()}
+        for k in d:
+            s = d[k]
+            for l in s:
+                sd["{}{}".format(k, l)] = s[l]
+        return sd

comfy/model_detection.py

@@ -1,7 +1,9 @@
 import comfy.supported_models
 import comfy.supported_models_base
+import comfy.utils
 import math
 import logging
+import torch
 
 def count_blocks(state_dict_keys, prefix_string):
     count = 0
@@ -39,7 +41,9 @@ def detect_unet_config(state_dict, key_prefix):
         unet_config["in_channels"] = state_dict['{}x_embedder.proj.weight'.format(key_prefix)].shape[1]
         patch_size = state_dict['{}x_embedder.proj.weight'.format(key_prefix)].shape[2]
         unet_config["patch_size"] = patch_size
-        unet_config["out_channels"] = state_dict['{}final_layer.linear.weight'.format(key_prefix)].shape[0] // (patch_size * patch_size)
+        final_layer = '{}final_layer.linear.weight'.format(key_prefix)
+        if final_layer in state_dict:
+            unet_config["out_channels"] = state_dict[final_layer].shape[0] // (patch_size * patch_size)
 
         unet_config["depth"] = state_dict['{}x_embedder.proj.weight'.format(key_prefix)].shape[0] // 64
         unet_config["input_size"] = None
@@ -101,6 +105,19 @@ def detect_unet_config(state_dict, key_prefix):
         unet_config["audio_model"] = "dit1.0"
         return unet_config
 
+    if '{}double_layers.0.attn.w1q.weight'.format(key_prefix) in state_dict_keys: #aura flow dit
+        unet_config = {}
+        unet_config["max_seq"] = state_dict['{}positional_encoding'.format(key_prefix)].shape[1]
+        unet_config["cond_seq_dim"] = state_dict['{}cond_seq_linear.weight'.format(key_prefix)].shape[1]
+        double_layers = count_blocks(state_dict_keys, '{}double_layers.'.format(key_prefix) + '{}.')
+        single_layers = count_blocks(state_dict_keys, '{}single_layers.'.format(key_prefix) + '{}.')
+        unet_config["n_double_layers"] = double_layers
+        unet_config["n_layers"] = double_layers + single_layers
+        return unet_config
+
+    if '{}input_blocks.0.0.weight'.format(key_prefix) not in state_dict_keys:
+        return None
+
     unet_config = {
         "use_checkpoint": False,
         "image_size": 32,
@@ -235,6 +252,8 @@ def model_config_from_unet_config(unet_config, state_dict=None):
 
 def model_config_from_unet(state_dict, unet_key_prefix, use_base_if_no_match=False):
     unet_config = detect_unet_config(state_dict, unet_key_prefix)
+    if unet_config is None:
+        return None
     model_config = model_config_from_unet_config(unet_config, state_dict)
     if model_config is None and use_base_if_no_match:
         return comfy.supported_models_base.BASE(unet_config)
@@ -244,6 +263,8 @@ def model_config_from_unet(state_dict, unet_key_prefix, use_base_if_no_match=Fal
 def unet_prefix_from_state_dict(state_dict):
     if "model.model.postprocess_conv.weight" in state_dict: #audio models
         unet_key_prefix = "model.model."
+    elif "model.double_layers.0.attn.w1q.weight" in state_dict: #aura flow
+        unet_key_prefix = "model."
     else:
         unet_key_prefix = "model.diffusion_model."
     return unet_key_prefix
@@ -431,3 +452,47 @@ def model_config_from_diffusers_unet(state_dict):
     if unet_config is not None:
         return model_config_from_unet_config(unet_config)
     return None
+
+def convert_diffusers_mmdit(state_dict, output_prefix=""):
+    out_sd = {}
+
+    if 'transformer_blocks.0.attn.add_q_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)
+    elif 'joint_transformer_blocks.0.attn.add_k_proj.weight' in state_dict: #AuraFlow
+        num_joint = count_blocks(state_dict, 'joint_transformer_blocks.{}.')
+        num_single = count_blocks(state_dict, 'single_transformer_blocks.{}.')
+        sd_map = comfy.utils.auraflow_to_diffusers({"n_double_layers": num_joint, "n_layers": num_joint + num_single}, output_prefix=output_prefix)
+    else:
+        return None
+
+    for k in sd_map:
+        weight = state_dict.get(k, None)
+        if weight is not None:
+            t = sd_map[k]
+
+            if not isinstance(t, str):
+                if len(t) > 2:
+                    fun = t[2]
+                else:
+                    fun = lambda a: a
+                offset = t[1]
+                if offset is not None:
+                    old_weight = out_sd.get(t[0], None)
+                    if old_weight is None:
+                        old_weight = torch.empty_like(weight)
+                        old_weight = old_weight.repeat([3] + [1] * (len(old_weight.shape) - 1))
+
+                    w = old_weight.narrow(offset[0], offset[1], offset[2])
+                else:
+                    old_weight = weight
+                    w = weight
+                w[:] = fun(weight)
+                t = t[0]
+                out_sd[t] = old_weight
+            else:
+                out_sd[t] = weight
+            state_dict.pop(k)
+
+    return out_sd

comfy/model_patcher.py

@@ -51,6 +51,18 @@ def set_model_options_patch_replace(model_options, patch, name, block_name, numb
     model_options["transformer_options"] = to
     return model_options
 
+def set_model_options_post_cfg_function(model_options, post_cfg_function, disable_cfg1_optimization=False):
+    model_options["sampler_post_cfg_function"] = model_options.get("sampler_post_cfg_function", []) + [post_cfg_function]
+    if disable_cfg1_optimization:
+        model_options["disable_cfg1_optimization"] = True
+    return model_options
+
+def set_model_options_pre_cfg_function(model_options, pre_cfg_function, disable_cfg1_optimization=False):
+    model_options["sampler_pre_cfg_function"] = model_options.get("sampler_pre_cfg_function", []) + [pre_cfg_function]
+    if disable_cfg1_optimization:
+        model_options["disable_cfg1_optimization"] = True
+    return model_options
+
 class ModelPatcher:
     def __init__(self, model, load_device, offload_device, size=0, current_device=None, weight_inplace_update=False):
         self.size = size
@@ -122,9 +134,10 @@ class ModelPatcher:
             self.model_options["disable_cfg1_optimization"] = True
 
     def set_model_sampler_post_cfg_function(self, post_cfg_function, disable_cfg1_optimization=False):
-        self.model_options["sampler_post_cfg_function"] = self.model_options.get("sampler_post_cfg_function", []) + [post_cfg_function]
+        self.model_options = set_model_options_post_cfg_function(self.model_options, post_cfg_function, disable_cfg1_optimization)
-        if disable_cfg1_optimization:
+
-            self.model_options["disable_cfg1_optimization"] = True
+    def set_model_sampler_pre_cfg_function(self, pre_cfg_function, disable_cfg1_optimization=False):
+        self.model_options = set_model_options_pre_cfg_function(self.model_options, pre_cfg_function, disable_cfg1_optimization)
 
     def set_model_unet_function_wrapper(self, unet_wrapper_function: UnetWrapperFunction):
         self.model_options["model_function_wrapper"] = unet_wrapper_function
@@ -210,16 +223,19 @@ class ModelPatcher:
         model_sd = self.model.state_dict()
         for k in patches:
             offset = None
+            function = None
             if isinstance(k, str):
                 key = k
             else:
                 offset = k[1]
                 key = k[0]
+                if len(k) > 2:
+                    function = k[2]
 
             if key in model_sd:
                 p.add(k)
                 current_patches = self.patches.get(key, [])
-                current_patches.append((strength_patch, patches[k], strength_model, offset))
+                current_patches.append((strength_patch, patches[k], strength_model, offset, function))
                 self.patches[key] = current_patches
 
         self.patches_uuid = uuid.uuid4()
@@ -347,6 +363,9 @@ class ModelPatcher:
             v = p[1]
             strength_model = p[2]
             offset = p[3]
+            function = p[4]
+            if function is None:
+                function = lambda a: a
 
             old_weight = None
             if offset is not None:
@@ -371,7 +390,7 @@ class ModelPatcher:
                     if w1.shape != weight.shape:
                         logging.warning("WARNING SHAPE MISMATCH {} WEIGHT NOT MERGED {} != {}".format(key, w1.shape, weight.shape))
                     else:
-                        weight += strength * comfy.model_management.cast_to_device(w1, weight.device, weight.dtype)
+                        weight += function(strength * comfy.model_management.cast_to_device(w1, weight.device, weight.dtype))
             elif patch_type == "lora": #lora/locon
                 mat1 = comfy.model_management.cast_to_device(v[0], weight.device, torch.float32)
                 mat2 = comfy.model_management.cast_to_device(v[1], weight.device, torch.float32)
@@ -389,9 +408,9 @@ class ModelPatcher:
                 try:
                     lora_diff = torch.mm(mat1.flatten(start_dim=1), mat2.flatten(start_dim=1)).reshape(weight.shape)
                     if dora_scale is not None:
-                        weight = weight_decompose(dora_scale, weight, lora_diff, alpha, strength)
+                        weight = function(weight_decompose(dora_scale, weight, lora_diff, alpha, strength))
                     else:
-                        weight += ((strength * alpha) * lora_diff).type(weight.dtype)
+                        weight += function(((strength * alpha) * lora_diff).type(weight.dtype))
                 except Exception as e:
                     logging.error("ERROR {} {} {}".format(patch_type, key, e))
             elif patch_type == "lokr":
@@ -435,9 +454,9 @@ class ModelPatcher:
                 try:
                     lora_diff = torch.kron(w1, w2).reshape(weight.shape)
                     if dora_scale is not None:
-                        weight = weight_decompose(dora_scale, weight, lora_diff, alpha, strength)
+                        weight = function(weight_decompose(dora_scale, weight, lora_diff, alpha, strength))
                     else:
-                        weight += ((strength * alpha) * lora_diff).type(weight.dtype)
+                        weight += function(((strength * alpha) * lora_diff).type(weight.dtype))
                 except Exception as e:
                     logging.error("ERROR {} {} {}".format(patch_type, key, e))
             elif patch_type == "loha":
@@ -472,9 +491,9 @@ class ModelPatcher:
                 try:
                     lora_diff = (m1 * m2).reshape(weight.shape)
                     if dora_scale is not None:
-                        weight = weight_decompose(dora_scale, weight, lora_diff, alpha, strength)
+                        weight = function(weight_decompose(dora_scale, weight, lora_diff, alpha, strength))
                     else:
-                        weight += ((strength * alpha) * lora_diff).type(weight.dtype)
+                        weight += function(((strength * alpha) * lora_diff).type(weight.dtype))
                 except Exception as e:
                     logging.error("ERROR {} {} {}".format(patch_type, key, e))
             elif patch_type == "glora":
@@ -493,9 +512,9 @@ class ModelPatcher:
                 try:
                     lora_diff = (torch.mm(b2, b1) + torch.mm(torch.mm(weight.flatten(start_dim=1), a2), a1)).reshape(weight.shape)
                     if dora_scale is not None:
-                        weight = weight_decompose(dora_scale, weight, lora_diff, alpha, strength)
+                        weight = function(weight_decompose(dora_scale, weight, lora_diff, alpha, strength))
                     else:
-                        weight += ((strength * alpha) * lora_diff).type(weight.dtype)
+                        weight += function(((strength * alpha) * lora_diff).type(weight.dtype))
                 except Exception as e:
                     logging.error("ERROR {} {} {}".format(patch_type, key, e))
             else:

comfy/model_sampling.py

@@ -59,8 +59,9 @@ class ModelSamplingDiscrete(torch.nn.Module):
         beta_schedule = sampling_settings.get("beta_schedule", "linear")
         linear_start = sampling_settings.get("linear_start", 0.00085)
         linear_end = sampling_settings.get("linear_end", 0.012)
+        timesteps = sampling_settings.get("timesteps", 1000)
 
-        self._register_schedule(given_betas=None, beta_schedule=beta_schedule, timesteps=1000, linear_start=linear_start, linear_end=linear_end, cosine_s=8e-3)
+        self._register_schedule(given_betas=None, beta_schedule=beta_schedule, timesteps=timesteps, linear_start=linear_start, linear_end=linear_end, cosine_s=8e-3)
         self.sigma_data = 1.0
 
     def _register_schedule(self, given_betas=None, beta_schedule="linear", timesteps=1000,
@@ -190,11 +191,12 @@ class ModelSamplingDiscreteFlow(torch.nn.Module):
         else:
             sampling_settings = {}
 
-        self.set_parameters(shift=sampling_settings.get("shift", 1.0))
+        self.set_parameters(shift=sampling_settings.get("shift", 1.0), multiplier=sampling_settings.get("multiplier", 1000))
 
-    def set_parameters(self, shift=1.0, timesteps=1000):
+    def set_parameters(self, shift=1.0, timesteps=1000, multiplier=1000):
         self.shift = shift
-        ts = self.sigma(torch.arange(1, timesteps + 1, 1))
+        self.multiplier = multiplier
+        ts = self.sigma((torch.arange(1, timesteps + 1, 1) / timesteps) * multiplier)
         self.register_buffer('sigmas', ts)
 
     @property
@@ -206,10 +208,10 @@ class ModelSamplingDiscreteFlow(torch.nn.Module):
         return self.sigmas[-1]
 
     def timestep(self, sigma):
-        return sigma * 1000
+        return sigma * self.multiplier
 
     def sigma(self, timestep):
-        return time_snr_shift(self.shift, timestep / 1000)
+        return time_snr_shift(self.shift, timestep / self.multiplier)
 
     def percent_to_sigma(self, percent):
         if percent <= 0.0:

comfy/samplers.py

@@ -6,6 +6,8 @@ from comfy import model_management
 import math
 import logging
 import comfy.sampler_helpers
+import scipy
+import numpy
 
 def get_area_and_mult(conds, x_in, timestep_in):
     dims = tuple(x_in.shape[2:])
@@ -275,6 +277,12 @@ def sampling_function(model, x, timestep, uncond, cond, cond_scale, model_option
 
     conds = [cond, uncond_]
     out = calc_cond_batch(model, conds, x, timestep, model_options)
+
+    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)
+
     return cfg_function(model, out[0], out[1], cond_scale, x, timestep, model_options=model_options, cond=cond, uncond=uncond_)
 
 
@@ -305,13 +313,18 @@ def simple_scheduler(model_sampling, steps):
 def ddim_scheduler(model_sampling, steps):
     s = model_sampling
     sigs = []
-    ss = max(len(s.sigmas) // steps, 1)
     x = 1
+    if math.isclose(float(s.sigmas[x]), 0, abs_tol=0.00001):
+        steps += 1
+        sigs = []
+    else:
+        sigs = [0.0]
+
+    ss = max(len(s.sigmas) // steps, 1)
     while x < len(s.sigmas):
         sigs += [float(s.sigmas[x])]
         x += ss
     sigs = sigs[::-1]
-    sigs += [0.0]
     return torch.FloatTensor(sigs)
 
 def normal_scheduler(model_sampling, steps, sgm=False, floor=False):
@@ -319,15 +332,34 @@ def normal_scheduler(model_sampling, steps, sgm=False, floor=False):
     start = s.timestep(s.sigma_max)
     end = s.timestep(s.sigma_min)
 
+    append_zero = True
     if sgm:
         timesteps = torch.linspace(start, end, steps + 1)[:-1]
     else:
+        if math.isclose(float(s.sigma(end)), 0, abs_tol=0.00001):
+            steps += 1
+            append_zero = False
         timesteps = torch.linspace(start, end, steps)
 
     sigs = []
     for x in range(len(timesteps)):
         ts = timesteps[x]
-        sigs.append(s.sigma(ts))
+        sigs.append(float(s.sigma(ts)))
+
+    if append_zero:
+        sigs += [0.0]
+
+    return torch.FloatTensor(sigs)
+
+# Implemented based on: https://arxiv.org/abs/2407.12173
+def beta_scheduler(model_sampling, steps, alpha=0.6, beta=0.6):
+    total_timesteps = (len(model_sampling.sigmas) - 1)
+    ts = 1 - numpy.linspace(0, 1, steps, endpoint=False)
+    ts = numpy.rint(scipy.stats.beta.ppf(ts, alpha, beta) * total_timesteps)
+
+    sigs = []
+    for t in ts:
+        sigs += [float(model_sampling.sigmas[int(t)])]
     sigs += [0.0]
     return torch.FloatTensor(sigs)
 
@@ -537,9 +569,10 @@ class Sampler:
         sigma = float(sigmas[0])
         return math.isclose(max_sigma, sigma, rel_tol=1e-05) or sigma > max_sigma
 
-KSAMPLER_NAMES = ["euler", "euler_ancestral", "heun", "heunpp2","dpm_2", "dpm_2_ancestral",
+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_sde", "dpmpp_sde_gpu",
-                  "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm"]
+                  "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm",
+                  "ipndm", "ipndm_v", "deis"]
 
 class KSAMPLER(Sampler):
     def __init__(self, sampler_function, extra_options={}, inpaint_options={}):
@@ -696,7 +729,7 @@ def sample(model, noise, positive, negative, cfg, device, sampler, sigmas, model
     return cfg_guider.sample(noise, latent_image, sampler, sigmas, denoise_mask, callback, disable_pbar, seed)
 
 
-SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform"]
+SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform", "beta"]
 SAMPLER_NAMES = KSAMPLER_NAMES + ["ddim", "uni_pc", "uni_pc_bh2"]
 
 def calculate_sigmas(model_sampling, scheduler_name, steps):
@@ -712,6 +745,8 @@ def calculate_sigmas(model_sampling, scheduler_name, steps):
         sigmas = ddim_scheduler(model_sampling, steps)
     elif scheduler_name == "sgm_uniform":
         sigmas = normal_scheduler(model_sampling, steps, sgm=True)
+    elif scheduler_name == "beta":
+        sigmas = beta_scheduler(model_sampling, steps)
     else:
         logging.error("error invalid scheduler {}".format(scheduler_name))
     return sigmas

comfy/sd.py

@@ -19,8 +19,9 @@ from . import model_detection
 from . import sd1_clip
 from . import sd2_clip
 from . import sdxl_clip
-from . import sd3_clip
+import comfy.text_encoders.sd3_clip
-from . import sa_t5
+import comfy.text_encoders.sa_t5
+import comfy.text_encoders.aura_t5
 
 import comfy.model_patcher
 import comfy.lora
@@ -28,36 +29,6 @@ import comfy.t2i_adapter.adapter
 import comfy.supported_models_base
 import comfy.taesd.taesd
 
-def load_model_weights(model, sd):
-    m, u = model.load_state_dict(sd, strict=False)
-    m = set(m)
-    unexpected_keys = set(u)
-
-    k = list(sd.keys())
-    for x in k:
-        if x not in unexpected_keys:
-            w = sd.pop(x)
-            del w
-    if len(m) > 0:
-        logging.warning("missing {}".format(m))
-    return model
-
-def load_clip_weights(model, sd):
-    k = list(sd.keys())
-    for x in k:
-        if x.startswith("cond_stage_model.transformer.") and not x.startswith("cond_stage_model.transformer.text_model."):
-            y = x.replace("cond_stage_model.transformer.", "cond_stage_model.transformer.text_model.")
-            sd[y] = sd.pop(x)
-
-    if 'cond_stage_model.transformer.text_model.embeddings.position_ids' in sd:
-        ids = sd['cond_stage_model.transformer.text_model.embeddings.position_ids']
-        if ids.dtype == torch.float32:
-            sd['cond_stage_model.transformer.text_model.embeddings.position_ids'] = ids.round()
-
-    sd = comfy.utils.clip_text_transformers_convert(sd, "cond_stage_model.model.", "cond_stage_model.transformer.")
-    return load_model_weights(model, sd)
-
-
 def load_lora_for_models(model, clip, lora, strength_model, strength_clip):
     key_map = {}
     if model is not None:
@@ -130,7 +101,7 @@ class CLIP:
     def tokenize(self, text, return_word_ids=False):
         return self.tokenizer.tokenize_with_weights(text, return_word_ids)
 
-    def encode_from_tokens(self, tokens, return_pooled=False):
+    def encode_from_tokens(self, tokens, return_pooled=False, return_dict=False):
         self.cond_stage_model.reset_clip_options()
 
         if self.layer_idx is not None:
@@ -140,7 +111,15 @@ class CLIP:
             self.cond_stage_model.set_clip_options({"projected_pooled": False})
 
         self.load_model()
-        cond, pooled = self.cond_stage_model.encode_token_weights(tokens)
+        o = self.cond_stage_model.encode_token_weights(tokens)
+        cond, pooled = o[:2]
+        if return_dict:
+            out = {"cond": cond, "pooled_output": pooled}
+            if len(o) > 2:
+                for k in o[2]:
+                    out[k] = o[2][k]
+            return out
+
         if return_pooled:
             return cond, pooled
         return cond
@@ -236,7 +215,7 @@ class VAE:
                     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.0.weight_v" in sd:
+            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)
                 self.memory_used_decode = lambda shape, dtype: (1000 * shape[2] * 2048) * model_management.dtype_size(dtype)
@@ -298,6 +277,10 @@ class VAE:
             / 3.0)
         return output
 
+    def decode_tiled_1d(self, samples, tile_x=128, overlap=32):
+        decode_fn = lambda a: self.first_stage_model.decode(a.to(self.vae_dtype).to(self.device)).float()
+        return comfy.utils.tiled_scale_multidim(samples, decode_fn, tile=(tile_x,), overlap=overlap, upscale_amount=self.upscale_ratio, out_channels=self.output_channels, output_device=self.output_device)
+
     def encode_tiled_(self, pixel_samples, tile_x=512, tile_y=512, overlap = 64):
         steps = pixel_samples.shape[0] * comfy.utils.get_tiled_scale_steps(pixel_samples.shape[3], pixel_samples.shape[2], tile_x, tile_y, overlap)
         steps += pixel_samples.shape[0] * comfy.utils.get_tiled_scale_steps(pixel_samples.shape[3], pixel_samples.shape[2], tile_x // 2, tile_y * 2, overlap)
@@ -311,6 +294,10 @@ class VAE:
         samples /= 3.0
         return samples
 
+    def encode_tiled_1d(self, samples, tile_x=128 * 2048, overlap=32 * 2048):
+        encode_fn = lambda a: self.first_stage_model.encode((self.process_input(a)).to(self.vae_dtype).to(self.device)).float()
+        return comfy.utils.tiled_scale_multidim(samples, encode_fn, tile=(tile_x,), overlap=overlap, upscale_amount=(1/self.downscale_ratio), out_channels=self.latent_channels, output_device=self.output_device)
+
     def decode(self, samples_in):
         try:
             memory_used = self.memory_used_decode(samples_in.shape, self.vae_dtype)
@@ -325,7 +312,10 @@ class VAE:
                 pixel_samples[x:x+batch_number] = self.process_output(self.first_stage_model.decode(samples).to(self.output_device).float())
         except model_management.OOM_EXCEPTION as e:
             logging.warning("Warning: Ran out of memory when regular VAE decoding, retrying with tiled VAE decoding.")
-            pixel_samples = self.decode_tiled_(samples_in)
+            if len(samples_in.shape) == 3:
+                pixel_samples = self.decode_tiled_1d(samples_in)
+            else:
+                pixel_samples = self.decode_tiled_(samples_in)
 
         pixel_samples = pixel_samples.to(self.output_device).movedim(1,-1)
         return pixel_samples
@@ -351,7 +341,10 @@ class VAE:
 
         except model_management.OOM_EXCEPTION as e:
             logging.warning("Warning: Ran out of memory when regular VAE encoding, retrying with tiled VAE encoding.")
-            samples = self.encode_tiled_(pixel_samples)
+            if len(pixel_samples.shape) == 3:
+                samples = self.encode_tiled_1d(pixel_samples)
+            else:
+                samples = self.encode_tiled_(pixel_samples)
 
         return samples
 
@@ -418,25 +411,30 @@ def load_clip(ckpt_paths, embedding_directory=None, clip_type=CLIPType.STABLE_DI
             clip_target.clip = sd2_clip.SD2ClipModel
             clip_target.tokenizer = sd2_clip.SD2Tokenizer
         elif "encoder.block.23.layer.1.DenseReluDense.wi_1.weight" in clip_data[0]:
-            dtype_t5 = clip_data[0]["encoder.block.23.layer.1.DenseReluDense.wi_1.weight"].dtype
+            weight = clip_data[0]["encoder.block.23.layer.1.DenseReluDense.wi_1.weight"]
-            clip_target.clip = sd3_clip.sd3_clip(clip_l=False, clip_g=False, t5=True, dtype_t5=dtype_t5)
+            dtype_t5 = weight.dtype
-            clip_target.tokenizer = sd3_clip.SD3Tokenizer
+            if weight.shape[-1] == 4096:
+                clip_target.clip = comfy.text_encoders.sd3_clip.sd3_clip(clip_l=False, clip_g=False, t5=True, dtype_t5=dtype_t5)
+                clip_target.tokenizer = comfy.text_encoders.sd3_clip.SD3Tokenizer
+            elif weight.shape[-1] == 2048:
+                clip_target.clip = comfy.text_encoders.aura_t5.AuraT5Model
+                clip_target.tokenizer = comfy.text_encoders.aura_t5.AuraT5Tokenizer
         elif "encoder.block.0.layer.0.SelfAttention.k.weight" in clip_data[0]:
-            clip_target.clip = sa_t5.SAT5Model
+            clip_target.clip = comfy.text_encoders.sa_t5.SAT5Model
-            clip_target.tokenizer = sa_t5.SAT5Tokenizer
+            clip_target.tokenizer = comfy.text_encoders.sa_t5.SAT5Tokenizer
         else:
             clip_target.clip = sd1_clip.SD1ClipModel
             clip_target.tokenizer = sd1_clip.SD1Tokenizer
     elif len(clip_data) == 2:
         if clip_type == CLIPType.SD3:
-            clip_target.clip = sd3_clip.sd3_clip(clip_l=True, clip_g=True, t5=False)
+            clip_target.clip = comfy.text_encoders.sd3_clip.sd3_clip(clip_l=True, clip_g=True, t5=False)
-            clip_target.tokenizer = sd3_clip.SD3Tokenizer
+            clip_target.tokenizer = comfy.text_encoders.sd3_clip.SD3Tokenizer
         else:
             clip_target.clip = sdxl_clip.SDXLClipModel
             clip_target.tokenizer = sdxl_clip.SDXLTokenizer
     elif len(clip_data) == 3:
-        clip_target.clip = sd3_clip.SD3ClipModel
+        clip_target.clip = comfy.text_encoders.sd3_clip.SD3ClipModel
-        clip_target.tokenizer = sd3_clip.SD3Tokenizer
+        clip_target.tokenizer = comfy.text_encoders.sd3_clip.SD3Tokenizer
 
     clip = CLIP(clip_target, embedding_directory=embedding_directory)
     for c in clip_data:
@@ -495,13 +493,13 @@ def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, o
     load_device = model_management.get_torch_device()
 
     model_config = model_detection.model_config_from_unet(sd, diffusion_model_prefix)
+    if model_config is None:
+        raise RuntimeError("ERROR: Could not detect model type of: {}".format(ckpt_path))
+
     unet_dtype = model_management.unet_dtype(model_params=parameters, supported_dtypes=model_config.supported_inference_dtypes)
     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)
 
-    if model_config is None:
-        raise RuntimeError("ERROR: Could not detect model type of: {}".format(ckpt_path))
-
     if model_config.clip_vision_prefix is not None:
         if output_clipvision:
             clipvision = clip_vision.load_clipvision_from_sd(sd, model_config.clip_vision_prefix, True)
@@ -549,30 +547,40 @@ def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, o
     return (model_patcher, clip, vae, clipvision)
 
 
-def load_unet_state_dict(sd): #load unet in diffusers format
+def load_unet_state_dict(sd): #load unet in diffusers or regular format
+
+    #Allow loading unets from checkpoint files
+    diffusion_model_prefix = model_detection.unet_prefix_from_state_dict(sd)
+    temp_sd = comfy.utils.state_dict_prefix_replace(sd, {diffusion_model_prefix: ""}, filter_keys=True)
+    if len(temp_sd) > 0:
+        sd = temp_sd
+
     parameters = comfy.utils.calculate_parameters(sd)
     unet_dtype = model_management.unet_dtype(model_params=parameters)
     load_device = model_management.get_torch_device()
+    model_config = model_detection.model_config_from_unet(sd, "")
 
-    if "input_blocks.0.0.weight" in sd or 'clf.1.weight' in sd: #ldm or stable cascade
+    if model_config is not None:
-        model_config = model_detection.model_config_from_unet(sd, "")
-        if model_config is None:
-            return None
         new_sd = sd
+    else:
+        new_sd = model_detection.convert_diffusers_mmdit(sd, "")
+        if new_sd is not None: #diffusers mmdit
+            model_config = model_detection.model_config_from_unet(new_sd, "")
+            if model_config is None:
+                return None
+        else: #diffusers unet
+            model_config = model_detection.model_config_from_diffusers_unet(sd)
+            if model_config is None:
+                return None
 
-    else: #diffusers
+            diffusers_keys = comfy.utils.unet_to_diffusers(model_config.unet_config)
-        model_config = model_detection.model_config_from_diffusers_unet(sd)
-        if model_config is None:
-            return None
 
-        diffusers_keys = comfy.utils.unet_to_diffusers(model_config.unet_config)
+            new_sd = {}
-
+            for k in diffusers_keys:
-        new_sd = {}
+                if k in sd:
-        for k in diffusers_keys:
+                    new_sd[diffusers_keys[k]] = sd.pop(k)
-            if k in sd:
+                else:
-                new_sd[diffusers_keys[k]] = sd.pop(k)
+                    logging.warning("{} {}".format(diffusers_keys[k], k))
-            else:
-                logging.warning("{} {}".format(diffusers_keys[k], k))
 
     offload_device = model_management.unet_offload_device()
     unet_dtype = model_management.unet_dtype(model_params=parameters, supported_dtypes=model_config.supported_inference_dtypes)
@@ -607,4 +615,9 @@ def save_checkpoint(output_path, model, clip=None, vae=None, clip_vision=None, m
     for k in extra_keys:
         sd[k] = extra_keys[k]
 
+    for k in sd:
+        t = sd[k]
+        if not t.is_contiguous():
+            sd[k] = t.contiguous()
+
     comfy.utils.save_torch_file(sd, output_path, metadata=metadata)

comfy/sd1_clip.py

@@ -9,6 +9,7 @@ from . import model_management
 import comfy.clip_model
 import json
 import logging
+import numbers
 
 def gen_empty_tokens(special_tokens, length):
     start_token = special_tokens.get("start", None)
@@ -37,7 +38,9 @@ class ClipTokenWeightEncoder:
         if has_weights or sections == 0:
             to_encode.append(gen_empty_tokens(self.special_tokens, max_token_len))
 
-        out, pooled = self.encode(to_encode)
+        o = self.encode(to_encode)
+        out, pooled = o[:2]
+
         if pooled is not None:
             first_pooled = pooled[0:1].to(model_management.intermediate_device())
         else:
@@ -56,8 +59,20 @@ class ClipTokenWeightEncoder:
             output.append(z)
 
         if (len(output) == 0):
-            return out[-1:].to(model_management.intermediate_device()), first_pooled
+            r = (out[-1:].to(model_management.intermediate_device()), first_pooled)
-        return torch.cat(output, dim=-2).to(model_management.intermediate_device()), first_pooled
+        else:
+            r = (torch.cat(output, dim=-2).to(model_management.intermediate_device()), first_pooled)
+
+        if len(o) > 2:
+            extra = {}
+            for k in o[2]:
+                v = o[2][k]
+                if k == "attention_mask":
+                    v = v[:sections].flatten().unsqueeze(dim=0).to(model_management.intermediate_device())
+                extra[k] = v
+
+            r = r + (extra,)
+        return r
 
 class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
     """Uses the CLIP transformer encoder for text (from huggingface)"""
@@ -69,7 +84,7 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
     def __init__(self, version="openai/clip-vit-large-patch14", device="cpu", max_length=77,
                  freeze=True, layer="last", layer_idx=None, textmodel_json_config=None, dtype=None, model_class=comfy.clip_model.CLIPTextModel,
                  special_tokens={"start": 49406, "end": 49407, "pad": 49407}, layer_norm_hidden_state=True, enable_attention_masks=False, zero_out_masked=False,
-                 return_projected_pooled=True):  # clip-vit-base-patch32
+                 return_projected_pooled=True, return_attention_masks=False):  # clip-vit-base-patch32
         super().__init__()
         assert layer in self.LAYERS
 
@@ -95,6 +110,7 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
 
         self.layer_norm_hidden_state = layer_norm_hidden_state
         self.return_projected_pooled = return_projected_pooled
+        self.return_attention_masks = return_attention_masks
 
         if layer == "hidden":
             assert layer_idx is not None
@@ -130,10 +146,10 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
         for x in tokens:
             tokens_temp = []
             for y in x:
-                if isinstance(y, int):
+                if isinstance(y, numbers.Integral):
                     if y == token_dict_size: #EOS token
                         y = -1
-                    tokens_temp += [y]
+                    tokens_temp += [int(y)]
                 else:
                     if y.shape[0] == current_embeds.weight.shape[1]:
                         embedding_weights += [y]
@@ -168,7 +184,7 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
         tokens = torch.LongTensor(tokens).to(device)
 
         attention_mask = None
-        if self.enable_attention_masks:
+        if self.enable_attention_masks or self.zero_out_masked or self.return_attention_masks:
             attention_mask = torch.zeros_like(tokens)
             end_token = self.special_tokens.get("end", -1)
             for x in range(attention_mask.shape[0]):
@@ -177,7 +193,11 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
                     if tokens[x, y] == end_token:
                         break
 
-        outputs = self.transformer(tokens, attention_mask, intermediate_output=self.layer_idx, final_layer_norm_intermediate=self.layer_norm_hidden_state)
+        attention_mask_model = None
+        if self.enable_attention_masks:
+            attention_mask_model = attention_mask
+
+        outputs = self.transformer(tokens, attention_mask_model, intermediate_output=self.layer_idx, final_layer_norm_intermediate=self.layer_norm_hidden_state)
         self.transformer.set_input_embeddings(backup_embeds)
 
         if self.layer == "last":
@@ -185,7 +205,7 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
         else:
             z = outputs[1].float()
 
-        if self.zero_out_masked and attention_mask is not None:
+        if self.zero_out_masked:
             z *= attention_mask.unsqueeze(-1).float()
 
         pooled_output = None
@@ -195,6 +215,13 @@ class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
             elif outputs[2] is not None:
                 pooled_output = outputs[2].float()
 
+        extra = {}
+        if self.return_attention_masks:
+            extra["attention_mask"] = attention_mask
+
+        if len(extra) > 0:
+            return z, pooled_output, extra
+
         return z, pooled_output
 
     def encode(self, tokens):
@@ -359,7 +386,7 @@ def load_embed(embedding_name, embedding_directory, embedding_size, embed_key=No
     return embed_out
 
 class SDTokenizer:
-    def __init__(self, tokenizer_path=None, max_length=77, pad_with_end=True, embedding_directory=None, embedding_size=768, embedding_key='clip_l', tokenizer_class=CLIPTokenizer, has_start_token=True, pad_to_max_length=True, min_length=None):
+    def __init__(self, tokenizer_path=None, max_length=77, pad_with_end=True, embedding_directory=None, embedding_size=768, embedding_key='clip_l', tokenizer_class=CLIPTokenizer, has_start_token=True, pad_to_max_length=True, min_length=None, pad_token=None):
         if tokenizer_path is None:
             tokenizer_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "sd1_tokenizer")
         self.tokenizer = tokenizer_class.from_pretrained(tokenizer_path)
@@ -375,6 +402,14 @@ class SDTokenizer:
             self.tokens_start = 0
             self.start_token = None
             self.end_token = empty[0]
+
+        if pad_token is not None:
+            self.pad_token = pad_token
+        elif pad_with_end:
+            self.pad_token = self.end_token
+        else:
+            self.pad_token = 0
+
         self.pad_with_end = pad_with_end
         self.pad_to_max_length = pad_to_max_length
 
@@ -407,10 +442,6 @@ class SDTokenizer:
         Word id values are unique per word and embedding, where the id 0 is reserved for non word tokens.
         Returned list has the dimensions NxM where M is the input size of CLIP
         '''
-        if self.pad_with_end:
-            pad_token = self.end_token
-        else:
-            pad_token = 0
 
         text = escape_important(text)
         parsed_weights = token_weights(text, 1.0)
@@ -462,7 +493,7 @@ class SDTokenizer:
                     else:
                         batch.append((self.end_token, 1.0, 0))
                         if self.pad_to_max_length:
-                            batch.extend([(pad_token, 1.0, 0)] * (remaining_length))
+                            batch.extend([(self.pad_token, 1.0, 0)] * (remaining_length))
                     #start new batch
                     batch = []
                     if self.start_token is not None:
@@ -475,9 +506,9 @@ class SDTokenizer:
         #fill last batch
         batch.append((self.end_token, 1.0, 0))
         if self.pad_to_max_length:
-            batch.extend([(pad_token, 1.0, 0)] * (self.max_length - len(batch)))
+            batch.extend([(self.pad_token, 1.0, 0)] * (self.max_length - len(batch)))
         if self.min_length is not None and len(batch) < self.min_length:
-            batch.extend([(pad_token, 1.0, 0)] * (self.min_length - len(batch)))
+            batch.extend([(self.pad_token, 1.0, 0)] * (self.min_length - len(batch)))
 
         if not return_word_ids:
             batched_tokens = [[(t, w) for t, w,_ in x] for x in batched_tokens]
@@ -505,10 +536,16 @@ class SD1Tokenizer:
 
 
 class SD1ClipModel(torch.nn.Module):
-    def __init__(self, device="cpu", dtype=None, clip_name="l", clip_model=SDClipModel, **kwargs):
+    def __init__(self, device="cpu", dtype=None, clip_name="l", clip_model=SDClipModel, name=None, **kwargs):
         super().__init__()
-        self.clip_name = clip_name
+
-        self.clip = "clip_{}".format(self.clip_name)
+        if name is not None:
+            self.clip_name = name
+            self.clip = "{}".format(self.clip_name)
+        else:
+            self.clip_name = clip_name
+            self.clip = "clip_{}".format(self.clip_name)
+
         setattr(self, self.clip, clip_model(device=device, dtype=dtype, **kwargs))
 
         self.dtypes = set()
@@ -523,8 +560,8 @@ class SD1ClipModel(torch.nn.Module):
 
     def encode_token_weights(self, token_weight_pairs):
         token_weight_pairs = token_weight_pairs[self.clip_name]
-        out, pooled = getattr(self, self.clip).encode_token_weights(token_weight_pairs)
+        out = getattr(self, self.clip).encode_token_weights(token_weight_pairs)
-        return out, pooled
+        return out
 
     def load_sd(self, sd):
         return getattr(self, self.clip).load_sd(sd)

comfy/supported_models.py

@@ -5,8 +5,9 @@ from . import utils
 from . import sd1_clip
 from . import sd2_clip
 from . import sdxl_clip
-from . import sd3_clip
+import comfy.text_encoders.sd3_clip
-from . import sa_t5
+import comfy.text_encoders.sa_t5
+import comfy.text_encoders.aura_t5
 
 from . import supported_models_base
 from . import latent_formats
@@ -523,7 +524,7 @@ class SD3(supported_models_base.BASE):
             t5 = True
             dtype_t5 = state_dict[t5_key].dtype
 
-        return supported_models_base.ClipTarget(sd3_clip.SD3Tokenizer, sd3_clip.sd3_clip(clip_l=clip_l, clip_g=clip_g, t5=t5, dtype_t5=dtype_t5))
+        return supported_models_base.ClipTarget(comfy.text_encoders.sd3_clip.SD3Tokenizer, comfy.text_encoders.sd3_clip.sd3_clip(clip_l=clip_l, clip_g=clip_g, t5=t5, dtype_t5=dtype_t5))
 
 class StableAudio(supported_models_base.BASE):
     unet_config = {
@@ -543,17 +544,42 @@ class StableAudio(supported_models_base.BASE):
         seconds_total_sd = utils.state_dict_prefix_replace(state_dict, {"conditioner.conditioners.seconds_total.": ""}, filter_keys=True)
         return model_base.StableAudio1(self, seconds_start_embedder_weights=seconds_start_sd, seconds_total_embedder_weights=seconds_total_sd, device=device)
 
-
     def process_unet_state_dict(self, state_dict):
         for k in list(state_dict.keys()):
             if k.endswith(".cross_attend_norm.beta") or k.endswith(".ff_norm.beta") or k.endswith(".pre_norm.beta"): #These weights are all zero
                 state_dict.pop(k)
         return state_dict
 
+    def process_unet_state_dict_for_saving(self, state_dict):
+        replace_prefix = {"": "model.model."}
+        return utils.state_dict_prefix_replace(state_dict, replace_prefix)
+
     def clip_target(self, state_dict={}):
-        return supported_models_base.ClipTarget(sa_t5.SAT5Tokenizer, sa_t5.SAT5Model)
+        return supported_models_base.ClipTarget(comfy.text_encoders.sa_t5.SAT5Tokenizer, comfy.text_encoders.sa_t5.SAT5Model)
 
+class AuraFlow(supported_models_base.BASE):
+    unet_config = {
+        "cond_seq_dim": 2048,
+    }
 
-models = [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]
+    sampling_settings = {
+        "multiplier": 1.0,
+        "shift": 1.73,
+    }
+
+    unet_extra_config = {}
+    latent_format = latent_formats.SDXL
+
+    vae_key_prefix = ["vae."]
+    text_encoder_key_prefix = ["text_encoders."]
+
+    def get_model(self, state_dict, prefix="", device=None):
+        out = model_base.AuraFlow(self, device=device)
+        return out
+
+    def clip_target(self, state_dict={}):
+        return supported_models_base.ClipTarget(comfy.text_encoders.aura_t5.AuraT5Tokenizer, comfy.text_encoders.aura_t5.AuraT5Model)
+
+models = [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]
 
 models += [SVD_img2vid]

comfy/t2i_adapter/adapter.py

@@ -153,7 +153,13 @@ class Adapter(nn.Module):
                 features.append(None)
             features.append(x)
 
-        return features
+        features = features[::-1]
+
+        if self.xl:
+            return {"input": features[1:], "middle": features[:1]}
+        else:
+            return {"input": features}
+
 
 
 class LayerNorm(nn.LayerNorm):
@@ -290,4 +296,4 @@ class Adapter_light(nn.Module):
             features.append(None)
             features.append(x)
 
-        return features
+        return {"input": features[::-1]}

comfy/text_encoders/aura_t5.py

@@ -0,0 +1,22 @@
+from comfy import sd1_clip
+from .llama_tokenizer import LLAMATokenizer
+import comfy.text_encoders.t5
+import os
+
+class PT5XlModel(sd1_clip.SDClipModel):
+    def __init__(self, device="cpu", layer="last", layer_idx=None, dtype=None):
+        textmodel_json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "t5_pile_config_xl.json")
+        super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config=textmodel_json_config, dtype=dtype, special_tokens={"end": 2, "pad": 1}, model_class=comfy.text_encoders.t5.T5, enable_attention_masks=True, zero_out_masked=True)
+
+class PT5XlTokenizer(sd1_clip.SDTokenizer):
+    def __init__(self, embedding_directory=None):
+        tokenizer_path = os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "t5_pile_tokenizer"), "tokenizer.model")
+        super().__init__(tokenizer_path, pad_with_end=False, embedding_size=2048, embedding_key='pile_t5xl', tokenizer_class=LLAMATokenizer, has_start_token=False, pad_to_max_length=False, max_length=99999999, min_length=256, pad_token=1)
+
+class AuraT5Tokenizer(sd1_clip.SD1Tokenizer):
+    def __init__(self, embedding_directory=None):
+        super().__init__(embedding_directory=embedding_directory, clip_name="pile_t5xl", tokenizer=PT5XlTokenizer)
+
+class AuraT5Model(sd1_clip.SD1ClipModel):
+    def __init__(self, device="cpu", dtype=None, **kwargs):
+        super().__init__(device=device, dtype=dtype, name="pile_t5xl", clip_model=PT5XlModel, **kwargs)

comfy/text_encoders/llama_tokenizer.py

@@ -0,0 +1,22 @@
+import os
+
+class LLAMATokenizer:
+    @staticmethod
+    def from_pretrained(path):
+        return LLAMATokenizer(path)
+
+    def __init__(self, tokenizer_path):
+        import sentencepiece
+        self.tokenizer = sentencepiece.SentencePieceProcessor(model_file=tokenizer_path)
+        self.end = self.tokenizer.eos_id()
+
+    def get_vocab(self):
+        out = {}
+        for i in range(self.tokenizer.get_piece_size()):
+            out[self.tokenizer.id_to_piece(i)] = i
+        return out
+
+    def __call__(self, string):
+        out = self.tokenizer.encode(string)
+        out += [self.end]
+        return {"input_ids": out}

comfy/text_encoders/sa_t5.py

@@ -1,12 +1,12 @@
 from comfy import sd1_clip
 from transformers import T5TokenizerFast
-import comfy.t5
+import comfy.text_encoders.t5
 import os
 
 class T5BaseModel(sd1_clip.SDClipModel):
     def __init__(self, device="cpu", layer="last", layer_idx=None, dtype=None):
         textmodel_json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "t5_config_base.json")
-        super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config=textmodel_json_config, dtype=dtype, special_tokens={"end": 1, "pad": 0}, model_class=comfy.t5.T5, enable_attention_masks=True, zero_out_masked=True)
+        super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config=textmodel_json_config, dtype=dtype, special_tokens={"end": 1, "pad": 0}, model_class=comfy.text_encoders.t5.T5, enable_attention_masks=True, zero_out_masked=True)
 
 class T5BaseTokenizer(sd1_clip.SDTokenizer):
     def __init__(self, embedding_directory=None):
@@ -19,4 +19,4 @@ class SAT5Tokenizer(sd1_clip.SD1Tokenizer):
 
 class SAT5Model(sd1_clip.SD1ClipModel):
     def __init__(self, device="cpu", dtype=None, **kwargs):
-        super().__init__(device=device, dtype=dtype, clip_name="t5base", clip_model=T5BaseModel, **kwargs)
+        super().__init__(device=device, dtype=dtype, name="t5base", clip_model=T5BaseModel, **kwargs)

comfy/text_encoders/sd3_clip.py

@@ -1,7 +1,7 @@
 from comfy import sd1_clip
 from comfy import sdxl_clip
 from transformers import T5TokenizerFast
-import comfy.t5
+import comfy.text_encoders.t5
 import torch
 import os
 import comfy.model_management
@@ -10,7 +10,7 @@ import logging
 class T5XXLModel(sd1_clip.SDClipModel):
     def __init__(self, device="cpu", layer="last", layer_idx=None, dtype=None):
         textmodel_json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "t5_config_xxl.json")
-        super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config=textmodel_json_config, dtype=dtype, special_tokens={"end": 1, "pad": 0}, model_class=comfy.t5.T5)
+        super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config=textmodel_json_config, dtype=dtype, special_tokens={"end": 1, "pad": 0}, model_class=comfy.text_encoders.t5.T5)
 
 class T5XXLTokenizer(sd1_clip.SDTokenizer):
     def __init__(self, embedding_directory=None):

comfy/text_encoders/t5.py

@@ -13,29 +13,36 @@ class T5LayerNorm(torch.nn.Module):
         x = x * torch.rsqrt(variance + self.variance_epsilon)
         return self.weight.to(device=x.device, dtype=x.dtype) * x
 
+activations = {
+    "gelu_pytorch_tanh": lambda a: torch.nn.functional.gelu(a, approximate="tanh"),
+    "relu": torch.nn.functional.relu,
+}
+
 class T5DenseActDense(torch.nn.Module):
-    def __init__(self, model_dim, ff_dim, dtype, device, operations):
+    def __init__(self, model_dim, ff_dim, ff_activation, dtype, device, operations):
         super().__init__()
         self.wi = operations.Linear(model_dim, ff_dim, bias=False, dtype=dtype, device=device)
         self.wo = operations.Linear(ff_dim, model_dim, bias=False, dtype=dtype, device=device)
         # self.dropout = nn.Dropout(config.dropout_rate)
+        self.act = activations[ff_activation]
 
     def forward(self, x):
-        x = torch.nn.functional.relu(self.wi(x))
+        x = self.act(self.wi(x))
         # x = self.dropout(x)
         x = self.wo(x)
         return x
 
 class T5DenseGatedActDense(torch.nn.Module):
-    def __init__(self, model_dim, ff_dim, dtype, device, operations):
+    def __init__(self, model_dim, ff_dim, ff_activation, dtype, device, operations):
         super().__init__()
         self.wi_0 = operations.Linear(model_dim, ff_dim, bias=False, dtype=dtype, device=device)
         self.wi_1 = operations.Linear(model_dim, ff_dim, bias=False, dtype=dtype, device=device)
         self.wo = operations.Linear(ff_dim, model_dim, bias=False, dtype=dtype, device=device)
         # self.dropout = nn.Dropout(config.dropout_rate)
+        self.act = activations[ff_activation]
 
     def forward(self, x):
-        hidden_gelu = torch.nn.functional.gelu(self.wi_0(x), approximate="tanh")
+        hidden_gelu = self.act(self.wi_0(x))
         hidden_linear = self.wi_1(x)
         x = hidden_gelu * hidden_linear
         # x = self.dropout(x)
@@ -43,12 +50,12 @@ class T5DenseGatedActDense(torch.nn.Module):
         return x
 
 class T5LayerFF(torch.nn.Module):
-    def __init__(self, model_dim, ff_dim, ff_activation, dtype, device, operations):
+    def __init__(self, model_dim, ff_dim, ff_activation, gated_act, dtype, device, operations):
         super().__init__()
-        if ff_activation == "gelu_pytorch_tanh":
+        if gated_act:
-            self.DenseReluDense = T5DenseGatedActDense(model_dim, ff_dim, dtype, device, operations)
+            self.DenseReluDense = T5DenseGatedActDense(model_dim, ff_dim, ff_activation, dtype, device, operations)
-        elif ff_activation == "relu":
+        else:
-            self.DenseReluDense = T5DenseActDense(model_dim, ff_dim, dtype, device, operations)
+            self.DenseReluDense = T5DenseActDense(model_dim, ff_dim, ff_activation, dtype, device, operations)
 
         self.layer_norm = T5LayerNorm(model_dim, dtype=dtype, device=device, operations=operations)
         # self.dropout = nn.Dropout(config.dropout_rate)
@@ -171,11 +178,11 @@ class T5LayerSelfAttention(torch.nn.Module):
         return x, past_bias
 
 class T5Block(torch.nn.Module):
-    def __init__(self, model_dim, inner_dim, ff_dim, ff_activation, num_heads, relative_attention_bias, dtype, device, operations):
+    def __init__(self, model_dim, inner_dim, ff_dim, ff_activation, gated_act, num_heads, relative_attention_bias, dtype, device, operations):
         super().__init__()
         self.layer = torch.nn.ModuleList()
         self.layer.append(T5LayerSelfAttention(model_dim, inner_dim, ff_dim, num_heads, relative_attention_bias, dtype, device, operations))
-        self.layer.append(T5LayerFF(model_dim, ff_dim, ff_activation, dtype, device, operations))
+        self.layer.append(T5LayerFF(model_dim, ff_dim, ff_activation, gated_act, dtype, device, operations))
 
     def forward(self, x, mask=None, past_bias=None, optimized_attention=None):
         x, past_bias = self.layer[0](x, mask, past_bias, optimized_attention)
@@ -183,11 +190,11 @@ class T5Block(torch.nn.Module):
         return x, past_bias
 
 class T5Stack(torch.nn.Module):
-    def __init__(self, num_layers, model_dim, inner_dim, ff_dim, ff_activation, num_heads, dtype, device, operations):
+    def __init__(self, num_layers, model_dim, inner_dim, ff_dim, ff_activation, gated_act, num_heads, relative_attention, dtype, device, operations):
         super().__init__()
 
         self.block = torch.nn.ModuleList(
-            [T5Block(model_dim, inner_dim, ff_dim, ff_activation, num_heads, relative_attention_bias=(i == 0), dtype=dtype, device=device, operations=operations) for i in range(num_layers)]
+            [T5Block(model_dim, inner_dim, ff_dim, ff_activation, gated_act, num_heads, relative_attention_bias=((not relative_attention) or (i == 0)), dtype=dtype, device=device, operations=operations) for i in range(num_layers)]
         )
         self.final_layer_norm = T5LayerNorm(model_dim, dtype=dtype, device=device, operations=operations)
         # self.dropout = nn.Dropout(config.dropout_rate)
@@ -216,7 +223,7 @@ class T5(torch.nn.Module):
         self.num_layers = config_dict["num_layers"]
         model_dim = config_dict["d_model"]
 
-        self.encoder = T5Stack(self.num_layers, model_dim, model_dim, config_dict["d_ff"], config_dict["dense_act_fn"], config_dict["num_heads"], dtype, device, operations)
+        self.encoder = T5Stack(self.num_layers, model_dim, model_dim, config_dict["d_ff"], config_dict["dense_act_fn"], config_dict["is_gated_act"], config_dict["num_heads"], config_dict["model_type"] == "t5", dtype, device, operations)
         self.dtype = dtype
         self.shared = torch.nn.Embedding(config_dict["vocab_size"], model_dim, device=device)
 

comfy/text_encoders/t5_config_base.json

@@ -8,6 +8,7 @@
   "dense_act_fn": "relu",
   "initializer_factor": 1.0,
   "is_encoder_decoder": true,
+  "is_gated_act": false,
   "layer_norm_epsilon": 1e-06,
   "model_type": "t5",
   "num_decoder_layers": 12,

comfy/text_encoders/t5_config_xxl.json

@@ -8,6 +8,7 @@
   "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": 24,

comfy/text_encoders/t5_pile_config_xl.json

@@ -0,0 +1,22 @@
+{
+  "d_ff": 5120,
+  "d_kv": 64,
+  "d_model": 2048,
+  "decoder_start_token_id": 0,
+  "dropout_rate": 0.1,
+  "eos_token_id": 2,
+  "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": "umt5",
+  "num_decoder_layers": 24,
+  "num_heads": 32,
+  "num_layers": 24,
+  "output_past": true,
+  "pad_token_id": 1,
+  "relative_attention_num_buckets": 32,
+  "tie_word_embeddings": false,
+  "vocab_size": 32128
+}

comfy/utils.py

@@ -6,6 +6,7 @@ import safetensors.torch
 import numpy as np
 from PIL import Image
 import logging
+import itertools
 
 def load_torch_file(ckpt, safe_load=False, device=None):
     if device is None:
@@ -249,6 +250,158 @@ def unet_to_diffusers(unet_config):
 
     return diffusers_unet_map
 
+def swap_scale_shift(weight):
+    shift, scale = weight.chunk(2, dim=0)
+    new_weight = torch.cat([scale, shift], dim=0)
+    return new_weight
+
+MMDIT_MAP_BASIC = {
+    ("context_embedder.bias", "context_embedder.bias"),
+    ("context_embedder.weight", "context_embedder.weight"),
+    ("t_embedder.mlp.0.bias", "time_text_embed.timestep_embedder.linear_1.bias"),
+    ("t_embedder.mlp.0.weight", "time_text_embed.timestep_embedder.linear_1.weight"),
+    ("t_embedder.mlp.2.bias", "time_text_embed.timestep_embedder.linear_2.bias"),
+    ("t_embedder.mlp.2.weight", "time_text_embed.timestep_embedder.linear_2.weight"),
+    ("x_embedder.proj.bias", "pos_embed.proj.bias"),
+    ("x_embedder.proj.weight", "pos_embed.proj.weight"),
+    ("y_embedder.mlp.0.bias", "time_text_embed.text_embedder.linear_1.bias"),
+    ("y_embedder.mlp.0.weight", "time_text_embed.text_embedder.linear_1.weight"),
+    ("y_embedder.mlp.2.bias", "time_text_embed.text_embedder.linear_2.bias"),
+    ("y_embedder.mlp.2.weight", "time_text_embed.text_embedder.linear_2.weight"),
+    ("pos_embed", "pos_embed.pos_embed"),
+    ("final_layer.adaLN_modulation.1.bias", "norm_out.linear.bias", swap_scale_shift),
+    ("final_layer.adaLN_modulation.1.weight", "norm_out.linear.weight", swap_scale_shift),
+    ("final_layer.linear.bias", "proj_out.bias"),
+    ("final_layer.linear.weight", "proj_out.weight"),
+}
+
+MMDIT_MAP_BLOCK = {
+    ("context_block.adaLN_modulation.1.bias", "norm1_context.linear.bias"),
+    ("context_block.adaLN_modulation.1.weight", "norm1_context.linear.weight"),
+    ("context_block.attn.proj.bias", "attn.to_add_out.bias"),
+    ("context_block.attn.proj.weight", "attn.to_add_out.weight"),
+    ("context_block.mlp.fc1.bias", "ff_context.net.0.proj.bias"),
+    ("context_block.mlp.fc1.weight", "ff_context.net.0.proj.weight"),
+    ("context_block.mlp.fc2.bias", "ff_context.net.2.bias"),
+    ("context_block.mlp.fc2.weight", "ff_context.net.2.weight"),
+    ("x_block.adaLN_modulation.1.bias", "norm1.linear.bias"),
+    ("x_block.adaLN_modulation.1.weight", "norm1.linear.weight"),
+    ("x_block.attn.proj.bias", "attn.to_out.0.bias"),
+    ("x_block.attn.proj.weight", "attn.to_out.0.weight"),
+    ("x_block.mlp.fc1.bias", "ff.net.0.proj.bias"),
+    ("x_block.mlp.fc1.weight", "ff.net.0.proj.weight"),
+    ("x_block.mlp.fc2.bias", "ff.net.2.bias"),
+    ("x_block.mlp.fc2.weight", "ff.net.2.weight"),
+}
+
+def mmdit_to_diffusers(mmdit_config, output_prefix=""):
+    key_map = {}
+
+    depth = mmdit_config.get("depth", 0)
+    num_blocks = mmdit_config.get("num_blocks", depth)
+    for i in range(num_blocks):
+        block_from = "transformer_blocks.{}".format(i)
+        block_to = "{}joint_blocks.{}".format(output_prefix, i)
+
+        offset = depth * 64
+
+        for end in ("weight", "bias"):
+            k = "{}.attn.".format(block_from)
+            qkv = "{}.x_block.attn.qkv.{}".format(block_to, end)
+            key_map["{}to_q.{}".format(k, end)] = (qkv, (0, 0, offset))
+            key_map["{}to_k.{}".format(k, end)] = (qkv, (0, offset, offset))
+            key_map["{}to_v.{}".format(k, end)] = (qkv, (0, offset * 2, offset))
+
+            qkv = "{}.context_block.attn.qkv.{}".format(block_to, end)
+            key_map["{}add_q_proj.{}".format(k, end)] = (qkv, (0, 0, offset))
+            key_map["{}add_k_proj.{}".format(k, end)] = (qkv, (0, offset, offset))
+            key_map["{}add_v_proj.{}".format(k, end)] = (qkv, (0, offset * 2, offset))
+
+        for k in MMDIT_MAP_BLOCK:
+            key_map["{}.{}".format(block_from, k[1])] = "{}.{}".format(block_to, k[0])
+
+    map_basic = MMDIT_MAP_BASIC.copy()
+    map_basic.add(("joint_blocks.{}.context_block.adaLN_modulation.1.bias".format(depth - 1), "transformer_blocks.{}.norm1_context.linear.bias".format(depth - 1), swap_scale_shift))
+    map_basic.add(("joint_blocks.{}.context_block.adaLN_modulation.1.weight".format(depth - 1), "transformer_blocks.{}.norm1_context.linear.weight".format(depth - 1), swap_scale_shift))
+
+    for k in map_basic:
+        if len(k) > 2:
+            key_map[k[1]] = ("{}{}".format(output_prefix, k[0]), None, k[2])
+        else:
+            key_map[k[1]] = "{}{}".format(output_prefix, k[0])
+
+    return key_map
+
+
+def auraflow_to_diffusers(mmdit_config, output_prefix=""):
+    n_double_layers = mmdit_config.get("n_double_layers", 0)
+    n_layers = mmdit_config.get("n_layers", 0)
+
+    key_map = {}
+    for i in range(n_layers):
+        if i < n_double_layers:
+            index = i
+            prefix_from = "joint_transformer_blocks"
+            prefix_to = "{}double_layers".format(output_prefix)
+            block_map = {
+                            "attn.to_q.weight": "attn.w2q.weight",
+                            "attn.to_k.weight": "attn.w2k.weight",
+                            "attn.to_v.weight": "attn.w2v.weight",
+                            "attn.to_out.0.weight": "attn.w2o.weight",
+                            "attn.add_q_proj.weight": "attn.w1q.weight",
+                            "attn.add_k_proj.weight": "attn.w1k.weight",
+                            "attn.add_v_proj.weight": "attn.w1v.weight",
+                            "attn.to_add_out.weight": "attn.w1o.weight",
+                            "ff.linear_1.weight": "mlpX.c_fc1.weight",
+                            "ff.linear_2.weight": "mlpX.c_fc2.weight",
+                            "ff.out_projection.weight": "mlpX.c_proj.weight",
+                            "ff_context.linear_1.weight": "mlpC.c_fc1.weight",
+                            "ff_context.linear_2.weight": "mlpC.c_fc2.weight",
+                            "ff_context.out_projection.weight": "mlpC.c_proj.weight",
+                            "norm1.linear.weight": "modX.1.weight",
+                            "norm1_context.linear.weight": "modC.1.weight",
+                        }
+        else:
+            index = i - n_double_layers
+            prefix_from = "single_transformer_blocks"
+            prefix_to = "{}single_layers".format(output_prefix)
+
+            block_map = {
+                            "attn.to_q.weight": "attn.w1q.weight",
+                            "attn.to_k.weight": "attn.w1k.weight",
+                            "attn.to_v.weight": "attn.w1v.weight",
+                            "attn.to_out.0.weight": "attn.w1o.weight",
+                            "norm1.linear.weight": "modCX.1.weight",
+                            "ff.linear_1.weight": "mlp.c_fc1.weight",
+                            "ff.linear_2.weight": "mlp.c_fc2.weight",
+                            "ff.out_projection.weight": "mlp.c_proj.weight"
+                        }
+
+        for k in block_map:
+            key_map["{}.{}.{}".format(prefix_from, index, k)] = "{}.{}.{}".format(prefix_to, index, block_map[k])
+
+    MAP_BASIC = {
+        ("positional_encoding", "pos_embed.pos_embed"),
+        ("register_tokens", "register_tokens"),
+        ("t_embedder.mlp.0.weight", "time_step_proj.linear_1.weight"),
+        ("t_embedder.mlp.0.bias", "time_step_proj.linear_1.bias"),
+        ("t_embedder.mlp.2.weight", "time_step_proj.linear_2.weight"),
+        ("t_embedder.mlp.2.bias", "time_step_proj.linear_2.bias"),
+        ("cond_seq_linear.weight", "context_embedder.weight"),
+        ("init_x_linear.weight", "pos_embed.proj.weight"),
+        ("init_x_linear.bias", "pos_embed.proj.bias"),
+        ("final_linear.weight", "proj_out.weight"),
+        ("modF.1.weight", "norm_out.linear.weight", swap_scale_shift),
+    }
+
+    for k in MAP_BASIC:
+        if len(k) > 2:
+            key_map[k[1]] = ("{}{}".format(output_prefix, k[0]), None, k[2])
+        else:
+            key_map[k[1]] = "{}{}".format(output_prefix, k[0])
+
+    return key_map
+
 def repeat_to_batch_size(tensor, batch_size, dim=0):
     if tensor.shape[dim] > batch_size:
         return tensor.narrow(dim, 0, batch_size)
@@ -425,34 +578,52 @@ def get_tiled_scale_steps(width, height, tile_x, tile_y, overlap):
     return math.ceil((height / (tile_y - overlap))) * math.ceil((width / (tile_x - overlap)))
 
 @torch.inference_mode()
-def tiled_scale(samples, function, tile_x=64, tile_y=64, overlap = 8, upscale_amount = 4, out_channels = 3, output_device="cpu", pbar = None):
+def tiled_scale_multidim(samples, function, tile=(64, 64), overlap = 8, upscale_amount = 4, out_channels = 3, output_device="cpu", pbar = None):
-    output = torch.empty((samples.shape[0], out_channels, round(samples.shape[2] * upscale_amount), round(samples.shape[3] * upscale_amount)), device=output_device)
+    dims = len(tile)
+    output = torch.empty([samples.shape[0], out_channels] + list(map(lambda a: round(a * upscale_amount), samples.shape[2:])), device=output_device)
+
     for b in range(samples.shape[0]):
         s = samples[b:b+1]
-        out = torch.zeros((s.shape[0], out_channels, round(s.shape[2] * upscale_amount), round(s.shape[3] * upscale_amount)), device=output_device)
+        out = torch.zeros([s.shape[0], out_channels] + list(map(lambda a: round(a * upscale_amount), s.shape[2:])), device=output_device)
-        out_div = torch.zeros((s.shape[0], out_channels, round(s.shape[2] * upscale_amount), round(s.shape[3] * upscale_amount)), device=output_device)
+        out_div = torch.zeros([s.shape[0], out_channels] + list(map(lambda a: round(a * upscale_amount), s.shape[2:])), device=output_device)
-        for y in range(0, s.shape[2], tile_y - overlap):
-            for x in range(0, s.shape[3], tile_x - overlap):
-                x = max(0, min(s.shape[-1] - overlap, x))
-                y = max(0, min(s.shape[-2] - overlap, y))
-                s_in = s[:,:,y:y+tile_y,x:x+tile_x]
 
-                ps = function(s_in).to(output_device)
+        for it in itertools.product(*map(lambda a: range(0, a[0], a[1] - overlap), zip(s.shape[2:], tile))):
-                mask = torch.ones_like(ps)
+            s_in = s
-                feather = round(overlap * upscale_amount)
+            upscaled = []
-                for t in range(feather):
+
-                        mask[:,:,t:1+t,:] *= ((1.0/feather) * (t + 1))
+            for d in range(dims):
-                        mask[:,:,mask.shape[2] -1 -t: mask.shape[2]-t,:] *= ((1.0/feather) * (t + 1))
+                pos = max(0, min(s.shape[d + 2] - overlap, it[d]))
-                        mask[:,:,:,t:1+t] *= ((1.0/feather) * (t + 1))
+                l = min(tile[d], s.shape[d + 2] - pos)
-                        mask[:,:,:,mask.shape[3]- 1 - t: mask.shape[3]- t] *= ((1.0/feather) * (t + 1))
+                s_in = s_in.narrow(d + 2, pos, l)
-                out[:,:,round(y*upscale_amount):round((y+tile_y)*upscale_amount),round(x*upscale_amount):round((x+tile_x)*upscale_amount)] += ps * mask
+                upscaled.append(round(pos * upscale_amount))
-                out_div[:,:,round(y*upscale_amount):round((y+tile_y)*upscale_amount),round(x*upscale_amount):round((x+tile_x)*upscale_amount)] += mask
+            ps = function(s_in).to(output_device)
-                if pbar is not None:
+            mask = torch.ones_like(ps)
-                    pbar.update(1)
+            feather = round(overlap * upscale_amount)
+            for t in range(feather):
+                for d in range(2, dims + 2):
+                    m = mask.narrow(d, t, 1)
+                    m *= ((1.0/feather) * (t + 1))
+                    m = mask.narrow(d, mask.shape[d] -1 -t, 1)
+                    m *= ((1.0/feather) * (t + 1))
+
+            o = out
+            o_d = out_div
+            for d in range(dims):
+                o = o.narrow(d + 2, upscaled[d], mask.shape[d + 2])
+                o_d = o_d.narrow(d + 2, upscaled[d], mask.shape[d + 2])
+
+            o += ps * mask
+            o_d += mask
+
+            if pbar is not None:
+                pbar.update(1)
 
         output[b:b+1] = out/out_div
     return output
 
+def tiled_scale(samples, function, tile_x=64, tile_y=64, overlap = 8, upscale_amount = 4, out_channels = 3, output_device="cpu", pbar = None):
+    return tiled_scale_multidim(samples, function, (tile_y, tile_x), overlap, upscale_amount, out_channels, output_device, pbar)
+
 PROGRESS_BAR_ENABLED = True
 def set_progress_bar_enabled(enabled):
     global PROGRESS_BAR_ENABLED

comfy_extras/nodes_advanced_samplers.py

@@ -56,6 +56,57 @@ class SamplerLCMUpscale:
         sampler = comfy.samplers.KSAMPLER(sample_lcm_upscale, extra_options={"total_upscale": scale_ratio, "upscale_steps": scale_steps, "upscale_method": upscale_method})
         return (sampler, )
 
+from comfy.k_diffusion.sampling import to_d
+import comfy.model_patcher
+
+@torch.no_grad()
+def sample_euler_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 - denoised + temp[0], sigmas[i], denoised)
+        if callback is not None:
+            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
+        dt = sigmas[i + 1] - sigma_hat
+        x = x + d * dt
+    return x
+
+
+class SamplerEulerCFGpp:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required":
+                    {"version": (["regular", "alternative"],),}
+               }
+    RETURN_TYPES = ("SAMPLER",)
+    # CATEGORY = "sampling/custom_sampling/samplers"
+    CATEGORY = "_for_testing"
+
+    FUNCTION = "get_sampler"
+
+    def get_sampler(self, version):
+        if version == "alternative":
+            sampler = comfy.samplers.KSAMPLER(sample_euler_pp)
+        else:
+            sampler = comfy.samplers.ksampler("euler_cfg_pp")
+        return (sampler, )
+
 NODE_CLASS_MAPPINGS = {
     "SamplerLCMUpscale": SamplerLCMUpscale,
+    "SamplerEulerCFGpp": SamplerEulerCFGpp,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "SamplerEulerCFGpp": "SamplerEulerCFG++",
 }

comfy_extras/nodes_audio.py

@@ -3,6 +3,11 @@ import torch
 import comfy.model_management
 import folder_paths
 import os
+import io
+import json
+import struct
+import random
+from comfy.cli_args import args
 
 class EmptyLatentAudio:
     def __init__(self):
@@ -10,15 +15,16 @@ class EmptyLatentAudio:
 
     @classmethod
     def INPUT_TYPES(s):
-        return {"required": {}}
+        return {"required": {"seconds": ("FLOAT", {"default": 47.6, "min": 1.0, "max": 1000.0, "step": 0.1})}}
     RETURN_TYPES = ("LATENT",)
     FUNCTION = "generate"
 
-    CATEGORY = "_for_testing/audio"
+    CATEGORY = "latent/audio"
 
-    def generate(self):
+    def generate(self, seconds):
         batch_size = 1
-        latent = torch.zeros([batch_size, 64, 1024], device=self.device)
+        length = round((seconds * 44100 / 2048) / 2) * 2
+        latent = torch.zeros([batch_size, 64, length], device=self.device)
         return ({"samples":latent, "type": "audio"}, )
 
 class VAEEncodeAudio:
@@ -28,10 +34,16 @@ class VAEEncodeAudio:
     RETURN_TYPES = ("LATENT",)
     FUNCTION = "encode"
 
-    CATEGORY = "_for_testing/audio"
+    CATEGORY = "latent/audio"
 
     def encode(self, vae, audio):
-        t = vae.encode(audio["waveform"].movedim(1, -1))
+        sample_rate = audio["sample_rate"]
+        if 44100 != sample_rate:
+            waveform = torchaudio.functional.resample(audio["waveform"], sample_rate, 44100)
+        else:
+            waveform = audio["waveform"]
+
+        t = vae.encode(waveform.movedim(1, -1))
         return ({"samples":t}, )
 
 class VAEDecodeAudio:
@@ -41,18 +53,72 @@ class VAEDecodeAudio:
     RETURN_TYPES = ("AUDIO",)
     FUNCTION = "decode"
 
-    CATEGORY = "_for_testing/audio"
+    CATEGORY = "latent/audio"
 
     def decode(self, vae, samples):
         audio = vae.decode(samples["samples"]).movedim(-1, 1)
         return ({"waveform": audio, "sample_rate": 44100}, )
 
+
+def create_vorbis_comment_block(comment_dict, last_block):
+    vendor_string = b'ComfyUI'
+    vendor_length = len(vendor_string)
+
+    comments = []
+    for key, value in comment_dict.items():
+        comment = f"{key}={value}".encode('utf-8')
+        comments.append(struct.pack('<I', len(comment)) + comment)
+
+    user_comment_list_length = len(comments)
+    user_comments = b''.join(comments)
+
+    comment_data = struct.pack('<I', vendor_length) + vendor_string + struct.pack('<I', user_comment_list_length) + user_comments
+    if last_block:
+        id = b'\x84'
+    else:
+        id = b'\x04'
+    comment_block = id + struct.pack('>I', len(comment_data))[1:] + comment_data
+
+    return comment_block
+
+def insert_or_replace_vorbis_comment(flac_io, comment_dict):
+    if len(comment_dict) == 0:
+        return flac_io
+
+    flac_io.seek(4)
+
+    blocks = []
+    last_block = False
+
+    while not last_block:
+        header = flac_io.read(4)
+        last_block = (header[0] & 0x80) != 0
+        block_type = header[0] & 0x7F
+        block_length = struct.unpack('>I', b'\x00' + header[1:])[0]
+        block_data = flac_io.read(block_length)
+
+        if block_type == 4 or block_type == 1:
+            pass
+        else:
+            header = bytes([(header[0] & (~0x80))]) + header[1:]
+            blocks.append(header + block_data)
+
+    blocks.append(create_vorbis_comment_block(comment_dict, last_block=True))
+
+    new_flac_io = io.BytesIO()
+    new_flac_io.write(b'fLaC')
+    for block in blocks:
+        new_flac_io.write(block)
+
+    new_flac_io.write(flac_io.read())
+    return new_flac_io
+
+
 class SaveAudio:
     def __init__(self):
         self.output_dir = folder_paths.get_output_directory()
         self.type = "output"
         self.prefix_append = ""
-        self.compress_level = 4
 
     @classmethod
     def INPUT_TYPES(s):
@@ -66,17 +132,33 @@ class SaveAudio:
 
     OUTPUT_NODE = True
 
-    CATEGORY = "_for_testing/audio"
+    CATEGORY = "audio"
 
     def save_audio(self, audio, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
         filename_prefix += self.prefix_append
         full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir)
         results = list()
+
+        metadata = {}
+        if not args.disable_metadata:
+            if prompt is not None:
+                metadata["prompt"] = json.dumps(prompt)
+            if extra_pnginfo is not None:
+                for x in extra_pnginfo:
+                    metadata[x] = json.dumps(extra_pnginfo[x])
+
         for (batch_number, waveform) in enumerate(audio["waveform"]):
-            #TODO: metadata
             filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
             file = f"{filename_with_batch_num}_{counter:05}_.flac"
-            torchaudio.save(os.path.join(full_output_folder, file), waveform, audio["sample_rate"], format="FLAC")
+
+            buff = io.BytesIO()
+            torchaudio.save(buff, waveform, audio["sample_rate"], format="FLAC")
+
+            buff = insert_or_replace_vorbis_comment(buff, metadata)
+
+            with open(os.path.join(full_output_folder, file), 'wb') as f:
+                f.write(buff.getbuffer())
+
             results.append({
                 "filename": file,
                 "subfolder": subfolder,
@@ -86,14 +168,34 @@ class SaveAudio:
 
         return { "ui": { "audio": results } }
 
+class PreviewAudio(SaveAudio):
+    def __init__(self):
+        self.output_dir = folder_paths.get_temp_directory()
+        self.type = "temp"
+        self.prefix_append = "_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for x in range(5))
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required":
+                    {"audio": ("AUDIO", ), },
+                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
+                }
+
 class LoadAudio:
+    SUPPORTED_FORMATS = ('.wav', '.mp3', '.ogg', '.flac', '.aiff', '.aif')
+
     @classmethod
     def INPUT_TYPES(s):
         input_dir = folder_paths.get_input_directory()
-        files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
+        files = [
-        return {"required": {"audio": [sorted(files), ]}, }
+            f for f in os.listdir(input_dir)
+            if (os.path.isfile(os.path.join(input_dir, f))
+                and f.endswith(LoadAudio.SUPPORTED_FORMATS)
+            )
+        ]
+        return {"required": {"audio": (sorted(files), {"audio_upload": True})}}
 
-    CATEGORY = "_for_testing/audio"
+    CATEGORY = "audio"
 
     RETURN_TYPES = ("AUDIO", )
     FUNCTION = "load"
@@ -101,7 +203,6 @@ class LoadAudio:
     def load(self, audio):
         audio_path = folder_paths.get_annotated_filepath(audio)
         waveform, sample_rate = torchaudio.load(audio_path)
-        multiplier = 1.0
         audio = {"waveform": waveform.unsqueeze(0), "sample_rate": sample_rate}
         return (audio, )
 
@@ -125,4 +226,5 @@ NODE_CLASS_MAPPINGS = {
     "VAEDecodeAudio": VAEDecodeAudio,
     "SaveAudio": SaveAudio,
     "LoadAudio": LoadAudio,
+    "PreviewAudio": PreviewAudio,
 }

comfy_extras/nodes_controlnet.py

@@ -0,0 +1,37 @@
+
+UNION_CONTROLNET_TYPES = {"auto": -1,
+                          "openpose": 0,
+                          "depth": 1,
+                          "hed/pidi/scribble/ted": 2,
+                          "canny/lineart/anime_lineart/mlsd": 3,
+                          "normal": 4,
+                          "segment": 5,
+                          "tile": 6,
+                          "repaint": 7,
+                        }
+
+class SetUnionControlNetType:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {"control_net": ("CONTROL_NET", ),
+                             "type": (list(UNION_CONTROLNET_TYPES.keys()),)
+                             }}
+
+    CATEGORY = "conditioning/controlnet"
+    RETURN_TYPES = ("CONTROL_NET",)
+
+    FUNCTION = "set_controlnet_type"
+
+    def set_controlnet_type(self, control_net, type):
+        control_net = control_net.copy()
+        type_number = UNION_CONTROLNET_TYPES[type]
+        if type_number >= 0:
+            control_net.set_extra_arg("control_type", [type_number])
+        else:
+            control_net.set_extra_arg("control_type", [])
+
+        return (control_net,)
+
+NODE_CLASS_MAPPINGS = {
+    "SetUnionControlNetType": SetUnionControlNetType,
+}

comfy_extras/nodes_custom_sampler.py

@@ -111,6 +111,25 @@ class SDTurboScheduler:
         sigmas = torch.cat([sigmas, sigmas.new_zeros([1])])
         return (sigmas, )
 
+class BetaSamplingScheduler:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required":
+                    {"model": ("MODEL",),
+                     "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
+                     "alpha": ("FLOAT", {"default": 0.6, "min": 0.0, "max": 50.0, "step":0.01, "round": False}),
+                     "beta": ("FLOAT", {"default": 0.6, "min": 0.0, "max": 50.0, "step":0.01, "round": False}),
+                      }
+               }
+    RETURN_TYPES = ("SIGMAS",)
+    CATEGORY = "sampling/custom_sampling/schedulers"
+
+    FUNCTION = "get_sigmas"
+
+    def get_sigmas(self, model, steps, alpha, beta):
+        sigmas = comfy.samplers.beta_scheduler(model.get_model_object("model_sampling"), steps, alpha=alpha, beta=beta)
+        return (sigmas, )
+
 class VPScheduler:
     @classmethod
     def INPUT_TYPES(s):
@@ -293,6 +312,25 @@ class SamplerEulerAncestral:
         sampler = comfy.samplers.ksampler("euler_ancestral", {"eta": eta, "s_noise": s_noise})
         return (sampler, )
 
+class SamplerEulerAncestralCFGPP:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "eta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step":0.01, "round": False}),
+                "s_noise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step":0.01, "round": False}),
+            }}
+    RETURN_TYPES = ("SAMPLER",)
+    CATEGORY = "sampling/custom_sampling/samplers"
+
+    FUNCTION = "get_sampler"
+
+    def get_sampler(self, eta, s_noise):
+        sampler = comfy.samplers.ksampler(
+            "euler_ancestral_cfg_pp",
+            {"eta": eta, "s_noise": s_noise})
+        return (sampler, )
+
 class SamplerLMS:
     @classmethod
     def INPUT_TYPES(s):
@@ -619,9 +657,11 @@ NODE_CLASS_MAPPINGS = {
     "ExponentialScheduler": ExponentialScheduler,
     "PolyexponentialScheduler": PolyexponentialScheduler,
     "VPScheduler": VPScheduler,
+    "BetaSamplingScheduler": BetaSamplingScheduler,
     "SDTurboScheduler": SDTurboScheduler,
     "KSamplerSelect": KSamplerSelect,
     "SamplerEulerAncestral": SamplerEulerAncestral,
+    "SamplerEulerAncestralCFGPP": SamplerEulerAncestralCFGPP,
     "SamplerLMS": SamplerLMS,
     "SamplerDPMPP_3M_SDE": SamplerDPMPP_3M_SDE,
     "SamplerDPMPP_2M_SDE": SamplerDPMPP_2M_SDE,
@@ -639,3 +679,7 @@ NODE_CLASS_MAPPINGS = {
     "AddNoise": AddNoise,
     "SamplerCustomAdvanced": SamplerCustomAdvanced,
 }
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "SamplerEulerAncestralCFGPP": "SamplerEulerAncestralCFG++",
+}

comfy_extras/nodes_freelunch.py

@@ -34,7 +34,7 @@ class FreeU:
     RETURN_TYPES = ("MODEL",)
     FUNCTION = "patch"
 
-    CATEGORY = "model_patches"
+    CATEGORY = "model_patches/unet"
 
     def patch(self, model, b1, b2, s1, s2):
         model_channels = model.model.model_config.unet_config["model_channels"]
@@ -73,7 +73,7 @@ class FreeU_V2:
     RETURN_TYPES = ("MODEL",)
     FUNCTION = "patch"
 
-    CATEGORY = "model_patches"
+    CATEGORY = "model_patches/unet"
 
     def patch(self, model, b1, b2, s1, s2):
         model_channels = model.model.model_config.unet_config["model_channels"]

comfy_extras/nodes_gits.py

@@ -0,0 +1,369 @@
+# from https://github.com/zju-pi/diff-sampler/tree/main/gits-main
+import numpy as np
+import torch
+
+def loglinear_interp(t_steps, num_steps):
+    """
+    Performs log-linear interpolation of a given array of decreasing numbers.
+    """
+    xs = np.linspace(0, 1, len(t_steps))
+    ys = np.log(t_steps[::-1])
+
+    new_xs = np.linspace(0, 1, num_steps)
+    new_ys = np.interp(new_xs, xs, ys)
+
+    interped_ys = np.exp(new_ys)[::-1].copy()
+    return interped_ys
+
+NOISE_LEVELS = {
+    0.80: [
+        [14.61464119, 7.49001646, 0.02916753],
+        [14.61464119, 11.54541874, 6.77309084, 0.02916753],
+        [14.61464119, 11.54541874, 7.49001646, 3.07277966, 0.02916753],
+        [14.61464119, 11.54541874, 7.49001646, 5.85520077, 2.05039096, 0.02916753],
+        [14.61464119, 12.2308979, 8.75849152, 7.49001646, 5.85520077, 2.05039096, 0.02916753],
+        [14.61464119, 12.2308979, 8.75849152, 7.49001646, 5.85520077, 3.07277966, 1.56271636, 0.02916753],
+        [14.61464119, 12.96784878, 11.54541874, 8.75849152, 7.49001646, 5.85520077, 3.07277966, 1.56271636, 0.02916753],
+        [14.61464119, 13.76078796, 12.2308979, 10.90732002, 8.75849152, 7.49001646, 5.85520077, 3.07277966, 1.56271636, 0.02916753],
+        [14.61464119, 13.76078796, 12.96784878, 12.2308979, 10.90732002, 8.75849152, 7.49001646, 5.85520077, 3.07277966, 1.56271636, 0.02916753],
+        [14.61464119, 13.76078796, 12.96784878, 12.2308979, 10.90732002, 9.24142551, 8.30717278, 7.49001646, 5.85520077, 3.07277966, 1.56271636, 0.02916753],
+        [14.61464119, 13.76078796, 12.96784878, 12.2308979, 10.90732002, 9.24142551, 8.30717278, 7.49001646, 6.14220476, 4.86714602, 3.07277966, 1.56271636, 0.02916753],
+        [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.31284904, 9.24142551, 8.30717278, 7.49001646, 6.14220476, 4.86714602, 3.07277966, 1.56271636, 0.02916753],
+        [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.90732002, 10.31284904, 9.24142551, 8.30717278, 7.49001646, 6.14220476, 4.86714602, 3.07277966, 1.56271636, 0.02916753],
+        [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.90732002, 10.31284904, 9.24142551, 8.75849152, 8.30717278, 7.49001646, 6.14220476, 4.86714602, 3.07277966, 1.56271636, 0.02916753],
+        [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.90732002, 10.31284904, 9.24142551, 8.75849152, 8.30717278, 7.49001646, 6.14220476, 4.86714602, 3.1956799, 1.98035145, 0.86115354, 0.02916753],
+        [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.90732002, 10.31284904, 9.75859547, 9.24142551, 8.75849152, 8.30717278, 7.49001646, 6.14220476, 4.86714602, 3.1956799, 1.98035145, 0.86115354, 0.02916753],
+        [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.90732002, 10.31284904, 9.75859547, 9.24142551, 8.75849152, 8.30717278, 7.49001646, 6.77309084, 5.85520077, 4.65472794, 3.07277966, 1.84880662, 0.83188516, 0.02916753],
+        [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.90732002, 10.31284904, 9.75859547, 9.24142551, 8.75849152, 8.30717278, 7.88507891, 7.49001646, 6.77309084, 5.85520077, 4.65472794, 3.07277966, 1.84880662, 0.83188516, 0.02916753],
+        [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.90732002, 10.31284904, 9.75859547, 9.24142551, 8.75849152, 8.30717278, 7.88507891, 7.49001646, 6.77309084, 5.85520077, 4.86714602, 3.75677586, 2.84484982, 1.78698075, 0.803307, 0.02916753],
+    ],
+    0.85: [
+        [14.61464119, 7.49001646, 0.02916753],
+        [14.61464119, 7.49001646, 1.84880662, 0.02916753],
+        [14.61464119, 11.54541874, 6.77309084, 1.56271636, 0.02916753],
+        [14.61464119, 11.54541874, 7.11996698, 3.07277966, 1.24153244, 0.02916753],
+        [14.61464119, 11.54541874, 7.49001646, 5.09240818, 2.84484982, 0.95350921, 0.02916753],
+        [14.61464119, 12.2308979, 8.75849152, 7.49001646, 5.09240818, 2.84484982, 0.95350921, 0.02916753],
+        [14.61464119, 12.2308979, 8.75849152, 7.49001646, 5.58536053, 3.1956799, 1.84880662, 0.803307, 0.02916753],
+        [14.61464119, 12.96784878, 11.54541874, 8.75849152, 7.49001646, 5.58536053, 3.1956799, 1.84880662, 0.803307, 0.02916753],
+        [14.61464119, 12.96784878, 11.54541874, 8.75849152, 7.49001646, 6.14220476, 4.65472794, 3.07277966, 1.84880662, 0.803307, 0.02916753],
+        [14.61464119, 13.76078796, 12.2308979, 10.90732002, 8.75849152, 7.49001646, 6.14220476, 4.65472794, 3.07277966, 1.84880662, 0.803307, 0.02916753],
+        [14.61464119, 13.76078796, 12.2308979, 10.90732002, 9.24142551, 8.30717278, 7.49001646, 6.14220476, 4.65472794, 3.07277966, 1.84880662, 0.803307, 0.02916753],
+        [14.61464119, 13.76078796, 12.96784878, 12.2308979, 10.90732002, 9.24142551, 8.30717278, 7.49001646, 6.14220476, 4.65472794, 3.07277966, 1.84880662, 0.803307, 0.02916753],
+        [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.31284904, 9.24142551, 8.30717278, 7.49001646, 6.14220476, 4.65472794, 3.07277966, 1.84880662, 0.803307, 0.02916753],
+        [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.31284904, 9.24142551, 8.30717278, 7.49001646, 6.14220476, 4.86714602, 3.60512662, 2.6383388, 1.56271636, 0.72133851, 0.02916753],
+        [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.31284904, 9.24142551, 8.30717278, 7.49001646, 6.77309084, 5.85520077, 4.65472794, 3.46139455, 2.45070267, 1.56271636, 0.72133851, 0.02916753],
+        [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.31284904, 9.24142551, 8.75849152, 8.30717278, 7.49001646, 6.77309084, 5.85520077, 4.65472794, 3.46139455, 2.45070267, 1.56271636, 0.72133851, 0.02916753],
+        [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.90732002, 10.31284904, 9.24142551, 8.75849152, 8.30717278, 7.49001646, 6.77309084, 5.85520077, 4.65472794, 3.46139455, 2.45070267, 1.56271636, 0.72133851, 0.02916753],
+        [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.90732002, 10.31284904, 9.75859547, 9.24142551, 8.75849152, 8.30717278, 7.49001646, 6.77309084, 5.85520077, 4.65472794, 3.46139455, 2.45070267, 1.56271636, 0.72133851, 0.02916753],
+        [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.90732002, 10.31284904, 9.75859547, 9.24142551, 8.75849152, 8.30717278, 7.88507891, 7.49001646, 6.77309084, 5.85520077, 4.65472794, 3.46139455, 2.45070267, 1.56271636, 0.72133851, 0.02916753],
+    ],
+    0.90: [
+        [14.61464119, 6.77309084, 0.02916753],
+        [14.61464119, 7.49001646, 1.56271636, 0.02916753],
+        [14.61464119, 7.49001646, 3.07277966, 0.95350921, 0.02916753],
+        [14.61464119, 7.49001646, 4.86714602, 2.54230714, 0.89115214, 0.02916753],
+        [14.61464119, 11.54541874, 7.49001646, 4.86714602, 2.54230714, 0.89115214, 0.02916753],
+        [14.61464119, 11.54541874, 7.49001646, 5.09240818, 3.07277966, 1.61558151, 0.69515091, 0.02916753],
+        [14.61464119, 12.2308979, 8.75849152, 7.11996698, 4.86714602, 3.07277966, 1.61558151, 0.69515091, 0.02916753],
+        [14.61464119, 12.2308979, 8.75849152, 7.49001646, 5.85520077, 4.45427561, 2.95596409, 1.61558151, 0.69515091, 0.02916753],
+        [14.61464119, 12.2308979, 8.75849152, 7.49001646, 5.85520077, 4.45427561, 3.1956799, 2.19988537, 1.24153244, 0.57119018, 0.02916753],
+        [14.61464119, 12.96784878, 10.90732002, 8.75849152, 7.49001646, 5.85520077, 4.45427561, 3.1956799, 2.19988537, 1.24153244, 0.57119018, 0.02916753],
+        [14.61464119, 12.96784878, 11.54541874, 9.24142551, 8.30717278, 7.49001646, 5.85520077, 4.45427561, 3.1956799, 2.19988537, 1.24153244, 0.57119018, 0.02916753],
+        [14.61464119, 12.96784878, 11.54541874, 9.24142551, 8.30717278, 7.49001646, 6.14220476, 4.86714602, 3.75677586, 2.84484982, 1.84880662, 1.08895338, 0.52423614, 0.02916753],
+        [14.61464119, 13.76078796, 12.2308979, 10.90732002, 9.24142551, 8.30717278, 7.49001646, 6.14220476, 4.86714602, 3.75677586, 2.84484982, 1.84880662, 1.08895338, 0.52423614, 0.02916753],
+        [14.61464119, 13.76078796, 12.2308979, 10.90732002, 9.24142551, 8.30717278, 7.49001646, 6.44769001, 5.58536053, 4.45427561, 3.32507086, 2.45070267, 1.61558151, 0.95350921, 0.45573691, 0.02916753],
+        [14.61464119, 13.76078796, 12.96784878, 12.2308979, 10.90732002, 9.24142551, 8.30717278, 7.49001646, 6.44769001, 5.58536053, 4.45427561, 3.32507086, 2.45070267, 1.61558151, 0.95350921, 0.45573691, 0.02916753],
+        [14.61464119, 13.76078796, 12.96784878, 12.2308979, 10.90732002, 9.24142551, 8.30717278, 7.49001646, 6.77309084, 5.85520077, 4.86714602, 3.91689563, 3.07277966, 2.27973175, 1.56271636, 0.95350921, 0.45573691, 0.02916753],
+        [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.31284904, 9.24142551, 8.30717278, 7.49001646, 6.77309084, 5.85520077, 4.86714602, 3.91689563, 3.07277966, 2.27973175, 1.56271636, 0.95350921, 0.45573691, 0.02916753],
+        [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.31284904, 9.24142551, 8.75849152, 8.30717278, 7.49001646, 6.77309084, 5.85520077, 4.86714602, 3.91689563, 3.07277966, 2.27973175, 1.56271636, 0.95350921, 0.45573691, 0.02916753],
+        [14.61464119, 13.76078796, 12.96784878, 12.2308979, 11.54541874, 10.31284904, 9.24142551, 8.75849152, 8.30717278, 7.49001646, 6.77309084, 5.85520077, 5.09240818, 4.45427561, 3.60512662, 2.95596409, 2.19988537, 1.51179266, 0.89115214, 0.43325692, 0.02916753],
+    ],
+    0.95: [
+        [14.61464119, 6.77309084, 0.02916753],
+        [14.61464119, 6.77309084, 1.56271636, 0.02916753],
+        [14.61464119, 7.49001646, 2.84484982, 0.89115214, 0.02916753],
+        [14.61464119, 7.49001646, 4.86714602, 2.36326075, 0.803307, 0.02916753],
+        [14.61464119, 7.49001646, 4.86714602, 2.95596409, 1.56271636, 0.64427125, 0.02916753],
+        [14.61464119, 11.54541874, 7.49001646, 4.86714602, 2.95596409, 1.56271636, 0.64427125, 0.02916753],
+        [14.61464119, 11.54541874, 7.49001646, 4.86714602, 3.07277966, 1.91321158, 1.08895338, 0.50118381, 0.02916753],
+        [14.61464119, 11.54541874, 7.49001646, 5.85520077, 4.45427561, 3.07277966, 1.91321158, 1.08895338, 0.50118381, 0.02916753],
+        [14.61464119, 12.2308979, 8.75849152, 7.49001646, 5.85520077, 4.45427561, 3.07277966, 1.91321158, 1.08895338, 0.50118381, 0.02916753],
+        [14.61464119, 12.2308979, 8.75849152, 7.49001646, 5.85520077, 4.45427561, 3.1956799, 2.19988537, 1.41535246, 0.803307, 0.38853383, 0.02916753],
+        [14.61464119, 12.2308979, 8.75849152, 7.49001646, 5.85520077, 4.65472794, 3.46139455, 2.6383388, 1.84880662, 1.24153244, 0.72133851, 0.34370604, 0.02916753],
+        [14.61464119, 12.96784878, 10.90732002, 8.75849152, 7.49001646, 5.85520077, 4.65472794, 3.46139455, 2.6383388, 1.84880662, 1.24153244, 0.72133851, 0.34370604, 0.02916753],
+        [14.61464119, 12.96784878, 10.90732002, 8.75849152, 7.49001646, 6.14220476, 4.86714602, 3.75677586, 2.95596409, 2.19988537, 1.56271636, 1.05362725, 0.64427125, 0.32104823, 0.02916753],
+        [14.61464119, 12.96784878, 10.90732002, 8.75849152, 7.49001646, 6.44769001, 5.58536053, 4.65472794, 3.60512662, 2.95596409, 2.19988537, 1.56271636, 1.05362725, 0.64427125, 0.32104823, 0.02916753],
+        [14.61464119, 12.96784878, 11.54541874, 9.24142551, 8.30717278, 7.49001646, 6.44769001, 5.58536053, 4.65472794, 3.60512662, 2.95596409, 2.19988537, 1.56271636, 1.05362725, 0.64427125, 0.32104823, 0.02916753],
+        [14.61464119, 12.96784878, 11.54541874, 9.24142551, 8.30717278, 7.49001646, 6.44769001, 5.58536053, 4.65472794, 3.75677586, 3.07277966, 2.45070267, 1.78698075, 1.24153244, 0.83188516, 0.50118381, 0.22545385, 0.02916753],
+        [14.61464119, 12.96784878, 11.54541874, 9.24142551, 8.30717278, 7.49001646, 6.77309084, 5.85520077, 5.09240818, 4.45427561, 3.60512662, 2.95596409, 2.36326075, 1.72759056, 1.24153244, 0.83188516, 0.50118381, 0.22545385, 0.02916753],
+        [14.61464119, 13.76078796, 12.2308979, 10.90732002, 9.24142551, 8.30717278, 7.49001646, 6.77309084, 5.85520077, 5.09240818, 4.45427561, 3.60512662, 2.95596409, 2.36326075, 1.72759056, 1.24153244, 0.83188516, 0.50118381, 0.22545385, 0.02916753],
+        [14.61464119, 13.76078796, 12.2308979, 10.90732002, 9.24142551, 8.30717278, 7.49001646, 6.77309084, 5.85520077, 5.09240818, 4.45427561, 3.75677586, 3.07277966, 2.45070267, 1.91321158, 1.46270394, 1.05362725, 0.72133851, 0.43325692, 0.19894916, 0.02916753],
+    ],
+    1.00: [
+        [14.61464119, 1.56271636, 0.02916753],
+        [14.61464119, 6.77309084, 0.95350921, 0.02916753],
+        [14.61464119, 6.77309084, 2.36326075, 0.803307, 0.02916753],
+        [14.61464119, 7.11996698, 3.07277966, 1.56271636, 0.59516323, 0.02916753],
+        [14.61464119, 7.49001646, 4.86714602, 2.84484982, 1.41535246, 0.57119018, 0.02916753],
+        [14.61464119, 7.49001646, 4.86714602, 2.84484982, 1.61558151, 0.86115354, 0.38853383, 0.02916753],
+        [14.61464119, 11.54541874, 7.49001646, 4.86714602, 2.84484982, 1.61558151, 0.86115354, 0.38853383, 0.02916753],
+        [14.61464119, 11.54541874, 7.49001646, 4.86714602, 3.07277966, 1.98035145, 1.24153244, 0.72133851, 0.34370604, 0.02916753],
+        [14.61464119, 11.54541874, 7.49001646, 5.85520077, 4.45427561, 3.07277966, 1.98035145, 1.24153244, 0.72133851, 0.34370604, 0.02916753],
+        [14.61464119, 11.54541874, 7.49001646, 5.85520077, 4.45427561, 3.1956799, 2.27973175, 1.51179266, 0.95350921, 0.54755926, 0.25053367, 0.02916753],
+        [14.61464119, 11.54541874, 7.49001646, 5.85520077, 4.45427561, 3.1956799, 2.36326075, 1.61558151, 1.08895338, 0.72133851, 0.41087446, 0.17026083, 0.02916753],
+        [14.61464119, 11.54541874, 8.75849152, 7.49001646, 5.85520077, 4.45427561, 3.1956799, 2.36326075, 1.61558151, 1.08895338, 0.72133851, 0.41087446, 0.17026083, 0.02916753],
+        [14.61464119, 11.54541874, 8.75849152, 7.49001646, 5.85520077, 4.65472794, 3.60512662, 2.84484982, 2.12350607, 1.56271636, 1.08895338, 0.72133851, 0.41087446, 0.17026083, 0.02916753],
+        [14.61464119, 11.54541874, 8.75849152, 7.49001646, 5.85520077, 4.65472794, 3.60512662, 2.84484982, 2.19988537, 1.61558151, 1.162866, 0.803307, 0.50118381, 0.27464288, 0.09824532, 0.02916753],
+        [14.61464119, 11.54541874, 8.75849152, 7.49001646, 5.85520077, 4.65472794, 3.75677586, 3.07277966, 2.45070267, 1.84880662, 1.36964464, 1.01931262, 0.72133851, 0.45573691, 0.25053367, 0.09824532, 0.02916753],
+        [14.61464119, 11.54541874, 8.75849152, 7.49001646, 6.14220476, 5.09240818, 4.26497746, 3.46139455, 2.84484982, 2.19988537, 1.67050016, 1.24153244, 0.92192322, 0.64427125, 0.43325692, 0.25053367, 0.09824532, 0.02916753],
+        [14.61464119, 11.54541874, 8.75849152, 7.49001646, 6.14220476, 5.09240818, 4.26497746, 3.60512662, 2.95596409, 2.45070267, 1.91321158, 1.51179266, 1.12534678, 0.83188516, 0.59516323, 0.38853383, 0.22545385, 0.09824532, 0.02916753],
+        [14.61464119, 12.2308979, 9.24142551, 8.30717278, 7.49001646, 6.14220476, 5.09240818, 4.26497746, 3.60512662, 2.95596409, 2.45070267, 1.91321158, 1.51179266, 1.12534678, 0.83188516, 0.59516323, 0.38853383, 0.22545385, 0.09824532, 0.02916753],
+        [14.61464119, 12.2308979, 9.24142551, 8.30717278, 7.49001646, 6.77309084, 5.85520077, 5.09240818, 4.26497746, 3.60512662, 2.95596409, 2.45070267, 1.91321158, 1.51179266, 1.12534678, 0.83188516, 0.59516323, 0.38853383, 0.22545385, 0.09824532, 0.02916753],
+    ],
+    1.05: [
+        [14.61464119, 0.95350921, 0.02916753],
+        [14.61464119, 6.77309084, 0.89115214, 0.02916753],
+        [14.61464119, 6.77309084, 2.05039096, 0.72133851, 0.02916753],
+        [14.61464119, 6.77309084, 2.84484982, 1.28281462, 0.52423614, 0.02916753],
+        [14.61464119, 6.77309084, 3.07277966, 1.61558151, 0.803307, 0.34370604, 0.02916753],
+        [14.61464119, 7.49001646, 4.86714602, 2.84484982, 1.56271636, 0.803307, 0.34370604, 0.02916753],
+        [14.61464119, 7.49001646, 4.86714602, 2.84484982, 1.61558151, 0.95350921, 0.52423614, 0.22545385, 0.02916753],
+        [14.61464119, 7.49001646, 4.86714602, 3.07277966, 1.98035145, 1.24153244, 0.74807048, 0.41087446, 0.17026083, 0.02916753],
+        [14.61464119, 7.49001646, 4.86714602, 3.1956799, 2.27973175, 1.51179266, 0.95350921, 0.59516323, 0.34370604, 0.13792117, 0.02916753],
+        [14.61464119, 7.49001646, 5.09240818, 3.46139455, 2.45070267, 1.61558151, 1.08895338, 0.72133851, 0.45573691, 0.25053367, 0.09824532, 0.02916753],
+        [14.61464119, 11.54541874, 7.49001646, 5.09240818, 3.46139455, 2.45070267, 1.61558151, 1.08895338, 0.72133851, 0.45573691, 0.25053367, 0.09824532, 0.02916753],
+        [14.61464119, 11.54541874, 7.49001646, 5.85520077, 4.45427561, 3.1956799, 2.36326075, 1.61558151, 1.08895338, 0.72133851, 0.45573691, 0.25053367, 0.09824532, 0.02916753],
+        [14.61464119, 11.54541874, 7.49001646, 5.85520077, 4.45427561, 3.1956799, 2.45070267, 1.72759056, 1.24153244, 0.86115354, 0.59516323, 0.38853383, 0.22545385, 0.09824532, 0.02916753],
+        [14.61464119, 11.54541874, 7.49001646, 5.85520077, 4.65472794, 3.60512662, 2.84484982, 2.19988537, 1.61558151, 1.162866, 0.83188516, 0.59516323, 0.38853383, 0.22545385, 0.09824532, 0.02916753],
+        [14.61464119, 11.54541874, 7.49001646, 5.85520077, 4.65472794, 3.60512662, 2.84484982, 2.19988537, 1.67050016, 1.28281462, 0.95350921, 0.72133851, 0.52423614, 0.34370604, 0.19894916, 0.09824532, 0.02916753],
+        [14.61464119, 11.54541874, 7.49001646, 5.85520077, 4.65472794, 3.60512662, 2.95596409, 2.36326075, 1.84880662, 1.41535246, 1.08895338, 0.83188516, 0.61951244, 0.45573691, 0.32104823, 0.19894916, 0.09824532, 0.02916753],
+        [14.61464119, 11.54541874, 7.49001646, 5.85520077, 4.65472794, 3.60512662, 2.95596409, 2.45070267, 1.91321158, 1.51179266, 1.20157266, 0.95350921, 0.74807048, 0.57119018, 0.43325692, 0.29807833, 0.19894916, 0.09824532, 0.02916753],
+        [14.61464119, 11.54541874, 8.30717278, 7.11996698, 5.85520077, 4.65472794, 3.60512662, 2.95596409, 2.45070267, 1.91321158, 1.51179266, 1.20157266, 0.95350921, 0.74807048, 0.57119018, 0.43325692, 0.29807833, 0.19894916, 0.09824532, 0.02916753],
+        [14.61464119, 11.54541874, 8.30717278, 7.11996698, 5.85520077, 4.65472794, 3.60512662, 2.95596409, 2.45070267, 1.98035145, 1.61558151, 1.32549286, 1.08895338, 0.86115354, 0.69515091, 0.54755926, 0.41087446, 0.29807833, 0.19894916, 0.09824532, 0.02916753],
+    ],
+    1.10: [
+        [14.61464119, 0.89115214, 0.02916753],
+        [14.61464119, 2.36326075, 0.72133851, 0.02916753],
+        [14.61464119, 5.85520077, 1.61558151, 0.57119018, 0.02916753],
+        [14.61464119, 6.77309084, 2.45070267, 1.08895338, 0.45573691, 0.02916753],
+        [14.61464119, 6.77309084, 2.95596409, 1.56271636, 0.803307, 0.34370604, 0.02916753],
+        [14.61464119, 6.77309084, 3.07277966, 1.61558151, 0.89115214, 0.4783645, 0.19894916, 0.02916753],
+        [14.61464119, 6.77309084, 3.07277966, 1.84880662, 1.08895338, 0.64427125, 0.34370604, 0.13792117, 0.02916753],
+        [14.61464119, 7.49001646, 4.86714602, 2.84484982, 1.61558151, 0.95350921, 0.54755926, 0.27464288, 0.09824532, 0.02916753],
+        [14.61464119, 7.49001646, 4.86714602, 2.95596409, 1.91321158, 1.24153244, 0.803307, 0.4783645, 0.25053367, 0.09824532, 0.02916753],
+        [14.61464119, 7.49001646, 4.86714602, 3.07277966, 2.05039096, 1.41535246, 0.95350921, 0.64427125, 0.41087446, 0.22545385, 0.09824532, 0.02916753],
+        [14.61464119, 7.49001646, 4.86714602, 3.1956799, 2.27973175, 1.61558151, 1.12534678, 0.803307, 0.54755926, 0.36617002, 0.22545385, 0.09824532, 0.02916753],
+        [14.61464119, 7.49001646, 4.86714602, 3.32507086, 2.45070267, 1.72759056, 1.24153244, 0.89115214, 0.64427125, 0.45573691, 0.32104823, 0.19894916, 0.09824532, 0.02916753],
+        [14.61464119, 7.49001646, 5.09240818, 3.60512662, 2.84484982, 2.05039096, 1.51179266, 1.08895338, 0.803307, 0.59516323, 0.43325692, 0.29807833, 0.19894916, 0.09824532, 0.02916753],
+        [14.61464119, 7.49001646, 5.09240818, 3.60512662, 2.84484982, 2.12350607, 1.61558151, 1.24153244, 0.95350921, 0.72133851, 0.54755926, 0.41087446, 0.29807833, 0.19894916, 0.09824532, 0.02916753],
+        [14.61464119, 7.49001646, 5.85520077, 4.45427561, 3.1956799, 2.45070267, 1.84880662, 1.41535246, 1.08895338, 0.83188516, 0.64427125, 0.50118381, 0.36617002, 0.25053367, 0.17026083, 0.09824532, 0.02916753],
+        [14.61464119, 7.49001646, 5.85520077, 4.45427561, 3.1956799, 2.45070267, 1.91321158, 1.51179266, 1.20157266, 0.95350921, 0.74807048, 0.59516323, 0.45573691, 0.34370604, 0.25053367, 0.17026083, 0.09824532, 0.02916753],
+        [14.61464119, 7.49001646, 5.85520077, 4.45427561, 3.46139455, 2.84484982, 2.19988537, 1.72759056, 1.36964464, 1.08895338, 0.86115354, 0.69515091, 0.54755926, 0.43325692, 0.34370604, 0.25053367, 0.17026083, 0.09824532, 0.02916753],
+        [14.61464119, 11.54541874, 7.49001646, 5.85520077, 4.45427561, 3.46139455, 2.84484982, 2.19988537, 1.72759056, 1.36964464, 1.08895338, 0.86115354, 0.69515091, 0.54755926, 0.43325692, 0.34370604, 0.25053367, 0.17026083, 0.09824532, 0.02916753],
+        [14.61464119, 11.54541874, 7.49001646, 5.85520077, 4.45427561, 3.46139455, 2.84484982, 2.19988537, 1.72759056, 1.36964464, 1.08895338, 0.89115214, 0.72133851, 0.59516323, 0.4783645, 0.38853383, 0.29807833, 0.22545385, 0.17026083, 0.09824532, 0.02916753],
+    ], 
+    1.15: [
+        [14.61464119, 0.83188516, 0.02916753],
+        [14.61464119, 1.84880662, 0.59516323, 0.02916753],
+        [14.61464119, 5.85520077, 1.56271636, 0.52423614, 0.02916753],
+        [14.61464119, 5.85520077, 1.91321158, 0.83188516, 0.34370604, 0.02916753],
+        [14.61464119, 5.85520077, 2.45070267, 1.24153244, 0.59516323, 0.25053367, 0.02916753],
+        [14.61464119, 5.85520077, 2.84484982, 1.51179266, 0.803307, 0.41087446, 0.17026083, 0.02916753],
+        [14.61464119, 5.85520077, 2.84484982, 1.56271636, 0.89115214, 0.50118381, 0.25053367, 0.09824532, 0.02916753],
+        [14.61464119, 6.77309084, 3.07277966, 1.84880662, 1.12534678, 0.72133851, 0.43325692, 0.22545385, 0.09824532, 0.02916753],
+        [14.61464119, 6.77309084, 3.07277966, 1.91321158, 1.24153244, 0.803307, 0.52423614, 0.34370604, 0.19894916, 0.09824532, 0.02916753],
+        [14.61464119, 7.49001646, 4.86714602, 2.95596409, 1.91321158, 1.24153244, 0.803307, 0.52423614, 0.34370604, 0.19894916, 0.09824532, 0.02916753],
+        [14.61464119, 7.49001646, 4.86714602, 3.07277966, 2.05039096, 1.36964464, 0.95350921, 0.69515091, 0.4783645, 0.32104823, 0.19894916, 0.09824532, 0.02916753],
+        [14.61464119, 7.49001646, 4.86714602, 3.07277966, 2.12350607, 1.51179266, 1.08895338, 0.803307, 0.59516323, 0.43325692, 0.29807833, 0.19894916, 0.09824532, 0.02916753],
+        [14.61464119, 7.49001646, 4.86714602, 3.07277966, 2.12350607, 1.51179266, 1.08895338, 0.803307, 0.59516323, 0.45573691, 0.34370604, 0.25053367, 0.17026083, 0.09824532, 0.02916753],
+        [14.61464119, 7.49001646, 4.86714602, 3.07277966, 2.19988537, 1.61558151, 1.24153244, 0.95350921, 0.74807048, 0.59516323, 0.45573691, 0.34370604, 0.25053367, 0.17026083, 0.09824532, 0.02916753],
+        [14.61464119, 7.49001646, 4.86714602, 3.1956799, 2.45070267, 1.78698075, 1.32549286, 1.01931262, 0.803307, 0.64427125, 0.50118381, 0.38853383, 0.29807833, 0.22545385, 0.17026083, 0.09824532, 0.02916753],
+        [14.61464119, 7.49001646, 4.86714602, 3.1956799, 2.45070267, 1.78698075, 1.32549286, 1.01931262, 0.803307, 0.64427125, 0.52423614, 0.41087446, 0.32104823, 0.25053367, 0.19894916, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 7.49001646, 4.86714602, 3.1956799, 2.45070267, 1.84880662, 1.41535246, 1.12534678, 0.89115214, 0.72133851, 0.59516323, 0.4783645, 0.38853383, 0.32104823, 0.25053367, 0.19894916, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 7.49001646, 4.86714602, 3.1956799, 2.45070267, 1.84880662, 1.41535246, 1.12534678, 0.89115214, 0.72133851, 0.59516323, 0.50118381, 0.41087446, 0.34370604, 0.27464288, 0.22545385, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 7.49001646, 4.86714602, 3.1956799, 2.45070267, 1.84880662, 1.41535246, 1.12534678, 0.89115214, 0.72133851, 0.59516323, 0.50118381, 0.41087446, 0.34370604, 0.29807833, 0.25053367, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+    ],
+    1.20: [
+        [14.61464119, 0.803307, 0.02916753],
+        [14.61464119, 1.56271636, 0.52423614, 0.02916753],
+        [14.61464119, 2.36326075, 0.92192322, 0.36617002, 0.02916753],
+        [14.61464119, 2.84484982, 1.24153244, 0.59516323, 0.25053367, 0.02916753],
+        [14.61464119, 5.85520077, 2.05039096, 0.95350921, 0.45573691, 0.17026083, 0.02916753],
+        [14.61464119, 5.85520077, 2.45070267, 1.24153244, 0.64427125, 0.29807833, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 2.45070267, 1.36964464, 0.803307, 0.45573691, 0.25053367, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 2.84484982, 1.61558151, 0.95350921, 0.59516323, 0.36617002, 0.19894916, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 2.84484982, 1.67050016, 1.08895338, 0.74807048, 0.50118381, 0.32104823, 0.19894916, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 2.95596409, 1.84880662, 1.24153244, 0.83188516, 0.59516323, 0.41087446, 0.27464288, 0.17026083, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 3.07277966, 1.98035145, 1.36964464, 0.95350921, 0.69515091, 0.50118381, 0.36617002, 0.25053367, 0.17026083, 0.09824532, 0.02916753],
+        [14.61464119, 6.77309084, 3.46139455, 2.36326075, 1.56271636, 1.08895338, 0.803307, 0.59516323, 0.45573691, 0.34370604, 0.25053367, 0.17026083, 0.09824532, 0.02916753],
+        [14.61464119, 6.77309084, 3.46139455, 2.45070267, 1.61558151, 1.162866, 0.86115354, 0.64427125, 0.50118381, 0.38853383, 0.29807833, 0.22545385, 0.17026083, 0.09824532, 0.02916753],
+        [14.61464119, 7.49001646, 4.65472794, 3.07277966, 2.12350607, 1.51179266, 1.08895338, 0.83188516, 0.64427125, 0.50118381, 0.38853383, 0.29807833, 0.22545385, 0.17026083, 0.09824532, 0.02916753],
+        [14.61464119, 7.49001646, 4.65472794, 3.07277966, 2.12350607, 1.51179266, 1.08895338, 0.83188516, 0.64427125, 0.50118381, 0.41087446, 0.32104823, 0.25053367, 0.19894916, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 7.49001646, 4.65472794, 3.07277966, 2.12350607, 1.51179266, 1.08895338, 0.83188516, 0.64427125, 0.50118381, 0.41087446, 0.34370604, 0.27464288, 0.22545385, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 7.49001646, 4.65472794, 3.07277966, 2.19988537, 1.61558151, 1.20157266, 0.92192322, 0.72133851, 0.57119018, 0.45573691, 0.36617002, 0.29807833, 0.25053367, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 7.49001646, 4.65472794, 3.07277966, 2.19988537, 1.61558151, 1.24153244, 0.95350921, 0.74807048, 0.59516323, 0.4783645, 0.38853383, 0.32104823, 0.27464288, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 7.49001646, 4.65472794, 3.07277966, 2.19988537, 1.61558151, 1.24153244, 0.95350921, 0.74807048, 0.59516323, 0.50118381, 0.41087446, 0.34370604, 0.29807833, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+    ],
+    1.25: [
+        [14.61464119, 0.72133851, 0.02916753],
+        [14.61464119, 1.56271636, 0.50118381, 0.02916753],
+        [14.61464119, 2.05039096, 0.803307, 0.32104823, 0.02916753],
+        [14.61464119, 2.36326075, 0.95350921, 0.43325692, 0.17026083, 0.02916753],
+        [14.61464119, 2.84484982, 1.24153244, 0.59516323, 0.27464288, 0.09824532, 0.02916753],
+        [14.61464119, 3.07277966, 1.51179266, 0.803307, 0.43325692, 0.22545385, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 2.36326075, 1.24153244, 0.72133851, 0.41087446, 0.22545385, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 2.45070267, 1.36964464, 0.83188516, 0.52423614, 0.34370604, 0.19894916, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 2.84484982, 1.61558151, 0.98595673, 0.64427125, 0.43325692, 0.27464288, 0.17026083, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 2.84484982, 1.67050016, 1.08895338, 0.74807048, 0.52423614, 0.36617002, 0.25053367, 0.17026083, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 2.84484982, 1.72759056, 1.162866, 0.803307, 0.59516323, 0.45573691, 0.34370604, 0.25053367, 0.17026083, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 2.95596409, 1.84880662, 1.24153244, 0.86115354, 0.64427125, 0.4783645, 0.36617002, 0.27464288, 0.19894916, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 2.95596409, 1.84880662, 1.28281462, 0.92192322, 0.69515091, 0.52423614, 0.41087446, 0.32104823, 0.25053367, 0.19894916, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 2.95596409, 1.91321158, 1.32549286, 0.95350921, 0.72133851, 0.54755926, 0.43325692, 0.34370604, 0.27464288, 0.22545385, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 2.95596409, 1.91321158, 1.32549286, 0.95350921, 0.72133851, 0.57119018, 0.45573691, 0.36617002, 0.29807833, 0.25053367, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 2.95596409, 1.91321158, 1.32549286, 0.95350921, 0.74807048, 0.59516323, 0.4783645, 0.38853383, 0.32104823, 0.27464288, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 3.07277966, 2.05039096, 1.41535246, 1.05362725, 0.803307, 0.61951244, 0.50118381, 0.41087446, 0.34370604, 0.29807833, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 3.07277966, 2.05039096, 1.41535246, 1.05362725, 0.803307, 0.64427125, 0.52423614, 0.43325692, 0.36617002, 0.32104823, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 3.07277966, 2.05039096, 1.46270394, 1.08895338, 0.83188516, 0.66947293, 0.54755926, 0.45573691, 0.38853383, 0.34370604, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+    ],
+    1.30: [
+        [14.61464119, 0.72133851, 0.02916753],
+        [14.61464119, 1.24153244, 0.43325692, 0.02916753],
+        [14.61464119, 1.56271636, 0.59516323, 0.22545385, 0.02916753],
+        [14.61464119, 1.84880662, 0.803307, 0.36617002, 0.13792117, 0.02916753],
+        [14.61464119, 2.36326075, 1.01931262, 0.52423614, 0.25053367, 0.09824532, 0.02916753],
+        [14.61464119, 2.84484982, 1.36964464, 0.74807048, 0.41087446, 0.22545385, 0.09824532, 0.02916753],
+        [14.61464119, 3.07277966, 1.56271636, 0.89115214, 0.54755926, 0.34370604, 0.19894916, 0.09824532, 0.02916753],
+        [14.61464119, 3.07277966, 1.61558151, 0.95350921, 0.61951244, 0.41087446, 0.27464288, 0.17026083, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 2.45070267, 1.36964464, 0.83188516, 0.54755926, 0.36617002, 0.25053367, 0.17026083, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 2.45070267, 1.41535246, 0.92192322, 0.64427125, 0.45573691, 0.34370604, 0.25053367, 0.17026083, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 2.6383388, 1.56271636, 1.01931262, 0.72133851, 0.50118381, 0.36617002, 0.27464288, 0.19894916, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 2.84484982, 1.61558151, 1.05362725, 0.74807048, 0.54755926, 0.41087446, 0.32104823, 0.25053367, 0.19894916, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 2.84484982, 1.61558151, 1.08895338, 0.77538133, 0.57119018, 0.43325692, 0.34370604, 0.27464288, 0.22545385, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 2.84484982, 1.61558151, 1.08895338, 0.803307, 0.59516323, 0.45573691, 0.36617002, 0.29807833, 0.25053367, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 2.84484982, 1.61558151, 1.08895338, 0.803307, 0.59516323, 0.4783645, 0.38853383, 0.32104823, 0.27464288, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 2.84484982, 1.72759056, 1.162866, 0.83188516, 0.64427125, 0.50118381, 0.41087446, 0.34370604, 0.29807833, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 2.84484982, 1.72759056, 1.162866, 0.83188516, 0.64427125, 0.52423614, 0.43325692, 0.36617002, 0.32104823, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 2.84484982, 1.78698075, 1.24153244, 0.92192322, 0.72133851, 0.57119018, 0.45573691, 0.38853383, 0.34370604, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 2.84484982, 1.78698075, 1.24153244, 0.92192322, 0.72133851, 0.57119018, 0.4783645, 0.41087446, 0.36617002, 0.32104823, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+    ], 
+    1.35: [
+        [14.61464119, 0.69515091, 0.02916753],
+        [14.61464119, 0.95350921, 0.34370604, 0.02916753],
+        [14.61464119, 1.56271636, 0.57119018, 0.19894916, 0.02916753],
+        [14.61464119, 1.61558151, 0.69515091, 0.29807833, 0.09824532, 0.02916753],
+        [14.61464119, 1.84880662, 0.83188516, 0.43325692, 0.22545385, 0.09824532, 0.02916753],
+        [14.61464119, 2.45070267, 1.162866, 0.64427125, 0.36617002, 0.19894916, 0.09824532, 0.02916753],
+        [14.61464119, 2.84484982, 1.36964464, 0.803307, 0.50118381, 0.32104823, 0.19894916, 0.09824532, 0.02916753],
+        [14.61464119, 2.84484982, 1.41535246, 0.83188516, 0.54755926, 0.36617002, 0.25053367, 0.17026083, 0.09824532, 0.02916753],
+        [14.61464119, 2.84484982, 1.56271636, 0.95350921, 0.64427125, 0.45573691, 0.32104823, 0.22545385, 0.17026083, 0.09824532, 0.02916753],
+        [14.61464119, 2.84484982, 1.56271636, 0.95350921, 0.64427125, 0.45573691, 0.34370604, 0.25053367, 0.19894916, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 3.07277966, 1.61558151, 1.01931262, 0.72133851, 0.52423614, 0.38853383, 0.29807833, 0.22545385, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 3.07277966, 1.61558151, 1.01931262, 0.72133851, 0.52423614, 0.41087446, 0.32104823, 0.25053367, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 3.07277966, 1.61558151, 1.05362725, 0.74807048, 0.54755926, 0.43325692, 0.34370604, 0.27464288, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 3.07277966, 1.72759056, 1.12534678, 0.803307, 0.59516323, 0.45573691, 0.36617002, 0.29807833, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 3.07277966, 1.72759056, 1.12534678, 0.803307, 0.59516323, 0.4783645, 0.38853383, 0.32104823, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 2.45070267, 1.51179266, 1.01931262, 0.74807048, 0.57119018, 0.45573691, 0.36617002, 0.32104823, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 2.6383388, 1.61558151, 1.08895338, 0.803307, 0.61951244, 0.50118381, 0.41087446, 0.34370604, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 2.6383388, 1.61558151, 1.08895338, 0.803307, 0.64427125, 0.52423614, 0.43325692, 0.36617002, 0.32104823, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 5.85520077, 2.6383388, 1.61558151, 1.08895338, 0.803307, 0.64427125, 0.52423614, 0.45573691, 0.38853383, 0.34370604, 0.32104823, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+    ],
+    1.40: [
+        [14.61464119, 0.59516323, 0.02916753],
+        [14.61464119, 0.95350921, 0.34370604, 0.02916753],
+        [14.61464119, 1.08895338, 0.43325692, 0.13792117, 0.02916753],
+        [14.61464119, 1.56271636, 0.64427125, 0.27464288, 0.09824532, 0.02916753],
+        [14.61464119, 1.61558151, 0.803307, 0.43325692, 0.22545385, 0.09824532, 0.02916753],
+        [14.61464119, 2.05039096, 0.95350921, 0.54755926, 0.34370604, 0.19894916, 0.09824532, 0.02916753],
+        [14.61464119, 2.45070267, 1.24153244, 0.72133851, 0.43325692, 0.27464288, 0.17026083, 0.09824532, 0.02916753],
+        [14.61464119, 2.45070267, 1.24153244, 0.74807048, 0.50118381, 0.34370604, 0.25053367, 0.17026083, 0.09824532, 0.02916753],
+        [14.61464119, 2.45070267, 1.28281462, 0.803307, 0.52423614, 0.36617002, 0.27464288, 0.19894916, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 2.45070267, 1.28281462, 0.803307, 0.54755926, 0.38853383, 0.29807833, 0.22545385, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 2.84484982, 1.41535246, 0.86115354, 0.59516323, 0.43325692, 0.32104823, 0.25053367, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 2.84484982, 1.51179266, 0.95350921, 0.64427125, 0.45573691, 0.34370604, 0.27464288, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 2.84484982, 1.51179266, 0.95350921, 0.64427125, 0.4783645, 0.36617002, 0.29807833, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 2.84484982, 1.56271636, 0.98595673, 0.69515091, 0.52423614, 0.41087446, 0.34370604, 0.29807833, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 2.84484982, 1.56271636, 1.01931262, 0.72133851, 0.54755926, 0.43325692, 0.36617002, 0.32104823, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 2.84484982, 1.61558151, 1.05362725, 0.74807048, 0.57119018, 0.45573691, 0.38853383, 0.34370604, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 2.84484982, 1.61558151, 1.08895338, 0.803307, 0.61951244, 0.50118381, 0.41087446, 0.36617002, 0.32104823, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 2.84484982, 1.61558151, 1.08895338, 0.803307, 0.61951244, 0.50118381, 0.43325692, 0.38853383, 0.34370604, 0.32104823, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 2.84484982, 1.61558151, 1.08895338, 0.803307, 0.64427125, 0.52423614, 0.45573691, 0.41087446, 0.36617002, 0.34370604, 0.32104823, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+    ],
+    1.45: [
+        [14.61464119, 0.59516323, 0.02916753],
+        [14.61464119, 0.803307, 0.25053367, 0.02916753],
+        [14.61464119, 0.95350921, 0.34370604, 0.09824532, 0.02916753],
+        [14.61464119, 1.24153244, 0.54755926, 0.25053367, 0.09824532, 0.02916753],
+        [14.61464119, 1.56271636, 0.72133851, 0.36617002, 0.19894916, 0.09824532, 0.02916753],
+        [14.61464119, 1.61558151, 0.803307, 0.45573691, 0.27464288, 0.17026083, 0.09824532, 0.02916753],
+        [14.61464119, 1.91321158, 0.95350921, 0.57119018, 0.36617002, 0.25053367, 0.17026083, 0.09824532, 0.02916753],
+        [14.61464119, 2.19988537, 1.08895338, 0.64427125, 0.41087446, 0.27464288, 0.19894916, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 2.45070267, 1.24153244, 0.74807048, 0.50118381, 0.34370604, 0.25053367, 0.19894916, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 2.45070267, 1.24153244, 0.74807048, 0.50118381, 0.36617002, 0.27464288, 0.22545385, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 2.45070267, 1.28281462, 0.803307, 0.54755926, 0.41087446, 0.32104823, 0.25053367, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 2.45070267, 1.28281462, 0.803307, 0.57119018, 0.43325692, 0.34370604, 0.27464288, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 2.45070267, 1.28281462, 0.83188516, 0.59516323, 0.45573691, 0.36617002, 0.29807833, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 2.45070267, 1.28281462, 0.83188516, 0.59516323, 0.45573691, 0.36617002, 0.32104823, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 2.84484982, 1.51179266, 0.95350921, 0.69515091, 0.52423614, 0.41087446, 0.34370604, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 2.84484982, 1.51179266, 0.95350921, 0.69515091, 0.52423614, 0.43325692, 0.36617002, 0.32104823, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 2.84484982, 1.56271636, 0.98595673, 0.72133851, 0.54755926, 0.45573691, 0.38853383, 0.34370604, 0.32104823, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 2.84484982, 1.56271636, 1.01931262, 0.74807048, 0.57119018, 0.4783645, 0.41087446, 0.36617002, 0.34370604, 0.32104823, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 2.84484982, 1.56271636, 1.01931262, 0.74807048, 0.59516323, 0.50118381, 0.43325692, 0.38853383, 0.36617002, 0.34370604, 0.32104823, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+    ],
+    1.50: [
+        [14.61464119, 0.54755926, 0.02916753],
+        [14.61464119, 0.803307, 0.25053367, 0.02916753],
+        [14.61464119, 0.86115354, 0.32104823, 0.09824532, 0.02916753],
+        [14.61464119, 1.24153244, 0.54755926, 0.25053367, 0.09824532, 0.02916753],
+        [14.61464119, 1.56271636, 0.72133851, 0.36617002, 0.19894916, 0.09824532, 0.02916753],
+        [14.61464119, 1.61558151, 0.803307, 0.45573691, 0.27464288, 0.17026083, 0.09824532, 0.02916753],
+        [14.61464119, 1.61558151, 0.83188516, 0.52423614, 0.34370604, 0.25053367, 0.17026083, 0.09824532, 0.02916753],
+        [14.61464119, 1.84880662, 0.95350921, 0.59516323, 0.38853383, 0.27464288, 0.19894916, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 1.84880662, 0.95350921, 0.59516323, 0.41087446, 0.29807833, 0.22545385, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 1.84880662, 0.95350921, 0.61951244, 0.43325692, 0.32104823, 0.25053367, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 2.19988537, 1.12534678, 0.72133851, 0.50118381, 0.36617002, 0.27464288, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 2.19988537, 1.12534678, 0.72133851, 0.50118381, 0.36617002, 0.29807833, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 2.36326075, 1.24153244, 0.803307, 0.57119018, 0.43325692, 0.34370604, 0.29807833, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 2.36326075, 1.24153244, 0.803307, 0.57119018, 0.43325692, 0.34370604, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 2.36326075, 1.24153244, 0.803307, 0.59516323, 0.45573691, 0.36617002, 0.32104823, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 2.36326075, 1.24153244, 0.803307, 0.59516323, 0.45573691, 0.38853383, 0.34370604, 0.32104823, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 2.45070267, 1.32549286, 0.86115354, 0.64427125, 0.50118381, 0.41087446, 0.36617002, 0.34370604, 0.32104823, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 2.45070267, 1.36964464, 0.92192322, 0.69515091, 0.54755926, 0.45573691, 0.41087446, 0.36617002, 0.34370604, 0.32104823, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+        [14.61464119, 2.45070267, 1.41535246, 0.95350921, 0.72133851, 0.57119018, 0.4783645, 0.43325692, 0.38853383, 0.36617002, 0.34370604, 0.32104823, 0.29807833, 0.27464288, 0.25053367, 0.22545385, 0.19894916, 0.17026083, 0.13792117, 0.09824532, 0.02916753],
+    ],
+}
+
+class GITSScheduler:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required":
+                    {"coeff": ("FLOAT", {"default": 1.20, "min": 0.80, "max": 1.50, "step": 0.05}),
+                     "steps": ("INT", {"default": 10, "min": 2, "max": 1000}),
+                     "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
+                      }
+               }
+    RETURN_TYPES = ("SIGMAS",)
+    CATEGORY = "sampling/custom_sampling/schedulers"
+
+    FUNCTION = "get_sigmas"
+
+    def get_sigmas(self, coeff, steps, denoise):
+        total_steps = steps
+        if denoise < 1.0:
+            if denoise <= 0.0:
+                return (torch.FloatTensor([]),)
+            total_steps = round(steps * denoise)
+
+        if steps <= 20:
+            sigmas = NOISE_LEVELS[round(coeff, 2)][steps-2][:]
+        else:
+            sigmas = NOISE_LEVELS[round(coeff, 2)][-1][:]
+            sigmas = loglinear_interp(sigmas, steps + 1)
+
+        sigmas = sigmas[-(total_steps + 1):]
+        sigmas[-1] = 0
+        return (torch.FloatTensor(sigmas), )
+
+NODE_CLASS_MAPPINGS = {
+    "GITSScheduler": GITSScheduler,
+}

comfy_extras/nodes_hypertile.py

@@ -32,7 +32,7 @@ class HyperTile:
     RETURN_TYPES = ("MODEL",)
     FUNCTION = "patch"
 
-    CATEGORY = "model_patches"
+    CATEGORY = "model_patches/unet"
 
     def patch(self, model, tile_size, swap_size, max_depth, scale_depth):
         model_channels = model.model.model_config.unet_config["model_channels"]

comfy_extras/nodes_model_advanced.py

@@ -144,7 +144,7 @@ class ModelSamplingSD3:
 
     CATEGORY = "advanced/model"
 
-    def patch(self, model, shift):
+    def patch(self, model, shift, multiplier=1000):
         m = model.clone()
 
         sampling_base = comfy.model_sampling.ModelSamplingDiscreteFlow
@@ -154,10 +154,22 @@ class ModelSamplingSD3:
             pass
 
         model_sampling = ModelSamplingAdvanced(model.model.model_config)
-        model_sampling.set_parameters(shift=shift)
+        model_sampling.set_parameters(shift=shift, multiplier=multiplier)
         m.add_object_patch("model_sampling", model_sampling)
         return (m, )
 
+class ModelSamplingAuraFlow(ModelSamplingSD3):
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": { "model": ("MODEL",),
+                              "shift": ("FLOAT", {"default": 1.73, "min": 0.0, "max": 100.0, "step":0.01}),
+                              }}
+
+    FUNCTION = "patch_aura"
+
+    def patch_aura(self, model, shift):
+        return self.patch(model, shift, multiplier=1.0)
+
 class ModelSamplingContinuousEDM:
     @classmethod
     def INPUT_TYPES(s):
@@ -271,5 +283,6 @@ NODE_CLASS_MAPPINGS = {
     "ModelSamplingContinuousV": ModelSamplingContinuousV,
     "ModelSamplingStableCascade": ModelSamplingStableCascade,
     "ModelSamplingSD3": ModelSamplingSD3,
+    "ModelSamplingAuraFlow": ModelSamplingAuraFlow,
     "RescaleCFG": RescaleCFG,
 }

comfy_extras/nodes_model_merging_model_specific.py

@@ -52,9 +52,32 @@ class ModelMergeSDXL(comfy_extras.nodes_model_merging.ModelMergeBlocks):
 
         return {"required": arg_dict}
 
+class ModelMergeSD3_2B(comfy_extras.nodes_model_merging.ModelMergeBlocks):
+    CATEGORY = "advanced/model_merging/model_specific"
+
+    @classmethod
+    def INPUT_TYPES(s):
+        arg_dict = { "model1": ("MODEL",),
+                              "model2": ("MODEL",)}
+
+        argument = ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01})
+
+        arg_dict["pos_embed."] = argument
+        arg_dict["x_embedder."] = argument
+        arg_dict["context_embedder."] = argument
+        arg_dict["y_embedder."] = argument
+        arg_dict["t_embedder."] = argument
+
+        for i in range(24):
+            arg_dict["joint_blocks.{}.".format(i)] = argument
+
+        arg_dict["final_layer."] = argument
+
+        return {"required": arg_dict}
 
 NODE_CLASS_MAPPINGS = {
     "ModelMergeSD1": ModelMergeSD1,
     "ModelMergeSD2": ModelMergeSD1, #SD1 and SD2 have the same blocks
     "ModelMergeSDXL": ModelMergeSDXL,
+    "ModelMergeSD3_2B": ModelMergeSD3_2B,
 }

comfy_extras/nodes_pag.py

@@ -19,7 +19,7 @@ class PerturbedAttentionGuidance:
     RETURN_TYPES = ("MODEL",)
     FUNCTION = "patch"
 
-    CATEGORY = "_for_testing"
+    CATEGORY = "model_patches/unet"
 
     def patch(self, model, scale):
         unet_block = "middle"

comfy_extras/nodes_sd3.py

@@ -80,8 +80,23 @@ class CLIPTextEncodeSD3:
         return ([[cond, {"pooled_output": pooled}]], )
 
 
+class ControlNetApplySD3(nodes.ControlNetApplyAdvanced):
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {"positive": ("CONDITIONING", ),
+                             "negative": ("CONDITIONING", ),
+                             "control_net": ("CONTROL_NET", ),
+                             "vae": ("VAE", ),
+                             "image": ("IMAGE", ),
+                             "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
+                             "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
+                             "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001})
+                             }}
+    CATEGORY = "conditioning/controlnet"
+
 NODE_CLASS_MAPPINGS = {
     "TripleCLIPLoader": TripleCLIPLoader,
     "EmptySD3LatentImage": EmptySD3LatentImage,
     "CLIPTextEncodeSD3": CLIPTextEncodeSD3,
+    "ControlNetApplySD3": ControlNetApplySD3,
 }

custom_nodes/example_node.py.example

@@ -12,9 +12,9 @@ class Example:
     Attributes
     ----------
     RETURN_TYPES (`tuple`): 
-        The type of each element in the output tulple.
+        The type of each element in the output tuple.
     RETURN_NAMES (`tuple`):
-        Optional: The name of each output in the output tulple.
+        Optional: The name of each output in the output tuple.
     FUNCTION (`str`):
         The name of the entry-point method. For example, if `FUNCTION = "execute"` then it will run Example().execute()
     OUTPUT_NODE ([`bool`]):
@@ -44,7 +44,7 @@ class Example:
                     * Key field_name (`string`): Name of a entry-point method's argument
                     * Value field_config (`tuple`):
                         + First value is a string indicate the type of field or a list for selection.
-                        + Secound value is a config for type "INT", "STRING" or "FLOAT".
+                        + Second value is a config for type "INT", "STRING" or "FLOAT".
         """
         return {
             "required": {
@@ -62,7 +62,7 @@ class Example:
                     "min": 0.0,
                     "max": 10.0,
                     "step": 0.01,
-                    "round": 0.001, #The value represeting the precision to round to, will be set to the step value by default. Can be set to False to disable rounding.
+                    "round": 0.001, #The value representing the precision to round to, will be set to the step value by default. Can be set to False to disable rounding.
                     "display": "number",
                     "lazy": True
                 }),
@@ -127,6 +127,16 @@ class Example:
 # Set the web directory, any .js file in that directory will be loaded by the frontend as a frontend extension
 # WEB_DIRECTORY = "./somejs"
 
+
+# Add custom API routes, using router
+from aiohttp import web
+from server import PromptServer
+
+@PromptServer.instance.routes.get("/hello")
+async def get_hello(request):
+    return web.json_response("hello")
+
+
 # A dictionary that contains all nodes you want to export with their names
 # NOTE: names should be globally unique
 NODE_CLASS_MAPPINGS = {

folder_paths.py

@@ -1,10 +1,13 @@
 import os
 import time
 import logging
+from typing import Set, List, Dict, Tuple
 
-supported_pt_extensions = set(['.ckpt', '.pt', '.bin', '.pth', '.safetensors', '.pkl'])
+supported_pt_extensions: Set[str] = set(['.ckpt', '.pt', '.bin', '.pth', '.safetensors', '.pkl'])
 
-folder_names_and_paths = {}
+SupportedFileExtensionsType = Set[str]
+ScanPathType = List[str]
+folder_names_and_paths: Dict[str, Tuple[ScanPathType, SupportedFileExtensionsType]] = {}
 
 base_path = os.path.dirname(os.path.realpath(__file__))
 models_dir = os.path.join(base_path, "models")
@@ -26,7 +29,7 @@ folder_names_and_paths["gligen"] = ([os.path.join(models_dir, "gligen")], suppor
 
 folder_names_and_paths["upscale_models"] = ([os.path.join(models_dir, "upscale_models")], supported_pt_extensions)
 
-folder_names_and_paths["custom_nodes"] = ([os.path.join(base_path, "custom_nodes")], [])
+folder_names_and_paths["custom_nodes"] = ([os.path.join(base_path, "custom_nodes")], set())
 
 folder_names_and_paths["hypernetworks"] = ([os.path.join(models_dir, "hypernetworks")], supported_pt_extensions)
 

main.py

@@ -5,6 +5,8 @@ import os
 import importlib.util
 import folder_paths
 import time
+from comfy.cli_args import args
+
 
 def execute_prestartup_script():
     def execute_script(script_path):
@@ -18,6 +20,9 @@ def execute_prestartup_script():
             print(f"Failed to execute startup-script: {script_path} / {e}")
         return False
 
+    if args.disable_all_custom_nodes:
+        return
+
     node_paths = folder_paths.get_folder_paths("custom_nodes")
     for custom_node_path in node_paths:
         possible_modules = os.listdir(custom_node_path)
@@ -53,7 +58,6 @@ import shutil
 import threading
 import gc
 
-from comfy.cli_args import args
 import logging
 
 if os.name == "nt":
@@ -76,7 +80,7 @@ import yaml
 import execution
 import server
 from server import BinaryEventTypes
-from nodes import init_custom_nodes
+import nodes
 import comfy.model_management
 
 def cuda_malloc_warning():
@@ -214,7 +218,7 @@ if __name__ == "__main__":
         for config_path in itertools.chain(*args.extra_model_paths_config):
             load_extra_path_config(config_path)
 
-    init_custom_nodes()
+    nodes.init_extra_nodes(init_custom_nodes=not args.disable_all_custom_nodes)
 
     cuda_malloc_warning()
 

node_helpers.py

@@ -1,3 +1,7 @@
+import hashlib
+
+from comfy.cli_args import args
+
 from PIL import ImageFile, UnidentifiedImageError
 
 def conditioning_set_values(conditioning, values={}):
@@ -22,3 +26,12 @@ def pillow(fn, arg):
         if prev_value is not None:
             ImageFile.LOAD_TRUNCATED_IMAGES = prev_value
         return x
+
+def hasher():
+    hashfuncs = {
+        "md5": hashlib.md5,
+        "sha1": hashlib.sha1,
+        "sha256": hashlib.sha256,
+        "sha512": hashlib.sha512
+    }
+    return hashfuncs[args.default_hashing_function]

nodes.py

@@ -55,8 +55,9 @@ class CLIPTextEncode:
 
     def encode(self, clip, text):
         tokens = clip.tokenize(text)
-        cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)
+        output = clip.encode_from_tokens(tokens, return_pooled=True, return_dict=True)
-        return ([[cond, {"pooled_output": pooled}]], )
+        cond = output.pop("cond")
+        return ([[cond, output]], )
 
 class ConditioningCombine:
     @classmethod
@@ -232,8 +233,9 @@ class ConditioningZeroOut:
         c = []
         for t in conditioning:
             d = t[1].copy()
-            if "pooled_output" in d:
+            pooled_output = d.get("pooled_output", None)
-                d["pooled_output"] = torch.zeros_like(d["pooled_output"])
+            if pooled_output is not None:
+                d["pooled_output"] = torch.zeros_like(pooled_output)
             n = [torch.zeros_like(t[0]), d]
             c.append(n)
         return (c, )
@@ -746,7 +748,7 @@ class ControlNetApply:
     RETURN_TYPES = ("CONDITIONING",)
     FUNCTION = "apply_controlnet"
 
-    CATEGORY = "conditioning"
+    CATEGORY = "conditioning/controlnet"
 
     def apply_controlnet(self, conditioning, control_net, image, strength):
         if strength == 0:
@@ -781,9 +783,9 @@ class ControlNetApplyAdvanced:
     RETURN_NAMES = ("positive", "negative")
     FUNCTION = "apply_controlnet"
 
-    CATEGORY = "conditioning"
+    CATEGORY = "conditioning/controlnet"
 
-    def apply_controlnet(self, positive, negative, control_net, image, strength, start_percent, end_percent):
+    def apply_controlnet(self, positive, negative, control_net, image, strength, start_percent, end_percent, vae=None):
         if strength == 0:
             return (positive, negative)
 
@@ -800,7 +802,7 @@ class ControlNetApplyAdvanced:
                 if prev_cnet in cnets:
                     c_net = cnets[prev_cnet]
                 else:
-                    c_net = control_net.copy().set_cond_hint(control_hint, strength, (start_percent, end_percent))
+                    c_net = control_net.copy().set_cond_hint(control_hint, strength, (start_percent, end_percent), vae)
                     c_net.set_previous_controlnet(prev_cnet)
                     cnets[prev_cnet] = c_net
 
@@ -1887,7 +1889,30 @@ NODE_DISPLAY_NAME_MAPPINGS = {
 
 EXTENSION_WEB_DIRS = {}
 
-def load_custom_node(module_path, ignore=set()):
+
+def get_module_name(module_path: str) -> str:
+    """
+    Returns the module name based on the given module path.
+    Examples:
+        get_module_name("C:/Users/username/ComfyUI/custom_nodes/my_custom_node.py") -> "my_custom_node"
+        get_module_name("C:/Users/username/ComfyUI/custom_nodes/my_custom_node") -> "my_custom_node"
+        get_module_name("C:/Users/username/ComfyUI/custom_nodes/my_custom_node/") -> "my_custom_node"
+        get_module_name("C:/Users/username/ComfyUI/custom_nodes/my_custom_node/__init__.py") -> "my_custom_node"
+        get_module_name("C:/Users/username/ComfyUI/custom_nodes/my_custom_node/__init__") -> "my_custom_node"
+        get_module_name("C:/Users/username/ComfyUI/custom_nodes/my_custom_node/__init__/") -> "my_custom_node"
+        get_module_name("C:/Users/username/ComfyUI/custom_nodes/my_custom_node.disabled") -> "custom_nodes
+    Args:
+        module_path (str): The path of the module.
+    Returns:
+        str: The module name.
+    """
+    base_path = os.path.basename(module_path)
+    if os.path.isfile(module_path):
+        base_path = os.path.splitext(base_path)[0]
+    return base_path
+
+
+def load_custom_node(module_path: str, ignore=set(), module_parent="custom_nodes") -> bool:
     module_name = os.path.basename(module_path)
     if os.path.isfile(module_path):
         sp = os.path.splitext(module_path)
@@ -1911,9 +1936,10 @@ def load_custom_node(module_path, ignore=set()):
                 EXTENSION_WEB_DIRS[module_name] = web_dir
 
         if hasattr(module, "NODE_CLASS_MAPPINGS") and getattr(module, "NODE_CLASS_MAPPINGS") is not None:
-            for name in module.NODE_CLASS_MAPPINGS:
+            for name, node_cls in module.NODE_CLASS_MAPPINGS.items():
                 if name not in ignore:
-                    NODE_CLASS_MAPPINGS[name] = module.NODE_CLASS_MAPPINGS[name]
+                    NODE_CLASS_MAPPINGS[name] = node_cls
+                    node_cls.RELATIVE_PYTHON_MODULE = "{}.{}".format(module_parent, get_module_name(module_path))
             if hasattr(module, "NODE_DISPLAY_NAME_MAPPINGS") and getattr(module, "NODE_DISPLAY_NAME_MAPPINGS") is not None:
                 NODE_DISPLAY_NAME_MAPPINGS.update(module.NODE_DISPLAY_NAME_MAPPINGS)
             return True
@@ -1925,7 +1951,16 @@ def load_custom_node(module_path, ignore=set()):
         logging.warning(f"Cannot import {module_path} module for custom nodes: {e}")
         return False
 
-def load_custom_nodes():
+def init_external_custom_nodes():
+    """
+    Initializes the external custom nodes.
+
+    This function loads custom nodes from the specified folder paths and imports them into the application.
+    It measures the import times for each custom node and logs the results.
+
+    Returns:
+        None
+    """
     base_node_names = set(NODE_CLASS_MAPPINGS.keys())
     node_paths = folder_paths.get_folder_paths("custom_nodes")
     node_import_times = []
@@ -1939,7 +1974,7 @@ def load_custom_nodes():
             if os.path.isfile(module_path) and os.path.splitext(module_path)[1] != ".py": continue
             if module_path.endswith(".disabled"): continue
             time_before = time.perf_counter()
-            success = load_custom_node(module_path, base_node_names)
+            success = load_custom_node(module_path, base_node_names, module_parent="custom_nodes")
             node_import_times.append((time.perf_counter() - time_before, module_path, success))
 
     if len(node_import_times) > 0:
@@ -1952,7 +1987,16 @@ def load_custom_nodes():
             logging.info("{:6.1f} seconds{}: {}".format(n[0], import_message, n[1]))
         logging.info("")
 
-def init_custom_nodes():
+def init_builtin_extra_nodes():
+    """
+    Initializes the built-in extra nodes in ComfyUI.
+
+    This function loads the extra node files located in the "comfy_extras" directory and imports them into ComfyUI.
+    If any of the extra node files fail to import, a warning message is logged.
+
+    Returns:
+        None
+    """
     extras_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras")
     extras_files = [
         "nodes_latent.py",
@@ -1991,14 +2035,25 @@ def init_custom_nodes():
         "nodes_webcam.py",
         "nodes_audio.py",
         "nodes_sd3.py",
+        "nodes_gits.py",
+        "nodes_controlnet.py",
     ]
 
     import_failed = []
     for node_file in extras_files:
-        if not load_custom_node(os.path.join(extras_dir, node_file)):
+        if not load_custom_node(os.path.join(extras_dir, node_file), module_parent="comfy_extras"):
             import_failed.append(node_file)
 
-    load_custom_nodes()
+    return import_failed
+
+
+def init_extra_nodes(init_custom_nodes=True):
+    import_failed = init_builtin_extra_nodes()
+
+    if init_custom_nodes:
+        init_external_custom_nodes()
+    else:
+        logging.info("Skipping loading of custom nodes")
 
     if len(import_failed) > 0:
         logging.warning("WARNING: some comfy_extras/ nodes did not import correctly. This may be because they are missing some dependencies.\n")

pytest.ini

@@ -2,5 +2,8 @@
 markers = 
   inference: mark as inference test (deselect with '-m "not inference"')
   execution: mark as execution test (deselect with '-m "not execution"')
-testpaths = tests
+testpaths =
+  tests
+  tests-unit
 addopts = -s
+pythonpath = .

requirements.txt

@@ -3,7 +3,9 @@ torchsde
 torchvision
 torchaudio
 einops
-transformers>=4.25.1
+transformers>=4.28.1
+tokenizers>=0.13.3
+sentencepiece
 safetensors>=0.4.2
 aiohttp
 pyyaml
@@ -15,3 +17,4 @@ psutil
 #non essential dependencies:
 kornia>=0.7.1
 spandrel
+soundfile

server.py

@@ -12,6 +12,7 @@ import json
 import glob
 import struct
 import ssl
+import hashlib
 from PIL import Image, ImageOps
 from PIL.PngImagePlugin import PngInfo
 from io import BytesIO
@@ -24,9 +25,11 @@ import mimetypes
 from comfy.cli_args import args
 import comfy.utils
 import comfy.model_management
-
+import node_helpers
+from app.frontend_management import FrontendManager
 from app.user_manager import UserManager
 
+
 class BinaryEventTypes:
     PREVIEW_IMAGE = 1
     UNENCODED_PREVIEW_IMAGE = 2
@@ -82,8 +85,12 @@ class PromptServer():
         max_upload_size = round(args.max_upload_size * 1024 * 1024)
         self.app = web.Application(client_max_size=max_upload_size, middlewares=middlewares)
         self.sockets = dict()
-        self.web_root = os.path.join(os.path.dirname(
+        self.web_root = (
-            os.path.realpath(__file__)), "web")
+            FrontendManager.init_frontend(args.front_end_version)
+            if args.front_end_root is None
+            else args.front_end_root
+        )
+        logging.info(f"[Prompt Server] web root: {self.web_root}")
         routes = web.RouteTableDef()
         self.routes = routes
         self.last_node_id = None
@@ -154,9 +161,25 @@ class PromptServer():
 
             return type_dir, dir_type
 
+        def compare_image_hash(filepath, image):
+            hasher = node_helpers.hasher()
+            
+            # function to compare hashes of two images to see if it already exists, fix to #3465
+            if os.path.exists(filepath):
+                a = hasher()
+                b = hasher()
+                with open(filepath, "rb") as f:
+                    a.update(f.read())
+                    b.update(image.file.read())
+                    image.file.seek(0)
+                    f.close()
+                return a.hexdigest() == b.hexdigest()
+            return False
+
         def image_upload(post, image_save_function=None):
             image = post.get("image")
             overwrite = post.get("overwrite")
+            image_is_duplicate = False
 
             image_upload_type = post.get("type")
             upload_dir, image_upload_type = get_dir_by_type(image_upload_type)
@@ -183,15 +206,19 @@ class PromptServer():
                 else:
                     i = 1
                     while os.path.exists(filepath):
+                        if compare_image_hash(filepath, image): #compare hash to prevent saving of duplicates with same name, fix for #3465
+                            image_is_duplicate = True
+                            break
                         filename = f"{split[0]} ({i}){split[1]}"
                         filepath = os.path.join(full_output_folder, filename)
                         i += 1
 
-                if image_save_function is not None:
+                if not image_is_duplicate:
-                    image_save_function(image, post, filepath)
+                    if image_save_function is not None:
-                else:
+                        image_save_function(image, post, filepath)
-                    with open(filepath, "wb") as f:
+                    else:
-                        f.write(image.file.read())
+                        with open(filepath, "wb") as f:
+                            f.write(image.file.read())
 
                 return web.json_response({"name" : filename, "subfolder": subfolder, "type": image_upload_type})
             else:
@@ -398,6 +425,7 @@ class PromptServer():
             info['name'] = node_class
             info['display_name'] = nodes.NODE_DISPLAY_NAME_MAPPINGS[node_class] if node_class in nodes.NODE_DISPLAY_NAME_MAPPINGS.keys() else node_class
             info['description'] = obj_class.DESCRIPTION if hasattr(obj_class,'DESCRIPTION') else ''
+            info['python_module'] = getattr(obj_class, "RELATIVE_PYTHON_MODULE", "nodes")
             info['category'] = 'sd'
             if hasattr(obj_class, 'OUTPUT_NODE') and obj_class.OUTPUT_NODE == True:
                 info['output_node'] = True
@@ -531,6 +559,19 @@ class PromptServer():
 
     def add_routes(self):
         self.user_manager.add_routes(self.routes)
+
+        # Prefix every route with /api for easier matching for delegation.
+        # This is very useful for frontend dev server, which need to forward
+        # everything except serving of static files.
+        # Currently both the old endpoints without prefix and new endpoints with
+        # prefix are supported.
+        api_routes = web.RouteTableDef()
+        for route in self.routes:
+            # Custom nodes might add extra static routes. Only process non-static
+            # routes to add /api prefix.
+            if isinstance(route, web.RouteDef):
+                api_routes.route(route.method, "/api" + route.path)(route.handler, **route.kwargs)
+        self.app.add_routes(api_routes)
         self.app.add_routes(self.routes)
 
         for name, dir in nodes.EXTENSION_WEB_DIRS.items():

tests-ui/utils/setup.js

@@ -72,6 +72,7 @@ export function mockApi(config = {}) {
 		storeUserData: jest.fn((file, data) => {
 			userData[file] = data;
 		}),
+		listUserData: jest.fn(() => [])
 	};
 	jest.mock("../../web/scripts/api", () => ({
 		get api() {

tests-unit/README.md

@@ -0,0 +1,8 @@
+# Pytest Unit Tests
+
+## Install test dependencies
+
+`pip install -r tests-units/requirements.txt`
+
+## Run tests
+`pytest tests-units/`

tests-unit/app_test/frontend_manager_test.py

@@ -0,0 +1,100 @@
+import argparse
+import pytest
+from requests.exceptions import HTTPError
+
+from app.frontend_management import (
+    FrontendManager,
+    FrontEndProvider,
+    Release,
+)
+from comfy.cli_args import DEFAULT_VERSION_STRING
+
+
+@pytest.fixture
+def mock_releases():
+    return [
+        Release(
+            id=1,
+            tag_name="1.0.0",
+            name="Release 1.0.0",
+            prerelease=False,
+            created_at="2022-01-01T00:00:00Z",
+            published_at="2022-01-01T00:00:00Z",
+            body="Release notes for 1.0.0",
+            assets=[{"name": "dist.zip", "url": "https://example.com/dist.zip"}],
+        ),
+        Release(
+            id=2,
+            tag_name="2.0.0",
+            name="Release 2.0.0",
+            prerelease=False,
+            created_at="2022-02-01T00:00:00Z",
+            published_at="2022-02-01T00:00:00Z",
+            body="Release notes for 2.0.0",
+            assets=[{"name": "dist.zip", "url": "https://example.com/dist.zip"}],
+        ),
+    ]
+
+
+@pytest.fixture
+def mock_provider(mock_releases):
+    provider = FrontEndProvider(
+        owner="test-owner",
+        repo="test-repo",
+    )
+    provider.all_releases = mock_releases
+    provider.latest_release = mock_releases[1]
+    FrontendManager.PROVIDERS = [provider]
+    return provider
+
+
+def test_get_release(mock_provider, mock_releases):
+    version = "1.0.0"
+    release = mock_provider.get_release(version)
+    assert release == mock_releases[0]
+
+
+def test_get_release_latest(mock_provider, mock_releases):
+    version = "latest"
+    release = mock_provider.get_release(version)
+    assert release == mock_releases[1]
+
+
+def test_get_release_invalid_version(mock_provider):
+    version = "invalid"
+    with pytest.raises(ValueError):
+        mock_provider.get_release(version)
+
+
+def test_init_frontend_default():
+    version_string = DEFAULT_VERSION_STRING
+    frontend_path = FrontendManager.init_frontend(version_string)
+    assert frontend_path == FrontendManager.DEFAULT_FRONTEND_PATH
+
+
+def test_init_frontend_invalid_version():
+    version_string = "test-owner/test-repo@1.100.99"
+    with pytest.raises(HTTPError):
+        FrontendManager.init_frontend_unsafe(version_string)
+
+
+def test_init_frontend_invalid_provider():
+    version_string = "invalid/invalid@latest"
+    with pytest.raises(HTTPError):
+        FrontendManager.init_frontend_unsafe(version_string)
+
+
+def test_parse_version_string():
+    version_string = "owner/repo@1.0.0"
+    repo_owner, repo_name, version = FrontendManager.parse_version_string(
+        version_string
+    )
+    assert repo_owner == "owner"
+    assert repo_name == "repo"
+    assert version == "1.0.0"
+
+
+def test_parse_version_string_invalid():
+    version_string = "invalid"
+    with pytest.raises(argparse.ArgumentTypeError):
+        FrontendManager.parse_version_string(version_string)

tests-unit/requirements.txt

@@ -0,0 +1 @@
+pytest>=7.8.0

web/extensions/core/colorPalette.js

@@ -63,6 +63,10 @@ const colorPalettes = {
 				"border-color": "#4e4e4e",
 				"tr-even-bg-color": "#222",
 				"tr-odd-bg-color": "#353535",
+				"content-bg": "#4e4e4e",
+				"content-fg": "#fff",
+				"content-hover-bg": "#222",
+				"content-hover-fg": "#fff"
 			}
 		},
 	},
@@ -120,6 +124,10 @@ const colorPalettes = {
 				"border-color": "#888",
 				"tr-even-bg-color": "#f9f9f9",
 				"tr-odd-bg-color": "#fff",
+				"content-bg": "#e0e0e0",
+				"content-fg": "#222",
+				"content-hover-bg": "#adadad",
+				"content-hover-fg": "#222"
 			}
 		},
 	},
@@ -176,6 +184,10 @@ const colorPalettes = {
 				"border-color": "#657b83", // Base00
 				"tr-even-bg-color": "#002b36",
 				"tr-odd-bg-color": "#073642",
+				"content-bg": "#657b83",
+				"content-fg": "#fdf6e3",
+				"content-hover-bg": "#002b36",
+				"content-hover-fg": "#fdf6e3"
 			}
 		},
 	},
@@ -244,7 +256,11 @@ const colorPalettes = {
 			  "error-text": "#ff4444",
 			  "border-color": "#6e7581",
 			  "tr-even-bg-color": "#2b2f38",
-			  "tr-odd-bg-color": "#242730"
+			  "tr-odd-bg-color": "#242730",
+			  "content-bg": "#6e7581",
+			  "content-fg": "#fff",
+			  "content-hover-bg": "#2b2f38",
+			  "content-hover-fg": "#fff"
 			}
 		},
 	},
@@ -313,7 +329,11 @@ const colorPalettes = {
 			  "error-text": "#ff4444",
 			  "border-color": "#545d70",
 			  "tr-even-bg-color": "#2e3440",
-			  "tr-odd-bg-color": "#161b22"
+			  "tr-odd-bg-color": "#161b22",
+			  "content-bg": "#545d70",
+			  "content-fg": "#e5eaf0",
+			  "content-hover-bg": "#2e3440",
+			  "content-hover-fg": "#e5eaf0"
 			}
 		},
 	},
@@ -382,7 +402,11 @@ const colorPalettes = {
 			  "error-text": "#ff4444",
 			  "border-color": "#30363d",
 			  "tr-even-bg-color": "#161b22",
-			  "tr-odd-bg-color": "#13171d"
+			  "tr-odd-bg-color": "#13171d",
+			  "content-bg": "#30363d",
+			  "content-fg": "#e5eaf0",
+			  "content-hover-bg": "#161b22",
+			  "content-hover-fg": "#e5eaf0"
 			}
 		},
 	}

web/extensions/core/undoRedo.js

@@ -1,177 +0,0 @@
-import { app } from "../../scripts/app.js";
-import { api } from "../../scripts/api.js"
-
-const MAX_HISTORY = 50;
-
-let undo = [];
-let redo = [];
-let activeState = null;
-let isOurLoad = false;
-function checkState() {
-	const currentState = app.graph.serialize();
-	if (!graphEqual(activeState, currentState)) {
-		undo.push(activeState);
-		if (undo.length > MAX_HISTORY) {
-			undo.shift();
-		}
-		activeState = clone(currentState);
-		redo.length = 0;
-		api.dispatchEvent(new CustomEvent("graphChanged", { detail: activeState }));
-	}
-}
-
-const loadGraphData = app.loadGraphData;
-app.loadGraphData = async function () {
-	const v = await loadGraphData.apply(this, arguments);
-	if (isOurLoad) {
-		isOurLoad = false;
-	} else {
-		checkState();
-	}
-	return v;
-};
-
-function clone(obj) {
-	try {
-		if (typeof structuredClone !== "undefined") {
-			return structuredClone(obj);
-		}
-	} catch (error) {
-		// structuredClone is stricter than using JSON.parse/stringify so fallback to that
-	}
-
-	return JSON.parse(JSON.stringify(obj));
-}
-
-function graphEqual(a, b, root = true) {
-	if (a === b) return true;
-
-	if (typeof a == "object" && a && typeof b == "object" && b) {
-		const keys = Object.getOwnPropertyNames(a);
-
-		if (keys.length != Object.getOwnPropertyNames(b).length) {
-			return false;
-		}
-
-		for (const key of keys) {
-			let av = a[key];
-			let bv = b[key];
-			if (root && key === "nodes") {
-				// Nodes need to be sorted as the order changes when selecting nodes
-				av = [...av].sort((a, b) => a.id - b.id);
-				bv = [...bv].sort((a, b) => a.id - b.id);
-			}
-			if (!graphEqual(av, bv, false)) {
-				return false;
-			}
-		}
-
-		return true;
-	}
-
-	return false;
-}
-
-const undoRedo = async (e) => {
-	const updateState = async (source, target) => {
-		const prevState = source.pop();
-		if (prevState) {
-			target.push(activeState);
-			isOurLoad = true;
-			await app.loadGraphData(prevState, false);
-			activeState = prevState;
-		}
-	}
-	if (e.ctrlKey || e.metaKey) {
-		if (e.key === "y") {
-			updateState(redo, undo);
-			return true;
-		} else if (e.key === "z") {
-			updateState(undo, redo);
-			return true;
-		}
-	}
-};
-
-const bindInput = (activeEl) => {
-	if (activeEl && activeEl.tagName !== "CANVAS" && activeEl.tagName !== "BODY") {
-		for (const evt of ["change", "input", "blur"]) {
-			if (`on${evt}` in activeEl) {
-				const listener = () => {
-					checkState();
-					activeEl.removeEventListener(evt, listener);
-				};
-				activeEl.addEventListener(evt, listener);
-				return true;
-			}
-		}
-	}
-};
-
-let keyIgnored = false;
-window.addEventListener(
-	"keydown",
-	(e) => {
-		requestAnimationFrame(async () => {
-			let activeEl;
-			// If we are auto queue in change mode then we do want to trigger on inputs
-			if (!app.ui.autoQueueEnabled || app.ui.autoQueueMode === "instant") {
-				activeEl = document.activeElement;
-				if (activeEl?.tagName === "INPUT" || activeEl?.type === "textarea") {
-					// Ignore events on inputs, they have their native history
-					return;
-				}
-			}
-		
-			keyIgnored = e.key === "Control" || e.key === "Shift" || e.key === "Alt" || e.key === "Meta";
-			if (keyIgnored) return;
-
-			// Check if this is a ctrl+z ctrl+y
-			if (await undoRedo(e)) return;
-
-			// If our active element is some type of input then handle changes after they're done
-			if (bindInput(activeEl)) return;
-			checkState();
-		});
-	},
-	true
-);
-
-window.addEventListener("keyup", (e) => {
-	if (keyIgnored) {
-		keyIgnored = false;
-		checkState();
-	}
-});
-
-// Handle clicking DOM elements (e.g. widgets)
-window.addEventListener("mouseup", () => {
-	checkState();
-});
-
-// Handle prompt queue event for dynamic widget changes
-api.addEventListener("promptQueued", () => {
-	checkState();
-});
-
-// Handle litegraph clicks
-const processMouseUp = LGraphCanvas.prototype.processMouseUp;
-LGraphCanvas.prototype.processMouseUp = function (e) {
-	const v = processMouseUp.apply(this, arguments);
-	checkState();
-	return v;
-};
-const processMouseDown = LGraphCanvas.prototype.processMouseDown;
-LGraphCanvas.prototype.processMouseDown = function (e) {
-	const v = processMouseDown.apply(this, arguments);
-	checkState();
-	return v;
-};
-
-// Handle litegraph context menu for COMBO widgets
-const close = LiteGraph.ContextMenu.prototype.close;
-LiteGraph.ContextMenu.prototype.close = function(e) {
-	const v = close.apply(this, arguments);
-	checkState();
-	return v;
-}

web/extensions/core/uploadAudio.js

@@ -0,0 +1,186 @@
+import { app } from "../../scripts/app.js"
+import { api } from "../../scripts/api.js"
+
+function splitFilePath(path) {
+  const folder_separator = path.lastIndexOf("/")
+  if (folder_separator === -1) {
+    return ["", path]
+  }
+  return [
+    path.substring(0, folder_separator),
+    path.substring(folder_separator + 1)
+  ]
+}
+
+function getResourceURL(subfolder, filename, type = "input") {
+  const params = [
+    "filename=" + encodeURIComponent(filename),
+    "type=" + type,
+    "subfolder=" + subfolder,
+    app.getRandParam().substring(1)
+  ].join("&")
+
+  return `/view?${params}`
+}
+
+async function uploadFile(
+  audioWidget,
+  audioUIWidget,
+  file,
+  updateNode,
+  pasted = false
+) {
+  try {
+    // Wrap file in formdata so it includes filename
+    const body = new FormData()
+    body.append("image", file)
+    if (pasted) body.append("subfolder", "pasted")
+    const resp = await api.fetchApi("/upload/image", {
+      method: "POST",
+      body
+    })
+
+    if (resp.status === 200) {
+      const data = await resp.json()
+      // Add the file to the dropdown list and update the widget value
+      let path = data.name
+      if (data.subfolder) path = data.subfolder + "/" + path
+
+      if (!audioWidget.options.values.includes(path)) {
+        audioWidget.options.values.push(path)
+      }
+
+      if (updateNode) {
+        audioUIWidget.element.src = api.apiURL(
+          getResourceURL(...splitFilePath(path))
+        )
+        audioWidget.value = path
+      }
+    } else {
+      alert(resp.status + " - " + resp.statusText)
+    }
+  } catch (error) {
+    alert(error)
+  }
+}
+
+// AudioWidget MUST be registered first, as AUDIOUPLOAD depends on AUDIO_UI to be
+// present.
+app.registerExtension({
+  name: "Comfy.AudioWidget",
+  async beforeRegisterNodeDef(nodeType, nodeData) {
+    if (["LoadAudio", "SaveAudio", "PreviewAudio"].includes(nodeType.comfyClass)) {
+      nodeData.input.required.audioUI = ["AUDIO_UI"]
+    }
+  },
+  getCustomWidgets() {
+    return {
+      AUDIO_UI(node, inputName) {
+        const audio = document.createElement("audio")
+        audio.controls = true
+        audio.classList.add("comfy-audio")
+        audio.setAttribute("name", "media")
+
+        const audioUIWidget = node.addDOMWidget(
+          inputName,
+          /* name=*/ "audioUI",
+          audio
+        )
+        // @ts-ignore
+        // TODO: Sort out the DOMWidget type.
+        audioUIWidget.serialize = false
+
+        const isOutputNode = node.constructor.nodeData.output_node
+        if (isOutputNode) {
+          // Hide the audio widget when there is no audio initially.
+          audioUIWidget.element.classList.add("empty-audio-widget")
+          // Populate the audio widget UI on node execution.
+          const onExecuted = node.onExecuted
+          node.onExecuted = function(message) {
+            onExecuted?.apply(this, arguments)
+            const audios = message.audio
+            if (!audios) return
+            const audio = audios[0]
+            audioUIWidget.element.src = api.apiURL(
+              getResourceURL(audio.subfolder, audio.filename, audio.type)
+            )
+            audioUIWidget.element.classList.remove("empty-audio-widget")
+          }
+        }
+        return { widget: audioUIWidget }
+      }
+    }
+  },
+  onNodeOutputsUpdated(nodeOutputs) {
+    for (const [nodeId, output] of Object.entries(nodeOutputs)) {
+      const node = app.graph.getNodeById(Number.parseInt(nodeId));
+      if ("audio" in output) {
+        const audioUIWidget = node.widgets.find((w) => w.name === "audioUI");
+        const audio = output.audio[0];
+        audioUIWidget.element.src = api.apiURL(getResourceURL(audio.subfolder, audio.filename, audio.type));
+        audioUIWidget.element.classList.remove("empty-audio-widget");
+      }
+    }
+  },
+})
+
+app.registerExtension({
+  name: "Comfy.UploadAudio",
+  async beforeRegisterNodeDef(nodeType, nodeData) {
+    if (nodeData?.input?.required?.audio?.[1]?.audio_upload === true) {
+      nodeData.input.required.upload = ["AUDIOUPLOAD"]
+    }
+  },
+  getCustomWidgets() {
+    return {
+      AUDIOUPLOAD(node, inputName) {
+        // The widget that allows user to select file.
+        const audioWidget = node.widgets.find(w => w.name === "audio")
+        const audioUIWidget = node.widgets.find(w => w.name === "audioUI")
+
+        const onAudioWidgetUpdate = () => {
+          audioUIWidget.element.src = api.apiURL(
+            getResourceURL(...splitFilePath(audioWidget.value))
+          )
+        }
+        // Initially load default audio file to audioUIWidget.
+        if (audioWidget.value) {
+          onAudioWidgetUpdate()
+        }
+        audioWidget.callback = onAudioWidgetUpdate
+
+        // Load saved audio file widget values if restoring from workflow
+        const onGraphConfigured = node.onGraphConfigured;
+        node.onGraphConfigured = function() {
+          onGraphConfigured?.apply(this, arguments)
+          if (audioWidget.value) {
+            onAudioWidgetUpdate()
+          }
+        }
+
+        const fileInput = document.createElement("input")
+        fileInput.type = "file"
+        fileInput.accept = "audio/*"
+        fileInput.style.display = "none"
+        fileInput.onchange = () => {
+          if (fileInput.files.length) {
+            uploadFile(audioWidget, audioUIWidget, fileInput.files[0], true)
+          }
+        }
+        // The widget to pop up the upload dialog.
+        const uploadWidget = node.addWidget(
+          "button",
+          inputName,
+          /* value=*/ "",
+          () => {
+            fileInput.click()
+          }
+        )
+        uploadWidget.label = "choose file to upload"
+        uploadWidget.serialize = false
+
+        return { widget: uploadWidget }
+      }
+    }
+  }
+})

web/index.html

@@ -5,6 +5,7 @@
 		<title>ComfyUI</title>
 		<meta name="viewport" content="width=device-width, initial-scale=1.0, user-scalable=no">
 		<link rel="stylesheet" type="text/css" href="./lib/litegraph.css" />
+		<link rel="stylesheet" type="text/css" href="./lib/materialdesignicons.min.css" />
 		<link rel="stylesheet" type="text/css" href="./style.css" />
 		<link rel="stylesheet" type="text/css" href="./user.css" />
 		<script type="text/javascript" src="./lib/litegraph.core.js"></script>

web/jsconfig.json

@@ -4,7 +4,9 @@
 		"paths": {
 			"/*": ["./*"]
 		},
-		"lib": ["DOM", "ES2022"]
+		"lib": ["DOM", "ES2022", "DOM.Iterable"],
+		"target": "ES2015",
+		"module": "es2020"
 	},
 	"include": ["."]
 }

web/scripts/api.js

@@ -394,7 +394,7 @@ class ComfyApi extends EventTarget {
 	 * Gets a user data file for the current user
 	 * @param { string } file The name of the userdata file to load
 	 * @param { RequestInit } [options]
-	 * @returns { Promise<unknown> } The fetch response object
+	 * @returns { Promise<Response> } The fetch response object
 	 */
 	async getUserData(file, options) {
 		return this.fetchApi(`/userdata/${encodeURIComponent(file)}`, options);
@@ -404,18 +404,75 @@ class ComfyApi extends EventTarget {
 	 * Stores a user data file for the current user
 	 * @param { string } file The name of the userdata file to save
 	 * @param { unknown } data The data to save to the file
-	 * @param { RequestInit & { stringify?: boolean, throwOnError?: boolean } } [options]
+	 * @param { RequestInit & { overwrite?: boolean, stringify?: boolean, throwOnError?: boolean } } [options]
-	 * @returns { Promise<void> }
+	 * @returns { Promise<Response> }
 	 */
-	async storeUserData(file, data, options = { stringify: true, throwOnError: true }) {
+	async storeUserData(file, data, options = { overwrite: true, stringify: true, throwOnError: true }) {
-		const resp = await this.fetchApi(`/userdata/${encodeURIComponent(file)}`, {
+		const resp = await this.fetchApi(`/userdata/${encodeURIComponent(file)}?overwrite=${options?.overwrite}`, {
 			method: "POST",
 			body: options?.stringify ? JSON.stringify(data) : data,
 			...options,
 		});
-		if (resp.status !== 200) {
+		if (resp.status !== 200 && options?.throwOnError !== false) {
 			throw new Error(`Error storing user data file '${file}': ${resp.status} ${(await resp).statusText}`);
 		}
+		return resp;
+	}
+
+	/**
+	 * Deletes a user data file for the current user
+	 * @param { string } file The name of the userdata file to delete
+	 */
+	async deleteUserData(file) {
+		const resp = await this.fetchApi(`/userdata/${encodeURIComponent(file)}`, {
+			method: "DELETE",
+		});
+		if (resp.status !== 204) {
+			throw new Error(`Error removing user data file '${file}': ${resp.status} ${(resp).statusText}`);
+		}
+	}
+
+	/**
+	 * Move a user data file for the current user
+	 * @param { string } source The userdata file to move
+	 * @param { string } dest The destination for the file
+	 */
+	async moveUserData(source, dest, options = { overwrite: false }) {
+		const resp = await this.fetchApi(`/userdata/${encodeURIComponent(source)}/move/${encodeURIComponent(dest)}?overwrite=${options?.overwrite}`, {
+			method: "POST",
+		});
+		return resp;
+	}
+
+	/**
+	 * @overload
+	 * Lists user data files for the current user
+	 * @param { string } dir The directory in which to list files
+	 * @param { boolean } [recurse] If the listing should be recursive
+	 * @param { true } [split] If the paths should be split based on the os path separator
+	 * @returns { Promise<string[][]>> } The list of split file paths in the format [fullPath, ...splitPath]
+	 */
+	/**
+	 * @overload
+	 * Lists user data files for the current user
+	 * @param { string } dir The directory in which to list files
+	 * @param { boolean } [recurse] If the listing should be recursive
+	 * @param { false | undefined } [split] If the paths should be split based on the os path separator
+	 * @returns { Promise<string[]>> } The list of files
+	 */
+	async listUserData(dir, recurse, split) {
+		const resp = await this.fetchApi(
+			`/userdata?${new URLSearchParams({
+				recurse,
+				dir,
+				split,
+			})}`
+		);
+		if (resp.status === 404) return [];
+		if (resp.status !== 200) {
+			throw new Error(`Error getting user data list '${dir}': ${resp.status} ${resp.statusText}`);
+		}
+		return resp.json();
 	}
 }
 

web/scripts/app.js

@@ -3,11 +3,13 @@ import { ComfyWidgets, initWidgets } from "./widgets.js";
 import { ComfyUI, $el } from "./ui.js";
 import { api } from "./api.js";
 import { defaultGraph } from "./defaultGraph.js";
-import { getPngMetadata, getWebpMetadata, importA1111, getLatentMetadata } from "./pnginfo.js";
+import { getPngMetadata, getWebpMetadata, getFlacMetadata, importA1111, getLatentMetadata } from "./pnginfo.js";
 import { addDomClippingSetting } from "./domWidget.js";
-import { createImageHost, calculateImageGrid } from "./ui/imagePreview.js"
+import { createImageHost, calculateImageGrid } from "./ui/imagePreview.js";
-
+import { ComfyAppMenu } from "./ui/menu/index.js";
-export const ANIM_PREVIEW_WIDGET = "$$comfy_animation_preview"
+import { getStorageValue, setStorageValue } from "./utils.js";
+import { ComfyWorkflowManager } from "./workflows.js";
+export const ANIM_PREVIEW_WIDGET = "$$comfy_animation_preview";
 
 function sanitizeNodeName(string) {
 	let entityMap = {
@@ -52,6 +54,12 @@ export class ComfyApp {
 	constructor() {
 		this.ui = new ComfyUI(this);
 		this.logging = new ComfyLogging(this);
+		this.workflowManager = new ComfyWorkflowManager(this);
+		this.bodyTop = $el("div.comfyui-body-top", { parent: document.body });
+		this.bodyLeft = $el("div.comfyui-body-left", { parent: document.body });
+		this.bodyRight = $el("div.comfyui-body-right", { parent: document.body });
+		this.bodyBottom = $el("div.comfyui-body-bottom", { parent: document.body });		  
+		this.menu = new ComfyAppMenu(this);
 
 		/**
 		 * List of extensions that are registered with the app
@@ -63,7 +71,7 @@ export class ComfyApp {
 		 * Stores the execution output data for each node
 		 * @type {Record<string, any>}
 		 */
-		this.nodeOutputs = {};
+		this._nodeOutputs = {};
 
 		/**
 		 * Stores the preview image data for each node
@@ -78,6 +86,15 @@ export class ComfyApp {
 		this.shiftDown = false;
 	}
 
+	get nodeOutputs() {
+		return this._nodeOutputs;
+	}
+
+	set nodeOutputs(value) {
+		this._nodeOutputs = value;
+		this.#invokeExtensions("onNodeOutputsUpdated", value);
+	}
+
 	getPreviewFormatParam() {
 		let preview_format = this.ui.settings.getSettingValue("Comfy.PreviewFormat");
 		if(preview_format)
@@ -1067,7 +1084,7 @@ export class ComfyApp {
 			if (e.type == "keydown" && !e.repeat) {
 
 				// Ctrl + M mute/unmute
-				if (e.key === 'm' && e.ctrlKey) {
+				if (e.key === 'm' && (e.metaKey || e.ctrlKey)) {
 					if (this.selected_nodes) {
 						for (var i in this.selected_nodes) {
 							if (this.selected_nodes[i].mode === 2) { // never
@@ -1081,7 +1098,7 @@ export class ComfyApp {
 				}
 
 				// Ctrl + B bypass
-				if (e.key === 'b' && e.ctrlKey) {
+				if (e.key === 'b' && (e.metaKey || e.ctrlKey)) {
 					if (this.selected_nodes) {
 						for (var i in this.selected_nodes) {
 							if (this.selected_nodes[i].mode === 4) { // never
@@ -1313,11 +1330,15 @@ export class ComfyApp {
 		});
 
 		api.addEventListener("progress", ({ detail }) => {
+			if (this.workflowManager.activePrompt?.workflow 
+				&& this.workflowManager.activePrompt.workflow !== this.workflowManager.activeWorkflow) return;
 			this.progress = detail;
 			this.graph.setDirtyCanvas(true, false);
 		});
 
 		api.addEventListener("executing", ({ detail }) => {
+			if (this.workflowManager.activePrompt ?.workflow
+				&& this.workflowManager.activePrompt.workflow !== this.workflowManager.activeWorkflow) return;
 			this.progress = null;
 			this.runningNodeId = detail;
 			this.graph.setDirtyCanvas(true, false);
@@ -1325,6 +1346,8 @@ export class ComfyApp {
 		});
 
 		api.addEventListener("executed", ({ detail }) => {
+			if (this.workflowManager.activePrompt ?.workflow
+				&& this.workflowManager.activePrompt.workflow !== this.workflowManager.activeWorkflow) return;
 			const output = this.nodeOutputs[detail.display_node];
 			if (detail.merge && output) {
 				for (const k in detail.output ?? {}) {
@@ -1433,6 +1456,11 @@ export class ComfyApp {
 	    });
 	
 	    await Promise.all(extensionPromises);
+		try {
+			this.menu.workflows.registerExtension(this);
+		} catch (error) {
+			console.error(error);
+		}
 	}
 
 	async #migrateSettings() {
@@ -1520,15 +1548,17 @@ export class ComfyApp {
 	 */
 	async setup() {
 		await this.#setUser();
-		await this.ui.settings.load();
-		await this.#loadExtensions();
 
 		// Create and mount the LiteGraph in the DOM
 		const mainCanvas = document.createElement("canvas")
 		mainCanvas.style.touchAction = "none"
 		const canvasEl = (this.canvasEl = Object.assign(mainCanvas, { id: "graph-canvas" }));
 		canvasEl.tabIndex = "1";
-		document.body.prepend(canvasEl);
+		document.body.append(canvasEl);
+		this.resizeCanvas();
+
+		await Promise.all([this.workflowManager.loadWorkflows(), this.ui.settings.load()]);
+		await this.#loadExtensions();
 
 		addDomClippingSetting();
 		this.#addProcessMouseHandler();
@@ -1541,7 +1571,7 @@ export class ComfyApp {
 
 		this.#addAfterConfigureHandler();
 
-		const canvas = (this.canvas = new LGraphCanvas(canvasEl, this.graph));
+		this.canvas = new LGraphCanvas(canvasEl, this.graph);
 		this.ctx = canvasEl.getContext("2d");
 
 		LiteGraph.release_link_on_empty_shows_menu = true;
@@ -1549,19 +1579,14 @@ export class ComfyApp {
 
 		this.graph.start();
 
-		function resizeCanvas() {
-			// Limit minimal scale to 1, see https://github.com/comfyanonymous/ComfyUI/pull/845
-			const scale = Math.max(window.devicePixelRatio, 1);
-			const { width, height } = canvasEl.getBoundingClientRect();
-			canvasEl.width = Math.round(width * scale);
-			canvasEl.height = Math.round(height * scale);
-			canvasEl.getContext("2d").scale(scale, scale);
-			canvas.draw(true, true);
-		}
-
 		// Ensure the canvas fills the window
-		resizeCanvas();
+		this.resizeCanvas();
-		window.addEventListener("resize", resizeCanvas);
+		window.addEventListener("resize", () => this.resizeCanvas());
+		const ro = new ResizeObserver(() => this.resizeCanvas());
+		ro.observe(this.bodyTop);
+		ro.observe(this.bodyLeft);
+		ro.observe(this.bodyRight);
+		ro.observe(this.bodyBottom);
 
 		await this.#invokeExtensionsAsync("init");
 		await this.registerNodes();
@@ -1573,7 +1598,8 @@ export class ComfyApp {
 			const loadWorkflow = async (json) => {
 				if (json) {
 					const workflow = JSON.parse(json);
-					await this.loadGraphData(workflow);
+					const workflowName = getStorageValue("Comfy.PreviousWorkflow");
+					await this.loadGraphData(workflow, true, true, workflowName);
 					return true;
 				}
 			};
@@ -1609,6 +1635,19 @@ export class ComfyApp {
 		await this.#invokeExtensionsAsync("setup");
 	}
 
+	resizeCanvas() {
+		// Limit minimal scale to 1, see https://github.com/comfyanonymous/ComfyUI/pull/845
+		const scale = Math.max(window.devicePixelRatio, 1);
+		
+		// Clear fixed width and height while calculating rect so it uses 100% instead
+		this.canvasEl.height = this.canvasEl.width = "";
+		const { width, height } = this.canvasEl.getBoundingClientRect();
+		this.canvasEl.width = Math.round(width * scale);
+		this.canvasEl.height = Math.round(height * scale);
+		this.canvasEl.getContext("2d").scale(scale, scale);
+		this.canvas?.draw(true, true);
+	}
+
 	/**
 	 * Registers nodes with the graph
 	 */
@@ -1795,12 +1834,29 @@ export class ComfyApp {
 		});
 	}
 
+	async changeWorkflow(callback, workflow = null) {
+		try {
+			this.workflowManager.activeWorkflow?.changeTracker?.store()
+		} catch (error) {
+			console.error(error);
+		}
+		await callback();
+		try {
+			this.workflowManager.setWorkflow(workflow);
+			this.workflowManager.activeWorkflow?.track()
+		} catch (error) {
+			console.error(error);
+		}
+	}
+
 	/**
 	 * Populates the graph with the specified workflow data
 	 * @param {*} graphData A serialized graph object
 	 * @param { boolean } clean If the graph state, e.g. images, should be cleared
+	 * @param { boolean } restore_view If the graph position should be restored
+	 * @param { import("./workflows.js").ComfyWorkflowInstance | null } workflow The workflow
 	 */
-	async loadGraphData(graphData, clean = true, restore_view = true) {
+	async loadGraphData(graphData, clean = true, restore_view = true, workflow = null) {
 		if (clean !== false) {
 			this.clean();
 		}
@@ -1819,6 +1875,12 @@ export class ComfyApp {
 			graphData = structuredClone(graphData);
 		}
 	
+		try {
+			this.workflowManager.setWorkflow(workflow);
+		} catch (error) {
+			console.error(error);
+		}
+
 		const missingNodeTypes = [];
 		await this.#invokeExtensionsAsync("beforeConfigureGraph", graphData, missingNodeTypes);
 		for (let n of graphData.nodes) {
@@ -1840,6 +1902,11 @@ export class ComfyApp {
 				this.canvas.ds.offset = graphData.extra.ds.offset;
 				this.canvas.ds.scale = graphData.extra.ds.scale;
 			}
+			
+			try {
+				this.workflowManager.activeWorkflow?.track()
+			} catch (error) {
+			}
 		} catch (error) {
 			let errorHint = [];
 			// Try extracting filename to see if it was caused by an extension script
@@ -1899,6 +1966,14 @@ export class ComfyApp {
 							if (widget.value.startsWith("sample_")) {
 								widget.value = widget.value.slice(7);
 							}
+							if (widget.value === "euler_pp" || widget.value === "euler_ancestral_pp") {
+								widget.value = widget.value.slice(0, -3);
+								for (let w of node.widgets) {
+									if (w.name == "cfg") {
+										w.value *= 2.0;
+									}
+								}
+							}
 						}
 					}
 					if (node.type == "KSampler" || node.type == "KSamplerAdvanced" || node.type == "PrimitiveNode") {
@@ -1927,14 +2002,17 @@ export class ComfyApp {
 			this.showMissingNodesError(missingNodeTypes);
 		}
 		await this.#invokeExtensionsAsync("afterConfigureGraph", missingNodeTypes);
+		requestAnimationFrame(() => {
+			this.graph.setDirtyCanvas(true, true);
+		});
 	}
 
 	/**
 	 * Converts the current graph workflow for sending to the API
 	 * @returns The workflow and node links
 	 */
-	async graphToPrompt() {
+	async graphToPrompt(graph = this.graph, clean = true) {
-		for (const outerNode of this.graph.computeExecutionOrder(false)) {
+		for (const outerNode of graph.computeExecutionOrder(false)) {
 			if (outerNode.widgets) {
 				for (const widget of outerNode.widgets) {
 					// Allow widgets to run callbacks before a prompt has been queued
@@ -1954,10 +2032,10 @@ export class ComfyApp {
 			}
 		}
 
-		const workflow = this.graph.serialize();
+		const workflow = graph.serialize();
 		const output = {};
 		// Process nodes in order of execution
-		for (const outerNode of this.graph.computeExecutionOrder(false)) {
+		for (const outerNode of graph.computeExecutionOrder(false)) {
 			const skipNode = outerNode.mode === 2 || outerNode.mode === 4;
 			const innerNodes = (!skipNode && outerNode.getInnerNodes) ? outerNode.getInnerNodes() : [outerNode];
 			for (const node of innerNodes) {
@@ -2049,13 +2127,14 @@ export class ComfyApp {
 		}
 
 		// Remove inputs connected to removed nodes
-
+		if(clean) {
-		for (const o in output) {
+			for (const o in output) {
-			for (const i in output[o].inputs) {
+				for (const i in output[o].inputs) {
-				if (Array.isArray(output[o].inputs[i])
+					if (Array.isArray(output[o].inputs[i])
-					&& output[o].inputs[i].length === 2
+						&& output[o].inputs[i].length === 2
-					&& !output[output[o].inputs[i][0]]) {
+						&& !output[output[o].inputs[i][0]]) {
-					delete output[o].inputs[i];
+						delete output[o].inputs[i];
+					}
 				}
 			}
 		}
@@ -2123,6 +2202,14 @@ export class ComfyApp {
 						this.lastNodeErrors = res.node_errors;
 						if (this.lastNodeErrors.length > 0) {
 							this.canvas.draw(true, true);
+						} else {
+							try {
+								this.workflowManager.storePrompt({
+									id: res.prompt_id,
+									nodes: Object.keys(p.output)
+								});
+							} catch (error) {
+							}
 						}
 					} catch (error) {
 						const formattedError = this.#formatPromptError(error)
@@ -2155,6 +2242,7 @@ export class ComfyApp {
 			this.#processingQueue = false;
 		}
 		api.dispatchEvent(new CustomEvent("promptQueued", { detail: { number, batchCount } }));
+		return !this.lastNodeErrors;
 	}
 
 	showErrorOnFileLoad(file) {
@@ -2170,14 +2258,24 @@ export class ComfyApp {
 	 * @param {File} file
 	 */
 	async handleFile(file) {
+		const removeExt = f => {
+			if(!f) return f;
+			const p = f.lastIndexOf(".");
+			if(p === -1) return f;
+			return f.substring(0, p);
+		};
+
+		const fileName = removeExt(file.name);
 		if (file.type === "image/png") {
 			const pngInfo = await getPngMetadata(file);
 			if (pngInfo?.workflow) {
-				await this.loadGraphData(JSON.parse(pngInfo.workflow));
+				await this.loadGraphData(JSON.parse(pngInfo.workflow), true, true, fileName);
 			} else if (pngInfo?.prompt) {
-				this.loadApiJson(JSON.parse(pngInfo.prompt));
+				this.loadApiJson(JSON.parse(pngInfo.prompt), fileName);
 			} else if (pngInfo?.parameters) {
-				importA1111(this.graph, pngInfo.parameters);
+				this.changeWorkflow(() => {
+					importA1111(this.graph, pngInfo.parameters);
+				}, fileName)
 			} else {
 				this.showErrorOnFileLoad(file);
 			}
@@ -2188,9 +2286,22 @@ export class ComfyApp {
 			const prompt = pngInfo?.prompt || pngInfo?.Prompt;
 
 			if (workflow) {
-				this.loadGraphData(JSON.parse(workflow));
+				this.loadGraphData(JSON.parse(workflow), true, true, fileName);
 			} else if (prompt) {
-				this.loadApiJson(JSON.parse(prompt));
+				this.loadApiJson(JSON.parse(prompt), fileName);
+			} else {
+				this.showErrorOnFileLoad(file);
+			}
+		} else if (file.type === "audio/flac" || file.type === "audio/x-flac") {
+			const pngInfo = await getFlacMetadata(file);
+			// Support loading workflows from that webp custom node.
+			const workflow = pngInfo?.workflow;
+			const prompt = pngInfo?.prompt;
+
+			if (workflow) {
+				this.loadGraphData(JSON.parse(workflow), true, true, fileName);
+			} else if (prompt) {
+				this.loadApiJson(JSON.parse(prompt), fileName);
 			} else {
 				this.showErrorOnFileLoad(file);
 			}
@@ -2201,16 +2312,16 @@ export class ComfyApp {
 				if (jsonContent?.templates) {
 					this.loadTemplateData(jsonContent);
 				} else if(this.isApiJson(jsonContent)) {
-					this.loadApiJson(jsonContent);
+					this.loadApiJson(jsonContent, fileName);
 				} else {
-					await this.loadGraphData(jsonContent);
+					await this.loadGraphData(jsonContent, true, true, fileName);
 				}
 			};
 			reader.readAsText(file);
 		} else if (file.name?.endsWith(".latent") || file.name?.endsWith(".safetensors")) {
 			const info = await getLatentMetadata(file);
 			if (info.workflow) {
-				await this.loadGraphData(JSON.parse(info.workflow));
+				await this.loadGraphData(JSON.parse(info.workflow), true, true, fileName);
 			} else if (info.prompt) {
 				this.loadApiJson(JSON.parse(info.prompt));
 			} else {
@@ -2225,7 +2336,7 @@ export class ComfyApp {
 		return Object.values(data).every((v) => v.class_type);
 	}
 
-	loadApiJson(apiData) {
+	loadApiJson(apiData, fileName) {
 		const missingNodeTypes = Object.values(apiData).filter((n) => !LiteGraph.registered_node_types[n.class_type]);
 		if (missingNodeTypes.length) {
 			this.showMissingNodesError(missingNodeTypes.map(t => t.class_type), false);
@@ -2239,41 +2350,42 @@ export class ComfyApp {
 			const node = LiteGraph.createNode(data.class_type);
 			node.id = isNaN(+id) ? id : +id;
 			node.title = data._meta?.title ?? node.title
-			graph.add(node);
+			app.graph.add(node);
 		}
 
-		for (const id of ids) {
+		this.changeWorkflow(() => {
-			const data = apiData[id];
+			for (const id of ids) {
-			const node = app.graph.getNodeById(id);
+				const data = apiData[id];
-			for (const input in data.inputs ?? {}) {
+				const node = app.graph.getNodeById(id);
-				const value = data.inputs[input];
+				for (const input in data.inputs ?? {}) {
-				if (value instanceof Array) {
+					const value = data.inputs[input];
-					const [fromId, fromSlot] = value;
+					if (value instanceof Array) {
-					const fromNode = app.graph.getNodeById(fromId);
+						const [fromId, fromSlot] = value;
-					let toSlot = node.inputs?.findIndex((inp) => inp.name === input);
+						const fromNode = app.graph.getNodeById(fromId);
-					if (toSlot == null || toSlot === -1) {
+						let toSlot = node.inputs?.findIndex((inp) => inp.name === input);
-						try {
+						if (toSlot == null || toSlot === -1) {
-							// Target has no matching input, most likely a converted widget
+							try {
-							const widget = node.widgets?.find((w) => w.name === input);
+								// Target has no matching input, most likely a converted widget
-							if (widget && node.convertWidgetToInput?.(widget)) {
+								const widget = node.widgets?.find((w) => w.name === input);
-								toSlot = node.inputs?.length - 1;
+								if (widget && node.convertWidgetToInput?.(widget)) {
-							}
+									toSlot = node.inputs?.length - 1;
-						} catch (error) {}
+								}
-					}
+							} catch (error) {}
-					if (toSlot != null || toSlot !== -1) {
+						}
-						fromNode.connect(fromSlot, node, toSlot);
+						if (toSlot != null || toSlot !== -1) {
-					}
+							fromNode.connect(fromSlot, node, toSlot);
-				} else {
+						}
-					const widget = node.widgets?.find((w) => w.name === input);
+					} else {
-					if (widget) {
+						const widget = node.widgets?.find((w) => w.name === input);
-						widget.value = value;
+						if (widget) {
-						widget.callback?.(value);
+							widget.value = value;
+							widget.callback?.(value);
+						}
 					}
 				}
 			}
-		}
+			app.graph.arrange();
-
+		}, fileName);
-		app.graph.arrange();
 	}
 
 	/**

web/scripts/changeTracker.js

@@ -0,0 +1,254 @@
+// @ts-check
+
+import { api } from "./api.js";
+import { clone } from "./utils.js";
+
+export class ChangeTracker {
+	static MAX_HISTORY = 50;
+	#app;
+	undo = [];
+	redo = [];
+	activeState = null;
+	isOurLoad = false;
+	/** @type { import("./workflows").ComfyWorkflow | null } */
+	workflow;
+
+	ds;
+	nodeOutputs;
+
+	get app() {
+		return this.#app ?? this.workflow.manager.app;
+	}
+
+	constructor(workflow) {
+		this.workflow = workflow;
+	}
+
+	#setApp(app) {
+		this.#app = app;
+	}
+
+	store() {
+		this.ds = { scale: this.app.canvas.ds.scale, offset: [...this.app.canvas.ds.offset] };
+	}
+
+	restore() {
+		if (this.ds) {
+			this.app.canvas.ds.scale = this.ds.scale;
+			this.app.canvas.ds.offset = this.ds.offset;
+		}
+		if (this.nodeOutputs) {
+			this.app.nodeOutputs = this.nodeOutputs;
+		}
+	}
+
+	checkState() {
+		if (!this.app.graph) return;
+
+		const currentState = this.app.graph.serialize();
+		if (!this.activeState) {
+			this.activeState = clone(currentState);
+			return;
+		}
+		if (!ChangeTracker.graphEqual(this.activeState, currentState)) {
+			this.undo.push(this.activeState);
+			if (this.undo.length > ChangeTracker.MAX_HISTORY) {
+				this.undo.shift();
+			}
+			this.activeState = clone(currentState);
+			this.redo.length = 0;
+			this.workflow.unsaved = true;
+			api.dispatchEvent(new CustomEvent("graphChanged", { detail: this.activeState }));
+		}
+	}
+
+	async updateState(source, target) {
+		const prevState = source.pop();
+		if (prevState) {
+			target.push(this.activeState);
+			this.isOurLoad = true;
+			await this.app.loadGraphData(prevState, false, false, this.workflow);
+			this.activeState = prevState;
+		}
+	}
+
+	async undoRedo(e) {
+		if (e.ctrlKey || e.metaKey) {
+			if (e.key === "y") {
+				this.updateState(this.redo, this.undo);
+				return true;
+			} else if (e.key === "z") {
+				this.updateState(this.undo, this.redo);
+				return true;
+			}
+		}
+	}
+
+	/** @param { import("./app.js").ComfyApp } app */
+	static init(app) {
+		const changeTracker = () => app.workflowManager.activeWorkflow?.changeTracker ?? globalTracker;
+		globalTracker.#setApp(app);
+
+		const loadGraphData = app.loadGraphData;
+		app.loadGraphData = async function () {
+			const v = await loadGraphData.apply(this, arguments);
+			const ct = changeTracker();
+			if (ct.isOurLoad) {
+				ct.isOurLoad = false;
+			} else {
+				ct.checkState();
+			}
+			return v;
+		};
+
+		let keyIgnored = false;
+		window.addEventListener(
+			"keydown",
+			(e) => {
+				requestAnimationFrame(async () => {
+					let activeEl;
+					// If we are auto queue in change mode then we do want to trigger on inputs
+					if (!app.ui.autoQueueEnabled || app.ui.autoQueueMode === "instant") {
+						activeEl = document.activeElement;
+						if (activeEl?.tagName === "INPUT" || activeEl?.["type"] === "textarea") {
+							// Ignore events on inputs, they have their native history
+							return;
+						}
+					}
+
+					keyIgnored = e.key === "Control" || e.key === "Shift" || e.key === "Alt" || e.key === "Meta";
+					if (keyIgnored) return;
+
+					// Check if this is a ctrl+z ctrl+y
+					if (await changeTracker().undoRedo(e)) return;
+
+					// If our active element is some type of input then handle changes after they're done
+					if (ChangeTracker.bindInput(activeEl)) return;
+					changeTracker().checkState();
+				});
+			},
+			true
+		);
+
+		window.addEventListener("keyup", (e) => {
+			if (keyIgnored) {
+				keyIgnored = false;
+				changeTracker().checkState();
+			}
+		});
+
+		// Handle clicking DOM elements (e.g. widgets)
+		window.addEventListener("mouseup", () => {
+			changeTracker().checkState();
+		});
+
+		// Handle prompt queue event for dynamic widget changes
+		api.addEventListener("promptQueued", () => {
+			changeTracker().checkState();
+		});
+
+		// Handle litegraph clicks
+		const processMouseUp = LGraphCanvas.prototype.processMouseUp;
+		LGraphCanvas.prototype.processMouseUp = function (e) {
+			const v = processMouseUp.apply(this, arguments);
+			changeTracker().checkState();
+			return v;
+		};
+		const processMouseDown = LGraphCanvas.prototype.processMouseDown;
+		LGraphCanvas.prototype.processMouseDown = function (e) {
+			const v = processMouseDown.apply(this, arguments);
+			changeTracker().checkState();
+			return v;
+		};
+
+		// Handle litegraph context menu for COMBO widgets
+		const close = LiteGraph.ContextMenu.prototype.close;
+		LiteGraph.ContextMenu.prototype.close = function (e) {
+			const v = close.apply(this, arguments);
+			changeTracker().checkState();
+			return v;
+		};
+
+		// Detects nodes being added via the node search dialog
+		const onNodeAdded = LiteGraph.LGraph.prototype.onNodeAdded;
+		LiteGraph.LGraph.prototype.onNodeAdded = function () {
+			const v = onNodeAdded?.apply(this, arguments);
+			if (!app?.configuringGraph) {
+				const ct = changeTracker();
+				if (!ct.isOurLoad) {
+					ct.checkState();
+				}
+			}
+			return v;
+		};
+
+		// Store node outputs
+		api.addEventListener("executed", ({ detail }) => {
+			const prompt = app.workflowManager.queuedPrompts[detail.prompt_id];
+			if (!prompt?.workflow) return;
+			const nodeOutputs = (prompt.workflow.changeTracker.nodeOutputs ??= {});
+			const output = nodeOutputs[detail.node];
+			if (detail.merge && output) {
+				for (const k in detail.output ?? {}) {
+					const v = output[k];
+					if (v instanceof Array) {
+						output[k] = v.concat(detail.output[k]);
+					} else {
+						output[k] = detail.output[k];
+					}
+				}
+			} else {
+				nodeOutputs[detail.node] = detail.output;
+			}
+		});
+	}
+
+	static bindInput(app, activeEl) {
+		if (activeEl && activeEl.tagName !== "CANVAS" && activeEl.tagName !== "BODY") {
+			for (const evt of ["change", "input", "blur"]) {
+				if (`on${evt}` in activeEl) {
+					const listener = () => {
+						app.workflowManager.activeWorkflow.changeTracker.checkState();
+						activeEl.removeEventListener(evt, listener);
+					};
+					activeEl.addEventListener(evt, listener);
+					return true;
+				}
+			}
+		}
+	}
+
+	static graphEqual(a, b, path = "") {
+		if (a === b) return true;
+
+		if (typeof a == "object" && a && typeof b == "object" && b) {
+			const keys = Object.getOwnPropertyNames(a);
+
+			if (keys.length != Object.getOwnPropertyNames(b).length) {
+				return false;
+			}
+
+			for (const key of keys) {
+				let av = a[key];
+				let bv = b[key];
+				if (!path && key === "nodes") {
+					// Nodes need to be sorted as the order changes when selecting nodes
+					av = [...av].sort((a, b) => a.id - b.id);
+					bv = [...bv].sort((a, b) => a.id - b.id);
+				} else if (path === "extra.ds") {
+					// Ignore view changes
+					continue;
+				}
+				if (!ChangeTracker.graphEqual(av, bv, path + (path ? "." : "") + key)) {
+					return false;
+				}
+			}
+
+			return true;
+		}
+
+		return false;
+	}
+}
+
+const globalTracker = new ChangeTracker({});

web/scripts/domWidget.js

@@ -34,8 +34,8 @@ function getClipPath(node, element) {
 		}
 
 		const widgetRect = element.getBoundingClientRect();
-		const clipX = intersection[0] - widgetRect.x / scale + "px";
+		const clipX = elRect.left + intersection[0] - widgetRect.x / scale + "px";
-		const clipY = intersection[1] - widgetRect.y / scale + "px";
+		const clipY = elRect.top + intersection[1] - widgetRect.y / scale + "px";
 		const clipWidth = intersection[2] + "px";
 		const clipHeight = intersection[3] + "px";
 		const path = `polygon(0% 0%, 0% 100%, ${clipX} 100%, ${clipX} ${clipY}, calc(${clipX} + ${clipWidth}) ${clipY}, calc(${clipX} + ${clipWidth}) calc(${clipY} + ${clipHeight}), ${clipX} calc(${clipY} + ${clipHeight}), ${clipX} 100%, 100% 100%, 100% 0%)`;
@@ -210,6 +210,8 @@ LGraphNode.prototype.addDOMWidget = function (name, type, element, options) {
 	if (!element.parentElement) {
 		document.body.append(element);
 	}
+	element.hidden = true;
+	element.style.display = "none";
 	
 	let mouseDownHandler;
 	if (element.blur) {
@@ -254,15 +256,15 @@ LGraphNode.prototype.addDOMWidget = function (name, type, element, options) {
 			const transform = new DOMMatrix()
 				.scaleSelf(elRect.width / ctx.canvas.width, elRect.height / ctx.canvas.height)
 				.multiplySelf(ctx.getTransform())
-				.translateSelf(margin, margin + y);
+				.translateSelf(margin, margin + y );
 
 			const scale = new DOMMatrix().scaleSelf(transform.a, transform.d);
 
 			Object.assign(element.style, {
 				transformOrigin: "0 0",
 				transform: scale,
-				left: `${transform.a + transform.e}px`,
+				left: `${transform.a + transform.e + elRect.left}px`,
-				top: `${transform.d + transform.f}px`,
+				top: `${transform.d + transform.f + elRect.top}px`,
 				width: `${widgetWidth - margin * 2}px`,
 				height: `${(widget.computedHeight ?? 50) - margin * 2}px`,
 				position: "absolute",

web/scripts/pnginfo.js

@@ -49,7 +49,7 @@ export function getPngMetadata(file) {
 
 function parseExifData(exifData) {
 	// Check for the correct TIFF header (0x4949 for little-endian or 0x4D4D for big-endian)
-	const isLittleEndian = new Uint16Array(exifData.slice(0, 2))[0] === 0x4949;
+	const isLittleEndian = String.fromCharCode(...exifData.slice(0, 2)) === "II";
 
 	// Function to read 16-bit and 32-bit integers from binary data
 	function readInt(offset, isLittleEndian, length) {
@@ -134,6 +134,7 @@ export function getWebpMetadata(file) {
 						let index = value.indexOf(':');
 						txt_chunks[value.slice(0, index)] = value.slice(index + 1);
 					}
+					break;
 				}
 
 				offset += 8 + chunk_length;
@@ -163,6 +164,78 @@ export function getLatentMetadata(file) {
 	});
 }
 
+
+function getString(dataView, offset, length) {
+	let string = '';
+	for (let i = 0; i < length; i++) {
+		string += String.fromCharCode(dataView.getUint8(offset + i));
+	}
+	return string;
+}
+
+// Function to parse the Vorbis Comment block
+function parseVorbisComment(dataView) {
+	let offset = 0;
+	const vendorLength = dataView.getUint32(offset, true);
+	offset += 4;
+	const vendorString = getString(dataView, offset, vendorLength);
+	offset += vendorLength;
+
+	const userCommentListLength = dataView.getUint32(offset, true);
+	offset += 4;
+	const comments = {};
+	for (let i = 0; i < userCommentListLength; i++) {
+		const commentLength = dataView.getUint32(offset, true);
+		offset += 4;
+		const comment = getString(dataView, offset, commentLength);
+		offset += commentLength;
+
+		const [key, value] = comment.split('=');
+
+		comments[key] = value;
+	}
+
+	return comments;
+}
+
+// Function to read a FLAC file and parse Vorbis comments
+export function getFlacMetadata(file) {
+	return new Promise((r) => {
+		const reader = new FileReader();
+		reader.onload = function(event) {
+			const arrayBuffer = event.target.result;
+			const dataView = new DataView(arrayBuffer);
+
+			// Verify the FLAC signature
+			const signature = String.fromCharCode(...new Uint8Array(arrayBuffer, 0, 4));
+			if (signature !== 'fLaC') {
+				console.error('Not a valid FLAC file');
+				return;
+			}
+
+			// Parse metadata blocks
+			let offset = 4;
+			let vorbisComment = null;
+			while (offset < dataView.byteLength) {
+				const isLastBlock = dataView.getUint8(offset) & 0x80;
+				const blockType = dataView.getUint8(offset) & 0x7F;
+				const blockSize = dataView.getUint32(offset, false) & 0xFFFFFF;
+				offset += 4;
+
+				if (blockType === 4) { // Vorbis Comment block type
+					vorbisComment = parseVorbisComment(new DataView(arrayBuffer, offset, blockSize));
+				}
+
+				offset += blockSize;
+				if (isLastBlock) break;
+			}
+
+			r(vorbisComment);
+		};
+		reader.readAsArrayBuffer(file);
+	});
+}
+
 export async function importA1111(graph, parameters) {
 	const p = parameters.lastIndexOf("\nSteps:");
 	if (p > -1) {

web/scripts/ui.js

@@ -6,17 +6,22 @@ import { ComfySettingsDialog } from "./ui/settings.js";
 export const ComfyDialog = _ComfyDialog;
 
 /**
- * 
+ * @template { string | (keyof HTMLElementTagNameMap) } K
- * @param { string } tag HTML Element Tag and optional classes e.g. div.class1.class2
+ * @typedef { K extends keyof HTMLElementTagNameMap ? HTMLElementTagNameMap[K] : HTMLElement } ElementType 
- * @param { string | Element | Element[] | {
+ */
+
+/**
+ * @template { string | (keyof HTMLElementTagNameMap) } K 
+ * @param { K } tag HTML Element Tag and optional classes e.g. div.class1.class2
+ * @param { string | Element | Element[] | ({
  * 	 parent?: Element,
- *   $?: (el: Element) => void, 
+ *   $?: (el: ElementType<K>) => void, 
  *   dataset?: DOMStringMap,
- *   style?: CSSStyleDeclaration,
+ *   style?: Partial<CSSStyleDeclaration>,
  * 	 for?: string
- * } | undefined } propsOrChildren 
+ * } & Omit<Partial<ElementType<K>>, "style">) | undefined } [propsOrChildren]
- * @param { Element[] | undefined } [children]
+ * @param { string | Element | Element[] | undefined } [children]
- * @returns 
+ * @returns { ElementType<K> }
  */
 export function $el(tag, propsOrChildren, children) {
 	const split = tag.split(".");
@@ -54,7 +59,7 @@ export function $el(tag, propsOrChildren, children) {
 
 			Object.assign(element, propsOrChildren);
 			if (children) {
-				element.append(...(children instanceof Array ? children : [children]));
+				element.append(...(children instanceof Array ? children.filter(Boolean) : [children]));
 			}
 
 			if (parent) {
@@ -102,6 +107,8 @@ function dragElement(dragEl, settings) {
 	}
 
 	function positionElement() {
+		if(dragEl.style.display === "none") return;
+
 		const halfWidth = document.body.clientWidth / 2;
 		const anchorRight = newPosX + dragEl.clientWidth / 2 > halfWidth;
 
@@ -191,6 +198,8 @@ function dragElement(dragEl, settings) {
 		document.onmouseup = null;
 		document.onmousemove = null;
 	}
+
+	return restorePos;
 }
 
 class ComfyList {
@@ -371,7 +380,7 @@ export class ComfyUI {
 		const fileInput = $el("input", {
 			id: "comfy-file-input",
 			type: "file",
-			accept: ".json,image/png,.latent,.safetensors,image/webp",
+			accept: ".json,image/png,.latent,.safetensors,image/webp,audio/flac",
 			style: {display: "none"},
 			parent: document.body,
 			onchange: () => {
@@ -379,6 +388,8 @@ export class ComfyUI {
 			},
 		});
 
+		this.loadFile = () => fileInput.click();
+
 		const autoQueueModeEl = toggleSwitch(
 			"autoQueueMode",
 			[
@@ -628,10 +639,10 @@ export class ComfyUI {
 			name: "Enable Dev mode Options",
 			type: "boolean",
 			defaultValue: false,
-			onChange: function(value) { document.getElementById("comfy-dev-save-api-button").style.display = value ? "block" : "none"},
+			onChange: function(value) { document.getElementById("comfy-dev-save-api-button").style.display = value ? "flex" : "none"},
 		});
 
-		dragElement(this.menuContainer, this.settings);
+		this.restoreMenuPosition = dragElement(this.menuContainer, this.settings);
 
 		this.setStatus({exec_info: {queue_remaining: "X"}});
 	}

web/scripts/ui/components/asyncDialog.js

@@ -0,0 +1,64 @@
+import { ComfyDialog } from "../dialog.js";
+import { $el } from "../../ui.js";
+
+export class ComfyAsyncDialog extends ComfyDialog {
+	#resolve;
+
+	constructor(actions) {
+		super(
+			"dialog.comfy-dialog.comfyui-dialog",
+			actions?.map((opt) => {
+				if (typeof opt === "string") {
+					opt = { text: opt };
+				}
+				return $el("button.comfyui-button", {
+					type: "button",
+					textContent: opt.text,
+					onclick: () => this.close(opt.value ?? opt.text),
+				});
+			})
+		);
+	}
+
+	show(html) {
+		this.element.addEventListener("close", () => {
+			this.close();
+		});
+
+		super.show(html);
+
+		return new Promise((resolve) => {
+			this.#resolve = resolve;
+		});
+	}
+
+	showModal(html) {
+		this.element.addEventListener("close", () => {
+			this.close();
+		});
+
+		super.show(html);
+		this.element.showModal();
+
+		return new Promise((resolve) => {
+			this.#resolve = resolve;
+		});
+	}
+
+	close(result = null) {
+		this.#resolve(result);
+		this.element.close();
+		super.close();
+	}
+
+	static async prompt({ title = null, message, actions }) {
+		const dialog = new ComfyAsyncDialog(actions);
+		const content = [$el("span", message)];
+		if (title) {
+			content.unshift($el("h3", title));
+		}
+		const res = await dialog.showModal(content);
+		dialog.element.remove();
+		return res;
+	}
+}

web/scripts/ui/components/button.js

@@ -0,0 +1,163 @@
+// @ts-check
+
+import { $el } from "../../ui.js";
+import { applyClasses, toggleElement } from "../utils.js";
+import { prop } from "../../utils.js";
+
+/**
+ * @typedef {{
+ *    icon?: string;
+ *    overIcon?: string;
+ *    iconSize?: number;
+ *    content?: string | HTMLElement;
+ *    tooltip?: string;
+ *    enabled?: boolean;
+ *    action?: (e: Event, btn: ComfyButton) => void,
+ *    classList?: import("../utils.js").ClassList,
+ * 	  visibilitySetting?: { id: string, showValue: any },
+ * 	  app?: import("../../app.js").ComfyApp
+ * }} ComfyButtonProps
+ */
+export class ComfyButton {
+	#over = 0;
+	#popupOpen = false;
+	isOver = false;	
+	iconElement = $el("i.mdi");
+	contentElement = $el("span");
+	/**
+	 * @type {import("./popup.js").ComfyPopup}
+	 */
+	popup;
+
+	/**
+	 * @param {ComfyButtonProps} opts
+	 */
+	constructor({
+		icon,
+		overIcon,
+		iconSize,
+		content,
+		tooltip,
+		action,
+		classList = "comfyui-button",
+		visibilitySetting,
+		app,
+		enabled = true,
+	}) {
+		this.element = $el("button", {
+			onmouseenter: () => {
+				this.isOver = true;
+				if(this.overIcon) {
+					this.updateIcon();
+				}
+			},
+			onmouseleave: () => {
+				this.isOver = false;
+				if(this.overIcon) {
+					this.updateIcon();
+				}
+			}
+
+		}, [this.iconElement, this.contentElement]);
+
+		this.icon = prop(this, "icon", icon, toggleElement(this.iconElement, { onShow: this.updateIcon }));
+		this.overIcon = prop(this, "overIcon", overIcon, () => {
+			if(this.isOver) {
+				this.updateIcon();
+			}
+		});
+		this.iconSize = prop(this, "iconSize", iconSize, this.updateIcon);
+		this.content = prop(
+			this,
+			"content",
+			content,
+			toggleElement(this.contentElement, {
+				onShow: (el, v) => {
+					if (typeof v === "string") {
+						el.textContent = v;
+					} else {
+						el.replaceChildren(v);
+					}
+				},
+			})
+		);
+
+		this.tooltip = prop(this, "tooltip", tooltip, (v) => {
+			if (v) {
+				this.element.title = v;
+			} else {
+				this.element.removeAttribute("title");
+			}
+		});
+		this.classList = prop(this, "classList", classList, this.updateClasses);
+		this.hidden = prop(this, "hidden", false, this.updateClasses);
+		this.enabled = prop(this, "enabled", enabled, () => {
+			this.updateClasses();
+			this.element.disabled = !this.enabled;
+		});
+		this.action = prop(this, "action", action);
+		this.element.addEventListener("click", (e) => {
+			if (this.popup) {
+				// we are either a touch device or triggered by click not hover
+				if (!this.#over) {
+					this.popup.toggle();
+				}
+			}
+			this.action?.(e, this);
+		});
+
+		if (visibilitySetting?.id) {
+			const settingUpdated = () => {
+				this.hidden = app.ui.settings.getSettingValue(visibilitySetting.id) !== visibilitySetting.showValue;
+			};
+			app.ui.settings.addEventListener(visibilitySetting.id + ".change", settingUpdated);
+			settingUpdated();
+		}
+	}
+
+	updateIcon = () => (this.iconElement.className = `mdi mdi-${(this.isOver && this.overIcon) || this.icon}${this.iconSize ? " mdi-" + this.iconSize + "px" : ""}`);
+	updateClasses = () => {
+		const internalClasses = [];
+		if (this.hidden) {
+			internalClasses.push("hidden");
+		}
+		if (!this.enabled) {
+			internalClasses.push("disabled");
+		}
+		if (this.popup) {
+			if (this.#popupOpen) {
+				internalClasses.push("popup-open");
+			} else {
+				internalClasses.push("popup-closed");
+			}
+		}
+		applyClasses(this.element, this.classList, ...internalClasses);
+	};
+
+	/**
+	 *
+	 * @param { import("./popup.js").ComfyPopup } popup
+	 * @param { "click" | "hover" } mode
+	 */
+	withPopup(popup, mode = "click") {
+		this.popup = popup;
+
+		if (mode === "hover") {
+			for (const el of [this.element, this.popup.element]) {
+				el.addEventListener("mouseenter", () => {
+					this.popup.open = !!++this.#over;
+				});
+				el.addEventListener("mouseleave", () => {
+					this.popup.open = !!--this.#over;
+				});
+			}
+		}
+
+		popup.addEventListener("change", () => {
+			this.#popupOpen = popup.open;
+			this.updateClasses();
+		});
+
+		return this;
+	}
+}

web/scripts/ui/components/buttonGroup.js

@@ -0,0 +1,45 @@
+// @ts-check
+
+import { $el } from "../../ui.js";
+import { ComfyButton } from "./button.js";
+import { prop } from "../../utils.js";
+
+export class ComfyButtonGroup {
+	element = $el("div.comfyui-button-group");
+
+	/** @param {Array<ComfyButton | HTMLElement>} buttons */
+	constructor(...buttons) {
+		this.buttons = prop(this, "buttons", buttons, () => this.update());
+	}
+
+	/**
+	 * @param {ComfyButton} button
+	 * @param {number} index
+	 */
+	insert(button, index) {
+		this.buttons.splice(index, 0, button);
+		this.update();
+	}
+
+	/** @param {ComfyButton} button */
+	append(button) {
+		this.buttons.push(button);
+		this.update();
+	}
+
+	/** @param {ComfyButton|number} indexOrButton */
+	remove(indexOrButton) {
+		if (typeof indexOrButton !== "number") {
+			indexOrButton = this.buttons.indexOf(indexOrButton);
+		}
+		if (indexOrButton > -1) {
+			const r = this.buttons.splice(indexOrButton, 1);
+			this.update();
+			return r;
+		}
+	}
+
+	update() {
+		this.element.replaceChildren(...this.buttons.map((b) => b["element"] ?? b));
+	}
+}

web/scripts/ui/components/popup.js

@@ -0,0 +1,128 @@
+// @ts-check
+
+import { prop } from "../../utils.js";
+import { $el } from "../../ui.js";
+import { applyClasses } from "../utils.js";
+
+export class ComfyPopup extends EventTarget {
+	element = $el("div.comfyui-popup");
+
+	/**
+	 * @param {{
+	 *     target: HTMLElement,
+	 *      container?: HTMLElement,
+	 *      classList?: import("../utils.js").ClassList,
+	 * 		ignoreTarget?: boolean,
+	 * 		closeOnEscape?: boolean,
+	 * 		position?: "absolute" | "relative",
+	 * 		horizontal?: "left" | "right"
+	 * }} param0
+	 * @param  {...HTMLElement} children
+	 */
+	constructor(
+		{
+			target,
+			container = document.body,
+			classList = "",
+			ignoreTarget = true,
+			closeOnEscape = true,
+			position = "absolute",
+			horizontal = "left",
+		},
+		...children
+	) {
+		super();
+		this.target = target;
+		this.ignoreTarget = ignoreTarget;
+		this.container = container;
+		this.position = position;
+		this.closeOnEscape = closeOnEscape;
+		this.horizontal = horizontal;
+
+		container.append(this.element);
+
+		this.children = prop(this, "children", children, () => {
+			this.element.replaceChildren(...this.children);
+			this.update();
+		});
+		this.classList = prop(this, "classList", classList, () => applyClasses(this.element, this.classList, "comfyui-popup", horizontal));
+		this.open = prop(this, "open", false, (v, o) => {
+			if (v === o) return;
+			if (v) {
+				this.#show();
+			} else {
+				this.#hide();
+			}
+		});
+	}
+
+	toggle() {
+		this.open = !this.open;
+	}
+
+	#hide() {
+		this.element.classList.remove("open");
+		window.removeEventListener("resize", this.update);
+		window.removeEventListener("click", this.#clickHandler, { capture: true });
+		window.removeEventListener("keydown", this.#escHandler, { capture: true });
+
+		this.dispatchEvent(new CustomEvent("close"));
+		this.dispatchEvent(new CustomEvent("change"));
+	}
+
+	#show() {
+		this.element.classList.add("open");
+		this.update();
+
+		window.addEventListener("resize", this.update);
+		window.addEventListener("click", this.#clickHandler, { capture: true });
+		if (this.closeOnEscape) {
+			window.addEventListener("keydown", this.#escHandler, { capture: true });
+		}
+
+		this.dispatchEvent(new CustomEvent("open"));
+		this.dispatchEvent(new CustomEvent("change"));
+	}
+
+	#escHandler = (e) => {
+		if (e.key === "Escape") {
+			this.open = false;
+			e.preventDefault();
+			e.stopImmediatePropagation();
+		}
+	};
+
+	#clickHandler = (e) => {
+		/** @type {any} */
+		const target = e.target;
+		if (!this.element.contains(target) && this.ignoreTarget && !this.target.contains(target)) {
+			this.open = false;
+		}
+	};
+
+	update = () => {
+		const rect = this.target.getBoundingClientRect();
+		this.element.style.setProperty("--bottom", "unset");
+		if (this.position === "absolute") {
+			if (this.horizontal === "left") {
+				this.element.style.setProperty("--left", rect.left + "px");
+			} else {
+				this.element.style.setProperty("--left", rect.right - this.element.clientWidth + "px");
+			}
+			this.element.style.setProperty("--top", rect.bottom + "px");
+			this.element.style.setProperty("--limit", rect.bottom + "px");
+		} else {
+			this.element.style.setProperty("--left", 0 + "px");
+			this.element.style.setProperty("--top", rect.height + "px");
+			this.element.style.setProperty("--limit", rect.height + "px");
+		}
+
+		const thisRect = this.element.getBoundingClientRect();
+		if (thisRect.height < 30) {
+			// Move up instead
+			this.element.style.setProperty("--top", "unset");
+			this.element.style.setProperty("--bottom", rect.height + 5 + "px");
+			this.element.style.setProperty("--limit", rect.height + 5 + "px");
+		}
+	};
+}

web/scripts/ui/components/splitButton.js

@@ -0,0 +1,43 @@
+// @ts-check
+
+import { $el } from "../../ui.js";
+import { ComfyButton } from "./button.js";
+import { prop } from "../../utils.js";
+import { ComfyPopup } from "./popup.js";
+
+export class ComfySplitButton {
+	/**
+	 *  @param {{
+	 * 		primary: ComfyButton,
+	 * 		mode?: "hover" | "click",
+	 * 		horizontal?: "left" | "right",
+	 * 		position?: "relative" | "absolute"
+	 *  }} param0
+	 *  @param {Array<ComfyButton> | Array<HTMLElement>} items
+	 */
+	constructor({ primary, mode, horizontal = "left", position = "relative" }, ...items) {
+		this.arrow = new ComfyButton({
+			icon: "chevron-down",
+		});
+		this.element = $el("div.comfyui-split-button" + (mode === "hover" ? ".hover" : ""), [
+			$el("div.comfyui-split-primary", primary.element),
+			$el("div.comfyui-split-arrow", this.arrow.element),
+		]);
+		this.popup = new ComfyPopup({
+			target: this.element,
+			container: position === "relative" ? this.element : document.body,
+			classList: "comfyui-split-button-popup" + (mode === "hover" ? " hover" : ""),
+			closeOnEscape: mode === "click",
+			position,
+			horizontal,
+		});
+
+		this.arrow.withPopup(this.popup, mode);
+
+		this.items = prop(this, "items", items, () => this.update());
+	}
+
+	update() {
+		this.popup.element.replaceChildren(...this.items.map((b) => b.element ?? b));
+	}
+}

web/scripts/ui/dialog.js

@@ -1,20 +1,26 @@
 import { $el } from "../ui.js";
 
-export class ComfyDialog {
+export class ComfyDialog extends EventTarget {
-	constructor() {
+	#buttons;
-		this.element = $el("div.comfy-modal", { parent: document.body }, [
+
+	constructor(type = "div", buttons = null) {
+		super();
+		this.#buttons = buttons;
+		this.element = $el(type + ".comfy-modal", { parent: document.body }, [
 			$el("div.comfy-modal-content", [$el("p", { $: (p) => (this.textElement = p) }), ...this.createButtons()]),
 		]);
 	}
 
 	createButtons() {
-		return [
+		return (
-			$el("button", {
+			this.#buttons ?? [
-				type: "button",
+				$el("button", {
-				textContent: "Close",
+					type: "button",
-				onclick: () => this.close(),
+					textContent: "Close",
-			}),
+					onclick: () => this.close(),
-		];
+				}),
+			]
+		);
 	}
 
 	close() {
@@ -25,7 +31,7 @@ export class ComfyDialog {
 		if (typeof html === "string") {
 			this.textElement.innerHTML = html;
 		} else {
-			this.textElement.replaceChildren(html);
+			this.textElement.replaceChildren(...(html instanceof Array ? html : [html]));
 		}
 		this.element.style.display = "flex";
 	}

web/scripts/ui/menu/index.js

@@ -0,0 +1,302 @@
+// @ts-check
+
+import { $el } from "../../ui.js";
+import { downloadBlob } from "../../utils.js";
+import { ComfyButton } from "../components/button.js";
+import { ComfyButtonGroup } from "../components/buttonGroup.js";
+import { ComfySplitButton } from "../components/splitButton.js";
+import { ComfyViewHistoryButton } from "./viewHistory.js";
+import { ComfyQueueButton } from "./queueButton.js";
+import { ComfyWorkflowsMenu } from "./workflows.js";
+import { ComfyViewQueueButton } from "./viewQueue.js";
+import { getInteruptButton } from "./interruptButton.js";
+
+const collapseOnMobile = (t) => {
+	(t.element ?? t).classList.add("comfyui-menu-mobile-collapse");
+	return t;
+};
+const showOnMobile = (t) => {
+	(t.element ?? t).classList.add("lt-lg-show");
+	return t;
+};
+
+export class ComfyAppMenu {
+	#sizeBreak = "lg";
+	#lastSizeBreaks = {
+		lg: null,
+		md: null,
+		sm: null,
+		xs: null,
+	};
+	#sizeBreaks = Object.keys(this.#lastSizeBreaks);
+	#cachedInnerSize = null;
+	#cacheTimeout = null;
+
+	/**
+	 * @param { import("../../app.js").ComfyApp } app
+	 */
+	constructor(app) {
+		this.app = app;
+
+		this.workflows = new ComfyWorkflowsMenu(app);
+		const getSaveButton = (t) =>
+			new ComfyButton({
+				icon: "content-save",
+				tooltip: "Save the current workflow",
+				action: () => app.workflowManager.activeWorkflow.save(),
+				content: t,
+			});
+
+		this.logo = $el("h1.comfyui-logo.nlg-hide", { title: "ComfyUI" }, "ComfyUI");
+		this.saveButton = new ComfySplitButton(
+			{
+				primary: getSaveButton(),
+				mode: "hover",
+				position: "absolute",
+			},
+			getSaveButton("Save"),
+			new ComfyButton({
+				icon: "content-save-edit",
+				content: "Save As",
+				tooltip: "Save the current graph as a new workflow",
+				action: () => app.workflowManager.activeWorkflow.save(true),
+			}),
+			new ComfyButton({
+				icon: "download",
+				content: "Export",
+				tooltip: "Export the current workflow as JSON",
+				action: () => this.exportWorkflow("workflow", "workflow"),
+			}),
+			new ComfyButton({
+				icon: "api",
+				content: "Export (API Format)",
+				tooltip: "Export the current workflow as JSON for use with the ComfyUI API",
+				action: () => this.exportWorkflow("workflow_api", "output"),
+				visibilitySetting: { id: "Comfy.DevMode", showValue: true },
+				app,
+			})
+		);
+		this.actionsGroup = new ComfyButtonGroup(
+			new ComfyButton({
+				icon: "refresh",
+				content: "Refresh",
+				tooltip: "Refresh widgets in nodes to find new models or files",
+				action: () => app.refreshComboInNodes(),
+			}),
+			new ComfyButton({
+				icon: "clipboard-edit-outline",
+				content: "Clipspace",
+				tooltip: "Open Clipspace window",
+				action: () => app["openClipspace"](),
+			}),
+			new ComfyButton({
+				icon: "fit-to-page-outline",
+				content: "Reset View",
+				tooltip: "Reset the canvas view",
+				action: () => app.resetView(),
+			}),
+			new ComfyButton({
+				icon: "cancel",
+				content: "Clear",
+				tooltip: "Clears current workflow",
+				action: () => {
+					if (!app.ui.settings.getSettingValue("Comfy.ConfirmClear", true) || confirm("Clear workflow?")) {
+						app.clean();
+						app.graph.clear();
+					}
+				},
+			})
+		);
+		this.settingsGroup = new ComfyButtonGroup(
+			new ComfyButton({
+				icon: "cog",
+				content: "Settings",
+				tooltip: "Open settings",
+				action: () => {
+					app.ui.settings.show();
+				},
+			})
+		);
+		this.viewGroup = new ComfyButtonGroup(
+			new ComfyViewHistoryButton(app).element,
+			new ComfyViewQueueButton(app).element,
+			getInteruptButton("nlg-hide").element
+		);
+		this.mobileMenuButton = new ComfyButton({
+			icon: "menu",
+			action: (_, btn) => {
+				btn.icon = this.element.classList.toggle("expanded") ? "menu-open" : "menu";
+				window.dispatchEvent(new Event("resize"));
+			},
+			classList: "comfyui-button comfyui-menu-button",
+		});
+
+		this.element = $el("nav.comfyui-menu.lg", { style: { display: "none" } }, [
+			this.logo,
+			this.workflows.element,
+			this.saveButton.element,
+			collapseOnMobile(this.actionsGroup).element,
+			$el("section.comfyui-menu-push"),
+			collapseOnMobile(this.settingsGroup).element,
+			collapseOnMobile(this.viewGroup).element,
+
+			getInteruptButton("lt-lg-show").element,
+			new ComfyQueueButton(app).element,
+			showOnMobile(this.mobileMenuButton).element,
+		]);
+
+		let resizeHandler;
+		this.menuPositionSetting = app.ui.settings.addSetting({
+			id: "Comfy.UseNewMenu",
+			defaultValue: "Disabled",
+			name: "[Beta] Use new menu and workflow management. Note: On small screens the menu will always be at the top.",
+			type: "combo",
+			options: ["Disabled", "Top", "Bottom"],
+			onChange: async (v) => {
+				if (v && v !== "Disabled") {
+					if (!resizeHandler) {
+						resizeHandler = () => {
+							this.calculateSizeBreak();
+						};
+						window.addEventListener("resize", resizeHandler);
+					}
+					this.updatePosition(v);
+				} else {
+					if (resizeHandler) {
+						window.removeEventListener("resize", resizeHandler);
+						resizeHandler = null;
+					}
+					document.body.style.removeProperty("display");
+					app.ui.menuContainer.style.removeProperty("display");
+					this.element.style.display = "none";
+					app.ui.restoreMenuPosition();
+				}
+				window.dispatchEvent(new Event("resize"));
+			},
+		});
+	}
+
+	updatePosition(v) {
+		document.body.style.display = "grid";
+		this.app.ui.menuContainer.style.display = "none";
+		this.element.style.removeProperty("display");
+		this.position = v;
+		if (v === "Bottom") {
+			this.app.bodyBottom.append(this.element);
+		} else {
+			this.app.bodyTop.prepend(this.element);
+		}
+		this.calculateSizeBreak();
+	}
+
+	updateSizeBreak(idx, prevIdx, direction) {
+		const newSize = this.#sizeBreaks[idx];
+		if (newSize === this.#sizeBreak) return;
+		this.#cachedInnerSize = null;
+		clearTimeout(this.#cacheTimeout);
+
+		this.#sizeBreak = this.#sizeBreaks[idx];
+		for (let i = 0; i < this.#sizeBreaks.length; i++) {
+			const sz = this.#sizeBreaks[i];
+			if (sz === this.#sizeBreak) {
+				this.element.classList.add(sz);
+			} else {
+				this.element.classList.remove(sz);
+			}
+			if (i < idx) {
+				this.element.classList.add("lt-" + sz);
+			} else {
+				this.element.classList.remove("lt-" + sz);
+			}
+		}
+
+		if (idx) {
+			// We're on a small screen, force the menu at the top
+			if (this.position !== "Top") {
+				this.updatePosition("Top");
+			}
+		} else if (this.position != this.menuPositionSetting.value) {
+			// Restore user position
+			this.updatePosition(this.menuPositionSetting.value);
+		}
+
+		// Allow multiple updates, but prevent bouncing
+		if (!direction) {
+			direction = prevIdx - idx;
+		} else if (direction != prevIdx - idx) {
+			return;
+		}
+		this.calculateSizeBreak(direction);
+	}
+
+	calculateSizeBreak(direction = 0) {
+		let idx = this.#sizeBreaks.indexOf(this.#sizeBreak);
+		const currIdx = idx;
+		const innerSize = this.calculateInnerSize(idx);
+		if (window.innerWidth >= this.#lastSizeBreaks[this.#sizeBreaks[idx - 1]]) {
+			if (idx > 0) {
+				idx--;
+			}
+		} else if (innerSize > this.element.clientWidth) {
+			this.#lastSizeBreaks[this.#sizeBreak] = Math.max(window.innerWidth, innerSize);
+			// We need to shrink
+			if (idx < this.#sizeBreaks.length - 1) {
+				idx++;
+			}
+		}
+
+		this.updateSizeBreak(idx, currIdx, direction);
+	}
+
+	calculateInnerSize(idx) {
+		// Cache the inner size to prevent too much calculation when resizing the window
+		clearTimeout(this.#cacheTimeout);
+		if (this.#cachedInnerSize) {
+			// Extend cache time
+			this.#cacheTimeout = setTimeout(() => (this.#cachedInnerSize = null), 100);
+		} else {
+			let innerSize = 0;
+			let count = 1;
+			for (const c of this.element.children) {
+				if (c.classList.contains("comfyui-menu-push")) continue; // ignore right push
+				if (idx && c.classList.contains("comfyui-menu-mobile-collapse")) continue; // ignore collapse items
+				innerSize += c.clientWidth;
+				count++;
+			}
+			innerSize += 8 * count;
+			this.#cachedInnerSize = innerSize;
+			this.#cacheTimeout = setTimeout(() => (this.#cachedInnerSize = null), 100);
+		}
+		return this.#cachedInnerSize;
+	}
+
+	/**
+	 * @param {string} defaultName
+	 */
+	getFilename(defaultName) {
+		if (this.app.ui.settings.getSettingValue("Comfy.PromptFilename", true)) {
+			defaultName = prompt("Save workflow as:", defaultName);
+			if (!defaultName) return;
+			if (!defaultName.toLowerCase().endsWith(".json")) {
+				defaultName += ".json";
+			}
+		}
+		return defaultName;
+	}
+
+	/**
+	 * @param {string} [filename]
+	 * @param { "workflow" | "output" } [promptProperty]
+	 */
+	async exportWorkflow(filename, promptProperty) {
+		if (this.app.workflowManager.activeWorkflow?.path) {
+			filename = this.app.workflowManager.activeWorkflow.name;
+		}
+		const p = await this.app.graphToPrompt();
+		const json = JSON.stringify(p[promptProperty], null, 2);
+		const blob = new Blob([json], { type: "application/json" });
+		const file = this.getFilename(filename);
+		if (!file) return;
+		downloadBlob(file, blob);
+	}
+}

web/scripts/ui/menu/interruptButton.js

@@ -0,0 +1,23 @@
+// @ts-check
+
+import { api } from "../../api.js";
+import { ComfyButton } from "../components/button.js";
+
+export function getInteruptButton(visibility) {
+	const btn = new ComfyButton({
+		icon: "close",
+		tooltip: "Cancel current generation",
+		enabled: false,
+		action: () => {
+			api.interrupt();
+		},
+		classList: ["comfyui-button", "comfyui-interrupt-button", visibility],
+	});
+
+	api.addEventListener("status", ({ detail }) => {
+		const sz = detail?.exec_info?.queue_remaining;
+		btn.enabled = sz > 0;
+	});
+
+	return btn;
+}

web/scripts/ui/menu/menu.css

@@ -0,0 +1,705 @@
+.relative {
+	position: relative;
+}
+.hidden {
+	display: none !important;
+}
+.mdi.rotate270::before {
+	transform: rotate(270deg);
+}
+
+/* Generic */
+.comfyui-button {
+	display: flex;
+	align-items: center;
+	gap: 0.5em;
+	cursor: pointer;
+	border: none;
+	border-radius: 4px;
+	padding: 4px 8px;
+	box-sizing: border-box;
+	margin: 0;
+	transition: box-shadow 0.1s;
+}
+
+.comfyui-button:active {
+	box-shadow: inset 1px 1px 10px rgba(0, 0, 0, 0.5);
+}
+.comfyui-button:disabled {
+	opacity: 0.5;
+	cursor: not-allowed;
+}
+.primary .comfyui-button,
+.primary.comfyui-button {
+	background-color: var(--primary-bg) !important;
+	color: var(--primary-fg) !important;
+}
+
+.primary .comfyui-button:not(:disabled):hover,
+.primary.comfyui-button:not(:disabled):hover {
+	background-color: var(--primary-hover-bg) !important;
+	color: var(--primary-hover-fg) !important;
+}
+
+/* Popup */
+.comfyui-popup {
+	position: absolute;
+	left: var(--left);
+	right: var(--right);
+	top: var(--top);
+	bottom: var(--bottom);
+	z-index: 2000;
+	max-height: calc(100vh - var(--limit) - 10px);
+	box-shadow: 3px 3px 5px 0px rgba(0, 0, 0, 0.3);
+}
+
+.comfyui-popup:not(.open) {
+	display: none;
+}
+
+.comfyui-popup.right.open {
+	border-top-left-radius: 4px;
+	border-bottom-right-radius: 4px;
+	border-bottom-left-radius: 4px;
+	overflow: hidden;
+}
+/* Split button */
+.comfyui-split-button {
+	position: relative;
+	display: flex;
+}
+
+.comfyui-split-primary {
+	flex: auto;
+}
+
+.comfyui-split-primary .comfyui-button {
+	border-top-right-radius: 0;
+	border-bottom-right-radius: 0;
+	border-right: 1px solid var(--comfy-menu-bg);
+	width: 100%;
+}
+
+.comfyui-split-arrow .comfyui-button {
+	border-top-left-radius: 0;
+	border-bottom-left-radius: 0;
+	padding-left: 2px;
+	padding-right: 2px;
+}
+
+.comfyui-split-button-popup {
+	white-space: nowrap;
+	background-color: var(--content-bg);
+	color: var(--content-fg);
+	display: flex;
+	flex-direction: column;
+	overflow: auto;
+}
+
+.comfyui-split-button-popup.hover {
+	z-index: 2001;
+}
+.comfyui-split-button-popup > .comfyui-button {
+	border: none;
+	background-color: transparent;
+	color: var(--fg-color);
+	padding: 8px 12px 8px 8px;
+}
+
+.comfyui-split-button-popup > .comfyui-button:not(:disabled):hover {
+	background-color: var(--comfy-input-bg);
+}
+
+/* Button group */
+.comfyui-button-group {
+	display: flex;
+	border-radius: 4px;
+	overflow: hidden;
+}
+
+.comfyui-button-group > .comfyui-button,
+.comfyui-button-group > .comfyui-button-wrapper > .comfyui-button {
+	padding: 4px 10px;
+	border-radius: 0;
+}
+
+/* Menu */
+.comfyui-menu {
+	width: 100vw;
+	background: var(--comfy-menu-bg);
+	color: var(--fg-color);
+	font-family: Arial, Helvetica, sans-serif;
+	font-size: 0.8em;
+	display: flex;
+	padding: 4px 8px;
+	align-items: center;
+	gap: 8px;
+	box-sizing: border-box;
+	z-index: 1000;
+	order: 0;
+	grid-column: 1/-1;
+	overflow: auto;
+	max-height: 90vh;
+}
+
+.comfyui-menu>* {
+	flex-shrink: 0;
+}
+.comfyui-menu .mdi::before {
+	font-size: 18px;
+}
+
+.comfyui-menu .comfyui-button {
+	background: var(--comfy-input-bg);
+	color: var(--fg-color);
+	white-space: nowrap;
+}
+
+.comfyui-menu .comfyui-button:not(:disabled):hover {
+	background: var(--border-color);
+	color: var(--content-fg);
+}
+
+.comfyui-menu .comfyui-split-button-popup > .comfyui-button {
+	border-radius: 0;
+	background-color: transparent;
+}
+
+.comfyui-menu .comfyui-split-button-popup > .comfyui-button:not(:disabled):hover {
+	background-color: var(--comfy-input-bg);
+}
+
+.comfyui-menu .comfyui-split-button-popup.left {
+	border-top-right-radius: 4px;
+	border-bottom-left-radius: 4px;
+	border-bottom-right-radius: 4px;
+}
+
+.comfyui-menu .comfyui-button.popup-open {
+	background-color: var(--content-bg);
+	color: var(--content-fg);
+}
+
+.comfyui-menu-push {
+	margin-left: -0.8em;
+	flex: auto;
+}
+
+.comfyui-logo {
+	font-size: 1.2em;
+	margin: 0;
+	user-select: none;
+	cursor: default;
+}
+
+/* Workflows */
+.comfyui-workflows-button {
+	flex-direction: row-reverse;
+	max-width: 200px;
+	position: relative;
+	z-index: 0;
+}
+
+.comfyui-workflows-button.popup-open {
+	border-bottom-left-radius: 0;
+	border-bottom-right-radius: 0;
+}
+.comfyui-workflows-button.unsaved {
+	font-style: italic;
+}
+.comfyui-workflows-button-progress {
+	position: absolute;
+	top: 0;
+	left: 0;
+	background-color: green;
+	height: 100%;
+	border-radius: 4px;
+	z-index: -1;
+}
+
+.comfyui-workflows-button > span {
+	flex: auto;
+	text-align: left;
+	overflow: hidden;
+}
+.comfyui-workflows-button-inner {
+	display: flex;
+	align-items: center;
+	gap: 7px;
+	width: 150px;
+}
+.comfyui-workflows-label {
+	overflow: hidden;
+	text-overflow: ellipsis;
+	direction: rtl;
+	flex: auto;
+	position: relative;
+}
+
+.comfyui-workflows-button.unsaved .comfyui-workflows-label {
+	padding-left: 8px;
+}
+
+.comfyui-workflows-button.unsaved .comfyui-workflows-label:after {
+	content: "*";
+	position: absolute;
+	top: 0;
+	left: 0;
+}
+.comfyui-workflows-button-inner .mdi-graph::before {
+	transform: rotate(-90deg);
+}
+
+.comfyui-workflows-popup {
+	font-family: Arial, Helvetica, sans-serif;
+	font-size: 0.8em;
+	padding: 10px;
+	overflow: auto;
+	background-color: var(--content-bg);
+	color: var(--content-fg);
+	border-top-right-radius: 4px;
+	border-bottom-right-radius: 4px;
+	border-bottom-left-radius: 4px;
+	z-index: 400;
+}
+
+.comfyui-workflows-panel {
+	min-height: 150px;
+}
+
+.comfyui-workflows-panel .lds-ring {
+	transform: translate(-50%);
+	position: absolute;
+	left: 50%;
+	top: 75px;
+}
+
+.comfyui-workflows-panel h3 {
+	margin: 10px 0 10px 0;
+	font-size: 11px;
+	opacity: 0.8;
+}
+
+.comfyui-workflows-panel section header {
+	display: flex;
+	justify-content: space-between;
+	align-items: center;
+}
+.comfy-ui-workflows-search .mdi {
+	position: relative;
+	top: 2px;
+	pointer-events: none;
+}
+.comfy-ui-workflows-search input {
+	background-color: var(--comfy-input-bg);
+	color: var(--input-text);
+	border: none;
+	border-radius: 4px;
+	padding: 4px 10px;
+	margin-left: -24px;
+	text-indent: 18px;
+}
+.comfy-ui-workflows-search input:placeholder-shown {
+	width: 10px;
+}
+.comfy-ui-workflows-search input:placeholder-shown:focus {
+	width: auto;
+}
+.comfyui-workflows-actions {
+	display: flex;
+	gap: 10px;
+	margin-bottom: 10px;
+}
+
+.comfyui-workflows-actions .comfyui-button {
+	background: var(--comfy-input-bg);
+	color: var(--input-text);
+}
+
+.comfyui-workflows-actions .comfyui-button:not(:disabled):hover {
+	background: var(--primary-bg);
+	color: var(--primary-fg);
+}
+
+.comfyui-workflows-favorites,
+.comfyui-workflows-open {
+	border-bottom: 1px solid var(--comfy-input-bg);
+	padding-bottom: 5px;
+	margin-bottom: 5px;
+}
+
+.comfyui-workflows-open .active {
+	font-weight: bold;
+}
+
+.comfyui-workflows-favorites:empty {
+	display: none;
+}
+
+.comfyui-workflows-tree {
+	padding: 0;
+	margin: 0;
+}
+
+.comfyui-workflows-tree:empty::after {
+	content: "No saved workflows";
+	display: block;
+	text-align: center;
+}
+.comfyui-workflows-tree > ul {
+	padding: 0;
+}
+
+.comfyui-workflows-tree > ul ul {
+	margin: 0;
+	padding: 0 0 0 25px;
+}
+
+.comfyui-workflows-tree:not(.filtered) .closed > ul {
+	display: none;
+}
+
+.comfyui-workflows-tree li,
+.comfyui-workflows-tree-file {
+	--item-height: 32px;
+	list-style-type: none;
+	height: var(--item-height);
+	display: flex;
+	align-items: center;
+	gap: 5px;
+	cursor: pointer;
+	user-select: none;
+}
+
+.comfyui-workflows-tree-file.active::before,
+.comfyui-workflows-tree li:hover::before,
+.comfyui-workflows-tree-file:hover::before {
+	content: "";
+	position: absolute;
+	width: 100%;
+	left: 0;
+	height: var(--item-height);
+	background-color: var(--content-hover-bg);
+	color: var(--content-hover-fg);
+	z-index: -1;
+}
+
+.comfyui-workflows-tree-file.active::before {
+	background-color: var(--primary-bg);
+	color: var(--primary-fg);
+}
+
+.comfyui-workflows-tree-file.running:not(:hover)::before {
+	content: "";
+	position: absolute;
+	width: var(--progress, 0);
+	left: 0;
+	height: var(--item-height);
+	background-color: green;
+	z-index: -1;
+}
+
+.comfyui-workflows-tree-file.unsaved span {
+	font-style: italic;
+}
+
+.comfyui-workflows-tree-file span {
+	flex: auto;
+}
+
+.comfyui-workflows-tree-file span + .comfyui-workflows-file-action {
+	margin-left: 10px;
+}
+
+.comfyui-workflows-tree-file .comfyui-workflows-file-action {
+	background-color: transparent;
+	color: var(--fg-color);
+	padding: 2px 4px;
+}
+
+.lg ~ .comfyui-workflows-popup .comfyui-workflows-tree-file:not(:hover) .comfyui-workflows-file-action {
+	opacity: 0;
+}
+
+.comfyui-workflows-tree-file .comfyui-workflows-file-action:hover {
+	background-color: var(--primary-bg);
+	color: var(--primary-fg);
+}
+
+.comfyui-workflows-tree-file .comfyui-workflows-file-action-primary {
+	background-color: transparent;
+	color: var(--fg-color);
+	padding: 2px 4px;
+	margin: 0 -4px;
+}
+
+.comfyui-workflows-file-action-favorite .mdi-star {
+	color: orange;
+}
+
+/* View List */
+.comfyui-view-list-popup {
+	padding: 10px;
+	background-color: var(--content-bg);
+	color: var(--content-fg);
+	min-width: 170px;
+	min-height: 435px;
+	display: flex;
+	flex-direction: column;
+	align-items: center;
+	box-sizing: border-box;
+}
+.comfyui-view-list-popup h3 {
+	margin: 0 0 5px 0;
+}
+.comfyui-view-list-items {
+	width: 100%;
+	background: var(--comfy-menu-bg);
+	border-radius: 5px;
+	display: flex;
+	justify-content: center;
+	flex: auto;
+	align-items: center;
+	flex-direction: column;
+}
+.comfyui-view-list-items section {
+	max-height: 400px;
+	overflow: auto;
+	width: 100%;
+	display: grid;
+	grid-template-columns: auto auto auto;
+	align-items: center;
+	justify-content: center;
+	gap: 5px;
+	padding: 5px 0;
+}
+.comfyui-view-list-items section + section {
+	border-top: 1px solid var(--border-color);
+	margin-top: 10px;
+	padding-top: 5px;
+}
+.comfyui-view-list-items section h5 {
+	grid-column: 1 / 4;
+	text-align: center;
+	margin: 5px;
+}
+.comfyui-view-list-items span {
+	text-align: center;
+	padding: 0 2px;
+}
+.comfyui-view-list-popup header {
+	margin-bottom: 10px;
+	display: flex;
+	gap: 5px;
+}
+.comfyui-view-list-popup header .comfyui-button {
+	border: 1px solid transparent;
+}
+.comfyui-view-list-popup header .comfyui-button:not(:disabled):hover {
+	border: 1px solid var(--comfy-menu-bg);
+}
+/* Queue button */
+.comfyui-queue-button .comfyui-split-primary .comfyui-button {
+	padding-right: 12px;
+}
+.comfyui-queue-count {
+	margin-left: 5px;
+	border-radius: 10px;
+	background-color: rgb(8, 80, 153);
+	padding: 2px 4px;
+	font-size: 10px;
+	min-width: 1em;
+	display: inline-block;
+}
+/* Queue options*/
+.comfyui-queue-options {
+	padding: 10px;
+	font-family: Arial, Helvetica, sans-serif;
+	font-size: 12px;
+	display: flex;
+	gap: 10px;
+}
+
+.comfyui-queue-batch {
+	display: flex;
+	flex-direction: column;
+	border-right: 1px solid var(--comfy-menu-bg);
+	padding-right: 10px;
+	gap: 5px;
+}
+
+.comfyui-queue-batch input {
+	width: 145px;
+}
+
+.comfyui-queue-batch .comfyui-queue-batch-value {
+	width: 70px;
+}
+
+.comfyui-queue-mode {
+	display: flex;
+	flex-direction: column;
+}
+
+.comfyui-queue-mode span {
+	font-weight: bold;
+	margin-bottom: 2px;
+}
+
+.comfyui-queue-mode label {
+	display: flex;
+	flex-direction: row-reverse;
+	justify-content: start;
+	gap: 5px;
+	padding: 2px 0;
+}
+
+.comfyui-queue-mode label input {
+	padding: 0;
+	margin: 0;
+}
+
+/** Send to workflow widget selection dialog */
+.comfy-widget-selection-dialog {
+	border: none;
+}
+
+.comfy-widget-selection-dialog div {
+	color: var(--fg-color);
+	font-family: Arial, Helvetica, sans-serif;
+}
+
+.comfy-widget-selection-dialog h2 {
+	margin-top: 0;
+}
+
+.comfy-widget-selection-dialog section {
+	width: fit-content;
+	display: flex;
+	flex-direction: column;
+}
+
+.comfy-widget-selection-item {
+	display: flex;
+	gap: 10px;
+	align-items: center;
+}
+
+.comfy-widget-selection-item span {
+	margin-right: auto;
+}
+
+.comfy-widget-selection-item span::before {
+	content: '#' attr(data-id);
+	opacity: 0.5;
+	margin-right: 5px;
+}
+
+.comfy-modal .comfy-widget-selection-item button {
+	font-size: 1em;
+}
+
+/***** Responsive *****/
+.lg.comfyui-menu .lt-lg-show {
+	display: none !important;
+}
+.comfyui-menu:not(.lg) .nlg-hide {
+	display: none !important;
+}
+/** Large screen */
+.lg.comfyui-menu>.comfyui-menu-mobile-collapse .comfyui-button span,
+.lg.comfyui-menu>.comfyui-menu-mobile-collapse.comfyui-button span {
+	display: none;
+}
+.lg.comfyui-menu>.comfyui-menu-mobile-collapse .comfyui-popup .comfyui-button span {
+	display: unset;
+}
+
+/** Non large screen */
+.lt-lg.comfyui-menu {
+	flex-wrap: wrap;
+}
+
+.lt-lg.comfyui-menu > *:not(.comfyui-menu-mobile-collapse) {
+	order: 1;
+}
+
+.lt-lg.comfyui-menu > .comfyui-menu-mobile-collapse {
+	order: 9999;
+	width: 100%;
+}
+
+.comfyui-body-bottom .lt-lg.comfyui-menu > .comfyui-menu-mobile-collapse {
+	order: -1;
+}
+
+.comfyui-body-bottom .lt-lg.comfyui-menu > .comfyui-menu-button {
+	top: unset;
+	bottom: 4px;
+}
+
+.lt-lg.comfyui-menu > .comfyui-menu-mobile-collapse.comfyui-button-group {
+	flex-wrap: wrap;
+}
+
+.lt-lg.comfyui-menu > .comfyui-menu-mobile-collapse .comfyui-button,
+.lt-lg.comfyui-menu > .comfyui-menu-mobile-collapse.comfyui-button {
+	padding: 10px;
+}
+.lt-lg.comfyui-menu > .comfyui-menu-mobile-collapse .comfyui-button,
+.lt-lg.comfyui-menu > .comfyui-menu-mobile-collapse .comfyui-button-wrapper {
+	width: 100%;
+}
+
+.lt-lg.comfyui-menu > .comfyui-menu-mobile-collapse .comfyui-popup {
+	position: static;
+	background-color: var(--comfy-input-bg);
+	max-width: unset;
+	max-height: 50vh;
+	overflow: auto;
+}
+
+.lt-lg.comfyui-menu:not(.expanded) > .comfyui-menu-mobile-collapse {
+	display: none;
+}
+
+.lt-lg .comfyui-queue-button {
+	margin-right: 44px;
+}
+
+.lt-lg .comfyui-menu-button {
+	position: absolute;
+	top: 4px;
+	right: 8px;
+}
+
+.lt-lg.comfyui-menu > .comfyui-menu-mobile-collapse .comfyui-view-list-popup {
+	border-radius: 0;
+}
+
+.lt-lg.comfyui-menu .comfyui-workflows-popup {
+	width: 100vw;
+}
+
+/** Small */
+.lt-md .comfyui-workflows-button-inner {
+	width: unset !important;
+}
+.lt-md  .comfyui-workflows-label {
+	display: none;
+}
+
+/** Extra small */
+.lt-sm .comfyui-queue-button {
+	margin-right: 0;
+	width: 100%;
+}
+.lt-sm .comfyui-queue-button .comfyui-button {
+	justify-content: center;
+}
+.lt-sm .comfyui-interrupt-button {
+	margin-right: 45px;
+}
+.comfyui-body-bottom .lt-sm.comfyui-menu > .comfyui-menu-button{
+	bottom: 41px;
+}

web/scripts/ui/menu/queueButton.js

@@ -0,0 +1,93 @@
+// @ts-check
+
+import { ComfyButton } from "../components/button.js";
+import { $el } from "../../ui.js";
+import { api } from "../../api.js";
+import { ComfySplitButton } from "../components/splitButton.js";
+import { ComfyQueueOptions } from "./queueOptions.js";
+import { prop } from "../../utils.js";
+
+export class ComfyQueueButton {
+	element = $el("div.comfyui-queue-button");
+	#internalQueueSize = 0;
+
+	queuePrompt = async (e) => {
+		this.#internalQueueSize += this.queueOptions.batchCount;
+		// Hold shift to queue front, event is undefined when auto-queue is enabled
+		await this.app.queuePrompt(e?.shiftKey ? -1 : 0, this.queueOptions.batchCount);
+	};
+
+	constructor(app) {
+		this.app = app;
+		this.queueSizeElement = $el("span.comfyui-queue-count", {
+			textContent: "?",
+		});
+
+		const queue = new ComfyButton({
+			content: $el("div", [
+				$el("span", {
+					textContent: "Queue",
+				}),
+				this.queueSizeElement,
+			]),
+			icon: "play",
+			classList: "comfyui-button",
+			action: this.queuePrompt,
+		});
+
+		this.queueOptions = new ComfyQueueOptions(app);
+
+		const btn = new ComfySplitButton(
+			{
+				primary: queue,
+				mode: "click",
+				position: "absolute",
+				horizontal: "right",
+			},
+			this.queueOptions.element
+		);
+		btn.element.classList.add("primary");
+		this.element.append(btn.element);
+
+		this.autoQueueMode = prop(this, "autoQueueMode", "", () => {
+			switch (this.autoQueueMode) {
+				case "instant":
+					queue.icon = "infinity";
+					break;
+				case "change":
+					queue.icon = "auto-mode";
+					break;
+				default:
+					queue.icon = "play";
+					break;
+			}
+		});
+
+		this.queueOptions.addEventListener("autoQueueMode", (e) => (this.autoQueueMode = e["detail"]));
+
+		api.addEventListener("graphChanged", () => {
+			if (this.autoQueueMode === "change") {
+				if (this.#internalQueueSize) {
+					this.graphHasChanged = true;
+				} else {
+					this.graphHasChanged = false;
+					this.queuePrompt();
+				}
+			}
+		});
+
+		api.addEventListener("status", ({ detail }) => {
+			this.#internalQueueSize = detail?.exec_info?.queue_remaining;
+			if (this.#internalQueueSize != null) {
+				this.queueSizeElement.textContent = this.#internalQueueSize > 99 ? "99+" : this.#internalQueueSize + "";
+				this.queueSizeElement.title = `${this.#internalQueueSize} prompts in queue`;
+				if (!this.#internalQueueSize && !app.lastExecutionError) {
+					if (this.autoQueueMode === "instant" || (this.autoQueueMode === "change" && this.graphHasChanged)) {
+						this.graphHasChanged = false;
+						this.queuePrompt();
+					}
+				}
+			}
+		});
+	}
+}

web/scripts/ui/menu/queueOptions.js

@@ -0,0 +1,77 @@
+// @ts-check
+
+import { $el } from "../../ui.js";
+import { prop } from "../../utils.js";
+
+export class ComfyQueueOptions extends EventTarget {
+	element = $el("div.comfyui-queue-options");
+
+	constructor(app) {
+		super();
+		this.app = app;
+
+		this.batchCountInput = $el("input", {
+			className: "comfyui-queue-batch-value",
+			type: "number",
+			min: "1",
+			value: "1",
+			oninput: () => (this.batchCount = +this.batchCountInput.value),
+		});
+
+		this.batchCountRange = $el("input", {
+			type: "range",
+			min: "1",
+			max: "100",
+			value: "1",
+			oninput: () => (this.batchCount = +this.batchCountRange.value),
+		});
+
+		this.element.append(
+			$el("div.comfyui-queue-batch", [
+				$el(
+					"label",
+					{
+						textContent: "Batch count: ",
+					},
+					this.batchCountInput
+				),
+				this.batchCountRange,
+			])
+		);
+
+		const createOption = (text, value, checked = false) =>
+			$el(
+				"label",
+				{ textContent: text },
+				$el("input", {
+					type: "radio",
+					name: "AutoQueueMode",
+					checked,
+					value,
+					oninput: (e) => (this.autoQueueMode = e.target["value"]),
+				})
+			);
+
+		this.autoQueueEl = $el("div.comfyui-queue-mode", [
+			$el("span", "Auto Queue:"),
+			createOption("Disabled", "", true),
+			createOption("Instant", "instant"),
+			createOption("On Change", "change"),
+		]);
+
+		this.element.append(this.autoQueueEl);
+
+		this.batchCount = prop(this, "batchCount", 1, () => {
+			this.batchCountInput.value = this.batchCount + "";
+			this.batchCountRange.value = this.batchCount + "";
+		});
+
+		this.autoQueueMode = prop(this, "autoQueueMode", "Disabled", () => {
+			this.dispatchEvent(
+				new CustomEvent("autoQueueMode", {
+					detail: this.autoQueueMode,
+				})
+			);
+		});
+	}
+}